agent_loop.py

"""Minimal Anthropic agent loop, transport-agnostic.

This is intentionally hand-rolled — not built on claude-agent-sdk — for two
reasons:

1. We need full control over what goes into the system prompt and tool
   list. claude-agent-sdk loads ~/.claude.json on every call (figma,
   pencil, PubMed, etc.) which contaminates the benchmark context.
2. We need access to ``cache_creation_input_tokens`` and
   ``cache_read_input_tokens`` directly from the API response, with no
   layer in between.

The loop is the same for every transport. What differs is the
``tool_runner`` callback: for the CLI transport it shells out to ``aws``
via the safety whitelist; for the MCP transport it forwards the call to
an mcp ClientSession.

Output is a single ``AgentResult`` per task containing every signal we
need to score and aggregate later.
"""

from __future__ import annotations

import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any

import anthropic
from anthropic.types import Message

# A tool call request from the model: name, opaque input dict, unique id.
# The runner returns either a string result or an error string.
ToolRunner = Callable[[str, dict[str, Any]], Awaitable[tuple[str, bool]]]


@dataclass
class ToolCallRecord:
    name: str
    input: dict[str, Any]
    output: str
    is_error: bool
    duration_ms: int


@dataclass
class AgentResult:
    """Everything we capture from one (task, transport, run) execution."""

    final_text: str
    stop_reason: str
    tool_calls: list[ToolCallRecord]
    # Aggregated usage across all turns of this run.
    input_tokens: int = 0
    output_tokens: int = 0
    cache_creation_input_tokens: int = 0
    cache_read_input_tokens: int = 0
    num_turns: int = 0
    wall_clock_ms: int = 0
    error: str | None = None
    model_replies: list[Message] = field(default_factory=list)


def _extract_text(msg: Message) -> str:
    parts: list[str] = []
    for block in msg.content:
        if block.type == "text":
            parts.append(block.text)
    return "".join(parts)


async def run_agent(
    *,
    client: anthropic.Anthropic,
    model: str,
    system: str,
    tools: list[dict[str, Any]],
    user_prompt: str,
    tool_runner: ToolRunner,
    max_turns: int = 10,
    max_tokens: int = 4096,
) -> AgentResult:
    """Run a single agent task and return aggregated metrics.

    Args:
        client: Anthropic client (already configured with API key).
        model: Model name, e.g. ``"claude-sonnet-4-6"``.
        system: System prompt text. No cache_control set — we deliberately
            measure uncached behaviour on the first pass.
        tools: List of tool specs in the Anthropic API shape (each is a
            dict with ``name``, ``description``, ``input_schema``).
        user_prompt: The task prompt the agent must answer.
        tool_runner: Async callable that takes ``(tool_name, input)`` and
            returns ``(output_text, is_error)``.
        max_turns: Hard cap on model -> tools -> model -> ... rounds. The
            cap counts assistant turns, not API calls.
        max_tokens: Per-turn output token cap.

    Returns:
        An ``AgentResult`` with the final answer text, captured tool
        calls, and aggregated token / time / cost-relevant metrics.
    """
    started = time.monotonic()
    result = AgentResult(final_text="", stop_reason="", tool_calls=[])
    messages: list[dict[str, Any]] = [{"role": "user", "content": user_prompt}]

    try:
        for _turn in range(max_turns):
            resp = client.messages.create(
                model=model,
                max_tokens=max_tokens,
                system=system,
                tools=tools,  # type: ignore[arg-type]
                messages=messages,  # type: ignore[arg-type]
            )
            result.num_turns += 1
            result.model_replies.append(resp)

            u = resp.usage
            result.input_tokens += u.input_tokens or 0
            result.output_tokens += u.output_tokens or 0
            result.cache_creation_input_tokens += u.cache_creation_input_tokens or 0
            result.cache_read_input_tokens += u.cache_read_input_tokens or 0
            result.stop_reason = resp.stop_reason or ""

            # Append the assistant turn verbatim so the model sees its
            # own tool_use blocks on the next round.
            messages.append({"role": "assistant", "content": resp.content})

            if resp.stop_reason == "end_turn":
                result.final_text = _extract_text(resp)
                break

            if resp.stop_reason == "tool_use":
                tool_results: list[dict[str, Any]] = []
                for block in resp.content:
                    if block.type != "tool_use":
                        continue
                    call_started = time.monotonic()
                    output, is_error = await tool_runner(block.name, dict(block.input))
                    call_ms = int((time.monotonic() - call_started) * 1000)
                    result.tool_calls.append(
                        ToolCallRecord(
                            name=block.name,
                            input=dict(block.input),
                            output=output,
                            is_error=is_error,
                            duration_ms=call_ms,
                        )
                    )
                    tool_results.append(
                        {
                            "type": "tool_result",
                            "tool_use_id": block.id,
                            "content": output,
                            "is_error": is_error,
                        }
                    )
                messages.append({"role": "user", "content": tool_results})
                continue

            # Any other stop_reason: capture whatever text we have and
            # exit. This covers max_tokens, refusal, pause_turn, etc.
            result.final_text = _extract_text(resp)
            break
        else:
            result.error = f"max_turns ({max_turns}) exceeded"

    except anthropic.APIError as exc:
        result.error = f"AnthropicAPIError: {exc}"
    except Exception as exc:  # noqa: BLE001
        result.error = f"{type(exc).__name__}: {exc}"

    result.wall_clock_ms = int((time.monotonic() - started) * 1000)
    return result