cleveragents-core/tools/controller/mcp/_builder_base.py

"""Shared base for per-role response-builder MCP servers.

Each builder MCP holds per-attempt state in a single module-level
``BuilderState`` instance (per-attempt subprocess model — see plan
v9). Tools mutate that state; ``finalize()`` validates and emits
canonical JSON.

This module is helper code, not an MCP server itself. The role-specific
modules import these helpers, declare their Pydantic output type, and
wire up tools via FastMCP.
"""
from __future__ import annotations

import json
import sys
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, TypeVar

from pydantic import BaseModel, ValidationError

T_Output = TypeVar("T_Output", bound=BaseModel)


@dataclass
class BuilderState:
    """Mutable accumulator for an in-flight response build.

    Holds:
    - ``started``: True after ``{role}_start`` was called (every
      builder enforces ``start`` is the first tool).
    - ``finalized``: True after a successful ``{role}_finalize``.
      Further mutating calls are refused.
    - ``identity``: workflow_id / attempt_id / pr_number — bound at
      start, surfaced in every event log line so per-PR debugging
      via ``grep "pr=N"`` finds builder activity.
    - ``fields``: working dict of model-shaped data the role-specific
      module fills in. Role modules know their model class; this
      class is intentionally untyped here so it can be shared.
    - ``audit``: ordered list of every tool call (name + args summary)
      for ``controller-cli tail-events`` and operator forensics.
    """

    started: bool = False
    finalized: bool = False
    identity: dict[str, Any] = field(default_factory=dict)
    fields: dict[str, Any] = field(default_factory=dict)
    audit: list[dict[str, Any]] = field(default_factory=list)
    started_at: datetime | None = None
    _lock: threading.Lock = field(default_factory=threading.Lock)

    def require_started(self) -> None:
        if not self.started:
            raise BuilderError("call {role}_start(...) before any other tool")

    def require_not_finalized(self) -> None:
        if self.finalized:
            raise BuilderError("response already finalized; no further mutations allowed")

    def reset_for_new_attempt(self) -> None:
        """Wipe per-attempt accumulator so a fresh ``*_start`` call sees
        a clean slate. Called by every ``{role}_start`` body when the
        arriving attempt_id differs from the stored identity's.

        Required because the opencode.json-registered local MCP servers
        are reused across multiple OpenCode sessions (each session ==
        one controller attempt). Without reset the second attempt's
        ``_start`` would either raise "already started" or inherit the
        first attempt's accumulated fields/identity.

        Force-reset semantics: if the prior attempt was interrupted
        (timeout / lost-lock / OpenCode hang) it can leave
        ``started=True, finalized=False`` indefinitely. Refusing to
        reset would permanently wedge the MCP for the rest of the
        OpenCode server's lifetime. Log a WARNING when this happens
        so operators see abandoned attempts but proceed with reset.

        Intra-session double-``_start`` (same attempt_id called twice)
        is detected separately by callers AFTER reset, via the
        ``_STATE.started`` check.
        """
        if self.started and not self.finalized:
            import logging
            logging.getLogger(__name__).warning(
                "builder reset_for_new_attempt: prior attempt was not "
                "finalized (identity=%s); force-resetting state. The "
                "previous attempt likely hit a timeout or lost lock.",
                self.identity,
            )
        with self._lock:
            self.started = False
            self.finalized = False
            self.identity = {}
            self.fields = {}
            self.audit = []
            self.started_at = None

    def record(self, tool: str, args: dict[str, Any]) -> None:
        """Append an audit entry. Args are str-coerced to avoid
        carrying large blobs (e.g. raw_log_excerpt) into the audit
        log; the controller_events sink will summarize."""
        with self._lock:
            self.audit.append(
                {
                    "tool": tool,
                    "ts": datetime.now(timezone.utc).isoformat(),
                    "args_summary": _summarize_args(args),
                }
            )


def _summarize_args(args: dict[str, Any]) -> dict[str, Any]:
    """Truncate any large string args so the audit log stays small.
    Caps each str field at 200 chars; lists/dicts cap at counts only."""
    out: dict[str, Any] = {}
    for k, v in args.items():
        if isinstance(v, str):
            out[k] = v if len(v) <= 200 else f"{v[:200]}…({len(v)} chars)"
        elif isinstance(v, (list, tuple)):
            out[k] = f"<list len={len(v)}>"
        elif isinstance(v, dict):
            out[k] = f"<dict keys={len(v)}>"
        else:
            out[k] = v
    return out


class BuilderError(ValueError):
    """Raised when a builder tool refuses a call due to an invariant
    violation. The MCP wrapper converts this to the
    ``{"error": str}`` envelope the agent's tool-result parser sees."""


