feat(token-id-capture): resolve each call's parent at request time

fern/versions/latest/pages/training-tutorials/external-agent-harnesses.mdx

@@ -70,7 +70,11 @@ This failure is silent. Check `n_calls` on the first run before relying on the r
 
 ## What you get back
 
-Each rollout's model calls are stitched into a single Responses payload with contiguous `output` items. Each item's `prompt_token_ids` contains the running sequence. Its `generation_token_ids` contains the tokens sampled by the policy at that step. Gym replaces the rollout's `response.output` with these items, so a trainer reads `response.output` the same way for native agents and external harnesses.
+Each captured call records one parent-resolution result. `ROOT` means that the request contains no previous model-authored output. `RESOLVED` means that exactly one committed call matches and its request context was verified. `UNRESOLVED` means that the request appears to continue prior model output, but Gym cannot prove which committed call produced it.
+
+Each rollout's resolved model calls are stitched into a single Responses payload with contiguous `output` items. Each item's `prompt_token_ids` contains the running sequence. Its `generation_token_ids` contains the tokens sampled by the policy at that step. Gym replaces the rollout's `response.output` with these items, so a trainer reads `response.output` the same way for native agents and external harnesses.
+
+An unresolved call begins a separate incomplete fragment. Gym never crosses that boundary with token-prefix inference. The call and its descendants remain available for diagnostics or a consumer that explicitly supports partial trajectories, but single-response delivery sets `mask_sample: true`. Prefix inference is reserved for records written before parent-resolution metadata existed.
 
 The loss mask follows from this structure instead of being sent separately. Prompt positions provide context. Generation positions are trainable.
 
@@ -87,7 +91,8 @@ Gym attaches a metrics dictionary to each rollout under `_ng_token_capture`. Agg
 | `delivered_fraction` | 1.0 | Sampled tokens were captured but not delivered. |
 | `quarantined_calls` | 0 | Two calls could not be told apart, so neither was used. |
 | `empty_generation_calls` | 0 | The output budget or a content filter is truncating generations. |
-| `mask_sample` | absent | The rollout lost a call and must not be trained on. |
+| `unresolved_parent_calls` | 0 | A request appeared to continue prior model output, but no exact committed parent was proven. |
+| `mask_sample` | absent | The rollout lost a call, split at an unresolved parent, or otherwise cannot be trained safely. |
 
 Pay particular attention to `n_calls`. A value of exactly 1 means the agentic path was never exercised, even though every other key can look correct.
 
@@ -99,7 +104,7 @@ Gym defines the record shape and builds each record. The training framework cont
 
 ### The interfaces
 
-Gym describes the sink and source as structural protocols. Framework adapters do not inherit from these definitions or import them at runtime. They implement the same method signatures, and Gym consumes the resulting objects by that method shape. The definitions below are the reference contract.
+Gym describes the sink, source, and lineage resolver as structural protocols. Framework adapters do not inherit from these definitions or import them at runtime. They implement the same method signatures, and Gym consumes the resulting objects by that method shape. The definitions below are the reference contract.
 
 ```python
 class TokenSink(Protocol):
@@ -111,9 +116,14 @@ class TokenSource(Protocol):
     async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: ...
     async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: ...
     async def close(self) -> None: ...
+
+class LineageStore(Protocol):
+    async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: ...
+    def is_process_shared(self) -> bool: ...
+    async def close(self) -> None: ...
 ```
 
-`put` must make the record durable before it returns. A reader that runs after the rollout must see every acknowledged record. This guarantee allows the consumer to freeze one complete view after the harness finishes.
+`TokenSink` remains the writer-side interface. `put` is the only publication boundary: it must make the complete token record and its compact continuation lookup metadata durable before returning. `LineageStore` is the worker-side read client over those committed entries; it does not publish a second record. A later request may resolve an entry only after the corresponding `put` has returned.
 
 `mark_incomplete` is the durable signal that a rollout lost a call. The model call still succeeds. A sink that drops this signal makes an incomplete rollout look complete.
 
@@ -123,9 +133,11 @@ class TokenSource(Protocol):
 
 `close` releases client resources. Gym closes clients it constructs, but it does not close a caller-installed source.
 
