Add producer-owned training captures for harness sessions

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.

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.

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.
 

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,22 +344,23 @@ def _create_once(
         start_agent: bool = True,
     ) -> OpenCodeSession:
         import logging
+
         _log = logging.getLogger(__name__)
 
         oc_task = OpenCodeTask.coerce(task)
         sandbox_timeout = int(self._config.agent_timeout_s) + 300
 
         _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)]