EvoScientist-Multi/EvoScientist/cli/_agent_loader.py

"""Background MCP/agent load lifecycle shared by CLI and TUI surfaces.

Holds no references to Rich, prompt_toolkit, or Textual — UI-specific
rendering and thread-hopping plug in via callbacks.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from typing import Any, Generic, TypeVar

_logger = logging.getLogger(__name__)

ProgressEvent = str  # "start" | "success" | "error"
ProgressState = str  # "pending" | "ok" | "error"

AgentT = TypeVar("AgentT")
ProgressCallback = Callable[[ProgressEvent, str, str], None]
SuccessCallback = Callable[[AgentT], None]
FailureCallback = Callable[[BaseException], None]


class MCPProgressTracker:
    """Per-server MCP load progress state.

    Reads and writes are GIL-atomic but iteration must go through
    :meth:`snapshot` — events can fire from a worker thread while the
    main thread renders.
    """

    __slots__ = ("progress",)

    def __init__(self) -> None:
        self.progress: dict[str, tuple[ProgressState, str]] = {}

    def prime(self) -> None:
        """Seed a ``pending`` entry for every configured server.

        Keeps the UI's "N / M" denominator stable from the first render.
        """
        try:
            from ..mcp import load_mcp_config

            cfg = load_mcp_config() or {}
            self.progress = dict.fromkeys(cfg, ("pending", ""))
        except Exception:
            self.progress = {}

    def record(
        self, event: ProgressEvent, server: str, detail: str
    ) -> ProgressState | None:
        """Apply an event and return the new state, or ``None`` if unknown."""
        if event == "start":
            self.progress.setdefault(server, ("pending", ""))
            return "pending"
        if event == "success":
            self.progress[server] = ("ok", detail)
            return "ok"
        if event == "error":
            self.progress[server] = ("error", detail)
            return "error"
        return None

    def snapshot(self) -> list[tuple[ProgressState, str]]:
        return list(self.progress.values())

    def totals(self) -> tuple[int, int]:
        """``(done, total)`` — done excludes ``pending``."""
        snap = self.snapshot()
        total = len(snap)
        done = sum(1 for state, _ in snap if state != "pending")
        return done, total


class BackgroundAgentLoader(Generic[AgentT]):
    """Owns the background ``_load_agent`` task and its generation token.

    Each :meth:`start` bumps an internal id; callbacks from a superseded
    load (the old worker thread keeps running after cancel, since
    ``asyncio.to_thread`` can't preempt arbitrary Python code) compare
    against it and drop silently.

    ``on_progress`` fires on the **worker thread**; UI callers hop
    threads inside it if needed.  ``on_success`` / ``on_failure`` fire
    on the event loop when the task completes.
    """

    def __init__(
        self,
        loader_fn: Callable[..., AgentT],
        *,
        on_progress: ProgressCallback | None = None,
        on_success: SuccessCallback | None = None,
        on_failure: FailureCallback | None = None,
    ) -> None:
        self._loader_fn = loader_fn
        self._on_progress = on_progress
        self._on_success = on_success
        self._on_failure = on_failure
        self.agent: AgentT | None = None
        self._task: asyncio.Task[AgentT] | None = None
        self._load_id: int = 0

    @property
    def task(self) -> asyncio.Task[AgentT] | None:
        return self._task

    @property
    def is_pending(self) -> bool:
        return self.agent is None and self._task is not None and not self._task.done()

    @property
    def needs_restart(self) -> bool:
        """True when no load is in flight and no agent is ready.

        Callers that want auto-retry behavior (e.g. TUI on the next
        user send after a failure) check this before :meth:`start`.
        """
        return self.agent is None and (self._task is None or self._task.done())

    def start(self, **loader_kwargs: Any) -> None:
        prev = self._task
        if prev is not None and not prev.done():
            prev.cancel()
        self._load_id += 1
        load_id = self._load_id
        self.agent = None

        def _gated_progress(event: str, server: str, detail: str) -> None:
            if load_id != self._load_id:
                return
            if self._on_progress is None:
                return
            try:
                self._on_progress(event, server, detail)
            except Exception:
                _logger.debug("MCP progress callback raised", exc_info=True)

        self._task = asyncio.create_task(
            asyncio.to_thread(
                self._loader_fn,
                on_mcp_progress=_gated_progress,
                **loader_kwargs,
            )
        )
        self._task.add_done_callback(lambda task, lid=load_id: self._on_done(task, lid))

    def adopt(self, agent: AgentT) -> None:
        """Install an externally-built agent and supersede any in-flight load.

        Used by ``/model`` (and any other caller that constructs a
        replacement agent directly): bumps the generation token so a
        late-arriving background load can't clobber ``self.agent`` via
        the done-callback, cancels the in-flight wrapper, and seats the
        new agent immediately.
        """
        prev = self._task
        if prev is not None and not prev.done():
            prev.cancel()
        self._load_id += 1
        self._task = None
        self.agent = agent

    async def await_ready(self) -> AgentT:
        """Return the loaded agent; re-raises on load failure.

        Idempotent.  State transitions (setting ``self.agent``, calling
        ``on_success`` / ``on_failure``) are handled exclusively by
        :meth:`_on_done`, which fires before this ``await`` resumes
        (asyncio guarantees done-callbacks run in registration order).
        """
        if self.agent is not None:
            return self.agent
        if self._task is None:
            raise RuntimeError(
                "BackgroundAgentLoader.await_ready called before start()"
            )
        await self._task
        if self.agent is None:
            raise RuntimeError("BackgroundAgentLoader completed without an agent")
        return self.agent

    def _on_done(self, task: asyncio.Task[AgentT], load_id: int) -> None:
        if load_id != self._load_id:
            return
        if task.cancelled():
            return
        try:
            self.agent = task.result()
        except Exception as exc:
            # Keep ``_task`` set so a later ``await_ready`` re-raises the
            # real exception instead of the "before start()" sentinel.
            self.agent = None
            if self._on_failure is not None:
                self._on_failure(exc)
            return
        if self._on_success is not None:
            self._on_success(self.agent)