feat(token-id-capture): Claude Code external-harness example

nemo_gym/base_responses_api_agent.py

@@ -26,7 +26,11 @@
     BaseVerifyResponse,
 )
 from nemo_gym.config_types import ROLLOUT_PATH_PREFIX
-from nemo_gym.global_config import OBSERVABILITY_ENABLED_KEY_NAME, get_first_server_config_dict
+from nemo_gym.global_config import (
+    OBSERVABILITY_ENABLED_KEY_NAME,
+    TOKEN_ID_CAPTURE_ENABLED_KEY_NAME,
+    get_first_server_config_dict,
+)
 from nemo_gym.openai_utils import (
     NeMoGymResponse,
     NeMoGymResponseCreateParamsNonStreaming,
@@ -43,7 +47,13 @@
 
 
 class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig):
-    pass
+    # Whether this agent's rollouts participate in training token capture. Native agents receive
+    # token ids inline on the model response and leave this off; opaque external harnesses (whose
+    # returned output carries no token ids) set it true so their model calls are correlated and
+    # captured into the token store, then rebuilt into a token-bearing response.output. The run-level
+    # token_id_capture_enabled switch still gates the capture infrastructure; this scopes which
+    # agents use it.
+    token_id_capture: bool = False
 
 
 class BaseResponsesAPIAgent(BaseServer):
@@ -79,21 +89,33 @@ async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyRespo
 
         return app
 
-    def _model_call_capture_enabled(self) -> bool:
-        # Fail closed: an agent whose client carries no usable global config runs uncorrelated
-        # rather than erroring on every model call.
+    def _capture_correlation_enabled(self) -> bool:
+        """Whether the per-rollout ``/ng-rollout/<id>`` correlation prefix should be applied.
+
+        Two independent capture paths consume the same prefix:
+        - Eval model-call capture (``observability_enabled``), which applies to every agent.
+        - Training token capture (``token_id_capture_enabled``), which applies only to agents
+          that opt in with the per-agent ``token_id_capture`` flag -- native agents carry token
+          ids inline and do not need the store, so they do not emit the prefix for token capture.
+
+        Fail closed: an agent whose client carries no usable global config runs uncorrelated
+        rather than erroring on every model call.
+        """
         global_config = getattr(self.server_client, "global_config_dict", None)
         if not isinstance(global_config, Mapping):
             return False
-        return bool(global_config.get(OBSERVABILITY_ENABLED_KEY_NAME, False))
+        token_capture = bool(global_config.get(TOKEN_ID_CAPTURE_ENABLED_KEY_NAME, False)) and bool(
+            getattr(self.config, "token_id_capture", False)
+        )
+        return bool(global_config.get(OBSERVABILITY_ENABLED_KEY_NAME, False) or token_capture)
 
     def rollout_id_from_run(self, body: Any) -> Optional[str]:
         """Per-rollout capture id for a run-request (its task/rollout indices).
 
-        None when model-call capture (observability) is disabled or the body carries no indices,
-        so callers apply no correlation prefix in either case.
+        None when neither capture path is enabled or the body carries no indices, so callers apply
+        no correlation prefix in either case.
         """
-        if not self._model_call_capture_enabled():
+        if not self._capture_correlation_enabled():
             return None
         return maybe_rollout_id_from_run_body(body)
 

nemo_gym/base_responses_api_model.py

@@ -67,6 +67,17 @@
     BaseServer,
     SimpleServer,
 )
+from nemo_gym.token_id_capture import (
+    CaptureContext,
+    TokenIdCaptureConfig,
+    capture_tokens,
+    reset_token_sink,
+    set_token_sink,
+)
+
+# The read route and its store factory need Gym's server stack, so they are not
+# re-exported from the leaf package (see nemo_gym/token_id_capture/__init__.py).
+from nemo_gym.token_id_capture.routes import install_token_capture_routes, make_token_store
 
 
 logger = logging.getLogger(__name__)
@@ -76,6 +87,47 @@
 _ANTHROPIC_CONVERTER = AnthropicConverter()
 
 
