iop/scripts/agent_benchmark/agy_iop.py

"""Closed agy-to-IOP adapter for the comparison benchmark.

The adapter deliberately recognises only a documented agy transport contract.
In particular, it never inherits an ambient Gemini configuration: absent or
unknown transport support is an implementation gap before a caller process is
constructed.  The module is standard-library-only and has no network calls.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable
from urllib.parse import urlsplit

from scripts.agent_benchmark.connectivity import (
    ISSUE_RESUME_CODES,
    CallerCapability,
    ConnectivityResult,
    ConnectivityIssue,
    EffectiveBinding,
    RequestedEffectiveBinding,
    classify_issues,
    make_result,
)
from scripts.agent_benchmark.lifecycle import (
    COMPLETION_EXIT_AFTER_IDLE,
    SUBMISSION_STDIN_ONCE,
    InvocationResult,
    InvocationSpec,
    LifecycleMetricError,
    ParsedMetric,
    SupervisorLocator,
    duration_metric,
    env_pairs,
    exact_value_redactor,
    is_reported_number,
    run_invocation,
)
from scripts.agent_benchmark.manifest import MatrixCell, TOKEN_RE, Timeout
from scripts.agent_benchmark.workspace import PreparedWorkspace


AGY_CALLER = "agy"
AGY_KNOWN_VERSION = "1.1.11"
AGY_PROVIDER_ENV = "AGY_PROVIDER"
AGY_ENDPOINT_ENV = "AGY_OPENAI_BASE_URL"
AGY_AUTH_ENV = "AGY_OPENAI_API_KEY"
_VERSION_RE = re.compile(r"(?:agy\s+)?(\d+\.\d+\.\d+)", re.IGNORECASE)
_SAFE_EVENT_FIELDS = ("type", "subtype", "model", "effort", "route_kind", "route_id")
_DOCUMENTED_OPTIONS = ("--print", "--output-format", "--sandbox", "--model", "--effort")
_DOCUMENTED_ENVIRONMENT = (AGY_PROVIDER_ENV, AGY_ENDPOINT_ENV, AGY_AUTH_ENV)
AGY_SAFE_METRIC_LABELS = ("metric:duration_ms",)
# agy reports whole durations in milliseconds.  Only these subtypes are
# converted, and a model-stage duration is marked as overlapping because it is
# reported inside the same window as the total.
AGY_DURATION_METRICS = {
    "duration_ms": ("total_duration", False),
    "model_duration_ms": ("model_duration", True),
    "queue_duration_ms": ("queue_duration", False),
}
_AGY_METRIC_KEYS = frozenset({"type", "subtype", "value", "model"})
_IDENTITY_RE = re.compile(r"sha256:[0-9a-f]{64}\Z")


def _exact_token_present(text: str, token: str) -> bool:
    """Match one documented help token, never a prefix or a suffix."""
    return re.search(r"(?<![A-Za-z0-9_-])" + re.escape(token) + r"(?![A-Za-z0-9_-])", text) is not None


@dataclass(frozen=True)
class AgyRuntimeInputs:
    """Private runtime values supplied by a preflight owner, never persisted."""

    binary: str
    endpoint: str
    credential: str


@dataclass(frozen=True)
class AgyRuntimeObservation:
    """Secret-free IOP configuration proof for exactly one benchmark cell."""

    cell_id: str
    route_kind: str
    route_id: str
    endpoint_identity: str
    credential_identity: str
    config_identity: str


@dataclass(frozen=True)
class _ValidatedAgyRuntime:
    """Private launch values admitted only after IOP configuration validation."""

    binary: str
    endpoint: str
    credential: str
    observation: AgyRuntimeObservation


@dataclass(frozen=True)
class AgyDocumentedCapabilities:
    """Exact, known-version tokens parsed from public agy help output."""

    options: tuple[str, ...]
    environment: tuple[str, ...]
    output_formats: tuple[str, ...]


@dataclass(frozen=True)
class AgyCapability:
    """Versioned, documented capability observation; no caller is launched."""

    version: str | None
    iop_transport_supported: bool
    endpoint_supported: bool
    auth_supported: bool
    protocol_supported: bool
    stream_supported: bool
    route_kinds: tuple[str, ...]
    efforts: tuple[str, ...]


@dataclass(frozen=True)
class AgyPreflightResult:
    """Closed outcome before invocation construction."""

    capability: AgyCapability
    binding: RequestedEffectiveBinding
    issues: tuple[ConnectivityIssue, ...]
    status: str
    runtime: _ValidatedAgyRuntime | None


class AgyAdapterError(Exception):
    """Raised when a caller launch is requested without a proven transport."""


def _issue(code: str) -> ConnectivityIssue:
    return ConnectivityIssue(code, ISSUE_RESUME_CODES[code])


def _requested_binding(cell: MatrixCell) -> RequestedEffectiveBinding:
    return RequestedEffectiveBinding(
        cell.id,
        cell.caller,
        cell.iop.route_kind,
        cell.iop.route_id,
        cell.iop.request_model,
        cell.iop.requested_effort,
    )


def parse_documented_agy_capabilities(help_output: str) -> AgyDocumentedCapabilities:
    """Parse only complete documented tokens from a known agy help surface."""
    if not isinstance(help_output, str):
        return AgyDocumentedCapabilities((), (), ())
    return AgyDocumentedCapabilities(
        tuple(token for token in _DOCUMENTED_OPTIONS if _exact_token_present(help_output, token)),
        tuple(token for token in _DOCUMENTED_ENVIRONMENT if _exact_token_present(help_output, token)),
        ("stream-json",) if _exact_token_present(help_output, "stream-json") else (),
    )


def inspect_agy_iop_capability(version_output: str, help_output: str) -> AgyCapability:
    """Inspect only public, versioned help text for the closed IOP transport.

    A version string is accepted only when it names the known agy release and
    every required transport variable is documented.  This prevents a new or
    partially documented client from silently inheriting ambient provider state.
    """
    if not isinstance(version_output, str) or not isinstance(help_output, str):
        return AgyCapability(None, False, False, False, False, False, (), ())
    matched = _VERSION_RE.fullmatch(version_output.strip())
    version = matched.group(1) if matched else None
    known_version = version == AGY_KNOWN_VERSION
    documented = parse_documented_agy_capabilities(help_output)
    endpoint_supported = AGY_ENDPOINT_ENV in documented.environment
    auth_supported = AGY_AUTH_ENV in documented.environment
    protocol_supported = known_version and all(
        option in documented.options for option in _DOCUMENTED_OPTIONS
    ) and AGY_PROVIDER_ENV in documented.environment
    stream_supported = "stream-json" in documented.output_formats
    supported = endpoint_supported and auth_supported and protocol_supported and stream_supported
    if not supported:
        return AgyCapability(
            version, False, endpoint_supported, auth_supported, protocol_supported, stream_supported, (), ()
        )
    return AgyCapability(
        version,
        True,
        True,
        True,
        True,
        True,
        ("direct", "execution_preset"),
        ("high", "low", "medium"),
    )


def _runtime_identity(label: str, value: str) -> str:
    return "sha256:" + hashlib.sha256(
        b"agy-iop-runtime-v1\0" + label.encode("ascii") + b"\0" + value.encode("utf-8")
    ).hexdigest()


def _validate_config_owner_observation(
    cell: MatrixCell, observation: AgyRuntimeObservation
) -> None:
    """Validate an opaque observation supplied by the independent config owner."""
    if not isinstance(observation, AgyRuntimeObservation):
        raise AgyAdapterError("agy runtime observation is invalid")
    if (observation.cell_id, observation.route_kind, observation.route_id) != (
        cell.id, cell.iop.route_kind, cell.iop.route_id,
    ):
        raise AgyAdapterError("agy IOP runtime observation mismatch")
    if not all(_IDENTITY_RE.fullmatch(value) for value in (
        observation.endpoint_identity,
        observation.credential_identity,
        observation.config_identity,
    )):
        raise AgyAdapterError("agy IOP runtime observation identity is invalid")


def validate_agy_iop_runtime(
    cell: MatrixCell, runtime: AgyRuntimeInputs, observation: AgyRuntimeObservation
) -> _ValidatedAgyRuntime:
    """Admit runtime-only launch values only when IOP identity proof is exact."""
    if not isinstance(cell, MatrixCell) or cell.caller != AGY_CALLER:
        raise AgyAdapterError("agy runtime validation requires an agy matrix cell")
    if not isinstance(runtime, AgyRuntimeInputs) or not isinstance(observation, AgyRuntimeObservation):
        raise AgyAdapterError("agy runtime inputs are invalid")
    if not all(isinstance(value, str) and value for value in (runtime.binary, runtime.endpoint, runtime.credential)):
        raise AgyAdapterError("agy runtime values are unavailable")
    endpoint = urlsplit(runtime.endpoint)
    if endpoint.scheme not in ("http", "https") or not endpoint.netloc or endpoint.query or endpoint.fragment:
        raise AgyAdapterError("agy IOP endpoint is invalid")
    _validate_config_owner_observation(cell, observation)
    if observation.endpoint_identity != _runtime_identity("endpoint", runtime.endpoint):
        raise AgyAdapterError("agy IOP runtime observation mismatch")
    if observation.credential_identity != _runtime_identity("credential", runtime.credential):
        raise AgyAdapterError("agy IOP runtime observation mismatch")
    return _ValidatedAgyRuntime(runtime.binary, runtime.endpoint, runtime.credential, observation)


def preflight_agy_iop(
    cell: MatrixCell,
    capability: AgyCapability,
    runtime: AgyRuntimeInputs,
    observation: AgyRuntimeObservation,
) -> AgyPreflightResult:
    """Classify only registration and implementation gaps without launching agy."""
    if not isinstance(cell, MatrixCell) or cell.caller != AGY_CALLER:
        raise AgyAdapterError("agy preflight requires an agy matrix cell")
    if not isinstance(capability, AgyCapability) or not isinstance(runtime, AgyRuntimeInputs):
        raise AgyAdapterError("agy preflight inputs are invalid")

    issues: list[ConnectivityIssue] = []
    validated_runtime: _ValidatedAgyRuntime | None = None
    if runtime.credential:
        try:
            validated_runtime = validate_agy_iop_runtime(cell, runtime, observation)
        except AgyAdapterError:
            issues.append(_issue("endpoint_incompatible"))
    if not runtime.credential:
        issues.append(_issue("credential_missing"))
    if not cell.iop.request_model:
        issues.append(_issue("model_missing"))
    if not runtime.endpoint:
        issues.append(_issue("endpoint_incompatible"))
    if not capability.endpoint_supported:
        issues.append(_issue("endpoint_incompatible"))
    if not capability.auth_supported:
        issues.append(_issue("auth_incompatible"))
    if not capability.protocol_supported:
        issues.append(_issue("protocol_incompatible"))
    if not capability.stream_supported:
        issues.append(_issue("stream_incompatible"))
    if capability.iop_transport_supported and cell.iop.route_kind not in capability.route_kinds:
        issues.append(_issue("protocol_incompatible"))
    elif capability.iop_transport_supported and cell.iop.requested_effort not in capability.efforts:
        issues.append(_issue("effort_unsupported"))

    # Preserve connectivity.py's canonical issue order without leaking values.
    unique = {item.code: item for item in issues}
    ordered = tuple(
        unique[code]
        for code in (
            "credential_missing", "model_missing", "route_missing", "effort_unsupported",
            "endpoint_incompatible", "auth_incompatible", "protocol_incompatible", "stream_incompatible",
        )
        if code in unique
    )
    status = classify_issues(ordered)
    return AgyPreflightResult(capability, _requested_binding(cell), ordered, status, validated_runtime)


def build_agy_invocation(
    cell: MatrixCell,
    prepared: PreparedWorkspace,
    task_payload: bytes,
    timeout: Timeout,
    preflight: AgyPreflightResult,
) -> InvocationSpec:
    """Build one isolated stdin-only agy invocation after a ready preflight."""
    if preflight.status != "ready" or not preflight.capability.iop_transport_supported or preflight.runtime is None:
        raise AgyAdapterError("agy IOP transport is not proven")
    runtime = preflight.runtime
    if not runtime.binary or not Path(runtime.binary).is_file():
        raise AgyAdapterError("agy binary is unavailable")
    if not isinstance(prepared, PreparedWorkspace) or not prepared.workspace_dir:
        raise AgyAdapterError("prepared workspace is unavailable")
    if not isinstance(task_payload, bytes) or not task_payload:
        raise AgyAdapterError("agy task payload is unavailable")

    # The child receives a minimal environment and explicit IOP-only provider
    # settings.  No parent agy/Gemini config or session variable is inherited.
    environment = {
        "PATH": os.environ.get("PATH", "/usr/bin:/bin"),
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
        "TZ": "UTC",
        AGY_PROVIDER_ENV: "iop-openai",
        AGY_ENDPOINT_ENV: runtime.endpoint,
        AGY_AUTH_ENV: runtime.credential,
    }
    return InvocationSpec(
        argv=(
            runtime.binary,
            "--print",
            "--sandbox",
            "--output-format", "stream-json",
            "--model", cell.iop.request_model,
            "--effort", cell.iop.requested_effort,
        ),
        cwd=prepared.workspace_dir,
        env=env_pairs(environment),
        env_allowlist=(AGY_PROVIDER_ENV, AGY_ENDPOINT_ENV, AGY_AUTH_ENV),
        submission_mode=SUBMISSION_STDIN_ONCE,
        completion_mode=COMPLETION_EXIT_AFTER_IDLE,
        timeout=timeout,
        evidence_dir=prepared.attempt_root,
        task_payload=task_payload,
        control_dir=str(Path(prepared.attempt_root) / "agy-control"),
    )


def _safe_identifier(value: Any) -> str | None:
    return value if isinstance(value, str) and TOKEN_RE.fullmatch(value) else None


def redact_agy_event(raw_line: str, sensitive_values: tuple[str, ...] = ()) -> str:
    """Return a canonical allowlisted event projection, never caller content."""
    try:
        parsed = json.loads(raw_line)
    except (TypeError, json.JSONDecodeError):
        return '{"event":"unparseable"}'
    if not isinstance(parsed, dict):
        return '{"event":"unparseable"}'
    safe: dict[str, str] = {}
    for field in _SAFE_EVENT_FIELDS:
        value = _safe_identifier(parsed.get(field))
        if value is not None and value not in sensitive_values:
            safe[field] = value
    if "type" not in safe or "subtype" not in safe:
        return '{"event":"unparseable"}'
    return json.dumps(safe, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


class AgyEventParser:
    """Strict stream-json parser bound to exactly one requested IOP cell."""

    def __init__(self, cell: MatrixCell) -> None:
        self._cell = cell
        self._observed_binding: RequestedEffectiveBinding | None = None
        self._binding_invalid = False

    def __call__(self, stream: str, raw_line: str) -> str | ParsedMetric | None:
        return self.parse(stream, raw_line)

    def parse(self, stream: str, raw_line: str) -> str | ParsedMetric | None:
        if stream != "stdout":
            return None
        try:
            event = json.loads(raw_line)
        except (TypeError, json.JSONDecodeError):
            return "malformed"
        if not isinstance(event, dict):
            return "malformed"
        event_type = event.get("type")
        subtype = event.get("subtype")
        if (event_type, subtype) == ("iop", "effective_binding"):
            self._observe_effective_binding(event)
            return None
        if event_type == "metric":
            return self._observe_duration(event, subtype)
        if event_type == "result" and subtype == "error":
            # Quota/provider errors can never be interpreted as finish/idle.
            return "quota_error" if event.get("reason") == "quota" else "malformed"
        terminal = (
            "finish" if (event_type, subtype) == ("result", "success")
            else "idle" if (event_type, subtype) == ("system", "idle")
            else None
        )
        if terminal is None or not self._matches_exact_binding(event):
            return "malformed"
        return terminal

    def _observe_duration(self, event: dict[str, Any], subtype: Any) -> str | ParsedMetric:
        """Convert one allowlisted agy duration losslessly, or fail closed."""
        mapped = AGY_DURATION_METRICS.get(subtype) if isinstance(subtype, str) else None
        value = event.get("value")
        if (
            mapped is None
            or not set(event) <= _AGY_METRIC_KEYS
            or not is_reported_number(value)
            or ("model" in event and event["model"] != self._cell.iop.request_model)
        ):
            return "malformed"
        name, overlap = mapped
        try:
            return duration_metric(
                name, value, reported_unit="ms",
                model=self._cell.iop.request_model, overlap=overlap,
            )
        except LifecycleMetricError:
            return "malformed"

    def _matches_exact_binding(self, event: dict[str, Any]) -> bool:
        expected = self._cell.iop
        return (
            event.get("model") == expected.request_model
            and event.get("effort") == expected.requested_effort
            and event.get("route_kind") == expected.route_kind
            and event.get("route_id") == expected.route_id
        )

    def _observe_effective_binding(self, event: dict[str, Any]) -> None:
        expected_keys = {"type", "subtype", "route_kind", "route_id", "model", "effort", "stages"}
        if set(event) != expected_keys or self._observed_binding is not None:
            self._binding_invalid = True
            return
        values = tuple(event[key] for key in ("route_kind", "route_id", "model", "effort"))
        stages = event.get("stages")
        if not all(isinstance(value, str) and TOKEN_RE.fullmatch(value) for value in values) or not isinstance(stages, list):
            self._binding_invalid = True
            return
        parsed_stages: list[EffectiveBinding] = []
        for stage in stages:
            if not isinstance(stage, dict) or set(stage) != {"stage", "model", "effort"}:
                self._binding_invalid = True
                return
            if not isinstance(stage["stage"], str) or not isinstance(stage["model"], str):
                self._binding_invalid = True
                return
            if stage["effort"] is not None and not isinstance(stage["effort"], str):
                self._binding_invalid = True
                return
            parsed_stages.append(EffectiveBinding(stage["stage"], stage["model"], stage["effort"]))
        self._observed_binding = RequestedEffectiveBinding(
            self._cell.id, self._cell.caller,
            self._cell.iop.route_kind, self._cell.iop.route_id,
            self._cell.iop.request_model, self._cell.iop.requested_effort,
            values[0], values[1], values[2], values[3], tuple(parsed_stages),
        )

    def observed_result(self, capability: AgyCapability, lifecycle: InvocationResult) -> ConnectivityResult:
        """Report ready only for successful lifecycle-owned explicit evidence."""
        requested = _requested_binding(self._cell)
        closed_gap = (_issue("stream_incompatible"),)
        caller_capability = CallerCapability(AGY_CALLER, capability.route_kinds, capability.efforts)
        if (
            not isinstance(lifecycle, InvocationResult)
            or not lifecycle.success
            or not lifecycle.finish_then_idle_then_quiet
            or self._binding_invalid
            or self._observed_binding is None
        ):
            return make_result(self._cell, caller_capability, requested, closed_gap)
        try:
            return make_result(self._cell, caller_capability, self._observed_binding)
        except Exception:
            return make_result(self._cell, caller_capability, requested, closed_gap)


def run_agy_invocation(
    spec: InvocationSpec,
    parser: AgyEventParser,
    preflight: AgyPreflightResult,
    on_started: Callable[[SupervisorLocator], None],
) -> InvocationResult:
    """Run a prepared agy call with structural output redaction only."""
    if (
        not isinstance(preflight, AgyPreflightResult)
        or preflight.status != "ready"
        or not preflight.capability.iop_transport_supported
        or preflight.runtime is None
    ):
        raise AgyAdapterError("agy IOP transport is not proven")
    runtime = preflight.runtime
    sensitive = (runtime.endpoint, runtime.credential)
    structural = lambda line: redact_agy_event(line, sensitive)
    # exact replacement is retained as a final defence for non-JSON stderr.
    exact = exact_value_redactor(sensitive)
    return run_invocation(
        spec,
        parse_event=parser,
        on_started=on_started,
        redact=lambda line: line if line in AGY_SAFE_METRIC_LABELS else structural(exact(line)),
    )