From 9410d3f1ac4fa5ea8b8365937884ec58b806f32b Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 02:41:09 +0800 Subject: [PATCH 1/6] Add turn training export core --- agent_base/training_export.py | 794 ++++++++++++++++++++++++++++++++++ 1 file changed, 794 insertions(+) create mode 100644 agent_base/training_export.py diff --git a/agent_base/training_export.py b/agent_base/training_export.py new file mode 100644 index 0000000..5bc7c35 --- /dev/null +++ b/agent_base/training_export.py @@ -0,0 +1,794 @@ +"""Immutable turn-level training export for single-agent ResearchHarness runs.""" + +from __future__ import annotations + +import ctypes +from dataclasses import dataclass +import hashlib +import json +import math +import os +from pathlib import Path, PurePosixPath +import re +import shutil +import tempfile +from typing import Any, Mapping, Sequence + + +UNIT_SCHEMA = "researchharness.turn-training-unit.v1" +SUMMARY_SCHEMA = "researchharness.rollout-summary.v1" +MANIFEST_SCHEMA = "researchharness.turn-training-export-manifest.v1" +EXPORT_DIR_NAME = "turn-training-export-v1" +CANONICALIZATION = "canonical-json-v1" +_IDENTITY_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:@+-]{0,255}") +_SHA256_RE = re.compile(r"[0-9a-f]{64}") +_PROHIBITED_AGENT_KEYS = { + "subagent", + "agent_id", + "parent_agent_id", + "parent_run_id", + "delegation_id", +} +_ORACLE_KEYS = { + "gold_patch", + "hidden_test", + "hidden_tests", + "verifier_internal", + "verifier_internals", + "resolved", + "reward", + "reference_answer", + "oracle", +} +_ORACLE_PATTERNS = ( + re.compile(r"\bgold[ _-]*patch\b", re.IGNORECASE), + re.compile(r"\bhidden[ _-]*(?:tests?|verifier)\b", re.IGNORECASE), + re.compile(r"\bverifier[ _-]*internals?\b", re.IGNORECASE), + re.compile(r"\bfinal[ _-]*(?:resolved|reward)\b", re.IGNORECASE), + re.compile(r"\bother[ _-]*rollout(?:'s)?[ _-]*(?:answer|output)\b", re.IGNORECASE), +) +_VALIDATOR_NAMES = ( + "TRACE_CAPTURE_CLOSURE", + "EXACT_REQUEST_MATCH", + "EXACT_RESPONSE_MATCH", + "TOOL_CALL_STRUCTURE", + "COMPACTION_CLOSURE", + "IDENTITY_JOIN", + "HASH_INTEGRITY", + "NO_MULTI_AGENT_V1", + "ORACLE_LEAK_SCAN", +) + + +@dataclass(frozen=True, slots=True) +class TrainingExportConfig: + output_dir: Path + task_id: str + task_revision_id: str + rollout_id: str + rollout_index: int + run_id: str + harness_revision: str + scaffold_version: str = "" + tool_protocol_version: str = "openai-chat-completions-native-tools-v1" + tokenizer_revision: str | None = None + final_patch_path: Path | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "output_dir", Path(self.output_dir)) + if self.output_dir.name != EXPORT_DIR_NAME: + raise ValueError(f"output_dir must end with {EXPORT_DIR_NAME!r}") + if self.final_patch_path is not None: + object.__setattr__(self, "final_patch_path", Path(self.final_patch_path)) + for label in ( + "task_id", + "task_revision_id", + "rollout_id", + "run_id", + "harness_revision", + ): + value = str(getattr(self, label)) + if _IDENTITY_RE.fullmatch(value) is None: + raise ValueError(f"{label} is not a safe, non-empty identity") + if isinstance(self.rollout_index, bool) or not isinstance(self.rollout_index, int) or self.rollout_index < 0: + raise ValueError("rollout_index must be a non-negative integer") + + def identity(self) -> dict[str, Any]: + return { + "task_id": self.task_id, + "task_revision_id": self.task_revision_id, + "rollout_id": self.rollout_id, + "rollout_index": self.rollout_index, + "run_id": self.run_id, + } + + +def canonical_json_bytes(value: Any) -> bytes: + _reject_nonfinite(value) + return json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + + +def _reject_nonfinite(value: Any) -> None: + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("canonical-json-v1 rejects NaN and Infinity") + if isinstance(value, Mapping): + for key, nested in value.items(): + if not isinstance(key, str): + raise ValueError("canonical-json-v1 object keys must be strings") + _reject_nonfinite(nested) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for nested in value: + _reject_nonfinite(nested) + + +def _hash_bytes(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def _hash_json(domain: str, value: Any, *, prefixed: bool = True) -> str: + digest = _hash_bytes( + canonical_json_bytes({"domain": domain, "payload": value}) + ) + return f"sha256:{digest}" if prefixed else digest + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _write_bytes(path: Path, payload: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("xb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + + +def _write_json(path: Path, value: Any) -> None: + _write_bytes(path, canonical_json_bytes(value) + b"\n") + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value + + +def _read_trace(path: Path) -> list[dict[str, Any]]: + _require_regular_source(path, label="trace") + rows: list[dict[str, Any]] = [] + for line_number, raw_line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not raw_line.strip(): + continue + value = json.loads(raw_line) + if not isinstance(value, dict): + raise ValueError(f"trace line {line_number} is not an object") + rows.append(value) + return rows + + +def _require_regular_source(path: Path, *, label: str) -> Path: + source = Path(path) + if source.is_symlink() or not source.is_file(): + raise ValueError(f"{label} must be an ordinary file") + return source + + +def _safe_ref_path(raw_path: object) -> PurePosixPath: + if not isinstance(raw_path, str) or not raw_path or "\\" in raw_path or "\x00" in raw_path: + raise ValueError("artifact ref path must be non-empty canonical POSIX text") + path = PurePosixPath(raw_path) + if ( + path.is_absolute() + or any(part in {"", ".", ".."} for part in path.parts) + or path.as_posix() != raw_path + ): + raise ValueError("artifact ref path escapes the export bundle") + return path + + +def _file_ref(root: Path, path: Path) -> dict[str, Any]: + relative = path.relative_to(root).as_posix() + _safe_ref_path(relative) + return {"path": relative, "sha256": _sha256_file(path)} + + +def _resolve_ref(root: Path, value: object) -> Path: + if not isinstance(value, Mapping) or set(value) != {"path", "sha256"}: + raise ValueError("artifact ref must contain exactly path and sha256") + relative = _safe_ref_path(value.get("path")) + digest = value.get("sha256") + if not isinstance(digest, str) or _SHA256_RE.fullmatch(digest) is None: + raise ValueError("artifact ref sha256 is invalid") + path = root.joinpath(*relative.parts) + if path.is_symlink() or not path.is_file() or path.stat().st_nlink != 1: + raise ValueError("artifact ref target must be one ordinary, non-hardlinked file") + if _sha256_file(path) != digest: + raise ValueError("artifact ref digest mismatch") + return path + + +def _artifact_json(root: Path, relative: str, value: Any) -> dict[str, Any]: + path = root / relative + _write_json(path, value) + return _file_ref(root, path) + + +def _oracle_findings(value: Any, *, location: str = "request") -> list[str]: + findings: list[str] = [] + if isinstance(value, Mapping): + for key, nested in value.items(): + if str(key).casefold() in _ORACLE_KEYS: + findings.append(f"{location}.{key}:forbidden-key") + findings.extend(_oracle_findings(str(key), location=f"{location}.{key}")) + findings.extend(_oracle_findings(nested, location=f"{location}.{key}")) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, nested in enumerate(value): + findings.extend(_oracle_findings(nested, location=f"{location}[{index}]")) + elif isinstance(value, str): + for pattern in _ORACLE_PATTERNS: + if pattern.search(value): + findings.append(f"{location}:{pattern.pattern}") + return sorted(set(findings)) + + +def _provider_request(payload: Mapping[str, Any], *, compaction: bool = False) -> dict[str, Any]: + if compaction: + response = payload.get("summary_response") + if isinstance(response, Mapping) and isinstance(response.get("provider_request"), Mapping): + return dict(response["provider_request"]) + return {} + if isinstance(payload.get("provider_request"), Mapping): + return dict(payload["provider_request"]) + return {} + + +def _normalized_request(provider_request: Mapping[str, Any]) -> dict[str, Any]: + sampling: dict[str, Any] = {} + for source, target in ( + ("temperature", "temperature"), + ("top_p", "top_p"), + ("presence_penalty", "presence_penalty"), + ("max_tokens", "max_output_tokens"), + ("max_completion_tokens", "max_output_tokens"), + ): + if source in provider_request: + sampling[target] = provider_request[source] + reserved = { + "model", + "messages", + "tools", + "temperature", + "top_p", + "presence_penalty", + "max_tokens", + "max_completion_tokens", + } + provider_fields = { + key: value for key, value in provider_request.items() if key not in reserved + } + return { + "model": provider_request.get("model", ""), + "messages": provider_request.get("messages", []), + "tools": provider_request.get("tools", []), + "sampling": sampling, + "provider_fields": provider_fields, + } + + +def _raw_response(payload: Mapping[str, Any], *, compaction: bool = False) -> dict[str, Any]: + if compaction: + response = payload.get("summary_response") + if isinstance(response, Mapping): + raw = response.get("raw_provider_response") + return dict(raw) if isinstance(raw, Mapping) else {} + return {} + raw = payload.get("raw_provider_response") + if isinstance(raw, Mapping): + return dict(raw) + return {} + + +def _response_source(payload: Mapping[str, Any], *, compaction: bool = False) -> dict[str, Any]: + value = payload.get("summary_response" if compaction else "response") + return dict(value) if isinstance(value, Mapping) else {} + + +def _normalized_response(response: Mapping[str, Any]) -> dict[str, Any]: + tool_calls = response.get("tool_calls") + if not isinstance(tool_calls, list): + tool_calls = [] + content = response.get("content") + if content is None: + content = "" + return { + "message": { + "role": "assistant", + "content": content, + "reasoning_content": response.get("reasoning_content"), + "tool_calls": tool_calls, + }, + "finish_reason": response.get("finish_reason"), + "usage": response.get("usage") if isinstance(response.get("usage"), Mapping) else {}, + "provider_status": response.get("status", "error"), + **({"error": response.get("error")} if response.get("error") else {}), + } + + +def _base_provenance(config: TrainingExportConfig, payload: Mapping[str, Any], findings: list[str]) -> dict[str, Any]: + response = payload.get("response") + if not isinstance(response, Mapping): + response = payload.get("summary_response") if isinstance(payload.get("summary_response"), Mapping) else {} + return { + "harness_revision": config.harness_revision, + "scaffold_version": config.scaffold_version, + "tool_protocol_version": config.tool_protocol_version, + "requested_model": payload.get("requested_model") or response.get("requested_model") or payload.get("model_name"), + "provider_model": payload.get("provider_model") or response.get("provider_model"), + "model_revision": None, + "tokenizer_revision": config.tokenizer_revision, + "provider_retry_count": int(payload.get("provider_retry_count", response.get("provider_retry_count", 0)) or 0), + "oracle_isolation": { + "status": "leak" if findings else "clear", + "findings": findings, + }, + } + + +def _finish_unit(unit: dict[str, Any]) -> dict[str, Any]: + request_sha = _hash_json("researchharness.request.v1", unit["request"]) + response_sha = _hash_json("researchharness.response.v1", unit["response"]) + identity_preimage = { + key: unit[key] + for key in ( + "unit_type", + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "run_id", + "turn_index", + ) + } + identity_preimage.update( + { + "request_sha256": request_sha, + "response_sha256": response_sha, + "trace_event_index": unit["source"]["trace_event_index"], + } + ) + unit["unit_id"] = _hash_json("researchharness.unit-id.v1", identity_preimage) + unit["integrity"] = { + "canonicalization": CANONICALIZATION, + "request_sha256": request_sha, + "response_sha256": response_sha, + } + unit["integrity"]["unit_payload_sha256"] = _hash_json( + "researchharness.unit-payload.v1", unit + ) + return unit + + +def _unit_common(config: TrainingExportConfig, row: Mapping[str, Any], unit_type: str) -> dict[str, Any]: + return { + "schema_version": UNIT_SCHEMA, + "unit_type": unit_type, + **config.identity(), + "turn_index": int(row.get("turn_index", 0) or 0), + } + + +def _build_main_unit(config: TrainingExportConfig, root: Path, row: Mapping[str, Any], trace_ref: Mapping[str, Any]) -> dict[str, Any]: + payload = row.get("payload") if isinstance(row.get("payload"), Mapping) else {} + provider_request = _provider_request(payload) + raw_response = _raw_response(payload) + event_index = int(row.get("event_index", 0) or 0) + request_ref = _artifact_json(root, f"artifacts/provider/main-{event_index:06d}.request.json", provider_request) + raw_ref = _artifact_json(root, f"artifacts/provider/main-{event_index:06d}.response.json", raw_response) + request = _normalized_request(provider_request) + response_source = _response_source(payload) + response = _normalized_response(response_source) + findings = _oracle_findings(provider_request) + unit = _unit_common(config, row, "main_agent_turn") + unit.update( + { + "capture_status": "ok" if response_source.get("status") == "ok" and not findings else "error", + "request": request, + "response": response, + "source": { + "trace_ref": dict(trace_ref), + "trace_event_index": event_index, + "capture_type": "llm_call", + "provider_request_ref": request_ref, + "raw_provider_response_ref": raw_ref, + }, + "provenance": _base_provenance(config, payload, findings), + } + ) + return _finish_unit(unit) + + +def _build_compaction_unit(config: TrainingExportConfig, root: Path, row: Mapping[str, Any], trace_ref: Mapping[str, Any]) -> dict[str, Any]: + payload = row.get("payload") if isinstance(row.get("payload"), Mapping) else {} + event_index = int(row.get("event_index", 0) or 0) + provider_request = _provider_request(payload, compaction=True) + raw_response = _raw_response(payload, compaction=True) + request_ref = _artifact_json(root, f"artifacts/provider/compaction-{event_index:06d}.request.json", provider_request) + raw_ref = _artifact_json(root, f"artifacts/provider/compaction-{event_index:06d}.response.json", raw_response) + pre_ref = _artifact_json(root, f"artifacts/compaction/{event_index:06d}.pre.json", payload.get("pre_messages", [])) + post_ref = _artifact_json(root, f"artifacts/compaction/{event_index:06d}.post.json", payload.get("post_messages", [])) + response_source = _response_source(payload, compaction=True) + response = _normalized_response(response_source) + findings = _oracle_findings(provider_request) + summary_text = str(payload.get("summary_text", "") or "").strip() + successful = ( + payload.get("status") == "ok" + and response_source.get("status") == "ok" + and bool(summary_text) + and not response["message"]["tool_calls"] + and not findings + ) + existing_memory = str(payload.get("existing_memory_text", "") or "") + unit = _unit_common(config, row, "memory_compaction") + unit.update( + { + "capture_status": "ok" if successful else "error", + "request": _normalized_request(provider_request), + "response": response, + "compaction": { + "trigger_reason": payload.get("trigger_reason", ""), + "prior_token_estimate": int(payload.get("prior_token_estimate", 0) or 0), + "new_token_estimate": int(payload.get("new_token_estimate", 0) or 0), + "compacted_group_count": int(payload.get("compacted_group_count", 0) or 0), + "kept_group_count": int(payload.get("kept_group_count", 0) or 0), + "existing_memory_text_sha256": "sha256:" + _hash_bytes(existing_memory.encode("utf-8")), + "pre_messages_ref": pre_ref, + "post_messages_ref": post_ref, + }, + "source": { + "trace_ref": dict(trace_ref), + "trace_event_index": event_index, + "capture_type": "compaction", + "provider_request_ref": request_ref, + "raw_provider_response_ref": raw_ref, + }, + "provenance": _base_provenance(config, payload, findings), + } + ) + return _finish_unit(unit) + + +def _termination(value: Mapping[str, Any] | str) -> dict[str, Any]: + if isinstance(value, str): + reason = value + submitted = value == "result" + return {"status": "completed" if submitted else "stopped", "reason": reason, "submitted": submitted, "clean_termination": submitted} + reason = str(value.get("reason", value.get("termination", ""))) + submitted = bool(value.get("submitted", reason == "result")) + return { + "status": str(value.get("status", "completed" if submitted else "stopped")), + "reason": reason, + "submitted": submitted, + "clean_termination": bool(value.get("clean_termination", submitted)), + } + + +def _all_files(root: Path) -> list[Path]: + paths: list[Path] = [] + for path in sorted(root.rglob("*"), key=lambda item: item.relative_to(root).as_posix()): + if path.is_symlink(): + raise ValueError("export bundle may not contain symlinks") + if path.is_file(): + if path.stat().st_nlink != 1: + raise ValueError("export bundle may not contain hardlinked files") + _safe_ref_path(path.relative_to(root).as_posix()) + paths.append(path) + elif not path.is_dir(): + raise ValueError("export bundle contains a non-ordinary entry") + return paths + + +def _copy_source(source: Path, destination: Path) -> None: + _require_regular_source(source, label="source artifact") + destination.parent.mkdir(parents=True, exist_ok=True) + with source.open("rb") as reader, destination.open("xb") as writer: + shutil.copyfileobj(reader, writer) + writer.flush() + os.fsync(writer.fileno()) + + +def export_turn_training_bundle( + config: TrainingExportConfig, + *, + trace_path: Path, + termination: Mapping[str, Any] | str, +) -> Path: + if not isinstance(config, TrainingExportConfig): + raise TypeError("config must be TrainingExportConfig") + trace_source = _require_regular_source(Path(trace_path), label="trace") + target = config.output_dir + parent = target.parent + parent.mkdir(parents=True, exist_ok=True) + if target.exists() or target.is_symlink(): + raise FileExistsError(f"immutable export already exists: {target}") + lock_path = parent / f".{target.name}.publish.lock" + lock_fd = os.open(lock_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + temp = Path(tempfile.mkdtemp(prefix=f".{target.name}.tmp-", dir=parent)) + try: + os.close(lock_fd) + trace_copy = temp / "artifacts" / "trace" / trace_source.name + _copy_source(trace_source, trace_copy) + trace_ref = _file_ref(temp, trace_copy) + rows = _read_trace(trace_copy) + if any(row.get("run_id") != config.run_id for row in rows): + raise ValueError("trace run_id differs from frozen training identity") + units: list[dict[str, Any]] = [] + for row in rows: + capture_type = row.get("capture_type") + if capture_type == "llm_call": + units.append(_build_main_unit(config, temp, row, trace_ref)) + elif capture_type == "compaction": + units.append(_build_compaction_unit(config, temp, row, trace_ref)) + units_path = temp / "turn-training-units.v1.jsonl" + _write_bytes(units_path, b"".join(canonical_json_bytes(unit) + b"\n" for unit in units)) + + final_patch_ref: dict[str, Any] | None = None + if config.final_patch_path is not None: + patch_path = temp / "agent.patch" + _copy_source(config.final_patch_path, patch_path) + final_patch_ref = _file_ref(temp, patch_path) + summary = { + "schema_version": SUMMARY_SCHEMA, + **config.identity(), + "main_agent_turn_count": sum(unit["unit_type"] == "main_agent_turn" for unit in units), + "memory_compaction_count": sum(unit["unit_type"] == "memory_compaction" for unit in units), + "termination": _termination(termination), + "final_patch_ref": final_patch_ref, + "trace_ref": trace_ref, + } + _write_json(temp / "rollout-summary.v1.json", summary) + + payload_files = [ + path for path in _all_files(temp) if path.name not in {"manifest.json", "SHA256SUMS"} + ] + manifest = { + "schema_version": MANIFEST_SCHEMA, + **config.identity(), + "canonicalization": CANONICALIZATION, + "validators": [{"name": name, "status": "pass"} for name in _VALIDATOR_NAMES], + "oracle_leak_unit_count": sum( + unit["provenance"]["oracle_isolation"]["status"] == "leak" for unit in units + ), + "files": [_file_ref(temp, path) for path in payload_files], + } + _write_json(temp / "manifest.json", manifest) + checksum_files = [path for path in _all_files(temp) if path.name != "SHA256SUMS"] + checksum_text = "".join( + f"{_sha256_file(path)} {path.relative_to(temp).as_posix()}\n" + for path in checksum_files + ) + _write_bytes(temp / "SHA256SUMS", checksum_text.encode("utf-8")) + validate_turn_training_bundle(temp) + if target.exists() or target.is_symlink(): + raise FileExistsError(f"immutable export appeared during publish: {target}") + os.rename(temp, target) + directory_fd = os.open(parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + return target + except Exception: + shutil.rmtree(temp, ignore_errors=True) + raise + finally: + try: + os.close(lock_fd) + except OSError: + pass + try: + lock_path.unlink() + except FileNotFoundError: + pass + + +def _walk_refs(value: Any) -> list[Mapping[str, Any]]: + refs: list[Mapping[str, Any]] = [] + if isinstance(value, Mapping): + if set(value) == {"path", "sha256"}: + refs.append(value) + else: + for nested in value.values(): + refs.extend(_walk_refs(nested)) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for nested in value: + refs.extend(_walk_refs(nested)) + return refs + + +def _contains_prohibited_key(value: Any) -> bool: + if isinstance(value, Mapping): + return bool(set(value) & _PROHIBITED_AGENT_KEYS) or any(_contains_prohibited_key(nested) for nested in value.values()) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return any(_contains_prohibited_key(nested) for nested in value) + return False + + +def _load_units(path: Path) -> list[dict[str, Any]]: + units: list[dict[str, Any]] = [] + for line in path.read_bytes().splitlines(keepends=True): + if not line.strip(): + continue + if not line.endswith(b"\n"): + raise ValueError("training JSONL line lacks trailing newline") + value = json.loads(line) + if not isinstance(value, dict) or line != canonical_json_bytes(value) + b"\n": + raise ValueError("training JSONL is not canonical-json-v1") + units.append(value) + return units + + +def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: + root = Path(bundle_path) + if root.is_symlink() or not root.is_dir(): + raise ValueError("training export must be an ordinary directory") + files = _all_files(root) + relative_files = {path.relative_to(root).as_posix(): path for path in files} + required = { + "turn-training-units.v1.jsonl", + "rollout-summary.v1.json", + "manifest.json", + "SHA256SUMS", + } + if not required <= set(relative_files): + raise ValueError("training export is incomplete") + checksum_lines = (root / "SHA256SUMS").read_text(encoding="utf-8").splitlines() + observed_checksums: dict[str, str] = {} + for line in checksum_lines: + if len(line) < 67 or line[64:66] != " ": + raise ValueError("SHA256SUMS line is malformed") + digest, raw_path = line[:64], line[66:] + path = _safe_ref_path(raw_path).as_posix() + if _SHA256_RE.fullmatch(digest) is None or path in observed_checksums: + raise ValueError("SHA256SUMS contains invalid or duplicate entries") + observed_checksums[path] = digest + expected_checksum_paths = set(relative_files) - {"SHA256SUMS"} + if set(observed_checksums) != expected_checksum_paths: + raise ValueError("SHA256SUMS file set is not exact") + for relative, digest in observed_checksums.items(): + if _sha256_file(relative_files[relative]) != digest: + raise ValueError("SHA256SUMS digest mismatch") + + manifest = _read_json(root / "manifest.json") + summary = _read_json(root / "rollout-summary.v1.json") + units = _load_units(root / "turn-training-units.v1.jsonl") + if manifest.get("schema_version") != MANIFEST_SCHEMA or summary.get("schema_version") != SUMMARY_SCHEMA: + raise ValueError("training export schema mismatch") + manifest_refs = manifest.get("files") + if not isinstance(manifest_refs, list): + raise ValueError("manifest files must be a list") + manifest_paths: list[str] = [] + for ref in manifest_refs: + path = _resolve_ref(root, ref) + manifest_paths.append(path.relative_to(root).as_posix()) + expected_manifest_paths = set(relative_files) - {"manifest.json", "SHA256SUMS"} + if len(manifest_paths) != len(set(manifest_paths)) or set(manifest_paths) != expected_manifest_paths: + raise ValueError("manifest file closure is not exact") + identity_fields = ("task_id", "task_revision_id", "rollout_id", "rollout_index", "run_id") + identity = {key: summary.get(key) for key in identity_fields} + if any(manifest.get(key) != value for key, value in identity.items()): + raise ValueError("manifest identity differs from summary") + + trace_path = _resolve_ref(root, summary.get("trace_ref")) + trace_rows = _read_trace(trace_path) + captures = [row for row in trace_rows if row.get("capture_type") in {"llm_call", "compaction"}] + if len(captures) != len(units): + raise ValueError("trace capture closure differs from unit count") + capture_by_event = {int(row.get("event_index", 0) or 0): row for row in captures} + if len(capture_by_event) != len(captures): + raise ValueError("trace capture event identities are not unique") + for unit in units: + if unit.get("schema_version") != UNIT_SCHEMA or unit.get("unit_type") not in {"main_agent_turn", "memory_compaction"}: + raise ValueError("training unit schema/type is invalid") + if _contains_prohibited_key(unit): + raise ValueError("V1 unit contains prohibited multi-agent fields") + if any(unit.get(key) != value for key, value in identity.items()): + raise ValueError("training unit identity differs from summary") + for ref in _walk_refs(unit): + _resolve_ref(root, ref) + event_index = unit.get("source", {}).get("trace_event_index") + row = capture_by_event.get(event_index) + expected_capture = "llm_call" if unit.get("unit_type") == "main_agent_turn" else "compaction" + if row is None or row.get("capture_type") != expected_capture: + raise ValueError("training unit source does not bind its trace capture") + integrity = unit.get("integrity") + if not isinstance(integrity, dict) or integrity.get("canonicalization") != CANONICALIZATION: + raise ValueError("training unit integrity is invalid") + if integrity.get("request_sha256") != _hash_json("researchharness.request.v1", unit.get("request")): + raise ValueError("training unit request digest mismatch") + if integrity.get("response_sha256") != _hash_json("researchharness.response.v1", unit.get("response")): + raise ValueError("training unit response digest mismatch") + payload_digest = integrity.get("unit_payload_sha256") + payload_copy = json.loads(json.dumps(unit, ensure_ascii=False)) + payload_copy["integrity"].pop("unit_payload_sha256", None) + if payload_digest != _hash_json("researchharness.unit-payload.v1", payload_copy): + raise ValueError("training unit payload digest mismatch") + identity_preimage = { + key: unit[key] + for key in ( + "unit_type", + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "run_id", + "turn_index", + ) + } + identity_preimage.update( + { + "request_sha256": integrity["request_sha256"], + "response_sha256": integrity["response_sha256"], + "trace_event_index": event_index, + } + ) + if unit.get("unit_id") != _hash_json("researchharness.unit-id.v1", identity_preimage): + raise ValueError("training unit_id mismatch") + provider_request = _read_json(_resolve_ref(root, unit["source"]["provider_request_ref"])) + payload = row.get("payload") if isinstance(row.get("payload"), Mapping) else {} + expected_request = _provider_request(payload, compaction=expected_capture == "compaction") + if provider_request != expected_request or unit.get("request") != _normalized_request(provider_request): + raise ValueError("exported request differs from provider-boundary trace") + raw_response = _read_json(_resolve_ref(root, unit["source"]["raw_provider_response_ref"])) + if raw_response != _raw_response(payload, compaction=expected_capture == "compaction"): + raise ValueError("raw provider response differs from trace") + response_source = _response_source(payload, compaction=expected_capture == "compaction") + if unit.get("response") != _normalized_response(response_source): + raise ValueError("exported response differs from accepted trace response") + findings = _oracle_findings(provider_request) + oracle = unit.get("provenance", {}).get("oracle_isolation", {}) + if oracle != {"status": "leak" if findings else "clear", "findings": findings}: + raise ValueError("oracle leak marker differs from model-visible request") + if findings and unit.get("capture_status") != "error": + raise ValueError("oracle-leaking unit is not fail-closed") + if expected_capture == "compaction": + compaction = unit.get("compaction") + if not isinstance(compaction, Mapping): + raise ValueError("memory compaction unit lacks compaction metadata") + pre = json.loads(_resolve_ref(root, compaction.get("pre_messages_ref")).read_text(encoding="utf-8")) + post = json.loads(_resolve_ref(root, compaction.get("post_messages_ref")).read_text(encoding="utf-8")) + if pre != payload.get("pre_messages", []) or post != payload.get("post_messages", []): + raise ValueError("compaction pre/post artifact differs from trace") + + if summary.get("main_agent_turn_count") != sum(unit.get("unit_type") == "main_agent_turn" for unit in units): + raise ValueError("rollout summary main-agent count mismatch") + if summary.get("memory_compaction_count") != sum(unit.get("unit_type") == "memory_compaction" for unit in units): + raise ValueError("rollout summary compaction count mismatch") + validators = manifest.get("validators") + expected_validators = [{"name": name, "status": "pass"} for name in _VALIDATOR_NAMES] + if validators != expected_validators: + raise ValueError("manifest validator closure is not canonical") + leak_count = sum(unit.get("provenance", {}).get("oracle_isolation", {}).get("status") == "leak" for unit in units) + if manifest.get("oracle_leak_unit_count") != leak_count: + raise ValueError("manifest oracle leak count mismatch") + return summary + + +__all__ = [ + "EXPORT_DIR_NAME", + "TrainingExportConfig", + "canonical_json_bytes", + "export_turn_training_bundle", + "validate_turn_training_bundle", +] From fcb2cedeeaf49a54b52b6e51efb06bd9249a972d Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 03:00:29 +0800 Subject: [PATCH 2/6] Add turn training export contract tests --- tests/test_turn_training_export.py | 477 +++++++++++++++++++++++++++++ 1 file changed, 477 insertions(+) create mode 100644 tests/test_turn_training_export.py diff --git a/tests/test_turn_training_export.py b/tests/test_turn_training_export.py new file mode 100644 index 0000000..a7b6429 --- /dev/null +++ b/tests/test_turn_training_export.py @@ -0,0 +1,477 @@ +import json +import os +import tempfile +import types +import unittest +import sys +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import agent_base.training_export as export_module +from agent_base.training_export import ( + EXPORT_DIR_NAME, + TrainingExportConfig, + canonical_json_bytes, + export_turn_training_bundle, + validate_turn_training_bundle, +) + + +RUN_ID = "run-001" +MODEL = "relay/gpt-test" + + +def _tool_schema(): + return { + "type": "function", + "function": { + "name": "Read", + "description": "Read one workspace file.", + "parameters": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, + }, + } + + +def _accepted_response(request, *, response_id, content, tool_calls, retry_count=0): + raw = { + "id": response_id, + "model": MODEL, + "choices": [ + { + "finish_reason": "tool_calls" if tool_calls else "stop", + "message": {"role": "assistant", "content": content, "tool_calls": tool_calls}, + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 20}, + } + return { + "status": "ok", + "finish_reason": "tool_calls" if tool_calls else "stop", + "content": content, + "reasoning_content": None, + "tool_calls": tool_calls, + "usage": {"prompt_tokens": 100, "completion_tokens": 20}, + "provider_request": request, + "raw_provider_response": raw, + "provider_retry_count": retry_count, + "requested_model": MODEL, + "provider_model": MODEL, + } + + +def _main_payload(request, response): + return { + "model_name": MODEL, + "request_messages": request["messages"], + "native_tools": request.get("tools", []), + "provider_request": request, + "raw_provider_response": response["raw_provider_response"], + "provider_retry_count": response["provider_retry_count"], + "requested_model": response["requested_model"], + "provider_model": response["provider_model"], + "response": response, + } + + +def _trace_rows(): + initial = [ + {"role": "system", "content": "Fix the repository using tools."}, + {"role": "user", "content": "Inspect src/auth.py and fix token expiry."}, + ] + call = { + "id": "call-read-1", + "type": "function", + "function": {"name": "Read", "arguments": '{"file_path":"src/auth.py"}'}, + } + first_request = { + "model": MODEL, + "messages": initial, + "max_tokens": 512, + "temperature": 0.2, + "top_p": 0.9, + "presence_penalty": 0.0, + "tools": [_tool_schema()], + "tool_choice": "auto", + "parallel_tool_calls": True, + } + first_response = _accepted_response( + first_request, + response_id="resp-tool", + content=None, + tool_calls=[call], + retry_count=1, + ) + pre_messages = [ + *initial, + {"role": "assistant", "content": "", "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "call-read-1", "content": "TOKEN_TTL = 60"}, + ] + summary_messages = [ + {"role": "system", "content": "Compress history without tools."}, + {"role": "user", "content": "Summarize the auth task."}, + ] + summary_request = { + "model": MODEL, + "messages": summary_messages, + "max_tokens": 128, + "temperature": 0.2, + "top_p": 0.9, + "presence_penalty": 0.0, + } + summary_text = "Goal: fix token expiry. Evidence: TOKEN_TTL is 60." + summary_response = _accepted_response( + summary_request, + response_id="resp-memory", + content=summary_text, + tool_calls=[], + ) + post_messages = [ + *initial, + {"role": "user", "content": "Runtime memory summary from earlier turns.\n\n" + summary_text}, + {"role": "assistant", "content": "", "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "call-read-1", "content": "TOKEN_TTL = 60"}, + ] + second_request = { + "model": MODEL, + "messages": post_messages, + "max_tokens": 512, + "temperature": 0.2, + "top_p": 0.9, + "presence_penalty": 0.0, + "tools": [_tool_schema()], + "tool_choice": "auto", + "parallel_tool_calls": True, + } + second_response = _accepted_response( + second_request, + response_id="resp-final", + content="Updated the token expiry handling.", + tool_calls=[], + ) + return [ + { + "run_id": RUN_ID, + "event_index": 1, + "turn_index": 1, + "capture_type": "llm_call", + "payload": _main_payload(first_request, first_response), + }, + { + "run_id": RUN_ID, + "event_index": 2, + "turn_index": 1, + "capture_type": "compaction", + "payload": { + "trigger_reason": "usage", + "status": "ok", + "prior_token_estimate": 9000, + "new_token_estimate": 2200, + "compacted_group_count": 4, + "kept_group_count": 1, + "existing_memory_text": "", + "summary_request": summary_messages, + "summary_response": summary_response, + "summary_text": summary_text, + "pre_messages": pre_messages, + "post_messages": post_messages, + }, + }, + { + "run_id": RUN_ID, + "event_index": 3, + "turn_index": 2, + "capture_type": "llm_call", + "payload": _main_payload(second_request, second_response), + }, + ] + + +def _write_trace(path, rows=None): + rows = _trace_rows() if rows is None else rows + path.write_text( + "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows), + encoding="utf-8", + ) + return path + + +def _config(tmp_path, *, name="case", patch_file=True): + case = tmp_path / name + case.mkdir() + patch_path = case / "agent.patch" + if patch_file: + patch_path.write_text("diff --git a/a.py b/a.py\n", encoding="utf-8") + return TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-auth", + task_revision_id="task-auth@base-abc", + rollout_id="r0", + rollout_index=0, + run_id=RUN_ID, + harness_revision="72e8a447", + scaffold_version="react-v1", + tool_protocol_version="native-tools-v1", + tokenizer_revision="tok-rev1", + final_patch_path=patch_path if patch_file else None, + ) + + +def _units(bundle): + return [ + json.loads(line) + for line in (bundle / "turn-training-units.v1.jsonl").read_text(encoding="utf-8").splitlines() + ] + + +def _write_json(path, value): + path.write_bytes(canonical_json_bytes(value) + b"\n") + + +def _write_units(bundle, units): + (bundle / "turn-training-units.v1.jsonl").write_bytes( + b"".join(canonical_json_bytes(unit) + b"\n" for unit in units) + ) + + +def _reseal_summary(bundle, summary): + payload = json.loads(json.dumps(summary)) + payload.pop("integrity", None) + summary["integrity"] = { + "canonicalization": export_module.CANONICALIZATION, + "summary_payload_sha256": export_module._hash_json( + "researchharness.summary-payload.v1", payload + ), + } + _write_json(bundle / "rollout-summary.v1.json", summary) + + +def _reseal_bundle(bundle): + manifest_path = bundle / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + payload_files = [ + path + for path in export_module._all_files(bundle) + if path.name not in {"manifest.json", "SHA256SUMS"} + ] + manifest["files"] = [export_module._file_ref(bundle, path) for path in payload_files] + _write_json(manifest_path, manifest) + checksum_files = [ + path for path in export_module._all_files(bundle) if path.name != "SHA256SUMS" + ] + (bundle / "SHA256SUMS").write_text( + "".join( + f"{export_module._sha256_file(path)} {path.relative_to(bundle).as_posix()}\n" + for path in checksum_files + ), + encoding="utf-8", + ) + + +def _export(tmp_path, *, name="case", rows=None, patch_file=True): + config = _config(tmp_path, name=name, patch_file=patch_file) + trace = _write_trace(tmp_path / name / "trace.jsonl", rows) + bundle = export_turn_training_bundle( + config, + trace_path=trace, + termination={ + "status": "completed", + "reason": "result", + "submitted": True, + "clean_termination": True, + }, + ) + return config, bundle + + +class FakeMessage: + def __init__(self, content): + self.content = content + self.tool_calls = None + self.reasoning_content = None + + def model_dump(self): + return {"role": "assistant", "content": self.content, "tool_calls": None} + + +class FakeUsage: + def model_dump(self): + return {"prompt_tokens": 10, "completion_tokens": 2} + + +class FakeResponse: + def __init__(self, response_id, content): + self.id = response_id + self.model = MODEL + self.usage = FakeUsage() + self.choices = [ + types.SimpleNamespace(finish_reason="stop", message=FakeMessage(content)) + ] + + def model_dump(self): + return { + "id": self.id, + "model": self.model, + "choices": [ + { + "finish_reason": "stop", + "message": {"role": "assistant", "content": self.choices[0].message.content}, + } + ], + "usage": self.usage.model_dump(), + } + + +class FakeClient: + def __init__(self, responses): + self.responses = list(responses) + self.requests = [] + self.chat = types.SimpleNamespace(completions=types.SimpleNamespace(create=self.create)) + + def with_options(self, **kwargs): + return self + + def create(self, **kwargs): + self.requests.append(kwargs) + return self.responses.pop(0) + + +class TrainingExportTests(unittest.TestCase): + def setUp(self): + self._temp = tempfile.TemporaryDirectory() + self.tmp_path = Path(self._temp.name) + + def tearDown(self): + self._temp.cleanup() + + def test_contract_smoke_closes_tool_round_compaction_and_integrity(self): + config, bundle = _export(self.tmp_path) + summary = validate_turn_training_bundle(bundle) + units = _units(bundle) + self.assertEqual(bundle, config.output_dir) + self.assertEqual(summary["main_agent_turn_count"], 2) + self.assertEqual(summary["memory_compaction_count"], 1) + self.assertEqual(summary["capture_error_count"], 0) + self.assertEqual([unit["unit_type"] for unit in units], [ + "main_agent_turn", "memory_compaction", "main_agent_turn" + ]) + self.assertEqual(units[0]["response"]["message"]["tool_calls"][0]["id"], "call-read-1") + self.assertEqual(units[0]["provenance"]["provider_retry_count"], 1) + self.assertEqual(units[0]["provenance"]["requested_model"], MODEL) + self.assertEqual(units[0]["provenance"]["provider_model"], MODEL) + self.assertTrue(any( + message.get("role") == "tool" and message.get("tool_call_id") == "call-read-1" + for message in units[2]["request"]["messages"] + )) + self.assertEqual(units[1]["response"]["message"]["content"], _trace_rows()[1]["payload"]["summary_text"]) + raw_ref = units[0]["source"]["raw_provider_response_ref"] + self.assertEqual(json.loads((bundle / raw_ref["path"]).read_text())["id"], "resp-tool") + self.assertEqual(summary["final_patch_ref"]["path"], "agent.patch") + self.assertNotIn("resolved", summary) + self.assertNotIn("reward", summary) + + def test_missing_provider_boundary_is_audited_as_error(self): + row = _trace_rows()[0] + row["payload"].pop("provider_request") + row["payload"].pop("raw_provider_response") + row["payload"]["provider_retry_count"] = 2 + row["payload"]["response"] = { + "status": "error", + "error": "provider unavailable", + "content": "", + "tool_calls": [], + "provider_retry_count": 2, + "requested_model": MODEL, + "provider_model": None, + } + _, bundle = _export(self.tmp_path, name="provider-error", rows=[row], patch_file=False) + self.assertEqual(_units(bundle)[0]["capture_status"], "error") + self.assertEqual(validate_turn_training_bundle(bundle)["capture_error_count"], 1) + + def test_failed_compaction_is_preserved_as_error_unit(self): + rows = _trace_rows() + failed = rows[1] + failed["payload"].update( + status="error", + summary_response={"status": "error", "error": "summary unavailable", "tool_calls": []}, + summary_text="", + post_messages=failed["payload"]["pre_messages"], + ) + _, bundle = _export(self.tmp_path, name="compact-error", rows=[rows[0], failed], patch_file=False) + self.assertEqual([unit["capture_status"] for unit in _units(bundle)], ["ok", "error"]) + + def test_successful_compaction_must_match_next_real_request(self): + rows = _trace_rows() + rows[2]["payload"]["provider_request"]["messages"] = [{"role": "user", "content": "wrong"}] + config = _config(self.tmp_path, name="bad-closure", patch_file=False) + trace = _write_trace(self.tmp_path / "bad-closure/trace.jsonl", rows) + with self.assertRaisesRegex(ValueError, "next provider request"): + export_turn_training_bundle(config, trace_path=trace, termination="result") + self.assertFalse(config.output_dir.exists()) + + def test_oracle_request_is_marked_error_for_downstream_fail_closed(self): + rows = _trace_rows() + rows[0]["payload"]["provider_request"]["messages"][1]["content"] = "Apply the GOLD PATCH." + _, bundle = _export(self.tmp_path, name="oracle", rows=rows, patch_file=False) + units = _units(bundle) + manifest = json.loads((bundle / "manifest.json").read_text()) + self.assertEqual(units[0]["capture_status"], "error") + self.assertEqual(units[0]["provenance"]["oracle_isolation"]["status"], "leak") + self.assertEqual(manifest["oracle_leak_unit_count"], 2) + + def test_existing_target_and_wrong_directory_name_are_rejected(self): + config, bundle = _export(self.tmp_path) + before = (bundle / "SHA256SUMS").read_bytes() + with self.assertRaises(FileExistsError): + export_turn_training_bundle(config, trace_path=self.tmp_path / "case/trace.jsonl", termination="result") + self.assertEqual((bundle / "SHA256SUMS").read_bytes(), before) + with self.assertRaisesRegex(ValueError, "output_dir"): + TrainingExportConfig( + output_dir=self.tmp_path / "wrong", + task_id="t", task_revision_id="r", rollout_id="r0", rollout_index=0, + run_id="run", harness_revision="rev", + ) + + def test_structured_target_tamper_fails_even_after_resealing(self): + _, bundle = _export(self.tmp_path) + units = _units(bundle) + units[0]["response"]["message"]["tool_calls"][0]["id"] = "tampered" + export_module._finish_unit(units[0]) + _write_units(bundle, units) + _reseal_bundle(bundle) + with self.assertRaisesRegex(ValueError, "accepted trace response"): + validate_turn_training_bundle(bundle) + + def test_summary_counts_are_recomputed(self): + _, bundle = _export(self.tmp_path) + summary = json.loads((bundle / "rollout-summary.v1.json").read_text()) + summary["main_agent_turn_count"] = 99 + _reseal_summary(bundle, summary) + _reseal_bundle(bundle) + with self.assertRaisesRegex(ValueError, "main-agent count"): + validate_turn_training_bundle(bundle) + + def test_prohibited_agent_field_and_unsafe_ref_fail_closed(self): + _, agents_bundle = _export(self.tmp_path, name="agents") + units = _units(agents_bundle) + units[0]["request"]["agent_id"] = "invented" + export_module._finish_unit(units[0]) + _write_units(agents_bundle, units) + _reseal_bundle(agents_bundle) + with self.assertRaisesRegex(ValueError, "multi-agent"): + validate_turn_training_bundle(agents_bundle) + + _, path_bundle = _export(self.tmp_path, name="unsafe") + summary = json.loads((path_bundle / "rollout-summary.v1.json").read_text()) + summary["trace_ref"]["path"] = "../trace.jsonl" + _reseal_summary(path_bundle, summary) + _reseal_bundle(path_bundle) + with self.assertRaisesRegex(ValueError, "path"): + validate_turn_training_bundle(path_bundle) + +if __name__ == "__main__": + unittest.main() From b8bbf5d12ef0223186c73e3ffb60ed6e92b092a6 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 03:39:14 +0800 Subject: [PATCH 3/6] Complete immutable turn training export --- .github/workflows/ci.yml | 12 + README.md | 1 + agent_base/react_agent.py | 204 ++++- agent_base/trace_utils.py | 14 +- agent_base/training_export.py | 830 ++++++++++++++++-- api/openai_server.py | 100 +++ docs/turn-training-export-v1.md | 124 +++ researchharness/runtime.py | 114 +++ tests/test_openai_api_checks.py | 5 +- tests/test_turn_training_export.py | 456 +++++++++- tests/test_turn_training_export_api.py | 154 ++++ .../test_turn_training_export_entrypoints.py | 95 ++ 12 files changed, 2017 insertions(+), 92 deletions(-) create mode 100644 docs/turn-training-export-v1.md create mode 100644 tests/test_turn_training_export_api.py create mode 100644 tests/test_turn_training_export_entrypoints.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f33441b..056bb5e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,6 +59,18 @@ jobs: - name: Run Python import API checks run: python tests/test_python_api_tools.py + - name: Run OpenAI-compatible API checks + run: python tests/test_openai_api_checks.py + + - name: Run turn-level training export contract checks + run: python tests/test_turn_training_export.py + + - name: Run turn-level training export API checks + run: python tests/test_turn_training_export_api.py + + - name: Run turn-level training export entrypoint checks + run: python tests/test_turn_training_export_entrypoints.py + main-agent-api: name: Main agent API smoke runs-on: ubuntu-latest diff --git a/README.md b/README.md index ddf6eee..6e1a823 100644 --- a/README.md +++ b/README.md @@ -257,6 +257,7 @@ Start here if you are reading the codebase for the first time. - [docs/tutorial_en.md](docs/tutorial_en.md): detailed English tutorial - [docs/tutorial_zh.md](docs/tutorial_zh.md): detailed Chinese tutorial +- [docs/turn-training-export-v1.md](docs/turn-training-export-v1.md): immutable per-provider-call training export contract and usage - [tests/](tests): tool checks and end-to-end agent tests - [tests/example_files/](tests/example_files): fixed local fixtures diff --git a/agent_base/react_agent.py b/agent_base/react_agent.py index 6324f74..d6d491c 100644 --- a/agent_base/react_agent.py +++ b/agent_base/react_agent.py @@ -9,6 +9,7 @@ import threading from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, Type +from uuid import uuid4 from openai import OpenAI, APIError, APIConnectionError, APITimeoutError import tiktoken @@ -20,6 +21,7 @@ from agent_base.prompt import composed_system_prompt from agent_base.session_state import AgentSessionState, CompactionRecord, persist_session_state, resolve_session_state_path from agent_base.trace_utils import FlatTraceWriter +from agent_base.training_export import TrainingExportConfig, export_turn_training_bundle from agent_base.tools.custom import build_custom_tool_map from agent_base.tools.tooling import ToolBase, normalize_workspace_root from agent_base.tools.tool_extra import StrReplaceEditor @@ -348,6 +350,14 @@ def llm_call_trace_payload( "native_tools": safe_jsonable(list(native_tools)), "response": safe_jsonable(response), } + if isinstance(response, dict): + if isinstance(response.get("provider_request"), dict): + payload["provider_request"] = safe_jsonable(response["provider_request"]) + if isinstance(response.get("raw_provider_response"), dict): + payload["raw_provider_response"] = safe_jsonable(response["raw_provider_response"]) + payload["provider_retry_count"] = int(response.get("provider_retry_count", 0) or 0) + payload["requested_model"] = str(response.get("requested_model", model_name) or model_name) + payload["provider_model"] = response.get("provider_model") if image_aging and int(image_aging.get("omitted_image_count", 0) or 0) > 0: payload["image_aging"] = safe_jsonable(image_aging) return payload @@ -720,6 +730,8 @@ def __init__( custom_tools: Optional[Sequence[Any]] = None, max_rounds: Optional[int] = None, max_runtime_seconds: Optional[int] = None, + training_export: Optional[TrainingExportConfig] = None, + run_id: Optional[str] = None, ): if not isinstance(llm, dict): raise ValueError("llm must be a dict configuration.") @@ -769,6 +781,15 @@ def __init__( else: self.llm_generate_cfg.pop("omit_generate_params", None) self.trace_dir = Path(trace_dir) if trace_dir else None + if training_export is not None and not isinstance(training_export, TrainingExportConfig): + raise ValueError("training_export must be a TrainingExportConfig") + if training_export is not None and self.trace_dir is None: + raise ValueError("turn-level training export requires trace_dir") + if training_export is not None and run_id is not None and run_id != training_export.run_id: + raise ValueError("run_id differs from the frozen training export run_id") + self.training_export = training_export + self.run_id = training_export.run_id if training_export is not None else run_id + self._training_export_consumed = False self.trace_path: Optional[Path] = None self.session_state_path: Optional[Path] = None self.role_prompt = self.resolve_role_prompt(role_prompt) @@ -818,12 +839,23 @@ def _call_chat_completion( base_sleep_time = 1 last_error = "unknown llm error" + attempts_made = 0 + last_provider_request: dict[str, Any] = {} + last_raw_provider_response: dict[str, Any] = {} + provider_attempts: list[dict[str, Any]] = [] for attempt in range(max_tries): remaining = remaining_runtime_seconds(runtime_deadline) if remaining is not None and remaining <= 0: last_error = "agent runtime limit reached before llm call could complete" break + attempt_record: dict[str, Any] = { + "attempt_index": attempt, + "provider_request": {}, + "raw_provider_response": {}, + } + provider_returned = False try: + attempts_made = attempt + 1 if debug_enabled(): print(f"--- Attempting to call the service, try {attempt + 1}/{max_tries} ---") request_timeout = ( @@ -861,8 +893,23 @@ def _call_chat_completion( request_kwargs["tools"] = self._native_tools request_kwargs["tool_choice"] = "auto" request_kwargs["parallel_tool_calls"] = True + last_provider_request = safe_jsonable(request_kwargs) + attempt_record["provider_request"] = last_provider_request + last_raw_provider_response = {} with llm_hard_timeout(request_timeout): chat_response = request_client.chat.completions.create(**request_kwargs) + provider_returned = True + last_raw_provider_response = { + "model": getattr(chat_response, "model", None), + "response_type": type(chat_response).__name__, + } + response_dump = getattr(chat_response, "model_dump", None) + if not callable(response_dump): + raise ValueError("provider response does not expose model_dump()") + dumped_response = safe_jsonable(response_dump()) + if not isinstance(dumped_response, dict): + raise ValueError("provider model_dump() must return an object") + last_raw_provider_response = dumped_response choice = chat_response.choices[0] message = choice.message content = message.content @@ -874,6 +921,11 @@ def _call_chat_completion( if assistant_has_meaningful_text(content) or tool_calls: if debug_enabled(): print("--- Service call successful, received a valid response ---") + attempt_record.update( + status="accepted", + raw_provider_response=last_raw_provider_response, + ) + provider_attempts.append(safe_jsonable(attempt_record)) return { "status": "ok", "finish_reason": choice.finish_reason, @@ -882,16 +934,51 @@ def _call_chat_completion( "reasoning_content": reasoning_content, "raw_message": raw_message, "usage": usage, + "provider_request": last_provider_request, + "raw_provider_response": last_raw_provider_response, + "provider_retry_count": attempt, + "provider_attempts": provider_attempts, + "requested_model": self.model, + "provider_model": getattr(chat_response, "model", None), } else: last_error = "empty response from llm api" + attempt_record.update( + status="rejected_empty", + raw_provider_response=last_raw_provider_response, + error=last_error, + ) + provider_attempts.append(safe_jsonable(attempt_record)) if debug_enabled(): print(f"Warning: Attempt {attempt + 1} received an empty response.") except (APIError, APIConnectionError, APITimeoutError, LLMHardTimeoutError) as e: last_error = str(e) + attempt_record.update( + status="error", + error_type=type(e).__name__, + error=last_error, + raw_provider_response=last_raw_provider_response, + ) + provider_attempts.append(safe_jsonable(attempt_record)) if debug_enabled(): print(f"Error: Attempt {attempt + 1} failed with an API or network error: {e}") + except Exception as e: + if not provider_returned: + raise + last_error = f"invalid provider response: {type(e).__name__}: {e}" + attempt_record.update( + status="parse_error", + error_type=type(e).__name__, + error=last_error, + raw_provider_response=last_raw_provider_response, + ) + provider_attempts.append(safe_jsonable(attempt_record)) + if debug_enabled(): + print( + f"Error: Attempt {attempt + 1} returned an invalid provider response: " + f"{type(e).__name__}: {e}" + ) if attempt < max_tries - 1: sleep_time = base_sleep_time * (2 ** attempt) + random.uniform(0, 1) @@ -910,7 +997,16 @@ def _call_chat_completion( if debug_enabled(): print("Error: All retry attempts have been exhausted. The call has failed.") - return {"status": "error", "error": f"llm api error: {last_error}"} + return { + "status": "error", + "error": f"llm api error: {last_error}", + "provider_request": last_provider_request, + "raw_provider_response": last_raw_provider_response, + "provider_retry_count": max(0, attempts_made - 1), + "provider_attempts": provider_attempts, + "requested_model": self.model, + "provider_model": last_raw_provider_response.get("model"), + } def call_llm_api(self, msgs, max_tries=10, runtime_deadline: Optional[float] = None) -> dict[str, Any]: return self._call_chat_completion( @@ -1014,6 +1110,10 @@ def _run_session( """Internal execution path with trace data for tests and debugging.""" if not isinstance(prompt, str) or not prompt.strip(): raise ValueError("prompt must be a non-empty string.") + if self.training_export is not None and self._training_export_consumed: + raise RuntimeError("one immutable training export config may be consumed by only one session") + if self.training_export is not None: + self._training_export_consumed = True prompt_text = prompt.strip() resolved_workspace_root = normalize_workspace_root( @@ -1076,6 +1176,7 @@ def _run_session( model_name=self.model, workspace_root=resolved_workspace_root, on_event=event_callback, + run_id=self.run_id, ) self.trace_path = trace_writer.path self.session_state_path = resolve_session_state_path(self.trace_path) if self.trace_path else None @@ -1110,6 +1211,21 @@ def finalize(result_text: str, termination: str, *, role: str = "runtime", error error=error, ) persist_state(termination=termination, error=error) + training_export_path = "" + if self.training_export is not None: + if self.trace_path is None: + raise RuntimeError("training export requires a durable trace path") + training_export_path = str( + export_turn_training_bundle( + self.training_export, + trace_path=self.trace_path, + termination={ + "reason": termination, + "submitted": termination == "result", + "clean_termination": not bool(error), + }, + ) + ) return { "prompt": prompt_text, "messages": messages, @@ -1117,6 +1233,7 @@ def finalize(result_text: str, termination: str, *, role: str = "runtime", error "termination": termination, "trace_path": str(self.trace_path) if self.trace_path else "", "session_state_path": str(self.session_state_path) if self.session_state_path else "", + "training_export_path": training_export_path, } def interruption_requested() -> bool: @@ -1530,22 +1647,29 @@ def resolve_agent_class_for_role_prompt_files(role_prompt_files: Sequence[str]) def _parse_cli_args( argv: list[str], -) -> tuple[ - str, - Optional[str], - Optional[str], - str, - list[str], - list[str], - Optional[bool], - list[str], - list[str], - dict[str, Any], -]: +) -> tuple[Any, ...]: parser = argparse.ArgumentParser(description="Run the local agent directly from agent_base.react_agent.") parser.add_argument("prompt", nargs="*", help="Prompt text.") parser.add_argument("--prompt-file", help="Optional UTF-8 text file containing the prompt.") parser.add_argument("--trace-dir", help="Optional directory where the run trace JSONL should be created.") + parser.add_argument("--training-export-dir", help="Exact output directory; its final name must be turn-training-export-v1.") + parser.add_argument("--task-id", help="Stable semantic task identity for turn-level export.") + parser.add_argument("--task-revision-id", help="Task/environment/verifier revision identity.") + parser.add_argument("--rollout-id", help="Stable rollout identity inside the task.") + parser.add_argument("--rollout-index", type=int, help="Non-negative rollout index.") + parser.add_argument("--run-id", help="Optional frozen run identity; generated when omitted.") + parser.add_argument("--harness-revision", help="Exact ResearchHarness source revision.") + parser.add_argument("--scaffold-version", default="", help="Optional scaffold version provenance.") + parser.add_argument( + "--tool-protocol-version", + default="openai-chat-completions-native-tools-v1", + help="Tool protocol provenance for exported units.", + ) + parser.add_argument("--tokenizer-revision", help="Optional tokenizer revision provenance.") + parser.add_argument( + "--final-patch-path", + help="Optional explicit final patch path. Generic tasks leave this unset; no patch is guessed.", + ) parser.add_argument( "--workspace-root", help="Optional workspace root for local file tools, Bash, and TerminalStart.", @@ -1621,6 +1745,50 @@ def _parse_cli_args( omit_generate_params = ( normalize_omit_generate_params(args.omit_generate_params) if args.omit_generate_params else None ) + training_export: Optional[TrainingExportConfig] = None + resolved_trace_dir = args.trace_dir + training_requested = args.training_export_dir is not None or any( + value is not None + for value in ( + args.task_id, + args.task_revision_id, + args.rollout_id, + args.rollout_index, + args.harness_revision, + args.final_patch_path, + ) + ) + if training_requested: + missing = [ + flag + for flag, value in ( + ("--training-export-dir", args.training_export_dir), + ("--task-id", args.task_id), + ("--task-revision-id", args.task_revision_id), + ("--rollout-id", args.rollout_id), + ("--rollout-index", args.rollout_index), + ("--harness-revision", args.harness_revision), + ) + if value is None + ] + if missing: + raise ValueError("turn-level training export requires: " + ", ".join(missing)) + output_dir = Path(args.training_export_dir) + resolved_run_id = args.run_id if args.run_id is not None else uuid4().hex + resolved_trace_dir = resolved_trace_dir or str(output_dir.parent / "trace") + training_export = TrainingExportConfig( + output_dir=output_dir, + task_id=str(args.task_id), + task_revision_id=str(args.task_revision_id), + rollout_id=str(args.rollout_id), + rollout_index=int(args.rollout_index), + run_id=resolved_run_id, + harness_revision=str(args.harness_revision), + scaffold_version=str(args.scaffold_version), + tool_protocol_version=str(args.tool_protocol_version), + tokenizer_revision=args.tokenizer_revision, + final_patch_path=Path(args.final_patch_path) if args.final_patch_path else None, + ) prompt_text = "" if args.prompt_file: @@ -1633,7 +1801,7 @@ def _parse_cli_args( role_prompt = read_role_prompt_files(args.role_prompt_files) return ( prompt_text, - args.trace_dir, + resolved_trace_dir, args.workspace_root, role_prompt, list(args.role_prompt_files), @@ -1643,6 +1811,7 @@ def _parse_cli_args( resolve_extra_tool_names(args.extra_tools), llm_extra_body, omit_generate_params, + training_export, ) @@ -1662,7 +1831,11 @@ def main(argv: Optional[list[str]] = None) -> int: extra_tools, llm_extra_body, omit_generate_params, + training_export, ) = _parse_cli_args(argv or sys.argv[1:]) + chat_enabled = chat_arg if chat_arg is not None else (sys.stdin.isatty() and sys.stdout.isatty()) + if training_export is not None and chat_enabled: + raise ValueError("V1 immutable training export does not support --chat; run each session as a distinct rollout") agent_cls = resolve_agent_class_for_role_prompt_files(role_prompt_files) forbidden_tools = set(getattr(agent_cls, "forbidden_tool_names", set())) forbidden_requested_tools = sorted(set(tool_names) & forbidden_tools) @@ -1680,6 +1853,8 @@ def main(argv: Optional[list[str]] = None) -> int: llm=default_llm_config(extra_body=llm_extra_body, omit_generate_params=omit_generate_params), trace_dir=trace_dir, role_prompt=role_prompt or None, + training_export=training_export, + run_id=training_export.run_id if training_export is not None else None, ) resolved_workspace_root = normalize_workspace_root(workspace_root) initial_content_parts: list[dict[str, Any]] = [] @@ -1705,7 +1880,6 @@ def main(argv: Optional[list[str]] = None) -> int: event_callback=printer.handle_event, initial_content_parts=initial_content_parts or None, ) - chat_enabled = chat_arg if chat_arg is not None else (sys.stdin.isatty() and sys.stdout.isatty()) messages = session.get("messages", []) while chat_enabled: try: diff --git a/agent_base/trace_utils.py b/agent_base/trace_utils.py index cfcbffa..21c1c58 100644 --- a/agent_base/trace_utils.py +++ b/agent_base/trace_utils.py @@ -1,6 +1,7 @@ import argparse import datetime from pathlib import Path +import re from typing import Any, Callable, Optional from uuid import uuid4 @@ -28,6 +29,9 @@ ] +_RUN_ID_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}") + + class FlatTraceWriter: def __init__( self, @@ -36,11 +40,19 @@ def __init__( model_name: str, workspace_root: str | Path, on_event: Optional[Callable[[dict[str, Any]], None]] = None, + run_id: Optional[str] = None, ): self.model_name = model_name self.workspace_root = str(workspace_root) self.on_event = on_event - self.run_id = uuid4().hex + if run_id is not None and not isinstance(run_id, str): + raise ValueError("run_id must be a string") + resolved_run_id = run_id if run_id is not None else uuid4().hex + if _RUN_ID_RE.fullmatch(resolved_run_id) is None: + raise ValueError( + "run_id must be 1-128 characters using only letters, digits, '.', '_', or '-'" + ) + self.run_id = resolved_run_id self.path = resolve_trace_path(trace_dir, run_id=self.run_id) if trace_dir else None self.event_index = 0 diff --git a/agent_base/training_export.py b/agent_base/training_export.py index 5bc7c35..4930a1b 100644 --- a/agent_base/training_export.py +++ b/agent_base/training_export.py @@ -4,6 +4,7 @@ import ctypes from dataclasses import dataclass +import errno import hashlib import json import math @@ -11,6 +12,7 @@ from pathlib import Path, PurePosixPath import re import shutil +import stat import tempfile from typing import Any, Mapping, Sequence @@ -21,6 +23,7 @@ EXPORT_DIR_NAME = "turn-training-export-v1" CANONICALIZATION = "canonical-json-v1" _IDENTITY_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:@+-]{0,255}") +_RUN_ID_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}") _SHA256_RE = re.compile(r"[0-9a-f]{64}") _PROHIBITED_AGENT_KEYS = { "subagent", @@ -87,11 +90,18 @@ def __post_init__(self) -> None: "run_id", "harness_revision", ): - value = str(getattr(self, label)) - if _IDENTITY_RE.fullmatch(value) is None: + value = getattr(self, label) + pattern = _RUN_ID_RE if label == "run_id" else _IDENTITY_RE + if not isinstance(value, str) or pattern.fullmatch(value) is None: raise ValueError(f"{label} is not a safe, non-empty identity") if isinstance(self.rollout_index, bool) or not isinstance(self.rollout_index, int) or self.rollout_index < 0: raise ValueError("rollout_index must be a non-negative integer") + if not isinstance(self.scaffold_version, str): + raise ValueError("scaffold_version must be a string") + if not isinstance(self.tool_protocol_version, str) or not self.tool_protocol_version: + raise ValueError("tool_protocol_version must be a non-empty string") + if self.tokenizer_revision is not None and not isinstance(self.tokenizer_revision, str): + raise ValueError("tokenizer_revision must be a string or null") def identity(self) -> dict[str, Any]: return { @@ -158,8 +168,16 @@ def _write_json(path: Path, value: Any) -> None: _write_bytes(path, canonical_json_bytes(value) + b"\n") +def _read_canonical_json(path: Path) -> Any: + raw = path.read_bytes() + value = json.loads(raw) + if raw != canonical_json_bytes(value) + b"\n": + raise ValueError(f"JSON artifact is not canonical-json-v1: {path}") + return value + + def _read_json(path: Path) -> dict[str, Any]: - value = json.loads(path.read_text(encoding="utf-8")) + value = _read_canonical_json(path) if not isinstance(value, dict): raise ValueError(f"expected a JSON object: {path}") return value @@ -174,7 +192,13 @@ def _read_trace(path: Path) -> list[dict[str, Any]]: value = json.loads(raw_line) if not isinstance(value, dict): raise ValueError(f"trace line {line_number} is not an object") + _reject_nonfinite(value) rows.append(value) + event_indexes = [row.get("event_index") for row in rows] + if any(isinstance(index, bool) or not isinstance(index, int) or index < 1 for index in event_indexes): + raise ValueError("trace event_index must be a positive integer") + if event_indexes != sorted(event_indexes) or len(event_indexes) != len(set(event_indexes)): + raise ValueError("trace event_index must be strictly increasing") return rows @@ -182,6 +206,9 @@ def _require_regular_source(path: Path, *, label: str) -> Path: source = Path(path) if source.is_symlink() or not source.is_file(): raise ValueError(f"{label} must be an ordinary file") + if source.stat().st_nlink != 1: + raise ValueError(f"{label} must not be hardlinked") + _reject_symlink_ancestors(source.parent) return source @@ -326,6 +353,307 @@ def _normalized_response(response: Mapping[str, Any]) -> dict[str, Any]: } +def _normalized_raw_provider_response(raw_response: Mapping[str, Any]) -> dict[str, Any]: + choices = raw_response.get("choices") + if not isinstance(choices, list) or not choices or not isinstance(choices[0], Mapping): + raise ValueError("raw provider response must contain a structured first choice") + choice = choices[0] + message = choice.get("message") + if not isinstance(message, Mapping): + raise ValueError("raw provider response first choice must contain a message") + raw_tool_calls = message.get("tool_calls") + if raw_tool_calls is None: + raw_tool_calls = [] + if not isinstance(raw_tool_calls, list): + raise ValueError("raw provider response tool_calls must be a list or null") + tool_calls: list[dict[str, Any]] = [] + for raw_tool_call in raw_tool_calls: + if not isinstance(raw_tool_call, Mapping): + raise ValueError("raw provider response tool call must be an object") + function = raw_tool_call.get("function") + if not isinstance(function, Mapping): + raise ValueError("raw provider response tool call function must be an object") + tool_calls.append( + { + "id": raw_tool_call.get("id"), + "type": raw_tool_call.get("type"), + "function": { + "name": function.get("name"), + "arguments": function.get("arguments"), + }, + } + ) + content = message.get("content") + if content is None: + content = "" + usage = raw_response.get("usage") + return { + "message": { + "role": "assistant", + "content": content, + "reasoning_content": message.get("reasoning_content"), + "tool_calls": tool_calls, + }, + "finish_reason": choice.get("finish_reason"), + "usage": dict(usage) if isinstance(usage, Mapping) else {}, + "provider_status": "ok", + } + + +def _has_exact_provider_capture(payload: Mapping[str, Any], *, compaction: bool) -> bool: + if compaction: + response = payload.get("summary_response") + return bool( + isinstance(response, Mapping) + and isinstance(response.get("provider_request"), Mapping) + and response.get("provider_request") + and isinstance(response.get("raw_provider_response"), Mapping) + and response.get("raw_provider_response") + ) + return bool( + isinstance(payload.get("provider_request"), Mapping) + and payload.get("provider_request") + and isinstance(payload.get("raw_provider_response"), Mapping) + and payload.get("raw_provider_response") + ) + + +def _main_capture_ok(payload: Mapping[str, Any], findings: Sequence[str]) -> bool: + response = _response_source(payload) + tool_calls = response.get("tool_calls") + return bool( + response.get("status") == "ok" + and _has_exact_provider_capture(payload, compaction=False) + and not ( + response.get("finish_reason") == "length" + and isinstance(tool_calls, list) + and tool_calls + ) + and not findings + ) + + +def _compaction_capture_ok(payload: Mapping[str, Any], findings: Sequence[str]) -> bool: + response = _response_source(payload, compaction=True) + summary_text = str(payload.get("summary_text", "") or "").strip() + response_text = str(response.get("content", "") or "").strip() + tool_calls = response.get("tool_calls") + provider_request = _provider_request(payload, compaction=True) + summary_request = payload.get("summary_request") + return bool( + payload.get("status") == "ok" + and response.get("status") == "ok" + and _has_exact_provider_capture(payload, compaction=True) + and summary_text + and summary_text == response_text + and isinstance(tool_calls, list) + and not tool_calls + and isinstance(summary_request, list) + and summary_request == provider_request.get("messages") + and not provider_request.get("tools") + and not findings + ) + + +def _validate_request_shape(request: Mapping[str, Any], *, capture_ok: bool) -> None: + if set(request) != {"model", "messages", "tools", "sampling", "provider_fields"}: + raise ValueError("normalized request has an unexpected schema") + if not isinstance(request.get("model"), str): + raise ValueError("request.model must be a string") + if capture_ok and not request.get("model"): + raise ValueError("successful capture must preserve a non-empty provider model request") + if not isinstance(request.get("messages"), list): + raise ValueError("request.messages must be a list") + if not isinstance(request.get("tools"), list): + raise ValueError("request.tools must be a list") + if not isinstance(request.get("sampling"), Mapping): + raise ValueError("request.sampling must be an object") + if not isinstance(request.get("provider_fields"), Mapping): + raise ValueError("request.provider_fields must be an object") + + +def _validate_tool_calls(response: Mapping[str, Any]) -> None: + message = response.get("message") + if not isinstance(message, Mapping): + raise ValueError("response.message must be an object") + tool_calls = message.get("tool_calls") + if not isinstance(tool_calls, list): + raise ValueError("response.message.tool_calls must be a list") + seen_ids: set[str] = set() + for tool_call in tool_calls: + if not isinstance(tool_call, Mapping) or tool_call.get("type") != "function": + raise ValueError("tool call must be a structured function call") + call_id = tool_call.get("id") + function = tool_call.get("function") + if not isinstance(call_id, str) or not call_id or call_id in seen_ids: + raise ValueError("tool call IDs must be non-empty and unique in response order") + if not isinstance(function, Mapping): + raise ValueError("tool call function must be an object") + if not isinstance(function.get("name"), str) or not function.get("name"): + raise ValueError("tool call function name must be non-empty") + if not isinstance(function.get("arguments"), str): + raise ValueError("tool call arguments must preserve the provider string") + seen_ids.add(call_id) + + +def _validate_provider_attempts( + response_source: Mapping[str, Any], + *, + provider_request: Mapping[str, Any], + raw_response: Mapping[str, Any], +) -> None: + attempts = response_source.get("provider_attempts") + status = response_source.get("status") + if not isinstance(attempts, list) or not attempts: + if status == "ok": + raise ValueError("successful provider capture must preserve every provider attempt") + if provider_request or raw_response: + raise ValueError("failed provider capture with boundary evidence must preserve attempts") + return + retry_count = _nonnegative_int( + response_source.get("provider_retry_count", len(attempts) - 1), + label="provider_retry_count", + ) + if retry_count != len(attempts) - 1: + raise ValueError("provider retry count differs from captured attempt closure") + allowed_statuses = {"accepted", "rejected_empty", "error", "parse_error"} + for index, attempt in enumerate(attempts): + if not isinstance(attempt, Mapping) or attempt.get("attempt_index") != index: + raise ValueError("provider attempts must be complete and canonically indexed") + if attempt.get("status") not in allowed_statuses: + raise ValueError("provider attempt has an invalid status") + if not isinstance(attempt.get("provider_request"), Mapping): + raise ValueError("provider attempt request must be an object") + if not isinstance(attempt.get("raw_provider_response"), Mapping): + raise ValueError("provider attempt raw response must be an object") + if attempt.get("provider_request") != provider_request: + raise ValueError("provider retry request differs from the final provider request") + if attempt.get("status") == "accepted" and index != len(attempts) - 1: + raise ValueError("only the final provider attempt may be accepted") + if attempt.get("status") == "rejected_empty": + try: + rejected_response = _normalized_raw_provider_response( + attempt["raw_provider_response"] + ) + except ValueError as exc: + raise ValueError("rejected-empty attempt lacks a parseable response") from exc + rejected_message = rejected_response["message"] + rejected_content = rejected_message.get("content") + if ( + (isinstance(rejected_content, str) and rejected_content.strip()) + or (not isinstance(rejected_content, str) and rejected_content) + or rejected_message.get("tool_calls") + ): + raise ValueError("rejected-empty attempt contains an acceptable response") + final_attempt = attempts[-1] + if final_attempt.get("provider_request") != provider_request: + raise ValueError("final provider attempt request differs from accepted boundary") + if final_attempt.get("raw_provider_response") != raw_response: + raise ValueError("final provider attempt response differs from accepted boundary") + if status == "ok" and final_attempt.get("status") != "accepted": + raise ValueError("successful provider capture must end in one accepted attempt") + if status != "ok" and final_attempt.get("status") == "accepted": + raise ValueError("failed provider capture cannot end in an accepted attempt") + + +def _tool_result_ids(messages: object) -> list[str]: + if not isinstance(messages, Sequence) or isinstance(messages, (str, bytes, bytearray)): + return [] + result: list[str] = [] + for message in messages: + if not isinstance(message, Mapping) or message.get("role") != "tool": + continue + tool_call_id = message.get("tool_call_id") + if isinstance(tool_call_id, str) and tool_call_id: + result.append(tool_call_id) + return result + + +def _contains_ordered_sequence(values: Sequence[str], expected: Sequence[str]) -> bool: + if not expected: + return True + cursor = 0 + for value in values: + if value == expected[cursor]: + cursor += 1 + if cursor == len(expected): + return True + return False + + +def _validate_tool_result_closure( + *, + trace_rows: Sequence[Mapping[str, Any]], + capture_row: Mapping[str, Any], + response: Mapping[str, Any], +) -> None: + message = response.get("message") + tool_calls = message.get("tool_calls") if isinstance(message, Mapping) else [] + if not isinstance(tool_calls, list) or not tool_calls: + return + if response.get("finish_reason") == "length": + return + expected_ids = [str(tool_call.get("id", "")) for tool_call in tool_calls] + capture_event_index = capture_row.get("event_index") + capture_position = next( + ( + index + for index, row in enumerate(trace_rows) + if row.get("event_index") == capture_event_index + ), + None, + ) + if capture_position is None: + raise ValueError("tool-call capture is missing from trace") + next_main_position = next( + ( + index + for index in range(capture_position + 1, len(trace_rows)) + if trace_rows[index].get("capture_type") == "llm_call" + ), + None, + ) + boundary = next_main_position if next_main_position is not None else len(trace_rows) + intervening = trace_rows[capture_position + 1 : boundary] + traced_result_ids: list[str] = [] + for row in intervening: + if row.get("role") == "tool": + ids = row.get("tool_call_ids") + if isinstance(ids, list): + traced_result_ids.extend( + str(tool_call_id) + for tool_call_id in ids + if isinstance(tool_call_id, str) and tool_call_id + ) + trace_proves_execution = _contains_ordered_sequence(traced_result_ids, expected_ids) + compaction_pre_proof = False + compacted_next_input_proof = False + for row in intervening: + if row.get("capture_type") != "compaction": + continue + payload = row.get("payload") + if not isinstance(payload, Mapping): + continue + if _contains_ordered_sequence(_tool_result_ids(payload.get("pre_messages")), expected_ids): + compaction_pre_proof = True + if payload.get("status") == "ok": + compacted_next_input_proof = True + next_input_proof = False + if next_main_position is not None: + next_payload = trace_rows[next_main_position].get("payload") + if isinstance(next_payload, Mapping): + next_input_proof = _contains_ordered_sequence( + _tool_result_ids(_provider_request(next_payload).get("messages")), + expected_ids, + ) + if not (trace_proves_execution or next_input_proof or compaction_pre_proof): + raise ValueError("tool-call target lacks matching tool-result closure") + if next_main_position is not None and not ( + next_input_proof or compacted_next_input_proof + ): + raise ValueError("tool results did not enter the next provider request") + + def _base_provenance(config: TrainingExportConfig, payload: Mapping[str, Any], findings: list[str]) -> dict[str, Any]: response = payload.get("response") if not isinstance(response, Mapping): @@ -338,7 +666,7 @@ def _base_provenance(config: TrainingExportConfig, payload: Mapping[str, Any], f "provider_model": payload.get("provider_model") or response.get("provider_model"), "model_revision": None, "tokenizer_revision": config.tokenizer_revision, - "provider_retry_count": int(payload.get("provider_retry_count", response.get("provider_retry_count", 0)) or 0), + "provider_retry_count": _nonnegative_int(payload.get("provider_retry_count", response.get("provider_retry_count", 0)), label="provider_retry_count"), "oracle_isolation": { "status": "leak" if findings else "clear", "findings": findings, @@ -365,7 +693,6 @@ def _finish_unit(unit: dict[str, Any]) -> dict[str, Any]: { "request_sha256": request_sha, "response_sha256": response_sha, - "trace_event_index": unit["source"]["trace_event_index"], } ) unit["unit_id"] = _hash_json("researchharness.unit-id.v1", identity_preimage) @@ -380,12 +707,19 @@ def _finish_unit(unit: dict[str, Any]) -> dict[str, Any]: return unit +def _nonnegative_int(value: Any, *, label: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{label} must be a non-negative integer") + return value + + def _unit_common(config: TrainingExportConfig, row: Mapping[str, Any], unit_type: str) -> dict[str, Any]: + turn_index = _nonnegative_int(row.get("turn_index"), label="turn_index") return { "schema_version": UNIT_SCHEMA, "unit_type": unit_type, **config.identity(), - "turn_index": int(row.get("turn_index", 0) or 0), + "turn_index": turn_index, } @@ -403,7 +737,7 @@ def _build_main_unit(config: TrainingExportConfig, root: Path, row: Mapping[str, unit = _unit_common(config, row, "main_agent_turn") unit.update( { - "capture_status": "ok" if response_source.get("status") == "ok" and not findings else "error", + "capture_status": "ok" if _main_capture_ok(payload, findings) else "error", "request": request, "response": response, "source": { @@ -431,14 +765,7 @@ def _build_compaction_unit(config: TrainingExportConfig, root: Path, row: Mappin response_source = _response_source(payload, compaction=True) response = _normalized_response(response_source) findings = _oracle_findings(provider_request) - summary_text = str(payload.get("summary_text", "") or "").strip() - successful = ( - payload.get("status") == "ok" - and response_source.get("status") == "ok" - and bool(summary_text) - and not response["message"]["tool_calls"] - and not findings - ) + successful = _compaction_capture_ok(payload, findings) existing_memory = str(payload.get("existing_memory_text", "") or "") unit = _unit_common(config, row, "memory_compaction") unit.update( @@ -448,10 +775,10 @@ def _build_compaction_unit(config: TrainingExportConfig, root: Path, row: Mappin "response": response, "compaction": { "trigger_reason": payload.get("trigger_reason", ""), - "prior_token_estimate": int(payload.get("prior_token_estimate", 0) or 0), - "new_token_estimate": int(payload.get("new_token_estimate", 0) or 0), - "compacted_group_count": int(payload.get("compacted_group_count", 0) or 0), - "kept_group_count": int(payload.get("kept_group_count", 0) or 0), + "prior_token_estimate": _nonnegative_int(payload.get("prior_token_estimate", 0), label="prior_token_estimate"), + "new_token_estimate": _nonnegative_int(payload.get("new_token_estimate", 0), label="new_token_estimate"), + "compacted_group_count": _nonnegative_int(payload.get("compacted_group_count", 0), label="compacted_group_count"), + "kept_group_count": _nonnegative_int(payload.get("kept_group_count", 0), label="kept_group_count"), "existing_memory_text_sha256": "sha256:" + _hash_bytes(existing_memory.encode("utf-8")), "pre_messages_ref": pre_ref, "post_messages_ref": post_ref, @@ -474,25 +801,40 @@ def _termination(value: Mapping[str, Any] | str) -> dict[str, Any]: reason = value submitted = value == "result" return {"status": "completed" if submitted else "stopped", "reason": reason, "submitted": submitted, "clean_termination": submitted} - reason = str(value.get("reason", value.get("termination", ""))) - submitted = bool(value.get("submitted", reason == "result")) + if not isinstance(value, Mapping): + raise TypeError("termination must be a string or object") + reason = value.get("reason", value.get("termination", "")) + status = value.get("status") + if not isinstance(reason, str): + raise ValueError("termination.reason must be a string") + if status is not None and not isinstance(status, str): + raise ValueError("termination.status must be a string") + for field in ("submitted", "clean_termination"): + if field in value and not isinstance(value[field], bool): + raise ValueError(f"termination.{field} must be a boolean") + submitted = value.get("submitted", reason == "result") return { - "status": str(value.get("status", "completed" if submitted else "stopped")), + "status": status or ("completed" if submitted else "stopped"), "reason": reason, "submitted": submitted, - "clean_termination": bool(value.get("clean_termination", submitted)), + "clean_termination": value.get("clean_termination", submitted), } def _all_files(root: Path) -> list[Path]: paths: list[Path] = [] + casefolded: set[str] = set() for path in sorted(root.rglob("*"), key=lambda item: item.relative_to(root).as_posix()): + relative = path.relative_to(root).as_posix() + if relative.casefold() in casefolded: + raise ValueError(f"export bundle contains a case-colliding path: {relative}") + casefolded.add(relative.casefold()) if path.is_symlink(): raise ValueError("export bundle may not contain symlinks") if path.is_file(): if path.stat().st_nlink != 1: raise ValueError("export bundle may not contain hardlinked files") - _safe_ref_path(path.relative_to(root).as_posix()) + _safe_ref_path(relative) paths.append(path) elif not path.is_dir(): raise ValueError("export bundle contains a non-ordinary entry") @@ -502,12 +844,52 @@ def _all_files(root: Path) -> list[Path]: def _copy_source(source: Path, destination: Path) -> None: _require_regular_source(source, label="source artifact") destination.parent.mkdir(parents=True, exist_ok=True) - with source.open("rb") as reader, destination.open("xb") as writer: + flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0) + source_fd = os.open(source, flags) + source_stat = os.fstat(source_fd) + if not stat.S_ISREG(source_stat.st_mode) or source_stat.st_nlink != 1: + os.close(source_fd) + raise ValueError("source artifact changed into a non-ordinary or hardlinked file") + with os.fdopen(source_fd, "rb") as reader, destination.open("xb") as writer: shutil.copyfileobj(reader, writer) writer.flush() os.fsync(writer.fileno()) +def _reject_symlink_ancestors(path: Path) -> None: + current = path.absolute() + while current != current.parent: + if current.is_symlink(): + raise ValueError(f"symlink path component is forbidden: {current}") + current = current.parent + + +def _rename_no_replace(source: Path, target: Path) -> None: + if os.name == "posix": + libc = ctypes.CDLL(None, use_errno=True) + renameat2 = getattr(libc, "renameat2", None) + if renameat2 is not None: + renameat2.argtypes = [ + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_uint, + ] + renameat2.restype = ctypes.c_int + if renameat2(-100, os.fsencode(source), -100, os.fsencode(target), 1) == 0: + return + error = ctypes.get_errno() + if error == 17: + raise FileExistsError(f"immutable export already exists: {target}") + raise OSError(error, os.strerror(error), str(target)) + raise OSError( + errno.ENOTSUP, + "atomic no-replace publish is unavailable; refusing non-atomic export", + str(target), + ) + + def export_turn_training_bundle( config: TrainingExportConfig, *, @@ -519,14 +901,17 @@ def export_turn_training_bundle( trace_source = _require_regular_source(Path(trace_path), label="trace") target = config.output_dir parent = target.parent + _reject_symlink_ancestors(parent) parent.mkdir(parents=True, exist_ok=True) + _reject_symlink_ancestors(parent) if target.exists() or target.is_symlink(): raise FileExistsError(f"immutable export already exists: {target}") lock_path = parent / f".{target.name}.publish.lock" lock_fd = os.open(lock_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) - temp = Path(tempfile.mkdtemp(prefix=f".{target.name}.tmp-", dir=parent)) + temp: Path | None = None try: os.close(lock_fd) + temp = Path(tempfile.mkdtemp(prefix=f".{target.name}.tmp-", dir=parent)) trace_copy = temp / "artifacts" / "trace" / trace_source.name _copy_source(trace_source, trace_copy) trace_ref = _file_ref(temp, trace_copy) @@ -553,10 +938,18 @@ def export_turn_training_bundle( **config.identity(), "main_agent_turn_count": sum(unit["unit_type"] == "main_agent_turn" for unit in units), "memory_compaction_count": sum(unit["unit_type"] == "memory_compaction" for unit in units), + "capture_error_count": sum(unit["capture_status"] == "error" for unit in units), + "unit_set_sha256": _hash_json( + "researchharness.unit-set.v1", [unit["unit_id"] for unit in units] + ), "termination": _termination(termination), "final_patch_ref": final_patch_ref, "trace_ref": trace_ref, } + summary["integrity"] = { + "canonicalization": CANONICALIZATION, + "summary_payload_sha256": _hash_json("researchharness.summary-payload.v1", summary), + } _write_json(temp / "rollout-summary.v1.json", summary) payload_files = [ @@ -566,6 +959,12 @@ def export_turn_training_bundle( "schema_version": MANIFEST_SCHEMA, **config.identity(), "canonicalization": CANONICALIZATION, + "provenance": { + "harness_revision": config.harness_revision, + "scaffold_version": config.scaffold_version, + "tool_protocol_version": config.tool_protocol_version, + "tokenizer_revision": config.tokenizer_revision, + }, "validators": [{"name": name, "status": "pass"} for name in _VALIDATOR_NAMES], "oracle_leak_unit_count": sum( unit["provenance"]["oracle_isolation"]["status"] == "leak" for unit in units @@ -580,17 +979,17 @@ def export_turn_training_bundle( ) _write_bytes(temp / "SHA256SUMS", checksum_text.encode("utf-8")) validate_turn_training_bundle(temp) - if target.exists() or target.is_symlink(): - raise FileExistsError(f"immutable export appeared during publish: {target}") - os.rename(temp, target) - directory_fd = os.open(parent, os.O_RDONLY) - try: - os.fsync(directory_fd) - finally: - os.close(directory_fd) + _rename_no_replace(temp, target) + if os.name == "posix": + directory_fd = os.open(parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) return target except Exception: - shutil.rmtree(temp, ignore_errors=True) + if temp is not None: + shutil.rmtree(temp, ignore_errors=True) raise finally: try: @@ -603,23 +1002,12 @@ def export_turn_training_bundle( pass -def _walk_refs(value: Any) -> list[Mapping[str, Any]]: - refs: list[Mapping[str, Any]] = [] - if isinstance(value, Mapping): - if set(value) == {"path", "sha256"}: - refs.append(value) - else: - for nested in value.values(): - refs.extend(_walk_refs(nested)) - elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): - for nested in value: - refs.extend(_walk_refs(nested)) - return refs - - def _contains_prohibited_key(value: Any) -> bool: if isinstance(value, Mapping): - return bool(set(value) & _PROHIBITED_AGENT_KEYS) or any(_contains_prohibited_key(nested) for nested in value.values()) + prohibited = any(str(key).casefold() in _PROHIBITED_AGENT_KEYS for key in value) + return prohibited or any( + _contains_prohibited_key(nested) for nested in value.values() + ) if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): return any(_contains_prohibited_key(nested) for nested in value) return False @@ -653,7 +1041,10 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: } if not required <= set(relative_files): raise ValueError("training export is incomplete") - checksum_lines = (root / "SHA256SUMS").read_text(encoding="utf-8").splitlines() + checksum_text = (root / "SHA256SUMS").read_text(encoding="utf-8") + if checksum_text and not checksum_text.endswith("\n"): + raise ValueError("SHA256SUMS must end with a newline") + checksum_lines = checksum_text.splitlines() observed_checksums: dict[str, str] = {} for line in checksum_lines: if len(line) < 67 or line[64:66] != " ": @@ -675,6 +1066,43 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: units = _load_units(root / "turn-training-units.v1.jsonl") if manifest.get("schema_version") != MANIFEST_SCHEMA or summary.get("schema_version") != SUMMARY_SCHEMA: raise ValueError("training export schema mismatch") + expected_manifest_keys = { + "schema_version", + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "run_id", + "canonicalization", + "provenance", + "validators", + "oracle_leak_unit_count", + "files", + } + expected_summary_keys = { + "schema_version", + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "run_id", + "main_agent_turn_count", + "memory_compaction_count", + "capture_error_count", + "unit_set_sha256", + "termination", + "final_patch_ref", + "trace_ref", + "integrity", + } + if _contains_prohibited_key(manifest) or _contains_prohibited_key(summary): + raise ValueError("V1 bundle metadata contains prohibited multi-agent fields") + if set(manifest) != expected_manifest_keys: + raise ValueError("manifest schema contains unexpected fields") + if set(summary) != expected_summary_keys: + raise ValueError("rollout summary schema contains unexpected fields") + if manifest.get("canonicalization") != CANONICALIZATION: + raise ValueError("manifest canonicalization is invalid") manifest_refs = manifest.get("files") if not isinstance(manifest_refs, list): raise ValueError("manifest files must be a list") @@ -685,36 +1113,137 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: expected_manifest_paths = set(relative_files) - {"manifest.json", "SHA256SUMS"} if len(manifest_paths) != len(set(manifest_paths)) or set(manifest_paths) != expected_manifest_paths: raise ValueError("manifest file closure is not exact") + if manifest_paths != sorted(manifest_paths): + raise ValueError("manifest file inventory is not canonically ordered") identity_fields = ("task_id", "task_revision_id", "rollout_id", "rollout_index", "run_id") identity = {key: summary.get(key) for key in identity_fields} + for key in ("task_id", "task_revision_id", "rollout_id", "run_id"): + pattern = _RUN_ID_RE if key == "run_id" else _IDENTITY_RE + if not isinstance(identity[key], str) or pattern.fullmatch(identity[key]) is None: + raise ValueError(f"rollout summary {key} is invalid") + _nonnegative_int(identity["rollout_index"], label="rollout_index") if any(manifest.get(key) != value for key, value in identity.items()): raise ValueError("manifest identity differs from summary") + manifest_provenance = manifest.get("provenance") + provenance_fields = ( + "harness_revision", + "scaffold_version", + "tool_protocol_version", + "tokenizer_revision", + ) + if not isinstance(manifest_provenance, Mapping) or set(manifest_provenance) != set(provenance_fields): + raise ValueError("manifest provenance is invalid") trace_path = _resolve_ref(root, summary.get("trace_ref")) + schema_referenced_paths = { + "turn-training-units.v1.jsonl", + "rollout-summary.v1.json", + trace_path.relative_to(root).as_posix(), + } trace_rows = _read_trace(trace_path) + if any(row.get("run_id") != identity["run_id"] for row in trace_rows): + raise ValueError("trace run_id differs from rollout identity") captures = [row for row in trace_rows if row.get("capture_type") in {"llm_call", "compaction"}] if len(captures) != len(units): raise ValueError("trace capture closure differs from unit count") capture_by_event = {int(row.get("event_index", 0) or 0): row for row in captures} if len(capture_by_event) != len(captures): raise ValueError("trace capture event identities are not unique") - for unit in units: + expected_capture_order = [ + (row["capture_type"], row["event_index"]) for row in captures + ] + actual_capture_order = [ + ( + unit.get("source", {}).get("capture_type"), + unit.get("source", {}).get("trace_event_index"), + ) + for unit in units + ] + if actual_capture_order != expected_capture_order: + raise ValueError("trace capture order/closure differs from training units") + unit_ids: set[str] = set() + main_turn_indexes: list[int] = [] + for unit_position, unit in enumerate(units): if unit.get("schema_version") != UNIT_SCHEMA or unit.get("unit_type") not in {"main_agent_turn", "memory_compaction"}: raise ValueError("training unit schema/type is invalid") - if _contains_prohibited_key(unit): + expected_unit_keys = { + "schema_version", + "unit_type", + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "run_id", + "turn_index", + "capture_status", + "request", + "response", + "source", + "provenance", + "unit_id", + "integrity", + } + if unit.get("unit_type") == "memory_compaction": + expected_unit_keys.add("compaction") + unit_metadata = { + key: value + for key, value in unit.items() + if key not in {"request", "response"} + } + request_keys = unit.get("request", {}).keys() if isinstance(unit.get("request"), Mapping) else () + if _contains_prohibited_key(unit_metadata) or any( + str(key).casefold() in _PROHIBITED_AGENT_KEYS for key in request_keys + ): raise ValueError("V1 unit contains prohibited multi-agent fields") + if set(unit) != expected_unit_keys: + raise ValueError("training unit schema contains unexpected fields") + unit_id = unit.get("unit_id") + if not isinstance(unit_id, str) or unit_id in unit_ids: + raise ValueError("training unit_id is missing or duplicated") + unit_ids.add(unit_id) if any(unit.get(key) != value for key, value in identity.items()): raise ValueError("training unit identity differs from summary") - for ref in _walk_refs(unit): - _resolve_ref(root, ref) - event_index = unit.get("source", {}).get("trace_event_index") + turn_index = _nonnegative_int(unit.get("turn_index"), label="turn_index") + if unit.get("unit_type") == "main_agent_turn": + main_turn_indexes.append(turn_index) + source = unit.get("source") + if not isinstance(source, Mapping) or source.get("trace_ref") != summary.get("trace_ref"): + raise ValueError("training unit trace_ref differs from rollout trace") + if set(source) != { + "trace_ref", + "trace_event_index", + "capture_type", + "provider_request_ref", + "raw_provider_response_ref", + }: + raise ValueError("training unit source schema is invalid") + for source_ref_name in ( + "trace_ref", + "provider_request_ref", + "raw_provider_response_ref", + ): + source_path = _resolve_ref(root, source.get(source_ref_name)) + schema_referenced_paths.add(source_path.relative_to(root).as_posix()) + event_index = _nonnegative_int( + source.get("trace_event_index"), + label="trace_event_index", + ) + if event_index < 1: + raise ValueError("trace_event_index must be positive") row = capture_by_event.get(event_index) expected_capture = "llm_call" if unit.get("unit_type") == "main_agent_turn" else "compaction" - if row is None or row.get("capture_type") != expected_capture: + if row is None or row.get("capture_type") != expected_capture or source.get("capture_type") != expected_capture: raise ValueError("training unit source does not bind its trace capture") integrity = unit.get("integrity") if not isinstance(integrity, dict) or integrity.get("canonicalization") != CANONICALIZATION: raise ValueError("training unit integrity is invalid") + if set(integrity) != { + "canonicalization", + "request_sha256", + "response_sha256", + "unit_payload_sha256", + }: + raise ValueError("training unit integrity schema is invalid") if integrity.get("request_sha256") != _hash_json("researchharness.request.v1", unit.get("request")): raise ValueError("training unit request digest mismatch") if integrity.get("response_sha256") != _hash_json("researchharness.response.v1", unit.get("response")): @@ -740,7 +1269,6 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: { "request_sha256": integrity["request_sha256"], "response_sha256": integrity["response_sha256"], - "trace_event_index": event_index, } ) if unit.get("unit_id") != _hash_json("researchharness.unit-id.v1", identity_preimage): @@ -754,33 +1282,199 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: if raw_response != _raw_response(payload, compaction=expected_capture == "compaction"): raise ValueError("raw provider response differs from trace") response_source = _response_source(payload, compaction=expected_capture == "compaction") - if unit.get("response") != _normalized_response(response_source): + expected_response = _normalized_response(response_source) + if unit.get("response") != expected_response: raise ValueError("exported response differs from accepted trace response") + if response_source.get("status") == "ok": + normalized_raw_response = _normalized_raw_provider_response(raw_response) + if expected_response != normalized_raw_response: + raise ValueError("accepted response differs from raw provider response") + _validate_tool_calls(expected_response) + _validate_provider_attempts( + response_source, + provider_request=provider_request, + raw_response=raw_response, + ) + if expected_capture == "llm_call": + _validate_tool_result_closure( + trace_rows=trace_rows, + capture_row=row, + response=expected_response, + ) findings = _oracle_findings(provider_request) - oracle = unit.get("provenance", {}).get("oracle_isolation", {}) + provenance = unit.get("provenance") + if not isinstance(provenance, Mapping): + raise ValueError("training unit provenance is missing") + if set(provenance) != { + "harness_revision", + "scaffold_version", + "tool_protocol_version", + "requested_model", + "provider_model", + "model_revision", + "tokenizer_revision", + "provider_retry_count", + "oracle_isolation", + }: + raise ValueError("training unit provenance schema is invalid") + if any(provenance.get(key) != manifest_provenance.get(key) for key in provenance_fields): + raise ValueError("training unit provenance differs from manifest") + oracle = provenance.get("oracle_isolation", {}) if oracle != {"status": "leak" if findings else "clear", "findings": findings}: raise ValueError("oracle leak marker differs from model-visible request") - if findings and unit.get("capture_status") != "error": - raise ValueError("oracle-leaking unit is not fail-closed") + source_response = payload.get("response") + if not isinstance(source_response, Mapping): + source_response = payload.get("summary_response") if isinstance(payload.get("summary_response"), Mapping) else {} + expected_retry_count = _nonnegative_int( + payload.get("provider_retry_count", source_response.get("provider_retry_count", 0)), + label="provider_retry_count", + ) + expected_requested_model = ( + payload.get("requested_model") + or source_response.get("requested_model") + or payload.get("model_name") + ) + expected_provider_model = payload.get("provider_model") or source_response.get("provider_model") + if ( + provenance.get("provider_retry_count") != expected_retry_count + or provenance.get("requested_model") != expected_requested_model + or provenance.get("provider_model") != expected_provider_model + or provenance.get("model_revision") is not None + ): + raise ValueError("training unit provider provenance differs from trace") + expected_status = ( + _compaction_capture_ok(payload, findings) + if expected_capture == "compaction" + else _main_capture_ok(payload, findings) + ) + _validate_request_shape(unit["request"], capture_ok=expected_status) + if unit.get("capture_status") != ("ok" if expected_status else "error"): + raise ValueError("capture_status differs from provider-boundary evidence") if expected_capture == "compaction": compaction = unit.get("compaction") if not isinstance(compaction, Mapping): raise ValueError("memory compaction unit lacks compaction metadata") - pre = json.loads(_resolve_ref(root, compaction.get("pre_messages_ref")).read_text(encoding="utf-8")) - post = json.loads(_resolve_ref(root, compaction.get("post_messages_ref")).read_text(encoding="utf-8")) + if set(compaction) != { + "trigger_reason", + "prior_token_estimate", + "new_token_estimate", + "compacted_group_count", + "kept_group_count", + "existing_memory_text_sha256", + "pre_messages_ref", + "post_messages_ref", + }: + raise ValueError("memory compaction metadata schema is invalid") + pre = _read_canonical_json(_resolve_ref(root, compaction.get("pre_messages_ref"))) + post = _read_canonical_json(_resolve_ref(root, compaction.get("post_messages_ref"))) + for compaction_ref_name in ("pre_messages_ref", "post_messages_ref"): + compaction_path = _resolve_ref(root, compaction.get(compaction_ref_name)) + schema_referenced_paths.add(compaction_path.relative_to(root).as_posix()) + if not isinstance(pre, list) or not isinstance(post, list): + raise ValueError("compaction pre/post artifacts must be message lists") if pre != payload.get("pre_messages", []) or post != payload.get("post_messages", []): raise ValueError("compaction pre/post artifact differs from trace") - - if summary.get("main_agent_turn_count") != sum(unit.get("unit_type") == "main_agent_turn" for unit in units): + if provider_request and payload.get("summary_request") != provider_request.get("messages"): + raise ValueError( + "compaction summary_request differs from the provider request messages" + ) + existing_memory = str(payload.get("existing_memory_text", "") or "") + expected_compaction = { + "trigger_reason": payload.get("trigger_reason", ""), + "prior_token_estimate": _nonnegative_int(payload.get("prior_token_estimate", 0), label="prior_token_estimate"), + "new_token_estimate": _nonnegative_int(payload.get("new_token_estimate", 0), label="new_token_estimate"), + "compacted_group_count": _nonnegative_int(payload.get("compacted_group_count", 0), label="compacted_group_count"), + "kept_group_count": _nonnegative_int(payload.get("kept_group_count", 0), label="kept_group_count"), + "existing_memory_text_sha256": "sha256:" + _hash_bytes(existing_memory.encode("utf-8")), + } + for key, expected_value in expected_compaction.items(): + if compaction.get(key) != expected_value: + raise ValueError(f"compaction {key} differs from trace") + if expected_status: + next_main_capture = next( + ( + later + for later in captures[unit_position + 1 :] + if later.get("capture_type") == "llm_call" + ), + None, + ) + if next_main_capture is not None: + from agent_base.react_agent import prepare_messages_for_llm + + prepared_post, _ = prepare_messages_for_llm(post) + next_payload = ( + next_main_capture.get("payload") + if isinstance(next_main_capture.get("payload"), Mapping) + else {} + ) + next_provider_request = _provider_request(next_payload) + if next_provider_request.get("messages") != prepared_post: + raise ValueError( + "successful compaction post_messages differ from the next provider request" + ) + + if main_turn_indexes != sorted(set(main_turn_indexes)): + raise ValueError("main-agent turn_index values must be strictly increasing") + summary_main_count = _nonnegative_int( + summary.get("main_agent_turn_count"), + label="main_agent_turn_count", + ) + summary_compaction_count = _nonnegative_int( + summary.get("memory_compaction_count"), + label="memory_compaction_count", + ) + summary_error_count = _nonnegative_int( + summary.get("capture_error_count"), + label="capture_error_count", + ) + if summary_main_count != sum(unit.get("unit_type") == "main_agent_turn" for unit in units): raise ValueError("rollout summary main-agent count mismatch") - if summary.get("memory_compaction_count") != sum(unit.get("unit_type") == "memory_compaction" for unit in units): + if summary_compaction_count != sum(unit.get("unit_type") == "memory_compaction" for unit in units): raise ValueError("rollout summary compaction count mismatch") + if summary_error_count != sum(unit.get("capture_status") == "error" for unit in units): + raise ValueError("rollout summary capture-error count mismatch") + if summary.get("unit_set_sha256") != _hash_json( + "researchharness.unit-set.v1", [unit["unit_id"] for unit in units] + ): + raise ValueError("rollout summary unit-set digest mismatch") + summary_integrity = summary.get("integrity") + if not isinstance(summary_integrity, Mapping) or summary_integrity.get("canonicalization") != CANONICALIZATION: + raise ValueError("rollout summary integrity is invalid") + if set(summary_integrity) != { + "canonicalization", + "summary_payload_sha256", + }: + raise ValueError("rollout summary integrity schema is invalid") + summary_payload = json.loads(json.dumps(summary, ensure_ascii=False)) + summary_payload.pop("integrity", None) + if summary_integrity.get("summary_payload_sha256") != _hash_json( + "researchharness.summary-payload.v1", summary_payload + ): + raise ValueError("rollout summary payload digest mismatch") + termination = summary.get("termination") + if not isinstance(termination, Mapping): + raise ValueError("rollout termination is missing") + if _termination(termination) != dict(termination): + raise ValueError("rollout termination is not canonical") + final_patch_ref = summary.get("final_patch_ref") + if final_patch_ref is not None: + final_patch_path = _resolve_ref(root, final_patch_ref) + schema_referenced_paths.add(final_patch_path.relative_to(root).as_posix()) + if set(manifest_paths) != schema_referenced_paths: + raise ValueError("manifest contains unreferenced or missing payload artifacts") + if "resolved" in summary or "reward" in summary: + raise ValueError("Harness export cannot contain resolved or reward labels") validators = manifest.get("validators") expected_validators = [{"name": name, "status": "pass"} for name in _VALIDATOR_NAMES] if validators != expected_validators: raise ValueError("manifest validator closure is not canonical") leak_count = sum(unit.get("provenance", {}).get("oracle_isolation", {}).get("status") == "leak" for unit in units) - if manifest.get("oracle_leak_unit_count") != leak_count: + manifest_leak_count = _nonnegative_int( + manifest.get("oracle_leak_unit_count"), + label="oracle_leak_unit_count", + ) + if manifest_leak_count != leak_count: raise ValueError("manifest oracle leak count mismatch") return summary diff --git a/api/openai_server.py b/api/openai_server.py index b3ce615..49032c8 100644 --- a/api/openai_server.py +++ b/api/openai_server.py @@ -28,6 +28,7 @@ resolve_extra_tool_names, ) from agent_base.provider_compat import normalize_omit_generate_params +from agent_base.training_export import EXPORT_DIR_NAME, TrainingExportConfig from agent_base.tools.tooling import normalize_workspace_root from agent_base.utils import append_jsonl, image_input_content_parts, read_role_prompt_files, safe_jsonable @@ -53,6 +54,7 @@ ) REQUEST_LLM_EXTRA_BODY_FIELD = "llm-extra-body" REQUEST_OMIT_GENERATE_PARAMS_FIELD = "omit-generate-params" +REQUEST_TRAINING_EXPORT_FIELD = "training-export" INPUT_WRAPPER_SYSTEM_PROMPT = """You are the ResearchHarness input wrapper. @@ -441,6 +443,88 @@ def request_omit_generate_params(payload: dict[str, Any]) -> tuple[str, ...] | N raise OpenAICompatError(400, f"{REQUEST_OMIT_GENERATE_PARAMS_FIELD}: {exc}") from exc +def request_training_export( + payload: dict[str, Any], + *, + run_root: Path, + workspace_root: Path, + run_id: str, +) -> TrainingExportConfig | None: + raw_value = payload.get(REQUEST_TRAINING_EXPORT_FIELD) + if raw_value is None: + return None + if not isinstance(raw_value, dict): + raise OpenAICompatError(400, f"{REQUEST_TRAINING_EXPORT_FIELD} must be a JSON object.") + allowed = { + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "harness_revision", + "scaffold_version", + "tool_protocol_version", + "tokenizer_revision", + "final_patch_path", + } + unknown = sorted(set(raw_value) - allowed) + if unknown: + raise OpenAICompatError( + 400, + f"{REQUEST_TRAINING_EXPORT_FIELD} contains unsupported fields: {', '.join(unknown)}", + ) + required = ( + "task_id", + "task_revision_id", + "rollout_id", + "rollout_index", + "harness_revision", + ) + missing = [name for name in required if raw_value.get(name) is None] + if missing: + raise OpenAICompatError( + 400, + f"{REQUEST_TRAINING_EXPORT_FIELD} requires: {', '.join(missing)}", + ) + for name in ("task_id", "task_revision_id", "rollout_id", "harness_revision"): + if not isinstance(raw_value.get(name), str) or not raw_value[name].strip(): + raise OpenAICompatError(400, f"training-export.{name} must be a non-empty string.") + rollout_index = raw_value.get("rollout_index") + if isinstance(rollout_index, bool) or not isinstance(rollout_index, int) or rollout_index < 0: + raise OpenAICompatError(400, "training-export.rollout_index must be a non-negative integer.") + for name in ("scaffold_version", "tool_protocol_version", "tokenizer_revision"): + if name in raw_value and raw_value[name] is not None and not isinstance(raw_value[name], str): + raise OpenAICompatError(400, f"training-export.{name} must be a string.") + final_patch_path: Path | None = None + if raw_value.get("final_patch_path") is not None: + if not isinstance(raw_value["final_patch_path"], str): + raise OpenAICompatError(400, "training-export.final_patch_path must be a string.") + raw_patch = str(raw_value["final_patch_path"]).strip() + patch = Path(raw_patch) + if not raw_patch or patch.is_absolute(): + raise OpenAICompatError(400, "training-export.final_patch_path must be a relative workspace path.") + final_patch_path = workspace_root / patch + try: + final_patch_path.resolve(strict=False).relative_to(workspace_root.resolve()) + except ValueError as exc: + raise OpenAICompatError(400, "training-export.final_patch_path escapes the workspace.") from exc + try: + return TrainingExportConfig( + output_dir=run_root / EXPORT_DIR_NAME, + task_id=raw_value["task_id"], + task_revision_id=raw_value["task_revision_id"], + rollout_id=raw_value["rollout_id"], + rollout_index=rollout_index, + run_id=run_id, + harness_revision=raw_value["harness_revision"], + scaffold_version=raw_value.get("scaffold_version", ""), + tool_protocol_version=raw_value.get("tool_protocol_version", "openai-chat-completions-native-tools-v1"), + tokenizer_revision=raw_value.get("tokenizer_revision"), + final_patch_path=final_patch_path, + ) + except (TypeError, ValueError) as exc: + raise OpenAICompatError(400, f"invalid {REQUEST_TRAINING_EXPORT_FIELD}: {exc}") from exc + + def resolve_request_workspace_root(payload: dict[str, Any], default_workspace_root: Path) -> tuple[Path, dict[str, Any]]: for alias in REQUEST_WORKSPACE_ROOT_ALIAS_FIELDS: if alias in payload: @@ -491,6 +575,13 @@ def append_api_event(trace_dir: Path, event: str, payload: dict[str, Any]) -> No def run_chat_completion(payload: dict[str, Any], config: ServerConfig) -> dict[str, Any]: payload = validate_chat_payload(payload) + if payload.get(REQUEST_TRAINING_EXPORT_FIELD) is not None and ( + config.input_wrapper or config.output_wrapper + ): + raise OpenAICompatError( + 400, + "V1 training export cannot be combined with input/output wrappers because wrapper model calls are outside the two V1 unit types.", + ) requested_model_label, backend_model = resolve_api_model_selection(str(payload.get("model") or API_MODEL_ALIAS)) payload["model"] = requested_model_label llm_extra_body = request_llm_extra_body(payload) @@ -504,6 +595,12 @@ def run_chat_completion(payload: dict[str, Any], config: ServerConfig) -> dict[s if workspace_meta["source"] == "default": agent_workspace.mkdir(parents=True, exist_ok=False) trace_dir.mkdir(parents=True, exist_ok=False) + training_export = request_training_export( + payload, + run_root=run_root, + workspace_root=agent_workspace, + run_id=run_id, + ) append_api_event( trace_dir, "workspace_selection", @@ -544,6 +641,8 @@ def run_chat_completion(payload: dict[str, Any], config: ServerConfig) -> dict[s llm=llm_config, trace_dir=str(trace_dir), role_prompt=config.role_prompt or None, + training_export=training_export, + run_id=run_id, ) if config.input_wrapper: @@ -589,6 +688,7 @@ def run_chat_completion(payload: dict[str, Any], config: ServerConfig) -> dict[s "termination": session.get("termination", ""), "result_text": agent_result_text, "trace_path": session.get("trace_path", ""), + "training_export_path": session.get("training_export_path", ""), }, ) diff --git a/docs/turn-training-export-v1.md b/docs/turn-training-export-v1.md new file mode 100644 index 0000000..3cef853 --- /dev/null +++ b/docs/turn-training-export-v1.md @@ -0,0 +1,124 @@ +# Turn-level training export V1 + +ResearchHarness can publish one immutable, self-validating training bundle for +one completed rollout. Each real provider call becomes one training unit, so a +downstream data job does not need to reconstruct model rounds from UI events. + +V1 has exactly two unit types: `main_agent_turn` and `memory_compaction`. +Compaction is an isolated semantic model call, not a sub-agent. + +## Bundle + +The output path must end in `turn-training-export-v1`: + +```text +turn-training-export-v1/ + turn-training-units.v1.jsonl + rollout-summary.v1.json + agent.patch # only when explicitly supplied + artifacts/trace/ + artifacts/provider/ + artifacts/compaction/ + manifest.json + SHA256SUMS +``` + +The Harness writes a temporary sibling, validates it, then publishes it with +no-replace semantics. Existing exports are never overwritten. Exact final +provider kwargs and full raw provider responses are retained as hashed +artifacts. Failed calls remain auditable with `capture_status="error"` and are +not SFT candidates. + +The normalized target is derived and revalidated against the raw provider +response. V1 applies one explicit canonical transform: an assistant +`message.content` of JSON `null` becomes the empty string `""`; all other +content, reasoning content, finish reason, usage, tool-call order, IDs, names, +and argument strings must match the raw first choice. Every retry attempt is +kept in the trace, and the final attempt must close against the exported +provider request and raw response. + +For tool rounds, each non-truncated tool-call ID must close against its tool +result and, when another model call occurs, that result must be present in the +next real provider input (directly or through a validated compaction boundary). +For compaction, `summary_request` must exactly equal the messages sent to the +provider, and successful `post_messages` must match the next real model input +after the documented image-aging transform. + +The Harness does not write `resolved`, reward, advantage, or verifier outcome. +Downstream code joins those outcomes by +`task_id + task_revision_id + rollout_id`. + +## CLI + +```bash +python -m agent_base.react_agent \ + "Fix the authentication regression." \ + --workspace-root /work/repo \ + --training-export-dir /outputs/r0/turn-training-export-v1 \ + --task-id auth-regression \ + --task-revision-id auth-regression@base-abc123 \ + --rollout-id r0 \ + --rollout-index 0 \ + --run-id run-auth-r0 \ + --harness-revision 72e8a44737aadabc \ + --final-patch-path /outputs/r0/agent.patch +``` + +`--run-id` is generated when omitted. `--trace-dir` defaults to a sibling +`trace/`. The patch is optional and never guessed. V1 export cannot be combined +with CLI chat mode because one immutable config represents one rollout. + +## Python API + +`create_agent(...)` and `run_agent(...)` accept the same identity fields: + +```python +from researchharness import run_agent +result = run_agent( + "Fix the authentication regression.", + workspace_root="/work/repo", + training_export_dir="/outputs/r0/turn-training-export-v1", + task_id="auth-regression", + task_revision_id="auth-regression@base-abc123", + rollout_id="r0", + rollout_index=0, + run_id="run-auth-r0", + harness_revision="72e8a44737aadabc", + final_patch_path="/outputs/r0/agent.patch", +) +``` + +## OpenAI-compatible API + +Add a `training-export` object to the request: + +```json +{ + "model": "RH", + "messages": [{"role": "user", "content": "Fix the auth regression."}], + "training-export": { + "task_id": "auth-regression", + "task_revision_id": "auth-regression@base-abc123", + "rollout_id": "r0", + "rollout_index": 0, + "harness_revision": "72e8a44737aadabc", + "final_patch_path": "agent.patch" + } +} +``` + +The API uses one outer `run_id` for the agent trace and export, and writes the +bundle under `//turn-training-export-v1`. API patch paths +are relative to the selected workspace. Training export fails closed when input +or output LLM wrappers are enabled because wrapper calls are outside V1's two +unit types. + +## Validation + +`validate_turn_training_bundle(path)` checks trace/order closure, exact request +and response matching, structured tool calls, compaction pre/post and next-input +continuity, frozen identity, canonical hashes and refs, oracle leakage, +multi-agent fields, path traversal, symlinks, hardlinks, and overwrite safety. + +Only units with `capture_status="ok"` are SFT candidates. Selection policy is a +downstream responsibility. diff --git a/researchharness/runtime.py b/researchharness/runtime.py index 08e4f2c..2cfe860 100644 --- a/researchharness/runtime.py +++ b/researchharness/runtime.py @@ -4,9 +4,11 @@ from pathlib import Path from typing import Any, Optional, Sequence +from uuid import uuid4 from agent_base.react_agent import ALL_TOOL_MAP, MultiTurnReactAgent, default_llm_config, default_tool_names from agent_base.provider_compat import normalize_omit_generate_params +from agent_base.training_export import TrainingExportConfig from agent_base.tools.custom import build_custom_tool_map, tool from agent_base.tools.tooling import ToolBase from agent_base.utils import load_default_dotenvs, read_role_prompt_files, require_required_env @@ -124,6 +126,69 @@ def _read_role_prompt_blocks(role_prompt_files: Optional[str | Path | Sequence[s return read_role_prompt_files(paths) +def _resolve_training_export( + *, + training_export_dir: Optional[str | Path], + task_id: Optional[str], + task_revision_id: Optional[str], + rollout_id: Optional[str], + rollout_index: Optional[int], + run_id: Optional[str], + harness_revision: Optional[str], + scaffold_version: str, + tool_protocol_version: str, + tokenizer_revision: Optional[str], + final_patch_path: Optional[str | Path], + trace_dir: Optional[str], +) -> tuple[Optional[TrainingExportConfig], Optional[str], Optional[str]]: + requested = training_export_dir is not None or any( + value is not None + for value in ( + task_id, + task_revision_id, + rollout_id, + rollout_index, + harness_revision, + final_patch_path, + ) + ) + if not requested: + return None, trace_dir, run_id + missing = [ + name + for name, value in ( + ("training_export_dir", training_export_dir), + ("task_id", task_id), + ("task_revision_id", task_revision_id), + ("rollout_id", rollout_id), + ("rollout_index", rollout_index), + ("harness_revision", harness_revision), + ) + if value is None + ] + if missing: + raise ValueError("turn-level training export requires: " + ", ".join(missing)) + if run_id is not None and not isinstance(run_id, str): + raise ValueError("run_id must be a string") + resolved_run_id = run_id if run_id is not None else uuid4().hex + output_dir = Path(training_export_dir) # type: ignore[arg-type] + resolved_trace_dir = trace_dir or str(output_dir.parent / "trace") + config = TrainingExportConfig( + output_dir=output_dir, + task_id=task_id, # type: ignore[arg-type] + task_revision_id=task_revision_id, # type: ignore[arg-type] + rollout_id=rollout_id, # type: ignore[arg-type] + rollout_index=rollout_index, # type: ignore[arg-type] + run_id=resolved_run_id, + harness_revision=harness_revision, # type: ignore[arg-type] + scaffold_version=scaffold_version, + tool_protocol_version=tool_protocol_version, + tokenizer_revision=tokenizer_revision, + final_patch_path=Path(final_patch_path) if final_patch_path is not None else None, + ) + return config, resolved_trace_dir, resolved_run_id + + def create_agent( *, model_name: Optional[str] = None, @@ -149,6 +214,17 @@ def create_agent( role_prompt_files: Optional[str | Path | Sequence[str | Path]] = None, tools: Optional[Sequence[Any]] = None, extra_tools: Optional[Sequence[str]] = None, + training_export_dir: Optional[str | Path] = None, + task_id: Optional[str] = None, + task_revision_id: Optional[str] = None, + rollout_id: Optional[str] = None, + rollout_index: Optional[int] = None, + run_id: Optional[str] = None, + harness_revision: Optional[str] = None, + scaffold_version: str = "", + tool_protocol_version: str = "openai-chat-completions-native-tools-v1", + tokenizer_revision: Optional[str] = None, + final_patch_path: Optional[str | Path] = None, require_env: bool = True, ) -> MultiTurnReactAgent: """Create a configured single-agent ResearchHarness runtime.""" @@ -162,6 +238,20 @@ def create_agent( role_blocks.append(_read_role_prompt_blocks(role_prompt_files)) resolved_role_prompt = "\n\n".join(block for block in role_blocks if block.strip()) function_list, custom_tools = _resolve_tools(tools, extra_tools) + training_export, trace_dir, resolved_run_id = _resolve_training_export( + training_export_dir=training_export_dir, + task_id=task_id, + task_revision_id=task_revision_id, + rollout_id=rollout_id, + rollout_index=rollout_index, + run_id=run_id, + harness_revision=harness_revision, + scaffold_version=scaffold_version, + tool_protocol_version=tool_protocol_version, + tokenizer_revision=tokenizer_revision, + final_patch_path=final_patch_path, + trace_dir=trace_dir, + ) llm = _apply_llm_overrides( default_llm_config(model_name=model_name), api_key=api_key, @@ -188,6 +278,8 @@ def create_agent( custom_tools=custom_tools, max_rounds=max_rounds, max_runtime_seconds=max_runtime_seconds, + training_export=training_export, + run_id=resolved_run_id, ) @@ -218,6 +310,17 @@ def run_agent( images: Optional[str | Path | Sequence[str | Path]] = None, tools: Optional[Sequence[Any]] = None, extra_tools: Optional[Sequence[str]] = None, + training_export_dir: Optional[str | Path] = None, + task_id: Optional[str] = None, + task_revision_id: Optional[str] = None, + rollout_id: Optional[str] = None, + rollout_index: Optional[int] = None, + run_id: Optional[str] = None, + harness_revision: Optional[str] = None, + scaffold_version: str = "", + tool_protocol_version: str = "openai-chat-completions-native-tools-v1", + tokenizer_revision: Optional[str] = None, + final_patch_path: Optional[str | Path] = None, require_env: bool = True, ) -> str: """Run ResearchHarness once and return the final assistant text.""" @@ -245,6 +348,17 @@ def run_agent( role_prompt_files=role_prompt_files, tools=tools, extra_tools=extra_tools, + training_export_dir=training_export_dir, + task_id=task_id, + task_revision_id=task_revision_id, + rollout_id=rollout_id, + rollout_index=rollout_index, + run_id=run_id, + harness_revision=harness_revision, + scaffold_version=scaffold_version, + tool_protocol_version=tool_protocol_version, + tokenizer_revision=tokenizer_revision, + final_patch_path=final_patch_path, require_env=require_env, ) return agent.run(prompt, images=images) diff --git a/tests/test_openai_api_checks.py b/tests/test_openai_api_checks.py index 2c5258a..ee45319 100644 --- a/tests/test_openai_api_checks.py +++ b/tests/test_openai_api_checks.py @@ -146,8 +146,11 @@ def call_llm_api(self, msgs, max_tries=10, runtime_deadline=None): fake_seen: dict[str, str] = {} class FakeAPIAgent: - def __init__(self, function_list, llm, trace_dir, role_prompt=None): + def __init__(self, function_list, llm, trace_dir, role_prompt=None, training_export=None, run_id=None): self.trace_dir = Path(trace_dir) + self.training_export = training_export + self.run_id = run_id + fake_seen.setdefault("run_ids", []).append(run_id) fake_seen["trace_dir"] = str(self.trace_dir) fake_seen.setdefault("trace_dirs", []).append(str(self.trace_dir)) fake_seen.setdefault("function_lists", []).append(list(function_list or [])) diff --git a/tests/test_turn_training_export.py b/tests/test_turn_training_export.py index a7b6429..33f985e 100644 --- a/tests/test_turn_training_export.py +++ b/tests/test_turn_training_export.py @@ -4,6 +4,7 @@ import types import unittest import sys +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from unittest.mock import patch @@ -49,6 +50,39 @@ def _accepted_response(request, *, response_id, content, tool_calls, retry_count ], "usage": {"prompt_tokens": 100, "completion_tokens": 20}, } + attempts = [] + for attempt_index in range(retry_count): + attempts.append( + { + "attempt_index": attempt_index, + "provider_request": request, + "raw_provider_response": { + "id": f"{response_id}-rejected-{attempt_index}", + "model": MODEL, + "choices": [ + { + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [], + }, + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 0}, + }, + "status": "rejected_empty", + "error": "empty response from llm api", + } + ) + attempts.append( + { + "attempt_index": retry_count, + "provider_request": request, + "raw_provider_response": raw, + "status": "accepted", + } + ) return { "status": "ok", "finish_reason": "tool_calls" if tool_calls else "stop", @@ -59,6 +93,7 @@ def _accepted_response(request, *, response_id, content, tool_calls, retry_count "provider_request": request, "raw_provider_response": raw, "provider_retry_count": retry_count, + "provider_attempts": attempts, "requested_model": MODEL, "provider_model": MODEL, } @@ -289,13 +324,27 @@ def _export(tmp_path, *, name="case", rows=None, patch_file=True): class FakeMessage: - def __init__(self, content): + def __init__(self, content, tool_calls=None): self.content = content - self.tool_calls = None + self.tool_calls = list(tool_calls or []) self.reasoning_content = None def model_dump(self): - return {"role": "assistant", "content": self.content, "tool_calls": None} + return { + "role": "assistant", + "content": self.content, + "tool_calls": [ + { + "id": tool_call.id, + "type": "function", + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments, + }, + } + for tool_call in self.tool_calls + ], + } class FakeUsage: @@ -304,12 +353,16 @@ def model_dump(self): class FakeResponse: - def __init__(self, response_id, content): + def __init__(self, response_id, content, tool_calls=None): self.id = response_id self.model = MODEL self.usage = FakeUsage() + message = FakeMessage(content, tool_calls=tool_calls) self.choices = [ - types.SimpleNamespace(finish_reason="stop", message=FakeMessage(content)) + types.SimpleNamespace( + finish_reason="tool_calls" if message.tool_calls else "stop", + message=message, + ) ] def model_dump(self): @@ -318,14 +371,49 @@ def model_dump(self): "model": self.model, "choices": [ { - "finish_reason": "stop", - "message": {"role": "assistant", "content": self.choices[0].message.content}, + "finish_reason": self.choices[0].finish_reason, + "message": self.choices[0].message.model_dump(), } ], "usage": self.usage.model_dump(), } +class FakeMalformedResponse: + def __init__(self, response_id): + self.id = response_id + self.model = MODEL + self.usage = FakeUsage() + self.choices = [] + + def model_dump(self): + return { + "id": self.id, + "model": self.model, + "choices": [], + "usage": self.usage.model_dump(), + } + + +class FakeResponseWithoutDump: + def __init__(self, content): + self.model = MODEL + self.usage = FakeUsage() + self.choices = [ + types.SimpleNamespace( + finish_reason="stop", + message=FakeMessage(content), + ) + ] + + +def _fake_tool_call(call_id, name, arguments): + return types.SimpleNamespace( + id=call_id, + function=types.SimpleNamespace(name=name, arguments=arguments), + ) + + class FakeClient: def __init__(self, responses): self.responses = list(responses) @@ -413,6 +501,35 @@ def test_successful_compaction_must_match_next_real_request(self): export_turn_training_bundle(config, trace_path=trace, termination="result") self.assertFalse(config.output_dir.exists()) + def test_compaction_next_input_uses_real_image_aging_transform(self): + from agent_base.react_agent import prepare_messages_for_llm + + rows = _trace_rows()[1:] + rows[0]["event_index"] = 1 + rows[1]["event_index"] = 2 + post_messages = rows[0]["payload"]["post_messages"] + post_messages[1]["content"] = [ + {"type": "text", "text": "Inspect this screenshot."}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + }, + ] + aged_messages, aging = prepare_messages_for_llm(post_messages) + self.assertEqual(aging["omitted_image_count"], 1) + rows[1]["payload"]["provider_request"]["messages"] = aged_messages + _, bundle = _export( + self.tmp_path, + name="image-aging-compaction", + rows=rows, + patch_file=False, + ) + units = _units(bundle) + self.assertEqual(validate_turn_training_bundle(bundle)["memory_compaction_count"], 1) + serialized_next_input = json.dumps(units[1]["request"]["messages"]) + self.assertIn("omitted", serialized_next_input) + self.assertNotIn("image_url", serialized_next_input) + def test_oracle_request_is_marked_error_for_downstream_fail_closed(self): rows = _trace_rows() rows[0]["payload"]["provider_request"]["messages"][1]["content"] = "Apply the GOLD PATCH." @@ -436,6 +553,26 @@ def test_existing_target_and_wrong_directory_name_are_rejected(self): run_id="run", harness_revision="rev", ) + def test_concurrent_publish_has_exactly_one_winner(self): + config = _config(self.tmp_path, name="concurrent", patch_file=False) + trace = _write_trace(self.tmp_path / "concurrent/trace.jsonl", [_trace_rows()[2]]) + + def publish(): + try: + return export_turn_training_bundle( + config, + trace_path=trace, + termination="result", + ) + except FileExistsError: + return None + + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(lambda _: publish(), range(2))) + self.assertEqual(sum(result == config.output_dir for result in results), 1) + self.assertEqual(sum(result is None for result in results), 1) + validate_turn_training_bundle(config.output_dir) + def test_structured_target_tamper_fails_even_after_resealing(self): _, bundle = _export(self.tmp_path) units = _units(bundle) @@ -446,6 +583,100 @@ def test_structured_target_tamper_fails_even_after_resealing(self): with self.assertRaisesRegex(ValueError, "accepted trace response"): validate_turn_training_bundle(bundle) + def test_raw_provider_response_must_match_normalized_target(self): + rows = _trace_rows() + payload = rows[0]["payload"] + divergent_raw = json.loads(json.dumps(payload["raw_provider_response"])) + divergent_raw["choices"][0]["finish_reason"] = "stop" + divergent_raw["choices"][0]["message"]["tool_calls"] = [] + payload["raw_provider_response"] = divergent_raw + payload["response"]["raw_provider_response"] = divergent_raw + payload["response"]["provider_attempts"][-1]["raw_provider_response"] = divergent_raw + config = _config(self.tmp_path, name="raw-divergence", patch_file=False) + trace = _write_trace(self.tmp_path / "raw-divergence/trace.jsonl", rows) + with self.assertRaisesRegex(ValueError, "raw provider response"): + export_turn_training_bundle(config, trace_path=trace, termination="result") + self.assertFalse(config.output_dir.exists()) + + def test_compaction_summary_request_must_match_real_provider_request(self): + rows = _trace_rows() + rows[1]["payload"]["summary_request"] = [ + {"role": "user", "content": "not the request sent to the provider"} + ] + config = _config(self.tmp_path, name="summary-divergence", patch_file=False) + trace = _write_trace(self.tmp_path / "summary-divergence/trace.jsonl", rows) + with self.assertRaisesRegex(ValueError, "summary_request"): + export_turn_training_bundle(config, trace_path=trace, termination="result") + self.assertFalse(config.output_dir.exists()) + + def test_tool_call_requires_matching_result_in_next_model_input(self): + rows = [_trace_rows()[0], _trace_rows()[2]] + rows[1]["event_index"] = 2 + request = rows[1]["payload"]["provider_request"] + request["messages"] = [ + {"role": "system", "content": "Fix the repository using tools."}, + {"role": "user", "content": "Inspect src/auth.py and fix token expiry."}, + ] + config = _config(self.tmp_path, name="missing-tool-result", patch_file=False) + trace = _write_trace(self.tmp_path / "missing-tool-result/trace.jsonl", rows) + with self.assertRaisesRegex(ValueError, "tool-result"): + export_turn_training_bundle(config, trace_path=trace, termination="result") + + def test_unit_id_does_not_depend_on_trace_event_position(self): + row_one = _trace_rows()[2] + row_one["event_index"] = 1 + row_two = json.loads(json.dumps(row_one)) + row_two["event_index"] = 17 + _, first_bundle = _export( + self.tmp_path, + name="unit-id-one", + rows=[row_one], + patch_file=False, + ) + _, second_bundle = _export( + self.tmp_path, + name="unit-id-two", + rows=[row_two], + patch_file=False, + ) + self.assertEqual(_units(first_bundle)[0]["unit_id"], _units(second_bundle)[0]["unit_id"]) + + def test_provider_owned_path_sha_mapping_is_not_an_artifact_ref(self): + rows = _trace_rows() + provider_owned = { + "path": "provider-owned-id", + "sha256": "0" * 64, + } + rows[0]["payload"]["provider_request"]["extra_body"] = { + "provider_metadata": provider_owned + } + _, bundle = _export( + self.tmp_path, + name="provider-path-object", + rows=rows, + patch_file=False, + ) + self.assertEqual(validate_turn_training_bundle(bundle)["capture_error_count"], 0) + + def test_outcome_and_agent_metadata_fields_are_rejected(self): + _, unit_bundle = _export(self.tmp_path, name="unit-outcome") + units = _units(unit_bundle) + units[0]["reward"] = 1.0 + export_module._finish_unit(units[0]) + _write_units(unit_bundle, units) + _reseal_bundle(unit_bundle) + with self.assertRaisesRegex(ValueError, "schema"): + validate_turn_training_bundle(unit_bundle) + + _, manifest_bundle = _export(self.tmp_path, name="manifest-agent") + manifest_path = manifest_bundle / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["agent_id"] = "invented" + _write_json(manifest_path, manifest) + _reseal_bundle(manifest_bundle) + with self.assertRaisesRegex(ValueError, "multi-agent"): + validate_turn_training_bundle(manifest_bundle) + def test_summary_counts_are_recomputed(self): _, bundle = _export(self.tmp_path) summary = json.loads((bundle / "rollout-summary.v1.json").read_text()) @@ -473,5 +704,216 @@ def test_prohibited_agent_field_and_unsafe_ref_fail_closed(self): with self.assertRaisesRegex(ValueError, "path"): validate_turn_training_bundle(path_bundle) + def test_symlink_and_hardlink_bundle_nodes_are_rejected(self): + for link_kind in ("symlink", "hardlink"): + with self.subTest(link_kind=link_kind): + _, bundle = _export(self.tmp_path, name=f"link-{link_kind}") + linked = bundle / "artifacts" / f"{link_kind}.patch" + if link_kind == "symlink": + linked.symlink_to(bundle / "agent.patch") + else: + os.link(bundle / "agent.patch", linked) + with self.assertRaisesRegex(ValueError, "link"): + validate_turn_training_bundle(bundle) + + def test_runtime_provider_boundary_retry_exports_one_accepted_unit(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "runtime" + case.mkdir() + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-runtime", + task_revision_id="task-runtime@1", + rollout_id="r0", + rollout_index=0, + run_id="runtime-run", + harness_revision="test-rev", + ) + agent = MultiTurnReactAgent( + function_list=[], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 2, + "temperature": 0.2, + "top_p": 0.9, + "presence_penalty": 0.0, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="runtime-run", + ) + client = FakeClient( + [ + FakeResponse("empty-first", ""), + FakeResponse("accepted-second", "done"), + ] + ) + agent._llm_client = client + agent._llm_api_base = "http://fake" + with patch("agent_base.react_agent.time.sleep", return_value=None): + session = agent._run_session("Finish the task.", workspace_root=str(case)) + self.assertEqual(session["termination"], "result") + bundle = Path(session["training_export_path"]) + unit = _units(bundle)[0] + self.assertEqual(unit["capture_status"], "ok") + self.assertEqual(unit["provenance"]["provider_retry_count"], 1) + self.assertEqual(unit["request"]["messages"], client.requests[-1]["messages"]) + raw_ref = unit["source"]["raw_provider_response_ref"] + self.assertEqual( + json.loads((bundle / raw_ref["path"]).read_text())["id"], + "accepted-second", + ) + summary = validate_turn_training_bundle(bundle) + self.assertEqual(summary["run_id"], "runtime-run") + trace_rows = [ + json.loads(line) + for line in (bundle / summary["trace_ref"]["path"]).read_text().splitlines() + ] + llm_capture = next(row for row in trace_rows if row.get("capture_type") == "llm_call") + self.assertEqual( + [item["status"] for item in llm_capture["payload"]["response"]["provider_attempts"]], + ["rejected_empty", "accepted"], + ) + + def test_malformed_provider_response_is_retried_and_exported_as_error(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "malformed-runtime" + case.mkdir() + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-malformed", + task_revision_id="task-malformed@1", + rollout_id="r0", + rollout_index=0, + run_id="malformed-run", + harness_revision="test-rev", + ) + agent = MultiTurnReactAgent( + function_list=[], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 2, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="malformed-run", + ) + agent._llm_client = FakeClient( + [ + FakeResponseWithoutDump("unverifiable"), + FakeMalformedResponse("malformed-2"), + ] + ) + agent._llm_api_base = "http://fake" + with patch("agent_base.react_agent.time.sleep", return_value=None): + session = agent._run_session("Finish the task.", workspace_root=str(case)) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + self.assertEqual(session["termination"], "llm api error") + self.assertEqual(summary["capture_error_count"], 1) + trace_rows = [ + json.loads(line) + for line in (bundle / summary["trace_ref"]["path"]).read_text().splitlines() + ] + llm_capture = next(row for row in trace_rows if row.get("capture_type") == "llm_call") + self.assertEqual( + [item["status"] for item in llm_capture["payload"]["response"]["provider_attempts"]], + ["parse_error", "parse_error"], + ) + + def test_runtime_parallel_tool_calls_close_into_next_provider_request(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "parallel-runtime" + case.mkdir() + (case / "a.txt").write_text("alpha\n", encoding="utf-8") + (case / "b.txt").write_text("beta\n", encoding="utf-8") + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-parallel", + task_revision_id="task-parallel@1", + rollout_id="r0", + rollout_index=0, + run_id="parallel-run", + harness_revision="test-rev", + ) + calls = [ + _fake_tool_call("call-a", "Read", '{"file_path":"a.txt"}'), + _fake_tool_call("call-b", "Read", '{"file_path":"b.txt"}'), + ] + agent = MultiTurnReactAgent( + function_list=["Read"], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=3, + training_export=config, + run_id="parallel-run", + ) + client = FakeClient( + [ + FakeResponse("tool-round", None, tool_calls=calls), + FakeResponse("final-round", "done"), + ] + ) + agent._llm_client = client + agent._llm_api_base = "http://fake" + session = agent._run_session("Read both files.", workspace_root=str(case)) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + self.assertEqual(summary["main_agent_turn_count"], 2) + next_messages = client.requests[1]["messages"] + self.assertEqual( + [ + message["tool_call_id"] + for message in next_messages + if message.get("role") == "tool" + ], + ["call-a", "call-b"], + ) + + def test_nonfinite_json_and_invalid_identity_types_are_rejected(self): + with self.assertRaisesRegex(ValueError, "NaN"): + canonical_json_bytes({"value": float("nan")}) + with self.assertRaisesRegex(ValueError, "rollout_index"): + TrainingExportConfig( + output_dir=self.tmp_path / EXPORT_DIR_NAME, + task_id="t", task_revision_id="r", rollout_id="r0", rollout_index=True, + run_id="run", harness_revision="rev", + ) + with self.assertRaisesRegex(ValueError, "task_id"): + TrainingExportConfig( + output_dir=self.tmp_path / EXPORT_DIR_NAME, + task_id=7, task_revision_id="r", rollout_id="r0", rollout_index=0, + run_id="run", harness_revision="rev", + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_turn_training_export_api.py b/tests/test_turn_training_export_api.py new file mode 100644 index 0000000..745fc76 --- /dev/null +++ b/tests/test_turn_training_export_api.py @@ -0,0 +1,154 @@ +import tempfile +import unittest +from pathlib import Path +import sys + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from api import openai_server +from api.openai_server import ( + OpenAICompatError, + REQUEST_TRAINING_EXPORT_FIELD, + ServerConfig, + request_training_export, + run_chat_completion, +) +from agent_base.training_export import EXPORT_DIR_NAME + + +def _training_payload(**overrides): + value = { + "task_id": "task-api", + "task_revision_id": "task-api@1", + "rollout_id": "r0", + "rollout_index": 0, + "harness_revision": "test-rev", + } + value.update(overrides) + return value + + +class FakeAPIAgent: + seen = [] + + def __init__( + self, + function_list, + llm, + trace_dir, + role_prompt=None, + training_export=None, + run_id=None, + ): + self.trace_dir = Path(trace_dir) + self.training_export = training_export + self.run_id = run_id + self.__class__.seen.append(self) + + def _run_session(self, prompt, workspace_root=None, initial_content_parts=None): + return { + "result_text": "done", + "termination": "result", + "trace_path": str(self.trace_dir / "trace.jsonl"), + "training_export_path": str(self.training_export.output_dir) + if self.training_export is not None + else "", + } + + +def _fake_llm_config(model_name=None, extra_body=None, omit_generate_params=None): + return { + "model": str(model_name or "fake-model"), + "api_key": "fake", + "api_base": "http://fake.invalid/v1", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + }, + } + + +class TrainingExportAPITests(unittest.TestCase): + def setUp(self): + self._temp = tempfile.TemporaryDirectory() + self.root = Path(self._temp.name) + FakeAPIAgent.seen.clear() + + def tearDown(self): + self._temp.cleanup() + + def test_request_parser_freezes_outer_run_identity_and_fixed_directory(self): + workspace = self.root / "workspace" + workspace.mkdir() + patch_path = workspace / "agent.patch" + patch_path.write_text("diff --git a/a b/a\n", encoding="utf-8") + config = request_training_export( + { + REQUEST_TRAINING_EXPORT_FIELD: _training_payload( + final_patch_path="agent.patch" + ) + }, + run_root=self.root / "run", + workspace_root=workspace, + run_id="run-api-1", + ) + self.assertIsNotNone(config) + self.assertEqual(config.run_id, "run-api-1") + self.assertEqual(config.output_dir, self.root / "run" / EXPORT_DIR_NAME) + self.assertEqual(config.final_patch_path, workspace / "agent.patch") + + def test_parser_rejects_unknown_fields_and_bool_rollout_index(self): + for value in ( + _training_payload(unknown=True), + _training_payload(rollout_index=True), + _training_payload(task_id=7), + ): + with self.subTest(value=value): + with self.assertRaises(OpenAICompatError): + request_training_export( + {REQUEST_TRAINING_EXPORT_FIELD: value}, + run_root=self.root / "run", + workspace_root=self.root, + run_id="run-api-1", + ) + + def test_training_export_with_wrapper_fails_before_any_model_call(self): + payload = { + "model": "RH", + "messages": [{"role": "user", "content": "hello"}], + REQUEST_TRAINING_EXPORT_FIELD: _training_payload(), + } + with self.assertRaisesRegex(OpenAICompatError, "cannot be combined"): + run_chat_completion( + payload, + ServerConfig(api_runs_dir=self.root / "runs", input_wrapper=True), + ) + self.assertEqual(FakeAPIAgent.seen, []) + + def test_api_passes_same_run_id_to_agent_trace_and_export(self): + previous_agent = openai_server.MultiTurnReactAgent + previous_config = openai_server.default_llm_config + openai_server.MultiTurnReactAgent = FakeAPIAgent + openai_server.default_llm_config = _fake_llm_config + try: + response = run_chat_completion( + { + "model": "RH", + "messages": [{"role": "user", "content": "hello"}], + REQUEST_TRAINING_EXPORT_FIELD: _training_payload(), + }, + ServerConfig(api_runs_dir=self.root / "runs"), + ) + finally: + openai_server.MultiTurnReactAgent = previous_agent + openai_server.default_llm_config = previous_config + self.assertEqual(response["choices"][0]["message"]["content"], "done") + self.assertEqual(len(FakeAPIAgent.seen), 1) + agent = FakeAPIAgent.seen[0] + self.assertEqual(agent.run_id, agent.training_export.run_id) + self.assertEqual(agent.training_export.output_dir.name, EXPORT_DIR_NAME) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_turn_training_export_entrypoints.py b/tests/test_turn_training_export_entrypoints.py new file mode 100644 index 0000000..6fd8d50 --- /dev/null +++ b/tests/test_turn_training_export_entrypoints.py @@ -0,0 +1,95 @@ +import tempfile +import unittest +from pathlib import Path +import sys + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from agent_base.react_agent import _parse_cli_args +from agent_base.training_export import EXPORT_DIR_NAME +from researchharness.runtime import _resolve_training_export + + +class TrainingExportEntrypointTests(unittest.TestCase): + def setUp(self): + self._temp = tempfile.TemporaryDirectory() + self.root = Path(self._temp.name) + + def tearDown(self): + self._temp.cleanup() + + def test_cli_builds_one_frozen_config_and_sibling_trace(self): + output = self.root / "r0" / EXPORT_DIR_NAME + parsed = _parse_cli_args( + [ + "fix", + "auth", + "--training-export-dir", + str(output), + "--task-id", + "task-auth", + "--task-revision-id", + "task-auth@1", + "--rollout-id", + "r0", + "--rollout-index", + "0", + "--run-id", + "run-cli", + "--harness-revision", + "test-rev", + ] + ) + self.assertEqual(parsed[0], "fix auth") + self.assertEqual(parsed[1], str(output.parent / "trace")) + config = parsed[-1] + self.assertEqual(config.output_dir, output) + self.assertEqual(config.run_id, "run-cli") + + def test_cli_partial_configuration_fails_closed(self): + with self.assertRaisesRegex(ValueError, "requires"): + _parse_cli_args( + [ + "fix", + "--training-export-dir", + str(self.root / EXPORT_DIR_NAME), + ] + ) + + def test_python_resolver_rejects_bool_index_and_preserves_no_export(self): + with self.assertRaisesRegex(ValueError, "rollout_index"): + _resolve_training_export( + training_export_dir=self.root / EXPORT_DIR_NAME, + task_id="task", + task_revision_id="task@1", + rollout_id="r0", + rollout_index=True, + run_id="run-python", + harness_revision="rev", + scaffold_version="", + tool_protocol_version="native-v1", + tokenizer_revision=None, + final_patch_path=None, + trace_dir=None, + ) + config, trace_dir, run_id = _resolve_training_export( + training_export_dir=None, + task_id=None, + task_revision_id=None, + rollout_id=None, + rollout_index=None, + run_id="ordinary-run", + harness_revision=None, + scaffold_version="", + tool_protocol_version="native-v1", + tokenizer_revision=None, + final_patch_path=None, + trace_dir="trace", + ) + self.assertIsNone(config) + self.assertEqual(trace_dir, "trace") + self.assertEqual(run_id, "ordinary-run") + + +if __name__ == "__main__": + unittest.main() From fec242b450085e24c5abeea61b0612e62ed49009 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 03:42:28 +0800 Subject: [PATCH 4/6] Harden training export closure validation --- agent_base/training_export.py | 93 +++++++++++++++++++++++++++--- tests/test_turn_training_export.py | 69 ++++++++++++++++++++++ 2 files changed, 155 insertions(+), 7 deletions(-) diff --git a/agent_base/training_export.py b/agent_base/training_export.py index 4930a1b..f015782 100644 --- a/agent_base/training_export.py +++ b/agent_base/training_export.py @@ -32,6 +32,15 @@ "parent_run_id", "delegation_id", } +_PROHIBITED_OUTCOME_KEYS = { + "resolved", + "reward", + "advantage", + "verifier_outcome", + "verifier_result", + "verifier_reward", + "verifier_verdict", +} _ORACLE_KEYS = { "gold_patch", "hidden_test", @@ -517,6 +526,8 @@ def _validate_provider_attempts( if retry_count != len(attempts) - 1: raise ValueError("provider retry count differs from captured attempt closure") allowed_statuses = {"accepted", "rejected_empty", "error", "parse_error"} + from agent_base.react_agent import assistant_has_meaningful_text + for index, attempt in enumerate(attempts): if not isinstance(attempt, Mapping) or attempt.get("attempt_index") != index: raise ValueError("provider attempts must be complete and canonically indexed") @@ -526,6 +537,13 @@ def _validate_provider_attempts( raise ValueError("provider attempt request must be an object") if not isinstance(attempt.get("raw_provider_response"), Mapping): raise ValueError("provider attempt raw response must be an object") + attempt_metadata = { + key: value + for key, value in attempt.items() + if key not in {"provider_request", "raw_provider_response"} + } + if _contains_prohibited_outcome_key(attempt_metadata): + raise ValueError("provider attempt contains prohibited outcome metadata") if attempt.get("provider_request") != provider_request: raise ValueError("provider retry request differs from the final provider request") if attempt.get("status") == "accepted" and index != len(attempts) - 1: @@ -540,8 +558,7 @@ def _validate_provider_attempts( rejected_message = rejected_response["message"] rejected_content = rejected_message.get("content") if ( - (isinstance(rejected_content, str) and rejected_content.strip()) - or (not isinstance(rejected_content, str) and rejected_content) + assistant_has_meaningful_text(rejected_content) or rejected_message.get("tool_calls") ): raise ValueError("rejected-empty attempt contains an acceptable response") @@ -556,12 +573,51 @@ def _validate_provider_attempts( raise ValueError("failed provider capture cannot end in an accepted attempt") -def _tool_result_ids(messages: object) -> list[str]: +def _message_tool_calls(message: Mapping[str, Any]) -> list[dict[str, Any]] | None: + raw_tool_calls = message.get("tool_calls") + if not isinstance(raw_tool_calls, list): + return None + tool_calls: list[dict[str, Any]] = [] + for raw_tool_call in raw_tool_calls: + if not isinstance(raw_tool_call, Mapping): + return None + function = raw_tool_call.get("function") + if not isinstance(function, Mapping): + return None + tool_calls.append( + { + "id": raw_tool_call.get("id"), + "type": raw_tool_call.get("type"), + "function": { + "name": function.get("name"), + "arguments": function.get("arguments"), + }, + } + ) + return tool_calls + + +def _tool_result_ids_after_target( + messages: object, + target_tool_calls: Sequence[Mapping[str, Any]], +) -> list[str]: if not isinstance(messages, Sequence) or isinstance(messages, (str, bytes, bytearray)): return [] + target_position: int | None = None + for index, message in enumerate(messages): + if not isinstance(message, Mapping) or message.get("role") != "assistant": + continue + if _message_tool_calls(message) == list(target_tool_calls): + target_position = index + if target_position is None: + return [] result: list[str] = [] - for message in messages: - if not isinstance(message, Mapping) or message.get("role") != "tool": + for message in messages[target_position + 1 :]: + if not isinstance(message, Mapping): + continue + if message.get("role") == "assistant": + break + if message.get("role") != "tool": continue tool_call_id = message.get("tool_call_id") if isinstance(tool_call_id, str) and tool_call_id: @@ -634,7 +690,13 @@ def _validate_tool_result_closure( payload = row.get("payload") if not isinstance(payload, Mapping): continue - if _contains_ordered_sequence(_tool_result_ids(payload.get("pre_messages")), expected_ids): + if _contains_ordered_sequence( + _tool_result_ids_after_target( + payload.get("pre_messages"), + tool_calls, + ), + expected_ids, + ): compaction_pre_proof = True if payload.get("status") == "ok": compacted_next_input_proof = True @@ -643,7 +705,10 @@ def _validate_tool_result_closure( next_payload = trace_rows[next_main_position].get("payload") if isinstance(next_payload, Mapping): next_input_proof = _contains_ordered_sequence( - _tool_result_ids(_provider_request(next_payload).get("messages")), + _tool_result_ids_after_target( + _provider_request(next_payload).get("messages"), + tool_calls, + ), expected_ids, ) if not (trace_proves_execution or next_input_proof or compaction_pre_proof): @@ -1013,6 +1078,20 @@ def _contains_prohibited_key(value: Any) -> bool: return False +def _contains_prohibited_outcome_key(value: Any) -> bool: + if isinstance(value, Mapping): + prohibited = any( + str(key).casefold().replace("-", "_") in _PROHIBITED_OUTCOME_KEYS + for key in value + ) + return prohibited or any( + _contains_prohibited_outcome_key(nested) for nested in value.values() + ) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return any(_contains_prohibited_outcome_key(nested) for nested in value) + return False + + def _load_units(path: Path) -> list[dict[str, Any]]: units: list[dict[str, Any]] = [] for line in path.read_bytes().splitlines(keepends=True): diff --git a/tests/test_turn_training_export.py b/tests/test_turn_training_export.py index 33f985e..4180064 100644 --- a/tests/test_turn_training_export.py +++ b/tests/test_turn_training_export.py @@ -622,6 +622,75 @@ def test_tool_call_requires_matching_result_in_next_model_input(self): with self.assertRaisesRegex(ValueError, "tool-result"): export_turn_training_bundle(config, trace_path=trace, termination="result") + def test_tool_result_closure_binds_to_current_reused_call_id(self): + def rows_with_reused_id(include_current_result): + rows = [_trace_rows()[0], _trace_rows()[2]] + rows[1]["event_index"] = 2 + call = rows[0]["payload"]["response"]["tool_calls"][0] + initial = rows[0]["payload"]["provider_request"]["messages"] + old_assistant = {"role": "assistant", "content": "", "tool_calls": [call]} + old_result = { + "role": "tool", + "tool_call_id": call["id"], + "content": "old result", + } + initial.extend([old_assistant, old_result]) + next_messages = [ + *initial, + {"role": "assistant", "content": "", "tool_calls": [call]}, + ] + if include_current_result: + next_messages.append( + { + "role": "tool", + "tool_call_id": call["id"], + "content": "new result", + } + ) + rows[1]["payload"]["provider_request"]["messages"] = next_messages + return rows + + config = _config(self.tmp_path, name="reused-id-missing", patch_file=False) + trace = _write_trace( + self.tmp_path / "reused-id-missing/trace.jsonl", + rows_with_reused_id(False), + ) + with self.assertRaisesRegex(ValueError, "tool-result"): + export_turn_training_bundle(config, trace_path=trace, termination="result") + + _, bundle = _export( + self.tmp_path, + name="reused-id-closed", + rows=rows_with_reused_id(True), + patch_file=False, + ) + self.assertEqual(validate_turn_training_bundle(bundle)["capture_error_count"], 0) + + def test_provider_attempt_metadata_cannot_carry_outcome_labels(self): + rows = _trace_rows() + rows[0]["payload"]["response"]["provider_attempts"][0]["details"] = { + "reward": 1.0 + } + config = _config(self.tmp_path, name="attempt-outcome", patch_file=False) + trace = _write_trace(self.tmp_path / "attempt-outcome/trace.jsonl", rows) + with self.assertRaisesRegex(ValueError, "outcome"): + export_turn_training_bundle(config, trace_path=trace, termination="result") + + def test_rejected_empty_attempt_uses_runtime_text_semantics(self): + rows = _trace_rows() + rows[0]["payload"]["response"]["provider_attempts"][0][ + "raw_provider_response" + ]["choices"][0]["message"]["content"] = [ + {"type": "text", "text": " "} + ] + _, bundle = _export( + self.tmp_path, + name="structured-empty-retry", + rows=rows, + patch_file=False, + ) + self.assertEqual(validate_turn_training_bundle(bundle)["capture_error_count"], 0) + def test_unit_id_does_not_depend_on_trace_event_position(self): row_one = _trace_rows()[2] row_one["event_index"] = 1 From 3a009fac002f485564b624fe6c1800b02856fc06 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 03:43:45 +0800 Subject: [PATCH 5/6] Preserve provider boundary failures in export --- agent_base/react_agent.py | 14 ++++++++++---- tests/test_turn_training_export.py | 12 ++++++++---- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/agent_base/react_agent.py b/agent_base/react_agent.py index d6d491c..fa42652 100644 --- a/agent_base/react_agent.py +++ b/agent_base/react_agent.py @@ -853,6 +853,7 @@ def _call_chat_completion( "provider_request": {}, "raw_provider_response": {}, } + provider_call_started = False provider_returned = False try: attempts_made = attempt + 1 @@ -896,6 +897,7 @@ def _call_chat_completion( last_provider_request = safe_jsonable(request_kwargs) attempt_record["provider_request"] = last_provider_request last_raw_provider_response = {} + provider_call_started = True with llm_hard_timeout(request_timeout): chat_response = request_client.chat.completions.create(**request_kwargs) provider_returned = True @@ -964,11 +966,15 @@ def _call_chat_completion( if debug_enabled(): print(f"Error: Attempt {attempt + 1} failed with an API or network error: {e}") except Exception as e: - if not provider_returned: + if not provider_call_started: raise - last_error = f"invalid provider response: {type(e).__name__}: {e}" + last_error = ( + f"invalid provider response: {type(e).__name__}: {e}" + if provider_returned + else f"provider call failed: {type(e).__name__}: {e}" + ) attempt_record.update( - status="parse_error", + status="parse_error" if provider_returned else "error", error_type=type(e).__name__, error=last_error, raw_provider_response=last_raw_provider_response, @@ -976,7 +982,7 @@ def _call_chat_completion( provider_attempts.append(safe_jsonable(attempt_record)) if debug_enabled(): print( - f"Error: Attempt {attempt + 1} returned an invalid provider response: " + f"Error: Attempt {attempt + 1} failed at the provider boundary: " f"{type(e).__name__}: {e}" ) diff --git a/tests/test_turn_training_export.py b/tests/test_turn_training_export.py index 4180064..b06fb24 100644 --- a/tests/test_turn_training_export.py +++ b/tests/test_turn_training_export.py @@ -425,7 +425,10 @@ def with_options(self, **kwargs): def create(self, **kwargs): self.requests.append(kwargs) - return self.responses.pop(0) + response = self.responses.pop(0) + if isinstance(response, BaseException): + raise response + return response class TrainingExportTests(unittest.TestCase): @@ -853,7 +856,7 @@ def test_runtime_provider_boundary_retry_exports_one_accepted_unit(self): ["rejected_empty", "accepted"], ) - def test_malformed_provider_response_is_retried_and_exported_as_error(self): + def test_provider_boundary_failures_are_retried_and_exported_as_error(self): from agent_base.react_agent import MultiTurnReactAgent case = self.tmp_path / "malformed-runtime" @@ -876,7 +879,7 @@ def test_malformed_provider_response_is_retried_and_exported_as_error(self): "generate_cfg": { "max_input_tokens": 32768, "max_output_tokens": 128, - "max_retries": 2, + "max_retries": 3, }, }, trace_dir=str(case / "trace"), @@ -887,6 +890,7 @@ def test_malformed_provider_response_is_retried_and_exported_as_error(self): ) agent._llm_client = FakeClient( [ + ValueError("unsupported request field"), FakeResponseWithoutDump("unverifiable"), FakeMalformedResponse("malformed-2"), ] @@ -905,7 +909,7 @@ def test_malformed_provider_response_is_retried_and_exported_as_error(self): llm_capture = next(row for row in trace_rows if row.get("capture_type") == "llm_call") self.assertEqual( [item["status"] for item in llm_capture["payload"]["response"]["provider_attempts"]], - ["parse_error", "parse_error"], + ["error", "parse_error", "parse_error"], ) def test_runtime_parallel_tool_calls_close_into_next_provider_request(self): From 4e330c6df040bba778ad7ad27f259158bfb6c922 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 04:02:05 +0800 Subject: [PATCH 6/6] Preserve interrupted and duplicate rollout exports --- agent_base/context_compact.py | 30 +- agent_base/react_agent.py | 126 ++++++- agent_base/training_export.py | 249 ++++++++++-- api/openai_server.py | 19 +- docs/turn-training-export-v1.md | 16 +- researchharness/runtime.py | 9 +- tests/test_turn_training_export.py | 355 +++++++++++++++++- tests/test_turn_training_export_api.py | 6 + .../test_turn_training_export_entrypoints.py | 31 ++ 9 files changed, 786 insertions(+), 55 deletions(-) diff --git a/agent_base/context_compact.py b/agent_base/context_compact.py index 5dae404..3b84904 100644 --- a/agent_base/context_compact.py +++ b/agent_base/context_compact.py @@ -133,11 +133,31 @@ def compact_messages( ), }, ] - summary_reply = llm_caller( - summary_request, - runtime_deadline=runtime_deadline, - max_output_tokens=model_profile.compact_summary_max_tokens, - ) + try: + summary_reply = llm_caller( + summary_request, + runtime_deadline=runtime_deadline, + max_output_tokens=model_profile.compact_summary_max_tokens, + ) + except KeyboardInterrupt: + return CompactionOutcome( + status="error", + compacted_messages=safe_messages, + prior_token_estimate=prior_token_estimate, + existing_memory_text=existing_memory_text, + summary_request=summary_request, + summary_response={ + "status": "error", + "error": "context compaction provider call interrupted", + "interrupted": True, + "tool_calls": [], + }, + pre_messages=safe_messages, + post_messages=safe_messages, + error="context compaction provider call interrupted", + compacted_group_count=len(compacted_groups), + kept_group_count=len(recent_groups), + ) if not isinstance(summary_reply, dict) or summary_reply.get("status") != "ok": error = summary_reply.get("error", "context compaction summary call failed") if isinstance(summary_reply, dict) else str(summary_reply) return CompactionOutcome( diff --git a/agent_base/react_agent.py b/agent_base/react_agent.py index fa42652..f141a5d 100644 --- a/agent_base/react_agent.py +++ b/agent_base/react_agent.py @@ -1,6 +1,7 @@ import argparse from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from dataclasses import replace import json import os import re @@ -21,7 +22,13 @@ from agent_base.prompt import composed_system_prompt from agent_base.session_state import AgentSessionState, CompactionRecord, persist_session_state, resolve_session_state_path from agent_base.trace_utils import FlatTraceWriter -from agent_base.training_export import TrainingExportConfig, export_turn_training_bundle +from agent_base.training_export import ( + TrainingExportConfig, + TrainingExportReservation, + export_turn_training_bundle, + reserve_turn_training_export, + validate_final_patch_source, +) from agent_base.tools.custom import build_custom_tool_map from agent_base.tools.tooling import ToolBase, normalize_workspace_root from agent_base.tools.tool_extra import StrReplaceEditor @@ -954,6 +961,26 @@ def _call_chat_completion( if debug_enabled(): print(f"Warning: Attempt {attempt + 1} received an empty response.") + except KeyboardInterrupt as e: + last_error = "provider call interrupted" + attempt_record.update( + status="error", + error_type=type(e).__name__, + error=last_error, + raw_provider_response=last_raw_provider_response, + ) + provider_attempts.append(safe_jsonable(attempt_record)) + return { + "status": "error", + "error": f"llm api error: {last_error}", + "interrupted": True, + "provider_request": last_provider_request, + "raw_provider_response": last_raw_provider_response, + "provider_retry_count": attempt, + "provider_attempts": provider_attempts, + "requested_model": self.model, + "provider_model": last_raw_provider_response.get("model"), + } except (APIError, APIConnectionError, APITimeoutError, LLMHardTimeoutError) as e: last_error = str(e) attempt_record.update( @@ -966,12 +993,14 @@ def _call_chat_completion( if debug_enabled(): print(f"Error: Attempt {attempt + 1} failed with an API or network error: {e}") except Exception as e: - if not provider_call_started: - raise last_error = ( f"invalid provider response: {type(e).__name__}: {e}" if provider_returned - else f"provider call failed: {type(e).__name__}: {e}" + else ( + f"provider call failed: {type(e).__name__}: {e}" + if provider_call_started + else f"provider setup failed: {type(e).__name__}: {e}" + ) ) attempt_record.update( status="parse_error" if provider_returned else "error", @@ -1082,6 +1111,21 @@ def run( images: Optional[str | Path | Sequence[str | Path]] = None, ) -> str: """Run the agent on one prompt and return only the final result text.""" + return str( + self.run_session( + prompt, + workspace_root=workspace_root, + images=images, + )["result_text"] + ) + + def run_session( + self, + prompt: str, + workspace_root: Optional[str] = None, + images: Optional[str | Path | Sequence[str | Path]] = None, + ) -> dict[str, Any]: + """Run one prompt and return result text plus trace/export metadata.""" resolved_workspace_root = normalize_workspace_root( workspace_root if workspace_root is not None else self.workspace_root ) @@ -1102,7 +1146,7 @@ def run( run_prompt, workspace_root=str(resolved_workspace_root), initial_content_parts=initial_content_parts or None, - )["result_text"] + ) def _run_session( self, @@ -1113,14 +1157,40 @@ def _run_session( prior_messages: Optional[Sequence[dict[str, Any]]] = None, interrupt_event: Optional[threading.Event] = None, ) -> dict: - """Internal execution path with trace data for tests and debugging.""" + """Run one session while reserving its immutable export before side effects.""" if not isinstance(prompt, str) or not prompt.strip(): raise ValueError("prompt must be a non-empty string.") if self.training_export is not None and self._training_export_consumed: raise RuntimeError("one immutable training export config may be consumed by only one session") + reservation: TrainingExportReservation | None = None if self.training_export is not None: + reservation = reserve_turn_training_export(self.training_export) self._training_export_consumed = True + try: + return self._run_session_impl( + prompt, + workspace_root=workspace_root, + event_callback=event_callback, + initial_content_parts=initial_content_parts, + prior_messages=prior_messages, + interrupt_event=interrupt_event, + training_export_reservation=reservation, + ) + finally: + if reservation is not None: + reservation.release() + def _run_session_impl( + self, + prompt: str, + workspace_root: Optional[str] = None, + event_callback: Optional[Callable[[dict[str, Any]], None]] = None, + initial_content_parts: Optional[Sequence[dict[str, Any]]] = None, + prior_messages: Optional[Sequence[dict[str, Any]]] = None, + interrupt_event: Optional[threading.Event] = None, + training_export_reservation: TrainingExportReservation | None = None, + ) -> dict: + """Internal execution path with trace data for tests and debugging.""" prompt_text = prompt.strip() resolved_workspace_root = normalize_workspace_root( workspace_root if workspace_root is not None else self.workspace_root @@ -1209,6 +1279,22 @@ def persist_state(*, termination: str = "", error: str = "") -> None: persist_session_state(self.session_state_path, session_state) def finalize(result_text: str, termination: str, *, role: str = "runtime", error: str = "") -> dict[str, Any]: + export_config = self.training_export + if ( + export_config is not None + and export_config.final_patch_path is not None + ): + try: + validate_final_patch_source(export_config.final_patch_path) + except (OSError, ValueError) as exc: + patch_error = f"final patch unavailable: {exc}" + export_config = replace(export_config, final_patch_path=None) + if error: + error = f"{error}; {patch_error}" + else: + termination = "final patch unavailable" + error = patch_error + role = "runtime" trace_writer.append( role=role, text=result_text, @@ -1218,21 +1304,23 @@ def finalize(result_text: str, termination: str, *, role: str = "runtime", error ) persist_state(termination=termination, error=error) training_export_path = "" - if self.training_export is not None: + if export_config is not None: if self.trace_path is None: raise RuntimeError("training export requires a durable trace path") training_export_path = str( export_turn_training_bundle( - self.training_export, + export_config, trace_path=self.trace_path, termination={ "reason": termination, "submitted": termination == "result", "clean_termination": not bool(error), }, + _reservation=training_export_reservation, ) ) return { + "run_id": trace_writer.run_id, "prompt": prompt_text, "messages": messages, "result_text": result_text, @@ -1350,6 +1438,11 @@ def finalize_interrupted() -> dict[str, Any]: payload=compaction_trace_payload(trigger_reason=compact_reason, outcome=compact_outcome), ) persist_state(error=compact_outcome.error) + if ( + isinstance(compact_outcome.summary_response, dict) + and compact_outcome.summary_response.get("interrupted") is True + ): + return finalize_interrupted() if current_token_estimate > max_input_tokens: result_text = "No result found before the maximum input token limit." termination = f"input token limit reached: {current_token_estimate} > {max_input_tokens}" @@ -1361,9 +1454,15 @@ def finalize_interrupted() -> dict[str, Any]: try: llm_reply = self.call_llm_api(llm_request_messages, runtime_deadline=runtime_deadline) except KeyboardInterrupt: - return finalize_interrupted() - if interruption_requested(): - return finalize_interrupted() + llm_reply = { + "status": "error", + "error": "llm api error: provider call interrupted", + "interrupted": True, + "tool_calls": [], + "provider_retry_count": 0, + "requested_model": self.model, + "provider_model": None, + } trace_writer.append( role="runtime", text="", @@ -1380,6 +1479,11 @@ def finalize_interrupted() -> dict[str, Any]: session_state.last_input_tokens = input_tokens_from_usage( llm_reply.get("usage") if isinstance(llm_reply, dict) else None ) + if ( + isinstance(llm_reply, dict) + and llm_reply.get("interrupted") is True + ) or interruption_requested(): + return finalize_interrupted() assistant_content = llm_reply.get("content") if isinstance(llm_reply, dict) else None assistant_tool_calls = llm_reply.get("tool_calls", []) if isinstance(llm_reply, dict) else [] assistant_reasoning = llm_reply.get("reasoning_content") if isinstance(llm_reply, dict) else None diff --git a/agent_base/training_export.py b/agent_base/training_export.py index f015782..c4ae547 100644 --- a/agent_base/training_export.py +++ b/agent_base/training_export.py @@ -2,6 +2,7 @@ from __future__ import annotations +import argparse import ctypes from dataclasses import dataclass import errno @@ -13,6 +14,7 @@ import re import shutil import stat +import sys import tempfile from typing import Any, Mapping, Sequence @@ -122,6 +124,38 @@ def identity(self) -> dict[str, Any]: } +@dataclass(slots=True) +class TrainingExportReservation: + target: Path + lock_path: Path + lock_fd: int + lock_device: int + lock_inode: int + active: bool = True + + def release(self) -> None: + if not self.active: + return + self.active = False + try: + os.close(self.lock_fd) + except OSError: + pass + try: + observed = self.lock_path.lstat() + except FileNotFoundError: + return + if ( + observed.st_dev == self.lock_device + and observed.st_ino == self.lock_inode + and stat.S_ISREG(observed.st_mode) + ): + try: + self.lock_path.unlink() + except FileNotFoundError: + pass + + def canonical_json_bytes(value: Any) -> bytes: _reject_nonfinite(value) return json.dumps( @@ -542,6 +576,8 @@ def _validate_provider_attempts( for key, value in attempt.items() if key not in {"provider_request", "raw_provider_response"} } + if _contains_prohibited_key(attempt_metadata): + raise ValueError("provider attempt contains prohibited multi-agent metadata") if _contains_prohibited_outcome_key(attempt_metadata): raise ValueError("provider attempt contains prohibited outcome metadata") if attempt.get("provider_request") != provider_request: @@ -600,17 +636,20 @@ def _message_tool_calls(message: Mapping[str, Any]) -> list[dict[str, Any]] | No def _tool_result_ids_after_target( messages: object, target_tool_calls: Sequence[Mapping[str, Any]], + *, + prior_occurrence_count: int, ) -> list[str]: if not isinstance(messages, Sequence) or isinstance(messages, (str, bytes, bytearray)): return [] - target_position: int | None = None + target_positions: list[int] = [] for index, message in enumerate(messages): if not isinstance(message, Mapping) or message.get("role") != "assistant": continue if _message_tool_calls(message) == list(target_tool_calls): - target_position = index - if target_position is None: + target_positions.append(index) + if len(target_positions) <= prior_occurrence_count: return [] + target_position = target_positions[prior_occurrence_count] result: list[str] = [] for message in messages[target_position + 1 :]: if not isinstance(message, Mapping): @@ -625,6 +664,20 @@ def _tool_result_ids_after_target( return result +def _tool_call_occurrence_count( + messages: object, + target_tool_calls: Sequence[Mapping[str, Any]], +) -> int: + if not isinstance(messages, Sequence) or isinstance(messages, (str, bytes, bytearray)): + return 0 + return sum( + isinstance(message, Mapping) + and message.get("role") == "assistant" + and _message_tool_calls(message) == list(target_tool_calls) + for message in messages + ) + + def _contains_ordered_sequence(values: Sequence[str], expected: Sequence[str]) -> bool: if not expected: return True @@ -642,14 +695,24 @@ def _validate_tool_result_closure( trace_rows: Sequence[Mapping[str, Any]], capture_row: Mapping[str, Any], response: Mapping[str, Any], -) -> None: + allow_terminal_incomplete: bool = False, +) -> bool: message = response.get("message") tool_calls = message.get("tool_calls") if isinstance(message, Mapping) else [] if not isinstance(tool_calls, list) or not tool_calls: - return + return True if response.get("finish_reason") == "length": - return + return True expected_ids = [str(tool_call.get("id", "")) for tool_call in tool_calls] + capture_payload = ( + capture_row.get("payload") + if isinstance(capture_row.get("payload"), Mapping) + else {} + ) + prior_occurrence_count = _tool_call_occurrence_count( + _provider_request(capture_payload).get("messages"), + tool_calls, + ) capture_event_index = capture_row.get("event_index") capture_position = next( ( @@ -694,6 +757,7 @@ def _validate_tool_result_closure( _tool_result_ids_after_target( payload.get("pre_messages"), tool_calls, + prior_occurrence_count=prior_occurrence_count, ), expected_ids, ): @@ -708,15 +772,40 @@ def _validate_tool_result_closure( _tool_result_ids_after_target( _provider_request(next_payload).get("messages"), tool_calls, + prior_occurrence_count=prior_occurrence_count, ), expected_ids, ) if not (trace_proves_execution or next_input_proof or compaction_pre_proof): + if allow_terminal_incomplete and next_main_position is None: + return False raise ValueError("tool-call target lacks matching tool-result closure") if next_main_position is not None and not ( next_input_proof or compacted_next_input_proof ): raise ValueError("tool results did not enter the next provider request") + return True + + +def _trace_proves_nonclean_terminal( + trace_rows: Sequence[Mapping[str, Any]], + termination: Mapping[str, Any], +) -> bool: + if termination.get("clean_termination") is not False or termination.get("submitted") is not False: + return False + terminal_rows = [ + row + for row in trace_rows + if isinstance(row.get("termination"), str) and row.get("termination") + ] + if not terminal_rows: + return False + final_row = terminal_rows[-1] + return bool( + final_row.get("termination") == termination.get("reason") + and isinstance(final_row.get("error"), str) + and final_row.get("error") + ) def _base_provenance(config: TrainingExportConfig, payload: Mapping[str, Any], findings: list[str]) -> dict[str, Any]: @@ -788,7 +877,14 @@ def _unit_common(config: TrainingExportConfig, row: Mapping[str, Any], unit_type } -def _build_main_unit(config: TrainingExportConfig, root: Path, row: Mapping[str, Any], trace_ref: Mapping[str, Any]) -> dict[str, Any]: +def _build_main_unit( + config: TrainingExportConfig, + root: Path, + row: Mapping[str, Any], + trace_ref: Mapping[str, Any], + *, + tool_closure_ok: bool, +) -> dict[str, Any]: payload = row.get("payload") if isinstance(row.get("payload"), Mapping) else {} provider_request = _provider_request(payload) raw_response = _raw_response(payload) @@ -802,7 +898,11 @@ def _build_main_unit(config: TrainingExportConfig, root: Path, row: Mapping[str, unit = _unit_common(config, row, "main_agent_turn") unit.update( { - "capture_status": "ok" if _main_capture_ok(payload, findings) else "error", + "capture_status": ( + "ok" + if _main_capture_ok(payload, findings) and tool_closure_ok + else "error" + ), "request": request, "response": response, "source": { @@ -929,6 +1029,50 @@ def _reject_symlink_ancestors(path: Path) -> None: current = current.parent +def reserve_turn_training_export( + config: TrainingExportConfig, +) -> TrainingExportReservation: + if not isinstance(config, TrainingExportConfig): + raise TypeError("config must be TrainingExportConfig") + target = config.output_dir + parent = target.parent + _reject_symlink_ancestors(parent) + parent.mkdir(parents=True, exist_ok=True) + _reject_symlink_ancestors(parent) + if target.exists() or target.is_symlink(): + raise FileExistsError(f"immutable export already exists: {target}") + lock_path = parent / f".{target.name}.publish.lock" + flags = ( + os.O_WRONLY + | os.O_CREAT + | os.O_EXCL + | getattr(os, "O_CLOEXEC", 0) + | getattr(os, "O_NOFOLLOW", 0) + ) + try: + lock_fd = os.open(lock_path, flags, 0o600) + except FileExistsError as exc: + raise FileExistsError( + f"immutable export is already reserved: {target}" + ) from exc + lock_stat = os.fstat(lock_fd) + reservation = TrainingExportReservation( + target=target, + lock_path=lock_path, + lock_fd=lock_fd, + lock_device=lock_stat.st_dev, + lock_inode=lock_stat.st_ino, + ) + if target.exists() or target.is_symlink(): + reservation.release() + raise FileExistsError(f"immutable export already exists: {target}") + return reservation + + +def validate_final_patch_source(path: Path) -> Path: + return _require_regular_source(Path(path), label="final patch") + + def _rename_no_replace(source: Path, target: Path) -> None: if os.name == "posix": libc = ctypes.CDLL(None, use_errno=True) @@ -960,22 +1104,23 @@ def export_turn_training_bundle( *, trace_path: Path, termination: Mapping[str, Any] | str, + _reservation: TrainingExportReservation | None = None, ) -> Path: if not isinstance(config, TrainingExportConfig): raise TypeError("config must be TrainingExportConfig") - trace_source = _require_regular_source(Path(trace_path), label="trace") target = config.output_dir parent = target.parent - _reject_symlink_ancestors(parent) - parent.mkdir(parents=True, exist_ok=True) - _reject_symlink_ancestors(parent) - if target.exists() or target.is_symlink(): - raise FileExistsError(f"immutable export already exists: {target}") - lock_path = parent / f".{target.name}.publish.lock" - lock_fd = os.open(lock_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + reservation = _reservation or reserve_turn_training_export(config) + if ( + not isinstance(reservation, TrainingExportReservation) + or not reservation.active + or reservation.target != target + or reservation.lock_path != parent / f".{target.name}.publish.lock" + ): + raise ValueError("training export reservation is invalid or inactive") temp: Path | None = None try: - os.close(lock_fd) + trace_source = _require_regular_source(Path(trace_path), label="trace") temp = Path(tempfile.mkdtemp(prefix=f".{target.name}.tmp-", dir=parent)) trace_copy = temp / "artifacts" / "trace" / trace_source.name _copy_source(trace_source, trace_copy) @@ -983,11 +1128,35 @@ def export_turn_training_bundle( rows = _read_trace(trace_copy) if any(row.get("run_id") != config.run_id for row in rows): raise ValueError("trace run_id differs from frozen training identity") + normalized_termination = _termination(termination) + allow_terminal_incomplete = _trace_proves_nonclean_terminal( + rows, + normalized_termination, + ) units: list[dict[str, Any]] = [] for row in rows: capture_type = row.get("capture_type") if capture_type == "llm_call": - units.append(_build_main_unit(config, temp, row, trace_ref)) + payload = ( + row.get("payload") + if isinstance(row.get("payload"), Mapping) + else {} + ) + tool_closure_ok = _validate_tool_result_closure( + trace_rows=rows, + capture_row=row, + response=_normalized_response(_response_source(payload)), + allow_terminal_incomplete=allow_terminal_incomplete, + ) + units.append( + _build_main_unit( + config, + temp, + row, + trace_ref, + tool_closure_ok=tool_closure_ok, + ) + ) elif capture_type == "compaction": units.append(_build_compaction_unit(config, temp, row, trace_ref)) units_path = temp / "turn-training-units.v1.jsonl" @@ -1007,7 +1176,7 @@ def export_turn_training_bundle( "unit_set_sha256": _hash_json( "researchharness.unit-set.v1", [unit["unit_id"] for unit in units] ), - "termination": _termination(termination), + "termination": normalized_termination, "final_patch_ref": final_patch_ref, "trace_ref": trace_ref, } @@ -1057,14 +1226,7 @@ def export_turn_training_bundle( shutil.rmtree(temp, ignore_errors=True) raise finally: - try: - os.close(lock_fd) - except OSError: - pass - try: - lock_path.unlink() - except FileNotFoundError: - pass + reservation.release() def _contains_prohibited_key(value: Any) -> bool: @@ -1222,6 +1384,11 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: trace_rows = _read_trace(trace_path) if any(row.get("run_id") != identity["run_id"] for row in trace_rows): raise ValueError("trace run_id differs from rollout identity") + summary_termination = summary.get("termination") + allow_terminal_incomplete = bool( + isinstance(summary_termination, Mapping) + and _trace_proves_nonclean_terminal(trace_rows, summary_termination) + ) captures = [row for row in trace_rows if row.get("capture_type") in {"llm_call", "compaction"}] if len(captures) != len(units): raise ValueError("trace capture closure differs from unit count") @@ -1374,11 +1541,13 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: provider_request=provider_request, raw_response=raw_response, ) + tool_closure_ok = True if expected_capture == "llm_call": - _validate_tool_result_closure( + tool_closure_ok = _validate_tool_result_closure( trace_rows=trace_rows, capture_row=row, response=expected_response, + allow_terminal_incomplete=allow_terminal_incomplete, ) findings = _oracle_findings(provider_request) provenance = unit.get("provenance") @@ -1424,7 +1593,7 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: expected_status = ( _compaction_capture_ok(payload, findings) if expected_capture == "compaction" - else _main_capture_ok(payload, findings) + else _main_capture_ok(payload, findings) and tool_closure_ok ) _validate_request_shape(unit["request"], capture_ok=expected_status) if unit.get("capture_status") != ("ok" if expected_status else "error"): @@ -1561,7 +1730,29 @@ def validate_turn_training_bundle(bundle_path: Path) -> dict[str, Any]: __all__ = [ "EXPORT_DIR_NAME", "TrainingExportConfig", + "TrainingExportReservation", "canonical_json_bytes", "export_turn_training_bundle", + "reserve_turn_training_export", + "validate_final_patch_source", "validate_turn_training_bundle", ] + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Validate one immutable ResearchHarness turn-training export bundle." + ) + parser.add_argument("bundle", type=Path, help="Path to turn-training-export-v1") + args = parser.parse_args(argv) + try: + summary = validate_turn_training_bundle(args.bundle) + except (OSError, TypeError, ValueError, json.JSONDecodeError) as exc: + print(f"invalid training export: {exc}", file=sys.stderr) + return 1 + print(json.dumps(summary, ensure_ascii=False, sort_keys=True, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/api/openai_server.py b/api/openai_server.py index 49032c8..742d150 100644 --- a/api/openai_server.py +++ b/api/openai_server.py @@ -150,8 +150,15 @@ def positive_int(value: Any, name: str) -> int: return parsed -def make_chat_completion_response(*, request_id: str, model: str, content: str) -> dict[str, Any]: - return { +def make_chat_completion_response( + *, + request_id: str, + model: str, + content: str, + run_id: str = "", + training_export_path: str = "", +) -> dict[str, Any]: + response = { "id": request_id, "object": "chat.completion", "created": int(time.time()), @@ -164,6 +171,12 @@ def make_chat_completion_response(*, request_id: str, model: str, content: str) } ], } + if run_id: + response["researchharness"] = { + "run_id": run_id, + "training_export_path": training_export_path or None, + } + return response def validate_chat_payload(payload: Any) -> dict[str, Any]: @@ -731,6 +744,8 @@ def run_chat_completion(payload: dict[str, Any], config: ServerConfig) -> dict[s request_id=request_id, model=requested_model_label, content=final_text, + run_id=run_id, + training_export_path=str(session.get("training_export_path", "") or ""), ) diff --git a/docs/turn-training-export-v1.md b/docs/turn-training-export-v1.md index 3cef853..25b913c 100644 --- a/docs/turn-training-export-v1.md +++ b/docs/turn-training-export-v1.md @@ -85,9 +85,15 @@ result = run_agent( run_id="run-auth-r0", harness_revision="72e8a44737aadabc", final_patch_path="/outputs/r0/agent.patch", + return_session=True, ) +print(result["run_id"], result["training_export_path"]) ``` +The default remains backward-compatible and returns only final text. +`return_session=True` returns the final text together with `run_id`, +`trace_path`, and `training_export_path`. + ## OpenAI-compatible API Add a `training-export` object to the request: @@ -111,7 +117,9 @@ The API uses one outer `run_id` for the agent trace and export, and writes the bundle under `//turn-training-export-v1`. API patch paths are relative to the selected workspace. Training export fails closed when input or output LLM wrappers are enabled because wrapper calls are outside V1's two -unit types. +unit types. Successful responses expose the frozen identity and bundle location +under `response.researchharness.run_id` and +`response.researchharness.training_export_path`. ## Validation @@ -122,3 +130,9 @@ multi-agent fields, path traversal, symlinks, hardlinks, and overwrite safety. Only units with `capture_status="ok"` are SFT candidates. Selection policy is a downstream responsibility. + +Validate a completed bundle from the command line: + +```bash +python -m agent_base.training_export /outputs/r0/turn-training-export-v1 +``` diff --git a/researchharness/runtime.py b/researchharness/runtime.py index 2cfe860..ae7cdc1 100644 --- a/researchharness/runtime.py +++ b/researchharness/runtime.py @@ -321,9 +321,12 @@ def run_agent( tool_protocol_version: str = "openai-chat-completions-native-tools-v1", tokenizer_revision: Optional[str] = None, final_patch_path: Optional[str | Path] = None, + return_session: bool = False, require_env: bool = True, -) -> str: - """Run ResearchHarness once and return the final assistant text.""" +) -> str | dict[str, Any]: + """Run once; optionally return trace/export metadata with the final text.""" + if not isinstance(return_session, bool): + raise ValueError("return_session must be a boolean") agent = create_agent( model_name=model_name, api_key=api_key, @@ -361,6 +364,8 @@ def run_agent( final_patch_path=final_patch_path, require_env=require_env, ) + if return_session: + return agent.run_session(prompt, images=images) return agent.run(prompt, images=images) diff --git a/tests/test_turn_training_export.py b/tests/test_turn_training_export.py index b06fb24..8094b4d 100644 --- a/tests/test_turn_training_export.py +++ b/tests/test_turn_training_export.py @@ -1,6 +1,8 @@ +import io import json import os import tempfile +import threading import types import unittest import sys @@ -431,6 +433,22 @@ def create(self, **kwargs): return response +class FakeSetupFailureClient: + def with_options(self, **kwargs): + raise RuntimeError("client setup crash") + + +class FakeInterruptingClient(FakeClient): + def __init__(self, interrupt_event, responses): + super().__init__(responses) + self.interrupt_event = interrupt_event + + def create(self, **kwargs): + response = super().create(**kwargs) + self.interrupt_event.set() + return response + + class TrainingExportTests(unittest.TestCase): def setUp(self): self._temp = tempfile.TemporaryDirectory() @@ -626,7 +644,7 @@ def test_tool_call_requires_matching_result_in_next_model_input(self): export_turn_training_bundle(config, trace_path=trace, termination="result") def test_tool_result_closure_binds_to_current_reused_call_id(self): - def rows_with_reused_id(include_current_result): + def rows_with_reused_id(include_current_result, include_current_assistant=True): rows = [_trace_rows()[0], _trace_rows()[2]] rows[1]["event_index"] = 2 call = rows[0]["payload"]["response"]["tool_calls"][0] @@ -638,10 +656,11 @@ def rows_with_reused_id(include_current_result): "content": "old result", } initial.extend([old_assistant, old_result]) - next_messages = [ - *initial, - {"role": "assistant", "content": "", "tool_calls": [call]}, - ] + next_messages = [*initial] + if include_current_assistant: + next_messages.append( + {"role": "assistant", "content": "", "tool_calls": [call]} + ) if include_current_result: next_messages.append( { @@ -661,6 +680,22 @@ def rows_with_reused_id(include_current_result): with self.assertRaisesRegex(ValueError, "tool-result"): export_turn_training_bundle(config, trace_path=trace, termination="result") + omitted_config = _config( + self.tmp_path, + name="reused-id-current-omitted", + patch_file=False, + ) + omitted_trace = _write_trace( + self.tmp_path / "reused-id-current-omitted/trace.jsonl", + rows_with_reused_id(False, include_current_assistant=False), + ) + with self.assertRaisesRegex(ValueError, "tool-result"): + export_turn_training_bundle( + omitted_config, + trace_path=omitted_trace, + termination="result", + ) + _, bundle = _export( self.tmp_path, name="reused-id-closed", @@ -679,6 +714,26 @@ def test_provider_attempt_metadata_cannot_carry_outcome_labels(self): with self.assertRaisesRegex(ValueError, "outcome"): export_turn_training_bundle(config, trace_path=trace, termination="result") + agent_rows = _trace_rows() + agent_rows[0]["payload"]["response"]["provider_attempts"][0]["details"] = { + "parent_agent_id": "invented" + } + agent_config = _config( + self.tmp_path, + name="attempt-agent", + patch_file=False, + ) + agent_trace = _write_trace( + self.tmp_path / "attempt-agent/trace.jsonl", + agent_rows, + ) + with self.assertRaisesRegex(ValueError, "multi-agent"): + export_turn_training_bundle( + agent_config, + trace_path=agent_trace, + termination="result", + ) + def test_rejected_empty_attempt_uses_runtime_text_semantics(self): rows = _trace_rows() rows[0]["payload"]["response"]["provider_attempts"][0][ @@ -855,6 +910,33 @@ def test_runtime_provider_boundary_retry_exports_one_accepted_unit(self): [item["status"] for item in llm_capture["payload"]["response"]["provider_attempts"]], ["rejected_empty", "accepted"], ) + duplicate_agent = MultiTurnReactAgent( + function_list=[], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + }, + }, + trace_dir=str(case / "duplicate-trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="runtime-run", + ) + duplicate_client = FakeClient([FakeResponse("must-not-run", "done")]) + duplicate_agent._llm_client = duplicate_client + duplicate_agent._llm_api_base = "http://fake" + with self.assertRaises(FileExistsError): + duplicate_agent._run_session( + "Must fail before provider.", + workspace_root=str(case), + ) + self.assertEqual(duplicate_client.requests, []) def test_provider_boundary_failures_are_retried_and_exported_as_error(self): from agent_base.react_agent import MultiTurnReactAgent @@ -912,6 +994,256 @@ def test_provider_boundary_failures_are_retried_and_exported_as_error(self): ["error", "parse_error", "parse_error"], ) + def test_provider_setup_failure_is_exported_as_error_rollout(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "setup-failure-runtime" + case.mkdir() + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-setup-failure", + task_revision_id="task-setup-failure@1", + rollout_id="r0", + rollout_index=0, + run_id="setup-failure-run", + harness_revision="test-rev", + ) + agent = MultiTurnReactAgent( + function_list=[], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 2, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="setup-failure-run", + ) + agent._llm_client = FakeSetupFailureClient() + agent._llm_api_base = "http://fake" + with patch("agent_base.react_agent.time.sleep", return_value=None): + session = agent._run_session("Finish the task.", workspace_root=str(case)) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + self.assertEqual(session["termination"], "llm api error") + self.assertEqual(summary["capture_error_count"], 1) + trace_rows = [ + json.loads(line) + for line in (bundle / summary["trace_ref"]["path"]).read_text().splitlines() + ] + llm_capture = next(row for row in trace_rows if row.get("capture_type") == "llm_call") + attempts = llm_capture["payload"]["response"]["provider_attempts"] + self.assertEqual([item["status"] for item in attempts], ["error", "error"]) + self.assertTrue( + all(item["error_type"] == "RuntimeError" for item in attempts) + ) + + def test_missing_explicit_patch_marks_rollout_error_without_losing_units(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "missing-patch-runtime" + case.mkdir() + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-missing-patch", + task_revision_id="task-missing-patch@1", + rollout_id="r0", + rollout_index=0, + run_id="missing-patch-run", + harness_revision="test-rev", + final_patch_path=case / "missing.patch", + ) + agent = MultiTurnReactAgent( + function_list=[], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="missing-patch-run", + ) + agent._llm_client = FakeClient([FakeResponse("patch-result", "done")]) + agent._llm_api_base = "http://fake" + session = agent._run_session("Finish the task.", workspace_root=str(case)) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + self.assertEqual(session["termination"], "final patch unavailable") + self.assertFalse(summary["termination"]["clean_termination"]) + self.assertFalse(summary["termination"]["submitted"]) + self.assertIsNone(summary["final_patch_ref"]) + self.assertEqual(summary["main_agent_turn_count"], 1) + self.assertEqual(_units(bundle)[0]["capture_status"], "ok") + + def test_inflight_interrupt_preserves_returned_provider_response(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "inflight-interrupt" + case.mkdir() + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-inflight-interrupt", + task_revision_id="task-inflight-interrupt@1", + rollout_id="r0", + rollout_index=0, + run_id="inflight-interrupt-run", + harness_revision="test-rev", + ) + agent = MultiTurnReactAgent( + function_list=[], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="inflight-interrupt-run", + ) + interrupt_event = threading.Event() + agent._llm_client = FakeInterruptingClient( + interrupt_event, + [FakeResponse("real-but-interrupted", "completed provider response")], + ) + agent._llm_api_base = "http://fake" + session = agent._run_session( + "Finish the task.", + workspace_root=str(case), + interrupt_event=interrupt_event, + ) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + unit = _units(bundle)[0] + self.assertEqual(session["termination"], "interrupted") + self.assertEqual(summary["main_agent_turn_count"], 1) + self.assertEqual(unit["capture_status"], "ok") + raw = json.loads( + (bundle / unit["source"]["raw_provider_response_ref"]["path"]).read_text() + ) + self.assertEqual(raw["id"], "real-but-interrupted") + + def test_tool_interrupt_preserves_pending_target_as_error_unit(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "tool-interrupt" + case.mkdir() + (case / "a.txt").write_text("alpha\n", encoding="utf-8") + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-tool-interrupt", + task_revision_id="task-tool-interrupt@1", + rollout_id="r0", + rollout_index=0, + run_id="tool-interrupt-run", + harness_revision="test-rev", + ) + call = _fake_tool_call("call-interrupt", "Read", '{"file_path":"a.txt"}') + agent = MultiTurnReactAgent( + function_list=["Read"], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=2, + training_export=config, + run_id="tool-interrupt-run", + ) + agent._llm_client = FakeClient( + [FakeResponse("pending-tool", None, tool_calls=[call])] + ) + agent._llm_api_base = "http://fake" + with patch.object(agent, "custom_call_tool", side_effect=KeyboardInterrupt): + session = agent._run_session("Read a.txt.", workspace_root=str(case)) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + self.assertEqual(session["termination"], "interrupted") + self.assertEqual(summary["main_agent_turn_count"], 1) + self.assertEqual(summary["capture_error_count"], 1) + self.assertEqual(_units(bundle)[0]["capture_status"], "error") + + def test_compaction_interrupt_exports_failed_semantic_unit(self): + from agent_base.react_agent import MultiTurnReactAgent + + case = self.tmp_path / "compaction-interrupt" + case.mkdir() + (case / "a.txt").write_text("alpha\n", encoding="utf-8") + config = TrainingExportConfig( + output_dir=case / EXPORT_DIR_NAME, + task_id="task-compaction-interrupt", + task_revision_id="task-compaction-interrupt@1", + rollout_id="r0", + rollout_index=0, + run_id="compaction-interrupt-run", + harness_revision="test-rev", + ) + call = _fake_tool_call("call-before-memory", "Read", '{"file_path":"a.txt"}') + agent = MultiTurnReactAgent( + function_list=["Read"], + llm={ + "model": MODEL, + "api_base": "http://fake", + "api_key": "fake", + "generate_cfg": { + "max_input_tokens": 32768, + "max_output_tokens": 128, + "max_retries": 1, + "compact_trigger_tokens": 1, + "recent_history_budget_tokens": 1, + }, + }, + trace_dir=str(case / "trace"), + workspace_root=str(case), + max_rounds=3, + training_export=config, + run_id="compaction-interrupt-run", + ) + agent._llm_client = FakeClient( + [FakeResponse("tool-before-memory", None, tool_calls=[call])] + ) + agent._llm_api_base = "http://fake" + with patch.object(agent, "call_compaction_api", side_effect=KeyboardInterrupt): + session = agent._run_session("Read a.txt.", workspace_root=str(case)) + bundle = Path(session["training_export_path"]) + summary = validate_turn_training_bundle(bundle) + units = _units(bundle) + self.assertEqual(session["termination"], "interrupted") + self.assertEqual(summary["main_agent_turn_count"], 1) + self.assertEqual(summary["memory_compaction_count"], 1) + self.assertEqual( + [unit["capture_status"] for unit in units], + ["ok", "error"], + ) + def test_runtime_parallel_tool_calls_close_into_next_provider_request(self): from agent_base.react_agent import MultiTurnReactAgent @@ -988,5 +1320,18 @@ def test_nonfinite_json_and_invalid_identity_types_are_rejected(self): run_id="run", harness_revision="rev", ) + def test_validator_cli_prints_validated_rollout_summary(self): + _, bundle = _export( + self.tmp_path, + name="validator-cli", + patch_file=False, + ) + output = io.StringIO() + with patch("sys.stdout", output): + self.assertEqual(export_module.main([str(bundle)]), 0) + rendered = json.loads(output.getvalue()) + self.assertEqual(rendered["run_id"], RUN_ID) + self.assertEqual(rendered["main_agent_turn_count"], 2) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_turn_training_export_api.py b/tests/test_turn_training_export_api.py index 745fc76..375e621 100644 --- a/tests/test_turn_training_export_api.py +++ b/tests/test_turn_training_export_api.py @@ -148,6 +148,12 @@ def test_api_passes_same_run_id_to_agent_trace_and_export(self): agent = FakeAPIAgent.seen[0] self.assertEqual(agent.run_id, agent.training_export.run_id) self.assertEqual(agent.training_export.output_dir.name, EXPORT_DIR_NAME) + metadata = response["researchharness"] + self.assertEqual(metadata["run_id"], agent.run_id) + self.assertEqual( + metadata["training_export_path"], + str(agent.training_export.output_dir), + ) if __name__ == "__main__": diff --git a/tests/test_turn_training_export_entrypoints.py b/tests/test_turn_training_export_entrypoints.py index 6fd8d50..878f781 100644 --- a/tests/test_turn_training_export_entrypoints.py +++ b/tests/test_turn_training_export_entrypoints.py @@ -2,11 +2,13 @@ import unittest from pathlib import Path import sys +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from agent_base.react_agent import _parse_cli_args from agent_base.training_export import EXPORT_DIR_NAME +import researchharness.runtime as runtime_module from researchharness.runtime import _resolve_training_export @@ -90,6 +92,35 @@ def test_python_resolver_rejects_bool_index_and_preserves_no_export(self): self.assertEqual(trace_dir, "trace") self.assertEqual(run_id, "ordinary-run") + def test_public_python_api_can_return_export_metadata_without_breaking_text_default(self): + class FakeAgent: + def run_session(self, prompt, images=None): + return { + "run_id": "run-python", + "result_text": "done", + "trace_path": "/tmp/trace.jsonl", + "training_export_path": "/tmp/turn-training-export-v1", + } + + def run(self, prompt, images=None): + return "done" + + fake_agent = FakeAgent() + with patch.object(runtime_module, "create_agent", return_value=fake_agent): + session = runtime_module.run_agent( + "fix", + return_session=True, + require_env=False, + ) + self.assertEqual(session["run_id"], "run-python") + self.assertEqual( + session["training_export_path"], + "/tmp/turn-training-export-v1", + ) + with patch.object(runtime_module, "create_agent", return_value=fake_agent): + text = runtime_module.run_agent("fix", require_env=False) + self.assertEqual(text, "done") + if __name__ == "__main__": unittest.main()