+def _request_facts(body: Any) -> dict[str, Any]:
+    """What the harness asked for, read off the parsed request body.
+
+    Two things, both used to tell trajectory calls from side calls: the model the
+    harness requested (not the one the server served), and whether the request
+    declared tools. Reading it from the parsed request avoids touching the body a
+    second time in middleware.
+    """
+    if body is None:
+        return {}
+    getter = body.get if isinstance(body, dict) else lambda key, default=None: getattr(body, key, default)
+    model = getter("model", None)
+    tools = getter("tools", None)
+    facts: dict[str, Any] = {}
+    if isinstance(model, str) and model:
+        facts["requested_model"] = model
+    if tools is not None:
+        facts["has_tools"] = bool(tools)
+    return facts
+
+
+def _request_messages(body: Any) -> list[dict]:
+    """The conversation a request carries, across the three dialects.
+
+    Used only to identify which recorded call this request continues (see
+    ``token_id_capture.lineage``): the assistant turns in it are the ones we
+    produced. Chat and Anthropic both use ``messages``; Responses carries
+    ``input``, which is a string for a first turn and a list of items after that.
+    """
+    if body is None:
+        return []
+    getter = body.get if isinstance(body, dict) else lambda key, default=None: getattr(body, key, default)
+    messages = getter("messages", None)
+    if isinstance(messages, list):
+        return [m if isinstance(m, dict) else m.model_dump() for m in messages if m is not None]
+    items = getter("input", None)
+    if isinstance(items, list):
+        return [i if isinstance(i, dict) else i.model_dump() for i in items if i is not None]
+    return []
+
+
 class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig):
     pass
 
@@ -90,7 +142,12 @@ def setup_webserver(self) -> FastAPI:
 
         self.setup_session_middleware(app)
         capture_config = ModelCallCaptureConfig.model_validate(self.server_client.global_config_dict)
-        install_model_call_capture(app, capture_config, model_server_name=self.config.name)
+        install_model_call_capture(
+            app,
+            capture_config,
+            model_server_name=self.config.name,
+            global_config_dict=self.server_client.global_config_dict,
+        )
 
         app.post("/v1/chat/completions")(self.chat_completions_dispatch)
 