+`LineageResolution` has three semantic outcomes: `ROOT`, `RESOLVED`, and `UNRESOLVED`. Diagnostic reasons such as an ambiguous lookup or unavailable backend may accompany `UNRESOLVED`, but they do not change reconstruction behavior.
+
 ### Connecting a framework-owned transport
 
-The training framework owns the transport and its client configuration. Gym owns the model-server processes, so each server worker constructs a framework-provided `TokenSink` proxy from the configured class path. The configured class is a worker factory descriptor, not a shared Python object or a transport implementation owned by Gym.
+The training framework owns the transport and its client configuration. Gym owns the model-server processes, so each server worker constructs framework-provided `TokenSink` and `LineageStore` proxies from configured class paths. The configured classes are worker factory descriptors, not shared Python objects or transport implementations owned by Gym.
 
 ```yaml
 env:
@@ -136,10 +148,16 @@ env:
       sink_kwargs:
         endpoint: ${oc.env:MY_DATAPLANE_URL}
         shard: ${oc.select:cluster_shard,0}
+      lineage_store: my_pkg.sinks:MyDataPlaneLineageStore
+      lineage_store_kwargs:
+        endpoint: ${oc.env:MY_DATAPLANE_URL}
+        shard: ${oc.select:cluster_shard,0}
       rebuild_response: false
 ```
 
-Gym passes `sink_kwargs` to the constructor, so a sink can receive the required endpoint, client, or credentials instead of reading ambient state. Use `${oc.env:VAR}` for secrets instead of writing them into the config. Unsupported constructor arguments cause a startup error. A sink that does not implement `mark_incomplete` also causes a startup error because it could otherwise make a rollout with a missing call look complete.
+Gym passes each kwargs block to its constructor, so the clients can receive the required endpoint, shard, or credentials instead of reading ambient state. Use `${oc.env:VAR}` for secrets instead of writing them into the config. Unsupported constructor arguments cause a startup error. A sink that does not implement `mark_incomplete` also causes a startup error because it could otherwise make a rollout with a missing call look complete.
+
+The sink and lineage resolver must use the same backend namespace. The backend transaction behind `TokenSink.put` stores the `TokenEntry` and updates the continuation-key index together. Pointing the two clients at unrelated services does not satisfy the protocol even if both methods return successfully.
 
 `sink` replaces the file store, so a `dir` configured alongside it is not used. This condition produces a warning instead of an error because no data is lost. No capture files appear on disk.
 
@@ -163,9 +181,11 @@ Configure the sink instead of installing it from a launcher script. Programmatic
 
 `install_token_sink` sets a process global. A model server with `num_workers > 1` launches uvicorn with an app string and `workers=N`. Uvicorn spawns workers that re-import the app module instead of inheriting the launcher's memory. A sink installed by the parent process therefore does not exist in any worker. Capture then falls back to the file store. If no `dir` is set, the worker has no local destination.
 
-Gym constructs the configured sink inside each worker at app startup, which avoids this process-boundary problem.
+Gym constructs configured sink and resolver clients inside each worker at app startup. Every client must connect to the same process-shared backend. `LineageStore.is_process_shared()` must return `true`; startup rejects a process-local resolver when `num_workers > 1`.
 </Warning>
 
+With Gym's local file backend, all workers append and resolve from the same token JSONL under one per-rollout lock. The resolver incrementally indexes only newly appended entries in each worker. Freeze and retirement use the same lock. Retirement keeps the lock file and a tombstone, so an old worker cannot recreate a retired capture.
+
 ### Reading records back
 
 `TokenCaptureStore` implements both protocols and is the default. A reader beside the store uses the store as its `TokenSource`. This arrangement applies to `gym eval run` and to a trainer colocated with the model server. The store directory should therefore be node-local.

nemo_gym/base_responses_api_model.py

@@ -71,15 +71,18 @@
 from nemo_gym.token_id_capture import (
     CaptureContext,
     capture_tokens,
+    installed_lineage_store,
     installed_token_sink,
     register_call_intent,
     reset_token_sink,
+    resolve_parent,
     set_token_sink,
 )
 
 # The store factory needs Gym's server stack.
 # The leaf package does not re-export it.
 from nemo_gym.token_id_capture.config import token_id_capture_config
