cleveragents-core/tools/controller/worker/agent_runner.py

"""Production agent_runner adapter — spawns the role-builder MCP +
the OpenCode session + returns the parsed worker output.

The injected ``agent_runner`` parameter to ``run_one_attempt``
(``runner.py``) is what production wires to this module. Tests use a
synthetic agent_runner.

Lifecycle per attempt (plan v9 per-attempt subprocess model):
1. Spawn the matching role-builder MCP subprocess (one of
   ``tools/controller/mcp/{role}_builder.py``) with stdout captured
   so we can read the finalize'd JSON.
2. Tell the OpenCode session to run the role's agent, pointed at the
   spawned MCP server (path passed via env / args).
3. Wait for the OpenCode session to complete + the MCP subprocess
   to emit one JSON line on stdout.
4. Strict-parse against the V1 output contract for the role.
5. Cleanup: terminate the MCP subprocess (in case it didn't exit
   cleanly after finalize).

The actual OpenCode session call is parameterized as
``run_opencode_session`` so this module is testable without a live
OpenCode. Production wires it to ``tools/_opencode_worker.run_session_blocking``.
"""
from __future__ import annotations

import logging
import os
import signal
import subprocess
import sys
import threading
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from ..contracts.parse import ContractValidationError, strict_parse
from ..contracts.v1 import (
    ConflictResolverOutputV1,
    EstimatorOutputV1,
    ImplementerOutputV1,
    ReviewerOutputV1,
    SummarizerOutputV1,
)
from .runner import WorkerError, WorkerLostLock

logger = logging.getLogger(__name__)


# Module path of each role's MCP server (Python -m runnable).
ROLE_TO_MCP_MODULE: dict[str, str] = {
    "implementer": "tools.controller.mcp.implementer_builder",
    "reviewer": "tools.controller.mcp.reviewer_builder",
    "estimator": "tools.controller.mcp.estimator_builder",
    "conflict_resolver": "tools.controller.mcp.conflict_resolver_builder",
    "summarizer": "tools.controller.mcp.summarizer_builder",
}

# The Pydantic model each role's output is parsed against.
ROLE_TO_OUTPUT_MODEL: dict[str, type] = {
    "implementer": ImplementerOutputV1,
    "reviewer": ReviewerOutputV1,
    "estimator": EstimatorOutputV1,
    "conflict_resolver": ConflictResolverOutputV1,
    "summarizer": SummarizerOutputV1,
}


# Callable signature for the injected OpenCode session runner.
# Inputs: role, mcp_subprocess (so the session can spawn-as-child or
# communicate via stdio), input_payload (already validated).
# Returns: nothing — the MCP subprocess is what emits the canonical
# JSON on its own stdout. The session is just the driver.
#
# Production wires this to a thin wrapper around
# ``_opencode_worker.run_session_blocking``.
OpenCodeSessionRunner = Callable[..., None]


@dataclass
class MCPSpawnResult:
    """Captured MCP subprocess + its stdout reader thread."""

    process: subprocess.Popen
    stdout_buffer: list[str]  # mutable; appended by reader thread
    reader_thread: threading.Thread


def production_agent_runner(
    *,
    attempt_id: int,
    role: str,
    tier: int | None,
    input_payload: dict[str, Any],
    instance_id: str,
    lost_lock_check: Callable[[], bool],
    run_opencode_session: OpenCodeSessionRunner,
    python_executable: str = sys.executable,
    mcp_startup_timeout_s: float = 5.0,
    finalize_timeout_s: float = 30.0,
) -> dict[str, Any]:
    """The production agent_runner used by run_one_attempt.

    Maps the role → matching MCP module + output contract; spawns
    the MCP as a subprocess; calls the injected ``run_opencode_session``
    to drive the LLM session; reads the MCP's emitted JSON; strict-
    parses; returns the parsed dict.

    Tests can call this directly with a synthetic ``run_opencode_session``
    that pokes the MCP's stdin (or otherwise simulates a session). The
    integration with ``_opencode_worker.run_session_blocking`` is the
    production wiring.

    Raises:
        ValueError: unknown role.
        WorkerError: MCP didn't start / didn't emit valid JSON /
                     strict-parse failed twice.
        WorkerLostLock: ``lost_lock_check`` returned True during the run.
    """
    if role not in ROLE_TO_MCP_MODULE:
        raise ValueError(f"unknown role {role!r}; cannot select MCP")
    mcp_module = ROLE_TO_MCP_MODULE[role]
    model_class = ROLE_TO_OUTPUT_MODEL[role]

    # Production: per-attempt canonical output path. MCP's
    # finalize_and_emit writes the canonical JSON HERE rather than
    # to stdout (which is shared with the FastMCP JSON-RPC transport).
    import tempfile
    out_fd, out_path = tempfile.mkstemp(
        prefix=f"controller-canonical-{attempt_id}-", suffix=".json",
    )
    os.close(out_fd)

    spawn = _spawn_mcp_subprocess(
        mcp_module,
        python_executable=python_executable,
        startup_timeout_s=mcp_startup_timeout_s,
        canonical_output_path=out_path,
    )
    try:
        # Hand the OpenCode session the MCP we just spawned + the
        # input payload. The session's job: drive the LLM through
        # its tool-calling sequence; the MCP captures the assembled
        # state + emits canonical JSON via its finalize() tool.
        try:
            run_opencode_session(
                role=role, tier=tier, input_payload=input_payload,
                mcp_process=spawn.process,
                attempt_id=attempt_id, instance_id=instance_id,
                lost_lock_check=lost_lock_check,
            )
        except WorkerLostLock:
            raise
        except Exception as exc:
            # Wrap so the runner's outer handler classifies properly.
            raise WorkerError(
                f"OpenCode session raised: {exc}",
                outcome="worker-internal-error",
            ) from exc

        if lost_lock_check():
            raise WorkerLostLock("lost lock during MCP session")

        # Wait for the MCP to emit its JSON to the canonical output
        # file. Tear down the subprocess first (which flushes its
        # finalize write), then read the file.
        canonical = _wait_for_canonical_output(
            spawn, out_path, timeout_s=finalize_timeout_s,
        )
        if canonical is None:
            raise WorkerError(
                f"MCP for role {role!r} did not emit canonical output "
                f"within {finalize_timeout_s}s",
                outcome="worker-internal-error",
            )

        try:
            parsed = strict_parse(model_class, canonical)
        except ContractValidationError as exc:
            raise WorkerError(
                f"MCP emitted JSON that failed {model_class.__name__} "
                f"strict-parse: {exc}",
                outcome="contract-violation",
            ) from exc

        return parsed.model_dump()
    finally:
        _terminate_mcp_subprocess(spawn)
        # Clean up the canonical output file (best-effort).
        try:
            os.unlink(out_path)
        except OSError:
            pass


