diff --git a/docs/source/environments/harbor.md b/docs/source/environments/harbor.md
index e79bdae73b..349ef4d862 100644
--- a/docs/source/environments/harbor.md
+++ b/docs/source/environments/harbor.md
@@ -287,7 +287,14 @@ Hugging Face Sandbox on a Space.
**Runs.** Every run you may see, filterable by status. A run shows what the agent did as a timeline
(its prompt, thinking, each tool call with its input and output, terminal sessions, the final
answer), the result and reward, the trace checks, and downloads of the result JSON and, for a
-trainable rollout, the training contract. Tick two to four runs to compare them.
+trainable rollout, the **Training trace** and **Capture audit**. The training trace is the same
+validated `TrainingTrace` returned by `HarborSession.fetch_training_trace()`: selected agent calls,
+engine token IDs, logprobs and explicit loss masks. Load it with
+`TrainingTrace.model_validate_json(Path("run.training_trace.json").read_text())` after importing
+`Path` from `pathlib` and `TrainingTrace` from `openenv.core.harness`. Task rewards and the full
+rollout remain in Result JSON. Capture audit preserves the existing `contract.json` export,
+including excluded turns. Older runs without explicit masks can still download their audit but
+cannot produce the typed training trace. Tick two to four runs to compare them.
**Setup.** The server's endpoint, which sandboxes have working credentials, the deployment settings
below, the datasets and the agents.
diff --git a/envs/harbor_env/README.md b/envs/harbor_env/README.md
index c5765cd2c0..e1fc374d89 100644
--- a/envs/harbor_env/README.md
+++ b/envs/harbor_env/README.md
@@ -295,7 +295,14 @@ Hugging Face Sandbox on a Space.
**Runs.** Every run you may see, filterable by status. A run shows what the agent did as a timeline
(its prompt, thinking, each tool call with its input and output, terminal sessions, the final
answer), the result and reward, the trace checks, and downloads of the result JSON and, for a
-trainable rollout, the training contract. Tick two to four runs to compare them.
+trainable rollout, the **Training trace** and **Capture audit**. The training trace is the same
+validated `TrainingTrace` returned by `HarborSession.fetch_training_trace()`: selected agent calls,
+engine token IDs, logprobs and explicit loss masks. Load it with
+`TrainingTrace.model_validate_json(Path("run.training_trace.json").read_text())` after importing
+`Path` from `pathlib` and `TrainingTrace` from `openenv.core.harness`. Task rewards and the full
+rollout remain in Result JSON. Capture audit preserves the existing `contract.json` export,
+including excluded turns. Older runs without explicit masks can still download their audit but
+cannot produce the typed training trace. Tick two to four runs to compare them.
**Setup.** The server's endpoint, which sandboxes have working credentials, the deployment settings
below, the datasets and the agents.
diff --git a/envs/harbor_env/harness.py b/envs/harbor_env/harness.py
index 4eaa5de66f..4dc0c15dfd 100644
--- a/envs/harbor_env/harness.py
+++ b/envs/harbor_env/harness.py
@@ -43,8 +43,9 @@
VerifyResult,
)
from openenv.core.harness.capture.upstream import training_sampling
+from openenv.core.harness.training import TrainingTrace
from openenv.harbor.client import HarborEnv
-from openenv.harbor.contract import to_trace_entries
+from openenv.harbor.contract import to_trace_entries, to_training_trace
from openenv.harbor.models import HarborRolloutResult
logger = logging.getLogger(__name__)
@@ -322,6 +323,12 @@ def _warn_if_oversized(self) -> None:
getattr(result, "n_turns", -1),
)
+ def fetch_training_trace(self) -> TrainingTrace:
+ """Return producer-validated tokens and masks after the rollout finishes."""
+ if self.result is None:
+ raise RuntimeError("rollout produced no capture")
+ return to_training_trace(self.result)
+
def fetch_proxy_trace(self) -> list[dict[str, Any]]:
"""Per-turn captured records, in TRL's `TraceEntry` shape.
diff --git a/envs/opencode_env/harness.py b/envs/opencode_env/harness.py
index 8c282ce937..01b433221f 100644
--- a/envs/opencode_env/harness.py
+++ b/envs/opencode_env/harness.py
@@ -38,8 +38,10 @@
ResourceSession,
ResourceSessionFactory,
ToolResult,
+ TrainingTrace,
VerifyResult,
)
+from openenv.core.harness.capture.upstream import normalize_response, training_sampling
from .config import OpenCodeConfig
from .opencode_runtime import (
@@ -175,6 +177,66 @@ def fetch_trace(self) -> str:
"""Return the raw ``opencode run`` log (JSON-lines when ``run_format=json``)."""
return self.sandbox.read_text(agent_log_path(self.config))
+ def fetch_training_trace(self) -> TrainingTrace:
+ """Select agent calls in OpenEnv and export their exact captured tokens."""
+ from types import SimpleNamespace
+
+ from openenv.core.harness.capture.contract import to_training_trace
+ from openenv.core.harness.capture.export import export_session
+ from openenv.core.harness.capture.graph import RolloutGraph, TurnNode
+
+ if self._proxy_trace_path is None:
+ raise ValueError("training requires transparent_proxy capture")
+ records = self.fetch_proxy_trace()
+ if not records:
+ raise RuntimeError("rollout produced no capture")
+ graph = RolloutGraph()
+ for entry in records:
+ request = entry["request"]
+ response = normalize_response(entry["response"])
+ if (
+ response.get("upstream_status", 200) >= 400
+ or response.get("upstream_error")
+ or response.get("error")
+ ):
+ if entry.get("completion_token_ids"):
+ raise ValueError("failed upstream call contains sampled tokens")
+ continue
+ choices = response.get("choices") or []
+ if not choices:
+ raise ValueError("successful capture has no response choices")
+ choice = choices[0]
+ graph.add_turn(
+ TurnNode(
+ node_id=str(entry["turn"]),
+ prompt_ids=entry["prompt_token_ids"],
+ sampled_ids=entry["completion_token_ids"],
+ sampled_logprobs=(
+ entry["per_token_logps"]
+ if choice.get("logprobs", {}) is not None
+ else None
+ ),
+ request_messages=request.get("messages", []),
+ request_tools=request.get("tools"),
+ n_tools=len(request.get("tools") or []),
+ response_message=choice.get("message", {}),
+ finish_reason=entry.get("finish_reason"),
+ )
+ )
+ if not graph.nodes():
+ raise RuntimeError("rollout produced no successful model calls")
+ document = export_session(
+ SimpleNamespace(
+ graph=graph,
+ session_id="opencode",
+ metadata={},
+ findings=[],
+ purpose="train",
+ ),
+ include_messages=True,
+ )
+ return to_training_trace(graph, document)
+
def fetch_proxy_trace(self) -> list[dict[str, Any]]:
"""Return per-turn proxy-captured records (Mode B only).
@@ -211,6 +273,7 @@ def __init__(
config: OpenCodeConfig,
sandbox_backend: SandboxBackend,
mode: Literal["black_box", "transparent_proxy"] = "black_box",
+ sampling: dict[str, Any] | None = None,
verifier: Verifier | None = None,
install_timeout_s: int = 240,
setup_timeout_s: int = 300,
@@ -219,6 +282,9 @@ def __init__(
) -> None:
if mode not in {"black_box", "transparent_proxy"}:
raise ValueError(f"Unknown mode: {mode!r}")
+ self.sampling = training_sampling(sampling) if sampling is not None else None
+ if self.sampling is not None and mode != "transparent_proxy":
+ raise ValueError("training sampling requires transparent_proxy mode")
self._config = config
self._backend = sandbox_backend
self._mode = mode
@@ -262,7 +328,10 @@ def create(
backoff = self._create_backoff_s * (2**i)
_log.warning(
"factory.create attempt %d/%d failed (%r); retrying in %.1fs",
- i + 1, self._create_attempts, exc, backoff,
+ i + 1,
+ self._create_attempts,
+ exc,
+ backoff,
)
time.sleep(backoff)
raise last_exc
@@ -275,6 +344,7 @@ def _create_once(
start_agent: bool = True,
) -> OpenCodeSession:
import logging
+
_log = logging.getLogger(__name__)
oc_task = OpenCodeTask.coerce(task)
@@ -282,15 +352,15 @@ def _create_once(
_log.info(
"factory.create: creating sandbox timeout=%ds mode=%s",
- sandbox_timeout, self._mode,
+ sandbox_timeout,
+ self._mode,
)
sandbox = self._backend.create(
timeout_s=sandbox_timeout,
metadata={"episode_id": episode_id} if episode_id else None,
)
- sid = (
- getattr(sandbox, "sandbox_id", None)
- or getattr(getattr(sandbox, "raw", None), "sandbox_id", "?")
+ sid = getattr(sandbox, "sandbox_id", None) or getattr(
+ getattr(sandbox, "raw", None), "sandbox_id", "?"
)
_log.info("factory.create: sandbox=%s — bootstrapping…", sid)
# Any failure past here (bootstrap/proxy/agent) must tear the sandbox down.
@@ -303,7 +373,8 @@ def _create_once(
if self._mode == "transparent_proxy":
_log.info(
"factory.create: starting interception proxy on :%d → %s",
- _PROXY_PORT, self._config.base_url,
+ _PROXY_PORT,
+ self._config.base_url,
)
proxy_bg_job, base_url_override, proxy_trace_path = self._start_proxy(
sandbox
@@ -343,7 +414,9 @@ def _create_once(
try:
sandbox.kill() # best-effort: don't let a cleanup failure mask the root cause
except Exception:
- _log.exception("factory.create: sandbox.kill() during cleanup also failed")
+ _log.exception(
+ "factory.create: sandbox.kill() during cleanup also failed"
+ )
raise
# ------------------------------------------------------------------
@@ -533,6 +606,11 @@ def _start_proxy(
)
sandbox.write_text(f"{proxy_dir(self._config)}/__init__.py", "")
+ if self.sampling is not None:
+ sandbox.write_text(
+ proxy_source_path(self._config), _PROXY_SOURCE_PATH.read_text()
+ )
+
proxy_args = [
"python",
"interception.py",
@@ -545,6 +623,8 @@ def _start_proxy(
"--top-logprobs",
str(self._config.proxy_top_logprobs),
]
+ if self.sampling is not None:
+ proxy_args.extend(["--sampling", json.dumps(self.sampling)])
if self._config.proxy_max_tokens_cap is not None:
proxy_args.extend(
["--max-tokens-cap", str(self._config.proxy_max_tokens_cap)]
diff --git a/envs/opencode_env/sandbox/interception.py b/envs/opencode_env/sandbox/interception.py
index 131d410243..57f3977bc4 100644
--- a/envs/opencode_env/sandbox/interception.py
+++ b/envs/opencode_env/sandbox/interception.py
@@ -63,6 +63,7 @@ class ProxyConfig:
trace_path: str = "/tmp/opencode-proxy-trace.jsonl"
host: str = "127.0.0.1"
port: int = 7000
+ sampling: dict[str, Any] | None = None
top_logprobs: int = 5
request_timeout_s: float = 600.0
# Cap ``max_tokens`` before forwarding. OpenCode historically asks for very
@@ -89,6 +90,7 @@ class TurnRecord:
request: dict[str, Any]
response: dict[str, Any]
logprobs: list[dict[str, Any]] | None
+ prompt_token_ids: list[int]
completion_tokens: list[str]
completion_token_ids: list[int]
per_token_logps: list[float]
@@ -130,9 +132,7 @@ async def chat_completions(request: Request) -> Response:
try:
body = json.loads(raw_body)
except json.JSONDecodeError:
- return JSONResponse(
- status_code=400, content={"error": "invalid json body"}
- )
+ return JSONResponse(status_code=400, content={"error": "invalid json body"})
forwarded_body = _prepare_forwarded_body(body, cfg)
headers = {
@@ -178,6 +178,11 @@ def _prepare_forwarded_body(body: dict[str, Any], cfg: ProxyConfig) -> dict[str,
``max_completion_tokens``.
"""
forwarded = copy.deepcopy(body)
+ if cfg.sampling is not None:
+ forwarded.update(cfg.sampling)
+ forwarded["logprobs"] = True
+ forwarded["return_tokens_as_token_ids"] = True
+ forwarded["return_token_ids"] = True
forwarded.setdefault("logprobs", True)
forwarded.setdefault("top_logprobs", cfg.top_logprobs)
@@ -252,7 +257,11 @@ async def _proxy_unary(
record = _build_turn_record(
turn_idx=turn_idx,
request_body=forwarded_body,
- response_json=response_json,
+ response_json=(
+ {**response_json, "upstream_status": upstream_response.status_code}
+ if upstream_response.status_code >= 400
+ else response_json
+ ),
latency_s=latency,
)
trace_file.write(record.to_json() + "\n")
@@ -338,7 +347,7 @@ async def _stream() -> Any:
yield line + "\n"
if not line.startswith("data:"):
continue
- data = line[len("data:"):].strip()
+ data = line[len("data:") :].strip()
if data == "[DONE]":
continue
try:
@@ -364,8 +373,16 @@ async def _stream() -> Any:
def _accumulate_stream_chunk(chunk: dict[str, Any], acc: dict[str, Any]) -> None:
+ if chunk.get("prompt_token_ids") is not None:
+ acc["prompt_token_ids"] = chunk["prompt_token_ids"]
for choice in chunk.get("choices", []) or []:
idx = choice.get("index", 0)
+ if choice.get("prompt_token_ids") is not None:
+ acc["prompt_token_ids"] = choice["prompt_token_ids"]
+ if choice.get("token_ids") is not None:
+ acc.setdefault("token_ids_by_idx", {}).setdefault(idx, []).extend(
+ choice["token_ids"]
+ )
delta = choice.get("delta") or {}
content = delta.get("content")
if content:
@@ -381,7 +398,11 @@ def _accumulate_stream_chunk(chunk: dict[str, Any], acc: dict[str, Any]) -> None
tc_idx = tc.get("index", 0)
bucket = acc["tool_calls_by_idx"].setdefault(
(idx, tc_idx),
- {"id": None, "type": "function", "function": {"name": "", "arguments": ""}},
+ {
+ "id": None,
+ "type": "function",
+ "function": {"name": "", "arguments": ""},
+ },
)
if tc.get("id"):
bucket["id"] = tc["id"]
@@ -406,6 +427,7 @@ def _assemble_streamed_response(
| set(acc["finish_by_idx"])
| {k[0] for k in acc["tool_calls_by_idx"]}
| set(acc["logprobs_by_idx"])
+ | set(acc.get("token_ids_by_idx", {}))
| {0}
)
choices: list[dict[str, Any]] = []
@@ -431,6 +453,8 @@ def _assemble_streamed_response(
}
if acc["logprobs_by_idx"].get(idx):
choice["logprobs"] = {"content": acc["logprobs_by_idx"][idx]}
+ if idx in acc.get("token_ids_by_idx", {}):
+ choice["token_ids"] = acc["token_ids_by_idx"][idx]
choices.append(choice)
return {
"id": last_chunk.get("id", ""),
@@ -438,6 +462,7 @@ def _assemble_streamed_response(
"model": last_chunk.get("model", ""),
"choices": choices,
"usage": last_chunk.get("usage"),
+ "prompt_token_ids": acc.get("prompt_token_ids", []),
}
@@ -455,14 +480,16 @@ def _build_turn_record(
content_lp = logprobs_field.get("content") or []
tokens: list[str] = []
- token_ids: list[int] = []
+ token_ids: list[int] = list(choice.get("token_ids") or [])
per_token_logps: list[float] = []
for entry in content_lp:
tokens.append(entry.get("token", ""))
- # OpenAI returns no raw token ids; vLLM returns them as ``token_id``.
- token_id = entry.get("token_id")
- if token_id is not None:
- token_ids.append(int(token_id))
+ if not choice.get("token_ids"):
+ token_id = entry.get("token_id")
+ if token_id is None and str(entry.get("token", "")).startswith("token_id:"):
+ token_id = int(entry["token"].removeprefix("token_id:"))
+ if token_id is not None:
+ token_ids.append(int(token_id))
lp = entry.get("logprob")
if lp is not None:
per_token_logps.append(float(lp))
@@ -472,6 +499,9 @@ def _build_turn_record(
request=request_body,
response=response_json,
logprobs=content_lp,
+ prompt_token_ids=response_json.get("prompt_token_ids")
+ or choice.get("prompt_token_ids")
+ or [],
completion_tokens=tokens,
completion_token_ids=token_ids,
per_token_logps=per_token_logps,
@@ -487,8 +517,7 @@ def _strip_logprobs(response_json: dict[str, Any]) -> dict[str, Any]:
choices = out.get("choices")
if isinstance(choices, list):
out["choices"] = [
- {k: v for k, v in (ch or {}).items() if k != "logprobs"}
- for ch in choices
+ {k: v for k, v in (ch or {}).items() if k != "logprobs"} for ch in choices
]
return out
@@ -537,9 +566,7 @@ def start(self) -> None:
lifespan="on",
)
self._server = uvicorn.Server(config)
- self._thread = threading.Thread(
- target=self._run_server, daemon=True
- )
+ self._thread = threading.Thread(target=self._run_server, daemon=True)
self._thread.start()
# Wait for the server to accept connections.
deadline = time.time() + 10
@@ -617,6 +644,7 @@ def main() -> None:
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=7000)
parser.add_argument("--top-logprobs", type=int, default=5)
+ parser.add_argument("--sampling", type=json.loads, default=None)
parser.add_argument("--request-timeout", type=float, default=600.0)
parser.add_argument(
"--max-tokens-cap",
@@ -649,6 +677,7 @@ def main() -> None:
host=args.host,
port=args.port,
top_logprobs=args.top_logprobs,
+ sampling=args.sampling,
request_timeout_s=args.request_timeout,
max_tokens_cap=args.max_tokens_cap,
disable_thinking=args.disable_thinking,
diff --git a/src/openenv/core/harness/README.md b/src/openenv/core/harness/README.md
index 1acdbd50a7..a8ebb3ed86 100644
--- a/src/openenv/core/harness/README.md
+++ b/src/openenv/core/harness/README.md
@@ -108,3 +108,57 @@ HF Dataset Viewer renders `results.jsonl` as `split=train`.
- **Filtering** — `should_keep` is a `Callable[[EpisodeRecord], bool]`. Defaults
to keeping everything. The CLI filters `reward >= 0` unless `--keep-losses`
is set.
+
+## Training captures
+
+A loop-owning session can implement `TrainableSession.fetch_training_trace()` to
+return a `TrainingTrace`. Harbor and native OpenCode implement this API. Raw
+`fetch_proxy_trace()` remains available for diagnostics and older consumers.
+
+```python
+from openenv.core.harness import TrainingTrace, TrainingTurn
+
+trace = TrainingTrace(turns=[TrainingTurn(
+ node_id="call-0",
+ prompt_token_ids=[10, 11],
+ completion_token_ids=[20, 21, 22],
+ per_token_logps=[-0.2, -0.4, -0.1],
+ loss_mask=[0, 0, 1, 0, 1],
+)])
+restored = TrainingTrace.model_validate_json(trace.model_dump_json())
+```
+
+The mask covers prompt plus completion. Prompt positions are zero; completion
+positions select the sampled tokens that receive loss. A zero-masked turn stays
+available as context. Tokens, masks and finite nonpositive logprobs are validated
+when the object is built or deserialized. Missing masks and duplicate call IDs
+are errors. Tokens are never reconstructed from text.
+
+The environment selects agent calls and excludes auxiliary calls and discarded
+retries before returning this object. Shared graph nodes appear once. Selection
+uses the existing capture classifier; its tool-manifest and retry heuristics are
+not a substitute for explicit harness evidence. The raw graph remains available
+for auditing. Changing a task's supervision means changing its producer mask,
+not adding a turn-selection callback in the trainer.
+
+`verify()` remains the source of the task reward. `None` is ungraded and `0` is an
+incorrect answer. A single-call rollout produces a diagnostic warning, not a
+capture failure. Masks do not alter the verifier score. Trainers own advantages,
+loss weighting, packing and weight updates. Graph identity does not by itself
+normalize a rollout that expands into multiple training rows.
+
+For training, construct the session factory with the trainer's `sampling` policy.
+Harbor and native OpenCode apply it before inference. The trainer need not compare
+sampling metadata after generation. Native OpenCode requires `transparent_proxy`
+mode and an engine returning prompt IDs and sampled-token logprobs; its standalone
+proxy retains choice token IDs across streaming, with `token_id:N` as a fallback.
+Training export uses the same token/logprob pairing check as Harbor when logprob
+entries identify their tokens. A disagreement on an agent turn rejects the training
+trace; raw capture stays available for inspection. Plain token text is not
+re-tokenized to infer an ID.
+
+Migration from TRL selection hooks: use `fetch_training_trace()` instead of a raw
+trace, move whole-turn exclusions into the producer's `loss_mask`, and keep final
+answers eligible unless your task explicitly masks them. The old Harbor export
+continues to read legacy captures; the strict API requires explicit masks and
+node IDs. Evaluation-only captures cannot be used for training.
diff --git a/src/openenv/core/harness/__init__.py b/src/openenv/core/harness/__init__.py
index 0b8f138bcd..2bad4b6f80 100644
--- a/src/openenv/core/harness/__init__.py
+++ b/src/openenv/core/harness/__init__.py
@@ -56,9 +56,11 @@
ToolResult,
ToolTraceEntry,
TraceEntry,
+ TrainableSession,
VerifyResult,
)
from .tools import resolve_tool_conflicts
+from .training import TrainingTrace, TrainingTurn
__all__ = [
# Trainer-side rollout API (openenv.core.harness.rollout)
@@ -80,6 +82,9 @@
"ToolResult",
"ToolTraceEntry",
"TraceEntry",
+ "TrainableSession",
+ "TrainingTrace",
+ "TrainingTurn",
"VerifyResult",
"build_harness_rollout_func",
# Turn-based agentic harness API (RFC 005)
diff --git a/src/openenv/core/harness/capture/__init__.py b/src/openenv/core/harness/capture/__init__.py
index f34faed83b..f74314d168 100644
--- a/src/openenv/core/harness/capture/__init__.py
+++ b/src/openenv/core/harness/capture/__init__.py
@@ -13,7 +13,12 @@
cross-checked against the harness's own trace.
"""
-from .contract import measure_retokenization_skew, to_trace_entries, to_turn_records
+from .contract import (
+ measure_retokenization_skew,
+ to_trace_entries,
+ to_training_trace,
+ to_turn_records,
+)
from .detection import APIType, detect
from .graph import RolloutGraph, TurnNode
from .runner import CaptureServer
@@ -24,6 +29,7 @@
"CaptureServer",
"to_turn_records",
"to_trace_entries",
+ "to_training_trace",
"measure_retokenization_skew",
"APIType",
"detect",
diff --git a/src/openenv/core/harness/capture/contract.py b/src/openenv/core/harness/capture/contract.py
index b3a3539ed4..b673d60545 100644
--- a/src/openenv/core/harness/capture/contract.py
+++ b/src/openenv/core/harness/capture/contract.py
@@ -1,6 +1,6 @@
"""What a rollout hands a trainer.
-The mask-aware training contract is `to_trace_entries`: one engine prompt, sampled completion,
+The typed training contract is `to_training_trace`: one engine prompt, sampled completion,
behavior logprobs, and a loss mask over prompt+completion per retained model call. Prompt ids and
sampled ids are never reconstructed from text. Masks are authoritative after reconciliation.
@@ -21,6 +21,7 @@
import logging
from typing import Any
+from ..training import TrainingTrace
from .graph import RolloutGraph, TurnNode
from .validate import check_turn, validate_training_turn
@@ -204,6 +205,16 @@ def _completion_masks(
return masks
+def to_training_trace(graph: RolloutGraph, document: dict[str, Any]) -> TrainingTrace:
+ """Return validated agent calls once each, including shared graph prefixes."""
+ if any("[FATAL" in finding for finding in document.get("validation", [])):
+ raise ValueError("cannot train a capture with fatal validation findings")
+ _require_trainable(document)
+ if any(not _usable(node) for node in _agent_nodes(graph, document)):
+ raise ValueError("agent capture has missing or invalid engine tokens/logprobs")
+ return TrainingTrace.from_entries(to_trace_entries(graph, document))
+
+
def to_turn_records(
graph: RolloutGraph, document: dict[str, Any]
) -> list[tuple[list[int], list[int], list[float]]]:
diff --git a/src/openenv/core/harness/capture/validate.py b/src/openenv/core/harness/capture/validate.py
index 5500e880ce..6569300e32 100644
--- a/src/openenv/core/harness/capture/validate.py
+++ b/src/openenv/core/harness/capture/validate.py
@@ -344,29 +344,13 @@ def check_rollout(
report.add(FATAL, "root_count", f"expected 1 root, got {stats['n_roots']}")
if stats["n_turns"] == 1 and budget_stop_count <= 0:
- # A recorded proxy budget stop explains a one-turn rollout. In particular, a tool can
- # return enough data after the first call to exhaust the next prompt's context budget.
- # Its original verifier score and sampled tokens remain valid; rejecting it here would
- # turn a legitimate bounded attempt into an unintended retry.
- # ONE call for an entire agentic task. Capture is trivially self-consistent here (a single
- # turn has nothing to stitch to and no prefix to disagree with), so every other check in this
- # file passes and the rollout reads as clean. It is not: an agent that made one model call
- # and stopped did not attempt the task.
- #
- # Found the hard way. swe-agent, trae-agent, nemo-agent and antigravity-sdk each passed 5/5
- # while producing exactly one turn per task and solving 0/5, for four unrelated harness-side
- # reasons (litellm cost registry, a null `prompt_tokens_details`, a tool-less prompt format,
- # and an SDK loop that exits after the first tool call). The capture layer was right every
- # time and the rollouts were still worthless.
- #
- # FATAL because the whole point of this layer is refusing to hand over data we cannot stand
- # behind, and a one-turn agentic rollout is a harness failure wearing a clean capture.
+ # A single call can be a valid answer or a harness that stopped early.
+ # Keep the diagnostic, but let the verifier score the task outcome.
report.add(
- FATAL,
+ WARN,
"degenerate_rollout",
- "exactly 1 model call for the whole task: the agent stopped after its first "
- "response. Capture is self-consistent because there is nothing to stitch, so the "
- "other checks cannot see this. Read the trial's agent stdout for the real cause.",
+ "exactly 1 model call for the whole task; inspect the agent log if this "
+ "was unexpected. Valid captured tokens remain eligible for training.",
)
if stats["n_discarded"]:
diff --git a/src/openenv/core/harness/rollout.py b/src/openenv/core/harness/rollout.py
index 7c4408a412..64d11e734e 100644
--- a/src/openenv/core/harness/rollout.py
+++ b/src/openenv/core/harness/rollout.py
@@ -29,6 +29,7 @@
from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool
from ..env_server.types import State
from ..llm_client import LLMResponse
+from .training import TrainingTrace
Message = dict[str, Any]
RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"})
@@ -210,6 +211,15 @@ def wait_for_completion(self, timeout_s: float | None = ...) -> int: ...
def fetch_proxy_trace(self) -> list[TraceEntry]: ...
+@runtime_checkable
+class TrainableSession(Protocol):
+ """A loop-owning session that supplies validated, masked engine captures."""
+
+ def wait_for_completion(self, timeout_s: float | None = ...) -> int: ...
+
+ def fetch_training_trace(self) -> TrainingTrace: ...
+
+
class ResourceSession(ABC):
"""Per-rollout environment/resource session exposed to harnesses."""
@@ -826,6 +836,7 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]:
"ToolResult",
"ToolTraceEntry",
"TraceEntry",
+ "TrainableSession",
"VerifyResult",
"build_harness_rollout_func",
]
diff --git a/src/openenv/core/harness/training.py b/src/openenv/core/harness/training.py
new file mode 100644
index 0000000000..111091c0c7
--- /dev/null
+++ b/src/openenv/core/harness/training.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Validated token capture shared by environment producers and trainers."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+class TrainingTurn(BaseModel):
+ """One engine call. The required loss mask covers prompt plus completion."""
+
+ model_config = ConfigDict(strict=True, extra="forbid")
+
+ node_id: str = Field(min_length=1)
+ prompt_token_ids: list[int]
+ completion_token_ids: list[int]
+ per_token_logps: list[float]
+ loss_mask: list[int]
+ request: dict[str, Any] = Field(default_factory=dict)
+ response: dict[str, Any] = Field(default_factory=dict)
+ metadata: dict[str, Any] = Field(default_factory=dict)
+
+ @model_validator(mode="after")
+ def validate_capture(self) -> TrainingTurn:
+ from .capture.validate import validate_training_turn
+
+ validate_training_turn(
+ self.prompt_token_ids,
+ self.completion_token_ids,
+ self.per_token_logps,
+ self.loss_mask,
+ )
+ return self
+
+
+class TrainingTrace(BaseModel):
+ """Selected agent calls, each with an authoritative mask and unique identity.
+
+ Zero-masked turns remain available for context and usage accounting. Task rewards
+ come from the session's `verify()` method, independently of these masks.
+ """
+
+ model_config = ConfigDict(strict=True, extra="forbid")
+
+ schema_version: Literal[1] = 1
+ turns: list[TrainingTurn]
+
+ @model_validator(mode="after")
+ def unique_calls(self) -> TrainingTrace:
+ ids = [turn.node_id for turn in self.turns]
+ if len(set(ids)) != len(ids):
+ raise ValueError("training trace contains duplicate node_ids")
+ return self
+
+ @classmethod
+ def from_entries(cls, entries: list[dict[str, Any]]) -> TrainingTrace:
+ """Adapt producer-selected entries without inferring tokens or masks."""
+ return cls(
+ turns=[
+ TrainingTurn(
+ node_id=entry["metadata"]["node_id"],
+ prompt_token_ids=entry["prompt_token_ids"],
+ completion_token_ids=entry["completion_token_ids"],
+ per_token_logps=entry["per_token_logps"],
+ loss_mask=entry["loss_mask"],
+ request=entry.get("request", {}),
+ response=entry.get("response", {}),
+ metadata=entry.get("metadata", {}),
+ )
+ for entry in entries
+ ]
+ )
+
+ def to_trace_entries(self) -> list[dict[str, Any]]:
+ """Return diagnostic records for reward functions and trace viewers."""
+ return [
+ {
+ **turn.model_dump(exclude={"node_id"}),
+ "metadata": {**turn.metadata, "node_id": turn.node_id},
+ }
+ for turn in self.turns
+ ]
diff --git a/src/openenv/harbor/contract.py b/src/openenv/harbor/contract.py
index c11361005c..4418d59f8c 100644
--- a/src/openenv/harbor/contract.py
+++ b/src/openenv/harbor/contract.py
@@ -9,6 +9,7 @@
from typing import Any
from openenv.core.harness.capture.validate import validate_training_turn
+from openenv.core.harness.training import TrainingTrace, TrainingTurn
from .models import HarborRolloutResult
@@ -131,6 +132,54 @@ def to_trace_entries(result: HarborRolloutResult) -> list[dict[str, Any]]:
return entries
+def to_training_trace(result: HarborRolloutResult) -> TrainingTrace:
+ """Return the strict training API; legacy results need explicit producer masks."""
+ if result.rollout_type != "train":
+ raise ValueError("eval-only rollout has no exact-token training contract")
+ if any("[FATAL" in finding for finding in result.findings):
+ raise ValueError("cannot train a capture with fatal validation findings")
+ turns = []
+ for turn in result.turns:
+ if turn.role != "agent" or turn.discarded or not turn.completion_token_ids:
+ continue
+ if turn.loss_mask is None:
+ raise ValueError("training capture requires an explicit loss_mask")
+ if not turn.trainable and any(turn.loss_mask):
+ raise ValueError("ineligible capture contains supervised tokens")
+ turns.append(
+ TrainingTurn(
+ node_id=turn.node_id,
+ prompt_token_ids=turn.prompt_token_ids,
+ completion_token_ids=turn.completion_token_ids,
+ per_token_logps=turn.per_token_logps,
+ loss_mask=turn.loss_mask,
+ request={
+ "messages": turn.request_messages,
+ "tools": turn.request_tools,
+ },
+ response={
+ "choices": [
+ {
+ "message": {
+ "role": "assistant",
+ "content": turn.text,
+ "tool_calls": _openai_tool_calls(turn.tool_calls)
+ or None,
+ },
+ "finish_reason": turn.finish_reason,
+ }
+ ]
+ },
+ metadata={
+ "turn": turn.turn,
+ "sampling_params": turn.sampling_params,
+ "requested_sampling_params": turn.requested_sampling_params,
+ },
+ )
+ )
+ return TrainingTrace(turns=turns)
+
+
def export_training_contract(result: HarborRolloutResult) -> dict[str, Any]:
"""Export validated supervision and an explicit audit of excluded turns.
diff --git a/src/openenv/harbor/ui.py b/src/openenv/harbor/ui.py
index 037b1ed593..a1a8db995b 100644
--- a/src/openenv/harbor/ui.py
+++ b/src/openenv/harbor/ui.py
@@ -194,19 +194,9 @@ def _hub_problem(spec: str) -> str | None:
def _contract(r: dict[str, Any]) -> dict[str, Any] | None:
- """The training contract of a rollout: exactly what a trainer consumes, nothing else.
+ """Keep the legacy audit export, including excluded turns and task rewards.
- Per turn, `(prompt_token_ids, completion_token_ids, per_token_logps)` plus the reward. The
- logprobs are the load-bearing part and the reason this is a separate download: they are the
- behaviour policy's, recorded at sampling time, and cannot be recovered afterwards by re-running
- the prompt. Discarded turns are kept but flagged, because they were generated and billed and a
- trainer must be able to see them in order to exclude them deliberately.
-
- `None` for an eval rollout: a download named `contract.json` whose every `prompt_token_ids` is
- `[]` would look like a contract and contain none.
-
- Raises:
- `ValueError`: when the result has FATAL findings or an invalid mask; the exporter refuses it.
+ Eval rollouts return None; malformed captures raise ValueError.
"""
turns = r.get("turns") or []
if not turns or r.get("rollout_type", "train") == "eval":
@@ -217,6 +207,18 @@ def _contract(r: dict[str, Any]) -> dict[str, Any] | None:
return export_training_contract(HarborRolloutResult.model_validate(r))
+def _training_trace(r: dict[str, Any]) -> dict[str, Any] | None:
+ """Download the same validated token objects returned by HarborSession."""
+ if not r.get("turns") or r.get("rollout_type", "train") == "eval":
+ return None
+ from .contract import to_training_trace
+ from .models import HarborRolloutResult
+
+ return to_training_trace(HarborRolloutResult.model_validate(r)).model_dump(
+ mode="json"
+ )
+
+
# ── the page's own pieces ─────────────────────────────────────────────────────────────────────────
@@ -941,9 +943,13 @@ def hb_download(args: list[Any]) -> dict[str, Any]:
if not result:
return {"error": "This link has expired. Open the run again."}
name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(rec.get("id") or "rollout"))
- if kind == "contract":
+ if kind in {"contract", "training_trace"}:
try:
- contract = _contract(result)
+ contract = (
+ _training_trace(result)
+ if kind == "training_trace"
+ else _contract(result)
+ )
except (
ValueError
) as exc: # the exporter refuses a rollout with FATAL findings
@@ -953,7 +959,7 @@ def hb_download(args: list[Any]) -> dict[str, Any]:
"error": "An eval rollout has nothing to train on, so it has no contract."
}
return {
- "name": f"{name}.contract.json",
+ "name": f"{name}.{kind}.json",
"text": json.dumps(contract, indent=2),
}
return {
diff --git a/src/openenv/harbor/ui_pages.py b/src/openenv/harbor/ui_pages.py
index e4db236462..69bdcc7697 100644
--- a/src/openenv/harbor/ui_pages.py
+++ b/src/openenv/harbor/ui_pages.py
@@ -532,7 +532,10 @@ def _run_head(rec: dict[str, Any], token: str) -> str:
)
if trainable:
actions.append(
- f''
+ f''
+ )
+ actions.append(
+ f''
)
title = rec.get("task_title") or rec.get("task_name")
return (
diff --git a/tests/core/test_harness_rollout_backcompat.py b/tests/core/test_harness_rollout_backcompat.py
index 9cef426764..03baeded3e 100644
--- a/tests/core/test_harness_rollout_backcompat.py
+++ b/tests/core/test_harness_rollout_backcompat.py
@@ -28,6 +28,7 @@
"ToolResult",
"ToolTraceEntry",
"TraceEntry",
+ "TrainableSession",
"VerifyResult",
"build_harness_rollout_func",
]
diff --git a/tests/envs/test_capture_model_call_budget.py b/tests/envs/test_capture_model_call_budget.py
index 43ce1826f3..fbb71bf778 100644
--- a/tests/envs/test_capture_model_call_budget.py
+++ b/tests/envs/test_capture_model_call_budget.py
@@ -153,7 +153,7 @@ async def invalid(request):
assert session.graph.stats()["n_turns"] == 0
-def test_single_turn_without_recorded_budget_stop_still_fails(app_and_engine):
+def test_single_turn_without_recorded_budget_stop_warns(app_and_engine):
from openenv.core.harness.capture.export import export_session
app, _ = app_and_engine
@@ -162,7 +162,7 @@ def test_single_turn_without_recorded_budget_stop_still_fails(app_and_engine):
_chat(client, session.session_id)
document = export_session(session, capture_level="text")
assert any(
- "[FATAL] degenerate_rollout" in finding for finding in document["validation"]
+ "[WARN] degenerate_rollout" in finding for finding in document["validation"]
)
diff --git a/tests/envs/test_harbor_capture_level.py b/tests/envs/test_harbor_capture_level.py
index c7192f3f98..05cd3ad3eb 100644
--- a/tests/envs/test_harbor_capture_level.py
+++ b/tests/envs/test_harbor_capture_level.py
@@ -362,9 +362,7 @@ def test_an_eval_rollout_is_recorded_without_fatal_findings(monkeypatch):
"""`check_turn`'s FATALs (`no_prompt_ids`, `no_logprobs`) are the expected condition here.
Letting them fire would mark every eval turn unusable and teach everyone to ignore findings.
- Two calls, not one: `degenerate_rollout` is FATAL for a single-call agentic rollout and stays
- that way on this path, because an agent that made one call and stopped did not attempt the task
- whether or not its tokens were captured.
+ Two calls exercise the multi-turn capture contract without the single-call warning.
"""
replies = [reply(), reply()]
diff --git a/tests/envs/test_harbor_ui_training_contract.py b/tests/envs/test_harbor_ui_training_contract.py
index c7bf9122e8..3c6829bd7c 100644
--- a/tests/envs/test_harbor_ui_training_contract.py
+++ b/tests/envs/test_harbor_ui_training_contract.py
@@ -6,7 +6,7 @@
from harbor_env.harness import to_trace_entries
from openenv.harbor.contract import export_training_contract
from openenv.harbor.models import HarborRolloutResult, HarborTurn
-from openenv.harbor.ui import _contract
+from openenv.harbor.ui import _contract, _training_trace
def rollout():
@@ -69,6 +69,73 @@ def test_invalid_mask_cannot_be_hidden_by_the_download_path():
_contract(result.model_dump())
+def test_typed_download_matches_session_capture_and_keeps_audit_separate():
+ from unittest.mock import MagicMock
+
+ from harbor_env.harness import HarborSession
+ from openenv.core.harness import TrainingTrace
+
+ result = rollout()
+ result.turns += [
+ result.turns[0].model_copy(
+ update={"turn": 1, "node_id": "aux", "role": "auxiliary"}
+ ),
+ result.turns[0].model_copy(
+ update={"turn": 2, "node_id": "retry", "discarded": True}
+ ),
+ result.turns[0].model_copy(
+ update={
+ "turn": 3,
+ "node_id": "context",
+ "loss_mask": [0, 0, 0, 0],
+ "trainable": False,
+ }
+ ),
+ ]
+ session = HarborSession(
+ env=MagicMock(),
+ split="test",
+ task_index=0,
+ instruction="task",
+ harness="opencode",
+ sandbox="docker",
+ llm_url="http://unused",
+ model="test",
+ )
+ session.result = result
+ download = TrainingTrace.model_validate_json(
+ json.dumps(_training_trace(result.model_dump()))
+ )
+ assert download == session.fetch_training_trace()
+ assert [t.node_id for t in download.turns] == ["agent", "context"]
+ assert [t.loss_mask for t in download.turns] == [[0, 0, 1, 0], [0, 0, 0, 0]]
+ assert _contract(result.model_dump()) == export_training_contract(result)
+
+
+def test_typed_download_rejects_legacy_masks_and_eval_capture():
+ result = rollout()
+ result.turns[0].loss_mask = None
+ with pytest.raises(ValueError, match="explicit loss_mask"):
+ _training_trace(result.model_dump())
+ assert _contract(result.model_dump()) is not None
+ result.rollout_type = "eval"
+ assert _training_trace(result.model_dump()) is None
+
+
+def test_run_page_keeps_downloads_behind_the_same_grant():
+ from openenv.harbor import ui_pages
+
+ result = rollout().model_dump()
+ rec = {"id": "typed-download-test", "status": "completed", "result": result}
+ page = ui_pages.run_html(rec)
+ token = ui_pages.grant(rec["id"])
+ for kind in ["result", "contract", "training_trace"]:
+ assert f'data-dl="{kind}" data-grant="{token}"' in page
+ assert ui_pages.granted(token) == rec["id"]
+ result["rollout_type"] = "eval"
+ assert 'data-dl="training_trace"' not in ui_pages.run_html(rec)
+
+
def test_ui_validation_keeps_the_qualified_acp_profile(tmp_path, monkeypatch):
import importlib
from types import SimpleNamespace
diff --git a/tests/envs/test_training_trace.py b/tests/envs/test_training_trace.py
new file mode 100644
index 0000000000..01fd9907bb
--- /dev/null
+++ b/tests/envs/test_training_trace.py
@@ -0,0 +1,513 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+import json
+from unittest.mock import MagicMock
+
+import pytest
+from openenv.core.harness import TrainingTrace, TrainingTurn
+from openenv.core.harness.capture.contract import to_training_trace as graph_trace
+from openenv.core.harness.capture.graph import RolloutGraph, TurnNode
+from openenv.core.harness.capture.upstream import training_sampling
+from openenv.harbor.contract import to_training_trace
+from openenv.harbor.models import HarborRolloutResult, HarborTurn
+
+
+def turn(**changes):
+ return TrainingTurn(
+ **{
+ "node_id": "agent",
+ "prompt_token_ids": [1, 2],
+ "completion_token_ids": [3, 4, 5],
+ "per_token_logps": [-0.1, -0.2, -0.3],
+ "loss_mask": [0, 0, 1, 0, 1],
+ **changes,
+ }
+ )
+
+
+def test_partial_mask_and_zero_mask_survive_wire_round_trip():
+ original = TrainingTrace(turns=[turn(), turn(node_id="context", loss_mask=[0] * 5)])
+ restored = TrainingTrace.model_validate_json(original.model_dump_json())
+ assert restored == original
+ assert restored.turns[0].loss_mask == [0, 0, 1, 0, 1]
+ assert restored.turns[1].per_token_logps == [-0.1, -0.2, -0.3]
+
+
+@pytest.mark.parametrize(
+ "changes",
+ [
+ {"loss_mask": None},
+ {"loss_mask": [0, 0, 1]},
+ {"loss_mask": [1, 0, 1, 0, 1]},
+ {"loss_mask": [0, 0, True, 0, 1]},
+ {"completion_token_ids": [3, True, 5]},
+ {"per_token_logps": []},
+ {"per_token_logps": [-0.1, float("nan"), -0.3]},
+ {"per_token_logps": [-0.1, 0.2, -0.3]},
+ {"per_token_logps": [-0.1, True, -0.3]},
+ {"prompt_token_ids": []},
+ ],
+)
+def test_malformed_capture_fails_at_producer_boundary(changes):
+ with pytest.raises(ValueError):
+ turn(**changes)
+
+
+def test_missing_mask_and_unknown_version_rejected():
+ values = turn().model_dump(exclude={"loss_mask"})
+ with pytest.raises(ValueError):
+ TrainingTurn.model_validate(values)
+ with pytest.raises(ValueError):
+ TrainingTrace(schema_version=2, turns=[])
+
+
+def test_duplicate_call_is_not_supervised_twice():
+ with pytest.raises(ValueError, match="duplicate node_ids"):
+ TrainingTrace(turns=[turn(), turn()])
+
+
+def test_shared_graph_prefix_is_emitted_once_and_partial_masks_preserved():
+ graph = RolloutGraph()
+ for name, prompt, sampled in [
+ ("root", [1], [2, 3]),
+ ("left", [1, 2, 3, 8], [4]),
+ ("right", [1, 2, 3, 9], [5]),
+ ("aux", [90], [91]),
+ ]:
+ graph.add_turn(
+ TurnNode(
+ node_id=name,
+ prompt_ids=prompt,
+ sampled_ids=sampled,
+ sampled_logprobs=[-0.1] * len(sampled),
+ )
+ )
+ document = {
+ "sequences": [
+ {
+ "role": "agent",
+ "node_ids": ["root", "left"],
+ "loss_mask": [0, 1, 0, 0, 1],
+ },
+ {
+ "role": "agent",
+ "node_ids": ["root", "right"],
+ "loss_mask": [0, 1, 0, 0, 1],
+ },
+ {"role": "auxiliary", "node_ids": ["aux"], "loss_mask": [0, 1]},
+ ]
+ }
+ trace = graph_trace(graph, document)
+ assert [t.node_id for t in trace.turns] == ["root", "left", "right"]
+ assert sum(sum(t.loss_mask) for t in trace.turns) == 3
+ assert trace.turns[0].loss_mask == [0, 1, 0]
+
+
+def test_harbor_masks_do_not_change_reward_or_drop_zero_masked_agent_calls():
+ turns = [
+ HarborTurn(
+ turn=0, **turn().model_dump(exclude={"request", "response", "metadata"})
+ ),
+ HarborTurn(
+ turn=1,
+ **turn(node_id="context", loss_mask=[0] * 5).model_dump(
+ exclude={"request", "response", "metadata"}
+ ),
+ ),
+ HarborTurn(turn=2, role="auxiliary"),
+ HarborTurn(turn=3, discarded=True),
+ ]
+ for reward in [None, 0.0, 1.0]:
+ result = HarborRolloutResult(turns=turns, reward=reward)
+ trace = to_training_trace(result)
+ assert [t.node_id for t in trace.turns] == ["agent", "context"]
+ assert result.reward == reward
+ result.turns[0].loss_mask = None
+ with pytest.raises(ValueError, match="explicit loss_mask"):
+ to_training_trace(result)
+ result.turns[0].loss_mask = [0, 0, 1, 0, 1]
+
+
+def test_eval_and_fatal_capture_are_not_empty_successful_training():
+ with pytest.raises(ValueError, match="eval-only"):
+ to_training_trace(HarborRolloutResult(rollout_type="eval"))
+ with pytest.raises(ValueError, match="fatal"):
+ to_training_trace(HarborRolloutResult(findings=["[FATAL] bad capture"]))
+
+
+def test_native_stream_keeps_engine_prompt_ids_and_normalizes_token_ids():
+ from opencode_env.sandbox.interception import (
+ _accumulate_stream_chunk,
+ _assemble_streamed_response,
+ _build_turn_record,
+ )
+
+ acc = {
+ "content_by_idx": {},
+ "tool_calls_by_idx": {},
+ "finish_by_idx": {},
+ "logprobs_by_idx": {},
+ }
+ _accumulate_stream_chunk(
+ {
+ "prompt_token_ids": [1, 2],
+ "choices": [
+ {
+ "index": 0,
+ "delta": {"content": "answer"},
+ "logprobs": {"content": [{"token": "token_id:3", "logprob": -0.3}]},
+ }
+ ],
+ },
+ acc,
+ )
+ final = {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
+ _accumulate_stream_chunk(final, acc)
+ record = _build_turn_record(
+ turn_idx=0,
+ request_body={},
+ response_json=_assemble_streamed_response(final, acc),
+ latency_s=0.1,
+ )
+ assert record.prompt_token_ids == [1, 2]
+ assert record.completion_token_ids == [3]
+ assert record.per_token_logps == [-0.3]
+
+
+def test_native_policy_overrides_harness_before_generation():
+ from opencode_env.config import OpenCodeConfig
+ from opencode_env.harness import OpenCodeSessionFactory
+ from opencode_env.sandbox.interception import _prepare_forwarded_body, ProxyConfig
+
+ policy = training_sampling({"temperature": 0.8})
+ factory = OpenCodeSessionFactory(
+ config=OpenCodeConfig(base_url="http://unused"),
+ sandbox_backend=MagicMock(),
+ mode="transparent_proxy",
+ sampling=policy,
+ )
+ forwarded = _prepare_forwarded_body(
+ {"temperature": 1.0, "top_p": 0.2, "logprobs": False},
+ ProxyConfig(upstream_url="http://unused", sampling=factory.sampling),
+ )
+ assert all(forwarded[k] == v for k, v in policy.items())
+ assert forwarded["logprobs"] is True
+ assert forwarded["return_tokens_as_token_ids"] is True
+ assert forwarded["return_token_ids"] is True
+
+
+def test_native_session_exports_all_agent_roots_without_auxiliary_calls():
+ from opencode_env.config import OpenCodeConfig
+ from opencode_env.harness import OpenCodeSession
+ from opencode_env.task import OpenCodeTask
+
+ records = []
+ for i, tools in enumerate([[{"type": "function"}], None, [{"type": "function"}]]):
+ records.append(
+ {
+ "turn": i,
+ "prompt_token_ids": [10 + i],
+ "completion_token_ids": [20 + i],
+ "per_token_logps": [-0.1],
+ "request": {
+ "messages": [{"role": "user", "content": "task"}],
+ "tools": tools,
+ },
+ "response": {"choices": [{"message": {"content": "ok"}}]},
+ }
+ )
+ sandbox = MagicMock(
+ read_text=lambda path: "\n".join(json.dumps(r) for r in records)
+ )
+ session = OpenCodeSession(
+ sandbox=sandbox,
+ config=OpenCodeConfig(base_url="http://unused"),
+ task=OpenCodeTask.coerce("task"),
+ proxy_trace_path="trace",
+ )
+ trace = session.fetch_training_trace()
+ assert [t.node_id for t in trace.turns] == ["0", "2"]
+ assert all(t.loss_mask == [0, 1] for t in trace.turns)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize(
+ "choice_ids, token_text, reject_pairing",
+ [
+ (None, "token_id:20", False),
+ ([], "token_id:20", False),
+ ([20], "answer", False),
+ ([20], "token_id:20", False),
+ ([20], "token_id:21", True),
+ ],
+)
+def test_native_proxy_http_capture_reaches_training_contract(
+ monkeypatch, tmp_path, stream, choice_ids, token_text, reject_pairing
+):
+ from functools import partial
+
+ import httpx
+ from fastapi.testclient import TestClient
+ from opencode_env.config import OpenCodeConfig
+ from opencode_env.harness import OpenCodeSession
+ from opencode_env.sandbox import interception
+ from opencode_env.task import OpenCodeTask
+
+ policy = training_sampling({"temperature": 0.7})
+ logprobs = {"content": [{"token": token_text, "logprob": -0.25}]}
+ ids = {"token_ids": choice_ids} if choice_ids is not None else {}
+
+ def engine(request):
+ body = json.loads(request.content)
+ assert all(body[key] == value for key, value in policy.items())
+ assert body["logprobs"] is True
+ assert body["return_tokens_as_token_ids"] is True
+ assert body["return_token_ids"] is True
+ if body["messages"][0]["content"] == "fail":
+ return httpx.Response(400, json={"error": {"message": "rejected request"}})
+ if stream:
+ chunks = [
+ {
+ "prompt_token_ids": [10],
+ "choices": [
+ {
+ "index": 0,
+ "delta": {"content": "answer"},
+ "logprobs": logprobs,
+ **(ids if choice_ids else {}),
+ }
+ ],
+ },
+ {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
+ ]
+ if choice_ids == []:
+ chunks[-1]["choices"][0]["token_ids"] = []
+ content = (
+ "".join("data: " + json.dumps(chunk) + "\n\n" for chunk in chunks)
+ + "data: [DONE]\n\n"
+ )
+ return httpx.Response(
+ 200, text=content, headers={"content-type": "text/event-stream"}
+ )
+ return httpx.Response(
+ 200,
+ json={
+ "prompt_token_ids": [10],
+ "choices": [
+ {
+ "message": {"role": "assistant", "content": "answer"},
+ "logprobs": logprobs,
+ **ids,
+ "finish_reason": "stop",
+ }
+ ],
+ },
+ )
+
+ monkeypatch.setattr(
+ interception.httpx,
+ "AsyncClient",
+ partial(httpx.AsyncClient, transport=httpx.MockTransport(engine)),
+ )
+ path = tmp_path / "trace.jsonl"
+ config = interception.ProxyConfig(
+ upstream_url="http://engine", trace_path=str(path), sampling=policy
+ )
+ with TestClient(interception._build_app(config)) as client:
+ failed = client.post(
+ "/v1/chat/completions",
+ json={
+ "model": "test",
+ "stream": stream,
+ "messages": [{"role": "user", "content": "fail"}],
+ },
+ )
+ assert failed.status_code == 400
+ response = client.post(
+ "/v1/chat/completions",
+ json={
+ "model": "test",
+ "messages": [{"role": "user", "content": "task"}],
+ "tools": [
+ {
+ "type": "function",
+ "function": {"name": "bash", "parameters": {"type": "object"}},
+ }
+ ],
+ "stream": stream,
+ "temperature": 1.0,
+ "logprobs": False,
+ },
+ )
+ assert response.status_code == 200
+ assert "answer" in response.text
+ session = OpenCodeSession(
+ sandbox=MagicMock(read_text=lambda _: path.read_text()),
+ task=OpenCodeTask(instruction="task"),
+ config=OpenCodeConfig(base_url="http://engine"),
+ verifier=None,
+ proxy_trace_path=str(path),
+ )
+ raw = session.fetch_proxy_trace()
+ assert len(raw) == 2
+ if reject_pairing:
+ with pytest.raises(
+ ValueError, match="missing or invalid engine tokens/logprobs"
+ ):
+ session.fetch_training_trace()
+ assert session.fetch_proxy_trace() == raw
+ return
+ trace = session.fetch_training_trace()
+ assert len(trace.turns) == 1
+ assert trace.turns[0].prompt_token_ids == [10]
+ assert trace.turns[0].completion_token_ids == [20]
+ assert trace.turns[0].per_token_logps == [-0.25]
+ assert trace.turns[0].loss_mask == [0, 1]
+
+
+@pytest.mark.parametrize("response", [{}, {"choices": []}])
+def test_missing_successful_response_is_not_silently_skipped(response):
+ from opencode_env.config import OpenCodeConfig
+ from opencode_env.harness import OpenCodeSession
+ from opencode_env.task import OpenCodeTask
+
+ session = OpenCodeSession(
+ sandbox=MagicMock(
+ read_text=lambda _: json.dumps({"request": {}, "response": response})
+ ),
+ config=OpenCodeConfig(base_url="http://unused"),
+ task=OpenCodeTask.coerce("task"),
+ proxy_trace_path="trace",
+ )
+ with pytest.raises(ValueError, match="no response choices"):
+ session.fetch_training_trace()
+
+
+def test_all_failed_upstream_calls_remain_a_transport_failure():
+ from opencode_env.config import OpenCodeConfig
+ from opencode_env.harness import OpenCodeSession
+ from opencode_env.task import OpenCodeTask
+
+ session = OpenCodeSession(
+ sandbox=MagicMock(
+ read_text=lambda _: json.dumps(
+ {"request": {}, "response": {"error": "unavailable"}}
+ )
+ ),
+ config=OpenCodeConfig(base_url="http://unused"),
+ task=OpenCodeTask.coerce("task"),
+ proxy_trace_path="trace",
+ )
+ with pytest.raises(RuntimeError, match="no successful model calls"):
+ session.fetch_training_trace()
+
+
+def test_streamed_choice_ids_accumulate_across_chunks_without_reparsing_text():
+ from opencode_env.sandbox.interception import (
+ _accumulate_stream_chunk,
+ _assemble_streamed_response,
+ _build_turn_record,
+ )
+
+ acc = {
+ "content_by_idx": {},
+ "tool_calls_by_idx": {},
+ "finish_by_idx": {},
+ "logprobs_by_idx": {},
+ }
+ for token in [20, 21]:
+ _accumulate_stream_chunk(
+ {
+ "prompt_token_ids": [10],
+ "choices": [
+ {
+ "index": 0,
+ "token_ids": [token],
+ "delta": {"content": "token_id:literal"},
+ "logprobs": {
+ "content": [{"token": "token_id:literal", "logprob": -0.25}]
+ },
+ }
+ ],
+ },
+ acc,
+ )
+ record = _build_turn_record(
+ turn_idx=1,
+ request_body={},
+ response_json=_assemble_streamed_response({}, acc),
+ latency_s=0.1,
+ )
+ assert record.completion_token_ids == [20, 21]
+ assert record.per_token_logps == [-0.25, -0.25]
+
+
+def test_choice_token_logprob_mismatch_is_not_truncated():
+ from opencode_env.sandbox.interception import _build_turn_record
+
+ record = _build_turn_record(
+ turn_idx=1,
+ request_body={},
+ response_json={
+ "prompt_token_ids": [10],
+ "choices": [
+ {
+ "token_ids": [20, 21],
+ "logprobs": {"content": [{"token": "answer", "logprob": -0.25}]},
+ }
+ ],
+ },
+ latency_s=0.1,
+ )
+ assert record.completion_token_ids == [20, 21]
+ with pytest.raises(ValueError):
+ TrainingTurn(
+ node_id="call",
+ prompt_token_ids=record.prompt_token_ids,
+ completion_token_ids=record.completion_token_ids,
+ per_token_logps=record.per_token_logps,
+ loss_mask=[0, 1, 1],
+ )
+
+
+def test_mispaired_auxiliary_call_does_not_invalidate_agent_capture():
+ from opencode_env.config import OpenCodeConfig
+ from opencode_env.harness import OpenCodeSession
+ from opencode_env.task import OpenCodeTask
+
+ records = [
+ {
+ "turn": i,
+ "prompt_token_ids": [10 + i],
+ "completion_token_ids": [20],
+ "per_token_logps": [-0.25],
+ "request": {"messages": [], "tools": tools},
+ "response": {
+ "choices": [
+ {
+ "token_ids": [20],
+ "message": {"content": "answer"},
+ "logprobs": {"content": [{"token": token, "logprob": -0.25}]},
+ }
+ ]
+ },
+ }
+ for i, (tools, token) in enumerate(
+ [
+ (None, "token_id:21"),
+ ([{"type": "function"}], "token_id:20"),
+ ]
+ )
+ ]
+ session = OpenCodeSession(
+ sandbox=MagicMock(
+ read_text=lambda _: "\n".join(json.dumps(r) for r in records)
+ ),
+ config=OpenCodeConfig(base_url="http://unused"),
+ task=OpenCodeTask.coerce("task"),
+ proxy_trace_path="trace",
+ )
+ trace = session.fetch_training_trace()
+ assert [t.node_id for t in trace.turns] == ["1"]
+ assert trace.turns[0].per_token_logps == [-0.25]
+ assert session.fetch_proxy_trace() == records