iop/scripts/agent_benchmark/claude_iop.py

"""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,
    exact_value_redactor,
)
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",
    }
)


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"

    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) or message.get("stop_reason") != "end_turn":
            raise ClaudeIopProtocolError("invalid Claude assistant terminal")
        if _required_string(message, "model") != self.cell.iop.request_model:
            raise ClaudeIopProtocolError("Claude model binding mismatch")
        self._phase = "await_result"
        return "finish"

    def _consume_result(self, event: dict[str, Any]) -> str:
        if self._phase != "await_result":
            raise ClaudeIopProtocolError("duplicate or out-of-order Claude result")
        self._require_bound_session(event)
        if event.get("subtype") != "success":
            raise ClaudeIopProtocolError("invalid Claude result terminal")
        self._phase = "complete"
        return "idle"

    def __call__(self, stream: str, raw_line: str) -> str | 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 == "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"),
        )
        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=",
            ),
            cwd=str(cwd),
            env=env,
            env_allowlist=(
                "ANTHROPIC_BASE_URL", "ANTHROPIC_API_KEY",
                "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "CLAUDE_CODE_DISABLE_AUTOUPDATER",
            ),
            submission_mode=SUBMISSION_STDIN_ONCE,
            completion_mode=COMPLETION_EXIT_AFTER_IDLE,
            timeout=timeout,
            evidence_dir=evidence_path,
            task_payload=task.encode("utf-8"),
        )