+from nemo_gym.token_id_capture.lineage import FileLineageStore
 from nemo_gym.token_id_capture.store import make_token_store
 
 
@@ -90,6 +93,65 @@
 _ANTHROPIC_CONVERTER = AnthropicConverter()
 
 
+def _request_messages(body: Any) -> list[dict]:
+    """Return the conversation carried by any supported dialect.
+
+    Lineage uses model-authored turns to identify the parent call.
+    Chat and Anthropic use ``messages``.
+    Responses uses ``input``.
+    The request envelope is prepended as a pseudo-turn because it also shapes the prompt.
+    """
+    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):
+        turns = [m if isinstance(m, dict) else m.model_dump() for m in messages if m is not None]
+    else:
+        items = getter("input", None)
+        turns = (
+            [i if isinstance(i, dict) else i.model_dump() for i in items if i is not None]
+            if isinstance(items, list)
+            else []
+        )
+    return _request_envelope(getter) + turns
+
+
+# This role cannot collide with a dialect role.
+# It is not assistant-authored and does not affect the lookup fingerprint.
+_ENVELOPE_ROLE = "_ng_request_envelope"
+
+
+def _request_envelope(getter: Any) -> list[dict]:
+    """Return prompt-shaping request fields that are not turns.
+
+    Instructions and tools can be siblings of the message list.
+    The chat template renders both into the prompt.
+    Including them prevents prefix reuse across different request envelopes.
+    """
+    instructions = getter("instructions", None)
+    tools = getter("tools", None)
+    if not instructions and not tools:
+        return []
+    envelope = {"instructions": _plain(instructions), "tools": _plain(tools)}
+    return [{"role": _ENVELOPE_ROLE, "content": json.dumps(envelope, sort_keys=True, default=str)}]
+
+
+def _plain(value: Any) -> Any:
+    """Reduce a request field to plain data so equal schemas serialize equally.
+
+    Handlers can expose tools as dictionaries or Pydantic models.
+    Normalization keeps an unchanged schema stable across handlers.
+    """
+    if hasattr(value, "model_dump"):
+        return value.model_dump()
+    if isinstance(value, (list, tuple)):
+        return [_plain(item) for item in value]
+    if isinstance(value, dict):
+        return {key: _plain(item) for key, item in value.items()}
+    return value
+
+
 class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig):
     pass
 
@@ -109,6 +171,7 @@ def setup_webserver(self) -> FastAPI:
             capture_config,
             model_server_name=self.config.name,
             global_config_dict=self.server_client.global_config_dict,
+            num_workers=self.config.num_workers,
         )
 
         app.post("/v1/chat/completions")(self.chat_completions_dispatch)
