iop/scripts/agent_benchmark/agy_iop.py

550 lines
21 KiB
Python

"""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,
RequestedEffectiveBinding,
classify_issues,
make_result,
)
from scripts.agent_benchmark.lifecycle import (
COMPLETION_EXIT_AFTER_IDLE,
SUBMISSION_ARGV_TASK,
InvocationResult,
InvocationSpec,
LifecycleMetricError,
ParsedMetric,
SupervisorLocator,
TLS_CA_ENV_KEYS,
count_metric,
duration_metric,
env_pairs,
exact_value_redactor,
inherited_tls_ca_environment,
run_invocation,
)
from scripts.agent_benchmark.manifest import MatrixCell, Timeout
from scripts.agent_benchmark.workspace import PreparedWorkspace
AGY_CALLER = "agy"
AGY_KNOWN_VERSION = "1.1.12"
AGY_ENDPOINT_ENV = "GOOGLE_GEMINI_BASE_URL"
AGY_AUTH_ENV = "GEMINI_API_KEY"
AGY_SETTINGS_RELATIVE_PATH = Path(".gemini/antigravity-cli/settings.json")
_AGY_ISOLATED_SETTINGS = {
"enableTelemetry": False,
"modelProvider": "gemini",
"toolPermission": "always-proceed",
}
AGY_MODEL_LABELS = {
"gemini-3.6-flash": "Gemini 3.6 Flash",
"gemini-hybrid": "Gemini 3.6 Flash",
}
_VERSION_RE = re.compile(r"(?:agy\s+)?(\d+\.\d+\.\d+)", re.IGNORECASE)
_SAFE_EVENT_FIELDS = ("event", "state", "step_type", "status")
_DOCUMENTED_OPTIONS = ("--print", "--output-format", "--sandbox", "--model")
_AGY_USAGE_METRICS = {
"input_tokens": "input_tokens",
"cache_read_tokens": "cached_input_tokens",
"output_tokens": "output_tokens",
"thinking_tokens": "reasoning_tokens",
"total_tokens": "total_tokens",
}
_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)),
(),
("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.
The known release is pinned because its Gemini provider environment is not
printed by ``--help``. The public options and stream format still have to
match exactly; a changed release fails closed until re-qualified.
"""
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 = known_version
auth_supported = known_version
protocol_supported = known_version and all(
option in documented.options for option in _DOCUMENTED_OPTIONS
)
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)
route_path = f"/gemini/{cell.iop.route_id}"
if (
endpoint.scheme != "https"
or not endpoint.netloc
or endpoint.query
or endpoint.fragment
or endpoint.path.rstrip("/") != route_path
):
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 cell.iop.request_model not in AGY_MODEL_LABELS:
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 official agy print 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 Path(prepared.workspace_dir).is_dir()
or not Path(prepared.session_dir).is_dir()
or not Path(prepared.attempt_root).is_dir()
):
raise AgyAdapterError("prepared workspace is unavailable")
if not isinstance(task_payload, bytes) or not task_payload:
raise AgyAdapterError("agy task payload is unavailable")
try:
task_text = task_payload.decode("utf-8")
except UnicodeDecodeError as exc:
raise AgyAdapterError("agy task payload must be UTF-8") from exc
_stage_agy_provider_settings(prepared.session_dir)
# 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",
"HOME": prepared.session_dir,
AGY_ENDPOINT_ENV: runtime.endpoint,
AGY_AUTH_ENV: runtime.credential,
}
environment.update(inherited_tls_ca_environment())
return InvocationSpec(
argv=(
runtime.binary,
"--sandbox",
"--output-format", "stream-json",
"--model", AGY_MODEL_LABELS[cell.iop.request_model],
"--print", task_text,
),
cwd=prepared.workspace_dir,
env=env_pairs(environment),
env_allowlist=(AGY_ENDPOINT_ENV, AGY_AUTH_ENV, *TLS_CA_ENV_KEYS),
submission_mode=SUBMISSION_ARGV_TASK,
completion_mode=COMPLETION_EXIT_AFTER_IDLE,
timeout=timeout,
evidence_dir=prepared.attempt_root,
control_dir=str(Path(prepared.attempt_root) / "agy-control"),
)
def _stage_agy_provider_settings(session_dir: str) -> None:
"""Create the exact secret-free provider selector in one isolated HOME."""
try:
session = Path(session_dir).resolve(strict=True)
if not session.is_dir():
raise OSError("session is not a directory")
settings = session / AGY_SETTINGS_RELATIVE_PATH
settings.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
if settings.parent.resolve(strict=True) != session / AGY_SETTINGS_RELATIVE_PATH.parent:
raise OSError("settings directory escapes session")
payload = (
json.dumps(_AGY_ISOLATED_SETTINGS, sort_keys=True, separators=(",", ":")) + "\n"
).encode("utf-8")
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
if hasattr(os, "O_NOFOLLOW"):
flags |= os.O_NOFOLLOW
descriptor = os.open(settings, flags, 0o600)
with os.fdopen(descriptor, "wb") as handle:
handle.write(payload)
except (OSError, ValueError) as exc:
raise AgyAdapterError("agy isolated provider settings are unavailable") from exc
def _safe_identifier(value: Any) -> str | None:
return value if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:+-]{0,63}", 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"}'
event = _safe_identifier(parsed.get("event"))
if event is None or event in sensitive_values:
return '{"event":"unparseable"}'
safe: dict[str, str] = {"event": event}
payload = parsed.get(event)
if not isinstance(payload, dict):
payload = {}
for field in _SAFE_EVENT_FIELDS[1:]:
value = _safe_identifier(payload.get(field))
if value is not None and value not in sensitive_values:
safe[field] = value
return json.dumps(safe, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
class AgyEventParser:
"""Parse the official 1.1.12 stream while trusting binding only from config."""
def __init__(self, cell: MatrixCell, admitted_binding: RequestedEffectiveBinding) -> None:
if not isinstance(admitted_binding, RequestedEffectiveBinding):
raise AgyAdapterError("agy admitted binding is invalid")
if admitted_binding.cell_id != cell.id or admitted_binding.caller != cell.caller:
raise AgyAdapterError("agy admitted binding mismatch")
self._cell = cell
self._admitted_binding = admitted_binding
self._init_seen = False
self._result_seen = False
self._latest_usage: dict[str, Any] | None = None
def __call__(self, stream: str, raw_line: str) -> str | ParsedMetric | tuple[Any, ...] | None:
return self.parse(stream, raw_line)
def parse(self, stream: str, raw_line: str) -> str | ParsedMetric | tuple[Any, ...] | None:
if stream != "stdout":
return None
try:
item = json.loads(raw_line)
except (TypeError, json.JSONDecodeError):
return "malformed"
if not isinstance(item, dict):
return "malformed"
event = item.get("event")
payload = item.get(event) if isinstance(event, str) else None
if not isinstance(payload, dict):
return "malformed"
if event == "init":
if self._init_seen or self._result_seen:
return "malformed"
self._init_seen = True
return None
if event == "step_update":
if not self._init_seen or self._result_seen:
return "malformed"
usage = payload.get("usage")
if usage is not None:
if self._usage_metrics(usage) is None:
return "malformed"
self._latest_usage = usage
return None
if event != "result" or not self._init_seen or self._result_seen:
return "malformed"
self._result_seen = True
if payload.get("status") != "SUCCESS":
return "malformed"
metrics: list[ParsedMetric] = []
usage = payload.get("usage", self._latest_usage)
if usage is not None:
parsed_usage = self._usage_metrics(usage)
if parsed_usage is None:
return "malformed"
metrics.extend(parsed_usage)
try:
if "duration_seconds" in payload:
metrics.append(duration_metric(
"total_duration", payload["duration_seconds"], reported_unit="s",
model=self._cell.iop.request_model,
))
if "num_turns" in payload:
metrics.append(count_metric(
"model_calls", payload["num_turns"], model=self._cell.iop.request_model,
))
except LifecycleMetricError:
return "malformed"
return tuple(metrics) + ("finish", "idle")
def _usage_metrics(self, usage: Any) -> tuple[ParsedMetric, ...] | None:
if not isinstance(usage, dict) or not set(usage) <= set(_AGY_USAGE_METRICS):
return None
metrics: list[ParsedMetric] = []
try:
for wire_name, metric_name in _AGY_USAGE_METRICS.items():
if wire_name in usage:
metrics.append(count_metric(
metric_name, usage[wire_name], model=self._cell.iop.request_model,
))
except LifecycleMetricError:
return None
return tuple(metrics)
def observed_result(self, capability: AgyCapability, lifecycle: InvocationResult) -> ConnectivityResult:
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 not self._result_seen
):
return make_result(self._cell, caller_capability, requested, closed_gap)
try:
return make_result(self._cell, caller_capability, self._admitted_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: structural(exact(line)),
)