feat(token-id-capture): capture training tokens from external harnesses

nemo_gym/base_responses_api_agent.py

@@ -26,8 +26,12 @@
     BaseRunRequest,
     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.config_types import ROLLOUT_PATH_PREFIX, TOKEN_CAPTURE_PATH_SEGMENT
+from nemo_gym.global_config import (
+    OBSERVABILITY_ENABLED_KEY_NAME,
+    TOKEN_ID_CAPTURE_BLOCK,
+    get_first_server_config_dict,
+)
 from nemo_gym.openai_utils import (
     NeMoGymResponse,
     NeMoGymResponseCreateParamsNonStreaming,
@@ -46,6 +50,12 @@
 class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig):
     skip_verification: bool = False
     skip_verification_reward: float = 0.0
+    # Whether this agent's rollouts participate in training token capture.
+    # Native agents already receive token ids inline and normally leave this disabled.
+    # Opaque external harnesses enable it because their returned output has no token ids.
+    # The run-level ``token_id_capture.enabled`` setting gates the capture infrastructure.
+    # The run-level ``token_id_capture.all_agents`` setting overrides this agent-level choice.
+    token_id_capture: bool = False
 
 
 class BaseResponsesAPIAgent(BaseServer):
@@ -61,10 +71,11 @@ def setup_webserver(self) -> FastAPI:
         self.setup_session_middleware(app)
 
         app.post("/v1/responses")(self.responses)
-        # Prefixed twin of /v1/responses: a self-call made with url_path_for_run() lands here, and
-        # responses() recovers the rollout id from the path (see url_path_for_request) to correlate
-        # its model calls. Same handler, so unprefixed calls are unaffected.
+        # A self-call made with ``url_path_for_run`` lands on a prefixed twin.
+        # ``responses`` recovers the rollout id from the path.
+        # The same handler serves prefixed and unprefixed calls.
         app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses")(self.responses)
+        app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/{TOKEN_CAPTURE_PATH_SEGMENT}/v1/responses")(self.responses)
 
         run = self.run
 
@@ -81,59 +92,87 @@ async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyRespo
 
         return app
 
+    def _capture_correlation_enabled(self) -> bool:
+        """Return whether this agent needs rollout correlation.
+
+        Evaluation uses ``/ng-rollout/<id>/...`` for every agent.
+        Training capture uses ``/ng-rollout/<id>/training-token-capture/...``.
+        Training capture requires ``token_id_capture.enabled``.
+        It also requires the static agent flag or run-level ``all_agents``.
+        Missing global configuration disables correlation.
+        """
+        return self._model_call_capture_enabled() or self._token_id_capture_enabled()
+
     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.
+        """Whether evaluation model-call observability is enabled."""
         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))
 
+    def _token_id_capture_enabled(self) -> bool:
+        """Whether this agent explicitly opted into training-token capture."""
+        global_config = getattr(self.server_client, "global_config_dict", None)
+        if not isinstance(global_config, Mapping):
+            return False
+        block = global_config.get(TOKEN_ID_CAPTURE_BLOCK) or {}
+        if not isinstance(block, Mapping) or not block.get("enabled", False):
+            return False
+        return bool(block.get("all_agents", False)) or bool(
+            getattr(getattr(self, "config", None), "token_id_capture", False)
+        )
+
     def rollout_id_from_run(self, body: Any) -> Optional[str]:
-        """Per-rollout capture id for a run-request (its task/rollout indices).
+        """Return the capture id for a run request.
 
-        None when model-call capture (observability) is disabled or the body carries no indices,
-        so callers apply no correlation prefix in either case.
+        Return ``None`` when capture is disabled.
+        Return ``None`` when the body has no usable identity.
         """
-        if not self._model_call_capture_enabled():
+        if not self._capture_correlation_enabled():
             return None
         return maybe_rollout_id_from_run_body(body)
 
     def url_path_for_run(self, url_path: str, body: Any) -> str:
-        """A downstream url_path with the per-rollout capture-correlation prefix applied.
+        """Apply this run's capture path to a downstream URL path.
 
-        Returns ``/ng-rollout/<id><url_path>`` when observability is enabled and the run body
-        carries task/rollout indices; otherwise ``url_path`` unchanged. Use for calls made while
-        handling ``/run`` — both direct model-server calls and self-calls to ``/v1/responses``
-        (the prefixed self-call route carries the id into ``responses()``).
+        Evaluation uses ``/ng-rollout/<id>/...``.
+        Training capture uses ``/ng-rollout/<id>/training-token-capture/...``.
+        Calls without a rollout id remain unchanged.
         """
-        return f"{rollout_path_prefix(self.rollout_id_from_run(body))}{url_path}"
+        return (
+            f"{rollout_path_prefix(self.rollout_id_from_run(body), token_capture=self._token_id_capture_enabled())}"
+            f"{url_path}"
+        )
 
     def base_url_for_run(self, base_url: str, body: Any) -> str:
-        """A model-server base URL with the per-rollout capture-correlation prefix applied.
+        """Apply this run's capture path to a model-server root URL.
 
-        ``base_url_for_run`` is the base-URL counterpart of ``url_path_for_run`` for SDK-style
-        harnesses that configure a client once instead of prefixing each call: same gating, applied
-        to a server root URL (append the API-version suffix afterwards).
+        Append the API-version suffix after this method returns.
         """
-        return apply_rollout_prefix(base_url, self.rollout_id_from_run(body))
+        return apply_rollout_prefix(
+            base_url,
+            self.rollout_id_from_run(body),
+            token_capture=self._token_id_capture_enabled(),
+        )
 
     def url_path_for_request(self, url_path: str, request: Optional[Request]) -> str:
-        """Carry an inbound ``/ng-rollout/<id>`` self-call prefix onto a downstream url_path.
+        """Carry an inbound capture path onto a downstream URL path.
 
-        Agents whose model calls happen inside ``responses()`` receive the correlation id as the
-        ``rollout_id`` path parameter of the prefixed self-call route; this re-applies it to the
-        outgoing model call. Unprefixed requests pass through unchanged.
+        Prefixed self-calls expose the rollout id as a path parameter.
+        Training-capture requests preserve their dedicated path segment.
+        Unprefixed requests remain unchanged.
         """
         path_params = getattr(request, "path_params", None)
         rollout_id = path_params.get("rollout_id") if isinstance(path_params, Mapping) else None
-        return f"{rollout_path_prefix(rollout_id)}{url_path}"
+        request_path = getattr(getattr(request, "url", None), "path", "")
+        token_capture = f"/{TOKEN_CAPTURE_PATH_SEGMENT}/" in request_path
+        return f"{rollout_path_prefix(rollout_id, token_capture=token_capture)}{url_path}"
 
     def resolve_model_base_url(self, model_server_name: str, rollout_id: Optional[str] = None) -> str:
         """Resolve a model-server URL with an optional rollout prefix."""
         server_config = get_first_server_config_dict(self.server_client.global_config_dict, model_server_name)
         base_url = self.server_client._build_server_base_url(server_config)
-        return f"{apply_rollout_prefix(base_url, rollout_id)}/v1"
+        return f"{apply_rollout_prefix(base_url, rollout_id, token_capture=self._token_id_capture_enabled())}/v1"
 
     # TODO: right now there is no validation on the TypedDict NeMoGymResponseCreateParamsNonStreaming
     # We should explicitly add validation at this server level or we should explicitly not validate so that there is flexibility in this API.