@@ -210,12 +273,19 @@ async def _invoke_chat_completions(
         # chat_completions() signatures vary across servers: some take a leading `request`, some
         # only `body`. Dispatch on whichever this server declares so the shared dispatch works for
         # all of them.
+        # Resolve the parent from the received request before dispatch.
+        # Exact prefix supply and capture share this decision.
+        request_messages = _request_messages(params)
+        await resolve_parent(request_messages)
         await register_call_intent()
         if "request" in inspect.signature(self.chat_completions).parameters:
             completion = await self.chat_completions(request=request, body=params)
         else:
             completion = await self.chat_completions(body=params)
-        await capture_tokens(completion)
+        await capture_tokens(
+            completion,
+            request_messages=request_messages,
+        )
         return completion
 
     async def messages(self, request: Request, body: dict = Body()):
@@ -244,6 +314,10 @@ async def _invoke_responses(
         # responses() signatures vary across servers: some take a leading `request`, some only
         # `body`. Dispatch on whichever this server declares so the default messages() works for
         # all of them.
+        # Resolve the parent from the received request before dispatch.
+        # Exact prefix supply and capture share this decision.
+        request_messages = _request_messages(params)
+        await resolve_parent(request_messages)
         await register_call_intent()
         if "request" in inspect.signature(self.responses).parameters:
             response = await self.responses(request=request, body=params)
@@ -252,7 +326,10 @@ async def _invoke_responses(
         # Capture before streaming dispatch wraps the response.
         # Anthropic mapping drops the token fields.
         # The assembled response still carries them here for every dialect.
-        await capture_tokens(response)
+        await capture_tokens(
+            response,
+            request_messages=request_messages,
+        )
         return response
 
 
@@ -1068,6 +1145,7 @@ def __init__(
         model_server_name: str | None,
         token_store: Any = None,
         configured_sink: Any = None,
+        lineage_store: Any = None,
         token_capture_enabled: bool = False,
     ) -> None:
         self._app = app
@@ -1077,6 +1155,7 @@ def __init__(
         self._token_store = token_store
         # Built from token_id_capture.sink, once, in this process.
         self._configured_sink = configured_sink
+        self._lineage_store = lineage_store
         # Capture may have no destination in this process.
         # A framework may stage records from its inference worker.
         # This process still resolves the capture identity.
@@ -1138,7 +1217,12 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None
         sink_token = None
         if capture_wanted:
             sink_token = set_token_sink(
-                CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, token_sink=token_sink)
+                CaptureContext(
+                    rollout_id=rollout_id,
+                    model_call_id=model_call_id,
+                    token_sink=token_sink,
+                    lineage_store=self._lineage_store,
+                )
             )
 
         # Training-only capture has no evaluation record.
@@ -1308,6 +1392,7 @@ def install_model_call_capture(
     *,
     model_server_name: str | None = None,
     global_config_dict: Any = None,
+    num_workers: int | None = None,
 ) -> None:
     """Install model-call capture middleware.
 
@@ -1319,22 +1404,36 @@ def install_model_call_capture(
     That path provides a request-scoped token sink.
     The model server records token ids from its complete response.
     Consumers access records through ``TokenSource.freeze``.
-    There is no HTTP token reader.
     """
+    capture_settings = token_id_capture_config(global_config_dict) if global_config_dict is not None else None
     token_store = make_token_store(global_config_dict) if global_config_dict is not None else None
     # Build this sink at app startup.
     # Each uvicorn worker constructs its own sink.
     # Spawned workers do not inherit a launcher-installed sink.
-    configured_sink = (
-        token_id_capture_config(global_config_dict).build_sink() if global_config_dict is not None else None
-    )
-    owned_sinks = [sink for sink in (configured_sink, token_store) if sink is not None]
+    configured_sink = capture_settings.build_sink() if capture_settings is not None else None
+    configured_lineage = capture_settings.build_lineage_store() if capture_settings is not None else None
+    lineage_store = configured_lineage or installed_lineage_store()
+    default_lineage = None
+    if lineage_store is None and token_store is not None:
+        default_lineage = FileLineageStore(token_store.root)
+        lineage_store = default_lineage
+    if capture_settings is not None and capture_settings.enabled and (num_workers or 1) > 1:
+        if lineage_store is None or not lineage_store.is_process_shared():
+            raise ValueError(
+                "token_id_capture with num_workers > 1 requires a process-shared lineage resolver "
+                "over the same backend used by the token sink"
+            )
+    owned_endpoints = [
+        endpoint
+        for endpoint in (configured_sink, token_store, configured_lineage, default_lineage)
+        if endpoint is not None
+    ]
 
-    async def _close_token_sinks() -> None:
-        for sink in owned_sinks:
-            await sink.close()
+    async def _close_capture_endpoints() -> None:
+        for endpoint in owned_endpoints:
+            await endpoint.close()
 
-    if owned_sinks:
+    if owned_endpoints:
         original_lifespan = app.router.lifespan_context
 
         @asynccontextmanager
@@ -1343,7 +1442,7 @@ async def _capture_lifespan(application):
                 async with original_lifespan(application) as state:
                     yield state
             finally:
-                await _close_token_sinks()
+                await _close_capture_endpoints()
 
         app.router.lifespan_context = _capture_lifespan
     app.add_middleware(
@@ -1352,9 +1451,8 @@ async def _capture_lifespan(application):
         model_server_name=model_server_name,
         token_store=token_store,
         configured_sink=configured_sink,
-        token_capture_enabled=(
-            token_id_capture_config(global_config_dict).enabled if global_config_dict is not None else False
-        ),
+        lineage_store=lineage_store,
+        token_capture_enabled=capture_settings.enabled if capture_settings is not None else False,
     )