orchestrator/agent-team/agent_team/invoker.py

"""Real Claude invokers for the ``claude_invoke`` billing seam (design §3.1).

:mod:`agent_team.billing` owns mode selection and the subscription-mode env
hygiene, but delegates the actual SDK call to a pluggable invoker bound via
:func:`agent_team.billing.set_invoker`. This module supplies that invoker: a
single function matching the ``billing.Invoker`` signature
``(prompt, *, mode, **kw) -> ClaudeResult`` that dispatches on
:class:`~agent_team.billing.BillingMode`:

* ``SUBSCRIPTION`` — the R720 default. Runs the Claude Agent SDK headless over
  the subscription OAuth token (``CLAUDE_CODE_OAUTH_TOKEN``), mirroring the
  canonical pattern in ``security-review/run_headless.py``. ``billing`` has
  already popped any stray ``ANTHROPIC_API_KEY`` for the duration of the call,
  so we only assert the OAuth token is present.
* ``API`` — a thin metered call through the ``anthropic`` SDK.
* ``BEDROCK`` — the rare cross-family tiebreak path; not wired for P1, so it
  raises :class:`NotImplementedError` honestly (a later config-flip wires it).

Deferred-import rationale: neither ``claude_agent_sdk`` nor ``anthropic`` is
installed in the test/Mac scaffolding environment, so importing either at
module load would raise :class:`ModuleNotFoundError` and break a clean import.
Following the deferred-import pattern of
:func:`agent_team.graph.build_sqlite_checkpointer`, the SDK imports live inside
the functions that actually call them and raise a clear :class:`RuntimeError`
when the package is missing. The SDK callables are also injectable (``_query``,
``_client``) so the real path stays unit-testable without the SDKs installed,
mirroring how the codebase keeps SDK calls injectable (see
:func:`agent_team.billing.set_invoker` and
:func:`agent_team.resume_worker.build_resume_command`).
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any, Callable

from agent_team.billing import BillingMode, ClaudeResult, set_invoker

_LOG = logging.getLogger(__name__)

__all__ = [
    "API_MODEL",
    "api_invoker",
    "bind_invoker",
    "bind_subscription_invoker",
    "subscription_invoker",
]

# Metered model for the API path (the rare opt-in billing mode).
API_MODEL = "claude-sonnet-4-6"

# Default per-call agent budget for the headless subscription path, in USD.
_DEFAULT_BUDGET_USD = 2.0
# The agent-team's subscription Claude calls (clarify confidence/questions,
# planner) are SINGLE-SHOT reasoning→JSON completions, NOT agentic sessions. A
# 40-turn, tool-enabled session made each call take minutes and let Claude wander
# (use tools / explore) and return output the planner/clarifier couldn't parse →
# spurious parks. One turn + no tools = a fast, deterministic completion. A
# genuinely agentic caller (e.g. a future fixer) overrides max_turns/allowed_tools.
_DEFAULT_MAX_TURNS = 1


# --------------------------------------------------------------------------- #
# Subscription path (Claude Agent SDK, headless OAuth)
# --------------------------------------------------------------------------- #


def _require_agent_sdk() -> Any:
    """Import and return ``claude_agent_sdk`` or raise a clear RuntimeError.

    Deferred so this module imports cleanly where the SDK is absent (the
    Mac/test scaffold). Mirrors graph.build_sqlite_checkpointer.
    """
    try:
        import claude_agent_sdk
    except ImportError as exc:  # pragma: no cover - depends on optional dep
        raise RuntimeError(
            "claude_agent_sdk is unavailable; install it to use the "
            "subscription billing path (the R720 default). Tests inject a fake "
            "query via the _query parameter."
        ) from exc
    return claude_agent_sdk


async def _collect_subscription_text(
    prompt: str,
    *,
    max_turns: int,
    budget_usd: float,
    model: str | None,
    _query: Callable[..., Any] | None = None,
    _options_cls: Callable[..., Any] | None = None,
) -> tuple[str, dict[str, Any], list[Any]]:
    """Drive one headless Agent SDK turn; return (text, usage, raw_messages).

    ``_query``/``_options_cls`` default to the real ``claude_agent_sdk``
    callables (lazily imported) but are injectable so tests can supply a fake
    async ``query`` without the SDK installed. Text extraction mirrors
    ``run_headless.py``: prefer the terminal ``ResultMessage.result``, falling
    back to concatenated ``AssistantMessage`` text blocks.
    """
    if _query is None or _options_cls is None:
        sdk = _require_agent_sdk()
        if _query is None:
            _query = sdk.query
        if _options_cls is None:
            _options_cls = sdk.ClaudeAgentOptions

    opts = _options_cls(
        permission_mode="bypassPermissions",
        setting_sources=[],  # hermetic: ignore user/project/local config + CLAUDE.md
        model=model,
        max_turns=max_turns,
        max_budget_usd=budget_usd,
        # No tools: these are pure reasoning→JSON completions. Disallowing tools
        # keeps the call a single deterministic turn (no repo exploration / tool
        # loops that produce slow, unparseable output).
        allowed_tools=[],
    )

    texts: list[str] = []
    result_text: str | None = None
    messages: list[Any] = []
    usage: dict[str, Any] = {}
    try:
        async for msg in _query(prompt=prompt, options=opts):
            messages.append(msg)
            name = type(msg).__name__
            if name == "AssistantMessage":
                for block in getattr(msg, "content", []) or []:
                    text = getattr(block, "text", None)
                    if text:
                        texts.append(text)
            elif name == "ResultMessage":
                result_text = getattr(msg, "result", None)
                cost = getattr(msg, "total_cost_usd", None)
                if cost is not None:
                    usage["total_cost_usd"] = float(cost)
                sdk_usage = getattr(msg, "usage", None)
                if isinstance(sdk_usage, dict):
                    usage.update(sdk_usage)
                elif sdk_usage is not None:
                    usage["usage"] = sdk_usage
    except Exception as exc:  # noqa: BLE001 - see below; we re-raise non-turn errors
        # These calls run SINGLE-SHOT (max_turns=1, allowed_tools=[]). When the
        # model's one turn does not terminate in a final ResultMessage, the Agent
        # SDK raises a terminal "Reached maximum number of turns (N)" error mid-
        # iteration. For a reasoning→JSON completion that is NOT fatal: the model
        # has usually already emitted its JSON as AssistantMessage text, so we
        # salvage what we collected and let the caller parse it (an empty salvage
        # surfaces cleanly as a PlannerError, which the coordinator turns into a
        # failed task — never a crashed daemon). Any OTHER error is a real fault
        # (auth, network, SDK bug) and must propagate.
        if "maximum number of turns" not in str(exc).lower():
            raise
        if not (result_text or texts):
            # Nothing usable was produced before the cap; re-raise so the caller
            # sees a real failure rather than silently parsing an empty reply.
            raise
        _LOG.warning(
            "subscription call hit the single-shot turn cap; salvaging %d "
            "collected text block(s) instead of crashing: %s",
            len(texts),
            exc,
        )

    return (result_text or "\n".join(texts)), usage, messages


def subscription_invoker(
    prompt: str,
    *,
    mode: BillingMode,
    max_turns: int = _DEFAULT_MAX_TURNS,
    budget_usd: float = _DEFAULT_BUDGET_USD,
    model: str | None = None,
    _query: Callable[..., Any] | None = None,
    _options_cls: Callable[..., Any] | None = None,
    **kw: Any,
) -> ClaudeResult:
    """Invoke Claude headless over the subscription OAuth token (§3.1).

    Asserts ``CLAUDE_CODE_OAUTH_TOKEN`` is present (the metered key is already
    popped by :func:`agent_team.billing.claude_invoke` in subscription mode) and
    refuses to run without it, naming ``~/secrev.env`` as the source. The Agent
    SDK ``query()`` coroutine is bridged to this sync seam with
    :func:`asyncio.run`; the box path is synchronous, but we fail clearly rather
    than silently if invoked from inside a running event loop.

    ``_query``/``_options_cls`` are injection seams for tests; production leaves
    them ``None`` so the real ``claude_agent_sdk`` callables are used.
    """
    import os

    if not os.environ.get("CLAUDE_CODE_OAUTH_TOKEN"):
        raise RuntimeError(
            "subscription_invoker requires CLAUDE_CODE_OAUTH_TOKEN to be set "
            "(source ~/secrev.env). Refusing to run the subscription OAuth path "
            "without it."
        )

    coro = _collect_subscription_text(
        prompt,
        max_turns=max_turns,
        budget_usd=budget_usd,
        model=model,
        _query=_query,
        _options_cls=_options_cls,
    )

    try:
        asyncio.get_running_loop()
    except RuntimeError:
        text, usage, messages = asyncio.run(coro)
    else:  # pragma: no cover - the box path is synchronous
        coro.close()
        raise RuntimeError(
            "subscription_invoker cannot bridge asyncio.run from within a "
            "running event loop; call claude_invoke from synchronous code."
        )

    return ClaudeResult(text=text, mode=mode, usage=usage, raw=messages)


# --------------------------------------------------------------------------- #
# API path (anthropic SDK, metered)
# --------------------------------------------------------------------------- #


def _require_anthropic() -> Any:
    """Import and return the ``anthropic`` module or raise a clear RuntimeError."""
    try:
        import anthropic
    except ImportError as exc:  # pragma: no cover - depends on optional dep
        raise RuntimeError(
            "anthropic is unavailable; install it to use the API billing path. "
            "Tests inject a fake client via the _client parameter."
        ) from exc
    return anthropic


def _extract_text(message: Any) -> str:
    """Join the text blocks of an anthropic Messages response."""
    parts: list[str] = []
    for block in getattr(message, "content", []) or []:
        text = getattr(block, "text", None)
        if text:
            parts.append(text)
    return "".join(parts)


def api_invoker(
    prompt: str,
    *,
    mode: BillingMode,
    model: str = API_MODEL,
    max_tokens: int = 4096,
    _client: Any | None = None,
    **kw: Any,
) -> ClaudeResult:
    """Invoke Claude through the metered ``anthropic`` SDK (§3.1, API mode).

    ``_client`` is an injection seam for tests; production leaves it ``None`` so
    a real ``anthropic.Anthropic()`` is constructed (reading
    ``ANTHROPIC_API_KEY`` from the environment, as the SDK does by default).
    """
    if _client is None:
        anthropic = _require_anthropic()
        _client = anthropic.Anthropic()

    message = _client.messages.create(
        model=model,
        max_tokens=max_tokens,
        messages=[{"role": "user", "content": prompt}],
    )

    usage_obj = getattr(message, "usage", None)
    if usage_obj is None:
        usage: dict[str, Any] = {}
    elif isinstance(usage_obj, dict):
        usage = dict(usage_obj)
    elif hasattr(usage_obj, "model_dump"):
        usage = usage_obj.model_dump()
    else:
        usage = {
            "input_tokens": getattr(usage_obj, "input_tokens", None),
            "output_tokens": getattr(usage_obj, "output_tokens", None),
        }

    return ClaudeResult(
        text=_extract_text(message), mode=mode, usage=usage, raw=message
    )


# --------------------------------------------------------------------------- #
# Dispatch + binding
# --------------------------------------------------------------------------- #


def real_invoker(prompt: str, *, mode: BillingMode, **kw: Any) -> ClaudeResult:
    """Dispatch to the per-mode real invoker (the ``billing.Invoker`` contract).

    ``BEDROCK`` is the rare cross-family tiebreak path and is not wired for P1;
    it raises :class:`NotImplementedError` honestly. Flipping it on is a later
    config change, not a code rewrite of the seam.
    """
    if mode is BillingMode.SUBSCRIPTION:
        return subscription_invoker(prompt, mode=mode, **kw)
    if mode is BillingMode.API:
        return api_invoker(prompt, mode=mode, **kw)
    if mode is BillingMode.BEDROCK:
        raise NotImplementedError(
            "BEDROCK billing is the rare cross-family tiebreak path and is not "
            "wired for P1; enable it later via config-flip once the cross-account "
            "Bedrock transport is provisioned."
        )
    raise NotImplementedError(f"no invoker for billing mode {mode!r}")


def bind_invoker(invoker: Callable[..., ClaudeResult] | None = None) -> None:
    """Bind a real invoker into the billing seam in one line at startup.

    Defaults to :func:`real_invoker` (mode-dispatching). Not called at import
    time so importing this module has no global side effects.
    """
    set_invoker(invoker or real_invoker)


def bind_subscription_invoker() -> None:
    """Bind the mode-dispatching real invoker (subscription is the default mode).

    Convenience for the common R720 startup: one call wires
    :func:`agent_team.billing.claude_invoke` to the real Claude path.
    """
    set_invoker(real_invoker)