feat(token-id-capture): require a bearer token on the token read route

The read route serves a rollout's raw training tokens on the same server the
harness calls to generate. Inside a trusted cluster that is acceptable. Once the
harness runs in a sandbox whose only egress is that server it is not: the
harness could read its own training data, or another rollout's.

token_id_capture_read_token requires a bearer token on the route. Comparison is
constant-time, since that token is the only thing between an untrusted harness
and every rollout's training data.

Left unset the route stays open and logs a warning at startup, so an existing
deployment keeps working and the gap is visible rather than silent. Set it before
running a harness in a sandbox.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

nemo_gym/base_responses_api_model.py

@@ -69,6 +69,7 @@
 )
 from nemo_gym.token_id_capture import (
     CaptureContext,
+    TokenIdCaptureConfig,
     capture_tokens,
     reset_token_sink,
     set_token_sink,
@@ -1290,7 +1291,11 @@ def install_model_call_capture(
         token_store=token_store,
     )
     if token_store is not None:
-        install_token_capture_routes(app, token_store)
+        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/token_id_capture/config.py

@@ -56,6 +56,13 @@ class TokenIdCaptureConfig(BaseModel):
     # Shared fallback directory (also used by evaluation capture).
     model_call_capture_dir: Path | None = None
 
+    # Bearer token required by the token read route. The route serves a rollout's raw training
+    # tokens on the same app the harness calls to generate, which is fine inside a trusted cluster
+    # and not fine once the harness runs in a sandbox whose only egress is this server -- it could
+    # read its own training data, or another rollout's. Set this before the sandbox work; when unset
+    # the route stays open and logs a warning once.
+    token_id_capture_read_token: str | None = None
+
     @model_validator(mode="after")
     def _validate(self) -> "TokenIdCaptureConfig":
         if not self.token_id_capture_enabled:

nemo_gym/token_id_capture/routes.py

@@ -25,9 +25,10 @@
 
 import asyncio
 import logging
+from hmac import compare_digest
 from typing import Any
 
-from fastapi import APIRouter, HTTPException, Response
+from fastapi import APIRouter, Header, HTTPException, Response
 
 from nemo_gym.token_id_capture.config import TokenIdCaptureConfig
 from nemo_gym.token_id_capture.store import TokenCaptureStore
@@ -44,19 +45,33 @@ def make_token_store(global_config_dict: Any) -> TokenCaptureStore | None:
     return TokenCaptureStore(config.resolved_dir())
 
 
-def install_token_capture_routes(app: Any, store: TokenCaptureStore) -> None:
+def install_token_capture_routes(app: Any, store: TokenCaptureStore, read_token: str | None = None) -> None:
     """Register the token read route.
 
     The route serves a rollout's raw training tokens on the same app the harness
     calls to generate. That is acceptable inside a trusted cluster; it is not
     acceptable once the harness runs in a sandbox whose only egress is this
     server, because it could read its own training data -- or another rollout's.
+    ``read_token`` requires a bearer token; when it is unset the route stays open
+    and warns once, so existing deployments keep working and the gap is visible.
     """
+    if not read_token:
+        logger.warning(
+            "The token-capture read route is unauthenticated. Anything that can reach this model "
+            "server can read captured training tokens for any rollout. Set token_id_capture_read_token "
+            "before running an agent harness in a sandbox whose only egress is this server."
+        )
 
     router = APIRouter()
 
     @router.get("/ng-capture/tokens/{rollout_id}")
-    async def get_tokens(rollout_id: str) -> Response:
+    async def get_tokens(rollout_id: str, authorization: str | None = Header(default=None)) -> Response:
+        if read_token:
+            expected = f"Bearer {read_token}"
+            # Constant-time compare: the token is the only thing standing between an untrusted
+            # harness and every rollout's training data.
+            if authorization is None or not compare_digest(authorization, expected):
+                raise HTTPException(status_code=401, detail="token capture read route requires a bearer token")
         try:
             entries = await asyncio.to_thread(store.read_entries, rollout_id)
         except ValueError as exc:

tests/unit_tests/test_token_id_capture.py

@@ -656,3 +656,28 @@ def test_lineage_resolves_a_tool_using_turn_echoed_in_anthropic_shape():
     parent = lineage.resolve(next_request)
     assert parent is not None and parent.call_id == "call-1"
     assert parent.cum_tokens == [1, 2, 3]
+
+
+def test_read_route_requires_a_token_when_one_is_configured(tmp_path):
+    """The route serves raw training tokens on the same app the harness calls to
+    generate. Inside a trusted cluster that is fine; once the harness runs in a
+    sandbox whose only egress is this server it could read its own training data,
+    or another rollout's."""
+    config = dict(_both_enabled(tmp_path))
+    config["token_id_capture_read_token"] = "s3cret"  # pragma: allowlist secret
+    client = TestClient(_server(config).setup_webserver())
+    client.post("/ng-rollout/auth0-roll0/v1/responses", json={"input": "hi"})
+
+    assert client.get("/ng-capture/tokens/auth0-roll0").status_code == 401
+    assert client.get("/ng-capture/tokens/auth0-roll0", headers={"Authorization": "Bearer wrong"}).status_code == 401
+
+    ok = client.get("/ng-capture/tokens/auth0-roll0", headers={"Authorization": "Bearer s3cret"})
+    assert ok.status_code == 200
+    assert TokenEntry.model_validate_json(ok.text.splitlines()[0]).generation_token_ids == GTOKS
+
+
+def test_read_route_stays_open_when_no_token_is_configured(tmp_path):
+    """Existing deployments keep working; the gap is logged rather than enforced."""
+    client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver())
+    client.post("/ng-rollout/auth1-roll0/v1/responses", json={"input": "hi"})
+    assert client.get("/ng-capture/tokens/auth1-roll0").status_code == 200