@@ -192,8 +249,15 @@ async def _invoke_chat_completions(
         # only `body`. Dispatch on whichever this server declares so the shared dispatch works for
         # all of them.
         if "request" in inspect.signature(self.chat_completions).parameters:
-            return await self.chat_completions(request=request, body=params)
-        return await self.chat_completions(body=params)
+            completion = await self.chat_completions(request=request, body=params)
+        else:
+            completion = await self.chat_completions(body=params)
+        await capture_tokens(
+            completion,
+            request_messages=_request_messages(params),
+            request_facts=_request_facts(params),
+        )
+        return completion
 
     async def messages(self, request: Request, body: dict = Body()):
         """Default Anthropic Messages <-> Responses mapping shared by every Gym model server.
@@ -206,6 +270,7 @@ async def messages(self, request: Request, body: dict = Body()):
         """
         params = _ANTHROPIC_CONVERTER.anthropic_request_to_responses(body)
         response = await self._invoke_responses(request, params)
+        # Capture here: the Anthropic response returned below has already dropped token ids.
         model_name = body.get("model") or response.model
         anthropic_response = _ANTHROPIC_CONVERTER.responses_to_anthropic_response(response, model=model_name)
         if body.get("stream"):
@@ -222,8 +287,18 @@ async def _invoke_responses(
         # `body`. Dispatch on whichever this server declares so the default messages() works for
         # all of them.
         if "request" in inspect.signature(self.responses).parameters:
-            return await self.responses(request=request, body=params)
-        return await self.responses(body=params)
+            response = await self.responses(request=request, body=params)
+        else:
+            response = await self.responses(body=params)
+        # Capture here rather than at the route: the streaming dispatch returns a StreamingResponse
+        # and the Anthropic mapping drops the token fields, so this is the last point where the
+        # assembled response still carries them, for every dialect.
+        await capture_tokens(
+            response,
+            request_messages=_request_messages(params),
+            request_facts=_request_facts(params),
+        )
+        return response
 
 
 def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming:
@@ -1003,10 +1078,19 @@ class _CaptureMiddleware:
     prefix and forwards only.
     """
 
-    def __init__(self, app: Any, *, store: Optional[CaptureStore], model_server_name: Optional[str]) -> None:
+    def __init__(
+        self,
+        app: Any,
+        *,
+        store: Optional[CaptureStore],
+        model_server_name: Optional[str],
+        token_store: Any = None,
+    ) -> None:
         self._app = app
         self._store = store
         self._model_server_name = model_server_name
+        # When set, correlated+observed calls also record training tokens via a per-request sink.
+        self._token_store = token_store
 
     async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
         if scope.get("type") != "http":
@@ -1021,24 +1105,36 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None
             path = prefix_match.group("rest")
             scope = {**scope, "path": path, "raw_path": path.encode("utf-8")}
 
-        # Capture disabled: the prefix is already stripped (routing preserved), so just forward.
-        if self._store is None:
-            await self._app(scope, receive, send)
-            return
+        dialect = _OBSERVED_PATHS.get(path)
 
-        # Only explicitly correlated model calls are captured. An unprefixed call is forwarded
-        # unchanged rather than being mixed with unrelated calls under a shared fallback key.
-        if rollout_from_path is None:
+        # Nothing to capture: neither store is active, the call isn't correlated to a rollout, or the
+        # path isn't an observed model endpoint. The prefix is already stripped, so just forward.
+        # An unprefixed call is forwarded rather than mixed with unrelated calls under a shared key.
+        if (self._store is None and self._token_store is None) or rollout_from_path is None or dialect is None:
             await self._app(scope, receive, send)
             return
 
-        dialect = _OBSERVED_PATHS.get(path)
-        if dialect is None:
-            await self._app(scope, receive, send)  # not observed (or a stripped non-/v1 path)
-            return
-
         rollout_id = rollout_from_path
         model_call_id = uuid4().hex
+
+        # Hand the model server a per-request token sink keyed to this call. It records token ids
+        # from its complete response (the middleware can't -- token ids are dropped on the SSE wire).
+        sink_token = None
+        if self._token_store is not None:
+            sink_token = set_token_sink(
+                CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, store=self._token_store)
+            )
+
+        # Training-token capture only: no evaluation record, so skip the response buffering entirely
+        # and just forward with the sink live.
+        if self._store is None:
+            try:
+                await self._app(scope, receive, send)
+            finally:
+                if sink_token is not None:
+                    reset_token_sink(sink_token)
+            return
+
         request_body = bytearray()
 
         async def _receive() -> dict[str, Any]:
@@ -1120,6 +1216,10 @@ async def _flush_deferred_response() -> None:
             finally:
                 await _flush_deferred_response()
             raise
+        finally:
+            # The sink is only needed while the model server produces the response.
+            if sink_token is not None:
+                reset_token_sink(sink_token)
 
         completed_at = time.time()
         latency_ms = (time.perf_counter() - start) * 1000.0
@@ -1187,22 +1287,39 @@ def _parse_and_record() -> None:
 
 
 def install_model_call_capture(
-    app: Any, config: ModelCallCaptureConfig, *, model_server_name: Optional[str] = None
+    app: Any,
+    config: ModelCallCaptureConfig,
+    *,
+    model_server_name: Optional[str] = None,
+    global_config_dict: Any = None,
 ) -> None:
     """Install model-call capture middleware.
 
     Always installed so the ``/ng-rollout/<id>`` correlation prefix is stripped before routing
     regardless of whether capture is enabled (otherwise a default ``gym eval`` would 404 on every
-    prefixed model call). When capture is enabled the middleware additionally records each observed
-    call's request + response into a rollout-keyed CaptureStore while forwarding bytes downstream
-    unchanged (non-terminal SSE chunks are forwarded as they arrive; the terminal event follows the
-    durable capture write).
+    prefixed model call). When evaluation capture is enabled the middleware additionally records each
+    observed call's request + response into a rollout-keyed CaptureStore while forwarding bytes
+    downstream unchanged (non-terminal SSE chunks are forwarded as they arrive; the terminal event
+    follows the durable capture write).
+
+    Training-token capture is a separate, independently-gated concern that reuses the same
+    correlation point: when enabled, the middleware hands the model server a per-request token sink
+    (keyed by the same rollout id and model_call_id) and the server records token ids from its
+    complete response. The read route is registered only when that capture is enabled.
     """
+    token_store = make_token_store(global_config_dict) if global_config_dict is not None else None
     app.add_middleware(
         _CaptureMiddleware,
         store=make_capture_store(config),
         model_server_name=model_server_name,
+        token_store=token_store,
     )
+    if token_store is not None:
+        install_token_capture_routes(
+            app,
+            token_store,
+            read_token=TokenIdCaptureConfig.model_validate(global_config_dict or {}).token_id_capture_read_token,
+        )
 
 
 # --- Run-level capture helpers (rollout-collection side) ---

nemo_gym/global_config.py

@@ -89,6 +89,9 @@
 QUERY_KEY_NAME = "query"
 OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled"
 MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir"
+TOKEN_ID_CAPTURE_ENABLED_KEY_NAME = "token_id_capture_enabled"
+# Per-agent opt-in (on an agent's config block) for participating in training token capture.
+TOKEN_ID_CAPTURE_KEY_NAME = "token_id_capture"
 COMPONENT_NAME_KEY_NAME = "component_name"
 NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [
     CONFIG_PATHS_KEY_NAME,