iop/scripts/agent_benchmark/claude_iop.py

517 lines
22 KiB
Python

"""Secret-safe Claude Code adapter for the IOP benchmark lifecycle.
The generic lifecycle deliberately does not know Claude's command line or its
JSONL protocol. This module converts one immutable benchmark cell and one
prepared workspace into that closed boundary. It has no network dependency;
tests use a local fake executable.
"""
from __future__ import annotations
import json
import os
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable
from scripts.agent_benchmark.connectivity import (
CallerCapability,
ConnectivityResult,
EffectiveBinding,
RequestedEffectiveBinding,
make_result,
)
from scripts.agent_benchmark.lifecycle import (
COMPLETION_EXIT_AFTER_IDLE,
SUBMISSION_STDIN_ONCE,
InvocationSpec,
LifecycleMetricError,
ParsedMetric,
TLS_CA_ENV_KEYS,
count_metric,
duration_metric,
exact_value_redactor,
inherited_tls_ca_environment,
is_reported_number,
)
from scripts.agent_benchmark.manifest import MatrixCell, Timeout
from scripts.agent_benchmark.workspace import PreparedWorkspace
REDACTED = "[redacted]"
CLAUDE_ROUTE_KINDS = ("direct", "execution_preset")
# This is lexical order, as required by the closed connectivity capability
# tuple. The cell's requested effort is still passed through unchanged.
CLAUDE_EFFORTS = ("high", "low", "max", "medium", "xhigh")
_STRUCTURAL_SECRET_KEYS = frozenset(
{
"content", "text", "input", "arguments", "tool_input", "prompt", "query",
"result", "error", "errors", "error_message",
}
)
# The exact Claude result fields this adapter is allowed to observe. Anything
# outside these allowlists never becomes timing or usage evidence, and
# ``duration_api_ms`` is marked as overlapping because it is reported inside
# the same wall-clock window as ``duration_ms``.
_CLAUDE_DURATION_FIELDS = (
("duration_ms", "total_duration", False),
("duration_api_ms", "model_duration", True),
)
_CLAUDE_USAGE_FIELDS = {
"input_tokens": "input_tokens",
"output_tokens": "output_tokens",
"cache_read_input_tokens": "cached_input_tokens",
"cache_creation_input_tokens": "cache_write_tokens",
}
_CLAUDE_USAGE_BOOKKEEPING_FIELDS = {
"cache_creation", "inference_geo", "iterations", "output_tokens_details",
"server_tool_use", "service_tier", "speed",
}
_CLAUDE_USAGE_NESTED_COUNTS = {
"cache_creation": {"ephemeral_1h_input_tokens", "ephemeral_5m_input_tokens"},
"output_tokens_details": {"thinking_tokens"},
"server_tool_use": {"web_fetch_requests", "web_search_requests"},
}
class ClaudeIopError(Exception):
"""Base exception for configuration and protocol failures."""
class ClaudeIopValidationError(ClaudeIopError):
"""Raised when a runtime input or preflight shape is inadmissible."""
class ClaudeIopProtocolError(ClaudeIopError):
"""Raised for a malformed or contradictory claimed Claude terminal."""
@dataclass(frozen=True)
class ClaudeIopRuntime:
"""Runtime-only IOP inputs. Values are never written to durable evidence."""
binary: str
base_url: str
api_key: str
def _require_string(value: Any, name: str) -> str:
if not isinstance(value, str) or not value:
raise ClaudeIopValidationError(f"invalid {name}")
return value
def _exact_object(raw_line: str) -> dict[str, Any]:
if not isinstance(raw_line, str):
raise ClaudeIopProtocolError("invalid Claude JSONL line")
try:
value = json.loads(raw_line)
except (TypeError, ValueError) as exc:
raise ClaudeIopProtocolError("invalid Claude JSONL line") from exc
if not isinstance(value, dict):
raise ClaudeIopProtocolError("invalid Claude JSONL object")
return value
def _required_string(data: dict[str, Any], name: str) -> str:
value = data.get(name)
if not isinstance(value, str) or not value:
raise ClaudeIopProtocolError(f"missing Claude {name}")
return value
def _structural_redact(
value: Any,
sensitive_values: tuple[str, ...],
key: str = "",
*,
redact_terminal_message: bool = False,
) -> Any:
if key in _STRUCTURAL_SECRET_KEYS:
return REDACTED
if redact_terminal_message and key == "message":
return REDACTED
if isinstance(value, str):
for sensitive in sensitive_values:
if sensitive:
value = value.replace(sensitive, REDACTED)
return value
if isinstance(value, list):
return [
_structural_redact(
item, sensitive_values, redact_terminal_message=redact_terminal_message
)
for item in value
]
if isinstance(value, dict):
return {
str(name): _structural_redact(
item,
sensitive_values,
str(name),
redact_terminal_message=redact_terminal_message,
)
for name, item in value.items()
}
return value
def redact_claude_event(raw_line: str, sensitive_values: tuple[str, ...]) -> str:
"""Return canonical JSON without task/tool content or runtime secrets.
A malformed line is represented by a fixed marker so error reporting cannot
accidentally retain the raw malformed payload.
"""
try:
event = _exact_object(raw_line)
except ClaudeIopProtocolError:
return '{"type":"invalid_claude_json"}'
redacted = _structural_redact(
event,
sensitive_values,
redact_terminal_message=event.get("type") in ("error", "result"),
)
return json.dumps(redacted, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
def parse_preflight_binding(raw_line: str, cell: MatrixCell) -> RequestedEffectiveBinding:
"""Parse an explicit, secret-free IOP binding observation from JSONL.
The adapter never derives effective values from the requested cell. A
caller must provide all requested/effective scalars and every stage, or the
connectivity contract rejects the observation as incomplete.
"""
event = _exact_object(raw_line)
if event.get("type") != "system" or event.get("subtype") != "iop_binding":
raise ClaudeIopProtocolError("invalid Claude preflight event")
binding = event.get("binding")
if not isinstance(binding, dict) or set(binding) != {
"cell_id", "caller", "requested_route_kind", "requested_route_id",
"requested_model", "requested_effort", "effective_route_kind",
"effective_route_id", "effective_model", "effective_effort", "effective_bindings",
}:
raise ClaudeIopProtocolError("invalid Claude preflight binding")
raw_stages = binding["effective_bindings"]
if not isinstance(raw_stages, list):
raise ClaudeIopProtocolError("invalid Claude preflight stages")
stages: list[EffectiveBinding] = []
for item in raw_stages:
if not isinstance(item, dict) or set(item) != {"stage", "model", "effort"}:
raise ClaudeIopProtocolError("invalid Claude preflight stage")
stages.append(EffectiveBinding(item["stage"], item["model"], item["effort"]))
try:
result = RequestedEffectiveBinding(
binding["cell_id"], binding["caller"], binding["requested_route_kind"],
binding["requested_route_id"], binding["requested_model"],
binding["requested_effort"], binding["effective_route_kind"],
binding["effective_route_id"], binding["effective_model"],
binding["effective_effort"], tuple(stages),
)
# Validation is intentionally delegated to the one shared contract.
make_result(cell, claude_capability(), result)
except Exception as exc:
if isinstance(exc, ClaudeIopProtocolError):
raise
raise ClaudeIopProtocolError("invalid Claude preflight binding") from exc
return result
class ClaudeStreamParser:
"""Parse one fresh Claude Code stream without inventing terminal bindings."""
def __init__(self, cell: MatrixCell, session_id: str) -> None:
if not isinstance(cell, MatrixCell) or cell.caller != "claude":
raise ClaudeIopValidationError("Claude adapter requires a Claude matrix cell")
self.cell = cell
# The prepared workspace identity proves that this is a fresh caller
# invocation. Claude Code emits its own UUID in system/init, so it
# must be derived from that event rather than compared to this local
# opaque label.
self.prepared_session_id = _require_string(session_id, "session_id")
self.claude_session_id: str | None = None
self._phase = "await_init"
self._assistant_messages = 0
self._active_message_id: str | None = None
self._completed_message_ids: set[str] = set()
self._continuation_pending = False
self._api_error_seen = False
def _require_bound_session(self, event: dict[str, Any]) -> None:
if self.claude_session_id is None:
raise ClaudeIopProtocolError("missing Claude init")
if _required_string(event, "session_id") != self.claude_session_id:
raise ClaudeIopProtocolError("Claude session binding mismatch")
def _consume_init(self, event: dict[str, Any]) -> None:
if self._phase != "await_init":
raise ClaudeIopProtocolError("duplicate or out-of-order Claude init")
if _required_string(event, "model") != self.cell.iop.request_model:
raise ClaudeIopProtocolError("Claude model binding mismatch")
self.claude_session_id = _required_string(event, "session_id")
self._phase = "await_assistant"
def _consume_assistant(self, event: dict[str, Any]) -> str:
if self._phase != "await_assistant":
raise ClaudeIopProtocolError("duplicate or out-of-order Claude assistant")
self._require_bound_session(event)
message = event.get("message")
if not isinstance(message, dict):
raise ClaudeIopProtocolError("invalid Claude assistant event")
if event.get("is_api_error_message") is True:
if message.get("model") != "<synthetic>" or message.get("stop_reason") != "stop_sequence":
raise ClaudeIopProtocolError("invalid Claude API error event")
self._complete_active_message()
self._api_error_seen = True
self._phase = "await_error_result"
return None
if _required_string(message, "model") != self.cell.iop.request_model:
raise ClaudeIopProtocolError("Claude model binding mismatch")
stop_reason = message.get("stop_reason")
if stop_reason not in (None, "tool_use", "end_turn"):
raise ClaudeIopProtocolError("invalid Claude assistant stop reason")
message_id = message.get("id")
if message_id is not None and (not isinstance(message_id, str) or not message_id):
raise ClaudeIopProtocolError("invalid Claude assistant message id")
if self._continuation_pending:
if message_id != self._active_message_id:
self._complete_active_message()
self._continuation_pending = False
if stop_reason is None:
if message_id is None:
raise ClaudeIopProtocolError("unbound Claude assistant snapshot")
if message_id in self._completed_message_ids:
raise ClaudeIopProtocolError("duplicate Claude assistant message")
if self._active_message_id not in (None, message_id):
raise ClaudeIopProtocolError("overlapping Claude assistant messages")
self._active_message_id = message_id
return None
if message_id is None:
# Older fixture-shaped output did not include a message id. It is
# admissible only for the single final assistant event.
if stop_reason != "end_turn" or self._active_message_id is not None:
raise ClaudeIopProtocolError("unbound Claude assistant terminal")
self._assistant_messages += 1
else:
if message_id in self._completed_message_ids:
raise ClaudeIopProtocolError("duplicate Claude assistant message")
if self._active_message_id not in (None, message_id):
raise ClaudeIopProtocolError("overlapping Claude assistant messages")
self._active_message_id = message_id
self._complete_active_message()
if stop_reason == "tool_use":
self._phase = "await_tool_result"
return None
self._phase = "await_result"
return "finish"
def _consume_user(self, event: dict[str, Any]) -> None:
if self._phase not in ("await_assistant", "await_tool_result"):
raise ClaudeIopProtocolError("out-of-order Claude user event")
self._require_bound_session(event)
if self._active_message_id is not None:
# Claude Code may publish cumulative snapshots with the same
# assistant message id on both sides of one or more tool results.
# The next assistant snapshot (or the result terminal) determines
# whether this message continues or a new model call begins.
self._continuation_pending = True
elif self._phase != "await_tool_result":
raise ClaudeIopProtocolError("unexpected Claude user event")
self._phase = "await_assistant"
def _complete_active_message(self) -> None:
if self._active_message_id is None:
return
if self._active_message_id in self._completed_message_ids:
raise ClaudeIopProtocolError("duplicate Claude assistant message")
self._completed_message_ids.add(self._active_message_id)
self._active_message_id = None
self._continuation_pending = False
self._assistant_messages += 1
def _consume_result(self, event: dict[str, Any]) -> tuple[Any, ...]:
if self._phase == "await_error_result":
self._require_bound_session(event)
if (
not self._api_error_seen
or event.get("subtype") != "success"
or event.get("is_error") is not True
or event.get("terminal_reason") != "api_error"
):
raise ClaudeIopProtocolError("invalid Claude API error terminal")
self._phase = "complete"
# Do not emit finish/idle for an upstream API failure. The caller's
# non-zero exit is the lifecycle terminal; an unexpected zero exit
# still fails closed as missing terminal evidence.
return ()
if self._phase == "await_assistant" and self._active_message_id is not None:
if self._continuation_pending:
raise ClaudeIopProtocolError("Claude result followed an unresolved tool result")
self._complete_active_message()
elif self._phase != "await_result":
raise ClaudeIopProtocolError("duplicate or out-of-order Claude result")
self._require_bound_session(event)
if event.get("subtype") != "success" or event.get("is_error") is True:
raise ClaudeIopProtocolError("invalid Claude result terminal")
self._phase = "complete"
return (*self._observations(event), "idle")
def _observations(self, event: dict[str, Any]) -> tuple[ParsedMetric, ...]:
"""Convert only allowlisted reported Claude values into observations."""
model = self.cell.iop.request_model
observations: list[ParsedMetric] = []
try:
for field, name, overlap in _CLAUDE_DURATION_FIELDS:
if field in event:
if not is_reported_number(event[field]):
raise ClaudeIopProtocolError("invalid Claude duration observation")
observations.append(duration_metric(
name, event[field], reported_unit="ms",
model=model, overlap=overlap,
))
observations.append(count_metric(
"model_calls", self._assistant_messages, model=model
))
observations.extend(self._usage_observations(event, model))
except LifecycleMetricError as exc:
raise ClaudeIopProtocolError("invalid Claude usage observation") from exc
return tuple(observations)
@staticmethod
def _usage_observations(event: dict[str, Any], model: str) -> list[ParsedMetric]:
usage = event.get("usage")
if usage is None:
return []
allowed = set(_CLAUDE_USAGE_FIELDS) | _CLAUDE_USAGE_BOOKKEEPING_FIELDS
if not isinstance(usage, dict) or not set(usage) <= allowed:
raise ClaudeIopProtocolError("invalid Claude usage observation")
ClaudeStreamParser._validate_usage_bookkeeping(usage)
return [
count_metric(_CLAUDE_USAGE_FIELDS[field], value, model=model)
for field, value in sorted(usage.items()) if field in _CLAUDE_USAGE_FIELDS
]
@staticmethod
def _validate_usage_bookkeeping(usage: dict[str, Any]) -> None:
for field, nested_fields in _CLAUDE_USAGE_NESTED_COUNTS.items():
if field not in usage:
continue
value = usage[field]
if not isinstance(value, dict) or not set(value) <= nested_fields:
raise ClaudeIopProtocolError("invalid Claude usage observation")
if any(not is_reported_number(item) for item in value.values()):
raise ClaudeIopProtocolError("invalid Claude usage observation")
if "iterations" in usage and usage["iterations"] != []:
raise ClaudeIopProtocolError("invalid Claude usage observation")
for field in ("inference_geo", "service_tier", "speed"):
if field in usage and (
not isinstance(usage[field], str) or len(usage[field]) > 64
):
raise ClaudeIopProtocolError("invalid Claude usage observation")
def __call__(self, stream: str, raw_line: str) -> str | tuple[Any, ...] | None:
if stream != "stdout":
return None
event = _exact_object(raw_line)
event_type = event.get("type")
if event_type == "system" and event.get("subtype") == "init":
self._consume_init(event)
return None
if event_type == "assistant":
return self._consume_assistant(event)
if event_type == "user":
self._consume_user(event)
return None
if event_type == "result":
return self._consume_result(event)
# Informational events are deliberately ignored only after they have
# passed exact JSON-object decoding above.
return None
def claude_capability() -> CallerCapability:
return CallerCapability("claude", CLAUDE_ROUTE_KINDS, CLAUDE_EFFORTS)
def resolve_claude_binary(binary: str) -> str:
"""Resolve one executable before lifecycle process creation."""
candidate = _require_string(binary, "Claude binary")
path = Path(candidate)
resolved = str(path.resolve()) if path.parent != Path(".") else shutil.which(candidate)
if not resolved or not Path(resolved).is_file() or not os.access(resolved, os.X_OK):
raise ClaudeIopValidationError("Claude binary is unavailable")
return str(Path(resolved).resolve())
class ClaudeIopAdapter:
"""Build safe Claude invocations and parse their IOP-bound stream evidence."""
capability = claude_capability()
def __init__(self, cell: MatrixCell, workspace: PreparedWorkspace, runtime: ClaudeIopRuntime) -> None:
if not isinstance(cell, MatrixCell) or cell.caller != "claude":
raise ClaudeIopValidationError("Claude adapter requires a Claude matrix cell")
if not isinstance(workspace, PreparedWorkspace) or not workspace.session_is_fresh:
raise ClaudeIopValidationError("Claude adapter requires a fresh prepared workspace")
if not isinstance(runtime, ClaudeIopRuntime):
raise ClaudeIopValidationError("invalid Claude runtime")
self.cell = cell
self.workspace = workspace
self.runtime = runtime
self.binary = resolve_claude_binary(runtime.binary)
self.base_url = _require_string(runtime.base_url, "IOP base URL")
self.api_key = _require_string(runtime.api_key, "IOP API key")
def preflight(self, raw_line: str) -> ConnectivityResult:
return make_result(self.cell, self.capability, parse_preflight_binding(raw_line, self.cell))
def parser(self) -> ClaudeStreamParser:
return ClaudeStreamParser(self.cell, self.workspace.session_id)
def redactor(self, task: str) -> Callable[[str], str]:
_require_string(task, "task")
structural_values = (task, self.base_url, self.api_key)
exact = exact_value_redactor(structural_values)
def _redact(line: str) -> str:
return exact(redact_claude_event(line, structural_values))
return _redact
def invocation(self, task: str, evidence_dir: str | Path, timeout: Timeout) -> InvocationSpec:
_require_string(task, "task")
if not isinstance(timeout, Timeout):
raise ClaudeIopValidationError("invalid invocation timeout")
evidence_path = _require_string(str(evidence_dir), "evidence directory")
cwd = Path(self.workspace.workspace_dir)
if not cwd.is_dir():
raise ClaudeIopValidationError("prepared workspace is unavailable")
env = (
("PATH", os.environ.get("PATH", "/usr/bin:/bin")),
("ANTHROPIC_BASE_URL", self.base_url),
("ANTHROPIC_API_KEY", self.api_key),
("CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1"),
("CLAUDE_CODE_DISABLE_AUTOUPDATER", "1"),
) + inherited_tls_ca_environment()
return InvocationSpec(
argv=(
self.binary, "--bare", "--print", "--verbose",
"--input-format", "text", "--output-format", "stream-json",
"--model", self.cell.iop.request_model, "--effort", self.cell.iop.requested_effort,
"--no-session-persistence", "--permission-mode", "dontAsk",
"--tools", "Read,Write,Edit", "--allowedTools", "Read,Write,Edit",
),
cwd=str(cwd),
env=env,
env_allowlist=(
"ANTHROPIC_BASE_URL", "ANTHROPIC_API_KEY",
"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "CLAUDE_CODE_DISABLE_AUTOUPDATER",
*TLS_CA_ENV_KEYS,
),
submission_mode=SUBMISSION_STDIN_ONCE,
completion_mode=COMPLETION_EXIT_AFTER_IDLE,
timeout=timeout,
evidence_dir=evidence_path,
task_payload=task.encode("utf-8"),
)