diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py
index d69ce85f58..80712059e3 100644
--- a/nemo_gym/base_responses_api_model.py
+++ b/nemo_gym/base_responses_api_model.py
@@ -194,6 +194,9 @@ class BaseResponsesAPIModel(BaseServer):
class SimpleResponsesAPIModel(BaseResponsesAPIModel, SimpleServer):
+ async def _finalize_served_response(self, response: Any) -> None:
+ """Finalize capture from the response representation returned to the client."""
+
def setup_webserver(self) -> FastAPI:
app = FastAPI()
@@ -255,7 +258,9 @@ async def responses_dispatch(self, request: Request, body: dict = Body()):
_reject_external_capture_streaming(body)
if not body.get("stream"):
params = _validate_responses_params(body)
- return _orjson_dispatch_response(await self._invoke_responses(request, params))
+ response = await self._invoke_responses(request, params)
+ await self._finalize_served_response(response)
+ return _orjson_dispatch_response(response)
cleaned, ns_map = sanitize_streaming_responses_body(body)
try:
@@ -300,7 +305,9 @@ async def chat_completions_dispatch(self, request: Request, body: dict = Body())
_reject_external_capture_streaming(body)
if body.get("stream") is not True:
params = _validate_chat_params(body)
- return _orjson_dispatch_response(await self._invoke_chat_completions(request, params))
+ response = await self._invoke_chat_completions(request, params)
+ await self._finalize_served_response(response)
+ return _orjson_dispatch_response(response)
cleaned, include_usage = sanitize_streaming_chat_body(body)
params = _validate_chat_params(cleaned)
@@ -349,6 +356,7 @@ async def messages(self, request: Request, body: dict = Body()):
response = await self._invoke_responses(request, params)
model_name = body.get("model") or response.model
anthropic_response = _ANTHROPIC_CONVERTER.responses_to_anthropic_response(response, model=model_name)
+ await self._finalize_served_response(anthropic_response)
if body.get("stream"):
return StreamingResponse(
_ANTHROPIC_CONVERTER.anthropic_response_to_sse(anthropic_response),
diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py
index 47f0baf8aa..2efcfa41ce 100644
--- a/nemo_gym/token_id_capture/records.py
+++ b/nemo_gym/token_id_capture/records.py
@@ -264,11 +264,28 @@ def response_to_output_items(payload: dict) -> list[dict]:
Responses payloads already carry ``output``.
Chat payloads carry ``choices[*].message``.
+ Anthropic Messages payloads carry top-level assistant content.
Wrap each assistant message as a Responses ``message`` item.
"""
output = payload.get("output")
if isinstance(output, list) and output:
return [item for item in output if isinstance(item, dict)]
+ if payload.get("type") == "message" and payload.get("role") == "assistant":
+ reasoning_items: list[dict] = []
+ message_content: list[Any] = []
+ for block in payload.get("content") or []:
+ if not isinstance(block, dict) or block.get("type") not in {"thinking", "redacted_thinking"}:
+ message_content.append(block)
+ continue
+ reasoning_item: dict[str, Any] = {"type": "reasoning", "summary": []}
+ if block.get("type") == "thinking" and isinstance(block.get("thinking"), str):
+ reasoning_item["summary"] = [{"type": "summary_text", "text": block["thinking"]}]
+ elif block.get("type") == "redacted_thinking" and block.get("data") is not None:
+ reasoning_item["encrypted_content"] = block["data"]
+ reasoning_items.append(reasoning_item)
+ if message_content:
+ reasoning_items.append({"type": "message", "role": "assistant", "content": message_content})
+ return reasoning_items
items: list[dict] = []
for choice in payload.get("choices") or []:
message = (choice or {}).get("message") or {}
diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py
index 1e41b27a11..c71b6698e1 100644
--- a/nemo_gym/token_id_capture/sink.py
+++ b/nemo_gym/token_id_capture/sink.py
@@ -105,6 +105,9 @@ class CaptureContext:
# ``resolve_parent`` so the commit hook can publish the ledger row with
# the exact representation the next request will echo.
request_items: list[dict] | None = None
+ # The worker acknowledgement arrives on the internal Chat response.
+ # Finalization consumes it after conversion to the served API dialect.
+ external_commit_coords: dict[str, Any] | None = None
@property
def parent_call_id(self) -> str | None:
diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py
index 0c7e6813d5..ff1f607025 100644
--- a/responses_api_models/vllm_model/app.py
+++ b/responses_api_models/vllm_model/app.py
@@ -971,7 +971,7 @@ async def chat_completions(
)
if self._external_capture_enabled:
- await self._finalize_external_capture(chat_completion_dict)
+ self._prepare_external_capture(chat_completion_dict)
if self.config.return_token_id_information:
message_dict = choice_dict["message"]
@@ -1039,13 +1039,24 @@ async def chat_completions(
return NeMoGymChatCompletion.model_validate(chat_completion_dict)
+ def _prepare_external_capture(self, payload: Dict[str, Any]) -> None:
+ """Retain the worker acknowledgement until API conversion finishes."""
+ context = current_capture_context()
+ if context is None or not context.external_staging:
+ return
+ context.external_commit_coords = payload.pop(NG_COMMIT_COORDS_FIELD, None)
+ self._strip_capture_transport_fields(payload)
+
+ async def _finalize_served_response(self, response: Any) -> None:
+ """Publish external capture using the response representation served to the client."""
+ await self._finalize_external_capture(_jsonable(response))
+
async def _finalize_external_capture(self, payload: Dict[str, Any]) -> None:
"""Validate and record a response staged by the inference worker.
The worker returns commit coordinates only after ``StagingSink.stage`` succeeds.
This method validates those coordinates against the active call.
- It then records the call in the lineage store.
- Finally, it removes token data and commit coordinates from the served response.
+ It records fingerprints from the response representation served to the client.
"""
context = current_capture_context()
if context is None or not context.external_staging or context.lineage_store is None:
@@ -1053,11 +1064,10 @@ async def _finalize_external_capture(self, payload: Dict[str, Any]) -> None:
ledger = context.lineage_store
if not isinstance(ledger, CaptureLedger):
raise ValueError("external staging requires a CaptureLedger on the capture context")
- coords_payload = payload.pop(NG_COMMIT_COORDS_FIELD, None)
+ coords_payload = context.external_commit_coords
admission = context.capture_admission
if admission is None:
# UNRESOLVED — the ledger already carries this call's poison row.
- self._strip_capture_transport_fields(payload)
return
try:
if coords_payload is None:
@@ -1157,8 +1167,6 @@ async def _finalize_external_capture(self, payload: Dict[str, Any]) -> None:
context.rollout_id,
context.model_call_id,
)
- finally:
- self._strip_capture_transport_fields(payload)
@staticmethod
def _strip_capture_transport_fields(payload: Dict[str, Any]) -> None:
diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py
index 295dde02e9..26b83032a8 100644
--- a/responses_api_models/vllm_model/tests/test_app.py
+++ b/responses_api_models/vllm_model/tests/test_app.py
@@ -63,6 +63,9 @@
resolve_parent,
set_token_sink,
)
+from nemo_gym.token_id_capture.fingerprint import assistant_fingerprint
+from nemo_gym.token_id_capture.records import response_to_output_items
+from nemo_gym.token_id_capture.staging.records import CaptureAdmission
from responses_api_models.vllm_model.app import (
VLLMConverter,
VLLMModel,
@@ -5421,6 +5424,236 @@ def test_reasoning_stripped_history_still_supplies_the_real_tokens(
assert out["required_prefix_token_ids"] == real_tokens_including_reasoning
+ def test_responses_reasoning_echo_supplies_the_exact_captured_prefix(
+ self, monkeypatch: MonkeyPatch, tmp_path
+ ) -> None:
+ config = VLLMModelConfig(
+ host="0.0.0.0",
+ port=8081,
+ base_url="http://api.openai.com/v1",
+ api_key="dummy_key", # pragma: allowlist secret
+ model="dummy_model",
+ entrypoint="",
+ name="vllm_model",
+ return_token_id_information=True,
+ uses_reasoning_parser=True,
+ uses_interleaved_reasoning=False,
+ supply_prefix_token_ids=True,
+ )
+ capture_config = {
+ "token_id_capture": {
+ "enabled": True,
+ "dir": str(tmp_path),
+ "rebuild_response": False,
+ }
+ }
+ monkeypatch.setattr(nemo_gym.server_utils, "get_global_config_dict", MagicMock(return_value=capture_config))
+ model = VLLMModel(
+ config=config,
+ server_client=MagicMock(spec=ServerClient, global_config_dict=capture_config),
+ )
+ requests: list[dict[str, Any]] = []
+ first_cumulative_tokens = [10, 11, 20, 21, 22]
+
+ def completion(
+ response_id: str,
+ prompt_token_ids: list[int],
+ generation_token_ids: list[int],
+ content: str,
+ *,
+ reasoning: str | None = None,
+ ) -> dict[str, Any]:
+ message = {"role": "assistant", "content": content}
+ if reasoning is not None:
+ message["reasoning_content"] = reasoning
+ return {
+ "id": response_id,
+ "object": "chat.completion",
+ "created": 0,
+ "model": "dummy_model",
+ "prompt_token_ids": prompt_token_ids,
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "stop",
+ "token_ids": generation_token_ids,
+ "message": message,
+ "logprobs": {
+ "content": [
+ {
+ "token": f"token_id:{token_id}",
+ "logprob": -0.1,
+ "bytes": None,
+ "top_logprobs": [],
+ }
+ for token_id in generation_token_ids
+ ]
+ },
+ }
+ ],
+ }
+
+ async def mock_create_chat_completion(**kwargs):
+ requests.append(kwargs)
+ if len(requests) == 1:
+ return completion(
+ "chatcmpl-first",
+ [10, 11],
+ [20, 21, 22],
+ "answer",
+ reasoning="hidden reasoning",
+ )
+ assert kwargs["required_prefix_token_ids"] == first_cumulative_tokens
+ return completion("chatcmpl-second", first_cumulative_tokens + [30], [40], "done")
+
+ mock_client = MagicMock(spec=NeMoGymAsyncOpenAI)
+ mock_client.create_chat_completion = AsyncMock(side_effect=mock_create_chat_completion)
+ mock_client.create_tokenize = AsyncMock(
+ side_effect=AssertionError("inline generation prompt tokens must prove prefix supply")
+ )
+ model._clients = [mock_client]
+ client = TestClient(model.setup_webserver())
+ path = "/ng-rollout/reasoning-prefix/training-token-capture/v1/responses"
+ first_input = [{"role": "user", "content": "question"}]
+
+ first_response = client.post(path, json={"model": "dummy_model", "input": first_input})
+
+ assert first_response.status_code == 200
+ first_output = first_response.json()["output"]
+ assert [item["type"] for item in first_output] == ["reasoning", "message"]
+ assert first_output[0]["summary"][0]["text"] == "hidden reasoning"
+ assert first_output[1]["content"][0]["text"] == "answer"
+ echoed_output = [
+ {
+ key: value
+ for key, value in item.items()
+ if key
+ not in {
+ "prompt_token_ids",
+ "generation_token_ids",
+ "generation_log_probs",
+ "routed_experts",
+ }
+ }
+ for item in first_output
+ ]
+
+ second_response = client.post(
+ path,
+ json={
+ "model": "dummy_model",
+ "input": first_input + echoed_output + [{"role": "user", "content": "next"}],
+ },
+ )
+
+ assert second_response.status_code == 200, second_response.text
+ assert len(requests) == 2
+ assert "hidden reasoning" not in str(requests[1]["messages"])
+ assert requests[1]["required_prefix_token_ids"] == first_cumulative_tokens
+ first_entry, second_entry = TokenCaptureStore(tmp_path).read_entries("reasoning-prefix")
+ assert [item["type"] for item in first_entry.output_items] == ["reasoning", "message"]
+ assert first_entry.prompt_token_ids + first_entry.generation_token_ids == first_cumulative_tokens
+ assert second_entry.parent_call_id == first_entry.model_call_id
+ assert second_entry.prefix_requested is True
+ assert second_entry.prefix_supplied is True
+ assert mock_client.create_tokenize.await_count == 0
+
+ def test_external_capture_fingerprints_the_served_responses_shape(self, monkeypatch: MonkeyPatch) -> None:
+ """Fingerprint the Responses payload after reasoning is separated from the answer."""
+
+ server = self._server(monkeypatch, enabled=True)
+ ledger = InMemoryLineageStore()
+ digest = "1" * 64
+ extras_digest = "2" * 64
+ chain_hash = "3" * 64
+ cumulative_hash = "4" * 64
+ coords = {
+ "rollout_id": "reasoning-external",
+ "model_call_id": "call-1",
+ "prev_len": 0,
+ "delta_len": 5,
+ "cum_len": 5,
+ "weight_version": 0,
+ "digest": digest,
+ "extras_digest": extras_digest,
+ "staging_key": "stage-call-1",
+ "chain_hash": chain_hash,
+ "cumulative_hash": cumulative_hash,
+ }
+ context = CaptureContext(
+ rollout_id="reasoning-external",
+ model_call_id="call-1",
+ token_sink=None,
+ lineage_store=ledger,
+ external_staging=True,
+ capture_admission=CaptureAdmission(
+ rollout_id="reasoning-external",
+ model_call_id="call-1",
+ mode="text",
+ ),
+ request_items=[{"role": "user", "content": "question"}],
+ )
+ internal_chat_payload = {
+ "id": "chatcmpl-reasoning",
+ "ng_commit_coords": coords,
+ "choices": [
+ {
+ "message": {
+ "role": "assistant",
+ "content": "hidden reasoninganswer",
+ }
+ }
+ ],
+ }
+ served_response = {
+ "id": "chatcmpl-reasoning",
+ "object": "response",
+ "output": [
+ {
+ "id": "reasoning-1",
+ "type": "reasoning",
+ "summary": [{"type": "summary_text", "text": "hidden reasoning"}],
+ },
+ {
+ "id": "message-1",
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": "answer", "annotations": []}],
+ },
+ ],
+ }
+
+ token = set_token_sink(context)
+ try:
+ server._prepare_external_capture(internal_chat_payload)
+ asyncio.run(server._finalize_served_response(served_response))
+ finally:
+ reset_token_sink(token)
+
+ assert context.committed is True
+ assert context.external_commit_coords == coords
+ assert "ng_commit_coords" not in internal_chat_payload
+ manifest = asyncio.run(ledger.manifest("reasoning-external"))
+ record = manifest["records"][0]
+ assert record["output_fingerprint"] == assistant_fingerprint(served_response["output"])
+ assert record["output_fingerprint"] != assistant_fingerprint(
+ [{"role": "assistant", "content": "hidden reasoninganswer"}]
+ )
+ anthropic_items = response_to_output_items(
+ {
+ "id": "chatcmpl-reasoning",
+ "type": "message",
+ "role": "assistant",
+ "content": [
+ {"type": "thinking", "thinking": "hidden reasoning", "signature": ""},
+ {"type": "text", "text": "answer"},
+ ],
+ }
+ )
+ assert [item["type"] for item in anthropic_items] == ["reasoning", "message"]
+ assert assistant_fingerprint(anthropic_items) == record["output_fingerprint"]
+
class TestPrefixSupplyAccounting:
"""Distinguish requested prefixes from prefixes proven to be applied.