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 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],
    *,
    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 to
    EITHER ``$CONTROLLER_CANONICAL_OUTPUT_PATH`` (preferred, set by
    the production worker agent_runner) OR stdout (fallback for
    direct-call tests).

    Why the env-var path: when the MCP runs as a subprocess of the
    worker controller (production), its stdout is shared with the
    JSON-RPC transport that FastMCP uses. Writing canonical JSON to
    stdout would collide with the transport's framing. The env-var
    path gives us a clean side channel: subprocess writes the
    canonical JSON to a file; agent_runner reads it after the
    subprocess exits.

    For direct-call tests (no subprocess), the env var is unset and
    we fall back to stdout (captured by capsys).
    """
    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()
    # Mark BEFORE emitting so duplicate finalize calls return error.
    state.finalized = True
    out_path = os.environ.get("CONTROLLER_CANONICAL_OUTPUT_PATH")
    if out_path:
        # Production / subprocess: file-based clean channel.
        try:
            with open(out_path, "w", encoding="utf-8") as f:
                f.write(canonical)
                f.write("\n")
                f.flush()
                os.fsync(f.fileno())
        except OSError as exc:
            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()
    return ok(committed=True, output_bytes=len(canonical))


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