# ─── MCP subprocess lifecycle ────────────────────────────────────────


def _spawn_mcp_subprocess(
    mcp_module: str, *, python_executable: str, startup_timeout_s: float,
    canonical_output_path: str,
) -> MCPSpawnResult:
    """Spawn ``python -m {mcp_module}`` with stdout captured via a
    background reader thread.

    ``canonical_output_path`` is the file the MCP writes its canonical
    JSON to (via the CONTROLLER_CANONICAL_OUTPUT_PATH env var).
    Stdout is still captured (for transport/debug logging) but the
    authoritative output goes to the file.
    """
    repo_root = str(Path(__file__).resolve().parents[3])
    env = dict(os.environ)
    # Prepend repo_root to PYTHONPATH so ``-m tools.controller.mcp....`` works.
    existing_pp = env.get("PYTHONPATH", "")
    env["PYTHONPATH"] = (
        f"{repo_root}{os.pathsep}{existing_pp}" if existing_pp else repo_root
    )
    env["CONTROLLER_CANONICAL_OUTPUT_PATH"] = canonical_output_path
    proc = subprocess.Popen(
        [python_executable, "-m", mcp_module],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        env=env,
        cwd=repo_root,
    )

    stdout_buffer: list[str] = []
    reader_thread = threading.Thread(
        target=_drain_stream_into_list,
        args=(proc.stdout, stdout_buffer),
        name=f"mcp-stdout-{proc.pid}",
        daemon=True,
    )
    reader_thread.start()

    # Best-effort startup check: did the process die immediately?
    # Don't actually wait the full timeout — production OpenCode sessions
    # are long-running; we just verify the subprocess is alive.
    if proc.poll() is not None:
        raise WorkerError(
            f"MCP subprocess died immediately (returncode={proc.returncode})",
            outcome="worker-internal-error",
        )

    return MCPSpawnResult(
        process=proc, stdout_buffer=stdout_buffer,
        reader_thread=reader_thread,
    )


def _drain_stream_into_list(stream, target: list[str]) -> None:
    """Read lines from ``stream`` into ``target`` until EOF.

    Runs in a background thread. Each line (without trailing newline)
    is appended. Robust against the stream closing mid-read.
    """
    try:
        for line in iter(stream.readline, ""):
            if not line:
                break
            target.append(line.rstrip("\n"))
    except (ValueError, OSError):
        # Stream closed; thread exits.
        pass


def _wait_for_canonical_output(
    spawn: MCPSpawnResult, out_path: str, *, timeout_s: float,
) -> str | None:
    """Poll the canonical-output file. Returns its content on first
    non-empty read, or None on timeout.

    The MCP writes the file atomically (open + write + fsync + close)
    in its finalize_and_emit. We poll until the file exists + has
    content OR the subprocess has been gone long enough to indicate
    it's not going to write.
    """
    import time

    start = time.monotonic()
    while True:
        try:
            with open(out_path, "r", encoding="utf-8") as f:
                content = f.read().strip()
            if content:
                return content
        except FileNotFoundError:
            pass
        except OSError:
            pass

        # Subprocess gone + file still empty → give it one more brief
        # window for the final flush to land, then give up.
        if spawn.process.poll() is not None:
            time.sleep(0.1)
            try:
                with open(out_path, "r", encoding="utf-8") as f:
                    content = f.read().strip()
                if content:
                    return content
            except OSError:
                pass
            return None

        if time.monotonic() - start > timeout_s:
            return None
        time.sleep(0.05)


def _terminate_mcp_subprocess(spawn: MCPSpawnResult) -> None:
    """Send SIGTERM + 2s grace + SIGKILL to ensure cleanup. Drain
    the reader thread."""
    proc = spawn.process
    if proc.poll() is None:
        try:
            proc.terminate()
        except ProcessLookupError:
            pass
        try:
            proc.wait(timeout=2.0)
        except subprocess.TimeoutExpired:
            try:
                proc.kill()
            except ProcessLookupError:
                pass
            try:
                proc.wait(timeout=2.0)
            except subprocess.TimeoutExpired:
                pass
    # Reader thread will exit on its own when stdout closes.
    spawn.reader_thread.join(timeout=2.0)


__all__ = [
    "MCPSpawnResult",
    "OpenCodeSessionRunner",
    "ROLE_TO_MCP_MODULE",
    "ROLE_TO_OUTPUT_MODEL",
    "production_agent_runner",
]