Add producer-owned training captures for harness sessions

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 training_sampling
 
 from .config import OpenCodeConfig
 from .opencode_runtime import (
@@ -175,6 +177,48 @@ 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"]
+            choice = entry["response"]["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"],
+                    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"),
+                )
+            )
+        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 +255,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 +264,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 +310,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 +326,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 +355,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 +396,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 +588,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 +605,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)]

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)
 
@@ -338,7 +343,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 +369,12 @@ 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"]
         delta = choice.get("delta") or {}
         content = delta.get("content")
         if content:
@@ -381,7 +390,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"]
@@ -438,6 +451,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", []),
     }
 
 
@@ -461,6 +475,8 @@ def _build_turn_record(
         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 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")
@@ -472,6 +488,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 +506,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 +555,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 +633,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 +666,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,

src/openenv/core/harness/README.md

@@ -108,3 +108,53 @@ 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 normalizes `token_id:N` values into IDs and retains them across streaming.
+
+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.

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)