From 22861954d1aead90b478a6902b2ccdc7a246ef80 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Fri, 28 Aug 2026 22:29:12 -0700 Subject: [PATCH 01/74] feat: capture LLM trajectories across auth modes --- docs/agent-quickstart.md | 5 +- docs/getting-started.md | 21 +- docs/running-any-benchmark.md | 3 +- src/benchflow/cli/main.py | 8 +- src/benchflow/rollout/__init__.py | 60 ++ src/benchflow/trajectories/__init__.py | 5 +- .../trajectories/export_prime_sft.py | 16 + src/benchflow/trajectories/llm_capture.py | 578 +++++++++++++ .../trajectories/llm_capture_manifest.py | 111 +++ .../trajectories/native_capture_parsers.py | 735 +++++++++++++++++ src/benchflow/trajectories/results.py | 65 +- tests/trajectories/test_export_prime_sft.py | 18 + tests/trajectories/test_native_llm_capture.py | 771 ++++++++++++++++++ 13 files changed, 2380 insertions(+), 16 deletions(-) create mode 100644 src/benchflow/trajectories/llm_capture.py create mode 100644 src/benchflow/trajectories/llm_capture_manifest.py create mode 100644 src/benchflow/trajectories/native_capture_parsers.py create mode 100644 tests/trajectories/test_native_llm_capture.py diff --git a/docs/agent-quickstart.md b/docs/agent-quickstart.md index 8eb6485f0..410fd0533 100644 --- a/docs/agent-quickstart.md +++ b/docs/agent-quickstart.md @@ -138,8 +138,9 @@ and what it is: agent/ — agent-side logs trajectory/acp_trajectory.jsonl — the full agent trace (every ACP event: prompts, tool calls, outputs) - trajectory/llm_trajectory.jsonl — raw provider requests/responses captured - by the usage-tracking proxy + trajectory/llm_trajectory.jsonl — always-present LLM exchange log + trajectory/llm_trajectory.manifest.json + — source/fidelity/completeness for that log trainer/verifiers.jsonl — trainer-ready scored trajectory record trainer/atif.json — the trajectory in ATIF interchange format (omitted if the trajectory is empty) diff --git a/docs/getting-started.md b/docs/getting-started.md index 624500025..4b2e15b4c 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -202,7 +202,8 @@ Each run writes under `--jobs-dir` (default `jobs/`): prompts.json # prompts sent to the agent trajectory/ acp_trajectory.jsonl # full agent trace (ACP events) - llm_trajectory.jsonl # raw provider requests/responses (when the usage-tracking proxy captured exchanges) + llm_trajectory.jsonl # always present; LLM exchanges at the fidelity described below + llm_trajectory.manifest.json # capture source, fidelity, completeness, and errors trainer/ verifiers.jsonl # trainer-ready scored trajectory (Verifiers/ORS record) atif.json # ATIF trajectory record (omitted if the trajectory is empty) @@ -213,6 +214,24 @@ Each run writes under `--jobs-dir` (default `jobs/`): test-stdout.txt # verifier stdout ``` +`llm_trajectory.jsonl` is created when the rollout directory is initialized, +including for setup failures and tasks that make no model call. Its sidecar is +the source of truth for interpreting the JSONL: + +| Agent/auth path | Primary source | `capture_fidelity` | +|---|---|---| +| API key through the BenchFlow gateway (including Azure) | LiteLLM provider request/response capture | `provider_wire` | +| Claude Code subscription/OAuth | Claude Code raw API-body files correlated by local OTLP logs | `provider_wire` | +| Claude Code subscription/OAuth fallback | Claude Code native session JSONL | `agent_session` | +| Codex subscription/OAuth | Codex native session JSONL | `agent_session` | + +The manifest status is `complete`, `partial`, `no_model_call`, or +`capture_failed`. Reconstructed `agent_session` rows remain useful for audit and +viewer workflows, but trainer exports fail closed unless the manifest says the +capture is complete provider-wire data. Claude's own raw-body telemetry can +still contain provider-redacted extended-thinking blocks; BenchFlow also applies +its normal secret redaction before publishing the JSONL. + ### Reading results Exit code 0 means the pipeline completed — it is not a pass/fail signal. A diff --git a/docs/running-any-benchmark.md b/docs/running-any-benchmark.md index 3ad809203..da45a2777 100644 --- a/docs/running-any-benchmark.md +++ b/docs/running-any-benchmark.md @@ -217,7 +217,8 @@ Every layer terminates at the *same* output contract, written per rollout under | `results.jsonl` | Verifiers/Prime-RL shaped rollout row | | `rewards.jsonl` | The reward record for the rollout (ORS / OpenReward shape) | | `trajectory/acp_trajectory.jsonl` | Full agent trace as ACP events | -| `trajectory/llm_trajectory.jsonl` | Raw provider requests/responses (when captured) | +| `trajectory/llm_trajectory.jsonl` | Always-present LLM exchange log (possibly empty for `no_model_call`) | +| `trajectory/llm_trajectory.manifest.json` | Capture source, fidelity, completeness, and errors | | `trainer/verifiers.jsonl` | Trainer-ready scored trajectory (Verifiers record) | | `trainer/atif.json` | ATIF trajectory record | | `trainer/adp.jsonl` | ADP trajectory record | diff --git a/src/benchflow/cli/main.py b/src/benchflow/cli/main.py index 266e74ffa..f0e5a4972 100644 --- a/src/benchflow/cli/main.py +++ b/src/benchflow/cli/main.py @@ -286,9 +286,11 @@ def eval_run( "--usage-tracking", help=( "Telemetry-enforcement policy: auto, required, or off. The " - "LiteLLM proxy is always used for routable agents (usage, cost, " - "and llm_trajectory.jsonl are always captured); this flag only " - "controls whether trusted telemetry is required." + "LiteLLM proxy is always used for routable API-key agents; " + "native subscription agents use their own telemetry/session " + "surface. llm_trajectory.jsonl is always emitted, and its " + "manifest records capture fidelity. This flag only controls " + "whether trusted usage telemetry is required." ), ), ] = None, diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 8e73b0f96..0ca690a37 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -219,6 +219,10 @@ make_trajectory_sink, ) from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter +from benchflow.trajectories.llm_capture import ( + LLMTrajectoryCapture, + model_call_seen_from_evidence, +) from benchflow.trajectories.tree import RolloutNode, RolloutTree, Step from benchflow.usage_tracking import ( USAGE_SOURCE_AGENT_NATIVE_ACP, @@ -645,6 +649,7 @@ def __init__(self, config: RolloutConfig) -> None: self._started_at: datetime | None = None self._job_name: str | None = None self._rollout_name: str | None = None + self._llm_capture: LLMTrajectoryCapture | None = None self._agent_env: dict[str, str] = {} self._resolved_prompts: list[str] = [] self._agent_launch: str = "" @@ -938,6 +943,17 @@ async def setup(self) -> None: self._rollout_name, ) = _init_rollout(cfg.task_path, cfg.job_name, cfg.rollout_name, cfg.jobs_dir) + # Create the artifact contract as soon as the rollout directory exists. + # Even setup/start failures therefore leave a valid (possibly empty) + # JSONL plus a sidecar explaining why no exchange was captured. + self._llm_capture = LLMTrajectoryCapture( + self._rollout_dir, + agent=cfg.primary_agent, + model=cfg.primary_model, + session_id=self._rollout_name or "", + started_at=self._started_at, + ) + # C-axis overlay: deep-merge cfg.config_override into the task's resolved # config here at the rollout layer (not in the Task constructor), so only # this run's tasks are patched and every downstream read sees it. No-op @@ -958,6 +974,7 @@ async def setup(self) -> None: ), disallow=self._disallow_web_tools, ) + self._llm_capture.configure(self._agent_env) env_config = getattr(getattr(self._task, "config", None), "sandbox", None) task_skill_policy = resolve_task_skill_policy( task_path=cfg.task_path, @@ -1225,6 +1242,16 @@ async def install_agent(self) -> None: await self._planes.upload_subscription_auth( self._env, agent_name, cred_home ) + llm_capture = getattr(self, "_llm_capture", None) + if llm_capture is not None: + self._agent_env = await llm_capture.prepare_agent( + self._env, + agent=agent_name, + model=cfg.primary_model, + agent_env=self._agent_env, + credential_home=cred_home, + sandbox_user=cfg.sandbox_user, + ) await self._planes.apply_web_tool_policy( self._env, agent_name, @@ -2010,6 +2037,27 @@ async def cleanup(self) -> None: self._usage_runtime = None self._finalize_usage_metrics() + llm_capture = getattr(self, "_llm_capture", None) + if llm_capture is not None: + acp_events = list(getattr(self, "_trajectory", None) or []) + model_call_seen = model_call_seen_from_evidence( + getattr(self, "_usage_metrics", None), acp_events + ) + try: + await llm_capture.finalize( + self._env, + acp_events=acp_events, + model_call_seen=model_call_seen, + ) + except Exception as e: + logger.warning(f"LLM trajectory finalization failed: {e}") + try: + llm_capture.record_failure(e, model_call_seen=model_call_seen) + except Exception as record_error: + logger.warning( + "LLM trajectory failure manifest write failed: %s", + record_error, + ) self._enforce_required_usage_tracking() if self._environment is not None: @@ -2341,6 +2389,18 @@ async def connect_as(self, role: Role) -> None: disallow=disallow_web_tools, ) + llm_capture = getattr(self, "_llm_capture", None) + if llm_capture is not None: + cred_home = f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" + agent_env = await llm_capture.prepare_agent( + self._env, + agent=role.agent, + model=role.model, + agent_env=agent_env, + credential_home=cred_home, + sandbox_user=cfg.sandbox_user, + ) + self._agent_launch = agent_launch sf_entrypoint = self._session_factory_entrypoint(role.agent) diff --git a/src/benchflow/trajectories/__init__.py b/src/benchflow/trajectories/__init__.py index 3565aeadd..84b1ba5fd 100644 --- a/src/benchflow/trajectories/__init__.py +++ b/src/benchflow/trajectories/__init__.py @@ -1,7 +1,8 @@ """Trajectory capture and exchange schemas. -BenchFlow persists ACP-native trajectories plus LiteLLM callback-derived LLM -request/response exchanges. +BenchFlow persists ACP-native trajectories plus provider-wire or explicitly +lower-fidelity native-agent LLM request/response exchanges. The adjacent +``llm_trajectory.manifest.json`` states the capture source and completeness. Files ----- diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 4ea47a56c..7f09d873f 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -62,6 +62,7 @@ class PrimeSftExportStats: skipped_exchanges_provider_error: int = 0 skipped_no_assistant: int = 0 skipped_missing_tool_defs: int = 0 + skipped_insufficient_capture_fidelity: int = 0 skipped_terminal_error: int = 0 skipped_invalid: int = 0 tool_call_ids_rewritten: int = 0 @@ -81,6 +82,9 @@ def as_dict(self) -> dict[str, Any]: "skipped_exchanges_provider_error": self.skipped_exchanges_provider_error, "skipped_no_assistant": self.skipped_no_assistant, "skipped_missing_tool_defs": self.skipped_missing_tool_defs, + "skipped_insufficient_capture_fidelity": ( + self.skipped_insufficient_capture_fidelity + ), "skipped_terminal_error": self.skipped_terminal_error, "skipped_invalid": self.skipped_invalid, "tool_call_ids_rewritten": self.tool_call_ids_rewritten, @@ -1117,6 +1121,15 @@ def normalize_prime_sft_exchange( redact: bool = True, ) -> tuple[PrimeSftExchangeData | None, str | None]: """Normalize one raw LLM exchange through the Prime-SFT validator path.""" + metadata = exchange.get("metadata") + if isinstance(metadata, dict): + fidelity = metadata.get("capture_fidelity") + if fidelity is not None and fidelity != "provider_wire": + return None, "insufficient_capture_fidelity" + if metadata.get("request_complete") is False: + return None, "insufficient_capture_fidelity" + if metadata.get("response_complete") is False: + return None, "insufficient_capture_fidelity" messages, tool_defs, skip_reason = _exchange_to_messages_and_tools( exchange, redact=redact ) @@ -1234,6 +1247,9 @@ def convert_benchflow_rollouts_to_prime_sft_rows( if skip_reason == "missing_tool_defs": stats.skipped_missing_tool_defs += 1 continue + if skip_reason == "insufficient_capture_fidelity": + stats.skipped_insufficient_capture_fidelity += 1 + continue if row is None: stats.skipped_invalid += 1 continue diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py new file mode 100644 index 000000000..7611dfc4f --- /dev/null +++ b/src/benchflow/trajectories/llm_capture.py @@ -0,0 +1,578 @@ +"""Lifecycle orchestration for the always-present LLM trajectory artifact.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import shlex +import tempfile +from contextlib import suppress +from datetime import datetime +from pathlib import Path +from typing import Any + +from benchflow.agents.env import uses_native_subscription_auth +from benchflow.agents.registry import AGENTS +from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter +from benchflow.trajectories.llm_capture_manifest import ( + LLM_TRAJECTORY_FILENAME, + LLM_TRAJECTORY_SCHEMA_VERSION, + AuthMode, + CaptureFidelity, + CaptureSource, + CaptureStatus, + initialize_llm_trajectory_artifacts, + write_llm_trajectory_manifest, +) +from benchflow.trajectories.native_capture_parsers import ( + NativeParseResult, + parse_claude_raw_capture, + parse_claude_sessions, + parse_codex_sessions, + project_acp_trajectory, +) +from benchflow.trajectories.types import ( + redact_trajectory_obj, + redact_trajectory_text, +) + +logger = logging.getLogger(__name__) + +_REMOTE_CAPTURE_PREFIX = "/tmp/benchflow-llm-capture-" +_OTEL_SINK_SOURCE = r""" +import { createServer } from 'node:http'; +import { mkdirSync, writeFileSync } from 'node:fs'; +import { join } from 'node:path'; + +const [outputDir, portFile] = process.argv.slice(2); +mkdirSync(outputDir, { recursive: true }); +let sequence = 0; +const server = createServer((request, response) => { + const chunks = []; + let size = 0; + request.on('data', chunk => { + size += chunk.length; + if (size <= 64 * 1024 * 1024) chunks.push(chunk); + }); + request.on('end', () => { + if (size <= 64 * 1024 * 1024) { + const name = `${Date.now()}-${String(sequence++).padStart(6, '0')}.json`; + writeFileSync(join(outputDir, name), Buffer.concat(chunks)); + } + response.writeHead(200, { 'content-type': 'application/json' }); + response.end('{}'); + }); +}); +server.listen(0, '127.0.0.1', () => { + const address = server.address(); + writeFileSync(portFile, `${address.port}\n`); +}); +process.on('SIGTERM', () => server.close(() => process.exit(0))); +""".strip() + + +class LLMTrajectoryCapture: + """Own capture initialization, sandbox instrumentation, and finalization.""" + + def __init__( + self, + rollout_dir: Path, + *, + agent: str, + model: str | None, + session_id: str, + started_at: datetime, + ) -> None: + self.rollout_dir = rollout_dir + self.agent = agent + self.model = model + self.session_id = session_id + self.started_at = started_at + capture_suffix = hashlib.sha256(session_id.encode()).hexdigest()[:12] + self._remote_capture_root = f"{_REMOTE_CAPTURE_PREFIX}{capture_suffix}" + self.manifest = initialize_llm_trajectory_artifacts( + rollout_dir, + agent=agent, + model=model, + session_id=session_id, + started_at=started_at, + ) + self._native_agents: dict[str, str | None] = {} + self._credential_homes: set[str] = set() + self._collector_started = False + self._capture_root_prepared = False + self._preparation_errors: list[str] = [] + + @property + def trajectory_path(self) -> Path: + return self.rollout_dir / "trajectory" / LLM_TRAJECTORY_FILENAME + + def configure(self, agent_env: dict[str, str]) -> None: + self.manifest.auth_mode = _resolve_auth_mode( + self.agent, + self.model, + agent_env, + ) + write_llm_trajectory_manifest(self.rollout_dir, self.manifest) + + async def prepare_agent( + self, + env: Any, + *, + agent: str, + model: str | None, + agent_env: dict[str, str], + credential_home: str, + sandbox_user: str | None, + ) -> dict[str, str]: + """Return the agent environment augmented for native capture.""" + + prepared = dict(agent_env) + if not uses_native_subscription_auth(agent, model, prepared): + return prepared + self._native_agents[agent] = model + self._credential_homes.add(credential_home) + self.manifest.auth_mode = _resolve_auth_mode(agent, model, prepared) + if _is_claude_code_agent(agent): + raw_dir = f"{self._remote_capture_root}/raw" + prepared.update( + { + "CLAUDE_CODE_ENABLE_TELEMETRY": "1", + "OTEL_LOG_RAW_API_BODIES": f"file:{raw_dir}", + } + ) + try: + port = await self._ensure_otel_sink(env, sandbox_user=sandbox_user) + except Exception as exc: + warning = _sanitized_error(exc) + self._preparation_errors.append(warning) + logger.warning( + "Claude OTel correlation unavailable; raw/session fallback remains " + "enabled: %s", + warning, + ) + else: + prepared.update( + { + "OTEL_LOGS_EXPORTER": "otlp", + "OTEL_EXPORTER_OTLP_PROTOCOL": "http/json", + "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL": "http/json", + "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": ( + f"http://127.0.0.1:{port}/v1/logs" + ), + "OTEL_LOGS_EXPORT_INTERVAL": "500", + } + ) + write_llm_trajectory_manifest(self.rollout_dir, self.manifest) + return prepared + + async def finalize( + self, + env: Any, + *, + acp_events: list[dict[str, Any]], + model_call_seen: bool, + ) -> None: + """Publish the highest-fidelity available capture and terminal sidecar.""" + + self.manifest.finished_at = datetime.now() + if self.trajectory_path.stat().st_size > 0: + try: + exchange_count = _annotate_provider_wire_jsonl( + self.trajectory_path, + auth_mode=self.manifest.auth_mode, + ) + self._finish_manifest( + status=CaptureStatus.COMPLETE, + source=CaptureSource.LITELLM_PROXY, + fidelity=CaptureFidelity.PROVIDER_WIRE, + exchange_count=exchange_count, + request_complete=True, + response_complete=True, + ) + finally: + # Mixed-role scenes can prepare native telemetry before another + # role writes provider capture. Never strand raw bodies in a + # reusable externally-owned sandbox on this early-return path. + if self._native_agents and env is not None: + await self._cleanup_remote_capture(env) + return + + native_result: NativeParseResult | None = None + collection_errors: list[str] = list(self._preparation_errors) + if self._native_agents and env is not None: + try: + native_result = await self._collect_native_result(env) + except Exception as exc: + collection_errors.append(_sanitized_error(exc)) + logger.warning("Native LLM trajectory collection failed: %s", exc) + finally: + await self._cleanup_remote_capture(env) + + if native_result is None and acp_events: + native_result = project_acp_trajectory( + acp_events, + agent=self.agent, + session_id=self.session_id, + started_at=self.started_at, + auth_mode=self.manifest.auth_mode.value, + ) + if native_result is not None: + LiveLLMTrajectoryWriter(self.trajectory_path).reconcile( + native_result.trajectory + ) + errors = [*collection_errors, *native_result.errors] + status = ( + CaptureStatus.COMPLETE + if native_result.fidelity is CaptureFidelity.PROVIDER_WIRE + and not errors + else CaptureStatus.PARTIAL + ) + self._finish_manifest( + status=status, + source=native_result.source, + fidelity=native_result.fidelity, + exchange_count=len(native_result.trajectory.exchanges), + request_complete=native_result.request_complete, + response_complete=native_result.response_complete, + missing_fields=native_result.missing_fields, + errors=errors, + ) + return + + self._finish_manifest( + status=( + CaptureStatus.CAPTURE_FAILED + if model_call_seen + else CaptureStatus.NO_MODEL_CALL + ), + source=CaptureSource.NONE, + fidelity=CaptureFidelity.NONE, + exchange_count=0, + request_complete=False, + response_complete=False, + missing_fields=( + ["provider_request", "provider_response"] if model_call_seen else [] + ), + errors=( + collection_errors + or ( + ["model call observed but no LLM capture source was readable"] + if model_call_seen + else [] + ) + ), + ) + + def record_failure(self, error: object, *, model_call_seen: bool) -> None: + """Leave a terminal, truthful sidecar when finalization itself fails.""" + + self.manifest.finished_at = datetime.now() + _atomic_replace_text(self.trajectory_path, "") + self._finish_manifest( + status=( + CaptureStatus.CAPTURE_FAILED + if model_call_seen + else CaptureStatus.NO_MODEL_CALL + ), + source=CaptureSource.NONE, + fidelity=CaptureFidelity.NONE, + exchange_count=0, + request_complete=False, + response_complete=False, + missing_fields=( + ["provider_request", "provider_response"] if model_call_seen else [] + ), + errors=[_sanitized_error(error)], + ) + + async def _ensure_otel_sink(self, env: Any, *, sandbox_user: str | None) -> int: + if self._collector_started: + return await self._read_collector_port(env) + capture_owner = shlex.quote(sandbox_user or "root") + setup = await env.exec( + f"mkdir -p {self._remote_capture_root}/raw " + f"{self._remote_capture_root}/otel\n" + f"chown -R {capture_owner} {self._remote_capture_root}\n" + f"chmod 700 {self._remote_capture_root} " + f"{self._remote_capture_root}/raw {self._remote_capture_root}/otel", + user="root", + timeout_sec=10, + ) + if setup.return_code != 0: + detail = (setup.stderr or setup.stdout or "capture directory setup failed")[ + :300 + ] + raise RuntimeError(f"Claude capture directory setup failed: {detail}") + self._capture_root_prepared = True + with tempfile.TemporaryDirectory(prefix="benchflow-otel-sink-") as temporary: + source = Path(temporary) / "otel_sink.mjs" + source.write_text(_OTEL_SINK_SOURCE + "\n") + await env.upload_file( + source, + f"{self._remote_capture_root}/otel_sink.mjs", + mode="755", + ) + command = f""" +find {self._remote_capture_root} -maxdepth 1 -type f -name port -delete +find {self._remote_capture_root} -maxdepth 1 -type f -name pid -delete +node_bin=/opt/benchflow/node/bin/node +if ! test -x "$node_bin"; then + node_bin=$(command -v node || true) +fi +if test -z "$node_bin"; then + echo "node runtime not found" >&2 + exit 1 +fi +nohup "$node_bin" {self._remote_capture_root}/otel_sink.mjs \ + {self._remote_capture_root}/otel {self._remote_capture_root}/port \ + >{self._remote_capture_root}/collector.stdout \ + 2>{self._remote_capture_root}/collector.stderr {self._remote_capture_root}/pid +for attempt in $(seq 1 50); do + if test -s {self._remote_capture_root}/port; then + cat {self._remote_capture_root}/port + exit 0 + fi + sleep 0.1 +done +tail -c 300 {self._remote_capture_root}/collector.stderr >&2 2>/dev/null || true +exit 1 +""" + result = await env.exec( + command, + user=sandbox_user or "root", + timeout_sec=10, + ) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "collector did not start")[:300] + raise RuntimeError(f"Claude OTel sink failed to start: {detail}") + self._collector_started = True + return _parse_port(result.stdout) + + async def _read_collector_port(self, env: Any) -> int: + result = await env.exec( + f"cat {self._remote_capture_root}/port", + user="root", + timeout_sec=5, + ) + if result.return_code != 0: + raise RuntimeError("Claude OTel sink port file is unavailable") + return _parse_port(result.stdout) + + async def _collect_native_result(self, env: Any) -> NativeParseResult | None: + if self._collector_started: + await env.exec( + f"if test -s {self._remote_capture_root}/pid; then " + f"kill -TERM $(cat {self._remote_capture_root}/pid) " + f"2>/dev/null || true; fi", + user="root", + timeout_sec=5, + ) + with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: + local_root = Path(temporary) + capture_dir = local_root / "capture" + if self._capture_root_prepared: + await env.download_dir(self._remote_capture_root, capture_dir) + result = parse_claude_raw_capture( + capture_dir, + agent=self.agent, + session_id=self.session_id, + started_at=self.started_at, + ) + if result is not None: + return result + session_roots: list[tuple[str, Path]] = [] + for index, credential_home in enumerate(sorted(self._credential_homes)): + claude_local = local_root / f"home-{index}" / "claude-projects" + codex_local = local_root / f"home-{index}" / "codex-sessions" + if any( + _is_claude_code_agent(agent) for agent in self._native_agents + ) and await _download_optional_dir( + env, f"{credential_home}/.claude/projects", claude_local + ): + session_roots.append(("claude", claude_local)) + if "codex-acp" in self._native_agents and await _download_optional_dir( + env, f"{credential_home}/.codex/sessions", codex_local + ): + session_roots.append(("codex", codex_local)) + + for source, root in session_roots: + if source == "claude": + result = parse_claude_sessions( + root, + agent=self.agent, + session_id=self.session_id, + started_at=self.started_at, + ) + else: + result = parse_codex_sessions( + root, + agent=self.agent, + session_id=self.session_id, + started_at=self.started_at, + configured_model=self._native_agents.get("codex-acp"), + auth_mode=self.manifest.auth_mode.value, + ) + if result is not None: + return result + return None + + async def _cleanup_remote_capture(self, env: Any) -> None: + if not self._capture_root_prepared: + return + try: + result = await env.exec( + f"find {self._remote_capture_root} -depth -delete", + user="root", + timeout_sec=10, + ) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "unknown error")[:300] + logger.warning("Sandbox LLM capture cleanup failed: %s", detail) + except Exception as exc: + logger.warning("Sandbox LLM capture cleanup failed: %s", exc) + + def _finish_manifest( + self, + *, + status: CaptureStatus, + source: CaptureSource, + fidelity: CaptureFidelity, + exchange_count: int, + request_complete: bool, + response_complete: bool, + missing_fields: list[str] | None = None, + errors: list[str] | None = None, + ) -> None: + self.manifest.status = status + self.manifest.capture_source = source + self.manifest.capture_fidelity = fidelity + self.manifest.exchange_count = exchange_count + self.manifest.request_complete = request_complete + self.manifest.response_complete = response_complete + self.manifest.missing_fields = sorted(set(missing_fields or [])) + self.manifest.errors = [_sanitized_error(item) for item in errors or []] + write_llm_trajectory_manifest(self.rollout_dir, self.manifest) + + +async def _download_optional_dir(env: Any, remote: str, local: Path) -> bool: + probe = await env.exec( + f"test -d {shlex.quote(remote)}", + user="root", + timeout_sec=5, + ) + if probe.return_code != 0: + return False + local.parent.mkdir(parents=True, exist_ok=True) + await env.download_dir(remote, local) + return True + + +def _is_claude_code_agent(agent: str) -> bool: + config = AGENTS.get(agent) + subscription = config.subscription_auth if config is not None else None + return bool( + subscription is not None + and subscription.replaces_env == "ANTHROPIC_API_KEY" + ) + + +def _resolve_auth_mode( + agent: str, + model: str | None, + agent_env: dict[str, str], +) -> AuthMode: + if not uses_native_subscription_auth(agent, model, agent_env): + return AuthMode.API_KEY + if agent != "codex-acp": + return AuthMode.OAUTH_SUBSCRIPTION + raw_auth = agent_env.get("CODEX_AUTH_JSON") + if raw_auth is None and agent_env.get("_BENCHFLOW_SUBSCRIPTION_AUTH") == "1": + host_auth = Path.home() / ".codex" / "auth.json" + with suppress(OSError): + raw_auth = host_auth.read_text() + if raw_auth: + try: + auth = json.loads(raw_auth) + except json.JSONDecodeError: + auth = None + if isinstance(auth, dict): + auth_name = str(auth.get("auth_mode") or "").casefold() + if auth_name == "chatgpt": + return AuthMode.OAUTH_SUBSCRIPTION + if auth_name in {"api_key", "apikey", "api-key"} or auth.get( + "OPENAI_API_KEY" + ): + return AuthMode.API_KEY + if isinstance(auth.get("tokens"), dict): + return AuthMode.OAUTH_SUBSCRIPTION + return AuthMode.OAUTH_SUBSCRIPTION + + +def _parse_port(value: str) -> int: + try: + port = int(value.strip().splitlines()[-1]) + except (ValueError, IndexError) as exc: + raise RuntimeError("Claude OTel sink returned an invalid port") from exc + if not 1 <= port <= 65535: + raise RuntimeError("Claude OTel sink returned an out-of-range port") + return port + + +def _annotate_provider_wire_jsonl(path: Path, *, auth_mode: AuthMode) -> int: + records: list[dict[str, Any]] = [] + for line_number, line in enumerate(path.read_text().splitlines(), start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"invalid LLM trajectory JSONL at line {line_number}: {exc.msg}" + ) from exc + if not isinstance(record, dict): + raise ValueError( + f"invalid LLM trajectory JSONL at line {line_number}: expected object" + ) + metadata = record.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + record["metadata"] = metadata + metadata.setdefault("schema_version", LLM_TRAJECTORY_SCHEMA_VERSION) + metadata.setdefault("capture_source", CaptureSource.LITELLM_PROXY.value) + metadata.setdefault("capture_fidelity", CaptureFidelity.PROVIDER_WIRE.value) + metadata.setdefault("auth_mode", auth_mode.value) + metadata.setdefault("request_complete", True) + metadata.setdefault("response_complete", True) + metadata.setdefault("payload_redacted", True) + records.append(redact_trajectory_obj(record)) + payload = "".join(json.dumps(record, default=str) + "\n" for record in records) + _atomic_replace_text(path, payload) + return len(records) + + +def _sanitized_error(error: object) -> str: + text = redact_trajectory_text(str(error)).replace("\n", " ").strip() + return text[:500] or type(error).__name__ + + +def _atomic_replace_text(path: Path, payload: str) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(payload) + os.replace(temporary, path) + + +def model_call_seen_from_evidence( + usage_metrics: dict[str, Any] | None, + acp_events: list[dict[str, Any]], +) -> bool: + for value in (usage_metrics or {}).values(): + if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0: + return True + return any( + event.get("type") in {"agent_message", "agent_thought", "tool_call"} + for event in acp_events + ) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py new file mode 100644 index 000000000..219306a00 --- /dev/null +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -0,0 +1,111 @@ +"""Typed provenance contract for ``llm_trajectory.jsonl``. + +The JSONL filename is intentionally stable across every authentication path. +This sidecar records what the file can actually prove so downstream code never +confuses a native-agent reconstruction with provider-wire traffic. +""" + +from __future__ import annotations + +import json +import os +from datetime import datetime +from enum import StrEnum +from pathlib import Path + +from pydantic import BaseModel, Field + +LLM_TRAJECTORY_FILENAME = "llm_trajectory.jsonl" +LLM_TRAJECTORY_MANIFEST_FILENAME = "llm_trajectory.manifest.json" +LLM_TRAJECTORY_SCHEMA_VERSION = 2 + + +class CaptureStatus(StrEnum): + PENDING = "pending" + COMPLETE = "complete" + PARTIAL = "partial" + NO_MODEL_CALL = "no_model_call" + CAPTURE_FAILED = "capture_failed" + + +class CaptureFidelity(StrEnum): + PROVIDER_WIRE = "provider_wire" + AGENT_SESSION = "agent_session" + ACP_PROJECTION = "acp_projection" + NONE = "none" + + +class CaptureSource(StrEnum): + LITELLM_PROXY = "litellm_proxy" + CLAUDE_OTEL_RAW_BODY = "claude_otel_raw_body" + CLAUDE_NATIVE_SESSION = "claude_native_session" + CODEX_NATIVE_SESSION = "codex_native_session" + ACP_PROJECTION = "acp_projection" + NONE = "none" + + +class AuthMode(StrEnum): + API_KEY = "api_key" + OAUTH_SUBSCRIPTION = "oauth_subscription" + UNKNOWN = "unknown" + + +class LLMTrajectoryManifest(BaseModel): + """Machine-readable fidelity and lifecycle state for the JSONL artifact.""" + + schema_version: int = LLM_TRAJECTORY_SCHEMA_VERSION + status: CaptureStatus = CaptureStatus.PENDING + capture_source: CaptureSource = CaptureSource.NONE + capture_fidelity: CaptureFidelity = CaptureFidelity.NONE + auth_mode: AuthMode = AuthMode.UNKNOWN + agent: str + model: str | None = None + session_id: str = "" + exchange_count: int = 0 + request_complete: bool = False + response_complete: bool = False + payload_redacted: bool = True + started_at: datetime + finished_at: datetime | None = None + missing_fields: list[str] = Field(default_factory=list) + errors: list[str] = Field(default_factory=list) + + +def initialize_llm_trajectory_artifacts( + rollout_dir: Path, + *, + agent: str, + model: str | None, + session_id: str, + started_at: datetime, +) -> LLMTrajectoryManifest: + """Create the always-present empty JSONL and its initial sidecar.""" + + trajectory_dir = rollout_dir / "trajectory" + trajectory_dir.mkdir(parents=True, exist_ok=True) + trajectory_path = trajectory_dir / LLM_TRAJECTORY_FILENAME + if not trajectory_path.exists(): + _atomic_write_text(trajectory_path, "") + manifest = LLMTrajectoryManifest( + agent=agent, + model=model, + session_id=session_id, + started_at=started_at, + ) + write_llm_trajectory_manifest(rollout_dir, manifest) + return manifest + + +def write_llm_trajectory_manifest( + rollout_dir: Path, manifest: LLMTrajectoryManifest +) -> None: + path = rollout_dir / "trajectory" / LLM_TRAJECTORY_MANIFEST_FILENAME + payload = json.dumps(manifest.model_dump(mode="json"), indent=2, sort_keys=True) + _atomic_write_text(path, payload + "\n") + + +def _atomic_write_text(path: Path, payload: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(payload) + os.replace(temporary, path) diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py new file mode 100644 index 000000000..f7c6ff9e4 --- /dev/null +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -0,0 +1,735 @@ +"""Normalize Claude Code and Codex native capture surfaces into LLM exchanges.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from benchflow.trajectories.llm_capture_manifest import ( + LLM_TRAJECTORY_SCHEMA_VERSION, + CaptureFidelity, + CaptureSource, +) +from benchflow.trajectories.types import ( + LLMExchange, + LLMRequest, + LLMResponse, + Trajectory, +) + + +@dataclass(frozen=True) +class NativeParseResult: + trajectory: Trajectory + source: CaptureSource + fidelity: CaptureFidelity + request_complete: bool + response_complete: bool + missing_fields: list[str] = field(default_factory=list) + errors: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class _OtelEvent: + name: str + timestamp: datetime + attributes: dict[str, Any] + + +@dataclass(frozen=True) +class _RequestBodyEvent: + path: Path + timestamp: datetime + session_key: str | None + otel_referenced: bool + + +def parse_claude_raw_capture( + capture_dir: Path, + *, + agent: str, + session_id: str, + started_at: datetime, +) -> NativeParseResult | None: + """Pair Claude raw bodies using OTLP log order within each session.""" + + raw_dir = capture_dir / "raw" + if not raw_dir.is_dir(): + return None + bodies = _load_body_files(raw_dir) + events = _load_otel_events(capture_dir / "otel") + if not bodies or not events: + return None + + pending = _request_body_events(raw_dir, bodies, events) + exchanges: list[LLMExchange] = [] + errors: list[str] = [] + pairing_ambiguous = False + observed_responses: set[Path] = set() + for event in sorted(events, key=lambda item: item.timestamp): + event_session = _event_session_key(event.attributes) + refs = _body_references(event.attributes) + if "api_response" not in event.name: + continue + response_ref = next((ref for ref in refs if ".response.json" in ref), None) + if response_ref is None: + continue + response_path = _resolve_body_reference(raw_dir, response_ref, bodies) + if response_path is not None: + if response_path in observed_responses: + continue + observed_responses.add(response_path) + candidates = [ + request + for request in pending + if request.timestamp <= event.timestamp + and request.session_key in {None, event_session} + ] + if response_path is None or not candidates: + errors.append("unpaired Claude raw API response") + continue + if len(candidates) > 1: + pairing_ambiguous = True + request_event = candidates[0] + pending.remove(request_event) + request_path = request_event.path + request_timestamp = request_event.timestamp + request_body = bodies[request_path] + response_body = bodies[response_path] + response_timestamp = event.timestamp + request_id = _request_id(event.attributes) or response_path.name.removesuffix( + ".response.json" + ) + exchanges.append( + _exchange( + request_body=request_body, + response_body=response_body, + request_timestamp=request_timestamp, + response_timestamp=response_timestamp, + path="/v1/messages", + source=CaptureSource.CLAUDE_OTEL_RAW_BODY, + fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode="oauth_subscription", + request_complete=True, + response_complete=True, + extra_metadata={ + "provider_request_id": request_id, + "request_body_file": request_path.name, + "response_body_file": response_path.name, + "pairing": ( + "otel_session_fifo" + if request_event.otel_referenced + else "raw_file_mtime_fifo" + ), + "correlation_complete": len(candidates) == 1, + }, + ) + ) + + dangling = len(pending) + if dangling: + errors.append(f"{dangling} unpaired Claude raw API request(s)") + unseen_responses = sum( + path.name.endswith(".response.json") and path not in observed_responses + for path in bodies + ) + if unseen_responses: + errors.append( + f"{unseen_responses} Claude raw API response(s) lacked an OTLP body event" + ) + if not exchanges: + return None + if pairing_ambiguous: + errors.append( + "concurrent Claude requests lacked a provider correlation id; " + "request/response pairing is ambiguous" + ) + for exchange in exchanges: + exchange.metadata["request_complete"] = False + exchange.metadata["correlation_complete"] = False + finished_at = max(exchange.response.timestamp for exchange in exchanges) + return NativeParseResult( + trajectory=Trajectory( + session_id=session_id, + agent_name=agent, + started_at=started_at, + finished_at=finished_at, + exchanges=exchanges, + ), + source=CaptureSource.CLAUDE_OTEL_RAW_BODY, + fidelity=CaptureFidelity.PROVIDER_WIRE, + request_complete=not pairing_ambiguous, + response_complete=True, + errors=errors, + ) + + +def parse_claude_sessions( + sessions_dir: Path, + *, + agent: str, + session_id: str, + started_at: datetime, +) -> NativeParseResult | None: + """Reconstruct model turns from Claude Code's native session transcript.""" + + records = _read_jsonl_tree(sessions_dir) + if not records: + return None + messages: list[dict[str, Any]] = [] + exchanges: list[LLMExchange] = [] + assistant_group: list[dict[str, Any]] = [] + group_key: str | None = None + + def flush_group() -> None: + nonlocal assistant_group, group_key + if not assistant_group: + return + message = _merge_claude_assistant_group(assistant_group) + request_timestamp = _record_timestamp(assistant_group[0], started_at) + response_timestamp = _record_timestamp(assistant_group[-1], request_timestamp) + model = message.get("model") + request_body: dict[str, Any] = {"messages": list(messages)} + if isinstance(model, str) and model: + request_body["model"] = model + exchanges.append( + _exchange( + request_body=request_body, + response_body=message, + request_timestamp=request_timestamp, + response_timestamp=response_timestamp, + path="/v1/messages", + source=CaptureSource.CLAUDE_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + auth_mode="oauth_subscription", + request_complete=False, + response_complete=False, + extra_metadata={ + "provider_request_id": group_key, + "missing_fields": [ + "system_prompt", + "tool_definitions", + "provider_response_envelope", + "headers", + ], + }, + ) + ) + messages.append(_message_for_history(message)) + assistant_group = [] + group_key = None + + for record in records: + record_type = record.get("type") + message = record.get("message") + if not isinstance(message, dict): + continue + if record_type == "user": + flush_group() + messages.append(_message_for_history(message)) + continue + if record_type != "assistant": + continue + key = str( + record.get("requestId") + or message.get("id") + or record.get("uuid") + or f"assistant-{len(exchanges)}" + ) + if assistant_group and key != group_key: + flush_group() + group_key = key + assistant_group.append(record) + flush_group() + if not exchanges: + return None + return NativeParseResult( + trajectory=Trajectory( + session_id=session_id, + agent_name=agent, + started_at=started_at, + finished_at=max(exchange.response.timestamp for exchange in exchanges), + exchanges=exchanges, + ), + source=CaptureSource.CLAUDE_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + request_complete=False, + response_complete=False, + missing_fields=[ + "system_prompt", + "tool_definitions", + "provider_response_envelope", + "headers", + ], + ) + + +def parse_codex_sessions( + sessions_dir: Path, + *, + agent: str, + session_id: str, + started_at: datetime, + configured_model: str | None, + auth_mode: str = "oauth_subscription", +) -> NativeParseResult | None: + """Reconstruct Responses-style calls from Codex native session records.""" + + records = _read_jsonl_tree(sessions_dir) + if not records: + return None + history: list[dict[str, Any]] = [] + output: list[dict[str, Any]] = [] + exchanges: list[LLMExchange] = [] + output_started_at: datetime | None = None + last_timestamp = started_at + pending_usage: dict[str, Any] | None = None + model = configured_model or _codex_session_model(records) + + def flush_output(response_timestamp: datetime) -> None: + nonlocal output, output_started_at, pending_usage + if not output: + return + response_body: dict[str, Any] = {"status": "completed", "output": output} + if pending_usage: + response_body["usage"] = _normalize_codex_usage(pending_usage) + exchanges.append( + _exchange( + request_body={"model": model, "input": list(history)}, + response_body=response_body, + request_timestamp=output_started_at or response_timestamp, + response_timestamp=response_timestamp, + path="/v1/responses", + source=CaptureSource.CODEX_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + auth_mode=auth_mode, + request_complete=False, + response_complete=False, + extra_metadata={ + "missing_fields": [ + "instructions", + "tool_definitions", + "provider_response_envelope", + "headers", + ] + }, + ) + ) + history.extend(output) + output = [] + output_started_at = None + pending_usage = None + + for record in records: + timestamp = _record_timestamp(record, last_timestamp) + last_timestamp = timestamp + payload = record.get("payload") + if not isinstance(payload, dict): + continue + if record.get("type") == "event_msg" and payload.get("type") == "token_count": + usage = payload.get("info") + if isinstance(usage, dict): + last_usage = usage.get("last_token_usage") + pending_usage = last_usage if isinstance(last_usage, dict) else usage + flush_output(timestamp) + continue + if record.get("type") != "response_item": + continue + payload_type = str(payload.get("type") or "") + role = str(payload.get("role") or "") + is_input = (payload_type == "message" and role in {"user", "developer"}) or ( + payload_type.endswith("_call_output") + ) + if is_input: + flush_output(timestamp) + history.append(payload) + continue + if payload_type in { + "message", + "reasoning", + "function_call", + "custom_tool_call", + "web_search_call", + } or payload_type.endswith("_call"): + if output_started_at is None: + output_started_at = timestamp + output.append(payload) + flush_output(last_timestamp) + if not exchanges: + return None + return NativeParseResult( + trajectory=Trajectory( + session_id=session_id, + agent_name=agent, + started_at=started_at, + finished_at=max(exchange.response.timestamp for exchange in exchanges), + exchanges=exchanges, + ), + source=CaptureSource.CODEX_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + request_complete=False, + response_complete=False, + missing_fields=[ + "instructions", + "tool_definitions", + "provider_response_envelope", + "headers", + ], + ) + + +def project_acp_trajectory( + events: list[dict[str, Any]], + *, + agent: str, + session_id: str, + started_at: datetime, + auth_mode: str, +) -> NativeParseResult | None: + """Last-resort projection that remains explicit about its low fidelity.""" + + prompts: list[dict[str, Any]] = [] + exchanges: list[LLMExchange] = [] + for index, event in enumerate(events): + event_type = event.get("type") + if event_type == "user_message": + prompts.append({"role": "user", "content": event.get("text", "")}) + elif event_type in {"agent_message", "agent_thought"}: + timestamp = _record_timestamp(event, started_at) + response_body = { + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + {"type": "output_text", "text": event.get("text", "")} + ], + } + ], + } + exchanges.append( + _exchange( + request_body={"input": list(prompts)}, + response_body=response_body, + request_timestamp=timestamp, + response_timestamp=timestamp, + path="acp://session", + source=CaptureSource.ACP_PROJECTION, + fidelity=CaptureFidelity.ACP_PROJECTION, + auth_mode=auth_mode, + request_complete=False, + response_complete=False, + extra_metadata={ + "acp_event_index": index, + "missing_fields": ["provider_request", "provider_response"], + }, + ) + ) + prompts.append( + {"role": "assistant", "content": str(event.get("text") or "")} + ) + if not exchanges: + return None + return NativeParseResult( + trajectory=Trajectory( + session_id=session_id, + agent_name=agent, + started_at=started_at, + finished_at=max(exchange.response.timestamp for exchange in exchanges), + exchanges=exchanges, + ), + source=CaptureSource.ACP_PROJECTION, + fidelity=CaptureFidelity.ACP_PROJECTION, + request_complete=False, + response_complete=False, + missing_fields=["provider_request", "provider_response"], + ) + + +def _exchange( + *, + request_body: dict[str, Any], + response_body: dict[str, Any], + request_timestamp: datetime, + response_timestamp: datetime, + path: str, + source: CaptureSource, + fidelity: CaptureFidelity, + auth_mode: str, + request_complete: bool, + response_complete: bool, + extra_metadata: dict[str, Any] | None = None, +) -> LLMExchange: + metadata = { + "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, + "capture_source": source.value, + "capture_fidelity": fidelity.value, + "auth_mode": auth_mode, + "request_complete": request_complete, + "response_complete": response_complete, + "payload_redacted": True, + **(extra_metadata or {}), + } + return LLMExchange( + request=LLMRequest(timestamp=request_timestamp, path=path, body=request_body), + response=LLMResponse( + timestamp=response_timestamp, + status_code=200, + body=response_body, + ), + duration_ms=max( + 0.0, (response_timestamp - request_timestamp).total_seconds() * 1000 + ), + metadata=metadata, + ) + + +def _read_jsonl_tree(root: Path) -> list[dict[str, Any]]: + if not root.is_dir(): + return [] + records: list[dict[str, Any]] = [] + for path in sorted(root.rglob("*.jsonl"), key=lambda item: item.stat().st_mtime): + for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(record, dict): + records.append(record) + return records + + +def _load_body_files(root: Path) -> dict[Path, dict[str, Any]]: + bodies: dict[Path, dict[str, Any]] = {} + for path in root.rglob("*.json"): + try: + body = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + continue + if isinstance(body, dict): + bodies[path] = body + return bodies + + +def _request_body_events( + root: Path, + bodies: dict[Path, dict[str, Any]], + otel_events: list[_OtelEvent], +) -> list[_RequestBodyEvent]: + """Resolve request bodies even when Claude omits the first body log event. + + Claude writes every raw request before sending it, but its periodic OTLP exporter + can start after the first request and omit that request's ``api_request_body`` + record. File modification time is therefore a safe ordering fallback. Pairing + still fails closed if more than one request could match a response. + """ + + referenced: dict[Path, _RequestBodyEvent] = {} + for event in otel_events: + request_ref = next( + ( + ref + for ref in _body_references(event.attributes) + if ".request.json" in ref + ), + None, + ) + request_path = _resolve_body_reference(root, request_ref, bodies) + if request_path is None or request_path in referenced: + continue + referenced[request_path] = _RequestBodyEvent( + path=request_path, + timestamp=event.timestamp, + session_key=_event_session_key(event.attributes), + otel_referenced=True, + ) + + requests = list(referenced.values()) + for path in bodies: + if not path.name.endswith(".request.json") or path in referenced: + continue + requests.append( + _RequestBodyEvent( + path=path, + timestamp=datetime.fromtimestamp(path.stat().st_mtime, tz=UTC), + session_key=None, + otel_referenced=False, + ) + ) + return sorted(requests, key=lambda item: (item.timestamp, item.path.name)) + + +def _load_otel_events(root: Path) -> list[_OtelEvent]: + events: list[_OtelEvent] = [] + if not root.is_dir(): + return events + for path in root.rglob("*.json"): + try: + envelope = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + continue + for record in _walk_log_records(envelope): + attributes = { + str(item.get("key")): _otel_value(item.get("value")) + for item in record.get("attributes") or [] + if isinstance(item, dict) and item.get("key") + } + body = _otel_value(record.get("body")) + name = str( + attributes.get("event.name") + or attributes.get("event_name") + or body + or "" + ) + timestamp = _unix_nanos_timestamp( + record.get("timeUnixNano") or record.get("observedTimeUnixNano") + ) + events.append(_OtelEvent(name=name, timestamp=timestamp, attributes=attributes)) + return events + + +def _walk_log_records(value: Any): + if isinstance(value, dict): + records = value.get("logRecords") + if isinstance(records, list): + yield from (record for record in records if isinstance(record, dict)) + for nested in value.values(): + yield from _walk_log_records(nested) + elif isinstance(value, list): + for nested in value: + yield from _walk_log_records(nested) + + +def _otel_value(value: Any) -> Any: + if not isinstance(value, dict): + return value + for key in ("stringValue", "intValue", "doubleValue", "boolValue", "bytesValue"): + if key in value: + return value[key] + array = value.get("arrayValue") + if isinstance(array, dict): + return [_otel_value(item) for item in array.get("values") or []] + pairs = value.get("kvlistValue") + if isinstance(pairs, dict): + return { + str(item.get("key")): _otel_value(item.get("value")) + for item in pairs.get("values") or [] + if isinstance(item, dict) + } + return value + + +def _body_references(attributes: dict[str, Any]) -> list[str]: + refs: list[str] = [] + for value in attributes.values(): + values = value if isinstance(value, list) else [value] + for item in values: + text = str(item) + if ".request.json" in text or ".response.json" in text: + refs.append(text) + return refs + + +def _resolve_body_reference( + root: Path, reference: str | None, bodies: dict[Path, dict[str, Any]] +) -> Path | None: + if reference is None: + return None + candidate = Path(reference) + for path in bodies: + if path == candidate or path.name == candidate.name: + return path + relative = root / reference + return relative if relative in bodies else None + + +def _event_session_key(attributes: dict[str, Any]) -> str: + for key in ("session.id", "session_id", "sessionId", "conversation.id"): + value = attributes.get(key) + if value: + return str(value) + return "default" + + +def _request_id(attributes: dict[str, Any]) -> str | None: + for key in ("request.id", "request_id", "requestId", "message.id"): + value = attributes.get(key) + if value: + return str(value) + return None + + +def _merge_claude_assistant_group(records: list[dict[str, Any]]) -> dict[str, Any]: + merged: dict[str, Any] = {"role": "assistant", "content": []} + seen_content: set[str] = set() + for record in records: + message = record.get("message") + if not isinstance(message, dict): + continue + for key in ("id", "type", "role", "model", "stop_reason", "stop_sequence"): + if key in message: + merged[key] = message[key] + if isinstance(message.get("usage"), dict): + merged["usage"] = message["usage"] + content = message.get("content") + blocks = content if isinstance(content, list) else [{"type": "text", "text": content}] + for block in blocks: + marker = json.dumps(block, sort_keys=True, default=str) + if marker not in seen_content: + seen_content.add(marker) + merged["content"].append(block) + return merged + + +def _message_for_history(message: dict[str, Any]) -> dict[str, Any]: + return { + key: value + for key, value in message.items() + if key in {"role", "content"} + } + + +def _codex_session_model(records: list[dict[str, Any]]) -> str: + for record in records: + payload = record.get("payload") + if isinstance(payload, dict): + model = payload.get("model") + if isinstance(model, str) and model: + return model + return "unknown" + + +def _normalize_codex_usage(usage: dict[str, Any]) -> dict[str, Any]: + return { + "input_tokens": usage.get("input_tokens", 0), + "input_tokens_details": { + "cached_tokens": usage.get("cached_input_tokens", 0) + }, + "output_tokens": usage.get("output_tokens", 0), + "output_tokens_details": { + "reasoning_tokens": usage.get("reasoning_output_tokens", 0) + }, + "total_tokens": usage.get("total_tokens", 0), + } + + +def _record_timestamp(record: dict[str, Any], default: datetime) -> datetime: + value = record.get("timestamp") + if not isinstance(value, str): + return default + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return default + return parsed + + +def _unix_nanos_timestamp(value: Any) -> datetime: + try: + return datetime.fromtimestamp(int(value) / 1_000_000_000, tz=UTC) + except (TypeError, ValueError, OSError): + return datetime.now(tz=UTC) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 700093fb3..3f67f7df7 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -3,14 +3,16 @@ This is the canonical trainer-facing rollout surface. It intentionally lives beside, not inside, raw traces: -- ``trajectory/llm_trajectory.jsonl`` remains the provider HTTP audit log. +- ``trajectory/llm_trajectory.jsonl`` remains the LLM exchange audit log; its + manifest distinguishes provider-wire traffic from native-agent reconstruction. - ``trajectory/acp_trajectory.jsonl`` remains the ACP event audit log. - ``results.jsonl`` is a Verifiers/Prime-RL-shaped rollout record. The writer is fail-closed for training readiness but not for artifact emission: even errored or unstructured rollouts get one JSONL row with an ``error``. The trainer-shaped trajectory is produced only from healthy -``llm_trajectory.jsonl`` exchanges; ACP is never used as a training fallback. +complete provider-wire ``llm_trajectory.jsonl`` exchanges; native session and +ACP projections are never used as a training fallback. """ from __future__ import annotations @@ -29,6 +31,9 @@ prime_sft_last_user_training_window, validate_prime_sft_row, ) +from benchflow.trajectories.llm_capture_manifest import ( + LLM_TRAJECTORY_MANIFEST_FILENAME, +) from benchflow.trajectories.types import redact_trajectory_obj from benchflow.usage_tracking import USAGE_SOURCE_AGENT_NATIVE_ACP @@ -38,7 +43,6 @@ logger = logging.getLogger(__name__) - def _record_to_redacted_json_line(record: dict[str, Any]) -> str: redacted = redact_trajectory_obj(scrub_non_finite(record)) return json.dumps(redacted, default=str, allow_nan=False) @@ -165,6 +169,11 @@ def _llm_steps_from_trajectory( exchanges = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError as exc: return [], [], f"Invalid LLM trajectory JSONL: {exc}" + capture_manifest = _load_llm_capture_manifest(rollout_dir) + if capture_manifest is not None and not _capture_manifest_allows_training( + capture_manifest + ): + return [], [], None training_success_indices = _training_success_exchange_indices(exchanges) skipped_successful: list[str] = [] for exchange_idx, exchange in enumerate(exchanges): @@ -205,12 +214,17 @@ def _llm_steps_from_trajectory( if isinstance(response.get("body"), dict) else {} ) + exchange_metadata = exchange.get("metadata") + tracking_source = "litellm_callback" + if isinstance(exchange_metadata, dict) and isinstance( + exchange_metadata.get("capture_source"), str + ): + tracking_source = str(exchange_metadata["capture_source"]) extras = { "source": "llm_trajectory", - "tracking_source": "litellm_callback", + "tracking_source": tracking_source, "exchange_index": exchange_idx, } - exchange_metadata = exchange.get("metadata") if isinstance(exchange_metadata, dict): extras.update( { @@ -260,6 +274,29 @@ def _response_is_training_success(response: Any) -> bool: return not bool(body.get("incomplete_details")) +def _load_llm_capture_manifest(rollout_dir: Path) -> dict[str, Any] | None: + path = rollout_dir / "trajectory" / LLM_TRAJECTORY_MANIFEST_FILENAME + if not path.exists(): + # Backward compatibility for artifacts created before the sidecar contract. + return None + try: + manifest = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return {"status": "capture_failed", "capture_fidelity": "none"} + if not isinstance(manifest, dict): + return {"status": "capture_failed", "capture_fidelity": "none"} + return manifest + + +def _capture_manifest_allows_training(manifest: dict[str, Any]) -> bool: + return bool( + manifest.get("status") == "complete" + and manifest.get("capture_fidelity") == "provider_wire" + and manifest.get("request_complete") is True + and manifest.get("response_complete") is True + ) + + def _training_success_exchange_indices( exchanges: list[dict[str, Any]], ) -> set[int]: @@ -466,9 +503,21 @@ def build_rollout_results_record( export_error=effective_export_error, ) terminal_health_error = bool(error or verifier_error or partial_trajectory) + capture_manifest = _load_llm_capture_manifest(rollout_path) + audit_only_capture = bool( + capture_manifest is not None + and int(capture_manifest.get("exchange_count") or 0) > 0 + and not _capture_manifest_allows_training(capture_manifest) + ) native_subscription_without_llm = bool( - (agent_result or {}).get("usage_source") == USAGE_SOURCE_AGENT_NATIVE_ACP - and not (rollout_path / "trajectory" / "llm_trajectory.jsonl").exists() + ( + (agent_result or {}).get("usage_source") == USAGE_SOURCE_AGENT_NATIVE_ACP + or (capture_manifest or {}).get("auth_mode") == "oauth_subscription" + ) + and ( + capture_manifest is None + or capture_manifest.get("status") in {"no_model_call", "partial"} + ) and effective_export_error is None and not terminal_health_error ) @@ -490,6 +539,8 @@ def build_rollout_results_record( training_ready_reason = "invalid_prime_sft_row" elif effective_export_error: training_ready_reason = "export_error" + elif audit_only_capture: + training_ready_reason = "insufficient_capture_fidelity" elif not steps or not completion: training_ready_reason = "missing_healthy_structured_llm_trajectory" elif partial_trajectory: diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index 3d8c12981..550d5ee70 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -188,6 +188,24 @@ def test_anthropic_tool_use_content_preserved_as_tool_calls(tmp_path: Path) -> N assert stats.rows_with_tool_calls == 1 +def test_native_session_exchange_is_not_exported_for_training(tmp_path: Path) -> None: + """Guards this PR's fail-closed export boundary for reconstructed capture.""" + + exchange = _anthropic_exchange() + exchange["metadata"] = { + "capture_fidelity": "agent_session", + "request_complete": False, + "response_complete": False, + } + _write_rollout(tmp_path / "job" / "rollout-1", exchanges=[exchange]) + + rows, stats = convert_benchflow_rollouts_to_prime_sft_rows(tmp_path / "job") + + assert rows == [] + assert stats.skipped_insufficient_capture_fidelity == 1 + assert stats.skipped_invalid == 0 + + def test_skipped_provider_error_counts_rollouts_not_exchanges(tmp_path: Path) -> None: """Guards #828 greptile P1: an all-failed rollout counts as ONE rollout skip, with the exchange count surfaced separately.""" diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py new file mode 100644 index 000000000..e7adc3bee --- /dev/null +++ b/tests/trajectories/test_native_llm_capture.py @@ -0,0 +1,771 @@ +"""Regression coverage for uniform LLM trajectory capture across auth modes.""" + +from __future__ import annotations + +import json +import os +from datetime import UTC, datetime +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from benchflow.trajectories.llm_capture import ( + LLMTrajectoryCapture, + _download_optional_dir, +) +from benchflow.trajectories.llm_capture_manifest import ( + CaptureFidelity, + CaptureSource, + CaptureStatus, +) +from benchflow.trajectories.native_capture_parsers import ( + parse_claude_raw_capture, + parse_claude_sessions, + parse_codex_sessions, + project_acp_trajectory, +) +from benchflow.trajectories.results import build_rollout_results_record + +STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + + +def _otel_value(value: str) -> dict: + return {"stringValue": value} + + +def _otel_record(name: str, timestamp_ns: int, **attributes: str) -> dict: + return { + "timeUnixNano": str(timestamp_ns), + "body": _otel_value(name), + "attributes": [ + {"key": key, "value": _otel_value(value)} + for key, value in attributes.items() + ], + } + + +def test_claude_otel_raw_bodies_become_provider_wire_exchanges(tmp_path: Path) -> None: + """Guards this PR's exact Claude OAuth raw-body capture contract.""" + + capture = tmp_path / "capture" + raw = capture / "raw" + raw.mkdir(parents=True) + request_path = raw / "random.request.json" + response_path = raw / "req_123.response.json" + request_path.write_text( + json.dumps( + { + "model": "claude-opus-4-1", + "messages": [{"role": "user", "content": "Solve it"}], + } + ) + ) + response_path.write_text( + json.dumps( + { + "id": "msg_123", + "role": "assistant", + "content": [{"type": "text", "text": "Done"}], + "usage": {"input_tokens": 10, "output_tokens": 2}, + } + ) + ) + otel = { + "resourceLogs": [ + { + "scopeLogs": [ + { + "logRecords": [ + _otel_record( + "claude_code.api_request", + 1_777_000_000_000_000_000, + session_id="session-1", + body_ref=str(request_path), + ), + _otel_record( + "claude_code.api_response", + 1_777_000_000_500_000_000, + session_id="session-1", + body_ref=str(response_path), + request_id="req_123", + ), + ] + } + ] + } + ] + } + otel_dir = capture / "otel" + otel_dir.mkdir() + (otel_dir / "000.json").write_text(json.dumps(otel)) + + result = parse_claude_raw_capture( + capture, + agent="claude-agent-acp", + session_id="rollout-1", + started_at=STARTED_AT, + ) + + assert result is not None + assert result.source is CaptureSource.CLAUDE_OTEL_RAW_BODY + assert result.fidelity is CaptureFidelity.PROVIDER_WIRE + assert result.request_complete is True + assert result.response_complete is True + assert len(result.trajectory.exchanges) == 1 + exchange = result.trajectory.exchanges[0] + assert exchange.request.body["messages"][0]["content"] == "Solve it" + assert exchange.response.body["id"] == "msg_123" + assert exchange.duration_ms == 500 + assert exchange.metadata["provider_request_id"] == "req_123" + + +def test_claude_concurrent_raw_pairing_fails_closed_for_training( + tmp_path: Path, +) -> None: + """Guards this PR against claiming FIFO pairing for concurrent Claude calls.""" + + capture = tmp_path / "capture" + raw = capture / "raw" + raw.mkdir(parents=True) + paths = { + name: raw / name + for name in ( + "one.request.json", + "two.request.json", + "req_two.response.json", + "req_one.response.json", + ) + } + for name, path in paths.items(): + path.write_text(json.dumps({"marker": name})) + records = [ + _otel_record( + "claude_code.api_request", + 1_777_000_000_000_000_000, + session_id="shared", + body_ref=str(paths["one.request.json"]), + ), + _otel_record( + "claude_code.api_request", + 1_777_000_000_100_000_000, + session_id="shared", + body_ref=str(paths["two.request.json"]), + ), + _otel_record( + "claude_code.api_response", + 1_777_000_000_200_000_000, + session_id="shared", + body_ref=str(paths["req_two.response.json"]), + request_id="req_two", + ), + _otel_record( + "claude_code.api_response", + 1_777_000_000_300_000_000, + session_id="shared", + body_ref=str(paths["req_one.response.json"]), + request_id="req_one", + ), + ] + otel_dir = capture / "otel" + otel_dir.mkdir() + (otel_dir / "events.json").write_text( + json.dumps({"resourceLogs": [{"scopeLogs": [{"logRecords": records}]}]}) + ) + + result = parse_claude_raw_capture( + capture, + agent="claude-agent-acp", + session_id="rollout-1", + started_at=STARTED_AT, + ) + + assert result is not None + assert result.request_complete is False + assert any("pairing is ambiguous" in error for error in result.errors) + assert all( + exchange.metadata["correlation_complete"] is False + and exchange.metadata["request_complete"] is False + for exchange in result.trajectory.exchanges + ) + + +def test_claude_first_raw_request_survives_missing_otel_body_event( + tmp_path: Path, +) -> None: + """Guards this PR against Claude omitting its first request-body OTLP event.""" + + capture = tmp_path / "capture" + raw = capture / "raw" + raw.mkdir(parents=True) + request_one = raw / "one.request.json" + request_two = raw / "two.request.json" + response_one = raw / "req_one.response.json" + response_two = raw / "req_two.response.json" + for path in (request_one, request_two, response_one, response_two): + path.write_text(json.dumps({"marker": path.name})) + request_one_timestamp = 1_776_999_999_000_000_000 + os.utime( + request_one, + ns=(request_one_timestamp, request_one_timestamp), + ) + records = [ + _otel_record( + "claude_code.api_response_body", + 1_777_000_000_500_000_000, + session_id="shared", + body_ref=str(response_one), + request_id="req_one", + ), + _otel_record( + "claude_code.api_request_body", + 1_777_000_000_600_000_000, + session_id="shared", + body_ref=str(request_two), + ), + _otel_record( + "claude_code.api_response_body", + 1_777_000_001_000_000_000, + session_id="shared", + body_ref=str(response_two), + request_id="req_two", + ), + ] + otel_dir = capture / "otel" + otel_dir.mkdir() + (otel_dir / "events.json").write_text( + json.dumps({"resourceLogs": [{"scopeLogs": [{"logRecords": records}]}]}) + ) + + result = parse_claude_raw_capture( + capture, + agent="claude-agent-acp", + session_id="rollout-1", + started_at=STARTED_AT, + ) + + assert result is not None + assert result.request_complete is True + assert result.errors == [] + assert len(result.trajectory.exchanges) == 2 + first, second = result.trajectory.exchanges + assert first.request.body["marker"] == "one.request.json" + assert first.metadata["pairing"] == "raw_file_mtime_fifo" + assert first.metadata["correlation_complete"] is True + assert second.request.body["marker"] == "two.request.json" + assert second.metadata["pairing"] == "otel_session_fifo" + + +def test_claude_session_fallback_is_truthfully_lower_fidelity(tmp_path: Path) -> None: + """Guards this PR's Claude OAuth fallback when raw OTel is unavailable.""" + + session = tmp_path / "claude" / "session.jsonl" + _write_jsonl( + session, + [ + { + "type": "user", + "timestamp": "2026-08-28T12:00:00Z", + "message": {"role": "user", "content": "Inspect the repo"}, + }, + { + "type": "assistant", + "requestId": "req-1", + "timestamp": "2026-08-28T12:00:01Z", + "message": { + "role": "assistant", + "model": "claude-opus-4-1", + "content": [{"type": "text", "text": "I will inspect it"}], + "usage": {"input_tokens": 5, "output_tokens": 3}, + }, + }, + { + "type": "assistant", + "requestId": "req-1", + "timestamp": "2026-08-28T12:00:02Z", + "message": { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "tool-1", "name": "Read"} + ], + "usage": {"input_tokens": 5, "output_tokens": 4}, + }, + }, + ], + ) + + result = parse_claude_sessions( + session.parent, + agent="claude-agent-acp", + session_id="rollout-1", + started_at=STARTED_AT, + ) + + assert result is not None + assert result.fidelity is CaptureFidelity.AGENT_SESSION + assert result.request_complete is False + assert result.response_complete is False + assert len(result.trajectory.exchanges) == 1 + assert len(result.trajectory.exchanges[0].response.body["content"]) == 2 + assert "tool_definitions" in result.missing_fields + + +def test_codex_oauth_session_splits_calls_and_preserves_usage(tmp_path: Path) -> None: + """Guards this PR's Codex OAuth native-session trajectory reconstruction.""" + + session = tmp_path / "codex" / "session.jsonl" + _write_jsonl( + session, + [ + { + "type": "session_meta", + "timestamp": "2026-08-28T12:00:00Z", + "payload": {"model": "gpt-5.6"}, + }, + { + "type": "response_item", + "timestamp": "2026-08-28T12:00:01Z", + "payload": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Run tests"}], + }, + }, + { + "type": "response_item", + "timestamp": "2026-08-28T12:00:02Z", + "payload": { + "type": "function_call", + "name": "exec_command", + "call_id": "call-1", + "arguments": '{"cmd":"pytest"}', + }, + }, + { + "type": "event_msg", + "timestamp": "2026-08-28T12:00:03Z", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": { + "input_tokens": 100, + "cached_input_tokens": 20, + "output_tokens": 10, + "reasoning_output_tokens": 4, + "total_tokens": 110, + } + }, + }, + }, + { + "type": "response_item", + "timestamp": "2026-08-28T12:00:04Z", + "payload": { + "type": "function_call_output", + "call_id": "call-1", + "output": "passed", + }, + }, + { + "type": "response_item", + "timestamp": "2026-08-28T12:00:05Z", + "payload": { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "All green"}], + }, + }, + { + "type": "event_msg", + "timestamp": "2026-08-28T12:00:06Z", + "payload": { + "type": "token_count", + "info": { + "last_token_usage": { + "input_tokens": 50, + "output_tokens": 5, + "total_tokens": 55, + } + }, + }, + }, + ], + ) + + result = parse_codex_sessions( + session.parent, + agent="codex-acp", + session_id="rollout-1", + started_at=STARTED_AT, + configured_model="gpt-5.6", + ) + + assert result is not None + assert result.source is CaptureSource.CODEX_NATIVE_SESSION + assert result.fidelity is CaptureFidelity.AGENT_SESSION + assert len(result.trajectory.exchanges) == 2 + first, second = result.trajectory.exchanges + assert first.response.body["usage"]["input_tokens"] == 100 + assert first.response.body["output"][0]["name"] == "exec_command" + assert second.request.body["input"][-1]["type"] == "function_call_output" + assert second.response.body["output"][0]["role"] == "assistant" + + +@pytest.mark.asyncio +async def test_capture_always_emits_empty_jsonl_and_terminal_manifest( + tmp_path: Path, +) -> None: + """Guards this PR's always-present artifact invariant for zero-call runs.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.configure({"OPENAI_API_KEY": "test-key"}) + + await capture.finalize(None, acp_events=[], model_call_seen=False) + + trajectory = tmp_path / "trajectory" / "llm_trajectory.jsonl" + manifest_path = tmp_path / "trajectory" / "llm_trajectory.manifest.json" + manifest = json.loads(manifest_path.read_text()) + assert trajectory.exists() + assert trajectory.read_text() == "" + assert manifest["status"] == CaptureStatus.NO_MODEL_CALL + assert manifest["exchange_count"] == 0 + assert manifest["auth_mode"] == "api_key" + + +@pytest.mark.parametrize( + ("auth_json", "expected"), + [ + ('{"auth_mode":"apikey","OPENAI_API_KEY":"test-key"}', "api_key"), + ( + '{"auth_mode":"chatgpt","tokens":{"refresh_token":"test-token"}}', + "oauth_subscription", + ), + ], +) +def test_codex_native_auth_file_mode_is_not_assumed_to_be_oauth( + tmp_path: Path, + auth_json: str, + expected: str, +) -> None: + """Guards this PR's Codex auth provenance for API-key auth.json files.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.configure({"CODEX_AUTH_JSON": auth_json}) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["auth_mode"] == expected + + +@pytest.mark.asyncio +async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> None: + """Guards this PR's uniform metadata for existing API-key proxy capture.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="azure-foundry-openai/gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.configure({"OPENAI_API_KEY": "test-key"}) + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"input": "hello"}}, + "response": {"status_code": 200, "body": {"output": []}}, + } + ) + + "\n" + ) + + await capture.finalize(None, acp_events=[], model_call_seen=True) + + exchange = json.loads(capture.trajectory_path.read_text()) + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert exchange["metadata"]["capture_fidelity"] == "provider_wire" + assert exchange["metadata"]["auth_mode"] == "api_key" + assert manifest["status"] == "complete" + assert manifest["capture_source"] == "litellm_proxy" + + +@pytest.mark.asyncio +async def test_provider_capture_early_return_cleans_native_raw_bodies( + tmp_path: Path, +) -> None: + """Guards this PR against raw-body leakage from mixed-role rollouts.""" + + commands: list[str] = [] + + class CleanupEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + return SimpleNamespace(return_code=0, stdout="", stderr="") + + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture._native_agents["claude-agent-acp"] = "claude-sonnet-4-6" + capture._capture_root_prepared = True + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"input": "hello"}}, + "response": {"status_code": 200, "body": {"output": []}}, + } + ) + + "\n" + ) + + await capture.finalize(CleanupEnv(), acp_events=[], model_call_seen=True) + + assert any("find /tmp/benchflow-llm-capture-" in command for command in commands) + assert any("-depth -delete" in command for command in commands) + + +@pytest.mark.asyncio +async def test_claude_capture_setup_failure_degrades_without_aborting( + tmp_path: Path, +) -> None: + """Guards this PR against observability setup breaking Claude OAuth runs.""" + + commands: list[str] = [] + + class FailingCollectorEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + if "nohup" in command: + return SimpleNamespace( + return_code=1, + stdout="", + stderr="node runtime not found", + ) + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def upload_file(self, *_args, **_kwargs): + return None + + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + prepared = await capture.prepare_agent( + FailingCollectorEnv(), + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "oauth-test-token"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + + assert prepared["CLAUDE_CODE_ENABLE_TELEMETRY"] == "1" + assert prepared["OTEL_LOG_RAW_API_BODIES"].startswith("file:") + assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in prepared + assert any("/opt/benchflow/node/bin/node" in command for command in commands) + + +@pytest.mark.asyncio +async def test_optional_session_download_creates_docker_copy_parent( + tmp_path: Path, +) -> None: + """Guards this PR's Docker native-session download path.""" + + class DockerLikeEnv: + async def exec(self, *_args, **_kwargs): + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def download_dir(self, _remote, local): + assert Path(local).parent.is_dir() + Path(local).mkdir() + + destination = tmp_path / "missing-parent" / "sessions" + downloaded = await _download_optional_dir( + DockerLikeEnv(), "/home/agent/.claude/projects", destination + ) + + assert downloaded is True + assert destination.is_dir() + + +def test_capture_failure_repairs_invalid_jsonl_and_redacts_manifest_error( + tmp_path: Path, +) -> None: + """Guards this PR's valid-JSONL and secret-redaction failure invariant.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.trajectory_path.write_text("{not-json\n") + + capture.record_failure( + "authorization: Bearer sk-test-secret-value", model_call_seen=True + ) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert capture.trajectory_path.read_text() == "" + assert manifest["status"] == "capture_failed" + assert "sk-test-secret-value" not in json.dumps(manifest) + + +def test_acp_projection_retains_the_actual_auth_mode() -> None: + """Guards this PR against labeling API-key fallback rows as OAuth.""" + + result = project_acp_trajectory( + [{"type": "agent_message", "text": "finished"}], + agent="codex-acp", + session_id="rollout-1", + started_at=STARTED_AT, + auth_mode="api_key", + ) + + assert result is not None + assert result.trajectory.exchanges[0].metadata["auth_mode"] == "api_key" + + +def test_agent_session_capture_is_audit_only_not_training_ready(tmp_path: Path) -> None: + """Guards this PR against silently training on reconstructed OAuth payloads.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + (trajectory_dir / "llm_trajectory.jsonl").write_text( + json.dumps( + { + "request": { + "body": {"messages": [{"role": "user", "content": "hello"}]} + }, + "response": { + "status_code": 200, + "body": { + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + }, + }, + "metadata": { + "capture_fidelity": "agent_session", + "request_complete": False, + "response_complete": True, + }, + } + ) + + "\n" + ) + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "partial", + "capture_fidelity": "agent_session", + "auth_mode": "oauth_subscription", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + } + ) + ) + + row = build_rollout_results_record( + tmp_path, + task_name="task", + rollout_name="rollout", + agent="claude-agent-acp", + agent_name="Claude Code", + model="claude-opus-4-1", + n_tool_calls=0, + prompts=["hello"], + trajectory=[], + partial_trajectory=False, + rewards={"reward": 1.0}, + error=None, + verifier_error=None, + agent_result={"usage_source": "agent_native_acp", "total_tokens": 2}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is True + assert row["error"] is None + + +def test_corrupt_capture_manifest_fails_closed_for_training(tmp_path: Path) -> None: + """Guards this PR against treating a corrupt new sidecar as a legacy artifact.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + (trajectory_dir / "llm_trajectory.jsonl").write_text( + json.dumps( + { + "request": { + "body": {"messages": [{"role": "user", "content": "hello"}]} + }, + "response": { + "status_code": 200, + "body": { + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + }, + }, + "metadata": { + "capture_fidelity": "provider_wire", + "request_complete": True, + "response_complete": True, + }, + } + ) + + "\n" + ) + (trajectory_dir / "llm_trajectory.manifest.json").write_text("{broken") + + row = build_rollout_results_record( + tmp_path, + task_name="task", + rollout_name="rollout", + agent="claude-agent-acp", + agent_name="Claude Code", + model="claude-opus-4-1", + n_tool_calls=0, + prompts=["hello"], + trajectory=[], + partial_trajectory=False, + rewards={"reward": 1.0}, + error=None, + verifier_error=None, + agent_result={"total_tokens": 2}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == ( + "missing_healthy_structured_llm_trajectory" + ) + assert row["is_completed"] is False From 9bc7d36ea9acd9654600e92ea91c46e90d7e25ec Mon Sep 17 00:00:00 2001 From: Bingran You Date: Fri, 28 Aug 2026 22:32:44 -0700 Subject: [PATCH 02/74] chore: update RestrictedPython security fix --- uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/uv.lock b/uv.lock index 971e01e27..c1224bb31 100644 --- a/uv.lock +++ b/uv.lock @@ -3937,11 +3937,11 @@ wheels = [ [[package]] name = "restrictedpython" -version = "8.2" +version = "8.5" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/97/04/7314639eab9edd57b5c8f6e158a558f5a5a65453c37a51fb8f9ecc1d112e/restrictedpython-8.2.tar.gz", hash = "sha256:b835c2ff87d71b76f6d8b129096cc5a1e893bf01839ab802ba47a4129f1bdfe4", size = 449323, upload-time = "2026-05-29T06:27:57.285Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7c/3b/8e41f7cfabbb30b1013ebc7484303d6c87da2906ec432d69dea11d2f7d75/restrictedpython-8.5.tar.gz", hash = "sha256:4ed1269dbe3caa88db650d1af325198a952aeb1451eca05df0cfa65db4466215", size = 455879, upload-time = "2026-08-19T07:02:10.934Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/35/59/b2778bfafcc37fdfd87c02702bbfc5753c58c28c7eb82999b7b49ec6e0e7/restrictedpython-8.2-py3-none-any.whl", hash = "sha256:6ff69e2fe06674bfe88943f17f2ce8353a9bc743ee37ea7bec38ff8b4851afd1", size = 27176, upload-time = "2026-05-29T06:27:55.651Z" }, + { url = "https://files.pythonhosted.org/packages/58/57/16ce3c721f5a33317e4110575d5c9976c0c45f7fd96ca2e0adeab06e6026/restrictedpython-8.5-py3-none-any.whl", hash = "sha256:6c70e0a3af13e830d37225788cdc8ab5804a8df4b500c135086eaef34b5c01e0", size = 30962, upload-time = "2026-08-19T07:02:09.553Z" }, ] [[package]] From 77fcef6d6e650c59d82a389321ab0c9967f4a6fc Mon Sep 17 00:00:00 2001 From: Bingran You Date: Fri, 28 Aug 2026 22:34:49 -0700 Subject: [PATCH 03/74] style: format trajectory capture modules --- src/benchflow/trajectories/llm_capture.py | 9 ++++++--- .../trajectories/native_capture_parsers.py | 20 +++++++++---------- src/benchflow/trajectories/results.py | 1 + tests/trajectories/test_native_llm_capture.py | 4 +--- 4 files changed, 18 insertions(+), 16 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 7611dfc4f..6a96316bb 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -475,8 +475,7 @@ def _is_claude_code_agent(agent: str) -> bool: config = AGENTS.get(agent) subscription = config.subscription_auth if config is not None else None return bool( - subscription is not None - and subscription.replaces_env == "ANTHROPIC_API_KEY" + subscription is not None and subscription.replaces_env == "ANTHROPIC_API_KEY" ) @@ -570,7 +569,11 @@ def model_call_seen_from_evidence( acp_events: list[dict[str, Any]], ) -> bool: for value in (usage_metrics or {}).values(): - if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0: + if ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and value > 0 + ): return True return any( event.get("type") in {"agent_message", "agent_thought", "tool_call"} diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index f7c6ff9e4..4cd39bca8 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -588,7 +588,9 @@ def _load_otel_events(root: Path) -> list[_OtelEvent]: timestamp = _unix_nanos_timestamp( record.get("timeUnixNano") or record.get("observedTimeUnixNano") ) - events.append(_OtelEvent(name=name, timestamp=timestamp, attributes=attributes)) + events.append( + _OtelEvent(name=name, timestamp=timestamp, attributes=attributes) + ) return events @@ -676,7 +678,11 @@ def _merge_claude_assistant_group(records: list[dict[str, Any]]) -> dict[str, An if isinstance(message.get("usage"), dict): merged["usage"] = message["usage"] content = message.get("content") - blocks = content if isinstance(content, list) else [{"type": "text", "text": content}] + blocks = ( + content + if isinstance(content, list) + else [{"type": "text", "text": content}] + ) for block in blocks: marker = json.dumps(block, sort_keys=True, default=str) if marker not in seen_content: @@ -686,11 +692,7 @@ def _merge_claude_assistant_group(records: list[dict[str, Any]]) -> dict[str, An def _message_for_history(message: dict[str, Any]) -> dict[str, Any]: - return { - key: value - for key, value in message.items() - if key in {"role", "content"} - } + return {key: value for key, value in message.items() if key in {"role", "content"}} def _codex_session_model(records: list[dict[str, Any]]) -> str: @@ -706,9 +708,7 @@ def _codex_session_model(records: list[dict[str, Any]]) -> str: def _normalize_codex_usage(usage: dict[str, Any]) -> dict[str, Any]: return { "input_tokens": usage.get("input_tokens", 0), - "input_tokens_details": { - "cached_tokens": usage.get("cached_input_tokens", 0) - }, + "input_tokens_details": {"cached_tokens": usage.get("cached_input_tokens", 0)}, "output_tokens": usage.get("output_tokens", 0), "output_tokens_details": { "reasoning_tokens": usage.get("reasoning_output_tokens", 0) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 3f67f7df7..c463f897d 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -43,6 +43,7 @@ logger = logging.getLogger(__name__) + def _record_to_redacted_json_line(record: dict[str, Any]) -> str: redacted = redact_trajectory_obj(scrub_non_finite(record)) return json.dumps(redacted, default=str, allow_nan=False) diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index e7adc3bee..aa64ef265 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -290,9 +290,7 @@ def test_claude_session_fallback_is_truthfully_lower_fidelity(tmp_path: Path) -> "timestamp": "2026-08-28T12:00:02Z", "message": { "role": "assistant", - "content": [ - {"type": "tool_use", "id": "tool-1", "name": "Read"} - ], + "content": [{"type": "tool_use", "id": "tool-1", "name": "Read"}], "usage": {"input_tokens": 5, "output_tokens": 4}, }, }, From 4c3ae11b1f482ffd6077217fa52174728ed11f34 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Fri, 28 Aug 2026 23:26:24 -0700 Subject: [PATCH 04/74] fix: harden mixed-auth trajectory capture --- docs/getting-started.md | 14 +- src/benchflow/eval_artifacts.py | 16 + .../trajectories/export_prime_sft.py | 11 + src/benchflow/trajectories/export_trl_sft.py | 17 + src/benchflow/trajectories/llm_capture.py | 294 ++++++------ .../trajectories/llm_capture_manifest.py | 44 ++ .../trajectories/llm_capture_records.py | 451 ++++++++++++++++++ .../trajectories/native_capture_parsers.py | 29 +- src/benchflow/trajectories/results.py | 34 +- tests/test_eval_artifact_cli.py | 29 ++ tests/trajectories/test_export_prime_sft.py | 41 +- .../test_llm_capture_training_contract.py | 99 ++++ tests/trajectories/test_native_llm_capture.py | 396 ++++++++++----- 13 files changed, 1167 insertions(+), 308 deletions(-) create mode 100644 src/benchflow/trajectories/llm_capture_records.py create mode 100644 tests/trajectories/test_llm_capture_training_contract.py diff --git a/docs/getting-started.md b/docs/getting-started.md index 4b2e15b4c..6990b5e3c 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -226,11 +226,15 @@ the source of truth for interpreting the JSONL: | Codex subscription/OAuth | Codex native session JSONL | `agent_session` | The manifest status is `complete`, `partial`, `no_model_call`, or -`capture_failed`. Reconstructed `agent_session` rows remain useful for audit and -viewer workflows, but trainer exports fail closed unless the manifest says the -capture is complete provider-wire data. Claude's own raw-body telemetry can -still contain provider-redacted extended-thinking blocks; BenchFlow also applies -its normal secret redaction before publishing the JSONL. +`capture_failed`. Mixed-role rollouts merge API-key and native-subscription +exchanges into the same JSONL; `role_captures` records each prepared role's +agent, model, auth mode, source, fidelity, completeness, and exchange count. +Missing or ambiguously attributed roles make the rollout-level capture +`partial`. Reconstructed `agent_session` rows remain useful for audit and viewer +workflows, but trainer exports fail closed unless the manifest says the capture +is complete provider-wire data. Claude's own raw-body telemetry can still +contain provider-redacted extended-thinking blocks; BenchFlow also applies its +normal secret redaction before publishing the JSONL. ### Reading results diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index f3169aa2a..188320d23 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -15,6 +15,10 @@ PrimeSftTrajectoryJsonlError, load_llm_trajectory_jsonl, ) +from benchflow.trajectories.llm_capture_manifest import ( + capture_manifest_allows_training, + read_llm_trajectory_manifest, +) CanonicalizePolicy = Literal["none", "one-healthy-per-task"] RetryPolicy = Literal["default", "unscored-only"] @@ -202,6 +206,18 @@ def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, int]: rows = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError: return True, False, 0 + manifest = read_llm_trajectory_manifest(rollout_dir) + if manifest is not None: + try: + expected_rows = int(manifest.get("exchange_count") or 0) + except (TypeError, ValueError): + return True, False, len(rows) + if ( + not capture_manifest_allows_training(manifest) + or not rows + or expected_rows != len(rows) + ): + return True, False, len(rows) return True, True, len(rows) diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 7f09d873f..8645c40cb 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -18,6 +18,10 @@ from typing import Any, Literal, cast from benchflow._utils.json_safe import dumps_finite, scrub_non_finite +from benchflow.trajectories.llm_capture_manifest import ( + capture_manifest_allows_training, + read_llm_trajectory_manifest, +) from benchflow.trajectories.types import redact_trajectory_obj PrimeSftRowMode = Literal["rollout", "exchange"] @@ -1215,6 +1219,13 @@ def convert_benchflow_rollouts_to_prime_sft_rows( stats.skipped_reward += 1 continue + capture_manifest = read_llm_trajectory_manifest(rollout_dir) + if capture_manifest is not None and not capture_manifest_allows_training( + capture_manifest + ): + stats.skipped_insufficient_capture_fidelity += 1 + continue + trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) if not exchanges: diff --git a/src/benchflow/trajectories/export_trl_sft.py b/src/benchflow/trajectories/export_trl_sft.py index 8a51fd65e..f1be50f0a 100644 --- a/src/benchflow/trajectories/export_trl_sft.py +++ b/src/benchflow/trajectories/export_trl_sft.py @@ -24,6 +24,10 @@ normalize_prime_sft_exchange, validate_prime_sft_row, ) +from benchflow.trajectories.llm_capture_manifest import ( + capture_manifest_allows_training, + read_llm_trajectory_manifest, +) from benchflow.trajectories.types import redact_trajectory_obj TrlSftRowMode = Literal["rollout", "exchange"] @@ -42,6 +46,7 @@ class TrlSftExportStats: skipped_provider_error: int = 0 skipped_no_assistant: int = 0 skipped_missing_tools: int = 0 + skipped_insufficient_capture_fidelity: int = 0 skipped_terminal_error: int = 0 skipped_helper_calls: int = 0 skipped_invalid: int = 0 @@ -64,6 +69,9 @@ def as_dict(self) -> dict[str, Any]: "skipped_provider_error": self.skipped_provider_error, "skipped_no_assistant": self.skipped_no_assistant, "skipped_missing_tools": self.skipped_missing_tools, + "skipped_insufficient_capture_fidelity": ( + self.skipped_insufficient_capture_fidelity + ), "skipped_terminal_error": self.skipped_terminal_error, "skipped_helper_calls": self.skipped_helper_calls, "skipped_invalid": self.skipped_invalid, @@ -314,6 +322,12 @@ def convert_benchflow_rollouts_to_trl_sft_rows( if min_reward is not None and (reward is None or reward < min_reward): stats.skipped_reward += 1 continue + capture_manifest = read_llm_trajectory_manifest(rollout_dir) + if capture_manifest is not None and not capture_manifest_allows_training( + capture_manifest + ): + stats.skipped_insufficient_capture_fidelity += 1 + continue trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) if not exchanges: @@ -351,6 +365,9 @@ def convert_benchflow_rollouts_to_trl_sft_rows( if skip_reason == "missing_tool_defs": stats.skipped_missing_tools += 1 continue + if skip_reason == "insufficient_capture_fidelity": + stats.skipped_insufficient_capture_fidelity += 1 + continue if row is None: stats.skipped_invalid += 1 continue diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 6a96316bb..f7295c575 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -15,28 +15,35 @@ from benchflow.agents.env import uses_native_subscription_auth from benchflow.agents.registry import AGENTS -from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter from benchflow.trajectories.llm_capture_manifest import ( LLM_TRAJECTORY_FILENAME, - LLM_TRAJECTORY_SCHEMA_VERSION, AuthMode, CaptureFidelity, CaptureSource, CaptureStatus, + LLMRoleCapture, initialize_llm_trajectory_artifacts, write_llm_trajectory_manifest, ) +from benchflow.trajectories.llm_capture_records import ( + CaptureTarget as _CaptureTarget, +) +from benchflow.trajectories.llm_capture_records import ( + NativeCaptureBundle as _NativeCaptureBundle, +) +from benchflow.trajectories.llm_capture_records import ( + assemble_capture, + load_provider_wire_records, + role_captures_for_targets, + write_exchange_records, +) from benchflow.trajectories.native_capture_parsers import ( - NativeParseResult, parse_claude_raw_capture, parse_claude_sessions, parse_codex_sessions, project_acp_trajectory, ) -from benchflow.trajectories.types import ( - redact_trajectory_obj, - redact_trajectory_text, -) +from benchflow.trajectories.types import redact_trajectory_text logger = logging.getLogger(__name__) @@ -99,8 +106,7 @@ def __init__( session_id=session_id, started_at=started_at, ) - self._native_agents: dict[str, str | None] = {} - self._credential_homes: set[str] = set() + self._targets: dict[tuple[str, str | None, str], _CaptureTarget] = {} self._collector_started = False self._capture_root_prepared = False self._preparation_errors: list[str] = [] @@ -130,11 +136,20 @@ async def prepare_agent( """Return the agent environment augmented for native capture.""" prepared = dict(agent_env) - if not uses_native_subscription_auth(agent, model, prepared): + native = uses_native_subscription_auth(agent, model, prepared) + auth_mode = _resolve_auth_mode(agent, model, prepared) + target = _CaptureTarget( + agent=agent, + model=model, + credential_home=credential_home, + auth_mode=auth_mode, + native=native, + ) + self._targets[(agent, model, credential_home)] = target + self._refresh_manifest_auth_mode() + if not native: + write_llm_trajectory_manifest(self.rollout_dir, self.manifest) return prepared - self._native_agents[agent] = model - self._credential_homes.add(credential_home) - self.manifest.auth_mode = _resolve_auth_mode(agent, model, prepared) if _is_claude_code_agent(agent): raw_dir = f"{self._remote_capture_root}/raw" prepared.update( @@ -168,6 +183,16 @@ async def prepare_agent( write_llm_trajectory_manifest(self.rollout_dir, self.manifest) return prepared + def _native_targets(self) -> list[_CaptureTarget]: + return [target for target in self._targets.values() if target.native] + + def _refresh_manifest_auth_mode(self) -> None: + modes = {target.auth_mode for target in self._targets.values()} + if len(modes) == 1: + self.manifest.auth_mode = next(iter(modes)) + elif len(modes) > 1: + self.manifest.auth_mode = AuthMode.MIXED + async def finalize( self, env: Any, @@ -178,92 +203,84 @@ async def finalize( """Publish the highest-fidelity available capture and terminal sidecar.""" self.manifest.finished_at = datetime.now() + provider_records: list[dict[str, Any]] = [] if self.trajectory_path.stat().st_size > 0: try: - exchange_count = _annotate_provider_wire_jsonl( + provider_records = load_provider_wire_records( self.trajectory_path, - auth_mode=self.manifest.auth_mode, + targets=[ + target for target in self._targets.values() if not target.native + ], + fallback_agent=self.agent, + fallback_model=self.model, + fallback_auth=self.manifest.auth_mode, ) - self._finish_manifest( - status=CaptureStatus.COMPLETE, - source=CaptureSource.LITELLM_PROXY, - fidelity=CaptureFidelity.PROVIDER_WIRE, - exchange_count=exchange_count, - request_complete=True, - response_complete=True, - ) - finally: - # Mixed-role scenes can prepare native telemetry before another - # role writes provider capture. Never strand raw bodies in a - # reusable externally-owned sandbox on this early-return path. - if self._native_agents and env is not None: + except Exception: + if env is not None: await self._cleanup_remote_capture(env) - return + raise - native_result: NativeParseResult | None = None + native_bundles: list[_NativeCaptureBundle] = [] collection_errors: list[str] = list(self._preparation_errors) - if self._native_agents and env is not None: + native_targets = self._native_targets() + if native_targets and env is not None: try: - native_result = await self._collect_native_result(env) + native_bundles = await self._collect_native_results(env) except Exception as exc: collection_errors.append(_sanitized_error(exc)) logger.warning("Native LLM trajectory collection failed: %s", exc) finally: await self._cleanup_remote_capture(env) - if native_result is None and acp_events: - native_result = project_acp_trajectory( + if not provider_records and not native_bundles and acp_events: + projected = project_acp_trajectory( acp_events, agent=self.agent, session_id=self.session_id, started_at=self.started_at, auth_mode=self.manifest.auth_mode.value, ) - if native_result is not None: - LiveLLMTrajectoryWriter(self.trajectory_path).reconcile( - native_result.trajectory - ) - errors = [*collection_errors, *native_result.errors] - status = ( - CaptureStatus.COMPLETE - if native_result.fidelity is CaptureFidelity.PROVIDER_WIRE - and not errors - else CaptureStatus.PARTIAL - ) - self._finish_manifest( - status=status, - source=native_result.source, - fidelity=native_result.fidelity, - exchange_count=len(native_result.trajectory.exchanges), - request_complete=native_result.request_complete, - response_complete=native_result.response_complete, - missing_fields=native_result.missing_fields, - errors=errors, - ) - return + if projected is not None: + native_bundles.append( + _NativeCaptureBundle( + targets=( + tuple(native_targets) + or (self._fallback_target(native=True),) + ), + result=projected, + ) + ) + prepared_targets = list(self._targets.values()) + assembly = assemble_capture( + provider_records=provider_records, + native_bundles=native_bundles, + targets=prepared_targets, + collection_errors=collection_errors, + model_call_seen=model_call_seen, + fallback_auth=self.manifest.auth_mode, + ) + write_exchange_records(self.trajectory_path, assembly.records) + self.manifest.auth_mode = assembly.auth_mode self._finish_manifest( - status=( - CaptureStatus.CAPTURE_FAILED - if model_call_seen - else CaptureStatus.NO_MODEL_CALL - ), - source=CaptureSource.NONE, - fidelity=CaptureFidelity.NONE, - exchange_count=0, - request_complete=False, - response_complete=False, - missing_fields=( - ["provider_request", "provider_response"] if model_call_seen else [] - ), - errors=( - collection_errors - or ( - ["model call observed but no LLM capture source was readable"] - if model_call_seen - else [] - ) - ), + status=assembly.status, + source=assembly.source, + fidelity=assembly.fidelity, + exchange_count=len(assembly.records), + request_complete=assembly.request_complete, + response_complete=assembly.response_complete, + missing_fields=assembly.missing_fields, + errors=assembly.errors, + role_captures=assembly.role_captures, + ) + + def _fallback_target(self, *, native: bool) -> _CaptureTarget: + return _CaptureTarget( + agent=self.agent, + model=self.model, + credential_home="", + auth_mode=self.manifest.auth_mode, + native=native, ) def record_failure(self, error: object, *, model_call_seen: bool) -> None: @@ -286,6 +303,7 @@ def record_failure(self, error: object, *, model_call_seen: bool) -> None: ["provider_request", "provider_response"] if model_call_seen else [] ), errors=[_sanitized_error(error)], + role_captures=role_captures_for_targets(list(self._targets.values())), ) async def _ensure_otel_sink(self, env: Any, *, sandbox_user: str | None) -> int: @@ -362,7 +380,7 @@ async def _read_collector_port(self, env: Any) -> int: raise RuntimeError("Claude OTel sink port file is unavailable") return _parse_port(result.stdout) - async def _collect_native_result(self, env: Any) -> NativeParseResult | None: + async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: if self._collector_started: await env.exec( f"if test -s {self._remote_capture_root}/pid; then " @@ -371,54 +389,88 @@ async def _collect_native_result(self, env: Any) -> NativeParseResult | None: user="root", timeout_sec=5, ) + bundles: list[_NativeCaptureBundle] = [] + native_targets = self._native_targets() + claude_targets = tuple( + target for target in native_targets if _is_claude_code_agent(target.agent) + ) with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: local_root = Path(temporary) capture_dir = local_root / "capture" + raw_claude_captured = False if self._capture_root_prepared: await env.download_dir(self._remote_capture_root, capture_dir) result = parse_claude_raw_capture( capture_dir, - agent=self.agent, + agent=(claude_targets[0].agent if claude_targets else self.agent), session_id=self.session_id, started_at=self.started_at, ) if result is not None: - return result - session_roots: list[tuple[str, Path]] = [] - for index, credential_home in enumerate(sorted(self._credential_homes)): + bundles.append( + _NativeCaptureBundle(targets=claude_targets, result=result) + ) + raw_claude_captured = True + credential_homes = sorted( + {target.credential_home for target in native_targets} + ) + for index, credential_home in enumerate(credential_homes): + home_targets = tuple( + target + for target in native_targets + if target.credential_home == credential_home + ) + home_claude_targets = tuple( + target + for target in home_targets + if _is_claude_code_agent(target.agent) + ) + home_codex_targets = tuple( + target for target in home_targets if target.agent == "codex-acp" + ) claude_local = local_root / f"home-{index}" / "claude-projects" codex_local = local_root / f"home-{index}" / "codex-sessions" - if any( - _is_claude_code_agent(agent) for agent in self._native_agents - ) and await _download_optional_dir( - env, f"{credential_home}/.claude/projects", claude_local - ): - session_roots.append(("claude", claude_local)) - if "codex-acp" in self._native_agents and await _download_optional_dir( - env, f"{credential_home}/.codex/sessions", codex_local + if ( + home_claude_targets + and not raw_claude_captured + and await _download_optional_dir( + env, f"{credential_home}/.claude/projects", claude_local + ) ): - session_roots.append(("codex", codex_local)) - - for source, root in session_roots: - if source == "claude": + target = home_claude_targets[0] result = parse_claude_sessions( - root, - agent=self.agent, + claude_local, + agent=target.agent, session_id=self.session_id, started_at=self.started_at, ) - else: + if result is not None: + bundles.append( + _NativeCaptureBundle( + targets=home_claude_targets, + result=result, + ) + ) + if home_codex_targets and await _download_optional_dir( + env, f"{credential_home}/.codex/sessions", codex_local + ): + target = home_codex_targets[0] result = parse_codex_sessions( - root, - agent=self.agent, + codex_local, + agent=target.agent, session_id=self.session_id, started_at=self.started_at, - configured_model=self._native_agents.get("codex-acp"), - auth_mode=self.manifest.auth_mode.value, + configured_model=target.model, + auth_mode=target.auth_mode.value, ) - if result is not None: - return result - return None + if result is not None: + bundles.append( + _NativeCaptureBundle( + targets=home_codex_targets, + result=result, + ) + ) + return bundles async def _cleanup_remote_capture(self, env: Any) -> None: if not self._capture_root_prepared: @@ -446,6 +498,7 @@ def _finish_manifest( response_complete: bool, missing_fields: list[str] | None = None, errors: list[str] | None = None, + role_captures: list[LLMRoleCapture] | None = None, ) -> None: self.manifest.status = status self.manifest.capture_source = source @@ -455,6 +508,7 @@ def _finish_manifest( self.manifest.response_complete = response_complete self.manifest.missing_fields = sorted(set(missing_fields or [])) self.manifest.errors = [_sanitized_error(item) for item in errors or []] + self.manifest.role_captures = role_captures or [] write_llm_trajectory_manifest(self.rollout_dir, self.manifest) @@ -521,38 +575,6 @@ def _parse_port(value: str) -> int: return port -def _annotate_provider_wire_jsonl(path: Path, *, auth_mode: AuthMode) -> int: - records: list[dict[str, Any]] = [] - for line_number, line in enumerate(path.read_text().splitlines(), start=1): - if not line.strip(): - continue - try: - record = json.loads(line) - except json.JSONDecodeError as exc: - raise ValueError( - f"invalid LLM trajectory JSONL at line {line_number}: {exc.msg}" - ) from exc - if not isinstance(record, dict): - raise ValueError( - f"invalid LLM trajectory JSONL at line {line_number}: expected object" - ) - metadata = record.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - record["metadata"] = metadata - metadata.setdefault("schema_version", LLM_TRAJECTORY_SCHEMA_VERSION) - metadata.setdefault("capture_source", CaptureSource.LITELLM_PROXY.value) - metadata.setdefault("capture_fidelity", CaptureFidelity.PROVIDER_WIRE.value) - metadata.setdefault("auth_mode", auth_mode.value) - metadata.setdefault("request_complete", True) - metadata.setdefault("response_complete", True) - metadata.setdefault("payload_redacted", True) - records.append(redact_trajectory_obj(record)) - payload = "".join(json.dumps(record, default=str) + "\n" for record in records) - _atomic_replace_text(path, payload) - return len(records) - - def _sanitized_error(error: object) -> str: text = redact_trajectory_text(str(error)).replace("\n", " ").strip() return text[:500] or type(error).__name__ diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 219306a00..525ccfa6e 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -12,6 +12,7 @@ from datetime import datetime from enum import StrEnum from pathlib import Path +from typing import Any from pydantic import BaseModel, Field @@ -32,6 +33,7 @@ class CaptureFidelity(StrEnum): PROVIDER_WIRE = "provider_wire" AGENT_SESSION = "agent_session" ACP_PROJECTION = "acp_projection" + MIXED = "mixed" NONE = "none" @@ -41,15 +43,30 @@ class CaptureSource(StrEnum): CLAUDE_NATIVE_SESSION = "claude_native_session" CODEX_NATIVE_SESSION = "codex_native_session" ACP_PROJECTION = "acp_projection" + MIXED = "mixed" NONE = "none" class AuthMode(StrEnum): API_KEY = "api_key" OAUTH_SUBSCRIPTION = "oauth_subscription" + MIXED = "mixed" UNKNOWN = "unknown" +class LLMRoleCapture(BaseModel): + """Per prepared role provenance for mixed-auth/mixed-agent rollouts.""" + + agent: str + model: str | None = None + auth_mode: AuthMode + capture_source: CaptureSource + capture_fidelity: CaptureFidelity + exchange_count: int + request_complete: bool + response_complete: bool + + class LLMTrajectoryManifest(BaseModel): """Machine-readable fidelity and lifecycle state for the JSONL artifact.""" @@ -69,6 +86,7 @@ class LLMTrajectoryManifest(BaseModel): finished_at: datetime | None = None missing_fields: list[str] = Field(default_factory=list) errors: list[str] = Field(default_factory=list) + role_captures: list[LLMRoleCapture] = Field(default_factory=list) def initialize_llm_trajectory_artifacts( @@ -104,6 +122,32 @@ def write_llm_trajectory_manifest( _atomic_write_text(path, payload + "\n") +def read_llm_trajectory_manifest(rollout_dir: Path) -> dict[str, Any] | None: + """Read the sidecar, distinguishing absent legacy data from corruption.""" + + path = rollout_dir / "trajectory" / LLM_TRAJECTORY_MANIFEST_FILENAME + if not path.exists(): + return None + try: + manifest = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return {"status": "capture_failed", "capture_fidelity": "none"} + if not isinstance(manifest, dict): + return {"status": "capture_failed", "capture_fidelity": "none"} + return manifest + + +def capture_manifest_allows_training(manifest: dict[str, Any]) -> bool: + """Return whether the rollout-level capture is safe for training export.""" + + return bool( + manifest.get("status") == "complete" + and manifest.get("capture_fidelity") == "provider_wire" + and manifest.get("request_complete") is True + and manifest.get("response_complete") is True + ) + + def _atomic_write_text(path: Path, payload: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py new file mode 100644 index 000000000..02f5b6006 --- /dev/null +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -0,0 +1,451 @@ +"""Pure record assembly for uniform LLM trajectory capture.""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from benchflow.trajectories.llm_capture_manifest import ( + LLM_TRAJECTORY_SCHEMA_VERSION, + AuthMode, + CaptureFidelity, + CaptureSource, + CaptureStatus, + LLMRoleCapture, +) +from benchflow.trajectories.native_capture_parsers import NativeParseResult +from benchflow.trajectories.types import redact_trajectory_obj + + +@dataclass(frozen=True) +class CaptureTarget: + """One prepared agent role whose model calls should be attributable.""" + + agent: str + model: str | None + credential_home: str + auth_mode: AuthMode + native: bool + + +@dataclass(frozen=True) +class NativeCaptureBundle: + """A native parser result plus the roles it may describe.""" + + targets: tuple[CaptureTarget, ...] + result: NativeParseResult + + +@dataclass(frozen=True) +class CaptureAssembly: + """Terminal artifact state derived from every available capture source.""" + + records: list[dict[str, Any]] + status: CaptureStatus + source: CaptureSource + fidelity: CaptureFidelity + auth_mode: AuthMode + request_complete: bool + response_complete: bool + missing_fields: list[str] + errors: list[str] + role_captures: list[LLMRoleCapture] + + +def load_provider_wire_records( + path: Path, + *, + targets: list[CaptureTarget], + fallback_agent: str, + fallback_model: str | None, + fallback_auth: AuthMode, +) -> list[dict[str, Any]]: + """Load and attribute the shared LiteLLM JSONL without guessing roles.""" + + records: list[dict[str, Any]] = [] + for line_number, line in enumerate(path.read_text().splitlines(), start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"invalid LLM trajectory JSONL at line {line_number}: {exc.msg}" + ) from exc + if not isinstance(record, dict): + raise ValueError( + f"invalid LLM trajectory JSONL at line {line_number}: expected object" + ) + metadata = record.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + record["metadata"] = metadata + target = _target_for_model(targets, _record_model(record)) + attribution_complete = target is not None or not targets + metadata.update( + { + "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, + "capture_source": CaptureSource.LITELLM_PROXY.value, + "capture_fidelity": CaptureFidelity.PROVIDER_WIRE.value, + "auth_mode": ( + target.auth_mode.value + if target is not None + else fallback_auth.value + ), + "agent": ( + target.agent + if target is not None + else (fallback_agent if not targets else "mixed") + ), + "model": ( + target.model + if target is not None + else (_record_model(record) or fallback_model) + ), + "role_attribution_complete": attribution_complete, + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + } + ) + if not attribution_complete: + metadata["role_candidates"] = _role_candidates(targets) + records.append(redact_trajectory_obj(record)) + return records + + +def assemble_capture( + *, + provider_records: list[dict[str, Any]], + native_bundles: list[NativeCaptureBundle], + targets: list[CaptureTarget], + collection_errors: list[str], + model_call_seen: bool, + fallback_auth: AuthMode, +) -> CaptureAssembly: + """Merge all sources and derive one fail-closed terminal contract.""" + + native_records = [ + record for bundle in native_bundles for record in _native_bundle_records(bundle) + ] + records = _sort_exchange_records([*provider_records, *native_records]) + errors = list(collection_errors) + attribution_incomplete = any( + _record_metadata(record).get("role_attribution_complete") is False + for record in records + ) + if attribution_incomplete: + errors.append( + "one or more exchanges could not be attributed to a unique prepared role" + ) + + captured_targets = { + *_captured_targets(provider_records, targets), + *_captured_targets(native_records, targets), + } + missing_targets = [target for target in targets if target not in captured_targets] + errors.extend( + f"no capture was attributable to {target.agent} " + f"({target.model or 'unknown model'}, {target.auth_mode.value})" + for target in missing_targets + ) + role_captures = _role_captures(records, targets=targets) + auth_mode = _aggregate_auth_mode( + records, targets=targets, fallback_auth=fallback_auth + ) + + if not records: + if model_call_seen and not errors: + errors.append("model call observed but no LLM capture source was readable") + return CaptureAssembly( + records=[], + status=( + CaptureStatus.CAPTURE_FAILED + if model_call_seen + else CaptureStatus.NO_MODEL_CALL + ), + source=CaptureSource.NONE, + fidelity=CaptureFidelity.NONE, + auth_mode=auth_mode, + request_complete=False, + response_complete=False, + missing_fields=( + ["provider_request", "provider_response"] if model_call_seen else [] + ), + errors=errors, + role_captures=role_captures, + ) + + missing_fields = { + field for bundle in native_bundles for field in bundle.result.missing_fields + } + if attribution_incomplete: + missing_fields.add("role_attribution") + source = _aggregate_source(records) + fidelity = _aggregate_fidelity(records) + if missing_targets: + source = CaptureSource.MIXED + fidelity = CaptureFidelity.MIXED + request_complete = ( + not missing_targets + and all(_record_metadata_bool(record, "request_complete") for record in records) + and all(bundle.result.request_complete for bundle in native_bundles) + ) + response_complete = ( + not missing_targets + and all( + _record_metadata_bool(record, "response_complete") for record in records + ) + and all(bundle.result.response_complete for bundle in native_bundles) + ) + errors.extend(error for bundle in native_bundles for error in bundle.result.errors) + status = ( + CaptureStatus.COMPLETE + if fidelity is CaptureFidelity.PROVIDER_WIRE + and request_complete + and response_complete + and not errors + else CaptureStatus.PARTIAL + ) + return CaptureAssembly( + records=records, + status=status, + source=source, + fidelity=fidelity, + auth_mode=auth_mode, + request_complete=request_complete, + response_complete=response_complete, + missing_fields=sorted(missing_fields), + errors=errors, + role_captures=role_captures, + ) + + +def write_exchange_records(path: Path, records: list[dict[str, Any]]) -> None: + """Atomically replace the JSONL with redacted assembled records.""" + + payload = "".join( + json.dumps(redact_trajectory_obj(record), default=str) + "\n" + for record in records + ) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(payload) + os.replace(temporary, path) + + +def role_captures_for_targets(targets: list[CaptureTarget]) -> list[LLMRoleCapture]: + """Return explicit zero-exchange provenance for a failed finalization.""" + + return _role_captures([], targets=targets) + + +def _native_bundle_records(bundle: NativeCaptureBundle) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for line in bundle.result.trajectory.to_jsonl(redact_keys=True).splitlines(): + if not line.strip(): + continue + record = json.loads(line) + if not isinstance(record, dict): + continue + metadata = record.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + record["metadata"] = metadata + target = _target_for_model(list(bundle.targets), _record_model(record)) + if target is not None: + metadata.update( + { + "agent": target.agent, + "model": target.model or _record_model(record), + "auth_mode": target.auth_mode.value, + "role_attribution_complete": True, + } + ) + elif bundle.targets: + metadata.update( + { + "agent": "mixed", + "model": _record_model(record), + "auth_mode": AuthMode.MIXED.value, + "role_attribution_complete": False, + "role_candidates": _role_candidates(list(bundle.targets)), + } + ) + records.append(redact_trajectory_obj(record)) + return records + + +def _role_candidates(targets: list[CaptureTarget]) -> list[dict[str, Any]]: + return [ + { + "agent": target.agent, + "model": target.model, + "auth_mode": target.auth_mode.value, + } + for target in targets + ] + + +def _captured_targets( + records: list[dict[str, Any]], targets: list[CaptureTarget] +) -> set[CaptureTarget]: + captured: set[CaptureTarget] = set() + for record in records: + metadata = _record_metadata(record) + if metadata.get("role_attribution_complete") is not True: + continue + for target in targets: + if ( + metadata.get("agent") == target.agent + and metadata.get("auth_mode") == target.auth_mode.value + and _model_matches_target(metadata.get("model"), target.model) + ): + captured.add(target) + return captured + + +def _model_matches_target(value: Any, configured_model: str | None) -> bool: + if configured_model is None: + return True + if not isinstance(value, str): + return False + normalized_value = value.casefold() + normalized_target = configured_model.casefold() + return bool( + normalized_value == normalized_target + or normalized_value.endswith(f"/{normalized_target}") + or normalized_target.endswith(f"/{normalized_value}") + ) + + +def _record_model(record: dict[str, Any]) -> str | None: + request = record.get("request") + body = request.get("body") if isinstance(request, dict) else None + model = body.get("model") if isinstance(body, dict) else None + return model if isinstance(model, str) and model else None + + +def _target_for_model( + targets: list[CaptureTarget], model: str | None +) -> CaptureTarget | None: + if len(targets) == 1: + return targets[0] + if model is None: + return None + matches = [ + target + for target in targets + if target.model and _model_matches_target(model, target.model) + ] + return matches[0] if len(matches) == 1 else None + + +def _sort_exchange_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + def timestamp(record: dict[str, Any]) -> str: + request = record.get("request") + value = request.get("timestamp") if isinstance(request, dict) else None + return value if isinstance(value, str) else "" + + return sorted(records, key=timestamp) + + +def _record_metadata(record: dict[str, Any]) -> dict[str, Any]: + metadata = record.get("metadata") + return metadata if isinstance(metadata, dict) else {} + + +def _record_metadata_bool(record: dict[str, Any], key: str) -> bool: + return _record_metadata(record).get(key) is True + + +def _aggregate_source(records: list[dict[str, Any]]) -> CaptureSource: + values = { + CaptureSource(str(_record_metadata(record).get("capture_source"))) + for record in records + } + return next(iter(values)) if len(values) == 1 else CaptureSource.MIXED + + +def _aggregate_fidelity(records: list[dict[str, Any]]) -> CaptureFidelity: + values = { + CaptureFidelity(str(_record_metadata(record).get("capture_fidelity"))) + for record in records + } + return next(iter(values)) if len(values) == 1 else CaptureFidelity.MIXED + + +def _aggregate_auth_mode( + records: list[dict[str, Any]], + *, + targets: list[CaptureTarget], + fallback_auth: AuthMode, +) -> AuthMode: + values = { + AuthMode(str(_record_metadata(record).get("auth_mode"))) for record in records + } + values.update(target.auth_mode for target in targets) + if not values: + return fallback_auth + return next(iter(values)) if len(values) == 1 else AuthMode.MIXED + + +def _role_captures( + records: list[dict[str, Any]], *, targets: list[CaptureTarget] +) -> list[LLMRoleCapture]: + grouped: dict[ + tuple[str, str | None, AuthMode, CaptureSource, CaptureFidelity], + list[dict[str, Any]], + ] = {} + for record in records: + metadata = _record_metadata(record) + agent = str(metadata.get("agent") or "unknown") + model_value = metadata.get("model") or _record_model(record) + model = str(model_value) if model_value is not None else None + key = ( + agent, + model, + AuthMode(str(metadata.get("auth_mode"))), + CaptureSource(str(metadata.get("capture_source"))), + CaptureFidelity(str(metadata.get("capture_fidelity"))), + ) + grouped.setdefault(key, []).append(record) + captures = [ + LLMRoleCapture( + agent=agent, + model=model, + auth_mode=auth_mode, + capture_source=source, + capture_fidelity=fidelity, + exchange_count=len(group), + request_complete=all( + _record_metadata_bool(record, "request_complete") for record in group + ), + response_complete=all( + _record_metadata_bool(record, "response_complete") for record in group + ), + ) + for (agent, model, auth_mode, source, fidelity), group in sorted( + grouped.items(), key=lambda item: (item[0][0], item[0][1] or "") + ) + ] + captured_roles = { + (capture.agent, capture.model, capture.auth_mode) for capture in captures + } + captures.extend( + LLMRoleCapture( + agent=target.agent, + model=target.model, + auth_mode=target.auth_mode, + capture_source=CaptureSource.NONE, + capture_fidelity=CaptureFidelity.NONE, + exchange_count=0, + request_complete=False, + response_complete=False, + ) + for target in targets + if (target.agent, target.model, target.auth_mode) not in captured_roles + ) + return sorted(captures, key=lambda capture: (capture.agent, capture.model or "")) diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index 4cd39bca8..1d37bc6f0 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -176,7 +176,7 @@ def parse_claude_sessions( ) -> NativeParseResult | None: """Reconstruct model turns from Claude Code's native session transcript.""" - records = _read_jsonl_tree(sessions_dir) + records = _read_jsonl_tree(sessions_dir, started_at=started_at) if not records: return None messages: list[dict[str, Any]] = [] @@ -278,7 +278,7 @@ def parse_codex_sessions( ) -> NativeParseResult | None: """Reconstruct Responses-style calls from Codex native session records.""" - records = _read_jsonl_tree(sessions_dir) + records = _read_jsonl_tree(sessions_dir, started_at=started_at) if not records: return None history: list[dict[str, Any]] = [] @@ -488,17 +488,29 @@ def _exchange( ) -def _read_jsonl_tree(root: Path) -> list[dict[str, Any]]: +def _read_jsonl_tree( + root: Path, *, started_at: datetime | None = None +) -> list[dict[str, Any]]: if not root.is_dir(): return [] + boundary = started_at.timestamp() - 1.0 if started_at is not None else None records: list[dict[str, Any]] = [] for path in sorted(root.rglob("*.jsonl"), key=lambda item: item.stat().st_mtime): + if boundary is not None and path.stat().st_mtime < boundary: + continue for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): try: record = json.loads(line) except json.JSONDecodeError: continue if isinstance(record, dict): + timestamp = _optional_record_timestamp(record) + if ( + boundary is not None + and timestamp is not None + and timestamp.timestamp() < boundary + ): + continue records.append(record) return records @@ -718,14 +730,17 @@ def _normalize_codex_usage(usage: dict[str, Any]) -> dict[str, Any]: def _record_timestamp(record: dict[str, Any], default: datetime) -> datetime: + return _optional_record_timestamp(record) or default + + +def _optional_record_timestamp(record: dict[str, Any]) -> datetime | None: value = record.get("timestamp") if not isinstance(value, str): - return default + return None try: - parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + return datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError: - return default - return parsed + return None def _unix_nanos_timestamp(value: Any) -> datetime: diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index c463f897d..edf3a4ba1 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -32,7 +32,8 @@ validate_prime_sft_row, ) from benchflow.trajectories.llm_capture_manifest import ( - LLM_TRAJECTORY_MANIFEST_FILENAME, + capture_manifest_allows_training, + read_llm_trajectory_manifest, ) from benchflow.trajectories.types import redact_trajectory_obj from benchflow.usage_tracking import USAGE_SOURCE_AGENT_NATIVE_ACP @@ -170,8 +171,8 @@ def _llm_steps_from_trajectory( exchanges = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError as exc: return [], [], f"Invalid LLM trajectory JSONL: {exc}" - capture_manifest = _load_llm_capture_manifest(rollout_dir) - if capture_manifest is not None and not _capture_manifest_allows_training( + capture_manifest = read_llm_trajectory_manifest(rollout_dir) + if capture_manifest is not None and not capture_manifest_allows_training( capture_manifest ): return [], [], None @@ -275,29 +276,6 @@ def _response_is_training_success(response: Any) -> bool: return not bool(body.get("incomplete_details")) -def _load_llm_capture_manifest(rollout_dir: Path) -> dict[str, Any] | None: - path = rollout_dir / "trajectory" / LLM_TRAJECTORY_MANIFEST_FILENAME - if not path.exists(): - # Backward compatibility for artifacts created before the sidecar contract. - return None - try: - manifest = json.loads(path.read_text()) - except (OSError, json.JSONDecodeError): - return {"status": "capture_failed", "capture_fidelity": "none"} - if not isinstance(manifest, dict): - return {"status": "capture_failed", "capture_fidelity": "none"} - return manifest - - -def _capture_manifest_allows_training(manifest: dict[str, Any]) -> bool: - return bool( - manifest.get("status") == "complete" - and manifest.get("capture_fidelity") == "provider_wire" - and manifest.get("request_complete") is True - and manifest.get("response_complete") is True - ) - - def _training_success_exchange_indices( exchanges: list[dict[str, Any]], ) -> set[int]: @@ -504,11 +482,11 @@ def build_rollout_results_record( export_error=effective_export_error, ) terminal_health_error = bool(error or verifier_error or partial_trajectory) - capture_manifest = _load_llm_capture_manifest(rollout_path) + capture_manifest = read_llm_trajectory_manifest(rollout_path) audit_only_capture = bool( capture_manifest is not None and int(capture_manifest.get("exchange_count") or 0) > 0 - and not _capture_manifest_allows_training(capture_manifest) + and not capture_manifest_allows_training(capture_manifest) ) native_subscription_without_llm = bool( ( diff --git a/tests/test_eval_artifact_cli.py b/tests/test_eval_artifact_cli.py index 11676b16b..641c874d2 100644 --- a/tests/test_eval_artifact_cli.py +++ b/tests/test_eval_artifact_cli.py @@ -81,6 +81,35 @@ def _write_rollout(rollout_dir: Path, task_name: str = "task-a") -> None: _write_llm_trajectory(rollout_dir) +def test_health_rejects_empty_terminal_llm_capture_manifest(tmp_path: Path) -> None: + """Guards PR #1057 against validating no-call artifacts as healthy traces.""" + + job = tmp_path / "job" + rollout = job / "task-a__abc" + _write_rollout(rollout) + trajectory = rollout / "trajectory" + (trajectory / "llm_trajectory.jsonl").write_text("") + (trajectory / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "no_model_call", + "capture_fidelity": "none", + "exchange_count": 0, + "request_complete": False, + "response_complete": False, + } + ) + ) + + health = build_health_summary(job) + + assert health["missing_llm_trajectory"] == 0 + assert health["malformed_llm_trajectory"] == 1 + assert health["rows"][0]["has_llm_trajectory"] is True + assert health["rows"][0]["valid_llm_trajectory"] is False + assert health["rows"][0]["llm_trajectory_rows"] == 0 + + def test_eval_run_writes_manifest_health_and_canonical_artifacts( tmp_path: Path, monkeypatch ) -> None: diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index 550d5ee70..da91689c1 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -15,6 +15,9 @@ normalize_prime_sft_exchange, validate_prime_sft_jsonl, ) +from benchflow.trajectories.export_trl_sft import ( + convert_benchflow_rollouts_to_trl_sft_rows, +) def _write_rollout( @@ -189,7 +192,7 @@ def test_anthropic_tool_use_content_preserved_as_tool_calls(tmp_path: Path) -> N def test_native_session_exchange_is_not_exported_for_training(tmp_path: Path) -> None: - """Guards this PR's fail-closed export boundary for reconstructed capture.""" + """Guards PR #1057's fail-closed export boundary for reconstructed capture.""" exchange = _anthropic_exchange() exchange["metadata"] = { @@ -206,6 +209,42 @@ def test_native_session_exchange_is_not_exported_for_training(tmp_path: Path) -> assert stats.skipped_invalid == 0 +def test_partial_manifest_blocks_prime_and_trl_training_exports( + tmp_path: Path, +) -> None: + """Guards PR #1057 against exporting individually complete partial capture.""" + + exchange = _exchange(final=True) + exchange["metadata"] = { + "capture_fidelity": "provider_wire", + "request_complete": True, + "response_complete": True, + } + rollout = tmp_path / "job" / "rollout-1" + _write_rollout(rollout, exchanges=[exchange]) + (rollout / "trajectory" / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "partial", + "capture_fidelity": "provider_wire", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + } + ) + ) + + prime_rows, prime_stats = convert_benchflow_rollouts_to_prime_sft_rows( + tmp_path / "job" + ) + trl_rows, trl_stats = convert_benchflow_rollouts_to_trl_sft_rows(tmp_path / "job") + + assert prime_rows == [] + assert prime_stats.skipped_insufficient_capture_fidelity == 1 + assert trl_rows == [] + assert trl_stats.skipped_insufficient_capture_fidelity == 1 + + def test_skipped_provider_error_counts_rollouts_not_exchanges(tmp_path: Path) -> None: """Guards #828 greptile P1: an all-failed rollout counts as ONE rollout skip, with the exchange count surfaced separately.""" diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py new file mode 100644 index 000000000..6d2ac1855 --- /dev/null +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -0,0 +1,99 @@ +"""Training boundaries for the uniform LLM capture manifest.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from benchflow.trajectories.results import build_rollout_results_record + + +def _build_results_row(rollout_dir: Path, *, agent_result: dict) -> dict: + return build_rollout_results_record( + rollout_dir, + task_name="task", + rollout_name="rollout", + agent="claude-agent-acp", + agent_name="Claude Code", + model="claude-opus-4-1", + n_tool_calls=0, + prompts=["hello"], + trajectory=[], + partial_trajectory=False, + rewards={"reward": 1.0}, + error=None, + verifier_error=None, + agent_result=agent_result, + ) + + +def _write_exchange(trajectory_dir: Path, *, fidelity: str) -> None: + (trajectory_dir / "llm_trajectory.jsonl").write_text( + json.dumps( + { + "request": { + "body": {"messages": [{"role": "user", "content": "hello"}]} + }, + "response": { + "status_code": 200, + "body": { + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + }, + }, + "metadata": { + "capture_fidelity": fidelity, + "request_complete": fidelity == "provider_wire", + "response_complete": True, + }, + } + ) + + "\n" + ) + + +def test_agent_session_capture_is_audit_only_not_training_ready(tmp_path: Path) -> None: + """Guards PR #1057 against silently training on reconstructed OAuth payloads.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="agent_session") + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "partial", + "capture_fidelity": "agent_session", + "auth_mode": "oauth_subscription", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + } + ) + ) + + row = _build_results_row( + tmp_path, + agent_result={"usage_source": "agent_native_acp", "total_tokens": 2}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is True + assert row["error"] is None + + +def test_corrupt_capture_manifest_fails_closed_for_training(tmp_path: Path) -> None: + """Guards PR #1057 against treating a corrupt new sidecar as a legacy artifact.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="provider_wire") + (trajectory_dir / "llm_trajectory.manifest.json").write_text("{broken") + + row = _build_results_row(tmp_path, agent_result={"total_tokens": 2}) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == ( + "missing_healthy_structured_llm_trajectory" + ) + assert row["is_completed"] is False diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index aa64ef265..f243e3ee4 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -12,9 +12,12 @@ from benchflow.trajectories.llm_capture import ( LLMTrajectoryCapture, + _CaptureTarget, _download_optional_dir, + _NativeCaptureBundle, ) from benchflow.trajectories.llm_capture_manifest import ( + AuthMode, CaptureFidelity, CaptureSource, CaptureStatus, @@ -25,7 +28,6 @@ parse_codex_sessions, project_acp_trajectory, ) -from benchflow.trajectories.results import build_rollout_results_record STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) @@ -51,7 +53,7 @@ def _otel_record(name: str, timestamp_ns: int, **attributes: str) -> dict: def test_claude_otel_raw_bodies_become_provider_wire_exchanges(tmp_path: Path) -> None: - """Guards this PR's exact Claude OAuth raw-body capture contract.""" + """Guards PR #1057's exact Claude OAuth raw-body capture contract.""" capture = tmp_path / "capture" raw = capture / "raw" @@ -128,7 +130,7 @@ def test_claude_otel_raw_bodies_become_provider_wire_exchanges(tmp_path: Path) - def test_claude_concurrent_raw_pairing_fails_closed_for_training( tmp_path: Path, ) -> None: - """Guards this PR against claiming FIFO pairing for concurrent Claude calls.""" + """Guards PR #1057 against claiming FIFO pairing for concurrent Claude calls.""" capture = tmp_path / "capture" raw = capture / "raw" @@ -198,7 +200,7 @@ def test_claude_concurrent_raw_pairing_fails_closed_for_training( def test_claude_first_raw_request_survives_missing_otel_body_event( tmp_path: Path, ) -> None: - """Guards this PR against Claude omitting its first request-body OTLP event.""" + """Guards PR #1057 against Claude omitting its first request-body OTLP event.""" capture = tmp_path / "capture" raw = capture / "raw" @@ -262,7 +264,7 @@ def test_claude_first_raw_request_survives_missing_otel_body_event( def test_claude_session_fallback_is_truthfully_lower_fidelity(tmp_path: Path) -> None: - """Guards this PR's Claude OAuth fallback when raw OTel is unavailable.""" + """Guards PR #1057's Claude OAuth fallback when raw OTel is unavailable.""" session = tmp_path / "claude" / "session.jsonl" _write_jsonl( @@ -313,8 +315,82 @@ def test_claude_session_fallback_is_truthfully_lower_fidelity(tmp_path: Path) -> assert "tool_definitions" in result.missing_fields +def test_native_session_fallback_filters_records_before_rollout_start( + tmp_path: Path, +) -> None: + """Guards PR #1057 against importing sessions from a reused sandbox.""" + + sessions = tmp_path / "claude" + previous = sessions / "previous.jsonl" + current = sessions / "current.jsonl" + _write_jsonl( + previous, + [ + { + "type": "assistant", + "requestId": "old-request", + "timestamp": "2026-08-28T11:59:00Z", + "message": { + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "old answer"}], + }, + } + ], + ) + _write_jsonl( + current, + [ + { + "type": "assistant", + "requestId": "old-appended-request", + "timestamp": "2026-08-28T11:59:30Z", + "message": { + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "old appended answer"}], + }, + }, + { + "type": "user", + "timestamp": "2026-08-28T12:00:01Z", + "message": {"role": "user", "content": "current prompt"}, + }, + { + "type": "assistant", + "requestId": "current-request", + "timestamp": "2026-08-28T12:00:02Z", + "message": { + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "current answer"}], + }, + }, + ], + ) + previous_time = int((STARTED_AT.timestamp() - 10) * 1_000_000_000) + current_time = int((STARTED_AT.timestamp() + 10) * 1_000_000_000) + os.utime(previous, ns=(previous_time, previous_time)) + os.utime(current, ns=(current_time, current_time)) + + result = parse_claude_sessions( + sessions, + agent="claude-agent-acp", + session_id="rollout-1", + started_at=STARTED_AT, + ) + + assert result is not None + assert len(result.trajectory.exchanges) == 1 + exchange = result.trajectory.exchanges[0] + assert exchange.response.body["content"][0]["text"] == "current answer" + assert exchange.request.body["messages"] == [ + {"role": "user", "content": "current prompt"} + ] + + def test_codex_oauth_session_splits_calls_and_preserves_usage(tmp_path: Path) -> None: - """Guards this PR's Codex OAuth native-session trajectory reconstruction.""" + """Guards PR #1057's Codex OAuth native-session trajectory reconstruction.""" session = tmp_path / "codex" / "session.jsonl" _write_jsonl( @@ -418,7 +494,7 @@ def test_codex_oauth_session_splits_calls_and_preserves_usage(tmp_path: Path) -> async def test_capture_always_emits_empty_jsonl_and_terminal_manifest( tmp_path: Path, ) -> None: - """Guards this PR's always-present artifact invariant for zero-call runs.""" + """Guards PR #1057's always-present artifact invariant for zero-call runs.""" capture = LLMTrajectoryCapture( tmp_path, @@ -456,7 +532,7 @@ def test_codex_native_auth_file_mode_is_not_assumed_to_be_oauth( auth_json: str, expected: str, ) -> None: - """Guards this PR's Codex auth provenance for API-key auth.json files.""" + """Guards PR #1057's Codex auth provenance for API-key auth.json files.""" capture = LLMTrajectoryCapture( tmp_path, @@ -475,7 +551,7 @@ def test_codex_native_auth_file_mode_is_not_assumed_to_be_oauth( @pytest.mark.asyncio async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> None: - """Guards this PR's uniform metadata for existing API-key proxy capture.""" + """Guards PR #1057's uniform metadata for existing API-key proxy capture.""" capture = LLMTrajectoryCapture( tmp_path, @@ -511,7 +587,7 @@ async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> async def test_provider_capture_early_return_cleans_native_raw_bodies( tmp_path: Path, ) -> None: - """Guards this PR against raw-body leakage from mixed-role rollouts.""" + """Guards PR #1057 against raw-body leakage from mixed-role rollouts.""" commands: list[str] = [] @@ -527,7 +603,14 @@ async def exec(self, command, **_kwargs): session_id="rollout-1", started_at=STARTED_AT, ) - capture._native_agents["claude-agent-acp"] = "claude-sonnet-4-6" + target = _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + ) + capture._targets[(target.agent, target.model, target.credential_home)] = target capture._capture_root_prepared = True capture.trajectory_path.write_text( json.dumps( @@ -545,11 +628,178 @@ async def exec(self, command, **_kwargs): assert any("-depth -delete" in command for command in commands) +@pytest.mark.asyncio +async def test_mixed_auth_rollout_merges_provider_and_native_exchanges( + tmp_path: Path, +) -> None: + """Guards PR #1057 against dropping native rows from mixed-auth scenes.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-api", + session_id="rollout-1", + started_at=STARTED_AT, + ) + api_target = _CaptureTarget( + agent="codex-acp", + model="gpt-api", + credential_home="/home/agent", + auth_mode=AuthMode.API_KEY, + native=False, + ) + native_target = _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + ) + for target in (api_target, native_target): + capture._targets[(target.agent, target.model, target.credential_home)] = target + capture._refresh_manifest_auth_mode() + capture.trajectory_path.write_text( + json.dumps( + { + "request": { + "timestamp": "2026-08-28T12:00:01Z", + "body": {"model": "gpt-api", "input": "hello"}, + }, + "response": { + "timestamp": "2026-08-28T12:00:02Z", + "status_code": 200, + "body": {"output": []}, + }, + } + ) + + "\n" + ) + session = tmp_path / "native-session" / "session.jsonl" + _write_jsonl( + session, + [ + { + "type": "user", + "timestamp": "2026-08-28T12:00:03Z", + "message": {"role": "user", "content": "native prompt"}, + }, + { + "type": "assistant", + "requestId": "req-native", + "timestamp": "2026-08-28T12:00:04Z", + "message": { + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "native answer"}], + }, + }, + ], + ) + native_result = parse_claude_sessions( + session.parent, + agent=native_target.agent, + session_id="rollout-1", + started_at=STARTED_AT, + ) + assert native_result is not None + + async def collect_native(_env): + return [_NativeCaptureBundle((native_target,), native_result)] + + capture._collect_native_results = collect_native + + await capture.finalize(object(), acp_events=[], model_call_seen=True) + + rows = [ + json.loads(line) for line in capture.trajectory_path.read_text().splitlines() + ] + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert len(rows) == 2 + assert {row["metadata"]["agent"] for row in rows} == { + "codex-acp", + "claude-agent-acp", + } + assert manifest["status"] == "partial" + assert manifest["capture_source"] == "mixed" + assert manifest["capture_fidelity"] == "mixed" + assert manifest["auth_mode"] == "mixed" + assert manifest["exchange_count"] == 2 + assert {role["agent"] for role in manifest["role_captures"]} == { + "codex-acp", + "claude-agent-acp", + } + + +@pytest.mark.asyncio +async def test_mixed_auth_rollout_marks_missing_native_role_partial( + tmp_path: Path, +) -> None: + """Guards PR #1057 against hiding a missing OAuth role behind API rows.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-api", + session_id="rollout-1", + started_at=STARTED_AT, + ) + api_target = _CaptureTarget( + agent="codex-acp", + model="gpt-api", + credential_home="/home/agent", + auth_mode=AuthMode.API_KEY, + native=False, + ) + native_target = _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + ) + for target in (api_target, native_target): + capture._targets[(target.agent, target.model, target.credential_home)] = target + capture._refresh_manifest_auth_mode() + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"model": "gpt-api", "input": "hello"}}, + "response": {"status_code": 200, "body": {"output": []}}, + } + ) + + "\n" + ) + + async def collect_native(_env): + return [] + + capture._collect_native_results = collect_native + + await capture.finalize(object(), acp_events=[], model_call_seen=True) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["status"] == "partial" + assert manifest["capture_source"] == "mixed" + assert manifest["capture_fidelity"] == "mixed" + assert manifest["auth_mode"] == "mixed" + assert manifest["request_complete"] is False + assert manifest["response_complete"] is False + assert any("no capture was attributable" in error for error in manifest["errors"]) + roles = {role["agent"]: role for role in manifest["role_captures"]} + assert roles["codex-acp"]["exchange_count"] == 1 + assert roles["claude-agent-acp"]["exchange_count"] == 0 + assert roles["claude-agent-acp"]["capture_fidelity"] == "none" + + @pytest.mark.asyncio async def test_claude_capture_setup_failure_degrades_without_aborting( tmp_path: Path, ) -> None: - """Guards this PR against observability setup breaking Claude OAuth runs.""" + """Guards PR #1057 against observability setup breaking Claude OAuth runs.""" commands: list[str] = [] @@ -593,7 +843,7 @@ async def upload_file(self, *_args, **_kwargs): async def test_optional_session_download_creates_docker_copy_parent( tmp_path: Path, ) -> None: - """Guards this PR's Docker native-session download path.""" + """Guards PR #1057's Docker native-session download path.""" class DockerLikeEnv: async def exec(self, *_args, **_kwargs): @@ -615,7 +865,7 @@ async def download_dir(self, _remote, local): def test_capture_failure_repairs_invalid_jsonl_and_redacts_manifest_error( tmp_path: Path, ) -> None: - """Guards this PR's valid-JSONL and secret-redaction failure invariant.""" + """Guards PR #1057's valid-JSONL and secret-redaction failure invariant.""" capture = LLMTrajectoryCapture( tmp_path, @@ -639,7 +889,7 @@ def test_capture_failure_repairs_invalid_jsonl_and_redacts_manifest_error( def test_acp_projection_retains_the_actual_auth_mode() -> None: - """Guards this PR against labeling API-key fallback rows as OAuth.""" + """Guards PR #1057 against labeling API-key fallback rows as OAuth.""" result = project_acp_trajectory( [{"type": "agent_message", "text": "finished"}], @@ -651,119 +901,3 @@ def test_acp_projection_retains_the_actual_auth_mode() -> None: assert result is not None assert result.trajectory.exchanges[0].metadata["auth_mode"] == "api_key" - - -def test_agent_session_capture_is_audit_only_not_training_ready(tmp_path: Path) -> None: - """Guards this PR against silently training on reconstructed OAuth payloads.""" - - trajectory_dir = tmp_path / "trajectory" - trajectory_dir.mkdir() - (trajectory_dir / "llm_trajectory.jsonl").write_text( - json.dumps( - { - "request": { - "body": {"messages": [{"role": "user", "content": "hello"}]} - }, - "response": { - "status_code": 200, - "body": { - "role": "assistant", - "content": [{"type": "text", "text": "hi"}], - }, - }, - "metadata": { - "capture_fidelity": "agent_session", - "request_complete": False, - "response_complete": True, - }, - } - ) - + "\n" - ) - (trajectory_dir / "llm_trajectory.manifest.json").write_text( - json.dumps( - { - "status": "partial", - "capture_fidelity": "agent_session", - "auth_mode": "oauth_subscription", - "exchange_count": 1, - "request_complete": False, - "response_complete": True, - } - ) - ) - - row = build_rollout_results_record( - tmp_path, - task_name="task", - rollout_name="rollout", - agent="claude-agent-acp", - agent_name="Claude Code", - model="claude-opus-4-1", - n_tool_calls=0, - prompts=["hello"], - trajectory=[], - partial_trajectory=False, - rewards={"reward": 1.0}, - error=None, - verifier_error=None, - agent_result={"usage_source": "agent_native_acp", "total_tokens": 2}, - ) - - assert row["info"]["training_ready"] is False - assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" - assert row["is_completed"] is True - assert row["error"] is None - - -def test_corrupt_capture_manifest_fails_closed_for_training(tmp_path: Path) -> None: - """Guards this PR against treating a corrupt new sidecar as a legacy artifact.""" - - trajectory_dir = tmp_path / "trajectory" - trajectory_dir.mkdir() - (trajectory_dir / "llm_trajectory.jsonl").write_text( - json.dumps( - { - "request": { - "body": {"messages": [{"role": "user", "content": "hello"}]} - }, - "response": { - "status_code": 200, - "body": { - "role": "assistant", - "content": [{"type": "text", "text": "hi"}], - }, - }, - "metadata": { - "capture_fidelity": "provider_wire", - "request_complete": True, - "response_complete": True, - }, - } - ) - + "\n" - ) - (trajectory_dir / "llm_trajectory.manifest.json").write_text("{broken") - - row = build_rollout_results_record( - tmp_path, - task_name="task", - rollout_name="rollout", - agent="claude-agent-acp", - agent_name="Claude Code", - model="claude-opus-4-1", - n_tool_calls=0, - prompts=["hello"], - trajectory=[], - partial_trajectory=False, - rewards={"reward": 1.0}, - error=None, - verifier_error=None, - agent_result={"total_tokens": 2}, - ) - - assert row["info"]["training_ready"] is False - assert row["info"]["training_ready_reason"] == ( - "missing_healthy_structured_llm_trajectory" - ) - assert row["is_completed"] is False From f9053743319041657e640a88c5332fda7efeb6a5 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Fri, 28 Aug 2026 23:39:45 -0700 Subject: [PATCH 05/74] fix: isolate reused rollouts and scene roles --- docs/getting-started.md | 3 +- src/benchflow/rollout/__init__.py | 1 + src/benchflow/trajectories/llm_capture.py | 10 ++- .../trajectories/llm_capture_manifest.py | 4 +- .../trajectories/llm_capture_records.py | 31 ++++++-- tests/trajectories/test_native_llm_capture.py | 79 +++++++++++++++++++ 6 files changed, 116 insertions(+), 12 deletions(-) diff --git a/docs/getting-started.md b/docs/getting-started.md index 6990b5e3c..88459d081 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -228,7 +228,8 @@ the source of truth for interpreting the JSONL: The manifest status is `complete`, `partial`, `no_model_call`, or `capture_failed`. Mixed-role rollouts merge API-key and native-subscription exchanges into the same JSONL; `role_captures` records each prepared role's -agent, model, auth mode, source, fidelity, completeness, and exchange count. +scene role, agent, model, auth mode, source, fidelity, completeness, and exchange +count. Missing or ambiguously attributed roles make the rollout-level capture `partial`. Reconstructed `agent_session` rows remain useful for audit and viewer workflows, but trainer exports fail closed unless the manifest says the capture diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 0ca690a37..d71670f68 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -2399,6 +2399,7 @@ async def connect_as(self, role: Role) -> None: agent_env=agent_env, credential_home=cred_home, sandbox_user=cfg.sandbox_user, + role_name=role.name, ) self._agent_launch = agent_launch diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index f7295c575..2da1cc13e 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -106,7 +106,7 @@ def __init__( session_id=session_id, started_at=started_at, ) - self._targets: dict[tuple[str, str | None, str], _CaptureTarget] = {} + self._targets: dict[tuple[str, str, str | None, str], _CaptureTarget] = {} self._collector_started = False self._capture_root_prepared = False self._preparation_errors: list[str] = [] @@ -132,6 +132,7 @@ async def prepare_agent( agent_env: dict[str, str], credential_home: str, sandbox_user: str | None, + role_name: str | None = None, ) -> dict[str, str]: """Return the agent environment augmented for native capture.""" @@ -144,9 +145,11 @@ async def prepare_agent( credential_home=credential_home, auth_mode=auth_mode, native=native, + role=role_name or "primary", ) - self._targets[(agent, model, credential_home)] = target - self._refresh_manifest_auth_mode() + if role_name is not None: + self._targets[(role_name, agent, model, credential_home)] = target + self._refresh_manifest_auth_mode() if not native: write_llm_trajectory_manifest(self.rollout_dir, self.manifest) return prepared @@ -281,6 +284,7 @@ def _fallback_target(self, *, native: bool) -> _CaptureTarget: credential_home="", auth_mode=self.manifest.auth_mode, native=native, + role="primary", ) def record_failure(self, error: object, *, model_call_seen: bool) -> None: diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 525ccfa6e..8dc5d86f5 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -57,6 +57,7 @@ class AuthMode(StrEnum): class LLMRoleCapture(BaseModel): """Per prepared role provenance for mixed-auth/mixed-agent rollouts.""" + role: str = "agent" agent: str model: str | None = None auth_mode: AuthMode @@ -102,8 +103,7 @@ def initialize_llm_trajectory_artifacts( trajectory_dir = rollout_dir / "trajectory" trajectory_dir.mkdir(parents=True, exist_ok=True) trajectory_path = trajectory_dir / LLM_TRAJECTORY_FILENAME - if not trajectory_path.exists(): - _atomic_write_text(trajectory_path, "") + _atomic_write_text(trajectory_path, "") manifest = LLMTrajectoryManifest( agent=agent, model=model, diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 02f5b6006..c925d2621 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -29,6 +29,7 @@ class CaptureTarget: credential_home: str auth_mode: AuthMode native: bool + role: str = "agent" @dataclass(frozen=True) @@ -100,6 +101,11 @@ def load_provider_wire_records( if target is not None else (fallback_agent if not targets else "mixed") ), + "role": ( + target.role + if target is not None + else ("primary" if not targets else "mixed") + ), "model": ( target.model if target is not None @@ -259,6 +265,7 @@ def _native_bundle_records(bundle: NativeCaptureBundle) -> list[dict[str, Any]]: metadata.update( { "agent": target.agent, + "role": target.role, "model": target.model or _record_model(record), "auth_mode": target.auth_mode.value, "role_attribution_complete": True, @@ -268,6 +275,7 @@ def _native_bundle_records(bundle: NativeCaptureBundle) -> list[dict[str, Any]]: metadata.update( { "agent": "mixed", + "role": "mixed", "model": _record_model(record), "auth_mode": AuthMode.MIXED.value, "role_attribution_complete": False, @@ -282,6 +290,7 @@ def _role_candidates(targets: list[CaptureTarget]) -> list[dict[str, Any]]: return [ { "agent": target.agent, + "role": target.role, "model": target.model, "auth_mode": target.auth_mode.value, } @@ -300,6 +309,7 @@ def _captured_targets( for target in targets: if ( metadata.get("agent") == target.agent + and metadata.get("role") == target.role and metadata.get("auth_mode") == target.auth_mode.value and _model_matches_target(metadata.get("model"), target.model) ): @@ -396,15 +406,17 @@ def _role_captures( records: list[dict[str, Any]], *, targets: list[CaptureTarget] ) -> list[LLMRoleCapture]: grouped: dict[ - tuple[str, str | None, AuthMode, CaptureSource, CaptureFidelity], + tuple[str, str, str | None, AuthMode, CaptureSource, CaptureFidelity], list[dict[str, Any]], ] = {} for record in records: metadata = _record_metadata(record) + role = str(metadata.get("role") or "unknown") agent = str(metadata.get("agent") or "unknown") model_value = metadata.get("model") or _record_model(record) model = str(model_value) if model_value is not None else None key = ( + role, agent, model, AuthMode(str(metadata.get("auth_mode"))), @@ -414,6 +426,7 @@ def _role_captures( grouped.setdefault(key, []).append(record) captures = [ LLMRoleCapture( + role=role, agent=agent, model=model, auth_mode=auth_mode, @@ -427,15 +440,17 @@ def _role_captures( _record_metadata_bool(record, "response_complete") for record in group ), ) - for (agent, model, auth_mode, source, fidelity), group in sorted( - grouped.items(), key=lambda item: (item[0][0], item[0][1] or "") + for (role, agent, model, auth_mode, source, fidelity), group in sorted( + grouped.items(), key=lambda item: (item[0][0], item[0][1], item[0][2] or "") ) ] captured_roles = { - (capture.agent, capture.model, capture.auth_mode) for capture in captures + (capture.role, capture.agent, capture.model, capture.auth_mode) + for capture in captures } captures.extend( LLMRoleCapture( + role=target.role, agent=target.agent, model=target.model, auth_mode=target.auth_mode, @@ -446,6 +461,10 @@ def _role_captures( response_complete=False, ) for target in targets - if (target.agent, target.model, target.auth_mode) not in captured_roles + if (target.role, target.agent, target.model, target.auth_mode) + not in captured_roles + ) + return sorted( + captures, + key=lambda capture: (capture.role, capture.agent, capture.model or ""), ) - return sorted(captures, key=lambda capture: (capture.agent, capture.model or "")) diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index f243e3ee4..c6886f2e4 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -517,6 +517,34 @@ async def test_capture_always_emits_empty_jsonl_and_terminal_manifest( assert manifest["auth_mode"] == "api_key" +def test_capture_initialization_truncates_reused_rollout_jsonl(tmp_path: Path) -> None: + """Guards PR #1057 against stale rows when a rollout name is reused.""" + + first = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="first-session", + started_at=STARTED_AT, + ) + first.trajectory_path.write_text('{"stale":true}\n') + + second = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="second-session", + started_at=STARTED_AT, + ) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert second.trajectory_path.read_text() == "" + assert manifest["status"] == "pending" + assert manifest["session_id"] == "second-session" + + @pytest.mark.parametrize( ("auth_json", "expected"), [ @@ -795,6 +823,57 @@ async def collect_native(_env): assert roles["claude-agent-acp"]["capture_fidelity"] == "none" +@pytest.mark.asyncio +async def test_same_agent_model_roles_remain_independently_auditable( + tmp_path: Path, +) -> None: + """Guards PR #1057 against collapsing same-model scene roles.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-api", + session_id="rollout-1", + started_at=STARTED_AT, + ) + for role_name in ("coder", "reviewer"): + await capture.prepare_agent( + None, + agent="codex-acp", + model="gpt-api", + agent_env={"OPENAI_API_KEY": "test-key"}, + credential_home="/home/agent", + sandbox_user="agent", + role_name=role_name, + ) + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"model": "gpt-api", "input": "hello"}}, + "response": {"status_code": 200, "body": {"output": []}}, + } + ) + + "\n" + ) + + await capture.finalize(None, acp_events=[], model_call_seen=True) + + row = json.loads(capture.trajectory_path.read_text()) + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert row["metadata"]["role"] == "mixed" + assert {item["role"] for item in row["metadata"]["role_candidates"]} == { + "coder", + "reviewer", + } + assert manifest["status"] == "partial" + roles = {role["role"]: role for role in manifest["role_captures"]} + assert roles["coder"]["exchange_count"] == 0 + assert roles["reviewer"]["exchange_count"] == 0 + assert roles["mixed"]["exchange_count"] == 1 + + @pytest.mark.asyncio async def test_claude_capture_setup_failure_degrades_without_aborting( tmp_path: Path, From 744d43c58ed7a80f64a8d92de288ab4aca23e76f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Fri, 28 Aug 2026 23:54:10 -0700 Subject: [PATCH 06/74] fix: scope native session collection --- src/benchflow/trajectories/llm_capture.py | 65 ++++++-- .../trajectories/native_capture_parsers.py | 114 +++++++++----- tests/trajectories/test_native_llm_capture.py | 24 --- .../test_native_session_boundaries.py | 145 ++++++++++++++++++ 4 files changed, 271 insertions(+), 77 deletions(-) create mode 100644 tests/trajectories/test_native_session_boundaries.py diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 2da1cc13e..1c63a6258 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import hashlib import json import logging @@ -10,7 +11,7 @@ import tempfile from contextlib import suppress from datetime import datetime -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Any from benchflow.agents.env import uses_native_subscription_auth @@ -147,9 +148,13 @@ async def prepare_agent( native=native, role=role_name or "primary", ) - if role_name is not None: + primary_key = ("primary", agent, model, credential_home) + if role_name is None: + self._targets[primary_key] = target + else: + self._targets.pop(primary_key, None) self._targets[(role_name, agent, model, credential_home)] = target - self._refresh_manifest_auth_mode() + self._refresh_manifest_auth_mode() if not native: write_llm_trajectory_manifest(self.rollout_dir, self.manifest) return prepared @@ -437,8 +442,11 @@ async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: if ( home_claude_targets and not raw_claude_captured - and await _download_optional_dir( - env, f"{credential_home}/.claude/projects", claude_local + and await _download_recent_session_files( + env, + f"{credential_home}/.claude/projects", + claude_local, + started_at=self.started_at, ) ): target = home_claude_targets[0] @@ -455,8 +463,11 @@ async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: result=result, ) ) - if home_codex_targets and await _download_optional_dir( - env, f"{credential_home}/.codex/sessions", codex_local + if home_codex_targets and await _download_recent_session_files( + env, + f"{credential_home}/.codex/sessions", + codex_local, + started_at=self.started_at, ): target = home_codex_targets[0] result = parse_codex_sessions( @@ -516,16 +527,42 @@ def _finish_manifest( write_llm_trajectory_manifest(self.rollout_dir, self.manifest) -async def _download_optional_dir(env: Any, remote: str, local: Path) -> bool: - probe = await env.exec( - f"test -d {shlex.quote(remote)}", +async def _download_recent_session_files( + env: Any, + remote: str, + local: Path, + *, + started_at: datetime, +) -> bool: + boundary = started_at.timestamp() - 1.0 + remote_root = shlex.quote(remote) + result = await env.exec( + f"if test -d {remote_root}; then " + f"find {remote_root} -type f -name '*.jsonl' " + f"-newermt {shlex.quote(f'@{boundary}')} -printf '%P\\n' | head -n 1001; " + "fi", user="root", - timeout_sec=5, + timeout_sec=10, ) - if probe.return_code != 0: + if result.return_code != 0: + detail = (result.stderr or result.stdout or "session discovery failed")[:300] + raise RuntimeError(f"Native session discovery failed: {detail}") + relative_paths = [line for line in result.stdout.splitlines() if line] + if len(relative_paths) > 1000: + raise RuntimeError("Native session discovery exceeded the 1000-file limit") + if not relative_paths: return False - local.parent.mkdir(parents=True, exist_ok=True) - await env.download_dir(remote, local) + downloads: list[tuple[str, Path]] = [] + for value in relative_paths: + relative = PurePosixPath(value) + if relative.is_absolute() or ".." in relative.parts: + raise RuntimeError("Native session discovery returned an unsafe path") + destination = local.joinpath(*relative.parts) + destination.parent.mkdir(parents=True, exist_ok=True) + downloads.append((f"{remote}/{relative.as_posix()}", destination)) + await asyncio.gather( + *(env.download_file(source, destination) for source, destination in downloads) + ) return True diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index 1d37bc6f0..cee757bc9 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -176,9 +176,41 @@ def parse_claude_sessions( ) -> NativeParseResult | None: """Reconstruct model turns from Claude Code's native session transcript.""" - records = _read_jsonl_tree(sessions_dir, started_at=started_at) - if not records: + record_groups = _read_jsonl_files(sessions_dir, started_at=started_at) + exchanges = [ + exchange + for records in record_groups + for exchange in _parse_claude_session_records(records, started_at=started_at) + ] + if not exchanges: return None + exchanges.sort(key=lambda exchange: exchange.request.timestamp) + return NativeParseResult( + trajectory=Trajectory( + session_id=session_id, + agent_name=agent, + started_at=started_at, + finished_at=max(exchange.response.timestamp for exchange in exchanges), + exchanges=exchanges, + ), + source=CaptureSource.CLAUDE_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + request_complete=False, + response_complete=False, + missing_fields=[ + "system_prompt", + "tool_definitions", + "provider_response_envelope", + "headers", + ], + ) + + +def _parse_claude_session_records( + records: list[dict[str, Any]], *, started_at: datetime +) -> list[LLMExchange]: + """Parse one Claude session file without leaking history across sessions.""" + messages: list[dict[str, Any]] = [] exchanges: list[LLMExchange] = [] assistant_group: list[dict[str, Any]] = [] @@ -244,8 +276,34 @@ def flush_group() -> None: group_key = key assistant_group.append(record) flush_group() + return exchanges + + +def parse_codex_sessions( + sessions_dir: Path, + *, + agent: str, + session_id: str, + started_at: datetime, + configured_model: str | None, + auth_mode: str = "oauth_subscription", +) -> NativeParseResult | None: + """Reconstruct Responses-style calls from Codex native session records.""" + + record_groups = _read_jsonl_files(sessions_dir, started_at=started_at) + exchanges = [ + exchange + for records in record_groups + for exchange in _parse_codex_session_records( + records, + started_at=started_at, + configured_model=configured_model, + auth_mode=auth_mode, + ) + ] if not exchanges: return None + exchanges.sort(key=lambda exchange: exchange.request.timestamp) return NativeParseResult( trajectory=Trajectory( session_id=session_id, @@ -254,12 +312,12 @@ def flush_group() -> None: finished_at=max(exchange.response.timestamp for exchange in exchanges), exchanges=exchanges, ), - source=CaptureSource.CLAUDE_NATIVE_SESSION, + source=CaptureSource.CODEX_NATIVE_SESSION, fidelity=CaptureFidelity.AGENT_SESSION, request_complete=False, response_complete=False, missing_fields=[ - "system_prompt", + "instructions", "tool_definitions", "provider_response_envelope", "headers", @@ -267,20 +325,15 @@ def flush_group() -> None: ) -def parse_codex_sessions( - sessions_dir: Path, +def _parse_codex_session_records( + records: list[dict[str, Any]], *, - agent: str, - session_id: str, started_at: datetime, configured_model: str | None, - auth_mode: str = "oauth_subscription", -) -> NativeParseResult | None: - """Reconstruct Responses-style calls from Codex native session records.""" + auth_mode: str, +) -> list[LLMExchange]: + """Parse one Codex session file without leaking history across sessions.""" - records = _read_jsonl_tree(sessions_dir, started_at=started_at) - if not records: - return None history: list[dict[str, Any]] = [] output: list[dict[str, Any]] = [] exchanges: list[LLMExchange] = [] @@ -358,27 +411,7 @@ def flush_output(response_timestamp: datetime) -> None: output_started_at = timestamp output.append(payload) flush_output(last_timestamp) - if not exchanges: - return None - return NativeParseResult( - trajectory=Trajectory( - session_id=session_id, - agent_name=agent, - started_at=started_at, - finished_at=max(exchange.response.timestamp for exchange in exchanges), - exchanges=exchanges, - ), - source=CaptureSource.CODEX_NATIVE_SESSION, - fidelity=CaptureFidelity.AGENT_SESSION, - request_complete=False, - response_complete=False, - missing_fields=[ - "instructions", - "tool_definitions", - "provider_response_envelope", - "headers", - ], - ) + return exchanges def project_acp_trajectory( @@ -488,16 +521,17 @@ def _exchange( ) -def _read_jsonl_tree( +def _read_jsonl_files( root: Path, *, started_at: datetime | None = None -) -> list[dict[str, Any]]: +) -> list[list[dict[str, Any]]]: if not root.is_dir(): return [] boundary = started_at.timestamp() - 1.0 if started_at is not None else None - records: list[dict[str, Any]] = [] + record_groups: list[list[dict[str, Any]]] = [] for path in sorted(root.rglob("*.jsonl"), key=lambda item: item.stat().st_mtime): if boundary is not None and path.stat().st_mtime < boundary: continue + records: list[dict[str, Any]] = [] for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): try: record = json.loads(line) @@ -512,7 +546,9 @@ def _read_jsonl_tree( ): continue records.append(record) - return records + if records: + record_groups.append(records) + return record_groups def _load_body_files(root: Path) -> dict[Path, dict[str, Any]]: diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index c6886f2e4..826338aaf 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -13,7 +13,6 @@ from benchflow.trajectories.llm_capture import ( LLMTrajectoryCapture, _CaptureTarget, - _download_optional_dir, _NativeCaptureBundle, ) from benchflow.trajectories.llm_capture_manifest import ( @@ -918,29 +917,6 @@ async def upload_file(self, *_args, **_kwargs): assert any("/opt/benchflow/node/bin/node" in command for command in commands) -@pytest.mark.asyncio -async def test_optional_session_download_creates_docker_copy_parent( - tmp_path: Path, -) -> None: - """Guards PR #1057's Docker native-session download path.""" - - class DockerLikeEnv: - async def exec(self, *_args, **_kwargs): - return SimpleNamespace(return_code=0, stdout="", stderr="") - - async def download_dir(self, _remote, local): - assert Path(local).parent.is_dir() - Path(local).mkdir() - - destination = tmp_path / "missing-parent" / "sessions" - downloaded = await _download_optional_dir( - DockerLikeEnv(), "/home/agent/.claude/projects", destination - ) - - assert downloaded is True - assert destination.is_dir() - - def test_capture_failure_repairs_invalid_jsonl_and_redacts_manifest_error( tmp_path: Path, ) -> None: diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py new file mode 100644 index 000000000..72f57ee69 --- /dev/null +++ b/tests/trajectories/test_native_session_boundaries.py @@ -0,0 +1,145 @@ +"""Session-boundary regressions for native OAuth trajectory capture.""" + +from __future__ import annotations + +import json +from datetime import UTC, datetime +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from benchflow.trajectories.llm_capture import ( + LLMTrajectoryCapture, + _download_recent_session_files, +) +from benchflow.trajectories.native_capture_parsers import parse_codex_sessions + +STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) + + +def _write_jsonl(path: Path, records: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + + +def _codex_session(prompt: str, answer: str, second: int) -> list[dict]: + return [ + { + "type": "response_item", + "timestamp": f"2026-08-28T12:00:{second:02d}Z", + "payload": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": prompt}], + }, + }, + { + "type": "response_item", + "timestamp": f"2026-08-28T12:00:{second + 1:02d}Z", + "payload": { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": answer}], + }, + }, + { + "type": "event_msg", + "timestamp": f"2026-08-28T12:00:{second + 2:02d}Z", + "payload": { + "type": "token_count", + "info": {"input_tokens": 10, "output_tokens": 2}, + }, + }, + ] + + +@pytest.mark.asyncio +async def test_phase_api_registers_provisional_primary_oauth_target( + tmp_path: Path, +) -> None: + """Guards PR #1057's native capture when no scene role name is supplied.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + await capture.prepare_agent( + None, + agent="codex-acp", + model="gpt-5.6", + agent_env={ + "CODEX_AUTH_JSON": '{"auth_mode":"chatgpt","tokens":{"refresh_token":"test"}}' + }, + credential_home="/home/agent", + sandbox_user="agent", + ) + + targets = capture._native_targets() + assert len(targets) == 1 + assert targets[0].role == "primary" + + +@pytest.mark.asyncio +async def test_native_download_selects_recent_files_before_copying( + tmp_path: Path, +) -> None: + """Guards PR #1057 against copying an entire reused native-session tree.""" + + commands: list[str] = [] + downloads: list[tuple[str, Path]] = [] + + class DockerLikeEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + return SimpleNamespace( + return_code=0, + stdout="2026/08/28/session.jsonl\n", + stderr="", + ) + + async def download_file(self, remote, local): + destination = Path(local) + downloads.append((remote, destination)) + assert destination.parent.is_dir() + destination.write_text("{}\n") + + destination = tmp_path / "missing-parent" / "sessions" + downloaded = await _download_recent_session_files( + DockerLikeEnv(), + "/home/agent/.codex/sessions", + destination, + started_at=STARTED_AT, + ) + + assert downloaded is True + assert len(downloads) == 1 + assert downloads[0][0].endswith("/2026/08/28/session.jsonl") + assert downloads[0][1].is_file() + assert "-newermt" in commands[0] + + +def test_codex_session_files_do_not_share_request_history(tmp_path: Path) -> None: + """Guards PR #1057 against merging unrelated reconnect histories.""" + + sessions = tmp_path / "codex" + _write_jsonl(sessions / "one.jsonl", _codex_session("first prompt", "one", 1)) + _write_jsonl(sessions / "two.jsonl", _codex_session("second prompt", "two", 11)) + + result = parse_codex_sessions( + sessions, + agent="codex-acp", + session_id="rollout-1", + started_at=STARTED_AT, + configured_model="gpt-5.6", + ) + + assert result is not None + assert len(result.trajectory.exchanges) == 2 + first_input = result.trajectory.exchanges[0].request.body["input"] + second_input = result.trajectory.exchanges[1].request.body["input"] + assert [item["content"][0]["text"] for item in first_input] == ["first prompt"] + assert [item["content"][0]["text"] for item in second_input] == ["second prompt"] From ffe3dcab870e6a5c75f4879711cc9c8cc2fa94d2 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 00:08:25 -0700 Subject: [PATCH 07/74] fix: reset retried Claude capture state --- src/benchflow/trajectories/llm_capture.py | 41 ++++++++++++++---- .../test_native_session_boundaries.py | 42 +++++++++++++++++++ 2 files changed, 76 insertions(+), 7 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 1c63a6258..08b678ed6 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -318,8 +318,11 @@ def record_failure(self, error: object, *, model_call_seen: bool) -> None: async def _ensure_otel_sink(self, env: Any, *, sandbox_user: str | None) -> int: if self._collector_started: return await self._read_collector_port(env) + await self._stop_otel_sink(env) capture_owner = shlex.quote(sandbox_user or "root") setup = await env.exec( + f"find {self._remote_capture_root} -depth -mindepth 1 -delete " + "2>/dev/null || true\n" f"mkdir -p {self._remote_capture_root}/raw " f"{self._remote_capture_root}/otel\n" f"chown -R {capture_owner} {self._remote_capture_root}\n" @@ -379,6 +382,36 @@ async def _ensure_otel_sink(self, env: Any, *, sandbox_user: str | None) -> int: self._collector_started = True return _parse_port(result.stdout) + async def _stop_otel_sink(self, env: Any) -> None: + command = f""" +if ! test -s {self._remote_capture_root}/pid; then + exit 0 +fi +read -r old_pid < {self._remote_capture_root}/pid || true +case "$old_pid" in + ''|*[!0-9]*) exit 0 ;; +esac +old_command=$(ps -p "$old_pid" -o command= 2>/dev/null || true) +case "$old_command" in + *{self._remote_capture_root}/otel_sink.mjs*) ;; + *) exit 0 ;; +esac +kill -TERM "$old_pid" 2>/dev/null || true +for attempt in $(seq 1 20); do + if ! kill -0 "$old_pid" 2>/dev/null; then + exit 0 + fi + sleep 0.05 +done +echo "previous Claude telemetry collector did not stop" >&2 +exit 1 +""" + result = await env.exec(command, user="root", timeout_sec=5) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "collector did not stop")[:300] + raise RuntimeError(f"Claude OTel sink shutdown failed: {detail}") + self._collector_started = False + async def _read_collector_port(self, env: Any) -> int: result = await env.exec( f"cat {self._remote_capture_root}/port", @@ -391,13 +424,7 @@ async def _read_collector_port(self, env: Any) -> int: async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: if self._collector_started: - await env.exec( - f"if test -s {self._remote_capture_root}/pid; then " - f"kill -TERM $(cat {self._remote_capture_root}/pid) " - f"2>/dev/null || true; fi", - user="root", - timeout_sec=5, - ) + await self._stop_otel_sink(env) bundles: list[_NativeCaptureBundle] = [] native_targets = self._native_targets() claude_targets = tuple( diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 72f57ee69..89d79f31b 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -83,6 +83,48 @@ async def test_phase_api_registers_provisional_primary_oauth_target( assert targets[0].role == "primary" +@pytest.mark.asyncio +async def test_claude_capture_setup_clears_reused_raw_attempt( + tmp_path: Path, +) -> None: + """Guards PR #1057 against importing raw bodies from a crashed retry.""" + + commands: list[str] = [] + + class FailingCollectorEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + if "nohup" in command: + return SimpleNamespace(return_code=1, stdout="", stderr="no node") + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def upload_file(self, *_args, **_kwargs): + return None + + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="reused-rollout", + started_at=STARTED_AT, + ) + await capture.prepare_agent( + FailingCollectorEnv(), + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "test-token"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + + stop, setup = commands[:2] + assert 'case "$old_pid"' in stop + assert "otel_sink.mjs*" in stop + clear_at = setup.index("-mindepth 1 -delete") + create_at = setup.index("mkdir -p") + assert clear_at < create_at + + @pytest.mark.asyncio async def test_native_download_selects_recent_files_before_copying( tmp_path: Path, From effc13256a1c219122ec904cf0c131c02df093fc Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 00:26:14 -0700 Subject: [PATCH 08/74] fix: preserve provider rows across runtime switches --- src/benchflow/providers/litellm_runtime.py | 8 +++-- src/benchflow/rollout/__init__.py | 10 +++--- src/benchflow/trajectories/_llm_capture.py | 40 +++++++++++++++++---- tests/test_live_llm_trajectory.py | 15 ++++++++ tests/test_usage_litellm.py | 42 ++++++++++++++++++++++ 5 files changed, 100 insertions(+), 15 deletions(-) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index fc53f1186..8cd99a8ce 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -373,7 +373,9 @@ async def _stop_live_capture(self) -> None: with contextlib.suppress(Exception): await self._capture_live_records() - def _reconcile_live_capture(self) -> None: + def reconcile_live_capture(self) -> None: + """Publish this runtime's final snapshot without dropping prior runtimes.""" + writer = getattr(self, "_live_writer", None) if writer is not None: writer.reconcile(self.trajectory) @@ -435,7 +437,7 @@ async def stop(self) -> None: self.process.kill() await asyncio.to_thread(self.process.wait, 10) self._load_callback_log() - self._reconcile_live_capture() + self.reconcile_live_capture() with contextlib.suppress(Exception): shutil.rmtree(self.runtime_dir, ignore_errors=True) @@ -539,7 +541,7 @@ async def stop(self) -> None: timeout_sec=10, ) await self._load_callback_log() - self._reconcile_live_capture() + self.reconcile_live_capture() with contextlib.suppress(Exception): await self.sandbox.exec( f"rm -rf {shlex.quote(self.runtime_dir)}", timeout_sec=10 diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index d71670f68..97b3da2f2 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -218,7 +218,6 @@ _scrape_agent_trajectory, make_trajectory_sink, ) -from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter from benchflow.trajectories.llm_capture import ( LLMTrajectoryCapture, model_call_seen_from_evidence, @@ -2548,12 +2547,13 @@ def _write_llm_trajectory(self, usage_runtime: Any) -> None: """Persist captured provider HTTP exchanges as JSONL.""" if self._rollout_dir is None: return - trajectory = getattr(getattr(usage_runtime, "server", None), "trajectory", None) + server = getattr(usage_runtime, "server", None) + if server is None: + return + trajectory = server.trajectory if trajectory is None or not trajectory.exchanges: return - LiveLLMTrajectoryWriter( - self._rollout_dir / "trajectory" / "llm_trajectory.jsonl" - ).reconcile(trajectory) + server.reconcile_live_capture() def _reconcile_acp_tool_evidence(self, usage_runtime: Any) -> None: """Repair lossy ACP tool details from trusted provider capture.""" diff --git a/src/benchflow/trajectories/_llm_capture.py b/src/benchflow/trajectories/_llm_capture.py index 4cbc3981d..ab2ccbf01 100644 --- a/src/benchflow/trajectories/_llm_capture.py +++ b/src/benchflow/trajectories/_llm_capture.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os from pathlib import Path @@ -9,12 +10,14 @@ class LiveLLMTrajectoryWriter: - """Atomically publish redacted snapshots of completed LLM exchanges. + """Atomically publish cumulative redacted snapshots of LLM exchanges. LiteLLM's callback log is append-only, but the public BenchFlow artifact is - rewritten from a parsed snapshot. This keeps concurrent readers from ever - observing a partial JSON line and lets end-of-run reconciliation repair a - missed live poll without changing the trajectory schema. + rewritten from parsed snapshots. Rows already published by an earlier + provider runtime are retained when a scene switches agents or models. This + keeps concurrent readers from ever observing a partial JSON line and lets + end-of-run reconciliation repair a missed live poll without changing the + trajectory schema. """ def __init__(self, path: Path) -> None: @@ -23,17 +26,22 @@ def __init__(self, path: Path) -> None: self._tmp = self.path.with_suffix(self.path.suffix + ".tmp") self._tmp.unlink(missing_ok=True) self._last_payload: str | None = None + self._base_payload = self._valid_existing_payload() def write(self, trajectory: Trajectory | None) -> bool: """Publish *trajectory* when it is non-empty and changed.""" if trajectory is None or not trajectory.exchanges: return False - payload = trajectory.to_jsonl(redact_keys=True) - if payload == self._last_payload: + snapshot = trajectory.to_jsonl(redact_keys=True) + if snapshot == self._last_payload: + return False + payload = _join_jsonl(self._base_payload, snapshot) + if payload == self._valid_existing_payload(): + self._last_payload = snapshot return False self._tmp.write_text(payload) os.replace(self._tmp, self.path) - self._last_payload = payload + self._last_payload = snapshot return True def reconcile(self, trajectory: Trajectory | None) -> bool: @@ -44,3 +52,21 @@ def reconcile(self, trajectory: Trajectory | None) -> bool: captured the final exchange. """ return self.write(trajectory) + + def _valid_existing_payload(self) -> str: + if not self.path.is_file(): + return "" + try: + payload = self.path.read_text() + for line in payload.splitlines(): + if line.strip(): + json.loads(line) + except (OSError, json.JSONDecodeError): + return "" + return payload.strip() + + +def _join_jsonl(existing: str, snapshot: str) -> str: + """Join a prior-runtime prefix to the current runtime's full snapshot.""" + + return "\n".join(part.strip() for part in (existing, snapshot) if part.strip()) diff --git a/tests/test_live_llm_trajectory.py b/tests/test_live_llm_trajectory.py index 7ae3419cd..7b568c5c4 100644 --- a/tests/test_live_llm_trajectory.py +++ b/tests/test_live_llm_trajectory.py @@ -117,6 +117,21 @@ def test_writer_deduplicates_unchanged_snapshot_and_reconciles(tmp_path): assert len(path.read_text().splitlines()) == 2 +def test_writer_preserves_prior_runtime_rows_across_scene_switch(tmp_path): + """Guards PR #1057 against replacing an earlier API-key role's rows.""" + + path = tmp_path / "llm_trajectory.jsonl" + assert LiveLLMTrajectoryWriter(path).write(_trajectory(content="first")) is True + + second_writer = LiveLLMTrajectoryWriter(path) + second = _trajectory(content="first") + assert second_writer.write(second) is True + assert len(path.read_text().splitlines()) == 2 + + assert second_writer.reconcile(second) is False + assert len(path.read_text().splitlines()) == 2 + + def test_writer_does_not_create_empty_live_artifact(tmp_path): """Guards empty-artifact suppression from commit c86adfb.""" path = tmp_path / "llm_trajectory.jsonl" diff --git a/tests/test_usage_litellm.py b/tests/test_usage_litellm.py index b6da441fa..48a377972 100644 --- a/tests/test_usage_litellm.py +++ b/tests/test_usage_litellm.py @@ -12,6 +12,7 @@ from benchflow.providers.litellm_logging import extract_usage_from_trajectory from benchflow.providers.runtime import ProviderRuntime, extract_usage +from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter from benchflow.trajectories.types import ( LLMExchange, LLMRequest, @@ -272,6 +273,42 @@ def test_extract_usage_reads_litellm_runtime_trajectory(): assert usage["n_output_tokens"] == 2 +def test_rollout_final_reconcile_preserves_prior_provider_runtime(tmp_path): + """Guards PR #1057's cumulative writer across runtime-switch cleanup.""" + + from benchflow.rollout import Rollout + + path = tmp_path / "trajectory" / "llm_trajectory.jsonl" + first = _trajectory({"choices": [{"message": {"content": "same"}}]}) + assert LiveLLMTrajectoryWriter(path).write(first) is True + + class FakeLiveServer: + def __init__(self): + self.trajectory = _trajectory( + {"choices": [{"message": {"content": "same"}}]} + ) + self.writer = LiveLLMTrajectoryWriter(path) + self.reconciled = 0 + + def reconcile_live_capture(self): + self.reconciled += 1 + self.writer.reconcile(self.trajectory) + + server = FakeLiveServer() + runtime = ProviderRuntime( + kind="litellm", + agent_base_url="http://127.0.0.1:4000", + server=server, + ) + rollout = Rollout.__new__(Rollout) + rollout._rollout_dir = tmp_path + + rollout._write_llm_trajectory(runtime) + + assert server.reconciled == 1 + assert len(path.read_text().splitlines()) == 2 + + @pytest.mark.asyncio async def test_rollout_cleanup_extracts_usage_and_writes_llm_trajectory(tmp_path): from benchflow.rollout import Rollout, RolloutConfig @@ -288,6 +325,11 @@ def __init__(self): async def stop(self): return None + def reconcile_live_capture(self): + LiveLLMTrajectoryWriter( + tmp_path / "trajectory" / "llm_trajectory.jsonl" + ).reconcile(self.trajectory) + server = FakeServer() rollout = Rollout.__new__(Rollout) rollout._config = RolloutConfig(task_path=tmp_path / "task") From a324ccb5b8bacdf86d1d30add0b9d4b22774de56 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 00:42:04 -0700 Subject: [PATCH 09/74] fix: fail closed on incomplete capture artifacts --- src/benchflow/eval_artifacts.py | 15 +++----- .../trajectories/export_prime_sft.py | 7 ++-- src/benchflow/trajectories/export_trl_sft.py | 6 ++-- src/benchflow/trajectories/llm_capture.py | 4 ++- .../trajectories/llm_capture_manifest.py | 12 +++++-- src/benchflow/trajectories/results.py | 35 ++++++++++++------ tests/trajectories/test_export_prime_sft.py | 36 +++++++++++++++++++ .../test_llm_capture_training_contract.py | 28 +++++++++++++++ tests/trajectories/test_native_llm_capture.py | 4 +-- 9 files changed, 114 insertions(+), 33 deletions(-) diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index 188320d23..c4205e2d1 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -207,17 +207,10 @@ def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, int]: except PrimeSftTrajectoryJsonlError: return True, False, 0 manifest = read_llm_trajectory_manifest(rollout_dir) - if manifest is not None: - try: - expected_rows = int(manifest.get("exchange_count") or 0) - except (TypeError, ValueError): - return True, False, len(rows) - if ( - not capture_manifest_allows_training(manifest) - or not rows - or expected_rows != len(rows) - ): - return True, False, len(rows) + if manifest is not None and not capture_manifest_allows_training( + manifest, exchange_count=len(rows) + ): + return True, False, len(rows) return True, True, len(rows) diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 8645c40cb..797efb00f 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -1219,15 +1219,14 @@ def convert_benchflow_rollouts_to_prime_sft_rows( stats.skipped_reward += 1 continue + trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" + exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) capture_manifest = read_llm_trajectory_manifest(rollout_dir) if capture_manifest is not None and not capture_manifest_allows_training( - capture_manifest + capture_manifest, exchange_count=len(exchanges) ): stats.skipped_insufficient_capture_fidelity += 1 continue - - trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" - exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) if not exchanges: stats.skipped_no_trajectory += 1 continue diff --git a/src/benchflow/trajectories/export_trl_sft.py b/src/benchflow/trajectories/export_trl_sft.py index f1be50f0a..f105443c3 100644 --- a/src/benchflow/trajectories/export_trl_sft.py +++ b/src/benchflow/trajectories/export_trl_sft.py @@ -322,14 +322,14 @@ def convert_benchflow_rollouts_to_trl_sft_rows( if min_reward is not None and (reward is None or reward < min_reward): stats.skipped_reward += 1 continue + trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" + exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) capture_manifest = read_llm_trajectory_manifest(rollout_dir) if capture_manifest is not None and not capture_manifest_allows_training( - capture_manifest + capture_manifest, exchange_count=len(exchanges) ): stats.skipped_insufficient_capture_fidelity += 1 continue - trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" - exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) if not exchanges: stats.skipped_no_trajectory += 1 continue diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 08b678ed6..3f0faf1c1 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -169,10 +169,12 @@ async def prepare_agent( try: port = await self._ensure_otel_sink(env, sandbox_user=sandbox_user) except Exception as exc: + prepared.pop("CLAUDE_CODE_ENABLE_TELEMETRY", None) + prepared.pop("OTEL_LOG_RAW_API_BODIES", None) warning = _sanitized_error(exc) self._preparation_errors.append(warning) logger.warning( - "Claude OTel correlation unavailable; raw/session fallback remains " + "Claude OTel correlation unavailable; session fallback remains " "enabled: %s", warning, ) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 8dc5d86f5..ae716b375 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -137,14 +137,22 @@ def read_llm_trajectory_manifest(rollout_dir: Path) -> dict[str, Any] | None: return manifest -def capture_manifest_allows_training(manifest: dict[str, Any]) -> bool: - """Return whether the rollout-level capture is safe for training export.""" +def capture_manifest_allows_training( + manifest: dict[str, Any], *, exchange_count: int +) -> bool: + """Return whether manifest fidelity and JSONL cardinality allow training.""" + + expected_count = manifest.get("exchange_count") + if not isinstance(expected_count, int) or isinstance(expected_count, bool): + return False return bool( manifest.get("status") == "complete" and manifest.get("capture_fidelity") == "provider_wire" and manifest.get("request_complete") is True and manifest.get("response_complete") is True + and exchange_count > 0 + and expected_count == exchange_count ) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index edf3a4ba1..b6da1be73 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -20,6 +20,7 @@ import json import logging from contextlib import suppress +from dataclasses import dataclass from pathlib import Path from typing import Any, cast @@ -45,6 +46,14 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _LLMStepsResult: + steps: list[dict[str, Any]] + tool_defs: list[dict[str, Any]] + export_error: str | None + capture_contract_rejected: bool = False + + def _record_to_redacted_json_line(record: dict[str, Any]) -> str: redacted = redact_trajectory_obj(scrub_non_finite(record)) return json.dumps(redacted, default=str, allow_nan=False) @@ -161,21 +170,21 @@ def _llm_steps_from_trajectory( reward: float, is_truncated: bool, trajectory_id_prefix: str, -) -> tuple[list[dict[str, Any]], list[dict[str, Any]], str | None]: +) -> _LLMStepsResult: path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" steps: list[dict[str, Any]] = [] tool_defs: list[dict[str, Any]] = [] if not path.exists(): - return steps, tool_defs, None + return _LLMStepsResult(steps, tool_defs, None) try: exchanges = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError as exc: - return [], [], f"Invalid LLM trajectory JSONL: {exc}" + return _LLMStepsResult([], [], f"Invalid LLM trajectory JSONL: {exc}") capture_manifest = read_llm_trajectory_manifest(rollout_dir) if capture_manifest is not None and not capture_manifest_allows_training( - capture_manifest + capture_manifest, exchange_count=len(exchanges) ): - return [], [], None + return _LLMStepsResult([], [], None, capture_contract_rejected=True) training_success_indices = _training_success_exchange_indices(exchanges) skipped_successful: list[str] = [] for exchange_idx, exchange in enumerate(exchanges): @@ -248,13 +257,13 @@ def _llm_steps_from_trajectory( } steps.append(step) if skipped_successful: - return ( + return _LLMStepsResult( steps, tool_defs, "Successful LLM exchanges were omitted from results.jsonl: " + "; ".join(skipped_successful), ) - return steps, tool_defs, None + return _LLMStepsResult(steps, tool_defs, None) def _response_is_truncated(response_body: dict[str, Any]) -> bool: @@ -463,12 +472,15 @@ def build_rollout_results_record( else f"{task_name}__{rollout_name}" ) is_truncated = bool(partial_trajectory) - steps, tool_defs, llm_export_error = _llm_steps_from_trajectory( + llm_steps = _llm_steps_from_trajectory( rollout_path, reward=reward, is_truncated=is_truncated, trajectory_id_prefix=trajectory_id_prefix, ) + steps = llm_steps.steps + tool_defs = llm_steps.tool_defs + llm_export_error = llm_steps.export_error prompt, completion = _top_level_prompt_completion(steps, prompts) validation_error = _prime_sft_validation_error( prompt=prompt, @@ -483,10 +495,13 @@ def build_rollout_results_record( ) terminal_health_error = bool(error or verifier_error or partial_trajectory) capture_manifest = read_llm_trajectory_manifest(rollout_path) + reported_exchange_count = (capture_manifest or {}).get("exchange_count") audit_only_capture = bool( capture_manifest is not None - and int(capture_manifest.get("exchange_count") or 0) > 0 - and not capture_manifest_allows_training(capture_manifest) + and isinstance(reported_exchange_count, int) + and not isinstance(reported_exchange_count, bool) + and reported_exchange_count > 0 + and llm_steps.capture_contract_rejected ) native_subscription_without_llm = bool( ( diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index da91689c1..839058a9a 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -245,6 +245,42 @@ def test_partial_manifest_blocks_prime_and_trl_training_exports( assert trl_stats.skipped_insufficient_capture_fidelity == 1 +def test_manifest_count_mismatch_blocks_prime_and_trl_training_exports( + tmp_path: Path, +) -> None: + """Guards PR #1057 against exporting a truncated provider-wire JSONL file.""" + + exchange = _exchange(final=True) + exchange["metadata"] = { + "capture_fidelity": "provider_wire", + "request_complete": True, + "response_complete": True, + } + rollout = tmp_path / "job" / "rollout-1" + _write_rollout(rollout, exchanges=[exchange]) + (rollout / "trajectory" / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "complete", + "capture_fidelity": "provider_wire", + "exchange_count": 2, + "request_complete": True, + "response_complete": True, + } + ) + ) + + prime_rows, prime_stats = convert_benchflow_rollouts_to_prime_sft_rows( + tmp_path / "job" + ) + trl_rows, trl_stats = convert_benchflow_rollouts_to_trl_sft_rows(tmp_path / "job") + + assert prime_rows == [] + assert prime_stats.skipped_insufficient_capture_fidelity == 1 + assert trl_rows == [] + assert trl_stats.skipped_insufficient_capture_fidelity == 1 + + def test_skipped_provider_error_counts_rollouts_not_exchanges(tmp_path: Path) -> None: """Guards #828 greptile P1: an all-failed rollout counts as ONE rollout skip, with the exchange count surfaced separately.""" diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 6d2ac1855..296797312 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -97,3 +97,31 @@ def test_corrupt_capture_manifest_fails_closed_for_training(tmp_path: Path) -> N "missing_healthy_structured_llm_trajectory" ) assert row["is_completed"] is False + + +def test_manifest_count_mismatch_fails_closed_for_canonical_results( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on a truncated canonical trajectory.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="provider_wire") + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "complete", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "exchange_count": 2, + "request_complete": True, + "response_complete": True, + } + ) + ) + + row = _build_results_row(tmp_path, agent_result={"total_tokens": 2}) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is False diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index 826338aaf..44dd9fc78 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -911,8 +911,8 @@ async def upload_file(self, *_args, **_kwargs): sandbox_user="agent", ) - assert prepared["CLAUDE_CODE_ENABLE_TELEMETRY"] == "1" - assert prepared["OTEL_LOG_RAW_API_BODIES"].startswith("file:") + assert "CLAUDE_CODE_ENABLE_TELEMETRY" not in prepared + assert "OTEL_LOG_RAW_API_BODIES" not in prepared assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in prepared assert any("/opt/benchflow/node/bin/node" in command for command in commands) From ee19fbb50ae379dd7d1ee12af4d01e972e0b6a2d Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 01:08:38 -0700 Subject: [PATCH 10/74] fix: bind native capture to ACP sessions --- src/benchflow/rollout/__init__.py | 38 ++++- src/benchflow/trajectories/llm_capture.py | 141 +++++++++++++----- .../trajectories/llm_capture_manifest.py | 23 ++- .../trajectories/llm_capture_records.py | 18 ++- .../trajectories/native_capture_parsers.py | 1 + src/benchflow/trajectories/results.py | 9 +- .../test_llm_capture_training_contract.py | 66 ++++++++ tests/trajectories/test_native_llm_capture.py | 1 + .../test_native_session_boundaries.py | 84 ++++++++++- 9 files changed, 330 insertions(+), 51 deletions(-) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 97b3da2f2..7e393e7be 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1365,6 +1365,13 @@ async def connect(self) -> None: getattr(self, "_agent_cfg", None), ), ) + self._bind_llm_capture_session( + agent=cfg.primary_agent, + model=cfg.primary_model, + credential_home=( + f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" + ), + ) self._native_usage_checkpoint = None self._reapply_ask_user_handler() self._attach_trajectory_writer(rollout_dir) @@ -1391,6 +1398,28 @@ def _attach_trajectory_writer(self, rollout_dir: Path) -> None: TrajectoryWriter(traj_path), prior ) + def _bind_llm_capture_session( + self, + *, + agent: str, + model: str | None, + credential_home: str, + role_name: str | None = None, + ) -> None: + """Bind the current ACP ID to the native session files it owns.""" + + llm_capture = getattr(self, "_llm_capture", None) + native_session_id = getattr(self._session, "session_id", None) + if llm_capture is None or not isinstance(native_session_id, str): + return + llm_capture.bind_native_session( + agent=agent, + model=model, + credential_home=credential_home, + native_session_id=native_session_id, + role_name=role_name, + ) + async def disconnect(self) -> None: """Close the ACP client and clean up agent process, keeping the environment alive.""" if getattr(self, "_is_session_factory", False): @@ -2347,6 +2376,7 @@ async def connect_as(self, role: Role) -> None: needs_role_credentials = ( role_agent_differs or role.model != cfg.primary_model or bool(role.env) ) + cred_home = f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" if role_agent_differs: if cfg.skip_agent_install: agent_cfg = None @@ -2360,7 +2390,6 @@ async def connect_as(self, role: Role) -> None: else: agent_cfg = getattr(self, "_agent_cfg", None) if needs_role_credentials: - cred_home = f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" await self._planes.write_credential_files( self._env, role.agent, @@ -2390,7 +2419,6 @@ async def connect_as(self, role: Role) -> None: llm_capture = getattr(self, "_llm_capture", None) if llm_capture is not None: - cred_home = f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" agent_env = await llm_capture.prepare_agent( self._env, agent=role.agent, @@ -2445,6 +2473,12 @@ async def connect_as(self, role: Role) -> None: role.agent, getattr(self, "_task", None), agent_cfg ), ) + self._bind_llm_capture_session( + agent=role.agent, + model=role.model, + credential_home=cred_home, + role_name=role.name, + ) self._reapply_ask_user_handler() self._attach_trajectory_writer(rollout_dir) self._active_role = role diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 3f0faf1c1..b0c426bb0 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -7,9 +7,11 @@ import json import logging import os +import re import shlex import tempfile from contextlib import suppress +from dataclasses import replace from datetime import datetime from pathlib import Path, PurePosixPath from typing import Any @@ -49,6 +51,8 @@ logger = logging.getLogger(__name__) _REMOTE_CAPTURE_PREFIX = "/tmp/benchflow-llm-capture-" +_MAX_NATIVE_SESSION_FILES = 1000 +_SAFE_NATIVE_SESSION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _OTEL_SINK_SOURCE = r""" import { createServer } from 'node:http'; import { mkdirSync, writeFileSync } from 'node:fs'; @@ -148,12 +152,24 @@ async def prepare_agent( native=native, role=role_name or "primary", ) - primary_key = ("primary", agent, model, credential_home) + primary_key = _capture_target_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=None, + ) if role_name is None: self._targets[primary_key] = target else: self._targets.pop(primary_key, None) - self._targets[(role_name, agent, model, credential_home)] = target + self._targets[ + _capture_target_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=role_name, + ) + ] = target self._refresh_manifest_auth_mode() if not native: write_llm_trajectory_manifest(self.rollout_dir, self.manifest) @@ -196,6 +212,39 @@ async def prepare_agent( def _native_targets(self) -> list[_CaptureTarget]: return [target for target in self._targets.values() if target.native] + def bind_native_session( + self, + *, + agent: str, + model: str | None, + credential_home: str, + native_session_id: str, + role_name: str | None = None, + ) -> None: + """Bind an ACP session ID to its prepared native capture target.""" + + key = _capture_target_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=role_name, + ) + target = self._targets.get(key) + if target is None or not target.native: + return + if _SAFE_NATIVE_SESSION_ID.fullmatch(native_session_id) is None: + warning = "native ACP session identifier was unsafe for file scoping" + self._preparation_errors.append(warning) + logger.warning("Native LLM capture disabled: %s", warning) + return + session_ids = tuple(sorted({*target.native_session_ids, native_session_id})) + if len(session_ids) > _MAX_NATIVE_SESSION_FILES: + warning = "native ACP session count exceeded the capture limit" + self._preparation_errors.append(warning) + logger.warning("Native LLM capture disabled: %s", warning) + return + self._targets[key] = replace(target, native_session_ids=session_ids) + def _refresh_manifest_auth_mode(self) -> None: modes = {target.auth_mode for target in self._targets.values()} if len(modes) == 1: @@ -449,36 +498,18 @@ async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: _NativeCaptureBundle(targets=claude_targets, result=result) ) raw_claude_captured = True - credential_homes = sorted( - {target.credential_home for target in native_targets} - ) - for index, credential_home in enumerate(credential_homes): - home_targets = tuple( - target - for target in native_targets - if target.credential_home == credential_home - ) - home_claude_targets = tuple( - target - for target in home_targets - if _is_claude_code_agent(target.agent) - ) - home_codex_targets = tuple( - target for target in home_targets if target.agent == "codex-acp" - ) - claude_local = local_root / f"home-{index}" / "claude-projects" - codex_local = local_root / f"home-{index}" / "codex-sessions" - if ( - home_claude_targets - and not raw_claude_captured - and await _download_recent_session_files( + for index, target in enumerate(native_targets): + if _is_claude_code_agent(target.agent) and not raw_claude_captured: + claude_local = local_root / f"target-{index}" / "claude-projects" + downloaded = await _download_bound_session_files( env, - f"{credential_home}/.claude/projects", + f"{target.credential_home}/.claude/projects", claude_local, started_at=self.started_at, + session_ids=target.native_session_ids, ) - ): - target = home_claude_targets[0] + if not downloaded: + continue result = parse_claude_sessions( claude_local, agent=target.agent, @@ -488,17 +519,21 @@ async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: if result is not None: bundles.append( _NativeCaptureBundle( - targets=home_claude_targets, + targets=(target,), result=result, ) ) - if home_codex_targets and await _download_recent_session_files( - env, - f"{credential_home}/.codex/sessions", - codex_local, - started_at=self.started_at, - ): - target = home_codex_targets[0] + elif target.agent == "codex-acp": + codex_local = local_root / f"target-{index}" / "codex-sessions" + downloaded = await _download_bound_session_files( + env, + f"{target.credential_home}/.codex/sessions", + codex_local, + started_at=self.started_at, + session_ids=target.native_session_ids, + ) + if not downloaded: + continue result = parse_codex_sessions( codex_local, agent=target.agent, @@ -510,7 +545,7 @@ async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: if result is not None: bundles.append( _NativeCaptureBundle( - targets=home_codex_targets, + targets=(target,), result=result, ) ) @@ -556,19 +591,35 @@ def _finish_manifest( write_llm_trajectory_manifest(self.rollout_dir, self.manifest) -async def _download_recent_session_files( +async def _download_bound_session_files( env: Any, remote: str, local: Path, *, started_at: datetime, + session_ids: tuple[str, ...], ) -> bool: + if not session_ids: + return False + if any(_SAFE_NATIVE_SESSION_ID.fullmatch(value) is None for value in session_ids): + raise RuntimeError("Native session discovery received an unsafe session ID") boundary = started_at.timestamp() - 1.0 remote_root = shlex.quote(remote) + filename_patterns = [ + pattern + for session_id in session_ids + for pattern in (f"{session_id}.jsonl", f"*-{session_id}.jsonl") + ] + filename_filter = ( + r"\( " + + " -o ".join(f"-name {shlex.quote(pattern)}" for pattern in filename_patterns) + + r" \)" + ) result = await env.exec( f"if test -d {remote_root}; then " f"find {remote_root} -type f -name '*.jsonl' " - f"-newermt {shlex.quote(f'@{boundary}')} -printf '%P\\n' | head -n 1001; " + f"-newermt {shlex.quote(f'@{boundary}')} {filename_filter} " + f"-printf '%P\\n' | head -n {_MAX_NATIVE_SESSION_FILES + 1}; " "fi", user="root", timeout_sec=10, @@ -577,7 +628,7 @@ async def _download_recent_session_files( detail = (result.stderr or result.stdout or "session discovery failed")[:300] raise RuntimeError(f"Native session discovery failed: {detail}") relative_paths = [line for line in result.stdout.splitlines() if line] - if len(relative_paths) > 1000: + if len(relative_paths) > _MAX_NATIVE_SESSION_FILES: raise RuntimeError("Native session discovery exceeded the 1000-file limit") if not relative_paths: return False @@ -595,6 +646,16 @@ async def _download_recent_session_files( return True +def _capture_target_key( + *, + agent: str, + model: str | None, + credential_home: str, + role_name: str | None, +) -> tuple[str, str, str | None, str]: + return (role_name or "primary", agent, model, credential_home) + + def _is_claude_code_agent(agent: str) -> bool: config = AGENTS.get(agent) subscription = config.subscription_auth if config is not None else None diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index ae716b375..5e890c03e 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -14,7 +14,7 @@ from pathlib import Path from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, ValidationError LLM_TRAJECTORY_FILENAME = "llm_trajectory.jsonl" LLM_TRAJECTORY_MANIFEST_FILENAME = "llm_trajectory.manifest.json" @@ -156,6 +156,27 @@ def capture_manifest_allows_training( ) +def capture_manifest_has_oauth_role_capture(manifest: dict[str, Any]) -> bool: + """Return whether a mixed manifest contains captured OAuth role evidence.""" + + role_captures = manifest.get("role_captures") + if not isinstance(role_captures, list): + return False + for value in role_captures: + try: + role_capture = LLMRoleCapture.model_validate(value) + except ValidationError: + continue + if ( + role_capture.auth_mode is AuthMode.OAUTH_SUBSCRIPTION + and role_capture.capture_source is not CaptureSource.NONE + and role_capture.capture_fidelity is not CaptureFidelity.NONE + and role_capture.exchange_count > 0 + ): + return True + return False + + def _atomic_write_text(path: Path, payload: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index c925d2621..011cd7491 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -30,6 +30,7 @@ class CaptureTarget: auth_mode: AuthMode native: bool role: str = "agent" + native_session_ids: tuple[str, ...] = () @dataclass(frozen=True) @@ -260,7 +261,7 @@ def _native_bundle_records(bundle: NativeCaptureBundle) -> list[dict[str, Any]]: if not isinstance(metadata, dict): metadata = {} record["metadata"] = metadata - target = _target_for_model(list(bundle.targets), _record_model(record)) + target = _target_for_record(list(bundle.targets), record) if target is not None: metadata.update( { @@ -353,6 +354,21 @@ def _target_for_model( return matches[0] if len(matches) == 1 else None +def _target_for_record( + targets: list[CaptureTarget], record: dict[str, Any] +) -> CaptureTarget | None: + native_session_id = _record_metadata(record).get("native_session_id") + if isinstance(native_session_id, str): + session_matches = [ + target + for target in targets + if native_session_id in target.native_session_ids + ] + if len(session_matches) == 1: + return session_matches[0] + return _target_for_model(targets, _record_model(record)) + + def _sort_exchange_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: def timestamp(record: dict[str, Any]) -> str: request = record.get("request") diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index cee757bc9..6e3c0fd8d 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -117,6 +117,7 @@ def parse_claude_raw_capture( response_complete=True, extra_metadata={ "provider_request_id": request_id, + "native_session_id": event_session, "request_body_file": request_path.name, "response_body_file": response_path.name, "pairing": ( diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index b6da1be73..3335f23ae 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -34,6 +34,7 @@ ) from benchflow.trajectories.llm_capture_manifest import ( capture_manifest_allows_training, + capture_manifest_has_oauth_role_capture, read_llm_trajectory_manifest, ) from benchflow.trajectories.types import redact_trajectory_obj @@ -503,10 +504,14 @@ def build_rollout_results_record( and reported_exchange_count > 0 and llm_steps.capture_contract_rejected ) - native_subscription_without_llm = bool( + audit_capture_preserves_completion = bool( ( (agent_result or {}).get("usage_source") == USAGE_SOURCE_AGENT_NATIVE_ACP or (capture_manifest or {}).get("auth_mode") == "oauth_subscription" + or ( + capture_manifest is not None + and capture_manifest_has_oauth_role_capture(capture_manifest) + ) ) and ( capture_manifest is None @@ -545,7 +550,7 @@ def build_rollout_results_record( training_ready_reason = "verifier_error" else: training_ready_reason = "missing_healthy_structured_llm_trajectory" - if error_obj is None and not native_subscription_without_llm: + if error_obj is None and not audit_capture_preserves_completion: error_name = ( training_ready_reason if training_ready_reason diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 296797312..4fa8a0f3e 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -125,3 +125,69 @@ def test_manifest_count_mismatch_fails_closed_for_canonical_results( assert row["info"]["training_ready"] is False assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" assert row["is_completed"] is False + + +def test_mixed_oauth_audit_capture_preserves_successful_completion( + tmp_path: Path, +) -> None: + """Guards PR #1057 against turning successful mixed-auth runs into errors.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="provider_wire") + provider_row = json.loads((trajectory_dir / "llm_trajectory.jsonl").read_text()) + oauth_row = json.loads(json.dumps(provider_row)) + oauth_row["metadata"].update( + { + "capture_fidelity": "agent_session", + "auth_mode": "oauth_subscription", + "request_complete": False, + } + ) + (trajectory_dir / "llm_trajectory.jsonl").write_text( + json.dumps(provider_row) + "\n" + json.dumps(oauth_row) + "\n" + ) + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "partial", + "capture_fidelity": "mixed", + "auth_mode": "mixed", + "exchange_count": 2, + "request_complete": False, + "response_complete": True, + "role_captures": [ + { + "role": "coder", + "agent": "codex-acp", + "auth_mode": "api_key", + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + }, + { + "role": "reviewer", + "agent": "claude-agent-acp", + "auth_mode": "oauth_subscription", + "capture_source": "claude_native_session", + "capture_fidelity": "agent_session", + "exchange_count": 1, + "request_complete": False, + "response_complete": False, + }, + ], + } + ) + ) + + row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is True + assert row["error"] is None diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index 44dd9fc78..d57b00266 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -124,6 +124,7 @@ def test_claude_otel_raw_bodies_become_provider_wire_exchanges(tmp_path: Path) - assert exchange.response.body["id"] == "msg_123" assert exchange.duration_ms == 500 assert exchange.metadata["provider_request_id"] == "req_123" + assert exchange.metadata["native_session_id"] == "session-1" def test_claude_concurrent_raw_pairing_fails_closed_for_training( diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 89d79f31b..1cf800e6d 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -6,12 +6,14 @@ from datetime import UTC, datetime from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock import pytest +from benchflow.rollout import Rollout from benchflow.trajectories.llm_capture import ( LLMTrajectoryCapture, - _download_recent_session_files, + _download_bound_session_files, ) from benchflow.trajectories.native_capture_parsers import parse_codex_sessions @@ -129,17 +131,18 @@ async def upload_file(self, *_args, **_kwargs): async def test_native_download_selects_recent_files_before_copying( tmp_path: Path, ) -> None: - """Guards PR #1057 against copying an entire reused native-session tree.""" + """Guards PR #1057 against copying unrelated concurrent native sessions.""" commands: list[str] = [] downloads: list[tuple[str, Path]] = [] + native_session_id = "019effaf-3966-75d3-b61a-2916c84b0ac8" class DockerLikeEnv: async def exec(self, command, **_kwargs): commands.append(command) return SimpleNamespace( return_code=0, - stdout="2026/08/28/session.jsonl\n", + stdout=f"2026/08/28/rollout-now-{native_session_id}.jsonl\n", stderr="", ) @@ -150,18 +153,89 @@ async def download_file(self, remote, local): destination.write_text("{}\n") destination = tmp_path / "missing-parent" / "sessions" - downloaded = await _download_recent_session_files( + downloaded = await _download_bound_session_files( DockerLikeEnv(), "/home/agent/.codex/sessions", destination, started_at=STARTED_AT, + session_ids=(native_session_id,), ) assert downloaded is True assert len(downloads) == 1 - assert downloads[0][0].endswith("/2026/08/28/session.jsonl") + assert downloads[0][0].endswith( + f"/2026/08/28/rollout-now-{native_session_id}.jsonl" + ) assert downloads[0][1].is_file() assert "-newermt" in commands[0] + assert f"-name {native_session_id}.jsonl" in commands[0] + assert f"-name '*-{native_session_id}.jsonl'" in commands[0] + + +@pytest.mark.asyncio +async def test_native_capture_binds_only_returned_acp_session_ids( + tmp_path: Path, +) -> None: + """Guards PR #1057's exact ACP-to-native-session ownership boundary.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + await capture.prepare_agent( + None, + agent="codex-acp", + model="gpt-5.6", + agent_env={"CODEX_AUTH_JSON": '{"tokens":{"access_token":"test"}}'}, + credential_home="/home/agent", + sandbox_user="agent", + ) + + capture.bind_native_session( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + native_session_id="019effaf-3966-75d3-b61a-2916c84b0ac8", + ) + capture.bind_native_session( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + native_session_id="019effaf-3ab7-71f1-8ff3-fdecf66b551e", + ) + + target = capture._native_targets()[0] + assert target.native_session_ids == ( + "019effaf-3966-75d3-b61a-2916c84b0ac8", + "019effaf-3ab7-71f1-8ff3-fdecf66b551e", + ) + + +def test_rollout_binds_the_acp_session_after_connect() -> None: + """Guards PR #1057's rollout-to-capture session binding seam.""" + + rollout = object.__new__(Rollout) + capture = SimpleNamespace(bind_native_session=Mock()) + rollout._llm_capture = capture + rollout._session = SimpleNamespace(session_id="native-session-1") + + rollout._bind_llm_capture_session( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + role_name="solver", + ) + + capture.bind_native_session.assert_called_once_with( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + native_session_id="native-session-1", + role_name="solver", + ) def test_codex_session_files_do_not_share_request_history(tmp_path: Path) -> None: From 69ed43596cd383d94a0d4f191760b3110afc35d6 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 01:20:29 -0700 Subject: [PATCH 11/74] fix: fail closed on capture finalization --- src/benchflow/rollout/__init__.py | 8 +++ src/benchflow/trajectories/llm_capture.py | 13 +++-- tests/test_usage_litellm.py | 50 +++++++++++++++++++ .../test_llm_capture_training_contract.py | 43 ++++++++++++++++ .../test_native_session_boundaries.py | 33 ++++++++++++ 5 files changed, 144 insertions(+), 3 deletions(-) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 7e393e7be..31d89b5cc 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -2022,12 +2022,16 @@ async def cleanup(self) -> None: self._evolved_skills = None usage_runtime = getattr(self, "_usage_runtime", None) + provider_capture_errors: list[str] = [] if usage_runtime is not None: try: await self._planes.stop_provider_runtime(usage_runtime) self._usage_metrics = self._planes.extract_usage(usage_runtime) except Exception as e: logger.warning(f"Usage telemetry runtime stop failed: {e}") + provider_capture_errors.append( + "provider runtime stop or remote capture import failed" + ) self._usage_metrics = self._planes.extract_usage(None) # Snapshot any provider failure (401/403/429/503) now that captures # are imported (stop() populated the trajectory). This must happen @@ -2057,6 +2061,9 @@ async def cleanup(self) -> None: self._write_llm_trajectory(usage_runtime) except Exception as e: logger.warning(f"LLM trajectory write failed: {e}") + provider_capture_errors.append( + "provider live trajectory reconciliation failed" + ) try: self._reconcile_acp_tool_evidence(usage_runtime) except Exception as e: @@ -2076,6 +2083,7 @@ async def cleanup(self) -> None: self._env, acp_events=acp_events, model_call_seen=model_call_seen, + capture_errors=provider_capture_errors, ) except Exception as e: logger.warning(f"LLM trajectory finalization failed: {e}") diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index b0c426bb0..11b9a633e 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -258,6 +258,7 @@ async def finalize( *, acp_events: list[dict[str, Any]], model_call_seen: bool, + capture_errors: list[str] | None = None, ) -> None: """Publish the highest-fidelity available capture and terminal sidecar.""" @@ -280,11 +281,17 @@ async def finalize( raise native_bundles: list[_NativeCaptureBundle] = [] - collection_errors: list[str] = list(self._preparation_errors) + collection_errors = list( + dict.fromkeys([*self._preparation_errors, *(capture_errors or [])]) + ) native_targets = self._native_targets() - if native_targets and env is not None: + native_resources_exist = self._collector_started or self._capture_root_prepared + if (native_targets or native_resources_exist) and env is not None: try: - native_bundles = await self._collect_native_results(env) + if native_targets: + native_bundles = await self._collect_native_results(env) + elif self._collector_started: + await self._stop_otel_sink(env) except Exception as exc: collection_errors.append(_sanitized_error(exc)) logger.warning("Native LLM trajectory collection failed: %s", exc) diff --git a/tests/test_usage_litellm.py b/tests/test_usage_litellm.py index 48a377972..5a726cf29 100644 --- a/tests/test_usage_litellm.py +++ b/tests/test_usage_litellm.py @@ -355,3 +355,53 @@ def reconcile_live_capture(self): assert rollout._usage_metrics["usage_source"] == "provider_response" assert (tmp_path / "trajectory" / "llm_trajectory.jsonl").exists() + + +@pytest.mark.asyncio +async def test_rollout_cleanup_propagates_provider_stop_failure_to_capture(tmp_path): + """Guards PR #1057 against accepting a truncated live provider prefix.""" + + from benchflow.rollout import Rollout, RolloutConfig + + class FakeServer: + trajectory = _trajectory( + { + "model": "gpt-5.6", + "usage": {"input_tokens": 10, "output_tokens": 2}, + } + ) + + def reconcile_live_capture(self): + return None + + async def fail_stop(_runtime): + raise RuntimeError("remote callback import failed") + + capture = SimpleNamespace(finalize=AsyncMock(), record_failure=AsyncMock()) + rollout = Rollout.__new__(Rollout) + rollout._config = RolloutConfig(task_path=tmp_path / "task") + rollout._trajectory = [] + rollout._acp_client = None + rollout._agent_launch = "" + rollout._env = SimpleNamespace(stop=AsyncMock()) + rollout._environment = None + rollout._usage_runtime = ProviderRuntime( + kind="litellm", + agent_base_url="http://127.0.0.1:4000", + backend_model="gpt-5.6", + server=FakeServer(), + ) + rollout._planes = SimpleNamespace( + stop_provider_runtime=fail_stop, + extract_usage=extract_usage, + ) + rollout._rollout_dir = tmp_path + rollout._env_externally_owned = False + rollout._llm_capture = capture + + await rollout.cleanup() + + capture.finalize.assert_awaited_once() + assert capture.finalize.await_args.kwargs["capture_errors"] == [ + "provider runtime stop or remote capture import failed" + ] diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 4fa8a0f3e..30f9d13da 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -3,8 +3,12 @@ from __future__ import annotations import json +from datetime import UTC, datetime from pathlib import Path +import pytest + +from benchflow.trajectories.llm_capture import LLMTrajectoryCapture from benchflow.trajectories.results import build_rollout_results_record @@ -191,3 +195,42 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" assert row["is_completed"] is True assert row["error"] is None + + +@pytest.mark.asyncio +async def test_provider_finalization_error_rejects_complete_live_prefix( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on a truncated provider prefix.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=datetime(2026, 8, 29, tzinfo=UTC), + ) + capture.configure({"OPENAI_API_KEY": "test-key"}) + _write_exchange(tmp_path / "trajectory", fidelity="provider_wire") + + await capture.finalize( + None, + acp_events=[], + model_call_seen=True, + capture_errors=["provider runtime stop or remote capture import failed"], + ) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["status"] == "partial" + assert manifest["exchange_count"] == 1 + assert manifest["errors"] == [ + "provider runtime stop or remote capture import failed" + ] + row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + assert row["info"]["training_ready"] is False + assert row["is_completed"] is False diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 1cf800e6d..3eb1b12b4 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -214,6 +214,39 @@ async def test_native_capture_binds_only_returned_acp_session_ids( ) +@pytest.mark.asyncio +async def test_replaced_native_target_still_releases_collector(tmp_path: Path) -> None: + """Guards PR #1057 against leaking a replaced OAuth collector.""" + + commands: list[str] = [] + + class CleanupEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + return SimpleNamespace(return_code=0, stdout="", stderr="") + + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.configure({"ANTHROPIC_API_KEY": "test-key"}) + capture._collector_started = True + capture._capture_root_prepared = True + + await capture.finalize( + CleanupEnv(), + acp_events=[], + model_call_seen=False, + ) + + assert any("kill -TERM" in command for command in commands) + assert any("-depth -delete" in command for command in commands) + assert capture._collector_started is False + + def test_rollout_binds_the_acp_session_after_connect() -> None: """Guards PR #1057's rollout-to-capture session binding seam.""" From bee800110d992dc53b8b75a093ded5a84b1a6bb8 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 01:33:40 -0700 Subject: [PATCH 12/74] fix: track collector ownership before readiness --- src/benchflow/trajectories/llm_capture.py | 9 ++- .../test_native_session_boundaries.py | 55 +++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 11b9a633e..425a07e73 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -113,6 +113,7 @@ def __init__( ) self._targets: dict[tuple[str, str, str | None, str], _CaptureTarget] = {} self._collector_started = False + self._collector_owned = False self._capture_root_prepared = False self._preparation_errors: list[str] = [] @@ -285,12 +286,12 @@ async def finalize( dict.fromkeys([*self._preparation_errors, *(capture_errors or [])]) ) native_targets = self._native_targets() - native_resources_exist = self._collector_started or self._capture_root_prepared + native_resources_exist = self._collector_owned or self._capture_root_prepared if (native_targets or native_resources_exist) and env is not None: try: if native_targets: native_bundles = await self._collect_native_results(env) - elif self._collector_started: + elif self._collector_owned: await self._stop_otel_sink(env) except Exception as exc: collection_errors.append(_sanitized_error(exc)) @@ -429,6 +430,7 @@ async def _ensure_otel_sink(self, env: Any, *, sandbox_user: str | None) -> int: tail -c 300 {self._remote_capture_root}/collector.stderr >&2 2>/dev/null || true exit 1 """ + self._collector_owned = True result = await env.exec( command, user=sandbox_user or "root", @@ -469,6 +471,7 @@ async def _stop_otel_sink(self, env: Any) -> None: detail = (result.stderr or result.stdout or "collector did not stop")[:300] raise RuntimeError(f"Claude OTel sink shutdown failed: {detail}") self._collector_started = False + self._collector_owned = False async def _read_collector_port(self, env: Any) -> int: result = await env.exec( @@ -481,7 +484,7 @@ async def _read_collector_port(self, env: Any) -> int: return _parse_port(result.stdout) async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: - if self._collector_started: + if self._collector_owned: await self._stop_otel_sink(env) bundles: list[_NativeCaptureBundle] = [] native_targets = self._native_targets() diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 3eb1b12b4..e71c35826 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -127,6 +127,59 @@ async def upload_file(self, *_args, **_kwargs): assert clear_at < create_at +@pytest.mark.asyncio +async def test_claude_startup_timeout_still_releases_owned_collector( + tmp_path: Path, +) -> None: + """Guards PR #1057 against leaking a collector after startup timeout.""" + + commands: list[str] = [] + + class StartupTimeoutEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + if "nohup" in command: + return SimpleNamespace( + return_code=1, + stdout="", + stderr="collector port was not observed", + ) + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def upload_file(self, *_args, **_kwargs): + return None + + async def download_dir(self, _remote, local): + Path(local).mkdir(parents=True) + + env = StartupTimeoutEnv() + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + await capture.prepare_agent( + env, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "test-token"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + + assert capture._collector_started is False + assert capture._collector_owned is True + + await capture.finalize(env, acp_events=[], model_call_seen=False) + + stop_commands = [command for command in commands if "read -r old_pid" in command] + assert len(stop_commands) == 2 + assert capture._collector_owned is False + assert any("-depth -delete" in command for command in commands) + + @pytest.mark.asyncio async def test_native_download_selects_recent_files_before_copying( tmp_path: Path, @@ -234,6 +287,7 @@ async def exec(self, command, **_kwargs): ) capture.configure({"ANTHROPIC_API_KEY": "test-key"}) capture._collector_started = True + capture._collector_owned = True capture._capture_root_prepared = True await capture.finalize( @@ -245,6 +299,7 @@ async def exec(self, command, **_kwargs): assert any("kill -TERM" in command for command in commands) assert any("-depth -delete" in command for command in commands) assert capture._collector_started is False + assert capture._collector_owned is False def test_rollout_binds_the_acp_session_after_connect() -> None: From 2f1172588456ac7d750a1f28aac5baeb3f62f63e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 01:44:43 -0700 Subject: [PATCH 13/74] fix: harden provisional capture cleanup --- src/benchflow/trajectories/llm_capture.py | 56 ++++++-- .../test_native_session_boundaries.py | 135 ++++++++++++++++++ 2 files changed, 176 insertions(+), 15 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 425a07e73..ebdb198f0 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -112,6 +112,7 @@ def __init__( started_at=started_at, ) self._targets: dict[tuple[str, str, str | None, str], _CaptureTarget] = {} + self._provisional_target_key: tuple[str, str, str | None, str] | None = None self._collector_started = False self._collector_owned = False self._capture_root_prepared = False @@ -160,9 +161,17 @@ async def prepare_agent( role_name=None, ) if role_name is None: + if ( + self._provisional_target_key is not None + and self._provisional_target_key != primary_key + ): + self._targets.pop(self._provisional_target_key, None) self._targets[primary_key] = target + self._provisional_target_key = primary_key else: - self._targets.pop(primary_key, None) + if self._provisional_target_key is not None: + self._targets.pop(self._provisional_target_key, None) + self._provisional_target_key = None self._targets[ _capture_target_key( agent=agent, @@ -231,7 +240,11 @@ def bind_native_session( role_name=role_name, ) target = self._targets.get(key) - if target is None or not target.native: + if target is None: + return + if key == self._provisional_target_key: + self._provisional_target_key = None + if not target.native: return if _SAFE_NATIVE_SESSION_ID.fullmatch(native_session_id) is None: warning = "native ACP session identifier was unsafe for file scoping" @@ -287,6 +300,7 @@ async def finalize( ) native_targets = self._native_targets() native_resources_exist = self._collector_owned or self._capture_root_prepared + cleanup_failed = False if (native_targets or native_resources_exist) and env is not None: try: if native_targets: @@ -297,7 +311,12 @@ async def finalize( collection_errors.append(_sanitized_error(exc)) logger.warning("Native LLM trajectory collection failed: %s", exc) finally: - await self._cleanup_remote_capture(env) + try: + await self._cleanup_remote_capture(env) + except Exception as exc: + collection_errors.append(_sanitized_error(exc)) + logger.error("Sandbox LLM capture cleanup failed: %s", exc) + cleanup_failed = True if not provider_records and not native_bundles and acp_events: projected = project_acp_trajectory( @@ -330,7 +349,9 @@ async def finalize( write_exchange_records(self.trajectory_path, assembly.records) self.manifest.auth_mode = assembly.auth_mode self._finish_manifest( - status=assembly.status, + status=( + CaptureStatus.CAPTURE_FAILED if cleanup_failed else assembly.status + ), source=assembly.source, fidelity=assembly.fidelity, exchange_count=len(assembly.records), @@ -564,17 +585,22 @@ async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: async def _cleanup_remote_capture(self, env: Any) -> None: if not self._capture_root_prepared: return - try: - result = await env.exec( - f"find {self._remote_capture_root} -depth -delete", - user="root", - timeout_sec=10, - ) - if result.return_code != 0: - detail = (result.stderr or result.stdout or "unknown error")[:300] - logger.warning("Sandbox LLM capture cleanup failed: %s", detail) - except Exception as exc: - logger.warning("Sandbox LLM capture cleanup failed: %s", exc) + result = await env.exec( + "for attempt in 1 2 3; do\n" + f" if ! test -e {self._remote_capture_root} || " + f"find {self._remote_capture_root} -depth -delete; then\n" + " exit 0\n" + " fi\n" + " sleep 0.1\n" + "done\n" + "exit 1", + user="root", + timeout_sec=10, + ) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "unknown error")[:300] + raise RuntimeError(f"Sandbox LLM capture cleanup failed: {detail}") + self._capture_root_prepared = False def _finish_manifest( self, diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index e71c35826..2259759e6 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -85,6 +85,89 @@ async def test_phase_api_registers_provisional_primary_oauth_target( assert targets[0].role == "primary" +@pytest.mark.asyncio +async def test_first_named_role_removes_actual_unbound_provisional_target( + tmp_path: Path, +) -> None: + """Guards PR #1057 against retaining a different provisional primary.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-primary", + session_id="rollout-1", + started_at=STARTED_AT, + ) + await capture.prepare_agent( + None, + agent="codex-acp", + model="gpt-primary", + agent_env={ + "CODEX_AUTH_JSON": '{"auth_mode":"chatgpt","tokens":{"refresh_token":"test"}}' + }, + credential_home="/home/agent", + sandbox_user="agent", + ) + await capture.prepare_agent( + None, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"ANTHROPIC_API_KEY": "test-key"}, + credential_home="/home/agent", + sandbox_user="agent", + role_name="reviewer", + ) + + targets = list(capture._targets.values()) + assert [(target.role, target.agent) for target in targets] == [ + ("reviewer", "claude-agent-acp") + ] + assert capture._provisional_target_key is None + + +@pytest.mark.asyncio +async def test_executed_primary_is_not_removed_by_later_named_role( + tmp_path: Path, +) -> None: + """Guards PR #1057 against deleting an activated primary capture target.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-primary", + session_id="rollout-1", + started_at=STARTED_AT, + ) + await capture.prepare_agent( + None, + agent="codex-acp", + model="gpt-primary", + agent_env={"OPENAI_API_KEY": "test-key"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + capture.bind_native_session( + agent="codex-acp", + model="gpt-primary", + credential_home="/home/agent", + native_session_id="primary-session", + ) + await capture.prepare_agent( + None, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"ANTHROPIC_API_KEY": "test-key"}, + credential_home="/home/agent", + sandbox_user="agent", + role_name="reviewer", + ) + + assert {target.role for target in capture._targets.values()} == { + "primary", + "reviewer", + } + + @pytest.mark.asyncio async def test_claude_capture_setup_clears_reused_raw_attempt( tmp_path: Path, @@ -302,6 +385,58 @@ async def exec(self, command, **_kwargs): assert capture._collector_owned is False +@pytest.mark.asyncio +async def test_raw_capture_cleanup_failure_marks_manifest_failed( + tmp_path: Path, +) -> None: + """Guards PR #1057 against hiding retained raw provider bodies.""" + + commands: list[str] = [] + + class CleanupFailureEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + return SimpleNamespace( + return_code=1, + stdout="", + stderr="capture directory deletion failed", + ) + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.configure({"OPENAI_API_KEY": "test-key"}) + capture._capture_root_prepared = True + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"model": "gpt-5.6", "input": "hello"}}, + "response": {"status_code": 200, "body": {"output": []}}, + } + ) + + "\n" + ) + + await capture.finalize( + CleanupFailureEnv(), + acp_events=[], + model_call_seen=True, + ) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert "for attempt in 1 2 3" in commands[0] + assert manifest["status"] == "capture_failed" + assert manifest["exchange_count"] == 1 + assert any("cleanup failed" in error for error in manifest["errors"]) + assert capture._capture_root_prepared is True + + def test_rollout_binds_the_acp_session_after_connect() -> None: """Guards PR #1057's rollout-to-capture session binding seam.""" From cedef525ed6181fd33ad559365b97a5ff528c154 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 01:54:42 -0700 Subject: [PATCH 14/74] fix: preserve native sessions across reconnects --- src/benchflow/trajectories/llm_capture.py | 26 ++++++---- .../test_native_session_boundaries.py | 48 +++++++++++++++++++ 2 files changed, 66 insertions(+), 8 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index ebdb198f0..7a605aec0 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -154,6 +154,23 @@ async def prepare_agent( native=native, role=role_name or "primary", ) + target_key = _capture_target_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=role_name, + ) + previous_target = self._targets.get(target_key) + if ( + previous_target is not None + and previous_target.native + and target.native + and previous_target.auth_mode is target.auth_mode + ): + target = replace( + target, + native_session_ids=previous_target.native_session_ids, + ) primary_key = _capture_target_key( agent=agent, model=model, @@ -172,14 +189,7 @@ async def prepare_agent( if self._provisional_target_key is not None: self._targets.pop(self._provisional_target_key, None) self._provisional_target_key = None - self._targets[ - _capture_target_key( - agent=agent, - model=model, - credential_home=credential_home, - role_name=role_name, - ) - ] = target + self._targets[target_key] = target self._refresh_manifest_auth_mode() if not native: write_llm_trajectory_manifest(self.rollout_dir, self.manifest) diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 2259759e6..036667a26 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -350,6 +350,54 @@ async def test_native_capture_binds_only_returned_acp_session_ids( ) +@pytest.mark.asyncio +async def test_native_role_reprepare_preserves_prior_session_ids( + tmp_path: Path, +) -> None: + """Guards PR #1057's prepare-bind reconnect sequence across rounds.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + prepare_kwargs = { + "agent": "codex-acp", + "model": "gpt-5.6", + "agent_env": { + "CODEX_AUTH_JSON": '{"auth_mode":"chatgpt","tokens":{"refresh_token":"test"}}' + }, + "credential_home": "/home/agent", + "sandbox_user": "agent", + "role_name": "solver", + } + await capture.prepare_agent(None, **prepare_kwargs) + capture.bind_native_session( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + native_session_id="019effaf-3966-75d3-b61a-2916c84b0ac8", + role_name="solver", + ) + + await capture.prepare_agent(None, **prepare_kwargs) + capture.bind_native_session( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + native_session_id="019effaf-3ab7-71f1-8ff3-fdecf66b551e", + role_name="solver", + ) + + target = capture._native_targets()[0] + assert target.native_session_ids == ( + "019effaf-3966-75d3-b61a-2916c84b0ac8", + "019effaf-3ab7-71f1-8ff3-fdecf66b551e", + ) + + @pytest.mark.asyncio async def test_replaced_native_target_still_releases_collector(tmp_path: Path) -> None: """Guards PR #1057 against leaking a replaced OAuth collector.""" From 4b671b559327fe1a40a24257b3d765a7058268b6 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 02:13:28 -0700 Subject: [PATCH 15/74] fix: isolate multi-target native capture --- src/benchflow/trajectories/llm_capture.py | 242 +++++++++++----- tests/trajectories/test_native_llm_capture.py | 7 +- .../test_native_session_boundaries.py | 266 +++++++++++++++++- 3 files changed, 441 insertions(+), 74 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 7a605aec0..edb7eefa8 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -11,7 +11,7 @@ import shlex import tempfile from contextlib import suppress -from dataclasses import replace +from dataclasses import dataclass, replace from datetime import datetime from pathlib import Path, PurePosixPath from typing import Any @@ -41,6 +41,7 @@ write_exchange_records, ) from benchflow.trajectories.native_capture_parsers import ( + NativeParseResult, parse_claude_raw_capture, parse_claude_sessions, parse_codex_sessions, @@ -85,6 +86,14 @@ """.strip() +@dataclass(frozen=True) +class _NativeCollection: + """Native bundles retained alongside isolated collection errors.""" + + bundles: tuple[_NativeCaptureBundle, ...] = () + errors: tuple[str, ...] = () + + class LLMTrajectoryCapture: """Own capture initialization, sandbox instrumentation, and finalization.""" @@ -117,6 +126,7 @@ def __init__( self._collector_owned = False self._capture_root_prepared = False self._preparation_errors: list[str] = [] + self._otel_setup_error: str | None = None @property def trajectory_path(self) -> Path: @@ -208,13 +218,14 @@ async def prepare_agent( prepared.pop("CLAUDE_CODE_ENABLE_TELEMETRY", None) prepared.pop("OTEL_LOG_RAW_API_BODIES", None) warning = _sanitized_error(exc) - self._preparation_errors.append(warning) + self._otel_setup_error = warning logger.warning( "Claude OTel correlation unavailable; session fallback remains " "enabled: %s", warning, ) else: + self._otel_setup_error = None prepared.update( { "OTEL_LOGS_EXPORTER": "otlp", @@ -305,8 +316,11 @@ async def finalize( raise native_bundles: list[_NativeCaptureBundle] = [] + preparation_errors = [*self._preparation_errors] + if self._otel_setup_error is not None: + preparation_errors.append(self._otel_setup_error) collection_errors = list( - dict.fromkeys([*self._preparation_errors, *(capture_errors or [])]) + dict.fromkeys([*preparation_errors, *(capture_errors or [])]) ) native_targets = self._native_targets() native_resources_exist = self._collector_owned or self._capture_root_prepared @@ -314,7 +328,9 @@ async def finalize( if (native_targets or native_resources_exist) and env is not None: try: if native_targets: - native_bundles = await self._collect_native_results(env) + collection = await self._collect_native_results(env) + native_bundles.extend(collection.bundles) + collection_errors.extend(collection.errors) elif self._collector_owned: await self._stop_otel_sink(env) except Exception as exc: @@ -514,83 +530,157 @@ async def _read_collector_port(self, env: Any) -> int: raise RuntimeError("Claude OTel sink port file is unavailable") return _parse_port(result.stdout) - async def _collect_native_results(self, env: Any) -> list[_NativeCaptureBundle]: - if self._collector_owned: - await self._stop_otel_sink(env) + async def _collect_native_results(self, env: Any) -> _NativeCollection: bundles: list[_NativeCaptureBundle] = [] + errors: list[str] = [] + if self._collector_owned: + try: + await self._stop_otel_sink(env) + except Exception as exc: + errors.append(_sanitized_error(exc)) + logger.warning("Claude OTel collector shutdown failed: %s", exc) + native_targets = self._native_targets() claude_targets = tuple( target for target in native_targets if _is_claude_code_agent(target.agent) ) with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: local_root = Path(temporary) - capture_dir = local_root / "capture" - raw_claude_captured = False - if self._capture_root_prepared: - await env.download_dir(self._remote_capture_root, capture_dir) - result = parse_claude_raw_capture( - capture_dir, - agent=(claude_targets[0].agent if claude_targets else self.agent), - session_id=self.session_id, - started_at=self.started_at, - ) - if result is not None: - bundles.append( - _NativeCaptureBundle(targets=claude_targets, result=result) - ) - raw_claude_captured = True + raw_claude_session_ids = await self._collect_claude_raw_capture( + env, + local_root=local_root, + claude_targets=claude_targets, + bundles=bundles, + errors=errors, + ) for index, target in enumerate(native_targets): - if _is_claude_code_agent(target.agent) and not raw_claude_captured: - claude_local = local_root / f"target-{index}" / "claude-projects" - downloaded = await _download_bound_session_files( - env, - f"{target.credential_home}/.claude/projects", - claude_local, - started_at=self.started_at, - session_ids=target.native_session_ids, - ) - if not downloaded: - continue - result = parse_claude_sessions( - claude_local, - agent=target.agent, - session_id=self.session_id, - started_at=self.started_at, - ) - if result is not None: - bundles.append( - _NativeCaptureBundle( - targets=(target,), - result=result, - ) + try: + if _is_claude_code_agent(target.agent): + fallback_ids = tuple( + session_id + for session_id in target.native_session_ids + if session_id not in raw_claude_session_ids ) - elif target.agent == "codex-acp": - codex_local = local_root / f"target-{index}" / "codex-sessions" - downloaded = await _download_bound_session_files( - env, - f"{target.credential_home}/.codex/sessions", - codex_local, - started_at=self.started_at, - session_ids=target.native_session_ids, - ) - if not downloaded: - continue - result = parse_codex_sessions( - codex_local, - agent=target.agent, - session_id=self.session_id, - started_at=self.started_at, - configured_model=target.model, - auth_mode=target.auth_mode.value, - ) - if result is not None: - bundles.append( - _NativeCaptureBundle( - targets=(target,), - result=result, - ) + bundle = await self._collect_claude_session_fallback( + env, + local_root=local_root, + index=index, + target=target, + session_ids=fallback_ids, + ) + elif target.agent == "codex-acp": + bundle = await self._collect_codex_session( + env, + local_root=local_root, + index=index, + target=target, ) - return bundles + else: + bundle = None + if bundle is not None: + bundles.append(bundle) + except Exception as exc: + warning = ( + f"native capture failed for role {target.role}: " + f"{_sanitized_error(exc)}" + ) + errors.append(warning) + logger.warning("%s", warning) + return _NativeCollection(bundles=tuple(bundles), errors=tuple(errors)) + + async def _collect_claude_raw_capture( + self, + env: Any, + *, + local_root: Path, + claude_targets: tuple[_CaptureTarget, ...], + bundles: list[_NativeCaptureBundle], + errors: list[str], + ) -> set[str]: + if not self._capture_root_prepared: + return set() + capture_dir = local_root / "capture" + try: + await env.download_dir(self._remote_capture_root, capture_dir) + result = parse_claude_raw_capture( + capture_dir, + agent=(claude_targets[0].agent if claude_targets else self.agent), + session_id=self.session_id, + started_at=self.started_at, + ) + except Exception as exc: + errors.append(_sanitized_error(exc)) + logger.warning("Claude raw LLM capture collection failed: %s", exc) + return set() + if result is None: + return set() + bundles.append(_NativeCaptureBundle(targets=claude_targets, result=result)) + return _native_session_ids(result) + + async def _collect_claude_session_fallback( + self, + env: Any, + *, + local_root: Path, + index: int, + target: _CaptureTarget, + session_ids: tuple[str, ...], + ) -> _NativeCaptureBundle | None: + if not session_ids: + return None + local = local_root / f"target-{index}" / "claude-projects" + downloaded = await _download_bound_session_files( + env, + f"{target.credential_home}/.claude/projects", + local, + started_at=self.started_at, + session_ids=session_ids, + ) + if not downloaded: + return None + result = parse_claude_sessions( + local, + agent=target.agent, + session_id=self.session_id, + started_at=self.started_at, + ) + return ( + _NativeCaptureBundle(targets=(target,), result=result) + if result is not None + else None + ) + + async def _collect_codex_session( + self, + env: Any, + *, + local_root: Path, + index: int, + target: _CaptureTarget, + ) -> _NativeCaptureBundle | None: + local = local_root / f"target-{index}" / "codex-sessions" + downloaded = await _download_bound_session_files( + env, + f"{target.credential_home}/.codex/sessions", + local, + started_at=self.started_at, + session_ids=target.native_session_ids, + ) + if not downloaded: + return None + result = parse_codex_sessions( + local, + agent=target.agent, + session_id=self.session_id, + started_at=self.started_at, + configured_model=target.model, + auth_mode=target.auth_mode.value, + ) + return ( + _NativeCaptureBundle(targets=(target,), result=result) + if result is not None + else None + ) async def _cleanup_remote_capture(self, env: Any) -> None: if not self._capture_root_prepared: @@ -702,6 +792,16 @@ def _capture_target_key( return (role_name or "primary", agent, model, credential_home) +def _native_session_ids(result: NativeParseResult) -> set[str]: + """Return exact ACP session IDs represented by parsed native exchanges.""" + + return { + value + for exchange in result.trajectory.exchanges + if isinstance((value := exchange.metadata.get("native_session_id")), str) + } + + def _is_claude_code_agent(agent: str) -> bool: config = AGENTS.get(agent) subscription = config.subscription_auth if config is not None else None diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index d57b00266..7a5c25341 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -14,6 +14,7 @@ LLMTrajectoryCapture, _CaptureTarget, _NativeCaptureBundle, + _NativeCollection, ) from benchflow.trajectories.llm_capture_manifest import ( AuthMode, @@ -732,7 +733,9 @@ async def test_mixed_auth_rollout_merges_provider_and_native_exchanges( assert native_result is not None async def collect_native(_env): - return [_NativeCaptureBundle((native_target,), native_result)] + return _NativeCollection( + bundles=(_NativeCaptureBundle((native_target,), native_result),) + ) capture._collect_native_results = collect_native @@ -801,7 +804,7 @@ async def test_mixed_auth_rollout_marks_missing_native_role_partial( ) async def collect_native(_env): - return [] + return _NativeCollection() capture._collect_native_results = collect_native diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 036667a26..7f7c52dac 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -11,11 +11,28 @@ import pytest from benchflow.rollout import Rollout +from benchflow.trajectories import llm_capture as llm_capture_module from benchflow.trajectories.llm_capture import ( LLMTrajectoryCapture, + _CaptureTarget, _download_bound_session_files, + _NativeCaptureBundle, +) +from benchflow.trajectories.llm_capture_manifest import ( + AuthMode, + CaptureFidelity, + CaptureSource, +) +from benchflow.trajectories.native_capture_parsers import ( + NativeParseResult, + parse_codex_sessions, +) +from benchflow.trajectories.types import ( + LLMExchange, + LLMRequest, + LLMResponse, + Trajectory, ) -from benchflow.trajectories.native_capture_parsers import parse_codex_sessions STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) @@ -56,6 +73,41 @@ def _codex_session(prompt: str, answer: str, second: int) -> list[dict]: ] +def _native_result( + *, + session_id: str, + source: CaptureSource, + fidelity: CaptureFidelity, + model: str, +) -> NativeParseResult: + exchange = LLMExchange( + request=LLMRequest(body={"model": model, "input": "hello"}), + response=LLMResponse(body={"output": []}), + metadata={ + "capture_source": source.value, + "capture_fidelity": fidelity.value, + "auth_mode": AuthMode.OAUTH_SUBSCRIPTION.value, + "native_session_id": session_id, + "request_complete": fidelity is CaptureFidelity.PROVIDER_WIRE, + "response_complete": True, + }, + ) + return NativeParseResult( + trajectory=Trajectory( + session_id="rollout-1", + agent_name="agent", + exchanges=[exchange], + ), + source=source, + fidelity=fidelity, + request_complete=fidelity is CaptureFidelity.PROVIDER_WIRE, + response_complete=True, + missing_fields=( + [] if fidelity is CaptureFidelity.PROVIDER_WIRE else ["request"] + ), + ) + + @pytest.mark.asyncio async def test_phase_api_registers_provisional_primary_oauth_target( tmp_path: Path, @@ -263,6 +315,64 @@ async def download_dir(self, _remote, local): assert any("-depth -delete" in command for command in commands) +@pytest.mark.asyncio +async def test_claude_capture_setup_error_is_cleared_after_recovery( + tmp_path: Path, +) -> None: + """Guards PR #1057 against retaining a recovered setup failure.""" + + class RecoveringCollectorEnv: + def __init__(self) -> None: + self.launch_attempts = 0 + + async def exec(self, command, **_kwargs): + if "nohup" not in command: + return SimpleNamespace(return_code=0, stdout="", stderr="") + self.launch_attempts += 1 + if self.launch_attempts == 1: + return SimpleNamespace( + return_code=1, + stdout="", + stderr="collector port was not observed", + ) + return SimpleNamespace(return_code=0, stdout="43123\n", stderr="") + + async def upload_file(self, *_args, **_kwargs): + return None + + env = RecoveringCollectorEnv() + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + first_env = await capture.prepare_agent( + env, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "test-token"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + assert "OTEL_LOGS_EXPORTER" not in first_env + assert capture._otel_setup_error is not None + + recovered_env = await capture.prepare_agent( + env, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "test-token"}, + credential_home="/home/agent", + sandbox_user="agent", + role_name="solver", + ) + + assert recovered_env["OTEL_LOGS_EXPORTER"] == "otlp" + assert capture._otel_setup_error is None + + @pytest.mark.asyncio async def test_native_download_selects_recent_files_before_copying( tmp_path: Path, @@ -398,6 +508,160 @@ async def test_native_role_reprepare_preserves_prior_session_ids( ) +@pytest.mark.asyncio +async def test_claude_fallback_collects_only_raw_uncovered_sessions( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Guards PR #1057's per-session Claude raw fallback coverage.""" + + first_id = "019effaf-3966-75d3-b61a-2916c84b0ac8" + second_id = "019effaf-3ab7-71f1-8ff3-fdecf66b551e" + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + targets = ( + _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + role="solver", + native_session_ids=(first_id,), + ), + _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + role="reviewer", + native_session_ids=(second_id,), + ), + ) + for target in targets: + capture._targets[ + (target.role, target.agent, target.model, target.credential_home) + ] = target + capture._capture_root_prepared = True + raw_result = _native_result( + session_id=first_id, + source=CaptureSource.CLAUDE_OTEL_RAW_BODY, + fidelity=CaptureFidelity.PROVIDER_WIRE, + model="claude-sonnet-4-6", + ) + fallback_result = _native_result( + session_id=second_id, + source=CaptureSource.CLAUDE_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + model="claude-sonnet-4-6", + ) + fallback_calls: list[tuple[str, ...]] = [] + + class CaptureEnv: + async def download_dir(self, _remote, local): + Path(local).mkdir(parents=True) + + async def download_bound(_env, _remote, local, *, started_at, session_ids): + del started_at + fallback_calls.append(session_ids) + Path(local).mkdir(parents=True) + return True + + monkeypatch.setattr( + llm_capture_module, + "parse_claude_raw_capture", + lambda *_args, **_kwargs: raw_result, + ) + monkeypatch.setattr( + llm_capture_module, + "_download_bound_session_files", + download_bound, + ) + monkeypatch.setattr( + llm_capture_module, + "parse_claude_sessions", + lambda *_args, **_kwargs: fallback_result, + ) + + collection = await capture._collect_native_results(CaptureEnv()) + + assert len(collection.bundles) == 2 + assert fallback_calls == [(second_id,)] + assert collection.errors == () + + +@pytest.mark.asyncio +async def test_later_native_target_failure_preserves_prior_bundle( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Guards PR #1057 against discarding earlier native role evidence.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-one", + session_id="rollout-1", + started_at=STARTED_AT, + ) + first = _CaptureTarget( + agent="codex-acp", + model="gpt-one", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + role="solver", + native_session_ids=("session-one",), + ) + second = _CaptureTarget( + agent="codex-acp", + model="gpt-two", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + role="reviewer", + native_session_ids=("session-two",), + ) + for target in (first, second): + capture._targets[ + (target.role, target.agent, target.model, target.credential_home) + ] = target + first_result = _native_result( + session_id="session-one", + source=CaptureSource.CODEX_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + model="gpt-one", + ) + + async def collect_codex(_env, *, local_root, index, target): + del local_root, index + if target is second: + raise RuntimeError("second target download failed") + return _NativeCaptureBundle(targets=(first,), result=first_result) + + monkeypatch.setattr(capture, "_collect_codex_session", collect_codex) + + await capture.finalize(object(), acp_events=[], model_call_seen=True) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + rows = [ + json.loads(line) for line in capture.trajectory_path.read_text().splitlines() + ] + assert len(rows) == 1 + assert rows[0]["metadata"]["role"] == "solver" + assert manifest["exchange_count"] == 1 + assert manifest["status"] == "partial" + assert any("role reviewer" in error for error in manifest["errors"]) + + @pytest.mark.asyncio async def test_replaced_native_target_still_releases_collector(tmp_path: Path) -> None: """Guards PR #1057 against leaking a replaced OAuth collector.""" From c1abd5249caba358eaf9a2c7593f243d39bdbaf9 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 02:28:22 -0700 Subject: [PATCH 16/74] fix: reconcile partial native captures --- src/benchflow/trajectories/llm_capture.py | 110 +++++++++------- .../trajectories/native_capture_parsers.py | 69 +++++++++- .../test_native_capture_resilience.py | 118 ++++++++++++++++++ .../test_native_session_boundaries.py | 99 +++++++++------ 4 files changed, 308 insertions(+), 88 deletions(-) create mode 100644 tests/trajectories/test_native_capture_resilience.py diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index edb7eefa8..d08a8a0e8 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -46,6 +46,7 @@ parse_claude_sessions, parse_codex_sessions, project_acp_trajectory, + retain_uncovered_claude_session_exchanges, ) from benchflow.trajectories.types import redact_trajectory_text @@ -312,6 +313,8 @@ async def finalize( ) except Exception: if env is not None: + if self._collector_owned: + await self._stop_otel_sink(env) await self._cleanup_remote_capture(env) raise @@ -510,6 +513,18 @@ async def _stop_otel_sink(self, env: Any) -> None: fi sleep 0.05 done +old_command=$(ps -p "$old_pid" -o command= 2>/dev/null || true) +case "$old_command" in + *{self._remote_capture_root}/otel_sink.mjs*) ;; + *) exit 0 ;; +esac +kill -KILL "$old_pid" 2>/dev/null || true +for attempt in $(seq 1 20); do + if ! kill -0 "$old_pid" 2>/dev/null; then + exit 0 + fi + sleep 0.05 +done echo "previous Claude telemetry collector did not stop" >&2 exit 1 """ @@ -546,7 +561,7 @@ async def _collect_native_results(self, env: Any) -> _NativeCollection: ) with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: local_root = Path(temporary) - raw_claude_session_ids = await self._collect_claude_raw_capture( + raw_claude_result = await self._collect_claude_raw_capture( env, local_root=local_root, claude_targets=claude_targets, @@ -556,18 +571,15 @@ async def _collect_native_results(self, env: Any) -> _NativeCollection: for index, target in enumerate(native_targets): try: if _is_claude_code_agent(target.agent): - fallback_ids = tuple( - session_id - for session_id in target.native_session_ids - if session_id not in raw_claude_session_ids - ) - bundle = await self._collect_claude_session_fallback( + target_bundles = await self._collect_claude_session_fallback( env, local_root=local_root, index=index, target=target, - session_ids=fallback_ids, + raw_result=raw_claude_result, ) + bundles.extend(target_bundles) + bundle = None elif target.agent == "codex-acp": bundle = await self._collect_codex_session( env, @@ -596,9 +608,9 @@ async def _collect_claude_raw_capture( claude_targets: tuple[_CaptureTarget, ...], bundles: list[_NativeCaptureBundle], errors: list[str], - ) -> set[str]: + ) -> NativeParseResult | None: if not self._capture_root_prepared: - return set() + return None capture_dir = local_root / "capture" try: await env.download_dir(self._remote_capture_root, capture_dir) @@ -611,11 +623,11 @@ async def _collect_claude_raw_capture( except Exception as exc: errors.append(_sanitized_error(exc)) logger.warning("Claude raw LLM capture collection failed: %s", exc) - return set() + return None if result is None: - return set() + return None bundles.append(_NativeCaptureBundle(targets=claude_targets, result=result)) - return _native_session_ids(result) + return result async def _collect_claude_session_fallback( self, @@ -624,31 +636,38 @@ async def _collect_claude_session_fallback( local_root: Path, index: int, target: _CaptureTarget, - session_ids: tuple[str, ...], - ) -> _NativeCaptureBundle | None: - if not session_ids: - return None - local = local_root / f"target-{index}" / "claude-projects" - downloaded = await _download_bound_session_files( - env, - f"{target.credential_home}/.claude/projects", - local, - started_at=self.started_at, - session_ids=session_ids, - ) - if not downloaded: - return None - result = parse_claude_sessions( - local, - agent=target.agent, - session_id=self.session_id, - started_at=self.started_at, - ) - return ( - _NativeCaptureBundle(targets=(target,), result=result) - if result is not None - else None - ) + raw_result: NativeParseResult | None, + ) -> tuple[_NativeCaptureBundle, ...]: + bundles: list[_NativeCaptureBundle] = [] + for session_index, native_session_id in enumerate(target.native_session_ids): + local = local_root / f"target-{index}" / f"claude-session-{session_index}" + downloaded = await _download_bound_session_files( + env, + f"{target.credential_home}/.claude/projects", + local, + started_at=self.started_at, + session_ids=(native_session_id,), + ) + if not downloaded: + continue + result = parse_claude_sessions( + local, + agent=target.agent, + session_id=self.session_id, + started_at=self.started_at, + ) + if result is None: + continue + uncovered = retain_uncovered_claude_session_exchanges( + raw_result, + result, + native_session_id=native_session_id, + ) + if uncovered is not None: + bundles.append( + _NativeCaptureBundle(targets=(target,), result=uncovered) + ) + return tuple(bundles) async def _collect_codex_session( self, @@ -685,6 +704,11 @@ async def _collect_codex_session( async def _cleanup_remote_capture(self, env: Any) -> None: if not self._capture_root_prepared: return + if self._collector_owned: + raise RuntimeError( + "Refusing to remove Claude capture ownership files while its " + "collector may still be running" + ) result = await env.exec( "for attempt in 1 2 3; do\n" f" if ! test -e {self._remote_capture_root} || " @@ -792,16 +816,6 @@ def _capture_target_key( return (role_name or "primary", agent, model, credential_home) -def _native_session_ids(result: NativeParseResult) -> set[str]: - """Return exact ACP session IDs represented by parsed native exchanges.""" - - return { - value - for exchange in result.trajectory.exchanges - if isinstance((value := exchange.metadata.get("native_session_id")), str) - } - - def _is_claude_code_agent(agent: str) -> bool: config = AGENTS.get(agent) subscription = config.subscription_auth if config is not None else None diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index 6e3c0fd8d..fde5a92dc 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -207,6 +207,48 @@ def parse_claude_sessions( ) +def retain_uncovered_claude_session_exchanges( + raw_result: NativeParseResult | None, + session_result: NativeParseResult, + *, + native_session_id: str, +) -> NativeParseResult | None: + """Keep transcript exchanges not already represented by provider-wire capture.""" + + raw_identities = { + identity + for exchange in ( + raw_result.trajectory.exchanges if raw_result is not None else [] + ) + if exchange.metadata.get("native_session_id") == native_session_id + for identity in _claude_exchange_identities(exchange) + } + uncovered: list[LLMExchange] = [] + for exchange in session_result.trajectory.exchanges: + exchange.metadata["native_session_id"] = native_session_id + identities = _claude_exchange_identities(exchange) + if identities and identities.intersection(raw_identities): + continue + uncovered.append(exchange) + if not uncovered: + return None + trajectory = session_result.trajectory.model_copy( + update={ + "exchanges": uncovered, + "finished_at": max(exchange.response.timestamp for exchange in uncovered), + } + ) + return NativeParseResult( + trajectory=trajectory, + source=session_result.source, + fidelity=session_result.fidelity, + request_complete=session_result.request_complete, + response_complete=session_result.response_complete, + missing_fields=list(session_result.missing_fields), + errors=list(session_result.errors), + ) + + def _parse_claude_session_records( records: list[dict[str, Any]], *, started_at: datetime ) -> list[LLMExchange]: @@ -498,6 +540,8 @@ def _exchange( response_complete: bool, extra_metadata: dict[str, Any] | None = None, ) -> LLMExchange: + request_timestamp = _utc_timestamp(request_timestamp) + response_timestamp = _utc_timestamp(response_timestamp) metadata = { "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, "capture_source": source.value, @@ -767,7 +811,7 @@ def _normalize_codex_usage(usage: dict[str, Any]) -> dict[str, Any]: def _record_timestamp(record: dict[str, Any], default: datetime) -> datetime: - return _optional_record_timestamp(record) or default + return _optional_record_timestamp(record) or _utc_timestamp(default) def _optional_record_timestamp(record: dict[str, Any]) -> datetime | None: @@ -775,11 +819,32 @@ def _optional_record_timestamp(record: dict[str, Any]) -> datetime | None: if not isinstance(value, str): return None try: - return datetime.fromisoformat(value.replace("Z", "+00:00")) + return _utc_timestamp(datetime.fromisoformat(value.replace("Z", "+00:00"))) except ValueError: return None +def _utc_timestamp(value: datetime) -> datetime: + """Normalize native timestamps, interpreting legacy naive values as UTC.""" + + if value.tzinfo is None: + return value.replace(tzinfo=UTC) + return value.astimezone(UTC) + + +def _claude_exchange_identities(exchange: LLMExchange) -> set[tuple[str, str]]: + """Return exact provider identifiers shared by raw and transcript records.""" + + identities: set[tuple[str, str]] = set() + request_id = exchange.metadata.get("provider_request_id") + if isinstance(request_id, str) and request_id: + identities.add(("request", request_id)) + response_id = exchange.response.body.get("id") + if isinstance(response_id, str) and response_id: + identities.add(("response", response_id)) + return identities + + def _unix_nanos_timestamp(value: Any) -> datetime: try: return datetime.fromtimestamp(int(value) / 1_000_000_000, tz=UTC) diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py new file mode 100644 index 000000000..8fc96459c --- /dev/null +++ b/tests/trajectories/test_native_capture_resilience.py @@ -0,0 +1,118 @@ +"""Failure-boundary regressions for native LLM trajectory capture.""" + +from __future__ import annotations + +import json +from datetime import UTC, datetime +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from benchflow.trajectories.llm_capture import LLMTrajectoryCapture, _CaptureTarget +from benchflow.trajectories.llm_capture_manifest import AuthMode +from benchflow.trajectories.native_capture_parsers import parse_codex_sessions + + +def test_native_parser_normalizes_fallback_and_record_timestamps( + tmp_path: Path, +) -> None: + """Guards PR #1057 against mixing naive fallback and aware record times.""" + + session = tmp_path / "sessions" / "rollout-session-one.jsonl" + session.parent.mkdir(parents=True) + session.write_text( + "\n".join( + json.dumps(record) + for record in ( + { + "type": "response_item", + "payload": { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "done"}], + }, + }, + { + "timestamp": "2026-08-29T12:00:01Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": {}}, + }, + ) + ) + + "\n" + ) + + result = parse_codex_sessions( + session.parent, + agent="codex-acp", + session_id="rollout-1", + started_at=datetime(2020, 1, 1), + configured_model="gpt-5.6", + ) + + assert result is not None + exchange = result.trajectory.exchanges[0] + assert exchange.request.timestamp.tzinfo is UTC + assert exchange.response.timestamp.tzinfo is UTC + assert exchange.duration_ms > 0 + + +@pytest.mark.asyncio +async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( + tmp_path: Path, +) -> None: + """Guards PR #1057 against leaking OTel on malformed mixed capture input.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="provider-agent", + model="provider-model", + session_id="rollout-1", + started_at=datetime(2026, 8, 29, 12, 0, tzinfo=UTC), + ) + provider = _CaptureTarget( + agent="provider-agent", + model="provider-model", + credential_home="/home/agent", + auth_mode=AuthMode.API_KEY, + native=False, + role="solver", + ) + native = _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + role="reviewer", + native_session_ids=("session-one",), + ) + for target in (provider, native): + capture._targets[ + (target.role, target.agent, target.model, target.credential_home) + ] = target + capture.trajectory_path.write_text("{malformed provider row\n") + capture._collector_owned = True + capture._capture_root_prepared = True + commands: list[str] = [] + + class RecordingEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + return SimpleNamespace(return_code=0, stdout="", stderr="") + + with pytest.raises(ValueError, match="invalid LLM trajectory JSONL"): + await capture.finalize( + RecordingEnv(), + acp_events=[], + model_call_seen=True, + ) + + stop_index = next(i for i, command in enumerate(commands) if "old_pid" in command) + cleanup_index = next( + i for i, command in enumerate(commands) if "for attempt in 1 2 3" in command + ) + assert stop_index < cleanup_index + assert capture._collector_owned is False + assert capture._capture_root_prepared is False diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 7f7c52dac..ffa2dff2b 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from dataclasses import replace from datetime import UTC, datetime from pathlib import Path from types import SimpleNamespace @@ -79,18 +80,26 @@ def _native_result( source: CaptureSource, fidelity: CaptureFidelity, model: str, + request_id: str | None = None, + response_id: str | None = None, ) -> NativeParseResult: + response_body: dict[str, object] = {"output": []} + if response_id is not None: + response_body["id"] = response_id + metadata: dict[str, object] = { + "capture_source": source.value, + "capture_fidelity": fidelity.value, + "auth_mode": AuthMode.OAUTH_SUBSCRIPTION.value, + "native_session_id": session_id, + "request_complete": fidelity is CaptureFidelity.PROVIDER_WIRE, + "response_complete": True, + } + if request_id is not None: + metadata["provider_request_id"] = request_id exchange = LLMExchange( request=LLMRequest(body={"model": model, "input": "hello"}), - response=LLMResponse(body={"output": []}), - metadata={ - "capture_source": source.value, - "capture_fidelity": fidelity.value, - "auth_mode": AuthMode.OAUTH_SUBSCRIPTION.value, - "native_session_id": session_id, - "request_complete": fidelity is CaptureFidelity.PROVIDER_WIRE, - "response_complete": True, - }, + response=LLMResponse(body=response_body), + metadata=metadata, ) return NativeParseResult( trajectory=Trajectory( @@ -509,14 +518,13 @@ async def test_native_role_reprepare_preserves_prior_session_ids( @pytest.mark.asyncio -async def test_claude_fallback_collects_only_raw_uncovered_sessions( +async def test_claude_fallback_collects_only_raw_uncovered_exchanges( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - """Guards PR #1057's per-session Claude raw fallback coverage.""" + """Guards PR #1057 against losing a partially raw-captured Claude session.""" first_id = "019effaf-3966-75d3-b61a-2916c84b0ac8" - second_id = "019effaf-3ab7-71f1-8ff3-fdecf66b551e" capture = LLMTrajectoryCapture( tmp_path, agent="claude-agent-acp", @@ -524,42 +532,53 @@ async def test_claude_fallback_collects_only_raw_uncovered_sessions( session_id="rollout-1", started_at=STARTED_AT, ) - targets = ( - _CaptureTarget( - agent="claude-agent-acp", - model="claude-sonnet-4-6", - credential_home="/home/agent", - auth_mode=AuthMode.OAUTH_SUBSCRIPTION, - native=True, - role="solver", - native_session_ids=(first_id,), - ), - _CaptureTarget( - agent="claude-agent-acp", - model="claude-sonnet-4-6", - credential_home="/home/agent", - auth_mode=AuthMode.OAUTH_SUBSCRIPTION, - native=True, - role="reviewer", - native_session_ids=(second_id,), - ), + target = _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + role="solver", + native_session_ids=(first_id,), ) - for target in targets: - capture._targets[ - (target.role, target.agent, target.model, target.credential_home) - ] = target + capture._targets[ + (target.role, target.agent, target.model, target.credential_home) + ] = target capture._capture_root_prepared = True raw_result = _native_result( session_id=first_id, source=CaptureSource.CLAUDE_OTEL_RAW_BODY, fidelity=CaptureFidelity.PROVIDER_WIRE, model="claude-sonnet-4-6", + request_id="request-covered", + response_id="message-covered", + ) + covered_fallback = _native_result( + session_id=first_id, + source=CaptureSource.CLAUDE_NATIVE_SESSION, + fidelity=CaptureFidelity.AGENT_SESSION, + model="claude-sonnet-4-6", + request_id="request-covered", + response_id="message-covered", ) - fallback_result = _native_result( - session_id=second_id, + missing_fallback = _native_result( + session_id=first_id, source=CaptureSource.CLAUDE_NATIVE_SESSION, fidelity=CaptureFidelity.AGENT_SESSION, model="claude-sonnet-4-6", + request_id="request-missing", + response_id="message-missing", + ) + fallback_result = replace( + covered_fallback, + trajectory=covered_fallback.trajectory.model_copy( + update={ + "exchanges": [ + *covered_fallback.trajectory.exchanges, + *missing_fallback.trajectory.exchanges, + ] + } + ), ) fallback_calls: list[tuple[str, ...]] = [] @@ -592,7 +611,11 @@ async def download_bound(_env, _remote, local, *, started_at, session_ids): collection = await capture._collect_native_results(CaptureEnv()) assert len(collection.bundles) == 2 - assert fallback_calls == [(second_id,)] + assert fallback_calls == [(first_id,)] + fallback_exchanges = collection.bundles[1].result.trajectory.exchanges + assert len(fallback_exchanges) == 1 + assert fallback_exchanges[0].metadata["provider_request_id"] == "request-missing" + assert fallback_exchanges[0].metadata["native_session_id"] == first_id assert collection.errors == () From fa6cbf5803d003902cb9496a737923b731bd774b Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 02:38:40 -0700 Subject: [PATCH 17/74] fix: reject empty legacy trajectory artifacts --- src/benchflow/eval_artifacts.py | 2 ++ tests/test_eval_artifact_cli.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index c4205e2d1..1f9f77fba 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -206,6 +206,8 @@ def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, int]: rows = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError: return True, False, 0 + if not rows: + return True, False, 0 manifest = read_llm_trajectory_manifest(rollout_dir) if manifest is not None and not capture_manifest_allows_training( manifest, exchange_count=len(rows) diff --git a/tests/test_eval_artifact_cli.py b/tests/test_eval_artifact_cli.py index 641c874d2..ead76e04b 100644 --- a/tests/test_eval_artifact_cli.py +++ b/tests/test_eval_artifact_cli.py @@ -110,6 +110,23 @@ def test_health_rejects_empty_terminal_llm_capture_manifest(tmp_path: Path) -> N assert health["rows"][0]["llm_trajectory_rows"] == 0 +def test_health_rejects_empty_llm_capture_without_manifest(tmp_path: Path) -> None: + """Guards PR #1057 against accepting an interrupted empty capture artifact.""" + + job = tmp_path / "job" + rollout = job / "task-a__abc" + _write_rollout(rollout) + (rollout / "trajectory" / "llm_trajectory.jsonl").write_text("") + + health = build_health_summary(job) + + assert health["missing_llm_trajectory"] == 0 + assert health["malformed_llm_trajectory"] == 1 + assert health["rows"][0]["has_llm_trajectory"] is True + assert health["rows"][0]["valid_llm_trajectory"] is False + assert health["rows"][0]["llm_trajectory_rows"] == 0 + + def test_eval_run_writes_manifest_health_and_canonical_artifacts( tmp_path: Path, monkeypatch ) -> None: From 72f478d4429b8574217e8521fbe1173f26c859b0 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 02:54:08 -0700 Subject: [PATCH 18/74] fix: attribute translated provider model aliases --- .../trajectories/llm_capture_records.py | 54 +++++++++++-- .../test_native_capture_resilience.py | 80 +++++++++++++++++++ 2 files changed, 126 insertions(+), 8 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 011cd7491..23e85a4f4 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import Any +from benchflow.providers.litellm_config import safe_model_alias from benchflow.trajectories.llm_capture_manifest import ( LLM_TRAJECTORY_SCHEMA_VERSION, AuthMode, @@ -85,7 +86,7 @@ def load_provider_wire_records( if not isinstance(metadata, dict): metadata = {} record["metadata"] = metadata - target = _target_for_model(targets, _record_model(record)) + target = _target_for_provider_record(targets, record) attribution_complete = target is not None or not targets metadata.update( { @@ -319,17 +320,27 @@ def _captured_targets( def _model_matches_target(value: Any, configured_model: str | None) -> bool: + return _model_match_score(value, configured_model) > 0 + + +def _model_match_score(value: Any, configured_model: str | None) -> int: if configured_model is None: - return True + return 1 if not isinstance(value, str): - return False + return 0 normalized_value = value.casefold() normalized_target = configured_model.casefold() - return bool( - normalized_value == normalized_target - or normalized_value.endswith(f"/{normalized_target}") - or normalized_target.endswith(f"/{normalized_value}") - ) + proxy_alias = safe_model_alias(configured_model).casefold() + if normalized_value == normalized_target or normalized_value in { + proxy_alias, + f"openai/{proxy_alias}", + }: + return 2 + if normalized_value.endswith(f"/{normalized_target}") or normalized_target.endswith( + f"/{normalized_value}" + ): + return 1 + return 0 def _record_model(record: dict[str, Any]) -> str | None: @@ -354,6 +365,33 @@ def _target_for_model( return matches[0] if len(matches) == 1 else None +def _target_for_provider_record( + targets: list[CaptureTarget], record: dict[str, Any] +) -> CaptureTarget | None: + if len(targets) == 1: + return targets[0] + candidates = {_record_model(record)} + metadata = _record_metadata(record) + candidates.update( + value + for key in ("model_group", "request_model", "provider_model") + if isinstance((value := metadata.get(key)), str) and value + ) + scored = [ + ( + max( + _model_match_score(candidate, target.model) for candidate in candidates + ), + target, + ) + for target in targets + if target.model + ] + best_score = max((score for score, _target in scored), default=0) + matches = [target for score, target in scored if score == best_score > 0] + return matches[0] if len(matches) == 1 else None + + def _target_for_record( targets: list[CaptureTarget], record: dict[str, Any] ) -> CaptureTarget | None: diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 8fc96459c..c9c2ab795 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -11,6 +11,7 @@ from benchflow.trajectories.llm_capture import LLMTrajectoryCapture, _CaptureTarget from benchflow.trajectories.llm_capture_manifest import AuthMode +from benchflow.trajectories.llm_capture_records import load_provider_wire_records from benchflow.trajectories.native_capture_parsers import parse_codex_sessions @@ -58,6 +59,85 @@ def test_native_parser_normalizes_fallback_and_record_timestamps( assert exchange.duration_ms > 0 +def test_provider_role_attribution_uses_proxy_model_aliases(tmp_path: Path) -> None: + """Guards PR #1057 against losing roles after LiteLLM model translation.""" + + targets = [ + _CaptureTarget( + agent="claude-agent-acp", + model="aws-bedrock/us.anthropic.claude-opus-4-8", + credential_home="/home/solver", + auth_mode=AuthMode.API_KEY, + native=False, + role="solver", + ), + _CaptureTarget( + agent="codex-acp", + model="azure-foundry-openai/gpt-5.5", + credential_home="/home/reviewer", + auth_mode=AuthMode.API_KEY, + native=False, + role="reviewer", + ), + _CaptureTarget( + agent="opencode", + model="openai/gpt-5.5", + credential_home="/home/critic", + auth_mode=AuthMode.API_KEY, + native=False, + role="critic", + ), + ] + trajectory = tmp_path / "llm_trajectory.jsonl" + trajectory.write_text( + "".join( + json.dumps(record) + "\n" + for record in ( + { + "request": { + "body": {"model": "bedrock/us.anthropic.claude-opus-4-8"} + }, + "response": {"status_code": 200, "body": {}}, + "metadata": { + "model_group": ( + "benchflow-aws-bedrock-us.anthropic.claude-opus-4-8" + ) + }, + }, + { + "request": {"body": {"model": "gpt-5.5"}}, + "response": {"status_code": 200, "body": {}}, + "metadata": { + "request_model": "benchflow-azure-foundry-openai-gpt-5.5" + }, + }, + { + "request": {"body": {"model": "gpt-5.5"}}, + "response": {"status_code": 200, "body": {}}, + "metadata": {"request_model": "benchflow-openai-gpt-5.5"}, + }, + ) + ) + ) + + records = load_provider_wire_records( + trajectory, + targets=targets, + fallback_agent="mixed", + fallback_model=None, + fallback_auth=AuthMode.API_KEY, + ) + + assert [record["metadata"]["role"] for record in records] == [ + "solver", + "reviewer", + "critic", + ] + assert all( + record["metadata"]["role_attribution_complete"] is True for record in records + ) + + @pytest.mark.asyncio async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( tmp_path: Path, From ae2ccb16b287fcc064fe5d7635e563c15f1cf630 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 03:03:25 -0700 Subject: [PATCH 19/74] fix: stamp provider route identity in capture --- src/benchflow/providers/litellm_logging.py | 8 ++++++ src/benchflow/providers/litellm_runtime.py | 5 ++++ .../trajectories/llm_capture_records.py | 8 +++++- tests/test_litellm_hardening.py | 7 +++++ tests/test_litellm_logging.py | 26 +++++++++++++++++++ .../test_native_capture_resilience.py | 12 +++++++-- 6 files changed, 63 insertions(+), 3 deletions(-) diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index cb0c7864e..b8c421bcf 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -259,6 +259,12 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - request_body[key] = value request_body = {k: v for k, v in request_body.items() if v is not None} return { + "benchflow_requested_model": os.environ.get( + "BENCHFLOW_LITELLM_REQUESTED_MODEL" + ), + "benchflow_model_alias": os.environ.get( + "BENCHFLOW_LITELLM_MODEL_ALIAS" + ), "request_model": kwargs.get("model"), "provider_model": litellm_params.get("model") or kwargs.get("model"), "model_group": metadata.get("model_group") if isinstance(metadata, dict) else None, @@ -447,6 +453,8 @@ def _exchange_metadata( metadata = { key: record.get(key) for key in ( + "benchflow_requested_model", + "benchflow_model_alias", "request_model", "provider_model", "model_group", diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 8cd99a8ce..fd5a79343 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -63,6 +63,7 @@ LITELLM_VERSION_SPEC = "litellm[proxy]==1.89.0" LITELLM_SANDBOX_ROOT = "/tmp/benchflow-litellm" _CALLBACK_MODULE = "benchflow_litellm_callback" +_LITELLM_REQUESTED_MODEL_ENV = "BENCHFLOW_LITELLM_REQUESTED_MODEL" _PATCH_MODULE = "benchflow_litellm_bedrock_patch" # The proxy is an internal single-route gateway — it must never register the @@ -857,6 +858,8 @@ async def _start_host_litellm( "PYTHONPATH": f"{runtime_dir}{os.pathsep}{env.get('PYTHONPATH', '')}", "LITELLM_MASTER_KEY": master_key, "BENCHFLOW_LITELLM_LOG_PATH": str(log_path), + LITELLM_MODEL_ALIAS_ENV: route.model_alias, + _LITELLM_REQUESTED_MODEL_ENV: route.requested_model, **_PROXY_DOCS_DISABLE_ENV, } ) @@ -1178,6 +1181,8 @@ async def _start_sandbox_litellm( "PYTHONPATH": f"{runtime_dir}:{env.get('PYTHONPATH', '')}", "LITELLM_MASTER_KEY": master_key, "BENCHFLOW_LITELLM_LOG_PATH": paths["log"], + LITELLM_MODEL_ALIAS_ENV: route.model_alias, + _LITELLM_REQUESTED_MODEL_ENV: route.requested_model, **_PROXY_DOCS_DISABLE_ENV, } ) diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 23e85a4f4..7bbd8c044 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -374,7 +374,13 @@ def _target_for_provider_record( metadata = _record_metadata(record) candidates.update( value - for key in ("model_group", "request_model", "provider_model") + for key in ( + "benchflow_requested_model", + "benchflow_model_alias", + "model_group", + "request_model", + "provider_model", + ) if isinstance((value := metadata.get(key)), str) and value ) scored = [ diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index c39ee1e82..9cdb1da5f 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -362,6 +362,13 @@ async def test_sandbox_litellm_launch_keeps_secrets_off_command_line(): launch_files = [k for k in sandbox.uploaded if k.endswith("launch_config.json")] assert launch_files, "launch_config.json should be uploaded" assert secret in sandbox.uploaded[launch_files[0]] + launch_config = json.loads(sandbox.uploaded[launch_files[0]]) + assert launch_config["env"]["BENCHFLOW_LITELLM_REQUESTED_MODEL"] == ( + "minimax/MiniMax-M3" + ) + assert launch_config["env"]["BENCHFLOW_LITELLM_MODEL_ALIAS"] == ( + "benchflow-minimax-MiniMax-M3" + ) assert set(sandbox.uploaded_modes.values()) == {"600"} # ...and the secret never appears on any exec command line (/proc exposure). assert all(secret not in call for call in sandbox.exec_calls) diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 298673c8d..c622c226e 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -112,6 +112,32 @@ def test_callback_record_preserves_logprob_request_fields(): assert record["request"]["body"]["top_logprobs"] == 1 +def test_callback_record_preserves_benchflow_route_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Guards PR #1057 against LiteLLM erasing provider-role identity.""" + + monkeypatch.setenv( + "BENCHFLOW_LITELLM_REQUESTED_MODEL", + "azure-foundry-openai/gpt-5.5", + ) + monkeypatch.setenv( + "BENCHFLOW_LITELLM_MODEL_ALIAS", + "benchflow-azure-foundry-openai-gpt-5.5", + ) + logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() + now = datetime.now() + + record = logger._base_record( + {"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}, + now, + now, + ) + + assert record["benchflow_requested_model"] == ("azure-foundry-openai/gpt-5.5") + assert record["benchflow_model_alias"] == ("benchflow-azure-foundry-openai-gpt-5.5") + + def test_callback_module_source_exposes_proxy_handler_instance(): source = callback_module_source() diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index c9c2ab795..314b339a6 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -108,13 +108,21 @@ def test_provider_role_attribution_uses_proxy_model_aliases(tmp_path: Path) -> N "request": {"body": {"model": "gpt-5.5"}}, "response": {"status_code": 200, "body": {}}, "metadata": { - "request_model": "benchflow-azure-foundry-openai-gpt-5.5" + "benchflow_requested_model": ("azure-foundry-openai/gpt-5.5"), + "benchflow_model_alias": ( + "benchflow-azure-foundry-openai-gpt-5.5" + ), + "request_model": "gpt-5.5", }, }, { "request": {"body": {"model": "gpt-5.5"}}, "response": {"status_code": 200, "body": {}}, - "metadata": {"request_model": "benchflow-openai-gpt-5.5"}, + "metadata": { + "benchflow_requested_model": "openai/gpt-5.5", + "benchflow_model_alias": "benchflow-openai-gpt-5.5", + "request_model": "gpt-5.5", + }, }, ) ) From 16d14bc3620489ce7d69506b703975b2df599d7c Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 03:14:46 -0700 Subject: [PATCH 20/74] fix: refresh stitched capture manifests --- src/benchflow/continue_run/orchestrator.py | 184 ++++++++++++++++++++- tests/continue_run/test_orchestrator.py | 128 ++++++++++++++ 2 files changed, 309 insertions(+), 3 deletions(-) diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 3b9a92310..dedde2144 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -27,6 +27,7 @@ import re from collections.abc import Awaitable, Callable from dataclasses import dataclass +from datetime import UTC, datetime from pathlib import Path from typing import Any, cast @@ -40,6 +41,17 @@ from benchflow.contracts import AgentProtocolError, SandboxStartupFailure from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.scenes import compile_scenes_to_steps +from benchflow.trajectories.llm_capture_manifest import ( + AuthMode, + CaptureFidelity, + CaptureSource, + CaptureStatus, + LLMRoleCapture, + LLMTrajectoryManifest, + capture_manifest_allows_training, + read_llm_trajectory_manifest, + write_llm_trajectory_manifest, +) from benchflow.trajectories.types import LLMExchange, redact_trajectory_text logger = logging.getLogger(__name__) @@ -274,16 +286,164 @@ def stitched_trajectory_lines( def write_stitched_trajectory( - rollout_dir: Path, original_llm_trajectory: Path, live_exchanges: list[LLMExchange] + rollout_dir: Path, + original_llm_trajectory: Path, + live_exchanges: list[LLMExchange], + *, + live_model: str | None = None, ) -> Path: """Write the stitched continuous trajectory into the new rollout folder.""" out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" out.parent.mkdir(parents=True, exist_ok=True) - lines = stitched_trajectory_lines(original_llm_trajectory, live_exchanges) - out.write_text("\n".join(lines) + ("\n" if lines else "")) + lines = stitched_trajectory_lines(original_llm_trajectory, []) + for exchange in live_exchanges: + payload = exchange.model_dump(mode="json") + metadata = payload.setdefault("metadata", {}) + metadata.update( + { + "agent": "openhands", + "role": "agent", + "model": live_model, + "auth_mode": AuthMode.API_KEY.value, + "capture_source": CaptureSource.LITELLM_PROXY.value, + "capture_fidelity": CaptureFidelity.PROVIDER_WIRE.value, + "request_complete": True, + "response_complete": True, + "role_attribution_complete": True, + "payload_redacted": True, + } + ) + lines.append(redact_trajectory_text(json.dumps(payload, default=str))) + rendered = "\n".join(lines) + ("\n" if lines else "") + temporary = out.with_suffix(out.suffix + ".tmp") + temporary.write_text(rendered) + os.replace(temporary, out) return out +def refresh_stitched_trajectory_manifest( + rollout_dir: Path, + source_rollout_dir: Path, + *, + original_model: str | None, + live_model: str | None, + n_recorded: int, + n_live: int, +) -> LLMTrajectoryManifest: + """Replace rollout-finalization provenance with the final stitched contract.""" + + current_raw = read_llm_trajectory_manifest(rollout_dir) + source_raw = read_llm_trajectory_manifest(source_rollout_dir) + try: + current = ( + LLMTrajectoryManifest.model_validate(current_raw) + if current_raw is not None + else None + ) + except ValueError: + current = None + try: + source = ( + LLMTrajectoryManifest.model_validate(source_raw) + if source_raw is not None + else None + ) + except ValueError: + source = None + + trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" + exchange_count = sum( + bool(line.strip()) for line in trajectory_path.read_text().splitlines() + ) + expected_count = n_recorded + n_live + count_matches = exchange_count == expected_count + source_allows_training = bool( + source_raw is not None + and capture_manifest_allows_training(source_raw, exchange_count=n_recorded) + ) + complete = source_allows_training and count_matches + + source_capture_source = source.capture_source if source else CaptureSource.NONE + source_fidelity = source.capture_fidelity if source else CaptureFidelity.NONE + source_auth = source.auth_mode if source else AuthMode.UNKNOWN + if n_live: + capture_source = ( + CaptureSource.LITELLM_PROXY + if source_capture_source is CaptureSource.LITELLM_PROXY + else CaptureSource.MIXED + ) + capture_fidelity = ( + CaptureFidelity.PROVIDER_WIRE + if source_fidelity is CaptureFidelity.PROVIDER_WIRE + else CaptureFidelity.MIXED + ) + auth_mode = ( + AuthMode.API_KEY if source_auth is AuthMode.API_KEY else AuthMode.MIXED + ) + else: + capture_source = source_capture_source + capture_fidelity = source_fidelity + auth_mode = source_auth + + request_complete = bool(source and source.request_complete and count_matches) + response_complete = bool(source and source.response_complete and count_matches) + errors = list(source.errors) if source and not complete else [] + missing_fields = list(source.missing_fields) if source and not complete else [] + if source is None: + errors.append("source LLM trajectory manifest is missing or malformed") + missing_fields.append("source_capture_provenance") + elif not source_allows_training: + errors.append("source LLM trajectory is not complete provider-wire capture") + if not count_matches: + errors.append( + "stitched LLM trajectory count mismatch: " + f"expected {expected_count}, found {exchange_count}" + ) + missing_fields.append("exchange_count") + + models = { + value + for value, present in ( + (original_model, n_recorded > 0), + (live_model, n_live > 0), + ) + if present and value + } + stitched_model = next(iter(models)) if len(models) == 1 else None + manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE if complete else CaptureStatus.PARTIAL, + capture_source=capture_source, + capture_fidelity=capture_fidelity, + auth_mode=auth_mode, + agent="openhands", + model=stitched_model, + session_id=current.session_id if current else rollout_dir.name, + exchange_count=exchange_count, + request_complete=request_complete, + response_complete=response_complete, + payload_redacted=source.payload_redacted if source else True, + started_at=current.started_at if current else datetime.now(UTC), + finished_at=datetime.now(UTC), + missing_fields=sorted(set(missing_fields)), + errors=errors, + role_captures=[ + LLMRoleCapture( + role="agent", + agent="openhands", + model=stitched_model, + auth_mode=auth_mode, + capture_source=capture_source, + capture_fidelity=capture_fidelity, + exchange_count=exchange_count, + request_complete=request_complete, + response_complete=response_complete, + ) + ], + ) + write_llm_trajectory_manifest(rollout_dir, manifest) + return manifest + + def _usage_int(usage: dict[str, Any], *keys: str) -> int: for key in keys: value = usage.get(key) @@ -561,6 +721,15 @@ async def continue_run( rollout_dir, run.path / "trajectory" / "llm_trajectory.jsonl", router.live_exchanges, + live_model=live_model, + ) + refresh_stitched_trajectory_manifest( + rollout_dir, + run.path, + original_model=run.model, + live_model=live_model, + n_recorded=run.n_recorded_exchanges, + n_live=len(router.live_exchanges), ) update_continued_metadata( rollout_dir, @@ -631,6 +800,15 @@ async def _write_artifacts_before_cleanup() -> None: rollout_dir, run.path / "trajectory" / "llm_trajectory.jsonl", live_exchanges, + live_model=live_model, + ) + refresh_stitched_trajectory_manifest( + rollout_dir, + run.path, + original_model=run.model, + live_model=live_model, + n_recorded=run.n_recorded_exchanges, + n_live=len(live_exchanges), ) update_continued_metadata( rollout_dir, diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index de67ade4d..c289610cd 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -12,6 +12,7 @@ build_agent_env, build_rollout_config, continued_rollout_name, + refresh_stitched_trajectory_manifest, resolve_task_path, select_proxy_mode, stitched_trajectory_lines, @@ -20,6 +21,16 @@ write_stitched_trajectory, ) from benchflow.continue_run.run_folder import RunFolderError, load_run_folder +from benchflow.trajectories.llm_capture_manifest import ( + AuthMode, + CaptureFidelity, + CaptureSource, + CaptureStatus, + LLMTrajectoryManifest, + capture_manifest_allows_training, + initialize_llm_trajectory_artifacts, + write_llm_trajectory_manifest, +) from ._helpers import completion, exchange, write_run_folder @@ -161,6 +172,123 @@ def test_write_stitched_trajectory_creates_file(tmp_path): assert len(out.read_text().strip().splitlines()) == 2 +def test_refresh_stitched_manifest_replaces_pre_stitch_finalization(tmp_path): + """Guards PR #1057 against stale continuation capture manifests.""" + + model = "openai/gpt-5.5" + source = write_run_folder( + tmp_path / "source", + exchanges=[exchange(completion(content="recorded"))], + model=model, + ) + source_manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="source", + exchange_count=1, + request_complete=True, + response_complete=True, + started_at="2026-08-29T00:00:00Z", + finished_at="2026-08-29T00:01:00Z", + ) + write_llm_trajectory_manifest(source, source_manifest) + + rollout = tmp_path / "continued" + initialize_llm_trajectory_artifacts( + rollout, + agent="openhands", + model=None, + session_id="continued", + started_at=source_manifest.finished_at, + ) + live = [exchange(completion(content="live"))] + out = write_stitched_trajectory( + rollout, + source / "trajectory" / "llm_trajectory.jsonl", + live, + live_model=model, + ) + manifest = refresh_stitched_trajectory_manifest( + rollout, + source, + original_model=model, + live_model=model, + n_recorded=1, + n_live=1, + ) + + assert manifest.status is CaptureStatus.COMPLETE + assert manifest.capture_fidelity is CaptureFidelity.PROVIDER_WIRE + assert manifest.exchange_count == 2 + assert capture_manifest_allows_training( + manifest.model_dump(mode="json"), exchange_count=2 + ) + live_row = json.loads(out.read_text().splitlines()[-1]) + assert live_row["metadata"]["capture_fidelity"] == "provider_wire" + assert live_row["metadata"]["model"] == model + + +def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path): + """Guards PR #1057 against promoting audit-only continuation prefixes.""" + + source = write_run_folder( + tmp_path / "source", + exchanges=[exchange(completion(content="recorded"))], + model="claude-sonnet-4-6", + ) + source_manifest = LLMTrajectoryManifest( + status=CaptureStatus.PARTIAL, + capture_source=CaptureSource.CLAUDE_NATIVE_SESSION, + capture_fidelity=CaptureFidelity.AGENT_SESSION, + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="source", + exchange_count=1, + request_complete=False, + response_complete=True, + started_at="2026-08-29T00:00:00Z", + finished_at="2026-08-29T00:01:00Z", + missing_fields=["provider_request"], + ) + write_llm_trajectory_manifest(source, source_manifest) + + rollout = tmp_path / "continued" + initialize_llm_trajectory_artifacts( + rollout, + agent="openhands", + model=None, + session_id="continued", + started_at=source_manifest.finished_at, + ) + write_stitched_trajectory( + rollout, + source / "trajectory" / "llm_trajectory.jsonl", + [exchange(completion(content="live"))], + live_model="openai/gpt-5.5", + ) + manifest = refresh_stitched_trajectory_manifest( + rollout, + source, + original_model="claude-sonnet-4-6", + live_model="openai/gpt-5.5", + n_recorded=1, + n_live=1, + ) + + assert manifest.status is CaptureStatus.PARTIAL + assert manifest.capture_source is CaptureSource.MIXED + assert manifest.capture_fidelity is CaptureFidelity.MIXED + assert manifest.auth_mode is AuthMode.MIXED + assert not capture_manifest_allows_training( + manifest.model_dump(mode="json"), exchange_count=2 + ) + + def test_summarize_llm_trajectory_usage_splits_recorded_and_live(tmp_path): """Guards the PR #648 continuation metadata fix for stitched token usage.""" traj = tmp_path / "llm_trajectory.jsonl" From 93bfdddcf7e16f0c31fbf9c325ae3388b970409f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 03:25:06 -0700 Subject: [PATCH 21/74] fix: require sidecars for schema-v2 captures --- src/benchflow/eval_artifacts.py | 6 +- .../trajectories/export_prime_sft.py | 7 +- src/benchflow/trajectories/export_trl_sft.py | 7 +- .../trajectories/llm_capture_manifest.py | 28 ++++++++ src/benchflow/trajectories/results.py | 7 +- tests/test_eval_artifact_cli.py | 70 ++++++++++++------- tests/trajectories/test_export_prime_sft.py | 25 +++++++ .../test_llm_capture_training_contract.py | 42 +++++++++-- 8 files changed, 149 insertions(+), 43 deletions(-) diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index 1f9f77fba..a97280218 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -16,7 +16,7 @@ load_llm_trajectory_jsonl, ) from benchflow.trajectories.llm_capture_manifest import ( - capture_manifest_allows_training, + capture_artifact_allows_training, read_llm_trajectory_manifest, ) @@ -209,9 +209,7 @@ def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, int]: if not rows: return True, False, 0 manifest = read_llm_trajectory_manifest(rollout_dir) - if manifest is not None and not capture_manifest_allows_training( - manifest, exchange_count=len(rows) - ): + if not capture_artifact_allows_training(manifest, exchanges=rows): return True, False, len(rows) return True, True, len(rows) diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 797efb00f..38a7e0d1a 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -19,7 +19,7 @@ from benchflow._utils.json_safe import dumps_finite, scrub_non_finite from benchflow.trajectories.llm_capture_manifest import ( - capture_manifest_allows_training, + capture_artifact_allows_training, read_llm_trajectory_manifest, ) from benchflow.trajectories.types import redact_trajectory_obj @@ -1222,8 +1222,9 @@ def convert_benchflow_rollouts_to_prime_sft_rows( trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) capture_manifest = read_llm_trajectory_manifest(rollout_dir) - if capture_manifest is not None and not capture_manifest_allows_training( - capture_manifest, exchange_count=len(exchanges) + if not capture_artifact_allows_training( + capture_manifest, + exchanges=exchanges, ): stats.skipped_insufficient_capture_fidelity += 1 continue diff --git a/src/benchflow/trajectories/export_trl_sft.py b/src/benchflow/trajectories/export_trl_sft.py index f105443c3..d83dcb405 100644 --- a/src/benchflow/trajectories/export_trl_sft.py +++ b/src/benchflow/trajectories/export_trl_sft.py @@ -25,7 +25,7 @@ validate_prime_sft_row, ) from benchflow.trajectories.llm_capture_manifest import ( - capture_manifest_allows_training, + capture_artifact_allows_training, read_llm_trajectory_manifest, ) from benchflow.trajectories.types import redact_trajectory_obj @@ -325,8 +325,9 @@ def convert_benchflow_rollouts_to_trl_sft_rows( trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) capture_manifest = read_llm_trajectory_manifest(rollout_dir) - if capture_manifest is not None and not capture_manifest_allows_training( - capture_manifest, exchange_count=len(exchanges) + if not capture_artifact_allows_training( + capture_manifest, + exchanges=exchanges, ): stats.skipped_insufficient_capture_fidelity += 1 continue diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 5e890c03e..19e9fa78d 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -9,6 +9,7 @@ import json import os +from collections.abc import Sequence from datetime import datetime from enum import StrEnum from pathlib import Path @@ -156,6 +157,33 @@ def capture_manifest_allows_training( ) +def capture_artifact_allows_training( + manifest: dict[str, Any] | None, + *, + exchanges: Sequence[dict[str, Any]], +) -> bool: + """Apply the sidecar contract while retaining genuine legacy JSONL support.""" + + if manifest is not None: + return capture_manifest_allows_training( + manifest, + exchange_count=len(exchanges), + ) + return not any(_exchange_requires_manifest(exchange) for exchange in exchanges) + + +def _exchange_requires_manifest(exchange: dict[str, Any]) -> bool: + metadata = exchange.get("metadata") + if not isinstance(metadata, dict): + return False + schema_version = metadata.get("schema_version") + return bool( + isinstance(schema_version, int) + and not isinstance(schema_version, bool) + and schema_version >= LLM_TRAJECTORY_SCHEMA_VERSION + ) + + def capture_manifest_has_oauth_role_capture(manifest: dict[str, Any]) -> bool: """Return whether a mixed manifest contains captured OAuth role evidence.""" diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 3335f23ae..112752524 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -33,7 +33,7 @@ validate_prime_sft_row, ) from benchflow.trajectories.llm_capture_manifest import ( - capture_manifest_allows_training, + capture_artifact_allows_training, capture_manifest_has_oauth_role_capture, read_llm_trajectory_manifest, ) @@ -182,8 +182,9 @@ def _llm_steps_from_trajectory( except PrimeSftTrajectoryJsonlError as exc: return _LLMStepsResult([], [], f"Invalid LLM trajectory JSONL: {exc}") capture_manifest = read_llm_trajectory_manifest(rollout_dir) - if capture_manifest is not None and not capture_manifest_allows_training( - capture_manifest, exchange_count=len(exchanges) + if not capture_artifact_allows_training( + capture_manifest, + exchanges=exchanges, ): return _LLMStepsResult([], [], None, capture_contract_rejected=True) training_success_indices = _training_success_exchange_indices(exchanges) diff --git a/tests/test_eval_artifact_cli.py b/tests/test_eval_artifact_cli.py index ead76e04b..6c0eae188 100644 --- a/tests/test_eval_artifact_cli.py +++ b/tests/test_eval_artifact_cli.py @@ -21,7 +21,12 @@ def _write_task(task_dir: Path) -> None: (task_dir / "task.toml").write_text('version = "1.0"\n', encoding="utf-8") -def _write_llm_trajectory(rollout_dir: Path, *, tool_calls: bool = True) -> None: +def _write_llm_trajectory( + rollout_dir: Path, + *, + tool_calls: bool = True, + schema_version: int | None = None, +) -> None: trajectory = rollout_dir / "trajectory" trajectory.mkdir(parents=True, exist_ok=True) message: dict[str, Any] = {"role": "assistant", "content": "done"} @@ -37,31 +42,31 @@ def _write_llm_trajectory(rollout_dir: Path, *, tool_calls: bool = True) -> None } ], } - (trajectory / "llm_trajectory.jsonl").write_text( - json.dumps( - { - "request": { - "body": { - "model": "m", - "messages": [{"role": "user", "content": "do it"}], - "tools": [ - { - "type": "function", - "function": { - "name": "finish", - "parameters": {"type": "object", "properties": {}}, - }, - } - ], + exchange = { + "request": { + "body": { + "model": "m", + "messages": [{"role": "user", "content": "do it"}], + "tools": [ + { + "type": "function", + "function": { + "name": "finish", + "parameters": {"type": "object", "properties": {}}, + }, } - }, - "response": { - "status_code": 200, - "body": {"choices": [{"message": message}]}, - }, + ], } - ) - + "\n", + }, + "response": { + "status_code": 200, + "body": {"choices": [{"message": message}]}, + }, + } + if schema_version is not None: + exchange["metadata"] = {"schema_version": schema_version} + (trajectory / "llm_trajectory.jsonl").write_text( + json.dumps(exchange) + "\n", encoding="utf-8", ) @@ -127,6 +132,23 @@ def test_health_rejects_empty_llm_capture_without_manifest(tmp_path: Path) -> No assert health["rows"][0]["llm_trajectory_rows"] == 0 +def test_health_rejects_sidecarless_schema_v2_llm_capture(tmp_path: Path) -> None: + """Guards PR #1057 against validating detached schema-v2 capture rows.""" + + job = tmp_path / "job" + rollout = job / "task-a__abc" + _write_rollout(rollout) + _write_llm_trajectory(rollout, schema_version=2) + + health = build_health_summary(job) + + assert health["missing_llm_trajectory"] == 0 + assert health["malformed_llm_trajectory"] == 1 + assert health["rows"][0]["has_llm_trajectory"] is True + assert health["rows"][0]["valid_llm_trajectory"] is False + assert health["rows"][0]["llm_trajectory_rows"] == 1 + + def test_eval_run_writes_manifest_health_and_canonical_artifacts( tmp_path: Path, monkeypatch ) -> None: diff --git a/tests/trajectories/test_export_prime_sft.py b/tests/trajectories/test_export_prime_sft.py index 839058a9a..c65c1bed2 100644 --- a/tests/trajectories/test_export_prime_sft.py +++ b/tests/trajectories/test_export_prime_sft.py @@ -245,6 +245,31 @@ def test_partial_manifest_blocks_prime_and_trl_training_exports( assert trl_stats.skipped_insufficient_capture_fidelity == 1 +def test_sidecarless_schema_v2_blocks_prime_and_trl_training_exports( + tmp_path: Path, +) -> None: + """Guards PR #1057 against exporting detached schema-v2 capture rows.""" + + exchange = _exchange(final=True) + exchange["metadata"] = { + "schema_version": 2, + "capture_fidelity": "provider_wire", + "request_complete": True, + "response_complete": True, + } + _write_rollout(tmp_path / "job" / "rollout-1", exchanges=[exchange]) + + prime_rows, prime_stats = convert_benchflow_rollouts_to_prime_sft_rows( + tmp_path / "job" + ) + trl_rows, trl_stats = convert_benchflow_rollouts_to_trl_sft_rows(tmp_path / "job") + + assert prime_rows == [] + assert prime_stats.skipped_insufficient_capture_fidelity == 1 + assert trl_rows == [] + assert trl_stats.skipped_insufficient_capture_fidelity == 1 + + def test_manifest_count_mismatch_blocks_prime_and_trl_training_exports( tmp_path: Path, ) -> None: diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 30f9d13da..51f3aa434 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -31,7 +31,19 @@ def _build_results_row(rollout_dir: Path, *, agent_result: dict) -> dict: ) -def _write_exchange(trajectory_dir: Path, *, fidelity: str) -> None: +def _write_exchange( + trajectory_dir: Path, + *, + fidelity: str, + schema_version: int | None = None, +) -> None: + metadata: dict[str, str | bool | int] = { + "capture_fidelity": fidelity, + "request_complete": fidelity == "provider_wire", + "response_complete": True, + } + if schema_version is not None: + metadata["schema_version"] = schema_version (trajectory_dir / "llm_trajectory.jsonl").write_text( json.dumps( { @@ -45,11 +57,7 @@ def _write_exchange(trajectory_dir: Path, *, fidelity: str) -> None: "content": [{"type": "text", "text": "hi"}], }, }, - "metadata": { - "capture_fidelity": fidelity, - "request_complete": fidelity == "provider_wire", - "response_complete": True, - }, + "metadata": metadata, } ) + "\n" @@ -103,6 +111,28 @@ def test_corrupt_capture_manifest_fails_closed_for_training(tmp_path: Path) -> N assert row["is_completed"] is False +def test_sidecarless_schema_v2_capture_fails_closed_for_canonical_results( + tmp_path: Path, +) -> None: + """Guards PR #1057 against treating detached schema-v2 rows as legacy.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange( + trajectory_dir, + fidelity="provider_wire", + schema_version=2, + ) + + row = _build_results_row(tmp_path, agent_result={"total_tokens": 2}) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == ( + "missing_healthy_structured_llm_trajectory" + ) + assert row["is_completed"] is False + + def test_manifest_count_mismatch_fails_closed_for_canonical_results( tmp_path: Path, ) -> None: From eee5952d12e3b6db9ac8db52ef7e57787cbea37e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 03:37:35 -0700 Subject: [PATCH 22/74] fix: fail closed on lost continuation calls --- src/benchflow/continue_run/orchestrator.py | 58 ++++++++++--- src/benchflow/continue_run/replay_proxy.py | 12 ++- src/benchflow/continue_run/sandbox_proxy.py | 90 +++++++++++++++++++-- tests/continue_run/test_orchestrator.py | 67 ++++++++++++++- tests/continue_run/test_replay_proxy.py | 55 +++++++++++++ 5 files changed, 264 insertions(+), 18 deletions(-) diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index dedde2144..29f8825b3 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -42,6 +42,7 @@ from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.scenes import compile_scenes_to_steps from benchflow.trajectories.llm_capture_manifest import ( + LLM_TRAJECTORY_SCHEMA_VERSION, AuthMode, CaptureFidelity, CaptureSource, @@ -270,15 +271,29 @@ def stitched_trajectory_lines( ) -> list[str]: """Build the continuous llm_trajectory: recorded prefix + live suffix. - The recorded prefix is taken verbatim from the source file (already redacted - and byte-identical to what the agent replayed); the live suffix is the - exchanges the proxy captured after the cut-point, redacted on the way out. + The recorded request/response payloads are preserved, while their metadata + is promoted to the current schema so a newly stitched artifact can never be + mistaken for sidecar-optional legacy data. The live suffix is redacted on + the way out. """ lines: list[str] = [] if original_llm_trajectory.is_file(): for raw in original_llm_trajectory.read_text().splitlines(): if raw.strip(): - lines.append(raw) + try: + payload = json.loads(raw) + except json.JSONDecodeError: + lines.append(raw) + continue + if not isinstance(payload, dict): + lines.append(raw) + continue + metadata = payload.setdefault("metadata", {}) + if not isinstance(metadata, dict): + metadata = {} + payload["metadata"] = metadata + metadata["schema_version"] = LLM_TRAJECTORY_SCHEMA_VERSION + lines.append(redact_trajectory_text(json.dumps(payload, default=str))) for exchange in live_exchanges: raw = json.dumps(exchange.model_dump(mode="json"), default=str) lines.append(redact_trajectory_text(raw)) @@ -307,6 +322,7 @@ def write_stitched_trajectory( "auth_mode": AuthMode.API_KEY.value, "capture_source": CaptureSource.LITELLM_PROXY.value, "capture_fidelity": CaptureFidelity.PROVIDER_WIRE.value, + "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, "request_complete": True, "response_complete": True, "role_attribution_complete": True, @@ -329,6 +345,8 @@ def refresh_stitched_trajectory_manifest( live_model: str | None, n_recorded: int, n_live: int, + live_attempt_count: int, + live_errors: list[str], ) -> LLMTrajectoryManifest: """Replace rollout-finalization provenance with the final stitched contract.""" @@ -355,13 +373,14 @@ def refresh_stitched_trajectory_manifest( exchange_count = sum( bool(line.strip()) for line in trajectory_path.read_text().splitlines() ) - expected_count = n_recorded + n_live + expected_count = n_recorded + live_attempt_count count_matches = exchange_count == expected_count source_allows_training = bool( source_raw is not None and capture_manifest_allows_training(source_raw, exchange_count=n_recorded) ) - complete = source_allows_training and count_matches + live_capture_complete = live_attempt_count == n_live and not live_errors + complete = source_allows_training and count_matches and live_capture_complete source_capture_source = source.capture_source if source else CaptureSource.NONE source_fidelity = source.capture_fidelity if source else CaptureFidelity.NONE @@ -385,10 +404,15 @@ def refresh_stitched_trajectory_manifest( capture_fidelity = source_fidelity auth_mode = source_auth - request_complete = bool(source and source.request_complete and count_matches) - response_complete = bool(source and source.response_complete and count_matches) + request_complete = bool( + source and source.request_complete and count_matches and live_capture_complete + ) + response_complete = bool( + source and source.response_complete and count_matches and live_capture_complete + ) errors = list(source.errors) if source and not complete else [] missing_fields = list(source.missing_fields) if source and not complete else [] + errors.extend(live_errors) if source is None: errors.append("source LLM trajectory manifest is missing or malformed") missing_fields.append("source_capture_provenance") @@ -400,12 +424,14 @@ def refresh_stitched_trajectory_manifest( f"expected {expected_count}, found {exchange_count}" ) missing_fields.append("exchange_count") + if not live_capture_complete: + missing_fields.append("live_provider_exchange") models = { value for value, present in ( (original_model, n_recorded > 0), - (live_model, n_live > 0), + (live_model, live_attempt_count > 0), ) if present and value } @@ -730,6 +756,8 @@ async def continue_run( live_model=live_model, n_recorded=run.n_recorded_exchanges, n_live=len(router.live_exchanges), + live_attempt_count=router.live_attempt_count, + live_errors=list(router.live_errors), ) update_continued_metadata( rollout_dir, @@ -809,6 +837,18 @@ async def _write_artifacts_before_cleanup() -> None: live_model=live_model, n_recorded=run.n_recorded_exchanges, n_live=len(live_exchanges), + live_attempt_count=( + replay_proxy.live_attempt_count if replay_proxy is not None else 0 + ), + live_errors=[ + *(replay_proxy.live_errors if replay_proxy is not None else []), + *( + ["continuation teardown failed before capture finalization"] + if isinstance(rollout._error, str) + and rollout._error.startswith("Continuation teardown warning:") + else [] + ), + ], ) update_continued_metadata( rollout_dir, diff --git a/src/benchflow/continue_run/replay_proxy.py b/src/benchflow/continue_run/replay_proxy.py index 0d730503b..01a196154 100644 --- a/src/benchflow/continue_run/replay_proxy.py +++ b/src/benchflow/continue_run/replay_proxy.py @@ -73,6 +73,8 @@ def __init__( self._lock = threading.Lock() self._cursor = 0 self.divergences = 0 + self.live_attempt_count = 0 + self.live_errors: list[str] = [] # Live-leg exchanges, in order, for stitching onto the recorded prefix. self.live_exchanges: list[LLMExchange] = [] @@ -115,9 +117,11 @@ def next_response(self, request_body: dict[str, Any]) -> ReplayResult: ) # Past the cut-point: live continuation. self._cursor += 1 + self.live_attempt_count += 1 forwarder = self._live_forwarder if forwarder is None: + self.live_errors.append("live continuation had no configured forwarder") logger.error( "recorded responses exhausted at turn %d and no live forwarder " "is configured — returning an error to the agent.", @@ -137,7 +141,13 @@ def next_response(self, request_body: dict[str, Any]) -> ReplayResult: }, ) - body = forwarder(request_body) + try: + body = forwarder(request_body) + except Exception: + self.live_errors.append( + "live provider request failed before a response could be captured" + ) + raise # Capture the live exchange so the caller can stitch a continuous # llm_trajectory.jsonl (recorded prefix + live suffix). self.live_exchanges.append( diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 5ad108ac8..d87a90038 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -39,6 +39,7 @@ def _sandbox_proxy_source() -> str: from __future__ import annotations import json +import os import sys import threading import time @@ -125,11 +126,27 @@ class ReplayState: upstream_api_key: str upstream_model: str live_log_path: str + state_path: str + port: int strict_divergence: bool = False cursor: int = 0 divergences: int = 0 + live_attempt_count: int = 0 + live_error_count: int = 0 lock: threading.Lock = field(default_factory=threading.Lock) + def _write_state(self): + payload = { + "port": self.port, + "live_attempt_count": self.live_attempt_count, + "live_error_count": self.live_error_count, + } + temporary = self.state_path + ".tmp" + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + handle.flush() + os.replace(temporary, self.state_path) + def _check_divergence(self, incoming, recorded_request): want = _n_messages(recorded_request) got = _n_messages(incoming) @@ -155,9 +172,22 @@ def next_response(self, request_body): response = exchange.get("response") or {} return "replay", int(response.get("status_code") or 200), dict(response.get("body") or {}) self.cursor += 1 + self.live_attempt_count += 1 + self._write_state() - status, body = self._forward_live(request_body) - self._append_live_exchange(request_body, status, body) + status, body, provider_observed = self._forward_live(request_body) + if provider_observed: + try: + self._append_live_exchange(request_body, status, body) + except Exception: + with self.lock: + self.live_error_count += 1 + self._write_state() + raise + else: + with self.lock: + self.live_error_count += 1 + self._write_state() return "live", status, body def _forward_live(self, request_body): @@ -177,17 +207,17 @@ def _forward_live(self, request_body): try: with urllib.request.urlopen(request, timeout=600) as response: raw = response.read().decode("utf-8") - return int(response.status), json.loads(raw or "{}") + return int(response.status), json.loads(raw or "{}"), True except urllib.error.HTTPError as exc: raw = exc.read().decode("utf-8", errors="replace") try: body = json.loads(raw or "{}") except json.JSONDecodeError: body = {"error": {"message": raw or str(exc)}} - return int(exc.code), body + return int(exc.code), body, True except Exception as exc: traceback.print_exc() - return 500, {"error": {"message": str(exc)}} + return 500, {"error": {"message": str(exc)}}, False def _append_live_exchange(self, request_body, status, body): row = { @@ -286,11 +316,12 @@ def main(): upstream_api_key=cfg["upstream_api_key"], upstream_model=cfg["upstream_model"], live_log_path=cfg["live_log_path"], + state_path=cfg["state_path"], + port=int(cfg["port"]), strict_divergence=bool(cfg.get("strict_divergence")), ) server = ReplayServer(("127.0.0.1", int(cfg["port"])), ReplayHandler, state) - with open(cfg["state_path"], "w", encoding="utf-8") as handle: - json.dump({"port": int(cfg["port"])}, handle) + state._write_state() server.serve_forever() @@ -341,6 +372,8 @@ class SandboxReplayProxy: stdout_path: str stderr_path: str live_exchanges: list[LLMExchange] = field(default_factory=list) + live_attempt_count: int = 0 + live_errors: list[str] = field(default_factory=list) @property def base_url(self) -> str: @@ -446,6 +479,7 @@ async def _wait_until_ready(self) -> None: async def stop(self) -> None: self.live_exchanges = await self._load_live_exchanges() + await self._load_live_state() with contextlib.suppress(Exception): await self.sandbox.exec( f"if [ -s {shlex.quote(self.pid_path)} ]; then " @@ -462,11 +496,53 @@ async def stop(self) -> None: async def _load_live_exchanges(self) -> list[LLMExchange]: text = await _read_remote_text(self.sandbox, self.live_log_path) exchanges: list[LLMExchange] = [] + malformed = 0 for raw in text.splitlines(): if not raw.strip(): continue try: exchanges.append(LLMExchange.model_validate_json(raw)) except Exception: + malformed += 1 continue + if malformed: + self.live_errors.append( + f"{malformed} sandbox live exchange record(s) were malformed" + ) return exchanges + + async def _load_live_state(self) -> None: + text = await _read_remote_text(self.sandbox, self.state_path) + try: + state = json.loads(text) + except (TypeError, json.JSONDecodeError): + state = None + if not isinstance(state, dict): + self.live_errors.append("sandbox live capture state was unavailable") + return + attempt_count = state.get("live_attempt_count") + error_count = state.get("live_error_count") + if ( + not isinstance(attempt_count, int) + or isinstance(attempt_count, bool) + or attempt_count < 0 + ): + self.live_errors.append("sandbox live attempt count was invalid") + return + self.live_attempt_count = attempt_count + if ( + not isinstance(error_count, int) + or isinstance(error_count, bool) + or error_count < 0 + ): + self.live_errors.append("sandbox live error count was invalid") + elif error_count: + self.live_errors.append( + f"{error_count} sandbox live provider request(s) failed before capture" + ) + if self.live_attempt_count != len(self.live_exchanges): + self.live_errors.append( + "sandbox live exchange recovery mismatch: " + f"attempted {self.live_attempt_count}, " + f"recovered {len(self.live_exchanges)}" + ) diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index c289610cd..b37fcf04e 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -156,7 +156,9 @@ def test_stitched_trajectory_recorded_prefix_plus_live_suffix(tmp_path): live = [exchange(completion(content="LIVE"))] lines = stitched_trajectory_lines(original, live) assert len(lines) == 3 - assert json.loads(lines[0]) == {"a": 1} + first = json.loads(lines[0]) + assert first["a"] == 1 + assert first["metadata"]["schema_version"] == 2 last = json.loads(lines[2]) assert last["response"]["body"]["choices"][0]["message"]["content"] == "LIVE" @@ -219,6 +221,8 @@ def test_refresh_stitched_manifest_replaces_pre_stitch_finalization(tmp_path): live_model=model, n_recorded=1, n_live=1, + live_attempt_count=1, + live_errors=[], ) assert manifest.status is CaptureStatus.COMPLETE @@ -230,6 +234,7 @@ def test_refresh_stitched_manifest_replaces_pre_stitch_finalization(tmp_path): live_row = json.loads(out.read_text().splitlines()[-1]) assert live_row["metadata"]["capture_fidelity"] == "provider_wire" assert live_row["metadata"]["model"] == model + assert live_row["metadata"]["schema_version"] == 2 def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path): @@ -278,6 +283,8 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) live_model="openai/gpt-5.5", n_recorded=1, n_live=1, + live_attempt_count=1, + live_errors=[], ) assert manifest.status is CaptureStatus.PARTIAL @@ -289,6 +296,64 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) ) +def test_refresh_stitched_manifest_rejects_missing_live_attempt(tmp_path): + """Guards PR #1057 against completing a lost continuation exchange.""" + + model = "openai/gpt-5.5" + source = write_run_folder( + tmp_path / "source", + exchanges=[exchange(completion(content="recorded"))], + model=model, + ) + source_manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="source", + exchange_count=1, + request_complete=True, + response_complete=True, + started_at="2026-08-29T00:00:00Z", + ) + write_llm_trajectory_manifest(source, source_manifest) + rollout = tmp_path / "continued" + initialize_llm_trajectory_artifacts( + rollout, + agent="openhands", + model=None, + session_id="continued", + started_at=source_manifest.started_at, + ) + stitched = write_stitched_trajectory( + rollout, + source / "trajectory" / "llm_trajectory.jsonl", + [], + live_model=model, + ) + + manifest = refresh_stitched_trajectory_manifest( + rollout, + source, + original_model=model, + live_model=model, + n_recorded=1, + n_live=0, + live_attempt_count=1, + live_errors=["live provider request failed before capture"], + ) + + assert manifest.status is CaptureStatus.PARTIAL + assert manifest.request_complete is False + assert manifest.response_complete is False + assert "live_provider_exchange" in manifest.missing_fields + assert any("count mismatch" in error for error in manifest.errors) + recorded_row = json.loads(stitched.read_text()) + assert recorded_row["metadata"]["schema_version"] == 2 + + def test_summarize_llm_trajectory_usage_splits_recorded_and_live(tmp_path): """Guards the PR #648 continuation metadata fix for stitched token usage.""" traj = tmp_path / "llm_trajectory.jsonl" diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index f7c1d55d5..411044207 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -13,6 +13,7 @@ ReplayRouter, completion_to_sse, ) +from benchflow.continue_run.sandbox_proxy import _sandbox_proxy_source from ._helpers import completion, exchange @@ -46,6 +47,8 @@ def forwarder(req): assert r3.body["choices"][0]["message"]["content"] == "LIVE" # the live exchange was captured for stitching assert len(router.live_exchanges) == 1 + assert router.live_attempt_count == 1 + assert router.live_errors == [] assert len(live) == 1 @@ -55,6 +58,58 @@ def test_exhausted_without_forwarder_returns_error(): result = router.next_response({"messages": [{}]}) assert result.source == "error" assert result.status == 503 + assert router.live_attempt_count == 1 + assert router.live_errors + + +def test_live_forwarder_failure_retains_unpaired_attempt() -> None: + """Guards PR #1057 against completing a lost host continuation call.""" + + def fail(_request): + raise RuntimeError("provider unavailable") + + router = ReplayRouter([], live_forwarder=fail) + + with pytest.raises(RuntimeError, match="provider unavailable"): + router.next_response({"messages": [{}]}) + + assert router.live_attempt_count == 1 + assert router.live_exchanges == [] + assert router.live_errors == [ + "live provider request failed before a response could be captured" + ] + + +def test_sandbox_forwarding_failure_is_not_logged_as_provider_exchange( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Guards PR #1057 against labeling a synthesized sandbox 500 provider-wire.""" + + namespace: dict[str, object] = {} + exec(_sandbox_proxy_source(), namespace) + state = namespace["ReplayState"]( + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(tmp_path / "live.jsonl"), + state_path=str(tmp_path / "state.json"), + port=61357, + ) + + def fail(*_args, **_kwargs): + raise TimeoutError("provider unavailable") + + monkeypatch.setattr(namespace["urllib"].request, "urlopen", fail) + source, status, body = state.next_response({"messages": [{"role": "user"}]}) + + assert source == "live" + assert status == 500 + assert "error" in body + assert not (tmp_path / "live.jsonl").exists() + capture_state = json.loads((tmp_path / "state.json").read_text()) + assert capture_state["live_attempt_count"] == 1 + assert capture_state["live_error_count"] == 1 def test_divergence_warns_by_default(): From 5426420ba6140b8758df9483de941dcfe9a5778d Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 03:56:09 -0700 Subject: [PATCH 23/74] fix: finalize continuation capture artifacts --- src/benchflow/continue_run/orchestrator.py | 66 ++++--- src/benchflow/continue_run/sandbox_proxy.py | 33 +++- src/benchflow/trajectories/results.py | 65 +++++++ tests/continue_run/test_orchestrator.py | 182 +++++++++++++++++++- tests/continue_run/test_replay_proxy.py | 65 ++++++- 5 files changed, 385 insertions(+), 26 deletions(-) diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 29f8825b3..690bb1690 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -53,7 +53,7 @@ read_llm_trajectory_manifest, write_llm_trajectory_manifest, ) -from benchflow.trajectories.types import LLMExchange, redact_trajectory_text +from benchflow.trajectories.types import LLMExchange, redact_trajectory_obj logger = logging.getLogger(__name__) @@ -293,10 +293,10 @@ def stitched_trajectory_lines( metadata = {} payload["metadata"] = metadata metadata["schema_version"] = LLM_TRAJECTORY_SCHEMA_VERSION - lines.append(redact_trajectory_text(json.dumps(payload, default=str))) + lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) for exchange in live_exchanges: - raw = json.dumps(exchange.model_dump(mode="json"), default=str) - lines.append(redact_trajectory_text(raw)) + payload = redact_trajectory_obj(exchange.model_dump(mode="json")) + lines.append(json.dumps(payload, default=str)) return lines @@ -329,7 +329,7 @@ def write_stitched_trajectory( "payload_redacted": True, } ) - lines.append(redact_trajectory_text(json.dumps(payload, default=str))) + lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) rendered = "\n".join(lines) + ("\n" if lines else "") temporary = out.with_suffix(out.suffix + ".tmp") temporary.write_text(rendered) @@ -370,9 +370,17 @@ def refresh_stitched_trajectory_manifest( source = None trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" - exchange_count = sum( - bool(line.strip()) for line in trajectory_path.read_text().splitlines() - ) + trajectory_lines = [ + line for line in trajectory_path.read_text().splitlines() if line.strip() + ] + exchange_count = len(trajectory_lines) + malformed_count = 0 + for line in trajectory_lines: + try: + LLMExchange.model_validate_json(line) + except ValueError: + malformed_count += 1 + rows_valid = malformed_count == 0 expected_count = n_recorded + live_attempt_count count_matches = exchange_count == expected_count source_allows_training = bool( @@ -380,7 +388,12 @@ def refresh_stitched_trajectory_manifest( and capture_manifest_allows_training(source_raw, exchange_count=n_recorded) ) live_capture_complete = live_attempt_count == n_live and not live_errors - complete = source_allows_training and count_matches and live_capture_complete + complete = ( + source_allows_training + and count_matches + and live_capture_complete + and rows_valid + ) source_capture_source = source.capture_source if source else CaptureSource.NONE source_fidelity = source.capture_fidelity if source else CaptureFidelity.NONE @@ -405,10 +418,18 @@ def refresh_stitched_trajectory_manifest( auth_mode = source_auth request_complete = bool( - source and source.request_complete and count_matches and live_capture_complete + source + and source.request_complete + and count_matches + and live_capture_complete + and rows_valid ) response_complete = bool( - source and source.response_complete and count_matches and live_capture_complete + source + and source.response_complete + and count_matches + and live_capture_complete + and rows_valid ) errors = list(source.errors) if source and not complete else [] missing_fields = list(source.missing_fields) if source and not complete else [] @@ -424,6 +445,11 @@ def refresh_stitched_trajectory_manifest( f"expected {expected_count}, found {exchange_count}" ) missing_fields.append("exchange_count") + if not rows_valid: + errors.append( + f"stitched LLM trajectory contains {malformed_count} malformed row(s)" + ) + missing_fields.append("valid_provider_exchange") if not live_capture_complete: missing_fields.append("live_provider_exchange") @@ -595,6 +621,9 @@ def update_continued_metadata( else usage_tracking.get("status", "off") ) result_path.write_text(json.dumps(result, indent=2) + "\n") + from benchflow.trajectories.results import refresh_rollout_results_jsonl + + refresh_rollout_results_jsonl(rollout_dir) def _host_proxy_binding(environment: str) -> tuple[str, str]: @@ -620,7 +649,7 @@ async def _safe_sandbox_continuation_teardown( replay_proxy: SandboxReplayProxy | None, provider_runtime: Any | None, stop_provider_runtime: Callable[[Any], Awaitable[None]], - before_cleanup: Callable[[], Awaitable[None]] | None = None, + before_cleanup: Callable[[list[str]], Awaitable[None]] | None = None, ) -> list[str]: """Stop continuation sidecars but always let Rollout write final artifacts.""" errors: list[str] = [] @@ -642,7 +671,7 @@ async def _capture(label: str, awaitable: Awaitable[Any]) -> None: rollout._error = "Continuation teardown warning: " + "; ".join(errors) if before_cleanup is not None: - await _capture("continuation artifact write", before_cleanup()) + await _capture("continuation artifact write", before_cleanup(errors)) await _capture("rollout cleanup", rollout.cleanup()) @@ -815,7 +844,9 @@ async def _continue_run_with_sandbox_proxy( live_exchanges: list[LLMExchange] = [] artifacts_written = False - async def _write_artifacts_before_cleanup() -> None: + async def _write_artifacts_before_cleanup( + teardown_errors: list[str] | None = None, + ) -> None: nonlocal artifacts_written, live_exchanges, result, rollout_dir if artifacts_written: return @@ -842,12 +873,7 @@ async def _write_artifacts_before_cleanup() -> None: ), live_errors=[ *(replay_proxy.live_errors if replay_proxy is not None else []), - *( - ["continuation teardown failed before capture finalization"] - if isinstance(rollout._error, str) - and rollout._error.startswith("Continuation teardown warning:") - else [] - ), + *(teardown_errors or []), ], ) update_continued_metadata( diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index d87a90038..009d5d5a7 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -27,6 +27,7 @@ SANDBOX_REPLAY_ROOT = "/tmp/benchflow-replay" DEFAULT_SANDBOX_REPLAY_PORT = 61357 +_CAPTURE_STATE_WRITE_FAILED = "BENCHFLOW_CAPTURE_STATE_WRITE_FAILED" def sandbox_replay_base_url(port: int = DEFAULT_SANDBOX_REPLAY_PORT) -> str: @@ -49,6 +50,8 @@ def _sandbox_proxy_source() -> str: from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +CAPTURE_STATE_WRITE_FAILED = "BENCHFLOW_CAPTURE_STATE_WRITE_FAILED" + def _n_messages(body): messages = body.get("messages") @@ -142,10 +145,23 @@ def _write_state(self): "live_error_count": self.live_error_count, } temporary = self.state_path + ".tmp" - with open(temporary, "w", encoding="utf-8") as handle: - json.dump(payload, handle) - handle.flush() - os.replace(temporary, self.state_path) + try: + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + handle.flush() + os.replace(temporary, self.state_path) + except Exception: + self.live_error_count += 1 + try: + os.unlink(self.state_path) + except OSError: + pass + try: + os.unlink(temporary) + except OSError: + pass + print(CAPTURE_STATE_WRITE_FAILED, file=sys.stderr, flush=True) + raise def _check_divergence(self, incoming, recorded_request): want = _n_messages(recorded_request) @@ -480,6 +496,7 @@ async def _wait_until_ready(self) -> None: async def stop(self) -> None: self.live_exchanges = await self._load_live_exchanges() await self._load_live_state() + await self._load_runtime_errors() with contextlib.suppress(Exception): await self.sandbox.exec( f"if [ -s {shlex.quote(self.pid_path)} ]; then " @@ -546,3 +563,11 @@ async def _load_live_state(self) -> None: f"attempted {self.live_attempt_count}, " f"recovered {len(self.live_exchanges)}" ) + + async def _load_runtime_errors(self) -> None: + stderr = await _read_remote_text(self.sandbox, self.stderr_path) + failures = stderr.count(_CAPTURE_STATE_WRITE_FAILED) + if failures: + self.live_errors.append( + f"sandbox live attempt journal failed {failures} time(s)" + ) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 112752524..9d6de8f22 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -654,6 +654,71 @@ def write_rollout_results_jsonl( return record +def refresh_rollout_results_jsonl(rollout_dir: str | Path) -> dict[str, Any] | None: + """Rebuild the trainer-facing row from finalized rollout artifacts. + + Continuation replaces the LLM trajectory and its manifest after Rollout's + normal result-building phase. Rebuilding here keeps ``results.jsonl`` in + lockstep with the final capture contract, model identity, and token usage. + """ + rollout_path = Path(rollout_dir) + result_path = rollout_path / "result.json" + if not result_path.is_file(): + return None + result = json.loads(result_path.read_text()) + if not isinstance(result, dict): + raise ValueError(f"Invalid rollout result object: {result_path}") + + prompts: list[str] = [] + prompts_path = rollout_path / "prompts.json" + if prompts_path.is_file(): + raw_prompts = json.loads(prompts_path.read_text()) + if isinstance(raw_prompts, list): + prompts = [item for item in raw_prompts if isinstance(item, str)] + + task_name = str(result.get("task_name") or rollout_path.name.split("__", 1)[0]) + rollout_name = str(result.get("rollout_name") or rollout_path.name) + agent = str(result.get("agent") or "unknown") + agent_name = str(result.get("agent_name") or agent) + n_tool_calls = result.get("n_tool_calls") + if not isinstance(n_tool_calls, int) or isinstance(n_tool_calls, bool): + n_tool_calls = 0 + record = write_rollout_results_jsonl( + rollout_path, + task_name=task_name, + rollout_name=rollout_name, + agent=agent, + agent_name=agent_name, + model=result.get("model") if isinstance(result.get("model"), str) else None, + n_tool_calls=n_tool_calls, + prompts=prompts, + trajectory=[], + partial_trajectory=bool(result.get("partial_trajectory")), + rewards=result.get("rewards") + if isinstance(result.get("rewards"), dict) + else None, + error=result.get("error") if isinstance(result.get("error"), str) else None, + verifier_error=( + result.get("verifier_error") + if isinstance(result.get("verifier_error"), str) + else None + ), + export_error=( + result.get("export_error") + if isinstance(result.get("export_error"), str) + else None + ), + timing=result.get("timing") if isinstance(result.get("timing"), dict) else None, + agent_result=( + result.get("agent_result") + if isinstance(result.get("agent_result"), dict) + else None + ), + ) + write_job_results_jsonl(rollout_path.parent) + return record + + def write_job_results_jsonl(job_dir: str | Path) -> Path | None: """Aggregate per-rollout ``results.jsonl`` files into ``job_dir/results.jsonl``.""" job_path = Path(job_dir) diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index b37fcf04e..eed97a457 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -476,9 +476,10 @@ async def stop_provider_runtime(runtime): rollout = FakeRollout() events: list[str] = [] - async def before_cleanup(): + async def before_cleanup(teardown_errors): events.append("artifact") assert rollout._error is not None + assert len(teardown_errors) == 2 errors = await _safe_sandbox_continuation_teardown( rollout=rollout, @@ -494,3 +495,182 @@ async def before_cleanup(): assert rollout._error is not None assert "proxy unavailable" in rollout._error assert "provider refused stop" in rollout._error + + +@pytest.mark.asyncio +async def test_sandbox_teardown_passes_errors_when_agent_already_failed(): + """Guards PR #1057 against hiding capture teardown behind an agent error.""" + + class FailingProxy: + async def stop(self): + raise RuntimeError("capture state unavailable") + + class FakeRollout: + _error = "agent failed first" + + async def cleanup(self): + return None + + observed: list[str] = [] + + async def before_cleanup(teardown_errors): + observed.extend(teardown_errors) + + async def stop_provider_runtime(_runtime): + return None + + rollout = FakeRollout() + errors = await _safe_sandbox_continuation_teardown( + rollout=rollout, + replay_proxy=FailingProxy(), + provider_runtime=None, + stop_provider_runtime=stop_provider_runtime, + before_cleanup=before_cleanup, + ) + + assert errors == observed + assert errors and "capture state unavailable" in errors[0] + assert rollout._error == "agent failed first" + + +def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): + """Guards PR #1057 against retaining the pre-stitch trainer row.""" + rollout = tmp_path / "job" / "demo-task__continued" + (rollout / "trajectory").mkdir(parents=True) + model = "openai/gpt-5.5" + row = exchange(completion(content="done")).model_dump(mode="json") + row["request"]["body"]["messages"] = [{"role": "user", "content": "Do the task."}] + row["metadata"] = { + "schema_version": 2, + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + } + trajectory_path = rollout / "trajectory" / "llm_trajectory.jsonl" + trajectory_path.write_text(json.dumps(row) + "\n") + manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="continued", + exchange_count=1, + request_complete=True, + response_complete=True, + started_at="2026-08-29T00:00:00Z", + finished_at="2026-08-29T00:01:00Z", + ) + write_llm_trajectory_manifest(rollout, manifest) + (rollout / "config.json").write_text(json.dumps({"model": None, "source": {}})) + (rollout / "prompts.json").write_text(json.dumps(["Do the task."])) + (rollout / "result.json").write_text( + json.dumps( + { + "task_name": "demo-task", + "rollout_name": "demo-task__continued", + "agent": "openhands", + "agent_name": "OpenHands", + "model": None, + "n_tool_calls": 0, + "partial_trajectory": False, + "rewards": {"reward": 1.0}, + "error": None, + "verifier_error": None, + "export_error": None, + "timing": {}, + "agent_result": {"total_tokens": 0, "usage_source": "unavailable"}, + "usage_tracking": {"requested": "off", "status": "off"}, + } + ) + ) + (rollout / "results.jsonl").write_text( + json.dumps({"info": {"training_ready": False, "model": None}}) + "\n" + ) + + update_continued_metadata( + rollout, + live_model=model, + usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), + environment="docker", + ) + + refreshed = json.loads((rollout / "results.jsonl").read_text()) + aggregated = json.loads((rollout.parent / "results.jsonl").read_text()) + assert refreshed["info"]["model"] == model + assert refreshed["info"]["training_ready"] is True + assert refreshed["token_usage"]["total_tokens"] == 2 + assert len(refreshed["trajectory"]) == 1 + assert aggregated == refreshed + + +def test_stitching_structurally_redacts_escaped_secret(tmp_path): + """Guards PR #1057 against string redaction corrupting stitched JSON.""" + secret = "ESCbearerSECRETtok123456" + source = tmp_path / "source.jsonl" + payload = exchange(completion(content="done")).model_dump(mode="json") + payload["request"]["body"]["authorization"] = f'Bearer {secret}\\"tail' + source.write_text(json.dumps(payload) + "\n") + + rendered = stitched_trajectory_lines(source, []) + + assert len(rendered) == 1 + restored = json.loads(rendered[0]) + assert secret not in rendered[0] + assert "***REDACTED***" in restored["request"]["body"]["authorization"] + assert restored["metadata"]["schema_version"] == 2 + + +def test_refresh_stitched_manifest_rejects_malformed_row(tmp_path): + """Guards PR #1057 against a complete sidecar for invalid stitched JSONL.""" + model = "openai/gpt-5.5" + source = write_run_folder( + tmp_path / "source", + exchanges=[exchange(completion(content="recorded"))], + model=model, + ) + source_manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="source", + exchange_count=1, + request_complete=True, + response_complete=True, + started_at="2026-08-29T00:00:00Z", + ) + write_llm_trajectory_manifest(source, source_manifest) + rollout = tmp_path / "continued" + initialize_llm_trajectory_artifacts( + rollout, + agent="openhands", + model=None, + session_id="continued", + started_at=source_manifest.started_at, + ) + stitched = rollout / "trajectory" / "llm_trajectory.jsonl" + stitched.write_text('{"request": broken}\n') + + manifest = refresh_stitched_trajectory_manifest( + rollout, + source, + original_model=model, + live_model=model, + n_recorded=1, + n_live=0, + live_attempt_count=0, + live_errors=[], + ) + + assert manifest.status is CaptureStatus.PARTIAL + assert manifest.request_complete is False + assert manifest.response_complete is False + assert "valid_provider_exchange" in manifest.missing_fields + assert any("malformed row" in error for error in manifest.errors) diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index 411044207..439ec1b51 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -13,7 +13,10 @@ ReplayRouter, completion_to_sse, ) -from benchflow.continue_run.sandbox_proxy import _sandbox_proxy_source +from benchflow.continue_run.sandbox_proxy import ( + SandboxReplayProxy, + _sandbox_proxy_source, +) from ._helpers import completion, exchange @@ -112,6 +115,66 @@ def fail(*_args, **_kwargs): assert capture_state["live_error_count"] == 1 +def test_sandbox_attempt_journal_failure_invalidates_stale_state( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Guards PR #1057 against completing an unjournaled sandbox call.""" + namespace: dict[str, object] = {} + exec(_sandbox_proxy_source(), namespace) + state_path = tmp_path / "state.json" + state_path.write_text(json.dumps({"live_attempt_count": 0, "live_error_count": 0})) + state = namespace["ReplayState"]( + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(tmp_path / "live.jsonl"), + state_path=str(state_path), + port=61357, + ) + + def fail_replace(*_args, **_kwargs): + raise OSError("disk full") + + monkeypatch.setattr(namespace["os"], "replace", fail_replace) + + with pytest.raises(OSError, match="disk full"): + state.next_response({"messages": [{"role": "user"}]}) + + assert state.live_attempt_count == 1 + assert state.live_error_count == 1 + assert not state_path.exists() + + +@pytest.mark.asyncio +async def test_sandbox_attempt_journal_marker_reaches_host( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Guards PR #1057 by surfacing the sandbox journal marker at teardown.""" + proxy = SandboxReplayProxy( + sandbox=object(), + runtime_dir="/tmp/runtime", + port=61357, + pid_path="/tmp/runtime/pid", + live_log_path="/tmp/runtime/live.jsonl", + state_path="/tmp/runtime/state.json", + stdout_path="/tmp/runtime/stdout.log", + stderr_path="/tmp/runtime/stderr.log", + ) + + async def read_remote_text(_sandbox, path, **_kwargs): + assert path == proxy.stderr_path + return "BENCHFLOW_CAPTURE_STATE_WRITE_FAILED\n" + + monkeypatch.setattr( + "benchflow.continue_run.sandbox_proxy._read_remote_text", read_remote_text + ) + + await proxy._load_runtime_errors() + + assert proxy.live_errors == ["sandbox live attempt journal failed 1 time(s)"] + + def test_divergence_warns_by_default(): recorded = [exchange(completion(content="a"), n_request_messages=3)] router = ReplayRouter(recorded) From 59883c661d4414cd09fbe1f3fc47b3f76dca8b5f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 04:07:22 -0700 Subject: [PATCH 24/74] fix: refinalize continuation after cleanup --- src/benchflow/continue_run/orchestrator.py | 14 +++++--- tests/continue_run/test_orchestrator.py | 38 ++++++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 690bb1690..5f81effde 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -650,6 +650,7 @@ async def _safe_sandbox_continuation_teardown( provider_runtime: Any | None, stop_provider_runtime: Callable[[Any], Awaitable[None]], before_cleanup: Callable[[list[str]], Awaitable[None]] | None = None, + after_cleanup: Callable[[list[str]], Awaitable[None]] | None = None, ) -> list[str]: """Stop continuation sidecars but always let Rollout write final artifacts.""" errors: list[str] = [] @@ -674,6 +675,8 @@ async def _capture(label: str, awaitable: Awaitable[Any]) -> None: await _capture("continuation artifact write", before_cleanup(errors)) await _capture("rollout cleanup", rollout.cleanup()) + if after_cleanup is not None: + await _capture("continuation artifact finalization", after_cleanup(errors)) return errors @@ -844,11 +847,13 @@ async def _continue_run_with_sandbox_proxy( live_exchanges: list[LLMExchange] = [] artifacts_written = False - async def _write_artifacts_before_cleanup( + async def _write_artifacts( teardown_errors: list[str] | None = None, + *, + force: bool = False, ) -> None: nonlocal artifacts_written, live_exchanges, result, rollout_dir - if artifacts_written: + if artifacts_written and not force: return result = _result_after_sandbox_teardown(rollout) if result is None: @@ -970,14 +975,15 @@ async def _write_artifacts_before_cleanup( replay_proxy=replay_proxy, provider_runtime=provider_runtime, stop_provider_runtime=stop_provider_runtime, - before_cleanup=_write_artifacts_before_cleanup, + before_cleanup=_write_artifacts, + after_cleanup=lambda errors: _write_artifacts(errors, force=True), ) if pending_acp_error is not None: rollout._error = rollout._classify_acp_error(pending_acp_error) logger.error(rollout._error) - await _write_artifacts_before_cleanup() + await _write_artifacts() if rollout_dir is None: rollout_dir = Path(rollout._rollout_dir or (output_dir / rollout_name)) diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index eed97a457..ace611005 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -533,6 +533,44 @@ async def stop_provider_runtime(_runtime): assert rollout._error == "agent failed first" +@pytest.mark.asyncio +async def test_sandbox_teardown_refinalizes_manifest_after_cleanup(tmp_path): + """Guards PR #1057 against cleanup overwriting teardown capture errors.""" + manifest_path = tmp_path / "manifest.json" + + class FailingProxy: + async def stop(self): + raise RuntimeError("capture stop failed") + + class FakeRollout: + _error = "agent failed first" + + async def cleanup(self): + manifest_path.write_text(json.dumps({"status": "complete"})) + + async def stop_provider_runtime(_runtime): + return None + + async def before_cleanup(teardown_errors): + assert teardown_errors + manifest_path.write_text(json.dumps({"status": "partial"})) + + async def after_cleanup(teardown_errors): + assert teardown_errors + manifest_path.write_text(json.dumps({"status": "partial"})) + + await _safe_sandbox_continuation_teardown( + rollout=FakeRollout(), + replay_proxy=FailingProxy(), + provider_runtime=None, + stop_provider_runtime=stop_provider_runtime, + before_cleanup=before_cleanup, + after_cleanup=after_cleanup, + ) + + assert json.loads(manifest_path.read_text())["status"] == "partial" + + def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): """Guards PR #1057 against retaining the pre-stitch trainer row.""" rollout = tmp_path / "job" / "demo-task__continued" From f6daaf2425e4446ea8e0436949e47b71c42a0d48 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 04:21:50 -0700 Subject: [PATCH 25/74] fix: quiesce capture and bind provider roles --- src/benchflow/continue_run/sandbox_proxy.py | 107 ++++++++++++++---- src/benchflow/providers/litellm_logging.py | 4 + src/benchflow/providers/litellm_runtime.py | 13 ++- src/benchflow/providers/runtime.py | 2 + src/benchflow/rollout/__init__.py | 1 + .../trajectories/llm_capture_records.py | 16 ++- tests/continue_run/test_replay_proxy.py | 50 ++++++++ tests/test_litellm_hardening.py | 3 + tests/test_litellm_logging.py | 4 + tests/test_litellm_runtime.py | 36 ++++++ .../test_native_capture_resilience.py | 52 +++++++++ 11 files changed, 267 insertions(+), 21 deletions(-) diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 009d5d5a7..6266b9962 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -137,6 +137,12 @@ class ReplayState: live_attempt_count: int = 0 live_error_count: int = 0 lock: threading.Lock = field(default_factory=threading.Lock) + condition: threading.Condition = field(init=False) + quiescing: bool = False + active_live_requests: int = 0 + + def __post_init__(self): + self.condition = threading.Condition(self.lock) def _write_state(self): payload = { @@ -177,7 +183,11 @@ def _check_divergence(self, incoming, recorded_request): print(message, file=sys.stderr, flush=True) def next_response(self, request_body): - with self.lock: + with self.condition: + if self.quiescing: + return "error", 503, { + "error": {"message": "replay proxy is quiescing"} + } if self.cursor < len(self.recorded): exchange = self.recorded[self.cursor] self._check_divergence( @@ -189,22 +199,44 @@ def next_response(self, request_body): return "replay", int(response.get("status_code") or 200), dict(response.get("body") or {}) self.cursor += 1 self.live_attempt_count += 1 - self._write_state() - - status, body, provider_observed = self._forward_live(request_body) - if provider_observed: + self.active_live_requests += 1 try: - self._append_live_exchange(request_body, status, body) + self._write_state() except Exception: + self.active_live_requests -= 1 + self.condition.notify_all() + raise + + try: + status, body, provider_observed = self._forward_live(request_body) + if provider_observed: + try: + self._append_live_exchange(request_body, status, body) + except Exception: + with self.lock: + self.live_error_count += 1 + self._write_state() + raise + else: with self.lock: self.live_error_count += 1 self._write_state() - raise - else: - with self.lock: - self.live_error_count += 1 - self._write_state() - return "live", status, body + return "live", status, body + finally: + with self.condition: + self.active_live_requests -= 1 + self.condition.notify_all() + + def quiesce(self, timeout=610): + deadline = time.monotonic() + timeout + with self.condition: + self.quiescing = True + while self.active_live_requests: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self.condition.wait(remaining) + return True def _forward_live(self, request_body): forwarded = dict(request_body) @@ -285,6 +317,13 @@ def do_GET(self): def do_POST(self): path = self.path.split("?", 1)[0] + if path == "/benchflow/quiesce": + quiesced = self.state.quiesce() + self._send_json( + 200 if quiesced else 503, + {"status": "quiesced" if quiesced else "quiesce_timeout"}, + ) + return if path not in ("/v1/chat/completions", "/chat/completions"): self._send_json(404, {"error": {"message": f"not found: {path}"}}) return @@ -494,22 +533,52 @@ async def _wait_until_ready(self) -> None: ) async def stop(self) -> None: + await self._quiesce() + await self._terminate() self.live_exchanges = await self._load_live_exchanges() await self._load_live_state() await self._load_runtime_errors() - with contextlib.suppress(Exception): - await self.sandbox.exec( - f"if [ -s {shlex.quote(self.pid_path)} ]; then " - f"kill -TERM $(cat {shlex.quote(self.pid_path)}) 2>/dev/null || true; " - "fi", - timeout_sec=10, - ) with contextlib.suppress(Exception): await self.sandbox.exec( f"rm -rf {shlex.quote(self.runtime_dir)}", timeout_sec=10, ) + async def _quiesce(self) -> None: + command = ( + "python3 - <<'PY'\n" + "import urllib.request\n" + "request = urllib.request.Request(\n" + f" 'http://127.0.0.1:{self.port}/benchflow/quiesce',\n" + " data=b'{}', method='POST')\n" + "urllib.request.urlopen(request, timeout=615).read()\n" + "PY" + ) + try: + result = await self.sandbox.exec(command, timeout_sec=620) + except Exception as exc: + self.live_errors.append(f"sandbox replay quiesce failed: {exc}") + return + if result.return_code != 0: + detail = (result.stderr or result.stdout or "unavailable").strip() + self.live_errors.append(f"sandbox replay quiesce failed: {detail}") + + async def _terminate(self) -> None: + command = ( + f"if [ -s {shlex.quote(self.pid_path)} ]; then " + f"pid=$(cat {shlex.quote(self.pid_path)}); " + 'kill -TERM "$pid" 2>/dev/null || true; i=0; ' + 'while kill -0 "$pid" 2>/dev/null; do ' + '[ "$i" -ge 100 ] && exit 1; i=$((i + 1)); sleep 0.1; done; fi' + ) + try: + result = await self.sandbox.exec(command, timeout_sec=15) + except Exception as exc: + self.live_errors.append(f"sandbox replay termination failed: {exc}") + return + if result.return_code != 0: + self.live_errors.append("sandbox replay termination did not quiesce") + async def _load_live_exchanges(self) -> list[LLMExchange]: text = await _read_remote_text(self.sandbox, self.live_log_path) exchanges: list[LLMExchange] = [] diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index b8c421bcf..80228a5a1 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -259,6 +259,8 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - request_body[key] = value request_body = {k: v for k, v in request_body.items() if v is not None} return { + "benchflow_agent": os.environ.get("BENCHFLOW_LITELLM_AGENT"), + "benchflow_role": os.environ.get("BENCHFLOW_LITELLM_ROLE"), "benchflow_requested_model": os.environ.get( "BENCHFLOW_LITELLM_REQUESTED_MODEL" ), @@ -453,6 +455,8 @@ def _exchange_metadata( metadata = { key: record.get(key) for key in ( + "benchflow_agent", + "benchflow_role", "benchflow_requested_model", "benchflow_model_alias", "request_model", diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index fd5a79343..3e9a73f2e 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -64,6 +64,8 @@ LITELLM_SANDBOX_ROOT = "/tmp/benchflow-litellm" _CALLBACK_MODULE = "benchflow_litellm_callback" _LITELLM_REQUESTED_MODEL_ENV = "BENCHFLOW_LITELLM_REQUESTED_MODEL" +_LITELLM_AGENT_ENV = "BENCHFLOW_LITELLM_AGENT" +_LITELLM_ROLE_ENV = "BENCHFLOW_LITELLM_ROLE" _PATCH_MODULE = "benchflow_litellm_bedrock_patch" # The proxy is an internal single-route gateway — it must never register the @@ -842,6 +844,7 @@ async def _start_host_litellm( environment: str, session_id: str, agent_name: str, + role_name: str | None = None, ) -> HostLiteLLMProcess: runtime_dir = Path(tempfile.mkdtemp(prefix="benchflow-litellm-")) log_path = runtime_dir / "callback.jsonl" @@ -860,6 +863,8 @@ async def _start_host_litellm( "BENCHFLOW_LITELLM_LOG_PATH": str(log_path), LITELLM_MODEL_ALIAS_ENV: route.model_alias, _LITELLM_REQUESTED_MODEL_ENV: route.requested_model, + _LITELLM_AGENT_ENV: agent_name, + _LITELLM_ROLE_ENV: role_name or "primary", **_PROXY_DOCS_DISABLE_ENV, } ) @@ -1162,6 +1167,7 @@ async def _start_sandbox_litellm( agent_env: dict[str, str], session_id: str, agent_name: str, + role_name: str | None = None, install_timeout_sec: int = 600, ) -> SandboxLiteLLMProcess: token = uuid4().hex[:16] @@ -1183,6 +1189,8 @@ async def _start_sandbox_litellm( "BENCHFLOW_LITELLM_LOG_PATH": paths["log"], LITELLM_MODEL_ALIAS_ENV: route.model_alias, _LITELLM_REQUESTED_MODEL_ENV: route.requested_model, + _LITELLM_AGENT_ENV: agent_name, + _LITELLM_ROLE_ENV: role_name or "primary", **_PROXY_DOCS_DISABLE_ENV, } ) @@ -1591,6 +1599,7 @@ async def ensure_litellm_runtime( required_skill_names: tuple[str, ...] = (), live_trajectory_path: Path | None = None, force_sandbox_local: bool = False, + role_name: str | None = None, ) -> tuple[dict[str, str], Any | None]: """Start/reuse LiteLLM and rewrite the agent env to talk to it. @@ -1659,7 +1668,7 @@ async def ensure_litellm_runtime( proxy_location = "sandbox" if sandbox_local else "host" config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" - f"{session_id}:{skill_gate_key}" + f"{session_id}:{role_name or 'primary'}:{skill_gate_key}" ) if runtime is not None and getattr(runtime, "kind", None) == "litellm": server = getattr(runtime, "server", None) @@ -1694,6 +1703,7 @@ async def ensure_litellm_runtime( agent_env=proxy_env, session_id=session_id, agent_name=agent, + role_name=role_name, install_timeout_sec=max(600, int(sandbox_setup_timeout)), ) else: @@ -1704,6 +1714,7 @@ async def ensure_litellm_runtime( environment=environment, session_id=session_id, agent_name=agent, + role_name=role_name, ) except BedrockPatchPreflightError: raise diff --git a/src/benchflow/providers/runtime.py b/src/benchflow/providers/runtime.py index e35d24b41..7e9e397c5 100644 --- a/src/benchflow/providers/runtime.py +++ b/src/benchflow/providers/runtime.py @@ -51,6 +51,7 @@ async def ensure_litellm_runtime( required_skill_names: tuple[str, ...] = (), live_trajectory_path: Path | None = None, force_sandbox_local: bool = False, + role_name: str | None = None, ) -> tuple[dict[str, str], ProviderRuntime | None]: from benchflow.providers.litellm_runtime import ( ensure_litellm_runtime as _ensure_litellm_runtime, @@ -69,6 +70,7 @@ async def ensure_litellm_runtime( required_skill_names=required_skill_names, live_trajectory_path=live_trajectory_path, force_sandbox_local=force_sandbox_local, + role_name=role_name, ) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 31d89b5cc..541da3360 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -2378,6 +2378,7 @@ async def connect_as(self, role: Role) -> None: required_skill_names=getattr(self, "_required_skill_names", ()), live_trajectory_path=rollout_dir / "trajectory" / "llm_trajectory.jsonl", force_sandbox_local=disallow_web_tools, + role_name=role.name, ) role_agent_differs = role.agent != cfg.primary_agent diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 7bbd8c044..147827f2c 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -370,8 +370,22 @@ def _target_for_provider_record( ) -> CaptureTarget | None: if len(targets) == 1: return targets[0] - candidates = {_record_model(record)} metadata = _record_metadata(record) + stamped_role = metadata.get("benchflow_role") + stamped_agent = metadata.get("benchflow_agent") + if isinstance(stamped_role, str) and stamped_role: + identity_matches = [ + target + for target in targets + if target.role == stamped_role + and ( + not isinstance(stamped_agent, str) + or not stamped_agent + or target.agent == stamped_agent + ) + ] + return identity_matches[0] if len(identity_matches) == 1 else None + candidates = {_record_model(record)} candidates.update( value for key in ( diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index 439ec1b51..fff84cdc5 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import threading import httpx import pytest @@ -146,6 +147,55 @@ def fail_replace(*_args, **_kwargs): assert not state_path.exists() +def test_sandbox_quiesce_waits_for_live_handler_before_snapshot(tmp_path) -> None: + """Guards PR #1057 against snapshotting before live calls are quiescent.""" + namespace: dict[str, object] = {} + exec(_sandbox_proxy_source(), namespace) + state = namespace["ReplayState"]( + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(tmp_path / "live.jsonl"), + state_path=str(tmp_path / "state.json"), + port=61357, + ) + forward_started = threading.Event() + release_forward = threading.Event() + responses: list[tuple] = [] + quiesced: list[bool] = [] + + def forward(_request): + forward_started.set() + assert release_forward.wait(timeout=5) + return 200, completion(content="done"), True + + state._forward_live = forward + worker = threading.Thread( + target=lambda: responses.append( + state.next_response({"messages": [{"role": "user"}]}) + ) + ) + worker.start() + assert forward_started.wait(timeout=5) + barrier = threading.Thread(target=lambda: quiesced.append(state.quiesce(timeout=5))) + barrier.start() + assert barrier.is_alive() + + release_forward.set() + worker.join(timeout=5) + barrier.join(timeout=5) + + assert quiesced == [True] + assert responses[0][0] == "live" + assert state.live_attempt_count == 1 + assert len((tmp_path / "live.jsonl").read_text().splitlines()) == 1 + late = state.next_response({"messages": [{"role": "user"}]}) + assert late[0] == "error" + assert late[1] == 503 + assert state.live_attempt_count == 1 + + @pytest.mark.asyncio async def test_sandbox_attempt_journal_marker_reaches_host( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index 9cdb1da5f..6e618486e 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -356,6 +356,7 @@ async def test_sandbox_litellm_launch_keeps_secrets_off_command_line(): }, session_id="s", agent_name="openhands", + role_name="solver", ) # launch_config is uploaded as a file (proxy needs the key)... @@ -369,6 +370,8 @@ async def test_sandbox_litellm_launch_keeps_secrets_off_command_line(): assert launch_config["env"]["BENCHFLOW_LITELLM_MODEL_ALIAS"] == ( "benchflow-minimax-MiniMax-M3" ) + assert launch_config["env"]["BENCHFLOW_LITELLM_AGENT"] == "openhands" + assert launch_config["env"]["BENCHFLOW_LITELLM_ROLE"] == "solver" assert set(sandbox.uploaded_modes.values()) == {"600"} # ...and the secret never appears on any exec command line (/proc exposure). assert all(secret not in call for call in sandbox.exec_calls) diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index c622c226e..02c892efe 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -125,6 +125,8 @@ def test_callback_record_preserves_benchflow_route_identity( "BENCHFLOW_LITELLM_MODEL_ALIAS", "benchflow-azure-foundry-openai-gpt-5.5", ) + monkeypatch.setenv("BENCHFLOW_LITELLM_AGENT", "codex-acp") + monkeypatch.setenv("BENCHFLOW_LITELLM_ROLE", "reviewer") logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() now = datetime.now() @@ -136,6 +138,8 @@ def test_callback_record_preserves_benchflow_route_identity( assert record["benchflow_requested_model"] == ("azure-foundry-openai/gpt-5.5") assert record["benchflow_model_alias"] == ("benchflow-azure-foundry-openai-gpt-5.5") + assert record["benchflow_agent"] == "codex-acp" + assert record["benchflow_role"] == "reviewer" def test_callback_module_source_exposes_proxy_handler_instance(): diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 21f934b77..b100020d5 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -310,6 +310,42 @@ async def fake_start(**kwargs): assert created[0].stopped is True +@pytest.mark.asyncio +async def test_runtime_is_not_reused_across_same_model_roles(monkeypatch): + """Guards PR #1057 by binding callback identity to one scene role.""" + created = [] + + async def fake_start(**kwargs): + server = FakeLiteLLMServer("http://127.0.0.1:4000", kwargs["route"]) + created.append((kwargs["role_name"], server)) + return server + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + env = {"OPENAI_API_KEY": "sk-openai"} + _updated, first = await ensure_litellm_runtime( + agent="opencode", + agent_env=env, + model="openai/gpt-4.1-mini", + runtime=None, + environment="local", + session_id="run-1", + role_name="solver", + ) + _updated, second = await ensure_litellm_runtime( + agent="opencode", + agent_env=env, + model="openai/gpt-4.1-mini", + runtime=first, + environment="local", + session_id="run-1", + role_name="reviewer", + ) + + assert second is not first + assert [role for role, _server in created] == ["solver", "reviewer"] + assert created[0][1].stopped is True + + @pytest.mark.asyncio async def test_required_usage_fails_when_litellm_lacks_provider_key(monkeypatch): monkeypatch.setattr(runtime_mod, "uses_native_subscription_auth", lambda *_: False) diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 314b339a6..01dd1847c 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -146,6 +146,58 @@ def test_provider_role_attribution_uses_proxy_model_aliases(tmp_path: Path) -> N ) +def test_provider_role_attribution_uses_runtime_identity_for_same_model( + tmp_path: Path, +) -> None: + """Guards PR #1057 for API-key scene roles sharing one model route.""" + model = "openai/gpt-5.5" + targets = [ + _CaptureTarget( + agent="opencode", + model=model, + credential_home=f"/home/{role}", + auth_mode=AuthMode.API_KEY, + native=False, + role=role, + ) + for role in ("solver", "reviewer") + ] + trajectory = tmp_path / "llm_trajectory.jsonl" + trajectory.write_text( + "".join( + json.dumps( + { + "request": {"body": {"model": "gpt-5.5"}}, + "response": {"status_code": 200, "body": {}}, + "metadata": { + "benchflow_agent": "opencode", + "benchflow_role": role, + "benchflow_requested_model": model, + }, + } + ) + + "\n" + for role in ("solver", "reviewer") + ) + ) + + records = load_provider_wire_records( + trajectory, + targets=targets, + fallback_agent="opencode", + fallback_model=model, + fallback_auth=AuthMode.API_KEY, + ) + + assert [record["metadata"]["role"] for record in records] == [ + "solver", + "reviewer", + ] + assert all( + record["metadata"]["role_attribution_complete"] is True for record in records + ) + + @pytest.mark.asyncio async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( tmp_path: Path, From 11e016827f4f86e18fe09091107d48a0281911f8 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 04:32:17 -0700 Subject: [PATCH 26/74] fix: journal late calls and refine role matching --- src/benchflow/continue_run/sandbox_proxy.py | 3 ++ .../trajectories/llm_capture_records.py | 7 ++- tests/continue_run/test_replay_proxy.py | 6 ++- .../test_native_capture_resilience.py | 51 +++++++++++++++++++ 4 files changed, 64 insertions(+), 3 deletions(-) diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 6266b9962..71a1339fa 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -185,6 +185,9 @@ def _check_divergence(self, incoming, recorded_request): def next_response(self, request_body): with self.condition: if self.quiescing: + self.live_attempt_count += 1 + self.live_error_count += 1 + self._write_state() return "error", 503, { "error": {"message": "replay proxy is quiescing"} } diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 147827f2c..33d80b992 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -374,7 +374,7 @@ def _target_for_provider_record( stamped_role = metadata.get("benchflow_role") stamped_agent = metadata.get("benchflow_agent") if isinstance(stamped_role, str) and stamped_role: - identity_matches = [ + targets = [ target for target in targets if target.role == stamped_role @@ -384,7 +384,10 @@ def _target_for_provider_record( or target.agent == stamped_agent ) ] - return identity_matches[0] if len(identity_matches) == 1 else None + if len(targets) == 1: + return targets[0] + if not targets: + return None candidates = {_record_model(record)} candidates.update( value diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index fff84cdc5..7a6e42bd3 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -193,7 +193,11 @@ def forward(_request): late = state.next_response({"messages": [{"role": "user"}]}) assert late[0] == "error" assert late[1] == 503 - assert state.live_attempt_count == 1 + assert state.live_attempt_count == 2 + assert state.live_error_count == 1 + capture_state = json.loads((tmp_path / "state.json").read_text()) + assert capture_state["live_attempt_count"] == 2 + assert capture_state["live_error_count"] == 1 @pytest.mark.asyncio diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 01dd1847c..c4eb78ac7 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -198,6 +198,57 @@ def test_provider_role_attribution_uses_runtime_identity_for_same_model( ) +def test_provider_role_attribution_intersects_repeated_role_with_model( + tmp_path: Path, +) -> None: + """Guards PR #1057 when scenes reuse a role name with different models.""" + targets = [ + _CaptureTarget( + agent="opencode", + model=model, + credential_home=f"/home/{index}", + auth_mode=AuthMode.API_KEY, + native=False, + role="solver", + ) + for index, model in enumerate(("openai/gpt-5.5", "openai/gpt-5.6")) + ] + trajectory = tmp_path / "llm_trajectory.jsonl" + trajectory.write_text( + "".join( + json.dumps( + { + "request": {"body": {"model": model.rsplit("/", 1)[-1]}}, + "response": {"status_code": 200, "body": {}}, + "metadata": { + "benchflow_agent": "opencode", + "benchflow_role": "solver", + "benchflow_requested_model": model, + }, + } + ) + + "\n" + for model in ("openai/gpt-5.5", "openai/gpt-5.6") + ) + ) + + records = load_provider_wire_records( + trajectory, + targets=targets, + fallback_agent="opencode", + fallback_model=None, + fallback_auth=AuthMode.API_KEY, + ) + + assert [record["metadata"]["model"] for record in records] == [ + "openai/gpt-5.5", + "openai/gpt-5.6", + ] + assert all( + record["metadata"]["role_attribution_complete"] is True for record in records + ) + + @pytest.mark.asyncio async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( tmp_path: Path, From 6fa31517bd0b1b31adbab368981202c3d696fd8c Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 05:04:08 -0700 Subject: [PATCH 27/74] fix: close capture trust and shutdown gaps --- docs/getting-started.md | 11 ++- src/benchflow/continue_run/orchestrator.py | 27 +++++- src/benchflow/continue_run/sandbox_proxy.py | 43 ++++++++- .../providers/litellm_capture_lifecycle.py | 89 +++++++++++++++++++ src/benchflow/providers/litellm_logging.py | 48 +++++++++- src/benchflow/providers/litellm_runtime.py | 89 +++++++++++++++---- src/benchflow/trajectories/llm_capture.py | 56 ++++++++++-- .../trajectories/llm_capture_manifest.py | 74 ++++++++++++++- .../trajectories/llm_capture_records.py | 14 +++ .../trajectories/native_capture_parsers.py | 7 +- tests/continue_run/test_replay_proxy.py | 67 ++++++++++++++ tests/test_litellm_hardening.py | 46 ++++++++++ .../test_llm_capture_training_contract.py | 51 ++++++++++- tests/trajectories/test_native_llm_capture.py | 18 +++- .../test_native_session_boundaries.py | 52 +++++++++++ 15 files changed, 644 insertions(+), 48 deletions(-) create mode 100644 src/benchflow/providers/litellm_capture_lifecycle.py diff --git a/docs/getting-started.md b/docs/getting-started.md index 88459d081..ff92ed73f 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -221,7 +221,7 @@ the source of truth for interpreting the JSONL: | Agent/auth path | Primary source | `capture_fidelity` | |---|---|---| | API key through the BenchFlow gateway (including Azure) | LiteLLM provider request/response capture | `provider_wire` | -| Claude Code subscription/OAuth | Claude Code raw API-body files correlated by local OTLP logs | `provider_wire` | +| Claude Code subscription/OAuth | Claude Code raw API-body files correlated by local OTLP logs | `agent_session` | | Claude Code subscription/OAuth fallback | Claude Code native session JSONL | `agent_session` | | Codex subscription/OAuth | Codex native session JSONL | `agent_session` | @@ -233,9 +233,12 @@ count. Missing or ambiguously attributed roles make the rollout-level capture `partial`. Reconstructed `agent_session` rows remain useful for audit and viewer workflows, but trainer exports fail closed unless the manifest says the capture -is complete provider-wire data. Claude's own raw-body telemetry can still -contain provider-redacted extended-thinking blocks; BenchFlow also applies its -normal secret redaction before publishing the JSONL. +is complete provider-wire data and every successful exchange has positive token +usage. Claude's raw-body surface is deliberately audit-only because Claude Code +writes it inside the agent-controlled sandbox; its API shape does not make its +custody provider-trusted. It can also contain provider-redacted +extended-thinking blocks. BenchFlow applies its normal secret redaction before +publishing the JSONL. ### Reading results diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 5f81effde..86a146acc 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -49,8 +49,9 @@ CaptureStatus, LLMRoleCapture, LLMTrajectoryManifest, - capture_manifest_allows_training, + capture_artifact_allows_training, read_llm_trajectory_manifest, + successful_exchanges_have_positive_usage, write_llm_trajectory_manifest, ) from benchflow.trajectories.types import LLMExchange, redact_trajectory_obj @@ -381,18 +382,37 @@ def refresh_stitched_trajectory_manifest( except ValueError: malformed_count += 1 rows_valid = malformed_count == 0 + parsed_rows: list[dict[str, Any]] = [] + if rows_valid: + parsed_rows = [json.loads(line) for line in trajectory_lines] expected_count = n_recorded + live_attempt_count count_matches = exchange_count == expected_count + source_rows: list[dict[str, Any]] = [] + try: + source_rows = [ + json.loads(line) + for line in (source_rollout_dir / "trajectory" / "llm_trajectory.jsonl") + .read_text() + .splitlines() + if line.strip() + ] + except (OSError, json.JSONDecodeError): + source_rows = [] source_allows_training = bool( source_raw is not None - and capture_manifest_allows_training(source_raw, exchange_count=n_recorded) + and len(source_rows) == n_recorded + and capture_artifact_allows_training(source_raw, exchanges=source_rows) ) live_capture_complete = live_attempt_count == n_live and not live_errors + usage_complete = bool( + parsed_rows and successful_exchanges_have_positive_usage(parsed_rows) + ) complete = ( source_allows_training and count_matches and live_capture_complete and rows_valid + and usage_complete ) source_capture_source = source.capture_source if source else CaptureSource.NONE @@ -450,6 +470,9 @@ def refresh_stitched_trajectory_manifest( f"stitched LLM trajectory contains {malformed_count} malformed row(s)" ) missing_fields.append("valid_provider_exchange") + if rows_valid and not usage_complete: + errors.append("stitched LLM trajectory lacks positive provider token usage") + missing_fields.append("token_usage") if not live_capture_complete: missing_fields.append("live_provider_exchange") diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 71a1339fa..b3987a891 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -140,6 +140,7 @@ class ReplayState: condition: threading.Condition = field(init=False) quiescing: bool = False active_live_requests: int = 0 + active_handlers: int = 0 def __post_init__(self): self.condition = threading.Condition(self.lock) @@ -234,13 +235,26 @@ def quiesce(self, timeout=610): deadline = time.monotonic() + timeout with self.condition: self.quiescing = True - while self.active_live_requests: + while self.active_live_requests or self.active_handlers > 1: remaining = deadline - time.monotonic() if remaining <= 0: return False self.condition.wait(remaining) return True + def begin_quiesce(self): + with self.condition: + self.quiescing = True + + def handler_started(self): + with self.condition: + self.active_handlers += 1 + + def handler_finished(self): + with self.condition: + self.active_handlers -= 1 + self.condition.notify_all() + def _forward_live(self, request_body): forwarded = dict(request_body) forwarded["model"] = self.upstream_model @@ -321,6 +335,13 @@ def do_GET(self): def do_POST(self): path = self.path.split("?", 1)[0] if path == "/benchflow/quiesce": + # Stop the accept loop before the barrier can report success. The + # server counts connections in ``process_request`` (before their + # worker threads start), so every already-accepted handler is in + # ``active_handlers`` even if it has not reached ``do_POST`` yet. + self.state.begin_quiesce() + self.server.shutdown() + self.server.server_close() quiesced = self.state.quiesce() self._send_json( 200 if quiesced else 503, @@ -361,10 +382,30 @@ def do_POST(self): class ReplayServer(ThreadingHTTPServer): + # ``/benchflow/quiesce`` closes the listener from its own worker thread. + # Do not let ``server_close`` try to join that current thread; the normal + # non-daemon handler lifecycle still keeps the process alive until the + # response and every already-accepted request have finished. + block_on_close = False + def __init__(self, address, handler, state): super().__init__(address, handler) self.state = state + def process_request(self, request, client_address): + self.state.handler_started() + try: + super().process_request(request, client_address) + except BaseException: + self.state.handler_finished() + raise + + def process_request_thread(self, request, client_address): + try: + super().process_request_thread(request, client_address) + finally: + self.state.handler_finished() + def main(): cfg = json.load(open(sys.argv[1], encoding="utf-8")) diff --git a/src/benchflow/providers/litellm_capture_lifecycle.py b/src/benchflow/providers/litellm_capture_lifecycle.py new file mode 100644 index 000000000..850ba3328 --- /dev/null +++ b/src/benchflow/providers/litellm_capture_lifecycle.py @@ -0,0 +1,89 @@ +"""Fail-closed capture accounting and graceful LiteLLM shutdown helpers.""" + +from __future__ import annotations + +import asyncio +import contextlib +import shlex +import subprocess +from typing import Any + +LITELLM_CAPTURE_STATE_ENV = "BENCHFLOW_LITELLM_CAPTURE_STATE_PATH" +PROVIDER_DRAIN_TIMEOUT_SEC = 610 + + +def capture_journal_error( + payload: dict[str, Any] | None, + *, + exchange_count: int, +) -> str | None: + """Validate durable accepted/terminal counts against imported JSONL rows.""" + + if payload is None: + return "LiteLLM capture attempt journal is missing" + attempt_count = payload.get("attempt_count") + terminal_count = payload.get("terminal_count") + if ( + not isinstance(attempt_count, int) + or isinstance(attempt_count, bool) + or attempt_count < 0 + or not isinstance(terminal_count, int) + or isinstance(terminal_count, bool) + or terminal_count < 0 + ): + return "LiteLLM capture attempt journal is malformed" + if attempt_count != terminal_count or terminal_count != exchange_count: + return ( + "LiteLLM capture did not drain every accepted provider request: " + f"attempts={attempt_count}, terminal={terminal_count}, " + f"exchanges={exchange_count}" + ) + return None + + +async def drain_host_process(process: subprocess.Popen[bytes]) -> str | None: + """Stop host acceptance and wait for the proxy to drain active requests.""" + + if process.poll() is not None: + return None + process.terminate() + try: + await asyncio.to_thread(process.wait, PROVIDER_DRAIN_TIMEOUT_SEC) + except subprocess.TimeoutExpired: + process.kill() + await asyncio.to_thread(process.wait, 10) + return "LiteLLM graceful shutdown timed out before provider requests drained" + return None + + +async def drain_sandbox_process( + sandbox: Any, + *, + pid_path: str, +) -> str | None: + """Stop sandbox acceptance and wait for the proxy to drain active requests.""" + + drain = await sandbox.exec( + ( + f"if [ ! -s {shlex.quote(pid_path)} ]; then exit 0; fi\n" + f"read -r pid < {shlex.quote(pid_path)}\n" + 'kill -TERM "$pid" 2>/dev/null || exit 0\n' + f"for attempt in $(seq 1 {PROVIDER_DRAIN_TIMEOUT_SEC * 10}); do\n" + ' if ! kill -0 "$pid" 2>/dev/null; then exit 0; fi\n' + ' if [ -r "/proc/$pid/stat" ] && ' + "[ \"$(awk '{print $3}' /proc/$pid/stat)\" = Z ]; then exit 0; fi\n" + " sleep 0.1\n" + "done\n" + "exit 124" + ), + timeout_sec=PROVIDER_DRAIN_TIMEOUT_SEC + 10, + ) + if drain.return_code == 0: + return None + with contextlib.suppress(Exception): + await sandbox.exec( + f"if [ -s {shlex.quote(pid_path)} ]; then " + f"kill -KILL $(cat {shlex.quote(pid_path)}) 2>/dev/null || true; fi", + timeout_sec=10, + ) + return "LiteLLM graceful shutdown timed out before provider requests drained" diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 80228a5a1..95b8bd20d 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -87,6 +87,7 @@ def callback_module_source() -> str: import json import os import re +import threading import time import traceback from datetime import datetime, timezone @@ -223,14 +224,50 @@ def _failure_traceback(detail: Any) -> str: class BenchFlowLiteLLMLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._capture_lock = threading.Lock() + self._attempt_count = 0 + self._terminal_count = 0 + with self._capture_lock: + self._write_state_locked() + + def _write_state_locked(self) -> None: + path = os.environ.get("BENCHFLOW_LITELLM_CAPTURE_STATE_PATH") + if not path: + return + payload = { + "attempt_count": self._attempt_count, + "terminal_count": self._terminal_count, + } + os.makedirs(os.path.dirname(path), exist_ok=True) + temporary = path + ".tmp" + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle, separators=(",", ":")) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + + def _journal_attempt(self) -> None: + with self._capture_lock: + self._attempt_count += 1 + self._write_state_locked() + def _write(self, payload: dict[str, Any]) -> None: path = os.environ.get("BENCHFLOW_LITELLM_LOG_PATH") if not path: return payload["logged_at"] = datetime.now(timezone.utc).isoformat() - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "a", encoding="utf-8") as handle: - handle.write(json.dumps(_jsonable(payload), separators=(",", ":")) + "\n") + with self._capture_lock: + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "a", encoding="utf-8") as handle: + handle.write( + json.dumps(_jsonable(payload), separators=(",", ":")) + "\n" + ) + handle.flush() + os.fsync(handle.fileno()) + self._terminal_count += 1 + self._write_state_locked() def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) -> dict[str, Any]: litellm_params = kwargs.get("litellm_params") or {} @@ -292,6 +329,11 @@ async def async_pre_call_hook( if not isinstance(data, dict): return None + # This hook is awaited before LiteLLM forwards the provider request. + # Persist the attempt first; the success/failure callback advances the + # terminal counter only after its full JSONL row is durable. + self._journal_attempt() + _gate_opencode_skill_catalog(data) cleaned = data diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 3e9a73f2e..f11308abb 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -36,6 +36,12 @@ preflight_sandbox_bedrock_patch, route_requires_bedrock_patch, ) +from benchflow.providers.litellm_capture_lifecycle import ( + LITELLM_CAPTURE_STATE_ENV, + capture_journal_error, + drain_host_process, + drain_sandbox_process, +) from benchflow.providers.litellm_config import ( LITELLM_MASTER_KEY_ENV, LITELLM_MODEL_ALIAS_ENV, @@ -410,12 +416,14 @@ def __init__( stderr_path: Path, session_id: str, agent_name: str, + capture_state_path: Path | None = None, ) -> None: self.route = route self.process = process self.runtime_dir = runtime_dir self.endpoint = endpoint self.log_path = log_path + self.capture_state_path = capture_state_path self.stdout_path = stdout_path self.stderr_path = stderr_path self.session_id = session_id @@ -430,19 +438,22 @@ async def is_running(self) -> bool: return self.process.poll() is None async def stop(self) -> None: + drain_error = await drain_host_process(self.process) await self._stop_live_capture() await _await_log_stable(self._log_size) - if self.process.poll() is None: - self.process.terminate() - try: - await asyncio.to_thread(self.process.wait, 10) - except subprocess.TimeoutExpired: - self.process.kill() - await asyncio.to_thread(self.process.wait, 10) self._load_callback_log() self.reconcile_live_capture() + trajectory = self.trajectory + assert trajectory is not None + journal_error = capture_journal_error( + self._load_capture_state(), + exchange_count=len(trajectory.exchanges), + ) with contextlib.suppress(Exception): shutil.rmtree(self.runtime_dir, ignore_errors=True) + error = drain_error or journal_error + if error is not None: + raise RuntimeError(error) def _log_size(self) -> int: try: @@ -479,6 +490,20 @@ def _load_callback_log(self) -> None: agent_name=self.agent_name, ) + def _load_capture_state(self) -> dict[str, Any] | None: + if self.capture_state_path is None: + trajectory = self.trajectory + assert trajectory is not None + return { + "attempt_count": len(trajectory.exchanges), + "terminal_count": len(trajectory.exchanges), + } + try: + payload = json.loads(self.capture_state_path.read_text()) + except (OSError, json.JSONDecodeError): + return None + return payload if isinstance(payload, dict) else None + def log_tail(self) -> str: chunks: list[str] = [] for label, path in (("stdout", self.stdout_path), ("stderr", self.stderr_path)): @@ -503,12 +528,14 @@ def __init__( stderr_path: str, session_id: str, agent_name: str, + capture_state_path: str | None = None, ) -> None: self.sandbox = sandbox self.route = route self.runtime_dir = runtime_dir self.endpoint = endpoint self.log_path = log_path + self.capture_state_path = capture_state_path self.pid_path = pid_path self.stdout_path = stdout_path self.stderr_path = stderr_path @@ -532,23 +559,27 @@ async def is_running(self) -> bool: return result.return_code == 0 and (result.stdout or "").strip() == "yes" async def stop(self) -> None: + drain_error = await drain_sandbox_process( + self.sandbox, + pid_path=self.pid_path, + ) await self._stop_live_capture() await _await_log_stable(self._remote_log_size) - with contextlib.suppress(Exception): - await self.sandbox.exec( - ( - f"if [ -s {shlex.quote(self.pid_path)} ]; then " - f"kill -TERM $(cat {shlex.quote(self.pid_path)}) 2>/dev/null || true; " - "fi" - ), - timeout_sec=10, - ) await self._load_callback_log() self.reconcile_live_capture() + trajectory = self.trajectory + assert trajectory is not None + journal_error = capture_journal_error( + await self._load_capture_state(), + exchange_count=len(trajectory.exchanges), + ) with contextlib.suppress(Exception): await self.sandbox.exec( f"rm -rf {shlex.quote(self.runtime_dir)}", timeout_sec=10 ) + error = drain_error or journal_error + if error is not None: + raise RuntimeError(error) async def _remote_log_size(self) -> int: with contextlib.suppress(Exception): @@ -635,6 +666,26 @@ async def _load_callback_log(self) -> None: agent_name=self.agent_name, ) + async def _load_capture_state(self) -> dict[str, Any] | None: + if self.capture_state_path is None: + trajectory = self.trajectory + assert trajectory is not None + return { + "attempt_count": len(trajectory.exchanges), + "terminal_count": len(trajectory.exchanges), + } + result = await self.sandbox.exec( + f"cat {shlex.quote(self.capture_state_path)} 2>/dev/null || true", + timeout_sec=15, + ) + if result.return_code != 0: + return None + try: + payload = json.loads(result.stdout or "") + except json.JSONDecodeError: + return None + return payload if isinstance(payload, dict) else None + async def log_tail(self) -> str: chunks: list[str] = [] for label, path in (("stdout", self.stdout_path), ("stderr", self.stderr_path)): @@ -848,6 +899,7 @@ async def _start_host_litellm( ) -> HostLiteLLMProcess: runtime_dir = Path(tempfile.mkdtemp(prefix="benchflow-litellm-")) log_path = runtime_dir / "callback.jsonl" + capture_state_path = runtime_dir / "capture_state.json" stdout_path = runtime_dir / "stdout.log" stderr_path = runtime_dir / "stderr.log" port = _find_free_port() @@ -861,6 +913,7 @@ async def _start_host_litellm( "PYTHONPATH": f"{runtime_dir}{os.pathsep}{env.get('PYTHONPATH', '')}", "LITELLM_MASTER_KEY": master_key, "BENCHFLOW_LITELLM_LOG_PATH": str(log_path), + LITELLM_CAPTURE_STATE_ENV: str(capture_state_path), LITELLM_MODEL_ALIAS_ENV: route.model_alias, _LITELLM_REQUESTED_MODEL_ENV: route.requested_model, _LITELLM_AGENT_ENV: agent_name, @@ -896,6 +949,7 @@ async def _start_host_litellm( runtime_dir=runtime_dir, endpoint=_agent_endpoint_for_environment(port, environment, bind), log_path=log_path, + capture_state_path=capture_state_path, stdout_path=stdout_path, stderr_path=stderr_path, session_id=session_id, @@ -1002,6 +1056,7 @@ async def _upload_runtime_files_to_sandbox( "stdout": f"{runtime_dir}/stdout.log", "stderr": f"{runtime_dir}/stderr.log", "log": f"{runtime_dir}/callback.jsonl", + "capture_state": f"{runtime_dir}/capture_state.json", "pid": f"{runtime_dir}/litellm.pid", "state": f"{runtime_dir}/state.json", "venv": f"{runtime_dir}/venv", @@ -1187,6 +1242,7 @@ async def _start_sandbox_litellm( "PYTHONPATH": f"{runtime_dir}:{env.get('PYTHONPATH', '')}", "LITELLM_MASTER_KEY": master_key, "BENCHFLOW_LITELLM_LOG_PATH": paths["log"], + LITELLM_CAPTURE_STATE_ENV: paths["capture_state"], LITELLM_MODEL_ALIAS_ENV: route.model_alias, _LITELLM_REQUESTED_MODEL_ENV: route.requested_model, _LITELLM_AGENT_ENV: agent_name, @@ -1260,6 +1316,7 @@ async def _start_sandbox_litellm( runtime_dir=runtime_dir, endpoint=endpoint, log_path=paths["log"], + capture_state_path=paths["capture_state"], pid_path=paths["pid"], stdout_path=paths["stdout"], stderr_path=paths["stderr"], diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index d08a8a0e8..0a3e6dbac 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -121,8 +121,16 @@ def __init__( session_id=session_id, started_at=started_at, ) - self._targets: dict[tuple[str, str, str | None, str], _CaptureTarget] = {} - self._provisional_target_key: tuple[str, str, str | None, str] | None = None + self._targets: dict[ + tuple[str, str, str | None, str, AuthMode], _CaptureTarget + ] = {} + self._active_target_keys: dict[ + tuple[str, str, str | None, str], + tuple[str, str, str | None, str, AuthMode], + ] = {} + self._provisional_target_key: ( + tuple[str, str, str | None, str, AuthMode] | None + ) = None self._collector_started = False self._collector_owned = False self._capture_root_prepared = False @@ -170,6 +178,7 @@ async def prepare_agent( model=model, credential_home=credential_home, role_name=role_name, + auth_mode=auth_mode, ) previous_target = self._targets.get(target_key) if ( @@ -187,20 +196,27 @@ async def prepare_agent( model=model, credential_home=credential_home, role_name=None, + auth_mode=auth_mode, + ) + base_key = _capture_target_base_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=role_name, ) if role_name is None: - if ( - self._provisional_target_key is not None - and self._provisional_target_key != primary_key - ): - self._targets.pop(self._provisional_target_key, None) self._targets[primary_key] = target self._provisional_target_key = primary_key else: if self._provisional_target_key is not None: - self._targets.pop(self._provisional_target_key, None) + removed_key = self._provisional_target_key + self._targets.pop(removed_key, None) + for active_base, active_key in list(self._active_target_keys.items()): + if active_key == removed_key: + self._active_target_keys.pop(active_base, None) self._provisional_target_key = None self._targets[target_key] = target + self._active_target_keys[base_key] = target_key self._refresh_manifest_auth_mode() if not native: write_llm_trajectory_manifest(self.rollout_dir, self.manifest) @@ -255,12 +271,15 @@ def bind_native_session( ) -> None: """Bind an ACP session ID to its prepared native capture target.""" - key = _capture_target_key( + base_key = _capture_target_base_key( agent=agent, model=model, credential_home=credential_home, role_name=role_name, ) + key = self._active_target_keys.get(base_key) + if key is None: + return target = self._targets.get(key) if target is None: return @@ -812,6 +831,25 @@ def _capture_target_key( model: str | None, credential_home: str, role_name: str | None, + auth_mode: AuthMode, +) -> tuple[str, str, str | None, str, AuthMode]: + return ( + *_capture_target_base_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=role_name, + ), + auth_mode, + ) + + +def _capture_target_base_key( + *, + agent: str, + model: str | None, + credential_home: str, + role_name: str | None, ) -> tuple[str, str, str | None, str]: return (role_name or "primary", agent, model, credential_home) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 19e9fa78d..a27925654 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -165,13 +165,81 @@ def capture_artifact_allows_training( """Apply the sidecar contract while retaining genuine legacy JSONL support.""" if manifest is not None: - return capture_manifest_allows_training( - manifest, - exchange_count=len(exchanges), + return bool( + capture_manifest_allows_training( + manifest, + exchange_count=len(exchanges), + ) + and successful_exchanges_have_positive_usage(exchanges) ) return not any(_exchange_requires_manifest(exchange) for exchange in exchanges) +def successful_exchanges_have_positive_usage( + exchanges: Sequence[dict[str, Any]], +) -> bool: + """Require positive token evidence for every successful provider response.""" + + successful = [ + exchange + for exchange in exchanges + if isinstance((response := exchange.get("response")), dict) + and isinstance(response.get("status_code"), int) + and 200 <= response["status_code"] < 300 + ] + return bool(successful) and all( + _exchange_has_positive_usage(exchange) for exchange in successful + ) + + +def _exchange_has_positive_usage(exchange: dict[str, Any]) -> bool: + response = exchange.get("response") + body = response.get("body") if isinstance(response, dict) else None + if not isinstance(body, dict): + return False + for container_name in ("usage", "usageMetadata"): + container = body.get(container_name) + if isinstance(container, dict) and _usage_payload_has_positive_tokens( + container + ): + return True + return False + + +def _usage_payload_has_positive_tokens(payload: dict[str, Any]) -> bool: + token_keys = { + "input_tokens", + "output_tokens", + "prompt_tokens", + "completion_tokens", + "total_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "inputTokens", + "outputTokens", + "totalTokens", + "promptTokenCount", + "candidatesTokenCount", + "totalTokenCount", + "cachedContentTokenCount", + "toolUsePromptTokenCount", + "thoughtsTokenCount", + "cached_tokens", + } + for key, value in payload.items(): + if key in token_keys: + if isinstance(value, bool): + continue + try: + if int(value) > 0: + return True + except (TypeError, ValueError): + continue + if isinstance(value, dict) and _usage_payload_has_positive_tokens(value): + return True + return False + + def _exchange_requires_manifest(exchange: dict[str, Any]) -> bool: metadata = exchange.get("metadata") if not isinstance(metadata, dict): diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 33d80b992..d2e551701 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -16,6 +16,7 @@ CaptureSource, CaptureStatus, LLMRoleCapture, + successful_exchanges_have_positive_usage, ) from benchflow.trajectories.native_capture_parsers import NativeParseResult from benchflow.trajectories.types import redact_trajectory_obj @@ -190,6 +191,18 @@ def assemble_capture( missing_fields = { field for bundle in native_bundles for field in bundle.result.missing_fields } + successful_records = [ + record + for record in records + if isinstance((response := record.get("response")), dict) + and isinstance(response.get("status_code"), int) + and 200 <= response["status_code"] < 300 + ] + usage_complete = not successful_records or successful_exchanges_have_positive_usage( + successful_records + ) + if not usage_complete: + missing_fields.add("token_usage") if attribution_incomplete: missing_fields.add("role_attribution") source = _aggregate_source(records) @@ -215,6 +228,7 @@ def assemble_capture( if fidelity is CaptureFidelity.PROVIDER_WIRE and request_complete and response_complete + and usage_complete and not errors else CaptureStatus.PARTIAL ) diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index fde5a92dc..0835acbfe 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -111,7 +111,10 @@ def parse_claude_raw_capture( response_timestamp=response_timestamp, path="/v1/messages", source=CaptureSource.CLAUDE_OTEL_RAW_BODY, - fidelity=CaptureFidelity.PROVIDER_WIRE, + # Claude writes this raw diagnostic surface from inside the + # agent sandbox. Its shape is exact, but its custody is not a + # trusted provider boundary, so it remains audit-only. + fidelity=CaptureFidelity.AGENT_SESSION, auth_mode="oauth_subscription", request_complete=True, response_complete=True, @@ -161,7 +164,7 @@ def parse_claude_raw_capture( exchanges=exchanges, ), source=CaptureSource.CLAUDE_OTEL_RAW_BODY, - fidelity=CaptureFidelity.PROVIDER_WIRE, + fidelity=CaptureFidelity.AGENT_SESSION, request_complete=not pairing_ambiguous, response_complete=True, errors=errors, diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index 7a6e42bd3..b15386311 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -200,6 +200,73 @@ def forward(_request): assert capture_state["live_error_count"] == 1 +def test_sandbox_quiesce_closes_listener_and_drains_accepted_handlers( + tmp_path, +) -> None: + """Guards PR #1057 against terminating a late rejected handler mid-journal.""" + + namespace: dict[str, object] = {} + exec(_sandbox_proxy_source(), namespace) + state = namespace["ReplayState"]( + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(tmp_path / "live.jsonl"), + state_path=str(tmp_path / "state.json"), + port=0, + ) + forward_started = threading.Event() + release_forward = threading.Event() + + def forward(_request): + forward_started.set() + assert release_forward.wait(timeout=5) + return 200, completion(content="done"), True + + state._forward_live = forward + server = namespace["ReplayServer"]( + ("127.0.0.1", 0), namespace["ReplayHandler"], state + ) + port = server.server_address[1] + server_thread = threading.Thread(target=server.serve_forever) + server_thread.start() + chat_responses: list[httpx.Response] = [] + chat_thread = threading.Thread( + target=lambda: chat_responses.append( + httpx.post( + f"http://127.0.0.1:{port}/v1/chat/completions", + json={"messages": [{"role": "user"}]}, + timeout=5, + ) + ) + ) + chat_thread.start() + assert forward_started.wait(timeout=5) + quiesce_responses: list[httpx.Response] = [] + quiesce_thread = threading.Thread( + target=lambda: quiesce_responses.append( + httpx.post( + f"http://127.0.0.1:{port}/benchflow/quiesce", + timeout=5, + ) + ) + ) + quiesce_thread.start() + assert quiesce_thread.is_alive() + + release_forward.set() + chat_thread.join(timeout=5) + quiesce_thread.join(timeout=5) + server_thread.join(timeout=5) + + assert chat_responses[0].status_code == 200 + assert quiesce_responses[0].status_code == 200 + assert state.active_handlers == 0 + with pytest.raises(httpx.ConnectError): + httpx.get(f"http://127.0.0.1:{port}/health", timeout=1) + + @pytest.mark.asyncio async def test_sandbox_attempt_journal_marker_reaches_host( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index 6e618486e..ad8578886 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -328,6 +328,10 @@ async def exec( return _ExecResult(0) if command.strip().startswith("rm -rf"): return _ExecResult(0) + if command.strip().startswith("cat") and "capture_state.json" in command: + return _ExecResult( + 0, stdout=json.dumps({"attempt_count": 1, "terminal_count": 1}) + ) if command.strip().startswith("cat") and "state.json" in command: if self._started: return _ExecResult(0, stdout=json.dumps({"pid": 4242, "port": 45999})) @@ -445,6 +449,41 @@ async def test_sandbox_litellm_stop_imports_usage_and_cleans_up(): assert any(call.strip().startswith("rm -rf") for call in sandbox.exec_calls) +@pytest.mark.asyncio +async def test_sandbox_litellm_stop_rejects_an_undrained_attempt_journal(): + """Guards PR #1057 against completing an accepted provider call lost at stop.""" + + route = resolve_litellm_route( + "minimax/MiniMax-M3", + {"MINIMAX_API_KEY": "k", "MINIMAX_BASE_URL": "https://api.minimax.io/v1"}, + ) + sandbox = _FakeSandbox() + original_exec = sandbox.exec + + async def mismatched_state(command, **kwargs): + if command.strip().startswith("cat") and "capture_state.json" in command: + return _ExecResult( + 0, stdout=json.dumps({"attempt_count": 2, "terminal_count": 1}) + ) + return await original_exec(command, **kwargs) + + sandbox.exec = mismatched_state + proc = await runtime_mod._start_sandbox_litellm( + sandbox=sandbox, + route=route, + master_key="sk-master", + agent_env={ + "MINIMAX_API_KEY": "k", + "MINIMAX_BASE_URL": "https://api.minimax.io/v1", + }, + session_id="s", + agent_name="openhands", + ) + + with pytest.raises(RuntimeError, match="did not drain every accepted"): + await proc.stop() + + @pytest.mark.asyncio async def test_sandbox_litellm_startup_failure_tears_down(): route = resolve_litellm_route( @@ -701,7 +740,9 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( logger = namespace["proxy_handler_instance"] log_path = tmp_path / "callback.jsonl" + state_path = tmp_path / "capture_state.json" monkeypatch.setenv("BENCHFLOW_LITELLM_LOG_PATH", str(log_path)) + monkeypatch.setenv("BENCHFLOW_LITELLM_CAPTURE_STATE_PATH", str(state_path)) response = { "model": "gpt-4.1-mini", @@ -718,6 +759,7 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( start = datetime(2026, 6, 4, 10, 0, 0) end = datetime(2026, 6, 4, 10, 0, 1) + await logger.async_pre_call_hook(None, None, kwargs, "acompletion") await logger.async_log_success_event(kwargs, response, start, end) text = log_path.read_text() @@ -732,6 +774,10 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( assert usage["usage_source"] == "provider_response" assert usage["n_input_tokens"] == 12 assert usage["n_output_tokens"] == 4 + assert json.loads(state_path.read_text()) == { + "attempt_count": 1, + "terminal_count": 1, + } def test_gemini_usage_metadata_is_detected_as_provider_usage(): diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 51f3aa434..dc01b5ce8 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -36,6 +36,7 @@ def _write_exchange( *, fidelity: str, schema_version: int | None = None, + include_usage: bool = True, ) -> None: metadata: dict[str, str | bool | int] = { "capture_fidelity": fidelity, @@ -44,6 +45,16 @@ def _write_exchange( } if schema_version is not None: metadata["schema_version"] = schema_version + response_body = { + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + } + if include_usage: + response_body["usage"] = { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + } (trajectory_dir / "llm_trajectory.jsonl").write_text( json.dumps( { @@ -52,10 +63,7 @@ def _write_exchange( }, "response": { "status_code": 200, - "body": { - "role": "assistant", - "content": [{"type": "text", "text": "hi"}], - }, + "body": response_body, }, "metadata": metadata, } @@ -161,6 +169,41 @@ def test_manifest_count_mismatch_fails_closed_for_canonical_results( assert row["is_completed"] is False +def test_provider_capture_without_positive_usage_is_not_training_ready( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on provider rows with zero token evidence.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange( + trajectory_dir, + fidelity="provider_wire", + include_usage=False, + ) + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "complete", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + } + ) + ) + + row = _build_results_row( + tmp_path, + agent_result={"usage_source": "unavailable", "total_tokens": 0}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is False + + def test_mixed_oauth_audit_capture_preserves_successful_completion( tmp_path: Path, ) -> None: diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index 7a5c25341..9b1c7180f 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -52,8 +52,8 @@ def _otel_record(name: str, timestamp_ns: int, **attributes: str) -> dict: } -def test_claude_otel_raw_bodies_become_provider_wire_exchanges(tmp_path: Path) -> None: - """Guards PR #1057's exact Claude OAuth raw-body capture contract.""" +def test_claude_otel_raw_bodies_remain_audit_only_exchanges(tmp_path: Path) -> None: + """Guards PR #1057 against trusting agent-writable Claude OAuth telemetry.""" capture = tmp_path / "capture" raw = capture / "raw" @@ -116,7 +116,7 @@ def test_claude_otel_raw_bodies_become_provider_wire_exchanges(tmp_path: Path) - assert result is not None assert result.source is CaptureSource.CLAUDE_OTEL_RAW_BODY - assert result.fidelity is CaptureFidelity.PROVIDER_WIRE + assert result.fidelity is CaptureFidelity.AGENT_SESSION assert result.request_complete is True assert result.response_complete is True assert len(result.trajectory.exchanges) == 1 @@ -594,7 +594,17 @@ async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> json.dumps( { "request": {"body": {"input": "hello"}}, - "response": {"status_code": 200, "body": {"output": []}}, + "response": { + "status_code": 200, + "body": { + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + }, + }, + }, } ) + "\n" diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index ffa2dff2b..3efc20ddd 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -229,6 +229,58 @@ async def test_executed_primary_is_not_removed_by_later_named_role( } +@pytest.mark.asyncio +async def test_repeated_role_preserves_auth_distinct_capture_targets( + tmp_path: Path, +) -> None: + """Guards PR #1057 against API auth replacing a bound OAuth role target.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + common = { + "env": None, + "agent": "codex-acp", + "model": "gpt-5.6", + "credential_home": "/home/agent", + "sandbox_user": "agent", + "role_name": "reviewer", + } + await capture.prepare_agent( + **common, + agent_env={ + "CODEX_AUTH_JSON": ( + '{"auth_mode":"chatgpt","tokens":{"refresh_token":"test"}}' + ) + }, + ) + capture.bind_native_session( + agent="codex-acp", + model="gpt-5.6", + credential_home="/home/agent", + native_session_id="oauth-session", + role_name="reviewer", + ) + await capture.prepare_agent( + **common, + agent_env={"OPENAI_API_KEY": "test-key"}, + ) + + targets = list(capture._targets.values()) + assert {target.auth_mode for target in targets} == { + AuthMode.OAUTH_SUBSCRIPTION, + AuthMode.API_KEY, + } + oauth_target = next( + target for target in targets if target.auth_mode is AuthMode.OAUTH_SUBSCRIPTION + ) + assert oauth_target.native_session_ids == ("oauth-session",) + + @pytest.mark.asyncio async def test_claude_capture_setup_clears_reused_raw_attempt( tmp_path: Path, From 767dcdefaeafb1c024b046ef94bfa2fa895271a6 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 05:15:56 -0700 Subject: [PATCH 28/74] fix: retire quiesce handler before acknowledgement --- src/benchflow/continue_run/sandbox_proxy.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index b3987a891..fe0245806 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -343,6 +343,11 @@ def do_POST(self): self.server.shutdown() self.server.server_close() quiesced = self.state.quiesce() + # The control request is the one handler deliberately excluded + # from the barrier above. Retire it from the accepted-handler + # accounting before publishing success so the response is also a + # reliable, race-free signal that the capture lifecycle drained. + self.server.release_current_handler() self._send_json( 200 if quiesced else 503, {"status": "quiesced" if quiesced else "quiesce_timeout"}, @@ -391,6 +396,7 @@ class ReplayServer(ThreadingHTTPServer): def __init__(self, address, handler, state): super().__init__(address, handler) self.state = state + self._handler_local = threading.local() def process_request(self, request, client_address): self.state.handler_started() @@ -401,9 +407,17 @@ def process_request(self, request, client_address): raise def process_request_thread(self, request, client_address): + self._handler_local.counted = True try: super().process_request_thread(request, client_address) finally: + if self._handler_local.counted: + self.state.handler_finished() + self._handler_local.counted = False + + def release_current_handler(self): + if getattr(self._handler_local, "counted", False): + self._handler_local.counted = False self.state.handler_finished() From 9c788d066a374e937a50f9ec747634b62b01bb2a Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 05:32:47 -0700 Subject: [PATCH 29/74] fix: invalidate stale LiteLLM attempt journals --- src/benchflow/providers/litellm_logging.py | 40 ++++++++++++++++++---- tests/test_litellm_hardening.py | 32 +++++++++++++++++ 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 95b8bd20d..7b67e5e63 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -240,13 +240,41 @@ def _write_state_locked(self) -> None: "attempt_count": self._attempt_count, "terminal_count": self._terminal_count, } - os.makedirs(os.path.dirname(path), exist_ok=True) temporary = path + ".tmp" - with open(temporary, "w", encoding="utf-8") as handle: - json.dump(payload, handle, separators=(",", ":")) - handle.flush() - os.fsync(handle.fileno()) - os.replace(temporary, path) + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle, separators=(",", ":")) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + except Exception: + # A previous, balanced journal must never survive a failed later + # update: cleanup could otherwise accept its stale counts and omit + # the rejected call. Both host and sandbox proxies execute this + # same embedded callback module. + self._invalidate_state_locked(path, temporary) + raise + + def _invalidate_state_locked(self, path: str, temporary: str) -> None: + try: + os.unlink(path) + except FileNotFoundError: + pass + except OSError: + # If removal itself is unavailable, corrupt the validity contract + # explicitly. Cleanup rejects this payload as a malformed journal. + try: + with open(path, "w", encoding="utf-8") as handle: + json.dump({"valid": False}, handle, separators=(",", ":")) + handle.flush() + os.fsync(handle.fileno()) + except OSError: + pass + try: + os.unlink(temporary) + except OSError: + pass def _journal_attempt(self) -> None: with self._capture_lock: diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index ad8578886..792866be4 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -780,6 +780,38 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( } +@pytest.mark.asyncio +async def test_embedded_logger_invalidates_stale_journal_after_attempt_write_failure( + tmp_path, monkeypatch +): + """Guards PR #1057 against accepting stale host or sandbox journal counts.""" + + namespace: dict[str, object] = {} + exec(callback_module_source(), namespace) + logger = namespace["proxy_handler_instance"] + state_path = tmp_path / "capture_state.json" + monkeypatch.setenv("BENCHFLOW_LITELLM_CAPTURE_STATE_PATH", str(state_path)) + + logger._journal_attempt() + logger._terminal_count = 1 + logger._write_state_locked() + assert json.loads(state_path.read_text()) == { + "attempt_count": 1, + "terminal_count": 1, + } + + def fail_replace(*_args, **_kwargs): + raise OSError("disk full") + + monkeypatch.setattr(namespace["os"], "replace", fail_replace) + + with pytest.raises(OSError, match="disk full"): + logger._journal_attempt() + + assert not state_path.exists() + assert not state_path.with_suffix(".json.tmp").exists() + + def test_gemini_usage_metadata_is_detected_as_provider_usage(): # LiteLLM normally normalizes to OpenAI shape, but a raw Gemini passthrough # reports usageMetadata; it must not silently degrade to 'unavailable'. From b709ec193900740dac431bf2406f129a2caf454e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 05:54:57 -0700 Subject: [PATCH 30/74] fix: enforce provider capture custody boundary --- docs/agent-quickstart.md | 7 ++ docs/continue-runs.md | 6 ++ docs/running-any-benchmark.md | 7 ++ src/benchflow/continue_run/orchestrator.py | 29 ++++++-- src/benchflow/providers/litellm_runtime.py | 8 ++- src/benchflow/providers/runtime.py | 3 + src/benchflow/rollout/__init__.py | 20 ++++++ src/benchflow/trajectories/llm_capture.py | 33 +++++++++ .../trajectories/llm_capture_records.py | 19 ++++- tests/continue_run/test_orchestrator.py | 63 ++++++++++++++++ tests/test_litellm_runtime.py | 28 ++++++++ .../test_native_capture_resilience.py | 72 ++++++++++++++++++- 12 files changed, 288 insertions(+), 7 deletions(-) diff --git a/docs/agent-quickstart.md b/docs/agent-quickstart.md index 410fd0533..cf6aa5c5e 100644 --- a/docs/agent-quickstart.md +++ b/docs/agent-quickstart.md @@ -148,6 +148,13 @@ and what it is: verifier/reward.txt — raw verifier reward verifier/test-stdout.txt — verifier stdout (and ctrf.json when the test emits a CTRF report) + +`llm_trajectory.jsonl` is an audit artifact for every run, not an unconditional +training claim. Only a `complete` manifest with `provider_wire` fidelity and +positive provider token usage is training-ready. Native OAuth/session capture, +and a sandbox-local proxy that shares root custody with the agent, stay +available for audit but are marked lower-fidelity and excluded from training. + Also note the job-level summary.json plus aggregated results.jsonl, verifiers.jsonl, and adp.jsonl in the job directory. The trainer/ files and the job-level aggregates are written by current BenchFlow releases; if they are missing, diff --git a/docs/continue-runs.md b/docs/continue-runs.md index e32b35ab6..00f6933b1 100644 --- a/docs/continue-runs.md +++ b/docs/continue-runs.md @@ -37,6 +37,12 @@ request/response pairs from the original run. `bench eval continue`: with a stitched `llm_trajectory.jsonl` (recorded prefix + live suffix) and `continued_from` provenance — a drop-in replacement for the timed-out entry. +The stitched manifest preserves capture custody. A live provider proxy across a +trusted host/non-root boundary can contribute `provider_wire` rows. If replay +and the untrusted agent share root custody in the sandbox, the live suffix is +still written for audit but is labeled `agent_session`; the stitched run cannot +be exported as training-ready provider evidence. + Because the agent rebuilds its own state by re-doing its own steps, no reverse-engineering of OpenHands internals is needed, and the result is a single continuous run rather than a fresh agent on a warm filesystem. diff --git a/docs/running-any-benchmark.md b/docs/running-any-benchmark.md index da45a2777..5e5580966 100644 --- a/docs/running-any-benchmark.md +++ b/docs/running-any-benchmark.md @@ -224,6 +224,13 @@ Every layer terminates at the *same* output contract, written per rollout under | `trainer/adp.jsonl` | ADP trajectory record | | `verifier/` | Raw verifier output (`reward.txt`, `ctrf.json`, stdout) | +The JSONL file is always an audit artifact; its sidecar determines whether it +is also training-ready. Training requires complete `provider_wire` capture with +positive provider token usage. Native OAuth/session capture and sandbox-local +capture that shares root custody with the agent are retained but marked +lower-fidelity, so exporters reject them rather than treating agent-writable +records as trusted provider evidence. + Hosted runs share this artifact contract too (see the `hosted_env.py` module docstring), with `source.type="hosted_env"` / `trajectory_source="hosted_env"` marking the lineage. diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 86a146acc..0441d494e 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -307,6 +307,7 @@ def write_stitched_trajectory( live_exchanges: list[LLMExchange], *, live_model: str | None = None, + live_capture_trusted: bool = True, ) -> Path: """Write the stitched continuous trajectory into the new rollout folder.""" out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" @@ -322,7 +323,11 @@ def write_stitched_trajectory( "model": live_model, "auth_mode": AuthMode.API_KEY.value, "capture_source": CaptureSource.LITELLM_PROXY.value, - "capture_fidelity": CaptureFidelity.PROVIDER_WIRE.value, + "capture_fidelity": ( + CaptureFidelity.PROVIDER_WIRE.value + if live_capture_trusted + else CaptureFidelity.AGENT_SESSION.value + ), "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, "request_complete": True, "response_complete": True, @@ -330,6 +335,8 @@ def write_stitched_trajectory( "payload_redacted": True, } ) + if not live_capture_trusted: + metadata["capture_custody"] = "agent_writable_sandbox" lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) rendered = "\n".join(lines) + ("\n" if lines else "") temporary = out.with_suffix(out.suffix + ".tmp") @@ -348,6 +355,7 @@ def refresh_stitched_trajectory_manifest( n_live: int, live_attempt_count: int, live_errors: list[str], + live_capture_trusted: bool = True, ) -> LLMTrajectoryManifest: """Replace rollout-finalization provenance with the final stitched contract.""" @@ -411,6 +419,7 @@ def refresh_stitched_trajectory_manifest( source_allows_training and count_matches and live_capture_complete + and (live_capture_trusted or n_live == 0) and rows_valid and usage_complete ) @@ -424,10 +433,13 @@ def refresh_stitched_trajectory_manifest( if source_capture_source is CaptureSource.LITELLM_PROXY else CaptureSource.MIXED ) - capture_fidelity = ( + live_fidelity = ( CaptureFidelity.PROVIDER_WIRE - if source_fidelity is CaptureFidelity.PROVIDER_WIRE - else CaptureFidelity.MIXED + if live_capture_trusted + else CaptureFidelity.AGENT_SESSION + ) + capture_fidelity = ( + live_fidelity if source_fidelity is live_fidelity else CaptureFidelity.MIXED ) auth_mode = ( AuthMode.API_KEY if source_auth is AuthMode.API_KEY else AuthMode.MIXED @@ -454,6 +466,8 @@ def refresh_stitched_trajectory_manifest( errors = list(source.errors) if source and not complete else [] missing_fields = list(source.missing_fields) if source and not complete else [] errors.extend(live_errors) + if n_live and not live_capture_trusted: + errors.append("sandbox replay capture shared root custody with the agent") if source is None: errors.append("source LLM trajectory manifest is missing or malformed") missing_fields.append("source_capture_provenance") @@ -888,6 +902,9 @@ async def _write_artifacts( run.path / "trajectory" / "llm_trajectory.jsonl", live_exchanges, live_model=live_model, + live_capture_trusted=bool( + config.sandbox_user and config.sandbox_user not in {"root", "0"} + ), ) refresh_stitched_trajectory_manifest( rollout_dir, @@ -903,6 +920,9 @@ async def _write_artifacts( *(replay_proxy.live_errors if replay_proxy is not None else []), *(teardown_errors or []), ], + live_capture_trusted=bool( + config.sandbox_user and config.sandbox_user not in {"root", "0"} + ), ) update_continued_metadata( rollout_dir, @@ -930,6 +950,7 @@ async def _write_artifacts( session_id=rollout_name, usage_tracking="required", sandbox=rollout.env, + sandbox_user=config.sandbox_user, ) replay_proxy = await SandboxReplayProxy.start( sandbox=rollout.env, diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index f11308abb..def1eb89b 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1657,6 +1657,7 @@ async def ensure_litellm_runtime( live_trajectory_path: Path | None = None, force_sandbox_local: bool = False, role_name: str | None = None, + sandbox_user: str | None = None, ) -> tuple[dict[str, str], Any | None]: """Start/reuse LiteLLM and rewrite the agent env to talk to it. @@ -1723,9 +1724,13 @@ async def ensure_litellm_runtime( sorted(set(required_skill_names)), separators=(",", ":") ) proxy_location = "sandbox" if sandbox_local else "host" + capture_trusted = not sandbox_local or bool( + sandbox_user and sandbox_user not in {"root", "0"} + ) config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" - f"{session_id}:{role_name or 'primary'}:{skill_gate_key}" + f"{session_id}:{role_name or 'primary'}:{sandbox_user or 'root'}:" + f"{skill_gate_key}" ) if runtime is not None and getattr(runtime, "kind", None) == "litellm": server = getattr(runtime, "server", None) @@ -1794,6 +1799,7 @@ async def ensure_litellm_runtime( server=server, config_key=config_key, master_key=master_key, + capture_trusted=capture_trusted, ) if live_trajectory_path is not None: server.start_live_capture(live_trajectory_path) diff --git a/src/benchflow/providers/runtime.py b/src/benchflow/providers/runtime.py index 7e9e397c5..5c757d724 100644 --- a/src/benchflow/providers/runtime.py +++ b/src/benchflow/providers/runtime.py @@ -31,6 +31,7 @@ class ProviderRuntime: server: LiteLLMProcess | None = None config_key: str | None = None master_key: str | None = None + capture_trusted: bool = True @property def base_url(self) -> str: @@ -52,6 +53,7 @@ async def ensure_litellm_runtime( live_trajectory_path: Path | None = None, force_sandbox_local: bool = False, role_name: str | None = None, + sandbox_user: str | None = None, ) -> tuple[dict[str, str], ProviderRuntime | None]: from benchflow.providers.litellm_runtime import ( ensure_litellm_runtime as _ensure_litellm_runtime, @@ -71,6 +73,7 @@ async def ensure_litellm_runtime( live_trajectory_path=live_trajectory_path, force_sandbox_local=force_sandbox_local, role_name=role_name, + sandbox_user=sandbox_user, ) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 541da3360..59a1e2f70 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1322,7 +1322,19 @@ async def connect(self) -> None: required_skill_names=getattr(self, "_required_skill_names", ()), live_trajectory_path=rollout_dir / "trajectory" / "llm_trajectory.jsonl", force_sandbox_local=getattr(self, "_disallow_web_tools", False), + sandbox_user=cfg.sandbox_user, ) + llm_capture = getattr(self, "_llm_capture", None) + if llm_capture is not None: + credential_home = ( + f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" + ) + llm_capture.bind_provider_capture_trust( + agent=cfg.primary_agent, + model=cfg.primary_model, + credential_home=credential_home, + trusted=getattr(self._usage_runtime, "capture_trusted", True), + ) sf_entrypoint = self._session_factory_entrypoint(cfg.primary_agent) self._is_session_factory = sf_entrypoint is not None if sf_entrypoint is not None: @@ -2379,6 +2391,7 @@ async def connect_as(self, role: Role) -> None: live_trajectory_path=rollout_dir / "trajectory" / "llm_trajectory.jsonl", force_sandbox_local=disallow_web_tools, role_name=role.name, + sandbox_user=cfg.sandbox_user, ) role_agent_differs = role.agent != cfg.primary_agent @@ -2437,6 +2450,13 @@ async def connect_as(self, role: Role) -> None: sandbox_user=cfg.sandbox_user, role_name=role.name, ) + llm_capture.bind_provider_capture_trust( + agent=role.agent, + model=role.model, + credential_home=cred_home, + trusted=getattr(self._usage_runtime, "capture_trusted", True), + role_name=role.name, + ) self._agent_launch = agent_launch diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 0a3e6dbac..563a47ed5 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -181,6 +181,8 @@ async def prepare_agent( auth_mode=auth_mode, ) previous_target = self._targets.get(target_key) + if previous_target is not None and not previous_target.provider_capture_trusted: + target = replace(target, provider_capture_trusted=False) if ( previous_target is not None and previous_target.native @@ -300,6 +302,37 @@ def bind_native_session( return self._targets[key] = replace(target, native_session_ids=session_ids) + def bind_provider_capture_trust( + self, + *, + agent: str, + model: str | None, + credential_home: str, + trusted: bool, + role_name: str | None = None, + ) -> None: + """Bind the gateway custody boundary to its prepared capture target.""" + + base_key = _capture_target_base_key( + agent=agent, + model=model, + credential_home=credential_home, + role_name=role_name, + ) + key = self._active_target_keys.get(base_key) + if key is None: + return + target = self._targets.get(key) + if target is None or target.native: + return + self._targets[key] = replace( + target, + # Once any provider rows for this target were collected under + # shared root custody, later trusted placement cannot recover the + # target-level training claim without per-row custody evidence. + provider_capture_trusted=(target.provider_capture_trusted and trusted), + ) + def _refresh_manifest_auth_mode(self) -> None: modes = {target.auth_mode for target in self._targets.values()} if len(modes) == 1: diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index d2e551701..64e59675f 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -33,6 +33,7 @@ class CaptureTarget: native: bool role: str = "agent" native_session_ids: tuple[str, ...] = () + provider_capture_trusted: bool = True @dataclass(frozen=True) @@ -89,11 +90,18 @@ def load_provider_wire_records( record["metadata"] = metadata target = _target_for_provider_record(targets, record) attribution_complete = target is not None or not targets + capture_trusted = ( + target.provider_capture_trusted if target is not None else True + ) metadata.update( { "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, "capture_source": CaptureSource.LITELLM_PROXY.value, - "capture_fidelity": CaptureFidelity.PROVIDER_WIRE.value, + "capture_fidelity": ( + CaptureFidelity.PROVIDER_WIRE.value + if capture_trusted + else CaptureFidelity.AGENT_SESSION.value + ), "auth_mode": ( target.auth_mode.value if target is not None @@ -120,6 +128,8 @@ def load_provider_wire_records( "payload_redacted": True, } ) + if not capture_trusted: + metadata["capture_custody"] = "agent_writable_sandbox" if not attribution_complete: metadata["role_candidates"] = _role_candidates(targets) records.append(redact_trajectory_obj(record)) @@ -142,6 +152,13 @@ def assemble_capture( ] records = _sort_exchange_records([*provider_records, *native_records]) errors = list(collection_errors) + if any( + _record_metadata(record).get("capture_custody") == "agent_writable_sandbox" + for record in provider_records + ): + errors.append( + "sandbox-local LiteLLM capture shared root custody with the agent" + ) attribution_incomplete = any( _record_metadata(record).get("role_attribution_complete") is False for record in records diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index ace611005..769c5b804 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -296,6 +296,69 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) ) +def test_root_sandbox_live_suffix_is_retained_but_audit_only(tmp_path): + """Guards PR #1057 against trusting root-writable continuation capture.""" + + model = "openai/gpt-5.5" + source = write_run_folder( + tmp_path / "source", + exchanges=[exchange(completion(content="recorded"))], + model=model, + ) + source_manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="source", + exchange_count=1, + request_complete=True, + response_complete=True, + started_at="2026-08-29T00:00:00Z", + finished_at="2026-08-29T00:01:00Z", + ) + write_llm_trajectory_manifest(source, source_manifest) + + rollout = tmp_path / "continued" + initialize_llm_trajectory_artifacts( + rollout, + agent="openhands", + model=None, + session_id="continued", + started_at=source_manifest.finished_at, + ) + out = write_stitched_trajectory( + rollout, + source / "trajectory" / "llm_trajectory.jsonl", + [exchange(completion(content="live"))], + live_model=model, + live_capture_trusted=False, + ) + manifest = refresh_stitched_trajectory_manifest( + rollout, + source, + original_model=model, + live_model=model, + n_recorded=1, + n_live=1, + live_attempt_count=1, + live_errors=[], + live_capture_trusted=False, + ) + + live_row = json.loads(out.read_text().splitlines()[-1]) + assert live_row["metadata"]["capture_fidelity"] == "agent_session" + assert live_row["metadata"]["capture_custody"] == "agent_writable_sandbox" + assert manifest.status is CaptureStatus.PARTIAL + assert manifest.capture_fidelity is CaptureFidelity.MIXED + assert any("shared root custody" in error for error in manifest.errors) + assert not capture_manifest_allows_training( + manifest.model_dump(mode="json"), exchange_count=2 + ) + + def test_refresh_stitched_manifest_rejects_missing_live_attempt(tmp_path): """Guards PR #1057 against completing a lost continuation exchange.""" diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index b100020d5..cffa9be4a 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -60,6 +60,7 @@ async def fake_start(**kwargs): assert provider_runtime is not None assert provider_runtime.kind == "litellm" assert provider_runtime.backend_model == "bedrock/us.anthropic.claude-opus-4-8" + assert provider_runtime.capture_trusted is True assert updated["OPENAI_BASE_URL"] == "http://host.docker.internal:32123/v1" assert updated["OPENAI_API_KEY"] == provider_runtime.master_key assert updated[LITELLM_MODEL_ALIAS_ENV] == ( @@ -160,10 +161,37 @@ async def fake_sandbox_start(**kwargs): assert starts[0]["sandbox"] is sandbox assert provider_runtime is not None assert provider_runtime.base_url == "http://127.0.0.1:45678" + assert provider_runtime.capture_trusted is False assert updated["LLM_BASE_URL"] == "http://127.0.0.1:45678/v1" assert updated["LLM_MODEL"].startswith("openai/benchflow-aws-bedrock") +@pytest.mark.asyncio +async def test_non_root_agent_keeps_sandbox_gateway_capture_trusted(monkeypatch): + """Guards PR #1057's root/non-root provider-capture custody boundary.""" + + async def fake_sandbox_start(**kwargs): + return FakeLiteLLMServer("http://127.0.0.1:45678", kwargs["route"]) + + monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) + _, provider_runtime = await ensure_litellm_runtime( + agent="openhands", + agent_env={ + "AWS_BEARER_TOKEN_BEDROCK": "token", + "AWS_REGION": "us-west-2", + }, + model="aws-bedrock/us.anthropic.claude-opus-4-8", + runtime=None, + environment="daytona", + session_id="run-non-root", + sandbox=SimpleNamespace(), + sandbox_user="agent", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is True + + @pytest.mark.asyncio async def test_apple_container_uses_sandbox_local_litellm(monkeypatch): """Guards PR #936 against handing the VM a host-loopback model endpoint.""" diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index c4eb78ac7..ff2fd7ff2 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -10,7 +10,11 @@ import pytest from benchflow.trajectories.llm_capture import LLMTrajectoryCapture, _CaptureTarget -from benchflow.trajectories.llm_capture_manifest import AuthMode +from benchflow.trajectories.llm_capture_manifest import ( + AuthMode, + CaptureFidelity, + CaptureStatus, +) from benchflow.trajectories.llm_capture_records import load_provider_wire_records from benchflow.trajectories.native_capture_parsers import parse_codex_sessions @@ -249,6 +253,72 @@ def test_provider_role_attribution_intersects_repeated_role_with_model( ) +@pytest.mark.asyncio +async def test_root_sandbox_provider_capture_is_retained_but_audit_only( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on root-agent-writable proxy capture.""" + + model = "openai/gpt-5.5" + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model=model, + session_id="rollout-root", + started_at=datetime(2026, 8, 29, 12, 0, tzinfo=UTC), + ) + await capture.prepare_agent( + None, + agent="codex-acp", + model=model, + agent_env={"OPENAI_API_KEY": "test-key"}, + credential_home="/root", + sandbox_user=None, + ) + capture.bind_provider_capture_trust( + agent="codex-acp", + model=model, + credential_home="/root", + trusted=False, + ) + capture.trajectory_path.write_text( + json.dumps( + { + "request": { + "body": { + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "solve"}], + } + }, + "response": { + "status_code": 200, + "body": { + "choices": [ + {"message": {"role": "assistant", "content": "done"}} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1}, + }, + }, + "metadata": { + "benchflow_agent": "codex-acp", + "benchflow_role": "primary", + "benchflow_requested_model": model, + }, + } + ) + + "\n" + ) + + await capture.finalize(None, acp_events=[], model_call_seen=True) + + record = json.loads(capture.trajectory_path.read_text()) + assert record["metadata"]["capture_fidelity"] == "agent_session" + assert record["metadata"]["capture_custody"] == "agent_writable_sandbox" + assert capture.manifest.status is CaptureStatus.PARTIAL + assert capture.manifest.capture_fidelity is CaptureFidelity.AGENT_SESSION + assert any("shared root custody" in error for error in capture.manifest.errors) + + @pytest.mark.asyncio async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( tmp_path: Path, From 89e720db1443ed1df264a5fd96e4a17bad0d851a Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 06:15:48 -0700 Subject: [PATCH 31/74] refactor: isolate continuation trajectory artifacts --- src/benchflow/continue_run/orchestrator.py | 297 +----------------- .../continue_run/trajectory_artifacts.py | 290 +++++++++++++++++ src/benchflow/providers/litellm_runtime.py | 8 +- .../trajectories/llm_capture_manifest.py | 8 + tests/continue_run/test_orchestrator.py | 8 +- 5 files changed, 321 insertions(+), 290 deletions(-) create mode 100644 src/benchflow/continue_run/trajectory_artifacts.py diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 0441d494e..f20e030ec 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -27,7 +27,6 @@ import re from collections.abc import Awaitable, Callable from dataclasses import dataclass -from datetime import UTC, datetime from pathlib import Path from typing import Any, cast @@ -38,23 +37,17 @@ SandboxReplayProxy, sandbox_replay_base_url, ) +from benchflow.continue_run.trajectory_artifacts import ( + refresh_stitched_trajectory_manifest, + write_stitched_trajectory, +) from benchflow.contracts import AgentProtocolError, SandboxStartupFailure from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.scenes import compile_scenes_to_steps from benchflow.trajectories.llm_capture_manifest import ( - LLM_TRAJECTORY_SCHEMA_VERSION, - AuthMode, - CaptureFidelity, - CaptureSource, - CaptureStatus, - LLMRoleCapture, - LLMTrajectoryManifest, - capture_artifact_allows_training, - read_llm_trajectory_manifest, - successful_exchanges_have_positive_usage, - write_llm_trajectory_manifest, + provider_capture_has_trusted_custody, ) -from benchflow.trajectories.types import LLMExchange, redact_trajectory_obj +from benchflow.trajectories.types import LLMExchange logger = logging.getLogger(__name__) @@ -267,272 +260,6 @@ def __call__(self, request_body: dict[str, Any]) -> dict[str, Any]: return dump() if callable(dump) else dict(response) -def stitched_trajectory_lines( - original_llm_trajectory: Path, live_exchanges: list[LLMExchange] -) -> list[str]: - """Build the continuous llm_trajectory: recorded prefix + live suffix. - - The recorded request/response payloads are preserved, while their metadata - is promoted to the current schema so a newly stitched artifact can never be - mistaken for sidecar-optional legacy data. The live suffix is redacted on - the way out. - """ - lines: list[str] = [] - if original_llm_trajectory.is_file(): - for raw in original_llm_trajectory.read_text().splitlines(): - if raw.strip(): - try: - payload = json.loads(raw) - except json.JSONDecodeError: - lines.append(raw) - continue - if not isinstance(payload, dict): - lines.append(raw) - continue - metadata = payload.setdefault("metadata", {}) - if not isinstance(metadata, dict): - metadata = {} - payload["metadata"] = metadata - metadata["schema_version"] = LLM_TRAJECTORY_SCHEMA_VERSION - lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) - for exchange in live_exchanges: - payload = redact_trajectory_obj(exchange.model_dump(mode="json")) - lines.append(json.dumps(payload, default=str)) - return lines - - -def write_stitched_trajectory( - rollout_dir: Path, - original_llm_trajectory: Path, - live_exchanges: list[LLMExchange], - *, - live_model: str | None = None, - live_capture_trusted: bool = True, -) -> Path: - """Write the stitched continuous trajectory into the new rollout folder.""" - out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" - out.parent.mkdir(parents=True, exist_ok=True) - lines = stitched_trajectory_lines(original_llm_trajectory, []) - for exchange in live_exchanges: - payload = exchange.model_dump(mode="json") - metadata = payload.setdefault("metadata", {}) - metadata.update( - { - "agent": "openhands", - "role": "agent", - "model": live_model, - "auth_mode": AuthMode.API_KEY.value, - "capture_source": CaptureSource.LITELLM_PROXY.value, - "capture_fidelity": ( - CaptureFidelity.PROVIDER_WIRE.value - if live_capture_trusted - else CaptureFidelity.AGENT_SESSION.value - ), - "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, - "request_complete": True, - "response_complete": True, - "role_attribution_complete": True, - "payload_redacted": True, - } - ) - if not live_capture_trusted: - metadata["capture_custody"] = "agent_writable_sandbox" - lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) - rendered = "\n".join(lines) + ("\n" if lines else "") - temporary = out.with_suffix(out.suffix + ".tmp") - temporary.write_text(rendered) - os.replace(temporary, out) - return out - - -def refresh_stitched_trajectory_manifest( - rollout_dir: Path, - source_rollout_dir: Path, - *, - original_model: str | None, - live_model: str | None, - n_recorded: int, - n_live: int, - live_attempt_count: int, - live_errors: list[str], - live_capture_trusted: bool = True, -) -> LLMTrajectoryManifest: - """Replace rollout-finalization provenance with the final stitched contract.""" - - current_raw = read_llm_trajectory_manifest(rollout_dir) - source_raw = read_llm_trajectory_manifest(source_rollout_dir) - try: - current = ( - LLMTrajectoryManifest.model_validate(current_raw) - if current_raw is not None - else None - ) - except ValueError: - current = None - try: - source = ( - LLMTrajectoryManifest.model_validate(source_raw) - if source_raw is not None - else None - ) - except ValueError: - source = None - - trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" - trajectory_lines = [ - line for line in trajectory_path.read_text().splitlines() if line.strip() - ] - exchange_count = len(trajectory_lines) - malformed_count = 0 - for line in trajectory_lines: - try: - LLMExchange.model_validate_json(line) - except ValueError: - malformed_count += 1 - rows_valid = malformed_count == 0 - parsed_rows: list[dict[str, Any]] = [] - if rows_valid: - parsed_rows = [json.loads(line) for line in trajectory_lines] - expected_count = n_recorded + live_attempt_count - count_matches = exchange_count == expected_count - source_rows: list[dict[str, Any]] = [] - try: - source_rows = [ - json.loads(line) - for line in (source_rollout_dir / "trajectory" / "llm_trajectory.jsonl") - .read_text() - .splitlines() - if line.strip() - ] - except (OSError, json.JSONDecodeError): - source_rows = [] - source_allows_training = bool( - source_raw is not None - and len(source_rows) == n_recorded - and capture_artifact_allows_training(source_raw, exchanges=source_rows) - ) - live_capture_complete = live_attempt_count == n_live and not live_errors - usage_complete = bool( - parsed_rows and successful_exchanges_have_positive_usage(parsed_rows) - ) - complete = ( - source_allows_training - and count_matches - and live_capture_complete - and (live_capture_trusted or n_live == 0) - and rows_valid - and usage_complete - ) - - source_capture_source = source.capture_source if source else CaptureSource.NONE - source_fidelity = source.capture_fidelity if source else CaptureFidelity.NONE - source_auth = source.auth_mode if source else AuthMode.UNKNOWN - if n_live: - capture_source = ( - CaptureSource.LITELLM_PROXY - if source_capture_source is CaptureSource.LITELLM_PROXY - else CaptureSource.MIXED - ) - live_fidelity = ( - CaptureFidelity.PROVIDER_WIRE - if live_capture_trusted - else CaptureFidelity.AGENT_SESSION - ) - capture_fidelity = ( - live_fidelity if source_fidelity is live_fidelity else CaptureFidelity.MIXED - ) - auth_mode = ( - AuthMode.API_KEY if source_auth is AuthMode.API_KEY else AuthMode.MIXED - ) - else: - capture_source = source_capture_source - capture_fidelity = source_fidelity - auth_mode = source_auth - - request_complete = bool( - source - and source.request_complete - and count_matches - and live_capture_complete - and rows_valid - ) - response_complete = bool( - source - and source.response_complete - and count_matches - and live_capture_complete - and rows_valid - ) - errors = list(source.errors) if source and not complete else [] - missing_fields = list(source.missing_fields) if source and not complete else [] - errors.extend(live_errors) - if n_live and not live_capture_trusted: - errors.append("sandbox replay capture shared root custody with the agent") - if source is None: - errors.append("source LLM trajectory manifest is missing or malformed") - missing_fields.append("source_capture_provenance") - elif not source_allows_training: - errors.append("source LLM trajectory is not complete provider-wire capture") - if not count_matches: - errors.append( - "stitched LLM trajectory count mismatch: " - f"expected {expected_count}, found {exchange_count}" - ) - missing_fields.append("exchange_count") - if not rows_valid: - errors.append( - f"stitched LLM trajectory contains {malformed_count} malformed row(s)" - ) - missing_fields.append("valid_provider_exchange") - if rows_valid and not usage_complete: - errors.append("stitched LLM trajectory lacks positive provider token usage") - missing_fields.append("token_usage") - if not live_capture_complete: - missing_fields.append("live_provider_exchange") - - models = { - value - for value, present in ( - (original_model, n_recorded > 0), - (live_model, live_attempt_count > 0), - ) - if present and value - } - stitched_model = next(iter(models)) if len(models) == 1 else None - manifest = LLMTrajectoryManifest( - status=CaptureStatus.COMPLETE if complete else CaptureStatus.PARTIAL, - capture_source=capture_source, - capture_fidelity=capture_fidelity, - auth_mode=auth_mode, - agent="openhands", - model=stitched_model, - session_id=current.session_id if current else rollout_dir.name, - exchange_count=exchange_count, - request_complete=request_complete, - response_complete=response_complete, - payload_redacted=source.payload_redacted if source else True, - started_at=current.started_at if current else datetime.now(UTC), - finished_at=datetime.now(UTC), - missing_fields=sorted(set(missing_fields)), - errors=errors, - role_captures=[ - LLMRoleCapture( - role="agent", - agent="openhands", - model=stitched_model, - auth_mode=auth_mode, - capture_source=capture_source, - capture_fidelity=capture_fidelity, - exchange_count=exchange_count, - request_complete=request_complete, - response_complete=response_complete, - ) - ], - ) - write_llm_trajectory_manifest(rollout_dir, manifest) - return manifest - - def _usage_int(usage: dict[str, Any], *keys: str) -> int: for key in keys: value = usage.get(key) @@ -875,6 +602,10 @@ async def _continue_run_with_sandbox_proxy( rollout_name=rollout_name, ) rollout = await Rollout.create(config) + live_capture_trusted = provider_capture_has_trusted_custody( + sandbox_local=True, + sandbox_user=config.sandbox_user, + ) replay_proxy: SandboxReplayProxy | None = None provider_runtime: Any | None = None result: Any | None = None @@ -902,9 +633,7 @@ async def _write_artifacts( run.path / "trajectory" / "llm_trajectory.jsonl", live_exchanges, live_model=live_model, - live_capture_trusted=bool( - config.sandbox_user and config.sandbox_user not in {"root", "0"} - ), + live_capture_trusted=live_capture_trusted, ) refresh_stitched_trajectory_manifest( rollout_dir, @@ -920,9 +649,7 @@ async def _write_artifacts( *(replay_proxy.live_errors if replay_proxy is not None else []), *(teardown_errors or []), ], - live_capture_trusted=bool( - config.sandbox_user and config.sandbox_user not in {"root", "0"} - ), + live_capture_trusted=live_capture_trusted, ) update_continued_metadata( rollout_dir, diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py new file mode 100644 index 000000000..4dd6b24d7 --- /dev/null +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -0,0 +1,290 @@ +"""Pure trajectory stitching and provenance reconciliation for continuation.""" + +from __future__ import annotations + +import json +import os +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from benchflow.trajectories.llm_capture_manifest import ( + LLM_TRAJECTORY_SCHEMA_VERSION, + AuthMode, + CaptureFidelity, + CaptureSource, + CaptureStatus, + LLMRoleCapture, + LLMTrajectoryManifest, + capture_artifact_allows_training, + read_llm_trajectory_manifest, + successful_exchanges_have_positive_usage, + write_llm_trajectory_manifest, +) +from benchflow.trajectories.types import LLMExchange, redact_trajectory_obj + + +def stitched_trajectory_lines( + original_llm_trajectory: Path, live_exchanges: list[LLMExchange] +) -> list[str]: + """Build the continuous llm_trajectory: recorded prefix + live suffix. + + The recorded request/response payloads are preserved, while their metadata + is promoted to the current schema so a newly stitched artifact can never be + mistaken for sidecar-optional legacy data. The live suffix is redacted on + the way out. + """ + lines: list[str] = [] + if original_llm_trajectory.is_file(): + for raw in original_llm_trajectory.read_text().splitlines(): + if raw.strip(): + try: + payload = json.loads(raw) + except json.JSONDecodeError: + lines.append(raw) + continue + if not isinstance(payload, dict): + lines.append(raw) + continue + metadata = payload.setdefault("metadata", {}) + if not isinstance(metadata, dict): + metadata = {} + payload["metadata"] = metadata + metadata["schema_version"] = LLM_TRAJECTORY_SCHEMA_VERSION + lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) + for exchange in live_exchanges: + payload = redact_trajectory_obj(exchange.model_dump(mode="json")) + lines.append(json.dumps(payload, default=str)) + return lines + + +def write_stitched_trajectory( + rollout_dir: Path, + original_llm_trajectory: Path, + live_exchanges: list[LLMExchange], + *, + live_model: str | None = None, + live_capture_trusted: bool = True, +) -> Path: + """Write the stitched continuous trajectory into the new rollout folder.""" + out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" + out.parent.mkdir(parents=True, exist_ok=True) + lines = stitched_trajectory_lines(original_llm_trajectory, []) + for exchange in live_exchanges: + payload = exchange.model_dump(mode="json") + metadata = payload.setdefault("metadata", {}) + metadata.update( + { + "agent": "openhands", + "role": "agent", + "model": live_model, + "auth_mode": AuthMode.API_KEY.value, + "capture_source": CaptureSource.LITELLM_PROXY.value, + "capture_fidelity": ( + CaptureFidelity.PROVIDER_WIRE.value + if live_capture_trusted + else CaptureFidelity.AGENT_SESSION.value + ), + "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, + "request_complete": True, + "response_complete": True, + "role_attribution_complete": True, + "payload_redacted": True, + } + ) + if not live_capture_trusted: + metadata["capture_custody"] = "agent_writable_sandbox" + lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) + rendered = "\n".join(lines) + ("\n" if lines else "") + temporary = out.with_suffix(out.suffix + ".tmp") + temporary.write_text(rendered) + os.replace(temporary, out) + return out + + +def refresh_stitched_trajectory_manifest( + rollout_dir: Path, + source_rollout_dir: Path, + *, + original_model: str | None, + live_model: str | None, + n_recorded: int, + n_live: int, + live_attempt_count: int, + live_errors: list[str], + live_capture_trusted: bool = True, +) -> LLMTrajectoryManifest: + """Replace rollout-finalization provenance with the final stitched contract.""" + + current_raw = read_llm_trajectory_manifest(rollout_dir) + source_raw = read_llm_trajectory_manifest(source_rollout_dir) + try: + current = ( + LLMTrajectoryManifest.model_validate(current_raw) + if current_raw is not None + else None + ) + except ValueError: + current = None + try: + source = ( + LLMTrajectoryManifest.model_validate(source_raw) + if source_raw is not None + else None + ) + except ValueError: + source = None + + trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" + trajectory_lines = [ + line for line in trajectory_path.read_text().splitlines() if line.strip() + ] + exchange_count = len(trajectory_lines) + malformed_count = 0 + for line in trajectory_lines: + try: + LLMExchange.model_validate_json(line) + except ValueError: + malformed_count += 1 + rows_valid = malformed_count == 0 + parsed_rows: list[dict[str, Any]] = [] + if rows_valid: + parsed_rows = [json.loads(line) for line in trajectory_lines] + expected_count = n_recorded + live_attempt_count + count_matches = exchange_count == expected_count + source_rows: list[dict[str, Any]] = [] + try: + source_rows = [ + json.loads(line) + for line in (source_rollout_dir / "trajectory" / "llm_trajectory.jsonl") + .read_text() + .splitlines() + if line.strip() + ] + except (OSError, json.JSONDecodeError): + source_rows = [] + source_allows_training = bool( + source_raw is not None + and len(source_rows) == n_recorded + and capture_artifact_allows_training(source_raw, exchanges=source_rows) + ) + live_capture_complete = live_attempt_count == n_live and not live_errors + usage_complete = bool( + parsed_rows and successful_exchanges_have_positive_usage(parsed_rows) + ) + complete = ( + source_allows_training + and count_matches + and live_capture_complete + and (live_capture_trusted or n_live == 0) + and rows_valid + and usage_complete + ) + + source_capture_source = source.capture_source if source else CaptureSource.NONE + source_fidelity = source.capture_fidelity if source else CaptureFidelity.NONE + source_auth = source.auth_mode if source else AuthMode.UNKNOWN + if n_live: + capture_source = ( + CaptureSource.LITELLM_PROXY + if source_capture_source is CaptureSource.LITELLM_PROXY + else CaptureSource.MIXED + ) + live_fidelity = ( + CaptureFidelity.PROVIDER_WIRE + if live_capture_trusted + else CaptureFidelity.AGENT_SESSION + ) + capture_fidelity = ( + live_fidelity if source_fidelity is live_fidelity else CaptureFidelity.MIXED + ) + auth_mode = ( + AuthMode.API_KEY if source_auth is AuthMode.API_KEY else AuthMode.MIXED + ) + else: + capture_source = source_capture_source + capture_fidelity = source_fidelity + auth_mode = source_auth + + request_complete = bool( + source + and source.request_complete + and count_matches + and live_capture_complete + and rows_valid + ) + response_complete = bool( + source + and source.response_complete + and count_matches + and live_capture_complete + and rows_valid + ) + errors = list(source.errors) if source and not complete else [] + missing_fields = list(source.missing_fields) if source and not complete else [] + errors.extend(live_errors) + if n_live and not live_capture_trusted: + errors.append("sandbox replay capture shared root custody with the agent") + if source is None: + errors.append("source LLM trajectory manifest is missing or malformed") + missing_fields.append("source_capture_provenance") + elif not source_allows_training: + errors.append("source LLM trajectory is not complete provider-wire capture") + if not count_matches: + errors.append( + "stitched LLM trajectory count mismatch: " + f"expected {expected_count}, found {exchange_count}" + ) + missing_fields.append("exchange_count") + if not rows_valid: + errors.append( + f"stitched LLM trajectory contains {malformed_count} malformed row(s)" + ) + missing_fields.append("valid_provider_exchange") + if rows_valid and not usage_complete: + errors.append("stitched LLM trajectory lacks positive provider token usage") + missing_fields.append("token_usage") + if not live_capture_complete: + missing_fields.append("live_provider_exchange") + + models = { + value + for value, present in ( + (original_model, n_recorded > 0), + (live_model, live_attempt_count > 0), + ) + if present and value + } + stitched_model = next(iter(models)) if len(models) == 1 else None + manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE if complete else CaptureStatus.PARTIAL, + capture_source=capture_source, + capture_fidelity=capture_fidelity, + auth_mode=auth_mode, + agent="openhands", + model=stitched_model, + session_id=current.session_id if current else rollout_dir.name, + exchange_count=exchange_count, + request_complete=request_complete, + response_complete=response_complete, + payload_redacted=source.payload_redacted if source else True, + started_at=current.started_at if current else datetime.now(UTC), + finished_at=datetime.now(UTC), + missing_fields=sorted(set(missing_fields)), + errors=errors, + role_captures=[ + LLMRoleCapture( + role="agent", + agent="openhands", + model=stitched_model, + auth_mode=auth_mode, + capture_source=capture_source, + capture_fidelity=capture_fidelity, + exchange_count=exchange_count, + request_complete=request_complete, + response_complete=response_complete, + ) + ], + ) + write_llm_trajectory_manifest(rollout_dir, manifest) + return manifest diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index def1eb89b..293237ad3 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -58,6 +58,9 @@ ) from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter +from benchflow.trajectories.llm_capture_manifest import ( + provider_capture_has_trusted_custody, +) from benchflow.trajectories.types import Trajectory from benchflow.usage_tracking import UsageTrackingConfig, usage_unavailable @@ -1724,8 +1727,9 @@ async def ensure_litellm_runtime( sorted(set(required_skill_names)), separators=(",", ":") ) proxy_location = "sandbox" if sandbox_local else "host" - capture_trusted = not sandbox_local or bool( - sandbox_user and sandbox_user not in {"root", "0"} + capture_trusted = provider_capture_has_trusted_custody( + sandbox_local=sandbox_local, + sandbox_user=sandbox_user, ) config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index a27925654..00bd16422 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -55,6 +55,14 @@ class AuthMode(StrEnum): UNKNOWN = "unknown" +def provider_capture_has_trusted_custody( + *, sandbox_local: bool, sandbox_user: str | None +) -> bool: + """Return whether an agent cannot rewrite its gateway's provider evidence.""" + + return not sandbox_local or bool(sandbox_user and sandbox_user not in {"root", "0"}) + + class LLMRoleCapture(BaseModel): """Per prepared role provenance for mixed-auth/mixed-agent rollouts.""" diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 769c5b804..b503d6d75 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -12,15 +12,17 @@ build_agent_env, build_rollout_config, continued_rollout_name, - refresh_stitched_trajectory_manifest, resolve_task_path, select_proxy_mode, - stitched_trajectory_lines, summarize_llm_trajectory_usage, update_continued_metadata, - write_stitched_trajectory, ) from benchflow.continue_run.run_folder import RunFolderError, load_run_folder +from benchflow.continue_run.trajectory_artifacts import ( + refresh_stitched_trajectory_manifest, + stitched_trajectory_lines, + write_stitched_trajectory, +) from benchflow.trajectories.llm_capture_manifest import ( AuthMode, CaptureFidelity, From c1517616418a41de10898a40f8c1715c3fc109b1 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 06:41:19 -0700 Subject: [PATCH 32/74] fix: close native capture review gaps --- docs/getting-started.md | 3 +- src/benchflow/continue_run/orchestrator.py | 19 +- .../continue_run/trajectory_artifacts.py | 79 ++- src/benchflow/providers/litellm_runtime.py | 38 +- src/benchflow/trajectories/llm_capture.py | 563 +++--------------- .../trajectories/llm_capture_manifest.py | 11 +- .../trajectories/native_capture_collection.py | 462 ++++++++++++++ tests/continue_run/test_orchestrator.py | 13 + tests/test_litellm_runtime.py | 39 +- .../test_native_capture_resilience.py | 133 ++++- tests/trajectories/test_native_llm_capture.py | 18 +- .../test_native_session_boundaries.py | 46 +- 12 files changed, 876 insertions(+), 548 deletions(-) create mode 100644 src/benchflow/trajectories/native_capture_collection.py diff --git a/docs/getting-started.md b/docs/getting-started.md index ff92ed73f..686c1fc86 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -229,7 +229,8 @@ The manifest status is `complete`, `partial`, `no_model_call`, or `capture_failed`. Mixed-role rollouts merge API-key and native-subscription exchanges into the same JSONL; `role_captures` records each prepared role's scene role, agent, model, auth mode, source, fidelity, completeness, and exchange -count. +count. Continued runs retain the source role entries as the `recorded` leg and +append a separate `live` leg, so a model or auth switch remains auditable. Missing or ambiguously attributed roles make the rollout-level capture `partial`. Reconstructed `agent_session` rows remain useful for audit and viewer workflows, but trainer exports fail closed unless the manifest says the capture diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index f20e030ec..52273c95a 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -44,9 +44,6 @@ from benchflow.contracts import AgentProtocolError, SandboxStartupFailure from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.scenes import compile_scenes_to_steps -from benchflow.trajectories.llm_capture_manifest import ( - provider_capture_has_trusted_custody, -) from benchflow.trajectories.types import LLMExchange logger = logging.getLogger(__name__) @@ -602,10 +599,7 @@ async def _continue_run_with_sandbox_proxy( rollout_name=rollout_name, ) rollout = await Rollout.create(config) - live_capture_trusted = provider_capture_has_trusted_custody( - sandbox_local=True, - sandbox_user=config.sandbox_user, - ) + live_capture_trusted = False replay_proxy: SandboxReplayProxy | None = None provider_runtime: Any | None = None result: Any | None = None @@ -620,13 +614,22 @@ async def _write_artifacts( *, force: bool = False, ) -> None: - nonlocal artifacts_written, live_exchanges, result, rollout_dir + nonlocal \ + artifacts_written, \ + live_capture_trusted, \ + live_exchanges, \ + result, \ + rollout_dir if artifacts_written and not force: return result = _result_after_sandbox_teardown(rollout) if result is None: return rollout_dir = Path(rollout._rollout_dir or (output_dir / rollout_name)) + live_capture_trusted = bool( + provider_runtime is not None + and getattr(provider_runtime, "capture_trusted", False) + ) live_exchanges = replay_proxy.live_exchanges if replay_proxy is not None else [] stitched_path = write_stitched_trajectory( rollout_dir, diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py index 4dd6b24d7..78cd8a62d 100644 --- a/src/benchflow/continue_run/trajectory_artifacts.py +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -256,6 +256,17 @@ def refresh_stitched_trajectory_manifest( if present and value } stitched_model = next(iter(models)) if len(models) == 1 else None + role_captures = _continuation_role_captures( + source=source, + original_model=original_model, + live_model=live_model, + n_recorded=n_recorded, + n_live=n_live, + live_attempt_count=live_attempt_count, + live_capture_complete=live_capture_complete, + live_capture_trusted=live_capture_trusted, + rows_valid=rows_valid, + ) manifest = LLMTrajectoryManifest( status=CaptureStatus.COMPLETE if complete else CaptureStatus.PARTIAL, capture_source=capture_source, @@ -272,19 +283,63 @@ def refresh_stitched_trajectory_manifest( finished_at=datetime.now(UTC), missing_fields=sorted(set(missing_fields)), errors=errors, - role_captures=[ + role_captures=role_captures, + ) + write_llm_trajectory_manifest(rollout_dir, manifest) + return manifest + + +def _continuation_role_captures( + *, + source: LLMTrajectoryManifest | None, + original_model: str | None, + live_model: str | None, + n_recorded: int, + n_live: int, + live_attempt_count: int, + live_capture_complete: bool, + live_capture_trusted: bool, + rows_valid: bool, +) -> list[LLMRoleCapture]: + """Preserve source-role provenance and append a distinct live leg.""" + + captures = [ + capture.model_copy(update={"leg": "recorded"}) + for capture in (source.role_captures if source else []) + ] + if source is not None and n_recorded and not captures: + captures.append( + LLMRoleCapture( + role="agent", + leg="recorded", + agent=source.agent, + model=source.model or original_model, + auth_mode=source.auth_mode, + capture_source=source.capture_source, + capture_fidelity=source.capture_fidelity, + exchange_count=n_recorded, + request_complete=source.request_complete, + response_complete=source.response_complete, + ) + ) + if live_attempt_count or n_live: + live_complete = live_capture_complete and rows_valid + captures.append( LLMRoleCapture( role="agent", + leg="live", agent="openhands", - model=stitched_model, - auth_mode=auth_mode, - capture_source=capture_source, - capture_fidelity=capture_fidelity, - exchange_count=exchange_count, - request_complete=request_complete, - response_complete=response_complete, + model=live_model, + auth_mode=AuthMode.API_KEY, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=( + CaptureFidelity.PROVIDER_WIRE + if live_capture_trusted + else CaptureFidelity.AGENT_SESSION + ), + exchange_count=n_live, + request_complete=live_complete, + response_complete=live_complete, ) - ], - ) - write_llm_trajectory_manifest(rollout_dir, manifest) - return manifest + ) + return captures diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 293237ad3..89bf13bc2 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -58,9 +58,6 @@ ) from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter -from benchflow.trajectories.llm_capture_manifest import ( - provider_capture_has_trusted_custody, -) from benchflow.trajectories.types import Trajectory from benchflow.usage_tracking import UsageTrackingConfig, usage_unavailable @@ -1727,9 +1724,10 @@ async def ensure_litellm_runtime( sorted(set(required_skill_names)), separators=(",", ":") ) proxy_location = "sandbox" if sandbox_local else "host" - capture_trusted = provider_capture_has_trusted_custody( + capture_trusted = await _provider_capture_has_verified_custody( sandbox_local=sandbox_local, sandbox_user=sandbox_user, + sandbox=sandbox, ) config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" @@ -1741,6 +1739,9 @@ async def ensure_litellm_runtime( if getattr(runtime, "config_key", None) == config_key and server is not None: is_running = await server.is_running() if is_running: + runtime.capture_trusted = bool( + getattr(runtime, "capture_trusted", False) and capture_trusted + ) if live_trajectory_path is not None: server.start_live_capture(live_trajectory_path) return ( @@ -1819,6 +1820,35 @@ async def ensure_litellm_runtime( ) +async def _provider_capture_has_verified_custody( + *, sandbox_local: bool, sandbox_user: str | None, sandbox: Any | None +) -> bool: + """Verify that the sandbox agent's effective UID cannot rewrite capture data.""" + + if not sandbox_local: + return True + if sandbox is None or sandbox_user in {None, "", "root", "0"}: + return False + try: + result = await sandbox.exec( + f"id -u -- {shlex.quote(sandbox_user)}", + user="root", + timeout_sec=10, + ) + except Exception as exc: + logger.warning("Provider capture custody UID check failed: %s", exc) + return False + if result.return_code != 0: + logger.warning("Provider capture custody UID check returned non-zero") + return False + try: + effective_uid = int(result.stdout.strip()) + except (AttributeError, ValueError): + logger.warning("Provider capture custody UID check returned invalid output") + return False + return effective_uid != 0 + + async def stop_litellm_runtime(runtime: Any | None) -> None: if runtime is None: return diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index 563a47ed5..c5dc28f69 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -2,22 +2,17 @@ from __future__ import annotations -import asyncio import hashlib import json import logging import os -import re -import shlex -import tempfile from contextlib import suppress -from dataclasses import dataclass, replace +from dataclasses import replace from datetime import datetime -from pathlib import Path, PurePosixPath +from pathlib import Path from typing import Any from benchflow.agents.env import uses_native_subscription_auth -from benchflow.agents.registry import AGENTS from benchflow.trajectories.llm_capture_manifest import ( LLM_TRAJECTORY_FILENAME, AuthMode, @@ -40,59 +35,19 @@ role_captures_for_targets, write_exchange_records, ) -from benchflow.trajectories.native_capture_parsers import ( - NativeParseResult, - parse_claude_raw_capture, - parse_claude_sessions, - parse_codex_sessions, - project_acp_trajectory, - retain_uncovered_claude_session_exchanges, +from benchflow.trajectories.native_capture_collection import ( + MAX_NATIVE_SESSION_FILES, + ClaudeOtelCollector, + NativeSessionCollector, + is_claude_code_agent, + native_session_id_is_safe, + sanitized_capture_error, ) -from benchflow.trajectories.types import redact_trajectory_text +from benchflow.trajectories.native_capture_parsers import project_acp_trajectory logger = logging.getLogger(__name__) _REMOTE_CAPTURE_PREFIX = "/tmp/benchflow-llm-capture-" -_MAX_NATIVE_SESSION_FILES = 1000 -_SAFE_NATIVE_SESSION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") -_OTEL_SINK_SOURCE = r""" -import { createServer } from 'node:http'; -import { mkdirSync, writeFileSync } from 'node:fs'; -import { join } from 'node:path'; - -const [outputDir, portFile] = process.argv.slice(2); -mkdirSync(outputDir, { recursive: true }); -let sequence = 0; -const server = createServer((request, response) => { - const chunks = []; - let size = 0; - request.on('data', chunk => { - size += chunk.length; - if (size <= 64 * 1024 * 1024) chunks.push(chunk); - }); - request.on('end', () => { - if (size <= 64 * 1024 * 1024) { - const name = `${Date.now()}-${String(sequence++).padStart(6, '0')}.json`; - writeFileSync(join(outputDir, name), Buffer.concat(chunks)); - } - response.writeHead(200, { 'content-type': 'application/json' }); - response.end('{}'); - }); -}); -server.listen(0, '127.0.0.1', () => { - const address = server.address(); - writeFileSync(portFile, `${address.port}\n`); -}); -process.on('SIGTERM', () => server.close(() => process.exit(0))); -""".strip() - - -@dataclass(frozen=True) -class _NativeCollection: - """Native bundles retained alongside isolated collection errors.""" - - bundles: tuple[_NativeCaptureBundle, ...] = () - errors: tuple[str, ...] = () class LLMTrajectoryCapture: @@ -114,6 +69,13 @@ def __init__( self.started_at = started_at capture_suffix = hashlib.sha256(session_id.encode()).hexdigest()[:12] self._remote_capture_root = f"{_REMOTE_CAPTURE_PREFIX}{capture_suffix}" + self._otel_collector = ClaudeOtelCollector(self._remote_capture_root) + self._native_collector = NativeSessionCollector( + agent=agent, + session_id=session_id, + started_at=started_at, + otel=self._otel_collector, + ) self.manifest = initialize_llm_trajectory_artifacts( rollout_dir, agent=agent, @@ -131,9 +93,6 @@ def __init__( self._provisional_target_key: ( tuple[str, str, str | None, str, AuthMode] | None ) = None - self._collector_started = False - self._collector_owned = False - self._capture_root_prepared = False self._preparation_errors: list[str] = [] self._otel_setup_error: str | None = None @@ -223,7 +182,7 @@ async def prepare_agent( if not native: write_llm_trajectory_manifest(self.rollout_dir, self.manifest) return prepared - if _is_claude_code_agent(agent): + if is_claude_code_agent(agent): raw_dir = f"{self._remote_capture_root}/raw" prepared.update( { @@ -232,11 +191,11 @@ async def prepare_agent( } ) try: - port = await self._ensure_otel_sink(env, sandbox_user=sandbox_user) + port = await self._otel_collector.ensure(env, sandbox_user=sandbox_user) except Exception as exc: prepared.pop("CLAUDE_CODE_ENABLE_TELEMETRY", None) prepared.pop("OTEL_LOG_RAW_API_BODIES", None) - warning = _sanitized_error(exc) + warning = sanitized_capture_error(exc) self._otel_setup_error = warning logger.warning( "Claude OTel correlation unavailable; session fallback remains " @@ -289,13 +248,13 @@ def bind_native_session( self._provisional_target_key = None if not target.native: return - if _SAFE_NATIVE_SESSION_ID.fullmatch(native_session_id) is None: + if not native_session_id_is_safe(native_session_id): warning = "native ACP session identifier was unsafe for file scoping" self._preparation_errors.append(warning) logger.warning("Native LLM capture disabled: %s", warning) return session_ids = tuple(sorted({*target.native_session_ids, native_session_id})) - if len(session_ids) > _MAX_NATIVE_SESSION_FILES: + if len(session_ids) > MAX_NATIVE_SESSION_FILES: warning = "native ACP session count exceeded the capture limit" self._preparation_errors.append(warning) logger.warning("Native LLM capture disabled: %s", warning) @@ -351,6 +310,12 @@ async def finalize( """Publish the highest-fidelity available capture and terminal sidecar.""" self.manifest.finished_at = datetime.now() + preparation_errors = [*self._preparation_errors] + if self._otel_setup_error is not None: + preparation_errors.append(self._otel_setup_error) + collection_errors = list( + dict.fromkeys([*preparation_errors, *(capture_errors or [])]) + ) provider_records: list[dict[str, Any]] = [] if self.trajectory_path.stat().st_size > 0: try: @@ -363,39 +328,37 @@ async def finalize( fallback_model=self.model, fallback_auth=self.manifest.auth_mode, ) - except Exception: - if env is not None: - if self._collector_owned: - await self._stop_otel_sink(env) - await self._cleanup_remote_capture(env) - raise + except Exception as exc: + warning = ( + f"provider capture parse failed: {sanitized_capture_error(exc)}" + ) + collection_errors.append(warning) + logger.warning("%s", warning) native_bundles: list[_NativeCaptureBundle] = [] - preparation_errors = [*self._preparation_errors] - if self._otel_setup_error is not None: - preparation_errors.append(self._otel_setup_error) - collection_errors = list( - dict.fromkeys([*preparation_errors, *(capture_errors or [])]) - ) native_targets = self._native_targets() - native_resources_exist = self._collector_owned or self._capture_root_prepared + native_resources_exist = ( + self._otel_collector.owned or self._otel_collector.root_prepared + ) cleanup_failed = False if (native_targets or native_resources_exist) and env is not None: try: if native_targets: - collection = await self._collect_native_results(env) + collection = await self._native_collector.collect( + env, targets=native_targets + ) native_bundles.extend(collection.bundles) collection_errors.extend(collection.errors) - elif self._collector_owned: - await self._stop_otel_sink(env) + elif self._otel_collector.owned: + await self._otel_collector.stop(env) except Exception as exc: - collection_errors.append(_sanitized_error(exc)) + collection_errors.append(sanitized_capture_error(exc)) logger.warning("Native LLM trajectory collection failed: %s", exc) finally: try: - await self._cleanup_remote_capture(env) + await self._otel_collector.cleanup(env) except Exception as exc: - collection_errors.append(_sanitized_error(exc)) + collection_errors.append(sanitized_capture_error(exc)) logger.error("Sandbox LLM capture cleanup failed: %s", exc) cleanup_failed = True @@ -457,326 +420,46 @@ def record_failure(self, error: object, *, model_call_seen: bool) -> None: """Leave a terminal, truthful sidecar when finalization itself fails.""" self.manifest.finished_at = datetime.now() - _atomic_replace_text(self.trajectory_path, "") + exchange_count = _valid_jsonl_row_count(self.trajectory_path) + if exchange_count is None: + _atomic_replace_text(self.trajectory_path, "") + exchange_count = 0 + rows_preserved = exchange_count > 0 self._finish_manifest( status=( CaptureStatus.CAPTURE_FAILED - if model_call_seen + if model_call_seen or rows_preserved else CaptureStatus.NO_MODEL_CALL ), - source=CaptureSource.NONE, - fidelity=CaptureFidelity.NONE, - exchange_count=0, + source=( + self.manifest.capture_source if rows_preserved else CaptureSource.NONE + ), + fidelity=( + self.manifest.capture_fidelity + if rows_preserved + else CaptureFidelity.NONE + ), + exchange_count=exchange_count, request_complete=False, response_complete=False, missing_fields=( - ["provider_request", "provider_response"] if model_call_seen else [] - ), - errors=[_sanitized_error(error)], - role_captures=role_captures_for_targets(list(self._targets.values())), - ) - - async def _ensure_otel_sink(self, env: Any, *, sandbox_user: str | None) -> int: - if self._collector_started: - return await self._read_collector_port(env) - await self._stop_otel_sink(env) - capture_owner = shlex.quote(sandbox_user or "root") - setup = await env.exec( - f"find {self._remote_capture_root} -depth -mindepth 1 -delete " - "2>/dev/null || true\n" - f"mkdir -p {self._remote_capture_root}/raw " - f"{self._remote_capture_root}/otel\n" - f"chown -R {capture_owner} {self._remote_capture_root}\n" - f"chmod 700 {self._remote_capture_root} " - f"{self._remote_capture_root}/raw {self._remote_capture_root}/otel", - user="root", - timeout_sec=10, - ) - if setup.return_code != 0: - detail = (setup.stderr or setup.stdout or "capture directory setup failed")[ - :300 - ] - raise RuntimeError(f"Claude capture directory setup failed: {detail}") - self._capture_root_prepared = True - with tempfile.TemporaryDirectory(prefix="benchflow-otel-sink-") as temporary: - source = Path(temporary) / "otel_sink.mjs" - source.write_text(_OTEL_SINK_SOURCE + "\n") - await env.upload_file( - source, - f"{self._remote_capture_root}/otel_sink.mjs", - mode="755", - ) - command = f""" -find {self._remote_capture_root} -maxdepth 1 -type f -name port -delete -find {self._remote_capture_root} -maxdepth 1 -type f -name pid -delete -node_bin=/opt/benchflow/node/bin/node -if ! test -x "$node_bin"; then - node_bin=$(command -v node || true) -fi -if test -z "$node_bin"; then - echo "node runtime not found" >&2 - exit 1 -fi -nohup "$node_bin" {self._remote_capture_root}/otel_sink.mjs \ - {self._remote_capture_root}/otel {self._remote_capture_root}/port \ - >{self._remote_capture_root}/collector.stdout \ - 2>{self._remote_capture_root}/collector.stderr {self._remote_capture_root}/pid -for attempt in $(seq 1 50); do - if test -s {self._remote_capture_root}/port; then - cat {self._remote_capture_root}/port - exit 0 - fi - sleep 0.1 -done -tail -c 300 {self._remote_capture_root}/collector.stderr >&2 2>/dev/null || true -exit 1 -""" - self._collector_owned = True - result = await env.exec( - command, - user=sandbox_user or "root", - timeout_sec=10, - ) - if result.return_code != 0: - detail = (result.stderr or result.stdout or "collector did not start")[:300] - raise RuntimeError(f"Claude OTel sink failed to start: {detail}") - self._collector_started = True - return _parse_port(result.stdout) - - async def _stop_otel_sink(self, env: Any) -> None: - command = f""" -if ! test -s {self._remote_capture_root}/pid; then - exit 0 -fi -read -r old_pid < {self._remote_capture_root}/pid || true -case "$old_pid" in - ''|*[!0-9]*) exit 0 ;; -esac -old_command=$(ps -p "$old_pid" -o command= 2>/dev/null || true) -case "$old_command" in - *{self._remote_capture_root}/otel_sink.mjs*) ;; - *) exit 0 ;; -esac -kill -TERM "$old_pid" 2>/dev/null || true -for attempt in $(seq 1 20); do - if ! kill -0 "$old_pid" 2>/dev/null; then - exit 0 - fi - sleep 0.05 -done -old_command=$(ps -p "$old_pid" -o command= 2>/dev/null || true) -case "$old_command" in - *{self._remote_capture_root}/otel_sink.mjs*) ;; - *) exit 0 ;; -esac -kill -KILL "$old_pid" 2>/dev/null || true -for attempt in $(seq 1 20); do - if ! kill -0 "$old_pid" 2>/dev/null; then - exit 0 - fi - sleep 0.05 -done -echo "previous Claude telemetry collector did not stop" >&2 -exit 1 -""" - result = await env.exec(command, user="root", timeout_sec=5) - if result.return_code != 0: - detail = (result.stderr or result.stdout or "collector did not stop")[:300] - raise RuntimeError(f"Claude OTel sink shutdown failed: {detail}") - self._collector_started = False - self._collector_owned = False - - async def _read_collector_port(self, env: Any) -> int: - result = await env.exec( - f"cat {self._remote_capture_root}/port", - user="root", - timeout_sec=5, - ) - if result.return_code != 0: - raise RuntimeError("Claude OTel sink port file is unavailable") - return _parse_port(result.stdout) - - async def _collect_native_results(self, env: Any) -> _NativeCollection: - bundles: list[_NativeCaptureBundle] = [] - errors: list[str] = [] - if self._collector_owned: - try: - await self._stop_otel_sink(env) - except Exception as exc: - errors.append(_sanitized_error(exc)) - logger.warning("Claude OTel collector shutdown failed: %s", exc) - - native_targets = self._native_targets() - claude_targets = tuple( - target for target in native_targets if _is_claude_code_agent(target.agent) - ) - with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: - local_root = Path(temporary) - raw_claude_result = await self._collect_claude_raw_capture( - env, - local_root=local_root, - claude_targets=claude_targets, - bundles=bundles, - errors=errors, - ) - for index, target in enumerate(native_targets): - try: - if _is_claude_code_agent(target.agent): - target_bundles = await self._collect_claude_session_fallback( - env, - local_root=local_root, - index=index, - target=target, - raw_result=raw_claude_result, - ) - bundles.extend(target_bundles) - bundle = None - elif target.agent == "codex-acp": - bundle = await self._collect_codex_session( - env, - local_root=local_root, - index=index, - target=target, - ) - else: - bundle = None - if bundle is not None: - bundles.append(bundle) - except Exception as exc: - warning = ( - f"native capture failed for role {target.role}: " - f"{_sanitized_error(exc)}" - ) - errors.append(warning) - logger.warning("%s", warning) - return _NativeCollection(bundles=tuple(bundles), errors=tuple(errors)) - - async def _collect_claude_raw_capture( - self, - env: Any, - *, - local_root: Path, - claude_targets: tuple[_CaptureTarget, ...], - bundles: list[_NativeCaptureBundle], - errors: list[str], - ) -> NativeParseResult | None: - if not self._capture_root_prepared: - return None - capture_dir = local_root / "capture" - try: - await env.download_dir(self._remote_capture_root, capture_dir) - result = parse_claude_raw_capture( - capture_dir, - agent=(claude_targets[0].agent if claude_targets else self.agent), - session_id=self.session_id, - started_at=self.started_at, - ) - except Exception as exc: - errors.append(_sanitized_error(exc)) - logger.warning("Claude raw LLM capture collection failed: %s", exc) - return None - if result is None: - return None - bundles.append(_NativeCaptureBundle(targets=claude_targets, result=result)) - return result - - async def _collect_claude_session_fallback( - self, - env: Any, - *, - local_root: Path, - index: int, - target: _CaptureTarget, - raw_result: NativeParseResult | None, - ) -> tuple[_NativeCaptureBundle, ...]: - bundles: list[_NativeCaptureBundle] = [] - for session_index, native_session_id in enumerate(target.native_session_ids): - local = local_root / f"target-{index}" / f"claude-session-{session_index}" - downloaded = await _download_bound_session_files( - env, - f"{target.credential_home}/.claude/projects", - local, - started_at=self.started_at, - session_ids=(native_session_id,), - ) - if not downloaded: - continue - result = parse_claude_sessions( - local, - agent=target.agent, - session_id=self.session_id, - started_at=self.started_at, - ) - if result is None: - continue - uncovered = retain_uncovered_claude_session_exchanges( - raw_result, - result, - native_session_id=native_session_id, - ) - if uncovered is not None: - bundles.append( - _NativeCaptureBundle(targets=(target,), result=uncovered) + sorted( + { + *self.manifest.missing_fields, + *( + ["provider_request", "provider_response"] + if model_call_seen or rows_preserved + else [] + ), + } ) - return tuple(bundles) - - async def _collect_codex_session( - self, - env: Any, - *, - local_root: Path, - index: int, - target: _CaptureTarget, - ) -> _NativeCaptureBundle | None: - local = local_root / f"target-{index}" / "codex-sessions" - downloaded = await _download_bound_session_files( - env, - f"{target.credential_home}/.codex/sessions", - local, - started_at=self.started_at, - session_ids=target.native_session_ids, - ) - if not downloaded: - return None - result = parse_codex_sessions( - local, - agent=target.agent, - session_id=self.session_id, - started_at=self.started_at, - configured_model=target.model, - auth_mode=target.auth_mode.value, - ) - return ( - _NativeCaptureBundle(targets=(target,), result=result) - if result is not None - else None - ) - - async def _cleanup_remote_capture(self, env: Any) -> None: - if not self._capture_root_prepared: - return - if self._collector_owned: - raise RuntimeError( - "Refusing to remove Claude capture ownership files while its " - "collector may still be running" - ) - result = await env.exec( - "for attempt in 1 2 3; do\n" - f" if ! test -e {self._remote_capture_root} || " - f"find {self._remote_capture_root} -depth -delete; then\n" - " exit 0\n" - " fi\n" - " sleep 0.1\n" - "done\n" - "exit 1", - user="root", - timeout_sec=10, + ), + errors=[*self.manifest.errors, sanitized_capture_error(error)], + role_captures=( + self.manifest.role_captures + or role_captures_for_targets(list(self._targets.values())) + ), ) - if result.return_code != 0: - detail = (result.stderr or result.stdout or "unknown error")[:300] - raise RuntimeError(f"Sandbox LLM capture cleanup failed: {detail}") - self._capture_root_prepared = False def _finish_manifest( self, @@ -798,66 +481,11 @@ def _finish_manifest( self.manifest.request_complete = request_complete self.manifest.response_complete = response_complete self.manifest.missing_fields = sorted(set(missing_fields or [])) - self.manifest.errors = [_sanitized_error(item) for item in errors or []] + self.manifest.errors = [sanitized_capture_error(item) for item in errors or []] self.manifest.role_captures = role_captures or [] write_llm_trajectory_manifest(self.rollout_dir, self.manifest) -async def _download_bound_session_files( - env: Any, - remote: str, - local: Path, - *, - started_at: datetime, - session_ids: tuple[str, ...], -) -> bool: - if not session_ids: - return False - if any(_SAFE_NATIVE_SESSION_ID.fullmatch(value) is None for value in session_ids): - raise RuntimeError("Native session discovery received an unsafe session ID") - boundary = started_at.timestamp() - 1.0 - remote_root = shlex.quote(remote) - filename_patterns = [ - pattern - for session_id in session_ids - for pattern in (f"{session_id}.jsonl", f"*-{session_id}.jsonl") - ] - filename_filter = ( - r"\( " - + " -o ".join(f"-name {shlex.quote(pattern)}" for pattern in filename_patterns) - + r" \)" - ) - result = await env.exec( - f"if test -d {remote_root}; then " - f"find {remote_root} -type f -name '*.jsonl' " - f"-newermt {shlex.quote(f'@{boundary}')} {filename_filter} " - f"-printf '%P\\n' | head -n {_MAX_NATIVE_SESSION_FILES + 1}; " - "fi", - user="root", - timeout_sec=10, - ) - if result.return_code != 0: - detail = (result.stderr or result.stdout or "session discovery failed")[:300] - raise RuntimeError(f"Native session discovery failed: {detail}") - relative_paths = [line for line in result.stdout.splitlines() if line] - if len(relative_paths) > _MAX_NATIVE_SESSION_FILES: - raise RuntimeError("Native session discovery exceeded the 1000-file limit") - if not relative_paths: - return False - downloads: list[tuple[str, Path]] = [] - for value in relative_paths: - relative = PurePosixPath(value) - if relative.is_absolute() or ".." in relative.parts: - raise RuntimeError("Native session discovery returned an unsafe path") - destination = local.joinpath(*relative.parts) - destination.parent.mkdir(parents=True, exist_ok=True) - downloads.append((f"{remote}/{relative.as_posix()}", destination)) - await asyncio.gather( - *(env.download_file(source, destination) for source, destination in downloads) - ) - return True - - def _capture_target_key( *, agent: str, @@ -887,14 +515,6 @@ def _capture_target_base_key( return (role_name or "primary", agent, model, credential_home) -def _is_claude_code_agent(agent: str) -> bool: - config = AGENTS.get(agent) - subscription = config.subscription_auth if config is not None else None - return bool( - subscription is not None and subscription.replaces_env == "ANTHROPIC_API_KEY" - ) - - def _resolve_auth_mode( agent: str, model: str | None, @@ -927,27 +547,28 @@ def _resolve_auth_mode( return AuthMode.OAUTH_SUBSCRIPTION -def _parse_port(value: str) -> int: - try: - port = int(value.strip().splitlines()[-1]) - except (ValueError, IndexError) as exc: - raise RuntimeError("Claude OTel sink returned an invalid port") from exc - if not 1 <= port <= 65535: - raise RuntimeError("Claude OTel sink returned an out-of-range port") - return port - - -def _sanitized_error(error: object) -> str: - text = redact_trajectory_text(str(error)).replace("\n", " ").strip() - return text[:500] or type(error).__name__ - - def _atomic_replace_text(path: Path, payload: str) -> None: temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(payload) os.replace(temporary, path) +def _valid_jsonl_row_count(path: Path) -> int | None: + """Count valid JSON objects, distinguishing an empty artifact from corruption.""" + + count = 0 + try: + for line in path.read_text().splitlines(): + if not line.strip(): + continue + if not isinstance(json.loads(line), dict): + return None + count += 1 + except (OSError, json.JSONDecodeError): + return None + return count + + def model_call_seen_from_evidence( usage_metrics: dict[str, Any] | None, acp_events: list[dict[str, Any]], diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 00bd16422..891cb25ab 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -13,7 +13,7 @@ from datetime import datetime from enum import StrEnum from pathlib import Path -from typing import Any +from typing import Any, Literal from pydantic import BaseModel, Field, ValidationError @@ -55,18 +55,11 @@ class AuthMode(StrEnum): UNKNOWN = "unknown" -def provider_capture_has_trusted_custody( - *, sandbox_local: bool, sandbox_user: str | None -) -> bool: - """Return whether an agent cannot rewrite its gateway's provider evidence.""" - - return not sandbox_local or bool(sandbox_user and sandbox_user not in {"root", "0"}) - - class LLMRoleCapture(BaseModel): """Per prepared role provenance for mixed-auth/mixed-agent rollouts.""" role: str = "agent" + leg: Literal["recorded", "live"] | None = None agent: str model: str | None = None auth_mode: AuthMode diff --git a/src/benchflow/trajectories/native_capture_collection.py b/src/benchflow/trajectories/native_capture_collection.py new file mode 100644 index 000000000..7c8db0182 --- /dev/null +++ b/src/benchflow/trajectories/native_capture_collection.py @@ -0,0 +1,462 @@ +"""Sandbox-owned native telemetry and session collection components.""" + +from __future__ import annotations + +import asyncio +import logging +import re +import shlex +import tempfile +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path, PurePosixPath +from typing import Any + +from benchflow.agents.registry import AGENTS +from benchflow.trajectories.llm_capture_records import ( + CaptureTarget, + NativeCaptureBundle, +) +from benchflow.trajectories.native_capture_parsers import ( + NativeParseResult, + parse_claude_raw_capture, + parse_claude_sessions, + parse_codex_sessions, + retain_uncovered_claude_session_exchanges, +) +from benchflow.trajectories.types import redact_trajectory_text + +logger = logging.getLogger(__name__) + +MAX_NATIVE_SESSION_FILES = 1000 +_SAFE_NATIVE_SESSION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_OTEL_SINK_SOURCE = r""" +import { createServer } from 'node:http'; +import { mkdirSync, writeFileSync } from 'node:fs'; +import { join } from 'node:path'; + +const [outputDir, portFile] = process.argv.slice(2); +mkdirSync(outputDir, { recursive: true }); +let sequence = 0; +const server = createServer((request, response) => { + const chunks = []; + let size = 0; + request.on('data', chunk => { + size += chunk.length; + if (size <= 64 * 1024 * 1024) chunks.push(chunk); + }); + request.on('end', () => { + if (size <= 64 * 1024 * 1024) { + const name = `${Date.now()}-${String(sequence++).padStart(6, '0')}.json`; + writeFileSync(join(outputDir, name), Buffer.concat(chunks)); + } + response.writeHead(200, { 'content-type': 'application/json' }); + response.end('{}'); + }); +}); +server.listen(0, '127.0.0.1', () => { + const address = server.address(); + writeFileSync(portFile, `${address.port}\n`); +}); +process.on('SIGTERM', () => server.close(() => process.exit(0))); +""".strip() + + +@dataclass(frozen=True) +class NativeCollection: + """Native bundles retained alongside isolated collection errors.""" + + bundles: tuple[NativeCaptureBundle, ...] = () + errors: tuple[str, ...] = () + + +@dataclass +class ClaudeOtelCollector: + """Own the sandbox OTel sink and its private capture directory.""" + + remote_root: str + started: bool = False + owned: bool = False + root_prepared: bool = False + + async def ensure(self, env: Any, *, sandbox_user: str | None) -> int: + if self.started: + return await self.read_port(env) + await self.stop(env) + capture_owner = shlex.quote(sandbox_user or "root") + setup = await env.exec( + f"find {self.remote_root} -depth -mindepth 1 -delete " + "2>/dev/null || true\n" + f"mkdir -p {self.remote_root}/raw {self.remote_root}/otel\n" + f"chown -R {capture_owner} {self.remote_root}\n" + f"chmod 700 {self.remote_root} " + f"{self.remote_root}/raw {self.remote_root}/otel", + user="root", + timeout_sec=10, + ) + if setup.return_code != 0: + detail = (setup.stderr or setup.stdout or "capture directory setup failed")[ + :300 + ] + raise RuntimeError(f"Claude capture directory setup failed: {detail}") + self.root_prepared = True + with tempfile.TemporaryDirectory(prefix="benchflow-otel-sink-") as temporary: + source = Path(temporary) / "otel_sink.mjs" + source.write_text(_OTEL_SINK_SOURCE + "\n") + await env.upload_file( + source, f"{self.remote_root}/otel_sink.mjs", mode="755" + ) + command = f""" +find {self.remote_root} -maxdepth 1 -type f -name port -delete +find {self.remote_root} -maxdepth 1 -type f -name pid -delete +node_bin=/opt/benchflow/node/bin/node +if ! test -x "$node_bin"; then + node_bin=$(command -v node || true) +fi +if test -z "$node_bin"; then + echo "node runtime not found" >&2 + exit 1 +fi +nohup "$node_bin" {self.remote_root}/otel_sink.mjs \ + {self.remote_root}/otel {self.remote_root}/port \ + >{self.remote_root}/collector.stdout \ + 2>{self.remote_root}/collector.stderr {self.remote_root}/pid +for attempt in $(seq 1 50); do + if test -s {self.remote_root}/port; then + cat {self.remote_root}/port + exit 0 + fi + sleep 0.1 +done +tail -c 300 {self.remote_root}/collector.stderr >&2 2>/dev/null || true +exit 1 +""" + self.owned = True + result = await env.exec(command, user=sandbox_user or "root", timeout_sec=10) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "collector did not start")[:300] + raise RuntimeError(f"Claude OTel sink failed to start: {detail}") + self.started = True + return _parse_port(result.stdout) + + async def stop(self, env: Any) -> None: + command = f""" +if ! test -s {self.remote_root}/pid; then + exit 0 +fi +read -r old_pid < {self.remote_root}/pid || true +case "$old_pid" in + ''|*[!0-9]*) exit 0 ;; +esac +old_command=$(ps -p "$old_pid" -o command= 2>/dev/null || true) +case "$old_command" in + *{self.remote_root}/otel_sink.mjs*) ;; + *) exit 0 ;; +esac +kill -TERM "$old_pid" 2>/dev/null || true +for attempt in $(seq 1 20); do + if ! kill -0 "$old_pid" 2>/dev/null; then + exit 0 + fi + sleep 0.05 +done +old_command=$(ps -p "$old_pid" -o command= 2>/dev/null || true) +case "$old_command" in + *{self.remote_root}/otel_sink.mjs*) ;; + *) exit 0 ;; +esac +kill -KILL "$old_pid" 2>/dev/null || true +for attempt in $(seq 1 20); do + if ! kill -0 "$old_pid" 2>/dev/null; then + exit 0 + fi + sleep 0.05 +done +echo "previous Claude telemetry collector did not stop" >&2 +exit 1 +""" + result = await env.exec(command, user="root", timeout_sec=5) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "collector did not stop")[:300] + raise RuntimeError(f"Claude OTel sink shutdown failed: {detail}") + self.started = False + self.owned = False + + async def read_port(self, env: Any) -> int: + result = await env.exec( + f"cat {self.remote_root}/port", user="root", timeout_sec=5 + ) + if result.return_code != 0: + raise RuntimeError("Claude OTel sink port file is unavailable") + return _parse_port(result.stdout) + + async def cleanup(self, env: Any) -> None: + if not self.root_prepared: + return + if self.owned: + raise RuntimeError( + "Refusing to remove Claude capture ownership files while its " + "collector may still be running" + ) + result = await env.exec( + "for attempt in 1 2 3; do\n" + f" if ! test -e {self.remote_root} || " + f"find {self.remote_root} -depth -delete; then\n" + " exit 0\n" + " fi\n" + " sleep 0.1\n" + "done\n" + "exit 1", + user="root", + timeout_sec=10, + ) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "unknown error")[:300] + raise RuntimeError(f"Sandbox LLM capture cleanup failed: {detail}") + self.root_prepared = False + + +@dataclass +class NativeSessionCollector: + """Collect session-bound Claude/Codex evidence after the agent exits.""" + + agent: str + session_id: str + started_at: datetime + otel: ClaudeOtelCollector + + async def collect( + self, env: Any, *, targets: list[CaptureTarget] + ) -> NativeCollection: + bundles: list[NativeCaptureBundle] = [] + errors: list[str] = [] + if self.otel.owned: + try: + await self.otel.stop(env) + except Exception as exc: + errors.append(sanitized_capture_error(exc)) + logger.warning("Claude OTel collector shutdown failed: %s", exc) + + claude_targets = tuple( + target for target in targets if is_claude_code_agent(target.agent) + ) + with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: + local_root = Path(temporary) + raw_claude_result = await self._collect_claude_raw_capture( + env, + local_root=local_root, + claude_targets=claude_targets, + bundles=bundles, + errors=errors, + ) + for index, target in enumerate(targets): + try: + if is_claude_code_agent(target.agent): + bundles.extend( + await self._collect_claude_session_fallback( + env, + local_root=local_root, + index=index, + target=target, + raw_result=raw_claude_result, + ) + ) + bundle = None + elif target.agent == "codex-acp": + bundle = await self._collect_codex_session( + env, local_root=local_root, index=index, target=target + ) + else: + bundle = None + if bundle is not None: + bundles.append(bundle) + except Exception as exc: + warning = ( + f"native capture failed for role {target.role}: " + f"{sanitized_capture_error(exc)}" + ) + errors.append(warning) + logger.warning("%s", warning) + return NativeCollection(bundles=tuple(bundles), errors=tuple(errors)) + + async def _collect_claude_raw_capture( + self, + env: Any, + *, + local_root: Path, + claude_targets: tuple[CaptureTarget, ...], + bundles: list[NativeCaptureBundle], + errors: list[str], + ) -> NativeParseResult | None: + if not self.otel.root_prepared: + return None + capture_dir = local_root / "capture" + try: + await env.download_dir(self.otel.remote_root, capture_dir) + result = parse_claude_raw_capture( + capture_dir, + agent=(claude_targets[0].agent if claude_targets else self.agent), + session_id=self.session_id, + started_at=self.started_at, + ) + except Exception as exc: + errors.append(sanitized_capture_error(exc)) + logger.warning("Claude raw LLM capture collection failed: %s", exc) + return None + if result is None: + return None + bundles.append(NativeCaptureBundle(targets=claude_targets, result=result)) + return result + + async def _collect_claude_session_fallback( + self, + env: Any, + *, + local_root: Path, + index: int, + target: CaptureTarget, + raw_result: NativeParseResult | None, + ) -> tuple[NativeCaptureBundle, ...]: + bundles: list[NativeCaptureBundle] = [] + for session_index, native_session_id in enumerate(target.native_session_ids): + local = local_root / f"target-{index}" / f"claude-session-{session_index}" + downloaded = await download_bound_session_files( + env, + f"{target.credential_home}/.claude/projects", + local, + started_at=self.started_at, + session_ids=(native_session_id,), + ) + if not downloaded: + continue + result = parse_claude_sessions( + local, + agent=target.agent, + session_id=self.session_id, + started_at=self.started_at, + ) + if result is None: + continue + uncovered = retain_uncovered_claude_session_exchanges( + raw_result, result, native_session_id=native_session_id + ) + if uncovered is not None: + bundles.append(NativeCaptureBundle(targets=(target,), result=uncovered)) + return tuple(bundles) + + async def _collect_codex_session( + self, + env: Any, + *, + local_root: Path, + index: int, + target: CaptureTarget, + ) -> NativeCaptureBundle | None: + local = local_root / f"target-{index}" / "codex-sessions" + downloaded = await download_bound_session_files( + env, + f"{target.credential_home}/.codex/sessions", + local, + started_at=self.started_at, + session_ids=target.native_session_ids, + ) + if not downloaded: + return None + result = parse_codex_sessions( + local, + agent=target.agent, + session_id=self.session_id, + started_at=self.started_at, + configured_model=target.model, + auth_mode=target.auth_mode.value, + ) + return ( + NativeCaptureBundle(targets=(target,), result=result) + if result is not None + else None + ) + + +async def download_bound_session_files( + env: Any, + remote: str, + local: Path, + *, + started_at: datetime, + session_ids: tuple[str, ...], +) -> bool: + """Download only explicitly bound, rollout-fresh native session files.""" + + if not session_ids: + return False + if any(not native_session_id_is_safe(value) for value in session_ids): + raise RuntimeError("Native session discovery received an unsafe session ID") + boundary = started_at.timestamp() - 1.0 + remote_root = shlex.quote(remote) + filename_patterns = [ + pattern + for session_id in session_ids + for pattern in (f"{session_id}.jsonl", f"*-{session_id}.jsonl") + ] + filename_filter = ( + r"\( " + + " -o ".join(f"-name {shlex.quote(pattern)}" for pattern in filename_patterns) + + r" \)" + ) + result = await env.exec( + f"if test -d {remote_root}; then " + f"find {remote_root} -type f -name '*.jsonl' " + f"-newermt {shlex.quote(f'@{boundary}')} {filename_filter} " + f"-printf '%P\\n' | head -n {MAX_NATIVE_SESSION_FILES + 1}; " + "fi", + user="root", + timeout_sec=10, + ) + if result.return_code != 0: + detail = (result.stderr or result.stdout or "session discovery failed")[:300] + raise RuntimeError(f"Native session discovery failed: {detail}") + relative_paths = [line for line in result.stdout.splitlines() if line] + if len(relative_paths) > MAX_NATIVE_SESSION_FILES: + raise RuntimeError("Native session discovery exceeded the 1000-file limit") + if not relative_paths: + return False + downloads: list[tuple[str, Path]] = [] + for value in relative_paths: + relative = PurePosixPath(value) + if relative.is_absolute() or ".." in relative.parts: + raise RuntimeError("Native session discovery returned an unsafe path") + destination = local.joinpath(*relative.parts) + destination.parent.mkdir(parents=True, exist_ok=True) + downloads.append((f"{remote}/{relative.as_posix()}", destination)) + await asyncio.gather( + *(env.download_file(source, destination) for source, destination in downloads) + ) + return True + + +def native_session_id_is_safe(value: str) -> bool: + return _SAFE_NATIVE_SESSION_ID.fullmatch(value) is not None + + +def is_claude_code_agent(agent: str) -> bool: + config = AGENTS.get(agent) + subscription = config.subscription_auth if config is not None else None + return bool( + subscription is not None and subscription.replaces_env == "ANTHROPIC_API_KEY" + ) + + +def sanitized_capture_error(error: object) -> str: + text = redact_trajectory_text(str(error)).replace("\n", " ").strip() + return text[:500] or type(error).__name__ + + +def _parse_port(value: str) -> int: + try: + port = int(value.strip().splitlines()[-1]) + except (ValueError, IndexError) as exc: + raise RuntimeError("Claude OTel sink returned an invalid port") from exc + if not 1 <= port <= 65535: + raise RuntimeError("Claude OTel sink returned an out-of-range port") + return port diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index b503d6d75..0dd4ab5ef 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -293,6 +293,19 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) assert manifest.capture_source is CaptureSource.MIXED assert manifest.capture_fidelity is CaptureFidelity.MIXED assert manifest.auth_mode is AuthMode.MIXED + assert [capture.leg for capture in manifest.role_captures] == [ + "recorded", + "live", + ] + recorded, live = manifest.role_captures + assert recorded.agent == "claude-agent-acp" + assert recorded.model == "claude-sonnet-4-6" + assert recorded.auth_mode is AuthMode.OAUTH_SUBSCRIPTION + assert recorded.capture_fidelity is CaptureFidelity.AGENT_SESSION + assert live.agent == "openhands" + assert live.model == "openai/gpt-5.5" + assert live.auth_mode is AuthMode.API_KEY + assert live.capture_fidelity is CaptureFidelity.PROVIDER_WIRE assert not capture_manifest_allows_training( manifest.model_dump(mode="json"), exchange_count=2 ) diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index cffa9be4a..c787b3420 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -174,6 +174,12 @@ async def fake_sandbox_start(**kwargs): return FakeLiteLLMServer("http://127.0.0.1:45678", kwargs["route"]) monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) + + class NonRootSandbox: + async def exec(self, command, **_kwargs): + assert command == "id -u -- agent" + return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") + _, provider_runtime = await ensure_litellm_runtime( agent="openhands", agent_env={ @@ -184,7 +190,7 @@ async def fake_sandbox_start(**kwargs): runtime=None, environment="daytona", session_id="run-non-root", - sandbox=SimpleNamespace(), + sandbox=NonRootSandbox(), sandbox_user="agent", ) @@ -192,6 +198,37 @@ async def fake_sandbox_start(**kwargs): assert provider_runtime.capture_trusted is True +@pytest.mark.asyncio +async def test_uid_zero_alias_keeps_sandbox_gateway_capture_audit_only(monkeypatch): + """Guards PR #1057 against trusting a non-root name that resolves to UID 0.""" + + async def fake_sandbox_start(**kwargs): + return FakeLiteLLMServer("http://127.0.0.1:45678", kwargs["route"]) + + class RootAliasSandbox: + async def exec(self, command, **_kwargs): + assert command == "id -u -- agent" + return SimpleNamespace(return_code=0, stdout="0\n", stderr="") + + monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) + _, provider_runtime = await ensure_litellm_runtime( + agent="openhands", + agent_env={ + "AWS_BEARER_TOKEN_BEDROCK": "token", + "AWS_REGION": "us-west-2", + }, + model="aws-bedrock/us.anthropic.claude-opus-4-8", + runtime=None, + environment="daytona", + session_id="run-root-alias", + sandbox=RootAliasSandbox(), + sandbox_user="agent", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is False + + @pytest.mark.asyncio async def test_apple_container_uses_sandbox_local_litellm(monkeypatch): """Guards PR #936 against handing the VM a host-loopback model endpoint.""" diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index ff2fd7ff2..62f273f55 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -9,13 +9,18 @@ import pytest +from benchflow.trajectories import llm_capture as llm_capture_module from benchflow.trajectories.llm_capture import LLMTrajectoryCapture, _CaptureTarget from benchflow.trajectories.llm_capture_manifest import ( AuthMode, CaptureFidelity, CaptureStatus, ) -from benchflow.trajectories.llm_capture_records import load_provider_wire_records +from benchflow.trajectories.llm_capture_records import ( + NativeCaptureBundle, + load_provider_wire_records, +) +from benchflow.trajectories.native_capture_collection import NativeCollection from benchflow.trajectories.native_capture_parsers import parse_codex_sessions @@ -320,10 +325,10 @@ async def test_root_sandbox_provider_capture_is_retained_but_audit_only( @pytest.mark.asyncio -async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( +async def test_malformed_provider_capture_preserves_native_evidence_and_cleanup( tmp_path: Path, ) -> None: - """Guards PR #1057 against leaking OTel on malformed mixed capture input.""" + """Guards PR #1057 against malformed provider rows suppressing native evidence.""" capture = LLMTrajectoryCapture( tmp_path, @@ -341,8 +346,8 @@ async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( role="solver", ) native = _CaptureTarget( - agent="claude-agent-acp", - model="claude-sonnet-4-6", + agent="codex-acp", + model="gpt-5.6", credential_home="/home/agent", auth_mode=AuthMode.OAUTH_SUBSCRIPTION, native=True, @@ -354,26 +359,126 @@ async def test_malformed_provider_capture_stops_owned_collector_before_cleanup( (target.role, target.agent, target.model, target.credential_home) ] = target capture.trajectory_path.write_text("{malformed provider row\n") - capture._collector_owned = True - capture._capture_root_prepared = True + capture._otel_collector.owned = True + capture._otel_collector.root_prepared = True commands: list[str] = [] + session = tmp_path / "native-session" / "rollout-session-one.jsonl" + session.parent.mkdir(parents=True) + session.write_text( + json.dumps( + { + "timestamp": "2026-08-29T12:00:01Z", + "type": "response_item", + "payload": { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "native answer"}], + }, + } + ) + + "\n" + ) + native_result = parse_codex_sessions( + session.parent, + agent=native.agent, + session_id="rollout-1", + started_at=datetime(2026, 8, 29, 11, 59, tzinfo=UTC), + configured_model=native.model, + ) + assert native_result is not None + class RecordingEnv: async def exec(self, command, **_kwargs): commands.append(command) return SimpleNamespace(return_code=0, stdout="", stderr="") - with pytest.raises(ValueError, match="invalid LLM trajectory JSONL"): - await capture.finalize( - RecordingEnv(), - acp_events=[], - model_call_seen=True, + async def collect_native(env, *, targets): + assert targets == [native] + await capture._otel_collector.stop(env) + return NativeCollection( + bundles=(NativeCaptureBundle(targets=(native,), result=native_result),) ) + capture._native_collector.collect = collect_native + await capture.finalize(RecordingEnv(), acp_events=[], model_call_seen=True) + stop_index = next(i for i, command in enumerate(commands) if "old_pid" in command) cleanup_index = next( i for i, command in enumerate(commands) if "for attempt in 1 2 3" in command ) assert stop_index < cleanup_index - assert capture._collector_owned is False - assert capture._capture_root_prepared is False + row = json.loads(capture.trajectory_path.read_text()) + assert row["metadata"]["agent"] == "codex-acp" + assert capture.manifest.status is CaptureStatus.PARTIAL + assert any("provider capture parse failed" in e for e in capture.manifest.errors) + assert capture._otel_collector.owned is False + assert capture._otel_collector.root_prepared is False + + +@pytest.mark.asyncio +async def test_manifest_write_failure_preserves_already_assembled_rows( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Guards PR #1057 against erasing rows after a sidecar write failure.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="openai/gpt-5.6", + session_id="rollout-1", + started_at=datetime(2026, 8, 29, 12, 0, tzinfo=UTC), + ) + target = _CaptureTarget( + agent="codex-acp", + model="openai/gpt-5.6", + credential_home="/home/agent", + auth_mode=AuthMode.API_KEY, + native=False, + role="agent", + ) + capture._targets[("agent", target.agent, target.model, target.credential_home)] = ( + target + ) + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"model": "gpt-5.6", "input": "hello"}}, + "response": { + "status_code": 200, + "body": { + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + }, + "metadata": {"benchflow_requested_model": "openai/gpt-5.6"}, + } + ) + + "\n" + ) + write_manifest = llm_capture_module.write_llm_trajectory_manifest + calls = 0 + + def fail_once(rollout_dir, manifest): + nonlocal calls + calls += 1 + if calls == 1: + raise OSError("manifest disk write failed") + write_manifest(rollout_dir, manifest) + + monkeypatch.setattr(llm_capture_module, "write_llm_trajectory_manifest", fail_once) + with pytest.raises(OSError, match="manifest disk write failed"): + await capture.finalize(None, acp_events=[], model_call_seen=True) + + assembled = capture.trajectory_path.read_text() + capture.record_failure("manifest disk write failed", model_call_seen=True) + + assert capture.trajectory_path.read_text() == assembled + assert len(assembled.splitlines()) == 1 + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["status"] == "capture_failed" + assert manifest["capture_source"] == "litellm_proxy" + assert manifest["capture_fidelity"] == "provider_wire" + assert manifest["exchange_count"] == 1 diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index 9b1c7180f..82fff3090 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -14,7 +14,6 @@ LLMTrajectoryCapture, _CaptureTarget, _NativeCaptureBundle, - _NativeCollection, ) from benchflow.trajectories.llm_capture_manifest import ( AuthMode, @@ -22,6 +21,7 @@ CaptureSource, CaptureStatus, ) +from benchflow.trajectories.native_capture_collection import NativeCollection from benchflow.trajectories.native_capture_parsers import ( parse_claude_raw_capture, parse_claude_sessions, @@ -650,7 +650,7 @@ async def exec(self, command, **_kwargs): native=True, ) capture._targets[(target.agent, target.model, target.credential_home)] = target - capture._capture_root_prepared = True + capture._otel_collector.root_prepared = True capture.trajectory_path.write_text( json.dumps( { @@ -742,12 +742,13 @@ async def test_mixed_auth_rollout_merges_provider_and_native_exchanges( ) assert native_result is not None - async def collect_native(_env): - return _NativeCollection( + async def collect_native(_env, *, targets): + assert targets == [native_target] + return NativeCollection( bundles=(_NativeCaptureBundle((native_target,), native_result),) ) - capture._collect_native_results = collect_native + capture._native_collector.collect = collect_native await capture.finalize(object(), acp_events=[], model_call_seen=True) @@ -813,10 +814,11 @@ async def test_mixed_auth_rollout_marks_missing_native_role_partial( + "\n" ) - async def collect_native(_env): - return _NativeCollection() + async def collect_native(_env, *, targets): + assert targets == [native_target] + return NativeCollection() - capture._collect_native_results = collect_native + capture._native_collector.collect = collect_native await capture.finalize(object(), acp_events=[], model_call_seen=True) diff --git a/tests/trajectories/test_native_session_boundaries.py b/tests/trajectories/test_native_session_boundaries.py index 3efc20ddd..7ba1106e0 100644 --- a/tests/trajectories/test_native_session_boundaries.py +++ b/tests/trajectories/test_native_session_boundaries.py @@ -12,11 +12,12 @@ import pytest from benchflow.rollout import Rollout -from benchflow.trajectories import llm_capture as llm_capture_module +from benchflow.trajectories import ( + native_capture_collection as native_capture_module, +) from benchflow.trajectories.llm_capture import ( LLMTrajectoryCapture, _CaptureTarget, - _download_bound_session_files, _NativeCaptureBundle, ) from benchflow.trajectories.llm_capture_manifest import ( @@ -24,6 +25,9 @@ CaptureFidelity, CaptureSource, ) +from benchflow.trajectories.native_capture_collection import ( + download_bound_session_files, +) from benchflow.trajectories.native_capture_parsers import ( NativeParseResult, parse_codex_sessions, @@ -365,14 +369,14 @@ async def download_dir(self, _remote, local): sandbox_user="agent", ) - assert capture._collector_started is False - assert capture._collector_owned is True + assert capture._otel_collector.started is False + assert capture._otel_collector.owned is True await capture.finalize(env, acp_events=[], model_call_seen=False) stop_commands = [command for command in commands if "read -r old_pid" in command] assert len(stop_commands) == 2 - assert capture._collector_owned is False + assert capture._otel_collector.owned is False assert any("-depth -delete" in command for command in commands) @@ -460,7 +464,7 @@ async def download_file(self, remote, local): destination.write_text("{}\n") destination = tmp_path / "missing-parent" / "sessions" - downloaded = await _download_bound_session_files( + downloaded = await download_bound_session_files( DockerLikeEnv(), "/home/agent/.codex/sessions", destination, @@ -596,7 +600,7 @@ async def test_claude_fallback_collects_only_raw_uncovered_exchanges( capture._targets[ (target.role, target.agent, target.model, target.credential_home) ] = target - capture._capture_root_prepared = True + capture._otel_collector.root_prepared = True raw_result = _native_result( session_id=first_id, source=CaptureSource.CLAUDE_OTEL_RAW_BODY, @@ -645,22 +649,22 @@ async def download_bound(_env, _remote, local, *, started_at, session_ids): return True monkeypatch.setattr( - llm_capture_module, + native_capture_module, "parse_claude_raw_capture", lambda *_args, **_kwargs: raw_result, ) monkeypatch.setattr( - llm_capture_module, - "_download_bound_session_files", + native_capture_module, + "download_bound_session_files", download_bound, ) monkeypatch.setattr( - llm_capture_module, + native_capture_module, "parse_claude_sessions", lambda *_args, **_kwargs: fallback_result, ) - collection = await capture._collect_native_results(CaptureEnv()) + collection = await capture._native_collector.collect(CaptureEnv(), targets=[target]) assert len(collection.bundles) == 2 assert fallback_calls == [(first_id,)] @@ -720,7 +724,9 @@ async def collect_codex(_env, *, local_root, index, target): raise RuntimeError("second target download failed") return _NativeCaptureBundle(targets=(first,), result=first_result) - monkeypatch.setattr(capture, "_collect_codex_session", collect_codex) + monkeypatch.setattr( + capture._native_collector, "_collect_codex_session", collect_codex + ) await capture.finalize(object(), acp_events=[], model_call_seen=True) @@ -756,9 +762,9 @@ async def exec(self, command, **_kwargs): started_at=STARTED_AT, ) capture.configure({"ANTHROPIC_API_KEY": "test-key"}) - capture._collector_started = True - capture._collector_owned = True - capture._capture_root_prepared = True + capture._otel_collector.started = True + capture._otel_collector.owned = True + capture._otel_collector.root_prepared = True await capture.finalize( CleanupEnv(), @@ -768,8 +774,8 @@ async def exec(self, command, **_kwargs): assert any("kill -TERM" in command for command in commands) assert any("-depth -delete" in command for command in commands) - assert capture._collector_started is False - assert capture._collector_owned is False + assert capture._otel_collector.started is False + assert capture._otel_collector.owned is False @pytest.mark.asyncio @@ -797,7 +803,7 @@ async def exec(self, command, **_kwargs): started_at=STARTED_AT, ) capture.configure({"OPENAI_API_KEY": "test-key"}) - capture._capture_root_prepared = True + capture._otel_collector.root_prepared = True capture.trajectory_path.write_text( json.dumps( { @@ -821,7 +827,7 @@ async def exec(self, command, **_kwargs): assert manifest["status"] == "capture_failed" assert manifest["exchange_count"] == 1 assert any("cleanup failed" in error for error in manifest["errors"]) - assert capture._capture_root_prepared is True + assert capture._otel_collector.root_prepared is True def test_rollout_binds_the_acp_session_after_connect() -> None: From 604351459f55f38aeb60027e03c35855ea6550fa Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 07:01:01 -0700 Subject: [PATCH 33/74] fix: preserve canonical provider requests --- src/benchflow/providers/litellm_logging.py | 26 +++++++--- .../trajectories/llm_capture_records.py | 8 ++- tests/test_litellm_hardening.py | 11 +++- tests/test_litellm_logging.py | 51 +++++++++++++++++++ tests/trajectories/test_native_llm_capture.py | 41 +++++++++++++++ 5 files changed, 128 insertions(+), 9 deletions(-) diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 7b67e5e63..6f038fb34 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -301,13 +301,23 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - litellm_params = kwargs.get("litellm_params") or {} optional_params = kwargs.get("optional_params") or {} metadata = kwargs.get("metadata") or litellm_params.get("metadata") or {} - request_body = { - "model": kwargs.get("model"), - "messages": kwargs.get("messages"), - "input": kwargs.get("input"), - "tools": optional_params.get("tools") or kwargs.get("tools"), - "stream": optional_params.get("stream") or kwargs.get("stream"), - } + proxy_request = litellm_params.get("proxy_server_request") or {} + proxy_body = ( + proxy_request.get("body") if isinstance(proxy_request, dict) else None + ) + request_complete = isinstance(proxy_body, dict) + # LiteLLM preserves the complete incoming provider request body here. + # Start from that canonical payload so ordinary/future sampling and + # output parameters are not silently lost by a hard-coded allowlist. + request_body = dict(proxy_body) if request_complete else {} + for key in ("model", "messages", "input"): + if kwargs.get(key) is not None: + request_body[key] = kwargs[key] + for key in ("tools", "stream"): + if key in optional_params: + request_body[key] = optional_params[key] + elif kwargs.get(key) is not None: + request_body[key] = kwargs[key] for key in ("reasoning_effort", "thinking", "output_config"): value = optional_params.get(key) if value is None: @@ -336,6 +346,7 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - "provider_model": litellm_params.get("model") or kwargs.get("model"), "model_group": metadata.get("model_group") if isinstance(metadata, dict) else None, "call_type": kwargs.get("call_type") or litellm_params.get("call_type"), + "request_complete": request_complete, "input_shape": { "has_messages": bool(kwargs.get("messages")), "has_input": kwargs.get("input") is not None, @@ -541,6 +552,7 @@ def _exchange_metadata( agent_name=agent_name, request_body=request_body, ) + metadata["request_complete"] = record.get("request_complete") is True return metadata diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 64e59675f..0c1f1b17a 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -93,6 +93,7 @@ def load_provider_wire_records( capture_trusted = ( target.provider_capture_trusted if target is not None else True ) + request_complete = metadata.get("request_complete") is True metadata.update( { "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, @@ -123,7 +124,7 @@ def load_provider_wire_records( else (_record_model(record) or fallback_model) ), "role_attribution_complete": attribution_complete, - "request_complete": True, + "request_complete": request_complete, "response_complete": True, "payload_redacted": True, } @@ -208,6 +209,11 @@ def assemble_capture( missing_fields = { field for bundle in native_bundles for field in bundle.result.missing_fields } + if any( + not _record_metadata_bool(record, "request_complete") + for record in provider_records + ): + missing_fields.add("provider_request") successful_records = [ record for record in records diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index 792866be4..bd44f12b5 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -752,7 +752,15 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( kwargs = { "model": "benchflow-openai-gpt-4.1-mini", "messages": [{"role": "user", "content": "hi"}], - "litellm_params": {"model": "openai/gpt-4.1-mini"}, + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "proxy_server_request": { + "body": { + "model": "benchflow-openai-gpt-4.1-mini", + "messages": [{"role": "user", "content": "hi"}], + } + }, + }, "optional_params": {}, "call_type": "acompletion", } @@ -774,6 +782,7 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( assert usage["usage_source"] == "provider_response" assert usage["n_input_tokens"] == 12 assert usage["n_output_tokens"] == 4 + assert trajectory.exchanges[0].metadata["request_complete"] is True assert json.loads(state_path.read_text()) == { "attempt_count": 1, "terminal_count": 1, diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 02c892efe..87cb3577e 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -149,6 +149,55 @@ def test_callback_module_source_exposes_proxy_handler_instance(): assert "proxy_handler_instance = BenchFlowLiteLLMLogger()" in source +def test_callback_preserves_complete_proxy_request_body(): + """Guards PR #1057 against claiming completeness for an allowlisted request.""" + + logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() + now = datetime.now() + proxy_body = { + "model": "benchflow-openai-gpt-5.6", + "messages": [{"role": "user", "content": "hi"}], + "temperature": 0.25, + "max_tokens": 321, + "top_p": 0.8, + "stop": ["DONE"], + "tool_choice": "required", + "response_format": {"type": "json_object"}, + } + + record = logger._base_record( + { + "model": "benchflow-openai-gpt-5.6", + "messages": proxy_body["messages"], + "litellm_params": { + "model": "openai/gpt-5.6", + "proxy_server_request": {"body": proxy_body}, + }, + "optional_params": {"stream": False}, + }, + now, + now, + ) + + assert record["request_complete"] is True + assert record["request"]["body"] == {**proxy_body, "stream": False} + + +def test_callback_marks_request_incomplete_without_proxy_body(): + """Guards PR #1057 against promoting reconstructed callback parameters.""" + + logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() + now = datetime.now() + + record = logger._base_record( + {"model": "gpt-5.6", "messages": [{"role": "user", "content": "hi"}]}, + now, + now, + ) + + assert record["request_complete"] is False + + @pytest.mark.asyncio async def test_callback_pre_call_hook_strips_chat_input_compat_field(): namespace: dict[str, object] = {} @@ -345,6 +394,7 @@ def test_opencode_callback_import_preserves_call_metadata_and_purpose(): "provider_model": "openai/glm-5.1", "model_group": "benchflow-glm-5.1", "call_type": "completion", + "request_complete": True, "input_shape": { "has_messages": True, "has_input": True, @@ -462,6 +512,7 @@ def test_opencode_callback_import_preserves_call_metadata_and_purpose(): "n_messages": 2, }, "call_purpose": "agent", + "request_complete": True, } assert [exchange.metadata["call_purpose"] for exchange in trajectory.exchanges] == [ "agent", diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index 82fff3090..effa27173 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -605,6 +605,7 @@ async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> }, }, }, + "metadata": {"request_complete": True}, } ) + "\n" @@ -622,6 +623,46 @@ async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> assert manifest["capture_source"] == "litellm_proxy" +@pytest.mark.asyncio +async def test_provider_jsonl_without_canonical_request_is_partial( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on an allowlisted callback request.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="openai/gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.configure({"OPENAI_API_KEY": "test-key"}) + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"model": "gpt-5.6", "input": "hello"}}, + "response": { + "status_code": 200, + "body": { + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + }, + } + ) + + "\n" + ) + + await capture.finalize(None, acp_events=[], model_call_seen=True) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["status"] == "partial" + assert manifest["request_complete"] is False + assert "provider_request" in manifest["missing_fields"] + + @pytest.mark.asyncio async def test_provider_capture_early_return_cleans_native_raw_bodies( tmp_path: Path, From 5b962c665ce14ef2971bafcb54673bb46309e6ee Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 07:28:33 -0700 Subject: [PATCH 34/74] fix: preserve provider capture fidelity --- src/benchflow/continue_run/replay_proxy.py | 33 +++- src/benchflow/continue_run/sandbox_proxy.py | 47 +++-- src/benchflow/providers/litellm_logging.py | 186 +++++++++++++++--- .../trajectories/llm_capture_records.py | 6 +- tests/continue_run/test_replay_proxy.py | 82 ++++++++ tests/test_litellm_hardening.py | 14 ++ tests/test_litellm_logging.py | 77 +++++++- .../test_native_capture_resilience.py | 8 +- tests/trajectories/test_native_llm_capture.py | 14 +- 9 files changed, 411 insertions(+), 56 deletions(-) diff --git a/src/benchflow/continue_run/replay_proxy.py b/src/benchflow/continue_run/replay_proxy.py index 01a196154..6e2be2e08 100644 --- a/src/benchflow/continue_run/replay_proxy.py +++ b/src/benchflow/continue_run/replay_proxy.py @@ -75,8 +75,21 @@ def __init__( self.divergences = 0 self.live_attempt_count = 0 self.live_errors: list[str] = [] - # Live-leg exchanges, in order, for stitching onto the recorded prefix. - self.live_exchanges: list[LLMExchange] = [] + # Provider calls can complete out of order. Retain the attempt sequence + # assigned at dispatch so stitching follows agent request order rather + # than response arrival order. + self._live_exchanges: list[tuple[int, LLMExchange]] = [] + + @property + def live_exchanges(self) -> list[LLMExchange]: + """Return captured live exchanges in provider-attempt order.""" + with self._lock: + return [ + exchange + for _, exchange in sorted( + self._live_exchanges, key=lambda item: item[0] + ) + ] @property def exhausted(self) -> bool: @@ -118,6 +131,7 @@ def next_response(self, request_body: dict[str, Any]) -> ReplayResult: # Past the cut-point: live continuation. self._cursor += 1 self.live_attempt_count += 1 + live_attempt = self.live_attempt_count forwarder = self._live_forwarder if forwarder is None: @@ -150,12 +164,17 @@ def next_response(self, request_body: dict[str, Any]) -> ReplayResult: raise # Capture the live exchange so the caller can stitch a continuous # llm_trajectory.jsonl (recorded prefix + live suffix). - self.live_exchanges.append( - LLMExchange( - request=LLMRequest(body=request_body), - response=LLMResponse(status_code=200, body=body), + with self._lock: + self._live_exchanges.append( + ( + live_attempt, + LLMExchange( + request=LLMRequest(body=request_body), + response=LLMResponse(status_code=200, body=body), + metadata={"continuation_attempt": live_attempt}, + ), + ) ) - ) return ReplayResult(source="live", status=200, body=body) diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index fe0245806..87d2fd2d8 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -203,6 +203,7 @@ def next_response(self, request_body): return "replay", int(response.get("status_code") or 200), dict(response.get("body") or {}) self.cursor += 1 self.live_attempt_count += 1 + live_attempt = self.live_attempt_count self.active_live_requests += 1 try: self._write_state() @@ -215,7 +216,10 @@ def next_response(self, request_body): status, body, provider_observed = self._forward_live(request_body) if provider_observed: try: - self._append_live_exchange(request_body, status, body) + with self.lock: + self._append_live_exchange( + request_body, status, body, live_attempt + ) except Exception: with self.lock: self.live_error_count += 1 @@ -284,10 +288,11 @@ def _forward_live(self, request_body): traceback.print_exc() return 500, {"error": {"message": str(exc)}}, False - def _append_live_exchange(self, request_body, status, body): + def _append_live_exchange(self, request_body, status, body, live_attempt): row = { "request": {"body": request_body}, "response": {"status_code": status, "body": body}, + "metadata": {"continuation_attempt": live_attempt}, } with open(self.live_log_path, "a", encoding="utf-8") as handle: handle.write(json.dumps(row) + "\n") @@ -443,6 +448,33 @@ def main(): """ +def _ordered_live_exchange_log(text: str) -> tuple[list[LLMExchange], int]: + """Parse sandbox live rows and restore their assigned attempt order.""" + sequenced: list[tuple[int, LLMExchange]] = [] + malformed = 0 + seen: set[int] = set() + for raw in text.splitlines(): + if not raw.strip(): + continue + try: + exchange = LLMExchange.model_validate_json(raw) + except Exception: + malformed += 1 + continue + attempt = exchange.metadata.get("continuation_attempt") + if ( + not isinstance(attempt, int) + or isinstance(attempt, bool) + or attempt <= 0 + or attempt in seen + ): + malformed += 1 + continue + seen.add(attempt) + sequenced.append((attempt, exchange)) + return [exchange for _, exchange in sorted(sequenced)], malformed + + async def _upload_text(sandbox: Any, text: str, target_path: str, suffix: str) -> None: with tempfile.NamedTemporaryFile("w", suffix=suffix, delete=False) as tmp: tmp.write(text) @@ -639,16 +671,7 @@ async def _terminate(self) -> None: async def _load_live_exchanges(self) -> list[LLMExchange]: text = await _read_remote_text(self.sandbox, self.live_log_path) - exchanges: list[LLMExchange] = [] - malformed = 0 - for raw in text.splitlines(): - if not raw.strip(): - continue - try: - exchanges.append(LLMExchange.model_validate_json(raw)) - except Exception: - malformed += 1 - continue + exchanges, malformed = _ordered_live_exchange_log(text) if malformed: self.live_errors.append( f"{malformed} sandbox live exchange record(s) were malformed" diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 6f038fb34..5e266a8b0 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -183,6 +183,80 @@ def _jsonable(value: Any) -> Any: return str(value) +_SENSITIVE_FIELD_SUFFIXES = ( + "api_key", + "access_key", + "account_key", + "secret_key", + "private_key", + "access_token", + "refresh_token", + "session_token", + "client_secret", + "token", + "secret", + "credential", + "password", + "passwd", + "credentials", +) +_SENSITIVE_FIELD_NAMES = { + "authorization", + "proxy_authorization", + "x_api_key", + "x_goog_api_key", + "api_key", + "apikey", +} +_URL_CREDENTIAL_RE = re.compile( + r"([?&](?:" + r"api_?key|access_key|access_token|refresh_token|session_token|" + r"session_id|sessionid|client_secret|aws_secret_access_key|account_key|" + r"secret|password_hash|password|passwd|" + r"signature|x-amz-signature|x-amz-credential|x-amz-security-token|sas" + r")=)[^&\s\"'*]+", + re.IGNORECASE, +) +_AZURE_SAS_RE = re.compile(r"([?&]sig=)[^&\s\"'*]{16,}", re.IGNORECASE) +_AUTHORIZATION_RE = re.compile( + r"\b(Bearer|Token|Basic)\s+(?!\*\*\*REDACTED\*\*\*)[^\s,;\"']+", + re.IGNORECASE, +) + + +def _is_sensitive_field(key: Any) -> bool: + if not isinstance(key, str): + return False + normalized = re.sub(r"[^a-z0-9]+", "_", key.lower()).strip("_") + return normalized in _SENSITIVE_FIELD_NAMES or normalized.endswith( + _SENSITIVE_FIELD_SUFFIXES + ) + + +def _redact_storage_text(value: str) -> str: + value = _URL_CREDENTIAL_RE.sub(r"\1***REDACTED***", value) + value = _AZURE_SAS_RE.sub(r"\1***REDACTED***", value) + return _AUTHORIZATION_RE.sub(r"\1 ***REDACTED***", value) + + +def _redact_for_storage(value: Any) -> Any: + # Redact callback payloads before the durable append-only write. + if isinstance(value, dict): + return { + str(key): ( + "***REDACTED***" + if _is_sensitive_field(key) + else _redact_for_storage(item) + ) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [_redact_for_storage(item) for item in value] + if isinstance(value, str): + return _redact_storage_text(value) + return value + + def _iso(value: Any) -> str: if isinstance(value, datetime): return value.isoformat() @@ -229,6 +303,7 @@ def __init__(self) -> None: self._capture_lock = threading.Lock() self._attempt_count = 0 self._terminal_count = 0 + self._provider_requests: dict[str, dict[str, Any]] = {} with self._capture_lock: self._write_state_locked() @@ -286,52 +361,105 @@ def _write(self, payload: dict[str, Any]) -> None: if not path: return payload["logged_at"] = datetime.now(timezone.utc).isoformat() + durable_payload = _redact_for_storage(_jsonable(payload)) with self._capture_lock: os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "a", encoding="utf-8") as handle: handle.write( - json.dumps(_jsonable(payload), separators=(",", ":")) + "\n" + json.dumps(durable_payload, separators=(",", ":")) + "\n" ) handle.flush() os.fsync(handle.fileno()) self._terminal_count += 1 self._write_state_locked() + @staticmethod + def _call_id(kwargs: dict[str, Any]) -> str | None: + direct = kwargs.get("litellm_call_id") + if isinstance(direct, str) and direct: + return direct + litellm_params = kwargs.get("litellm_params") or {} + nested = ( + litellm_params.get("litellm_call_id") + if isinstance(litellm_params, dict) + else None + ) + return nested if isinstance(nested, str) and nested else None + + def log_pre_api_call(self, model, messages, kwargs) -> None: + # Retain LiteLLM's post-transform body until its terminal callback. + if not isinstance(kwargs, dict): + return + call_id = self._call_id(kwargs) + additional_args = kwargs.get("additional_args") or {} + provider_body = ( + additional_args.get("complete_input_dict") + if isinstance(additional_args, dict) + else None + ) + jsonable_body = _jsonable(provider_body) + if call_id is None or not isinstance(jsonable_body, dict): + return + with self._capture_lock: + self._provider_requests[call_id] = _redact_for_storage(jsonable_body) + def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) -> dict[str, Any]: litellm_params = kwargs.get("litellm_params") or {} optional_params = kwargs.get("optional_params") or {} metadata = kwargs.get("metadata") or litellm_params.get("metadata") or {} + additional_args = kwargs.get("additional_args") or {} + callback_provider_body = ( + additional_args.get("complete_input_dict") + if isinstance(additional_args, dict) + else None + ) + call_id = self._call_id(kwargs) + with self._capture_lock: + provider_body = ( + self._provider_requests.pop(call_id, None) + if call_id is not None + else None + ) + if provider_body is None: + provider_body = callback_provider_body + jsonable_provider_body = _jsonable(provider_body) + request_complete = isinstance(jsonable_provider_body, dict) proxy_request = litellm_params.get("proxy_server_request") or {} proxy_body = ( proxy_request.get("body") if isinstance(proxy_request, dict) else None ) - request_complete = isinstance(proxy_body, dict) - # LiteLLM preserves the complete incoming provider request body here. - # Start from that canonical payload so ordinary/future sampling and - # output parameters are not silently lost by a hard-coded allowlist. - request_body = dict(proxy_body) if request_complete else {} - for key in ("model", "messages", "input"): - if kwargs.get(key) is not None: - request_body[key] = kwargs[key] - for key in ("tools", "stream"): - if key in optional_params: - request_body[key] = optional_params[key] - elif kwargs.get(key) is not None: - request_body[key] = kwargs[key] - for key in ("reasoning_effort", "thinking", "output_config"): - value = optional_params.get(key) - if value is None: - value = kwargs.get(key) - if value is None: - value = litellm_params.get(key) - if value is not None: - request_body[key] = value - for key in ("logprobs", "top_logprobs"): - value = optional_params.get(key) - if value is None: - value = kwargs.get(key) - if value is not None: - request_body[key] = value + if request_complete: + # ``Logging.pre_call(..., complete_input_dict=...)`` is populated + # after provider-specific transformation and is the body handed to + # the HTTP client. Proxy ingress is only an audit fallback: routes + # may drop or rewrite its fields before the provider sees them. + request_body = dict(jsonable_provider_body) + request_capture_source = "litellm_pre_api_call_complete_input_dict" + else: + request_body = dict(proxy_body) if isinstance(proxy_body, dict) else {} + for key in ("model", "messages", "input"): + if kwargs.get(key) is not None: + request_body[key] = kwargs[key] + for key in ("tools", "stream"): + if key in optional_params: + request_body[key] = optional_params[key] + elif kwargs.get(key) is not None: + request_body[key] = kwargs[key] + for key in ("reasoning_effort", "thinking", "output_config"): + value = optional_params.get(key) + if value is None: + value = kwargs.get(key) + if value is None: + value = litellm_params.get(key) + if value is not None: + request_body[key] = value + for key in ("logprobs", "top_logprobs"): + value = optional_params.get(key) + if value is None: + value = kwargs.get(key) + if value is not None: + request_body[key] = value + request_capture_source = "proxy_ingress_reconstruction" request_body = {k: v for k, v in request_body.items() if v is not None} return { "benchflow_agent": os.environ.get("BENCHFLOW_LITELLM_AGENT"), @@ -347,6 +475,7 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - "model_group": metadata.get("model_group") if isinstance(metadata, dict) else None, "call_type": kwargs.get("call_type") or litellm_params.get("call_type"), "request_complete": request_complete, + "request_capture_source": request_capture_source, "input_shape": { "has_messages": bool(kwargs.get("messages")), "has_input": kwargs.get("input") is not None, @@ -545,6 +674,7 @@ def _exchange_metadata( "model_group", "call_type", "input_shape", + "request_capture_source", ) if record.get(key) is not None } diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 0c1f1b17a..1c178428c 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -94,13 +94,17 @@ def load_provider_wire_records( target.provider_capture_trusted if target is not None else True ) request_complete = metadata.get("request_complete") is True + provider_request_observed = ( + metadata.get("request_capture_source") + == "litellm_pre_api_call_complete_input_dict" + ) metadata.update( { "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, "capture_source": CaptureSource.LITELLM_PROXY.value, "capture_fidelity": ( CaptureFidelity.PROVIDER_WIRE.value - if capture_trusted + if capture_trusted and provider_request_observed else CaptureFidelity.AGENT_SESSION.value ), "auth_mode": ( diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index b15386311..a1214a8a0 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -16,6 +16,7 @@ ) from benchflow.continue_run.sandbox_proxy import ( SandboxReplayProxy, + _ordered_live_exchange_log, _sandbox_proxy_source, ) @@ -84,6 +85,43 @@ def fail(_request): ] +def test_host_live_exchanges_preserve_attempt_order_under_concurrency() -> None: + """Guards PR #1057 against stitching host calls in completion order.""" + + first_started = threading.Event() + release_first = threading.Event() + + def forward(request): + request_id = request["request_id"] + if request_id == 1: + first_started.set() + assert release_first.wait(timeout=5) + return completion(content=f"live-{request_id}") + + router = ReplayRouter([], live_forwarder=forward) + results: list[object] = [] + first = threading.Thread( + target=lambda: results.append(router.next_response({"request_id": 1})) + ) + second = threading.Thread( + target=lambda: results.append(router.next_response({"request_id": 2})) + ) + + first.start() + assert first_started.wait(timeout=5) + second.start() + second.join(timeout=5) + assert not second.is_alive() + release_first.set() + first.join(timeout=5) + + assert [row.request.body["request_id"] for row in router.live_exchanges] == [1, 2] + assert [row.metadata["continuation_attempt"] for row in router.live_exchanges] == [ + 1, + 2, + ] + + def test_sandbox_forwarding_failure_is_not_logged_as_provider_exchange( tmp_path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -200,6 +238,50 @@ def forward(_request): assert capture_state["live_error_count"] == 1 +def test_sandbox_live_exchange_recovery_restores_attempt_order(tmp_path) -> None: + """Guards PR #1057 against stitching sandbox calls in completion order.""" + + namespace: dict[str, object] = {} + exec(_sandbox_proxy_source(), namespace) + live_log = tmp_path / "live.jsonl" + state = namespace["ReplayState"]( + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(live_log), + state_path=str(tmp_path / "state.json"), + port=61357, + ) + first_started = threading.Event() + release_first = threading.Event() + + def forward(request): + request_id = request["request_id"] + if request_id == 1: + first_started.set() + assert release_first.wait(timeout=5) + return 200, completion(content=f"live-{request_id}"), True + + state._forward_live = forward + first = threading.Thread(target=lambda: state.next_response({"request_id": 1})) + second = threading.Thread(target=lambda: state.next_response({"request_id": 2})) + + first.start() + assert first_started.wait(timeout=5) + second.start() + second.join(timeout=5) + assert not second.is_alive() + release_first.set() + first.join(timeout=5) + + raw_rows = [json.loads(line) for line in live_log.read_text().splitlines()] + assert [row["metadata"]["continuation_attempt"] for row in raw_rows] == [2, 1] + exchanges, malformed = _ordered_live_exchange_log(live_log.read_text()) + assert malformed == 0 + assert [row.request.body["request_id"] for row in exchanges] == [1, 2] + + def test_sandbox_quiesce_closes_listener_and_drains_accepted_handlers( tmp_path, ) -> None: diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index bd44f12b5..f41ad7ace 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -750,6 +750,7 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( "usage": {"prompt_tokens": 12, "completion_tokens": 4, "total_tokens": 16}, } kwargs = { + "litellm_call_id": "round-trip-provider-request", "model": "benchflow-openai-gpt-4.1-mini", "messages": [{"role": "user", "content": "hi"}], "litellm_params": { @@ -768,6 +769,19 @@ async def test_embedded_callback_logger_round_trips_to_provider_usage( end = datetime(2026, 6, 4, 10, 0, 1) await logger.async_pre_call_hook(None, None, kwargs, "acompletion") + logger.log_pre_api_call( + "openai/gpt-4.1-mini", + kwargs["messages"], + { + "litellm_call_id": kwargs["litellm_call_id"], + "additional_args": { + "complete_input_dict": { + "model": "gpt-4.1-mini", + "messages": kwargs["messages"], + } + }, + }, + ) await logger.async_log_success_event(kwargs, response, start, end) text = log_path.read_text() diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 87cb3577e..5c1cf195f 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -149,8 +149,8 @@ def test_callback_module_source_exposes_proxy_handler_instance(): assert "proxy_handler_instance = BenchFlowLiteLLMLogger()" in source -def test_callback_preserves_complete_proxy_request_body(): - """Guards PR #1057 against claiming completeness for an allowlisted request.""" +def test_callback_preserves_post_transform_provider_request_body(): + """Guards PR #1057 against labeling proxy ingress as provider wire.""" logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() now = datetime.now() @@ -163,10 +163,24 @@ def test_callback_preserves_complete_proxy_request_body(): "stop": ["DONE"], "tool_choice": "required", "response_format": {"type": "json_object"}, + "input": "removed before provider dispatch", } + provider_body = {key: value for key, value in proxy_body.items() if key != "input"} + provider_body["model"] = "gpt-5.6-deployment" + provider_body["stream"] = False + call_id = "call-provider-body" + logger.log_pre_api_call( + provider_body["model"], + provider_body["messages"], + { + "litellm_call_id": call_id, + "additional_args": {"complete_input_dict": provider_body}, + }, + ) record = logger._base_record( { + "litellm_call_id": call_id, "model": "benchflow-openai-gpt-5.6", "messages": proxy_body["messages"], "litellm_params": { @@ -180,22 +194,73 @@ def test_callback_preserves_complete_proxy_request_body(): ) assert record["request_complete"] is True - assert record["request"]["body"] == {**proxy_body, "stream": False} + assert record["request_capture_source"] == ( + "litellm_pre_api_call_complete_input_dict" + ) + assert record["request"]["body"] == provider_body + assert "input" not in record["request"]["body"] -def test_callback_marks_request_incomplete_without_proxy_body(): - """Guards PR #1057 against promoting reconstructed callback parameters.""" +def test_callback_marks_proxy_ingress_request_incomplete(): + """Guards PR #1057 against promoting reconstructed proxy-ingress parameters.""" logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() now = datetime.now() record = logger._base_record( - {"model": "gpt-5.6", "messages": [{"role": "user", "content": "hi"}]}, + { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "hi"}], + "litellm_params": { + "proxy_server_request": { + "body": {"model": "alias", "temperature": 0.25} + } + }, + }, now, now, ) assert record["request_complete"] is False + assert record["request_capture_source"] == "proxy_ingress_reconstruction" + assert record["request"]["body"]["temperature"] == 0.25 + + +def test_callback_redacts_secrets_before_durable_journal( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Guards PR #1057 against persisting raw secrets before finalization.""" + + logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() + log_path = tmp_path / "callback.jsonl" + monkeypatch.setenv("BENCHFLOW_LITELLM_LOG_PATH", str(log_path)) + api_key = "provider-key-without-a-recognizable-prefix" + access_token = "oauth-token-without-a-recognizable-prefix" + sas = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG" + + logger._write( + { + "event": "success", + "request": { + "body": { + "api_key": api_key, + "nested": {"access_token": access_token}, + "download_url": f"https://blob.invalid/x?sig={sas}&sp=r", + "note": f"Authorization: Bearer {access_token}", + } + }, + } + ) + + raw = log_path.read_text() + assert api_key not in raw + assert access_token not in raw + assert sas not in raw + row = json.loads(raw) + body = row["request"]["body"] + assert body["api_key"] == "***REDACTED***" + assert body["nested"]["access_token"] == "***REDACTED***" + assert "sig=***REDACTED***" in body["download_url"] @pytest.mark.asyncio diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 62f273f55..be69c74eb 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -451,7 +451,13 @@ async def test_manifest_write_failure_preserves_already_assembled_rows( "usage": {"input_tokens": 1, "output_tokens": 1}, }, }, - "metadata": {"benchflow_requested_model": "openai/gpt-5.6"}, + "metadata": { + "benchflow_requested_model": "openai/gpt-5.6", + "request_complete": True, + "request_capture_source": ( + "litellm_pre_api_call_complete_input_dict" + ), + }, } ) + "\n" diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index effa27173..bd659539d 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -605,7 +605,12 @@ async def test_provider_jsonl_gets_complete_fidelity_metadata(tmp_path: Path) -> }, }, }, - "metadata": {"request_complete": True}, + "metadata": { + "request_complete": True, + "request_capture_source": ( + "litellm_pre_api_call_complete_input_dict" + ), + }, } ) + "\n" @@ -659,6 +664,7 @@ async def test_provider_jsonl_without_canonical_request_is_partial( (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() ) assert manifest["status"] == "partial" + assert manifest["capture_fidelity"] == "agent_session" assert manifest["request_complete"] is False assert "provider_request" in manifest["missing_fields"] @@ -750,6 +756,12 @@ async def test_mixed_auth_rollout_merges_provider_and_native_exchanges( "status_code": 200, "body": {"output": []}, }, + "metadata": { + "request_complete": True, + "request_capture_source": ( + "litellm_pre_api_call_complete_input_dict" + ), + }, } ) + "\n" From 87166fbe34807f72abac0ba8b0f819f48c304fe7 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 07:49:59 -0700 Subject: [PATCH 35/74] fix: fail closed on capture provenance --- docs/agent-quickstart.md | 7 +- docs/continue-runs.md | 13 +-- docs/getting-started.md | 5 +- src/benchflow/continue_run/orchestrator.py | 10 +-- .../continue_run/trajectory_artifacts.py | 63 +++++++------- src/benchflow/providers/litellm_logging.py | 86 ++----------------- .../trajectories/llm_capture_manifest.py | 1 + src/benchflow/trajectories/types.py | 33 +++++++ tests/continue_run/test_orchestrator.py | 30 ++++--- tests/test_litellm_logging.py | 43 ++++++++++ 10 files changed, 159 insertions(+), 132 deletions(-) diff --git a/docs/agent-quickstart.md b/docs/agent-quickstart.md index cf6aa5c5e..78f8656a8 100644 --- a/docs/agent-quickstart.md +++ b/docs/agent-quickstart.md @@ -151,9 +151,10 @@ and what it is: `llm_trajectory.jsonl` is an audit artifact for every run, not an unconditional training claim. Only a `complete` manifest with `provider_wire` fidelity and -positive provider token usage is training-ready. Native OAuth/session capture, -and a sandbox-local proxy that shares root custody with the agent, stay -available for audit but are marked lower-fidelity and excluded from training. +positive provider token usage is training-ready. Native OAuth/session capture +and continuation replay-proxy ingress stay available for audit but are marked +lower-fidelity and excluded from training. A sandbox-local proxy that shares +root custody with the agent records that additional custody limitation. Also note the job-level summary.json plus aggregated results.jsonl, verifiers.jsonl, and adp.jsonl in the job directory. The trainer/ files and the job-level diff --git a/docs/continue-runs.md b/docs/continue-runs.md index 00f6933b1..84d3314be 100644 --- a/docs/continue-runs.md +++ b/docs/continue-runs.md @@ -37,11 +37,14 @@ request/response pairs from the original run. `bench eval continue`: with a stitched `llm_trajectory.jsonl` (recorded prefix + live suffix) and `continued_from` provenance — a drop-in replacement for the timed-out entry. -The stitched manifest preserves capture custody. A live provider proxy across a -trusted host/non-root boundary can contribute `provider_wire` rows. If replay -and the untrusted agent share root custody in the sandbox, the live suffix is -still written for audit but is labeled `agent_session`; the stitched run cannot -be exported as training-ready provider evidence. +The stitched manifest preserves both capture boundary and custody. The live +suffix is observed at replay-proxy ingress, before the forwarder and LiteLLM can +replace, filter, or transform the provider request. It is therefore always +labeled `replay_proxy` / `agent_session` with an incomplete provider request, +even when the proxy is host-owned. If replay and the untrusted agent also share +root custody in the sandbox, the manifest records that additional limitation. +The suffix remains useful for audit and continuity, but any continued run with +a live suffix is excluded from training-ready provider evidence. Because the agent rebuilds its own state by re-doing its own steps, no reverse-engineering of OpenHands internals is needed, and the result is a single diff --git a/docs/getting-started.md b/docs/getting-started.md index 686c1fc86..eaa36b90f 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -224,13 +224,16 @@ the source of truth for interpreting the JSONL: | Claude Code subscription/OAuth | Claude Code raw API-body files correlated by local OTLP logs | `agent_session` | | Claude Code subscription/OAuth fallback | Claude Code native session JSONL | `agent_session` | | Codex subscription/OAuth | Codex native session JSONL | `agent_session` | +| Continued-run live suffix | Replay-proxy ingress before provider transformation | `agent_session` | The manifest status is `complete`, `partial`, `no_model_call`, or `capture_failed`. Mixed-role rollouts merge API-key and native-subscription exchanges into the same JSONL; `role_captures` records each prepared role's scene role, agent, model, auth mode, source, fidelity, completeness, and exchange count. Continued runs retain the source role entries as the `recorded` leg and -append a separate `live` leg, so a model or auth switch remains auditable. +append a separate `live` leg, so a model or auth switch remains auditable. The +live continuation request is captured before provider transformation and is +therefore audit-only even when its replay proxy is host-owned. Missing or ambiguously attributed roles make the rollout-level capture `partial`. Reconstructed `agent_session` rows remain useful for audit and viewer workflows, but trainer exports fail closed unless the manifest says the capture diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 52273c95a..1a5d44004 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -599,7 +599,7 @@ async def _continue_run_with_sandbox_proxy( rollout_name=rollout_name, ) rollout = await Rollout.create(config) - live_capture_trusted = False + live_capture_host_owned = False replay_proxy: SandboxReplayProxy | None = None provider_runtime: Any | None = None result: Any | None = None @@ -616,7 +616,7 @@ async def _write_artifacts( ) -> None: nonlocal \ artifacts_written, \ - live_capture_trusted, \ + live_capture_host_owned, \ live_exchanges, \ result, \ rollout_dir @@ -626,7 +626,7 @@ async def _write_artifacts( if result is None: return rollout_dir = Path(rollout._rollout_dir or (output_dir / rollout_name)) - live_capture_trusted = bool( + live_capture_host_owned = bool( provider_runtime is not None and getattr(provider_runtime, "capture_trusted", False) ) @@ -636,7 +636,7 @@ async def _write_artifacts( run.path / "trajectory" / "llm_trajectory.jsonl", live_exchanges, live_model=live_model, - live_capture_trusted=live_capture_trusted, + live_capture_host_owned=live_capture_host_owned, ) refresh_stitched_trajectory_manifest( rollout_dir, @@ -652,7 +652,7 @@ async def _write_artifacts( *(replay_proxy.live_errors if replay_proxy is not None else []), *(teardown_errors or []), ], - live_capture_trusted=live_capture_trusted, + live_capture_host_owned=live_capture_host_owned, ) update_continued_metadata( rollout_dir, diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py index 78cd8a62d..98bf4819c 100644 --- a/src/benchflow/continue_run/trajectory_artifacts.py +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -64,9 +64,15 @@ def write_stitched_trajectory( live_exchanges: list[LLMExchange], *, live_model: str | None = None, - live_capture_trusted: bool = True, + live_capture_host_owned: bool = True, ) -> Path: - """Write the stitched continuous trajectory into the new rollout folder.""" + """Write a stitched trajectory without overstating the live request boundary. + + Continuation captures the agent-facing request at replay-proxy ingress. The + live forwarder and LiteLLM can still replace, filter, or transform it before + provider dispatch, so the suffix remains useful for audit/replay but is never + provider-wire training evidence. + """ out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" out.parent.mkdir(parents=True, exist_ok=True) lines = stitched_trajectory_lines(original_llm_trajectory, []) @@ -79,21 +85,21 @@ def write_stitched_trajectory( "role": "agent", "model": live_model, "auth_mode": AuthMode.API_KEY.value, - "capture_source": CaptureSource.LITELLM_PROXY.value, - "capture_fidelity": ( - CaptureFidelity.PROVIDER_WIRE.value - if live_capture_trusted - else CaptureFidelity.AGENT_SESSION.value - ), + "capture_source": CaptureSource.REPLAY_PROXY.value, + "capture_fidelity": CaptureFidelity.AGENT_SESSION.value, "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, - "request_complete": True, + "request_complete": False, "response_complete": True, "role_attribution_complete": True, "payload_redacted": True, + "request_capture_source": "replay_proxy_ingress", + "capture_custody": ( + "host_owned" + if live_capture_host_owned + else "agent_writable_sandbox" + ), } ) - if not live_capture_trusted: - metadata["capture_custody"] = "agent_writable_sandbox" lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) rendered = "\n".join(lines) + ("\n" if lines else "") temporary = out.with_suffix(out.suffix + ".tmp") @@ -112,7 +118,7 @@ def refresh_stitched_trajectory_manifest( n_live: int, live_attempt_count: int, live_errors: list[str], - live_capture_trusted: bool = True, + live_capture_host_owned: bool = True, ) -> LLMTrajectoryManifest: """Replace rollout-finalization provenance with the final stitched contract.""" @@ -176,7 +182,7 @@ def refresh_stitched_trajectory_manifest( source_allows_training and count_matches and live_capture_complete - and (live_capture_trusted or n_live == 0) + and n_live == 0 and rows_valid and usage_complete ) @@ -186,15 +192,9 @@ def refresh_stitched_trajectory_manifest( source_auth = source.auth_mode if source else AuthMode.UNKNOWN if n_live: capture_source = ( - CaptureSource.LITELLM_PROXY - if source_capture_source is CaptureSource.LITELLM_PROXY - else CaptureSource.MIXED - ) - live_fidelity = ( - CaptureFidelity.PROVIDER_WIRE - if live_capture_trusted - else CaptureFidelity.AGENT_SESSION + CaptureSource.REPLAY_PROXY if n_recorded == 0 else CaptureSource.MIXED ) + live_fidelity = CaptureFidelity.AGENT_SESSION capture_fidelity = ( live_fidelity if source_fidelity is live_fidelity else CaptureFidelity.MIXED ) @@ -212,6 +212,7 @@ def refresh_stitched_trajectory_manifest( and count_matches and live_capture_complete and rows_valid + and n_live == 0 ) response_complete = bool( source @@ -223,7 +224,13 @@ def refresh_stitched_trajectory_manifest( errors = list(source.errors) if source and not complete else [] missing_fields = list(source.missing_fields) if source and not complete else [] errors.extend(live_errors) - if n_live and not live_capture_trusted: + if n_live: + errors.append( + "live continuation request was captured at replay-proxy ingress, " + "before provider transformation" + ) + missing_fields.append("live_provider_request") + if n_live and not live_capture_host_owned: errors.append("sandbox replay capture shared root custody with the agent") if source is None: errors.append("source LLM trajectory manifest is missing or malformed") @@ -264,7 +271,6 @@ def refresh_stitched_trajectory_manifest( n_live=n_live, live_attempt_count=live_attempt_count, live_capture_complete=live_capture_complete, - live_capture_trusted=live_capture_trusted, rows_valid=rows_valid, ) manifest = LLMTrajectoryManifest( @@ -298,7 +304,6 @@ def _continuation_role_captures( n_live: int, live_attempt_count: int, live_capture_complete: bool, - live_capture_trusted: bool, rows_valid: bool, ) -> list[LLMRoleCapture]: """Preserve source-role provenance and append a distinct live leg.""" @@ -331,14 +336,10 @@ def _continuation_role_captures( agent="openhands", model=live_model, auth_mode=AuthMode.API_KEY, - capture_source=CaptureSource.LITELLM_PROXY, - capture_fidelity=( - CaptureFidelity.PROVIDER_WIRE - if live_capture_trusted - else CaptureFidelity.AGENT_SESSION - ), + capture_source=CaptureSource.REPLAY_PROXY, + capture_fidelity=CaptureFidelity.AGENT_SESSION, exchange_count=n_live, - request_complete=live_complete, + request_complete=False, response_complete=live_complete, ) ) diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 5e266a8b0..144e64760 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -13,6 +13,7 @@ LLMRequest, LLMResponse, Trajectory, + canonical_redaction_source, ) from benchflow.usage_tracking import usage_unavailable @@ -81,7 +82,8 @@ def callback_module_source() -> str: """Return the Python module written next to LiteLLM config.yaml.""" - return r""" + return ( + r""" from __future__ import annotations import json @@ -95,6 +97,9 @@ def callback_module_source() -> str: import litellm from litellm.integrations.custom_logger import CustomLogger +""" + + canonical_redaction_source() + + r""" _skill_catalog_gate_passed = False @@ -183,80 +188,6 @@ def _jsonable(value: Any) -> Any: return str(value) -_SENSITIVE_FIELD_SUFFIXES = ( - "api_key", - "access_key", - "account_key", - "secret_key", - "private_key", - "access_token", - "refresh_token", - "session_token", - "client_secret", - "token", - "secret", - "credential", - "password", - "passwd", - "credentials", -) -_SENSITIVE_FIELD_NAMES = { - "authorization", - "proxy_authorization", - "x_api_key", - "x_goog_api_key", - "api_key", - "apikey", -} -_URL_CREDENTIAL_RE = re.compile( - r"([?&](?:" - r"api_?key|access_key|access_token|refresh_token|session_token|" - r"session_id|sessionid|client_secret|aws_secret_access_key|account_key|" - r"secret|password_hash|password|passwd|" - r"signature|x-amz-signature|x-amz-credential|x-amz-security-token|sas" - r")=)[^&\s\"'*]+", - re.IGNORECASE, -) -_AZURE_SAS_RE = re.compile(r"([?&]sig=)[^&\s\"'*]{16,}", re.IGNORECASE) -_AUTHORIZATION_RE = re.compile( - r"\b(Bearer|Token|Basic)\s+(?!\*\*\*REDACTED\*\*\*)[^\s,;\"']+", - re.IGNORECASE, -) - - -def _is_sensitive_field(key: Any) -> bool: - if not isinstance(key, str): - return False - normalized = re.sub(r"[^a-z0-9]+", "_", key.lower()).strip("_") - return normalized in _SENSITIVE_FIELD_NAMES or normalized.endswith( - _SENSITIVE_FIELD_SUFFIXES - ) - - -def _redact_storage_text(value: str) -> str: - value = _URL_CREDENTIAL_RE.sub(r"\1***REDACTED***", value) - value = _AZURE_SAS_RE.sub(r"\1***REDACTED***", value) - return _AUTHORIZATION_RE.sub(r"\1 ***REDACTED***", value) - - -def _redact_for_storage(value: Any) -> Any: - # Redact callback payloads before the durable append-only write. - if isinstance(value, dict): - return { - str(key): ( - "***REDACTED***" - if _is_sensitive_field(key) - else _redact_for_storage(item) - ) - for key, item in value.items() - } - if isinstance(value, (list, tuple)): - return [_redact_for_storage(item) for item in value] - if isinstance(value, str): - return _redact_storage_text(value) - return value - - def _iso(value: Any) -> str: if isinstance(value, datetime): return value.isoformat() @@ -361,7 +292,7 @@ def _write(self, payload: dict[str, Any]) -> None: if not path: return payload["logged_at"] = datetime.now(timezone.utc).isoformat() - durable_payload = _redact_for_storage(_jsonable(payload)) + durable_payload = redact_trajectory_obj(_jsonable(payload)) with self._capture_lock: os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "a", encoding="utf-8") as handle: @@ -401,7 +332,7 @@ def log_pre_api_call(self, model, messages, kwargs) -> None: if call_id is None or not isinstance(jsonable_body, dict): return with self._capture_lock: - self._provider_requests[call_id] = _redact_for_storage(jsonable_body) + self._provider_requests[call_id] = redact_trajectory_obj(jsonable_body) def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) -> dict[str, Any]: litellm_params = kwargs.get("litellm_params") or {} @@ -618,6 +549,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti proxy_handler_instance = BenchFlowLiteLLMLogger() """ + ) def _parse_time(value: Any) -> datetime: diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 891cb25ab..3dfbe1e0c 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -40,6 +40,7 @@ class CaptureFidelity(StrEnum): class CaptureSource(StrEnum): LITELLM_PROXY = "litellm_proxy" + REPLAY_PROXY = "replay_proxy" CLAUDE_OTEL_RAW_BODY = "claude_otel_raw_body" CLAUDE_NATIVE_SESSION = "claude_native_session" CODEX_NATIVE_SESSION = "codex_native_session" diff --git a/src/benchflow/trajectories/types.py b/src/benchflow/trajectories/types.py index 219b8ebcf..c6d5b3f6d 100644 --- a/src/benchflow/trajectories/types.py +++ b/src/benchflow/trajectories/types.py @@ -557,6 +557,39 @@ def _redact_field(key: Any, value: Any) -> Any: return redact_trajectory_text(value) +def canonical_redaction_source() -> str: + """Render the canonical JSON-object redactor as a standalone source fragment. + + LiteLLM callbacks also run in isolated sandbox virtualenvs where BenchFlow is + not installed. Generate their stdlib-only redactor from these exact pattern + objects and function definitions instead of maintaining a second, inevitably + narrower implementation beside the callback. + """ + import inspect + + pattern_specs = [ + (pattern.pattern, pattern.flags, replacement, category) + for pattern, replacement, category in _REDACTION_PATTERNS + ] + definitions = ( + redact_trajectory_text, + redact_trajectory_text_with_count, + redact_trajectory_text_with_categories, + redact_trajectory_obj, + _redact_field, + ) + return "\n".join( + [ + f"_REDACTION_PATTERN_SPECS = {pattern_specs!r}", + "_REDACTION_PATTERNS = [", + " (re.compile(pattern, flags), replacement, category)", + " for pattern, flags, replacement, category in _REDACTION_PATTERN_SPECS", + "]", + *(inspect.getsource(definition) for definition in definitions), + ] + ) + + def redact_acp_trajectory_jsonl(trajectory: list[dict[str, Any]]) -> str: """Serialize an ACP trajectory list to redacted JSONL. diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 0dd4ab5ef..ab0944876 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -176,8 +176,8 @@ def test_write_stitched_trajectory_creates_file(tmp_path): assert len(out.read_text().strip().splitlines()) == 2 -def test_refresh_stitched_manifest_replaces_pre_stitch_finalization(tmp_path): - """Guards PR #1057 against stale continuation capture manifests.""" +def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_path): + """Guards PR #1057 against promoting continuation ingress to provider wire.""" model = "openai/gpt-5.5" source = write_run_folder( @@ -227,14 +227,21 @@ def test_refresh_stitched_manifest_replaces_pre_stitch_finalization(tmp_path): live_errors=[], ) - assert manifest.status is CaptureStatus.COMPLETE - assert manifest.capture_fidelity is CaptureFidelity.PROVIDER_WIRE + assert manifest.status is CaptureStatus.PARTIAL + assert manifest.capture_source is CaptureSource.MIXED + assert manifest.capture_fidelity is CaptureFidelity.MIXED assert manifest.exchange_count == 2 - assert capture_manifest_allows_training( + assert manifest.request_complete is False + assert "live_provider_request" in manifest.missing_fields + assert not capture_manifest_allows_training( manifest.model_dump(mode="json"), exchange_count=2 ) live_row = json.loads(out.read_text().splitlines()[-1]) - assert live_row["metadata"]["capture_fidelity"] == "provider_wire" + assert live_row["metadata"]["capture_source"] == "replay_proxy" + assert live_row["metadata"]["capture_fidelity"] == "agent_session" + assert live_row["metadata"]["request_complete"] is False + assert live_row["metadata"]["response_complete"] is True + assert live_row["metadata"]["request_capture_source"] == "replay_proxy_ingress" assert live_row["metadata"]["model"] == model assert live_row["metadata"]["schema_version"] == 2 @@ -291,7 +298,7 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) assert manifest.status is CaptureStatus.PARTIAL assert manifest.capture_source is CaptureSource.MIXED - assert manifest.capture_fidelity is CaptureFidelity.MIXED + assert manifest.capture_fidelity is CaptureFidelity.AGENT_SESSION assert manifest.auth_mode is AuthMode.MIXED assert [capture.leg for capture in manifest.role_captures] == [ "recorded", @@ -305,7 +312,10 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) assert live.agent == "openhands" assert live.model == "openai/gpt-5.5" assert live.auth_mode is AuthMode.API_KEY - assert live.capture_fidelity is CaptureFidelity.PROVIDER_WIRE + assert live.capture_source is CaptureSource.REPLAY_PROXY + assert live.capture_fidelity is CaptureFidelity.AGENT_SESSION + assert live.request_complete is False + assert live.response_complete is True assert not capture_manifest_allows_training( manifest.model_dump(mode="json"), exchange_count=2 ) @@ -349,7 +359,7 @@ def test_root_sandbox_live_suffix_is_retained_but_audit_only(tmp_path): source / "trajectory" / "llm_trajectory.jsonl", [exchange(completion(content="live"))], live_model=model, - live_capture_trusted=False, + live_capture_host_owned=False, ) manifest = refresh_stitched_trajectory_manifest( rollout, @@ -360,7 +370,7 @@ def test_root_sandbox_live_suffix_is_retained_but_audit_only(tmp_path): n_live=1, live_attempt_count=1, live_errors=[], - live_capture_trusted=False, + live_capture_host_owned=False, ) live_row = json.loads(out.read_text().splitlines()[-1]) diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 5c1cf195f..1affd9244 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -236,7 +236,15 @@ def test_callback_redacts_secrets_before_durable_journal( monkeypatch.setenv("BENCHFLOW_LITELLM_LOG_PATH", str(log_path)) api_key = "provider-key-without-a-recognizable-prefix" access_token = "oauth-token-without-a-recognizable-prefix" + anthropic_key = "sk-ant-api03-" + "A" * 40 + google_key = "AIzaSy" + "B" * 33 sas = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG" + pem = ( + "-----BEGIN PRIVATE KEY-----\n" + "PEMsecretMaterial1234567890\n" + "-----END PRIVATE KEY-----" + ) + userinfo_password = "userinfo-password-without-prefix" logger._write( { @@ -247,6 +255,11 @@ def test_callback_redacts_secrets_before_durable_journal( "nested": {"access_token": access_token}, "download_url": f"https://blob.invalid/x?sig={sas}&sp=r", "note": f"Authorization: Bearer {access_token}", + "provider_error": ( + f"anthropic={anthropic_key} google={google_key} " + f"api_key={api_key} private={pem} " + f"url=https://user:{userinfo_password}@db.invalid/path" + ), } }, } @@ -255,7 +268,11 @@ def test_callback_redacts_secrets_before_durable_journal( raw = log_path.read_text() assert api_key not in raw assert access_token not in raw + assert anthropic_key not in raw + assert google_key not in raw assert sas not in raw + assert "PEMsecretMaterial1234567890" not in raw + assert userinfo_password not in raw row = json.loads(raw) body = row["request"]["body"] assert body["api_key"] == "***REDACTED***" @@ -263,6 +280,32 @@ def test_callback_redacts_secrets_before_durable_journal( assert "sig=***REDACTED***" in body["download_url"] +def test_callback_uses_exact_canonical_object_redactor() -> None: + """Guards PR #1057 against callback redaction drifting from trajectories.""" + + from benchflow.trajectories.types import redact_trajectory_obj + + callback_redactor = _callback_namespace()["redact_trajectory_obj"] + payload = { + "api_key": "prefixlessSecretValue1234567890", + "message": ( + "sk-ant-api03-" + "C" * 40 + " " + "https://admin:password@db.invalid/x?access_token=token1234567890" + ), + "nested": [ + { + "private_key": ( + "-----BEGIN PRIVATE KEY-----\n" + "CANONICALpemMaterial12345\n" + "-----END PRIVATE KEY-----" + ) + } + ], + } + + assert callback_redactor(payload) == redact_trajectory_obj(payload) + + @pytest.mark.asyncio async def test_callback_pre_call_hook_strips_chat_input_compat_field(): namespace: dict[str, object] = {} From f0746ab1ac5e465df0ee3738ead156ee4578c908 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 08:20:16 -0700 Subject: [PATCH 36/74] fix: preserve audit continuation results --- src/benchflow/acp/container_transport.py | 2 +- src/benchflow/providers/litellm_logging.py | 16 +- src/benchflow/providers/litellm_runtime.py | 16 + src/benchflow/publish/redact.py | 2 +- src/benchflow/sandbox/process/_base.py | 2 +- src/benchflow/trajectories/_export_common.py | 2 +- .../trajectories/llm_capture_manifest.py | 26 ++ .../trajectories/native_capture_collection.py | 2 +- src/benchflow/trajectories/redaction.py | 383 ++++++++++++++++ src/benchflow/trajectories/results.py | 5 + src/benchflow/trajectories/types.py | 430 +----------------- tests/continue_run/test_orchestrator.py | 45 +- tests/test_litellm_hardening.py | 37 ++ tests/test_litellm_logging.py | 1 + .../test_llm_capture_training_contract.py | 55 +++ tests/trajectories/test_redaction.py | 10 +- 16 files changed, 605 insertions(+), 429 deletions(-) create mode 100644 src/benchflow/trajectories/redaction.py diff --git a/src/benchflow/acp/container_transport.py b/src/benchflow/acp/container_transport.py index f5a784c2d..4bbeeb9d3 100644 --- a/src/benchflow/acp/container_transport.py +++ b/src/benchflow/acp/container_transport.py @@ -7,7 +7,7 @@ from benchflow.sandbox.process import LiveProcess from benchflow.sandbox.process._base import _ANSI_CSI_RE, _ANSI_OSC_RE -from benchflow.trajectories.types import redact_trajectory_text +from benchflow.trajectories.redaction import redact_trajectory_text from .transport import Transport, decode_json_rpc_message diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 144e64760..7486a9f58 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -13,7 +13,6 @@ LLMRequest, LLMResponse, Trajectory, - canonical_redaction_source, ) from benchflow.usage_tracking import usage_unavailable @@ -82,8 +81,7 @@ def callback_module_source() -> str: """Return the Python module written next to LiteLLM config.yaml.""" - return ( - r""" + return r""" from __future__ import annotations import json @@ -97,9 +95,14 @@ def callback_module_source() -> str: import litellm from litellm.integrations.custom_logger import CustomLogger -""" - + canonical_redaction_source() - + r""" + +try: + from benchflow_trajectory_redaction import redact_trajectory_obj +except ModuleNotFoundError: + # Direct source execution in BenchFlow tests/development. Production proxy + # runtimes always receive the packaged standalone module beside this file; + # isolated sandboxes fail loudly if that deployment artifact is absent. + from benchflow.trajectories.redaction import redact_trajectory_obj _skill_catalog_gate_passed = False @@ -549,7 +552,6 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti proxy_handler_instance = BenchFlowLiteLLMLogger() """ - ) def _parse_time(value: Any) -> datetime: diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 89bf13bc2..d76490f82 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -18,6 +18,7 @@ import tempfile from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass +from importlib.resources import files from pathlib import Path from typing import TYPE_CHECKING, Any, NoReturn, cast from uuid import uuid4 @@ -69,6 +70,7 @@ LITELLM_VERSION_SPEC = "litellm[proxy]==1.89.0" LITELLM_SANDBOX_ROOT = "/tmp/benchflow-litellm" _CALLBACK_MODULE = "benchflow_litellm_callback" +_REDACTION_MODULE = "benchflow_trajectory_redaction" _LITELLM_REQUESTED_MODEL_ENV = "BENCHFLOW_LITELLM_REQUESTED_MODEL" _LITELLM_AGENT_ENV = "BENCHFLOW_LITELLM_AGENT" _LITELLM_ROLE_ENV = "BENCHFLOW_LITELLM_ROLE" @@ -823,10 +825,12 @@ def _write_runtime_files( ) -> tuple[Path, Path, Path]: runtime_dir.mkdir(parents=True, exist_ok=True) callback_path = runtime_dir / f"{_CALLBACK_MODULE}.py" + redaction_path = runtime_dir / f"{_REDACTION_MODULE}.py" patch_path = runtime_dir / f"{_PATCH_MODULE}.py" sitecustomize_path = runtime_dir / "sitecustomize.py" config_path = runtime_dir / "config.yaml" callback_path.write_text(callback_module_source()) + redaction_path.write_text(_redaction_module_source()) patch_source = Path(__file__).with_name("litellm_bedrock_patch.py").read_text() patch_path.write_text(patch_source) sitecustomize_path.write_text(f"import {_PATCH_MODULE}\n") @@ -834,6 +838,16 @@ def _write_runtime_files( return config_path, callback_path, patch_path +def _redaction_module_source() -> str: + """Read the packaged canonical redactor for isolated proxy runtimes.""" + + return ( + files("benchflow.trajectories") + .joinpath("redaction.py") + .read_text(encoding="utf-8") + ) + + # How long to wait for the *host* per-run LiteLLM proxy to become healthy. # litellm's cold start runs tens of seconds, and when many runs launch in parallel # (max-parallel sweeps) the proxies cold-start simultaneously and contend for CPU, @@ -1050,6 +1064,7 @@ async def _upload_runtime_files_to_sandbox( paths = { "config": f"{runtime_dir}/config.yaml", "callback": f"{runtime_dir}/{_CALLBACK_MODULE}.py", + "redaction": f"{runtime_dir}/{_REDACTION_MODULE}.py", "patch": f"{runtime_dir}/{_PATCH_MODULE}.py", "sitecustomize": f"{runtime_dir}/sitecustomize.py", "launcher": f"{runtime_dir}/launcher.py", @@ -1070,6 +1085,7 @@ async def _upload_runtime_files_to_sandbox( sandbox, yaml.safe_dump(config, sort_keys=False), paths["config"], ".yaml" ) await _upload_text(sandbox, callback_module_source(), paths["callback"], ".py") + await _upload_text(sandbox, _redaction_module_source(), paths["redaction"], ".py") await _upload_text( sandbox, Path(__file__).with_name("litellm_bedrock_patch.py").read_text(), diff --git a/src/benchflow/publish/redact.py b/src/benchflow/publish/redact.py index 530ef1f2c..f9495c34c 100644 --- a/src/benchflow/publish/redact.py +++ b/src/benchflow/publish/redact.py @@ -8,7 +8,7 @@ from collections.abc import Mapping, Sequence from typing import Any, NamedTuple -from benchflow.trajectories.types import ( +from benchflow.trajectories.redaction import ( REDACTION_CATEGORY_API_KEY, REDACTION_CATEGORY_BEARER_TOKEN, REDACTION_CATEGORY_CREDENTIAL_FIELD, diff --git a/src/benchflow/sandbox/process/_base.py b/src/benchflow/sandbox/process/_base.py index 974102213..b4a05e910 100644 --- a/src/benchflow/sandbox/process/_base.py +++ b/src/benchflow/sandbox/process/_base.py @@ -229,7 +229,7 @@ async def readline(self) -> bytes: msg = f"Process closed stdout (rc={rc}): {hint}" stderr_snippet: str | None = None if stderr_text: - from benchflow.trajectories.types import redact_trajectory_text + from benchflow.trajectories.redaction import redact_trajectory_text stderr_snippet = redact_trajectory_text(stderr_text)[:_DIAG_TRUNCATE] msg += f"\nstderr: {stderr_snippet}" diff --git a/src/benchflow/trajectories/_export_common.py b/src/benchflow/trajectories/_export_common.py index ed6687c1e..88828fe8d 100644 --- a/src/benchflow/trajectories/_export_common.py +++ b/src/benchflow/trajectories/_export_common.py @@ -22,7 +22,7 @@ from pathlib import Path from typing import Any -from benchflow.trajectories.types import redact_trajectory_text +from benchflow.trajectories.redaction import redact_trajectory_text logger = logging.getLogger(__name__) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 3dfbe1e0c..3d36772fd 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -275,6 +275,32 @@ def capture_manifest_has_oauth_role_capture(manifest: dict[str, Any]) -> bool: return False +def capture_manifest_has_replay_capture(manifest: dict[str, Any]) -> bool: + """Return whether a manifest contains an audit-only continuation suffix.""" + + if ( + manifest.get("capture_source") == CaptureSource.REPLAY_PROXY.value + and isinstance(manifest.get("exchange_count"), int) + and not isinstance(manifest.get("exchange_count"), bool) + and manifest["exchange_count"] > 0 + ): + return True + role_captures = manifest.get("role_captures") + if not isinstance(role_captures, list): + return False + for value in role_captures: + try: + role_capture = LLMRoleCapture.model_validate(value) + except ValidationError: + continue + if ( + role_capture.capture_source is CaptureSource.REPLAY_PROXY + and role_capture.exchange_count > 0 + ): + return True + return False + + def _atomic_write_text(path: Path, payload: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") diff --git a/src/benchflow/trajectories/native_capture_collection.py b/src/benchflow/trajectories/native_capture_collection.py index 7c8db0182..dedd0d6ed 100644 --- a/src/benchflow/trajectories/native_capture_collection.py +++ b/src/benchflow/trajectories/native_capture_collection.py @@ -24,7 +24,7 @@ parse_codex_sessions, retain_uncovered_claude_session_exchanges, ) -from benchflow.trajectories.types import redact_trajectory_text +from benchflow.trajectories.redaction import redact_trajectory_text logger = logging.getLogger(__name__) diff --git a/src/benchflow/trajectories/redaction.py b/src/benchflow/trajectories/redaction.py new file mode 100644 index 000000000..29b27b078 --- /dev/null +++ b/src/benchflow/trajectories/redaction.py @@ -0,0 +1,383 @@ +"""Canonical stdlib-only secret redaction for trajectory and proxy artifacts.""" + +from __future__ import annotations + +import re +from typing import Any + +# Human-facing redaction categories. Every canonical pattern below is tagged +# with the kind of secret its rule actually detects, so the upload preview can +# tell contributors WHAT was masked ("2 API keys, 1 bearer token") instead of +# only how many values. The taxonomy names the real rules — there is no +# PII/email detection, so no such category exists. +REDACTION_CATEGORY_API_KEY = "API key" +REDACTION_CATEGORY_BEARER_TOKEN = "bearer token" +REDACTION_CATEGORY_PRIVATE_KEY = "private key block" +REDACTION_CATEGORY_PASSWORD = "password" +REDACTION_CATEGORY_URL_CREDENTIAL = "URL credential" +REDACTION_CATEGORY_CREDENTIAL_FIELD = "credential-bearing field value" + +# Quote atom for the header/key-value carriers below: a run of 0-8 optional +# backslashes followed by an optional quote. This makes every carrier fire on raw +# text (``"k": "v"``), Python dict-repr (``'k': 'v'``), AND json.dumps-escaped +# JSON at any nesting depth (``\"k\": \"v\"``, ``\\\"k\\\": …``) — the escaped +# form still appears when a serialized artifact is redacted as text (the +# ADP/ATIF exporters, ``results.py``), or when a provider error body embedded in +# ``error.message`` is itself JSON (#830). (``Trajectory.to_jsonl`` and the ACP +# writer now redact record *values* before ``json.dumps`` via +# ``redact_trajectory_obj`` so they can never emit a corrupt escape.) The +# ``{0,8}`` cap keeps it ReDoS-bounded. +_ESCQ = r'(?:\\{0,8}["\']|)' +# Secret value class for the carriers: stops at a quote (plain, or the backslash +# of an escaped closing quote), whitespace, and the ``,`` ``}`` ``&`` delimiters +# so a redacted value can never swallow a sibling JSON field or URL query param; +# excludes ``*`` so an already-redacted value isn't re-matched. +_SECVAL = r"[^\"'\s,}\\*&]+" +# Name-prefix class for the env/JSON carriers, LENGTH-CAPPED. An uncapped +# ``[A-Za-z0-9_]*`` before a required marker backtracks O(n²) on a long +# alphanumeric run (e.g. a base64 image field), so a benign trajectory could +# stall the redactor for tens of seconds. Real env-var/header names are short; +# ``{0,64}`` makes each match attempt O(1) → overall O(n) (#830 ReDoS fix). +_NAME = r"[A-Za-z0-9_]{0,64}" + +_REDACTION_PATTERNS: list[tuple[re.Pattern[str], str, str]] = [ + # PEM private-key blocks (GCP/Vertex service-account JSON, RSA keys). Redact + # the whole armored block incl. the base64 body; the carriers below stop at + # the first space and would leave the key material. Lazy + length-capped body + # keeps it bounded (#830). + ( + re.compile( + r"-----BEGIN [A-Z0-9 ]{0,40}PRIVATE KEY-----" + r"[\s\S]{0,8192}?" + r"-----END [A-Z0-9 ]{0,40}PRIVATE KEY-----" + ), + "***REDACTED***", + REDACTION_CATEGORY_PRIVATE_KEY, + ), + # --- Token families: redacted WHOLE (prefix included) so the v0.5 leak audit + # greps (``AIzaSy``/``dtn_``…) see no live-key shape (#537/#585). --- + # Anthropic: sk-ant-api03-... + ( + re.compile(r"sk-ant-[a-zA-Z0-9_-]{12,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # BenchFlow proxy master key (sk-benchflow-), OpenAI service + # account (sk-svcacct-), OpenRouter (sk-or-v1-): rare labels that won't appear + # in a normal kebab slug, so the hyphen-permitting entropy class is safe here. + ( + re.compile(r"sk-(?:benchflow|svcacct|or-v1)-[A-Za-z0-9_-]{12,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # OpenAI org-scoped sk-proj-/sk-admin-: `proj`/`admin` ARE common identifier + # words, so a bare hyphen class would redact kebab slugs like + # `sk-proj-refactor-auth`. Real keys are high-entropy base64url and always + # carry an uppercase char; gate on a lookahead for one so lowercase slugs + # survive (#830). Before generic sk- so it wins. + ( + re.compile( + r"sk-(?:proj|admin)-(?=[A-Za-z0-9_-]{0,200}[A-Z])[A-Za-z0-9_-]{12,}" + ), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # OpenAI / generic sk- (alphanumeric only — widening to include `-` would + # match common slugs like `task-sk-us-east-1-...`) + (re.compile(r"sk-[a-zA-Z0-9]{12,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), + # Google AI / Gemini: AIzaSy... (≥20 char suffix avoids matching `AIzaSy` + # alone). Prefix is redacted too so the audit grep for `AIzaSy` is clean. + ( + re.compile(r"AIzaSy[A-Za-z0-9_-]{20,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # AWS access keys: AKIA/ASIA + exactly 16 chars; length anchor avoids + # matching English words like "ASIAPACIFIC". + ( + re.compile(r"(?:AKIA|ASIA)[A-Z0-9]{16}(?![A-Z0-9])"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # Daytona SDK tokens: dtn_... — ≥16 char suffix avoids short ids (`dtn_v2`). + ( + re.compile(r"dtn_[A-Za-z0-9_]{16,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # GitHub tokens: PATs / OAuth / app / refresh (ghp_/gho_/ghu_/ghs_/ghr_) and + # fine-grained PATs (github_pat_...). ≥20 char suffix avoids short slugs. + ( + re.compile(r"(?:ghp|gho|ghu|ghs|ghr)_[A-Za-z0-9]{20,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + ( + re.compile(r"github_pat_[A-Za-z0-9_]{20,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # Slack tokens: bot/user/app/refresh/legacy (xoxb-/xoxp-/xoxa-/xoxr-/xoxs-). + ( + re.compile(r"xox[baprs]-[A-Za-z0-9-]{10,}"), + "***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # Other provider key families with distinctive prefixes — caught here for the + # bare-token-in-prose form that no NAME=value carrier sees (a provider + # exception that echoes the key). Groq gsk_, xAI xai-, Replicate r8_, + # HuggingFace hf_, Fireworks fw_. ≥20-char anchors avoid short ids (#830). + (re.compile(r"gsk_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), + (re.compile(r"xai-[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), + (re.compile(r"r8_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), + (re.compile(r"hf_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), + (re.compile(r"fw_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), + # JSON Web Tokens (session/bearer creds): three base64url segments split by + # dots. The ``eyJ`` prefix is base64url of ``{"`` — a JWT header. The 3rd + # (signature) segment requires ≥20 chars (real HS256/RS256 sigs are 43+) so + # short dotted base64/method-chains that merely start ``eyJ`` aren't redacted; + # a left boundary keeps it from starting mid-identifier. Segment upper bounds + # keep it ReDoS-bounded (#830). + ( + re.compile( + r"(?` — bare `sig` is excluded from the curated list above + # (a `?sig=v2` scheme-version flag is a common non-secret), so gate on a + # high-entropy value length: a real SAS signature is a 40+ char base64 HMAC, + # while a version flag is short (#830). Excludes `*` like the curated params + # above so the 14-char `***REDACTED***` marker can never re-match (#1008). + ( + re.compile(r"([?&]sig=)[^&\s\"'*]{16,}", re.IGNORECASE), + r"\1***REDACTED***", + REDACTION_CATEGORY_URL_CREDENTIAL, + ), + # --- Header / key-value secret carriers --- + # Match `name: value` / `name=value` in raw, JSON, Python dict-repr + # (single-quoted), AND json.dumps-escaped (`\"name\": \"value\"`) forms — see + # _ESCQ. The name is kept; only the value is dropped (#585/#830). No leading + # boundary on the hyphenated header names — the hyphen keeps them from matching + # inside underscore variable names. + ( + re.compile( + rf"((?:{_ESCQ}(?:x-api-key|x-goog-api-key|api-key){_ESCQ})\s*[:=]\s*{_ESCQ})" + rf"{_SECVAL}", + re.IGNORECASE, + ), + r"\1***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # Underscore form `api_key`. No leading boundary: namespaced env dumps like + # `GEMINI_API_KEY=secret` / JSON keys such as `"openai_api_key"` must redact + # too (#585). The `\1` capture preserves the matched name+separator. + ( + re.compile( + rf"({_ESCQ}api_key{_ESCQ}\s*[:=]\s*{_ESCQ}){_SECVAL}", re.IGNORECASE + ), + r"\1***REDACTED***", + REDACTION_CATEGORY_API_KEY, + ), + # `master_key`/`master-key` and `private_key`/`private-key` carriers — the + # BenchFlow proxy master key and GCP/Vertex service-account private-key field + # as labelled values (the PEM body itself is scrubbed by the block rule above). + ( + re.compile( + rf"({_ESCQ}(?:master|private)[_-]key{_ESCQ}\s*[:=]\s*{_ESCQ}){_SECVAL}", + re.IGNORECASE, + ), + r"\1***REDACTED***", + REDACTION_CATEGORY_CREDENTIAL_FIELD, + ), + # Bearer-token env vars whose NAME has BEARER_TOKEN in the MIDDLE, e.g. + # `AWS_BEARER_TOKEN_BEDROCK=` — the secret-suffix rule below anchors its marker + # at the name end, so a TOKEN in the middle slips through. `BEARER_TOKEN` is an + # unambiguous secret signal, safe as a substring (#830). Names length-capped. + ( + re.compile( + rf"(?` isn't double-redacted into `***REDACTED*** ***...`. + ( + re.compile( + rf"(? str: + """Apply all secret-redaction patterns to *text*. + + Token families (Anthropic/OpenAI/Google/AWS/Daytona/GitHub/Slack) are + redacted whole, prefix included, so the secret-leak audit greps see no + live-key shape. Header/key-value carriers keep the field name but drop the + value, in both JSON (``"x-api-key": "v"``) and raw-text (``x-api-key: v``) + forms; this includes generic ``*TOKEN*``/``*SECRET*`` carriers so a + ``GITHUB_TOKEN=...`` env dump is scrubbed even without a known prefix. + """ + return redact_trajectory_text_with_count(text)[0] + + +def redact_trajectory_text_with_count(text: str) -> tuple[str, int]: + """Apply the canonical patterns and report how many matches were replaced.""" + text, categories = redact_trajectory_text_with_categories(text) + return text, sum(categories.values()) + + +def redact_trajectory_text_with_categories(text: str) -> tuple[str, dict[str, int]]: + """Apply the canonical patterns and report replacements per category. + + Categories name the kind of secret each rule actually detects (API key, + bearer token, private key block, URL credential, credential-bearing field + value) so upload previews can itemize what was masked without inventing + detection that does not exist. The total replacement count is the sum of + the returned values. + """ + categories: dict[str, int] = {} + for pattern, replacement, category in _REDACTION_PATTERNS: + text, count = pattern.subn(replacement, text) + if count: + categories[category] = categories.get(category, 0) + count + return text, categories + + +def redact_trajectory_obj(obj: Any) -> Any: + """Recursively redact secrets in the string leaves of a JSON-serializable value. + + Prefer this over running :func:`redact_trajectory_text` on an + already-``json.dumps``-ed record. Redacting the serialized *text* can split a + ``\\`` escape — a secret sitting next to one leaves a lone ``\\X`` — which + corrupts the JSON so the trajectory file no longer parses and + ``bench train convert`` fails. Redacting the *values* before serializing keeps + escaping valid while preserving coverage: free-text leaves are scrubbed by the + token-family and inline ``name: value`` carriers, and structured fields (e.g. + a ``{"x-api-key": ""}`` header) are scrubbed in their key context + by :func:`_redact_field` so the carriers that key off a field name still fire. + """ + if isinstance(obj, dict): + return {key: _redact_field(key, value) for key, value in obj.items()} + if isinstance(obj, list): + return [redact_trajectory_obj(item) for item in obj] + if isinstance(obj, str): + return redact_trajectory_text(obj) + return obj + + +def _redact_field(key: Any, value: Any) -> Any: + """Redact a dict field's *value*, keeping the field name as carrier context. + + The ``name: value`` carriers (``x-api-key``, ``api_key``, ``authorization``, + ``*_TOKEN`` …) strip *prefixless* secret values only when they can see the + field name; in a parsed object the name and value are separate. So redact the + value inside a ``{"": ""}`` probe and read it back. The probe + round-trips through json so escaping stays valid; if redacting the serialized + probe would corrupt it (the escape hazard this module guards), fall back to + redacting the bare value, which still catches the prefixed token families. + """ + if isinstance(value, (dict, list)): + return redact_trajectory_obj(value) + if not isinstance(value, str): + return value + if isinstance(key, str): + import json + + probe = json.dumps({key: value}) + redacted = redact_trajectory_text(probe) + if redacted == probe: + return value + try: + restored = json.loads(redacted) + except ValueError: + pass + else: + if isinstance(restored, dict) and isinstance(restored.get(key), str): + return restored[key] + return redact_trajectory_text(value) diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index 9d6de8f22..b980ddabb 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -35,6 +35,7 @@ from benchflow.trajectories.llm_capture_manifest import ( capture_artifact_allows_training, capture_manifest_has_oauth_role_capture, + capture_manifest_has_replay_capture, read_llm_trajectory_manifest, ) from benchflow.trajectories.types import redact_trajectory_obj @@ -513,6 +514,10 @@ def build_rollout_results_record( capture_manifest is not None and capture_manifest_has_oauth_role_capture(capture_manifest) ) + or ( + capture_manifest is not None + and capture_manifest_has_replay_capture(capture_manifest) + ) ) and ( capture_manifest is None diff --git a/src/benchflow/trajectories/types.py b/src/benchflow/trajectories/types.py index c6d5b3f6d..dc07481d8 100644 --- a/src/benchflow/trajectories/types.py +++ b/src/benchflow/trajectories/types.py @@ -1,12 +1,29 @@ """Trajectory types — raw LLM API request/response pairs captured from providers.""" -import re from dataclasses import dataclass from datetime import datetime from typing import Any from pydantic import BaseModel, Field +from benchflow.trajectories import redaction as _redaction + +# Keep the historical ``benchflow.trajectories.types`` redaction imports +# working for downstream callers. The implementation lives in the standalone, +# stdlib-only module so the exact same code can be deployed beside LiteLLM. +REDACTION_CATEGORY_API_KEY = _redaction.REDACTION_CATEGORY_API_KEY +REDACTION_CATEGORY_BEARER_TOKEN = _redaction.REDACTION_CATEGORY_BEARER_TOKEN +REDACTION_CATEGORY_CREDENTIAL_FIELD = _redaction.REDACTION_CATEGORY_CREDENTIAL_FIELD +REDACTION_CATEGORY_PASSWORD = _redaction.REDACTION_CATEGORY_PASSWORD +REDACTION_CATEGORY_PRIVATE_KEY = _redaction.REDACTION_CATEGORY_PRIVATE_KEY +REDACTION_CATEGORY_URL_CREDENTIAL = _redaction.REDACTION_CATEGORY_URL_CREDENTIAL +redact_trajectory_obj = _redaction.redact_trajectory_obj +redact_trajectory_text = _redaction.redact_trajectory_text +redact_trajectory_text_with_categories = ( + _redaction.redact_trajectory_text_with_categories +) +redact_trajectory_text_with_count = _redaction.redact_trajectory_text_with_count + _USAGE_KEYS = { "input_tokens", "output_tokens", @@ -179,417 +196,6 @@ def _exchange_token_usage(exchange: "LLMExchange") -> TokenUsage: ) -# Human-facing redaction categories. Every canonical pattern below is tagged -# with the kind of secret its rule actually detects, so the upload preview can -# tell contributors WHAT was masked ("2 API keys, 1 bearer token") instead of -# only how many values. The taxonomy names the real rules — there is no -# PII/email detection, so no such category exists. -REDACTION_CATEGORY_API_KEY = "API key" -REDACTION_CATEGORY_BEARER_TOKEN = "bearer token" -REDACTION_CATEGORY_PRIVATE_KEY = "private key block" -REDACTION_CATEGORY_PASSWORD = "password" -REDACTION_CATEGORY_URL_CREDENTIAL = "URL credential" -REDACTION_CATEGORY_CREDENTIAL_FIELD = "credential-bearing field value" - -# Quote atom for the header/key-value carriers below: a run of 0-8 optional -# backslashes followed by an optional quote. This makes every carrier fire on raw -# text (``"k": "v"``), Python dict-repr (``'k': 'v'``), AND json.dumps-escaped -# JSON at any nesting depth (``\"k\": \"v\"``, ``\\\"k\\\": …``) — the escaped -# form still appears when a serialized artifact is redacted as text (the -# ADP/ATIF exporters, ``results.py``), or when a provider error body embedded in -# ``error.message`` is itself JSON (#830). (``Trajectory.to_jsonl`` and the ACP -# writer now redact record *values* before ``json.dumps`` via -# ``redact_trajectory_obj`` so they can never emit a corrupt escape.) The -# ``{0,8}`` cap keeps it ReDoS-bounded. -_ESCQ = r'(?:\\{0,8}["\']|)' -# Secret value class for the carriers: stops at a quote (plain, or the backslash -# of an escaped closing quote), whitespace, and the ``,`` ``}`` ``&`` delimiters -# so a redacted value can never swallow a sibling JSON field or URL query param; -# excludes ``*`` so an already-redacted value isn't re-matched. -_SECVAL = r"[^\"'\s,}\\*&]+" -# Name-prefix class for the env/JSON carriers, LENGTH-CAPPED. An uncapped -# ``[A-Za-z0-9_]*`` before a required marker backtracks O(n²) on a long -# alphanumeric run (e.g. a base64 image field), so a benign trajectory could -# stall the redactor for tens of seconds. Real env-var/header names are short; -# ``{0,64}`` makes each match attempt O(1) → overall O(n) (#830 ReDoS fix). -_NAME = r"[A-Za-z0-9_]{0,64}" - -_REDACTION_PATTERNS: list[tuple[re.Pattern[str], str, str]] = [ - # PEM private-key blocks (GCP/Vertex service-account JSON, RSA keys). Redact - # the whole armored block incl. the base64 body; the carriers below stop at - # the first space and would leave the key material. Lazy + length-capped body - # keeps it bounded (#830). - ( - re.compile( - r"-----BEGIN [A-Z0-9 ]{0,40}PRIVATE KEY-----" - r"[\s\S]{0,8192}?" - r"-----END [A-Z0-9 ]{0,40}PRIVATE KEY-----" - ), - "***REDACTED***", - REDACTION_CATEGORY_PRIVATE_KEY, - ), - # --- Token families: redacted WHOLE (prefix included) so the v0.5 leak audit - # greps (``AIzaSy``/``dtn_``…) see no live-key shape (#537/#585). --- - # Anthropic: sk-ant-api03-... - ( - re.compile(r"sk-ant-[a-zA-Z0-9_-]{12,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # BenchFlow proxy master key (sk-benchflow-), OpenAI service - # account (sk-svcacct-), OpenRouter (sk-or-v1-): rare labels that won't appear - # in a normal kebab slug, so the hyphen-permitting entropy class is safe here. - ( - re.compile(r"sk-(?:benchflow|svcacct|or-v1)-[A-Za-z0-9_-]{12,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # OpenAI org-scoped sk-proj-/sk-admin-: `proj`/`admin` ARE common identifier - # words, so a bare hyphen class would redact kebab slugs like - # `sk-proj-refactor-auth`. Real keys are high-entropy base64url and always - # carry an uppercase char; gate on a lookahead for one so lowercase slugs - # survive (#830). Before generic sk- so it wins. - ( - re.compile( - r"sk-(?:proj|admin)-(?=[A-Za-z0-9_-]{0,200}[A-Z])[A-Za-z0-9_-]{12,}" - ), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # OpenAI / generic sk- (alphanumeric only — widening to include `-` would - # match common slugs like `task-sk-us-east-1-...`) - (re.compile(r"sk-[a-zA-Z0-9]{12,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), - # Google AI / Gemini: AIzaSy... (≥20 char suffix avoids matching `AIzaSy` - # alone). Prefix is redacted too so the audit grep for `AIzaSy` is clean. - ( - re.compile(r"AIzaSy[A-Za-z0-9_-]{20,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # AWS access keys: AKIA/ASIA + exactly 16 chars; length anchor avoids - # matching English words like "ASIAPACIFIC". - ( - re.compile(r"(?:AKIA|ASIA)[A-Z0-9]{16}(?![A-Z0-9])"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # Daytona SDK tokens: dtn_... — ≥16 char suffix avoids short ids (`dtn_v2`). - ( - re.compile(r"dtn_[A-Za-z0-9_]{16,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # GitHub tokens: PATs / OAuth / app / refresh (ghp_/gho_/ghu_/ghs_/ghr_) and - # fine-grained PATs (github_pat_...). ≥20 char suffix avoids short slugs. - ( - re.compile(r"(?:ghp|gho|ghu|ghs|ghr)_[A-Za-z0-9]{20,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - ( - re.compile(r"github_pat_[A-Za-z0-9_]{20,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # Slack tokens: bot/user/app/refresh/legacy (xoxb-/xoxp-/xoxa-/xoxr-/xoxs-). - ( - re.compile(r"xox[baprs]-[A-Za-z0-9-]{10,}"), - "***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # Other provider key families with distinctive prefixes — caught here for the - # bare-token-in-prose form that no NAME=value carrier sees (a provider - # exception that echoes the key). Groq gsk_, xAI xai-, Replicate r8_, - # HuggingFace hf_, Fireworks fw_. ≥20-char anchors avoid short ids (#830). - (re.compile(r"gsk_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), - (re.compile(r"xai-[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), - (re.compile(r"r8_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), - (re.compile(r"hf_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), - (re.compile(r"fw_[A-Za-z0-9]{20,}"), "***REDACTED***", REDACTION_CATEGORY_API_KEY), - # JSON Web Tokens (session/bearer creds): three base64url segments split by - # dots. The ``eyJ`` prefix is base64url of ``{"`` — a JWT header. The 3rd - # (signature) segment requires ≥20 chars (real HS256/RS256 sigs are 43+) so - # short dotted base64/method-chains that merely start ``eyJ`` aren't redacted; - # a left boundary keeps it from starting mid-identifier. Segment upper bounds - # keep it ReDoS-bounded (#830). - ( - re.compile( - r"(?` — bare `sig` is excluded from the curated list above - # (a `?sig=v2` scheme-version flag is a common non-secret), so gate on a - # high-entropy value length: a real SAS signature is a 40+ char base64 HMAC, - # while a version flag is short (#830). Excludes `*` like the curated params - # above so the 14-char `***REDACTED***` marker can never re-match (#1008). - ( - re.compile(r"([?&]sig=)[^&\s\"'*]{16,}", re.IGNORECASE), - r"\1***REDACTED***", - REDACTION_CATEGORY_URL_CREDENTIAL, - ), - # --- Header / key-value secret carriers --- - # Match `name: value` / `name=value` in raw, JSON, Python dict-repr - # (single-quoted), AND json.dumps-escaped (`\"name\": \"value\"`) forms — see - # _ESCQ. The name is kept; only the value is dropped (#585/#830). No leading - # boundary on the hyphenated header names — the hyphen keeps them from matching - # inside underscore variable names. - ( - re.compile( - rf"((?:{_ESCQ}(?:x-api-key|x-goog-api-key|api-key){_ESCQ})\s*[:=]\s*{_ESCQ})" - rf"{_SECVAL}", - re.IGNORECASE, - ), - r"\1***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # Underscore form `api_key`. No leading boundary: namespaced env dumps like - # `GEMINI_API_KEY=secret` / JSON keys such as `"openai_api_key"` must redact - # too (#585). The `\1` capture preserves the matched name+separator. - ( - re.compile( - rf"({_ESCQ}api_key{_ESCQ}\s*[:=]\s*{_ESCQ}){_SECVAL}", re.IGNORECASE - ), - r"\1***REDACTED***", - REDACTION_CATEGORY_API_KEY, - ), - # `master_key`/`master-key` and `private_key`/`private-key` carriers — the - # BenchFlow proxy master key and GCP/Vertex service-account private-key field - # as labelled values (the PEM body itself is scrubbed by the block rule above). - ( - re.compile( - rf"({_ESCQ}(?:master|private)[_-]key{_ESCQ}\s*[:=]\s*{_ESCQ}){_SECVAL}", - re.IGNORECASE, - ), - r"\1***REDACTED***", - REDACTION_CATEGORY_CREDENTIAL_FIELD, - ), - # Bearer-token env vars whose NAME has BEARER_TOKEN in the MIDDLE, e.g. - # `AWS_BEARER_TOKEN_BEDROCK=` — the secret-suffix rule below anchors its marker - # at the name end, so a TOKEN in the middle slips through. `BEARER_TOKEN` is an - # unambiguous secret signal, safe as a substring (#830). Names length-capped. - ( - re.compile( - rf"(?` isn't double-redacted into `***REDACTED*** ***...`. - ( - re.compile( - rf"(? str: - """Apply all secret-redaction patterns to *text*. - - Token families (Anthropic/OpenAI/Google/AWS/Daytona/GitHub/Slack) are - redacted whole, prefix included, so the secret-leak audit greps see no - live-key shape. Header/key-value carriers keep the field name but drop the - value, in both JSON (``"x-api-key": "v"``) and raw-text (``x-api-key: v``) - forms; this includes generic ``*TOKEN*``/``*SECRET*`` carriers so a - ``GITHUB_TOKEN=...`` env dump is scrubbed even without a known prefix. - """ - return redact_trajectory_text_with_count(text)[0] - - -def redact_trajectory_text_with_count(text: str) -> tuple[str, int]: - """Apply the canonical patterns and report how many matches were replaced.""" - text, categories = redact_trajectory_text_with_categories(text) - return text, sum(categories.values()) - - -def redact_trajectory_text_with_categories(text: str) -> tuple[str, dict[str, int]]: - """Apply the canonical patterns and report replacements per category. - - Categories name the kind of secret each rule actually detects (API key, - bearer token, private key block, URL credential, credential-bearing field - value) so upload previews can itemize what was masked without inventing - detection that does not exist. The total replacement count is the sum of - the returned values. - """ - categories: dict[str, int] = {} - for pattern, replacement, category in _REDACTION_PATTERNS: - text, count = pattern.subn(replacement, text) - if count: - categories[category] = categories.get(category, 0) + count - return text, categories - - -def redact_trajectory_obj(obj: Any) -> Any: - """Recursively redact secrets in the string leaves of a JSON-serializable value. - - Prefer this over running :func:`redact_trajectory_text` on an - already-``json.dumps``-ed record. Redacting the serialized *text* can split a - ``\\`` escape — a secret sitting next to one leaves a lone ``\\X`` — which - corrupts the JSON so the trajectory file no longer parses and - ``bench train convert`` fails. Redacting the *values* before serializing keeps - escaping valid while preserving coverage: free-text leaves are scrubbed by the - token-family and inline ``name: value`` carriers, and structured fields (e.g. - a ``{"x-api-key": ""}`` header) are scrubbed in their key context - by :func:`_redact_field` so the carriers that key off a field name still fire. - """ - if isinstance(obj, dict): - return {key: _redact_field(key, value) for key, value in obj.items()} - if isinstance(obj, list): - return [redact_trajectory_obj(item) for item in obj] - if isinstance(obj, str): - return redact_trajectory_text(obj) - return obj - - -def _redact_field(key: Any, value: Any) -> Any: - """Redact a dict field's *value*, keeping the field name as carrier context. - - The ``name: value`` carriers (``x-api-key``, ``api_key``, ``authorization``, - ``*_TOKEN`` …) strip *prefixless* secret values only when they can see the - field name; in a parsed object the name and value are separate. So redact the - value inside a ``{"": ""}`` probe and read it back. The probe - round-trips through json so escaping stays valid; if redacting the serialized - probe would corrupt it (the escape hazard this module guards), fall back to - redacting the bare value, which still catches the prefixed token families. - """ - if isinstance(value, (dict, list)): - return redact_trajectory_obj(value) - if not isinstance(value, str): - return value - if isinstance(key, str): - import json - - probe = json.dumps({key: value}) - redacted = redact_trajectory_text(probe) - if redacted == probe: - return value - try: - restored = json.loads(redacted) - except ValueError: - pass - else: - if isinstance(restored, dict) and isinstance(restored.get(key), str): - return restored[key] - return redact_trajectory_text(value) - - -def canonical_redaction_source() -> str: - """Render the canonical JSON-object redactor as a standalone source fragment. - - LiteLLM callbacks also run in isolated sandbox virtualenvs where BenchFlow is - not installed. Generate their stdlib-only redactor from these exact pattern - objects and function definitions instead of maintaining a second, inevitably - narrower implementation beside the callback. - """ - import inspect - - pattern_specs = [ - (pattern.pattern, pattern.flags, replacement, category) - for pattern, replacement, category in _REDACTION_PATTERNS - ] - definitions = ( - redact_trajectory_text, - redact_trajectory_text_with_count, - redact_trajectory_text_with_categories, - redact_trajectory_obj, - _redact_field, - ) - return "\n".join( - [ - f"_REDACTION_PATTERN_SPECS = {pattern_specs!r}", - "_REDACTION_PATTERNS = [", - " (re.compile(pattern, flags), replacement, category)", - " for pattern, flags, replacement, category in _REDACTION_PATTERN_SPECS", - "]", - *(inspect.getsource(definition) for definition in definitions), - ] - ) - - def redact_acp_trajectory_jsonl(trajectory: list[dict[str, Any]]) -> str: """Serialize an ACP trajectory list to redacted JSONL. diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index ab0944876..363aea935 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -28,6 +28,7 @@ CaptureFidelity, CaptureSource, CaptureStatus, + LLMRoleCapture, LLMTrajectoryManifest, capture_manifest_allows_training, initialize_llm_trajectory_artifacts, @@ -660,7 +661,7 @@ async def after_cleanup(teardown_errors): def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): - """Guards PR #1057 against retaining the pre-stitch trainer row.""" + """Guards PR #1057 against retaining stale or incomplete continuation rows.""" rollout = tmp_path / "job" / "demo-task__continued" (rollout / "trajectory").mkdir(parents=True) model = "openai/gpt-5.5" @@ -733,6 +734,48 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): assert len(refreshed["trajectory"]) == 1 assert aggregated == refreshed + replay_manifest = manifest.model_copy( + update={ + "status": CaptureStatus.PARTIAL, + "capture_source": CaptureSource.MIXED, + "capture_fidelity": CaptureFidelity.MIXED, + "request_complete": False, + "role_captures": [ + LLMRoleCapture( + role="agent", + leg="live", + agent="openhands", + model=model, + auth_mode=AuthMode.API_KEY, + capture_source=CaptureSource.REPLAY_PROXY, + capture_fidelity=CaptureFidelity.AGENT_SESSION, + exchange_count=1, + request_complete=False, + response_complete=True, + ) + ], + } + ) + write_llm_trajectory_manifest(rollout, replay_manifest) + (rollout / "results.jsonl").write_text( + json.dumps({"info": {"training_ready": True}, "is_completed": False}) + "\n" + ) + + update_continued_metadata( + rollout, + live_model=model, + usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), + environment="docker", + ) + + replay_refreshed = json.loads((rollout / "results.jsonl").read_text()) + assert replay_refreshed["info"]["training_ready"] is False + assert replay_refreshed["info"]["training_ready_reason"] == ( + "insufficient_capture_fidelity" + ) + assert replay_refreshed["is_completed"] is True + assert replay_refreshed["error"] is None + def test_stitching_structurally_redacts_escaped_secret(tmp_path): """Guards PR #1057 against string redaction corrupting stitched JSON.""" diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index f41ad7ace..729cd2f66 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -382,6 +382,13 @@ async def test_sandbox_litellm_launch_keeps_secrets_off_command_line(): # config.yaml uses os.environ/ refs, so the secret is not inlined there either. config_files = [k for k in sandbox.uploaded if k.endswith("config.yaml")] assert config_files and secret not in sandbox.uploaded[config_files[0]] + redaction_files = [ + k for k in sandbox.uploaded if k.endswith("benchflow_trajectory_redaction.py") + ] + assert len(redaction_files) == 1 + assert ( + sandbox.uploaded[redaction_files[0]] == runtime_mod._redaction_module_source() + ) launch_command = next(call for call in sandbox.exec_calls if "launcher.py" in call) assert f"rc=$?; rm -f {launch_files[0]}; exit $rc" in launch_command @@ -632,6 +639,36 @@ def test_bedrock_patch_preflight_passes_when_runtime_files_on_pythonpath(tmp_pat assert result.returncode == 0, result.stdout + result.stderr +def test_runtime_packages_canonical_redactor_verbatim(tmp_path): + """Guards PR #1057 against reflective callback source reconstruction.""" + + import os + import subprocess + import sys + + runtime_mod._write_runtime_files(tmp_path, config={"model_list": []}) + + packaged = tmp_path / "benchflow_trajectory_redaction.py" + assert packaged.read_text() == runtime_mod._redaction_module_source() + assert "def redact_trajectory_obj" in packaged.read_text() + env = dict(os.environ) + env["PYTHONPATH"] = str(tmp_path) + result = subprocess.run( + [ + sys.executable, + "-c", + "import benchflow_litellm_callback as callback; " + "assert callback.redact_trajectory_obj.__module__ == " + "'benchflow_trajectory_redaction'", + ], + env=env, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_bedrock_patch_preflight_fails_closed_when_patch_not_loaded(tmp_path): """THE regression test for issue #602's fail-open (fixed in PR #668): when the patch never loads (sitecustomize missing from PYTHONPATH — the exact silent-failure diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 1affd9244..999d4f187 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -146,6 +146,7 @@ def test_callback_module_source_exposes_proxy_handler_instance(): source = callback_module_source() assert "class BenchFlowLiteLLMLogger" in source + assert "from benchflow_trajectory_redaction import redact_trajectory_obj" in source assert "proxy_handler_instance = BenchFlowLiteLLMLogger()" in source diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index dc01b5ce8..845b2c837 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -270,6 +270,61 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( assert row["error"] is None +def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> None: + """Guards PR #1057 against marking successful continuations incomplete.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="agent_session") + replay_row = json.loads((trajectory_dir / "llm_trajectory.jsonl").read_text()) + replay_row["metadata"].update( + { + "capture_source": "replay_proxy", + "request_capture_source": "replay_proxy_ingress", + "auth_mode": "api_key", + "request_complete": False, + } + ) + (trajectory_dir / "llm_trajectory.jsonl").write_text(json.dumps(replay_row) + "\n") + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "partial", + "capture_source": "mixed", + "capture_fidelity": "mixed", + "auth_mode": "api_key", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + "role_captures": [ + { + "role": "agent", + "leg": "live", + "agent": "openhands", + "model": "openai/gpt-5.6", + "auth_mode": "api_key", + "capture_source": "replay_proxy", + "capture_fidelity": "agent_session", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + } + ], + } + ) + ) + + row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is True + assert row["error"] is None + + @pytest.mark.asyncio async def test_provider_finalization_error_rejects_complete_live_prefix( tmp_path: Path, diff --git a/tests/trajectories/test_redaction.py b/tests/trajectories/test_redaction.py index ca316aa25..4cb196eea 100644 --- a/tests/trajectories/test_redaction.py +++ b/tests/trajectories/test_redaction.py @@ -10,14 +10,16 @@ import pytest +from benchflow.trajectories.redaction import ( + redact_trajectory_text, + redact_trajectory_text_with_count, +) from benchflow.trajectories.types import ( LLMExchange, LLMRequest, LLMResponse, Trajectory, redact_acp_trajectory_jsonl, - redact_trajectory_text, - redact_trajectory_text_with_count, ) # Fake token fixtures assembled from split literals so the full token string @@ -1015,7 +1017,7 @@ def test_redact_trajectory_obj_preserves_key_context_for_structured_secret(): def test_canonical_text_redaction_reports_a_category(text, expected_category): """Guards the redaction-transparency feature from PR #1022: each canonical rule tags its replacements with the kind of secret it actually detects.""" - from benchflow.trajectories.types import redact_trajectory_text_with_categories + from benchflow.trajectories.redaction import redact_trajectory_text_with_categories redacted, categories = redact_trajectory_text_with_categories(text) assert categories == {expected_category: 1} @@ -1025,7 +1027,7 @@ def test_canonical_text_redaction_reports_a_category(text, expected_category): def test_text_categories_sum_to_the_backward_compatible_count(): """Guards the redaction-transparency feature from PR #1022: the categorized scan and the count scan agree, so ``redaction_replacements`` stays backward-compatible.""" - from benchflow.trajectories.types import ( + from benchflow.trajectories.redaction import ( redact_trajectory_text_with_categories, redact_trajectory_text_with_count, ) From 206c13d792deeb2c9bf45f6f009dffc635737fef Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 08:43:25 -0700 Subject: [PATCH 37/74] test: split native capture resilience cases --- .../test_native_capture_resilience.py | 92 ++++++++++++++++++- tests/trajectories/test_native_llm_capture.py | 86 ----------------- 2 files changed, 91 insertions(+), 87 deletions(-) diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index be69c74eb..34bc7600b 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -21,7 +21,12 @@ load_provider_wire_records, ) from benchflow.trajectories.native_capture_collection import NativeCollection -from benchflow.trajectories.native_capture_parsers import parse_codex_sessions +from benchflow.trajectories.native_capture_parsers import ( + parse_codex_sessions, + project_acp_trajectory, +) + +STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) def test_native_parser_normalizes_fallback_and_record_timestamps( @@ -488,3 +493,88 @@ def fail_once(rollout_dir, manifest): assert manifest["capture_source"] == "litellm_proxy" assert manifest["capture_fidelity"] == "provider_wire" assert manifest["exchange_count"] == 1 + + +@pytest.mark.asyncio +async def test_claude_capture_setup_failure_degrades_without_aborting( + tmp_path: Path, +) -> None: + """Guards PR #1057 against observability setup breaking Claude OAuth runs.""" + + commands: list[str] = [] + + class FailingCollectorEnv: + async def exec(self, command, **_kwargs): + commands.append(command) + if "nohup" in command: + return SimpleNamespace( + return_code=1, + stdout="", + stderr="node runtime not found", + ) + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def upload_file(self, *_args, **_kwargs): + return None + + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + prepared = await capture.prepare_agent( + FailingCollectorEnv(), + agent="claude-agent-acp", + model="claude-sonnet-4-6", + agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "oauth-test-token"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + + assert "CLAUDE_CODE_ENABLE_TELEMETRY" not in prepared + assert "OTEL_LOG_RAW_API_BODIES" not in prepared + assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in prepared + assert any("/opt/benchflow/node/bin/node" in command for command in commands) + + +def test_capture_failure_repairs_invalid_jsonl_and_redacts_manifest_error( + tmp_path: Path, +) -> None: + """Guards PR #1057's valid-JSONL and secret-redaction failure invariant.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + capture.trajectory_path.write_text("{not-json\n") + + capture.record_failure( + "authorization: Bearer sk-test-secret-value", model_call_seen=True + ) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert capture.trajectory_path.read_text() == "" + assert manifest["status"] == "capture_failed" + assert "sk-test-secret-value" not in json.dumps(manifest) + + +def test_acp_projection_retains_the_actual_auth_mode() -> None: + """Guards PR #1057 against labeling API-key fallback rows as OAuth.""" + + result = project_acp_trajectory( + [{"type": "agent_message", "text": "finished"}], + agent="codex-acp", + session_id="rollout-1", + started_at=STARTED_AT, + auth_mode="api_key", + ) + + assert result is not None + assert result.trajectory.exchanges[0].metadata["auth_mode"] == "api_key" diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index bd659539d..f937c029c 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -26,7 +26,6 @@ parse_claude_raw_capture, parse_claude_sessions, parse_codex_sessions, - project_acp_trajectory, ) STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) @@ -940,88 +939,3 @@ async def test_same_agent_model_roles_remain_independently_auditable( assert roles["coder"]["exchange_count"] == 0 assert roles["reviewer"]["exchange_count"] == 0 assert roles["mixed"]["exchange_count"] == 1 - - -@pytest.mark.asyncio -async def test_claude_capture_setup_failure_degrades_without_aborting( - tmp_path: Path, -) -> None: - """Guards PR #1057 against observability setup breaking Claude OAuth runs.""" - - commands: list[str] = [] - - class FailingCollectorEnv: - async def exec(self, command, **_kwargs): - commands.append(command) - if "nohup" in command: - return SimpleNamespace( - return_code=1, - stdout="", - stderr="node runtime not found", - ) - return SimpleNamespace(return_code=0, stdout="", stderr="") - - async def upload_file(self, *_args, **_kwargs): - return None - - capture = LLMTrajectoryCapture( - tmp_path, - agent="claude-agent-acp", - model="claude-sonnet-4-6", - session_id="rollout-1", - started_at=STARTED_AT, - ) - prepared = await capture.prepare_agent( - FailingCollectorEnv(), - agent="claude-agent-acp", - model="claude-sonnet-4-6", - agent_env={"CLAUDE_CODE_OAUTH_TOKEN": "oauth-test-token"}, - credential_home="/home/agent", - sandbox_user="agent", - ) - - assert "CLAUDE_CODE_ENABLE_TELEMETRY" not in prepared - assert "OTEL_LOG_RAW_API_BODIES" not in prepared - assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in prepared - assert any("/opt/benchflow/node/bin/node" in command for command in commands) - - -def test_capture_failure_repairs_invalid_jsonl_and_redacts_manifest_error( - tmp_path: Path, -) -> None: - """Guards PR #1057's valid-JSONL and secret-redaction failure invariant.""" - - capture = LLMTrajectoryCapture( - tmp_path, - agent="codex-acp", - model="gpt-5.6", - session_id="rollout-1", - started_at=STARTED_AT, - ) - capture.trajectory_path.write_text("{not-json\n") - - capture.record_failure( - "authorization: Bearer sk-test-secret-value", model_call_seen=True - ) - - manifest = json.loads( - (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() - ) - assert capture.trajectory_path.read_text() == "" - assert manifest["status"] == "capture_failed" - assert "sk-test-secret-value" not in json.dumps(manifest) - - -def test_acp_projection_retains_the_actual_auth_mode() -> None: - """Guards PR #1057 against labeling API-key fallback rows as OAuth.""" - - result = project_acp_trajectory( - [{"type": "agent_message", "text": "finished"}], - agent="codex-acp", - session_id="rollout-1", - started_at=STARTED_AT, - auth_mode="api_key", - ) - - assert result is not None - assert result.trajectory.exchanges[0].metadata["auth_mode"] == "api_key" From 750de4af7f8df11d9d4da47a4af0c079b3264f3e Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 09:06:06 -0700 Subject: [PATCH 38/74] fix: fail closed on unverified capture states --- .../continue_run/trajectory_artifacts.py | 9 +- src/benchflow/providers/litellm_logging.py | 5 + src/benchflow/providers/litellm_runtime.py | 100 ++++++++++++++++-- .../trajectories/llm_capture_manifest.py | 60 +++++++++++ .../trajectories/llm_capture_records.py | 21 +++- src/benchflow/trajectories/results.py | 18 ++-- tests/continue_run/test_orchestrator.py | 2 + tests/test_litellm_logging.py | 2 + tests/test_litellm_runtime.py | 61 ++++++++++- .../test_llm_capture_training_contract.py | 18 ++++ .../test_native_capture_resilience.py | 67 ++++++++++++ 11 files changed, 332 insertions(+), 31 deletions(-) diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py index 98bf4819c..c8de4d4f2 100644 --- a/src/benchflow/continue_run/trajectory_artifacts.py +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -9,7 +9,9 @@ from typing import Any from benchflow.trajectories.llm_capture_manifest import ( + CONTINUATION_SOURCE_AUDIT_ERROR, LLM_TRAJECTORY_SCHEMA_VERSION, + REPLAY_PROXY_INGRESS_AUDIT_ERROR, AuthMode, CaptureFidelity, CaptureSource, @@ -225,10 +227,7 @@ def refresh_stitched_trajectory_manifest( missing_fields = list(source.missing_fields) if source and not complete else [] errors.extend(live_errors) if n_live: - errors.append( - "live continuation request was captured at replay-proxy ingress, " - "before provider transformation" - ) + errors.append(REPLAY_PROXY_INGRESS_AUDIT_ERROR) missing_fields.append("live_provider_request") if n_live and not live_capture_host_owned: errors.append("sandbox replay capture shared root custody with the agent") @@ -236,7 +235,7 @@ def refresh_stitched_trajectory_manifest( errors.append("source LLM trajectory manifest is missing or malformed") missing_fields.append("source_capture_provenance") elif not source_allows_training: - errors.append("source LLM trajectory is not complete provider-wire capture") + errors.append(CONTINUATION_SOURCE_AUDIT_ERROR) if not count_matches: errors.append( "stitched LLM trajectory count mismatch: " diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 7486a9f58..8b6043f0e 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -617,6 +617,11 @@ def _exchange_metadata( request_body=request_body, ) metadata["request_complete"] = record.get("request_complete") is True + # A success callback proves LiteLLM received a provider response. Failure + # callbacks can also represent local DNS/connect/timeout failures after + # ``pre_api_call``; without explicit provider-response evidence they must + # fail closed even when the transformed request was captured completely. + metadata["response_complete"] = record.get("event") == "success" return metadata diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index d76490f82..0d6bd8f6c 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1740,11 +1740,7 @@ async def ensure_litellm_runtime( sorted(set(required_skill_names)), separators=(",", ":") ) proxy_location = "sandbox" if sandbox_local else "host" - capture_trusted = await _provider_capture_has_verified_custody( - sandbox_local=sandbox_local, - sandbox_user=sandbox_user, - sandbox=sandbox, - ) + capture_trusted = not sandbox_local config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" f"{session_id}:{role_name or 'primary'}:{sandbox_user or 'root'}:" @@ -1755,6 +1751,13 @@ async def ensure_litellm_runtime( if getattr(runtime, "config_key", None) == config_key and server is not None: is_running = await server.is_running() if is_running: + if sandbox_local: + capture_trusted = await _provider_capture_has_verified_custody( + sandbox_local=True, + sandbox_user=sandbox_user, + sandbox=sandbox, + runtime_dir=getattr(server, "runtime_dir", None), + ) runtime.capture_trusted = bool( getattr(runtime, "capture_trusted", False) and capture_trusted ) @@ -1811,6 +1814,14 @@ async def ensure_litellm_runtime( ), ) + if sandbox_local: + capture_trusted = await _provider_capture_has_verified_custody( + sandbox_local=True, + sandbox_user=sandbox_user, + sandbox=sandbox, + runtime_dir=getattr(server, "runtime_dir", None), + ) + from benchflow.providers.runtime import ProviderRuntime new_runtime = ProviderRuntime( @@ -1837,13 +1848,17 @@ async def ensure_litellm_runtime( async def _provider_capture_has_verified_custody( - *, sandbox_local: bool, sandbox_user: str | None, sandbox: Any | None + *, + sandbox_local: bool, + sandbox_user: str | None, + sandbox: Any | None, + runtime_dir: str | None, ) -> bool: - """Verify that the sandbox agent's effective UID cannot rewrite capture data.""" + """Prove the agent cannot read or mutate a root-owned runtime artifact.""" if not sandbox_local: return True - if sandbox is None or sandbox_user in {None, "", "root", "0"}: + if sandbox is None or sandbox_user in {None, "", "root", "0"} or not runtime_dir: return False try: result = await sandbox.exec( @@ -1862,7 +1877,74 @@ async def _provider_capture_has_verified_custody( except (AttributeError, ValueError): logger.warning("Provider capture custody UID check returned invalid output") return False - return effective_uid != 0 + if effective_uid == 0: + return False + + probe_path = f"{runtime_dir}/.benchflow-custody-{uuid4().hex}" + probe_value = secrets.token_hex(32) + quoted_path = shlex.quote(probe_path) + quoted_value = shlex.quote(probe_value) + try: + created = await sandbox.exec( + ( + f"umask 077; printf %s {quoted_value} > {quoted_path}; " + f"chown 0:0 {quoted_path}; chmod 600 {quoted_path}" + ), + user="root", + timeout_sec=10, + ) + if created.return_code != 0: + logger.warning("Provider capture custody probe creation failed") + return False + access = await sandbox.exec( + ( + f"if cat {quoted_path} >/dev/null 2>&1; then exit 0; fi; " + f"if printf compromised >> {quoted_path} 2>/dev/null; " + "then exit 0; fi; " + f"if rm -f {quoted_path} 2>/dev/null; then exit 0; fi; " + "if command -v sudo >/dev/null 2>&1 && " + f"sudo -n cat {quoted_path} >/dev/null 2>&1; then exit 0; fi; " + "if command -v doas >/dev/null 2>&1 && " + f"doas -n cat {quoted_path} >/dev/null 2>&1; then exit 0; fi; " + "if id -G 2>/dev/null | tr ' ' '\\n' | grep -qx 0; " + "then exit 0; fi; " + "effective_caps=$(awk '/^CapEff:/ {print $2}' " + "/proc/self/status 2>/dev/null); " + 'if [ -n "$effective_caps" ] && ' + "[ \"$effective_caps\" != '0000000000000000' ]; then exit 0; fi; " + "for privilege_socket in /var/run/docker.sock " + "/run/containerd/containerd.sock /run/podman/podman.sock; do " + 'if [ -S "$privilege_socket" ] && ' + '[ -r "$privilege_socket" ] && ' + '[ -w "$privilege_socket" ]; then exit 0; fi; done; ' + "exit 1" + ), + user=sandbox_user, + timeout_sec=10, + ) + if access.return_code != 1: + logger.warning( + "Provider capture custody probe found agent-accessible root data" + ) + return False + verified = await sandbox.exec( + ( + f'test "$(cat {quoted_path})" = {quoted_value} && ' + f"test \"$(stat -c '%u:%a' {quoted_path})\" = '0:600'" + ), + user="root", + timeout_sec=10, + ) + if verified.return_code != 0: + logger.warning("Provider capture custody probe integrity check failed") + return False + return True + except Exception as exc: + logger.warning("Provider capture custody artifact probe failed: %s", exc) + return False + finally: + with contextlib.suppress(Exception): + await sandbox.exec(f"rm -f {quoted_path}", user="root", timeout_sec=10) async def stop_litellm_runtime(runtime: Any | None) -> None: diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 3d36772fd..5708ee078 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -20,6 +20,13 @@ LLM_TRAJECTORY_FILENAME = "llm_trajectory.jsonl" LLM_TRAJECTORY_MANIFEST_FILENAME = "llm_trajectory.manifest.json" LLM_TRAJECTORY_SCHEMA_VERSION = 2 +REPLAY_PROXY_INGRESS_AUDIT_ERROR = ( + "live continuation request was captured at replay-proxy ingress, " + "before provider transformation" +) +CONTINUATION_SOURCE_AUDIT_ERROR = ( + "source LLM trajectory is not complete provider-wire capture" +) class CaptureStatus(StrEnum): @@ -301,6 +308,59 @@ def capture_manifest_has_replay_capture(manifest: dict[str, Any]) -> bool: return False +def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> bool: + """Accept only expected, internally complete audit-only capture states.""" + + if manifest.get("status") not in { + CaptureStatus.NO_MODEL_CALL.value, + CaptureStatus.PARTIAL.value, + }: + return False + raw_errors = manifest.get("errors", []) + if not isinstance(raw_errors, list) or not all( + isinstance(error, str) for error in raw_errors + ): + return False + errors = set(raw_errors) + has_oauth_capture = bool( + manifest.get("auth_mode") == AuthMode.OAUTH_SUBSCRIPTION.value + or capture_manifest_has_oauth_role_capture(manifest) + ) + if not capture_manifest_has_replay_capture(manifest): + return has_oauth_capture and not errors + + allowed_errors = {REPLAY_PROXY_INGRESS_AUDIT_ERROR} + if has_oauth_capture: + allowed_errors.add(CONTINUATION_SOURCE_AUDIT_ERROR) + if ( + REPLAY_PROXY_INGRESS_AUDIT_ERROR not in errors + or not errors.issubset(allowed_errors) + or manifest.get("response_complete") is not True + ): + return False + role_captures = manifest.get("role_captures") + if not isinstance(role_captures, list): + return False + replay_captures: list[LLMRoleCapture] = [] + for value in role_captures: + try: + role_capture = LLMRoleCapture.model_validate(value) + except ValidationError: + return False + if role_capture.capture_source is CaptureSource.REPLAY_PROXY: + replay_captures.append(role_capture) + return bool( + replay_captures + and all( + capture.capture_fidelity is CaptureFidelity.AGENT_SESSION + and capture.exchange_count > 0 + and capture.request_complete is False + and capture.response_complete is True + for capture in replay_captures + ) + ) + + def _atomic_write_text(path: Path, payload: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 1c178428c..70f8e1c64 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -98,13 +98,25 @@ def load_provider_wire_records( metadata.get("request_capture_source") == "litellm_pre_api_call_complete_input_dict" ) + response = record.get("response") + response_status = ( + response.get("status_code") if isinstance(response, dict) else None + ) + response_complete_value = metadata.get("response_complete") + response_complete = response_complete_value is True or ( + response_complete_value is None + and isinstance(response_status, int) + and 200 <= response_status < 300 + ) metadata.update( { "schema_version": LLM_TRAJECTORY_SCHEMA_VERSION, "capture_source": CaptureSource.LITELLM_PROXY.value, "capture_fidelity": ( CaptureFidelity.PROVIDER_WIRE.value - if capture_trusted and provider_request_observed + if capture_trusted + and provider_request_observed + and response_complete else CaptureFidelity.AGENT_SESSION.value ), "auth_mode": ( @@ -129,7 +141,7 @@ def load_provider_wire_records( ), "role_attribution_complete": attribution_complete, "request_complete": request_complete, - "response_complete": True, + "response_complete": response_complete, "payload_redacted": True, } ) @@ -218,6 +230,11 @@ def assemble_capture( for record in provider_records ): missing_fields.add("provider_request") + if any( + not _record_metadata_bool(record, "response_complete") + for record in provider_records + ): + missing_fields.add("provider_response") successful_records = [ record for record in records diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index b980ddabb..cc6ddfdda 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -34,8 +34,7 @@ ) from benchflow.trajectories.llm_capture_manifest import ( capture_artifact_allows_training, - capture_manifest_has_oauth_role_capture, - capture_manifest_has_replay_capture, + capture_manifest_preserves_audit_completion, read_llm_trajectory_manifest, ) from benchflow.trajectories.types import redact_trajectory_obj @@ -508,21 +507,16 @@ def build_rollout_results_record( ) audit_capture_preserves_completion = bool( ( - (agent_result or {}).get("usage_source") == USAGE_SOURCE_AGENT_NATIVE_ACP - or (capture_manifest or {}).get("auth_mode") == "oauth_subscription" - or ( - capture_manifest is not None - and capture_manifest_has_oauth_role_capture(capture_manifest) + ( + capture_manifest is None + and (agent_result or {}).get("usage_source") + == USAGE_SOURCE_AGENT_NATIVE_ACP ) or ( capture_manifest is not None - and capture_manifest_has_replay_capture(capture_manifest) + and capture_manifest_preserves_audit_completion(capture_manifest) ) ) - and ( - capture_manifest is None - or capture_manifest.get("status") in {"no_model_call", "partial"} - ) and effective_export_error is None and not terminal_health_error ) diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 363aea935..22bc458df 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -24,6 +24,7 @@ write_stitched_trajectory, ) from benchflow.trajectories.llm_capture_manifest import ( + REPLAY_PROXY_INGRESS_AUDIT_ERROR, AuthMode, CaptureFidelity, CaptureSource, @@ -740,6 +741,7 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): "capture_source": CaptureSource.MIXED, "capture_fidelity": CaptureFidelity.MIXED, "request_complete": False, + "errors": [REPLAY_PROXY_INGRESS_AUDIT_ERROR], "role_captures": [ LLMRoleCapture( role="agent", diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 999d4f187..77f3969d6 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -622,6 +622,7 @@ def test_opencode_callback_import_preserves_call_metadata_and_purpose(): }, "call_purpose": "agent", "request_complete": True, + "response_complete": True, } assert [exchange.metadata["call_purpose"] for exchange in trajectory.exchanges] == [ "agent", @@ -731,6 +732,7 @@ def test_litellm_failure_records_become_error_exchanges(): assert trajectory.exchanges[0].response.status_code == 500 assert trajectory.exchanges[0].response.body["error"]["message"] == "bad key" + assert trajectory.exchanges[0].metadata["response_complete"] is False usage = extract_usage_from_trajectory(trajectory, fallback_model="openai/gpt-4") assert usage["usage_source"] == "unavailable" diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index c787b3420..643bffd35 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -20,6 +20,7 @@ class FakeLiteLLMServer: def __init__(self, base_url: str, route): self._base_url = base_url self.route = route + self.runtime_dir = "/tmp/benchflow-litellm/test-runtime" self.stopped = False self.trajectory = None @@ -176,9 +177,24 @@ async def fake_sandbox_start(**kwargs): monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) class NonRootSandbox: - async def exec(self, command, **_kwargs): - assert command == "id -u -- agent" - return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") + async def exec(self, command, **kwargs): + if command == "id -u -- agent": + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") + if command.startswith("umask 077"): + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") + if kwargs["user"] == "agent": + assert "sudo -n cat" in command + assert "CapEff" in command + assert "/var/run/docker.sock" in command + return SimpleNamespace(return_code=1, stdout="", stderr="") + if command.startswith('test "$(cat'): + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") + assert command.startswith("rm -f") + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") _, provider_runtime = await ensure_litellm_runtime( agent="openhands", @@ -198,6 +214,45 @@ async def exec(self, command, **_kwargs): assert provider_runtime.capture_trusted is True +@pytest.mark.asyncio +async def test_passwordless_privilege_keeps_sandbox_capture_audit_only(monkeypatch): + """Guards PR #1057 against trusting users that can regain root access.""" + + async def fake_sandbox_start(**kwargs): + return FakeLiteLLMServer("http://127.0.0.1:45678", kwargs["route"]) + + monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) + + class PrivilegedSandbox: + async def exec(self, command, **kwargs): + if command == "id -u -- agent": + return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") + if command.startswith("umask 077"): + return SimpleNamespace(return_code=0, stdout="", stderr="") + if kwargs["user"] == "agent": + assert "sudo -n cat" in command + return SimpleNamespace(return_code=0, stdout="", stderr="") + assert command.startswith("rm -f") + return SimpleNamespace(return_code=0, stdout="", stderr="") + + _, provider_runtime = await ensure_litellm_runtime( + agent="openhands", + agent_env={ + "AWS_BEARER_TOKEN_BEDROCK": "token", + "AWS_REGION": "us-west-2", + }, + model="aws-bedrock/us.anthropic.claude-opus-4-8", + runtime=None, + environment="daytona", + session_id="run-privileged-agent", + sandbox=PrivilegedSandbox(), + sandbox_user="agent", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is False + + @pytest.mark.asyncio async def test_uid_zero_alias_keeps_sandbox_gateway_capture_audit_only(monkeypatch): """Guards PR #1057 against trusting a non-root name that resolves to UID 0.""" diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 845b2c837..aaaf85a7a 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -9,6 +9,9 @@ import pytest from benchflow.trajectories.llm_capture import LLMTrajectoryCapture +from benchflow.trajectories.llm_capture_manifest import ( + REPLAY_PROXY_INGRESS_AUDIT_ERROR, +) from benchflow.trajectories.results import build_rollout_results_record @@ -296,6 +299,7 @@ def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> "exchange_count": 1, "request_complete": False, "response_complete": True, + "errors": [REPLAY_PROXY_INGRESS_AUDIT_ERROR], "role_captures": [ { "role": "agent", @@ -324,6 +328,20 @@ def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> assert row["is_completed"] is True assert row["error"] is None + manifest_path = trajectory_dir / "llm_trajectory.manifest.json" + unhealthy_manifest = json.loads(manifest_path.read_text()) + unhealthy_manifest["errors"].append("live attempt journal mismatch") + manifest_path.write_text(json.dumps(unhealthy_manifest)) + + unhealthy_row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + + assert unhealthy_row["info"]["training_ready"] is False + assert unhealthy_row["is_completed"] is False + assert unhealthy_row["error"]["error"] == "missing_llm_trajectory" + @pytest.mark.asyncio async def test_provider_finalization_error_rejects_complete_live_prefix( diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 34bc7600b..07d616248 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -18,6 +18,7 @@ ) from benchflow.trajectories.llm_capture_records import ( NativeCaptureBundle, + assemble_capture, load_provider_wire_records, ) from benchflow.trajectories.native_capture_collection import NativeCollection @@ -160,6 +161,72 @@ def test_provider_role_attribution_uses_proxy_model_aliases(tmp_path: Path) -> N ) +def test_local_provider_failure_downgrades_mixed_capture(tmp_path: Path) -> None: + """Guards PR #1057 against treating local LiteLLM failures as responses.""" + + trajectory_path = tmp_path / "llm_trajectory.jsonl" + rows = [ + { + "request": {"body": {"model": "gpt-5.6", "input": "first"}}, + "response": { + "status_code": 200, + "body": { + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + }, + "metadata": { + "benchflow_requested_model": "openai/gpt-5.6", + "request_complete": True, + "response_complete": True, + "request_capture_source": ("litellm_pre_api_call_complete_input_dict"), + }, + }, + { + "request": {"body": {"model": "gpt-5.6", "input": "second"}}, + "response": { + "status_code": 500, + "body": {"error": {"message": "connection failed locally"}}, + }, + "metadata": { + "benchflow_requested_model": "openai/gpt-5.6", + "request_complete": True, + "response_complete": False, + "request_capture_source": ("litellm_pre_api_call_complete_input_dict"), + }, + }, + ] + trajectory_path.write_text("".join(json.dumps(row) + "\n" for row in rows)) + target = _CaptureTarget( + agent="codex-acp", + model="openai/gpt-5.6", + credential_home="/home/agent", + auth_mode=AuthMode.API_KEY, + native=False, + ) + + provider_records = load_provider_wire_records( + trajectory_path, + targets=[target], + fallback_agent="codex-acp", + fallback_model="openai/gpt-5.6", + fallback_auth=AuthMode.API_KEY, + ) + assembly = assemble_capture( + provider_records=provider_records, + native_bundles=[], + targets=[target], + collection_errors=[], + model_call_seen=True, + fallback_auth=AuthMode.API_KEY, + ) + + assert assembly.status is CaptureStatus.PARTIAL + assert assembly.fidelity is CaptureFidelity.MIXED + assert assembly.response_complete is False + assert "provider_response" in assembly.missing_fields + + def test_provider_role_attribution_uses_runtime_identity_for_same_model( tmp_path: Path, ) -> None: From 58c1e04044df7815653c214386a9dfb22831745c Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 09:29:02 -0700 Subject: [PATCH 39/74] fix: reject incomplete audit capture metadata --- .../trajectories/llm_capture_manifest.py | 28 ++++++++++- tests/continue_run/test_orchestrator.py | 1 + .../test_llm_capture_training_contract.py | 48 ++++++++++++++----- 3 files changed, 65 insertions(+), 12 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 5708ee078..3a0d157c4 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -27,6 +27,18 @@ CONTINUATION_SOURCE_AUDIT_ERROR = ( "source LLM trajectory is not complete provider-wire capture" ) +_OAUTH_AUDIT_MISSING_FIELDS = frozenset( + { + "headers", + "instructions", + "provider_request", + "provider_response", + "provider_response_envelope", + "system_prompt", + "tool_definitions", + } +) +_REPLAY_AUDIT_MISSING_FIELD = "live_provider_request" class CaptureStatus(StrEnum): @@ -322,19 +334,33 @@ def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> boo ): return False errors = set(raw_errors) + raw_missing_fields = manifest.get("missing_fields", []) + if not isinstance(raw_missing_fields, list) or not all( + isinstance(field, str) for field in raw_missing_fields + ): + return False + missing_fields = set(raw_missing_fields) has_oauth_capture = bool( manifest.get("auth_mode") == AuthMode.OAUTH_SUBSCRIPTION.value or capture_manifest_has_oauth_role_capture(manifest) ) if not capture_manifest_has_replay_capture(manifest): - return has_oauth_capture and not errors + return bool( + has_oauth_capture + and not errors + and missing_fields.issubset(_OAUTH_AUDIT_MISSING_FIELDS) + ) allowed_errors = {REPLAY_PROXY_INGRESS_AUDIT_ERROR} + allowed_missing_fields = {_REPLAY_AUDIT_MISSING_FIELD} if has_oauth_capture: allowed_errors.add(CONTINUATION_SOURCE_AUDIT_ERROR) + allowed_missing_fields.update(_OAUTH_AUDIT_MISSING_FIELDS) if ( REPLAY_PROXY_INGRESS_AUDIT_ERROR not in errors or not errors.issubset(allowed_errors) + or _REPLAY_AUDIT_MISSING_FIELD not in missing_fields + or not missing_fields.issubset(allowed_missing_fields) or manifest.get("response_complete") is not True ): return False diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 22bc458df..e4b39eae6 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -741,6 +741,7 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): "capture_source": CaptureSource.MIXED, "capture_fidelity": CaptureFidelity.MIXED, "request_complete": False, + "missing_fields": ["live_provider_request"], "errors": [REPLAY_PROXY_INGRESS_AUDIT_ERROR], "role_captures": [ LLMRoleCapture( diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index aaaf85a7a..def53eacf 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -236,6 +236,12 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( "exchange_count": 2, "request_complete": False, "response_complete": True, + "missing_fields": [ + "headers", + "provider_response_envelope", + "system_prompt", + "tool_definitions", + ], "role_captures": [ { "role": "coder", @@ -272,6 +278,20 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( assert row["is_completed"] is True assert row["error"] is None + manifest_path = trajectory_dir / "llm_trajectory.manifest.json" + unhealthy_manifest = json.loads(manifest_path.read_text()) + unhealthy_manifest["missing_fields"].append("token_usage") + manifest_path.write_text(json.dumps(unhealthy_manifest)) + + unhealthy_row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + + assert unhealthy_row["info"]["training_ready"] is False + assert unhealthy_row["is_completed"] is False + assert unhealthy_row["error"]["error"] == "missing_llm_trajectory" + def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> None: """Guards PR #1057 against marking successful continuations incomplete.""" @@ -299,6 +319,7 @@ def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> "exchange_count": 1, "request_complete": False, "response_complete": True, + "missing_fields": ["live_provider_request"], "errors": [REPLAY_PROXY_INGRESS_AUDIT_ERROR], "role_captures": [ { @@ -329,18 +350,23 @@ def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> assert row["error"] is None manifest_path = trajectory_dir / "llm_trajectory.manifest.json" - unhealthy_manifest = json.loads(manifest_path.read_text()) - unhealthy_manifest["errors"].append("live attempt journal mismatch") - manifest_path.write_text(json.dumps(unhealthy_manifest)) - - unhealthy_row = _build_results_row( - tmp_path, - agent_result={"usage_source": "provider_response", "total_tokens": 2}, - ) + healthy_manifest = json.loads(manifest_path.read_text()) + for field, value in ( + ("missing_fields", "token_usage"), + ("errors", "live attempt journal mismatch"), + ): + unhealthy_manifest = json.loads(json.dumps(healthy_manifest)) + unhealthy_manifest[field].append(value) + manifest_path.write_text(json.dumps(unhealthy_manifest)) + + unhealthy_row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) - assert unhealthy_row["info"]["training_ready"] is False - assert unhealthy_row["is_completed"] is False - assert unhealthy_row["error"]["error"] == "missing_llm_trajectory" + assert unhealthy_row["info"]["training_ready"] is False + assert unhealthy_row["is_completed"] is False + assert unhealthy_row["error"]["error"] == "missing_llm_trajectory" @pytest.mark.asyncio From ad2f66645574af1451860310389632634a613654 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 09:49:10 -0700 Subject: [PATCH 40/74] fix: validate audit completion per role --- .../trajectories/llm_capture_manifest.py | 173 +++++++++++------- .../test_llm_capture_training_contract.py | 49 ++++- 2 files changed, 154 insertions(+), 68 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 3a0d157c4..f654dadb1 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -68,6 +68,16 @@ class CaptureSource(StrEnum): NONE = "none" +_OAUTH_AUDIT_CAPTURE_SOURCES = frozenset( + { + CaptureSource.CLAUDE_OTEL_RAW_BODY, + CaptureSource.CLAUDE_NATIVE_SESSION, + CaptureSource.CODEX_NATIVE_SESSION, + CaptureSource.ACP_PROJECTION, + } +) + + class AuthMode(StrEnum): API_KEY = "api_key" OAUTH_SUBSCRIPTION = "oauth_subscription" @@ -273,53 +283,6 @@ def _exchange_requires_manifest(exchange: dict[str, Any]) -> bool: ) -def capture_manifest_has_oauth_role_capture(manifest: dict[str, Any]) -> bool: - """Return whether a mixed manifest contains captured OAuth role evidence.""" - - role_captures = manifest.get("role_captures") - if not isinstance(role_captures, list): - return False - for value in role_captures: - try: - role_capture = LLMRoleCapture.model_validate(value) - except ValidationError: - continue - if ( - role_capture.auth_mode is AuthMode.OAUTH_SUBSCRIPTION - and role_capture.capture_source is not CaptureSource.NONE - and role_capture.capture_fidelity is not CaptureFidelity.NONE - and role_capture.exchange_count > 0 - ): - return True - return False - - -def capture_manifest_has_replay_capture(manifest: dict[str, Any]) -> bool: - """Return whether a manifest contains an audit-only continuation suffix.""" - - if ( - manifest.get("capture_source") == CaptureSource.REPLAY_PROXY.value - and isinstance(manifest.get("exchange_count"), int) - and not isinstance(manifest.get("exchange_count"), bool) - and manifest["exchange_count"] > 0 - ): - return True - role_captures = manifest.get("role_captures") - if not isinstance(role_captures, list): - return False - for value in role_captures: - try: - role_capture = LLMRoleCapture.model_validate(value) - except ValidationError: - continue - if ( - role_capture.capture_source is CaptureSource.REPLAY_PROXY - and role_capture.exchange_count > 0 - ): - return True - return False - - def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> bool: """Accept only expected, internally complete audit-only capture states.""" @@ -340,11 +303,38 @@ def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> boo ): return False missing_fields = set(raw_missing_fields) - has_oauth_capture = bool( - manifest.get("auth_mode") == AuthMode.OAUTH_SUBSCRIPTION.value - or capture_manifest_has_oauth_role_capture(manifest) + role_captures = _validated_role_captures(manifest) + if role_captures is None or not _role_captures_match_manifest( + manifest, role_captures + ): + return False + has_oauth_capture = any( + _is_oauth_audit_role_capture(capture) for capture in role_captures ) - if not capture_manifest_has_replay_capture(manifest): + if manifest.get("status") == CaptureStatus.NO_MODEL_CALL.value: + return bool( + manifest.get("auth_mode") == AuthMode.OAUTH_SUBSCRIPTION.value + and not errors + and not missing_fields + and all( + capture.auth_mode is AuthMode.OAUTH_SUBSCRIPTION + and capture.capture_source is CaptureSource.NONE + and capture.capture_fidelity is CaptureFidelity.NONE + and capture.exchange_count == 0 + and capture.request_complete is False + and capture.response_complete is False + for capture in role_captures + ) + ) + if not role_captures or not all( + _role_capture_preserves_audit_completion(capture) for capture in role_captures + ): + return False + has_replay_capture = any( + capture.capture_source is CaptureSource.REPLAY_PROXY + for capture in role_captures + ) + if not has_replay_capture: return bool( has_oauth_capture and not errors @@ -356,34 +346,83 @@ def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> boo if has_oauth_capture: allowed_errors.add(CONTINUATION_SOURCE_AUDIT_ERROR) allowed_missing_fields.update(_OAUTH_AUDIT_MISSING_FIELDS) - if ( + return not ( REPLAY_PROXY_INGRESS_AUDIT_ERROR not in errors or not errors.issubset(allowed_errors) or _REPLAY_AUDIT_MISSING_FIELD not in missing_fields or not missing_fields.issubset(allowed_missing_fields) or manifest.get("response_complete") is not True + ) + + +def _validated_role_captures( + manifest: dict[str, Any], +) -> list[LLMRoleCapture] | None: + raw_role_captures = manifest.get("role_captures", []) + if not isinstance(raw_role_captures, list) or any( + not isinstance(value, dict) + or not isinstance(value.get("exchange_count"), int) + or isinstance(value.get("exchange_count"), bool) + or value["exchange_count"] < 0 + or not isinstance(value.get("request_complete"), bool) + or not isinstance(value.get("response_complete"), bool) + for value in raw_role_captures ): - return False - role_captures = manifest.get("role_captures") - if not isinstance(role_captures, list): - return False - replay_captures: list[LLMRoleCapture] = [] - for value in role_captures: - try: - role_capture = LLMRoleCapture.model_validate(value) - except ValidationError: - return False - if role_capture.capture_source is CaptureSource.REPLAY_PROXY: - replay_captures.append(role_capture) + return None + try: + return [LLMRoleCapture.model_validate(value) for value in raw_role_captures] + except ValidationError: + return None + + +def _is_oauth_audit_role_capture(capture: LLMRoleCapture) -> bool: return bool( - replay_captures - and all( + capture.auth_mode is AuthMode.OAUTH_SUBSCRIPTION + and capture.capture_source in _OAUTH_AUDIT_CAPTURE_SOURCES + and capture.capture_fidelity + in {CaptureFidelity.AGENT_SESSION, CaptureFidelity.ACP_PROJECTION} + and capture.exchange_count > 0 + ) + + +def _role_captures_match_manifest( + manifest: dict[str, Any], role_captures: list[LLMRoleCapture] +) -> bool: + exchange_count = manifest.get("exchange_count") + if ( + not isinstance(exchange_count, int) + or isinstance(exchange_count, bool) + or exchange_count < 0 + or sum(capture.exchange_count for capture in role_captures) != exchange_count + ): + return False + if exchange_count == 0: + return True + active_auth_modes = { + capture.auth_mode for capture in role_captures if capture.exchange_count > 0 + } + expected_auth_mode = ( + next(iter(active_auth_modes)) if len(active_auth_modes) == 1 else AuthMode.MIXED + ) + return manifest.get("auth_mode") == expected_auth_mode.value + + +def _role_capture_preserves_audit_completion(capture: LLMRoleCapture) -> bool: + if capture.capture_source is CaptureSource.REPLAY_PROXY: + return bool( capture.capture_fidelity is CaptureFidelity.AGENT_SESSION and capture.exchange_count > 0 and capture.request_complete is False and capture.response_complete is True - for capture in replay_captures ) + if _is_oauth_audit_role_capture(capture): + return True + return bool( + capture.capture_source is not CaptureSource.NONE + and capture.capture_fidelity is not CaptureFidelity.NONE + and capture.exchange_count > 0 + and capture.request_complete is True + and capture.response_complete is True ) diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index def53eacf..20cc8950e 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -85,11 +85,24 @@ def test_agent_session_capture_is_audit_only_not_training_ready(tmp_path: Path) json.dumps( { "status": "partial", + "capture_source": "codex_native_session", "capture_fidelity": "agent_session", "auth_mode": "oauth_subscription", "exchange_count": 1, "request_complete": False, "response_complete": True, + "role_captures": [ + { + "role": "agent", + "agent": "codex-acp", + "auth_mode": "oauth_subscription", + "capture_source": "codex_native_session", + "capture_fidelity": "agent_session", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + } + ], } ) ) @@ -279,7 +292,8 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( assert row["error"] is None manifest_path = trajectory_dir / "llm_trajectory.manifest.json" - unhealthy_manifest = json.loads(manifest_path.read_text()) + healthy_manifest = json.loads(manifest_path.read_text()) + unhealthy_manifest = json.loads(json.dumps(healthy_manifest)) unhealthy_manifest["missing_fields"].append("token_usage") manifest_path.write_text(json.dumps(unhealthy_manifest)) @@ -292,6 +306,39 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( assert unhealthy_row["is_completed"] is False assert unhealthy_row["error"]["error"] == "missing_llm_trajectory" + cross_role_manifest = json.loads(json.dumps(healthy_manifest)) + cross_role_manifest["missing_fields"] = [ + "headers", + "provider_response", + "provider_response_envelope", + "system_prompt", + "tool_definitions", + ] + cross_role_manifest["role_captures"][0]["response_complete"] = False + manifest_path.write_text(json.dumps(cross_role_manifest)) + + cross_role_row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + + assert cross_role_row["info"]["training_ready"] is False + assert cross_role_row["is_completed"] is False + assert cross_role_row["error"]["error"] == "missing_llm_trajectory" + + omitted_role_manifest = json.loads(json.dumps(healthy_manifest)) + omitted_role_manifest["role_captures"].pop(0) + manifest_path.write_text(json.dumps(omitted_role_manifest)) + + omitted_role_row = _build_results_row( + tmp_path, + agent_result={"usage_source": "provider_response", "total_tokens": 2}, + ) + + assert omitted_role_row["info"]["training_ready"] is False + assert omitted_role_row["is_completed"] is False + assert omitted_role_row["error"]["error"] == "missing_llm_trajectory" + def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> None: """Guards PR #1057 against marking successful continuations incomplete.""" From f5a38de8feabd8bfc64a2143386a0229e98587cb Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 10:07:38 -0700 Subject: [PATCH 41/74] fix: preserve repeated replay completion --- .../trajectories/llm_capture_manifest.py | 8 +++- tests/continue_run/test_orchestrator.py | 38 +++++++++++++++++++ .../test_llm_capture_training_contract.py | 14 +++++++ 3 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index f654dadb1..add10fe9a 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -343,8 +343,14 @@ def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> boo allowed_errors = {REPLAY_PROXY_INGRESS_AUDIT_ERROR} allowed_missing_fields = {_REPLAY_AUDIT_MISSING_FIELD} - if has_oauth_capture: + has_recorded_replay_capture = any( + capture.capture_source is CaptureSource.REPLAY_PROXY + and capture.leg == "recorded" + for capture in role_captures + ) + if has_oauth_capture or has_recorded_replay_capture: allowed_errors.add(CONTINUATION_SOURCE_AUDIT_ERROR) + if has_oauth_capture: allowed_missing_fields.update(_OAUTH_AUDIT_MISSING_FIELDS) return not ( REPLAY_PROXY_INGRESS_AUDIT_ERROR not in errors diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index e4b39eae6..1d371697d 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -24,6 +24,7 @@ write_stitched_trajectory, ) from benchflow.trajectories.llm_capture_manifest import ( + CONTINUATION_SOURCE_AUDIT_ERROR, REPLAY_PROXY_INGRESS_AUDIT_ERROR, AuthMode, CaptureFidelity, @@ -32,6 +33,7 @@ LLMRoleCapture, LLMTrajectoryManifest, capture_manifest_allows_training, + capture_manifest_preserves_audit_completion, initialize_llm_trajectory_artifacts, write_llm_trajectory_manifest, ) @@ -247,6 +249,42 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p assert live_row["metadata"]["model"] == model assert live_row["metadata"]["schema_version"] == 2 + repeated_rollout = tmp_path / "continued-twice" + initialize_llm_trajectory_artifacts( + repeated_rollout, + agent="openhands", + model=None, + session_id="continued-twice", + started_at=manifest.finished_at, + ) + write_stitched_trajectory( + repeated_rollout, + rollout / "trajectory" / "llm_trajectory.jsonl", + [exchange(completion(content="live-again"))], + live_model=model, + ) + repeated_manifest = refresh_stitched_trajectory_manifest( + repeated_rollout, + rollout, + original_model=model, + live_model=model, + n_recorded=2, + n_live=1, + live_attempt_count=1, + live_errors=[], + ) + + assert repeated_manifest.exchange_count == 3 + assert CONTINUATION_SOURCE_AUDIT_ERROR in repeated_manifest.errors + assert [capture.leg for capture in repeated_manifest.role_captures] == [ + "recorded", + "recorded", + "live", + ] + assert capture_manifest_preserves_audit_completion( + repeated_manifest.model_dump(mode="json") + ) + def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path): """Guards PR #1057 against promoting audit-only continuation prefixes.""" diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 20cc8950e..a9889d2d5 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -10,7 +10,9 @@ from benchflow.trajectories.llm_capture import LLMTrajectoryCapture from benchflow.trajectories.llm_capture_manifest import ( + CONTINUATION_SOURCE_AUDIT_ERROR, REPLAY_PROXY_INGRESS_AUDIT_ERROR, + capture_manifest_preserves_audit_completion, ) from benchflow.trajectories.results import build_rollout_results_record @@ -398,6 +400,18 @@ def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> manifest_path = trajectory_dir / "llm_trajectory.manifest.json" healthy_manifest = json.loads(manifest_path.read_text()) + repeated_manifest = json.loads(json.dumps(healthy_manifest)) + repeated_manifest["exchange_count"] = 2 + repeated_manifest["errors"].append(CONTINUATION_SOURCE_AUDIT_ERROR) + repeated_manifest["role_captures"][0]["leg"] = "recorded" + live_role = json.loads(json.dumps(repeated_manifest["role_captures"][0])) + live_role["leg"] = "live" + repeated_manifest["role_captures"].append(live_role) + assert capture_manifest_preserves_audit_completion(repeated_manifest) is True + + repeated_manifest["role_captures"][0]["leg"] = "live" + assert capture_manifest_preserves_audit_completion(repeated_manifest) is False + for field, value in ( ("missing_fields", "token_usage"), ("errors", "live attempt journal mismatch"), From 9bf15d8f31fb0fd3b89882f4c10608e4383e48f9 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 10:34:30 -0700 Subject: [PATCH 42/74] fix repeated continuation model provenance --- .../continue_run/trajectory_artifacts.py | 13 +++---- tests/continue_run/test_orchestrator.py | 35 +++++++++++++------ 2 files changed, 29 insertions(+), 19 deletions(-) diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py index c8de4d4f2..220a35709 100644 --- a/src/benchflow/continue_run/trajectory_artifacts.py +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -253,15 +253,6 @@ def refresh_stitched_trajectory_manifest( if not live_capture_complete: missing_fields.append("live_provider_exchange") - models = { - value - for value, present in ( - (original_model, n_recorded > 0), - (live_model, live_attempt_count > 0), - ) - if present and value - } - stitched_model = next(iter(models)) if len(models) == 1 else None role_captures = _continuation_role_captures( source=source, original_model=original_model, @@ -272,6 +263,10 @@ def refresh_stitched_trajectory_manifest( live_capture_complete=live_capture_complete, rows_valid=rows_valid, ) + active_models = { + capture.model for capture in role_captures if capture.exchange_count > 0 + } + stitched_model = next(iter(active_models)) if len(active_models) == 1 else None manifest = LLMTrajectoryManifest( status=CaptureStatus.COMPLETE if complete else CaptureStatus.PARTIAL, capture_source=capture_source, diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 1d371697d..44a36a3c0 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -183,11 +183,12 @@ def test_write_stitched_trajectory_creates_file(tmp_path): def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_path): """Guards PR #1057 against promoting continuation ingress to provider wire.""" - model = "openai/gpt-5.5" + recorded_model = "openai/gpt-5.4" + continued_model = "openai/gpt-5.5" source = write_run_folder( tmp_path / "source", exchanges=[exchange(completion(content="recorded"))], - model=model, + model=recorded_model, ) source_manifest = LLMTrajectoryManifest( status=CaptureStatus.COMPLETE, @@ -195,7 +196,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p capture_fidelity=CaptureFidelity.PROVIDER_WIRE, auth_mode=AuthMode.API_KEY, agent="openhands", - model=model, + model=recorded_model, session_id="source", exchange_count=1, request_complete=True, @@ -218,13 +219,13 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p rollout, source / "trajectory" / "llm_trajectory.jsonl", live, - live_model=model, + live_model=continued_model, ) manifest = refresh_stitched_trajectory_manifest( rollout, source, - original_model=model, - live_model=model, + original_model=recorded_model, + live_model=continued_model, n_recorded=1, n_live=1, live_attempt_count=1, @@ -235,6 +236,11 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p assert manifest.capture_source is CaptureSource.MIXED assert manifest.capture_fidelity is CaptureFidelity.MIXED assert manifest.exchange_count == 2 + assert manifest.model is None + assert [capture.model for capture in manifest.role_captures] == [ + recorded_model, + continued_model, + ] assert manifest.request_complete is False assert "live_provider_request" in manifest.missing_fields assert not capture_manifest_allows_training( @@ -246,7 +252,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p assert live_row["metadata"]["request_complete"] is False assert live_row["metadata"]["response_complete"] is True assert live_row["metadata"]["request_capture_source"] == "replay_proxy_ingress" - assert live_row["metadata"]["model"] == model + assert live_row["metadata"]["model"] == continued_model assert live_row["metadata"]["schema_version"] == 2 repeated_rollout = tmp_path / "continued-twice" @@ -261,13 +267,16 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p repeated_rollout, rollout / "trajectory" / "llm_trajectory.jsonl", [exchange(completion(content="live-again"))], - live_model=model, + live_model=continued_model, ) repeated_manifest = refresh_stitched_trajectory_manifest( repeated_rollout, rollout, - original_model=model, - live_model=model, + # update_continued_metadata() rewrites the first continuation's + # run-level config to the live model, even though its prefix retains + # recorded_model. The aggregate must therefore come from role captures. + original_model=continued_model, + live_model=continued_model, n_recorded=2, n_live=1, live_attempt_count=1, @@ -275,12 +284,18 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p ) assert repeated_manifest.exchange_count == 3 + assert repeated_manifest.model is None assert CONTINUATION_SOURCE_AUDIT_ERROR in repeated_manifest.errors assert [capture.leg for capture in repeated_manifest.role_captures] == [ "recorded", "recorded", "live", ] + assert [capture.model for capture in repeated_manifest.role_captures] == [ + recorded_model, + continued_model, + continued_model, + ] assert capture_manifest_preserves_audit_completion( repeated_manifest.model_dump(mode="json") ) From 3885e0daa80e6e78ac923eb16873043b5751d955 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 10:58:03 -0700 Subject: [PATCH 43/74] require redacted payloads for training --- .../trajectories/llm_capture_manifest.py | 1 + .../test_llm_capture_training_contract.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index add10fe9a..a016a1228 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -183,6 +183,7 @@ def capture_manifest_allows_training( and manifest.get("capture_fidelity") == "provider_wire" and manifest.get("request_complete") is True and manifest.get("response_complete") is True + and manifest.get("payload_redacted") is True and exchange_count > 0 and expected_count == exchange_count ) diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index a9889d2d5..8771e9ecf 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -12,6 +12,7 @@ from benchflow.trajectories.llm_capture_manifest import ( CONTINUATION_SOURCE_AUDIT_ERROR, REPLAY_PROXY_INGRESS_AUDIT_ERROR, + capture_manifest_allows_training, capture_manifest_preserves_audit_completion, ) from benchflow.trajectories.results import build_rollout_results_record @@ -187,6 +188,32 @@ def test_manifest_count_mismatch_fails_closed_for_canonical_results( assert row["is_completed"] is False +def test_unredacted_provider_capture_fails_closed_for_training( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on explicitly unredacted payloads.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="provider_wire", schema_version=2) + manifest = { + "status": "complete", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + "payload_redacted": False, + } + (trajectory_dir / "llm_trajectory.manifest.json").write_text(json.dumps(manifest)) + + assert not capture_manifest_allows_training(manifest, exchange_count=1) + row = _build_results_row(tmp_path, agent_result={"total_tokens": 2}) + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is False + + def test_provider_capture_without_positive_usage_is_not_training_ready( tmp_path: Path, ) -> None: From 8dad07af53c4249cf33ef971b6d9782968a934d6 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 11:25:00 -0700 Subject: [PATCH 44/74] Harden provider capture artifact custody --- src/benchflow/providers/litellm_logging.py | 29 ++++++-- src/benchflow/providers/litellm_runtime.py | 78 +++++++++++++++------- tests/test_litellm_hardening.py | 21 ++++++ tests/test_litellm_logging.py | 25 +++++++ tests/test_litellm_runtime.py | 65 +++++++++++++++--- 5 files changed, 180 insertions(+), 38 deletions(-) diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 8b6043f0e..90281df9d 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -108,6 +108,27 @@ def callback_module_source() -> str: _skill_catalog_gate_passed = False +def _secure_parent(path: str) -> None: + parent = os.path.dirname(path) + if not parent: + return + os.makedirs(parent, mode=0o700, exist_ok=True) + os.chmod(parent, 0o700) + + +def _secure_text_file(path: str, *, append: bool): + _secure_parent(path) + flags = os.O_WRONLY | os.O_CREAT | (os.O_APPEND if append else os.O_TRUNC) + flags |= getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0) + fd = os.open(path, flags, 0o600) + try: + os.fchmod(fd, 0o600) + return os.fdopen(fd, "a" if append else "w", encoding="utf-8") + except BaseException: + os.close(fd) + raise + + def _required_skill_names() -> tuple[str, ...]: raw = os.environ.get("BENCHFLOW_REQUIRED_SKILL_NAMES_JSON", "") if not raw: @@ -251,8 +272,7 @@ def _write_state_locked(self) -> None: } temporary = path + ".tmp" try: - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(temporary, "w", encoding="utf-8") as handle: + with _secure_text_file(temporary, append=False) as handle: json.dump(payload, handle, separators=(",", ":")) handle.flush() os.fsync(handle.fileno()) @@ -274,7 +294,7 @@ def _invalidate_state_locked(self, path: str, temporary: str) -> None: # If removal itself is unavailable, corrupt the validity contract # explicitly. Cleanup rejects this payload as a malformed journal. try: - with open(path, "w", encoding="utf-8") as handle: + with _secure_text_file(path, append=False) as handle: json.dump({"valid": False}, handle, separators=(",", ":")) handle.flush() os.fsync(handle.fileno()) @@ -297,8 +317,7 @@ def _write(self, payload: dict[str, Any]) -> None: payload["logged_at"] = datetime.now(timezone.utc).isoformat() durable_payload = redact_trajectory_obj(_jsonable(payload)) with self._capture_lock: - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "a", encoding="utf-8") as handle: + with _secure_text_file(path, append=True) as handle: handle.write( json.dumps(durable_payload, separators=(",", ":")) + "\n" ) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 0d6bd8f6c..6cc208c9e 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -824,6 +824,7 @@ def _write_runtime_files( config: dict[str, object], ) -> tuple[Path, Path, Path]: runtime_dir.mkdir(parents=True, exist_ok=True) + runtime_dir.chmod(0o700) callback_path = runtime_dir / f"{_CALLBACK_MODULE}.py" redaction_path = runtime_dir / f"{_REDACTION_MODULE}.py" patch_path = runtime_dir / f"{_PATCH_MODULE}.py" @@ -835,6 +836,14 @@ def _write_runtime_files( patch_path.write_text(patch_source) sitecustomize_path.write_text(f"import {_PATCH_MODULE}\n") config_path.write_text(yaml.safe_dump(config, sort_keys=False)) + for path in ( + callback_path, + redaction_path, + patch_path, + sitecustomize_path, + config_path, + ): + path.chmod(0o600) return config_path, callback_path, patch_path @@ -1006,6 +1015,8 @@ def _sandbox_launcher_source() -> str: import subprocess import sys +os.umask(0o077) + # Read launch config from a file (argv[1]), never the command line: provider # keys live in cfg["env"], and a shared sandbox exposes exec argv via /proc. cfg = json.loads(open(sys.argv[1], encoding="utf-8").read()) @@ -1078,7 +1089,15 @@ async def _upload_runtime_files_to_sandbox( "launch_config": f"{runtime_dir}/launch_config.json", "preflight": f"{runtime_dir}/bedrock_patch_preflight.py", } - result = await sandbox.exec(f"mkdir -p {shlex.quote(runtime_dir)}", timeout_sec=20) + quoted_runtime_dir = shlex.quote(runtime_dir) + result = await sandbox.exec( + ( + f"umask 077; mkdir -p {quoted_runtime_dir}; " + f"chown 0:0 {quoted_runtime_dir}; chmod 700 {quoted_runtime_dir}" + ), + user="root", + timeout_sec=20, + ) if result.return_code != 0: raise RuntimeError(_exec_details("prepare LiteLLM runtime directory", result)) await _upload_text( @@ -1224,10 +1243,13 @@ async def _terminate_sandbox_litellm( await sandbox.exec( f"if [ -s {shlex.quote(pid_path)} ]; then " f"kill -TERM $(cat {shlex.quote(pid_path)}) 2>/dev/null || true; fi", + user="root", timeout_sec=10, ) with contextlib.suppress(Exception): - await sandbox.exec(f"rm -rf {shlex.quote(runtime_dir)}", timeout_sec=10) + await sandbox.exec( + f"rm -rf {shlex.quote(runtime_dir)}", user="root", timeout_sec=10 + ) async def _start_sandbox_litellm( @@ -1284,13 +1306,16 @@ async def _start_sandbox_litellm( # the sandbox filesystem for the life of the run. command = ( f"rm -f {shlex.quote(paths['state'])} {shlex.quote(paths['pid'])} " - f"{shlex.quote(paths['log'])} && " + f"{shlex.quote(paths['log'])} {shlex.quote(paths['capture_state'])} && " + f"umask 077 && : > {shlex.quote(paths['log'])} && " + f"chown 0:0 {shlex.quote(paths['log'])} && " + f"chmod 600 {shlex.quote(paths['log'])} && " f"{shlex.quote(python)} {shlex.quote(paths['launcher'])} " f"{shlex.quote(paths['launch_config'])}; " f"rc=$?; rm -f {shlex.quote(paths['launch_config'])}; exit $rc" ) try: - result = await sandbox.exec(command, timeout_sec=20) + result = await sandbox.exec(command, user="root", timeout_sec=20) if result.return_code != 0: raise RuntimeError(_exec_details("start sandbox LiteLLM", result)) state = await _wait_for_sandbox_state( @@ -1854,7 +1879,7 @@ async def _provider_capture_has_verified_custody( sandbox: Any | None, runtime_dir: str | None, ) -> bool: - """Prove the agent cannot read or mutate a root-owned runtime artifact.""" + """Prove the agent cannot read or mutate the real provider capture files.""" if not sandbox_local: return True @@ -1880,32 +1905,37 @@ async def _provider_capture_has_verified_custody( if effective_uid == 0: return False - probe_path = f"{runtime_dir}/.benchflow-custody-{uuid4().hex}" - probe_value = secrets.token_hex(32) - quoted_path = shlex.quote(probe_path) - quoted_value = shlex.quote(probe_value) + quoted_runtime = shlex.quote(runtime_dir) + quoted_log = shlex.quote(f"{runtime_dir}/callback.jsonl") + quoted_state = shlex.quote(f"{runtime_dir}/capture_state.json") try: - created = await sandbox.exec( + secured = await sandbox.exec( ( - f"umask 077; printf %s {quoted_value} > {quoted_path}; " - f"chown 0:0 {quoted_path}; chmod 600 {quoted_path}" + f"test -d {quoted_runtime} && test ! -L {quoted_runtime} && " + f"test -f {quoted_log} && test ! -L {quoted_log} && " + f"test -f {quoted_state} && test ! -L {quoted_state} && " + f"chown 0:0 {quoted_runtime} {quoted_log} {quoted_state} && " + f"chmod 700 {quoted_runtime} && " + f"chmod 600 {quoted_log} {quoted_state}" ), user="root", timeout_sec=10, ) - if created.return_code != 0: - logger.warning("Provider capture custody probe creation failed") + if secured.return_code != 0: + logger.warning("Provider capture custody artifact hardening failed") return False access = await sandbox.exec( ( - f"if cat {quoted_path} >/dev/null 2>&1; then exit 0; fi; " - f"if printf compromised >> {quoted_path} 2>/dev/null; " - "then exit 0; fi; " - f"if rm -f {quoted_path} 2>/dev/null; then exit 0; fi; " + f"for artifact in {quoted_log} {quoted_state}; do " + 'if cat "$artifact" >/dev/null 2>&1; then exit 0; fi; ' + 'if [ -r "$artifact" ] || [ -w "$artifact" ]; then exit 0; fi; ' + "done; " + f"if [ -r {quoted_runtime} ] || [ -w {quoted_runtime} ] || " + f"[ -x {quoted_runtime} ]; then exit 0; fi; " "if command -v sudo >/dev/null 2>&1 && " - f"sudo -n cat {quoted_path} >/dev/null 2>&1; then exit 0; fi; " + f"sudo -n cat {quoted_state} >/dev/null 2>&1; then exit 0; fi; " "if command -v doas >/dev/null 2>&1 && " - f"doas -n cat {quoted_path} >/dev/null 2>&1; then exit 0; fi; " + f"doas -n cat {quoted_state} >/dev/null 2>&1; then exit 0; fi; " "if id -G 2>/dev/null | tr ' ' '\\n' | grep -qx 0; " "then exit 0; fi; " "effective_caps=$(awk '/^CapEff:/ {print $2}' " @@ -1929,8 +1959,9 @@ async def _provider_capture_has_verified_custody( return False verified = await sandbox.exec( ( - f'test "$(cat {quoted_path})" = {quoted_value} && ' - f"test \"$(stat -c '%u:%a' {quoted_path})\" = '0:600'" + f"test \"$(stat -c '%u:%a' {quoted_runtime})\" = '0:700' && " + f"test \"$(stat -c '%u:%a' {quoted_log})\" = '0:600' && " + f"test \"$(stat -c '%u:%a' {quoted_state})\" = '0:600'" ), user="root", timeout_sec=10, @@ -1942,9 +1973,6 @@ async def _provider_capture_has_verified_custody( except Exception as exc: logger.warning("Provider capture custody artifact probe failed: %s", exc) return False - finally: - with contextlib.suppress(Exception): - await sandbox.exec(f"rm -f {quoted_path}", user="root", timeout_sec=10) async def stop_litellm_runtime(runtime: Any | None) -> None: diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index 729cd2f66..b31bf4f39 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -295,6 +295,7 @@ def __init__( self.uploaded_modes: dict[str, str | None] = {} self.exec_calls: list[str] = [] self.exec_timeouts: list[int | None] = [] + self.exec_users: list[str | None] = [] self.fail_launch = fail_launch self.fail_preflight = fail_preflight self.log_content = log_content @@ -309,6 +310,7 @@ async def exec( ) -> _ExecResult: self.exec_calls.append(command) self.exec_timeouts.append(timeout_sec) + self.exec_users.append(user) if "stat -c %s" in command: return _ExecResult(0, stdout=str(len(self.log_content))) if "urllib.request" in command: @@ -391,6 +393,8 @@ async def test_sandbox_litellm_launch_keeps_secrets_off_command_line(): ) launch_command = next(call for call in sandbox.exec_calls if "launcher.py" in call) assert f"rc=$?; rm -f {launch_files[0]}; exit $rc" in launch_command + launch_index = sandbox.exec_calls.index(launch_command) + assert sandbox.exec_users[launch_index] == "root" assert proc.base_url == "http://127.0.0.1:45999" assert await proc.is_running() is True @@ -669,6 +673,23 @@ def test_runtime_packages_canonical_redactor_verbatim(tmp_path): assert result.returncode == 0, result.stdout + result.stderr +def test_runtime_files_are_private_even_with_permissive_umask(tmp_path): + """Guards PR #1057 against permissive host runtime-file permissions.""" + + import os + + runtime_dir = tmp_path / "runtime" + previous_umask = os.umask(0) + try: + runtime_mod._write_runtime_files(runtime_dir, config={"model_list": []}) + finally: + os.umask(previous_umask) + + assert runtime_dir.stat().st_mode & 0o777 == 0o700 + for path in runtime_dir.iterdir(): + assert path.stat().st_mode & 0o777 == 0o600, path + + def test_bedrock_patch_preflight_fails_closed_when_patch_not_loaded(tmp_path): """THE regression test for issue #602's fail-open (fixed in PR #668): when the patch never loads (sitecustomize missing from PYTHONPATH — the exact silent-failure diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 77f3969d6..640cd4b29 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -2,6 +2,7 @@ import asyncio import json +import os from datetime import datetime import pytest @@ -150,6 +151,30 @@ def test_callback_module_source_exposes_proxy_handler_instance(): assert "proxy_handler_instance = BenchFlowLiteLLMLogger()" in source +def test_callback_forces_private_artifacts_under_permissive_umask( + tmp_path, monkeypatch +): + """Guards PR #1057 against permissive umasks exposing provider capture.""" + + namespace = _callback_namespace() + runtime_dir = tmp_path / "runtime" + log_path = runtime_dir / "callback.jsonl" + state_path = runtime_dir / "capture_state.json" + monkeypatch.setenv("BENCHFLOW_LITELLM_LOG_PATH", str(log_path)) + monkeypatch.setenv("BENCHFLOW_LITELLM_CAPTURE_STATE_PATH", str(state_path)) + + previous_umask = os.umask(0) + try: + logger = namespace["BenchFlowLiteLLMLogger"]() + logger._write({"event": "failure"}) + finally: + os.umask(previous_umask) + + assert runtime_dir.stat().st_mode & 0o777 == 0o700 + assert log_path.stat().st_mode & 0o777 == 0o600 + assert state_path.stat().st_mode & 0o777 == 0o600 + + def test_callback_preserves_post_transform_provider_request_body(): """Guards PR #1057 against labeling proxy ingress as provider wire.""" diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 643bffd35..e566fb55a 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -181,20 +181,26 @@ async def exec(self, command, **kwargs): if command == "id -u -- agent": assert kwargs["user"] == "root" return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") - if command.startswith("umask 077"): + if command.startswith("test -d"): assert kwargs["user"] == "root" + assert "callback.jsonl" in command + assert "capture_state.json" in command + assert "chmod 700" in command + assert "chmod 600" in command return SimpleNamespace(return_code=0, stdout="", stderr="") if kwargs["user"] == "agent": + assert "callback.jsonl" in command + assert "capture_state.json" in command assert "sudo -n cat" in command assert "CapEff" in command assert "/var/run/docker.sock" in command return SimpleNamespace(return_code=1, stdout="", stderr="") - if command.startswith('test "$(cat'): + if command.startswith('test "$(stat -c'): assert kwargs["user"] == "root" + assert "callback.jsonl" in command + assert "capture_state.json" in command return SimpleNamespace(return_code=0, stdout="", stderr="") - assert command.startswith("rm -f") - assert kwargs["user"] == "root" - return SimpleNamespace(return_code=0, stdout="", stderr="") + raise AssertionError(f"unexpected custody command: {command}") _, provider_runtime = await ensure_litellm_runtime( agent="openhands", @@ -227,13 +233,14 @@ class PrivilegedSandbox: async def exec(self, command, **kwargs): if command == "id -u -- agent": return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") - if command.startswith("umask 077"): + if command.startswith("test -d"): return SimpleNamespace(return_code=0, stdout="", stderr="") if kwargs["user"] == "agent": + assert "callback.jsonl" in command + assert "capture_state.json" in command assert "sudo -n cat" in command return SimpleNamespace(return_code=0, stdout="", stderr="") - assert command.startswith("rm -f") - return SimpleNamespace(return_code=0, stdout="", stderr="") + raise AssertionError(f"unexpected custody command: {command}") _, provider_runtime = await ensure_litellm_runtime( agent="openhands", @@ -253,6 +260,48 @@ async def exec(self, command, **kwargs): assert provider_runtime.capture_trusted is False +@pytest.mark.asyncio +async def test_agent_access_to_real_capture_artifact_keeps_capture_audit_only( + monkeypatch, +): + """Guards PR #1057 against trusting a writable real callback artifact.""" + + async def fake_sandbox_start(**kwargs): + return FakeLiteLLMServer("http://127.0.0.1:45678", kwargs["route"]) + + monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) + + class WritableArtifactSandbox: + async def exec(self, command, **kwargs): + if command == "id -u -- agent": + return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") + if command.startswith("test -d"): + return SimpleNamespace(return_code=0, stdout="", stderr="") + if kwargs["user"] == "agent": + assert "callback.jsonl" in command + assert "capture_state.json" in command + # Exit 0 is the probe's contract for any real-artifact access. + return SimpleNamespace(return_code=0, stdout="", stderr="") + raise AssertionError(f"unexpected custody command: {command}") + + _, provider_runtime = await ensure_litellm_runtime( + agent="openhands", + agent_env={ + "AWS_BEARER_TOKEN_BEDROCK": "token", + "AWS_REGION": "us-west-2", + }, + model="aws-bedrock/us.anthropic.claude-opus-4-8", + runtime=None, + environment="daytona", + session_id="run-writable-artifact", + sandbox=WritableArtifactSandbox(), + sandbox_user="agent", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is False + + @pytest.mark.asyncio async def test_uid_zero_alias_keeps_sandbox_gateway_capture_audit_only(monkeypatch): """Guards PR #1057 against trusting a non-root name that resolves to UID 0.""" From 914e5ff1ed3f6f8ecd17891238c6a5c12f2bf4a4 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 11:49:36 -0700 Subject: [PATCH 45/74] Keep agent-visible provider credentials audit-only --- src/benchflow/providers/litellm_runtime.py | 20 ++++- .../trajectories/llm_capture_records.py | 20 ++++- tests/test_litellm_credential_custody.py | 78 +++++++++++++++++++ .../test_native_capture_resilience.py | 68 ++++++++++++++++ 4 files changed, 182 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm_credential_custody.py diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 6cc208c9e..c70bcd6dd 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1400,6 +1400,8 @@ def _provider_secret_env_names() -> set[str]: "GEMINI_API_KEY", "GOOGLE_API_KEY", "GOOGLE_GENERATIVE_AI_API_KEY", + "GOOGLE_APPLICATION_CREDENTIALS", + "GOOGLE_APPLICATION_CREDENTIALS_JSON", "AWS_BEARER_TOKEN_BEDROCK", "AZURE_API_KEY", } @@ -1409,6 +1411,15 @@ def _provider_secret_env_names() -> set[str]: return names +def _provider_credentials_have_proxy_only_custody(route: LiteLLMRoute) -> bool: + """Whether no provider credential file is deliberately exposed to the agent.""" + + from benchflow.agents.providers import PROVIDERS + + provider = PROVIDERS.get(route.provider_name) + return provider is None or not provider.credential_files + + def _provider_model_id(entry: object) -> str | None: if not isinstance(entry, Mapping): return None @@ -1765,7 +1776,8 @@ async def ensure_litellm_runtime( sorted(set(required_skill_names)), separators=(",", ":") ) proxy_location = "sandbox" if sandbox_local else "host" - capture_trusted = not sandbox_local + credentials_isolated = _provider_credentials_have_proxy_only_custody(route) + capture_trusted = not sandbox_local and credentials_isolated config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" f"{session_id}:{role_name or 'primary'}:{sandbox_user or 'root'}:" @@ -1777,12 +1789,13 @@ async def ensure_litellm_runtime( is_running = await server.is_running() if is_running: if sandbox_local: - capture_trusted = await _provider_capture_has_verified_custody( + artifact_custody = await _provider_capture_has_verified_custody( sandbox_local=True, sandbox_user=sandbox_user, sandbox=sandbox, runtime_dir=getattr(server, "runtime_dir", None), ) + capture_trusted = credentials_isolated and artifact_custody runtime.capture_trusted = bool( getattr(runtime, "capture_trusted", False) and capture_trusted ) @@ -1840,12 +1853,13 @@ async def ensure_litellm_runtime( ) if sandbox_local: - capture_trusted = await _provider_capture_has_verified_custody( + artifact_custody = await _provider_capture_has_verified_custody( sandbox_local=True, sandbox_user=sandbox_user, sandbox=sandbox, runtime_dir=getattr(server, "runtime_dir", None), ) + capture_trusted = credentials_isolated and artifact_custody from benchflow.providers.runtime import ProviderRuntime diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index 70f8e1c64..e22075461 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -60,6 +60,16 @@ class CaptureAssembly: role_captures: list[LLMRoleCapture] +def _untrusted_provider_custody(target: CaptureTarget | None) -> str: + if target is not None and target.model: + from benchflow.agents.providers import find_provider + + provider = find_provider(target.model) + if provider is not None and provider[1].credential_files: + return "agent_accessible_provider_credentials" + return "agent_writable_sandbox" + + def load_provider_wire_records( path: Path, *, @@ -146,7 +156,7 @@ def load_provider_wire_records( } ) if not capture_trusted: - metadata["capture_custody"] = "agent_writable_sandbox" + metadata["capture_custody"] = _untrusted_provider_custody(target) if not attribution_complete: metadata["role_candidates"] = _role_candidates(targets) records.append(redact_trajectory_obj(record)) @@ -176,6 +186,14 @@ def assemble_capture( errors.append( "sandbox-local LiteLLM capture shared root custody with the agent" ) + if any( + _record_metadata(record).get("capture_custody") + == "agent_accessible_provider_credentials" + for record in provider_records + ): + errors.append( + "provider credentials were available to the agent outside LiteLLM" + ) attribution_incomplete = any( _record_metadata(record).get("role_attribution_complete") is False for record in records diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py new file mode 100644 index 000000000..8cb5e0116 --- /dev/null +++ b/tests/test_litellm_credential_custody.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from benchflow.providers import litellm_runtime as runtime_mod +from benchflow.providers.runtime import ensure_litellm_runtime + + +@pytest.mark.asyncio +async def test_vertex_adc_provider_capture_remains_audit_only_on_host(monkeypatch): + """Guards PR #1057 against trusting agent-accessible Vertex credentials.""" + + async def fake_start(**_kwargs): + return SimpleNamespace(base_url="http://127.0.0.1:4000") + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + + updated, provider_runtime = await ensure_litellm_runtime( + agent="codex-acp", + agent_env={ + "GOOGLE_APPLICATION_CREDENTIALS_JSON": '{"type":"authorized_user"}', + "GOOGLE_APPLICATION_CREDENTIALS": ( + "/home/agent/.config/gcloud/application_default_credentials.json" + ), + "GOOGLE_CLOUD_PROJECT": "project", + "GOOGLE_CLOUD_LOCATION": "global", + }, + model="google-vertex/gemini-2.5-flash", + runtime=None, + environment="docker", + session_id="run-vertex-adc", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is False + assert "GOOGLE_APPLICATION_CREDENTIALS_JSON" not in updated + assert "GOOGLE_APPLICATION_CREDENTIALS" not in updated + + +@pytest.mark.asyncio +async def test_vertex_adc_stays_audit_only_with_verified_sandbox_artifacts(monkeypatch): + """Guards PR #1057 against equating file custody with ADC isolation.""" + + async def fake_start(**_kwargs): + return SimpleNamespace( + base_url="http://127.0.0.1:4000", + runtime_dir="/tmp/benchflow-litellm/test-runtime", + ) + + async def fake_artifact_custody(**_kwargs): + return True + + monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_start) + monkeypatch.setattr( + runtime_mod, + "_provider_capture_has_verified_custody", + fake_artifact_custody, + ) + + _, provider_runtime = await ensure_litellm_runtime( + agent="codex-acp", + agent_env={ + "GOOGLE_APPLICATION_CREDENTIALS_JSON": '{"type":"authorized_user"}', + "GOOGLE_CLOUD_PROJECT": "project", + "GOOGLE_CLOUD_LOCATION": "global", + }, + model="google-vertex/gemini-2.5-flash", + runtime=None, + environment="daytona", + session_id="run-vertex-adc-sandbox", + sandbox=SimpleNamespace(), + sandbox_user="agent", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is False diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 07d616248..dc4416e81 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -396,6 +396,74 @@ async def test_root_sandbox_provider_capture_is_retained_but_audit_only( assert any("shared root custody" in error for error in capture.manifest.errors) +@pytest.mark.asyncio +async def test_vertex_provider_capture_is_audit_only_when_agent_has_adc( + tmp_path: Path, +) -> None: + """Guards PR #1057 against training on bypassable Vertex ADC capture.""" + + model = "google-vertex/gemini-2.5-flash" + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model=model, + session_id="rollout-vertex", + started_at=datetime(2026, 8, 29, 12, 0, tzinfo=UTC), + ) + await capture.prepare_agent( + None, + agent="codex-acp", + model=model, + agent_env={"GOOGLE_APPLICATION_CREDENTIALS_JSON": "test-adc"}, + credential_home="/home/agent", + sandbox_user="agent", + ) + capture.bind_provider_capture_trust( + agent="codex-acp", + model=model, + credential_home="/home/agent", + trusted=False, + ) + capture.trajectory_path.write_text( + json.dumps( + { + "request": {"body": {"model": model, "contents": []}}, + "response": { + "status_code": 200, + "body": { + "candidates": [{"content": {"parts": [{"text": "done"}]}}], + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 1, + }, + }, + }, + "metadata": { + "benchflow_agent": "codex-acp", + "benchflow_role": "primary", + "benchflow_requested_model": model, + "request_complete": True, + "response_complete": True, + "request_capture_source": ( + "litellm_pre_api_call_complete_input_dict" + ), + }, + } + ) + + "\n" + ) + + await capture.finalize(None, acp_events=[], model_call_seen=True) + + record = json.loads(capture.trajectory_path.read_text()) + assert record["metadata"]["capture_fidelity"] == "agent_session" + assert record["metadata"]["capture_custody"] == ( + "agent_accessible_provider_credentials" + ) + assert capture.manifest.status is CaptureStatus.PARTIAL + assert any("outside LiteLLM" in error for error in capture.manifest.errors) + + @pytest.mark.asyncio async def test_malformed_provider_capture_preserves_native_evidence_and_cleanup( tmp_path: Path, From 3d6d1986901bcab821fadea1441f74bd18f3486b Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 12:18:30 -0700 Subject: [PATCH 46/74] Strip alternate credentials from proxy agents --- src/benchflow/providers/litellm_runtime.py | 55 +++++++++++++++++-- tests/test_litellm_credential_custody.py | 64 ++++++++++++++++++++++ 2 files changed, 115 insertions(+), 4 deletions(-) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index c70bcd6dd..ed4263e2b 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1389,8 +1389,36 @@ def _missing_required_env(route: LiteLLMRoute, env: dict[str, str]) -> list[str] return missing -def _provider_secret_env_names() -> set[str]: - """Upstream provider credentials the proxy owns and the agent must not see.""" +_PROVIDER_CREDENTIAL_ENV_EXACT = frozenset( + { + "GOOGLE_APPLICATION_CREDENTIALS", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "CODEX_AUTH_JSON", + } +) +_PROVIDER_CREDENTIAL_ENV_SUFFIXES = ( + "_API_KEY", + "_AUTH_TOKEN", + "_ACCESS_TOKEN", + "_OAUTH_TOKEN", + "_BEARER_TOKEN", + "_AUTH_JSON", + "_CREDENTIALS_JSON", + "_SECRET_ACCESS_KEY", + "_SESSION_TOKEN", +) + + +def _looks_like_provider_credential_env(name: str) -> bool: + return name in _PROVIDER_CREDENTIAL_ENV_EXACT or name.endswith( + _PROVIDER_CREDENTIAL_ENV_SUFFIXES + ) + + +def _provider_secret_env_names(env: Mapping[str, object] | None = None) -> set[str]: + """Upstream provider credential aliases the agent must never receive.""" from benchflow.agents.providers import PROVIDERS names = { @@ -1404,10 +1432,29 @@ def _provider_secret_env_names() -> set[str]: "GOOGLE_APPLICATION_CREDENTIALS_JSON", "AWS_BEARER_TOKEN_BEDROCK", "AZURE_API_KEY", + "CODEX_API_KEY", + "CODEX_ACCESS_TOKEN", + "CODEX_AUTH_JSON", + "CLAUDE_CODE_OAUTH_TOKEN", + "CLAUDE_OAUTH_TOKEN", + "LLM_API_KEY", + "BENCHFLOW_PROVIDER_API_KEY", } for cfg in PROVIDERS.values(): if cfg.auth_env: names.add(cfg.auth_env) + for cfg in AGENTS.values(): + names.update(cfg.requires_env) + names.update(file.env_source for file in cfg.credential_files) + names.update( + destination + for source, destination in cfg.env_mapping.items() + if source == "BENCHFLOW_PROVIDER_API_KEY" + ) + if cfg.subscription_auth is not None: + names.add(cfg.subscription_auth.replaces_env) + if env is not None: + names.update(name for name in env if _looks_like_provider_credential_env(name)) return names @@ -1502,7 +1549,7 @@ def _assert_proxy_isolated(agent: str, env: dict[str, str], *, master_key: str) """ leaked = sorted( name - for name in _provider_secret_env_names() + for name in _provider_secret_env_names(env) if env.get(name) and env.get(name) != master_key ) if leaked: @@ -1563,7 +1610,7 @@ def _wire_litellm_agent_env( # proxy process holds them via its own env. (In sandbox-local mode the proxy # shares the agent's sandbox, so this reduces — but cannot fully remove — key # visibility.) - for secret_key in _provider_secret_env_names(): + for secret_key in _provider_secret_env_names(updated): updated.pop(secret_key, None) for endpoint_key in _PROVIDER_ENDPOINT_ENV_NAMES: updated.pop(endpoint_key, None) diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 8cb5e0116..5463b43f7 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -1,10 +1,12 @@ from __future__ import annotations +import json from types import SimpleNamespace import pytest from benchflow.providers import litellm_runtime as runtime_mod +from benchflow.providers.litellm_config import resolve_litellm_route from benchflow.providers.runtime import ensure_litellm_runtime @@ -76,3 +78,65 @@ async def fake_artifact_custody(**_kwargs): assert provider_runtime is not None assert provider_runtime.capture_trusted is False + + +def test_proxy_strips_every_supported_alternate_credential_alias(): + """Guards PR #1057 against alternate credentials bypassing the proxy.""" + + master_key = "sk-benchflow-master" + route = resolve_litellm_route( + "openai/gpt-5.6-luna", + {"OPENAI_API_KEY": "sk-provider"}, + ) + raw_credentials = { + "OPENAI_API_KEY": "sk-provider", + "CODEX_API_KEY": "sk-codex", + "CODEX_ACCESS_TOKEN": "codex-access", + "CODEX_AUTH_JSON": '{"tokens":{"access_token":"codex-json"}}', + "CLAUDE_CODE_OAUTH_TOKEN": "claude-code-oauth", + "CLAUDE_OAUTH_TOKEN": "claude-oauth", + "ANTHROPIC_AUTH_TOKEN": "anthropic-auth", + "AWS_ACCESS_KEY_ID": "aws-access", + "AWS_SECRET_ACCESS_KEY": "aws-secret", + "AWS_SESSION_TOKEN": "aws-session", + "CUSTOM_API_KEY": "custom-provider-key", + } + + agent_env = { + **raw_credentials, + "DEFAULT_AUTH_REQUEST": json.dumps( + { + "methodId": "api-key", + "_meta": {"api-key": {"apiKey": raw_credentials["CODEX_API_KEY"]}}, + } + ), + } + updated = runtime_mod._wire_litellm_agent_env( + agent="codex-acp", + agent_env=agent_env, + route=route, + base_url="http://127.0.0.1:4000", + master_key=master_key, + ) + + for name, value in raw_credentials.items(): + assert updated.get(name) != value, name + assert value not in json.dumps(updated), name + assert updated["OPENAI_API_KEY"] == master_key + assert updated["BENCHFLOW_PROVIDER_API_KEY"] == master_key + runtime_mod._assert_proxy_isolated( + "codex-acp", + updated, + master_key=master_key, + ) + + +def test_proxy_isolation_guard_detects_unregistered_api_key_alias(): + """Guards PR #1057 against custom API-key aliases escaping the guard.""" + + with pytest.raises(RuntimeError, match="CUSTOM_API_KEY"): + runtime_mod._assert_proxy_isolated( + "custom-agent", + {"CUSTOM_API_KEY": "raw-provider-key"}, + master_key="sk-benchflow-master", + ) From 142910a44fdbce17e94eae92d48054f80c9d660f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 12:52:58 -0700 Subject: [PATCH 47/74] Harden proxy credential and continuation journals --- src/benchflow/agents/codex_config.py | 5 +- src/benchflow/continue_run/sandbox_proxy.py | 25 ++++++-- src/benchflow/providers/litellm_runtime.py | 8 +-- src/benchflow/trajectories/redaction.py | 11 ++++ tests/continue_run/test_replay_proxy.py | 68 +++++++++++++++++++++ tests/test_litellm_credential_custody.py | 40 ++++++++++++ 6 files changed, 145 insertions(+), 12 deletions(-) diff --git a/src/benchflow/agents/codex_config.py b/src/benchflow/agents/codex_config.py index 7dc27e69f..1ee9b1279 100644 --- a/src/benchflow/agents/codex_config.py +++ b/src/benchflow/agents/codex_config.py @@ -55,7 +55,10 @@ def apply_codex_provider_config( provider = dict(provider) if isinstance(provider, dict) else {} provider.setdefault("name", provider_name) provider["base_url"] = base_url - provider.setdefault("env_key", "OPENAI_API_KEY") + if strict: + provider["env_key"] = "OPENAI_API_KEY" + else: + provider.setdefault("env_key", "OPENAI_API_KEY") provider.setdefault("wire_api", "responses") provider.setdefault("supports_websockets", False) diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 87d2fd2d8..38b92371d 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -23,6 +23,7 @@ from typing import Any from uuid import uuid4 +from benchflow.trajectories.redaction import canonical_redaction_module_source from benchflow.trajectories.types import LLMExchange SANDBOX_REPLAY_ROOT = "/tmp/benchflow-replay" @@ -50,6 +51,11 @@ def _sandbox_proxy_source() -> str: from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +try: + from benchflow_trajectory_redaction import redact_trajectory_obj +except ImportError: + from benchflow.trajectories.redaction import redact_trajectory_obj + CAPTURE_STATE_WRITE_FAILED = "BENCHFLOW_CAPTURE_STATE_WRITE_FAILED" @@ -289,11 +295,13 @@ def _forward_live(self, request_body): return 500, {"error": {"message": str(exc)}}, False def _append_live_exchange(self, request_body, status, body, live_attempt): - row = { - "request": {"body": request_body}, - "response": {"status_code": status, "body": body}, - "metadata": {"continuation_attempt": live_attempt}, - } + row = redact_trajectory_obj( + { + "request": {"body": request_body}, + "response": {"status_code": status, "body": body}, + "metadata": {"continuation_attempt": live_attempt}, + } + ) with open(self.live_log_path, "a", encoding="utf-8") as handle: handle.write(json.dumps(row) + "\n") handle.flush() @@ -540,6 +548,7 @@ async def start( runtime_dir = f"{SANDBOX_REPLAY_ROOT}/{token}" paths = { "script": f"{runtime_dir}/replay_proxy.py", + "redaction": f"{runtime_dir}/benchflow_trajectory_redaction.py", "config": f"{runtime_dir}/config.json", "state": f"{runtime_dir}/state.json", "pid": f"{runtime_dir}/replay.pid", @@ -572,6 +581,12 @@ async def start( "live_log_path": paths["live_log"], } await _upload_text(sandbox, _sandbox_proxy_source(), paths["script"], ".py") + await _upload_text( + sandbox, + canonical_redaction_module_source(), + paths["redaction"], + ".py", + ) await _upload_text(sandbox, json.dumps(config), paths["config"], ".json") command = ( diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index ed4263e2b..618710c5f 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -18,7 +18,6 @@ import tempfile from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass -from importlib.resources import files from pathlib import Path from typing import TYPE_CHECKING, Any, NoReturn, cast from uuid import uuid4 @@ -59,6 +58,7 @@ ) from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter +from benchflow.trajectories.redaction import canonical_redaction_module_source from benchflow.trajectories.types import Trajectory from benchflow.usage_tracking import UsageTrackingConfig, usage_unavailable @@ -850,11 +850,7 @@ def _write_runtime_files( def _redaction_module_source() -> str: """Read the packaged canonical redactor for isolated proxy runtimes.""" - return ( - files("benchflow.trajectories") - .joinpath("redaction.py") - .read_text(encoding="utf-8") - ) + return canonical_redaction_module_source() # How long to wait for the *host* per-run LiteLLM proxy to become healthy. diff --git a/src/benchflow/trajectories/redaction.py b/src/benchflow/trajectories/redaction.py index 29b27b078..ceffa1ff1 100644 --- a/src/benchflow/trajectories/redaction.py +++ b/src/benchflow/trajectories/redaction.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from importlib.resources import files from typing import Any # Human-facing redaction categories. Every canonical pattern below is tagged @@ -293,6 +294,16 @@ ] +def canonical_redaction_module_source() -> str: + """Read this packaged redactor for an isolated stdlib-only runtime.""" + + return ( + files("benchflow.trajectories") + .joinpath("redaction.py") + .read_text(encoding="utf-8") + ) + + def redact_trajectory_text(text: str) -> str: """Apply all secret-redaction patterns to *text*. diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index a1214a8a0..cbc16572b 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -4,6 +4,7 @@ import json import threading +from types import SimpleNamespace import httpx import pytest @@ -19,6 +20,7 @@ _ordered_live_exchange_log, _sandbox_proxy_source, ) +from benchflow.trajectories.redaction import canonical_redaction_module_source from ._helpers import completion, exchange @@ -154,6 +156,72 @@ def fail(*_args, **_kwargs): assert capture_state["live_error_count"] == 1 +def test_sandbox_live_exchange_is_redacted_before_journaling(tmp_path) -> None: + """Guards PR #1057 against persisting raw continuation secrets.""" + + namespace: dict[str, object] = {} + exec(_sandbox_proxy_source(), namespace) + live_log = tmp_path / "live.jsonl" + state = namespace["ReplayState"]( + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(live_log), + state_path=str(tmp_path / "state.json"), + port=61357, + ) + request_secret = "sk-ant-api03-requestsecret123456" + response_secret = "sk-ant-api03-responsesecret123456" + state._forward_live = lambda _request: ( + 401, + {"error": {"message": response_secret}}, + True, + ) + + _source, status, response = state.next_response( + {"messages": [{"role": "user", "content": request_secret}]} + ) + + assert status == 401 + assert response["error"]["message"] == response_secret + raw = live_log.read_text() + assert request_secret not in raw + assert response_secret not in raw + assert raw.count("***REDACTED***") == 2 + + +@pytest.mark.asyncio +async def test_sandbox_replay_uploads_canonical_redactor() -> None: + """Guards PR #1057 against launching continuation without its redactor.""" + + class FakeSandbox: + def __init__(self) -> None: + self.uploaded: dict[str, str] = {} + + async def exec(self, _command, **_kwargs): + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def upload_file(self, source, target): + self.uploaded[target] = source.read_text() + + sandbox = FakeSandbox() + proxy = await SandboxReplayProxy.start( + sandbox=sandbox, + recorded=[], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + ) + + redaction_path = f"{proxy.runtime_dir}/benchflow_trajectory_redaction.py" + assert sandbox.uploaded[redaction_path] == canonical_redaction_module_source() + assert ( + "from benchflow_trajectory_redaction import" + in sandbox.uploaded[f"{proxy.runtime_dir}/replay_proxy.py"] + ) + + def test_sandbox_attempt_journal_failure_invalidates_stale_state( tmp_path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 5463b43f7..650d5bf85 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -131,6 +131,46 @@ def test_proxy_strips_every_supported_alternate_credential_alias(): ) +def test_proxy_rebinds_custom_codex_provider_to_master_key(): + """Guards PR #1057 against stripping a custom Codex provider's env key.""" + + master_key = "sk-benchflow-master" + route = resolve_litellm_route( + "openai/gpt-5.6-luna", + {"OPENAI_API_KEY": "sk-provider"}, + ) + updated = runtime_mod._wire_litellm_agent_env( + agent="codex-acp", + agent_env={ + "CODEX_API_KEY": "sk-codex-provider", + "CODEX_CONFIG": json.dumps( + { + "model_provider": "custom", + "model_providers": { + "custom": { + "env_key": "CODEX_API_KEY", + "wire_api": "responses", + } + }, + } + ), + }, + route=route, + base_url="http://127.0.0.1:4000", + master_key=master_key, + ) + + provider = json.loads(updated["CODEX_CONFIG"])["model_providers"]["custom"] + assert provider["env_key"] == "OPENAI_API_KEY" + assert updated["OPENAI_API_KEY"] == master_key + assert "CODEX_API_KEY" not in updated + runtime_mod._assert_proxy_isolated( + "codex-acp", + updated, + master_key=master_key, + ) + + def test_proxy_isolation_guard_detects_unregistered_api_key_alias(): """Guards PR #1057 against custom API-key aliases escaping the guard.""" From 0aa771ce7b6fc6ad297f7e06246e453ff583fbe0 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 13:21:12 -0700 Subject: [PATCH 48/74] Sanitize Codex config and package redactor data --- hatch_build.py | 44 ++++++++++++++++++++++++ pyproject.toml | 4 +++ src/benchflow/agents/codex_config.py | 39 +++++++++++++-------- src/benchflow/trajectories/redaction.py | 11 +++--- tests/test_litellm_credential_custody.py | 28 ++++++++++++--- tests/test_redaction_resource.py | 21 +++++++++++ 6 files changed, 125 insertions(+), 22 deletions(-) create mode 100644 hatch_build.py create mode 100644 tests/test_redaction_resource.py diff --git a/hatch_build.py b/hatch_build.py new file mode 100644 index 000000000..f19c974a6 --- /dev/null +++ b/hatch_build.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Any + +from hatchling.builders.hooks.plugin.interface import BuildHookInterface + +_REDACTOR_SOURCE = "src/benchflow/trajectories/redaction.py" +_REDACTOR_RESOURCE = "benchflow/trajectories/resources/canonical_redaction.py.txt" + + +class CustomBuildHook(BuildHookInterface): + PLUGIN_NAME = "custom" + + def initialize(self, _version: str, build_data: dict[str, Any]) -> None: + if self.target_name != "wheel": + return + source = Path(self.root, _REDACTOR_SOURCE) + with tempfile.NamedTemporaryFile( + prefix="benchflow-redactor-", + suffix=".txt", + delete=False, + ) as generated: + generated.write(source.read_bytes()) + self._generated_path = Path(generated.name) + build_data["force_include"][str(self._generated_path)] = _REDACTOR_RESOURCE + + def finalize( + self, + _version: str, + _build_data: dict[str, Any], + _artifact_path: str, + ) -> None: + self._discard_generated() + + def clean(self, _versions: list[str]) -> None: + self._discard_generated() + + def _discard_generated(self) -> None: + generated = getattr(self, "_generated_path", None) + if generated is not None: + generated.unlink(missing_ok=True) + self._generated_path = None diff --git a/pyproject.toml b/pyproject.toml index 2c1753f80..78c172bd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -146,12 +146,16 @@ build-backend = "hatchling.build" only-include = [ "src", "tests", + "hatch_build.py", "README.md", "CHANGELOG.md", "LICENSE", "pyproject.toml", ] +[tool.hatch.build.hooks.custom] +path = "hatch_build.py" + [tool.pytest.ini_options] asyncio_mode = "auto" addopts = "-m 'not live and not integration'" diff --git a/src/benchflow/agents/codex_config.py b/src/benchflow/agents/codex_config.py index 1ee9b1279..70ab85677 100644 --- a/src/benchflow/agents/codex_config.py +++ b/src/benchflow/agents/codex_config.py @@ -44,23 +44,34 @@ def apply_codex_provider_config( raise ValueError(f"{CODEX_CONFIG_ENV} must decode to a JSON object") return - provider_id = ( - agent_env.get(CODEX_MODEL_PROVIDER_ENV) - or config.get("model_provider") - or codex_provider_id(provider_name) - ) - providers = config.get("model_providers") - providers = {} if not isinstance(providers, dict) else dict(providers) - provider = providers.get(provider_id) - provider = dict(provider) if isinstance(provider, dict) else {} - provider.setdefault("name", provider_name) - provider["base_url"] = base_url if strict: - provider["env_key"] = "OPENAI_API_KEY" + provider_id = codex_provider_id(provider_name) + config = {} + providers: dict[str, Any] = {} + provider: dict[str, Any] = { + "name": provider_name, + "base_url": base_url, + "env_key": "OPENAI_API_KEY", + "wire_api": "responses", + "supports_websockets": False, + } else: + provider_id = ( + agent_env.get(CODEX_MODEL_PROVIDER_ENV) + or config.get("model_provider") + or codex_provider_id(provider_name) + ) + providers_value = config.get("model_providers") + providers = ( + {} if not isinstance(providers_value, dict) else dict(providers_value) + ) + provider_value = providers.get(provider_id) + provider = dict(provider_value) if isinstance(provider_value, dict) else {} + provider.setdefault("name", provider_name) + provider["base_url"] = base_url provider.setdefault("env_key", "OPENAI_API_KEY") - provider.setdefault("wire_api", "responses") - provider.setdefault("supports_websockets", False) + provider.setdefault("wire_api", "responses") + provider.setdefault("supports_websockets", False) providers[provider_id] = provider config["model_providers"] = providers diff --git a/src/benchflow/trajectories/redaction.py b/src/benchflow/trajectories/redaction.py index ceffa1ff1..acabc783b 100644 --- a/src/benchflow/trajectories/redaction.py +++ b/src/benchflow/trajectories/redaction.py @@ -297,11 +297,14 @@ def canonical_redaction_module_source() -> str: """Read this packaged redactor for an isolated stdlib-only runtime.""" - return ( - files("benchflow.trajectories") - .joinpath("redaction.py") - .read_text(encoding="utf-8") + package = files("benchflow.trajectories") + resource = package.joinpath( + "resources", + "canonical_redaction.py.txt", ) + if resource.is_file(): + return resource.read_text(encoding="utf-8") + return package.joinpath("redaction.py").read_text(encoding="utf-8") def redact_trajectory_text(text: str) -> str: diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 650d5bf85..ede7d9b70 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -132,9 +132,10 @@ def test_proxy_strips_every_supported_alternate_credential_alias(): def test_proxy_rebinds_custom_codex_provider_to_master_key(): - """Guards PR #1057 against stripping a custom Codex provider's env key.""" + """Guards PR #1057 against retaining custom Codex provider credentials.""" master_key = "sk-benchflow-master" + literal_key = "literal-provider-credential" route = resolve_litellm_route( "openai/gpt-5.6-luna", {"OPENAI_API_KEY": "sk-provider"}, @@ -145,12 +146,18 @@ def test_proxy_rebinds_custom_codex_provider_to_master_key(): "CODEX_API_KEY": "sk-codex-provider", "CODEX_CONFIG": json.dumps( { + "top_level_secret": literal_key, "model_provider": "custom", "model_providers": { "custom": { "env_key": "CODEX_API_KEY", "wire_api": "responses", - } + "http_headers": {"Authorization": literal_key}, + }, + "unused": { + "env_key": "UNUSED_API_KEY", + "http_headers": {"Authorization": literal_key}, + }, }, } ), @@ -160,10 +167,23 @@ def test_proxy_rebinds_custom_codex_provider_to_master_key(): master_key=master_key, ) - provider = json.loads(updated["CODEX_CONFIG"])["model_providers"]["custom"] - assert provider["env_key"] == "OPENAI_API_KEY" + config = json.loads(updated["CODEX_CONFIG"]) + assert config == { + "model_providers": { + "benchflow-litellm": { + "name": "litellm", + "base_url": "http://127.0.0.1:4000/v1", + "env_key": "OPENAI_API_KEY", + "wire_api": "responses", + "supports_websockets": False, + } + }, + "model_provider": "benchflow-litellm", + "model": route.model_alias, + } assert updated["OPENAI_API_KEY"] == master_key assert "CODEX_API_KEY" not in updated + assert literal_key not in json.dumps(updated) runtime_mod._assert_proxy_isolated( "codex-acp", updated, diff --git a/tests/test_redaction_resource.py b/tests/test_redaction_resource.py new file mode 100644 index 000000000..e538c0743 --- /dev/null +++ b/tests/test_redaction_resource.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import tomllib +from pathlib import Path + +from benchflow.trajectories.redaction import canonical_redaction_module_source + + +def test_wheel_packages_canonical_redactor_as_data() -> None: + """Guards PR #1057 against source-stripped installs losing the redactor.""" + + root = Path(__file__).parents[1] + config = tomllib.loads((root / "pyproject.toml").read_text()) + build = config["tool"]["hatch"]["build"] + + assert build["hooks"]["custom"] == {"path": "hatch_build.py"} + assert "hatch_build.py" in build["targets"]["sdist"]["only-include"] + assert ( + canonical_redaction_module_source() + == (root / "src/benchflow/trajectories/redaction.py").read_text() + ) From e139e62e372c35636d73ee29b0b364f63e4d5908 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 13:55:49 -0700 Subject: [PATCH 49/74] Preserve direct Codex provider settings --- src/benchflow/agents/codex_config.py | 15 ++++-- src/benchflow/providers/litellm_runtime.py | 38 +++++++++++++- tests/examples/test_codex_custom_provider.sh | 31 ++++++----- tests/test_litellm_runtime.py | 34 ++++++++++++ tests/test_resolve_env_helpers.py | 54 ++++++++++++++++++++ 5 files changed, 152 insertions(+), 20 deletions(-) diff --git a/src/benchflow/agents/codex_config.py b/src/benchflow/agents/codex_config.py index 70ab85677..8a31e1fd1 100644 --- a/src/benchflow/agents/codex_config.py +++ b/src/benchflow/agents/codex_config.py @@ -27,8 +27,14 @@ def apply_codex_provider_config( model: str | None, provider_name: str, strict: bool = False, + isolate: bool = False, ) -> None: - """Create or update Codex's model provider entry in ``agent_env``.""" + """Create or update Codex's model provider entry in ``agent_env``. + + ``isolate`` discards caller config and is reserved for a BenchFlow-owned + proxy boundary. Direct provider configuration keeps unrelated caller + settings while ``strict`` still validates JSON and pins the credential key. + """ raw_config = agent_env.get(CODEX_CONFIG_ENV) if not raw_config: config: dict[str, Any] = {} @@ -44,7 +50,7 @@ def apply_codex_provider_config( raise ValueError(f"{CODEX_CONFIG_ENV} must decode to a JSON object") return - if strict: + if isolate: provider_id = codex_provider_id(provider_name) config = {} providers: dict[str, Any] = {} @@ -69,7 +75,10 @@ def apply_codex_provider_config( provider = dict(provider_value) if isinstance(provider_value, dict) else {} provider.setdefault("name", provider_name) provider["base_url"] = base_url - provider.setdefault("env_key", "OPENAI_API_KEY") + if strict: + provider["env_key"] = "OPENAI_API_KEY" + else: + provider.setdefault("env_key", "OPENAI_API_KEY") provider.setdefault("wire_api", "responses") provider.setdefault("supports_websockets", False) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 618710c5f..82a052b63 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -17,9 +17,10 @@ import sys import tempfile from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import TYPE_CHECKING, Any, NoReturn, cast +from urllib.parse import urlsplit, urlunsplit from uuid import uuid4 import httpx @@ -84,6 +85,9 @@ # it's ignored) and `NO_DOCS=true` so the docs route is skipped regardless of the # inherited environment. _PROXY_DOCS_DISABLE_ENV = {"DOCS_URL": "", "NO_DOCS": "true"} +_CONTAINER_HOST_ALIASES = frozenset( + {"host.docker.internal", "gateway.docker.internal", "host.containers.internal"} +) _SKILL_CATALOG_GATE_AGENT_ENV = "BENCHFLOW_SKILL_CATALOG_GATE_AGENT" _REQUIRED_SKILL_NAMES_ENV = "BENCHFLOW_REQUIRED_SKILL_NAMES_JSON" # Live callback-log capture budgets. The reader tails the gateway's @@ -804,6 +808,35 @@ def _host_bind_address(environment: str) -> str: return address +def _route_for_host_proxy(route: LiteLLMRoute, environment: str) -> LiteLLMRoute: + """Translate container-view host aliases for a host-owned proxy. + + Before provider capture became mandatory, a Docker agent used names such as + ``host.docker.internal`` to reach a provider running on the host. The + always-on host LiteLLM process must instead use host loopback; those Docker + DNS aliases are often deliberately undefined in the host resolver. + """ + + if environment != "docker": + return route + api_base = route.litellm_params.get("api_base") + if not isinstance(api_base, str): + return route + parsed = urlsplit(api_base) + if ( + parsed.hostname not in _CONTAINER_HOST_ALIASES + or parsed.username is not None + or parsed.password is not None + ): + return route + netloc = "127.0.0.1" + if parsed.port is not None: + netloc = f"{netloc}:{parsed.port}" + params = dict(route.litellm_params) + params["api_base"] = urlunsplit(parsed._replace(netloc=netloc)) + return replace(route, litellm_params=params) + + def _agent_endpoint_for_environment( port: int, environment: str, bind: str ) -> LiteLLMEndpoint: @@ -1650,6 +1683,7 @@ def _wire_litellm_agent_env( model=route.model_alias, provider_name="litellm", strict=True, + isolate=True, ) return updated if agent == "opencode": @@ -1798,6 +1832,8 @@ async def ensure_litellm_runtime( "directly — register the provider/model or fix the route." ), ) + if not sandbox_local: + route = _route_for_host_proxy(route, environment) missing = _missing_required_env(route, agent_env) if missing: missing_text = ", ".join(missing) diff --git a/tests/examples/test_codex_custom_provider.sh b/tests/examples/test_codex_custom_provider.sh index c6c8ec63e..76034f57e 100644 --- a/tests/examples/test_codex_custom_provider.sh +++ b/tests/examples/test_codex_custom_provider.sh @@ -5,17 +5,16 @@ # to fail with a mocked 401. The test passes if codex-acp sends a request to # the local stub instead of api.openai.com. # -# Two deliberate choices keep this a *direct* routing check: -# --usage-tracking off — with the default (auto), benchflow stands up its -# own LiteLLM usage proxy in front of the agent (#587/#613). That proxy -# would intercept codex's call, fall back to codex's built-in default -# model, and 400 before forwarding upstream — so the stub would never -# see the request. Off sends provider traffic straight to the stub. -# (Usage-proxy forwarding for custom providers is a separate concern.) -# --model vllm/gpt-5.4 — codex-acp validates the model against its own +# Two deliberate choices isolate custom-provider routing: +# --usage-tracking off — disables the usage requirement. BenchFlow still +# routes API-key calls through its trajectory-capture proxy (PR #1057), +# so reaching the stub proves the configured provider remains the +# upstream destination after capture wiring. +# --model vllm/gpt-5.6-luna — codex-acp validates the model against its own # catalog at session/set_model, so a synthetic id like "mock-model" is -# rejected with -32603 before any HTTP request. gpt-5.4 is a real catalog -# id that is accepted but is NOT codex's built-in default (gpt-5.5), so +# rejected with -32603 before any HTTP request. gpt-5.6-luna is a real +# catalog id that is accepted but is NOT codex's built-in default +# (gpt-5.5), so # the model assertion below also proves codex sent the *configured* model # rather than silently falling back. The stub returns a mocked 401. # @@ -95,7 +94,7 @@ env -u CODEX_ACCESS_TOKEN -u CODEX_API_KEY -u OPENAI_BASE_URL -u OPENAI_API_KEY uv run bench eval run \ --tasks-dir "$TASK" \ --agent codex-acp \ - --model vllm/gpt-5.4 \ + --model vllm/gpt-5.6-luna \ --sandbox docker \ --usage-tracking off \ --jobs-dir "$JOBS_DIR" \ @@ -126,11 +125,11 @@ fi # The configured model must reach the wire. If codex silently fell back to its # built-in default (gpt-5.5 — the failure mode when the usage proxy intercepts), -# the request body's model field would not be gpt-5.4. Match the escaped model -# field exactly; codex mentions "gpt-5.5" in its system prompt, so a loose grep -# would give a false pass. -if ! grep -q '\\"model\\":\\"gpt-5\.4\\"' "$LOG_FILE"; then - echo "FAIL: stub did not receive the configured model gpt-5.4 (codex may have fallen back to its default)" +# the request body's model field would not be gpt-5.6-luna. Match the escaped +# model field exactly; codex mentions "gpt-5.5" in its system prompt, so a +# loose grep would give a false pass. +if ! grep -q '\\"model\\":\\"gpt-5\.6-luna\\"' "$LOG_FILE"; then + echo "FAIL: stub did not receive the configured model gpt-5.6-luna (codex may have fallen back to its default)" cat "$LOG_FILE" exit 1 fi diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index e566fb55a..6d3994357 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -402,6 +402,40 @@ async def fake_start(**kwargs): ) +@pytest.mark.asyncio +async def test_host_proxy_translates_container_view_provider_alias(monkeypatch): + """Guards PR #1057 against host capture losing Docker-host providers.""" + + starts = [] + + async def fake_start(**kwargs): + starts.append(kwargs) + return FakeLiteLLMServer("http://host.docker.internal:45678", kwargs["route"]) + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + + updated, provider_runtime = await ensure_litellm_runtime( + agent="codex-acp", + agent_env={ + "BENCHFLOW_PROVIDER_BASE_URL": ( + "http://host.docker.internal:18765/custom/v1?tenant=science" + ), + "BENCHFLOW_PROVIDER_API_KEY": "dummy-local-key", + }, + model="vllm/gpt-5.6-luna", + runtime=None, + environment="docker", + session_id="run-local-provider", + usage_tracking="off", + ) + + assert provider_runtime is not None + assert starts[0]["route"].litellm_params["api_base"] == ( + "http://127.0.0.1:18765/custom/v1?tenant=science" + ) + assert updated["OPENAI_BASE_URL"] == "http://host.docker.internal:45678/v1" + + @pytest.mark.asyncio async def test_pi_acp_proxy_preserves_provider_model_metadata(monkeypatch): """Guards PR #803: Pi metadata follows the LiteLLM alias in proxy mode.""" diff --git a/tests/test_resolve_env_helpers.py b/tests/test_resolve_env_helpers.py index e28cbe1e7..0141e181f 100644 --- a/tests/test_resolve_env_helpers.py +++ b/tests/test_resolve_env_helpers.py @@ -750,6 +750,60 @@ def test_codex_uses_azure_api_key_and_endpoint(self, monkeypatch): "supports_websockets": False, } + def test_codex_direct_provider_preserves_caller_settings(self, monkeypatch): + """Guards PR #1057 against minimalizing direct Codex provider config.""" + + monkeypatch.setenv("AZURE_API_KEY", "az-test") + monkeypatch.setenv( + "AZURE_API_ENDPOINT", "https://example-resource.openai.azure.com/" + ) + caller_config = { + "model_reasoning_effort": "xhigh", + "service_tier": "priority", + "model_provider": "enterprise", + "model_providers": { + "enterprise": { + "name": "Enterprise Azure", + "base_url": "https://stale.example/v1", + "env_key": "AZURE_API_KEY", + "wire_api": "chat", + "supports_websockets": True, + "http_headers": {"x-organization": "science"}, + }, + "unused": { + "name": "Unselected provider", + "base_url": "https://unused.example/v1", + }, + }, + } + + result = resolve_agent_env( + "codex-acp", + "azure-foundry-openai/gpt-5.5", + { + "CODEX_CONFIG": json.dumps(caller_config), + "MODEL_PROVIDER": "enterprise", + }, + ) + + config = json.loads(result["CODEX_CONFIG"]) + assert config["model_reasoning_effort"] == "xhigh" + assert config["service_tier"] == "priority" + assert config["model"] == "gpt-5.5" + assert config["model_provider"] == "enterprise" + assert ( + config["model_providers"]["unused"] + == caller_config["model_providers"]["unused"] + ) + assert config["model_providers"]["enterprise"] == { + "name": "Enterprise Azure", + "base_url": "https://example-resource.openai.azure.com/openai/v1", + "env_key": "OPENAI_API_KEY", + "wire_api": "chat", + "supports_websockets": True, + "http_headers": {"x-organization": "science"}, + } + def test_claude_uses_same_azure_key_on_anthropic_surface(self, monkeypatch): monkeypatch.setenv("AZURE_API_KEY", "az-test") monkeypatch.setenv( From 1fabaf05f07020ac945d890531512f84e8abd30d Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 16:59:21 -0700 Subject: [PATCH 50/74] refactor llm capture runtime custody --- hatch_build.py | 38 +- src/benchflow/agents/codex_config.py | 111 +++-- src/benchflow/agents/env.py | 5 +- src/benchflow/cli/train.py | 5 +- src/benchflow/continue_run/orchestrator.py | 3 + src/benchflow/continue_run/sandbox_proxy.py | 455 +----------------- .../continue_run/sandbox_replay_runtime.py | 427 ++++++++++++++++ .../continue_run/trajectory_artifacts.py | 56 ++- src/benchflow/eval_artifacts.py | 32 +- .../providers/litellm_capture_custody.py | 112 +++++ src/benchflow/providers/litellm_logging.py | 82 +++- src/benchflow/providers/litellm_runtime.py | 201 ++------ .../resources/provider_capture_custody.sh | 62 +++ src/benchflow/rollout/__init__.py | 18 +- src/benchflow/sandbox/files.py | 25 + src/benchflow/trajectories/export.py | 2 +- src/benchflow/trajectories/export_adp.py | 2 +- src/benchflow/trajectories/export_atif.py | 2 +- .../trajectories/export_prime_sft.py | 10 +- src/benchflow/trajectories/export_trl_sft.py | 10 +- src/benchflow/trajectories/io.py | 27 ++ src/benchflow/trajectories/llm_capture.py | 19 +- .../trajectories/llm_capture_manifest.py | 41 +- .../trajectories/llm_capture_records.py | 38 +- .../trajectories/native_capture_collection.py | 47 +- .../trajectories/native_capture_parsers.py | 1 - src/benchflow/trajectories/redaction.py | 32 +- src/benchflow/trajectories/results.py | 9 +- src/benchflow/trajectories/types.py | 18 +- tests/continue_run/test_orchestrator.py | 14 +- tests/continue_run/test_replay_proxy.py | 46 +- tests/test_eval_artifact_cli.py | 12 +- tests/test_litellm_hardening.py | 5 +- tests/test_litellm_logging.py | 32 ++ tests/test_litellm_runtime.py | 95 ++-- tests/test_live_llm_trajectory.py | 4 + tests/test_redaction_resource.py | 5 + tests/test_train_cli.py | 68 +++ tests/trajectories/test_native_llm_capture.py | 26 + 39 files changed, 1283 insertions(+), 914 deletions(-) create mode 100644 src/benchflow/continue_run/sandbox_replay_runtime.py create mode 100644 src/benchflow/providers/litellm_capture_custody.py create mode 100644 src/benchflow/providers/resources/provider_capture_custody.sh create mode 100644 src/benchflow/sandbox/files.py create mode 100644 src/benchflow/trajectories/io.py diff --git a/hatch_build.py b/hatch_build.py index f19c974a6..2dd7673a9 100644 --- a/hatch_build.py +++ b/hatch_build.py @@ -6,8 +6,16 @@ from hatchling.builders.hooks.plugin.interface import BuildHookInterface -_REDACTOR_SOURCE = "src/benchflow/trajectories/redaction.py" -_REDACTOR_RESOURCE = "benchflow/trajectories/resources/canonical_redaction.py.txt" +_RUNTIME_RESOURCES = ( + ( + "src/benchflow/trajectories/redaction.py", + "benchflow/trajectories/resources/canonical_redaction.py.txt", + ), + ( + "src/benchflow/continue_run/sandbox_replay_runtime.py", + "benchflow/continue_run/resources/sandbox_replay_runtime.py.txt", + ), +) class CustomBuildHook(BuildHookInterface): @@ -16,15 +24,18 @@ class CustomBuildHook(BuildHookInterface): def initialize(self, _version: str, build_data: dict[str, Any]) -> None: if self.target_name != "wheel": return - source = Path(self.root, _REDACTOR_SOURCE) - with tempfile.NamedTemporaryFile( - prefix="benchflow-redactor-", - suffix=".txt", - delete=False, - ) as generated: - generated.write(source.read_bytes()) - self._generated_path = Path(generated.name) - build_data["force_include"][str(self._generated_path)] = _REDACTOR_RESOURCE + self._generated_paths: list[Path] = [] + for source_name, resource_name in _RUNTIME_RESOURCES: + source = Path(self.root, source_name) + with tempfile.NamedTemporaryFile( + prefix="benchflow-runtime-resource-", + suffix=".txt", + delete=False, + ) as generated: + generated.write(source.read_bytes()) + generated_path = Path(generated.name) + self._generated_paths.append(generated_path) + build_data["force_include"][str(generated_path)] = resource_name def finalize( self, @@ -38,7 +49,6 @@ def clean(self, _versions: list[str]) -> None: self._discard_generated() def _discard_generated(self) -> None: - generated = getattr(self, "_generated_path", None) - if generated is not None: + for generated in getattr(self, "_generated_paths", []): generated.unlink(missing_ok=True) - self._generated_path = None + self._generated_paths = [] diff --git a/src/benchflow/agents/codex_config.py b/src/benchflow/agents/codex_config.py index 8a31e1fd1..eb986236f 100644 --- a/src/benchflow/agents/codex_config.py +++ b/src/benchflow/agents/codex_config.py @@ -20,21 +20,15 @@ def codex_provider_id(provider_name: str | None) -> str: return f"{_CODEX_PROVIDER_ID_PREFIX}{safe_name or 'provider'}" -def apply_codex_provider_config( +def apply_codex_custom_provider_config( agent_env: dict[str, str], *, base_url: str, model: str | None, provider_name: str, - strict: bool = False, - isolate: bool = False, ) -> None: - """Create or update Codex's model provider entry in ``agent_env``. + """Update caller-owned Codex config for a direct custom provider.""" - ``isolate`` discards caller config and is reserved for a BenchFlow-owned - proxy boundary. Direct provider configuration keeps unrelated caller - settings while ``strict`` still validates JSON and pins the credential key. - """ raw_config = agent_env.get(CODEX_CONFIG_ENV) if not raw_config: config: dict[str, Any] = {} @@ -42,45 +36,80 @@ def apply_codex_provider_config( try: config = json.loads(raw_config) except json.JSONDecodeError as exc: - if strict: - raise ValueError(f"{CODEX_CONFIG_ENV} must be valid JSON") from exc - return + raise ValueError(f"{CODEX_CONFIG_ENV} must be valid JSON") from exc if not isinstance(config, dict): - if strict: - raise ValueError(f"{CODEX_CONFIG_ENV} must decode to a JSON object") - return + raise ValueError(f"{CODEX_CONFIG_ENV} must decode to a JSON object") + + configured_provider_id = agent_env.get(CODEX_MODEL_PROVIDER_ENV) or config.get( + "model_provider" + ) + provider_id = ( + configured_provider_id + if isinstance(configured_provider_id, str) and configured_provider_id + else codex_provider_id(provider_name) + ) + providers_value = config.get("model_providers") + providers = {} if not isinstance(providers_value, dict) else dict(providers_value) + provider_value = providers.get(provider_id) + provider = dict(provider_value) if isinstance(provider_value, dict) else {} + provider.setdefault("name", provider_name) + provider["base_url"] = base_url + provider["env_key"] = "OPENAI_API_KEY" + provider.setdefault("wire_api", "responses") + provider.setdefault("supports_websockets", False) + + _write_codex_provider_config( + agent_env, + config=config, + providers=providers, + provider_id=provider_id, + provider=provider, + model=model, + base_url=base_url, + provider_name=provider_name, + ) + + +def apply_codex_proxy_config( + agent_env: dict[str, str], + *, + base_url: str, + model: str | None, + provider_name: str, +) -> None: + """Replace Codex config with one BenchFlow-owned proxy provider.""" - if isolate: - provider_id = codex_provider_id(provider_name) - config = {} - providers: dict[str, Any] = {} - provider: dict[str, Any] = { + provider_id = codex_provider_id(provider_name) + _write_codex_provider_config( + agent_env, + config={}, + providers={}, + provider_id=provider_id, + provider={ "name": provider_name, "base_url": base_url, "env_key": "OPENAI_API_KEY", "wire_api": "responses", "supports_websockets": False, - } - else: - provider_id = ( - agent_env.get(CODEX_MODEL_PROVIDER_ENV) - or config.get("model_provider") - or codex_provider_id(provider_name) - ) - providers_value = config.get("model_providers") - providers = ( - {} if not isinstance(providers_value, dict) else dict(providers_value) - ) - provider_value = providers.get(provider_id) - provider = dict(provider_value) if isinstance(provider_value, dict) else {} - provider.setdefault("name", provider_name) - provider["base_url"] = base_url - if strict: - provider["env_key"] = "OPENAI_API_KEY" - else: - provider.setdefault("env_key", "OPENAI_API_KEY") - provider.setdefault("wire_api", "responses") - provider.setdefault("supports_websockets", False) + }, + model=model, + base_url=base_url, + provider_name=provider_name, + ) + + +def _write_codex_provider_config( + agent_env: dict[str, str], + *, + config: dict[str, Any], + providers: dict[str, Any], + provider_id: str, + provider: dict[str, Any], + model: str | None, + base_url: str, + provider_name: str, +) -> None: + """Serialize one already-resolved Codex provider configuration.""" providers[provider_id] = provider config["model_providers"] = providers @@ -88,7 +117,7 @@ def apply_codex_provider_config( if model: config["model"] = model - agent_env[CODEX_MODEL_PROVIDER_ENV] = str(provider_id) + agent_env[CODEX_MODEL_PROVIDER_ENV] = provider_id agent_env[CODEX_CONFIG_ENV] = json.dumps(config, separators=(",", ":")) _apply_codex_default_auth_request( agent_env, diff --git a/src/benchflow/agents/env.py b/src/benchflow/agents/env.py index ebf83e15c..202813e04 100644 --- a/src/benchflow/agents/env.py +++ b/src/benchflow/agents/env.py @@ -23,7 +23,7 @@ from urllib.parse import urlparse from benchflow._dotenv import load_dotenv_env -from benchflow.agents.codex_config import apply_codex_provider_config +from benchflow.agents.codex_config import apply_codex_custom_provider_config from benchflow.agents.registry import AGENTS logger = logging.getLogger(__name__) @@ -615,12 +615,11 @@ def _configure_codex_custom_provider( if not base_url or not provider_model: return - apply_codex_provider_config( + apply_codex_custom_provider_config( agent_env, base_url=base_url, model=provider_model, provider_name=agent_env.get("BENCHFLOW_PROVIDER_NAME", "openai-compatible"), - strict=True, ) diff --git a/src/benchflow/cli/train.py b/src/benchflow/cli/train.py index d11004876..2987fa4ab 100644 --- a/src/benchflow/cli/train.py +++ b/src/benchflow/cli/train.py @@ -283,9 +283,11 @@ def train_validate( if require_llm_trajectory and ( health["missing_llm_trajectory"] or health["malformed_llm_trajectory"] + or health["non_training_grade_llm_trajectory"] ): raise ValueError( - "source jobs contain missing/malformed llm_trajectory.jsonl" + "source jobs contain missing, malformed, or non-training-grade " + "llm_trajectory.jsonl" ) if ( require_tool_calls @@ -304,6 +306,7 @@ def train_validate( "rows_with_tool_calls", "missing_llm_trajectory", "malformed_llm_trajectory", + "non_training_grade_llm_trajectory", ) } if source_canonical_selection is not None: diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 1a5d44004..501db3383 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -744,6 +744,9 @@ async def _write_artifacts( rollout._error = str(exc) logger.error("Run failed", exc_info=True) finally: + # Snapshot once before cleanup while sandbox capture files still exist. + # Re-render atomically after cleanup only to attach teardown errors that + # were unknowable during the first pass; the final call below is a no-op. await _safe_sandbox_continuation_teardown( rollout=rollout, replay_proxy=replay_proxy, diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 38b92371d..097182b46 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -19,10 +19,12 @@ import shlex import tempfile from dataclasses import dataclass, field +from importlib.resources import files from pathlib import Path from typing import Any from uuid import uuid4 +from benchflow.sandbox.files import upload_private_text from benchflow.trajectories.redaction import canonical_redaction_module_source from benchflow.trajectories.types import LLMExchange @@ -36,424 +38,17 @@ def sandbox_replay_base_url(port: int = DEFAULT_SANDBOX_REPLAY_PORT) -> str: return f"http://127.0.0.1:{port}/v1" -def _sandbox_proxy_source() -> str: - return r""" -from __future__ import annotations - -import json -import os -import sys -import threading -import time -import traceback -import urllib.error -import urllib.request -from dataclasses import dataclass, field -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer - -try: - from benchflow_trajectory_redaction import redact_trajectory_obj -except ImportError: - from benchflow.trajectories.redaction import redact_trajectory_obj - -CAPTURE_STATE_WRITE_FAILED = "BENCHFLOW_CAPTURE_STATE_WRITE_FAILED" - - -def _n_messages(body): - messages = body.get("messages") - return len(messages) if isinstance(messages, list) else 0 - - -def _completion_to_sse(body): - base_id = body.get("id") or f"chatcmpl-replay-{int(time.time() * 1000)}" - created = body.get("created") or int(time.time()) - model = body.get("model") or "replay" - choices = body.get("choices") or [{}] - choice = choices[0] if choices and isinstance(choices[0], dict) else {} - message = choice.get("message") or {} - finish_reason = choice.get("finish_reason") or "stop" - - def chunk(delta, finish=None): - return json.dumps( - { - "id": base_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], - } - ) - - payloads = [chunk({"role": "assistant"})] - content = message.get("content") - if content: - delta = {"content": content} - for key in ("reasoning_content", "thinking"): - if message.get(key): - delta[key] = message[key] - payloads.append(chunk(delta)) - - for i, tool_call in enumerate(message.get("tool_calls") or []): - if not isinstance(tool_call, dict): - continue - function = tool_call.get("function") or {} - payloads.append( - chunk( - { - "tool_calls": [ - { - "index": tool_call.get("index", i), - "id": tool_call.get("id"), - "type": tool_call.get("type", "function"), - "function": { - "name": function.get("name"), - "arguments": function.get("arguments", ""), - }, - } - ] - } - ) - ) - - final = { - "id": base_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}], - } - if isinstance(body.get("usage"), dict): - final["usage"] = body["usage"] - payloads.append(json.dumps(final)) - return payloads - - -@dataclass -class ReplayState: - recorded: list - upstream_url: str - upstream_api_key: str - upstream_model: str - live_log_path: str - state_path: str - port: int - strict_divergence: bool = False - cursor: int = 0 - divergences: int = 0 - live_attempt_count: int = 0 - live_error_count: int = 0 - lock: threading.Lock = field(default_factory=threading.Lock) - condition: threading.Condition = field(init=False) - quiescing: bool = False - active_live_requests: int = 0 - active_handlers: int = 0 - - def __post_init__(self): - self.condition = threading.Condition(self.lock) - - def _write_state(self): - payload = { - "port": self.port, - "live_attempt_count": self.live_attempt_count, - "live_error_count": self.live_error_count, - } - temporary = self.state_path + ".tmp" - try: - with open(temporary, "w", encoding="utf-8") as handle: - json.dump(payload, handle) - handle.flush() - os.replace(temporary, self.state_path) - except Exception: - self.live_error_count += 1 - try: - os.unlink(self.state_path) - except OSError: - pass - try: - os.unlink(temporary) - except OSError: - pass - print(CAPTURE_STATE_WRITE_FAILED, file=sys.stderr, flush=True) - raise - - def _check_divergence(self, incoming, recorded_request): - want = _n_messages(recorded_request) - got = _n_messages(incoming) - if want and got and want != got: - self.divergences += 1 - message = ( - f"replay divergence at turn {self.cursor}: agent sent {got} " - f"messages, recorded turn had {want}" - ) - if self.strict_divergence: - raise RuntimeError(message) - print(message, file=sys.stderr, flush=True) - - def next_response(self, request_body): - with self.condition: - if self.quiescing: - self.live_attempt_count += 1 - self.live_error_count += 1 - self._write_state() - return "error", 503, { - "error": {"message": "replay proxy is quiescing"} - } - if self.cursor < len(self.recorded): - exchange = self.recorded[self.cursor] - self._check_divergence( - request_body, - ((exchange.get("request") or {}).get("body") or {}), - ) - self.cursor += 1 - response = exchange.get("response") or {} - return "replay", int(response.get("status_code") or 200), dict(response.get("body") or {}) - self.cursor += 1 - self.live_attempt_count += 1 - live_attempt = self.live_attempt_count - self.active_live_requests += 1 - try: - self._write_state() - except Exception: - self.active_live_requests -= 1 - self.condition.notify_all() - raise +def sandbox_replay_runtime_source() -> str: + """Read the real replay runtime module for sandbox deployment.""" - try: - status, body, provider_observed = self._forward_live(request_body) - if provider_observed: - try: - with self.lock: - self._append_live_exchange( - request_body, status, body, live_attempt - ) - except Exception: - with self.lock: - self.live_error_count += 1 - self._write_state() - raise - else: - with self.lock: - self.live_error_count += 1 - self._write_state() - return "live", status, body - finally: - with self.condition: - self.active_live_requests -= 1 - self.condition.notify_all() - - def quiesce(self, timeout=610): - deadline = time.monotonic() + timeout - with self.condition: - self.quiescing = True - while self.active_live_requests or self.active_handlers > 1: - remaining = deadline - time.monotonic() - if remaining <= 0: - return False - self.condition.wait(remaining) - return True - - def begin_quiesce(self): - with self.condition: - self.quiescing = True - - def handler_started(self): - with self.condition: - self.active_handlers += 1 - - def handler_finished(self): - with self.condition: - self.active_handlers -= 1 - self.condition.notify_all() - - def _forward_live(self, request_body): - forwarded = dict(request_body) - forwarded["model"] = self.upstream_model - forwarded["stream"] = False - data = json.dumps(forwarded).encode("utf-8") - request = urllib.request.Request( - self.upstream_url.rstrip("/") + "/chat/completions", - data=data, - headers={ - "Content-Type": "application/json", - "Authorization": f"Bearer {self.upstream_api_key}", - }, - method="POST", - ) - try: - with urllib.request.urlopen(request, timeout=600) as response: - raw = response.read().decode("utf-8") - return int(response.status), json.loads(raw or "{}"), True - except urllib.error.HTTPError as exc: - raw = exc.read().decode("utf-8", errors="replace") - try: - body = json.loads(raw or "{}") - except json.JSONDecodeError: - body = {"error": {"message": raw or str(exc)}} - return int(exc.code), body, True - except Exception as exc: - traceback.print_exc() - return 500, {"error": {"message": str(exc)}}, False - - def _append_live_exchange(self, request_body, status, body, live_attempt): - row = redact_trajectory_obj( - { - "request": {"body": request_body}, - "response": {"status_code": status, "body": body}, - "metadata": {"continuation_attempt": live_attempt}, - } - ) - with open(self.live_log_path, "a", encoding="utf-8") as handle: - handle.write(json.dumps(row) + "\n") - handle.flush() - - -class ReplayHandler(BaseHTTPRequestHandler): - def log_message(self, fmt, *args): - print("replay-proxy " + fmt % args, file=sys.stderr, flush=True) - - @property - def state(self): - return self.server.state - - def _send_json(self, status, payload): - data = json.dumps(payload).encode("utf-8") - self.send_response(status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(data))) - self.end_headers() - self.wfile.write(data) - - def _send_sse(self, payloads): - self.close_connection = True - self.send_response(200) - self.send_header("Content-Type", "text/event-stream") - self.send_header("Cache-Control", "no-cache") - self.send_header("Connection", "close") - self.end_headers() - for payload in payloads: - self.wfile.write(f"data: {payload}\n\n".encode("utf-8")) - self.wfile.write(b"data: [DONE]\n\n") - self.wfile.flush() - - def do_GET(self): - path = self.path.split("?", 1)[0] - if path in ("/health", "/health/liveliness", "/v1/health"): - self._send_json(200, {"status": "ok"}) - return - if path in ("/v1/models", "/models"): - self._send_json(200, {"object": "list", "data": [{"id": "replay", "object": "model"}]}) - return - self._send_json(404, {"error": {"message": f"not found: {path}"}}) - - def do_POST(self): - path = self.path.split("?", 1)[0] - if path == "/benchflow/quiesce": - # Stop the accept loop before the barrier can report success. The - # server counts connections in ``process_request`` (before their - # worker threads start), so every already-accepted handler is in - # ``active_handlers`` even if it has not reached ``do_POST`` yet. - self.state.begin_quiesce() - self.server.shutdown() - self.server.server_close() - quiesced = self.state.quiesce() - # The control request is the one handler deliberately excluded - # from the barrier above. Retire it from the accepted-handler - # accounting before publishing success so the response is also a - # reliable, race-free signal that the capture lifecycle drained. - self.server.release_current_handler() - self._send_json( - 200 if quiesced else 503, - {"status": "quiesced" if quiesced else "quiesce_timeout"}, - ) - return - if path not in ("/v1/chat/completions", "/chat/completions"): - self._send_json(404, {"error": {"message": f"not found: {path}"}}) - return - try: - length = int(self.headers.get("Content-Length") or 0) - raw = self.rfile.read(length) if length else b"{}" - body = json.loads(raw or b"{}") - if not isinstance(body, dict): - raise ValueError("request body must be a JSON object") - except Exception as exc: - self._send_json(400, {"error": {"message": f"bad request: {exc}"}}) - return - - want_stream = bool(body.get("stream")) - try: - _, status, response = self.state.next_response(body) - except RuntimeError as exc: - self._send_json(409, {"error": {"message": str(exc), "type": "divergence"}}) - return - except Exception as exc: - traceback.print_exc() - self._send_json(500, {"error": {"message": str(exc)}}) - return - - if status >= 400 or not response.get("choices"): - self._send_json(status, response) - return - if want_stream: - self._send_sse(_completion_to_sse(response)) - else: - self._send_json(status, response) - - -class ReplayServer(ThreadingHTTPServer): - # ``/benchflow/quiesce`` closes the listener from its own worker thread. - # Do not let ``server_close`` try to join that current thread; the normal - # non-daemon handler lifecycle still keeps the process alive until the - # response and every already-accepted request have finished. - block_on_close = False - - def __init__(self, address, handler, state): - super().__init__(address, handler) - self.state = state - self._handler_local = threading.local() - - def process_request(self, request, client_address): - self.state.handler_started() - try: - super().process_request(request, client_address) - except BaseException: - self.state.handler_finished() - raise - - def process_request_thread(self, request, client_address): - self._handler_local.counted = True - try: - super().process_request_thread(request, client_address) - finally: - if self._handler_local.counted: - self.state.handler_finished() - self._handler_local.counted = False - - def release_current_handler(self): - if getattr(self._handler_local, "counted", False): - self._handler_local.counted = False - self.state.handler_finished() - - -def main(): - cfg = json.load(open(sys.argv[1], encoding="utf-8")) - state = ReplayState( - recorded=cfg["recorded"], - upstream_url=cfg["upstream_url"], - upstream_api_key=cfg["upstream_api_key"], - upstream_model=cfg["upstream_model"], - live_log_path=cfg["live_log_path"], - state_path=cfg["state_path"], - port=int(cfg["port"]), - strict_divergence=bool(cfg.get("strict_divergence")), + package = files("benchflow.continue_run") + resource = package.joinpath( + "resources", + "sandbox_replay_runtime.py.txt", ) - server = ReplayServer(("127.0.0.1", int(cfg["port"])), ReplayHandler, state) - state._write_state() - server.serve_forever() - - -if __name__ == "__main__": - main() -""" + if resource.is_file(): + return resource.read_text(encoding="utf-8") + return package.joinpath("sandbox_replay_runtime.py").read_text(encoding="utf-8") def _ordered_live_exchange_log(text: str) -> tuple[list[LLMExchange], int]: @@ -483,16 +78,6 @@ def _ordered_live_exchange_log(text: str) -> tuple[list[LLMExchange], int]: return [exchange for _, exchange in sorted(sequenced)], malformed -async def _upload_text(sandbox: Any, text: str, target_path: str, suffix: str) -> None: - with tempfile.NamedTemporaryFile("w", suffix=suffix, delete=False) as tmp: - tmp.write(text) - tmp_path = Path(tmp.name) - try: - await sandbox.upload_file(tmp_path, target_path) - finally: - tmp_path.unlink(missing_ok=True) - - async def _read_remote_text(sandbox: Any, path: str, *, timeout_sec: int = 15) -> str: download_file = getattr(sandbox, "download_file", None) if download_file is not None: @@ -580,14 +165,24 @@ async def start( "state_path": paths["state"], "live_log_path": paths["live_log"], } - await _upload_text(sandbox, _sandbox_proxy_source(), paths["script"], ".py") - await _upload_text( + await upload_private_text( + sandbox, + sandbox_replay_runtime_source(), + paths["script"], + suffix=".py", + ) + await upload_private_text( sandbox, canonical_redaction_module_source(), paths["redaction"], - ".py", + suffix=".py", + ) + await upload_private_text( + sandbox, + json.dumps(config), + paths["config"], + suffix=".json", ) - await _upload_text(sandbox, json.dumps(config), paths["config"], ".json") command = ( f"rm -f {shlex.quote(paths['state'])} {shlex.quote(paths['pid'])} " diff --git a/src/benchflow/continue_run/sandbox_replay_runtime.py b/src/benchflow/continue_run/sandbox_replay_runtime.py new file mode 100644 index 000000000..c0edb701b --- /dev/null +++ b/src/benchflow/continue_run/sandbox_replay_runtime.py @@ -0,0 +1,427 @@ +"""Stdlib-only replay proxy runtime uploaded into remote sandboxes.""" + +from __future__ import annotations + +import contextlib +import json +import os +import sys +import threading +import time +import traceback +import urllib.error +import urllib.request +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, cast + +try: + from benchflow_trajectory_redaction import redact_trajectory_obj +except ImportError: + from benchflow.trajectories.redaction import redact_trajectory_obj + +CAPTURE_STATE_WRITE_FAILED = "BENCHFLOW_CAPTURE_STATE_WRITE_FAILED" +# Provider calls can run for ten minutes; teardown must not snapshot a request +# while the provider timeout is still live. +PROVIDER_DRAIN_TIMEOUT_SEC = 610 + + +def _n_messages(body): + messages = body.get("messages") + return len(messages) if isinstance(messages, list) else 0 + + +def _completion_to_sse(body): + base_id = body.get("id") or f"chatcmpl-replay-{int(time.time() * 1000)}" + created = body.get("created") or int(time.time()) + model = body.get("model") or "replay" + choices = body.get("choices") or [{}] + choice = choices[0] if choices and isinstance(choices[0], dict) else {} + message = choice.get("message") or {} + finish_reason = choice.get("finish_reason") or "stop" + + def chunk(delta, finish=None): + return json.dumps( + { + "id": base_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + ) + + payloads = [chunk({"role": "assistant"})] + content = message.get("content") + if content: + delta = {"content": content} + for key in ("reasoning_content", "thinking"): + if message.get(key): + delta[key] = message[key] + payloads.append(chunk(delta)) + + for i, tool_call in enumerate(message.get("tool_calls") or []): + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function") or {} + payloads.append( + chunk( + { + "tool_calls": [ + { + "index": tool_call.get("index", i), + "id": tool_call.get("id"), + "type": tool_call.get("type", "function"), + "function": { + "name": function.get("name"), + "arguments": function.get("arguments", ""), + }, + } + ] + } + ) + ) + + final = { + "id": base_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}], + } + if isinstance(body.get("usage"), dict): + final["usage"] = body["usage"] + payloads.append(json.dumps(final)) + return payloads + + +@dataclass +class ReplayState: + recorded: list + upstream_url: str + upstream_api_key: str + upstream_model: str + live_log_path: str + state_path: str + port: int + strict_divergence: bool = False + cursor: int = 0 + divergences: int = 0 + live_attempt_count: int = 0 + live_error_count: int = 0 + lock: threading.Lock = field(default_factory=threading.Lock) + condition: threading.Condition = field(init=False) + quiescing: bool = False + active_live_requests: int = 0 + active_handlers: int = 0 + + def __post_init__(self): + self.condition = threading.Condition(self.lock) + + def _write_state(self): + payload = { + "port": self.port, + "live_attempt_count": self.live_attempt_count, + "live_error_count": self.live_error_count, + } + temporary = self.state_path + ".tmp" + try: + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + handle.flush() + os.replace(temporary, self.state_path) + except Exception: + self.live_error_count += 1 + with contextlib.suppress(OSError): + os.unlink(self.state_path) + with contextlib.suppress(OSError): + os.unlink(temporary) + print(CAPTURE_STATE_WRITE_FAILED, file=sys.stderr, flush=True) + raise + + def _check_divergence(self, incoming, recorded_request): + want = _n_messages(recorded_request) + got = _n_messages(incoming) + if want and got and want != got: + self.divergences += 1 + message = ( + f"replay divergence at turn {self.cursor}: agent sent {got} " + f"messages, recorded turn had {want}" + ) + if self.strict_divergence: + raise RuntimeError(message) + print(message, file=sys.stderr, flush=True) + + def next_response(self, request_body): + with self.condition: + if self.quiescing: + self.live_attempt_count += 1 + self.live_error_count += 1 + self._write_state() + return "error", 503, {"error": {"message": "replay proxy is quiescing"}} + if self.cursor < len(self.recorded): + exchange = self.recorded[self.cursor] + self._check_divergence( + request_body, + ((exchange.get("request") or {}).get("body") or {}), + ) + self.cursor += 1 + response = exchange.get("response") or {} + return ( + "replay", + int(response.get("status_code") or 200), + dict(response.get("body") or {}), + ) + self.cursor += 1 + self.live_attempt_count += 1 + live_attempt = self.live_attempt_count + self.active_live_requests += 1 + try: + self._write_state() + except Exception: + self.active_live_requests -= 1 + self.condition.notify_all() + raise + + try: + status, body, provider_observed = self._forward_live(request_body) + if provider_observed: + try: + with self.lock: + self._append_live_exchange( + request_body, status, body, live_attempt + ) + except Exception: + with self.lock: + self.live_error_count += 1 + self._write_state() + raise + else: + with self.lock: + self.live_error_count += 1 + self._write_state() + return "live", status, body + finally: + with self.condition: + self.active_live_requests -= 1 + self.condition.notify_all() + + def quiesce(self, timeout=PROVIDER_DRAIN_TIMEOUT_SEC): + deadline = time.monotonic() + timeout + with self.condition: + self.quiescing = True + while self.active_live_requests or self.active_handlers > 1: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self.condition.wait(remaining) + return True + + def begin_quiesce(self): + with self.condition: + self.quiescing = True + + def handler_started(self): + with self.condition: + self.active_handlers += 1 + + def handler_finished(self): + with self.condition: + self.active_handlers -= 1 + self.condition.notify_all() + + def _forward_live(self, request_body): + forwarded = dict(request_body) + forwarded["model"] = self.upstream_model + forwarded["stream"] = False + data = json.dumps(forwarded).encode("utf-8") + request = urllib.request.Request( + self.upstream_url.rstrip("/") + "/chat/completions", + data=data, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {self.upstream_api_key}", + }, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=600) as response: + raw = response.read().decode("utf-8") + return int(response.status), json.loads(raw or "{}"), True + except urllib.error.HTTPError as exc: + raw = exc.read().decode("utf-8", errors="replace") + try: + body = json.loads(raw or "{}") + except json.JSONDecodeError: + body = {"error": {"message": raw or str(exc)}} + return int(exc.code), body, True + except Exception as exc: + traceback.print_exc() + return 500, {"error": {"message": str(exc)}}, False + + def _append_live_exchange(self, request_body, status, body, live_attempt): + row = redact_trajectory_obj( + { + "request": {"body": request_body}, + "response": {"status_code": status, "body": body}, + "metadata": {"continuation_attempt": live_attempt}, + } + ) + with open(self.live_log_path, "a", encoding="utf-8") as handle: + handle.write(json.dumps(row) + "\n") + handle.flush() + + +class ReplayHandler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: Any) -> None: + print("replay-proxy " + format % args, file=sys.stderr, flush=True) + + @property + def replay_server(self) -> ReplayServer: + return cast("ReplayServer", self.server) + + @property + def state(self) -> ReplayState: + return self.replay_server.state + + def _send_json(self, status, payload): + data = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def _send_sse(self, payloads): + self.close_connection = True + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "close") + self.end_headers() + for payload in payloads: + self.wfile.write(f"data: {payload}\n\n".encode()) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + + def do_GET(self): + path = self.path.split("?", 1)[0] + if path in ("/health", "/health/liveliness", "/v1/health"): + self._send_json(200, {"status": "ok"}) + return + if path in ("/v1/models", "/models"): + self._send_json( + 200, {"object": "list", "data": [{"id": "replay", "object": "model"}]} + ) + return + self._send_json(404, {"error": {"message": f"not found: {path}"}}) + + def do_POST(self): + path = self.path.split("?", 1)[0] + if path == "/benchflow/quiesce": + # Stop the accept loop before the barrier can report success. The + # server counts connections in ``process_request`` (before their + # worker threads start), so every already-accepted handler is in + # ``active_handlers`` even if it has not reached ``do_POST`` yet. + self.state.begin_quiesce() + self.replay_server.shutdown() + self.replay_server.server_close() + quiesced = self.state.quiesce() + # The control request is the one handler deliberately excluded + # from the barrier above. Retire it from the accepted-handler + # accounting before publishing success so the response is also a + # reliable, race-free signal that the capture lifecycle drained. + self.replay_server.release_current_handler() + self._send_json( + 200 if quiesced else 503, + {"status": "quiesced" if quiesced else "quiesce_timeout"}, + ) + return + if path not in ("/v1/chat/completions", "/chat/completions"): + self._send_json(404, {"error": {"message": f"not found: {path}"}}) + return + try: + length = int(self.headers.get("Content-Length") or 0) + raw = self.rfile.read(length) if length else b"{}" + body = json.loads(raw or b"{}") + if not isinstance(body, dict): + raise ValueError("request body must be a JSON object") + except Exception as exc: + self._send_json(400, {"error": {"message": f"bad request: {exc}"}}) + return + + want_stream = bool(body.get("stream")) + try: + _, status, response = self.state.next_response(body) + except RuntimeError as exc: + self._send_json(409, {"error": {"message": str(exc), "type": "divergence"}}) + return + except Exception as exc: + traceback.print_exc() + self._send_json(500, {"error": {"message": str(exc)}}) + return + + if status >= 400 or not response.get("choices"): + self._send_json(status, response) + return + if want_stream: + self._send_sse(_completion_to_sse(response)) + else: + self._send_json(status, response) + + +class ReplayServer(ThreadingHTTPServer): + # ``/benchflow/quiesce`` closes the listener from its own worker thread. + # Do not let ``server_close`` try to join that current thread; the normal + # non-daemon handler lifecycle still keeps the process alive until the + # response and every already-accepted request have finished. + block_on_close = False + + def __init__(self, address, handler, state): + super().__init__(address, handler) + self.state = state + self._handler_local = threading.local() + + def process_request(self, request, client_address): + self.state.handler_started() + try: + super().process_request(request, client_address) + except BaseException: + self.state.handler_finished() + raise + + def process_request_thread(self, request, client_address): + self._handler_local.counted = True + try: + super().process_request_thread(request, client_address) + finally: + if self._handler_local.counted: + self.state.handler_finished() + self._handler_local.counted = False + + def release_current_handler(self): + if getattr(self._handler_local, "counted", False): + self._handler_local.counted = False + self.state.handler_finished() + + +def main(): + with open(sys.argv[1], encoding="utf-8") as config_file: + cfg = json.load(config_file) + state = ReplayState( + recorded=cfg["recorded"], + upstream_url=cfg["upstream_url"], + upstream_api_key=cfg["upstream_api_key"], + upstream_model=cfg["upstream_model"], + live_log_path=cfg["live_log_path"], + state_path=cfg["state_path"], + port=int(cfg["port"]), + strict_divergence=bool(cfg.get("strict_divergence")), + ) + server = ReplayServer(("127.0.0.1", int(cfg["port"])), ReplayHandler, state) + state._write_state() + server.serve_forever() + + +if __name__ == "__main__": + main() diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py index 220a35709..144385f8d 100644 --- a/src/benchflow/continue_run/trajectory_artifacts.py +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -3,11 +3,11 @@ from __future__ import annotations import json -import os from datetime import UTC, datetime from pathlib import Path from typing import Any +from benchflow.trajectories.io import atomic_write_text from benchflow.trajectories.llm_capture_manifest import ( CONTINUATION_SOURCE_AUDIT_ERROR, LLM_TRAJECTORY_SCHEMA_VERSION, @@ -18,23 +18,27 @@ CaptureStatus, LLMRoleCapture, LLMTrajectoryManifest, - capture_artifact_allows_training, read_llm_trajectory_manifest, + rollout_capture_is_training_grade, successful_exchanges_have_positive_usage, write_llm_trajectory_manifest, ) -from benchflow.trajectories.types import LLMExchange, redact_trajectory_obj +from benchflow.trajectories.redaction import ( + jsonl_payload_is_redacted, + redact_trajectory_obj, + redact_trajectory_obj_with_audit, +) +from benchflow.trajectories.types import LLMExchange + +CONTINUATION_AGENT = "openhands" -def stitched_trajectory_lines( - original_llm_trajectory: Path, live_exchanges: list[LLMExchange] -) -> list[str]: - """Build the continuous llm_trajectory: recorded prefix + live suffix. +def stitched_trajectory_lines(original_llm_trajectory: Path) -> list[str]: + """Build the schema-promoted recorded prefix of a continuation trajectory. The recorded request/response payloads are preserved, while their metadata is promoted to the current schema so a newly stitched artifact can never be - mistaken for sidecar-optional legacy data. The live suffix is redacted on - the way out. + mistaken for sidecar-optional legacy data. """ lines: list[str] = [] if original_llm_trajectory.is_file(): @@ -54,9 +58,6 @@ def stitched_trajectory_lines( payload["metadata"] = metadata metadata["schema_version"] = LLM_TRAJECTORY_SCHEMA_VERSION lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) - for exchange in live_exchanges: - payload = redact_trajectory_obj(exchange.model_dump(mode="json")) - lines.append(json.dumps(payload, default=str)) return lines @@ -77,13 +78,13 @@ def write_stitched_trajectory( """ out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" out.parent.mkdir(parents=True, exist_ok=True) - lines = stitched_trajectory_lines(original_llm_trajectory, []) + lines = stitched_trajectory_lines(original_llm_trajectory) for exchange in live_exchanges: payload = exchange.model_dump(mode="json") metadata = payload.setdefault("metadata", {}) metadata.update( { - "agent": "openhands", + "agent": CONTINUATION_AGENT, "role": "agent", "model": live_model, "auth_mode": AuthMode.API_KEY.value, @@ -93,7 +94,6 @@ def write_stitched_trajectory( "request_complete": False, "response_complete": True, "role_attribution_complete": True, - "payload_redacted": True, "request_capture_source": "replay_proxy_ingress", "capture_custody": ( "host_owned" @@ -102,11 +102,18 @@ def write_stitched_trajectory( ), } ) - lines.append(json.dumps(redact_trajectory_obj(payload), default=str)) + redacted, payload_redacted = redact_trajectory_obj_with_audit(payload) + if not isinstance(redacted, dict): + payload_redacted = False + redacted = {} + redacted_metadata = redacted.get("metadata") + if not isinstance(redacted_metadata, dict): + redacted_metadata = {} + redacted["metadata"] = redacted_metadata + redacted_metadata["payload_redacted"] = payload_redacted + lines.append(json.dumps(redacted, default=str)) rendered = "\n".join(lines) + ("\n" if lines else "") - temporary = out.with_suffix(out.suffix + ".tmp") - temporary.write_text(rendered) - os.replace(temporary, out) + atomic_write_text(out, rendered) return out @@ -174,7 +181,10 @@ def refresh_stitched_trajectory_manifest( source_allows_training = bool( source_raw is not None and len(source_rows) == n_recorded - and capture_artifact_allows_training(source_raw, exchanges=source_rows) + and rollout_capture_is_training_grade( + source_rollout_dir, + exchanges=source_rows, + ) ) live_capture_complete = live_attempt_count == n_live and not live_errors usage_complete = bool( @@ -272,13 +282,13 @@ def refresh_stitched_trajectory_manifest( capture_source=capture_source, capture_fidelity=capture_fidelity, auth_mode=auth_mode, - agent="openhands", + agent=CONTINUATION_AGENT, model=stitched_model, session_id=current.session_id if current else rollout_dir.name, exchange_count=exchange_count, request_complete=request_complete, response_complete=response_complete, - payload_redacted=source.payload_redacted if source else True, + payload_redacted=jsonl_payload_is_redacted(trajectory_path), started_at=current.started_at if current else datetime.now(UTC), finished_at=datetime.now(UTC), missing_fields=sorted(set(missing_fields)), @@ -327,7 +337,7 @@ def _continuation_role_captures( LLMRoleCapture( role="agent", leg="live", - agent="openhands", + agent=CONTINUATION_AGENT, model=live_model, auth_mode=AuthMode.API_KEY, capture_source=CaptureSource.REPLAY_PROXY, diff --git a/src/benchflow/eval_artifacts.py b/src/benchflow/eval_artifacts.py index a97280218..48eae2391 100644 --- a/src/benchflow/eval_artifacts.py +++ b/src/benchflow/eval_artifacts.py @@ -16,8 +16,7 @@ load_llm_trajectory_jsonl, ) from benchflow.trajectories.llm_capture_manifest import ( - capture_artifact_allows_training, - read_llm_trajectory_manifest, + rollout_capture_is_training_grade, ) CanonicalizePolicy = Literal["none", "one-healthy-per-task"] @@ -198,20 +197,18 @@ def _tool_call_count(result: dict[str, Any]) -> int: return 0 -def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, int]: +def _llm_trajectory_status(rollout_dir: Path) -> tuple[bool, bool, bool, int]: path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" if not path.is_file(): - return False, False, 0 + return False, False, False, 0 try: rows = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError: - return True, False, 0 - if not rows: - return True, False, 0 - manifest = read_llm_trajectory_manifest(rollout_dir) - if not capture_artifact_allows_training(manifest, exchanges=rows): - return True, False, len(rows) - return True, True, len(rows) + return True, False, False, 0 + training_grade = bool( + rows and rollout_capture_is_training_grade(rollout_dir, exchanges=rows) + ) + return True, True, training_grade, len(rows) def build_health_summary( @@ -226,6 +223,7 @@ def build_health_summary( "zero_tool_rows": 0, "missing_llm_trajectory": 0, "malformed_llm_trajectory": 0, + "non_training_grade_llm_trajectory": 0, } rollout_dirs = ( _iter_selected_rollouts(canonical_selection) @@ -248,11 +246,15 @@ def build_health_summary( counts["rows_with_tool_calls"] += 1 else: counts["zero_tool_rows"] += 1 - has_llm, valid_llm, llm_rows = _llm_trajectory_status(rollout_dir) + has_llm, well_formed_llm, training_grade_llm, llm_rows = _llm_trajectory_status( + rollout_dir + ) if not has_llm: counts["missing_llm_trajectory"] += 1 - elif not valid_llm: + elif not well_formed_llm: counts["malformed_llm_trajectory"] += 1 + elif not training_grade_llm: + counts["non_training_grade_llm_trajectory"] += 1 rows.append( { "task_id": result.get("task_name") or rollout_dir.name, @@ -261,7 +263,9 @@ def build_health_summary( "scored": scored, "tool_calls": tool_calls, "has_llm_trajectory": has_llm, - "valid_llm_trajectory": valid_llm, + "well_formed_llm_trajectory": well_formed_llm, + "valid_llm_trajectory": training_grade_llm, + "training_grade_llm_trajectory": training_grade_llm, "llm_trajectory_rows": llm_rows, "error": result.get("error"), "verifier_error": result.get("verifier_error"), diff --git a/src/benchflow/providers/litellm_capture_custody.py b/src/benchflow/providers/litellm_capture_custody.py new file mode 100644 index 000000000..bfe2ea987 --- /dev/null +++ b/src/benchflow/providers/litellm_capture_custody.py @@ -0,0 +1,112 @@ +"""Sandbox artifact-custody proof for provider-wire capture.""" + +from __future__ import annotations + +import contextlib +import logging +import secrets +import shlex +from importlib.resources import files +from typing import Any + +from benchflow.sandbox.files import upload_private_text + +logger = logging.getLogger(__name__) + + +def provider_capture_custody_probe_source() -> str: + """Read the packaged custody prober deployed beside sandbox runtimes.""" + + return ( + files("benchflow.providers") + .joinpath("resources", "provider_capture_custody.sh") + .read_text(encoding="utf-8") + ) + + +async def provider_capture_has_verified_custody( + *, + sandbox_user: str | None, + sandbox: Any | None, + runtime_dir: str | None, +) -> bool: + """Prove the agent cannot read or mutate provider capture files.""" + + if sandbox is None or sandbox_user in {None, "", "root", "0"} or not runtime_dir: + return False + try: + result = await sandbox.exec( + f"id -u -- {shlex.quote(sandbox_user)}", + user="root", + timeout_sec=10, + ) + except Exception as exc: + logger.warning("Provider capture custody UID check failed: %s", exc) + return False + if result.return_code != 0: + logger.warning("Provider capture custody UID check returned non-zero") + return False + try: + effective_uid = int(result.stdout.strip()) + except (AttributeError, ValueError): + logger.warning("Provider capture custody UID check returned invalid output") + return False + if effective_uid == 0: + return False + + probe_path = f"/tmp/benchflow-capture-custody-{secrets.token_hex(8)}.sh" + quoted_probe = shlex.quote(probe_path) + quoted_runtime = shlex.quote(runtime_dir) + try: + await upload_private_text( + sandbox, + provider_capture_custody_probe_source(), + probe_path, + suffix=".sh", + ) + prepared = await sandbox.exec( + f"chown 0:0 {quoted_probe} && chmod 755 {quoted_probe}", + user="root", + timeout_sec=10, + ) + if prepared.return_code != 0: + logger.warning("Provider capture custody probe upload hardening failed") + return False + + secured = await sandbox.exec( + f"{quoted_probe} harden {quoted_runtime}", + user="root", + timeout_sec=10, + ) + if secured.return_code != 0: + logger.warning("Provider capture custody artifact hardening failed") + return False + access = await sandbox.exec( + f"{quoted_probe} probe {quoted_runtime}", + user=sandbox_user, + timeout_sec=10, + ) + if access.return_code != 1: + logger.warning( + "Provider capture custody probe found agent-accessible root data" + ) + return False + verified = await sandbox.exec( + f"{quoted_probe} verify {quoted_runtime}", + user="root", + timeout_sec=10, + ) + if verified.return_code != 0: + logger.warning("Provider capture custody probe integrity check failed") + return False + return True + except Exception as exc: + logger.warning("Provider capture custody artifact probe failed: %s", exc) + return False + finally: + with contextlib.suppress(Exception): + await sandbox.exec( + f"rm -f {quoted_probe}", + user="root", + timeout_sec=10, + ) diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 90281df9d..cc8f52721 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -212,6 +212,52 @@ def _jsonable(value: Any) -> Any: return str(value) +def _reconstruct_proxy_ingress_body( + kwargs: dict[str, Any], + *, + litellm_params: dict[str, Any], + optional_params: dict[str, Any], + proxy_body: Any, +) -> dict[str, Any]: + # Build the explicitly incomplete fallback when provider input is absent. + + request_body = dict(proxy_body) if isinstance(proxy_body, dict) else {} + for key in ("model", "messages", "input"): + if kwargs.get(key) is not None: + request_body[key] = kwargs[key] + for key in ("tools", "stream"): + if key in optional_params: + request_body[key] = optional_params[key] + elif kwargs.get(key) is not None: + request_body[key] = kwargs[key] + for key in ("reasoning_effort", "thinking", "output_config"): + value = optional_params.get(key) + if value is None: + value = kwargs.get(key) + if value is None: + value = litellm_params.get(key) + if value is not None: + request_body[key] = value + for key in ("logprobs", "top_logprobs"): + value = optional_params.get(key) + if value is None: + value = kwargs.get(key) + if value is not None: + request_body[key] = value + return {key: value for key, value in request_body.items() if value is not None} + + +def _provider_request_path(call_type: Any) -> str: + # Map LiteLLM's provider call type to the captured upstream API path. + + normalized = str(call_type or "").casefold() + if normalized == "anthropic_messages": + return "/v1/messages" + if "responses" in normalized: + return "/v1/responses" + return "/v1/chat/completions" + + def _iso(value: Any) -> str: if isinstance(value, datetime): return value.isoformat() @@ -389,31 +435,15 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - request_body = dict(jsonable_provider_body) request_capture_source = "litellm_pre_api_call_complete_input_dict" else: - request_body = dict(proxy_body) if isinstance(proxy_body, dict) else {} - for key in ("model", "messages", "input"): - if kwargs.get(key) is not None: - request_body[key] = kwargs[key] - for key in ("tools", "stream"): - if key in optional_params: - request_body[key] = optional_params[key] - elif kwargs.get(key) is not None: - request_body[key] = kwargs[key] - for key in ("reasoning_effort", "thinking", "output_config"): - value = optional_params.get(key) - if value is None: - value = kwargs.get(key) - if value is None: - value = litellm_params.get(key) - if value is not None: - request_body[key] = value - for key in ("logprobs", "top_logprobs"): - value = optional_params.get(key) - if value is None: - value = kwargs.get(key) - if value is not None: - request_body[key] = value + request_body = _reconstruct_proxy_ingress_body( + kwargs, + litellm_params=litellm_params, + optional_params=optional_params, + proxy_body=proxy_body, + ) request_capture_source = "proxy_ingress_reconstruction" - request_body = {k: v for k, v in request_body.items() if v is not None} + request_body = {key: value for key, value in request_body.items() if value is not None} + call_type = kwargs.get("call_type") or litellm_params.get("call_type") return { "benchflow_agent": os.environ.get("BENCHFLOW_LITELLM_AGENT"), "benchflow_role": os.environ.get("BENCHFLOW_LITELLM_ROLE"), @@ -426,7 +456,7 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - "request_model": kwargs.get("model"), "provider_model": litellm_params.get("model") or kwargs.get("model"), "model_group": metadata.get("model_group") if isinstance(metadata, dict) else None, - "call_type": kwargs.get("call_type") or litellm_params.get("call_type"), + "call_type": call_type, "request_complete": request_complete, "request_capture_source": request_capture_source, "input_shape": { @@ -436,7 +466,7 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - }, "request": { "method": "POST", - "path": "/v1/messages" if kwargs.get("call_type") == "anthropic_messages" else "/v1/chat/completions", + "path": _provider_request_path(call_type), "body": request_body, }, "start_time": _iso(start_time), diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 82a052b63..148186d75 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -27,7 +27,7 @@ import yaml from benchflow._utils.text import describe_exception -from benchflow.agents.codex_config import apply_codex_provider_config +from benchflow.agents.codex_config import apply_codex_proxy_config from benchflow.agents.env import uses_native_subscription_auth from benchflow.agents.registry import AGENTS from benchflow.providers.litellm_bedrock_preflight import ( @@ -37,6 +37,9 @@ preflight_sandbox_bedrock_patch, route_requires_bedrock_patch, ) +from benchflow.providers.litellm_capture_custody import ( + provider_capture_has_verified_custody as _provider_capture_has_verified_custody, +) from benchflow.providers.litellm_capture_lifecycle import ( LITELLM_CAPTURE_STATE_ENV, capture_journal_error, @@ -57,6 +60,7 @@ extract_usage_from_trajectory, trajectory_from_litellm_callback_log, ) +from benchflow.sandbox.files import upload_private_text from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter from benchflow.trajectories.redaction import canonical_redaction_module_source @@ -422,7 +426,7 @@ def __init__( stderr_path: Path, session_id: str, agent_name: str, - capture_state_path: Path | None = None, + capture_state_path: Path, ) -> None: self.route = route self.process = process @@ -497,13 +501,6 @@ def _load_callback_log(self) -> None: ) def _load_capture_state(self) -> dict[str, Any] | None: - if self.capture_state_path is None: - trajectory = self.trajectory - assert trajectory is not None - return { - "attempt_count": len(trajectory.exchanges), - "terminal_count": len(trajectory.exchanges), - } try: payload = json.loads(self.capture_state_path.read_text()) except (OSError, json.JSONDecodeError): @@ -534,7 +531,7 @@ def __init__( stderr_path: str, session_id: str, agent_name: str, - capture_state_path: str | None = None, + capture_state_path: str, ) -> None: self.sandbox = sandbox self.route = route @@ -673,13 +670,6 @@ async def _load_callback_log(self) -> None: ) async def _load_capture_state(self) -> dict[str, Any] | None: - if self.capture_state_path is None: - trajectory = self.trajectory - assert trajectory is not None - return { - "attempt_count": len(trajectory.exchanges), - "terminal_count": len(trajectory.exchanges), - } result = await self.sandbox.exec( f"cat {shlex.quote(self.capture_state_path)} 2>/dev/null || true", timeout_sec=15, @@ -864,7 +854,7 @@ def _write_runtime_files( sitecustomize_path = runtime_dir / "sitecustomize.py" config_path = runtime_dir / "config.yaml" callback_path.write_text(callback_module_source()) - redaction_path.write_text(_redaction_module_source()) + redaction_path.write_text(canonical_redaction_module_source()) patch_source = Path(__file__).with_name("litellm_bedrock_patch.py").read_text() patch_path.write_text(patch_source) sitecustomize_path.write_text(f"import {_PATCH_MODULE}\n") @@ -880,12 +870,6 @@ def _write_runtime_files( return config_path, callback_path, patch_path -def _redaction_module_source() -> str: - """Read the packaged canonical redactor for isolated proxy runtimes.""" - - return canonical_redaction_module_source() - - # How long to wait for the *host* per-run LiteLLM proxy to become healthy. # litellm's cold start runs tens of seconds, and when many runs launch in parallel # (max-parallel sweeps) the proxies cold-start simultaneously and contend for CPU, @@ -1082,19 +1066,6 @@ def _sandbox_launcher_source() -> str: """ -async def _upload_text(sandbox: Any, text: str, target_path: str, suffix: str) -> None: - with tempfile.NamedTemporaryFile("w", suffix=suffix, delete=False) as tmp: - tmp.write(text) - tmp_path = Path(tmp.name) - try: - # Runtime files carry the provider environment and proxy master key; - # ``mode`` is part of the sandbox protocol, so every backend either - # honors it or fails loudly — never a silently world-readable secret. - await sandbox.upload_file(tmp_path, target_path, mode="600") - finally: - tmp_path.unlink(missing_ok=True) - - async def _upload_runtime_files_to_sandbox( sandbox: Any, *, @@ -1129,23 +1100,41 @@ async def _upload_runtime_files_to_sandbox( ) if result.return_code != 0: raise RuntimeError(_exec_details("prepare LiteLLM runtime directory", result)) - await _upload_text( - sandbox, yaml.safe_dump(config, sort_keys=False), paths["config"], ".yaml" + await upload_private_text( + sandbox, + yaml.safe_dump(config, sort_keys=False), + paths["config"], + suffix=".yaml", + ) + await upload_private_text( + sandbox, callback_module_source(), paths["callback"], suffix=".py" + ) + await upload_private_text( + sandbox, + canonical_redaction_module_source(), + paths["redaction"], + suffix=".py", ) - await _upload_text(sandbox, callback_module_source(), paths["callback"], ".py") - await _upload_text(sandbox, _redaction_module_source(), paths["redaction"], ".py") - await _upload_text( + await upload_private_text( sandbox, Path(__file__).with_name("litellm_bedrock_patch.py").read_text(), paths["patch"], - ".py", + suffix=".py", ) - await _upload_text( - sandbox, f"import {_PATCH_MODULE}\n", paths["sitecustomize"], ".py" + await upload_private_text( + sandbox, + f"import {_PATCH_MODULE}\n", + paths["sitecustomize"], + suffix=".py", + ) + await upload_private_text( + sandbox, _sandbox_launcher_source(), paths["launcher"], suffix=".py" ) - await _upload_text(sandbox, _sandbox_launcher_source(), paths["launcher"], ".py") - await _upload_text( - sandbox, BEDROCK_PATCH_PREFLIGHT_SOURCE, paths["preflight"], ".py" + await upload_private_text( + sandbox, + BEDROCK_PATCH_PREFLIGHT_SOURCE, + paths["preflight"], + suffix=".py", ) return paths @@ -1327,8 +1316,11 @@ async def _start_sandbox_litellm( "state": paths["state"], "env": env, } - await _upload_text( - sandbox, json.dumps(launch_config), paths["launch_config"], ".json" + await upload_private_text( + sandbox, + json.dumps(launch_config), + paths["launch_config"], + suffix=".json", ) # The launcher reads launch_config.json (provider env + master key) at # startup; unlink it immediately afterwards so the secret does not sit in @@ -1677,13 +1669,11 @@ def _wire_litellm_agent_env( updated["OPENAI_BASE_URL"] = openai_base_url updated["OPENAI_API_KEY"] = master_key updated[LITELLM_MODEL_VIA_ENV] = "1" - apply_codex_provider_config( + apply_codex_proxy_config( updated, base_url=openai_base_url, model=route.model_alias, provider_name="litellm", - strict=True, - isolate=True, ) return updated if agent == "opencode": @@ -1869,7 +1859,6 @@ async def ensure_litellm_runtime( if is_running: if sandbox_local: artifact_custody = await _provider_capture_has_verified_custody( - sandbox_local=True, sandbox_user=sandbox_user, sandbox=sandbox, runtime_dir=getattr(server, "runtime_dir", None), @@ -1933,7 +1922,6 @@ async def ensure_litellm_runtime( if sandbox_local: artifact_custody = await _provider_capture_has_verified_custody( - sandbox_local=True, sandbox_user=sandbox_user, sandbox=sandbox, runtime_dir=getattr(server, "runtime_dir", None), @@ -1965,109 +1953,6 @@ async def ensure_litellm_runtime( ) -async def _provider_capture_has_verified_custody( - *, - sandbox_local: bool, - sandbox_user: str | None, - sandbox: Any | None, - runtime_dir: str | None, -) -> bool: - """Prove the agent cannot read or mutate the real provider capture files.""" - - if not sandbox_local: - return True - if sandbox is None or sandbox_user in {None, "", "root", "0"} or not runtime_dir: - return False - try: - result = await sandbox.exec( - f"id -u -- {shlex.quote(sandbox_user)}", - user="root", - timeout_sec=10, - ) - except Exception as exc: - logger.warning("Provider capture custody UID check failed: %s", exc) - return False - if result.return_code != 0: - logger.warning("Provider capture custody UID check returned non-zero") - return False - try: - effective_uid = int(result.stdout.strip()) - except (AttributeError, ValueError): - logger.warning("Provider capture custody UID check returned invalid output") - return False - if effective_uid == 0: - return False - - quoted_runtime = shlex.quote(runtime_dir) - quoted_log = shlex.quote(f"{runtime_dir}/callback.jsonl") - quoted_state = shlex.quote(f"{runtime_dir}/capture_state.json") - try: - secured = await sandbox.exec( - ( - f"test -d {quoted_runtime} && test ! -L {quoted_runtime} && " - f"test -f {quoted_log} && test ! -L {quoted_log} && " - f"test -f {quoted_state} && test ! -L {quoted_state} && " - f"chown 0:0 {quoted_runtime} {quoted_log} {quoted_state} && " - f"chmod 700 {quoted_runtime} && " - f"chmod 600 {quoted_log} {quoted_state}" - ), - user="root", - timeout_sec=10, - ) - if secured.return_code != 0: - logger.warning("Provider capture custody artifact hardening failed") - return False - access = await sandbox.exec( - ( - f"for artifact in {quoted_log} {quoted_state}; do " - 'if cat "$artifact" >/dev/null 2>&1; then exit 0; fi; ' - 'if [ -r "$artifact" ] || [ -w "$artifact" ]; then exit 0; fi; ' - "done; " - f"if [ -r {quoted_runtime} ] || [ -w {quoted_runtime} ] || " - f"[ -x {quoted_runtime} ]; then exit 0; fi; " - "if command -v sudo >/dev/null 2>&1 && " - f"sudo -n cat {quoted_state} >/dev/null 2>&1; then exit 0; fi; " - "if command -v doas >/dev/null 2>&1 && " - f"doas -n cat {quoted_state} >/dev/null 2>&1; then exit 0; fi; " - "if id -G 2>/dev/null | tr ' ' '\\n' | grep -qx 0; " - "then exit 0; fi; " - "effective_caps=$(awk '/^CapEff:/ {print $2}' " - "/proc/self/status 2>/dev/null); " - 'if [ -n "$effective_caps" ] && ' - "[ \"$effective_caps\" != '0000000000000000' ]; then exit 0; fi; " - "for privilege_socket in /var/run/docker.sock " - "/run/containerd/containerd.sock /run/podman/podman.sock; do " - 'if [ -S "$privilege_socket" ] && ' - '[ -r "$privilege_socket" ] && ' - '[ -w "$privilege_socket" ]; then exit 0; fi; done; ' - "exit 1" - ), - user=sandbox_user, - timeout_sec=10, - ) - if access.return_code != 1: - logger.warning( - "Provider capture custody probe found agent-accessible root data" - ) - return False - verified = await sandbox.exec( - ( - f"test \"$(stat -c '%u:%a' {quoted_runtime})\" = '0:700' && " - f"test \"$(stat -c '%u:%a' {quoted_log})\" = '0:600' && " - f"test \"$(stat -c '%u:%a' {quoted_state})\" = '0:600'" - ), - user="root", - timeout_sec=10, - ) - if verified.return_code != 0: - logger.warning("Provider capture custody probe integrity check failed") - return False - return True - except Exception as exc: - logger.warning("Provider capture custody artifact probe failed: %s", exc) - return False - - async def stop_litellm_runtime(runtime: Any | None) -> None: if runtime is None: return diff --git a/src/benchflow/providers/resources/provider_capture_custody.sh b/src/benchflow/providers/resources/provider_capture_custody.sh new file mode 100644 index 000000000..fe48b36cd --- /dev/null +++ b/src/benchflow/providers/resources/provider_capture_custody.sh @@ -0,0 +1,62 @@ +#!/bin/sh +set -eu + +mode=$1 +runtime_dir=$2 +callback_log=$runtime_dir/callback.jsonl +capture_state=$runtime_dir/capture_state.json + +case "$mode" in + harden) + test -d "$runtime_dir" + test ! -L "$runtime_dir" + test -f "$callback_log" + test ! -L "$callback_log" + test -f "$capture_state" + test ! -L "$capture_state" + chown 0:0 "$runtime_dir" "$callback_log" "$capture_state" + chmod 700 "$runtime_dir" + chmod 600 "$callback_log" "$capture_state" + ;; + probe) + for artifact in "$callback_log" "$capture_state"; do + if cat "$artifact" >/dev/null 2>&1; then exit 0; fi + if [ -r "$artifact" ] || [ -w "$artifact" ]; then exit 0; fi + done + if [ -r "$runtime_dir" ] || [ -w "$runtime_dir" ] || [ -x "$runtime_dir" ]; then + exit 0 + fi + if command -v sudo >/dev/null 2>&1 && \ + sudo -n cat "$capture_state" >/dev/null 2>&1; then + exit 0 + fi + if command -v doas >/dev/null 2>&1 && \ + doas -n cat "$capture_state" >/dev/null 2>&1; then + exit 0 + fi + if id -G 2>/dev/null | tr ' ' '\n' | grep -qx 0; then exit 0; fi + effective_caps=$(awk '/^CapEff:/ {print $2}' /proc/self/status 2>/dev/null || true) + if [ -n "$effective_caps" ] && [ "$effective_caps" != 0000000000000000 ]; then + exit 0 + fi + for privilege_socket in \ + /var/run/docker.sock \ + /run/containerd/containerd.sock \ + /run/podman/podman.sock; do + if [ -S "$privilege_socket" ] && \ + [ -r "$privilege_socket" ] && \ + [ -w "$privilege_socket" ]; then + exit 0 + fi + done + exit 1 + ;; + verify) + test "$(stat -c '%u:%a' "$runtime_dir")" = 0:700 + test "$(stat -c '%u:%a' "$callback_log")" = 0:600 + test "$(stat -c '%u:%a' "$capture_state")" = 0:600 + ;; + *) + exit 2 + ;; +esac diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 59a1e2f70..0866688d6 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -230,6 +230,12 @@ ) logger = logging.getLogger(__name__) + + +def _sandbox_user_home(sandbox_user: str | None) -> str: + return f"/home/{sandbox_user}" if sandbox_user else "/root" + + _SETUP_COMMAND_LOCK_SAFE_RE = re.compile(r"[^A-Za-z0-9._-]+") # Lifecycle phases from verify() onward. The agent will not run again in this @@ -1221,7 +1227,7 @@ async def install_agent(self) -> None: workspace=self._agent_cwd, timeout_sec=cfg.sandbox_setup_timeout, ) - cred_home = f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" + cred_home = _sandbox_user_home(cfg.sandbox_user) await self._planes.write_credential_files( self._env, agent_name, @@ -1326,9 +1332,7 @@ async def connect(self) -> None: ) llm_capture = getattr(self, "_llm_capture", None) if llm_capture is not None: - credential_home = ( - f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" - ) + credential_home = _sandbox_user_home(cfg.sandbox_user) llm_capture.bind_provider_capture_trust( agent=cfg.primary_agent, model=cfg.primary_model, @@ -1380,9 +1384,7 @@ async def connect(self) -> None: self._bind_llm_capture_session( agent=cfg.primary_agent, model=cfg.primary_model, - credential_home=( - f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" - ), + credential_home=_sandbox_user_home(cfg.sandbox_user), ) self._native_usage_checkpoint = None self._reapply_ask_user_handler() @@ -2398,7 +2400,7 @@ async def connect_as(self, role: Role) -> None: needs_role_credentials = ( role_agent_differs or role.model != cfg.primary_model or bool(role.env) ) - cred_home = f"/home/{cfg.sandbox_user}" if cfg.sandbox_user else "/root" + cred_home = _sandbox_user_home(cfg.sandbox_user) if role_agent_differs: if cfg.skip_agent_install: agent_cfg = None diff --git a/src/benchflow/sandbox/files.py b/src/benchflow/sandbox/files.py new file mode 100644 index 000000000..bc82f90d2 --- /dev/null +++ b/src/benchflow/sandbox/files.py @@ -0,0 +1,25 @@ +"""Small host-to-sandbox file transfer primitives.""" + +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import Any + + +async def upload_private_text( + sandbox: Any, + text: str, + target_path: str, + *, + suffix: str, +) -> None: + """Upload text with an explicit owner-only mode and remove the host temp.""" + + with tempfile.NamedTemporaryFile("w", suffix=suffix, delete=False) as temporary: + temporary.write(text) + source_path = Path(temporary.name) + try: + await sandbox.upload_file(source_path, target_path, mode="600") + finally: + source_path.unlink(missing_ok=True) diff --git a/src/benchflow/trajectories/export.py b/src/benchflow/trajectories/export.py index 347c4f8ff..5e17efe33 100644 --- a/src/benchflow/trajectories/export.py +++ b/src/benchflow/trajectories/export.py @@ -33,7 +33,7 @@ from benchflow.adapters.ors import ORSAdapter from benchflow.rewards.protocol import VerifyResult from benchflow.trajectories._export_common import aggregate_rollout_jsonl -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj # Canonical artifact locations (see issue #385). ROLLOUT_ARTIFACT_RELPATH = "trainer/verifiers.jsonl" diff --git a/src/benchflow/trajectories/export_adp.py b/src/benchflow/trajectories/export_adp.py index ac23c54b6..62edb60c2 100644 --- a/src/benchflow/trajectories/export_adp.py +++ b/src/benchflow/trajectories/export_adp.py @@ -50,7 +50,7 @@ aggregate_rollout_jsonl, content_blocks_to_text, ) -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj logger = logging.getLogger(__name__) diff --git a/src/benchflow/trajectories/export_atif.py b/src/benchflow/trajectories/export_atif.py index d96dcf87d..8ffd96559 100644 --- a/src/benchflow/trajectories/export_atif.py +++ b/src/benchflow/trajectories/export_atif.py @@ -41,7 +41,7 @@ from benchflow._utils.json_safe import dumps_finite from benchflow.trajectories._export_common import ThoughtBuffer, content_blocks_to_text -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj ATIF_SCHEMA_VERSION = "ATIF-v1.7" diff --git a/src/benchflow/trajectories/export_prime_sft.py b/src/benchflow/trajectories/export_prime_sft.py index 38a7e0d1a..68bbe89e4 100644 --- a/src/benchflow/trajectories/export_prime_sft.py +++ b/src/benchflow/trajectories/export_prime_sft.py @@ -19,10 +19,9 @@ from benchflow._utils.json_safe import dumps_finite, scrub_non_finite from benchflow.trajectories.llm_capture_manifest import ( - capture_artifact_allows_training, - read_llm_trajectory_manifest, + rollout_capture_is_training_grade, ) -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj PrimeSftRowMode = Literal["rollout", "exchange"] @@ -1221,9 +1220,8 @@ def convert_benchflow_rollouts_to_prime_sft_rows( trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) - capture_manifest = read_llm_trajectory_manifest(rollout_dir) - if not capture_artifact_allows_training( - capture_manifest, + if not rollout_capture_is_training_grade( + rollout_dir, exchanges=exchanges, ): stats.skipped_insufficient_capture_fidelity += 1 diff --git a/src/benchflow/trajectories/export_trl_sft.py b/src/benchflow/trajectories/export_trl_sft.py index d83dcb405..0166a2ac9 100644 --- a/src/benchflow/trajectories/export_trl_sft.py +++ b/src/benchflow/trajectories/export_trl_sft.py @@ -25,10 +25,9 @@ validate_prime_sft_row, ) from benchflow.trajectories.llm_capture_manifest import ( - capture_artifact_allows_training, - read_llm_trajectory_manifest, + rollout_capture_is_training_grade, ) -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj TrlSftRowMode = Literal["rollout", "exchange"] TrlSftContextPolicy = Literal["full", "message-window"] @@ -324,9 +323,8 @@ def convert_benchflow_rollouts_to_trl_sft_rows( continue trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" exchanges = load_llm_trajectory_jsonl(trajectory_path, strict=True) - capture_manifest = read_llm_trajectory_manifest(rollout_dir) - if not capture_artifact_allows_training( - capture_manifest, + if not rollout_capture_is_training_grade( + rollout_dir, exchanges=exchanges, ): stats.skipped_insufficient_capture_fidelity += 1 diff --git a/src/benchflow/trajectories/io.py b/src/benchflow/trajectories/io.py new file mode 100644 index 000000000..37a4e893d --- /dev/null +++ b/src/benchflow/trajectories/io.py @@ -0,0 +1,27 @@ +"""Durable filesystem primitives shared by trajectory artifacts.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + + +def atomic_write_text(path: Path, payload: str) -> None: + """Replace *path* atomically after writing the complete text payload.""" + + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + "w", + dir=path.parent, + prefix=f".{path.name}.", + delete=False, + ) as temporary: + temporary.write(payload) + temporary.flush() + os.fsync(temporary.fileno()) + temporary_path = Path(temporary.name) + try: + os.replace(temporary_path, path) + finally: + temporary_path.unlink(missing_ok=True) diff --git a/src/benchflow/trajectories/llm_capture.py b/src/benchflow/trajectories/llm_capture.py index c5dc28f69..addbed349 100644 --- a/src/benchflow/trajectories/llm_capture.py +++ b/src/benchflow/trajectories/llm_capture.py @@ -5,7 +5,6 @@ import hashlib import json import logging -import os from contextlib import suppress from dataclasses import replace from datetime import datetime @@ -13,6 +12,7 @@ from typing import Any from benchflow.agents.env import uses_native_subscription_auth +from benchflow.trajectories.io import atomic_write_text from benchflow.trajectories.llm_capture_manifest import ( LLM_TRAJECTORY_FILENAME, AuthMode, @@ -44,6 +44,7 @@ sanitized_capture_error, ) from benchflow.trajectories.native_capture_parsers import project_acp_trajectory +from benchflow.trajectories.redaction import jsonl_payload_is_redacted logger = logging.getLogger(__name__) @@ -390,7 +391,9 @@ async def finalize( model_call_seen=model_call_seen, fallback_auth=self.manifest.auth_mode, ) - write_exchange_records(self.trajectory_path, assembly.records) + payload_redacted = write_exchange_records( + self.trajectory_path, assembly.records + ) self.manifest.auth_mode = assembly.auth_mode self._finish_manifest( status=( @@ -401,6 +404,7 @@ async def finalize( exchange_count=len(assembly.records), request_complete=assembly.request_complete, response_complete=assembly.response_complete, + payload_redacted=payload_redacted, missing_fields=assembly.missing_fields, errors=assembly.errors, role_captures=assembly.role_captures, @@ -422,7 +426,7 @@ def record_failure(self, error: object, *, model_call_seen: bool) -> None: self.manifest.finished_at = datetime.now() exchange_count = _valid_jsonl_row_count(self.trajectory_path) if exchange_count is None: - _atomic_replace_text(self.trajectory_path, "") + atomic_write_text(self.trajectory_path, "") exchange_count = 0 rows_preserved = exchange_count > 0 self._finish_manifest( @@ -442,6 +446,7 @@ def record_failure(self, error: object, *, model_call_seen: bool) -> None: exchange_count=exchange_count, request_complete=False, response_complete=False, + payload_redacted=jsonl_payload_is_redacted(self.trajectory_path), missing_fields=( sorted( { @@ -470,6 +475,7 @@ def _finish_manifest( exchange_count: int, request_complete: bool, response_complete: bool, + payload_redacted: bool, missing_fields: list[str] | None = None, errors: list[str] | None = None, role_captures: list[LLMRoleCapture] | None = None, @@ -480,6 +486,7 @@ def _finish_manifest( self.manifest.exchange_count = exchange_count self.manifest.request_complete = request_complete self.manifest.response_complete = response_complete + self.manifest.payload_redacted = payload_redacted self.manifest.missing_fields = sorted(set(missing_fields or [])) self.manifest.errors = [sanitized_capture_error(item) for item in errors or []] self.manifest.role_captures = role_captures or [] @@ -547,12 +554,6 @@ def _resolve_auth_mode( return AuthMode.OAUTH_SUBSCRIPTION -def _atomic_replace_text(path: Path, payload: str) -> None: - temporary = path.with_suffix(path.suffix + ".tmp") - temporary.write_text(payload) - os.replace(temporary, path) - - def _valid_jsonl_row_count(path: Path) -> int | None: """Count valid JSON objects, distinguishing an empty artifact from corruption.""" diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index a016a1228..cd8eaa977 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -8,7 +8,6 @@ from __future__ import annotations import json -import os from collections.abc import Sequence from datetime import datetime from enum import StrEnum @@ -17,6 +16,8 @@ from pydantic import BaseModel, Field, ValidationError +from benchflow.trajectories.io import atomic_write_text + LLM_TRAJECTORY_FILENAME = "llm_trajectory.jsonl" LLM_TRAJECTORY_MANIFEST_FILENAME = "llm_trajectory.manifest.json" LLM_TRAJECTORY_SCHEMA_VERSION = 2 @@ -114,7 +115,7 @@ class LLMTrajectoryManifest(BaseModel): exchange_count: int = 0 request_complete: bool = False response_complete: bool = False - payload_redacted: bool = True + payload_redacted: bool = False started_at: datetime finished_at: datetime | None = None missing_fields: list[str] = Field(default_factory=list) @@ -135,7 +136,7 @@ def initialize_llm_trajectory_artifacts( trajectory_dir = rollout_dir / "trajectory" trajectory_dir.mkdir(parents=True, exist_ok=True) trajectory_path = trajectory_dir / LLM_TRAJECTORY_FILENAME - _atomic_write_text(trajectory_path, "") + atomic_write_text(trajectory_path, "") manifest = LLMTrajectoryManifest( agent=agent, model=model, @@ -151,7 +152,7 @@ def write_llm_trajectory_manifest( ) -> None: path = rollout_dir / "trajectory" / LLM_TRAJECTORY_MANIFEST_FILENAME payload = json.dumps(manifest.model_dump(mode="json"), indent=2, sort_keys=True) - _atomic_write_text(path, payload + "\n") + atomic_write_text(path, payload + "\n") def read_llm_trajectory_manifest(rollout_dir: Path) -> dict[str, Any] | None: @@ -207,6 +208,19 @@ def capture_artifact_allows_training( return not any(_exchange_requires_manifest(exchange) for exchange in exchanges) +def rollout_capture_is_training_grade( + rollout_dir: Path, + *, + exchanges: Sequence[dict[str, Any]], +) -> bool: + """Apply the canonical training-admission contract for one rollout.""" + + return capture_artifact_allows_training( + read_llm_trajectory_manifest(rollout_dir), + exchanges=exchanges, + ) + + def successful_exchanges_have_positive_usage( exchanges: Sequence[dict[str, Any]], ) -> bool: @@ -353,12 +367,12 @@ def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> boo allowed_errors.add(CONTINUATION_SOURCE_AUDIT_ERROR) if has_oauth_capture: allowed_missing_fields.update(_OAUTH_AUDIT_MISSING_FIELDS) - return not ( - REPLAY_PROXY_INGRESS_AUDIT_ERROR not in errors - or not errors.issubset(allowed_errors) - or _REPLAY_AUDIT_MISSING_FIELD not in missing_fields - or not missing_fields.issubset(allowed_missing_fields) - or manifest.get("response_complete") is not True + return bool( + REPLAY_PROXY_INGRESS_AUDIT_ERROR in errors + and errors.issubset(allowed_errors) + and _REPLAY_AUDIT_MISSING_FIELD in missing_fields + and missing_fields.issubset(allowed_missing_fields) + and manifest.get("response_complete") is True ) @@ -431,10 +445,3 @@ def _role_capture_preserves_audit_completion(capture: LLMRoleCapture) -> bool: and capture.request_complete is True and capture.response_complete is True ) - - -def _atomic_write_text(path: Path, payload: str) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_suffix(path.suffix + ".tmp") - temporary.write_text(payload) - os.replace(temporary, path) diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index e22075461..b14f2cdf6 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -3,12 +3,12 @@ from __future__ import annotations import json -import os from dataclasses import dataclass from pathlib import Path from typing import Any from benchflow.providers.litellm_config import safe_model_alias +from benchflow.trajectories.io import atomic_write_text from benchflow.trajectories.llm_capture_manifest import ( LLM_TRAJECTORY_SCHEMA_VERSION, AuthMode, @@ -19,7 +19,7 @@ successful_exchanges_have_positive_usage, ) from benchflow.trajectories.native_capture_parsers import NativeParseResult -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj_with_audit @dataclass(frozen=True) @@ -152,14 +152,13 @@ def load_provider_wire_records( "role_attribution_complete": attribution_complete, "request_complete": request_complete, "response_complete": response_complete, - "payload_redacted": True, } ) if not capture_trusted: metadata["capture_custody"] = _untrusted_provider_custody(target) if not attribution_complete: metadata["role_candidates"] = _role_candidates(targets) - records.append(redact_trajectory_obj(record)) + records.append(record) return records @@ -308,16 +307,25 @@ def assemble_capture( ) -def write_exchange_records(path: Path, records: list[dict[str, Any]]) -> None: - """Atomically replace the JSONL with redacted assembled records.""" +def write_exchange_records(path: Path, records: list[dict[str, Any]]) -> bool: + """Atomically write records and return the verified redaction state.""" - payload = "".join( - json.dumps(redact_trajectory_obj(record), default=str) + "\n" - for record in records - ) - temporary = path.with_suffix(path.suffix + ".tmp") - temporary.write_text(payload) - os.replace(temporary, path) + rendered: list[str] = [] + payload_redacted = True + for record in records: + redacted, row_redacted = redact_trajectory_obj_with_audit(record) + if not isinstance(redacted, dict): + row_redacted = False + redacted = {} + metadata = redacted.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + redacted["metadata"] = metadata + metadata["payload_redacted"] = row_redacted + payload_redacted = payload_redacted and row_redacted + rendered.append(json.dumps(redacted, default=str) + "\n") + atomic_write_text(path, "".join(rendered)) + return payload_redacted def role_captures_for_targets(targets: list[CaptureTarget]) -> list[LLMRoleCapture]: @@ -328,7 +336,7 @@ def role_captures_for_targets(targets: list[CaptureTarget]) -> list[LLMRoleCaptu def _native_bundle_records(bundle: NativeCaptureBundle) -> list[dict[str, Any]]: records: list[dict[str, Any]] = [] - for line in bundle.result.trajectory.to_jsonl(redact_keys=True).splitlines(): + for line in bundle.result.trajectory.to_jsonl(redact_keys=False).splitlines(): if not line.strip(): continue record = json.loads(line) @@ -360,7 +368,7 @@ def _native_bundle_records(bundle: NativeCaptureBundle) -> list[dict[str, Any]]: "role_candidates": _role_candidates(list(bundle.targets)), } ) - records.append(redact_trajectory_obj(record)) + records.append(record) return records diff --git a/src/benchflow/trajectories/native_capture_collection.py b/src/benchflow/trajectories/native_capture_collection.py index dedd0d6ed..88ab0d773 100644 --- a/src/benchflow/trajectories/native_capture_collection.py +++ b/src/benchflow/trajectories/native_capture_collection.py @@ -243,12 +243,21 @@ async def collect( ) with tempfile.TemporaryDirectory(prefix="benchflow-native-llm-") as temporary: local_root = Path(temporary) - raw_claude_result = await self._collect_claude_raw_capture( - env, - local_root=local_root, - claude_targets=claude_targets, - bundles=bundles, - errors=errors, + try: + raw_claude_bundle = await self._collect_claude_raw_capture( + env, + local_root=local_root, + claude_targets=claude_targets, + ) + except Exception as exc: + warning = sanitized_capture_error(exc) + errors.append(warning) + logger.warning("Claude raw LLM capture collection failed: %s", exc) + raw_claude_bundle = None + if raw_claude_bundle is not None: + bundles.append(raw_claude_bundle) + raw_claude_result = ( + raw_claude_bundle.result if raw_claude_bundle is not None else None ) for index, target in enumerate(targets): try: @@ -286,28 +295,20 @@ async def _collect_claude_raw_capture( *, local_root: Path, claude_targets: tuple[CaptureTarget, ...], - bundles: list[NativeCaptureBundle], - errors: list[str], - ) -> NativeParseResult | None: + ) -> NativeCaptureBundle | None: if not self.otel.root_prepared: return None capture_dir = local_root / "capture" - try: - await env.download_dir(self.otel.remote_root, capture_dir) - result = parse_claude_raw_capture( - capture_dir, - agent=(claude_targets[0].agent if claude_targets else self.agent), - session_id=self.session_id, - started_at=self.started_at, - ) - except Exception as exc: - errors.append(sanitized_capture_error(exc)) - logger.warning("Claude raw LLM capture collection failed: %s", exc) - return None + await env.download_dir(self.otel.remote_root, capture_dir) + result = parse_claude_raw_capture( + capture_dir, + agent=(claude_targets[0].agent if claude_targets else self.agent), + session_id=self.session_id, + started_at=self.started_at, + ) if result is None: return None - bundles.append(NativeCaptureBundle(targets=claude_targets, result=result)) - return result + return NativeCaptureBundle(targets=claude_targets, result=result) async def _collect_claude_session_fallback( self, diff --git a/src/benchflow/trajectories/native_capture_parsers.py b/src/benchflow/trajectories/native_capture_parsers.py index 0835acbfe..d4b23722d 100644 --- a/src/benchflow/trajectories/native_capture_parsers.py +++ b/src/benchflow/trajectories/native_capture_parsers.py @@ -552,7 +552,6 @@ def _exchange( "auth_mode": auth_mode, "request_complete": request_complete, "response_complete": response_complete, - "payload_redacted": True, **(extra_metadata or {}), } return LLMExchange( diff --git a/src/benchflow/trajectories/redaction.py b/src/benchflow/trajectories/redaction.py index acabc783b..7e9b23642 100644 --- a/src/benchflow/trajectories/redaction.py +++ b/src/benchflow/trajectories/redaction.py @@ -2,8 +2,10 @@ from __future__ import annotations +import json import re from importlib.resources import files +from pathlib import Path from typing import Any # Human-facing redaction categories. Every canonical pattern below is tagged @@ -365,6 +367,34 @@ def redact_trajectory_obj(obj: Any) -> Any: return obj +def trajectory_obj_is_redacted(obj: Any) -> bool: + """Verify that a JSON-compatible value contains no canonical secret match.""" + + return redact_trajectory_obj(obj) == obj + + +def redact_trajectory_obj_with_audit(obj: Any) -> tuple[Any, bool]: + """Redact one value and independently verify the emitted value is clean.""" + + redacted = redact_trajectory_obj(obj) + return redacted, trajectory_obj_is_redacted(redacted) + + +def jsonl_payload_is_redacted(path: Path) -> bool: + """Verify every non-empty JSONL row is an object with no secret match.""" + + try: + for raw in path.read_text().splitlines(): + if not raw.strip(): + continue + row = json.loads(raw) + if not isinstance(row, dict) or not trajectory_obj_is_redacted(row): + return False + except (OSError, json.JSONDecodeError): + return False + return True + + def _redact_field(key: Any, value: Any) -> Any: """Redact a dict field's *value*, keeping the field name as carrier context. @@ -381,8 +411,6 @@ def _redact_field(key: Any, value: Any) -> Any: if not isinstance(value, str): return value if isinstance(key, str): - import json - probe = json.dumps({key: value}) redacted = redact_trajectory_text(probe) if redacted == probe: diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index cc6ddfdda..f05c704f2 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -33,11 +33,11 @@ validate_prime_sft_row, ) from benchflow.trajectories.llm_capture_manifest import ( - capture_artifact_allows_training, capture_manifest_preserves_audit_completion, read_llm_trajectory_manifest, + rollout_capture_is_training_grade, ) -from benchflow.trajectories.types import redact_trajectory_obj +from benchflow.trajectories.redaction import redact_trajectory_obj from benchflow.usage_tracking import USAGE_SOURCE_AGENT_NATIVE_ACP ROLLOUT_RESULTS_FILENAME = "results.jsonl" @@ -181,9 +181,8 @@ def _llm_steps_from_trajectory( exchanges = load_llm_trajectory_jsonl(path, strict=True) except PrimeSftTrajectoryJsonlError as exc: return _LLMStepsResult([], [], f"Invalid LLM trajectory JSONL: {exc}") - capture_manifest = read_llm_trajectory_manifest(rollout_dir) - if not capture_artifact_allows_training( - capture_manifest, + if not rollout_capture_is_training_grade( + rollout_dir, exchanges=exchanges, ): return _LLMStepsResult([], [], None, capture_contract_rejected=True) diff --git a/src/benchflow/trajectories/types.py b/src/benchflow/trajectories/types.py index dc07481d8..cce405745 100644 --- a/src/benchflow/trajectories/types.py +++ b/src/benchflow/trajectories/types.py @@ -6,23 +6,7 @@ from pydantic import BaseModel, Field -from benchflow.trajectories import redaction as _redaction - -# Keep the historical ``benchflow.trajectories.types`` redaction imports -# working for downstream callers. The implementation lives in the standalone, -# stdlib-only module so the exact same code can be deployed beside LiteLLM. -REDACTION_CATEGORY_API_KEY = _redaction.REDACTION_CATEGORY_API_KEY -REDACTION_CATEGORY_BEARER_TOKEN = _redaction.REDACTION_CATEGORY_BEARER_TOKEN -REDACTION_CATEGORY_CREDENTIAL_FIELD = _redaction.REDACTION_CATEGORY_CREDENTIAL_FIELD -REDACTION_CATEGORY_PASSWORD = _redaction.REDACTION_CATEGORY_PASSWORD -REDACTION_CATEGORY_PRIVATE_KEY = _redaction.REDACTION_CATEGORY_PRIVATE_KEY -REDACTION_CATEGORY_URL_CREDENTIAL = _redaction.REDACTION_CATEGORY_URL_CREDENTIAL -redact_trajectory_obj = _redaction.redact_trajectory_obj -redact_trajectory_text = _redaction.redact_trajectory_text -redact_trajectory_text_with_categories = ( - _redaction.redact_trajectory_text_with_categories -) -redact_trajectory_text_with_count = _redaction.redact_trajectory_text_with_count +from benchflow.trajectories.redaction import redact_trajectory_obj _USAGE_KEYS = { "input_tokens", diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 44a36a3c0..dc4b5bcb6 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -156,17 +156,14 @@ def test_live_forwarder_build_kwargs_resolves_route_offline(): assert kwargs["temperature"] == 0.5 -def test_stitched_trajectory_recorded_prefix_plus_live_suffix(tmp_path): +def test_stitched_trajectory_promotes_recorded_prefix(tmp_path): original = tmp_path / "orig.jsonl" original.write_text('{"a": 1}\n{"b": 2}\n') - live = [exchange(completion(content="LIVE"))] - lines = stitched_trajectory_lines(original, live) - assert len(lines) == 3 + lines = stitched_trajectory_lines(original) + assert len(lines) == 2 first = json.loads(lines[0]) assert first["a"] == 1 assert first["metadata"]["schema_version"] == 2 - last = json.loads(lines[2]) - assert last["response"]["body"]["choices"][0]["message"]["content"] == "LIVE" def test_write_stitched_trajectory_creates_file(tmp_path): @@ -201,6 +198,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p exchange_count=1, request_complete=True, response_complete=True, + payload_redacted=True, started_at="2026-08-29T00:00:00Z", finished_at="2026-08-29T00:01:00Z", ) @@ -396,6 +394,7 @@ def test_root_sandbox_live_suffix_is_retained_but_audit_only(tmp_path): exchange_count=1, request_complete=True, response_complete=True, + payload_redacted=True, started_at="2026-08-29T00:00:00Z", finished_at="2026-08-29T00:01:00Z", ) @@ -743,6 +742,7 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): exchange_count=1, request_complete=True, response_complete=True, + payload_redacted=True, started_at="2026-08-29T00:00:00Z", finished_at="2026-08-29T00:01:00Z", ) @@ -841,7 +841,7 @@ def test_stitching_structurally_redacts_escaped_secret(tmp_path): payload["request"]["body"]["authorization"] = f'Bearer {secret}\\"tail' source.write_text(json.dumps(payload) + "\n") - rendered = stitched_trajectory_lines(source, []) + rendered = stitched_trajectory_lines(source) assert len(rendered) == 1 restored = json.loads(rendered[0]) diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index cbc16572b..9812851d9 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -9,6 +9,7 @@ import httpx import pytest +from benchflow.continue_run import sandbox_replay_runtime as replay_runtime from benchflow.continue_run.replay_proxy import ( ReplayDivergenceError, ReplayProxy, @@ -18,7 +19,12 @@ from benchflow.continue_run.sandbox_proxy import ( SandboxReplayProxy, _ordered_live_exchange_log, - _sandbox_proxy_source, + sandbox_replay_runtime_source, +) +from benchflow.continue_run.sandbox_replay_runtime import ( + ReplayHandler, + ReplayServer, + ReplayState, ) from benchflow.trajectories.redaction import canonical_redaction_module_source @@ -129,9 +135,7 @@ def test_sandbox_forwarding_failure_is_not_logged_as_provider_exchange( ) -> None: """Guards PR #1057 against labeling a synthesized sandbox 500 provider-wire.""" - namespace: dict[str, object] = {} - exec(_sandbox_proxy_source(), namespace) - state = namespace["ReplayState"]( + state = ReplayState( recorded=[], upstream_url="https://provider.invalid/v1", upstream_api_key="test-key", @@ -144,7 +148,7 @@ def test_sandbox_forwarding_failure_is_not_logged_as_provider_exchange( def fail(*_args, **_kwargs): raise TimeoutError("provider unavailable") - monkeypatch.setattr(namespace["urllib"].request, "urlopen", fail) + monkeypatch.setattr(replay_runtime.urllib.request, "urlopen", fail) source, status, body = state.next_response({"messages": [{"role": "user"}]}) assert source == "live" @@ -159,10 +163,8 @@ def fail(*_args, **_kwargs): def test_sandbox_live_exchange_is_redacted_before_journaling(tmp_path) -> None: """Guards PR #1057 against persisting raw continuation secrets.""" - namespace: dict[str, object] = {} - exec(_sandbox_proxy_source(), namespace) live_log = tmp_path / "live.jsonl" - state = namespace["ReplayState"]( + state = ReplayState( recorded=[], upstream_url="https://provider.invalid/v1", upstream_api_key="test-key", @@ -202,7 +204,8 @@ def __init__(self) -> None: async def exec(self, _command, **_kwargs): return SimpleNamespace(return_code=0, stdout="", stderr="") - async def upload_file(self, source, target): + async def upload_file(self, source, target, *, mode): + assert mode == "600" self.uploaded[target] = source.read_text() sandbox = FakeSandbox() @@ -220,17 +223,18 @@ async def upload_file(self, source, target): "from benchflow_trajectory_redaction import" in sandbox.uploaded[f"{proxy.runtime_dir}/replay_proxy.py"] ) + assert sandbox.uploaded[f"{proxy.runtime_dir}/replay_proxy.py"] == ( + sandbox_replay_runtime_source() + ) def test_sandbox_attempt_journal_failure_invalidates_stale_state( tmp_path, monkeypatch: pytest.MonkeyPatch ) -> None: """Guards PR #1057 against completing an unjournaled sandbox call.""" - namespace: dict[str, object] = {} - exec(_sandbox_proxy_source(), namespace) state_path = tmp_path / "state.json" state_path.write_text(json.dumps({"live_attempt_count": 0, "live_error_count": 0})) - state = namespace["ReplayState"]( + state = ReplayState( recorded=[], upstream_url="https://provider.invalid/v1", upstream_api_key="test-key", @@ -243,7 +247,7 @@ def test_sandbox_attempt_journal_failure_invalidates_stale_state( def fail_replace(*_args, **_kwargs): raise OSError("disk full") - monkeypatch.setattr(namespace["os"], "replace", fail_replace) + monkeypatch.setattr(replay_runtime.os, "replace", fail_replace) with pytest.raises(OSError, match="disk full"): state.next_response({"messages": [{"role": "user"}]}) @@ -255,9 +259,7 @@ def fail_replace(*_args, **_kwargs): def test_sandbox_quiesce_waits_for_live_handler_before_snapshot(tmp_path) -> None: """Guards PR #1057 against snapshotting before live calls are quiescent.""" - namespace: dict[str, object] = {} - exec(_sandbox_proxy_source(), namespace) - state = namespace["ReplayState"]( + state = ReplayState( recorded=[], upstream_url="https://provider.invalid/v1", upstream_api_key="test-key", @@ -309,10 +311,8 @@ def forward(_request): def test_sandbox_live_exchange_recovery_restores_attempt_order(tmp_path) -> None: """Guards PR #1057 against stitching sandbox calls in completion order.""" - namespace: dict[str, object] = {} - exec(_sandbox_proxy_source(), namespace) live_log = tmp_path / "live.jsonl" - state = namespace["ReplayState"]( + state = ReplayState( recorded=[], upstream_url="https://provider.invalid/v1", upstream_api_key="test-key", @@ -355,9 +355,7 @@ def test_sandbox_quiesce_closes_listener_and_drains_accepted_handlers( ) -> None: """Guards PR #1057 against terminating a late rejected handler mid-journal.""" - namespace: dict[str, object] = {} - exec(_sandbox_proxy_source(), namespace) - state = namespace["ReplayState"]( + state = ReplayState( recorded=[], upstream_url="https://provider.invalid/v1", upstream_api_key="test-key", @@ -375,9 +373,7 @@ def forward(_request): return 200, completion(content="done"), True state._forward_live = forward - server = namespace["ReplayServer"]( - ("127.0.0.1", 0), namespace["ReplayHandler"], state - ) + server = ReplayServer(("127.0.0.1", 0), ReplayHandler, state) port = server.server_address[1] server_thread = threading.Thread(target=server.serve_forever) server_thread.start() diff --git a/tests/test_eval_artifact_cli.py b/tests/test_eval_artifact_cli.py index 6c0eae188..1ffe04c9b 100644 --- a/tests/test_eval_artifact_cli.py +++ b/tests/test_eval_artifact_cli.py @@ -109,9 +109,11 @@ def test_health_rejects_empty_terminal_llm_capture_manifest(tmp_path: Path) -> N health = build_health_summary(job) assert health["missing_llm_trajectory"] == 0 - assert health["malformed_llm_trajectory"] == 1 + assert health["malformed_llm_trajectory"] == 0 + assert health["non_training_grade_llm_trajectory"] == 1 assert health["rows"][0]["has_llm_trajectory"] is True assert health["rows"][0]["valid_llm_trajectory"] is False + assert health["rows"][0]["well_formed_llm_trajectory"] is True assert health["rows"][0]["llm_trajectory_rows"] == 0 @@ -126,9 +128,11 @@ def test_health_rejects_empty_llm_capture_without_manifest(tmp_path: Path) -> No health = build_health_summary(job) assert health["missing_llm_trajectory"] == 0 - assert health["malformed_llm_trajectory"] == 1 + assert health["malformed_llm_trajectory"] == 0 + assert health["non_training_grade_llm_trajectory"] == 1 assert health["rows"][0]["has_llm_trajectory"] is True assert health["rows"][0]["valid_llm_trajectory"] is False + assert health["rows"][0]["well_formed_llm_trajectory"] is True assert health["rows"][0]["llm_trajectory_rows"] == 0 @@ -143,9 +147,11 @@ def test_health_rejects_sidecarless_schema_v2_llm_capture(tmp_path: Path) -> Non health = build_health_summary(job) assert health["missing_llm_trajectory"] == 0 - assert health["malformed_llm_trajectory"] == 1 + assert health["malformed_llm_trajectory"] == 0 + assert health["non_training_grade_llm_trajectory"] == 1 assert health["rows"][0]["has_llm_trajectory"] is True assert health["rows"][0]["valid_llm_trajectory"] is False + assert health["rows"][0]["well_formed_llm_trajectory"] is True assert health["rows"][0]["llm_trajectory_rows"] == 1 diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index b31bf4f39..2d1c27139 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -389,7 +389,8 @@ async def test_sandbox_litellm_launch_keeps_secrets_off_command_line(): ] assert len(redaction_files) == 1 assert ( - sandbox.uploaded[redaction_files[0]] == runtime_mod._redaction_module_source() + sandbox.uploaded[redaction_files[0]] + == runtime_mod.canonical_redaction_module_source() ) launch_command = next(call for call in sandbox.exec_calls if "launcher.py" in call) assert f"rc=$?; rm -f {launch_files[0]}; exit $rc" in launch_command @@ -653,7 +654,7 @@ def test_runtime_packages_canonical_redactor_verbatim(tmp_path): runtime_mod._write_runtime_files(tmp_path, config={"model_list": []}) packaged = tmp_path / "benchflow_trajectory_redaction.py" - assert packaged.read_text() == runtime_mod._redaction_module_source() + assert packaged.read_text() == runtime_mod.canonical_redaction_module_source() assert "def redact_trajectory_obj" in packaged.read_text() env = dict(os.environ) env["PYTHONPATH"] = str(tmp_path) diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 640cd4b29..3ef5195c4 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -227,6 +227,38 @@ def test_callback_preserves_post_transform_provider_request_body(): assert "input" not in record["request"]["body"] +def test_callback_uses_responses_path_for_responses_provider_body() -> None: + """Guards PR #1057 against labeling Responses wire data as chat completions.""" + + logger = _callback_namespace()["BenchFlowLiteLLMLogger"]() + now = datetime.now() + call_id = "call-responses-body" + provider_body = {"model": "gpt-5.5", "input": "hello"} + logger.log_pre_api_call( + provider_body["model"], + None, + { + "litellm_call_id": call_id, + "additional_args": {"complete_input_dict": provider_body}, + }, + ) + + record = logger._base_record( + { + "litellm_call_id": call_id, + "model": "gpt-5.5", + "input": "hello", + "call_type": "aresponses", + }, + now, + now, + ) + + assert record["request_complete"] is True + assert record["request"]["path"] == "/v1/responses" + assert record["request"]["body"] == provider_body + + def test_callback_marks_proxy_ingress_request_incomplete(): """Guards PR #1057 against promoting reconstructed proxy-ingress parameters.""" diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 6d3994357..32ec8ad6c 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -167,6 +167,43 @@ async def fake_sandbox_start(**kwargs): assert updated["LLM_MODEL"].startswith("openai/benchflow-aws-bedrock") +class _CustodyProbeSandbox: + def __init__(self, *, agent_probe_return_code: int) -> None: + self.agent_probe_return_code = agent_probe_return_code + self.probe_path = "" + + async def upload_file(self, source, target, *, mode): + assert mode == "600" + assert "callback.jsonl" in source.read_text() + self.probe_path = target + + async def exec(self, command, **kwargs): + if command == "id -u -- agent": + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") + if command.startswith("chown 0:0"): + assert self.probe_path in command + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") + if f"{self.probe_path} harden " in command: + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") + if f"{self.probe_path} probe " in command: + assert kwargs["user"] == "agent" + return SimpleNamespace( + return_code=self.agent_probe_return_code, + stdout="", + stderr="", + ) + if f"{self.probe_path} verify " in command: + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") + if command.startswith("rm -f"): + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") + raise AssertionError(f"unexpected custody command: {command}") + + @pytest.mark.asyncio async def test_non_root_agent_keeps_sandbox_gateway_capture_trusted(monkeypatch): """Guards PR #1057's root/non-root provider-capture custody boundary.""" @@ -176,32 +213,6 @@ async def fake_sandbox_start(**kwargs): monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) - class NonRootSandbox: - async def exec(self, command, **kwargs): - if command == "id -u -- agent": - assert kwargs["user"] == "root" - return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") - if command.startswith("test -d"): - assert kwargs["user"] == "root" - assert "callback.jsonl" in command - assert "capture_state.json" in command - assert "chmod 700" in command - assert "chmod 600" in command - return SimpleNamespace(return_code=0, stdout="", stderr="") - if kwargs["user"] == "agent": - assert "callback.jsonl" in command - assert "capture_state.json" in command - assert "sudo -n cat" in command - assert "CapEff" in command - assert "/var/run/docker.sock" in command - return SimpleNamespace(return_code=1, stdout="", stderr="") - if command.startswith('test "$(stat -c'): - assert kwargs["user"] == "root" - assert "callback.jsonl" in command - assert "capture_state.json" in command - return SimpleNamespace(return_code=0, stdout="", stderr="") - raise AssertionError(f"unexpected custody command: {command}") - _, provider_runtime = await ensure_litellm_runtime( agent="openhands", agent_env={ @@ -212,7 +223,7 @@ async def exec(self, command, **kwargs): runtime=None, environment="daytona", session_id="run-non-root", - sandbox=NonRootSandbox(), + sandbox=_CustodyProbeSandbox(agent_probe_return_code=1), sandbox_user="agent", ) @@ -229,19 +240,6 @@ async def fake_sandbox_start(**kwargs): monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) - class PrivilegedSandbox: - async def exec(self, command, **kwargs): - if command == "id -u -- agent": - return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") - if command.startswith("test -d"): - return SimpleNamespace(return_code=0, stdout="", stderr="") - if kwargs["user"] == "agent": - assert "callback.jsonl" in command - assert "capture_state.json" in command - assert "sudo -n cat" in command - return SimpleNamespace(return_code=0, stdout="", stderr="") - raise AssertionError(f"unexpected custody command: {command}") - _, provider_runtime = await ensure_litellm_runtime( agent="openhands", agent_env={ @@ -252,7 +250,7 @@ async def exec(self, command, **kwargs): runtime=None, environment="daytona", session_id="run-privileged-agent", - sandbox=PrivilegedSandbox(), + sandbox=_CustodyProbeSandbox(agent_probe_return_code=0), sandbox_user="agent", ) @@ -271,19 +269,6 @@ async def fake_sandbox_start(**kwargs): monkeypatch.setattr(runtime_mod, "_start_sandbox_litellm", fake_sandbox_start) - class WritableArtifactSandbox: - async def exec(self, command, **kwargs): - if command == "id -u -- agent": - return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") - if command.startswith("test -d"): - return SimpleNamespace(return_code=0, stdout="", stderr="") - if kwargs["user"] == "agent": - assert "callback.jsonl" in command - assert "capture_state.json" in command - # Exit 0 is the probe's contract for any real-artifact access. - return SimpleNamespace(return_code=0, stdout="", stderr="") - raise AssertionError(f"unexpected custody command: {command}") - _, provider_runtime = await ensure_litellm_runtime( agent="openhands", agent_env={ @@ -294,7 +279,7 @@ async def exec(self, command, **kwargs): runtime=None, environment="daytona", session_id="run-writable-artifact", - sandbox=WritableArtifactSandbox(), + sandbox=_CustodyProbeSandbox(agent_probe_return_code=0), sandbox_user="agent", ) diff --git a/tests/test_live_llm_trajectory.py b/tests/test_live_llm_trajectory.py index 7b568c5c4..415a12e2a 100644 --- a/tests/test_live_llm_trajectory.py +++ b/tests/test_live_llm_trajectory.py @@ -157,6 +157,7 @@ async def test_host_proxy_mirrors_callback_before_stop(tmp_path, monkeypatch): stderr_path=tmp_path / "stderr.log", session_id="run", agent_name="opencode", + capture_state_path=tmp_path / "capture_state.json", ) process.start_live_capture(output_path) @@ -227,6 +228,7 @@ async def test_daytona_proxy_incrementally_mirrors_callback(tmp_path, monkeypatc stderr_path="/tmp/runtime/stderr", session_id="run", agent_name="opencode", + capture_state_path="/tmp/runtime/capture_state.json", ) process.start_live_capture(output_path) @@ -256,6 +258,7 @@ def _sandbox_process(sandbox) -> SandboxLiteLLMProcess: stderr_path="/tmp/runtime/stderr", session_id="run", agent_name="prime-agent", + capture_state_path="/tmp/runtime/capture_state.json", ) @@ -457,6 +460,7 @@ async def test_daytona_proxy_uses_transient_exec_for_callback_poll(tmp_path): stderr_path="/tmp/runtime/stderr", session_id="run", agent_name="openhands", + capture_state_path="/tmp/runtime/capture_state.json", ) chunk = await process._read_callback_chunk(0, 24 * 1024) diff --git a/tests/test_redaction_resource.py b/tests/test_redaction_resource.py index e538c0743..8bc797e5a 100644 --- a/tests/test_redaction_resource.py +++ b/tests/test_redaction_resource.py @@ -3,6 +3,7 @@ import tomllib from pathlib import Path +from benchflow.continue_run.sandbox_proxy import sandbox_replay_runtime_source from benchflow.trajectories.redaction import canonical_redaction_module_source @@ -19,3 +20,7 @@ def test_wheel_packages_canonical_redactor_as_data() -> None: canonical_redaction_module_source() == (root / "src/benchflow/trajectories/redaction.py").read_text() ) + assert ( + sandbox_replay_runtime_source() + == (root / "src/benchflow/continue_run/sandbox_replay_runtime.py").read_text() + ) diff --git a/tests/test_train_cli.py b/tests/test_train_cli.py index 67fae4328..2dfadc912 100644 --- a/tests/test_train_cli.py +++ b/tests/test_train_cli.py @@ -468,6 +468,74 @@ def test_train_validate_source_health_requirements(tmp_path: Path) -> None: payload = json.loads(result.output) assert payload["source_health"]["total_rows"] == 1 assert payload["source_health"]["rows_with_tool_calls"] == 1 + assert payload["source_health"]["non_training_grade_llm_trajectory"] == 0 + + +def test_train_validate_rejects_non_training_grade_source_trajectory( + tmp_path: Path, +) -> None: + """Guards PR #1057 against admitting audit-only source trajectories.""" + + jobs = tmp_path / "jobs" + rollout = jobs / "run" / "task-a__abc123" + _write_rollout(rollout) + trajectory = rollout / "trajectory" + (trajectory / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "schema_version": 2, + "status": "partial", + "capture_source": "claude_otel_raw_body", + "capture_fidelity": "agent_session", + "auth_mode": "oauth_subscription", + "agent": "claude-agent-acp", + "model": "claude-haiku-4-5-20251001", + "session_id": "session", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + "started_at": "2026-08-29T00:00:00Z", + "finished_at": "2026-08-29T00:00:01Z", + "missing_fields": [], + "errors": [], + "role_captures": [], + } + ), + encoding="utf-8", + ) + out = tmp_path / "train.jsonl" + conversion = runner.invoke( + app, + [ + "train", + "convert", + str(jobs), + "--out", + str(out), + "--expected-rows", + "0", + ], + ) + assert conversion.exit_code == 0, conversion.output + + result = runner.invoke( + app, + [ + "train", + "validate", + str(out), + "--source-jobs", + str(jobs), + "--expected-rows", + "0", + "--require-llm-trajectory", + ], + ) + + assert result.exit_code == 1 + assert "non-training-grade" in result.output + assert "llm_trajectory.jsonl" in result.output def test_train_convert_rejects_malformed_llm_jsonl(tmp_path: Path) -> None: diff --git a/tests/trajectories/test_native_llm_capture.py b/tests/trajectories/test_native_llm_capture.py index f937c029c..600e0724c 100644 --- a/tests/trajectories/test_native_llm_capture.py +++ b/tests/trajectories/test_native_llm_capture.py @@ -545,6 +545,32 @@ def test_capture_initialization_truncates_reused_rollout_jsonl(tmp_path: Path) - assert manifest["session_id"] == "second-session" +def test_capture_failure_derives_unredacted_state_from_preserved_rows( + tmp_path: Path, +) -> None: + """Guards PR #1057 against self-attesting an unscanned failure artifact.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="codex-acp", + model="gpt-5.6", + session_id="failed-session", + started_at=STARTED_AT, + ) + secret = "sk-ant-api03-" + "A" * 40 + capture.trajectory_path.write_text( + json.dumps({"request": {"body": {"secret": secret}}}) + "\n" + ) + + capture.record_failure(RuntimeError("finalization stopped"), model_call_seen=True) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["payload_redacted"] is False + assert manifest["status"] == CaptureStatus.CAPTURE_FAILED + + @pytest.mark.parametrize( ("auth_json", "expected"), [ From ebd0e04d060792aa52ec40cd3822f40633585d0f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 17:20:48 -0700 Subject: [PATCH 51/74] reject contradictory capture manifests --- .../trajectories/llm_capture_manifest.py | 2 + .../test_llm_capture_training_contract.py | 41 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index cd8eaa977..94a9e2b0a 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -185,6 +185,8 @@ def capture_manifest_allows_training( and manifest.get("request_complete") is True and manifest.get("response_complete") is True and manifest.get("payload_redacted") is True + and manifest.get("missing_fields") == [] + and manifest.get("errors") == [] and exchange_count > 0 and expected_count == exchange_count ) diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 8771e9ecf..824b85d8e 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -214,6 +214,47 @@ def test_unredacted_provider_capture_fails_closed_for_training( assert row["is_completed"] is False +@pytest.mark.parametrize( + ("contradictory_field", "value"), + [ + ("errors", ["capture journal missing"]), + ("missing_fields", ["provider_request"]), + ], +) +def test_complete_manifest_with_capture_gaps_fails_closed_for_training( + tmp_path: Path, + contradictory_field: str, + value: list[str], +) -> None: + """Guards PR #1057 against training on contradictory complete manifests.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="provider_wire", schema_version=2) + manifest = { + "status": "complete", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + "missing_fields": [], + "errors": [], + } + manifest[contradictory_field] = value + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps(manifest), + encoding="utf-8", + ) + + assert not capture_manifest_allows_training(manifest, exchange_count=1) + row = _build_results_row(tmp_path, agent_result={"total_tokens": 2}) + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is False + + def test_provider_capture_without_positive_usage_is_not_training_ready( tmp_path: Path, ) -> None: From e685ab66f8e8c89395ee00b319e09cb1d6db43d4 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 17:44:41 -0700 Subject: [PATCH 52/74] close llm capture provenance gaps --- src/benchflow/agents/registry.py | 12 +- src/benchflow/continue_run/orchestrator.py | 22 +++- src/benchflow/continue_run/replay_proxy.py | 10 +- src/benchflow/continue_run/sandbox_proxy.py | 13 +++ .../continue_run/sandbox_replay_runtime.py | 5 + .../continue_run/trajectory_artifacts.py | 47 +++++++- .../trajectories/llm_capture_manifest.py | 46 ++++++++ tests/continue_run/test_orchestrator.py | 103 +++++++++++++++++- tests/continue_run/test_replay_proxy.py | 62 +++++++++++ tests/test_opencode_family_proxy_tracking.py | 50 +++++++++ .../test_llm_capture_training_contract.py | 59 ++++++++++ 11 files changed, 413 insertions(+), 16 deletions(-) diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index 253f602b3..e10f41987 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -316,9 +316,9 @@ def _json_settings_merge(path: str, mutator: str) -> str: # routes through the proxy. def _opencode_family_proxy_wrapper_install(binary: str, config_path: str) -> str: """Install ``/opt/benchflow/bin/-proxy``: a thin wrapper that, in - proxy mode, registers the LiteLLM gateway alias under a dedicated - OpenCode provider, then execs the isolated agent binary. Idempotent - (preserves existing config); no-op outside proxy mode. + proxy mode, replaces the provider map with the dedicated LiteLLM gateway + provider, then execs the isolated agent binary. Non-provider settings are + preserved; no-op outside proxy mode. The gateway alias is registered under the ``{OPENCODE_PROXY_PROVIDER_ID}`` provider using ``@ai-sdk/openai-compatible`` — NOT the built-in ``openai`` @@ -357,7 +357,11 @@ def _opencode_family_proxy_wrapper_install(binary: str, config_path: str) -> str " const d = text ? JSON.parse(text) : {};", # Dedicated provider id (see docstring) → chat completions, not the # Responses API the built-in ``openai`` id is hard-coded to. - " const providers = d.provider ||= {};", + # An image can carry a pre-existing provider with a literal key or + # endpoint. Retaining it would let the agent bypass the capture + # proxy even after BenchFlow strips provider credentials from the + # process environment, so proxy mode owns the entire provider map. + " const providers = d.provider = {};", f" const prov = providers[{provider_id!r}] ||= {{}};", ' prov.npm = "@ai-sdk/openai-compatible";', ' prov.name ||= "BenchFlow Gateway";', diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 501db3383..08f7f8784 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -540,6 +540,7 @@ async def continue_run( rollout_dir, run.path / "trajectory" / "llm_trajectory.jsonl", router.live_exchanges, + recorded_consumed_count=router.recorded_consumed_count, live_model=live_model, ) refresh_stitched_trajectory_manifest( @@ -548,6 +549,7 @@ async def continue_run( original_model=run.model, live_model=live_model, n_recorded=run.n_recorded_exchanges, + n_recorded_consumed=router.recorded_consumed_count, n_live=len(router.live_exchanges), live_attempt_count=router.live_attempt_count, live_errors=list(router.live_errors), @@ -557,7 +559,7 @@ async def continue_run( live_model=live_model, usage=summarize_llm_trajectory_usage( stitched_path, - n_recorded=run.n_recorded_exchanges, + n_recorded=router.recorded_consumed_count, ), environment=run.environment, ) @@ -566,7 +568,7 @@ async def continue_run( rollout_dir=rollout_dir, rewards=getattr(result, "rewards", None), error=getattr(result, "error", None), - n_recorded=run.n_recorded_exchanges, + n_recorded=router.recorded_consumed_count, n_live=len(router.live_exchanges), divergences=router.divergences, ) @@ -635,6 +637,9 @@ async def _write_artifacts( rollout_dir, run.path / "trajectory" / "llm_trajectory.jsonl", live_exchanges, + recorded_consumed_count=( + replay_proxy.recorded_consumed_count if replay_proxy is not None else 0 + ), live_model=live_model, live_capture_host_owned=live_capture_host_owned, ) @@ -644,6 +649,9 @@ async def _write_artifacts( original_model=run.model, live_model=live_model, n_recorded=run.n_recorded_exchanges, + n_recorded_consumed=( + replay_proxy.recorded_consumed_count if replay_proxy is not None else 0 + ), n_live=len(live_exchanges), live_attempt_count=( replay_proxy.live_attempt_count if replay_proxy is not None else 0 @@ -659,7 +667,11 @@ async def _write_artifacts( live_model=live_model, usage=summarize_llm_trajectory_usage( stitched_path, - n_recorded=run.n_recorded_exchanges, + n_recorded=( + replay_proxy.recorded_consumed_count + if replay_proxy is not None + else 0 + ), ), environment=run.environment, ) @@ -768,7 +780,9 @@ async def _write_artifacts( rollout_dir=rollout_dir, rewards=getattr(result, "rewards", None), error=getattr(result, "error", None), - n_recorded=run.n_recorded_exchanges, + n_recorded=( + replay_proxy.recorded_consumed_count if replay_proxy is not None else 0 + ), n_live=len(live_exchanges), divergences=0, ) diff --git a/src/benchflow/continue_run/replay_proxy.py b/src/benchflow/continue_run/replay_proxy.py index 6e2be2e08..c2f362e9c 100644 --- a/src/benchflow/continue_run/replay_proxy.py +++ b/src/benchflow/continue_run/replay_proxy.py @@ -93,7 +93,15 @@ def live_exchanges(self) -> list[LLMExchange]: @property def exhausted(self) -> bool: - return self._cursor >= len(self._recorded) + with self._lock: + return self._cursor >= len(self._recorded) + + @property + def recorded_consumed_count(self) -> int: + """Return the recorded prefix actually served to the agent.""" + + with self._lock: + return min(self._cursor, len(self._recorded)) def _check_divergence( self, incoming: dict[str, Any], recorded_req: dict[str, Any] diff --git a/src/benchflow/continue_run/sandbox_proxy.py b/src/benchflow/continue_run/sandbox_proxy.py index 097182b46..ee85a7e80 100644 --- a/src/benchflow/continue_run/sandbox_proxy.py +++ b/src/benchflow/continue_run/sandbox_proxy.py @@ -109,6 +109,8 @@ class SandboxReplayProxy: state_path: str stdout_path: str stderr_path: str + recorded_exchange_count: int = 0 + recorded_consumed_count: int = 0 live_exchanges: list[LLMExchange] = field(default_factory=list) live_attempt_count: int = 0 live_errors: list[str] = field(default_factory=list) @@ -205,6 +207,7 @@ async def start( state_path=paths["state"], stdout_path=paths["stdout"], stderr_path=paths["stderr"], + recorded_exchange_count=len(recorded), ) try: await proxy._wait_until_ready() @@ -297,8 +300,18 @@ async def _load_live_state(self) -> None: if not isinstance(state, dict): self.live_errors.append("sandbox live capture state was unavailable") return + recorded_consumed_count = state.get("recorded_consumed_count") attempt_count = state.get("live_attempt_count") error_count = state.get("live_error_count") + if ( + not isinstance(recorded_consumed_count, int) + or isinstance(recorded_consumed_count, bool) + or recorded_consumed_count < 0 + or recorded_consumed_count > self.recorded_exchange_count + ): + self.live_errors.append("sandbox recorded replay count was invalid") + else: + self.recorded_consumed_count = recorded_consumed_count if ( not isinstance(attempt_count, int) or isinstance(attempt_count, bool) diff --git a/src/benchflow/continue_run/sandbox_replay_runtime.py b/src/benchflow/continue_run/sandbox_replay_runtime.py index c0edb701b..323ed5194 100644 --- a/src/benchflow/continue_run/sandbox_replay_runtime.py +++ b/src/benchflow/continue_run/sandbox_replay_runtime.py @@ -121,6 +121,7 @@ def __post_init__(self): def _write_state(self): payload = { "port": self.port, + "recorded_consumed_count": min(self.cursor, len(self.recorded)), "live_attempt_count": self.live_attempt_count, "live_error_count": self.live_error_count, } @@ -166,6 +167,10 @@ def next_response(self, request_body): ((exchange.get("request") or {}).get("body") or {}), ) self.cursor += 1 + # Journal the recorded prefix before returning the response. + # If this write fails, the handler returns 500 and the host + # cannot mistake an unjournaled replay for a completed prefix. + self._write_state() response = exchange.get("response") or {} return ( "replay", diff --git a/src/benchflow/continue_run/trajectory_artifacts.py b/src/benchflow/continue_run/trajectory_artifacts.py index 144385f8d..bf2cb9142 100644 --- a/src/benchflow/continue_run/trajectory_artifacts.py +++ b/src/benchflow/continue_run/trajectory_artifacts.py @@ -66,6 +66,7 @@ def write_stitched_trajectory( original_llm_trajectory: Path, live_exchanges: list[LLMExchange], *, + recorded_consumed_count: int, live_model: str | None = None, live_capture_host_owned: bool = True, ) -> Path: @@ -79,6 +80,15 @@ def write_stitched_trajectory( out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" out.parent.mkdir(parents=True, exist_ok=True) lines = stitched_trajectory_lines(original_llm_trajectory) + if ( + isinstance(recorded_consumed_count, bool) + or recorded_consumed_count < 0 + or recorded_consumed_count > len(lines) + ): + raise ValueError( + "recorded_consumed_count must identify an available trajectory prefix" + ) + lines = lines[:recorded_consumed_count] for exchange in live_exchanges: payload = exchange.model_dump(mode="json") metadata = payload.setdefault("metadata", {}) @@ -124,6 +134,7 @@ def refresh_stitched_trajectory_manifest( original_model: str | None, live_model: str | None, n_recorded: int, + n_recorded_consumed: int, n_live: int, live_attempt_count: int, live_errors: list[str], @@ -150,6 +161,13 @@ def refresh_stitched_trajectory_manifest( except ValueError: source = None + if ( + isinstance(n_recorded_consumed, bool) + or n_recorded_consumed < 0 + or n_recorded_consumed > n_recorded + ): + raise ValueError("n_recorded_consumed must be between zero and n_recorded") + trajectory_path = rollout_dir / "trajectory" / "llm_trajectory.jsonl" trajectory_lines = [ line for line in trajectory_path.read_text().splitlines() if line.strip() @@ -165,7 +183,7 @@ def refresh_stitched_trajectory_manifest( parsed_rows: list[dict[str, Any]] = [] if rows_valid: parsed_rows = [json.loads(line) for line in trajectory_lines] - expected_count = n_recorded + live_attempt_count + expected_count = n_recorded_consumed + live_attempt_count count_matches = exchange_count == expected_count source_rows: list[dict[str, Any]] = [] try: @@ -186,12 +204,14 @@ def refresh_stitched_trajectory_manifest( exchanges=source_rows, ) ) + recorded_prefix_complete = n_recorded_consumed == n_recorded live_capture_complete = live_attempt_count == n_live and not live_errors usage_complete = bool( parsed_rows and successful_exchanges_have_positive_usage(parsed_rows) ) complete = ( source_allows_training + and recorded_prefix_complete and count_matches and live_capture_complete and n_live == 0 @@ -225,6 +245,7 @@ def refresh_stitched_trajectory_manifest( and live_capture_complete and rows_valid and n_live == 0 + and recorded_prefix_complete ) response_complete = bool( source @@ -232,6 +253,7 @@ def refresh_stitched_trajectory_manifest( and count_matches and live_capture_complete and rows_valid + and recorded_prefix_complete ) errors = list(source.errors) if source and not complete else [] missing_fields = list(source.missing_fields) if source and not complete else [] @@ -246,6 +268,12 @@ def refresh_stitched_trajectory_manifest( missing_fields.append("source_capture_provenance") elif not source_allows_training: errors.append(CONTINUATION_SOURCE_AUDIT_ERROR) + if not recorded_prefix_complete: + errors.append( + "continuation stopped after consuming " + f"{n_recorded_consumed} of {n_recorded} recorded provider exchanges" + ) + missing_fields.append("recorded_replay_prefix") if not count_matches: errors.append( "stitched LLM trajectory count mismatch: " @@ -267,7 +295,7 @@ def refresh_stitched_trajectory_manifest( source=source, original_model=original_model, live_model=live_model, - n_recorded=n_recorded, + n_recorded=n_recorded_consumed, n_live=n_live, live_attempt_count=live_attempt_count, live_capture_complete=live_capture_complete, @@ -312,10 +340,17 @@ def _continuation_role_captures( ) -> list[LLMRoleCapture]: """Preserve source-role provenance and append a distinct live leg.""" - captures = [ - capture.model_copy(update={"leg": "recorded"}) - for capture in (source.role_captures if source else []) - ] + captures: list[LLMRoleCapture] = [] + remaining = n_recorded + for capture in source.role_captures if source else []: + retained = min(capture.exchange_count, remaining) + if retained: + captures.append( + capture.model_copy( + update={"leg": "recorded", "exchange_count": retained} + ) + ) + remaining -= retained if source is not None and n_recorded and not captures: captures.append( LLMRoleCapture( diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 94a9e2b0a..f57dfd452 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -205,6 +205,10 @@ def capture_artifact_allows_training( manifest, exchange_count=len(exchanges), ) + and all( + _exchange_matches_training_manifest(exchange, manifest) + for exchange in exchanges + ) and successful_exchanges_have_positive_usage(exchanges) ) return not any(_exchange_requires_manifest(exchange) for exchange in exchanges) @@ -300,6 +304,48 @@ def _exchange_requires_manifest(exchange: dict[str, Any]) -> bool: ) +def _exchange_matches_training_manifest( + exchange: dict[str, Any], manifest: dict[str, Any] +) -> bool: + """Require every canonical row to agree with its trusted sidecar.""" + + metadata = exchange.get("metadata") + if not isinstance(metadata, dict): + return False + if ( + metadata.get("schema_version") != LLM_TRAJECTORY_SCHEMA_VERSION + or metadata.get("capture_fidelity") != CaptureFidelity.PROVIDER_WIRE.value + or metadata.get("request_complete") is not True + or metadata.get("response_complete") is not True + or metadata.get("payload_redacted") is not True + ): + return False + + for field, mixed_value in ( + ("capture_source", CaptureSource.MIXED.value), + ("auth_mode", AuthMode.MIXED.value), + ): + manifest_value = manifest.get(field) + if not isinstance(manifest_value, str): + return False + allowed = {manifest_value} if manifest_value != mixed_value else set() + if manifest_value == mixed_value: + role_captures = manifest.get("role_captures") + if not isinstance(role_captures, list): + return False + allowed = { + capture.get(field) + for capture in role_captures + if isinstance(capture, dict) + and capture.get("capture_fidelity") + == CaptureFidelity.PROVIDER_WIRE.value + and isinstance(capture.get(field), str) + } + if not allowed or metadata.get(field) not in allowed: + return False + return True + + def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> bool: """Accept only expected, internally complete audit-only capture states.""" diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index dc4b5bcb6..3b5f65b8b 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -171,12 +171,102 @@ def test_write_stitched_trajectory_creates_file(tmp_path): original.write_text('{"a": 1}\n') rollout_dir = tmp_path / "rollout" out = write_stitched_trajectory( - rollout_dir, original, [exchange(completion(content="L"))] + rollout_dir, + original, + [exchange(completion(content="L"))], + recorded_consumed_count=1, ) assert out == rollout_dir / "trajectory" / "llm_trajectory.jsonl" assert len(out.read_text().strip().splitlines()) == 2 +def test_stitching_rejects_unconsumed_recorded_suffix(tmp_path): + """Guards PR #1057 against training on replay responses never consumed.""" + + model = "openai/gpt-5.5" + source = write_run_folder( + tmp_path / "source", + exchanges=[ + exchange(completion(content="consumed")), + exchange(completion(content="never-consumed")), + ], + model=model, + ) + trajectory_path = source / "trajectory" / "llm_trajectory.jsonl" + source_rows = [ + json.loads(line) for line in trajectory_path.read_text().splitlines() + ] + for row in source_rows: + row["metadata"].update( + { + "schema_version": 2, + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + } + ) + trajectory_path.write_text("".join(json.dumps(row) + "\n" for row in source_rows)) + source_manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="source", + exchange_count=2, + request_complete=True, + response_complete=True, + payload_redacted=True, + started_at="2026-08-29T00:00:00Z", + finished_at="2026-08-29T00:01:00Z", + ) + write_llm_trajectory_manifest(source, source_manifest) + + rollout = tmp_path / "continued" + initialize_llm_trajectory_artifacts( + rollout, + agent="openhands", + model=None, + session_id="continued", + started_at=source_manifest.finished_at, + ) + stitched = write_stitched_trajectory( + rollout, + trajectory_path, + [], + recorded_consumed_count=1, + live_model=model, + ) + manifest = refresh_stitched_trajectory_manifest( + rollout, + source, + original_model=model, + live_model=model, + n_recorded=2, + n_recorded_consumed=1, + n_live=0, + live_attempt_count=0, + live_errors=[], + ) + + stitched_rows = [json.loads(line) for line in stitched.read_text().splitlines()] + assert len(stitched_rows) == 1 + assert "consumed" in json.dumps(stitched_rows[0]) + assert "never-consumed" not in stitched.read_text() + assert manifest.status is CaptureStatus.PARTIAL + assert manifest.exchange_count == 1 + assert manifest.request_complete is False + assert manifest.response_complete is False + assert manifest.role_captures[0].exchange_count == 1 + assert "recorded_replay_prefix" in manifest.missing_fields + assert any("1 of 2 recorded" in error for error in manifest.errors) + assert CONTINUATION_SOURCE_AUDIT_ERROR not in manifest.errors + + def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_path): """Guards PR #1057 against promoting continuation ingress to provider wire.""" @@ -217,6 +307,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p rollout, source / "trajectory" / "llm_trajectory.jsonl", live, + recorded_consumed_count=1, live_model=continued_model, ) manifest = refresh_stitched_trajectory_manifest( @@ -225,6 +316,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p original_model=recorded_model, live_model=continued_model, n_recorded=1, + n_recorded_consumed=1, n_live=1, live_attempt_count=1, live_errors=[], @@ -265,6 +357,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p repeated_rollout, rollout / "trajectory" / "llm_trajectory.jsonl", [exchange(completion(content="live-again"))], + recorded_consumed_count=2, live_model=continued_model, ) repeated_manifest = refresh_stitched_trajectory_manifest( @@ -276,6 +369,7 @@ def test_refresh_stitched_manifest_rejects_replay_ingress_as_provider_wire(tmp_p original_model=continued_model, live_model=continued_model, n_recorded=2, + n_recorded_consumed=2, n_live=1, live_attempt_count=1, live_errors=[], @@ -336,6 +430,7 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) rollout, source / "trajectory" / "llm_trajectory.jsonl", [exchange(completion(content="live"))], + recorded_consumed_count=1, live_model="openai/gpt-5.5", ) manifest = refresh_stitched_trajectory_manifest( @@ -344,6 +439,7 @@ def test_refresh_stitched_manifest_keeps_lower_fidelity_prefix_partial(tmp_path) original_model="claude-sonnet-4-6", live_model="openai/gpt-5.5", n_recorded=1, + n_recorded_consumed=1, n_live=1, live_attempt_count=1, live_errors=[], @@ -412,6 +508,7 @@ def test_root_sandbox_live_suffix_is_retained_but_audit_only(tmp_path): rollout, source / "trajectory" / "llm_trajectory.jsonl", [exchange(completion(content="live"))], + recorded_consumed_count=1, live_model=model, live_capture_host_owned=False, ) @@ -421,6 +518,7 @@ def test_root_sandbox_live_suffix_is_retained_but_audit_only(tmp_path): original_model=model, live_model=model, n_recorded=1, + n_recorded_consumed=1, n_live=1, live_attempt_count=1, live_errors=[], @@ -473,6 +571,7 @@ def test_refresh_stitched_manifest_rejects_missing_live_attempt(tmp_path): rollout, source / "trajectory" / "llm_trajectory.jsonl", [], + recorded_consumed_count=1, live_model=model, ) @@ -482,6 +581,7 @@ def test_refresh_stitched_manifest_rejects_missing_live_attempt(tmp_path): original_model=model, live_model=model, n_recorded=1, + n_recorded_consumed=1, n_live=0, live_attempt_count=1, live_errors=["live provider request failed before capture"], @@ -889,6 +989,7 @@ def test_refresh_stitched_manifest_rejects_malformed_row(tmp_path): original_model=model, live_model=model, n_recorded=1, + n_recorded_consumed=1, n_live=0, live_attempt_count=0, live_errors=[], diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index 9812851d9..e987a5110 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -49,11 +49,13 @@ def forwarder(req): r1 = router.next_response({"messages": [{"role": "user"}]}) assert r1.source == "replay" assert r1.body["choices"][0]["message"]["content"] == "first" + assert router.recorded_consumed_count == 1 r2 = router.next_response({"messages": [{"role": "user"}]}) assert r2.source == "replay" assert r2.body["choices"][0]["message"]["content"] == "second" assert router.exhausted is True + assert router.recorded_consumed_count == 2 r3 = router.next_response({"messages": [{"role": "user"}]}) assert r3.source == "live" @@ -160,6 +162,28 @@ def fail(*_args, **_kwargs): assert capture_state["live_error_count"] == 1 +def test_sandbox_recorded_prefix_is_journaled_before_response(tmp_path) -> None: + """Guards PR #1057 against overstating a sandbox replay prefix.""" + + state = ReplayState( + recorded=[exchange(completion(content="recorded")).model_dump(mode="json")], + upstream_url="https://provider.invalid/v1", + upstream_api_key="test-key", + upstream_model="openai/test-model", + live_log_path=str(tmp_path / "live.jsonl"), + state_path=str(tmp_path / "state.json"), + port=61357, + ) + + source, status, _body = state.next_response({"messages": [{"role": "user"}]}) + + assert source == "replay" + assert status == 200 + capture_state = json.loads((tmp_path / "state.json").read_text()) + assert capture_state["recorded_consumed_count"] == 1 + assert capture_state["live_attempt_count"] == 0 + + def test_sandbox_live_exchange_is_redacted_before_journaling(tmp_path) -> None: """Guards PR #1057 against persisting raw continuation secrets.""" @@ -442,6 +466,44 @@ async def read_remote_text(_sandbox, path, **_kwargs): assert proxy.live_errors == ["sandbox live attempt journal failed 1 time(s)"] +@pytest.mark.asyncio +async def test_sandbox_replay_count_is_recovered_by_host( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Guards PR #1057 against stitching unconsumed sandbox responses.""" + + proxy = SandboxReplayProxy( + sandbox=object(), + runtime_dir="/tmp/runtime", + port=61357, + pid_path="/tmp/runtime/pid", + live_log_path="/tmp/runtime/live.jsonl", + state_path="/tmp/runtime/state.json", + stdout_path="/tmp/runtime/stdout.log", + stderr_path="/tmp/runtime/stderr.log", + recorded_exchange_count=3, + ) + + async def read_remote_text(_sandbox, path, **_kwargs): + assert path == proxy.state_path + return json.dumps( + { + "recorded_consumed_count": 2, + "live_attempt_count": 0, + "live_error_count": 0, + } + ) + + monkeypatch.setattr( + "benchflow.continue_run.sandbox_proxy._read_remote_text", read_remote_text + ) + + await proxy._load_live_state() + + assert proxy.recorded_consumed_count == 2 + assert proxy.live_errors == [] + + def test_divergence_warns_by_default(): recorded = [exchange(completion(content="a"), n_request_messages=3)] router = ReplayRouter(recorded) diff --git a/tests/test_opencode_family_proxy_tracking.py b/tests/test_opencode_family_proxy_tracking.py index d33d0ee49..5c0b0109a 100644 --- a/tests/test_opencode_family_proxy_tracking.py +++ b/tests/test_opencode_family_proxy_tracking.py @@ -12,7 +12,10 @@ """ import base64 +import json +import os import re +import subprocess import pytest @@ -58,6 +61,53 @@ def test_proxy_wrapper_wires_gateway_base_url(agent, wrapper_bin, cfg): assert "baseURL" in w +def test_proxy_wrapper_removes_preexisting_provider_credentials(tmp_path): + """Guards PR #1057 against an OpenCode config bypassing provider capture.""" + + wrapper = _wrapper_script("opencode", "opencode-proxy") + register_js = wrapper.split("<<'JSEOF'\n", 1)[1].split("\nJSEOF", 1)[0] + config = tmp_path / ".config" / "opencode" / "opencode.json" + config.parent.mkdir(parents=True) + config.write_text( + json.dumps( + { + "provider": { + "retained": { + "options": { + "apiKey": "literal-bypass-key", + "baseURL": "https://bypass.invalid/v1", + } + } + }, + "tools": {"webfetch": False}, + } + ), + encoding="utf-8", + ) + env = { + **os.environ, + "BENCHFLOW_AGENT_HOME": str(tmp_path), + "BENCHFLOW_LITELLM_MODEL_ALIAS": "benchflow-provider-model", + "OPENAI_BASE_URL": "http://127.0.0.1:4000/v1", + "OPENAI_API_KEY": "sk-benchflow-proxy-master-key", + } + + subprocess.run( + ["node"], + input=register_js, + text=True, + env=env, + check=True, + timeout=15, + ) + + updated = json.loads(config.read_text()) + assert set(updated["provider"]) == {OPENCODE_PROXY_PROVIDER_ID} + assert "literal-bypass-key" not in config.read_text() + assert "bypass.invalid" not in config.read_text() + assert updated["tools"] == {"webfetch": False} + + @pytest.mark.parametrize("agent,wrapper_bin,cfg", CASES) def test_proxy_wrapper_forces_chat_completions_sdk(agent, wrapper_bin, cfg): """The dedicated provider must use ``@ai-sdk/openai-compatible`` (chat diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 824b85d8e..af8446f5d 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -12,6 +12,7 @@ from benchflow.trajectories.llm_capture_manifest import ( CONTINUATION_SOURCE_AUDIT_ERROR, REPLAY_PROXY_INGRESS_AUDIT_ERROR, + capture_artifact_allows_training, capture_manifest_allows_training, capture_manifest_preserves_audit_completion, ) @@ -255,6 +256,64 @@ def test_complete_manifest_with_capture_gaps_fails_closed_for_training( assert row["is_completed"] is False +@pytest.mark.parametrize( + ("contradictory_field", "value"), + [ + ("capture_fidelity", "agent_session"), + ("request_complete", False), + ("response_complete", False), + ("payload_redacted", False), + ("capture_source", "codex_native_session"), + ("auth_mode", "oauth_subscription"), + ], +) +def test_schema_v2_row_cannot_contradict_training_manifest( + tmp_path: Path, contradictory_field: str, value: str | bool +) -> None: + """Guards PR #1057 against trusting a sidecar over contradictory rows.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange( + trajectory_dir, + fidelity="provider_wire", + schema_version=2, + ) + exchange = json.loads((trajectory_dir / "llm_trajectory.jsonl").read_text()) + exchange["metadata"].update( + { + "capture_source": "litellm_proxy", + "auth_mode": "api_key", + "payload_redacted": True, + } + ) + manifest = { + "status": "complete", + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + "missing_fields": [], + "errors": [], + } + + assert capture_artifact_allows_training(manifest, exchanges=[exchange]) + exchange["metadata"][contradictory_field] = value + assert not capture_artifact_allows_training(manifest, exchanges=[exchange]) + (trajectory_dir / "llm_trajectory.jsonl").write_text( + json.dumps(exchange) + "\n", encoding="utf-8" + ) + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps(manifest), encoding="utf-8" + ) + row = _build_results_row(tmp_path, agent_result={"total_tokens": 2}) + assert row["info"]["training_ready"] is False + assert row["is_completed"] is False + + def test_provider_capture_without_positive_usage_is_not_training_ready( tmp_path: Path, ) -> None: From 0e4ed7127705f1fb5c76de458bf6672debff2e7f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 17:54:13 -0700 Subject: [PATCH 53/74] preserve opencode manifest parity --- src/benchflow/acp/runtime.py | 13 +++++++ src/benchflow/agents/opencode_config.py | 38 ++++++++++++++++++++ src/benchflow/agents/registry.py | 12 +++---- tests/test_opencode_family_proxy_tracking.py | 29 +++++++++++++++ 4 files changed, 84 insertions(+), 8 deletions(-) create mode 100644 src/benchflow/agents/opencode_config.py diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 770262de1..9431001c6 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -28,6 +28,7 @@ from benchflow.acp.container_transport import ContainerTransport from benchflow.acp.selection import selected_acp_transport from benchflow.acp.types import McpServerSpec +from benchflow.agents.opencode_config import opencode_provider_reset_command from benchflow.agents.protocol import ACPSessionAdapter from benchflow.agents.providers import ( find_provider, @@ -72,6 +73,16 @@ _OPENHANDS_DISABLE_SUBAGENTS_ENV = "BENCHFLOW_OPENHANDS_DISABLE_SUBAGENTS" +def _harden_proxy_agent_launch( + agent: str, agent_launch: str, agent_env: dict[str, str] +) -> str: + """Insert agent-side config hardening after proxy selection.""" + + if agent != "opencode" or not agent_env.get("BENCHFLOW_LITELLM_MODEL_ALIAS"): + return agent_launch + return f"{opencode_provider_reset_command()} && {agent_launch}" + + def _acp_handshake_timeout_sec() -> float: """Effective timeout for the pre-prompt ACP handshake, in seconds. @@ -610,6 +621,8 @@ async def connect_acp( agent_launch = " ".join(parts) logger.info(f"Resolved agent path: {agent_launch}") + agent_launch = _harden_proxy_agent_launch(agent, agent_launch, agent_env) + if sandbox_user: agent_launch = build_priv_drop_cmd(agent_launch, sandbox_user) logger.info(f"Agent sandboxed as: {sandbox_user}") diff --git a/src/benchflow/agents/opencode_config.py b/src/benchflow/agents/opencode_config.py new file mode 100644 index 000000000..06cc8e376 --- /dev/null +++ b/src/benchflow/agents/opencode_config.py @@ -0,0 +1,38 @@ +"""OpenCode proxy-boundary configuration hardening.""" + +from __future__ import annotations + +import shlex + +OPENCODE_CONFIG_RELATIVE_PATH = ".config/opencode/opencode.json" +_OPENCODE_NODE = "/opt/benchflow/node/bin/node" + + +def opencode_provider_reset_source() -> str: + """Return stdlib Node.js that removes every pre-existing provider.""" + + return "\n".join( + [ + 'const fs = require("fs");', + 'const os = require("os");', + 'const path = require("path");', + 'const home = (process.env.BENCHFLOW_AGENT_HOME || "").trim() || os.homedir();', + f"const p = path.join(home, {OPENCODE_CONFIG_RELATIVE_PATH!r});", + "fs.mkdirSync(path.dirname(p), { recursive: true });", + 'const text = fs.existsSync(p) ? fs.readFileSync(p, "utf8").trim() : "";', + "const d = text ? JSON.parse(text) : {};", + # The manifest-owned wrapper adds the one BenchFlow gateway provider + # immediately after this pre-launch boundary. Removing the full map + # here prevents an image-baked literal key/endpoint from surviving. + "d.provider = {};", + 'const temporary = p + ".benchflow-" + process.pid + ".tmp";', + 'fs.writeFileSync(temporary, JSON.stringify(d, null, 2) + "\\n", { mode: 0o600 });', + "fs.renameSync(temporary, p);", + ] + ) + + +def opencode_provider_reset_command() -> str: + """Return the shell-safe reset command run immediately before OpenCode.""" + + return f"{_OPENCODE_NODE} -e {shlex.quote(opencode_provider_reset_source())}" diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index e10f41987..253f602b3 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -316,9 +316,9 @@ def _json_settings_merge(path: str, mutator: str) -> str: # routes through the proxy. def _opencode_family_proxy_wrapper_install(binary: str, config_path: str) -> str: """Install ``/opt/benchflow/bin/-proxy``: a thin wrapper that, in - proxy mode, replaces the provider map with the dedicated LiteLLM gateway - provider, then execs the isolated agent binary. Non-provider settings are - preserved; no-op outside proxy mode. + proxy mode, registers the LiteLLM gateway alias under a dedicated + OpenCode provider, then execs the isolated agent binary. Idempotent + (preserves existing config); no-op outside proxy mode. The gateway alias is registered under the ``{OPENCODE_PROXY_PROVIDER_ID}`` provider using ``@ai-sdk/openai-compatible`` — NOT the built-in ``openai`` @@ -357,11 +357,7 @@ def _opencode_family_proxy_wrapper_install(binary: str, config_path: str) -> str " const d = text ? JSON.parse(text) : {};", # Dedicated provider id (see docstring) → chat completions, not the # Responses API the built-in ``openai`` id is hard-coded to. - # An image can carry a pre-existing provider with a literal key or - # endpoint. Retaining it would let the agent bypass the capture - # proxy even after BenchFlow strips provider credentials from the - # process environment, so proxy mode owns the entire provider map. - " const providers = d.provider = {};", + " const providers = d.provider ||= {};", f" const prov = providers[{provider_id!r}] ||= {{}};", ' prov.npm = "@ai-sdk/openai-compatible";', ' prov.name ||= "BenchFlow Gateway";', diff --git a/tests/test_opencode_family_proxy_tracking.py b/tests/test_opencode_family_proxy_tracking.py index 5c0b0109a..fac978755 100644 --- a/tests/test_opencode_family_proxy_tracking.py +++ b/tests/test_opencode_family_proxy_tracking.py @@ -19,6 +19,8 @@ import pytest +from benchflow.acp.runtime import _harden_proxy_agent_launch +from benchflow.agents.opencode_config import opencode_provider_reset_source from benchflow.agents.registry import AGENTS, OPENCODE_PROXY_PROVIDER_ID # (agent name, proxy wrapper binary, agent config filename). MiMo is an @@ -92,6 +94,17 @@ def test_proxy_wrapper_removes_preexisting_provider_credentials(tmp_path): "OPENAI_API_KEY": "sk-benchflow-proxy-master-key", } + subprocess.run( + ["node", "-e", opencode_provider_reset_source()], + text=True, + env=env, + check=True, + timeout=15, + ) + sanitized = json.loads(config.read_text()) + assert sanitized["provider"] == {} + assert sanitized["tools"] == {"webfetch": False} + subprocess.run( ["node"], input=register_js, @@ -108,6 +121,22 @@ def test_proxy_wrapper_removes_preexisting_provider_credentials(tmp_path): assert updated["tools"] == {"webfetch": False} +def test_proxy_launch_resets_providers_immediately_before_manifest_wrapper(): + """Guards PR #1057 without drifting the external agent manifest.""" + + launch = "/opt/benchflow/bin/opencode-proxy acp" + hardened = _harden_proxy_agent_launch( + "opencode", + launch, + {"BENCHFLOW_LITELLM_MODEL_ALIAS": "benchflow-provider-model"}, + ) + + assert "opencode.json" in hardened + assert "d.provider = {}" in hardened + assert hardened.endswith("&& " + launch) + assert _harden_proxy_agent_launch("opencode", launch, {}) == launch + + @pytest.mark.parametrize("agent,wrapper_bin,cfg", CASES) def test_proxy_wrapper_forces_chat_completions_sdk(agent, wrapper_bin, cfg): """The dedicated provider must use ``@ai-sdk/openai-compatible`` (chat From e113d0e3797e61e9ffde858e2c64fdcb67780584 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 18:24:37 -0700 Subject: [PATCH 54/74] Harden proxy config and role capture admission --- src/benchflow/acp/runtime.py | 13 +- .../trajectories/llm_capture_manifest.py | 58 +++++++ tests/continue_run/test_orchestrator.py | 137 ++------------- tests/continue_run/test_training_metadata.py | 158 ++++++++++++++++++ tests/test_opencode_family_proxy_tracking.py | 41 +++++ .../test_llm_capture_training_contract.py | 87 ++++++++++ 6 files changed, 371 insertions(+), 123 deletions(-) create mode 100644 tests/continue_run/test_training_metadata.py diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 9431001c6..57efcc144 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -78,9 +78,18 @@ def _harden_proxy_agent_launch( ) -> str: """Insert agent-side config hardening after proxy selection.""" - if agent != "opencode" or not agent_env.get("BENCHFLOW_LITELLM_MODEL_ALIAS"): + if not agent_env.get("BENCHFLOW_LITELLM_MODEL_ALIAS"): return agent_launch - return f"{opencode_provider_reset_command()} && {agent_launch}" + if agent == "opencode": + return f"{opencode_provider_reset_command()} && {agent_launch}" + if agent == "mimo": + # MiMo's manifest-owned launcher replaces both canonical config files + # in proxy mode, but the CLI also honors this arbitrary alternate path. + # Do not let an image-baked config bypass the capture proxy. Direct mode + # deliberately retains the override as part of the caller's provider + # configuration. + return f"unset MIMOCODE_CONFIG && {agent_launch}" + return agent_launch def _acp_handshake_timeout_sec() -> float: diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index f57dfd452..9be618149 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -8,6 +8,7 @@ from __future__ import annotations import json +from collections import Counter from collections.abc import Sequence from datetime import datetime from enum import StrEnum @@ -178,9 +179,11 @@ def capture_manifest_allows_training( expected_count = manifest.get("exchange_count") if not isinstance(expected_count, int) or isinstance(expected_count, bool): return False + role_captures = _validated_role_captures(manifest) return bool( manifest.get("status") == "complete" + and manifest.get("capture_source") == CaptureSource.LITELLM_PROXY.value and manifest.get("capture_fidelity") == "provider_wire" and manifest.get("request_complete") is True and manifest.get("response_complete") is True @@ -189,6 +192,16 @@ def capture_manifest_allows_training( and manifest.get("errors") == [] and exchange_count > 0 and expected_count == exchange_count + and role_captures + and _role_captures_match_manifest(manifest, role_captures) + and all( + capture.capture_source is CaptureSource.LITELLM_PROXY + and capture.capture_fidelity is CaptureFidelity.PROVIDER_WIRE + and capture.exchange_count > 0 + and capture.request_complete is True + and capture.response_complete is True + for capture in role_captures + ) ) @@ -209,6 +222,7 @@ def capture_artifact_allows_training( _exchange_matches_training_manifest(exchange, manifest) for exchange in exchanges ) + and _exchanges_match_training_roles(exchanges, manifest) and successful_exchanges_have_positive_usage(exchanges) ) return not any(_exchange_requires_manifest(exchange) for exchange in exchanges) @@ -346,6 +360,50 @@ def _exchange_matches_training_manifest( return True +def _exchanges_match_training_roles( + exchanges: Sequence[dict[str, Any]], manifest: dict[str, Any] +) -> bool: + """Require row attribution and cardinality to match per-role provenance.""" + + role_captures = _validated_role_captures(manifest) + if not role_captures: + return False + expected: Counter[tuple[str, str, str | None, str, str, str]] = Counter() + for capture in role_captures: + expected[ + ( + capture.role, + capture.agent, + capture.model, + capture.auth_mode.value, + capture.capture_source.value, + capture.capture_fidelity.value, + ) + ] += capture.exchange_count + actual: Counter[tuple[str, str, str | None, str, str, str]] = Counter() + for exchange in exchanges: + metadata = exchange.get("metadata") + if not isinstance(metadata, dict): + return False + role = metadata.get("role") + agent = metadata.get("agent") + model = metadata.get("model") + auth_mode = metadata.get("auth_mode") + capture_source = metadata.get("capture_source") + capture_fidelity = metadata.get("capture_fidelity") + if not ( + isinstance(role, str) + and isinstance(agent, str) + and (model is None or isinstance(model, str)) + and isinstance(auth_mode, str) + and isinstance(capture_source, str) + and isinstance(capture_fidelity, str) + ): + return False + actual[(role, agent, model, auth_mode, capture_source, capture_fidelity)] += 1 + return actual == expected + + def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> bool: """Accept only expected, internally complete audit-only capture states.""" diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 3b5f65b8b..728272e78 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -25,7 +25,6 @@ ) from benchflow.trajectories.llm_capture_manifest import ( CONTINUATION_SOURCE_AUDIT_ERROR, - REPLAY_PROXY_INGRESS_AUDIT_ERROR, AuthMode, CaptureFidelity, CaptureSource, @@ -206,6 +205,9 @@ def test_stitching_rejects_unconsumed_recorded_suffix(tmp_path): "request_complete": True, "response_complete": True, "payload_redacted": True, + "role": "agent", + "agent": "openhands", + "model": model, } ) trajectory_path.write_text("".join(json.dumps(row) + "\n" for row in source_rows)) @@ -223,6 +225,19 @@ def test_stitching_rejects_unconsumed_recorded_suffix(tmp_path): payload_redacted=True, started_at="2026-08-29T00:00:00Z", finished_at="2026-08-29T00:01:00Z", + role_captures=[ + LLMRoleCapture( + role="agent", + agent="openhands", + model=model, + auth_mode=AuthMode.API_KEY, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + exchange_count=2, + request_complete=True, + response_complete=True, + ) + ], ) write_llm_trajectory_manifest(source, source_manifest) @@ -813,126 +828,6 @@ async def after_cleanup(teardown_errors): assert json.loads(manifest_path.read_text())["status"] == "partial" -def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): - """Guards PR #1057 against retaining stale or incomplete continuation rows.""" - rollout = tmp_path / "job" / "demo-task__continued" - (rollout / "trajectory").mkdir(parents=True) - model = "openai/gpt-5.5" - row = exchange(completion(content="done")).model_dump(mode="json") - row["request"]["body"]["messages"] = [{"role": "user", "content": "Do the task."}] - row["metadata"] = { - "schema_version": 2, - "capture_source": "litellm_proxy", - "capture_fidelity": "provider_wire", - "auth_mode": "api_key", - "request_complete": True, - "response_complete": True, - "payload_redacted": True, - } - trajectory_path = rollout / "trajectory" / "llm_trajectory.jsonl" - trajectory_path.write_text(json.dumps(row) + "\n") - manifest = LLMTrajectoryManifest( - status=CaptureStatus.COMPLETE, - capture_source=CaptureSource.LITELLM_PROXY, - capture_fidelity=CaptureFidelity.PROVIDER_WIRE, - auth_mode=AuthMode.API_KEY, - agent="openhands", - model=model, - session_id="continued", - exchange_count=1, - request_complete=True, - response_complete=True, - payload_redacted=True, - started_at="2026-08-29T00:00:00Z", - finished_at="2026-08-29T00:01:00Z", - ) - write_llm_trajectory_manifest(rollout, manifest) - (rollout / "config.json").write_text(json.dumps({"model": None, "source": {}})) - (rollout / "prompts.json").write_text(json.dumps(["Do the task."])) - (rollout / "result.json").write_text( - json.dumps( - { - "task_name": "demo-task", - "rollout_name": "demo-task__continued", - "agent": "openhands", - "agent_name": "OpenHands", - "model": None, - "n_tool_calls": 0, - "partial_trajectory": False, - "rewards": {"reward": 1.0}, - "error": None, - "verifier_error": None, - "export_error": None, - "timing": {}, - "agent_result": {"total_tokens": 0, "usage_source": "unavailable"}, - "usage_tracking": {"requested": "off", "status": "off"}, - } - ) - ) - (rollout / "results.jsonl").write_text( - json.dumps({"info": {"training_ready": False, "model": None}}) + "\n" - ) - - update_continued_metadata( - rollout, - live_model=model, - usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), - environment="docker", - ) - - refreshed = json.loads((rollout / "results.jsonl").read_text()) - aggregated = json.loads((rollout.parent / "results.jsonl").read_text()) - assert refreshed["info"]["model"] == model - assert refreshed["info"]["training_ready"] is True - assert refreshed["token_usage"]["total_tokens"] == 2 - assert len(refreshed["trajectory"]) == 1 - assert aggregated == refreshed - - replay_manifest = manifest.model_copy( - update={ - "status": CaptureStatus.PARTIAL, - "capture_source": CaptureSource.MIXED, - "capture_fidelity": CaptureFidelity.MIXED, - "request_complete": False, - "missing_fields": ["live_provider_request"], - "errors": [REPLAY_PROXY_INGRESS_AUDIT_ERROR], - "role_captures": [ - LLMRoleCapture( - role="agent", - leg="live", - agent="openhands", - model=model, - auth_mode=AuthMode.API_KEY, - capture_source=CaptureSource.REPLAY_PROXY, - capture_fidelity=CaptureFidelity.AGENT_SESSION, - exchange_count=1, - request_complete=False, - response_complete=True, - ) - ], - } - ) - write_llm_trajectory_manifest(rollout, replay_manifest) - (rollout / "results.jsonl").write_text( - json.dumps({"info": {"training_ready": True}, "is_completed": False}) + "\n" - ) - - update_continued_metadata( - rollout, - live_model=model, - usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), - environment="docker", - ) - - replay_refreshed = json.loads((rollout / "results.jsonl").read_text()) - assert replay_refreshed["info"]["training_ready"] is False - assert replay_refreshed["info"]["training_ready_reason"] == ( - "insufficient_capture_fidelity" - ) - assert replay_refreshed["is_completed"] is True - assert replay_refreshed["error"] is None - - def test_stitching_structurally_redacts_escaped_secret(tmp_path): """Guards PR #1057 against string redaction corrupting stitched JSON.""" secret = "ESCbearerSECRETtok123456" diff --git a/tests/continue_run/test_training_metadata.py b/tests/continue_run/test_training_metadata.py new file mode 100644 index 000000000..d7e3bbd2b --- /dev/null +++ b/tests/continue_run/test_training_metadata.py @@ -0,0 +1,158 @@ +"""Continuation metadata regeneration at the trainer-admission boundary.""" + +from __future__ import annotations + +import json + +from benchflow.continue_run.orchestrator import ( + summarize_llm_trajectory_usage, + update_continued_metadata, +) +from benchflow.trajectories.llm_capture_manifest import ( + REPLAY_PROXY_INGRESS_AUDIT_ERROR, + AuthMode, + CaptureFidelity, + CaptureSource, + CaptureStatus, + LLMRoleCapture, + LLMTrajectoryManifest, + write_llm_trajectory_manifest, +) + +from ._helpers import completion, exchange + + +def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): + """Guards PR #1057 against retaining stale or incomplete continuation rows.""" + rollout = tmp_path / "job" / "demo-task__continued" + (rollout / "trajectory").mkdir(parents=True) + model = "openai/gpt-5.5" + row = exchange(completion(content="done")).model_dump(mode="json") + row["request"]["body"]["messages"] = [{"role": "user", "content": "Do the task."}] + row["metadata"] = { + "schema_version": 2, + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + "role": "agent", + "agent": "openhands", + "model": model, + } + trajectory_path = rollout / "trajectory" / "llm_trajectory.jsonl" + trajectory_path.write_text(json.dumps(row) + "\n") + manifest = LLMTrajectoryManifest( + status=CaptureStatus.COMPLETE, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + auth_mode=AuthMode.API_KEY, + agent="openhands", + model=model, + session_id="continued", + exchange_count=1, + request_complete=True, + response_complete=True, + payload_redacted=True, + started_at="2026-08-29T00:00:00Z", + finished_at="2026-08-29T00:01:00Z", + role_captures=[ + LLMRoleCapture( + role="agent", + agent="openhands", + model=model, + auth_mode=AuthMode.API_KEY, + capture_source=CaptureSource.LITELLM_PROXY, + capture_fidelity=CaptureFidelity.PROVIDER_WIRE, + exchange_count=1, + request_complete=True, + response_complete=True, + ) + ], + ) + write_llm_trajectory_manifest(rollout, manifest) + (rollout / "config.json").write_text(json.dumps({"model": None, "source": {}})) + (rollout / "prompts.json").write_text(json.dumps(["Do the task."])) + (rollout / "result.json").write_text( + json.dumps( + { + "task_name": "demo-task", + "rollout_name": "demo-task__continued", + "agent": "openhands", + "agent_name": "OpenHands", + "model": None, + "n_tool_calls": 0, + "partial_trajectory": False, + "rewards": {"reward": 1.0}, + "error": None, + "verifier_error": None, + "export_error": None, + "timing": {}, + "agent_result": {"total_tokens": 0, "usage_source": "unavailable"}, + "usage_tracking": {"requested": "off", "status": "off"}, + } + ) + ) + (rollout / "results.jsonl").write_text( + json.dumps({"info": {"training_ready": False, "model": None}}) + "\n" + ) + + update_continued_metadata( + rollout, + live_model=model, + usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), + environment="docker", + ) + + refreshed = json.loads((rollout / "results.jsonl").read_text()) + aggregated = json.loads((rollout.parent / "results.jsonl").read_text()) + assert refreshed["info"]["model"] == model + assert refreshed["info"]["training_ready"] is True + assert refreshed["token_usage"]["total_tokens"] == 2 + assert len(refreshed["trajectory"]) == 1 + assert aggregated == refreshed + + replay_manifest = manifest.model_copy( + update={ + "status": CaptureStatus.PARTIAL, + "capture_source": CaptureSource.MIXED, + "capture_fidelity": CaptureFidelity.MIXED, + "request_complete": False, + "missing_fields": ["live_provider_request"], + "errors": [REPLAY_PROXY_INGRESS_AUDIT_ERROR], + "role_captures": [ + LLMRoleCapture( + role="agent", + leg="live", + agent="openhands", + model=model, + auth_mode=AuthMode.API_KEY, + capture_source=CaptureSource.REPLAY_PROXY, + capture_fidelity=CaptureFidelity.AGENT_SESSION, + exchange_count=1, + request_complete=False, + response_complete=True, + ) + ], + } + ) + write_llm_trajectory_manifest(rollout, replay_manifest) + (rollout / "results.jsonl").write_text( + json.dumps({"info": {"training_ready": True}, "is_completed": False}) + "\n" + ) + + update_continued_metadata( + rollout, + live_model=model, + usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), + environment="docker", + ) + + replay_refreshed = json.loads((rollout / "results.jsonl").read_text()) + assert replay_refreshed["info"]["training_ready"] is False + assert replay_refreshed["info"]["training_ready_reason"] == ( + "insufficient_capture_fidelity" + ) + assert replay_refreshed["is_completed"] is True + assert replay_refreshed["error"] is None diff --git a/tests/test_opencode_family_proxy_tracking.py b/tests/test_opencode_family_proxy_tracking.py index fac978755..675f077b6 100644 --- a/tests/test_opencode_family_proxy_tracking.py +++ b/tests/test_opencode_family_proxy_tracking.py @@ -137,6 +137,47 @@ def test_proxy_launch_resets_providers_immediately_before_manifest_wrapper(): assert _harden_proxy_agent_launch("opencode", launch, {}) == launch +def test_proxy_launch_unsets_mimo_alternate_config_before_manifest_launcher( + tmp_path, +): + """Guards PR #1057 against a MiMo alternate-config capture bypass.""" + + alternate = tmp_path / "alternate-mimocode.json" + alternate.write_text( + json.dumps( + { + "provider": { + "bypass": { + "options": { + "apiKey": "literal-bypass-key", + "baseURL": "https://bypass.invalid/v1", + } + } + } + } + ), + encoding="utf-8", + ) + launch = 'test -z "${MIMOCODE_CONFIG:-}"' + hardened = _harden_proxy_agent_launch( + "mimo", + launch, + {"BENCHFLOW_LITELLM_MODEL_ALIAS": "benchflow-provider-model"}, + ) + + subprocess.run( + ["sh", "-c", hardened], + env={**os.environ, "MIMOCODE_CONFIG": str(alternate)}, + check=True, + timeout=15, + ) + assert hardened == f"unset MIMOCODE_CONFIG && {launch}" + assert ( + _harden_proxy_agent_launch("mimo", launch, {"MIMOCODE_CONFIG": str(alternate)}) + == launch + ) + + @pytest.mark.parametrize("agent,wrapper_bin,cfg", CASES) def test_proxy_wrapper_forces_chat_completions_sdk(agent, wrapper_bin, cfg): """The dedicated provider must use ``@ai-sdk/openai-compatible`` (chat diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index af8446f5d..a5d3297e4 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -285,6 +285,8 @@ def test_schema_v2_row_cannot_contradict_training_manifest( "capture_source": "litellm_proxy", "auth_mode": "api_key", "payload_redacted": True, + "role": "agent", + "agent": "codex-acp", } ) manifest = { @@ -298,9 +300,34 @@ def test_schema_v2_row_cannot_contradict_training_manifest( "payload_redacted": True, "missing_fields": [], "errors": [], + "role_captures": [ + { + "role": "agent", + "agent": "codex-acp", + "auth_mode": "api_key", + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + } + ], } assert capture_artifact_allows_training(manifest, exchanges=[exchange]) + manifest_without_roles = json.loads(json.dumps(manifest)) + manifest_without_roles["role_captures"] = [] + assert not capture_artifact_allows_training( + manifest_without_roles, exchanges=[exchange] + ) + manifest_with_wrong_attribution = json.loads(json.dumps(manifest)) + manifest_with_wrong_attribution["role_captures"][0]["agent"] = "claude-agent-acp" + assert capture_manifest_allows_training( + manifest_with_wrong_attribution, exchange_count=1 + ) + assert not capture_artifact_allows_training( + manifest_with_wrong_attribution, exchanges=[exchange] + ) exchange["metadata"][contradictory_field] = value assert not capture_artifact_allows_training(manifest, exchanges=[exchange]) (trajectory_dir / "llm_trajectory.jsonl").write_text( @@ -314,6 +341,66 @@ def test_schema_v2_row_cannot_contradict_training_manifest( assert row["is_completed"] is False +@pytest.mark.parametrize( + ("role_field", "value"), + [ + ("capture_source", "codex_native_session"), + ("capture_fidelity", "agent_session"), + ("exchange_count", 0), + ("request_complete", False), + ("response_complete", False), + ], +) +def test_role_capture_cannot_contradict_training_manifest( + tmp_path: Path, role_field: str, value: str | int | bool +) -> None: + """Guards PR #1057 against training on contradictory role provenance.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange( + trajectory_dir, + fidelity="provider_wire", + schema_version=2, + ) + exchange = json.loads((trajectory_dir / "llm_trajectory.jsonl").read_text()) + exchange["metadata"].update( + { + "capture_source": "litellm_proxy", + "auth_mode": "api_key", + "payload_redacted": True, + } + ) + manifest = { + "status": "complete", + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "auth_mode": "api_key", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + "payload_redacted": True, + "missing_fields": [], + "errors": [], + "role_captures": [ + { + "role": "agent", + "agent": "codex-acp", + "auth_mode": "api_key", + "capture_source": "litellm_proxy", + "capture_fidelity": "provider_wire", + "exchange_count": 1, + "request_complete": True, + "response_complete": True, + } + ], + } + manifest["role_captures"][0][role_field] = value + + assert not capture_manifest_allows_training(manifest, exchange_count=1) + assert not capture_artifact_allows_training(manifest, exchanges=[exchange]) + + def test_provider_capture_without_positive_usage_is_not_training_ready( tmp_path: Path, ) -> None: From 33590d3269daa5c6ca268b78b70904304f46077f Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 18:49:06 -0700 Subject: [PATCH 55/74] Require attribution and isolate subscription auth --- src/benchflow/agents/credentials.py | 50 +++++++++++++++++++ src/benchflow/providers/litellm_runtime.py | 15 +++++- .../trajectories/llm_capture_manifest.py | 1 + tests/continue_run/test_orchestrator.py | 1 + tests/continue_run/test_training_metadata.py | 1 + tests/test_litellm_credential_custody.py | 44 ++++++++++++++++ tests/test_litellm_runtime.py | 4 +- .../test_llm_capture_training_contract.py | 2 + 8 files changed, 116 insertions(+), 2 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index 4570bebf5..a68ec3ead 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -5,6 +5,7 @@ - write_credential_files agent + provider credential files (cf. AgentConfig) - write_gemini_vertex_settings ~/.gemini/settings.json for Vertex backend - upload_subscription_auth host login files (e.g. ~/.claude/.credentials.json) + - isolate_subscription_auth_for_proxy remove stale native login in API mode The Gemini Vertex settings helper lives here (not in _agent_env.py) so the module has a single coherent role and zero horizontal imports between phase @@ -178,3 +179,52 @@ async def upload_subscription_auth( host_path, container_path, ) + + +async def isolate_subscription_auth_for_proxy( + env, + *, + agent: str, + cred_home: str, +) -> bool: + """Remove the primary native-login credential before an API proxy run. + + A reused image may already contain the CLI login detected by + ``SubscriptionAuth.detect_file``. API-key mode must not leave that alternate + provider route available to the agent. Return false unless the registry's + primary credential can be identified, removed, and verified absent. + """ + + agent_cfg = AGENTS.get(agent) + subscription_auth = agent_cfg.subscription_auth if agent_cfg else None + if subscription_auth is None: + return True + primary_files = [ + auth_file + for auth_file in subscription_auth.files + if auth_file.host_path == subscription_auth.detect_file + ] + if len(primary_files) != 1 or env is None: + logger.warning( + "Cannot prove proxy isolation for %s subscription credentials", agent + ) + return False + + path = primary_files[0].container_path.format(home=cred_home) + quoted_path = shlex.quote(path) + try: + result = await env.exec( + f"rm -f -- {quoted_path} && ! test -e {quoted_path} && ! test -L {quoted_path}", + user="root", + timeout_sec=10, + ) + except Exception as exc: + logger.warning("Failed to isolate %s subscription credential: %s", agent, exc) + return False + if result.return_code != 0: + logger.warning( + "Subscription credential remained accessible in proxy mode for %s", + agent, + ) + return False + return True diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 148186d75..d4cd8d6cc 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -28,6 +28,7 @@ from benchflow._utils.text import describe_exception from benchflow.agents.codex_config import apply_codex_proxy_config +from benchflow.agents.credentials import isolate_subscription_auth_for_proxy from benchflow.agents.env import uses_native_subscription_auth from benchflow.agents.registry import AGENTS from benchflow.providers.litellm_bedrock_preflight import ( @@ -1837,6 +1838,15 @@ async def ensure_litellm_runtime( ), ) + credential_home = agent_env.get("BENCHFLOW_AGENT_HOME", "").strip() or ( + f"/home/{sandbox_user}" if sandbox_user else "/root" + ) + subscription_credentials_isolated = await isolate_subscription_auth_for_proxy( + sandbox, + agent=agent, + cred_home=credential_home, + ) + master_key = ( agent_env.get(LITELLM_MASTER_KEY_ENV) or f"sk-benchflow-{secrets.token_urlsafe(24)}" @@ -1845,7 +1855,10 @@ async def ensure_litellm_runtime( sorted(set(required_skill_names)), separators=(",", ":") ) proxy_location = "sandbox" if sandbox_local else "host" - credentials_isolated = _provider_credentials_have_proxy_only_custody(route) + credentials_isolated = bool( + _provider_credentials_have_proxy_only_custody(route) + and subscription_credentials_isolated + ) capture_trusted = not sandbox_local and credentials_isolated config_key = ( f"{environment}:{proxy_location}:{route.config_key}:{agent}:" diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index 9be618149..adf4cab3b 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -332,6 +332,7 @@ def _exchange_matches_training_manifest( or metadata.get("request_complete") is not True or metadata.get("response_complete") is not True or metadata.get("payload_redacted") is not True + or metadata.get("role_attribution_complete") is not True ): return False diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index 728272e78..3d6db382e 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -205,6 +205,7 @@ def test_stitching_rejects_unconsumed_recorded_suffix(tmp_path): "request_complete": True, "response_complete": True, "payload_redacted": True, + "role_attribution_complete": True, "role": "agent", "agent": "openhands", "model": model, diff --git a/tests/continue_run/test_training_metadata.py b/tests/continue_run/test_training_metadata.py index d7e3bbd2b..1fa57044b 100644 --- a/tests/continue_run/test_training_metadata.py +++ b/tests/continue_run/test_training_metadata.py @@ -37,6 +37,7 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): "request_complete": True, "response_complete": True, "payload_redacted": True, + "role_attribution_complete": True, "role": "agent", "agent": "openhands", "model": model, diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index ede7d9b70..55d5c88cc 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -10,6 +10,50 @@ from benchflow.providers.runtime import ensure_litellm_runtime +class _SubscriptionIsolationSandbox: + def __init__(self, *, return_code: int = 0) -> None: + self.return_code = return_code + self.commands: list[str] = [] + + async def exec(self, command, **kwargs): + self.commands.append(command) + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=self.return_code, stdout="", stderr="") + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("return_code", "trusted"), [(0, True), (1, False)]) +async def test_claude_api_proxy_isolates_preexisting_subscription_credential( + monkeypatch, return_code: int, trusted: bool +) -> None: + """Guards PR #1057 against a reused Claude OAuth file bypassing capture.""" + + async def fake_start(**_kwargs): + return SimpleNamespace(base_url="http://127.0.0.1:4000") + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + sandbox = _SubscriptionIsolationSandbox(return_code=return_code) + + _, provider_runtime = await ensure_litellm_runtime( + agent="claude-agent-acp", + agent_env={"ANTHROPIC_API_KEY": "sk-ant-provider"}, + model="claude-sonnet-4-6", + runtime=None, + environment="docker", + session_id="run-claude-api-with-stale-oauth", + sandbox=sandbox, + sandbox_user="agent", + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is trusted + assert len(sandbox.commands) == 1 + assert "/home/agent/.claude/.credentials.json" in sandbox.commands[0] + assert "rm -f" in sandbox.commands[0] + assert "! test -e" in sandbox.commands[0] + assert "! test -L" in sandbox.commands[0] + + @pytest.mark.asyncio async def test_vertex_adc_provider_capture_remains_audit_only_on_host(monkeypatch): """Guards PR #1057 against trusting agent-accessible Vertex credentials.""" diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 32ec8ad6c..c3e95af67 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -61,7 +61,9 @@ async def fake_start(**kwargs): assert provider_runtime is not None assert provider_runtime.kind == "litellm" assert provider_runtime.backend_model == "bedrock/us.anthropic.claude-opus-4-8" - assert provider_runtime.capture_trusted is True + # Without a sandbox handle, the registry-known Codex subscription file + # cannot be proven absent, so provider rows remain audit-only. + assert provider_runtime.capture_trusted is False assert updated["OPENAI_BASE_URL"] == "http://host.docker.internal:32123/v1" assert updated["OPENAI_API_KEY"] == provider_runtime.master_key assert updated[LITELLM_MODEL_ALIAS_ENV] == ( diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index a5d3297e4..23314d1d2 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -263,6 +263,7 @@ def test_complete_manifest_with_capture_gaps_fails_closed_for_training( ("request_complete", False), ("response_complete", False), ("payload_redacted", False), + ("role_attribution_complete", False), ("capture_source", "codex_native_session"), ("auth_mode", "oauth_subscription"), ], @@ -285,6 +286,7 @@ def test_schema_v2_row_cannot_contradict_training_manifest( "capture_source": "litellm_proxy", "auth_mode": "api_key", "payload_redacted": True, + "role_attribution_complete": True, "role": "agent", "agent": "codex-acp", } From d8342c6ffa0a237c9f0bb2eb2f740de9112d1f82 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 19:17:59 -0700 Subject: [PATCH 56/74] Harden proxy auth config homes --- src/benchflow/agents/credentials.py | 56 +++++++++++++-- src/benchflow/agents/registry.py | 6 ++ src/benchflow/providers/litellm_runtime.py | 5 +- tests/test_litellm_credential_custody.py | 82 +++++++++++++++++++--- tests/test_subscription_auth.py | 2 + 5 files changed, 131 insertions(+), 20 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index a68ec3ead..38c59fc0d 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -22,7 +22,7 @@ import os import shlex import tempfile -from pathlib import Path +from pathlib import Path, PurePosixPath from benchflow.agents.registry import AGENTS @@ -185,14 +185,18 @@ async def isolate_subscription_auth_for_proxy( env, *, agent: str, + agent_env: dict[str, str], cred_home: str, ) -> bool: """Remove the primary native-login credential before an API proxy run. A reused image may already contain the CLI login detected by ``SubscriptionAuth.detect_file``. API-key mode must not leave that alternate - provider route available to the agent. Return false unless the registry's - primary credential can be identified, removed, and verified absent. + provider route available to the agent. A safe CLI config-home override is + scrubbed as well as removed before launch; an override outside the sandbox + user's home fails capture trust without allowing a root deletion there. + Return false unless every eligible credential can be identified, removed, + and verified absent. """ agent_cfg = AGENTS.get(agent) @@ -210,18 +214,56 @@ async def isolate_subscription_auth_for_proxy( ) return False - path = primary_files[0].container_path.format(home=cred_home) - quoted_path = shlex.quote(path) + primary_path = primary_files[0].container_path.format(home=cred_home) + paths = [primary_path] + paths_safe = True + if subscription_auth.config_dir_env: + override_value = agent_env.pop(subscription_auth.config_dir_env, "") + if override_value: + override_dir = PurePosixPath(override_value) + home = PurePosixPath(cred_home) + try: + override_dir.relative_to(home) + except ValueError: + paths_safe = False + if ( + not override_dir.is_absolute() + or ".." in override_dir.parts + or "\x00" in override_value + ): + paths_safe = False + if paths_safe: + paths.append(str(override_dir / PurePosixPath(primary_path).name)) + else: + logger.warning( + "Refusing root cleanup outside the sandbox home for %s", agent + ) + + paths = list(dict.fromkeys(paths)) + quoted_paths = [shlex.quote(path) for path in paths] + parent_paths: set[str] = set() + for path in paths: + parent = PurePosixPath(path).parent + parent_paths.update( + str(candidate) + for candidate in (*reversed(parent.parents), parent) + if str(candidate) != "/" + ) + command_parts = [ + *(f"! test -L {shlex.quote(parent)}" for parent in sorted(parent_paths)), + f"rm -f -- {' '.join(quoted_paths)}", + *(f"! test -e {path} && ! test -L {path}" for path in quoted_paths), + ] try: result = await env.exec( - f"rm -f -- {quoted_path} && ! test -e {quoted_path} && ! test -L {quoted_path}", + " && ".join(command_parts), user="root", timeout_sec=10, ) except Exception as exc: logger.warning("Failed to isolate %s subscription credential: %s", agent, exc) return False - if result.return_code != 0: + if result.return_code != 0 or not paths_safe: logger.warning( "Subscription credential remained accessible in proxy mode for %s", agent, diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index 253f602b3..4a19924ba 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -437,6 +437,10 @@ class SubscriptionAuth: replaces_env: str # The env var this substitutes, e.g. "ANTHROPIC_API_KEY" detect_file: str # Host path to check for login, e.g. "~/.claude/.credentials.json" files: list[HostAuthFile] = field(default_factory=list) # All files to copy + config_dir_env: str = "" + # Optional CLI config-home override. Proxy mode removes this variable so a + # reused image cannot select an alternate native-login file outside the + # registry-owned container path. @dataclass @@ -537,6 +541,7 @@ class AgentConfig: subscription_auth=SubscriptionAuth( replaces_env="ANTHROPIC_API_KEY", detect_file="~/.claude/.credentials.json", + config_dir_env="CLAUDE_CONFIG_DIR", files=[ HostAuthFile( "~/.claude/.credentials.json", "{home}/.claude/.credentials.json" @@ -643,6 +648,7 @@ class AgentConfig: subscription_auth=SubscriptionAuth( replaces_env="OPENAI_API_KEY", detect_file="~/.codex/auth.json", + config_dir_env="CODEX_HOME", files=[ HostAuthFile("~/.codex/auth.json", "{home}/.codex/auth.json"), ], diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index d4cd8d6cc..76c1fadaa 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1838,12 +1838,11 @@ async def ensure_litellm_runtime( ), ) - credential_home = agent_env.get("BENCHFLOW_AGENT_HOME", "").strip() or ( - f"/home/{sandbox_user}" if sandbox_user else "/root" - ) + credential_home = f"/home/{sandbox_user}" if sandbox_user else "/root" subscription_credentials_isolated = await isolate_subscription_auth_for_proxy( sandbox, agent=agent, + agent_env=agent_env, cred_home=credential_home, ) diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 55d5c88cc..2879d6c9b 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -5,6 +5,7 @@ import pytest +from benchflow.agents.credentials import isolate_subscription_auth_for_proxy from benchflow.providers import litellm_runtime as runtime_mod from benchflow.providers.litellm_config import resolve_litellm_route from benchflow.providers.runtime import ensure_litellm_runtime @@ -22,38 +23,99 @@ async def exec(self, command, **kwargs): @pytest.mark.asyncio -@pytest.mark.parametrize(("return_code", "trusted"), [(0, True), (1, False)]) -async def test_claude_api_proxy_isolates_preexisting_subscription_credential( - monkeypatch, return_code: int, trusted: bool +@pytest.mark.parametrize( + ("agent", "override_env", "default_auth_path", "return_code", "trusted"), + [ + ( + "claude-agent-acp", + "CLAUDE_CONFIG_DIR", + "/home/agent/.claude/.credentials.json", + 0, + True, + ), + ( + "codex-acp", + "CODEX_HOME", + "/home/agent/.codex/auth.json", + 0, + True, + ), + ( + "claude-agent-acp", + "CLAUDE_CONFIG_DIR", + "/home/agent/.claude/.credentials.json", + 1, + False, + ), + ], +) +async def test_api_proxy_isolates_preexisting_subscription_credential( + monkeypatch, + agent: str, + override_env: str, + default_auth_path: str, + return_code: int, + trusted: bool, ) -> None: - """Guards PR #1057 against a reused Claude OAuth file bypassing capture.""" + """Guards PR #1057 against reused native auth bypassing capture.""" async def fake_start(**_kwargs): return SimpleNamespace(base_url="http://127.0.0.1:4000") monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) sandbox = _SubscriptionIsolationSandbox(return_code=return_code) + agent_env = { + "ANTHROPIC_API_KEY": "sk-ant-provider", + "OPENAI_API_KEY": "sk-openai-provider", + override_env: f"/home/agent/custom-{agent}", + } - _, provider_runtime = await ensure_litellm_runtime( - agent="claude-agent-acp", - agent_env={"ANTHROPIC_API_KEY": "sk-ant-provider"}, - model="claude-sonnet-4-6", + updated, provider_runtime = await ensure_litellm_runtime( + agent=agent, + agent_env=agent_env, + model="claude-sonnet-4-6" if agent == "claude-agent-acp" else "gpt-5.6-sol", runtime=None, environment="docker", - session_id="run-claude-api-with-stale-oauth", + session_id=f"run-{agent}-api-with-stale-oauth", sandbox=sandbox, sandbox_user="agent", ) assert provider_runtime is not None assert provider_runtime.capture_trusted is trusted + assert override_env not in updated + assert override_env not in agent_env assert len(sandbox.commands) == 1 - assert "/home/agent/.claude/.credentials.json" in sandbox.commands[0] + assert default_auth_path in sandbox.commands[0] + assert ( + f"custom-{agent}/{default_auth_path.rsplit('/', 1)[-1]}" in sandbox.commands[0] + ) assert "rm -f" in sandbox.commands[0] assert "! test -e" in sandbox.commands[0] assert "! test -L" in sandbox.commands[0] +@pytest.mark.asyncio +async def test_proxy_auth_cleanup_rejects_override_outside_sandbox_home() -> None: + """Guards PR #1057 against root deletion through a config-home override.""" + + sandbox = _SubscriptionIsolationSandbox() + agent_env = {"CODEX_HOME": "/etc/custom-codex"} + + trusted = await isolate_subscription_auth_for_proxy( + sandbox, + agent="codex-acp", + agent_env=agent_env, + cred_home="/home/agent", + ) + + assert trusted is False + assert "CODEX_HOME" not in agent_env + assert len(sandbox.commands) == 1 + assert "/home/agent/.codex/auth.json" in sandbox.commands[0] + assert "/etc/custom-codex" not in sandbox.commands[0] + + @pytest.mark.asyncio async def test_vertex_adc_provider_capture_remains_audit_only_on_host(monkeypatch): """Guards PR #1057 against trusting agent-accessible Vertex credentials.""" diff --git a/tests/test_subscription_auth.py b/tests/test_subscription_auth.py index 45da6fb54..c57bc751f 100644 --- a/tests/test_subscription_auth.py +++ b/tests/test_subscription_auth.py @@ -20,6 +20,7 @@ def test_claude_subscription_auth(self): assert sa is not None assert sa.replaces_env == "ANTHROPIC_API_KEY" assert ".claude/.credentials.json" in sa.detect_file + assert sa.config_dir_env == "CLAUDE_CONFIG_DIR" assert len(sa.files) == 1 def test_codex_subscription_auth(self): @@ -28,6 +29,7 @@ def test_codex_subscription_auth(self): assert sa is not None assert sa.replaces_env == "OPENAI_API_KEY" assert ".codex/auth.json" in sa.detect_file + assert sa.config_dir_env == "CODEX_HOME" assert len(sa.files) == 1 def test_gemini_subscription_auth(self): From ee737184a1344d69c3779c868ac8efd9a2f29fba Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 19:35:43 -0700 Subject: [PATCH 57/74] Harden effective proxy auth homes --- src/benchflow/agents/credentials.py | 66 ++++++++++++++++------ src/benchflow/providers/litellm_runtime.py | 5 ++ tests/test_litellm_credential_custody.py | 57 +++++++++++++++++++ 3 files changed, 112 insertions(+), 16 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index 38c59fc0d..b54ac2648 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -214,30 +214,38 @@ async def isolate_subscription_auth_for_proxy( ) return False + home = PurePosixPath(cred_home) primary_path = primary_files[0].container_path.format(home=cred_home) + primary_relative_path = PurePosixPath(primary_path).relative_to(home) paths = [primary_path] paths_safe = True + + for home_env in ("HOME", "BENCHFLOW_AGENT_HOME"): + override_value = agent_env.get(home_env, "") + if override_value and override_value != cred_home: + override_home = _safe_proxy_auth_root( + override_value, + cred_home=cred_home, + agent=agent, + ) + if override_home is None: + paths_safe = False + else: + paths.append(str(override_home / primary_relative_path)) + agent_env[home_env] = cred_home + if subscription_auth.config_dir_env: override_value = agent_env.pop(subscription_auth.config_dir_env, "") if override_value: - override_dir = PurePosixPath(override_value) - home = PurePosixPath(cred_home) - try: - override_dir.relative_to(home) - except ValueError: + override_dir = _safe_proxy_auth_root( + override_value, + cred_home=cred_home, + agent=agent, + ) + if override_dir is None: paths_safe = False - if ( - not override_dir.is_absolute() - or ".." in override_dir.parts - or "\x00" in override_value - ): - paths_safe = False - if paths_safe: - paths.append(str(override_dir / PurePosixPath(primary_path).name)) else: - logger.warning( - "Refusing root cleanup outside the sandbox home for %s", agent - ) + paths.append(str(override_dir / PurePosixPath(primary_path).name)) paths = list(dict.fromkeys(paths)) quoted_paths = [shlex.quote(path) for path in paths] @@ -270,3 +278,29 @@ async def isolate_subscription_auth_for_proxy( ) return False return True + + +def _safe_proxy_auth_root( + value: str, + *, + cred_home: str, + agent: str, +) -> PurePosixPath | None: + """Return an in-home auth root, refusing unsafe root-owned deletion paths.""" + + candidate = PurePosixPath(value) + home = PurePosixPath(cred_home) + try: + candidate.relative_to(home) + except ValueError: + candidate_safe = False + else: + candidate_safe = bool( + candidate.is_absolute() + and ".." not in candidate.parts + and "\x00" not in value + ) + if not candidate_safe: + logger.warning("Refusing root cleanup outside the sandbox home for %s", agent) + return None + return candidate diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 76c1fadaa..51f6b5640 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1606,6 +1606,11 @@ def _litellm_proxy_env( *, agent: str, agent_env: dict[str, str], required_skill_names: tuple[str, ...] ) -> dict[str, str]: updated = dict(agent_env) + # These values describe the sandbox agent, not the LiteLLM process. A host + # proxy must retain its real host HOME; a sandbox-local proxy runs as root + # and likewise must not inherit the unprivileged agent's config root. + updated.pop("HOME", None) + updated.pop("BENCHFLOW_AGENT_HOME", None) updated.pop(_SKILL_CATALOG_GATE_AGENT_ENV, None) updated.pop(_REQUIRED_SKILL_NAMES_ENV, None) expected = sorted(set(required_skill_names)) diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 2879d6c9b..95e61f917 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -85,6 +85,8 @@ async def fake_start(**_kwargs): assert provider_runtime.capture_trusted is trusted assert override_env not in updated assert override_env not in agent_env + assert updated["HOME"] == "/home/agent" + assert updated["BENCHFLOW_AGENT_HOME"] == "/home/agent" assert len(sandbox.commands) == 1 assert default_auth_path in sandbox.commands[0] assert ( @@ -116,6 +118,61 @@ async def test_proxy_auth_cleanup_rejects_override_outside_sandbox_home() -> Non assert "/etc/custom-codex" not in sandbox.commands[0] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("custom_home", "trusted", "expected_in_command"), + [ + ("/root/custom-agent-home", True, True), + ("/tmp/custom-agent-home", False, False), + ], +) +async def test_root_proxy_auth_cleanup_handles_effective_agent_home( + monkeypatch, + custom_home: str, + trusted: bool, + expected_in_command: bool, +) -> None: + """Guards PR #1057 against root HOME redirects bypassing capture.""" + + proxy_agent_env: dict[str, str] = {} + + async def fake_start(**kwargs): + proxy_agent_env.update(kwargs["agent_env"]) + return SimpleNamespace(base_url="http://127.0.0.1:4000") + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + sandbox = _SubscriptionIsolationSandbox() + agent_env = { + "ANTHROPIC_API_KEY": "sk-ant-provider", + "HOME": custom_home, + "BENCHFLOW_AGENT_HOME": custom_home, + } + + updated, provider_runtime = await ensure_litellm_runtime( + agent="claude-agent-acp", + agent_env=agent_env, + model="claude-sonnet-4-6", + runtime=None, + environment="docker", + session_id="run-root-claude-api-with-stale-oauth", + sandbox=sandbox, + sandbox_user=None, + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is trusted + assert updated["HOME"] == "/root" + assert updated["BENCHFLOW_AGENT_HOME"] == "/root" + assert agent_env["HOME"] == "/root" + assert agent_env["BENCHFLOW_AGENT_HOME"] == "/root" + assert "HOME" not in proxy_agent_env + assert "BENCHFLOW_AGENT_HOME" not in proxy_agent_env + assert (f"{custom_home}/.claude/.credentials.json" in sandbox.commands[0]) is ( + expected_in_command + ) + assert "/root/.claude/.credentials.json" in sandbox.commands[0] + + @pytest.mark.asyncio async def test_vertex_adc_provider_capture_remains_audit_only_on_host(monkeypatch): """Guards PR #1057 against trusting agent-accessible Vertex credentials.""" From 42e3d99e309d7a92feebb8f58b88a3c0a1ee9edc Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 20:07:01 -0700 Subject: [PATCH 58/74] fix(auth): make proxy credential cleanup race-safe --- src/benchflow/agents/credentials.py | 134 ++++++++++++++++++++--- tests/test_litellm_credential_custody.py | 73 +++++++++++- 2 files changed, 185 insertions(+), 22 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index b54ac2648..2379356bd 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -29,6 +29,103 @@ logger = logging.getLogger(__name__) +_BENCHFLOW_NODE_BIN = "/opt/benchflow/node/bin/node" +_PROXY_AUTH_CLEANUP_JS = r""" +const fs = require("fs"); + +const sandboxUid = Number(process.argv[1]); +const targets = JSON.parse(process.argv[2]); +const constants = fs.constants; +const directoryFlags = + constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW; + +function processesForUid(uid) { + const matches = []; + for (const entry of fs.readdirSync("/proc")) { + if (!/^\d+$/.test(entry) || Number(entry) === process.pid) continue; + try { + const status = fs.readFileSync(`/proc/${entry}/status`, "utf8"); + const owner = /^Uid:\s+(\d+)/m.exec(status); + const state = /^State:\s+(\S+)/m.exec(status); + if (owner && Number(owner[1]) === uid && (!state || state[1] !== "Z")) { + matches.push(Number(entry)); + } + } catch (error) { + if (error.code !== "ENOENT") throw error; + } + } + return matches; +} + +function stopPriorAgentProcesses(uid) { + if (!Number.isInteger(uid) || uid < 0) return; + for (let attempt = 0; attempt < 20; attempt += 1) { + const pids = processesForUid(uid); + if (pids.length === 0) return; + for (const pid of pids) { + try { + process.kill(pid, "SIGKILL"); + } catch (error) { + if (error.code !== "ESRCH") throw error; + } + } + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 5); + } + throw new Error("sandbox user still has a live process after isolation"); +} + +function openDirectoryNoFollow(parts) { + let descriptor = fs.openSync("/", directoryFlags); + try { + for (const part of parts) { + const next = fs.openSync( + `/proc/self/fd/${descriptor}/${part}`, + directoryFlags, + ); + fs.closeSync(descriptor); + descriptor = next; + } + return descriptor; + } catch (error) { + fs.closeSync(descriptor); + throw error; + } +} + +function removeCredentialNoFollow(target) { + const parts = target.split("/").filter(Boolean); + const basename = parts.pop(); + let parent; + try { + parent = openDirectoryNoFollow(parts); + } catch (error) { + if (error.code === "ENOENT") return; + throw error; + } + try { + const descriptorPath = `/proc/self/fd/${parent}/${basename}`; + try { + fs.unlinkSync(descriptorPath); + } catch (error) { + if (error.code !== "ENOENT") throw error; + } + try { + fs.lstatSync(descriptorPath); + } catch (error) { + if (error.code === "ENOENT") return; + throw error; + } + throw new Error(`credential remained after deletion: ${target}`); + } finally { + fs.closeSync(parent); + } +} + +stopPriorAgentProcesses(sandboxUid); +for (const target of targets) removeCredentialNoFollow(target); +""".strip() + + def _owner_from_home(cred_home: str) -> str | None: """Return sandbox username for /home/ credential homes.""" parts = Path(cred_home).parts @@ -194,9 +291,11 @@ async def isolate_subscription_auth_for_proxy( ``SubscriptionAuth.detect_file``. API-key mode must not leave that alternate provider route available to the agent. A safe CLI config-home override is scrubbed as well as removed before launch; an override outside the sandbox - user's home fails capture trust without allowing a root deletion there. - Return false unless every eligible credential can be identified, removed, - and verified absent. + user's home fails capture trust without allowing a root deletion there. The + cleanup stops stale processes for a non-root sandbox user, then uses the + JavaScript runtime already required by every subscription-capable agent to + traverse with no-follow directory descriptors. Return false unless every + eligible credential can be identified, removed, and verified absent. """ agent_cfg = AGENTS.get(agent) @@ -248,23 +347,20 @@ async def isolate_subscription_auth_for_proxy( paths.append(str(override_dir / PurePosixPath(primary_path).name)) paths = list(dict.fromkeys(paths)) - quoted_paths = [shlex.quote(path) for path in paths] - parent_paths: set[str] = set() - for path in paths: - parent = PurePosixPath(path).parent - parent_paths.update( - str(candidate) - for candidate in (*reversed(parent.parents), parent) - if str(candidate) != "/" + sandbox_owner = _owner_from_home(cred_home) + uid_command = f"$(id -u -- {shlex.quote(sandbox_owner)})" if sandbox_owner else "-1" + cleanup_command = " ".join( + ( + f"{_BENCHFLOW_NODE_BIN} -e", + shlex.quote(_PROXY_AUTH_CLEANUP_JS), + "--", + uid_command, + shlex.quote(json.dumps(paths, separators=(",", ":"))), ) - command_parts = [ - *(f"! test -L {shlex.quote(parent)}" for parent in sorted(parent_paths)), - f"rm -f -- {' '.join(quoted_paths)}", - *(f"! test -e {path} && ! test -L {path}" for path in quoted_paths), - ] + ) try: result = await env.exec( - " && ".join(command_parts), + cleanup_command, user="root", timeout_sec=10, ) @@ -272,9 +368,11 @@ async def isolate_subscription_auth_for_proxy( logger.warning("Failed to isolate %s subscription credential: %s", agent, exc) return False if result.return_code != 0 or not paths_safe: + detail = (getattr(result, "stderr", "") or "").strip()[:500] logger.warning( - "Subscription credential remained accessible in proxy mode for %s", + "Subscription credential remained accessible in proxy mode for %s%s", agent, + f": {detail}" if detail else "", ) return False return True diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 95e61f917..1f162a483 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -1,11 +1,17 @@ from __future__ import annotations import json +import shutil +import subprocess +import sys from types import SimpleNamespace import pytest -from benchflow.agents.credentials import isolate_subscription_auth_for_proxy +from benchflow.agents.credentials import ( + _PROXY_AUTH_CLEANUP_JS, + isolate_subscription_auth_for_proxy, +) from benchflow.providers import litellm_runtime as runtime_mod from benchflow.providers.litellm_config import resolve_litellm_route from benchflow.providers.runtime import ensure_litellm_runtime @@ -92,9 +98,10 @@ async def fake_start(**_kwargs): assert ( f"custom-{agent}/{default_auth_path.rsplit('/', 1)[-1]}" in sandbox.commands[0] ) - assert "rm -f" in sandbox.commands[0] - assert "! test -e" in sandbox.commands[0] - assert "! test -L" in sandbox.commands[0] + assert "O_NOFOLLOW" in sandbox.commands[0] + assert "/proc/self/fd/" in sandbox.commands[0] + assert "SIGKILL" in sandbox.commands[0] + assert "rm -f" not in sandbox.commands[0] @pytest.mark.asyncio @@ -118,6 +125,64 @@ async def test_proxy_auth_cleanup_rejects_override_outside_sandbox_home() -> Non assert "/etc/custom-codex" not in sandbox.commands[0] +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("node") is None, + reason="the sandbox cleanup primitive requires Linux procfs and Node.js", +) +def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: + """Guards PR #1057 review r3888287847 against credential cleanup TOCTOU.""" + + outside = tmp_path / "outside" + outside.mkdir() + outside_auth = outside / "auth.json" + outside_auth.write_text("subscription-secret") + + sandbox_home = tmp_path / "home" / "agent" + sandbox_home.mkdir(parents=True) + auth_dir = sandbox_home / ".codex" + auth_dir.symlink_to(outside, target_is_directory=True) + credential_path = auth_dir / "auth.json" + + refused = subprocess.run( + [ + "node", + "-e", + _PROXY_AUTH_CLEANUP_JS, + "--", + "-1", + json.dumps([str(credential_path)]), + ], + check=False, + capture_output=True, + text=True, + ) + + assert refused.returncode != 0 + assert outside_auth.read_text() == "subscription-secret" + + auth_dir.unlink() + auth_dir.mkdir() + credential_path.symlink_to(outside_auth) + removed = subprocess.run( + [ + "node", + "-e", + _PROXY_AUTH_CLEANUP_JS, + "--", + "-1", + json.dumps([str(credential_path)]), + ], + check=False, + capture_output=True, + text=True, + ) + + assert removed.returncode == 0, removed.stderr + assert not credential_path.exists() + assert not credential_path.is_symlink() + assert outside_auth.read_text() == "subscription-secret" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("custom_home", "trusted", "expected_in_command"), From d4557e9721bddbb38b13653cfb041dd18a564aed Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 20:42:18 -0700 Subject: [PATCH 59/74] fix(auth): scope stale process isolation --- src/benchflow/agents/credentials.py | 61 +++++++++++++++++++----- tests/test_litellm_credential_custody.py | 51 ++++++++++++++++++-- 2 files changed, 96 insertions(+), 16 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index 2379356bd..bcf29984a 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -29,12 +29,14 @@ logger = logging.getLogger(__name__) -_BENCHFLOW_NODE_BIN = "/opt/benchflow/node/bin/node" +_BENCHFLOW_PREFIX = "/opt/benchflow" +_BENCHFLOW_NODE_BIN = f"{_BENCHFLOW_PREFIX}/node/bin/node" _PROXY_AUTH_CLEANUP_JS = r""" const fs = require("fs"); const sandboxUid = Number(process.argv[1]); -const targets = JSON.parse(process.argv[2]); +const agentProcessMarkers = JSON.parse(process.argv[2]); +const targets = JSON.parse(process.argv[3]); const constants = fs.constants; const directoryFlags = constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW; @@ -48,7 +50,11 @@ const owner = /^Uid:\s+(\d+)/m.exec(status); const state = /^State:\s+(\S+)/m.exec(status); if (owner && Number(owner[1]) === uid && (!state || state[1] !== "Z")) { - matches.push(Number(entry)); + const argv = fs + .readFileSync(`/proc/${entry}/cmdline`, "utf8") + .split("\0") + .filter(Boolean); + matches.push({pid: Number(entry), argv}); } } catch (error) { if (error.code !== "ENOENT") throw error; @@ -57,12 +63,22 @@ return matches; } -function stopPriorAgentProcesses(uid) { - if (!Number.isInteger(uid) || uid < 0) return; +function isManagedAgentProcess(processInfo, markers) { + return markers.some( + (marker) => processInfo.argv[0] === marker || processInfo.argv[1] === marker, + ); +} + +function quiescePriorAgentProcesses(uid, markers) { + if (!Number.isInteger(uid) || uid < 0) { + throw new Error("root sandbox agents cannot prove stale-process isolation"); + } for (let attempt = 0; attempt < 20; attempt += 1) { - const pids = processesForUid(uid); - if (pids.length === 0) return; - for (const pid of pids) { + const agents = processesForUid(uid).filter((processInfo) => + isManagedAgentProcess(processInfo, markers), + ); + if (agents.length === 0) break; + for (const {pid} of agents) { try { process.kill(pid, "SIGKILL"); } catch (error) { @@ -71,7 +87,16 @@ } Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 5); } - throw new Error("sandbox user still has a live process after isolation"); + const remaining = processesForUid(uid); + const staleAgents = remaining.filter((processInfo) => + isManagedAgentProcess(processInfo, markers), + ); + if (staleAgents.length !== 0) { + throw new Error("stale agent process survived credential isolation"); + } + if (remaining.length !== 0) { + throw new Error("unrelated sandbox-user process prevents trusted cleanup"); + } } function openDirectoryNoFollow(parts) { @@ -121,7 +146,7 @@ } } -stopPriorAgentProcesses(sandboxUid); +quiescePriorAgentProcesses(sandboxUid, agentProcessMarkers); for (const target of targets) removeCredentialNoFollow(target); """.strip() @@ -292,10 +317,13 @@ async def isolate_subscription_auth_for_proxy( provider route available to the agent. A safe CLI config-home override is scrubbed as well as removed before launch; an override outside the sandbox user's home fails capture trust without allowing a root deletion there. The - cleanup stops stale processes for a non-root sandbox user, then uses the + cleanup stops only stale ACP processes for a non-root sandbox user and + refuses trust when unrelated user processes are alive. It then uses the JavaScript runtime already required by every subscription-capable agent to - traverse with no-follow directory descriptors. Return false unless every - eligible credential can be identified, removed, and verified absent. + traverse with no-follow directory descriptors. Root-agent runs remain + audit-only because stale root processes cannot be safely distinguished. + Return false unless every eligible credential can be identified, removed, + and verified absent. """ agent_cfg = AGENTS.get(agent) @@ -348,13 +376,20 @@ async def isolate_subscription_auth_for_proxy( paths = list(dict.fromkeys(paths)) sandbox_owner = _owner_from_home(cred_home) + if sandbox_owner is None: + paths_safe = False uid_command = f"$(id -u -- {shlex.quote(sandbox_owner)})" if sandbox_owner else "-1" + agent_process_markers = ( + f"{_BENCHFLOW_PREFIX}/js-agents/bin/{agent}", + f"{_BENCHFLOW_PREFIX}/bin/{agent}", + ) cleanup_command = " ".join( ( f"{_BENCHFLOW_NODE_BIN} -e", shlex.quote(_PROXY_AUTH_CLEANUP_JS), "--", uid_command, + shlex.quote(json.dumps(agent_process_markers, separators=(",", ":"))), shlex.quote(json.dumps(paths, separators=(",", ":"))), ) ) diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 1f162a483..6de8639f8 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import shutil import subprocess import sys @@ -149,7 +150,8 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: "-e", _PROXY_AUTH_CLEANUP_JS, "--", - "-1", + "2147483646", + json.dumps(["/opt/benchflow/js-agents/bin/codex-acp"]), json.dumps([str(credential_path)]), ], check=False, @@ -169,7 +171,8 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: "-e", _PROXY_AUTH_CLEANUP_JS, "--", - "-1", + "2147483646", + json.dumps(["/opt/benchflow/js-agents/bin/codex-acp"]), json.dumps([str(credential_path)]), ], check=False, @@ -183,11 +186,53 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: assert outside_auth.read_text() == "subscription-secret" +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("node") is None, + reason="the sandbox cleanup primitive requires Linux procfs and Node.js", +) +def test_proxy_auth_cleanup_preserves_unrelated_user_processes(tmp_path) -> None: + """Guards PR #1057 review r3888337690 against killing task processes.""" + + credential_path = tmp_path / ".codex" / "auth.json" + credential_path.parent.mkdir() + credential_path.write_text("subscription-secret") + absent_agent_marker = "/opt/benchflow/js-agents/bin/not-this-process" + + unrelated = subprocess.Popen( + ["node", "-e", "setInterval(() => {}, 1000)", absent_agent_marker], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + refused = subprocess.run( + [ + "node", + "-e", + _PROXY_AUTH_CLEANUP_JS, + "--", + str(os.getuid()), + json.dumps([absent_agent_marker]), + json.dumps([str(credential_path)]), + ], + check=False, + capture_output=True, + text=True, + ) + + assert refused.returncode != 0 + assert "unrelated sandbox-user process" in refused.stderr + assert credential_path.read_text() == "subscription-secret" + assert unrelated.poll() is None + finally: + unrelated.terminate() + unrelated.wait(timeout=5) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("custom_home", "trusted", "expected_in_command"), [ - ("/root/custom-agent-home", True, True), + ("/root/custom-agent-home", False, True), ("/tmp/custom-agent-home", False, False), ], ) From 88fc1b5434ad743d4debd5a7b8c73dc3710cd303 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 21:10:38 -0700 Subject: [PATCH 60/74] fix(auth): gate every proxy agent on process isolation --- src/benchflow/agents/credentials.py | 136 ++++++++------------- src/benchflow/providers/litellm_runtime.py | 7 +- tests/test_litellm_credential_custody.py | 133 ++++++++++++++++---- tests/test_litellm_runtime.py | 3 + 4 files changed, 164 insertions(+), 115 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index bcf29984a..c45604102 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -5,7 +5,7 @@ - write_credential_files agent + provider credential files (cf. AgentConfig) - write_gemini_vertex_settings ~/.gemini/settings.json for Vertex backend - upload_subscription_auth host login files (e.g. ~/.claude/.credentials.json) - - isolate_subscription_auth_for_proxy remove stale native login in API mode + - isolate_agent_for_proxy_capture prove process/auth isolation in API mode The Gemini Vertex settings helper lives here (not in _agent_env.py) so the module has a single coherent role and zero horizontal imports between phase @@ -34,71 +34,11 @@ _PROXY_AUTH_CLEANUP_JS = r""" const fs = require("fs"); -const sandboxUid = Number(process.argv[1]); -const agentProcessMarkers = JSON.parse(process.argv[2]); -const targets = JSON.parse(process.argv[3]); +const targets = JSON.parse(process.argv[1]); const constants = fs.constants; const directoryFlags = constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW; -function processesForUid(uid) { - const matches = []; - for (const entry of fs.readdirSync("/proc")) { - if (!/^\d+$/.test(entry) || Number(entry) === process.pid) continue; - try { - const status = fs.readFileSync(`/proc/${entry}/status`, "utf8"); - const owner = /^Uid:\s+(\d+)/m.exec(status); - const state = /^State:\s+(\S+)/m.exec(status); - if (owner && Number(owner[1]) === uid && (!state || state[1] !== "Z")) { - const argv = fs - .readFileSync(`/proc/${entry}/cmdline`, "utf8") - .split("\0") - .filter(Boolean); - matches.push({pid: Number(entry), argv}); - } - } catch (error) { - if (error.code !== "ENOENT") throw error; - } - } - return matches; -} - -function isManagedAgentProcess(processInfo, markers) { - return markers.some( - (marker) => processInfo.argv[0] === marker || processInfo.argv[1] === marker, - ); -} - -function quiescePriorAgentProcesses(uid, markers) { - if (!Number.isInteger(uid) || uid < 0) { - throw new Error("root sandbox agents cannot prove stale-process isolation"); - } - for (let attempt = 0; attempt < 20; attempt += 1) { - const agents = processesForUid(uid).filter((processInfo) => - isManagedAgentProcess(processInfo, markers), - ); - if (agents.length === 0) break; - for (const {pid} of agents) { - try { - process.kill(pid, "SIGKILL"); - } catch (error) { - if (error.code !== "ESRCH") throw error; - } - } - Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 5); - } - const remaining = processesForUid(uid); - const staleAgents = remaining.filter((processInfo) => - isManagedAgentProcess(processInfo, markers), - ); - if (staleAgents.length !== 0) { - throw new Error("stale agent process survived credential isolation"); - } - if (remaining.length !== 0) { - throw new Error("unrelated sandbox-user process prevents trusted cleanup"); - } -} - function openDirectoryNoFollow(parts) { let descriptor = fs.openSync("/", directoryFlags); try { @@ -146,7 +86,6 @@ } } -quiescePriorAgentProcesses(sandboxUid, agentProcessMarkers); for (const target of targets) removeCredentialNoFollow(target); """.strip() @@ -159,6 +98,23 @@ def _owner_from_home(cred_home: str) -> str | None: return None +def _proxy_process_isolation_guard(cred_home: str) -> tuple[str, bool]: + """Return a root shell guard proving the agent UID has no live process.""" + + sandbox_owner = _owner_from_home(cred_home) + if sandbox_owner is None: + return "false", False + quoted_owner = shlex.quote(sandbox_owner) + return ( + "command -v pgrep >/dev/null 2>&1 && " + f"bf_agent_uid=$(id -u -- {quoted_owner}) && " + '[ "$bf_agent_uid" -ne 0 ] && ' + '{ pgrep -u "$bf_agent_uid" >/dev/null 2>&1; ' + 'bf_pgrep_rc=$?; [ "$bf_pgrep_rc" -eq 1 ]; }', + True, + ) + + async def upload_credential( env, path: str, @@ -303,39 +259,55 @@ async def upload_subscription_auth( ) -async def isolate_subscription_auth_for_proxy( +async def isolate_agent_for_proxy_capture( env, *, agent: str, agent_env: dict[str, str], cred_home: str, ) -> bool: - """Remove the primary native-login credential before an API proxy run. + """Prove process isolation and remove native auth before an API proxy run. A reused image may already contain the CLI login detected by ``SubscriptionAuth.detect_file``. API-key mode must not leave that alternate provider route available to the agent. A safe CLI config-home override is scrubbed as well as removed before launch; an override outside the sandbox - user's home fails capture trust without allowing a root deletion there. The - cleanup stops only stale ACP processes for a non-root sandbox user and - refuses trust when unrelated user processes are alive. It then uses the - JavaScript runtime already required by every subscription-capable agent to - traverse with no-follow directory descriptors. Root-agent runs remain - audit-only because stale root processes cannot be safely distinguished. - Return false unless every eligible credential can be identified, removed, - and verified absent. + user's home fails capture trust without allowing a root deletion there. + Every API-proxied agent first requires a non-root sandbox user with no live + processes; an existing agent or task process is preserved and makes capture + audit-only. Subscription-capable agents then use their already-required + JavaScript runtime to traverse with no-follow directory descriptors. + Root-agent runs remain audit-only because stale root processes cannot be + safely distinguished. Return false unless process isolation is proven and + every eligible credential can be identified, removed, and verified absent. """ agent_cfg = AGENTS.get(agent) subscription_auth = agent_cfg.subscription_auth if agent_cfg else None + if env is None: + logger.warning("Cannot prove proxy process isolation for %s", agent) + return False + + process_guard, processes_safe = _proxy_process_isolation_guard(cred_home) + if subscription_auth is None: - return True + try: + result = await env.exec( + process_guard, + user="root", + timeout_sec=10, + ) + except Exception as exc: + logger.warning("Failed to isolate %s agent processes: %s", agent, exc) + return False + return bool(result.return_code == 0 and processes_safe) + primary_files = [ auth_file for auth_file in subscription_auth.files if auth_file.host_path == subscription_auth.detect_file ] - if len(primary_files) != 1 or env is None: + if len(primary_files) != 1: logger.warning( "Cannot prove proxy isolation for %s subscription credentials", agent ) @@ -375,21 +347,11 @@ async def isolate_subscription_auth_for_proxy( paths.append(str(override_dir / PurePosixPath(primary_path).name)) paths = list(dict.fromkeys(paths)) - sandbox_owner = _owner_from_home(cred_home) - if sandbox_owner is None: - paths_safe = False - uid_command = f"$(id -u -- {shlex.quote(sandbox_owner)})" if sandbox_owner else "-1" - agent_process_markers = ( - f"{_BENCHFLOW_PREFIX}/js-agents/bin/{agent}", - f"{_BENCHFLOW_PREFIX}/bin/{agent}", - ) - cleanup_command = " ".join( + cleanup_command = f"{process_guard} && " + " ".join( ( f"{_BENCHFLOW_NODE_BIN} -e", shlex.quote(_PROXY_AUTH_CLEANUP_JS), "--", - uid_command, - shlex.quote(json.dumps(agent_process_markers, separators=(",", ":"))), shlex.quote(json.dumps(paths, separators=(",", ":"))), ) ) @@ -402,7 +364,7 @@ async def isolate_subscription_auth_for_proxy( except Exception as exc: logger.warning("Failed to isolate %s subscription credential: %s", agent, exc) return False - if result.return_code != 0 or not paths_safe: + if result.return_code != 0 or not paths_safe or not processes_safe: detail = (getattr(result, "stderr", "") or "").strip()[:500] logger.warning( "Subscription credential remained accessible in proxy mode for %s%s", diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 51f6b5640..3644944c0 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -28,7 +28,7 @@ from benchflow._utils.text import describe_exception from benchflow.agents.codex_config import apply_codex_proxy_config -from benchflow.agents.credentials import isolate_subscription_auth_for_proxy +from benchflow.agents.credentials import isolate_agent_for_proxy_capture from benchflow.agents.env import uses_native_subscription_auth from benchflow.agents.registry import AGENTS from benchflow.providers.litellm_bedrock_preflight import ( @@ -1844,7 +1844,7 @@ async def ensure_litellm_runtime( ) credential_home = f"/home/{sandbox_user}" if sandbox_user else "/root" - subscription_credentials_isolated = await isolate_subscription_auth_for_proxy( + agent_capture_isolated = await isolate_agent_for_proxy_capture( sandbox, agent=agent, agent_env=agent_env, @@ -1860,8 +1860,7 @@ async def ensure_litellm_runtime( ) proxy_location = "sandbox" if sandbox_local else "host" credentials_isolated = bool( - _provider_credentials_have_proxy_only_custody(route) - and subscription_credentials_isolated + _provider_credentials_have_proxy_only_custody(route) and agent_capture_isolated ) capture_trusted = not sandbox_local and credentials_isolated config_key = ( diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 6de8639f8..7c4fcf47b 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -11,7 +11,8 @@ from benchflow.agents.credentials import ( _PROXY_AUTH_CLEANUP_JS, - isolate_subscription_auth_for_proxy, + _proxy_process_isolation_guard, + isolate_agent_for_proxy_capture, ) from benchflow.providers import litellm_runtime as runtime_mod from benchflow.providers.litellm_config import resolve_litellm_route @@ -101,7 +102,8 @@ async def fake_start(**_kwargs): ) assert "O_NOFOLLOW" in sandbox.commands[0] assert "/proc/self/fd/" in sandbox.commands[0] - assert "SIGKILL" in sandbox.commands[0] + assert "pgrep -u" in sandbox.commands[0] + assert "SIGKILL" not in sandbox.commands[0] assert "rm -f" not in sandbox.commands[0] @@ -112,7 +114,7 @@ async def test_proxy_auth_cleanup_rejects_override_outside_sandbox_home() -> Non sandbox = _SubscriptionIsolationSandbox() agent_env = {"CODEX_HOME": "/etc/custom-codex"} - trusted = await isolate_subscription_auth_for_proxy( + trusted = await isolate_agent_for_proxy_capture( sandbox, agent="codex-acp", agent_env=agent_env, @@ -126,9 +128,105 @@ async def test_proxy_auth_cleanup_rejects_override_outside_sandbox_home() -> Non assert "/etc/custom-codex" not in sandbox.commands[0] +@pytest.mark.asyncio +@pytest.mark.parametrize("agent", ["opencode", "openhands", "pi-acp"]) +async def test_proxy_process_guard_applies_without_subscription_auth(agent) -> None: + """Guards PR #1057 review r3888399860 across every proxied API agent.""" + + sandbox = _SubscriptionIsolationSandbox() + + trusted = await isolate_agent_for_proxy_capture( + sandbox, + agent=agent, + agent_env={}, + cred_home="/home/agent", + ) + + assert trusted is True + assert len(sandbox.commands) == 1 + assert "pgrep -u" in sandbox.commands[0] + assert '"$bf_agent_uid" -ne 0' in sandbox.commands[0] + assert '"$bf_pgrep_rc" -eq 1' in sandbox.commands[0] + assert _PROXY_AUTH_CLEANUP_JS not in sandbox.commands[0] + + +@pytest.mark.asyncio +async def test_proxy_process_guard_rejects_root_agent_without_subscription_auth() -> ( + None +): + """Guards PR #1057 review r3888399860 against trusted root OpenCode runs.""" + + sandbox = _SubscriptionIsolationSandbox() + + trusted = await isolate_agent_for_proxy_capture( + sandbox, + agent="opencode", + agent_env={}, + cred_home="/root", + ) + + assert trusted is False + assert sandbox.commands == ["false"] + + +@pytest.mark.parametrize( + ("pgrep_return_code", "trusted"), + [(0, False), (1, True), (2, False)], +) +def test_proxy_process_guard_accepts_only_exact_no_match_exit( + tmp_path, pgrep_return_code: int, trusted: bool +) -> None: + """Guards PR #1057 review r3888399860 against process-probe errors.""" + + fake_id = tmp_path / "id" + fake_id.write_text("#!/bin/sh\nprintf '1000\\n'\n") + fake_id.chmod(0o755) + fake_pgrep = tmp_path / "pgrep" + fake_pgrep.write_text(f"#!/bin/sh\nexit {pgrep_return_code}\n") + fake_pgrep.chmod(0o755) + guard, owner_safe = _proxy_process_isolation_guard("/home/agent") + + result = subprocess.run( + ["sh", "-c", guard], + check=False, + env={"PATH": f"{tmp_path}:{os.environ['PATH']}"}, + ) + + assert owner_safe is True + assert (result.returncode == 0) is trusted + + +@pytest.mark.asyncio +async def test_root_opencode_proxy_capture_remains_audit_only(monkeypatch) -> None: + """Guards PR #1057 review r3888399860 at the runtime trust boundary.""" + + async def fake_start(**_kwargs): + return SimpleNamespace(base_url="http://127.0.0.1:4000") + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + sandbox = _SubscriptionIsolationSandbox() + + _, provider_runtime = await ensure_litellm_runtime( + agent="opencode", + agent_env={"OPENAI_API_KEY": "sk-provider"}, + model="openai/gpt-5.6-luna", + runtime=None, + environment="docker", + session_id="run-root-opencode", + sandbox=sandbox, + sandbox_user=None, + ) + + assert provider_runtime is not None + assert provider_runtime.capture_trusted is False + assert sandbox.commands == ["false"] + + @pytest.mark.skipif( - sys.platform != "linux" or shutil.which("node") is None, - reason="the sandbox cleanup primitive requires Linux procfs and Node.js", + sys.platform != "linux" + or shutil.which("node") is None + or shutil.which("pgrep") is None, + reason="the sandbox process guard requires Linux, Node.js, and pgrep", ) def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: """Guards PR #1057 review r3888287847 against credential cleanup TOCTOU.""" @@ -150,8 +248,6 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: "-e", _PROXY_AUTH_CLEANUP_JS, "--", - "2147483646", - json.dumps(["/opt/benchflow/js-agents/bin/codex-acp"]), json.dumps([str(credential_path)]), ], check=False, @@ -171,8 +267,6 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: "-e", _PROXY_AUTH_CLEANUP_JS, "--", - "2147483646", - json.dumps(["/opt/benchflow/js-agents/bin/codex-acp"]), json.dumps([str(credential_path)]), ], check=False, @@ -187,8 +281,10 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: @pytest.mark.skipif( - sys.platform != "linux" or shutil.which("node") is None, - reason="the sandbox cleanup primitive requires Linux procfs and Node.js", + sys.platform != "linux" + or shutil.which("node") is None + or shutil.which("pgrep") is None, + reason="the sandbox process guard requires Linux, Node.js, and pgrep", ) def test_proxy_auth_cleanup_preserves_unrelated_user_processes(tmp_path) -> None: """Guards PR #1057 review r3888337690 against killing task processes.""" @@ -196,31 +292,20 @@ def test_proxy_auth_cleanup_preserves_unrelated_user_processes(tmp_path) -> None credential_path = tmp_path / ".codex" / "auth.json" credential_path.parent.mkdir() credential_path.write_text("subscription-secret") - absent_agent_marker = "/opt/benchflow/js-agents/bin/not-this-process" - unrelated = subprocess.Popen( - ["node", "-e", "setInterval(() => {}, 1000)", absent_agent_marker], + ["node", "-e", "setInterval(() => {}, 1000)"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) try: refused = subprocess.run( - [ - "node", - "-e", - _PROXY_AUTH_CLEANUP_JS, - "--", - str(os.getuid()), - json.dumps([absent_agent_marker]), - json.dumps([str(credential_path)]), - ], + ["sh", "-c", '! pgrep -u "$(id -u)" >/dev/null 2>&1'], check=False, capture_output=True, text=True, ) assert refused.returncode != 0 - assert "unrelated sandbox-user process" in refused.stderr assert credential_path.read_text() == "subscription-secret" assert unrelated.poll() is None finally: diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index c3e95af67..f340a2301 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -180,6 +180,9 @@ async def upload_file(self, source, target, *, mode): self.probe_path = target async def exec(self, command, **kwargs): + if command.startswith("command -v pgrep"): + assert kwargs["user"] == "root" + return SimpleNamespace(return_code=0, stdout="", stderr="") if command == "id -u -- agent": assert kwargs["user"] == "root" return SimpleNamespace(return_code=0, stdout="1000\n", stderr="") From 7183e738c0f8536046c7d55cb36749b74d0aebe8 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 21:38:05 -0700 Subject: [PATCH 61/74] fix(auth): sanitize Claude credential settings --- src/benchflow/agents/credentials.py | 89 +++----- src/benchflow/agents/registry.py | 12 ++ .../agents/resources/proxy_auth_cleanup.js | 193 ++++++++++++++++++ tests/test_litellm_credential_custody.py | 118 +++++++++++ tests/test_registry_invariants.py | 10 + 5 files changed, 364 insertions(+), 58 deletions(-) create mode 100644 src/benchflow/agents/resources/proxy_auth_cleanup.js diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index c45604102..ae07067cf 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -22,6 +22,7 @@ import os import shlex import tempfile +from importlib.resources import files as resource_files from pathlib import Path, PurePosixPath from benchflow.agents.registry import AGENTS @@ -31,63 +32,12 @@ _BENCHFLOW_PREFIX = "/opt/benchflow" _BENCHFLOW_NODE_BIN = f"{_BENCHFLOW_PREFIX}/node/bin/node" -_PROXY_AUTH_CLEANUP_JS = r""" -const fs = require("fs"); - -const targets = JSON.parse(process.argv[1]); -const constants = fs.constants; -const directoryFlags = - constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW; - -function openDirectoryNoFollow(parts) { - let descriptor = fs.openSync("/", directoryFlags); - try { - for (const part of parts) { - const next = fs.openSync( - `/proc/self/fd/${descriptor}/${part}`, - directoryFlags, - ); - fs.closeSync(descriptor); - descriptor = next; - } - return descriptor; - } catch (error) { - fs.closeSync(descriptor); - throw error; - } -} - -function removeCredentialNoFollow(target) { - const parts = target.split("/").filter(Boolean); - const basename = parts.pop(); - let parent; - try { - parent = openDirectoryNoFollow(parts); - } catch (error) { - if (error.code === "ENOENT") return; - throw error; - } - try { - const descriptorPath = `/proc/self/fd/${parent}/${basename}`; - try { - fs.unlinkSync(descriptorPath); - } catch (error) { - if (error.code !== "ENOENT") throw error; - } - try { - fs.lstatSync(descriptorPath); - } catch (error) { - if (error.code === "ENOENT") return; - throw error; - } - throw new Error(`credential remained after deletion: ${target}`); - } finally { - fs.closeSync(parent); - } -} - -for (const target of targets) removeCredentialNoFollow(target); -""".strip() +_PROXY_AUTH_CLEANUP_JS = ( + resource_files("benchflow.agents") + .joinpath("resources", "proxy_auth_cleanup.js") + .read_text(encoding="utf-8") + .strip() +) def _owner_from_home(cred_home: str) -> str | None: @@ -276,7 +226,8 @@ async def isolate_agent_for_proxy_capture( Every API-proxied agent first requires a non-root sandbox user with no live processes; an existing agent or task process is preserved and makes capture audit-only. Subscription-capable agents then use their already-required - JavaScript runtime to traverse with no-follow directory descriptors. + JavaScript runtime to traverse with no-follow directory descriptors, remove + native login files, and sanitize registry-declared credential settings. Root-agent runs remain audit-only because stale root processes cannot be safely distinguished. Return false unless process isolation is proven and every eligible credential can be identified, removed, and verified absent. @@ -347,12 +298,34 @@ async def isolate_agent_for_proxy_capture( paths.append(str(override_dir / PurePosixPath(primary_path).name)) paths = list(dict.fromkeys(paths)) + settings_targets: list[dict[str, object]] = [] + if subscription_auth.proxy_settings_file: + settings_file = PurePosixPath(subscription_auth.proxy_settings_file) + if ( + not subscription_auth.proxy_settings_drop_keys + or settings_file.is_absolute() + or len(settings_file.parts) != 1 + or settings_file.name != subscription_auth.proxy_settings_file + ): + logger.warning("Cannot prove proxy settings isolation for %s", agent) + return False + settings_targets = [ + { + "path": str(PurePosixPath(credential_path).parent / settings_file), + "drop_keys": list(subscription_auth.proxy_settings_drop_keys), + } + for credential_path in paths + ] + settings_targets = list( + {str(target["path"]): target for target in settings_targets}.values() + ) cleanup_command = f"{process_guard} && " + " ".join( ( f"{_BENCHFLOW_NODE_BIN} -e", shlex.quote(_PROXY_AUTH_CLEANUP_JS), "--", shlex.quote(json.dumps(paths, separators=(",", ":"))), + shlex.quote(json.dumps(settings_targets, separators=(",", ":"))), ) ) try: diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index 4a19924ba..7216ebd0f 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -441,6 +441,11 @@ class SubscriptionAuth: # Optional CLI config-home override. Proxy mode removes this variable so a # reused image cannot select an alternate native-login file outside the # registry-owned container path. + proxy_settings_file: str = "" + proxy_settings_drop_keys: tuple[str, ...] = () + # Optional settings file stored beside the primary credential. Proxy mode + # atomically removes these top-level keys at every effective config root so + # a reused image cannot inject a direct-provider token, endpoint, or helper. @dataclass @@ -547,6 +552,13 @@ class AgentConfig: "~/.claude/.credentials.json", "{home}/.claude/.credentials.json" ), ], + proxy_settings_file="settings.json", + proxy_settings_drop_keys=( + "env", + "apiKeyHelper", + "awsAuthRefresh", + "awsCredentialExport", + ), ), disallow_web_tools_setup_cmd=_json_settings_merge( "$BENCHFLOW_AGENT_HOME/.claude/settings.json", diff --git a/src/benchflow/agents/resources/proxy_auth_cleanup.js b/src/benchflow/agents/resources/proxy_auth_cleanup.js new file mode 100644 index 000000000..9dd61d356 --- /dev/null +++ b/src/benchflow/agents/resources/proxy_auth_cleanup.js @@ -0,0 +1,193 @@ +"use strict"; + +const crypto = require("crypto"); +const fs = require("fs"); + +const credentialTargets = JSON.parse(process.argv[1]); +const settingsTargets = JSON.parse(process.argv[2]); +const constants = fs.constants; +const directoryFlags = + constants.O_RDONLY | constants.O_DIRECTORY | constants.O_NOFOLLOW; + +function openDirectoryNoFollow(parts) { + let descriptor = fs.openSync("/", directoryFlags); + try { + for (const part of parts) { + const next = fs.openSync( + `/proc/self/fd/${descriptor}/${part}`, + directoryFlags, + ); + fs.closeSync(descriptor); + descriptor = next; + } + return descriptor; + } catch (error) { + fs.closeSync(descriptor); + throw error; + } +} + +function splitTarget(target) { + if (typeof target !== "string" || !target.startsWith("/")) { + throw new Error("cleanup target must be an absolute path"); + } + const parts = target.split("/").filter(Boolean); + const basename = parts.pop(); + if (!basename || parts.includes("..")) { + throw new Error(`invalid cleanup target: ${target}`); + } + return {parts, basename}; +} + +function openParent(target) { + const {parts, basename} = splitTarget(target); + try { + return {descriptor: openDirectoryNoFollow(parts), basename}; + } catch (error) { + if (error.code === "ENOENT") return null; + throw error; + } +} + +function removeCredentialNoFollow(target) { + const parent = openParent(target); + if (parent === null) return; + const {descriptor, basename} = parent; + try { + const descriptorPath = `/proc/self/fd/${descriptor}/${basename}`; + try { + fs.unlinkSync(descriptorPath); + } catch (error) { + if (error.code !== "ENOENT") throw error; + } + try { + fs.lstatSync(descriptorPath); + } catch (error) { + if (error.code === "ENOENT") return; + throw error; + } + throw new Error(`credential remained after deletion: ${target}`); + } finally { + fs.closeSync(descriptor); + } +} + +function parseSettings(raw, target) { + let parsed; + try { + parsed = JSON.parse(raw); + } catch (error) { + throw new Error(`invalid JSON settings file ${target}: ${error.message}`); + } + if (parsed === null || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error(`settings file must contain a JSON object: ${target}`); + } + return parsed; +} + +function sanitizeSettingsNoFollow(spec) { + if ( + spec === null || + typeof spec !== "object" || + typeof spec.path !== "string" || + !Array.isArray(spec.drop_keys) || + spec.drop_keys.some((key) => typeof key !== "string" || !key) + ) { + throw new Error("invalid settings sanitization specification"); + } + + const parent = openParent(spec.path); + if (parent === null) return; + const {descriptor, basename} = parent; + const descriptorPath = `/proc/self/fd/${descriptor}/${basename}`; + let sourceDescriptor; + let tempPath; + try { + try { + sourceDescriptor = fs.openSync( + descriptorPath, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + } catch (error) { + if (error.code === "ENOENT") return; + throw error; + } + + const sourceStat = fs.fstatSync(sourceDescriptor); + if (!sourceStat.isFile()) { + throw new Error(`settings target is not a regular file: ${spec.path}`); + } + const settings = parseSettings( + fs.readFileSync(sourceDescriptor, "utf8"), + spec.path, + ); + fs.closeSync(sourceDescriptor); + sourceDescriptor = undefined; + + let changed = false; + for (const key of spec.drop_keys) { + if (Object.prototype.hasOwnProperty.call(settings, key)) { + delete settings[key]; + changed = true; + } + } + if (!changed) return; + + const tempName = `.benchflow-sanitized-${process.pid}-${crypto + .randomBytes(8) + .toString("hex")}`; + tempPath = `/proc/self/fd/${descriptor}/${tempName}`; + const tempDescriptor = fs.openSync( + tempPath, + constants.O_WRONLY | + constants.O_CREAT | + constants.O_EXCL | + constants.O_NOFOLLOW, + 0o600, + ); + try { + fs.fchownSync(tempDescriptor, sourceStat.uid, sourceStat.gid); + fs.fchmodSync(tempDescriptor, sourceStat.mode & 0o777); + fs.writeFileSync(tempDescriptor, `${JSON.stringify(settings, null, 2)}\n`); + fs.fsyncSync(tempDescriptor); + } finally { + fs.closeSync(tempDescriptor); + } + fs.renameSync(tempPath, descriptorPath); + tempPath = undefined; + + const verifiedDescriptor = fs.openSync( + descriptorPath, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + try { + const verified = parseSettings( + fs.readFileSync(verifiedDescriptor, "utf8"), + spec.path, + ); + for (const key of spec.drop_keys) { + if (Object.prototype.hasOwnProperty.call(verified, key)) { + throw new Error(`credential setting remained after sanitization: ${key}`); + } + } + } finally { + fs.closeSync(verifiedDescriptor); + } + } finally { + if (sourceDescriptor !== undefined) fs.closeSync(sourceDescriptor); + if (tempPath !== undefined) { + try { + fs.unlinkSync(tempPath); + } catch (error) { + if (error.code !== "ENOENT") throw error; + } + } + fs.closeSync(descriptor); + } +} + +if (!Array.isArray(credentialTargets) || !Array.isArray(settingsTargets)) { + throw new Error("cleanup arguments must be arrays"); +} +for (const target of credentialTargets) removeCredentialNoFollow(target); +for (const spec of settingsTargets) sanitizeSettingsNoFollow(spec); diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 7c4fcf47b..2d9d61558 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -100,6 +100,11 @@ async def fake_start(**_kwargs): assert ( f"custom-{agent}/{default_auth_path.rsplit('/', 1)[-1]}" in sandbox.commands[0] ) + if agent == "claude-agent-acp": + assert "/home/agent/.claude/settings.json" in sandbox.commands[0] + assert f"custom-{agent}/settings.json" in sandbox.commands[0] + for key in ("env", "apiKeyHelper", "awsAuthRefresh", "awsCredentialExport"): + assert key in sandbox.commands[0] assert "O_NOFOLLOW" in sandbox.commands[0] assert "/proc/self/fd/" in sandbox.commands[0] assert "pgrep -u" in sandbox.commands[0] @@ -249,6 +254,7 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: _PROXY_AUTH_CLEANUP_JS, "--", json.dumps([str(credential_path)]), + json.dumps([]), ], check=False, capture_output=True, @@ -268,6 +274,7 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: _PROXY_AUTH_CLEANUP_JS, "--", json.dumps([str(credential_path)]), + json.dumps([]), ], check=False, capture_output=True, @@ -280,6 +287,117 @@ def test_proxy_auth_cleanup_never_follows_credential_symlinks(tmp_path) -> None: assert outside_auth.read_text() == "subscription-secret" +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("node") is None, + reason="the sandbox settings sanitizer requires Linux procfs and Node.js", +) +def test_proxy_auth_cleanup_sanitizes_claude_credential_settings(tmp_path) -> None: + """Guards PR #1057 review r3888455412 against settings-based auth bypass.""" + + claude_dir = tmp_path / "home" / "agent" / ".claude" + claude_dir.mkdir(parents=True) + credential_path = claude_dir / ".credentials.json" + credential_path.write_text("subscription-secret") + settings_path = claude_dir / "settings.json" + settings_path.write_text( + json.dumps( + { + "env": { + "ANTHROPIC_AUTH_TOKEN": "literal-provider-secret", + "ANTHROPIC_BASE_URL": "https://api.anthropic.example", + }, + "apiKeyHelper": "/opt/provider-key-helper", + "awsAuthRefresh": "/opt/aws-login", + "awsCredentialExport": "/opt/aws-export", + "permissions": {"deny": ["WebSearch", "WebFetch"]}, + "theme": "dark", + } + ) + ) + settings_path.chmod(0o640) + before = settings_path.stat() + + sanitized = subprocess.run( + [ + "node", + "-e", + _PROXY_AUTH_CLEANUP_JS, + "--", + json.dumps([str(credential_path)]), + json.dumps( + [ + { + "path": str(settings_path), + "drop_keys": [ + "env", + "apiKeyHelper", + "awsAuthRefresh", + "awsCredentialExport", + ], + } + ] + ), + ], + check=False, + capture_output=True, + text=True, + ) + + assert sanitized.returncode == 0, sanitized.stderr + assert not credential_path.exists() + assert json.loads(settings_path.read_text()) == { + "permissions": {"deny": ["WebSearch", "WebFetch"]}, + "theme": "dark", + } + after = settings_path.stat() + assert after.st_mode & 0o777 == before.st_mode & 0o777 + assert (after.st_uid, after.st_gid) == (before.st_uid, before.st_gid) + + +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("node") is None, + reason="the sandbox settings sanitizer requires Linux procfs and Node.js", +) +@pytest.mark.parametrize("unsafe_kind", ["symlink", "malformed"]) +def test_proxy_auth_cleanup_fails_closed_for_unsafe_claude_settings( + tmp_path, unsafe_kind: str +) -> None: + """Guards PR #1057 review r3888455412 against unsafe settings rewrites.""" + + claude_dir = tmp_path / "home" / "agent" / ".claude" + claude_dir.mkdir(parents=True) + outside = tmp_path / "outside-settings.json" + outside.write_text('{"env":{"ANTHROPIC_AUTH_TOKEN":"outside-secret"}}') + settings_path = claude_dir / "settings.json" + if unsafe_kind == "symlink": + settings_path.symlink_to(outside) + else: + settings_path.write_text("{not-json") + + refused = subprocess.run( + [ + "node", + "-e", + _PROXY_AUTH_CLEANUP_JS, + "--", + json.dumps([]), + json.dumps([{"path": str(settings_path), "drop_keys": ["env"]}]), + ], + check=False, + capture_output=True, + text=True, + ) + + assert refused.returncode != 0 + if unsafe_kind == "symlink": + assert settings_path.is_symlink() + assert json.loads(outside.read_text())["env"]["ANTHROPIC_AUTH_TOKEN"] == ( + "outside-secret" + ) + else: + assert settings_path.read_text() == "{not-json" + + @pytest.mark.skipif( sys.platform != "linux" or shutil.which("node") is None diff --git a/tests/test_registry_invariants.py b/tests/test_registry_invariants.py index a6dc7d03f..34e4854a1 100644 --- a/tests/test_registry_invariants.py +++ b/tests/test_registry_invariants.py @@ -314,6 +314,16 @@ def test_agent_credential_and_subscription_auth(name, cfg): sa = cfg.subscription_auth assert sa.replaces_env, "SubscriptionAuth.replaces_env must be set" assert sa.detect_file, "SubscriptionAuth.detect_file must be set" + assert bool(sa.proxy_settings_file) is bool(sa.proxy_settings_drop_keys), ( + "Proxy settings filename and drop keys must be declared together" + ) + if sa.proxy_settings_file: + assert "/" not in sa.proxy_settings_file + assert sa.proxy_settings_file not in {".", ".."} + assert len(set(sa.proxy_settings_drop_keys)) == len( + sa.proxy_settings_drop_keys + ) + assert all(sa.proxy_settings_drop_keys) for f in sa.files: assert f.host_path, "HostAuthFile.host_path must be set" assert f.container_path, "HostAuthFile.container_path must be set" From b2ca3f3c516a2baeb15c4ce1a5940b3d17510d57 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 21:53:28 -0700 Subject: [PATCH 62/74] fix(auth): replace stale Pi providers in proxy mode --- src/benchflow/agents/pi_acp_launcher.py | 8 +++- tests/test_pi_acp_launcher.py | 51 +++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/benchflow/agents/pi_acp_launcher.py b/src/benchflow/agents/pi_acp_launcher.py index 93e5a347c..9e7f0e4a6 100644 --- a/src/benchflow/agents/pi_acp_launcher.py +++ b/src/benchflow/agents/pi_acp_launcher.py @@ -79,6 +79,7 @@ def setup_provider() -> None: provider_name = os.environ.get("BENCHFLOW_PROVIDER_NAME") or _derive_provider_name( model ) + proxy_mode = bool(os.environ.get("BENCHFLOW_LITELLM_MODEL_ALIAS")) if protocol == "openai-completions": if not base_url: @@ -117,8 +118,11 @@ def setup_provider() -> None: config_dir = Path.home() / ".pi" / "agent" config_dir.mkdir(parents=True, exist_ok=True) models_path = config_dir / "models.json" - # Merge with existing config so manually-added providers survive - if models_path.exists(): + # A proxy run has exclusive custody of model traffic. Replacing the + # provider map removes stale direct providers (and their literal API + # keys) that could otherwise let Pi bypass LiteLLM. Direct-provider + # runs retain the historical merge behavior for user-added providers. + if not proxy_mode and models_path.exists(): try: existing = json.loads(models_path.read_text()) existing.setdefault("providers", {}).update(config["providers"]) diff --git a/tests/test_pi_acp_launcher.py b/tests/test_pi_acp_launcher.py index 004031eb4..7af3fc2ec 100644 --- a/tests/test_pi_acp_launcher.py +++ b/tests/test_pi_acp_launcher.py @@ -32,6 +32,7 @@ def _pi_env(monkeypatch, tmp_path): "BENCHFLOW_PROVIDER_MODEL", "BENCHFLOW_PROVIDER_MODELS", "BENCHFLOW_PROVIDER_NAME", + "BENCHFLOW_LITELLM_MODEL_ALIAS", "ANTHROPIC_BASE_URL", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_MODEL", @@ -144,6 +145,56 @@ def test_merges_with_existing_providers(self, monkeypatch, tmp_path): assert "other" in config["providers"], "pre-existing provider must survive" assert "vllm" in config["providers"], "new provider must be added" + def test_proxy_mode_replaces_existing_providers(self, monkeypatch, tmp_path): + """Guards PR #1057 review r3888489750 against a Pi proxy escape. + + Proxy mode must discard every pre-existing provider and top-level field, + including literal API keys, so Pi can only route through LiteLLM. + """ + config_dir = tmp_path / ".pi" / "agent" + config_dir.mkdir(parents=True) + stale_key = "stale-direct-provider-key" + existing = { + "providers": { + "direct-anthropic": { + "baseUrl": "https://api.anthropic.com", + "api": "anthropic-messages", + "apiKey": stale_key, + "models": [{"id": "claude", "name": "claude"}], + }, + "direct-openai": { + "baseUrl": "https://api.openai.com/v1", + "api": "openai-completions", + "apiKey": "another-stale-key", + "models": [{"id": "gpt", "name": "gpt"}], + }, + }, + "unrelated": {"credential": "must-also-be-removed"}, + } + models_path = config_dir / "models.json" + models_path.write_text(json.dumps(existing)) + + monkeypatch.setenv("BENCHFLOW_PROVIDER_PROTOCOL", "openai-completions") + monkeypatch.setenv("BENCHFLOW_PROVIDER_BASE_URL", "http://127.0.0.1:4000/v1") + monkeypatch.setenv("BENCHFLOW_PROVIDER_API_KEY", "proxy-master-key") + monkeypatch.setenv("BENCHFLOW_PROVIDER_MODEL", "benchflow-model-alias") + monkeypatch.setenv("BENCHFLOW_PROVIDER_NAME", "litellm") + monkeypatch.setenv("BENCHFLOW_LITELLM_MODEL_ALIAS", "benchflow-model-alias") + + from benchflow.agents.pi_acp_launcher import setup_provider + + setup_provider() + + serialized = models_path.read_text() + config = json.loads(serialized) + assert set(config) == {"providers"} + assert set(config["providers"]) == {"litellm"} + assert config["providers"]["litellm"]["baseUrl"] == ("http://127.0.0.1:4000/v1") + assert config["providers"]["litellm"]["apiKey"] == "proxy-master-key" + assert stale_key not in serialized + assert "another-stale-key" not in serialized + assert "must-also-be-removed" not in serialized + def test_overwrites_corrupt_models_json(self, monkeypatch, tmp_path, capsys): config_dir = tmp_path / ".pi" / "agent" config_dir.mkdir(parents=True) From aa078e09c679f22c162c7fe19687eec470e4d5d3 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 21:58:30 -0700 Subject: [PATCH 63/74] fix(auth): scrub stale Pi config before proxy launch --- src/benchflow/agents/credentials.py | 57 +++++++++++++++++++----- src/benchflow/agents/pi_acp_launcher.py | 8 +--- tests/test_litellm_credential_custody.py | 57 +++++++++++++++++++++++- tests/test_pi_acp_launcher.py | 51 --------------------- 4 files changed, 105 insertions(+), 68 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index ae07067cf..274c14ca7 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -65,6 +65,24 @@ def _proxy_process_isolation_guard(cred_home: str) -> tuple[str, bool]: ) +def _proxy_auth_cleanup_command( + process_guard: str, + paths: list[str], + settings_targets: list[dict[str, object]], +) -> str: + """Build the no-follow cleanup command gated by process isolation.""" + + return f"{process_guard} && " + " ".join( + ( + f"{_BENCHFLOW_NODE_BIN} -e", + shlex.quote(_PROXY_AUTH_CLEANUP_JS), + "--", + shlex.quote(json.dumps(paths, separators=(",", ":"))), + shlex.quote(json.dumps(settings_targets, separators=(",", ":"))), + ) + ) + + async def upload_credential( env, path: str, @@ -242,16 +260,39 @@ async def isolate_agent_for_proxy_capture( process_guard, processes_safe = _proxy_process_isolation_guard(cred_home) if subscription_auth is None: + if agent == "pi-acp": + home = PurePosixPath(cred_home) + relative_models_path = PurePosixPath(".pi/agent/models.json") + paths = [str(home / relative_models_path)] + paths_safe = True + for home_env in ("HOME", "BENCHFLOW_AGENT_HOME"): + override_value = agent_env.get(home_env, "") + if override_value and override_value != cred_home: + override_home = _safe_proxy_auth_root( + override_value, + cred_home=cred_home, + agent=agent, + ) + if override_home is None: + paths_safe = False + else: + paths.append(str(override_home / relative_models_path)) + agent_env[home_env] = cred_home + paths = list(dict.fromkeys(paths)) + cleanup_command = _proxy_auth_cleanup_command(process_guard, paths, []) + else: + paths_safe = True + cleanup_command = process_guard try: result = await env.exec( - process_guard, + cleanup_command, user="root", timeout_sec=10, ) except Exception as exc: logger.warning("Failed to isolate %s agent processes: %s", agent, exc) return False - return bool(result.return_code == 0 and processes_safe) + return bool(result.return_code == 0 and paths_safe and processes_safe) primary_files = [ auth_file @@ -319,14 +360,10 @@ async def isolate_agent_for_proxy_capture( settings_targets = list( {str(target["path"]): target for target in settings_targets}.values() ) - cleanup_command = f"{process_guard} && " + " ".join( - ( - f"{_BENCHFLOW_NODE_BIN} -e", - shlex.quote(_PROXY_AUTH_CLEANUP_JS), - "--", - shlex.quote(json.dumps(paths, separators=(",", ":"))), - shlex.quote(json.dumps(settings_targets, separators=(",", ":"))), - ) + cleanup_command = _proxy_auth_cleanup_command( + process_guard, + paths, + settings_targets, ) try: result = await env.exec( diff --git a/src/benchflow/agents/pi_acp_launcher.py b/src/benchflow/agents/pi_acp_launcher.py index 9e7f0e4a6..93e5a347c 100644 --- a/src/benchflow/agents/pi_acp_launcher.py +++ b/src/benchflow/agents/pi_acp_launcher.py @@ -79,7 +79,6 @@ def setup_provider() -> None: provider_name = os.environ.get("BENCHFLOW_PROVIDER_NAME") or _derive_provider_name( model ) - proxy_mode = bool(os.environ.get("BENCHFLOW_LITELLM_MODEL_ALIAS")) if protocol == "openai-completions": if not base_url: @@ -118,11 +117,8 @@ def setup_provider() -> None: config_dir = Path.home() / ".pi" / "agent" config_dir.mkdir(parents=True, exist_ok=True) models_path = config_dir / "models.json" - # A proxy run has exclusive custody of model traffic. Replacing the - # provider map removes stale direct providers (and their literal API - # keys) that could otherwise let Pi bypass LiteLLM. Direct-provider - # runs retain the historical merge behavior for user-added providers. - if not proxy_mode and models_path.exists(): + # Merge with existing config so manually-added providers survive + if models_path.exists(): try: existing = json.loads(models_path.read_text()) existing.setdefault("providers", {}).update(config["providers"]) diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index 2d9d61558..c50acb3f1 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -152,7 +152,62 @@ async def test_proxy_process_guard_applies_without_subscription_auth(agent) -> N assert "pgrep -u" in sandbox.commands[0] assert '"$bf_agent_uid" -ne 0' in sandbox.commands[0] assert '"$bf_pgrep_rc" -eq 1' in sandbox.commands[0] - assert _PROXY_AUTH_CLEANUP_JS not in sandbox.commands[0] + if agent == "pi-acp": + assert _PROXY_AUTH_CLEANUP_JS in sandbox.commands[0] + assert "/home/agent/.pi/agent/models.json" in sandbox.commands[0] + else: + assert _PROXY_AUTH_CLEANUP_JS not in sandbox.commands[0] + + +@pytest.mark.asyncio +async def test_pi_proxy_cleanup_removes_every_safe_effective_models_config() -> None: + """Guards PR #1057 review r3888489750 against stale direct Pi routes.""" + + sandbox = _SubscriptionIsolationSandbox() + agent_env = { + "HOME": "/home/agent/custom-home", + "BENCHFLOW_AGENT_HOME": "/home/agent/custom-agent-home", + } + + trusted = await isolate_agent_for_proxy_capture( + sandbox, + agent="pi-acp", + agent_env=agent_env, + cred_home="/home/agent", + ) + + assert trusted is True + assert agent_env["HOME"] == "/home/agent" + assert agent_env["BENCHFLOW_AGENT_HOME"] == "/home/agent" + assert len(sandbox.commands) == 1 + command = sandbox.commands[0] + assert "/home/agent/.pi/agent/models.json" in command + assert "/home/agent/custom-home/.pi/agent/models.json" in command + assert "/home/agent/custom-agent-home/.pi/agent/models.json" in command + assert "O_NOFOLLOW" in command + assert "/proc/self/fd/" in command + + +@pytest.mark.asyncio +async def test_pi_proxy_cleanup_rejects_models_config_outside_sandbox_home() -> None: + """Guards PR #1057 review r3888489750 against unsafe root deletion.""" + + sandbox = _SubscriptionIsolationSandbox() + agent_env = {"HOME": "/etc/pi-home"} + + trusted = await isolate_agent_for_proxy_capture( + sandbox, + agent="pi-acp", + agent_env=agent_env, + cred_home="/home/agent", + ) + + assert trusted is False + assert agent_env["HOME"] == "/home/agent" + assert agent_env["BENCHFLOW_AGENT_HOME"] == "/home/agent" + assert len(sandbox.commands) == 1 + assert "/home/agent/.pi/agent/models.json" in sandbox.commands[0] + assert "/etc/pi-home" not in sandbox.commands[0] @pytest.mark.asyncio diff --git a/tests/test_pi_acp_launcher.py b/tests/test_pi_acp_launcher.py index 7af3fc2ec..004031eb4 100644 --- a/tests/test_pi_acp_launcher.py +++ b/tests/test_pi_acp_launcher.py @@ -32,7 +32,6 @@ def _pi_env(monkeypatch, tmp_path): "BENCHFLOW_PROVIDER_MODEL", "BENCHFLOW_PROVIDER_MODELS", "BENCHFLOW_PROVIDER_NAME", - "BENCHFLOW_LITELLM_MODEL_ALIAS", "ANTHROPIC_BASE_URL", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_MODEL", @@ -145,56 +144,6 @@ def test_merges_with_existing_providers(self, monkeypatch, tmp_path): assert "other" in config["providers"], "pre-existing provider must survive" assert "vllm" in config["providers"], "new provider must be added" - def test_proxy_mode_replaces_existing_providers(self, monkeypatch, tmp_path): - """Guards PR #1057 review r3888489750 against a Pi proxy escape. - - Proxy mode must discard every pre-existing provider and top-level field, - including literal API keys, so Pi can only route through LiteLLM. - """ - config_dir = tmp_path / ".pi" / "agent" - config_dir.mkdir(parents=True) - stale_key = "stale-direct-provider-key" - existing = { - "providers": { - "direct-anthropic": { - "baseUrl": "https://api.anthropic.com", - "api": "anthropic-messages", - "apiKey": stale_key, - "models": [{"id": "claude", "name": "claude"}], - }, - "direct-openai": { - "baseUrl": "https://api.openai.com/v1", - "api": "openai-completions", - "apiKey": "another-stale-key", - "models": [{"id": "gpt", "name": "gpt"}], - }, - }, - "unrelated": {"credential": "must-also-be-removed"}, - } - models_path = config_dir / "models.json" - models_path.write_text(json.dumps(existing)) - - monkeypatch.setenv("BENCHFLOW_PROVIDER_PROTOCOL", "openai-completions") - monkeypatch.setenv("BENCHFLOW_PROVIDER_BASE_URL", "http://127.0.0.1:4000/v1") - monkeypatch.setenv("BENCHFLOW_PROVIDER_API_KEY", "proxy-master-key") - monkeypatch.setenv("BENCHFLOW_PROVIDER_MODEL", "benchflow-model-alias") - monkeypatch.setenv("BENCHFLOW_PROVIDER_NAME", "litellm") - monkeypatch.setenv("BENCHFLOW_LITELLM_MODEL_ALIAS", "benchflow-model-alias") - - from benchflow.agents.pi_acp_launcher import setup_provider - - setup_provider() - - serialized = models_path.read_text() - config = json.loads(serialized) - assert set(config) == {"providers"} - assert set(config["providers"]) == {"litellm"} - assert config["providers"]["litellm"]["baseUrl"] == ("http://127.0.0.1:4000/v1") - assert config["providers"]["litellm"]["apiKey"] == "proxy-master-key" - assert stale_key not in serialized - assert "another-stale-key" not in serialized - assert "must-also-be-removed" not in serialized - def test_overwrites_corrupt_models_json(self, monkeypatch, tmp_path, capsys): config_dir = tmp_path / ".pi" / "agent" config_dir.mkdir(parents=True) From df513055ed4b77c3f7f6973fe825ea43fbb5075b Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 22:28:14 -0700 Subject: [PATCH 64/74] fix(capture): preserve role usage and block privilege regain --- .../providers/litellm_capture_custody.py | 8 +- .../resources/provider_capture_custody.sh | 2 + src/benchflow/rollout/__init__.py | 86 +++++++----- src/benchflow/rollout/_usage.py | 55 +++++++- src/benchflow/sandbox/lockdown.py | 12 +- tests/test_api_error_capture.py | 13 ++ tests/test_litellm_runtime.py | 4 +- tests/test_sdk_lockdown.py | 35 ++++- tests/test_trial_litellm_runtime.py | 12 +- tests/test_usage_litellm.py | 132 ++++++++++++++++++ 10 files changed, 314 insertions(+), 45 deletions(-) diff --git a/src/benchflow/providers/litellm_capture_custody.py b/src/benchflow/providers/litellm_capture_custody.py index bfe2ea987..116bb9622 100644 --- a/src/benchflow/providers/litellm_capture_custody.py +++ b/src/benchflow/providers/litellm_capture_custody.py @@ -10,6 +10,7 @@ from typing import Any from benchflow.sandbox.files import upload_private_text +from benchflow.sandbox.lockdown import build_priv_drop_cmd logger = logging.getLogger(__name__) @@ -82,8 +83,11 @@ async def provider_capture_has_verified_custody( logger.warning("Provider capture custody artifact hardening failed") return False access = await sandbox.exec( - f"{quoted_probe} probe {quoted_runtime}", - user=sandbox_user, + build_priv_drop_cmd( + f"{quoted_probe} probe {quoted_runtime}", + sandbox_user, + ), + user="root", timeout_sec=10, ) if access.return_code != 1: diff --git a/src/benchflow/providers/resources/provider_capture_custody.sh b/src/benchflow/providers/resources/provider_capture_custody.sh index fe48b36cd..0d64643de 100644 --- a/src/benchflow/providers/resources/provider_capture_custody.sh +++ b/src/benchflow/providers/resources/provider_capture_custody.sh @@ -19,6 +19,8 @@ case "$mode" in chmod 600 "$callback_log" "$capture_state" ;; probe) + no_new_privs=$(awk '/^NoNewPrivs:/ {print $2}' /proc/self/status 2>/dev/null || true) + if [ "$no_new_privs" != 1 ]; then exit 0; fi for artifact in "$callback_log" "$capture_state"; do if cat "$artifact" >/dev/null 2>&1; then exit 0; fi if [ -r "$artifact" ] || [ -w "$artifact" ]; then exit 0; fi diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 0866688d6..526b3688e 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -163,9 +163,12 @@ ProviderFailure as ProviderFailure, ) from benchflow.rollout._usage import _as_nonnegative_int as _as_nonnegative_int +from benchflow.rollout._usage import ( + _merge_provider_usage_metrics as _merge_provider_usage_metrics, +) from benchflow.rollout._usage import _native_acp_usage_delta as _native_acp_usage_delta from benchflow.rollout._usage import ( - _provider_api_failure_summary_from_runtime as _provider_api_failure_summary_from_runtime, + _provider_api_failure_summary_from_runtimes as _provider_api_failure_summary_from_runtimes, ) from benchflow.rollout._usage import ( _provider_auth_status_from_runtime as _provider_auth_status_from_runtime, @@ -677,6 +680,7 @@ def __init__(self, config: RolloutConfig) -> None: self._task_tmp: Path | None = None self._task_skill_policy: TaskSkillPolicy | None = None self._usage_runtime: Any = None + self._retired_usage_runtimes: list[Any] = [] self._usage_metrics: dict[str, Any] = self._planes.extract_usage(None) self._native_usage_metrics: dict[str, Any] = _zero_native_acp_usage_metrics() self._native_usage_checkpoint: dict[str, int | None] | None = None @@ -1312,14 +1316,12 @@ async def connect(self) -> None: rollout_dir = self._require_rollout_dir() t0 = datetime.now() - ( - self._agent_env, - self._usage_runtime, - ) = await self._planes.ensure_litellm_runtime( + previous_usage_runtime = getattr(self, "_usage_runtime", None) + self._agent_env, next_usage_runtime = await self._planes.ensure_litellm_runtime( agent=cfg.primary_agent, agent_env=self._agent_env, model=cfg.primary_model, - runtime=getattr(self, "_usage_runtime", None), + runtime=previous_usage_runtime, environment=cfg.environment, session_id=getattr(self, "_rollout_name", "") or "", usage_tracking=cfg.usage_tracking, @@ -1330,6 +1332,7 @@ async def connect(self) -> None: force_sandbox_local=getattr(self, "_disallow_web_tools", False), sandbox_user=cfg.sandbox_user, ) + self._adopt_usage_runtime(previous_usage_runtime, next_usage_runtime) llm_capture = getattr(self, "_llm_capture", None) if llm_capture is not None: credential_home = _sandbox_user_home(cfg.sandbox_user) @@ -2036,6 +2039,10 @@ async def cleanup(self) -> None: self._evolved_skills = None usage_runtime = getattr(self, "_usage_runtime", None) + retired_usage_runtimes = list(getattr(self, "_retired_usage_runtimes", [])) + all_usage_runtimes = [*retired_usage_runtimes] + if usage_runtime is not None: + all_usage_runtimes.append(usage_runtime) provider_capture_errors: list[str] = [] if usage_runtime is not None: try: @@ -2047,30 +2054,6 @@ async def cleanup(self) -> None: "provider runtime stop or remote capture import failed" ) self._usage_metrics = self._planes.extract_usage(None) - # Snapshot any provider failure (401/403/429/503) now that captures - # are imported (stop() populated the trajectory). This must happen - # before we drop the runtime reference below, and is read later by - # ACP-error classification — for Daytona the trajectory is empty - # until here (#546/#564). - # - # Coverage gap: only `self._usage_runtime` is scanned here. Bedrock - # auth failures flow through `self._provider_runtime`, whose server - # (BedrockProxyServer) exposes no `.trajectory`/`.exchanges`, so a - # fallback scan of it would always return None — useless, so it's - # not implemented. The direct-AWS-Bedrock case (remote sandbox, - # runtime=None) bypasses both proxies entirely and is out of scope. - self._provider_failure_cached = _provider_failure_from_runtime( - usage_runtime - ) - self._provider_auth_status_cached = ( - self._provider_failure_cached.status - if self._provider_failure_cached is not None - and self._provider_failure_cached.marker == "provider auth failed" - else None - ) - self._api_failure_summary_cached = ( - _provider_api_failure_summary_from_runtime(usage_runtime) - ) try: self._write_llm_trajectory(usage_runtime) except Exception as e: @@ -2085,6 +2068,32 @@ async def cleanup(self) -> None: finally: self._usage_runtime = None + # Role-specific gateways are stopped during connect_as() rotation. Keep + # their token/cost and failure evidence alongside the final gateway; + # the live trajectory writer already preserves their exchange rows. + prior_provider_metrics = [ + self._planes.extract_usage(completed_runtime) + for completed_runtime in retired_usage_runtimes + ] + self._usage_metrics = _merge_provider_usage_metrics( + [*prior_provider_metrics, getattr(self, "_usage_metrics", {})] + ) + self._retired_usage_runtimes = [] + self._provider_failure_cached = None + for completed_runtime in all_usage_runtimes: + failure = _provider_failure_from_runtime(completed_runtime) + if failure is not None: + self._provider_failure_cached = failure + self._provider_auth_status_cached = ( + self._provider_failure_cached.status + if self._provider_failure_cached is not None + and self._provider_failure_cached.marker == "provider auth failed" + else None + ) + self._api_failure_summary_cached = _provider_api_failure_summary_from_runtimes( + all_usage_runtimes + ) + self._finalize_usage_metrics() llm_capture = getattr(self, "_llm_capture", None) if llm_capture is not None: @@ -2130,6 +2139,17 @@ async def cleanup(self) -> None: self._phase = "cleaned" + def _adopt_usage_runtime(self, previous: Any, current: Any) -> None: + """Retain stopped role runtimes for final usage/failure aggregation.""" + + if previous is not None and previous is not current: + retired = getattr(self, "_retired_usage_runtimes", None) + if retired is None: + retired = [] + self._retired_usage_runtimes = retired + retired.append(previous) + self._usage_runtime = current + def _finalize_usage_metrics(self) -> None: """Prefer LiteLLM usage, otherwise use trusted native ACP usage.""" current_metrics = getattr( @@ -2379,11 +2399,12 @@ async def connect_as(self, role: Role) -> None: ), disallow=disallow_web_tools, ) - agent_env, self._usage_runtime = await self._planes.ensure_litellm_runtime( + previous_usage_runtime = getattr(self, "_usage_runtime", None) + agent_env, next_usage_runtime = await self._planes.ensure_litellm_runtime( agent=role.agent, agent_env=agent_env, model=role.model, - runtime=getattr(self, "_usage_runtime", None), + runtime=previous_usage_runtime, environment=cfg.environment, session_id=getattr(self, "_rollout_name", "") or "", usage_tracking=cfg.usage_tracking, @@ -2395,6 +2416,7 @@ async def connect_as(self, role: Role) -> None: role_name=role.name, sandbox_user=cfg.sandbox_user, ) + self._adopt_usage_runtime(previous_usage_runtime, next_usage_runtime) role_agent_differs = role.agent != cfg.primary_agent needs_role_credentials = ( diff --git a/src/benchflow/rollout/_usage.py b/src/benchflow/rollout/_usage.py index 8434e3176..dfb6d6acc 100644 --- a/src/benchflow/rollout/_usage.py +++ b/src/benchflow/rollout/_usage.py @@ -103,9 +103,19 @@ def _provider_api_failure_summary_from_runtime(runtime: Any) -> dict[str, Any] | counts, and the dominant failure's (subcategory, transient, fingerprint) classification. Reads only integer status codes (#546/#564). """ - server = getattr(runtime, "server", None) - trajectory = getattr(server, "trajectory", None) - exchanges = getattr(trajectory, "exchanges", None) or [] + return _provider_api_failure_summary_from_runtimes([runtime]) + + +def _provider_api_failure_summary_from_runtimes( + runtimes: list[Any], +) -> dict[str, Any] | None: + """Summarize HTTP failures across role-scoped provider runtimes.""" + + exchanges = [] + for runtime in runtimes: + server = getattr(runtime, "server", None) + trajectory = getattr(server, "trajectory", None) + exchanges.extend(getattr(trajectory, "exchanges", None) or []) total = 0 failed: dict[int, int] = {} last_failed_status: int | None = None @@ -138,6 +148,45 @@ def _provider_api_failure_summary_from_runtime(runtime: Any) -> dict[str, Any] | return summary +_PROVIDER_USAGE_COUNT_FIELDS = ( + "n_input_tokens", + "n_output_tokens", + "n_cache_read_tokens", + "n_cache_creation_tokens", + "total_tokens", +) + + +def _merge_provider_usage_metrics( + metrics: list[dict[str, Any]], +) -> dict[str, Any]: + """Sum trusted provider metrics from successively rotated gateways. + + Cost remains unknown if any contributing provider response was unpriced; + summing only the priced subset would silently under-report the rollout. + """ + + available = [ + item for item in metrics if item.get("usage_source") == "provider_response" + ] + if not available: + return usage_unavailable() + merged = usage_unavailable() + for field in _PROVIDER_USAGE_COUNT_FIELDS: + merged[field] = sum(_as_nonnegative_int(item.get(field)) for item in available) + costs = [item.get("cost_usd") for item in available] + numeric_costs = [float(cost) for cost in costs if isinstance(cost, int | float)] + if len(numeric_costs) == len(costs): + merged["cost_usd"] = round(sum(numeric_costs), 10) + merged["price_source"] = ( + "litellm" + if all(item.get("price_source") == "litellm" for item in available) + else None + ) + merged["usage_source"] = "provider_response" + return merged + + def classify_api_failure( summary: dict[str, Any] | None, *, diff --git a/src/benchflow/sandbox/lockdown.py b/src/benchflow/sandbox/lockdown.py index 9fbfde900..5fd6bfa2f 100644 --- a/src/benchflow/sandbox/lockdown.py +++ b/src/benchflow/sandbox/lockdown.py @@ -129,14 +129,18 @@ def _agent_egress_firewall_cmd(sandbox_user: str) -> str: def build_priv_drop_cmd(agent_launch: str, sandbox_user: str) -> str: """Build a shell command that drops to sandbox_user via setpriv or su. - setpriv (util-linux) execs directly; su -l is the fallback for Alpine/BusyBox. - No outer sh -c wrapper — DockerProcess wraps in bash -c already. + A capable util-linux ``setpriv`` also sets ``no_new_privs`` so setuid/file- + capability helpers cannot regain privilege after launch. ``su -l`` remains + the compatibility fallback; capture custody independently treats that path + as audit-only because it cannot prove the same kernel boundary. No outer + sh -c wrapper — DockerProcess wraps in bash -c already. """ inner = f"export HOME=/home/{sandbox_user} && {agent_launch}" quoted = shlex.quote(inner) return ( - f"if setpriv --help 2>&1 | grep -q reuid; then" - f" exec setpriv --reuid={sandbox_user} --regid={sandbox_user}" + f"if setpriv --help 2>&1 | grep -q reuid &&" + f" setpriv --help 2>&1 | grep -q no-new-privs; then" + f" exec setpriv --no-new-privs --reuid={sandbox_user} --regid={sandbox_user}" f" --init-groups -- bash -c {quoted};" f" else exec su -l {sandbox_user} -c {quoted};" f" fi" diff --git a/tests/test_api_error_capture.py b/tests/test_api_error_capture.py index 59cbfa8b8..3b37adf32 100644 --- a/tests/test_api_error_capture.py +++ b/tests/test_api_error_capture.py @@ -27,6 +27,7 @@ from benchflow.rollout._usage import ( _api_error_subcategory, _provider_api_failure_summary_from_runtime, + _provider_api_failure_summary_from_runtimes, classify_api_failure, ) @@ -116,6 +117,18 @@ def test_non_int_statuses_skipped(self): assert s["failed_requests"] == 1 assert s["subcategory"] == "auth" + def test_role_runtime_summaries_are_combined(self): + """Guards PR #1057 review r3888535158 beyond token accounting.""" + + summary = _provider_api_failure_summary_from_runtimes( + [_runtime([429]), _runtime([200, 500])] + ) + + assert summary is not None + assert summary["total_requests"] == 3 + assert summary["failed_requests"] == 2 + assert summary["status_counts"] == {"429": 1, "500": 1} + class TestClassifyApiFailure: def test_proxy_proven(self): diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index f340a2301..e317fa1c3 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -194,7 +194,9 @@ async def exec(self, command, **kwargs): assert kwargs["user"] == "root" return SimpleNamespace(return_code=0, stdout="", stderr="") if f"{self.probe_path} probe " in command: - assert kwargs["user"] == "agent" + assert kwargs["user"] == "root" + assert "setpriv --no-new-privs" in command + assert "--reuid=agent" in command return SimpleNamespace( return_code=self.agent_probe_return_code, stdout="", diff --git a/tests/test_sdk_lockdown.py b/tests/test_sdk_lockdown.py index ace503fb6..6cff79175 100644 --- a/tests/test_sdk_lockdown.py +++ b/tests/test_sdk_lockdown.py @@ -1,6 +1,9 @@ """Tests for path lockdown — _validate_locked_path, _resolve_locked_paths, lockdown_paths.""" +import os +import shutil import subprocess +import sys from unittest.mock import AsyncMock, MagicMock import pytest @@ -365,9 +368,39 @@ class TestPrivDropCommand: def test_contains_setpriv_and_su_fallback(self): cmd = build_priv_drop_cmd("my-agent --stdio", "agent") - assert "setpriv --reuid=agent --regid=agent --init-groups" in cmd + assert "setpriv --no-new-privs --reuid=agent --regid=agent --init-groups" in cmd assert "su -l agent -c" in cmd + def test_privilege_regain_is_blocked(self): + """Guards PR #1057 review r3888535162 against setuid regain.""" + + cmd = build_priv_drop_cmd("my-agent", "agent") + assert "grep -q no-new-privs" in cmd + assert "setpriv --no-new-privs" in cmd + + @pytest.mark.skipif( + sys.platform != "linux" + or getattr(os, "geteuid", lambda: -1)() != 0 + or shutil.which("setpriv") is None, + reason="real no_new_privs launch proof requires root and Linux setpriv", + ) + def test_linux_launch_really_sets_no_new_privs(self): + """Guards PR #1057 review r3888535162 at the kernel boundary.""" + + cmd = build_priv_drop_cmd( + "awk '/^NoNewPrivs:/ {print $2}' /proc/self/status", + "nobody", + ) + + result = subprocess.run( + ["bash", "-c", cmd], + check=True, + capture_output=True, + text=True, + ) + + assert result.stdout.strip() == "1" + def test_exec_prefix(self): """Both branches use exec to replace the shell (no lingering parent).""" cmd = build_priv_drop_cmd("my-agent", "agent") diff --git a/tests/test_trial_litellm_runtime.py b/tests/test_trial_litellm_runtime.py index 183c0ba30..149c33c58 100644 --- a/tests/test_trial_litellm_runtime.py +++ b/tests/test_trial_litellm_runtime.py @@ -61,6 +61,8 @@ async def fake_connect_acp(**kwargs): @pytest.mark.asyncio async def test_trial_connect_as_starts_litellm_for_role(tmp_path: Path): + """Guards PR #1057 review r3888535158 across role runtime rotation.""" + rollout = Rollout.__new__(Rollout) rollout._config = RolloutConfig( task_path=tmp_path / "task", @@ -73,7 +75,10 @@ async def test_trial_connect_as_starts_litellm_for_role(tmp_path: Path): rollout._rollout_name = "rollout" rollout._agent_cwd = "/workspace" rollout._env = SimpleNamespace() - rollout._usage_runtime = None + previous_runtime = SimpleNamespace(kind="litellm", name="primary") + next_runtime = SimpleNamespace(kind="litellm", name="reviewer") + rollout._usage_runtime = previous_runtime + rollout._retired_usage_runtimes = [] rollout._timing = {} rollout._disallow_web_tools = False rollout._agent_cfg = SimpleNamespace() @@ -90,9 +95,10 @@ async def test_trial_connect_as_starts_litellm_for_role(tmp_path: Path): async def fake_litellm(**kwargs): calls.append("litellm") assert kwargs["agent"] == "claude-agent-acp" + assert kwargs["runtime"] is previous_runtime env = dict(kwargs["agent_env"]) env["ANTHROPIC_BASE_URL"] = "http://host.docker.internal:4000" - return env, SimpleNamespace(kind="litellm") + return env, next_runtime async def fake_connect_acp(**kwargs): calls.append("acp") @@ -115,3 +121,5 @@ async def fake_connect_acp(**kwargs): await rollout.connect_as(role) assert calls == ["litellm", "acp"] + assert rollout._usage_runtime is next_runtime + assert rollout._retired_usage_runtimes == [previous_runtime] diff --git a/tests/test_usage_litellm.py b/tests/test_usage_litellm.py index 5a726cf29..3501f9fd9 100644 --- a/tests/test_usage_litellm.py +++ b/tests/test_usage_litellm.py @@ -273,6 +273,77 @@ def test_extract_usage_reads_litellm_runtime_trajectory(): assert usage["n_output_tokens"] == 2 +def test_merge_provider_usage_metrics_across_role_runtimes(): + """Guards PR #1057 review r3888535158 against role-rotation undercount.""" + + from benchflow.rollout._usage import _merge_provider_usage_metrics + + merged = _merge_provider_usage_metrics( + [ + { + "n_input_tokens": 10, + "n_output_tokens": 2, + "n_cache_read_tokens": 1, + "n_cache_creation_tokens": 0, + "total_tokens": 13, + "cost_usd": 0.01, + "usage_source": "provider_response", + "price_source": "litellm", + }, + { + "n_input_tokens": 20, + "n_output_tokens": 3, + "n_cache_read_tokens": 0, + "n_cache_creation_tokens": 2, + "total_tokens": 25, + "cost_usd": 0.02, + "usage_source": "provider_response", + "price_source": "litellm", + }, + ] + ) + + assert merged == { + "n_input_tokens": 30, + "n_output_tokens": 5, + "n_cache_read_tokens": 1, + "n_cache_creation_tokens": 2, + "total_tokens": 38, + "cost_usd": 0.03, + "usage_source": "provider_response", + "price_source": "litellm", + } + + +def test_merge_provider_usage_keeps_mixed_pricing_unknown(): + """Unpriced role traffic must not make aggregate cost look complete.""" + + from benchflow.rollout._usage import _merge_provider_usage_metrics + + priced = { + "n_input_tokens": 10, + "n_output_tokens": 2, + "total_tokens": 12, + "cost_usd": 0.01, + "usage_source": "provider_response", + "price_source": "litellm", + } + unpriced = { + "n_input_tokens": 4, + "n_output_tokens": 1, + "total_tokens": 5, + "cost_usd": None, + "usage_source": "provider_response", + "price_source": None, + } + + merged = _merge_provider_usage_metrics([priced, unpriced]) + + assert merged["total_tokens"] == 17 + assert merged["cost_usd"] is None + assert merged["price_source"] is None + + def test_rollout_final_reconcile_preserves_prior_provider_runtime(tmp_path): """Guards PR #1057's cumulative writer across runtime-switch cleanup.""" @@ -357,6 +428,67 @@ def reconcile_live_capture(self): assert (tmp_path / "trajectory" / "llm_trajectory.jsonl").exists() +@pytest.mark.asyncio +async def test_rollout_cleanup_accumulates_rotated_provider_usage(tmp_path): + """Guards PR #1057 review r3888535158 against final-role-only metrics.""" + + from benchflow.rollout import Rollout, RolloutConfig + + class FakeServer: + def __init__(self, input_tokens: int, output_tokens: int): + self.trajectory = _trajectory( + { + "model": "gpt-5.6", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + }, + } + ) + + async def stop(self): + return None + + def reconcile_live_capture(self): + LiveLLMTrajectoryWriter( + tmp_path / "trajectory" / "llm_trajectory.jsonl" + ).reconcile(self.trajectory) + + def runtime(input_tokens: int, output_tokens: int) -> ProviderRuntime: + return ProviderRuntime( + kind="litellm", + agent_base_url="http://127.0.0.1:4000", + backend_model="gpt-5.6", + server=FakeServer(input_tokens, output_tokens), + ) + + first = runtime(10, 2) + final = runtime(20, 3) + rollout = Rollout.__new__(Rollout) + rollout._config = RolloutConfig(task_path=tmp_path / "task") + rollout._trajectory = [] + rollout._acp_client = None + rollout._agent_launch = "" + rollout._env = SimpleNamespace(stop=AsyncMock()) + rollout._environment = None + rollout._retired_usage_runtimes = [first] + rollout._usage_runtime = final + rollout._planes = SimpleNamespace( + stop_provider_runtime=lambda provider_runtime: provider_runtime.server.stop(), + extract_usage=extract_usage, + ) + rollout._rollout_dir = tmp_path + rollout._env_externally_owned = False + + await rollout.cleanup() + + assert rollout._usage_metrics["usage_source"] == "provider_response" + assert rollout._usage_metrics["n_input_tokens"] == 30 + assert rollout._usage_metrics["n_output_tokens"] == 5 + assert rollout._usage_metrics["total_tokens"] == 35 + assert rollout._retired_usage_runtimes == [] + + @pytest.mark.asyncio async def test_rollout_cleanup_propagates_provider_stop_failure_to_capture(tmp_path): """Guards PR #1057 against accepting a truncated live provider prefix.""" From 200b1075a6ee06b245c3fca30217efc6f7a4b955 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 22:29:28 -0700 Subject: [PATCH 65/74] test(capture): use matching Linux user group --- tests/test_sdk_lockdown.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sdk_lockdown.py b/tests/test_sdk_lockdown.py index 6cff79175..868495cd4 100644 --- a/tests/test_sdk_lockdown.py +++ b/tests/test_sdk_lockdown.py @@ -389,7 +389,7 @@ def test_linux_launch_really_sets_no_new_privs(self): cmd = build_priv_drop_cmd( "awk '/^NoNewPrivs:/ {print $2}' /proc/self/status", - "nobody", + "daemon", ) result = subprocess.run( From c3ac7f2dc4f2bd3ef5bc33df82341173e33f3fed Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 22:44:07 -0700 Subject: [PATCH 66/74] fix: preserve zero-call OAuth completion --- .../trajectories/llm_capture_records.py | 15 +++-- .../test_native_capture_resilience.py | 56 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture_records.py b/src/benchflow/trajectories/llm_capture_records.py index b14f2cdf6..6a8e27cd3 100644 --- a/src/benchflow/trajectories/llm_capture_records.py +++ b/src/benchflow/trajectories/llm_capture_records.py @@ -207,11 +207,16 @@ def assemble_capture( *_captured_targets(native_records, targets), } missing_targets = [target for target in targets if target not in captured_targets] - errors.extend( - f"no capture was attributable to {target.agent} " - f"({target.model or 'unknown model'}, {target.auth_mode.value})" - for target in missing_targets - ) + # A prepared role is expected to have no attributable capture when the + # rollout never made a model call. Captured rows are themselves evidence + # of a call when the ACP-level signal is unavailable, so keep the + # fail-closed missing-role check for every non-empty assembly. + if model_call_seen or records: + errors.extend( + f"no capture was attributable to {target.agent} " + f"({target.model or 'unknown model'}, {target.auth_mode.value})" + for target in missing_targets + ) role_captures = _role_captures(records, targets=targets) auth_mode = _aggregate_auth_mode( records, targets=targets, fallback_auth=fallback_auth diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index dc4416e81..3e71739a6 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -15,6 +15,7 @@ AuthMode, CaptureFidelity, CaptureStatus, + capture_manifest_preserves_audit_completion, ) from benchflow.trajectories.llm_capture_records import ( NativeCaptureBundle, @@ -30,6 +31,61 @@ STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) +@pytest.mark.asyncio +async def test_zero_call_oauth_target_preserves_clean_audit_completion( + tmp_path: Path, +) -> None: + """Guards PR #1057 review r3888576789 for clean zero-call OAuth runs.""" + + capture = LLMTrajectoryCapture( + tmp_path, + agent="claude-agent-acp", + model="claude-sonnet-4-6", + session_id="rollout-1", + started_at=STARTED_AT, + ) + target = _CaptureTarget( + agent="claude-agent-acp", + model="claude-sonnet-4-6", + credential_home="/home/agent", + auth_mode=AuthMode.OAUTH_SUBSCRIPTION, + native=True, + ) + capture._targets[(target.agent, target.model, target.credential_home)] = target + capture._refresh_manifest_auth_mode() + + async def collect_native(_env, *, targets): + assert targets == [target] + return NativeCollection() + + capture._native_collector.collect = collect_native + + await capture.finalize(None, acp_events=[], model_call_seen=False) + + manifest = json.loads( + (tmp_path / "trajectory" / "llm_trajectory.manifest.json").read_text() + ) + assert manifest["status"] == CaptureStatus.NO_MODEL_CALL + assert manifest["errors"] == [] + assert manifest["missing_fields"] == [] + assert manifest["exchange_count"] == 0 + assert manifest["role_captures"] == [ + { + "role": "agent", + "agent": "claude-agent-acp", + "model": "claude-sonnet-4-6", + "auth_mode": "oauth_subscription", + "capture_source": "none", + "capture_fidelity": "none", + "exchange_count": 0, + "request_complete": False, + "response_complete": False, + "leg": None, + } + ] + assert capture_manifest_preserves_audit_completion(manifest) is True + + def test_native_parser_normalizes_fallback_and_record_timestamps( tmp_path: Path, ) -> None: From f071fd5d9650bad4cc79e67e6514195fa911475b Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 23:07:34 -0700 Subject: [PATCH 67/74] fix: aggregate mixed-auth token usage --- src/benchflow/models.py | 2 +- src/benchflow/rollout/__init__.py | 31 ++++++--- src/benchflow/rollout/_usage.py | 70 ++++++++++++++++++- src/benchflow/usage_tracking.py | 17 ++++- tests/test_metrics.py | 4 ++ tests/test_native_acp_usage.py | 100 +++++++++++++++++++++++++--- tests/test_trial_litellm_runtime.py | 3 + 7 files changed, 204 insertions(+), 23 deletions(-) diff --git a/src/benchflow/models.py b/src/benchflow/models.py index b3558e0c5..9f673abe6 100644 --- a/src/benchflow/models.py +++ b/src/benchflow/models.py @@ -90,7 +90,7 @@ class RolloutResult: or None when provider telemetry was unavailable. cost_usd: Provider cost estimate in USD, or None when unavailable. usage_source: Token telemetry source. One of "provider_response", - "agent_native_acp", or "unavailable". + "agent_native_acp", "mixed", or "unavailable". price_source: Pricing table version used for cost_usd, or None. usage_details: Optional source-specific telemetry details. error: Error description string, or None on success. diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 526b3688e..329f3a47e 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -77,6 +77,7 @@ from benchflow._utils.text import describe_exception from benchflow.acp.types import McpServerSpec from benchflow.agents.credentials import upload_credential +from benchflow.agents.env import uses_native_subscription_auth from benchflow.agents.registry import AGENTS from benchflow.contracts import ( AgentProtocolError, @@ -163,6 +164,9 @@ ProviderFailure as ProviderFailure, ) from benchflow.rollout._usage import _as_nonnegative_int as _as_nonnegative_int +from benchflow.rollout._usage import ( + _merge_provider_and_native_usage_metrics as _merge_provider_and_native_usage_metrics, +) from benchflow.rollout._usage import ( _merge_provider_usage_metrics as _merge_provider_usage_metrics, ) @@ -228,7 +232,6 @@ from benchflow.trajectories.tree import RolloutNode, RolloutTree, Step from benchflow.usage_tracking import ( USAGE_SOURCE_AGENT_NATIVE_ACP, - USAGE_SOURCE_PROVIDER_RESPONSE, is_token_usage_available, ) @@ -684,6 +687,7 @@ def __init__(self, config: RolloutConfig) -> None: self._usage_metrics: dict[str, Any] = self._planes.extract_usage(None) self._native_usage_metrics: dict[str, Any] = _zero_native_acp_usage_metrics() self._native_usage_checkpoint: dict[str, int | None] | None = None + self._current_usage_is_native_subscription = False # Provider failure snapshotted during cleanup, after the usage proxy # imports its captures (Daytona's SandboxUsageProxy only fills trajectory # on stop()). Read by _provider_failure() so ACP-error classification can @@ -1389,6 +1393,11 @@ async def connect(self) -> None: model=cfg.primary_model, credential_home=_sandbox_user_home(cfg.sandbox_user), ) + self._current_usage_is_native_subscription = uses_native_subscription_auth( + cfg.primary_agent, + cfg.primary_model, + self._agent_env, + ) self._native_usage_checkpoint = None self._reapply_ask_user_handler() self._attach_trajectory_writer(rollout_dir) @@ -1772,6 +1781,8 @@ def _commit_acp_execution( def _collect_native_acp_usage(self) -> None: """Accumulate ACP PromptResponse.usage deltas for native subscription runs.""" + if getattr(self, "_current_usage_is_native_subscription", True) is False: + return session = getattr(self, "_session", None) latest_fn = getattr(session, "latest_usage_totals", None) if not callable(latest_fn): @@ -2151,17 +2162,15 @@ def _adopt_usage_runtime(self, previous: Any, current: Any) -> None: self._usage_runtime = current def _finalize_usage_metrics(self) -> None: - """Prefer LiteLLM usage, otherwise use trusted native ACP usage.""" + """Combine every trusted token source used by this rollout.""" current_metrics = getattr( self, "_usage_metrics", {"usage_source": "unavailable"} ) - if current_metrics.get("usage_source") == USAGE_SOURCE_PROVIDER_RESPONSE: - return native_metrics = getattr(self, "_native_usage_metrics", None) - if isinstance(native_metrics, dict) and is_token_usage_available( - native_metrics - ): - self._usage_metrics = native_metrics + self._usage_metrics = _merge_provider_and_native_usage_metrics( + current_metrics, + native_metrics if isinstance(native_metrics, dict) else None, + ) def _enforce_required_usage_tracking(self) -> None: usage_cfg = self._config.usage_tracking.with_env_defaults() @@ -2532,6 +2541,12 @@ async def connect_as(self, role: Role) -> None: credential_home=cred_home, role_name=role.name, ) + self._current_usage_is_native_subscription = uses_native_subscription_auth( + role.agent, + role.model, + agent_env, + ) + self._native_usage_checkpoint = None self._reapply_ask_user_handler() self._attach_trajectory_writer(rollout_dir) self._active_role = role diff --git a/src/benchflow/rollout/_usage.py b/src/benchflow/rollout/_usage.py index dfb6d6acc..1a030d0a2 100644 --- a/src/benchflow/rollout/_usage.py +++ b/src/benchflow/rollout/_usage.py @@ -12,7 +12,12 @@ from dataclasses import dataclass from typing import Any -from benchflow.usage_tracking import usage_unavailable +from benchflow.usage_tracking import ( + USAGE_SOURCE_AGENT_NATIVE_ACP, + USAGE_SOURCE_MIXED, + USAGE_SOURCE_PROVIDER_RESPONSE, + usage_unavailable, +) @dataclass(frozen=True) @@ -187,6 +192,69 @@ def _merge_provider_usage_metrics( return merged +def _merge_provider_and_native_usage_metrics( + provider_metrics: dict[str, Any], + native_metrics: dict[str, Any] | None, +) -> dict[str, Any]: + """Combine trusted provider and native-ACP totals without hiding provenance. + + A mixed-auth scene has two independently trusted telemetry surfaces. Its + token counters are additive, but native ACP does not provide pricing, so a + rollout-wide cost must remain unknown. The priced provider component and + both counter sets remain available in ``usage_details.source_breakdown``. + """ + + provider_available = ( + provider_metrics.get("usage_source") == USAGE_SOURCE_PROVIDER_RESPONSE + ) + native_available = bool( + isinstance(native_metrics, dict) + and native_metrics.get("usage_source") == USAGE_SOURCE_AGENT_NATIVE_ACP + ) + if provider_metrics.get("usage_source") == USAGE_SOURCE_MIXED: + return provider_metrics + if not provider_available: + if native_available and isinstance(native_metrics, dict): + return dict(native_metrics) + return provider_metrics + if not native_available: + return provider_metrics + + assert isinstance(native_metrics, dict) + merged = usage_unavailable() + for field in _PROVIDER_USAGE_COUNT_FIELDS: + merged[field] = _as_nonnegative_int( + provider_metrics.get(field) + ) + _as_nonnegative_int(native_metrics.get(field)) + provider_breakdown: dict[str, Any] = { + field: _as_nonnegative_int(provider_metrics.get(field)) + for field in _PROVIDER_USAGE_COUNT_FIELDS + } + provider_breakdown.update( + cost_usd=provider_metrics.get("cost_usd"), + price_source=provider_metrics.get("price_source"), + ) + native_breakdown: dict[str, Any] = { + field: _as_nonnegative_int(native_metrics.get(field)) + for field in _PROVIDER_USAGE_COUNT_FIELDS + } + native_details: dict[str, Any] = dict(native_metrics.get("usage_details") or {}) + if native_details: + native_breakdown["usage_details"] = native_details + usage_details = dict(native_details) + usage_details["source_breakdown"] = { + USAGE_SOURCE_PROVIDER_RESPONSE: provider_breakdown, + USAGE_SOURCE_AGENT_NATIVE_ACP: native_breakdown, + } + merged.update( + usage_source=USAGE_SOURCE_MIXED, + cost_usd=None, + price_source=None, + usage_details=usage_details, + ) + return merged + + def classify_api_failure( summary: dict[str, Any] | None, *, diff --git a/src/benchflow/usage_tracking.py b/src/benchflow/usage_tracking.py index ef63566aa..013ad2928 100644 --- a/src/benchflow/usage_tracking.py +++ b/src/benchflow/usage_tracking.py @@ -7,20 +7,31 @@ from typing import Any, Literal, cast UsageTrackingMode = Literal["auto", "required", "off"] -UsageSource = Literal["provider_response", "agent_native_acp", "unavailable"] +UsageSource = Literal[ + "provider_response", + "agent_native_acp", + "mixed", + "unavailable", +] USAGE_TRACKING_ENV = "BENCHFLOW_USAGE_TRACKING" USAGE_SOURCE_PROVIDER_RESPONSE = "provider_response" USAGE_SOURCE_AGENT_NATIVE_ACP = "agent_native_acp" +USAGE_SOURCE_MIXED = "mixed" USAGE_SOURCE_UNAVAILABLE = "unavailable" TRUSTED_USAGE_SOURCES: frozenset[str] = frozenset( - {USAGE_SOURCE_PROVIDER_RESPONSE, USAGE_SOURCE_AGENT_NATIVE_ACP} + { + USAGE_SOURCE_PROVIDER_RESPONSE, + USAGE_SOURCE_AGENT_NATIVE_ACP, + USAGE_SOURCE_MIXED, + } ) _MODES: set[str] = {"auto", "required", "off"} _USAGE_SOURCES: set[str] = { USAGE_SOURCE_PROVIDER_RESPONSE, USAGE_SOURCE_AGENT_NATIVE_ACP, + USAGE_SOURCE_MIXED, USAGE_SOURCE_UNAVAILABLE, } _LEGACY_USAGE_PROXY_KEYS: frozenset[str] = frozenset( @@ -174,6 +185,8 @@ def to_result_metadata( endpoint_kind = "sandbox" if environment == "daytona" else "host" if usage_source == USAGE_SOURCE_AGENT_NATIVE_ACP: endpoint_kind = "agent_native" + elif usage_source == USAGE_SOURCE_MIXED: + endpoint_kind = "mixed" elif usage_source == USAGE_SOURCE_UNAVAILABLE: endpoint_kind = "none" return { diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 7c9731c26..3d2ffe11d 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -432,6 +432,7 @@ def test_usage_source_type_contract_tracks_trusted_sources(): from benchflow.usage_tracking import ( TRUSTED_USAGE_SOURCES, USAGE_SOURCE_AGENT_NATIVE_ACP, + USAGE_SOURCE_MIXED, USAGE_SOURCE_PROVIDER_RESPONSE, USAGE_SOURCE_UNAVAILABLE, UsageSource, @@ -441,14 +442,17 @@ def test_usage_source_type_contract_tracks_trusted_sources(): assert set(get_args(UsageSource)) == { USAGE_SOURCE_PROVIDER_RESPONSE, USAGE_SOURCE_AGENT_NATIVE_ACP, + USAGE_SOURCE_MIXED, USAGE_SOURCE_UNAVAILABLE, } assert { USAGE_SOURCE_PROVIDER_RESPONSE, USAGE_SOURCE_AGENT_NATIVE_ACP, + USAGE_SOURCE_MIXED, } == TRUSTED_USAGE_SOURCES assert normalize_usage_source(USAGE_SOURCE_AGENT_NATIVE_ACP) == ( USAGE_SOURCE_AGENT_NATIVE_ACP ) + assert normalize_usage_source(USAGE_SOURCE_MIXED) == USAGE_SOURCE_MIXED with pytest.raises(ValueError, match="usage_source must be one of"): normalize_usage_source("new_unregistered_source") diff --git a/tests/test_native_acp_usage.py b/tests/test_native_acp_usage.py index ebfea42f2..eb819b236 100644 --- a/tests/test_native_acp_usage.py +++ b/tests/test_native_acp_usage.py @@ -55,6 +55,7 @@ def test_rollout_native_acp_usage_uses_cumulative_deltas(): session = ACPSession("session-1") rollout = Rollout.__new__(Rollout) rollout._session = session + rollout._current_usage_is_native_subscription = True rollout._native_usage_checkpoint = None session.record_prompt_usage( @@ -93,17 +94,46 @@ def test_rollout_native_acp_usage_uses_cumulative_deltas(): } -def test_rollout_provider_usage_wins_over_native_acp_usage(): - """Guards PR #613 follow-up: LiteLLM provider telemetry remains authoritative.""" +def test_rollout_provider_acp_usage_is_not_double_counted_as_native(): + """Guards PR #1057 review r3888606859 against mixed-source duplication.""" + from benchflow.acp.session import ACPSession + from benchflow.rollout import Rollout, _zero_native_acp_usage_metrics + + session = ACPSession("session-1") + session.record_prompt_usage( + SimpleNamespace( + input_tokens=10, + output_tokens=4, + total_tokens=14, + cached_read_tokens=0, + cached_write_tokens=0, + thought_tokens=0, + ) + ) + rollout = Rollout.__new__(Rollout) + rollout._session = session + rollout._current_usage_is_native_subscription = False + rollout._native_usage_metrics = _zero_native_acp_usage_metrics() + rollout._native_usage_checkpoint = None + + rollout._collect_native_acp_usage() + + assert rollout._native_usage_metrics["usage_source"] == "unavailable" + assert rollout._native_usage_metrics["total_tokens"] == 0 + assert rollout._native_usage_checkpoint is None + + +def test_rollout_mixed_auth_sums_provider_and_native_acp_usage(): + """Guards PR #1057 review r3888606859 against dropping OAuth usage.""" from benchflow.rollout import Rollout rollout = Rollout.__new__(Rollout) rollout._usage_metrics = { "n_input_tokens": 100, "n_output_tokens": 20, - "n_cache_read_tokens": 0, - "n_cache_creation_tokens": 0, - "total_tokens": 120, + "n_cache_read_tokens": 4, + "n_cache_creation_tokens": 2, + "total_tokens": 126, "cost_usd": 0.01, "usage_source": "provider_response", "price_source": "litellm", @@ -111,19 +141,67 @@ def test_rollout_provider_usage_wins_over_native_acp_usage(): rollout._native_usage_metrics = { "n_input_tokens": 10, "n_output_tokens": 5, - "n_cache_read_tokens": 0, - "n_cache_creation_tokens": 0, - "total_tokens": 15, + "n_cache_read_tokens": 3, + "n_cache_creation_tokens": 1, + "total_tokens": 19, "cost_usd": None, "usage_source": "agent_native_acp", "price_source": None, - "usage_details": {"thought_tokens": 0}, + "usage_details": {"thought_tokens": 2}, } rollout._finalize_usage_metrics() - assert rollout._usage_metrics["usage_source"] == "provider_response" - assert rollout._usage_metrics["total_tokens"] == 120 + assert rollout._usage_metrics == { + "n_input_tokens": 110, + "n_output_tokens": 25, + "n_cache_read_tokens": 7, + "n_cache_creation_tokens": 3, + "total_tokens": 145, + "cost_usd": None, + "usage_source": "mixed", + "price_source": None, + "usage_details": { + "thought_tokens": 2, + "source_breakdown": { + "provider_response": { + "n_input_tokens": 100, + "n_output_tokens": 20, + "n_cache_read_tokens": 4, + "n_cache_creation_tokens": 2, + "total_tokens": 126, + "cost_usd": 0.01, + "price_source": "litellm", + }, + "agent_native_acp": { + "n_input_tokens": 10, + "n_output_tokens": 5, + "n_cache_read_tokens": 3, + "n_cache_creation_tokens": 1, + "total_tokens": 19, + "usage_details": {"thought_tokens": 2}, + }, + }, + }, + } + + first_finalization = rollout._usage_metrics + rollout._finalize_usage_metrics() + assert rollout._usage_metrics is first_finalization + + +def test_mixed_usage_metadata_reports_both_endpoint_kinds(): + """Guards PR #1057 review r3888606859 for explicit mixed provenance.""" + from benchflow.usage_tracking import UsageTrackingConfig + + metadata = UsageTrackingConfig(mode="required").to_result_metadata( + environment="docker", + status="enabled", + usage_source="mixed", + ) + + assert metadata["usage_source"] == "mixed" + assert metadata["endpoint_kind"] == "mixed" def test_required_usage_accepts_native_acp_usage(tmp_path): diff --git a/tests/test_trial_litellm_runtime.py b/tests/test_trial_litellm_runtime.py index 149c33c58..5e9971f51 100644 --- a/tests/test_trial_litellm_runtime.py +++ b/tests/test_trial_litellm_runtime.py @@ -79,6 +79,7 @@ async def test_trial_connect_as_starts_litellm_for_role(tmp_path: Path): next_runtime = SimpleNamespace(kind="litellm", name="reviewer") rollout._usage_runtime = previous_runtime rollout._retired_usage_runtimes = [] + rollout._native_usage_checkpoint = {"total_tokens": 99} rollout._timing = {} rollout._disallow_web_tools = False rollout._agent_cfg = SimpleNamespace() @@ -123,3 +124,5 @@ async def fake_connect_acp(**kwargs): assert calls == ["litellm", "acp"] assert rollout._usage_runtime is next_runtime assert rollout._retired_usage_runtimes == [previous_runtime] + assert rollout._current_usage_is_native_subscription is False + assert rollout._native_usage_checkpoint is None From 3f58679ca68393924a8740fd4d37355d22e66d7d Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 23:17:14 -0700 Subject: [PATCH 68/74] fix: honor primary role-scoped credentials --- src/benchflow/rollout/__init__.py | 2 +- src/benchflow/rollout/_config.py | 10 ++++++++++ tests/test_connect_as_env.py | 14 ++++++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 329f3a47e..4c5241997 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -983,7 +983,7 @@ async def setup(self) -> None: ) and cfg.primary_agent != "oracle" self._agent_env = _apply_web_policy( self._planes.resolve_agent_env( - cfg.primary_agent, cfg.primary_model, cfg.agent_env + cfg.primary_agent, cfg.primary_model, cfg.primary_env ), disallow=self._disallow_web_tools, ) diff --git a/src/benchflow/rollout/_config.py b/src/benchflow/rollout/_config.py index d88a47c76..f3f9d6733 100644 --- a/src/benchflow/rollout/_config.py +++ b/src/benchflow/rollout/_config.py @@ -432,3 +432,13 @@ def primary_reasoning_effort(self) -> str | None: """Reasoning effort for the first role of the first scene.""" role = self._primary_role return role.reasoning_effort if role else self.reasoning_effort + + @property + def primary_env(self) -> dict[str, str]: + """Config env merged with the first scene role's scoped overrides.""" + + role = self._primary_role + return { + **(self.agent_env or {}), + **((role.env or {}) if role is not None else {}), + } diff --git a/tests/test_connect_as_env.py b/tests/test_connect_as_env.py index 129e1d86b..084240063 100644 --- a/tests/test_connect_as_env.py +++ b/tests/test_connect_as_env.py @@ -32,6 +32,20 @@ def _make_config(agent_env=None, role_env=None): class TestConnectAsEnvMerge: """Verify connect_as() merges cfg.agent_env with role.env correctly.""" + def test_primary_setup_env_includes_role_scoped_credentials(self): + """Guards PR #1057 mixed-auth setup against requiring global secrets.""" + + config = _make_config( + agent_env={"GLOBAL": "yes", "SHARED": "config"}, + role_env={"CLAUDE_CODE_OAUTH_TOKEN": "oauth", "SHARED": "role"}, + ) + + assert config.primary_env == { + "GLOBAL": "yes", + "CLAUDE_CODE_OAUTH_TOKEN": "oauth", + "SHARED": "role", + } + @pytest.fixture() def _mock_trial(self, tmp_path): """Return a Rollout stub wired to capture the agent_env passed to connect_acp.""" From 1f18b77452b0dbaac3a8a3a1ec4c225a37013bd1 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 23:30:59 -0700 Subject: [PATCH 69/74] fix: isolate credentials across agent roles --- src/benchflow/agents/credentials.py | 227 ++++++++++++----------- src/benchflow/rollout/__init__.py | 31 +++- tests/test_litellm_credential_custody.py | 86 ++++++++- tests/test_session_factory_runtime.py | 44 +++++ 4 files changed, 274 insertions(+), 114 deletions(-) diff --git a/src/benchflow/agents/credentials.py b/src/benchflow/agents/credentials.py index 274c14ca7..cb32cd84c 100644 --- a/src/benchflow/agents/credentials.py +++ b/src/benchflow/agents/credentials.py @@ -70,9 +70,16 @@ def _proxy_auth_cleanup_command( paths: list[str], settings_targets: list[dict[str, object]], ) -> str: - """Build the no-follow cleanup command gated by process isolation.""" + """Build the no-follow cleanup command gated by process isolation. - return f"{process_guard} && " + " ".join( + JavaScript agents install BenchFlow's private Node runtime before this + command executes. Python-only agents need not pay that installation cost: + when the runtime is absent they remain trusted only if every possible + credential and settings target is already absent. A present target fails + closed instead of attempting a weaker shell rewrite. + """ + + cleanup = " ".join( ( f"{_BENCHFLOW_NODE_BIN} -e", shlex.quote(_PROXY_AUTH_CLEANUP_JS), @@ -81,6 +88,17 @@ def _proxy_auth_cleanup_command( shlex.quote(json.dumps(settings_targets, separators=(",", ":"))), ) ) + target_paths = [*paths, *(str(target["path"]) for target in settings_targets)] + absence_checks = " && ".join( + f"[ ! -e {shlex.quote(path)} ] && [ ! -L {shlex.quote(path)} ]" + for path in dict.fromkeys(target_paths) + ) + if not absence_checks: + absence_checks = "true" + return ( + f"{process_guard} && if [ -x {_BENCHFLOW_NODE_BIN} ]; then " + f"{cleanup}; else {absence_checks}; fi" + ) async def upload_credential( @@ -234,83 +252,25 @@ async def isolate_agent_for_proxy_capture( agent_env: dict[str, str], cred_home: str, ) -> bool: - """Prove process isolation and remove native auth before an API proxy run. - - A reused image may already contain the CLI login detected by - ``SubscriptionAuth.detect_file``. API-key mode must not leave that alternate - provider route available to the agent. A safe CLI config-home override is - scrubbed as well as removed before launch; an override outside the sandbox - user's home fails capture trust without allowing a root deletion there. - Every API-proxied agent first requires a non-root sandbox user with no live - processes; an existing agent or task process is preserved and makes capture - audit-only. Subscription-capable agents then use their already-required - JavaScript runtime to traverse with no-follow directory descriptors, remove - native login files, and sanitize registry-declared credential settings. - Root-agent runs remain audit-only because stale root processes cannot be - safely distinguished. Return false unless process isolation is proven and - every eligible credential can be identified, removed, and verified absent. + """Prove process isolation and remove every native route before API proxying. + + A shared sandbox can retain credentials from an earlier role. The current + agent's login file is therefore not a sufficient cleanup boundary: all + registry-declared subscription files, credential-bearing settings, and Pi's + generated provider map are removed through the same no-follow traversal. + Unsafe effective/config homes and any live sandbox-user process fail capture + trust closed. Root-agent runs remain audit-only because their processes + cannot be distinguished safely from orchestration. """ - agent_cfg = AGENTS.get(agent) - subscription_auth = agent_cfg.subscription_auth if agent_cfg else None if env is None: logger.warning("Cannot prove proxy process isolation for %s", agent) return False process_guard, processes_safe = _proxy_process_isolation_guard(cred_home) - - if subscription_auth is None: - if agent == "pi-acp": - home = PurePosixPath(cred_home) - relative_models_path = PurePosixPath(".pi/agent/models.json") - paths = [str(home / relative_models_path)] - paths_safe = True - for home_env in ("HOME", "BENCHFLOW_AGENT_HOME"): - override_value = agent_env.get(home_env, "") - if override_value and override_value != cred_home: - override_home = _safe_proxy_auth_root( - override_value, - cred_home=cred_home, - agent=agent, - ) - if override_home is None: - paths_safe = False - else: - paths.append(str(override_home / relative_models_path)) - agent_env[home_env] = cred_home - paths = list(dict.fromkeys(paths)) - cleanup_command = _proxy_auth_cleanup_command(process_guard, paths, []) - else: - paths_safe = True - cleanup_command = process_guard - try: - result = await env.exec( - cleanup_command, - user="root", - timeout_sec=10, - ) - except Exception as exc: - logger.warning("Failed to isolate %s agent processes: %s", agent, exc) - return False - return bool(result.return_code == 0 and paths_safe and processes_safe) - - primary_files = [ - auth_file - for auth_file in subscription_auth.files - if auth_file.host_path == subscription_auth.detect_file - ] - if len(primary_files) != 1: - logger.warning( - "Cannot prove proxy isolation for %s subscription credentials", agent - ) - return False - home = PurePosixPath(cred_home) - primary_path = primary_files[0].container_path.format(home=cred_home) - primary_relative_path = PurePosixPath(primary_path).relative_to(home) - paths = [primary_path] + roots = [home] paths_safe = True - for home_env in ("HOME", "BENCHFLOW_AGENT_HOME"): override_value = agent_env.get(home_env, "") if override_value and override_value != cred_home: @@ -322,44 +282,99 @@ async def isolate_agent_for_proxy_capture( if override_home is None: paths_safe = False else: - paths.append(str(override_home / primary_relative_path)) + roots.append(override_home) agent_env[home_env] = cred_home - if subscription_auth.config_dir_env: - override_value = agent_env.pop(subscription_auth.config_dir_env, "") - if override_value: - override_dir = _safe_proxy_auth_root( - override_value, - cred_home=cred_home, - agent=agent, + roots = list(dict.fromkeys(roots)) + paths = [str(root / ".pi/agent/models.json") for root in roots] + settings_targets: list[dict[str, object]] = [] + seen_auth: set[int] = set() + for registered_agent in AGENTS.values(): + subscription_auth = registered_agent.subscription_auth + if subscription_auth is None or id(subscription_auth) in seen_auth: + continue + seen_auth.add(id(subscription_auth)) + primary_files = [ + auth_file + for auth_file in subscription_auth.files + if auth_file.host_path == subscription_auth.detect_file + ] + if len(primary_files) != 1: + logger.warning( + "Cannot prove proxy isolation for registered %s credentials", + registered_agent.name, ) - if override_dir is None: + paths_safe = False + continue + + relative_files: list[PurePosixPath] = [] + for auth_file in subscription_auth.files: + try: + rendered = PurePosixPath( + auth_file.container_path.format(home=cred_home) + ) + relative_files.append(rendered.relative_to(home)) + except (KeyError, ValueError): + logger.warning( + "Unsafe registered credential path for %s", + registered_agent.name, + ) paths_safe = False - else: - paths.append(str(override_dir / PurePosixPath(primary_path).name)) + for root in roots: + paths.extend(str(root / relative_path) for relative_path in relative_files) + + try: + primary_rendered = PurePosixPath( + primary_files[0].container_path.format(home=cred_home) + ) + primary_relative = primary_rendered.relative_to(home) + except (KeyError, ValueError): + paths_safe = False + continue + settings_parents = [root / primary_relative.parent for root in roots] + + if subscription_auth.config_dir_env: + override_value = agent_env.pop(subscription_auth.config_dir_env, "") + if override_value: + override_dir = _safe_proxy_auth_root( + override_value, + cred_home=cred_home, + agent=agent, + ) + if override_dir is None: + paths_safe = False + else: + paths.extend( + str(override_dir / relative_path.name) + for relative_path in relative_files + ) + settings_parents.append(override_dir) + + if subscription_auth.proxy_settings_file: + settings_file = PurePosixPath(subscription_auth.proxy_settings_file) + if ( + not subscription_auth.proxy_settings_drop_keys + or settings_file.is_absolute() + or len(settings_file.parts) != 1 + or settings_file.name != subscription_auth.proxy_settings_file + ): + logger.warning( + "Cannot prove proxy settings isolation for %s", + registered_agent.name, + ) + paths_safe = False + continue + for parent in settings_parents: + target: dict[str, object] = { + "path": str(parent / settings_file), + "drop_keys": list(subscription_auth.proxy_settings_drop_keys), + } + settings_targets.append(target) paths = list(dict.fromkeys(paths)) - settings_targets: list[dict[str, object]] = [] - if subscription_auth.proxy_settings_file: - settings_file = PurePosixPath(subscription_auth.proxy_settings_file) - if ( - not subscription_auth.proxy_settings_drop_keys - or settings_file.is_absolute() - or len(settings_file.parts) != 1 - or settings_file.name != subscription_auth.proxy_settings_file - ): - logger.warning("Cannot prove proxy settings isolation for %s", agent) - return False - settings_targets = [ - { - "path": str(PurePosixPath(credential_path).parent / settings_file), - "drop_keys": list(subscription_auth.proxy_settings_drop_keys), - } - for credential_path in paths - ] - settings_targets = list( - {str(target["path"]): target for target in settings_targets}.values() - ) + settings_targets = list( + {str(target["path"]): target for target in settings_targets}.values() + ) cleanup_command = _proxy_auth_cleanup_command( process_guard, paths, @@ -372,13 +387,17 @@ async def isolate_agent_for_proxy_capture( timeout_sec=10, ) except Exception as exc: - logger.warning("Failed to isolate %s subscription credential: %s", agent, exc) + logger.warning("Failed to isolate %s proxy credentials: %s", agent, exc) return False if result.return_code != 0 or not paths_safe or not processes_safe: detail = (getattr(result, "stderr", "") or "").strip()[:500] logger.warning( - "Subscription credential remained accessible in proxy mode for %s%s", + "Proxy credential/process isolation failed for %s " + "(rc=%s, paths_safe=%s, processes_safe=%s)%s", agent, + result.return_code, + paths_safe, + processes_safe, f": {detail}" if detail else "", ) return False diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 4c5241997..ca28620a5 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1464,11 +1464,36 @@ async def disconnect(self) -> None: # Kill any lingering agent processes to prevent context bleed between scenes agent_pattern = _agent_process_kill_pattern(self._agent_launch) if self._env and agent_pattern: - with contextlib.suppress(Exception): - await self._env.exec( - f"pkill -f {shlex.quote(agent_pattern)} || true", + sandbox_user = getattr(getattr(self, "_config", None), "sandbox_user", None) + if sandbox_user: + quoted_user = shlex.quote(sandbox_user) + quoted_pattern = shlex.quote(agent_pattern) + termination_command = ( + f"bf_agent_uid=$(id -u -- {quoted_user}) || exit 1\n" + f'pkill -u "$bf_agent_uid" -f {quoted_pattern} ' + ">/dev/null 2>&1 || true\n" + "bf_wait=0\n" + f'while [ "$bf_wait" -lt 20 ] && pgrep -u "$bf_agent_uid" -f ' + f"{quoted_pattern} >/dev/null 2>&1; do\n" + " sleep 0.1\n" + " bf_wait=$((bf_wait + 1))\n" + "done\n" + f'! pgrep -u "$bf_agent_uid" -f {quoted_pattern} ' + ">/dev/null 2>&1" + ) + else: + termination_command = f"pkill -f {shlex.quote(agent_pattern)} || true" + try: + terminated = await self._env.exec( + termination_command, timeout_sec=10, ) + if getattr(terminated, "return_code", 0) != 0: + logger.warning("Agent process did not quiesce after disconnect") + except Exception as exc: + logger.warning( + "Agent process cleanup failed during disconnect: %s", exc + ) self._active_role = None self._session_tool_count = 0 self._session_traj_count = 0 diff --git a/tests/test_litellm_credential_custody.py b/tests/test_litellm_credential_custody.py index c50acb3f1..7ad624169 100644 --- a/tests/test_litellm_credential_custody.py +++ b/tests/test_litellm_credential_custody.py @@ -152,11 +152,41 @@ async def test_proxy_process_guard_applies_without_subscription_auth(agent) -> N assert "pgrep -u" in sandbox.commands[0] assert '"$bf_agent_uid" -ne 0' in sandbox.commands[0] assert '"$bf_pgrep_rc" -eq 1' in sandbox.commands[0] - if agent == "pi-acp": - assert _PROXY_AUTH_CLEANUP_JS in sandbox.commands[0] - assert "/home/agent/.pi/agent/models.json" in sandbox.commands[0] - else: - assert _PROXY_AUTH_CLEANUP_JS not in sandbox.commands[0] + assert _PROXY_AUTH_CLEANUP_JS in sandbox.commands[0] + assert "if [ -x /opt/benchflow/node/bin/node ]" in sandbox.commands[0] + assert "[ ! -e /home/agent/.codex/auth.json ]" in sandbox.commands[0] + assert "[ ! -L /home/agent/.codex/auth.json ]" in sandbox.commands[0] + assert "/home/agent/.pi/agent/models.json" in sandbox.commands[0] + assert "/home/agent/.claude/.credentials.json" in sandbox.commands[0] + assert "/home/agent/.codex/auth.json" in sandbox.commands[0] + assert "/home/agent/.gemini/oauth_creds.json" in sandbox.commands[0] + + +@pytest.mark.asyncio +async def test_proxy_cleanup_covers_credentials_from_previous_agent_roles() -> None: + """Guards PR #1057 mixed auth against cross-role subscription leakage.""" + + sandbox = _SubscriptionIsolationSandbox() + + trusted = await isolate_agent_for_proxy_capture( + sandbox, + agent="codex-acp", + agent_env={"AZURE_API_KEY": "provider-key"}, + cred_home="/home/agent", + ) + + assert trusted is True + command = sandbox.commands[0] + for credential_path in ( + "/home/agent/.claude/.credentials.json", + "/home/agent/.codex/auth.json", + "/home/agent/.gemini/oauth_creds.json", + "/home/agent/.gemini/settings.json", + "/home/agent/.gemini/google_accounts.json", + "/home/agent/.pi/agent/models.json", + ): + assert credential_path in command + assert "/home/agent/.claude/settings.json" in command @pytest.mark.asyncio @@ -226,7 +256,9 @@ async def test_proxy_process_guard_rejects_root_agent_without_subscription_auth( ) assert trusted is False - assert sandbox.commands == ["false"] + assert len(sandbox.commands) == 1 + assert sandbox.commands[0].startswith("false && ") + assert _PROXY_AUTH_CLEANUP_JS in sandbox.commands[0] @pytest.mark.parametrize( @@ -279,7 +311,9 @@ async def fake_start(**_kwargs): assert provider_runtime is not None assert provider_runtime.capture_trusted is False - assert sandbox.commands == ["false"] + assert len(sandbox.commands) == 1 + assert sandbox.commands[0].startswith("false && ") + assert _PROXY_AUTH_CLEANUP_JS in sandbox.commands[0] @pytest.mark.skipif( @@ -409,6 +443,44 @@ def test_proxy_auth_cleanup_sanitizes_claude_credential_settings(tmp_path) -> No assert (after.st_uid, after.st_gid) == (before.st_uid, before.st_gid) +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("node") is None, + reason="the sandbox credential cleanup requires Linux procfs and Node.js", +) +def test_proxy_auth_cleanup_removes_credentials_from_every_agent_role(tmp_path) -> None: + """Guards PR #1057 mixed auth against a previous role's native login.""" + + sandbox_home = tmp_path / "home" / "agent" + credential_paths = [ + sandbox_home / ".claude" / ".credentials.json", + sandbox_home / ".codex" / "auth.json", + sandbox_home / ".gemini" / "oauth_creds.json", + sandbox_home / ".gemini" / "settings.json", + sandbox_home / ".gemini" / "google_accounts.json", + sandbox_home / ".pi" / "agent" / "models.json", + ] + for credential_path in credential_paths: + credential_path.parent.mkdir(parents=True, exist_ok=True) + credential_path.write_text("subscription-secret") + + removed = subprocess.run( + [ + "node", + "-e", + _PROXY_AUTH_CLEANUP_JS, + "--", + json.dumps([str(path) for path in credential_paths]), + json.dumps([]), + ], + check=False, + capture_output=True, + text=True, + ) + + assert removed.returncode == 0, removed.stderr + assert all(not path.exists() for path in credential_paths) + + @pytest.mark.skipif( sys.platform != "linux" or shutil.which("node") is None, reason="the sandbox settings sanitizer requires Linux procfs and Node.js", diff --git a/tests/test_session_factory_runtime.py b/tests/test_session_factory_runtime.py index 23d08d3b7..7ea04d273 100644 --- a/tests/test_session_factory_runtime.py +++ b/tests/test_session_factory_runtime.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace import pytest @@ -339,6 +340,49 @@ async def test_disconnect_clears_session_factory_state(): assert rollout._phase == "installed" +@pytest.mark.asyncio +async def test_disconnect_waits_for_role_agent_process_to_quiesce(): + """Guards PR #1057 mixed auth against cross-role process overlap.""" + + class _RecordingEnv: + def __init__(self) -> None: + self.calls: list[tuple[str, dict]] = [] + + async def exec(self, command: str, **kwargs): + self.calls.append((command, kwargs)) + return SimpleNamespace(return_code=0) + + from benchflow.rollout import Rollout + + env = _RecordingEnv() + rollout = Rollout.__new__(Rollout) + rollout._acp_client = None + rollout._session = _FakeSession() + rollout._session_adapter = None + rollout._is_session_factory = True + rollout._agent_launch = "/opt/benchflow/bin/claude-agent-acp" + rollout._env = env + rollout._config = SimpleNamespace(sandbox_user="agent") + rollout._active_role = object() + rollout._session_tool_count = 0 + rollout._session_traj_count = 0 + rollout._trajectory = [] + rollout._partial_trajectory = False + rollout._trajectory_source = None + rollout._phase = "connected" + + await rollout.disconnect() + + assert len(env.calls) == 1 + command, kwargs = env.calls[0] + assert "id -u -- agent" in command + assert 'pkill -u "$bf_agent_uid" -f' in command + assert 'pgrep -u "$bf_agent_uid" -f' in command + assert 'while [ "$bf_wait" -lt 20 ]' in command + assert "sleep 0.1" in command + assert kwargs == {"timeout_sec": 10} + + @pytest.mark.asyncio async def test_steps_only_session_trajectory_sink_writes_steps(tmp_path): """The real streaming sink must work on a steps-only session-factory Session. From b21b1aaeed8af8b189f2852ffee1260687fde709 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sat, 29 Aug 2026 23:43:47 -0700 Subject: [PATCH 70/74] fix: quiesce native capture roles safely --- src/benchflow/rollout/__init__.py | 57 +++++++--------- src/benchflow/rollout/_setup.py | 67 +++++++++++++++++++ .../trajectories/native_capture_collection.py | 14 ++-- tests/test_agent_setup.py | 30 ++++++++- tests/test_session_factory_runtime.py | 5 +- .../test_native_capture_resilience.py | 35 +++++++++- 6 files changed, 167 insertions(+), 41 deletions(-) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index ca28620a5..68aa05f17 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -124,7 +124,7 @@ _agent_launch_with_web_policy as _agent_launch_with_web_policy, ) from benchflow.rollout._setup import ( - _agent_process_kill_pattern as _agent_process_kill_pattern, + _agent_process_termination_command as _agent_process_termination_command, ) from benchflow.rollout._setup import _apply_prompt_prefix as _apply_prompt_prefix from benchflow.rollout._setup import _apply_web_policy as _apply_web_policy @@ -1452,48 +1452,37 @@ async def disconnect(self) -> None: self._capture_partial_session_factory_trajectory() else: self._capture_partial_acp_trajectory() - if self._acp_client: - try: - await self._acp_client.close() - except Exception as e: - logger.warning(f"ACP client close failed: {e}") - self._acp_client = None - self._session = None - self._session_adapter = None - self._is_session_factory = False - # Kill any lingering agent processes to prevent context bleed between scenes - agent_pattern = _agent_process_kill_pattern(self._agent_launch) - if self._env and agent_pattern: - sandbox_user = getattr(getattr(self, "_config", None), "sandbox_user", None) - if sandbox_user: - quoted_user = shlex.quote(sandbox_user) - quoted_pattern = shlex.quote(agent_pattern) - termination_command = ( - f"bf_agent_uid=$(id -u -- {quoted_user}) || exit 1\n" - f'pkill -u "$bf_agent_uid" -f {quoted_pattern} ' - ">/dev/null 2>&1 || true\n" - "bf_wait=0\n" - f'while [ "$bf_wait" -lt 20 ] && pgrep -u "$bf_agent_uid" -f ' - f"{quoted_pattern} >/dev/null 2>&1; do\n" - " sleep 0.1\n" - " bf_wait=$((bf_wait + 1))\n" - "done\n" - f'! pgrep -u "$bf_agent_uid" -f {quoted_pattern} ' - ">/dev/null 2>&1" - ) - else: - termination_command = f"pkill -f {shlex.quote(agent_pattern)} || true" + # Terminate the complete agent tree while the ACP wrapper still owns its + # children. Closing the transport first can orphan the real CLI under + # PID 1, making it impossible to distinguish from an unrelated task + # process during the next role's proxy-custody check. + termination_command = _agent_process_termination_command( + self._agent_launch, + getattr(getattr(self, "_config", None), "sandbox_user", None), + ) + if self._env and termination_command: try: terminated = await self._env.exec( termination_command, timeout_sec=10, ) if getattr(terminated, "return_code", 0) != 0: - logger.warning("Agent process did not quiesce after disconnect") + logger.warning( + "Agent process tree did not quiesce after disconnect" + ) except Exception as exc: logger.warning( - "Agent process cleanup failed during disconnect: %s", exc + "Agent process-tree cleanup failed during disconnect: %s", exc ) + if self._acp_client: + try: + await self._acp_client.close() + except Exception as e: + logger.warning(f"ACP client close failed: {e}") + self._acp_client = None + self._session = None + self._session_adapter = None + self._is_session_factory = False self._active_role = None self._session_tool_count = 0 self._session_traj_count = 0 diff --git a/src/benchflow/rollout/_setup.py b/src/benchflow/rollout/_setup.py index 33c95c7e8..23e0aeaaa 100644 --- a/src/benchflow/rollout/_setup.py +++ b/src/benchflow/rollout/_setup.py @@ -186,6 +186,73 @@ def _agent_process_kill_pattern(agent_launch: str) -> str | None: return None +def _agent_process_termination_command( + agent_launch: str, sandbox_user: str | None +) -> str | None: + """Return a bounded command that terminates one agent process tree. + + ACP wrappers commonly spawn the real CLI as a child. Killing only the + wrapper can orphan that child under PID 1, which both leaks context across + roles and prevents the next provider-proxy role from proving exclusive + credential custody. For non-root sandbox users, freeze each discovered + descendant before walking deeper, then terminate only the recorded tree. + Unrelated processes owned by the same user are deliberately untouched. + """ + + pattern = _agent_process_kill_pattern(agent_launch) + if pattern is None: + return None + quoted_pattern = shlex.quote(pattern) + if not sandbox_user: + return f"pkill -f {quoted_pattern} || true" + + quoted_user = shlex.quote(sandbox_user) + return ( + f"bf_agent_uid=$(id -u -- {quoted_user}) || exit 1\n" + f'bf_frontier=$(pgrep -u "$bf_agent_uid" -f {quoted_pattern} ' + "2>/dev/null || true)\n" + "bf_pids=\n" + 'while [ -n "$bf_frontier" ]; do\n' + " bf_next=\n" + " for bf_pid in $bf_frontier; do\n" + " case \"$bf_pid\" in ''|*[!0-9]*) exit 1 ;; esac\n" + ' case " $bf_pids " in *" $bf_pid "*) continue ;; esac\n' + ' bf_pids="$bf_pids $bf_pid"\n' + ' kill -STOP "$bf_pid" 2>/dev/null || true\n' + ' bf_children=$(pgrep -u "$bf_agent_uid" -P "$bf_pid" ' + "2>/dev/null || true)\n" + ' bf_next="$bf_next $bf_children"\n' + " done\n" + ' bf_frontier="$bf_next"\n' + "done\n" + 'if [ -n "$bf_pids" ]; then\n' + " kill -TERM $bf_pids 2>/dev/null || true\n" + " kill -CONT $bf_pids 2>/dev/null || true\n" + "fi\n" + "bf_wait=0\n" + 'while [ "$bf_wait" -lt 20 ]; do\n' + " bf_alive=\n" + " for bf_pid in $bf_pids; do\n" + ' if kill -0 "$bf_pid" 2>/dev/null; then ' + 'bf_alive="$bf_alive $bf_pid"; fi\n' + " done\n" + ' [ -z "$bf_alive" ] && break\n' + " sleep 0.1\n" + " bf_wait=$((bf_wait + 1))\n" + "done\n" + 'if [ -n "$bf_alive" ]; then\n' + " kill -KILL $bf_alive 2>/dev/null || true\n" + " sleep 0.1\n" + "fi\n" + "bf_alive=\n" + "for bf_pid in $bf_pids; do\n" + ' if kill -0 "$bf_pid" 2>/dev/null; then ' + 'bf_alive="$bf_alive $bf_pid"; fi\n' + "done\n" + '[ -z "$bf_alive" ]' + ) + + def _configured_task_workdir(task: Any) -> str | None: """Return the task-declared sandbox workdir, if any.""" diff --git a/src/benchflow/trajectories/native_capture_collection.py b/src/benchflow/trajectories/native_capture_collection.py index 88ab0d773..099074d94 100644 --- a/src/benchflow/trajectories/native_capture_collection.py +++ b/src/benchflow/trajectories/native_capture_collection.py @@ -88,9 +88,10 @@ async def ensure(self, env: Any, *, sandbox_user: str | None) -> int: f"find {self.remote_root} -depth -mindepth 1 -delete " "2>/dev/null || true\n" f"mkdir -p {self.remote_root}/raw {self.remote_root}/otel\n" - f"chown -R {capture_owner} {self.remote_root}\n" - f"chmod 700 {self.remote_root} " - f"{self.remote_root}/raw {self.remote_root}/otel", + f"chown root:root {self.remote_root} {self.remote_root}/otel\n" + f"chown {capture_owner} {self.remote_root}/raw\n" + f"chmod 711 {self.remote_root}\n" + f"chmod 700 {self.remote_root}/raw {self.remote_root}/otel", user="root", timeout_sec=10, ) @@ -133,7 +134,12 @@ async def ensure(self, env: Any, *, sandbox_user: str | None) -> int: exit 1 """ self.owned = True - result = await env.exec(command, user=sandbox_user or "root", timeout_sec=10) + # The sink is BenchFlow infrastructure, not part of the agent process + # tree. Keep it root-owned so role teardown and the next API-proxy + # custody proof can require zero live sandbox-user processes. Claude + # still writes its raw-body files into the separately agent-owned + # ``raw`` directory and exports OTel to this loopback listener. + result = await env.exec(command, user="root", timeout_sec=10) if result.return_code != 0: detail = (result.stderr or result.stdout or "collector did not start")[:300] raise RuntimeError(f"Claude OTel sink failed to start: {detail}") diff --git a/tests/test_agent_setup.py b/tests/test_agent_setup.py index 39359d167..47ea79645 100644 --- a/tests/test_agent_setup.py +++ b/tests/test_agent_setup.py @@ -11,7 +11,10 @@ from benchflow.agents.install import apply_web_tool_policy, deploy_skills, install_agent from benchflow.agents.registry import AGENTS, AgentConfig from benchflow.models import AgentInstallError -from benchflow.rollout._setup import _agent_process_kill_pattern +from benchflow.rollout._setup import ( + _agent_process_kill_pattern, + _agent_process_termination_command, +) class LocalShellEnv: @@ -879,3 +882,28 @@ def test_agent_kill_pattern_targets_agent_not_python_services(launch, agent_argv def test_agent_kill_pattern_empty_launch_is_none(): assert _agent_process_kill_pattern("") is None assert _agent_process_kill_pattern(" ") is None + + +def test_agent_termination_command_tracks_descendants_before_killing_wrapper(): + """Guards PR #1057 mixed auth against orphaned native-agent children.""" + + command = _agent_process_termination_command( + "/opt/benchflow/bin/claude-agent-acp", "agent" + ) + + assert command is not None + assert "id -u -- agent" in command + assert 'pgrep -u "$bf_agent_uid" -f' in command + assert 'pgrep -u "$bf_agent_uid" -P "$bf_pid"' in command + assert 'kill -STOP "$bf_pid"' in command + assert "kill -TERM $bf_pids" in command + assert "kill -KILL $bf_alive" in command + assert 'while [ "$bf_wait" -lt 20 ]' in command + assert "pkill -u" not in command + + +def test_agent_termination_command_preserves_legacy_root_scope(): + command = _agent_process_termination_command("codex-acp", None) + + assert command is not None + assert command.startswith("pkill -f ") diff --git a/tests/test_session_factory_runtime.py b/tests/test_session_factory_runtime.py index 7ea04d273..ce7268d7e 100644 --- a/tests/test_session_factory_runtime.py +++ b/tests/test_session_factory_runtime.py @@ -376,8 +376,11 @@ async def exec(self, command: str, **kwargs): assert len(env.calls) == 1 command, kwargs = env.calls[0] assert "id -u -- agent" in command - assert 'pkill -u "$bf_agent_uid" -f' in command assert 'pgrep -u "$bf_agent_uid" -f' in command + assert 'pgrep -u "$bf_agent_uid" -P "$bf_pid"' in command + assert 'kill -STOP "$bf_pid"' in command + assert "kill -TERM $bf_pids" in command + assert "kill -KILL $bf_alive" in command assert 'while [ "$bf_wait" -lt 20 ]' in command assert "sleep 0.1" in command assert kwargs == {"timeout_sec": 10} diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 3e71739a6..0991dc49e 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -22,7 +22,10 @@ assemble_capture, load_provider_wire_records, ) -from benchflow.trajectories.native_capture_collection import NativeCollection +from benchflow.trajectories.native_capture_collection import ( + ClaudeOtelCollector, + NativeCollection, +) from benchflow.trajectories.native_capture_parsers import ( parse_codex_sessions, project_acp_trajectory, @@ -31,6 +34,36 @@ STARTED_AT = datetime(2026, 8, 28, 12, 0, tzinfo=UTC) +@pytest.mark.asyncio +async def test_claude_otel_collector_is_outside_agent_process_custody() -> None: + """Guards PR #1057 mixed auth against its sink blocking the next role.""" + + calls: list[tuple[str, dict]] = [] + + class RecordingEnv: + async def exec(self, command, **kwargs): + calls.append((command, kwargs)) + if "nohup" in command: + return SimpleNamespace(return_code=0, stdout="43123\n", stderr="") + return SimpleNamespace(return_code=0, stdout="", stderr="") + + async def upload_file(self, *_args, **_kwargs): + return None + + collector = ClaudeOtelCollector("/tmp/benchflow-capture-test") + port = await collector.ensure(RecordingEnv(), sandbox_user="agent") + + assert port == 43123 + setup_command, setup_kwargs = calls[1] + assert "chown agent /tmp/benchflow-capture-test/raw" in setup_command + assert "chown root:root /tmp/benchflow-capture-test" in setup_command + assert "chmod 711 /tmp/benchflow-capture-test" in setup_command + assert setup_kwargs["user"] == "root" + launch_command, launch_kwargs = calls[2] + assert "nohup" in launch_command + assert launch_kwargs["user"] == "root" + + @pytest.mark.asyncio async def test_zero_call_oauth_target_preserves_clean_audit_completion( tmp_path: Path, From 5646fe83836870a945bd4320650b57412ee769d5 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sun, 30 Aug 2026 00:00:06 -0700 Subject: [PATCH 71/74] fix: restore OAuth credentials on role reconnect --- src/benchflow/rollout/__init__.py | 16 ++++++++++------ tests/test_connect_as_env.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 68aa05f17..2da2917d3 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -2441,9 +2441,17 @@ async def connect_as(self, role: Role) -> None: ) self._adopt_usage_runtime(previous_usage_runtime, next_usage_runtime) + native_subscription = uses_native_subscription_auth( + role.agent, + role.model, + agent_env, + ) role_agent_differs = role.agent != cfg.primary_agent needs_role_credentials = ( - role_agent_differs or role.model != cfg.primary_model or bool(role.env) + role_agent_differs + or role.model != cfg.primary_model + or bool(role.env) + or native_subscription ) cred_home = _sandbox_user_home(cfg.sandbox_user) if role_agent_differs: @@ -2555,11 +2563,7 @@ async def connect_as(self, role: Role) -> None: credential_home=cred_home, role_name=role.name, ) - self._current_usage_is_native_subscription = uses_native_subscription_auth( - role.agent, - role.model, - agent_env, - ) + self._current_usage_is_native_subscription = native_subscription self._native_usage_checkpoint = None self._reapply_ask_user_handler() self._attach_trajectory_writer(rollout_dir) diff --git a/tests/test_connect_as_env.py b/tests/test_connect_as_env.py index 084240063..115374857 100644 --- a/tests/test_connect_as_env.py +++ b/tests/test_connect_as_env.py @@ -180,3 +180,34 @@ async def test_same_agent_different_model_refreshes_credentials(self, _mock_tria args, _kwargs = _mock_trial._planes.write_credential_files.await_args assert args[1] == "claude-agent-acp" assert args[4] == "other-model" + + @pytest.mark.asyncio + async def test_same_primary_oauth_role_reuploads_cleaned_credentials( + self, _mock_trial + ): + """Guards PR #1057 review r3888704109 for OAuth -> API -> OAuth.""" + + role = Role( + name="primary", + agent="claude-agent-acp", + model="claude-sonnet-4-6", + ) + _mock_trial._config.scenes[0].roles = [role] + native_env = {"_BENCHFLOW_SUBSCRIPTION_AUTH": "1"} + _mock_trial._planes.resolve_agent_env.side_effect = None + _mock_trial._planes.resolve_agent_env.return_value = dict(native_env) + _mock_trial._planes.ensure_litellm_runtime.return_value = ( + dict(native_env), + None, + ) + + await _mock_trial.connect_as(role) + + _mock_trial._planes.install_agent.assert_not_awaited() + _mock_trial._planes.write_credential_files.assert_awaited_once() + _mock_trial._planes.upload_subscription_auth.assert_awaited_once_with( + _mock_trial._env, + "claude-agent-acp", + "/home/agent", + ) + _mock_trial._planes.apply_web_tool_policy.assert_awaited_once() From bd768e05d1cda6d345ad678f7074f2fc57280c1a Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sun, 30 Aug 2026 00:27:09 -0700 Subject: [PATCH 72/74] fix: close proxy and continuation provenance gaps --- src/benchflow/acp/runtime.py | 18 +++- src/benchflow/agents/opencode_config.py | 25 ++++- src/benchflow/continue_run/orchestrator.py | 39 ++++++- tests/continue_run/test_training_metadata.py | 9 +- tests/test_opencode_family_proxy_tracking.py | 101 +++++++++++++++++++ 5 files changed, 184 insertions(+), 8 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 57efcc144..0a65a291e 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -81,7 +81,23 @@ def _harden_proxy_agent_launch( if not agent_env.get("BENCHFLOW_LITELLM_MODEL_ALIAS"): return agent_launch if agent == "opencode": - return f"{opencode_provider_reset_command()} && {agent_launch}" + # OpenCode merges several global filenames, ~/.opencode, project + # configs, an arbitrary file/directory, and inline JSON. Proxy mode + # must expose exactly the manifest-owned config that the wrapper + # registers below. Pin XDG to the same canonical agent home, disable + # project discovery, and remove every supported alternate env source + # before the process imports its configuration flags. + isolate_sources = ( + "unset OPENCODE_CONFIG OPENCODE_CONFIG_DIR " + "OPENCODE_CONFIG_CONTENT XDG_CONFIG_HOME && " + 'export HOME="${BENCHFLOW_AGENT_HOME:-$HOME}" && ' + 'export XDG_CONFIG_HOME="$HOME/.config" && ' + "export OPENCODE_DISABLE_PROJECT_CONFIG=1" + ) + return ( + f"{isolate_sources} && {opencode_provider_reset_command()} " + f"&& {agent_launch}" + ) if agent == "mimo": # MiMo's manifest-owned launcher replaces both canonical config files # in proxy mode, but the CLI also honors this arbitrary alternate path. diff --git a/src/benchflow/agents/opencode_config.py b/src/benchflow/agents/opencode_config.py index 06cc8e376..d6d679792 100644 --- a/src/benchflow/agents/opencode_config.py +++ b/src/benchflow/agents/opencode_config.py @@ -9,7 +9,7 @@ def opencode_provider_reset_source() -> str: - """Return stdlib Node.js that removes every pre-existing provider.""" + """Return stdlib Node.js that isolates OpenCode's file config sources.""" return "\n".join( [ @@ -17,8 +17,27 @@ def opencode_provider_reset_source() -> str: 'const os = require("os");', 'const path = require("path");', 'const home = (process.env.BENCHFLOW_AGENT_HOME || "").trim() || os.homedir();', + 'const globalDir = path.join(home, ".config", "opencode");', f"const p = path.join(home, {OPENCODE_CONFIG_RELATIVE_PATH!r});", - "fs.mkdirSync(path.dirname(p), { recursive: true });", + "fs.mkdirSync(globalDir, { recursive: true });", + "function removeConfigFile(target) {", + " try {", + " const stat = fs.lstatSync(target);", + " if (stat.isDirectory()) throw new Error(`config source is a directory: ${target}`);", + " fs.unlinkSync(target);", + " } catch (error) {", + ' if (error.code !== "ENOENT") throw error;', + " }", + "}", + # OpenCode merges all three global JSON names plus a legacy TOML + # file. Keep only the manifest-owned opencode.json boundary. + 'for (const name of ["config", "config.json", "opencode.jsonc"]) {', + " removeConfigFile(path.join(globalDir, name));", + "}", + # It also loads ~/.opencode independently of project-config flags. + 'for (const name of ["opencode.json", "opencode.jsonc"]) {', + ' removeConfigFile(path.join(home, ".opencode", name));', + "}", 'const text = fs.existsSync(p) ? fs.readFileSync(p, "utf8").trim() : "";', "const d = text ? JSON.parse(text) : {};", # The manifest-owned wrapper adds the one BenchFlow gateway provider @@ -33,6 +52,6 @@ def opencode_provider_reset_source() -> str: def opencode_provider_reset_command() -> str: - """Return the shell-safe reset command run immediately before OpenCode.""" + """Return the shell-safe file reset run immediately before OpenCode.""" return f"{_OPENCODE_NODE} -e {shlex.quote(opencode_provider_reset_source())}" diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 08f7f8784..cd9bfb934 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -44,6 +44,10 @@ from benchflow.contracts import AgentProtocolError, SandboxStartupFailure from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.scenes import compile_scenes_to_steps +from benchflow.trajectories.llm_capture_manifest import ( + LLMTrajectoryManifest, + read_llm_trajectory_manifest, +) from benchflow.trajectories.types import LLMExchange logger = logging.getLogger(__name__) @@ -332,10 +336,11 @@ def update_continued_metadata( HF-compatible artifacts still need the actual live model and token usage. The stitched LLM trajectory is authoritative for provider usage. """ + attributed_model = _continued_attributed_model(rollout_dir, fallback=live_model) config_path = rollout_dir / "config.json" if config_path.is_file(): config = json.loads(config_path.read_text()) - config["model"] = live_model + config["model"] = attributed_model config.setdefault("source", {})["usage_source"] = "stitched_llm_trajectory" config["usage_tracking"] = { "requested": "required", @@ -356,7 +361,7 @@ def update_continued_metadata( if not result_path.is_file(): return result = json.loads(result_path.read_text()) - result["model"] = live_model + result["model"] = attributed_model agent_result = result.setdefault("agent_result", {}) if isinstance(agent_result, dict): agent_result.update(usage.as_agent_result_patch()) @@ -387,6 +392,36 @@ def update_continued_metadata( refresh_rollout_results_jsonl(rollout_dir) +def _continued_attributed_model( + rollout_dir: Path, *, fallback: str | None +) -> str | None: + """Return the model represented by finalized continuation exchanges. + + Legacy/direct callers without a valid manifest retain the requested-model + fallback. Once continuation finalization has written a valid manifest, its + active role captures are authoritative: a replay-only run must not claim a + requested live model that never produced an exchange, and a mixed-model + run must not collapse to either model. + """ + + raw = read_llm_trajectory_manifest(rollout_dir) + if raw is None: + return fallback + try: + manifest = LLMTrajectoryManifest.model_validate(raw) + except ValueError: + return fallback + + active_models = { + capture.model + for capture in manifest.role_captures + if capture.exchange_count > 0 + } + if active_models: + return next(iter(active_models)) if len(active_models) == 1 else None + return manifest.model + + def _host_proxy_binding(environment: str) -> tuple[str, str]: """(bind_host, advertise_host) so a Docker agent can reach the host proxy. diff --git a/tests/continue_run/test_training_metadata.py b/tests/continue_run/test_training_metadata.py index 1fa57044b..33b045b15 100644 --- a/tests/continue_run/test_training_metadata.py +++ b/tests/continue_run/test_training_metadata.py @@ -23,7 +23,7 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): - """Guards PR #1057 against retaining stale or incomplete continuation rows.""" + """Guards PR #1057 review r3888738115 and stale continuation rows.""" rollout = tmp_path / "job" / "demo-task__continued" (rollout / "trajectory").mkdir(parents=True) model = "openai/gpt-5.5" @@ -99,15 +99,20 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): json.dumps({"info": {"training_ready": False, "model": None}}) + "\n" ) + requested_but_unused_model = "openai/gpt-5.6" update_continued_metadata( rollout, - live_model=model, + live_model=requested_but_unused_model, usage=summarize_llm_trajectory_usage(trajectory_path, n_recorded=0), environment="docker", ) + refreshed_config = json.loads((rollout / "config.json").read_text()) + refreshed_result = json.loads((rollout / "result.json").read_text()) refreshed = json.loads((rollout / "results.jsonl").read_text()) aggregated = json.loads((rollout.parent / "results.jsonl").read_text()) + assert refreshed_config["model"] == model + assert refreshed_result["model"] == model assert refreshed["info"]["model"] == model assert refreshed["info"]["training_ready"] is True assert refreshed["token_usage"]["total_tokens"] == 2 diff --git a/tests/test_opencode_family_proxy_tracking.py b/tests/test_opencode_family_proxy_tracking.py index 675f077b6..ba74377ac 100644 --- a/tests/test_opencode_family_proxy_tracking.py +++ b/tests/test_opencode_family_proxy_tracking.py @@ -15,6 +15,7 @@ import json import os import re +import shutil import subprocess import pytest @@ -133,10 +134,110 @@ def test_proxy_launch_resets_providers_immediately_before_manifest_wrapper(): assert "opencode.json" in hardened assert "d.provider = {}" in hardened + assert ( + "unset OPENCODE_CONFIG OPENCODE_CONFIG_DIR " + "OPENCODE_CONFIG_CONTENT XDG_CONFIG_HOME" + ) in hardened + for alternate in ( + "OPENCODE_CONFIG", + "OPENCODE_CONFIG_DIR", + "OPENCODE_CONFIG_CONTENT", + "XDG_CONFIG_HOME", + ): + assert alternate in hardened + assert "OPENCODE_DISABLE_PROJECT_CONFIG=1" in hardened assert hardened.endswith("&& " + launch) assert _harden_proxy_agent_launch("opencode", launch, {}) == launch +def test_proxy_launch_isolates_every_effective_opencode_config_source(tmp_path): + """Guards PR #1057 review r3888738113 against alternate config bypasses.""" + + home = tmp_path / "home" + global_dir = home / ".config" / "opencode" + dot_dir = home / ".opencode" + project_dir = tmp_path / "project" + alternate_dir = tmp_path / "alternate-dir" + for directory in (global_dir, dot_dir, project_dir, alternate_dir): + directory.mkdir(parents=True) + + bypass = json.dumps( + { + "provider": { + "bypass": { + "options": { + "apiKey": "literal-bypass-key", + "baseURL": "https://bypass.invalid/v1", + } + } + } + } + ) + canonical = global_dir / "opencode.json" + canonical.write_text( + json.dumps({"provider": json.loads(bypass)["provider"], "theme": "dark"}) + ) + for path in ( + global_dir / "config", + global_dir / "config.json", + global_dir / "opencode.jsonc", + dot_dir / "opencode.json", + dot_dir / "opencode.jsonc", + ): + path.write_text(bypass) + alternate_file = tmp_path / "alternate.json" + alternate_file.write_text(bypass) + (alternate_dir / "opencode.json").write_text(bypass) + (project_dir / "opencode.json").write_text(bypass) + + launch = ( + 'test -z "${OPENCODE_CONFIG:-}" && ' + 'test -z "${OPENCODE_CONFIG_DIR:-}" && ' + 'test -z "${OPENCODE_CONFIG_CONTENT:-}" && ' + 'test "$OPENCODE_DISABLE_PROJECT_CONFIG" = 1 && ' + 'test "$HOME" = "$BENCHFLOW_AGENT_HOME" && ' + 'test "$XDG_CONFIG_HOME" = "$HOME/.config"' + ) + hardened = _harden_proxy_agent_launch( + "opencode", + launch, + {"BENCHFLOW_LITELLM_MODEL_ALIAS": "benchflow-provider-model"}, + ) + node = shutil.which("node") + assert node is not None + executable_hardened = hardened.replace("/opt/benchflow/node/bin/node", node, 1) + subprocess.run( + ["sh", "-c", executable_hardened], + cwd=project_dir, + env={ + **os.environ, + "HOME": str(tmp_path / "alternate-home"), + "BENCHFLOW_AGENT_HOME": str(home), + "OPENCODE_CONFIG": str(alternate_file), + "OPENCODE_CONFIG_DIR": str(alternate_dir), + "OPENCODE_CONFIG_CONTENT": bypass, + "XDG_CONFIG_HOME": str(tmp_path / "alternate-xdg"), + }, + check=True, + timeout=15, + ) + + sanitized = json.loads(canonical.read_text()) + assert sanitized == {"provider": {}, "theme": "dark"} + for removed in ( + global_dir / "config", + global_dir / "config.json", + global_dir / "opencode.jsonc", + dot_dir / "opencode.json", + dot_dir / "opencode.jsonc", + ): + assert not removed.exists() + # Alternate/project sources are not destructively edited; the launch + # boundary makes them unreachable in proxy mode. + assert alternate_file.read_text() == bypass + assert (project_dir / "opencode.json").read_text() == bypass + + def test_proxy_launch_unsets_mimo_alternate_config_before_manifest_launcher( tmp_path, ): From b45e2aae5ce713479217ec16b60bc8ce6d8c88d7 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sun, 30 Aug 2026 00:44:41 -0700 Subject: [PATCH 73/74] fix: isolate all OpenCode config authorities --- src/benchflow/acp/runtime.py | 22 ++++++--- src/benchflow/agents/opencode_config.py | 8 +++- tests/test_opencode_family_proxy_tracking.py | 47 +++++++++++++++++++- 3 files changed, 67 insertions(+), 10 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 0a65a291e..562bf6057 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -81,17 +81,25 @@ def _harden_proxy_agent_launch( if not agent_env.get("BENCHFLOW_LITELLM_MODEL_ALIAS"): return agent_launch if agent == "opencode": - # OpenCode merges several global filenames, ~/.opencode, project - # configs, an arbitrary file/directory, and inline JSON. Proxy mode - # must expose exactly the manifest-owned config that the wrapper - # registers below. Pin XDG to the same canonical agent home, disable - # project discovery, and remove every supported alternate env source - # before the process imports its configuration flags. + # OpenCode merges several global filenames, ~/.opencode, project and + # system-managed configs, arbitrary env file/directory/content sources, + # remote configs from auth/account state, and test-only path overrides. + # Proxy mode must expose exactly the manifest-owned config that the + # wrapper registers below. Pin every path to the canonical agent home, + # replace auth/account state with empty in-memory state, disable project + # discovery, and remove every alternate source before module import. isolate_sources = ( "unset OPENCODE_CONFIG OPENCODE_CONFIG_DIR " - "OPENCODE_CONFIG_CONTENT XDG_CONFIG_HOME && " + "OPENCODE_CONFIG_CONTENT OPENCODE_AUTH_CONTENT OPENCODE_DB " + "OPENCODE_TEST_HOME OPENCODE_TEST_MANAGED_CONFIG_DIR " + "XDG_CONFIG_HOME && " 'export HOME="${BENCHFLOW_AGENT_HOME:-$HOME}" && ' 'export XDG_CONFIG_HOME="$HOME/.config" && ' + 'export OPENCODE_AUTH_CONTENT="{}" && ' + 'export OPENCODE_DB=":memory:" && ' + 'export OPENCODE_TEST_HOME="$HOME" && ' + "export OPENCODE_TEST_MANAGED_CONFIG_DIR=" + '"$HOME/.config/opencode/benchflow-managed-disabled" && ' "export OPENCODE_DISABLE_PROJECT_CONFIG=1" ) return ( diff --git a/src/benchflow/agents/opencode_config.py b/src/benchflow/agents/opencode_config.py index d6d679792..0bf8a183d 100644 --- a/src/benchflow/agents/opencode_config.py +++ b/src/benchflow/agents/opencode_config.py @@ -18,6 +18,9 @@ def opencode_provider_reset_source() -> str: 'const path = require("path");', 'const home = (process.env.BENCHFLOW_AGENT_HOME || "").trim() || os.homedir();', 'const globalDir = path.join(home, ".config", "opencode");', + "const managedDir = " + '(process.env.OPENCODE_TEST_MANAGED_CONFIG_DIR || "").trim() || ' + 'path.join(globalDir, "benchflow-managed-disabled");', f"const p = path.join(home, {OPENCODE_CONFIG_RELATIVE_PATH!r});", "fs.mkdirSync(globalDir, { recursive: true });", "function removeConfigFile(target) {", @@ -34,9 +37,12 @@ def opencode_provider_reset_source() -> str: 'for (const name of ["config", "config.json", "opencode.jsonc"]) {', " removeConfigFile(path.join(globalDir, name));", "}", - # It also loads ~/.opencode independently of project-config flags. + # It also loads ~/.opencode independently of project-config flags + # and a late system-managed directory. Runtime hardening redirects + # the latter to this agent-owned empty boundary. 'for (const name of ["opencode.json", "opencode.jsonc"]) {', ' removeConfigFile(path.join(home, ".opencode", name));', + " removeConfigFile(path.join(managedDir, name));", "}", 'const text = fs.existsSync(p) ? fs.readFileSync(p, "utf8").trim() : "";', "const d = text ? JSON.parse(text) : {};", diff --git a/tests/test_opencode_family_proxy_tracking.py b/tests/test_opencode_family_proxy_tracking.py index ba74377ac..323ae6697 100644 --- a/tests/test_opencode_family_proxy_tracking.py +++ b/tests/test_opencode_family_proxy_tracking.py @@ -136,12 +136,17 @@ def test_proxy_launch_resets_providers_immediately_before_manifest_wrapper(): assert "d.provider = {}" in hardened assert ( "unset OPENCODE_CONFIG OPENCODE_CONFIG_DIR " - "OPENCODE_CONFIG_CONTENT XDG_CONFIG_HOME" + "OPENCODE_CONFIG_CONTENT OPENCODE_AUTH_CONTENT OPENCODE_DB " + "OPENCODE_TEST_HOME OPENCODE_TEST_MANAGED_CONFIG_DIR XDG_CONFIG_HOME" ) in hardened for alternate in ( "OPENCODE_CONFIG", "OPENCODE_CONFIG_DIR", "OPENCODE_CONFIG_CONTENT", + "OPENCODE_AUTH_CONTENT", + "OPENCODE_DB", + "OPENCODE_TEST_HOME", + "OPENCODE_TEST_MANAGED_CONFIG_DIR", "XDG_CONFIG_HOME", ): assert alternate in hardened @@ -156,9 +161,18 @@ def test_proxy_launch_isolates_every_effective_opencode_config_source(tmp_path): home = tmp_path / "home" global_dir = home / ".config" / "opencode" dot_dir = home / ".opencode" + data_dir = home / ".local" / "share" / "opencode" + managed_dir = global_dir / "benchflow-managed-disabled" project_dir = tmp_path / "project" alternate_dir = tmp_path / "alternate-dir" - for directory in (global_dir, dot_dir, project_dir, alternate_dir): + for directory in ( + global_dir, + dot_dir, + data_dir, + managed_dir, + project_dir, + alternate_dir, + ): directory.mkdir(parents=True) bypass = json.dumps( @@ -183,17 +197,38 @@ def test_proxy_launch_isolates_every_effective_opencode_config_source(tmp_path): global_dir / "opencode.jsonc", dot_dir / "opencode.json", dot_dir / "opencode.jsonc", + managed_dir / "opencode.json", + managed_dir / "opencode.jsonc", ): path.write_text(bypass) + auth_file = data_dir / "auth.json" + auth_file.write_text( + json.dumps( + { + "https://bypass.invalid": { + "type": "wellknown", + "key": "BYPASS_TOKEN", + "token": "literal-bypass-key", + } + } + ) + ) alternate_file = tmp_path / "alternate.json" alternate_file.write_text(bypass) (alternate_dir / "opencode.json").write_text(bypass) (project_dir / "opencode.json").write_text(bypass) + alternate_database = tmp_path / "alternate-opencode.db" + alternate_database.write_text("bypass-account-database") launch = ( 'test -z "${OPENCODE_CONFIG:-}" && ' 'test -z "${OPENCODE_CONFIG_DIR:-}" && ' 'test -z "${OPENCODE_CONFIG_CONTENT:-}" && ' + 'test "$OPENCODE_AUTH_CONTENT" = "{}" && ' + 'test "$OPENCODE_DB" = ":memory:" && ' + 'test "$OPENCODE_TEST_HOME" = "$HOME" && ' + 'test "$OPENCODE_TEST_MANAGED_CONFIG_DIR" = ' + '"$HOME/.config/opencode/benchflow-managed-disabled" && ' 'test "$OPENCODE_DISABLE_PROJECT_CONFIG" = 1 && ' 'test "$HOME" = "$BENCHFLOW_AGENT_HOME" && ' 'test "$XDG_CONFIG_HOME" = "$HOME/.config"' @@ -216,6 +251,10 @@ def test_proxy_launch_isolates_every_effective_opencode_config_source(tmp_path): "OPENCODE_CONFIG": str(alternate_file), "OPENCODE_CONFIG_DIR": str(alternate_dir), "OPENCODE_CONFIG_CONTENT": bypass, + "OPENCODE_AUTH_CONTENT": auth_file.read_text(), + "OPENCODE_DB": str(alternate_database), + "OPENCODE_TEST_HOME": str(tmp_path / "alternate-test-home"), + "OPENCODE_TEST_MANAGED_CONFIG_DIR": str(tmp_path / "alternate-managed"), "XDG_CONFIG_HOME": str(tmp_path / "alternate-xdg"), }, check=True, @@ -230,12 +269,16 @@ def test_proxy_launch_isolates_every_effective_opencode_config_source(tmp_path): global_dir / "opencode.jsonc", dot_dir / "opencode.json", dot_dir / "opencode.jsonc", + managed_dir / "opencode.json", + managed_dir / "opencode.jsonc", ): assert not removed.exists() # Alternate/project sources are not destructively edited; the launch # boundary makes them unreachable in proxy mode. assert alternate_file.read_text() == bypass assert (project_dir / "opencode.json").read_text() == bypass + assert auth_file.read_text().find("literal-bypass-key") >= 0 + assert alternate_database.read_text() == "bypass-account-database" def test_proxy_launch_unsets_mimo_alternate_config_before_manifest_launcher( From 272557493b0f056502b1a90bd8bbfeb749ba9316 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Sun, 30 Aug 2026 00:57:27 -0700 Subject: [PATCH 74/74] fix: validate audit artifacts before completion --- .../trajectories/llm_capture_manifest.py | 16 ++-- src/benchflow/trajectories/results.py | 19 +++- tests/continue_run/test_training_metadata.py | 9 ++ .../test_llm_capture_training_contract.py | 87 +++++++++++++++++++ .../test_native_capture_resilience.py | 19 ++++ 5 files changed, 140 insertions(+), 10 deletions(-) diff --git a/src/benchflow/trajectories/llm_capture_manifest.py b/src/benchflow/trajectories/llm_capture_manifest.py index adf4cab3b..18d35941b 100644 --- a/src/benchflow/trajectories/llm_capture_manifest.py +++ b/src/benchflow/trajectories/llm_capture_manifest.py @@ -222,7 +222,7 @@ def capture_artifact_allows_training( _exchange_matches_training_manifest(exchange, manifest) for exchange in exchanges ) - and _exchanges_match_training_roles(exchanges, manifest) + and _exchanges_match_roles(exchanges, manifest) and successful_exchanges_have_positive_usage(exchanges) ) return not any(_exchange_requires_manifest(exchange) for exchange in exchanges) @@ -361,7 +361,7 @@ def _exchange_matches_training_manifest( return True -def _exchanges_match_training_roles( +def _exchanges_match_roles( exchanges: Sequence[dict[str, Any]], manifest: dict[str, Any] ) -> bool: """Require row attribution and cardinality to match per-role provenance.""" @@ -405,7 +405,11 @@ def _exchanges_match_training_roles( return actual == expected -def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> bool: +def capture_manifest_preserves_audit_completion( + manifest: dict[str, Any], + *, + exchanges: Sequence[dict[str, Any]] | None = None, +) -> bool: """Accept only expected, internally complete audit-only capture states.""" if manifest.get("status") not in { @@ -426,9 +430,9 @@ def capture_manifest_preserves_audit_completion(manifest: dict[str, Any]) -> boo return False missing_fields = set(raw_missing_fields) role_captures = _validated_role_captures(manifest) - if role_captures is None or not _role_captures_match_manifest( - manifest, role_captures - ): + if not role_captures or not _role_captures_match_manifest(manifest, role_captures): + return False + if exchanges is not None and not _exchanges_match_roles(exchanges, manifest): return False has_oauth_capture = any( _is_oauth_audit_role_capture(capture) for capture in role_captures diff --git a/src/benchflow/trajectories/results.py b/src/benchflow/trajectories/results.py index f05c704f2..9b3472dc1 100644 --- a/src/benchflow/trajectories/results.py +++ b/src/benchflow/trajectories/results.py @@ -20,7 +20,7 @@ import json import logging from contextlib import suppress -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any, cast @@ -53,6 +53,7 @@ class _LLMStepsResult: tool_defs: list[dict[str, Any]] export_error: str | None capture_contract_rejected: bool = False + exchanges: list[dict[str, Any]] = field(default_factory=list) def _record_to_redacted_json_line(record: dict[str, Any]) -> str: @@ -185,7 +186,13 @@ def _llm_steps_from_trajectory( rollout_dir, exchanges=exchanges, ): - return _LLMStepsResult([], [], None, capture_contract_rejected=True) + return _LLMStepsResult( + [], + [], + None, + capture_contract_rejected=True, + exchanges=exchanges, + ) training_success_indices = _training_success_exchange_indices(exchanges) skipped_successful: list[str] = [] for exchange_idx, exchange in enumerate(exchanges): @@ -263,8 +270,9 @@ def _llm_steps_from_trajectory( tool_defs, "Successful LLM exchanges were omitted from results.jsonl: " + "; ".join(skipped_successful), + exchanges=exchanges, ) - return _LLMStepsResult(steps, tool_defs, None) + return _LLMStepsResult(steps, tool_defs, None, exchanges=exchanges) def _response_is_truncated(response_body: dict[str, Any]) -> bool: @@ -513,7 +521,10 @@ def build_rollout_results_record( ) or ( capture_manifest is not None - and capture_manifest_preserves_audit_completion(capture_manifest) + and capture_manifest_preserves_audit_completion( + capture_manifest, + exchanges=llm_steps.exchanges, + ) ) ) and effective_export_error is None diff --git a/tests/continue_run/test_training_metadata.py b/tests/continue_run/test_training_metadata.py index 33b045b15..a08ddd2e1 100644 --- a/tests/continue_run/test_training_metadata.py +++ b/tests/continue_run/test_training_metadata.py @@ -144,6 +144,15 @@ def test_update_continued_metadata_rebuilds_trainer_results(tmp_path): } ) write_llm_trajectory_manifest(rollout, replay_manifest) + replay_row = json.loads(trajectory_path.read_text()) + replay_row["metadata"].update( + { + "capture_source": "replay_proxy", + "capture_fidelity": "agent_session", + "request_complete": False, + } + ) + trajectory_path.write_text(json.dumps(replay_row) + "\n") (rollout / "results.jsonl").write_text( json.dumps({"info": {"training_ready": True}, "is_completed": False}) + "\n" ) diff --git a/tests/trajectories/test_llm_capture_training_contract.py b/tests/trajectories/test_llm_capture_training_contract.py index 23314d1d2..bb6721e06 100644 --- a/tests/trajectories/test_llm_capture_training_contract.py +++ b/tests/trajectories/test_llm_capture_training_contract.py @@ -85,6 +85,17 @@ def test_agent_session_capture_is_audit_only_not_training_ready(tmp_path: Path) trajectory_dir = tmp_path / "trajectory" trajectory_dir.mkdir() _write_exchange(trajectory_dir, fidelity="agent_session") + exchange_path = trajectory_dir / "llm_trajectory.jsonl" + exchange = json.loads(exchange_path.read_text()) + exchange["metadata"].update( + { + "capture_source": "codex_native_session", + "auth_mode": "oauth_subscription", + "role": "agent", + "agent": "codex-acp", + } + ) + exchange_path.write_text(json.dumps(exchange) + "\n") (trajectory_dir / "llm_trajectory.manifest.json").write_text( json.dumps( { @@ -122,6 +133,67 @@ def test_agent_session_capture_is_audit_only_not_training_ready(tmp_path: Path) assert row["error"] is None +@pytest.mark.parametrize("corruption", ["truncated", "wrong_role"]) +def test_corrupt_audit_rows_fail_closed_for_completion( + tmp_path: Path, corruption: str +) -> None: + """Guards PR #1057 review r3888793594 against stale audit sidecars.""" + + trajectory_dir = tmp_path / "trajectory" + trajectory_dir.mkdir() + _write_exchange(trajectory_dir, fidelity="agent_session") + exchange_path = trajectory_dir / "llm_trajectory.jsonl" + exchange = json.loads(exchange_path.read_text()) + exchange["metadata"].update( + { + "capture_source": "codex_native_session", + "auth_mode": "oauth_subscription", + "role": "agent", + "agent": "codex-acp", + } + ) + if corruption == "truncated": + exchange_path.write_text("") + else: + exchange["metadata"]["role"] = "unprepared-role" + exchange_path.write_text(json.dumps(exchange) + "\n") + (trajectory_dir / "llm_trajectory.manifest.json").write_text( + json.dumps( + { + "status": "partial", + "capture_source": "codex_native_session", + "capture_fidelity": "agent_session", + "auth_mode": "oauth_subscription", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + "role_captures": [ + { + "role": "agent", + "agent": "codex-acp", + "auth_mode": "oauth_subscription", + "capture_source": "codex_native_session", + "capture_fidelity": "agent_session", + "exchange_count": 1, + "request_complete": False, + "response_complete": True, + } + ], + } + ) + ) + + row = _build_results_row( + tmp_path, + agent_result={"usage_source": "agent_native_acp", "total_tokens": 2}, + ) + + assert row["info"]["training_ready"] is False + assert row["info"]["training_ready_reason"] == "insufficient_capture_fidelity" + assert row["is_completed"] is False + assert row["error"]["error"] == "missing_llm_trajectory" + + def test_corrupt_capture_manifest_fails_closed_for_training(tmp_path: Path) -> None: """Guards PR #1057 against treating a corrupt new sidecar as a legacy artifact.""" @@ -447,12 +519,24 @@ def test_mixed_oauth_audit_capture_preserves_successful_completion( trajectory_dir.mkdir() _write_exchange(trajectory_dir, fidelity="provider_wire") provider_row = json.loads((trajectory_dir / "llm_trajectory.jsonl").read_text()) + provider_row["metadata"].update( + { + "capture_source": "litellm_proxy", + "auth_mode": "api_key", + "role": "coder", + "agent": "codex-acp", + } + ) oauth_row = json.loads(json.dumps(provider_row)) oauth_row["metadata"].update( { "capture_fidelity": "agent_session", + "capture_source": "claude_native_session", "auth_mode": "oauth_subscription", "request_complete": False, + "response_complete": False, + "role": "reviewer", + "agent": "claude-agent-acp", } ) (trajectory_dir / "llm_trajectory.jsonl").write_text( @@ -571,6 +655,9 @@ def test_mixed_replay_capture_preserves_successful_completion(tmp_path: Path) -> "request_capture_source": "replay_proxy_ingress", "auth_mode": "api_key", "request_complete": False, + "role": "agent", + "agent": "openhands", + "model": "openai/gpt-5.6", } ) (trajectory_dir / "llm_trajectory.jsonl").write_text(json.dumps(replay_row) + "\n") diff --git a/tests/trajectories/test_native_capture_resilience.py b/tests/trajectories/test_native_capture_resilience.py index 0991dc49e..ad877d8dc 100644 --- a/tests/trajectories/test_native_capture_resilience.py +++ b/tests/trajectories/test_native_capture_resilience.py @@ -119,6 +119,25 @@ async def collect_native(_env, *, targets): assert capture_manifest_preserves_audit_completion(manifest) is True +def test_zero_call_oauth_manifest_requires_a_prepared_role() -> None: + """Guards PR #1057 review r3888793598 against vacuous zero-call evidence.""" + + manifest = { + "status": "no_model_call", + "capture_source": "none", + "capture_fidelity": "none", + "auth_mode": "oauth_subscription", + "exchange_count": 0, + "request_complete": False, + "response_complete": False, + "errors": [], + "missing_fields": [], + "role_captures": [], + } + + assert capture_manifest_preserves_audit_completion(manifest, exchanges=[]) is False + + def test_native_parser_normalizes_fallback_and_record_timestamps( tmp_path: Path, ) -> None: