dry_run.py

"""One-shot smoke test of a single (task, transport) on the new agent loop.

This is the smoke test that validates the rewrite away from
claude-agent-sdk. It is intentionally dumb and verbose so failures are
obvious.

Usage:
    AWS_PROFILE=mcp-benchmark AWS_REGION=us-west-2 \\
        python -m src.dry_run [transport] [task_id]

    transport: "cli" (default) | "mcp"
    task_id:   default "ec2_running"
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
from pathlib import Path

import anthropic
from dotenv import load_dotenv

from .agent_loop import run_agent
from .tasks import get_task
from .tools_cli import AWS_CLI_TOOL_SPEC, run_aws_cli_tool
from .tools_cli_rich import AWS_CLI_RICH_TOOL_SPEC, run_aws_cli_rich_tool
from .tools_cli_renamed import CALL_AWS_RENAMED_SPEC, run_call_aws_renamed_tool
from .tools_cli_v2 import AWS_CLI_V2_TOOL_SPEC, run_aws_cli_v2_tool
from .tools_cli_with_fake_suggest import (
    CLI_WITH_FAKE_SUGGEST_TOOLS,
    run_cli_with_fake_suggest_tool,
)
from .tools_mcp import open_aws_api_mcp
from .verify import verify

REPO_ROOT = Path(__file__).resolve().parent.parent
GROUND_TRUTH_PATH = REPO_ROOT / "results" / "ground_truth.local.json"
MODEL = "claude-sonnet-4-6"

SYSTEM_PROMPT = """\
You are a read-only AWS inspector. Your job is to answer questions about \
an AWS account by calling the available AWS tool. You may call the tool \
as many times as needed but you must finish each task by emitting a \
single JSON value (object or array) in the exact shape requested. Do not \
include explanations alongside the final JSON — return only the JSON, \
optionally inside a ```json fenced block. Use --output json on aws CLI \
calls when it helps. Never attempt write operations; only describe/list/\
get commands are permitted, and any attempt outside the allowlist will \
be rejected.\
"""


def _system_prompt_with_context() -> str:
    """Augment the base system prompt with runtime context.

    Hypothesis: the awslabs MCP server effectively gives the model the
    current UTC date for free because every tool result includes the AWS
    HTTP response headers (which contain a `date` field). Raw CLI tools
    return only the response body, so the model has to guess at
    timestamps for CloudWatch / time-windowed queries.

    By injecting the same minimum context (date, account, region) into
    the CLI system prompt, we can isolate this effect from the protocol
    or the tool description.
    """
    from datetime import datetime, timezone
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    return SYSTEM_PROMPT + (
        f"\n\nRuntime context (provided by the runner, not by the tool):\n"
        f"- Current UTC time: {now}\n"
        f"- Default AWS region: {os.environ.get('AWS_REGION', 'us-west-2')}\n"
        f"- IAM principal: arn:aws:iam::{os.environ.get('AWS_ACCOUNT_ID', '<account>')}:user/mcp-benchmark\n"
        f"- This account is real and live; commands return real data."
    )


def _print_result(result, verdict) -> None:
    print()
    print("=== Result ===")
    print(f"Stop reason: {result.stop_reason}")
    print(f"Turns: {result.num_turns}")
    print(f"Tool calls: {len(result.tool_calls)}")
    for i, c in enumerate(result.tool_calls, 1):
        payload = c.input.get("command") or c.input.get("cli_command") or json.dumps(c.input)
        print(f"  {i}. {c.name} ({c.duration_ms}ms, error={c.is_error})")
        print(f"     {str(payload)[:200]}")
    print()
    print("Final answer:")
    print(result.final_text[:500])
    print()
    print(
        f"Tokens: input={result.input_tokens} "
        f"cache_read={result.cache_read_input_tokens} "
        f"cache_creation={result.cache_creation_input_tokens} "
        f"output={result.output_tokens}"
    )
    print(f"Wall clock: {result.wall_clock_ms}ms")
    if result.error:
        print(f"ERROR: {result.error}")
    print()
    print(f"Verification: ok={verdict.ok} reason={verdict.reason}")


async def run_cli(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    print(f"=== Dry run: task={task.id} transport=cli model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print()

    result = await run_agent(
        client=client,
        model=MODEL,
        system=SYSTEM_PROMPT,
        tools=[AWS_CLI_TOOL_SPEC],
        user_prompt=task.prompt,
        tool_runner=run_aws_cli_tool,
        max_turns=10,
    )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


async def run_cli_v2(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    print(f"=== Dry run: task={task.id} transport=cli-v2 model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print()

    result = await run_agent(
        client=client,
        model=MODEL,
        system=SYSTEM_PROMPT,
        tools=[AWS_CLI_V2_TOOL_SPEC],
        user_prompt=task.prompt,
        tool_runner=run_aws_cli_v2_tool,
        max_turns=15,
    )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


async def run_cli_ctx(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    sys_prompt = _system_prompt_with_context()
    print(f"=== Dry run: task={task.id} transport=cli-ctx model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print(f"System prompt extra:\n{sys_prompt[len(SYSTEM_PROMPT):]}")
    print()

    result = await run_agent(
        client=client,
        model=MODEL,
        system=sys_prompt,
        tools=[AWS_CLI_TOOL_SPEC],
        user_prompt=task.prompt,
        tool_runner=run_aws_cli_tool,
        max_turns=15,
    )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


async def run_cli_renamed(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    print(f"=== Dry run: task={task.id} transport=cli-renamed model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print()

    result = await run_agent(
        client=client,
        model=MODEL,
        system=SYSTEM_PROMPT,
        tools=[CALL_AWS_RENAMED_SPEC],
        user_prompt=task.prompt,
        tool_runner=run_call_aws_renamed_tool,
        max_turns=10,
    )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


async def run_cli_fake_suggest(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    print(f"=== Dry run: task={task.id} transport=cli-fake-suggest model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print()

    result = await run_agent(
        client=client,
        model=MODEL,
        system=SYSTEM_PROMPT,
        tools=CLI_WITH_FAKE_SUGGEST_TOOLS,
        user_prompt=task.prompt,
        tool_runner=run_cli_with_fake_suggest_tool,
        max_turns=10,
    )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


async def run_cli_rich(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    print(f"=== Dry run: task={task.id} transport=cli-rich model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print()

    result = await run_agent(
        client=client,
        model=MODEL,
        system=SYSTEM_PROMPT,
        tools=[AWS_CLI_RICH_TOOL_SPEC],
        user_prompt=task.prompt,
        tool_runner=run_aws_cli_rich_tool,
        max_turns=10,
    )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


async def run_mcp(task_id: str) -> int:
    truth_blob = json.loads(GROUND_TRUTH_PATH.read_text())
    task = get_task(task_id)
    client = anthropic.Anthropic()

    print(f"=== Dry run: task={task.id} transport=mcp model={MODEL} ===")
    print(f"AWS_PROFILE={os.environ.get('AWS_PROFILE')!r} AWS_REGION={os.environ.get('AWS_REGION')!r}")
    print()

    async with open_aws_api_mcp() as (tools, tool_runner):
        print(f"MCP server published {len(tools)} tools: {[t['name'] for t in tools]}")
        print()
        result = await run_agent(
            client=client,
            model=MODEL,
            system=SYSTEM_PROMPT,
            tools=tools,
            user_prompt=task.prompt,
            tool_runner=tool_runner,
            max_turns=10,
        )
    verdict = verify(task_id, result.final_text, truth_blob)
    _print_result(result, verdict)
    return 0 if verdict.ok else 2


def main() -> int:
    load_dotenv(REPO_ROOT / ".env")
    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("ERROR: ANTHROPIC_API_KEY missing in env / .env")
        return 1
    if not GROUND_TRUTH_PATH.exists():
        print(f"ERROR: {GROUND_TRUTH_PATH} not found. Run: python -m src.ground_truth")
        return 1

    transport = sys.argv[1] if len(sys.argv) > 1 else "cli"
    task_id = sys.argv[2] if len(sys.argv) > 2 else "ec2_running"

    if transport == "cli":
        return asyncio.run(run_cli(task_id))
    if transport == "cli-rich":
        return asyncio.run(run_cli_rich(task_id))
    if transport == "cli-fake-suggest":
        return asyncio.run(run_cli_fake_suggest(task_id))
    if transport == "cli-renamed":
        return asyncio.run(run_cli_renamed(task_id))
    if transport == "cli-ctx":
        return asyncio.run(run_cli_ctx(task_id))
    if transport == "cli-v2":
        return asyncio.run(run_cli_v2(task_id))
    if transport == "mcp":
        return asyncio.run(run_mcp(task_id))
    print(f"ERROR: unknown transport {transport!r}")
    return 1


if __name__ == "__main__":
    sys.exit(main())