def ok(**fields: Any) -> dict[str, Any]:
    """Standard success envelope for builder tools."""
    return {"status": "ok", **fields}


def err(msg: str, **fields: Any) -> dict[str, Any]:
    """Standard error envelope. Builder tools wrap their bodies in
    a try/except BuilderError so the agent sees structured failures."""
    return {"error": msg, **fields}


def call_with_invariants(
    state: BuilderState,
    tool: str,
    fn: Callable[[], dict[str, Any]],
    args: dict[str, Any],
) -> dict[str, Any]:
    """Run a tool body with the standard invariant guard rails.

    Order of checks:
    1. Audit-record the call.
    2. Run ``fn()`` (which may raise ``BuilderError``).
    3. On BuilderError → ``err(msg)``; on success → fn's return value
       (caller already wrapped it in ``ok(...)``).
    """
    state.record(tool, args)
    try:
        return fn()
    except BuilderError as e:
        return err(str(e), tool=tool)


def finalize_and_emit(
    state: BuilderState,
    model_class: type[T_Output],
    *,
    output_path: str | None = None,
    extra_required_check: Callable[[BuilderState], None] | None = None,
) -> dict[str, Any]:
    """Common finalize helper: validate state.fields against
    ``model_class``, mark state.finalized, emit canonical JSON.

    Path resolution (precedence: explicit > env > stdout):
    1. ``output_path`` arg (preferred — the controller's prompt tells
       the agent the per-attempt path; the agent passes it to
       ``{role}_finalize(output_path=...)``). This is the
       opencode.json-registered-MCP path, where opencode.json's static
       env can't inject a per-attempt path.
    2. ``$CONTROLLER_CANONICAL_OUTPUT_PATH`` env var (legacy — used
       when the controller's worker spawns the MCP itself as a
       subprocess and injects the env var per-attempt).
    3. stdout (fallback for direct-call tests).
    """
    import os
    state.require_started()
    state.require_not_finalized()
    if extra_required_check is not None:
        extra_required_check(state)
    try:
        model = model_class.model_validate(state.fields)
    except ValidationError as exc:
        return err(
            f"finalize failed schema validation for {model_class.__name__}: {exc}",
            validation_error=str(exc),
        )
    canonical = model.model_dump_json()
    out_path = output_path or os.environ.get("CONTROLLER_CANONICAL_OUTPUT_PATH")
    if out_path:
        # CA3 path-injection defense: refuse paths that look obviously
        # adversarial. Real defense lives at the workspace boundary
        # (worker validates owner/repo + workspace_root is under /tmp),
        # but a defense-in-depth check here surfaces bugs early.
        if "\x00" in out_path:
            return err(f"output_path contains NUL byte; refusing to write")
        if not os.path.isabs(out_path):
            return err(
                f"output_path must be absolute; got {out_path!r}"
            )
        # Production: file-based clean channel.
        parent = os.path.dirname(out_path)
        if parent and os.path.exists(parent) and not os.path.isdir(parent):
            return err(
                f"parent of output_path={out_path!r} exists but is a file, "
                "not a directory; refusing to write"
            )
        try:
            # Ensure the parent dir exists so the controller doesn't
            # have to pre-create it; same idempotent guarantee as
            # tempfile.mkstemp gave.
            if parent:
                os.makedirs(parent, exist_ok=True)
            # Atomic write: write to .tmp then os.replace so a concurrent
            # reader (the worker's _wait_for_canonical_output poller)
            # never sees a half-written file.
            tmp_path = f"{out_path}.tmp"
            with open(tmp_path, "w", encoding="utf-8") as f:
                f.write(canonical)
                f.write("\n")
                f.flush()
                os.fsync(f.fileno())
            os.replace(tmp_path, out_path)
        except OSError as exc:
            # Best-effort cleanup of the .tmp leftover.
            try:
                os.unlink(f"{out_path}.tmp")
            except OSError:
                pass
            return err(f"failed to write canonical output to {out_path}: {exc}")
    else:
        # Direct-call tests: emit to stdout (captured by capsys).
        sys.stdout.write(canonical)
        sys.stdout.write("\n")
        sys.stdout.flush()
    # Mark AFTER successful emission so a failed write allows the agent
    # to retry finalize within the same session (CA12).
    state.finalized = True
    return ok(
        committed=True, output_bytes=len(canonical),
        wrote_to=out_path or "stdout",
    )


__all__ = [
    "BuilderError",
    "BuilderState",
    "call_with_invariants",
    "err",
    "finalize_and_emit",
    "ok",
]