From 26527272fce31d4cf5e94139da6cd9e3664bf608 Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Thu, 23 Jul 2026 15:07:55 +0200
Subject: [PATCH 1/8] Define rollout observation and correlation contract

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 .../pages/model-server/model-call-capture.mdx |  40 ++-
 nemo_gym/base_responses_api_model.py          | 252 ++++++++++----
 nemo_gym/rollout_collection.py                |   6 +-
 nemo_gym/rollout_observability.py             | 328 ++++++++++++++++++
 .../test_base_responses_api_model.py          | 255 ++++++++++++--
 tests/unit_tests/test_rollout_collection.py   |   9 +-
 .../unit_tests/test_rollout_observability.py  | 244 +++++++++++++
 7 files changed, 1037 insertions(+), 97 deletions(-)
 create mode 100644 nemo_gym/rollout_observability.py
 create mode 100644 tests/unit_tests/test_rollout_observability.py

diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx
index 8f2a5c0ec3..fb8618a9f2 100644
--- a/fern/versions/latest/pages/model-server/model-call-capture.mdx
+++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx
@@ -41,6 +41,8 @@ Use `/ng-rollout/<rollout_id>` as the model-server URL prefix. SDKs append their
 the rollout prefix before routing.
 
 Correlation is caller-supplied. An unprefixed call is forwarded normally but is not captured.
+SDK-based harnesses must configure their `model_server` field for capture; calls sent directly to an
+external provider do not cross a Gym model server.
 
 Agents built on `SimpleResponsesAPIAgent` can use these helpers: `url_path_for_run(url_path, body)`
 prefixes a downstream call from the run request's task/rollout indices (only when
@@ -77,9 +79,10 @@ totals = aggregate_model_call_metrics(store, rollout_id)
 ```
 
 `ModelCallRecord` is an observability serialization model derived from captured HTTP exchanges. It
-contains a unique server-generated `model_call_id`, typed `model_ref`, wall-clock `started_at` and
-`completed_at`, a `call_index`, API dialect, token and cache usage, latency, error details, tool
-calls, reasoning content, and the captured request and response. `started_at` is recorded immediately
+contains a unique server-generated `model_call_id`, the protocol `response_id` when present, typed
+`model_ref`, wall-clock `started_at` and `completed_at`, a `call_index`, API dialect, token and cache
+usage, latency, error details, tool calls, reasoning content, and the captured request and response.
+`started_at` is recorded immediately
 before invoking the downstream ASGI application; `completed_at` is recorded when that invocation
 returns or raises, before capture parsing and persistence. Both are UTC Unix seconds for external
 trace correlation; durations use the monotonic latency fields. `call_index` reflects durable append
@@ -111,7 +114,32 @@ data does not mix with an earlier attempt. Before dispatch, the collector clears
 for that exact rollout-attempt id, including a kill-shaped attempt being redispatched.
 
 The attachment is additive: it does not replace or rewrite the existing response, reward,
-`NeMoGymResponse`, token-id, or log-prob fields. Downstream consumers can choose whether to read it.
+`NeMoGymResponse`, token-id, or log-prob fields. Downstream consumers can choose whether to read it,
+and aggregate-metrics requests exclude it.
+
+## Agent observations
+
+Supported Agent Servers may also attach `ng_agent_observations` when observability is enabled. It
+contains an unordered `records` list of typed agent invocations, tool-call intervals, explicit
+context-compaction events, and sandbox observations. `gaps` reports unavailable evidence. An
+invocation's `conversation` contains the ordered, normalized items exposed by that integration.
+
+Agent observations and model-call capture are separate evidence. Join an invocation's model-call
+references by `model_call_id`, or by the exact `(model_ref, response_id)` pair when the harness sees
+the protocol response ID. Do not infer ownership from timestamps, text, or list position. The full
+model request and response remain in `CaptureStore`; rollout attachments intentionally omit them.
+
+Tool timestamps are UTC Unix seconds when the source provides wall-clock time. `duration_ms` is the
+measured interval, and `timing_source` identifies executor, harness, or artifact-derived timing.
+
+Model-visible tool calls and results remain in `AgentInvocation.conversation` as `function_call` and
+`function_call_output` items. Execution timing and outcome, when observable, live in
+`ToolCallObservation` and join through `(invocation_id, tool_call_id)`. A tool observation may also
+reference its enclosing `sandbox_id`; concurrent calls retain independent timing and outcome.
+
+Each `SandboxObservation` covers one sandbox execution. Usage fields contain measured values only;
+configured limits are never reported as usage. Integrations emit only facts available at their
+execution boundary or in retained artifacts and report unavailable evidence in `gaps`.
 
 ## Limitations
 
@@ -119,3 +147,7 @@ The model-server boundary observes model HTTP requests and responses. It can rec
 reasoning, and tool results present in those payloads, but it does not observe the actual tool
 execution boundary, environment events, context compaction, semantic turns, or subagent structure.
 Those require instrumentation at the agent, tool, environment, or rollout layer.
+
+Sandbox CPU and memory observations describe the sandbox as a whole. They cannot be attributed to
+individual overlapping tool calls unless each call runs in a separately measured execution scope.
+Gym therefore records the shared sandbox relationship without estimating per-tool resource usage.
diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py
index 7d29549ad9..8a33ca380f 100644
--- a/nemo_gym/base_responses_api_model.py
+++ b/nemo_gym/base_responses_api_model.py
@@ -33,7 +33,6 @@
 import logging
 import os
 import re
-import threading
 import time
 from abc import abstractmethod
 from pathlib import Path
@@ -65,6 +64,7 @@
     synthesize_responses_sse,
     validate_streaming_responses_params,
 )
+from nemo_gym.rollout_observability import AgentObservationBundle, ObservationGap, join_model_call_observations
 from nemo_gym.server_utils import (
     BaseRunServerInstanceConfig,
     BaseServer,
@@ -228,7 +228,6 @@ class CaptureStore:
     def __init__(self, root: str | Path) -> None:
         self._root = Path(root)
         self._root.mkdir(parents=True, exist_ok=True)
-        self._lock = threading.Lock()
 
     @property
     def root(self) -> Path:
@@ -237,25 +236,32 @@ def root(self) -> Path:
     def path_for(self, rollout_id: str) -> Path:
         return self._root / f"{_validate_rollout_id(rollout_id)}.capture.jsonl"
 
+    def incomplete_path_for(self, rollout_id: str) -> Path:
+        return self._root / f"{_validate_rollout_id(rollout_id)}.capture.incomplete"
+
+    def mark_incomplete(self, rollout_id: str) -> None:
+        self.incomplete_path_for(rollout_id).touch(exist_ok=True)
+
+    def is_incomplete(self, rollout_id: str) -> bool:
+        return self.incomplete_path_for(rollout_id).exists()
+
     def record(self, rollout_id: str, exchange: dict[str, Any]) -> None:
         """Append one exchange and fsync (durable across a killed box).
 
-        ``flock`` serializes appends across worker processes (a model server may run with
-        ``num_workers > 1``, where the in-process lock can't coordinate); the in-process lock
-        serializes threads. This does blocking file IO + fsync, so callers run it off the event
-        loop (the capture middleware offloads it via ``asyncio.to_thread``).
+        ``flock`` serializes appends to the same rollout across worker processes and threads while
+        allowing independent rollouts to write concurrently. This does blocking file IO + fsync,
+        so callers run it off the event loop (the capture middleware uses ``asyncio.to_thread``).
         """
         line = orjson.dumps(exchange, default=str, option=orjson.OPT_APPEND_NEWLINE)
         path = self.path_for(rollout_id)
-        with self._lock:
-            with path.open("ab") as handle:
-                fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
-                try:
-                    handle.write(line)
-                    handle.flush()
-                    os.fsync(handle.fileno())
-                finally:
-                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+        with path.open("ab") as handle:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
+            try:
+                handle.write(line)
+                handle.flush()
+                os.fsync(handle.fileno())
+            finally:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
 
     def read(self, rollout_id: str) -> list[dict[str, Any]]:
         path = self.path_for(rollout_id)
@@ -263,19 +269,47 @@ def read(self, rollout_id: str) -> list[dict[str, Any]]:
             return []
         exchanges: list[dict[str, Any]] = []
         # Stream line-by-line; a capture can be large (token-ids / logprobs).
-        with self._lock:
-            with path.open("rb") as handle:
-                fcntl.flock(handle.fileno(), fcntl.LOCK_SH)
-                try:
-                    for line in handle:
-                        stripped = line.strip()
-                        if not stripped:
-                            continue
-                        exchanges.append(orjson.loads(stripped))
-                finally:
-                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+        with path.open("rb") as handle:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_SH)
+            try:
+                for line in handle:
+                    stripped = line.strip()
+                    if not stripped:
+                        continue
+                    exchanges.append(orjson.loads(stripped))
+            finally:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
         return exchanges
 
+    def read_available(self, rollout_id: str) -> tuple[list[tuple[int, dict[str, Any]]], int]:
+        """Read valid exchanges without letting one damaged line hide the rest."""
+        path = self.path_for(rollout_id)
+        if not path.exists():
+            return [], 0
+        exchanges: list[tuple[int, dict[str, Any]]] = []
+        invalid_count = 0
+        capture_index = 0
+        with path.open("rb") as handle:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_SH)
+            try:
+                for line in handle:
+                    stripped = line.strip()
+                    if not stripped:
+                        continue
+                    try:
+                        exchange = orjson.loads(stripped)
+                    except orjson.JSONDecodeError:
+                        invalid_count += 1
+                    else:
+                        if isinstance(exchange, dict):
+                            exchanges.append((capture_index, exchange))
+                        else:
+                            invalid_count += 1
+                    capture_index += 1
+            finally:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+        return exchanges, invalid_count
+
 
 def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]:
     """Per-rollout model-call capture id from a run-request's task/rollout indices.
@@ -308,6 +342,10 @@ def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -
 # --- Observability records derived from captured exchanges ---
 
 
+def _token_count(value: Any) -> Optional[int]:
+    return value if type(value) is int and value >= 0 else None
+
+
 def extract_token_stats(usage: Any) -> dict[str, Optional[int]]:
     """Normalize token totals across Responses, Chat Completions, and Anthropic Messages usage.
 
@@ -321,7 +359,7 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]:
     cache-creation (~1.25x) differently from base input, so cost-accurate consumers should weight
     ``cached_tokens`` and ``cache_creation_tokens`` separately rather than summing ``tokens_in``.
     """
-    if not usage:
+    if not isinstance(usage, Mapping):
         return {
             "tokens_in": None,
             "tokens_out": None,
@@ -329,28 +367,30 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]:
             "tokens_total": None,
             "cache_creation_tokens": None,
         }
-    tokens_in = usage.get("input_tokens")
+    tokens_in = _token_count(usage.get("input_tokens"))
     if tokens_in is None:
-        tokens_in = usage.get("prompt_tokens")
-    tokens_out = usage.get("output_tokens")
+        tokens_in = _token_count(usage.get("prompt_tokens"))
+    tokens_out = _token_count(usage.get("output_tokens"))
     if tokens_out is None:
-        tokens_out = usage.get("completion_tokens")
+        tokens_out = _token_count(usage.get("completion_tokens"))
     # Anthropic-native shape: top-level cache_* keys mean input_tokens excludes cached tokens.
-    cache_read = usage.get("cache_read_input_tokens")
-    cache_creation = usage.get("cache_creation_input_tokens")
+    cache_read = _token_count(usage.get("cache_read_input_tokens"))
+    cache_creation = _token_count(usage.get("cache_creation_input_tokens"))
     if cache_read is not None or cache_creation is not None:
         # A fully-cached response can omit input_tokens; use a 0 base so the folded prompt size is
         # preserved rather than dropped to null. (Top-level cache_* keys are Anthropic-only, so the
         # OpenAI/Responses path -- nested prompt_tokens_details.cached_tokens -- never enters here.)
         tokens_in = (tokens_in or 0) + (cache_read or 0) + (cache_creation or 0)
-    tokens_total = usage.get("total_tokens")
+    tokens_total = _token_count(usage.get("total_tokens"))
     if tokens_total is None and tokens_in is not None and tokens_out is not None:
         tokens_total = tokens_in + tokens_out
     details = usage.get("output_tokens_details") or usage.get("completion_tokens_details") or {}
+    if not isinstance(details, Mapping):
+        details = {}
     return {
         "tokens_in": tokens_in,
         "tokens_out": tokens_out,
-        "tokens_reasoning": details.get("reasoning_tokens"),
+        "tokens_reasoning": _token_count(details.get("reasoning_tokens")),
         "tokens_total": tokens_total,
         "cache_creation_tokens": cache_creation,
     }
@@ -358,12 +398,14 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]:
 
 def _cache_signal(usage: Any) -> tuple[Optional[bool], Optional[int]]:
     """Cache hit/miss + cached-token count, from usage cache fields (OpenAI / Anthropic)."""
-    if not usage:
+    if not isinstance(usage, Mapping):
         return None, None
     details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") or {}
-    cached = details.get("cached_tokens")
+    if not isinstance(details, Mapping):
+        details = {}
+    cached = _token_count(details.get("cached_tokens"))
     if cached is None:
-        cached = usage.get("cache_read_input_tokens")  # Anthropic
+        cached = _token_count(usage.get("cache_read_input_tokens"))  # Anthropic
     if cached is None:
         return None, None
     return cached > 0, cached
@@ -401,7 +443,7 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str,
             elif item.get("type") == "reasoning":
                 for summary in item.get("summary") or []:
                     text = summary.get("text") if isinstance(summary, dict) else None
-                    if text:
+                    if isinstance(text, str) and text:
                         reasoning.append(text)
         return tool_calls, ("\n".join(reasoning) or None)
 
@@ -409,19 +451,21 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str,
     if isinstance(choices, list):  # Chat Completions
         for choice in choices:
             message = choice.get("message") if isinstance(choice, dict) else None
-            if not message:
+            if not isinstance(message, Mapping):
                 continue
             for tc in message.get("tool_calls") or []:
                 if not isinstance(tc, dict):
                     continue
                 fn = tc.get("function") or {}
+                if not isinstance(fn, Mapping):
+                    fn = {}
                 tool_calls.append(
                     {"call_id": tc.get("id"), "name": fn.get("name"), "arguments": _as_arguments(fn.get("arguments"))}
                 )
             # vLLM and newer OpenAI-compatible servers emit `reasoning`; `reasoning_content` is the
             # older field. Accept either (reasoning_content wins when both are present).
             reasoning_text = message.get("reasoning_content") or message.get("reasoning")
-            if reasoning_text:
+            if isinstance(reasoning_text, str) and reasoning_text:
                 reasoning.append(reasoning_text)
         return tool_calls, ("\n".join(reasoning) or None)
 
@@ -434,7 +478,7 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str,
                 tool_calls.append(
                     {"call_id": block.get("id"), "name": block.get("name"), "arguments": block.get("input") or {}}
                 )
-            elif block.get("type") in ("thinking", "redacted_thinking") and block.get("thinking"):
+            elif block.get("type") in ("thinking", "redacted_thinking") and isinstance(block.get("thinking"), str):
                 reasoning.append(block["thinking"])
         return tool_calls, ("\n".join(reasoning) or None)
 
@@ -446,12 +490,15 @@ class ModelCallRecord(BaseModel):
 
     # Unique server-generated identity for each persisted call.
     model_call_id: Optional[str] = None
+    response_id: Optional[str] = None
 
     # Durable append order, not a causal or semantic order for concurrent calls.
     call_index: int
     model_ref: Optional[ModelServerRef] = None
+    model: Optional[str] = None
     dialect: Optional[str] = None
     status_code: Optional[int] = None
+    finish_reason: Optional[str] = None
 
     # Wall-clock bounds around the downstream ASGI invocation, as UTC Unix timestamps. These are
     # for external trace correlation; durations use the monotonic latency fields below.
@@ -469,6 +516,8 @@ class ModelCallRecord(BaseModel):
     # Model-call record.
     request: Optional[dict[str, Any]] = None
     response: Optional[dict[str, Any]] = None
+    request_raw: Optional[str] = None
+    response_raw: Optional[str] = None
     tool_calls: list[dict[str, Any]] = Field(default_factory=list)
 
     # Structured reasoning (not flattened into the response text).
@@ -490,20 +539,46 @@ class ModelCallRecord(BaseModel):
 
 def build_model_call_record(exchange: dict[str, Any], *, call_index: int) -> ModelCallRecord:
     """Map one captured exchange and its transport metadata into an observability record."""
-    response = exchange.get("response") or {}
+    raw_response = exchange.get("response")
+    response = raw_response if isinstance(raw_response, dict) else {}
     tokens = extract_token_stats(response.get("usage"))
     cache_hit, cached_tokens = _cache_signal(response.get("usage"))
     tool_calls, reasoning_content = _tool_calls_and_reasoning(response)
+    raw_request = exchange.get("request")
+    request = raw_request if isinstance(raw_request, dict) else {}
+    choices = response.get("choices")
+    first_choice = choices[0] if isinstance(choices, list) and choices and isinstance(choices[0], dict) else {}
+    incomplete_details = response.get("incomplete_details")
+    if not isinstance(incomplete_details, dict):
+        incomplete_details = {}
+    finish_reason = next(
+        (
+            value
+            for value in (
+                response.get("stop_reason"),
+                first_choice.get("finish_reason"),
+                incomplete_details.get("reason"),
+            )
+            if isinstance(value, str)
+        ),
+        None,
+    )
+    model = response.get("model") or request.get("model")
     return ModelCallRecord(
         model_call_id=exchange.get("model_call_id"),
+        response_id=response.get("id") if isinstance(response.get("id"), str) else None,
         call_index=call_index,
         model_ref=exchange.get("model_ref"),
+        model=model if isinstance(model, str) else None,
         dialect=exchange.get("dialect"),
         status_code=exchange.get("status_code"),
+        finish_reason=finish_reason,
         started_at=exchange.get("started_at"),
         completed_at=exchange.get("completed_at"),
-        request=exchange.get("request"),
-        response=response or None,
+        request=raw_request if isinstance(raw_request, dict) else None,
+        response=raw_response if isinstance(raw_response, dict) else None,
+        request_raw=exchange.get("request_raw") if isinstance(exchange.get("request_raw"), str) else None,
+        response_raw=exchange.get("response_raw") if isinstance(exchange.get("response_raw"), str) else None,
         tool_calls=tool_calls,
         reasoning_content=reasoning_content,
         cache_hit=cache_hit,
@@ -522,6 +597,18 @@ def read_model_call_records(store: CaptureStore, rollout_id: str) -> list[ModelC
     ]
 
 
+def read_available_model_call_records(store: CaptureStore, rollout_id: str) -> tuple[list[ModelCallRecord], int]:
+    """Read valid call records and count damaged records."""
+    exchanges, invalid_count = store.read_available(rollout_id)
+    calls = []
+    for index, exchange in exchanges:
+        try:
+            calls.append(build_model_call_record(exchange, call_index=index))
+        except Exception:
+            invalid_count += 1
+    return calls, invalid_count
+
+
 def aggregate_model_call_records(calls: list[ModelCallRecord]) -> dict[str, Any]:
     """Aggregate token and latency values from model-call records."""
 
@@ -719,11 +806,14 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An
     tool_calls: dict[int, dict[str, Any]] = {}
     usage: Optional[dict[str, Any]] = None
     model: Optional[str] = None
+    response_id: Optional[str] = None
     role = "assistant"
     finish_reason: Optional[str] = None
     saw_choice = False
     for chunk in events:
         model = chunk.get("model") or model
+        if isinstance(chunk.get("id"), str):
+            response_id = chunk["id"]
         if chunk.get("usage"):
             usage = chunk["usage"]
         for choice in chunk.get("choices") or []:
@@ -762,6 +852,8 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An
         "model": model,
         "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
     }
+    if response_id is not None:
+        result["id"] = response_id
     if usage:
         result["usage"] = usage
     return result
@@ -807,6 +899,7 @@ def _record(
     error_category: Optional[str],
     latency_ms: float,
     ttft_ms: Optional[float] = None,
+    response_raw: Optional[str] = None,
 ) -> None:
     """Append one exchange (success or failure). Best-effort: never raises."""
     request_body = None
@@ -837,9 +930,15 @@ def _record(
         }
         if request_raw is not None:
             exchange["request_raw"] = request_raw
+        if response_raw is not None:
+            exchange["response_raw"] = response_raw
         store.record(rollout_id, exchange)
     except Exception:
         logger.warning("Model-call capture failed for one %s call.", dialect, exc_info=True)
+        try:
+            store.mark_incomplete(rollout_id)
+        except Exception:
+            logger.warning("Could not mark rollout %s capture as incomplete.", rollout_id, exc_info=True)
 
 
 class _CaptureMiddleware:
@@ -956,10 +1055,11 @@ async def _flush_deferred_response() -> None:
                     started_at=started_at,
                     completed_at=completed_at,
                     response_body=None,
-                    status_code=None,
+                    status_code=state["status"],
                     error_category=_classify_exception(exc),
                     latency_ms=(time.perf_counter() - start) * 1000.0,
                     ttft_ms=state["ttft_ms"],
+                    response_raw=(bytes(state["body"]).decode("utf-8", errors="replace") if state["body"] else None),
                 )
             except Exception:
                 logger.warning("Model-call capture finalization failed.", exc_info=True)
@@ -986,6 +1086,8 @@ def _parse_and_record() -> None:
                     response_body = (
                         _reconstruct_streamed_response(body_bytes, dialect) if streaming else json.loads(body_bytes)
                     )
+                    if not isinstance(response_body, dict):
+                        response_body = None
                 except Exception:
                     response_body = None
             error_category = _classify_status(status) if status is not None else None
@@ -1000,6 +1102,11 @@ def _parse_and_record() -> None:
             # doesn't silently count as a success with null tokens in reliability/cost sums.
             if error_category is None and body_bytes and response_body is None:
                 error_category = "capture_parse_error"
+            response_raw = (
+                body_bytes.decode("utf-8", errors="replace")
+                if body_bytes and (streaming or response_body is None)
+                else None
+            )
             _record(
                 store,
                 dialect,
@@ -1014,6 +1121,7 @@ def _parse_and_record() -> None:
                 error_category=error_category,
                 latency_ms=latency_ms,
                 ttft_ms=ttft_ms,
+                response_raw=response_raw,
             )
 
         try:
@@ -1058,7 +1166,7 @@ def model_call_capture_dirs_from_config(global_config_dict: Any) -> list[Path]:
 def _store_for_rollout(rollout_id: str, capture_dirs: list[Path]) -> Optional[CaptureStore]:
     for directory in capture_dirs:
         store = CaptureStore(directory)
-        if store.path_for(rollout_id).exists():
+        if store.path_for(rollout_id).exists() or store.is_incomplete(rollout_id):
             return store
     return None
 
@@ -1078,34 +1186,60 @@ def clear_model_call_captures_for_rollouts(records: list[Any], capture_dirs: lis
             rollout_id = maybe_rollout_id_from_run_body(record)
             if rollout_id:
                 store.path_for(rollout_id).unlink(missing_ok=True)
+                store.incomplete_path_for(rollout_id).unlink(missing_ok=True)
 
 
-def merge_model_call_capture_into_record(
-    record: dict[str, Any], capture_dirs: list[Path], *, include_payloads: bool = False
-) -> dict[str, Any]:
+def merge_model_call_capture_into_record(record: dict[str, Any], capture_dirs: list[Path]) -> dict[str, Any]:
     """Attach captured model-call observability data to a rollout record in place.
 
     Keyed by the rollout id derived from the record's task/rollout/attempt indices, so the attached
     shape is identical for every agent harness. Adds
     ``ng_model_call_capture = {rollout_id, metrics, calls}`` where ``calls`` are derived observability
-    records. Raw request and response payloads remain in the capture store unless ``include_payloads``
-    is true. No-op when no capture exists. The harness output and reward are not modified.
+    records. Raw request and response payloads remain in the capture store. Capture/read/join
+    failures are attached as ``gaps``. The harness output and reward are not modified.
     """
     if not capture_dirs:
         return record
     rollout_id = maybe_rollout_id_from_run_body(record)
     if rollout_id is None:
         return record
+    gaps: list[ObservationGap] = []
     store = _store_for_rollout(rollout_id, capture_dirs)
     if store is None:
-        return record
-    calls = read_model_call_records(store, rollout_id)
-    if not calls:
-        return record
-    exclude = None if include_payloads else {"request", "response"}
-    record["ng_model_call_capture"] = {
+        calls = []
+        gaps.append(ObservationGap(code="model_call_capture_no_records"))
+    else:
+        try:
+            calls, invalid_count = read_available_model_call_records(store, rollout_id)
+            if invalid_count:
+                gaps.append(
+                    ObservationGap(
+                        code="model_call_capture_records_unreadable",
+                        detail=f"count={invalid_count}",
+                    )
+                )
+            if store.is_incomplete(rollout_id):
+                gaps.append(ObservationGap(code="model_call_capture_incomplete"))
+            elif not calls and not invalid_count:
+                gaps.append(ObservationGap(code="model_call_capture_no_records"))
+        except Exception:
+            logger.warning("Could not read model-call capture for rollout %s.", rollout_id, exc_info=True)
+            calls = []
+            gaps.append(ObservationGap(code="model_call_capture_unreadable"))
+    observations = record.get("ng_agent_observations")
+    if observations is not None and calls:
+        try:
+            bundle = AgentObservationBundle.model_validate(observations)
+            record["ng_agent_observations"] = join_model_call_observations(bundle, calls).model_dump(mode="json")
+        except Exception:
+            logger.warning("Could not join agent observations for rollout %s.", rollout_id, exc_info=True)
+            gaps.append(ObservationGap(code="agent_observation_join_failed"))
+    capture = {
         "rollout_id": rollout_id,
         "metrics": aggregate_model_call_records(calls),
-        "calls": [call.model_dump(exclude=exclude) for call in calls],
+        "calls": [call.model_dump(exclude={"request", "response", "request_raw", "response_raw"}) for call in calls],
     }
+    if gaps:
+        capture["gaps"] = [gap.model_dump(mode="json", exclude_none=True) for gap in gaps]
+    record["ng_model_call_capture"] = capture
     return record
diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py
index 2fb0bd15b1..17b7cee014 100644
--- a/nemo_gym/rollout_collection.py
+++ b/nemo_gym/rollout_collection.py
@@ -653,7 +653,11 @@ async def _fetch_agent_metrics(agent_name: str, agent_result_list: List[Dict]) -
             # Strip heavyweight fields before sending, but preserve response.usage
             stripped = []
             for r in agent_result_list:
-                entry = {k: v for k, v in r.items() if k not in ("response", "responses_create_params")}
+                entry = {
+                    k: v
+                    for k, v in r.items()
+                    if k not in ("response", "responses_create_params", "ng_agent_observations")
+                }
                 usage = (r.get("response") or {}).get("usage")
                 if usage:
                     entry["response"] = {"usage": usage}
diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py
new file mode 100644
index 0000000000..dd722f7b2a
--- /dev/null
+++ b/nemo_gym/rollout_observability.py
@@ -0,0 +1,328 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Small shared contract for observations exposed by Agent integrations."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from nemo_gym.config_types import ModelServerRef
+from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseInputItem
+
+
+if TYPE_CHECKING:
+    from nemo_gym.base_responses_api_model import ModelCallRecord
+
+
+class ObservationModel(BaseModel):
+    model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+
+class ModelCallRef(ObservationModel):
+    """Stable identifiers an Agent integration can observe for one model call."""
+
+    model_call_id: Optional[str] = None
+    model_ref: Optional[ModelServerRef] = None
+    response_id: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_join_key(self) -> "ModelCallRef":
+        if not self.model_call_id and not (self.model_ref is not None and self.response_id):
+            raise ValueError("model_call_id or both model_ref and response_id are required")
+        return self
+
+
+class AgentInvocation(ObservationModel):
+    """One root Agent or subagent conversation observed by a harness."""
+
+    kind: Literal["agent_invocation"] = "agent_invocation"
+    invocation_id: str
+    parent_invocation_id: Optional[str] = None
+    spawned_by_tool_call_id: Optional[str] = None
+    status: Literal["completed", "failed", "incomplete", "unknown"] = Field(
+        default="unknown", description="Harness-reported invocation outcome; unknown when not explicit."
+    )
+    model_calls: list[ModelCallRef] = Field(default_factory=list)
+    conversation: list[NeMoGymResponseInputItem] = Field(
+        default_factory=list,
+        description="Normalized conversation items supported by this producer; gaps describe unavailable evidence.",
+    )
+
+
+class ToolCallObservation(ObservationModel):
+    """Timing observed for one tool call at an Agent-owned boundary."""
+
+    kind: Literal["tool_call"] = "tool_call"
+    invocation_id: str
+    tool_call_id: str
+    sandbox_id: Optional[str] = Field(
+        default=None,
+        description="Enclosing sandbox instance, shared by concurrent calls; not per-call resource attribution.",
+    )
+    tool_name: Optional[str] = None
+    started_at: Optional[float] = None
+    completed_at: Optional[float] = None
+    duration_ms: Optional[float] = None
+    timing_source: Optional[Literal["executor", "artifact", "harness"]] = None
+    status: Literal["completed", "failed", "timeout", "cancelled", "incomplete", "unknown"] = "unknown"
+    error_type: Optional[str] = None
+
+
+class SandboxObservation(ObservationModel):
+    """Outcome and lifetime resource usage reported by a sandbox-owning harness."""
+
+    kind: Literal["sandbox"] = "sandbox"
+    role: Literal["agent", "verifier", "environment"]
+    provider: Optional[str] = None
+    sandbox_id: Optional[str] = None
+    outcome: Literal["completed", "failed", "timeout", "oom", "sandbox_error", "cancelled", "unknown"] = "unknown"
+    exit_code: Optional[int] = None
+    wall_time_s: Optional[float] = Field(default=None, ge=0)
+    cpu_time_s: Optional[float] = Field(
+        default=None,
+        ge=0,
+        description="Cumulative CPU time for the sandbox, never an allocation or per-tool estimate.",
+    )
+    peak_memory_mib: Optional[float] = Field(
+        default=None,
+        ge=0,
+        description="Measured sandbox high-water mark, never its configured memory limit.",
+    )
+    resource_usage_source: Optional[str] = None
+    error_type: Optional[str] = None
+
+
+class ContextCompactionObservation(ObservationModel):
+    """An explicit context-compaction event reported by the Agent harness."""
+
+    kind: Literal["context_compaction"] = "context_compaction"
+    invocation_id: str
+    observed_at: Optional[float] = None
+    trigger: Optional[str] = None
+    tokens_before: Optional[int] = None
+    tokens_after: Optional[int] = None
+    outcome: Literal["completed", "failed", "aborted", "unknown"] = "unknown"
+    summary: Optional[str] = None
+    first_kept_item_id: Optional[str] = None
+    before_model_call: Optional[ModelCallRef] = None
+    after_model_call: Optional[ModelCallRef] = None
+
+
+class ObservationGap(ObservationModel):
+    """A fact that the selected integration could not observe or join exactly."""
+
+    code: str
+    invocation_id: Optional[str] = None
+    detail: Optional[str] = None
+
+
+AgentObservationRecord = Annotated[
+    AgentInvocation | ToolCallObservation | ContextCompactionObservation | SandboxObservation,
+    Field(discriminator="kind"),
+]
+
+
+class AgentObservationBundle(ObservationModel):
+    """Normalized observations returned by one Agent Server for one rollout."""
+
+    source: str
+    records: list[AgentObservationRecord] = Field(
+        default_factory=list,
+        description="Unordered typed records; list position does not imply execution order.",
+    )
+    gaps: list[ObservationGap] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def validate_identity(self) -> "AgentObservationBundle":
+        invocation_ids = [record.invocation_id for record in self.records if isinstance(record, AgentInvocation)]
+        if len(invocation_ids) != len(set(invocation_ids)):
+            raise ValueError("invocation_id must be unique within an observation bundle")
+        return self
+
+
+def link_tool_calls_to_sandbox(bundle: AgentObservationBundle, sandbox_id: Optional[str]) -> None:
+    """Link unassigned tool calls to their shared enclosing sandbox."""
+    if sandbox_id is None:
+        return
+    for record in bundle.records:
+        if isinstance(record, ToolCallObservation) and record.sandbox_id is None:
+            record.sandbox_id = sandbox_id
+
+
+def join_model_call_observations(
+    bundle: AgentObservationBundle,
+    calls: Iterable[ModelCallRecord],
+) -> AgentObservationBundle:
+    """Resolve harness call references against captured model calls without guessing ownership."""
+
+    result = bundle.model_copy()
+    result.records = [
+        record.model_copy() if isinstance(record, (AgentInvocation, ContextCompactionObservation)) else record
+        for record in bundle.records
+    ]
+    invocations = [record for record in result.records if isinstance(record, AgentInvocation)]
+    compactions = [record for record in result.records if isinstance(record, ContextCompactionObservation)]
+    captured = list(calls)
+    by_call_id: dict[str, list[ModelCallRecord]] = {}
+    by_response: dict[tuple[str, str, str], list[ModelCallRecord]] = {}
+    for call in captured:
+        if call.model_call_id:
+            by_call_id.setdefault(call.model_call_id, []).append(call)
+        if call.model_ref is not None and call.response_id:
+            key = (call.model_ref.type, call.model_ref.name, call.response_id)
+            by_response.setdefault(key, []).append(call)
+
+    def matches(ref: ModelCallRef) -> list[ModelCallRecord]:
+        if ref.model_call_id:
+            candidates = by_call_id.get(ref.model_call_id, [])
+            return [
+                call
+                for call in candidates
+                if (ref.model_ref is None or ref.model_ref == call.model_ref)
+                and (ref.response_id is None or ref.response_id == call.response_id)
+            ]
+        assert ref.model_ref is not None and ref.response_id is not None
+        return by_response.get((ref.model_ref.type, ref.model_ref.name, ref.response_id), [])
+
+    def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef:
+        return ModelCallRef.model_validate(
+            {
+                "model_call_id": call.model_call_id,
+                "model_ref": call.model_ref,
+                "response_id": call.response_id,
+            }
+        )
+
+    join_codes = {
+        "model_call_reference_ambiguous",
+        "model_call_reference_conflict",
+        "model_call_reference_unmatched",
+    }
+    result.gaps = [
+        gap
+        for gap in bundle.gaps
+        if gap.code not in join_codes
+        and not (
+            gap.code == "model_call_ownership_unavailable"
+            and gap.invocation_id is None
+            and (gap.detail is None or gap.detail.startswith("capture:"))
+        )
+    ]
+
+    claimed: set[int] = set()
+    join_gaps: list[ObservationGap] = []
+    for invocation in invocations:
+        resolved: list[ModelCallRef] = []
+        for ref in invocation.model_calls:
+            candidates = matches(ref)
+            detail = ref.model_call_id or ref.response_id
+            if len(candidates) != 1:
+                join_gaps.append(
+                    ObservationGap(
+                        code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"),
+                        invocation_id=invocation.invocation_id,
+                        detail=detail,
+                    )
+                )
+                resolved.append(ref)
+                continue
+
+            call = candidates[0]
+            identity = id(call)
+            if identity in claimed:
+                join_gaps.append(
+                    ObservationGap(
+                        code="model_call_reference_conflict",
+                        invocation_id=invocation.invocation_id,
+                        detail=call.model_call_id or call.response_id,
+                    )
+                )
+                resolved.append(ref)
+                continue
+            claimed.add(identity)
+            resolved.append(canonical(ref, call))
+        invocation.model_calls = resolved
+
+    for compaction in compactions:
+        for field_name in ("before_model_call", "after_model_call"):
+            ref = getattr(compaction, field_name)
+            if ref is None:
+                continue
+            candidates = matches(ref)
+            if len(candidates) == 1:
+                setattr(compaction, field_name, canonical(ref, candidates[0]))
+            else:
+                join_gaps.append(
+                    ObservationGap(
+                        code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"),
+                        invocation_id=compaction.invocation_id,
+                        detail=f"{field_name}:{ref.model_call_id or ref.response_id}",
+                    )
+                )
+
+    result.gaps.extend(join_gaps)
+    for call in captured:
+        if id(call) not in claimed:
+            result.gaps.append(
+                ObservationGap(
+                    code="model_call_ownership_unavailable",
+                    detail=f"capture:{call.model_call_id or call.response_id or 'unknown'}:call_index={call.call_index}",
+                )
+            )
+    result.gaps = list({(gap.code, gap.invocation_id, gap.detail): gap for gap in result.gaps}.values())
+    return result
+
+
+def model_visible_tool_calls(
+    conversation: Iterable[NeMoGymResponseInputItem],
+) -> list[tuple[str, Optional[str], str]]:
+    """Return model-visible call IDs, tool names, and result statuses."""
+
+    def field(item: Any, name: str) -> Any:
+        return item.get(name) if isinstance(item, dict) else getattr(item, name, None)
+
+    items = list(conversation)
+    results = {
+        call_id: field(item, "status")
+        for item in items
+        if field(item, "type") == "function_call_output" and isinstance((call_id := field(item, "call_id")), str)
+    }
+    return [
+        (
+            call_id,
+            field(item, "name"),
+            (
+                results[call_id]
+                if results.get(call_id) in {"completed", "incomplete"}
+                else "unknown"
+                if call_id in results
+                else "incomplete"
+            ),
+        )
+        for item in items
+        if field(item, "type") == "function_call" and isinstance((call_id := field(item, "call_id")), str) and call_id
+    ]
+
+
+@dataclass(frozen=True, slots=True)
+class AgentEpisode:
+    """An Agent response and the observations available at its execution boundary."""
+
+    response: NeMoGymResponse
+    observations: AgentObservationBundle
+
+
+def response_with_observations(episode: AgentEpisode) -> NeMoGymResponse:
+    """Carry observations through a prefixed Agent self-call."""
+    return episode.response.model_copy(update={"ng_agent_observations": episode.observations.model_dump(mode="json")})
+
+
+def pop_response_observations(response: dict[str, Any]) -> Optional[AgentObservationBundle]:
+    """Remove and validate observations carried by an Agent self-call."""
+    observations = response.pop("ng_agent_observations", None)
+    return AgentObservationBundle.model_validate(observations) if observations is not None else None
diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py
index c1fd00304d..c401da9b9d 100644
--- a/tests/unit_tests/test_base_responses_api_model.py
+++ b/tests/unit_tests/test_base_responses_api_model.py
@@ -135,6 +135,7 @@ def test_build_model_call_record_from_exchange():
         "latency_ms": 18.4,
         "request": {"input": "hi"},
         "response": {
+            "id": "resp-1",
             "model": "m",
             "usage": {
                 "input_tokens": 10,
@@ -152,8 +153,10 @@ def test_build_model_call_record_from_exchange():
     }
     rec = build_model_call_record(exchange, call_index=3)
     assert rec.model_call_id == "call-1"
+    assert rec.response_id == "resp-1"
     assert rec.call_index == 3
     assert rec.model_ref is not None and rec.model_ref.name == "srv"
+    assert rec.model == "m"
     assert rec.dialect == "responses"
     assert rec.started_at == 100.0 and rec.completed_at == 100.02
     assert (rec.tokens_in, rec.tokens_out, rec.tokens_total, rec.tokens_reasoning) == (10, 5, 15, 3)
@@ -161,12 +164,19 @@ def test_build_model_call_record_from_exchange():
     assert rec.reasoning_content == "thinking..."
     assert rec.tool_calls == [{"call_id": "c1", "name": "calc", "arguments": {"x": 1}}]
     assert rec.latency_total_ms == 18.4
+    assert build_model_call_record({"response": {"id": 123}}, call_index=0).response_id is None
+    empty = build_model_call_record({"request": {}, "response": {}}, call_index=0)
+    assert empty.request == {}
+    assert empty.response == {}
     assert {
         "model_call_id",
+        "response_id",
         "call_index",
         "model_ref",
+        "model",
         "dialect",
         "status_code",
+        "finish_reason",
         "started_at",
         "completed_at",
         "tokens_in",
@@ -175,6 +185,8 @@ def test_build_model_call_record_from_exchange():
         "tokens_total",
         "request",
         "response",
+        "request_raw",
+        "response_raw",
         "tool_calls",
         "reasoning_content",
         "cache_hit",
@@ -186,6 +198,36 @@ def test_build_model_call_record_from_exchange():
     } <= type(rec).model_json_schema()["properties"].keys()
 
 
+@pytest.mark.parametrize(
+    "response,expected",
+    [
+        ({"choices": [{"finish_reason": "tool_calls"}]}, "tool_calls"),
+        ({"stop_reason": "end_turn"}, "end_turn"),
+        ({"incomplete_details": {"reason": "max_output_tokens"}}, "max_output_tokens"),
+        ({"status": "completed"}, None),
+        ({"choices": {"finish_reason": "invalid"}, "incomplete_details": []}, None),
+    ],
+)
+def test_build_model_call_record_normalizes_finish_reason(response, expected):
+    assert build_model_call_record({"response": response}, call_index=0).finish_reason == expected
+
+
+def test_build_model_call_record_tolerates_malformed_nested_shapes():
+    record = build_model_call_record(
+        {
+            "response": {
+                "usage": {"input_tokens": "invalid", "prompt_tokens_details": {"cached_tokens": []}},
+                "choices": [{"message": "invalid"}],
+                "output": [{"type": "reasoning", "summary": [{"text": {}}]}],
+            }
+        },
+        call_index=0,
+    )
+
+    assert record.tokens_total is None
+    assert record.tool_calls == []
+
+
 def test_capture_is_durable_before_stream_terminal_event_is_sent(tmp_path):
     import asyncio
 
@@ -232,6 +274,48 @@ async def send(message):
     assert durable_call_counts == [0, 1, 1]
 
 
+def test_capture_retains_partial_stream_when_downstream_raises(tmp_path):
+    import asyncio
+
+    from nemo_gym.base_responses_api_model import _CaptureMiddleware
+
+    store = CaptureStore(tmp_path)
+    partial = b'data: {"type":"response.output_text.delta","delta":"partial"}\n\n'
+
+    async def app(_scope, receive, send):
+        await receive()
+        await send(
+            {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]}
+        )
+        await send({"type": "http.response.body", "body": partial, "more_body": True})
+        raise RuntimeError("stream failed")
+
+    async def receive():
+        return {"type": "http.request", "body": b'{"input":"hi"}', "more_body": False}
+
+    async def send(_message):
+        pass
+
+    with pytest.raises(RuntimeError, match="stream failed"):
+        asyncio.run(
+            _CaptureMiddleware(app, store=store, model_server_name="srv")(
+                {
+                    "type": "http",
+                    "path": "/ng-rollout/partial/v1/responses",
+                    "raw_path": b"/ng-rollout/partial/v1/responses",
+                    "headers": [],
+                },
+                receive,
+                send,
+            )
+        )
+
+    [call] = read_model_call_records(store, "partial")
+    assert call.status_code == 200
+    assert call.error_category == "exception"
+    assert call.response_raw == partial.decode()
+
+
 def test_stream_error_events_are_terminal():
     from nemo_gym.base_responses_api_model import _consume_terminal_sse_event
 
@@ -343,7 +427,7 @@ async def _responses(body: dict = Body()) -> dict:
 
     [record] = read_model_call_records(CaptureStore(tmp_path), "invalid-request")
     assert record.request is None
-    assert "request_raw" not in record.model_dump()
+    assert record.request_raw == request_bytes.decode()
 
 
 def test_per_rollout_url_prefix_correlates_and_is_openai_compatible(tmp_path):
@@ -490,16 +574,15 @@ def _boom(_root):
     assert obs.make_capture_store(config) is None
 
 
-def test_record_swallows_store_failure():
+def test_record_swallows_store_failure_and_marks_capture_incomplete(tmp_path, monkeypatch):
     from nemo_gym.base_responses_api_model import _record
 
-    class _BadStore:
-        def record(self, *args, **kwargs):
-            raise RuntimeError("disk full")
+    store = CaptureStore(tmp_path)
+    monkeypatch.setattr(store, "record", MagicMock(side_effect=RuntimeError("disk full")))
 
     # Best-effort: a failing store must not raise out of _record.
     _record(
-        _BadStore(),
+        store,
         "chat",
         "srv",
         b"{}",
@@ -512,6 +595,7 @@ def record(self, *args, **kwargs):
         error_category=None,
         latency_ms=1.0,
     )
+    assert store.is_incomplete("r")
 
 
 def test_record_falls_back_to_raw_when_request_parser_raises(tmp_path, monkeypatch):
@@ -555,8 +639,11 @@ async def _r(body: dict = Body()) -> PlainTextResponse:
     assert r.status_code == 200 and r.text == "not json"  # response passed through unaltered
     records = CaptureStore(tmp_path).read("rnj")
     assert len(records) == 1 and records[0]["response"] is None  # non-JSON body -> None
+    assert records[0]["response_raw"] == "not json"
     # a 2xx whose body we couldn't parse is flagged, not silently counted as a clean success
     assert records[0]["error_category"] == "capture_parse_error"
+    [record] = read_model_call_records(CaptureStore(tmp_path), "rnj")
+    assert record.response_raw == "not json"
 
 
 def test_as_arguments():
@@ -835,11 +922,13 @@ async def gen():
 
     records = CaptureStore(tmp_path).read("3-0")
     assert len(records) == 1 and records[0]["response"] is not None  # reassembled, not dropped
+    assert records[0]["response_raw"] == r.text  # lossless SSE evidence is retained alongside the projection
 
     calls = read_model_call_records(CaptureStore(tmp_path), "3-0")
     assert len(calls) == 1
     call = calls[0]
     assert call.model_call_id
+    assert call.response_id == "msg_1"
     assert call.model_ref is not None and call.model_ref.name == "srv"
     assert call.started_at is not None and call.completed_at is not None
     assert call.started_at <= call.completed_at
@@ -850,13 +939,18 @@ async def gen():
     assert call.tool_calls == [{"call_id": "t1", "name": "calc", "arguments": {"x": 1}}]
     assert call.latency_ttft_ms is not None
     assert call.error_category is None
+    assert call.response_raw == r.text
 
 
 def test_reconstruct_chat_sse():
     from nemo_gym.base_responses_api_model import _reconstruct_streamed_response
 
     chunks = [
-        {"model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}]},
+        {
+            "id": "chatcmpl-1",
+            "model": "m",
+            "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}],
+        },
         {"choices": [{"index": 0, "delta": {"content": "lo", "reasoning": "hmm"}}]},  # vLLM `reasoning` alias
         {
             "choices": [
@@ -882,6 +976,7 @@ def test_reconstruct_chat_sse():
     raw = (b"".join(_sse("", c) for c in chunks) + b"data: [DONE]\n\n").replace(b"\n", b"\r\n")
     resp = _reconstruct_streamed_response(raw, "chat")
     msg = resp["choices"][0]["message"]
+    assert resp["id"] == "chatcmpl-1"
     assert msg["content"] == "Hello" and msg["reasoning_content"] == "hmm"
     assert msg["tool_calls"][0]["function"] == {"name": "f", "arguments": '{"a":1}'}
     assert resp["usage"]["total_tokens"] == 8
@@ -957,17 +1052,41 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path):
     from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record
 
     store = CaptureStore(tmp_path)
-    store.record(
-        "0-0",
-        _capture_exchange(
-            "responses",
-            "A",
-            {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
-            {"output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}]},
-        ),
+    exchange = _capture_exchange(
+        "responses",
+        "A",
+        {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
+        {
+            "id": "resp-A",
+            "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}],
+        },
     )
-
-    record = {"_ng_task_index": 0, "_ng_rollout_index": 0, "reward": 1.0, "response": {"harness": "A"}}
+    exchange["request_raw"] = "malformed request"
+    exchange["response_raw"] = "malformed response"
+    store.record("0-0", exchange)
+
+    record = {
+        "_ng_task_index": 0,
+        "_ng_rollout_index": 0,
+        "reward": 1.0,
+        "response": {"harness": "A"},
+        "ng_agent_observations": {
+            "source": "test",
+            "records": [
+                {
+                    "kind": "agent_invocation",
+                    "invocation_id": "root",
+                    "model_calls": [
+                        {
+                            "model_ref": {"type": "responses_api_models", "name": "A"},
+                            "response_id": "resp-A",
+                        }
+                    ],
+                }
+            ],
+            "gaps": [{"code": "model_call_ownership_unavailable"}],
+        },
+    }
     merge_model_call_capture_into_record(record, [tmp_path])
 
     capture = record["ng_model_call_capture"]
@@ -976,33 +1095,74 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path):
     assert capture["metrics"]["num_calls"] == 1
     attached_call = capture["calls"][0]
     assert attached_call["model_call_id"] == "call-A"
+    assert attached_call["response_id"] == "resp-A"
     assert attached_call["model_ref"] == {"type": "responses_api_models", "name": "A"}
+    assert attached_call["model"] == "m"
     assert attached_call["started_at"] == 100.0 and attached_call["completed_at"] == 100.01
     assert attached_call["tokens_in"] == 3
-    assert "request" not in attached_call and "response" not in attached_call
+    assert {"request", "response", "request_raw", "response_raw"}.isdisjoint(attached_call)
     assert record["response"] == {"harness": "A"} and record["reward"] == 1.0
+    [joined_ref] = record["ng_agent_observations"]["records"][0]["model_calls"]
+    assert joined_ref["model_call_id"] == "call-A"
+    assert record["ng_agent_observations"]["gaps"] == []
 
 
-def test_merge_capture_noop_without_capture(tmp_path):
-    from nemo_gym.base_responses_api_model import merge_model_call_capture_into_record
+def test_merge_capture_reports_missing_capture(tmp_path):
+    from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record
 
-    rec = {"_ng_task_index": 9, "_ng_rollout_index": 9, "reward": 1.0}
+    rec = {
+        "_ng_task_index": 9,
+        "_ng_rollout_index": 9,
+        "reward": 1.0,
+        "ng_agent_observations": {
+            "source": "test",
+            "gaps": [{"code": "model_call_ownership_unavailable"}],
+        },
+    }
     merge_model_call_capture_into_record(rec, [tmp_path])  # no capture file for 9-9
-    assert "ng_model_call_capture" not in rec
-    merge_model_call_capture_into_record(rec, [])  # no dirs
-    assert "ng_model_call_capture" not in rec
+    assert rec["ng_model_call_capture"]["calls"] == []
+    assert [gap["code"] for gap in rec["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"]
+    assert rec["ng_agent_observations"]["gaps"] == [{"code": "model_call_ownership_unavailable"}]
 
+    CaptureStore(tmp_path).path_for("8-8").touch()
+    empty = {"_ng_task_index": 8, "_ng_rollout_index": 8}
+    merge_model_call_capture_into_record(empty, [tmp_path])
+    assert [gap["code"] for gap in empty["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"]
 
-def test_merge_capture_surfaces_malformed_data_only_when_active(tmp_path):
+
+def test_merge_capture_preserves_valid_records_around_malformed_data(tmp_path):
     from nemo_gym.base_responses_api_model import merge_model_call_capture_into_record
 
     store = CaptureStore(tmp_path)
-    store.path_for("9-9").write_bytes(b"{not-json}\n")
-    record = {"_ng_task_index": 9, "_ng_rollout_index": 9}
+    first = _capture_exchange("responses", "A", {}, {"id": "resp-A"})
+    third = _capture_exchange("responses", "B", {}, {"id": "resp-B"})
+    store.path_for("9-8").write_bytes(orjson.dumps(first) + b"\n{not-json}\n" + orjson.dumps(third) + b"\n")
+    record = {"_ng_task_index": 9, "_ng_rollout_index": 8}
 
-    merge_model_call_capture_into_record(record, [])
-    with pytest.raises(orjson.JSONDecodeError):
-        merge_model_call_capture_into_record(record, [tmp_path])
+    merge_model_call_capture_into_record(record, [tmp_path])
+
+    assert [call["call_index"] for call in record["ng_model_call_capture"]["calls"]] == [0, 2]
+    assert [call["response_id"] for call in record["ng_model_call_capture"]["calls"]] == ["resp-A", "resp-B"]
+    assert record["ng_model_call_capture"]["gaps"] == [
+        {
+            "code": "model_call_capture_records_unreadable",
+            "detail": "count=1",
+        }
+    ]
+
+
+def test_merge_capture_reports_partial_write_loss(tmp_path):
+    from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record
+
+    store = CaptureStore(tmp_path)
+    store.record("4-2", _capture_exchange("responses", "A", {}, {"id": "resp-A"}))
+    store.mark_incomplete("4-2")
+    record = {"_ng_task_index": 4, "_ng_rollout_index": 2}
+
+    merge_model_call_capture_into_record(record, [tmp_path])
+
+    assert len(record["ng_model_call_capture"]["calls"]) == 1
+    assert [gap["code"] for gap in record["ng_model_call_capture"]["gaps"]] == ["model_call_capture_incomplete"]
 
 
 def test_clear_model_call_captures_for_rollouts_run_scoping(tmp_path, monkeypatch):
@@ -1012,11 +1172,12 @@ def test_clear_model_call_captures_for_rollouts_run_scoping(tmp_path, monkeypatc
     store = CaptureStore(tmp_path)
     store.record("0-0", {"dialect": "chat", "request": {}, "response": {}})
     store.record("1-0", {"dialect": "chat", "request": {}, "response": {}})
+    store.mark_incomplete("0-0")
     assert store.read("0-0") and store.read("1-0")
 
     # Clears only the rollout ids about to be (re)run; rows without indices are skipped, others stay.
     clear_model_call_captures_for_rollouts([{"_ng_task_index": 0, "_ng_rollout_index": 0}, {"no": "id"}], [tmp_path])
-    assert store.read("0-0") == [] and store.read("1-0")
+    assert store.read("0-0") == [] and not store.is_incomplete("0-0") and store.read("1-0")
     clear_model_call_captures_for_rollouts([{"_ng_task_index": 1, "_ng_rollout_index": 0}], [])  # no dirs -> no-op
     assert store.read("1-0")
 
@@ -1126,10 +1287,40 @@ def _write(i: int) -> None:
     for t in threads:
         t.join()
     rows = store.read("0-0")
-    assert len(rows) == 20  # flock + in-process lock: no lost or corrupted appends
+    assert len(rows) == 20  # flock prevents lost or corrupted appends
     assert sorted(r["request"]["i"] for r in rows) == list(range(20))
 
 
+def test_capture_store_does_not_serialize_independent_rollouts(tmp_path, monkeypatch):
+    import os
+    import threading
+
+    store = CaptureStore(tmp_path)
+    first_fsync = threading.Event()
+    release_first = threading.Event()
+    original_fsync = os.fsync
+
+    def _fsync(fd):
+        if threading.current_thread().name == "first-rollout":
+            first_fsync.set()
+            assert release_first.wait(timeout=5)
+        original_fsync(fd)
+
+    monkeypatch.setattr(os, "fsync", _fsync)
+    first = threading.Thread(name="first-rollout", target=store.record, args=("0-0", {"request": {}, "response": {}}))
+    second = threading.Thread(target=store.record, args=("1-0", {"request": {}, "response": {}}))
+    first.start()
+    assert first_fsync.wait(timeout=5)
+    second.start()
+    second.join(timeout=1)
+    try:
+        assert not second.is_alive()
+        assert store.read("1-0") == [{"request": {}, "response": {}}]
+    finally:
+        release_first.set()
+        first.join(timeout=5)
+
+
 def test_capture_store_read_waits_for_in_progress_append(tmp_path):
     import fcntl
     import threading
diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py
index ece7bacb1b..134d7d331c 100644
--- a/tests/unit_tests/test_rollout_collection.py
+++ b/tests/unit_tests/test_rollout_collection.py
@@ -994,7 +994,13 @@ def setup_server_client(self):
             {AGENT_REF_KEY_NAME: {"name": "my_agent"}, TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 1},
         ]
         results = [
-            {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 10}}},
+            {
+                TASK_INDEX_KEY_NAME: 0,
+                ROLLOUT_INDEX_KEY_NAME: 0,
+                "reward": 1.0,
+                "response": {"usage": {"tokens": 10}},
+                "ng_agent_observations": {"invocations": [{"conversation": ["large"]}]},
+            },
             {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 12}}},
             {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 8}}},
             {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 15}}},
@@ -1023,6 +1029,7 @@ def setup_server_client(self):
         )
         for item in sent_data:
             assert "responses_create_params" not in item
+            assert "ng_agent_observations" not in item
             assert "usage" in item["response"]
 
     async def test_call_aggregate_metrics_multiple_agents(self, tmp_path: Path) -> None:
diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py
new file mode 100644
index 0000000000..c5dadd658d
--- /dev/null
+++ b/tests/unit_tests/test_rollout_observability.py
@@ -0,0 +1,244 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from pydantic import ValidationError
+
+from nemo_gym.base_responses_api_model import ModelCallRecord
+from nemo_gym.config_types import ModelServerRef
+from nemo_gym.openai_utils import NeMoGymResponse
+from nemo_gym.rollout_observability import (
+    AgentEpisode,
+    AgentInvocation,
+    AgentObservationBundle,
+    ContextCompactionObservation,
+    ModelCallRef,
+    ObservationGap,
+    SandboxObservation,
+    ToolCallObservation,
+    join_model_call_observations,
+    link_tool_calls_to_sandbox,
+    pop_response_observations,
+    response_with_observations,
+)
+
+
+@pytest.mark.parametrize(
+    "value",
+    ({}, {"response_id": "resp-1"}, {"model_ref": {"name": "policy", "type": "responses_api_models"}}),
+)
+def test_model_call_ref_rejects_incomplete_join_keys(value: dict) -> None:
+    with pytest.raises(ValidationError, match="model_call_id or both model_ref and response_id"):
+        ModelCallRef.model_validate(value)
+
+
+def test_observations_round_trip_through_internal_response() -> None:
+    response = NeMoGymResponse.model_validate(
+        {
+            "id": "response",
+            "created_at": 0,
+            "model": "model",
+            "object": "response",
+            "output": [],
+            "parallel_tool_calls": True,
+            "tool_choice": "auto",
+            "tools": [],
+        }
+    )
+    carried = response_with_observations(
+        AgentEpisode(response=response, observations=AgentObservationBundle(source="test"))
+    ).model_dump(mode="json")
+
+    observations = pop_response_observations(carried)
+
+    assert observations == AgentObservationBundle(source="test")
+    assert "ng_agent_observations" not in carried
+
+
+def test_parallel_tool_calls_share_enclosing_sandbox_without_overwriting_existing_link() -> None:
+    bundle = AgentObservationBundle(
+        source="test",
+        records=[
+            ToolCallObservation(
+                invocation_id="root",
+                tool_call_id="call-1",
+                started_at=1.0,
+                completed_at=3.0,
+            ),
+            ToolCallObservation(
+                invocation_id="root",
+                tool_call_id="call-2",
+                started_at=2.0,
+                completed_at=4.0,
+            ),
+            ToolCallObservation(
+                invocation_id="child",
+                tool_call_id="call-3",
+                sandbox_id="other-sandbox",
+            ),
+        ],
+    )
+
+    link_tool_calls_to_sandbox(bundle, "shared-sandbox")
+
+    tools = [record for record in bundle.records if isinstance(record, ToolCallObservation)]
+    assert [tool.sandbox_id for tool in tools] == [
+        "shared-sandbox",
+        "shared-sandbox",
+        "other-sandbox",
+    ]
+
+
+def test_observation_bundle_rejects_duplicate_invocation_ids() -> None:
+    with pytest.raises(ValidationError, match="invocation_id must be unique"):
+        AgentObservationBundle(
+            source="test",
+            records=[AgentInvocation(invocation_id="root"), AgentInvocation(invocation_id="root")],
+        )
+
+
+def test_observation_models_reject_unknown_fields() -> None:
+    with pytest.raises(ValidationError, match="producer_extension"):
+        ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"})
+
+
+def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None:
+    model_ref = ModelServerRef(name="policy", type="responses_api_models")
+    bundle = AgentObservationBundle(
+        source="test",
+        records=[
+            AgentInvocation(
+                invocation_id="root",
+                model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")],
+            )
+        ],
+        gaps=[ObservationGap(code="model_call_ownership_unavailable")],
+    )
+    calls = [
+        ModelCallRecord(
+            model_call_id="call-1",
+            response_id="resp-1",
+            model_ref=model_ref,
+            call_index=0,
+        ),
+        ModelCallRecord(model_call_id="call-2", model_ref=model_ref, call_index=1),
+    ]
+
+    joined = join_model_call_observations(bundle, calls)
+
+    [invocation] = [record for record in joined.records if isinstance(record, AgentInvocation)]
+    [joined_call] = invocation.model_calls
+    assert joined_call.model_call_id == "call-1"
+    assert joined_call.model_ref == model_ref
+    assert joined_call.response_id == "resp-1"
+    ownership_gaps = [gap for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"]
+    assert [gap.detail for gap in ownership_gaps] == ["capture:call-2:call_index=1"]
+
+
+def test_join_model_calls_does_not_guess_ambiguous_response_ids() -> None:
+    model_ref = ModelServerRef(name="policy", type="responses_api_models")
+    bundle = AgentObservationBundle(
+        source="test",
+        records=[
+            AgentInvocation(
+                invocation_id="root",
+                model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")],
+            )
+        ],
+    )
+    calls = [
+        ModelCallRecord(model_call_id="call-1", response_id="resp-1", model_ref=model_ref, call_index=0),
+        ModelCallRecord(model_call_id="call-2", response_id="resp-1", model_ref=model_ref, call_index=1),
+    ]
+
+    joined = join_model_call_observations(bundle, calls)
+
+    [invocation] = [record for record in joined.records if isinstance(record, AgentInvocation)]
+    assert invocation.model_calls[0].model_call_id is None
+    assert "model_call_reference_ambiguous" in {gap.code for gap in joined.gaps}
+    assert [gap.detail for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"] == [
+        "capture:call-1:call_index=0",
+        "capture:call-2:call_index=1",
+    ]
+
+
+def test_join_model_calls_reports_conflicting_and_unmatched_references() -> None:
+    model_ref = ModelServerRef(name="policy", type="responses_api_models")
+    bundle = AgentObservationBundle(
+        source="test",
+        records=[
+            AgentInvocation(
+                invocation_id="root",
+                model_calls=[
+                    ModelCallRef(model_call_id="call-1"),
+                    ModelCallRef(model_ref=model_ref, response_id="resp-1"),
+                    ModelCallRef(model_call_id="missing"),
+                ],
+            )
+        ],
+        gaps=[
+            ObservationGap(
+                code="model_call_ownership_unavailable",
+                invocation_id="root",
+                detail="producer gap",
+            )
+        ],
+    )
+    calls = [
+        ModelCallRecord(
+            model_call_id="call-1",
+            response_id="resp-1",
+            model_ref=model_ref,
+            call_index=0,
+        )
+    ]
+
+    joined = join_model_call_observations(bundle, calls)
+    joined_again = join_model_call_observations(joined, calls)
+
+    assert joined_again.model_dump() == joined.model_dump()
+    assert {gap.code for gap in joined.gaps} == {
+        "model_call_ownership_unavailable",
+        "model_call_reference_conflict",
+        "model_call_reference_unmatched",
+    }
+    assert any(gap.detail == "producer gap" for gap in joined.gaps)
+
+
+def test_join_model_calls_resolves_compaction_boundaries() -> None:
+    model_ref = ModelServerRef(name="policy", type="responses_api_models")
+    bundle = AgentObservationBundle(
+        source="test",
+        records=[
+            ContextCompactionObservation(
+                invocation_id="root",
+                before_model_call=ModelCallRef(model_ref=model_ref, response_id="before"),
+                after_model_call=ModelCallRef(model_ref=model_ref, response_id="after"),
+            )
+        ],
+    )
+    calls = [
+        ModelCallRecord(
+            model_call_id="call-before",
+            response_id="before",
+            model_ref=model_ref,
+            call_index=0,
+        ),
+        ModelCallRecord(
+            model_call_id="call-after",
+            response_id="after",
+            model_ref=model_ref,
+            call_index=1,
+        ),
+    ]
+
+    joined = join_model_call_observations(bundle, calls)
+    [compaction] = [record for record in joined.records if isinstance(record, ContextCompactionObservation)]
+
+    assert compaction.before_model_call.model_call_id == "call-before"
+    assert compaction.after_model_call.model_call_id == "call-after"
+
+
+def test_sandbox_observation_rejects_negative_usage() -> None:
+    with pytest.raises(ValidationError):
+        SandboxObservation(role="agent", cpu_time_s=-1)

From b851fdf3f17bf51db8e21236a363926b8b94d80a Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Fri, 24 Jul 2026 18:29:35 +0200
Subject: [PATCH 2/8] Propagate rollout correlation across servers

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 nemo_gym/base_resources_server.py            |   2 +
 nemo_gym/base_responses_api_agent.py         |  16 +-
 nemo_gym/base_responses_api_model.py         |  34 +---
 nemo_gym/mcp_auto_exposure.py                |   2 +
 nemo_gym/rollout_correlation.py              |  87 ++++++++
 nemo_gym/server_utils.py                     |  30 ++-
 tests/unit_tests/test_rollout_correlation.py | 198 +++++++++++++++++++
 7 files changed, 332 insertions(+), 37 deletions(-)
 create mode 100644 nemo_gym/rollout_correlation.py
 create mode 100644 tests/unit_tests/test_rollout_correlation.py

diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py
index 1c4b6be35c..2eba9f0605 100644
--- a/nemo_gym/base_resources_server.py
+++ b/nemo_gym/base_resources_server.py
@@ -30,6 +30,7 @@
     NeMoGymResponseCreateParamsNonStreaming,
 )
 from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics
+from nemo_gym.rollout_correlation import RolloutContextMiddleware
 from nemo_gym.server_utils import BaseRunServerInstanceConfig, BaseServer, SimpleServer
 
 
@@ -110,6 +111,7 @@ def setup_webserver(self) -> FastAPI:
         app = FastAPI()
 
         self.setup_session_middleware(app)
+        app.add_middleware(RolloutContextMiddleware)
 
         app.post("/seed_session")(self.seed_session)
         app.post("/verify")(self.verify)
diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py
index 18fe0816ab..c6e76cd89f 100644
--- a/nemo_gym/base_responses_api_agent.py
+++ b/nemo_gym/base_responses_api_agent.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 from abc import abstractmethod
 from collections.abc import Mapping
+from functools import wraps
 from typing import Any, Optional
 
 from fastapi import Body, FastAPI, Request
@@ -24,7 +25,6 @@
     BaseRunRequest,
     BaseVerifyResponse,
 )
-from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body
 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.openai_utils import (
@@ -32,6 +32,7 @@
     NeMoGymResponseCreateParamsNonStreaming,
 )
 from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics
+from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body, rollout_context
 from nemo_gym.server_utils import (
     BaseRunServerInstanceConfig,
     BaseServer,
@@ -62,7 +63,18 @@ def setup_webserver(self) -> FastAPI:
         # 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.
         app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses")(self.responses)
-        app.post("/run")(self.run)
+
+        run = self.run
+
+        @wraps(run)
+        async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyResponse:
+            body = kwargs.get("body")
+            if body is None:
+                body = next((arg for arg in args if isinstance(arg, BaseRunRequest)), None)
+            with rollout_context(self.rollout_id_from_run(body)):
+                return await run(*args, **kwargs)
+
+        app.post("/run")(run_with_rollout_context)
         app.post("/aggregate_metrics")(self.aggregate_metrics)
 
         return app
diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py
index 8a33ca380f..daaa25f93f 100644
--- a/nemo_gym/base_responses_api_model.py
+++ b/nemo_gym/base_responses_api_model.py
@@ -47,11 +47,6 @@
 
 from nemo_gym.anthropic_converter import AnthropicConverter
 from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, ModelServerRef
-from nemo_gym.global_config import (
-    ATTEMPT_INDEX_KEY_NAME,
-    ROLLOUT_INDEX_KEY_NAME,
-    TASK_INDEX_KEY_NAME,
-)
 from nemo_gym.openai_utils import (
     NeMoGymChatCompletion,
     NeMoGymChatCompletionCreateParamsNonStreaming,
@@ -64,6 +59,7 @@
     synthesize_responses_sse,
     validate_streaming_responses_params,
 )
+from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body
 from nemo_gym.rollout_observability import AgentObservationBundle, ObservationGap, join_model_call_observations
 from nemo_gym.server_utils import (
     BaseRunServerInstanceConfig,
@@ -311,34 +307,6 @@ def read_available(self, rollout_id: str) -> tuple[list[tuple[int, dict[str, Any
         return exchanges, invalid_count
 
 
-def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]:
-    """Per-rollout model-call capture id from a run-request's task/rollout indices.
-
-    Reads the canonical row keys (``_ng_task_index`` / ``_ng_rollout_index``) that
-    rollout_collection ships to an agent's ``/run``. When a resume re-dispatch attempt is present
-    (``_ng_attempt_index`` > 0), an ``-a<n>`` suffix is appended so a retry's captured model calls
-    stay separable from the prior attempt; the first attempt (0) keeps the bare ``<task>-<rollout>``
-    key for backward compatibility.
-    """
-    if isinstance(body, BaseModel):
-        data = body.model_dump()
-    elif isinstance(body, Mapping):
-        data = body
-    else:
-        return None
-    task = data.get(TASK_INDEX_KEY_NAME)
-    rollout = data.get(ROLLOUT_INDEX_KEY_NAME)
-    if task is None or rollout is None:
-        return None
-    rollout_id = f"{task}-{rollout}"
-    attempt = data.get(ATTEMPT_INDEX_KEY_NAME)
-    if attempt is not None:
-        attempt_index = int(attempt)
-        if attempt_index > 0:
-            rollout_id = f"{rollout_id}-a{attempt_index}"
-    return rollout_id
-
-
 # --- Observability records derived from captured exchanges ---
 
 
diff --git a/nemo_gym/mcp_auto_exposure.py b/nemo_gym/mcp_auto_exposure.py
index 45de0aaaef..0e6aab4f6d 100644
--- a/nemo_gym/mcp_auto_exposure.py
+++ b/nemo_gym/mcp_auto_exposure.py
@@ -438,6 +438,8 @@ def harvest_tools(app: FastAPI, server: Any) -> dict[str, MCPTool]:
         cls = m.cls
         if f"{cls.__module__}.{cls.__name__}" == "starlette.middleware.sessions.SessionMiddleware":
             continue  # Gym's SessionMiddleware — replaced by a materialized session on direct dispatch
+        if f"{cls.__module__}.{cls.__name__}" == "nemo_gym.rollout_correlation.RolloutContextMiddleware":
+            continue  # Correlation prefixes are handled before resource routes; direct MCP dispatch has no prefix.
         dispatch = m.kwargs.get("dispatch")
         if dispatch is not None and getattr(dispatch, "__module__", None) in _GYM_MIDDLEWARE_MODULES:
             continue  # Gym's add_session_id / exception middleware
diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py
new file mode 100644
index 0000000000..e41a70b240
--- /dev/null
+++ b/nemo_gym/rollout_correlation.py
@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import re
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from nemo_gym.config_types import ROLLOUT_PATH_PREFIX
+from nemo_gym.global_config import (
+    ATTEMPT_INDEX_KEY_NAME,
+    ROLLOUT_INDEX_KEY_NAME,
+    TASK_INDEX_KEY_NAME,
+)
+
+
+_ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_rollout_id", default=None)
+
+
+def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]:
+    """Build the capture key stamped by rollout collection."""
+    if isinstance(body, BaseModel):
+        data = body.model_dump()
+    elif isinstance(body, Mapping):
+        data = body
+    else:
+        return None
+
+    task = data.get(TASK_INDEX_KEY_NAME)
+    rollout = data.get(ROLLOUT_INDEX_KEY_NAME)
+    if task is None or rollout is None:
+        return None
+
+    rollout_id = f"{task}-{rollout}"
+    attempt = data.get(ATTEMPT_INDEX_KEY_NAME)
+    if attempt is not None and int(attempt) > 0:
+        rollout_id = f"{rollout_id}-a{int(attempt)}"
+    return rollout_id
+
+
+def current_rollout_id() -> Optional[str]:
+    return _ROLLOUT_ID.get()
+
+
+@contextmanager
+def rollout_context(rollout_id: Optional[str]) -> Iterator[None]:
+    token = _ROLLOUT_ID.set(rollout_id)
+    try:
+        yield
+    finally:
+        _ROLLOUT_ID.reset(token)
+
+
+class RolloutContextMiddleware:
+    """Strip a rollout prefix and expose it to downstream Gym calls for this request."""
+
+    _PREFIX = re.compile(
+        rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P<rollout_id>[A-Za-z0-9][A-Za-z0-9._-]*)(?P<rest>/.*)$"
+    )
+
+    def __init__(self, app: Any) -> None:
+        self._app = app
+
+    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
+        match = self._PREFIX.match(scope.get("path", "")) if scope.get("type") == "http" else None
+        if match is None:
+            await self._app(scope, receive, send)
+            return
+
+        path = match.group("rest")
+        scope = {**scope, "path": path, "raw_path": path.encode()}
+        with rollout_context(match.group("rollout_id")):
+            await self._app(scope, receive, send)
diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py
index 03070f237f..d05f3100cd 100644
--- a/nemo_gym/server_utils.py
+++ b/nemo_gym/server_utils.py
@@ -64,6 +64,7 @@
     DRY_RUN_KEY_NAME,
     HEAD_SERVER_KEY_NAME,
     NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME,
+    OBSERVABILITY_ENABLED_KEY_NAME,
     RAY_HEAD_NODE_ADDRESS_KEY_NAME,
     GlobalConfigDictParser,
     GlobalConfigDictParserConfig,
@@ -71,6 +72,7 @@
     get_global_config_dict,
 )
 from nemo_gym.profiling import Profiler
+from nemo_gym.rollout_correlation import current_rollout_id, maybe_rollout_id_from_run_body
 
 
 _GLOBAL_AIOHTTP_CLIENT: Union[None, ClientSession] = None
@@ -330,10 +332,34 @@ async def request(
         server_config_dict = get_first_server_config_dict(self.global_config_dict, server_name)
         base_url = self._build_server_base_url(server_config_dict)
 
+        json_obj = kwargs.get("json")
         if "json" in kwargs:
-            json_obj = kwargs["json"]
             if isinstance(json_obj, BaseModel):
-                kwargs["json"] = json_obj.model_dump(exclude_unset=True)
+                json_obj = json_obj.model_dump(exclude_unset=True)
+                kwargs["json"] = json_obj
+
+        observability_enabled = self.global_config_dict.get(OBSERVABILITY_ENABLED_KEY_NAME, False)
+        server_entry = self.global_config_dict.get(server_name)
+        rollout_id = current_rollout_id()
+        if (
+            observability_enabled
+            and server_entry is not None
+            and "resources_servers" in server_entry
+            and url_path == "/verify"
+        ):
+            rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj)
+            if rollout_id is not None:
+                url_path = f"{rollout_path_prefix(rollout_id)}{url_path}"
+
+        if (
+            rollout_id is not None
+            and observability_enabled
+            and server_entry is not None
+            and "responses_api_models" in server_entry
+            and url_path.partition("?")[0] in {"/v1/responses", "/v1/chat/completions", "/v1/messages"}
+            and not url_path.startswith(f"/{ROLLOUT_PATH_PREFIX}/")
+        ):
+            url_path = f"{rollout_path_prefix(rollout_id)}{url_path}"
 
         return await request(method=method, url=f"{base_url}{url_path}", _internal=True, **kwargs)
 
diff --git a/tests/unit_tests/test_rollout_correlation.py b/tests/unit_tests/test_rollout_correlation.py
new file mode 100644
index 0000000000..25f18ba62f
--- /dev/null
+++ b/tests/unit_tests/test_rollout_correlation.py
@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from urllib.parse import urlsplit
+
+import orjson
+import pytest
+from fastapi import Body, FastAPI
+from fastapi.responses import JSONResponse
+from omegaconf import OmegaConf
+from pydantic import ConfigDict
+from starlette.testclient import TestClient
+
+import nemo_gym.server_utils
+from nemo_gym.base_resources_server import (
+    BaseResourcesServerConfig,
+    BaseRunRequest,
+    BaseVerifyRequest,
+    BaseVerifyResponse,
+    SimpleResourcesServer,
+)
+from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig, SimpleResponsesAPIAgent
+from nemo_gym.base_responses_api_model import (
+    CaptureStore,
+    ModelCallCaptureConfig,
+    install_model_call_capture,
+    merge_model_call_capture_into_record,
+)
+from nemo_gym.config_types import BaseServerConfig
+from nemo_gym.server_utils import ServerClient, get_response_json
+
+
+def _model_response(model: str, text: str = "") -> dict:
+    return {
+        "id": f"resp-{model}",
+        "created_at": 0.0,
+        "model": model,
+        "object": "response",
+        "output": [
+            {
+                "id": f"msg-{model}",
+                "type": "message",
+                "role": "assistant",
+                "status": "completed",
+                "content": [{"type": "output_text", "text": text, "annotations": []}],
+            }
+        ],
+        "parallel_tool_calls": True,
+        "tool_choice": "auto",
+        "tools": [],
+    }
+
+
+def _model_app(capture_dir, name: str) -> FastAPI:
+    app = FastAPI()
+
+    @app.post("/v1/responses")
+    async def responses(body: dict = Body()) -> JSONResponse:
+        return JSONResponse(_model_response(name, "[[A=B]]" if name == "judge" else "answer"))
+
+    install_model_call_capture(
+        app,
+        ModelCallCaptureConfig(observability_enabled=True, model_call_capture_dir=capture_dir),
+        model_server_name=name,
+    )
+    return app
+
+
+class _JudgeResourcesServer(SimpleResourcesServer):
+    async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse:
+        judge = await self.server_client.post(
+            server_name="judge",
+            url_path="/v1/responses",
+            json={"input": "grade"},
+        )
+        await get_response_json(judge)
+        return BaseVerifyResponse(**body.model_dump(), reward=1.0)
+
+
+class _AgentRunRequest(BaseRunRequest):
+    model_config = ConfigDict(extra="allow")
+
+
+class _Agent(SimpleResponsesAPIAgent):
+    async def responses(self, body):
+        raise NotImplementedError
+
+    async def run(self, body: _AgentRunRequest) -> BaseVerifyResponse:
+        policy = await self.server_client.post(
+            server_name="policy",
+            url_path="/v1/responses",
+            json=body.responses_create_params,
+        )
+        response = orjson.loads(await policy.read())
+        verify = await self.server_client.post(
+            server_name="resources",
+            url_path="/verify",
+            json={
+                "responses_create_params": body.responses_create_params.model_dump(),
+                "response": response,
+            },
+        )
+        return BaseVerifyResponse.model_validate(orjson.loads(await verify.read()))
+
+
+class _Response:
+    def __init__(self, response) -> None:
+        self.status = response.status_code
+        self.ok = response.is_success
+        self.cookies = response.cookies
+        self._content = response.content
+
+    async def read(self) -> bytes:
+        return self._content
+
+
+@pytest.mark.asyncio
+async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_capture(tmp_path, monkeypatch) -> None:
+    capture_dir = tmp_path / "captures"
+    config = OmegaConf.create(
+        {
+            "observability_enabled": True,
+            "policy": {"responses_api_models": {"model": {"host": "policy.test", "port": 80}}},
+            "judge": {"responses_api_models": {"model": {"host": "judge.test", "port": 80}}},
+            "resources": {"resources_servers": {"judge": {"host": "resources.test", "port": 80}}},
+            "agent": {"responses_api_agents": {"agent": {"host": "agent.test", "port": 80}}},
+        }
+    )
+    server_client = ServerClient(
+        head_server_config=BaseServerConfig(host="head.test", port=80),
+        global_config_dict=config,
+    )
+    resources = _JudgeResourcesServer(
+        config=BaseResourcesServerConfig(
+            host="resources.test",
+            port=80,
+            entrypoint="app.py",
+            name="resources",
+        ),
+        server_client=server_client,
+    )
+    agent = _Agent(
+        config=BaseResponsesAPIAgentConfig(
+            host="agent.test",
+            port=80,
+            entrypoint="app.py",
+            name="agent",
+        ),
+        server_client=server_client,
+    )
+    clients = {
+        "policy.test": TestClient(_model_app(capture_dir, "policy")),
+        "judge.test": TestClient(_model_app(capture_dir, "judge")),
+        "resources.test": TestClient(resources.setup_webserver()),
+        "agent.test": TestClient(agent.setup_webserver()),
+    }
+
+    async def dispatch(method: str, url: str, **kwargs):
+        parsed = urlsplit(url)
+        response = clients[parsed.hostname].request(method, parsed.path, json=kwargs.get("json"))
+        return _Response(response)
+
+    monkeypatch.setattr(nemo_gym.server_utils, "request", dispatch)
+
+    verify = await server_client.post(
+        server_name="agent",
+        url_path="/run",
+        json={
+            "_ng_task_index": 4,
+            "_ng_rollout_index": 2,
+            "responses_create_params": {"input": "solve"},
+        },
+    )
+    assert orjson.loads(await verify.read())["reward"] == 1.0
+
+    store = CaptureStore(capture_dir)
+    capture_path = store.path_for("4-2")
+    assert capture_path.is_file()
+    exchanges = store.read("4-2")
+    assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "judge"]
+    assert all(exchange.get("request") is not None or exchange.get("request_raw") for exchange in exchanges)
+    assert all(exchange.get("response") is not None or exchange.get("response_raw") for exchange in exchanges)
+
+    rollout = {"_ng_task_index": 4, "_ng_rollout_index": 2}
+    merge_model_call_capture_into_record(rollout, [capture_dir])
+    assert capture_path.is_file()
+    assert len(capture_path.read_bytes()) > 0

From 52279d7e1ad08c89bdbe869cb1e4b1bca38deb55 Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Mon, 27 Jul 2026 17:38:53 +0200
Subject: [PATCH 3/8] Tighten rollout observation foundation

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 nemo_gym/rollout_collection.py                |  8 +-
 nemo_gym/rollout_correlation.py               | 15 ++--
 nemo_gym/rollout_observability.py             | 61 ++-------------
 nemo_gym/server_utils.py                      | 12 +--
 tests/unit_tests/test_rollout_collection.py   |  2 +
 tests/unit_tests/test_rollout_correlation.py  | 40 +++++++++-
 .../unit_tests/test_rollout_observability.py  | 74 +++----------------
 7 files changed, 79 insertions(+), 133 deletions(-)

diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py
index 17b7cee014..d503687980 100644
--- a/nemo_gym/rollout_collection.py
+++ b/nemo_gym/rollout_collection.py
@@ -656,7 +656,13 @@ async def _fetch_agent_metrics(agent_name: str, agent_result_list: List[Dict]) -
                 entry = {
                     k: v
                     for k, v in r.items()
-                    if k not in ("response", "responses_create_params", "ng_agent_observations")
+                    if k
+                    not in (
+                        "response",
+                        "responses_create_params",
+                        "ng_agent_observations",
+                        "ng_model_call_capture",
+                    )
                 }
                 usage = (r.get("response") or {}).get("usage")
                 if usage:
diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py
index e41a70b240..c1d58a2694 100644
--- a/nemo_gym/rollout_correlation.py
+++ b/nemo_gym/rollout_correlation.py
@@ -33,20 +33,19 @@
 
 def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]:
     """Build the capture key stamped by rollout collection."""
-    if isinstance(body, BaseModel):
-        data = body.model_dump()
-    elif isinstance(body, Mapping):
-        data = body
-    else:
+    if not isinstance(body, (BaseModel, Mapping)):
         return None
 
-    task = data.get(TASK_INDEX_KEY_NAME)
-    rollout = data.get(ROLLOUT_INDEX_KEY_NAME)
+    def field(key: str) -> Any:
+        return body.get(key) if isinstance(body, Mapping) else getattr(body, key, None)
+
+    task = field(TASK_INDEX_KEY_NAME)
+    rollout = field(ROLLOUT_INDEX_KEY_NAME)
     if task is None or rollout is None:
         return None
 
     rollout_id = f"{task}-{rollout}"
-    attempt = data.get(ATTEMPT_INDEX_KEY_NAME)
+    attempt = field(ATTEMPT_INDEX_KEY_NAME)
     if attempt is not None and int(attempt) > 0:
         rollout_id = f"{rollout_id}-a{int(attempt)}"
     return rollout_id
diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py
index dd722f7b2a..2ea664de12 100644
--- a/nemo_gym/rollout_observability.py
+++ b/nemo_gym/rollout_observability.py
@@ -7,7 +7,7 @@
 
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional
+from typing import TYPE_CHECKING, Annotated, Literal, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
@@ -67,11 +67,17 @@ class ToolCallObservation(ObservationModel):
     tool_name: Optional[str] = None
     started_at: Optional[float] = None
     completed_at: Optional[float] = None
-    duration_ms: Optional[float] = None
+    duration_ms: Optional[float] = Field(default=None, ge=0)
     timing_source: Optional[Literal["executor", "artifact", "harness"]] = None
     status: Literal["completed", "failed", "timeout", "cancelled", "incomplete", "unknown"] = "unknown"
     error_type: Optional[str] = None
 
+    @model_validator(mode="after")
+    def validate_timing(self) -> "ToolCallObservation":
+        if self.started_at is not None and self.completed_at is not None and self.completed_at < self.started_at:
+            raise ValueError("completed_at must not precede started_at")
+        return self
+
 
 class SandboxObservation(ObservationModel):
     """Outcome and lifetime resource usage reported by a sandbox-owning harness."""
@@ -145,15 +151,6 @@ def validate_identity(self) -> "AgentObservationBundle":
         return self
 
 
-def link_tool_calls_to_sandbox(bundle: AgentObservationBundle, sandbox_id: Optional[str]) -> None:
-    """Link unassigned tool calls to their shared enclosing sandbox."""
-    if sandbox_id is None:
-        return
-    for record in bundle.records:
-        if isinstance(record, ToolCallObservation) and record.sandbox_id is None:
-            record.sandbox_id = sandbox_id
-
-
 def join_model_call_observations(
     bundle: AgentObservationBundle,
     calls: Iterable[ModelCallRecord],
@@ -278,51 +275,9 @@ def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef:
     return result
 
 
-def model_visible_tool_calls(
-    conversation: Iterable[NeMoGymResponseInputItem],
-) -> list[tuple[str, Optional[str], str]]:
-    """Return model-visible call IDs, tool names, and result statuses."""
-
-    def field(item: Any, name: str) -> Any:
-        return item.get(name) if isinstance(item, dict) else getattr(item, name, None)
-
-    items = list(conversation)
-    results = {
-        call_id: field(item, "status")
-        for item in items
-        if field(item, "type") == "function_call_output" and isinstance((call_id := field(item, "call_id")), str)
-    }
-    return [
-        (
-            call_id,
-            field(item, "name"),
-            (
-                results[call_id]
-                if results.get(call_id) in {"completed", "incomplete"}
-                else "unknown"
-                if call_id in results
-                else "incomplete"
-            ),
-        )
-        for item in items
-        if field(item, "type") == "function_call" and isinstance((call_id := field(item, "call_id")), str) and call_id
-    ]
-
-
 @dataclass(frozen=True, slots=True)
 class AgentEpisode:
     """An Agent response and the observations available at its execution boundary."""
 
     response: NeMoGymResponse
     observations: AgentObservationBundle
-
-
-def response_with_observations(episode: AgentEpisode) -> NeMoGymResponse:
-    """Carry observations through a prefixed Agent self-call."""
-    return episode.response.model_copy(update={"ng_agent_observations": episode.observations.model_dump(mode="json")})
-
-
-def pop_response_observations(response: dict[str, Any]) -> Optional[AgentObservationBundle]:
-    """Remove and validate observations carried by an Agent self-call."""
-    observations = response.pop("ng_agent_observations", None)
-    return AgentObservationBundle.model_validate(observations) if observations is not None else None
diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py
index d05f3100cd..b676c8f49a 100644
--- a/nemo_gym/server_utils.py
+++ b/nemo_gym/server_utils.py
@@ -341,14 +341,10 @@ async def request(
         observability_enabled = self.global_config_dict.get(OBSERVABILITY_ENABLED_KEY_NAME, False)
         server_entry = self.global_config_dict.get(server_name)
         rollout_id = current_rollout_id()
-        if (
-            observability_enabled
-            and server_entry is not None
-            and "resources_servers" in server_entry
-            and url_path == "/verify"
-        ):
-            rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj)
-            if rollout_id is not None:
+        if observability_enabled and server_entry is not None and "resources_servers" in server_entry:
+            if url_path == "/verify":
+                rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj)
+            if rollout_id is not None and not url_path.startswith(f"/{ROLLOUT_PATH_PREFIX}/"):
                 url_path = f"{rollout_path_prefix(rollout_id)}{url_path}"
 
         if (
diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py
index 134d7d331c..d3145eb403 100644
--- a/tests/unit_tests/test_rollout_collection.py
+++ b/tests/unit_tests/test_rollout_collection.py
@@ -1000,6 +1000,7 @@ def setup_server_client(self):
                 "reward": 1.0,
                 "response": {"usage": {"tokens": 10}},
                 "ng_agent_observations": {"invocations": [{"conversation": ["large"]}]},
+                "ng_model_call_capture": {"calls": [{"request": "large"}]},
             },
             {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 12}}},
             {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 8}}},
@@ -1030,6 +1031,7 @@ def setup_server_client(self):
         for item in sent_data:
             assert "responses_create_params" not in item
             assert "ng_agent_observations" not in item
+            assert "ng_model_call_capture" not in item
             assert "usage" in item["response"]
 
     async def test_call_aggregate_metrics_multiple_agents(self, tmp_path: Path) -> None:
diff --git a/tests/unit_tests/test_rollout_correlation.py b/tests/unit_tests/test_rollout_correlation.py
index 25f18ba62f..3954ece3fc 100644
--- a/tests/unit_tests/test_rollout_correlation.py
+++ b/tests/unit_tests/test_rollout_correlation.py
@@ -38,6 +38,7 @@
     merge_model_call_capture_into_record,
 )
 from nemo_gym.config_types import BaseServerConfig
+from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body
 from nemo_gym.server_utils import ServerClient, get_response_json
 
 
@@ -78,6 +79,19 @@ async def responses(body: dict = Body()) -> JSONResponse:
 
 
 class _JudgeResourcesServer(SimpleResourcesServer):
+    def setup_webserver(self) -> FastAPI:
+        app = super().setup_webserver()
+        app.post("/tool")(self.tool)
+        return app
+
+    async def tool(self, body: dict = Body()) -> dict:
+        tool_model = await self.server_client.post(
+            server_name="tool_model",
+            url_path="/v1/responses",
+            json={"input": body["input"]},
+        )
+        return await get_response_json(tool_model)
+
     async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse:
         judge = await self.server_client.post(
             server_name="judge",
@@ -103,6 +117,12 @@ async def run(self, body: _AgentRunRequest) -> BaseVerifyResponse:
             json=body.responses_create_params,
         )
         response = orjson.loads(await policy.read())
+        tool = await self.server_client.post(
+            server_name="resources",
+            url_path="/tool",
+            json={"input": "lookup"},
+        )
+        await get_response_json(tool)
         verify = await self.server_client.post(
             server_name="resources",
             url_path="/verify",
@@ -132,6 +152,7 @@ async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_captur
         {
             "observability_enabled": True,
             "policy": {"responses_api_models": {"model": {"host": "policy.test", "port": 80}}},
+            "tool_model": {"responses_api_models": {"model": {"host": "tool-model.test", "port": 80}}},
             "judge": {"responses_api_models": {"model": {"host": "judge.test", "port": 80}}},
             "resources": {"resources_servers": {"judge": {"host": "resources.test", "port": 80}}},
             "agent": {"responses_api_agents": {"agent": {"host": "agent.test", "port": 80}}},
@@ -161,6 +182,7 @@ async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_captur
     )
     clients = {
         "policy.test": TestClient(_model_app(capture_dir, "policy")),
+        "tool-model.test": TestClient(_model_app(capture_dir, "tool_model")),
         "judge.test": TestClient(_model_app(capture_dir, "judge")),
         "resources.test": TestClient(resources.setup_webserver()),
         "agent.test": TestClient(agent.setup_webserver()),
@@ -188,7 +210,7 @@ async def dispatch(method: str, url: str, **kwargs):
     capture_path = store.path_for("4-2")
     assert capture_path.is_file()
     exchanges = store.read("4-2")
-    assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "judge"]
+    assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "tool_model", "judge"]
     assert all(exchange.get("request") is not None or exchange.get("request_raw") for exchange in exchanges)
     assert all(exchange.get("response") is not None or exchange.get("response_raw") for exchange in exchanges)
 
@@ -196,3 +218,19 @@ async def dispatch(method: str, url: str, **kwargs):
     merge_model_call_capture_into_record(rollout, [capture_dir])
     assert capture_path.is_file()
     assert len(capture_path.read_bytes()) > 0
+
+
+def test_rollout_id_does_not_serialize_run_body() -> None:
+    class UndumpableRunRequest(_AgentRunRequest):
+        def model_dump(self, *args, **kwargs):
+            raise AssertionError("run body must not be serialized")
+
+    body = UndumpableRunRequest.model_validate(
+        {
+            "_ng_task_index": 4,
+            "_ng_rollout_index": 2,
+            "responses_create_params": {"input": "solve"},
+        }
+    )
+
+    assert maybe_rollout_id_from_run_body(body) == "4-2"
diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py
index c5dadd658d..daf42e5b5a 100644
--- a/tests/unit_tests/test_rollout_observability.py
+++ b/tests/unit_tests/test_rollout_observability.py
@@ -6,9 +6,7 @@
 
 from nemo_gym.base_responses_api_model import ModelCallRecord
 from nemo_gym.config_types import ModelServerRef
-from nemo_gym.openai_utils import NeMoGymResponse
 from nemo_gym.rollout_observability import (
-    AgentEpisode,
     AgentInvocation,
     AgentObservationBundle,
     ContextCompactionObservation,
@@ -17,9 +15,6 @@
     SandboxObservation,
     ToolCallObservation,
     join_model_call_observations,
-    link_tool_calls_to_sandbox,
-    pop_response_observations,
-    response_with_observations,
 )
 
 
@@ -32,63 +27,6 @@ def test_model_call_ref_rejects_incomplete_join_keys(value: dict) -> None:
         ModelCallRef.model_validate(value)
 
 
-def test_observations_round_trip_through_internal_response() -> None:
-    response = NeMoGymResponse.model_validate(
-        {
-            "id": "response",
-            "created_at": 0,
-            "model": "model",
-            "object": "response",
-            "output": [],
-            "parallel_tool_calls": True,
-            "tool_choice": "auto",
-            "tools": [],
-        }
-    )
-    carried = response_with_observations(
-        AgentEpisode(response=response, observations=AgentObservationBundle(source="test"))
-    ).model_dump(mode="json")
-
-    observations = pop_response_observations(carried)
-
-    assert observations == AgentObservationBundle(source="test")
-    assert "ng_agent_observations" not in carried
-
-
-def test_parallel_tool_calls_share_enclosing_sandbox_without_overwriting_existing_link() -> None:
-    bundle = AgentObservationBundle(
-        source="test",
-        records=[
-            ToolCallObservation(
-                invocation_id="root",
-                tool_call_id="call-1",
-                started_at=1.0,
-                completed_at=3.0,
-            ),
-            ToolCallObservation(
-                invocation_id="root",
-                tool_call_id="call-2",
-                started_at=2.0,
-                completed_at=4.0,
-            ),
-            ToolCallObservation(
-                invocation_id="child",
-                tool_call_id="call-3",
-                sandbox_id="other-sandbox",
-            ),
-        ],
-    )
-
-    link_tool_calls_to_sandbox(bundle, "shared-sandbox")
-
-    tools = [record for record in bundle.records if isinstance(record, ToolCallObservation)]
-    assert [tool.sandbox_id for tool in tools] == [
-        "shared-sandbox",
-        "shared-sandbox",
-        "other-sandbox",
-    ]
-
-
 def test_observation_bundle_rejects_duplicate_invocation_ids() -> None:
     with pytest.raises(ValidationError, match="invocation_id must be unique"):
         AgentObservationBundle(
@@ -102,6 +40,18 @@ def test_observation_models_reject_unknown_fields() -> None:
         ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"})
 
 
+@pytest.mark.parametrize(
+    "timing",
+    (
+        {"duration_ms": -1},
+        {"started_at": 2.0, "completed_at": 1.0},
+    ),
+)
+def test_tool_call_observation_rejects_invalid_timing(timing: dict) -> None:
+    with pytest.raises(ValidationError):
+        ToolCallObservation(invocation_id="root", tool_call_id="call-1", **timing)
+
+
 def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None:
     model_ref = ModelServerRef(name="policy", type="responses_api_models")
     bundle = AgentObservationBundle(

From 0d2836cc4cdeb66189401d8667d5b363c3d596aa Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Mon, 27 Jul 2026 19:06:15 +0200
Subject: [PATCH 4/8] Record agent invocation outcomes

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 nemo_gym/rollout_observability.py              | 2 ++
 tests/unit_tests/test_rollout_observability.py | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py
index 2ea664de12..162359e059 100644
--- a/nemo_gym/rollout_observability.py
+++ b/nemo_gym/rollout_observability.py
@@ -47,6 +47,8 @@ class AgentInvocation(ObservationModel):
     status: Literal["completed", "failed", "incomplete", "unknown"] = Field(
         default="unknown", description="Harness-reported invocation outcome; unknown when not explicit."
     )
+    duration_ms: Optional[float] = Field(default=None, ge=0)
+    error_type: Optional[str] = None
     model_calls: list[ModelCallRef] = Field(default_factory=list)
     conversation: list[NeMoGymResponseInputItem] = Field(
         default_factory=list,
diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py
index daf42e5b5a..2a09972cf5 100644
--- a/tests/unit_tests/test_rollout_observability.py
+++ b/tests/unit_tests/test_rollout_observability.py
@@ -52,6 +52,11 @@ def test_tool_call_observation_rejects_invalid_timing(timing: dict) -> None:
         ToolCallObservation(invocation_id="root", tool_call_id="call-1", **timing)
 
 
+def test_agent_invocation_rejects_negative_duration() -> None:
+    with pytest.raises(ValidationError):
+        AgentInvocation(invocation_id="root", duration_ms=-1)
+
+
 def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None:
     model_ref = ModelServerRef(name="policy", type="responses_api_models")
     bundle = AgentObservationBundle(

From 20f8fc763177f5f25087f266f6177b42aa9b1e0a Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Mon, 27 Jul 2026 19:34:30 +0200
Subject: [PATCH 5/8] Preserve upstream failure evidence

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 nemo_gym/base_responses_api_model.py          | 29 +++++++++++++++++--
 .../test_base_responses_api_model.py          | 23 +++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py
index daaa25f93f..377de77daa 100644
--- a/nemo_gym/base_responses_api_model.py
+++ b/nemo_gym/base_responses_api_model.py
@@ -698,6 +698,24 @@ def _classify_exception(exc: BaseException) -> str:
     return "exception"
 
 
+def _exception_http_details(exc: BaseException) -> tuple[Optional[int], bytes]:
+    response = getattr(exc, "response", None)
+    status = getattr(exc, "status", None)
+    if not isinstance(status, int):
+        status = getattr(exc, "status_code", None)
+    if not isinstance(status, int) and response is not None:
+        status = getattr(response, "status_code", None)
+
+    body = getattr(exc, "response_content", None)
+    if body is None and response is not None:
+        body = getattr(response, "content", None)
+        if body is None:
+            body = getattr(response, "text", None)
+    if isinstance(body, str):
+        body = body.encode()
+    return (status if isinstance(status, int) else None, bytes(body) if isinstance(body, (bytes, bytearray)) else b"")
+
+
 # --- SSE reconstruction: rebuild a final response object from a streamed body ---
 def _parse_sse_events(raw: bytes) -> list[dict[str, Any]]:
     """Parse an SSE byte stream into its JSON ``data:`` payloads (best-effort; non-JSON skipped)."""
@@ -1010,6 +1028,11 @@ async def _flush_deferred_response() -> None:
             await self._app(scope, _receive, _send)
         except Exception as exc:
             completed_at = time.time()
+            exception_status, exception_body = _exception_http_details(exc)
+            upstream_status = state["status"] or exception_status
+            upstream_body = bytes(state["body"]) or exception_body
+            error_category = _classify_status(upstream_status) if isinstance(upstream_status, int) else None
+            error_category = error_category or _classify_exception(exc)
             # Offload the blocking write+fsync so it never stalls the event loop.
             try:
                 await asyncio.to_thread(
@@ -1023,11 +1046,11 @@ async def _flush_deferred_response() -> None:
                     started_at=started_at,
                     completed_at=completed_at,
                     response_body=None,
-                    status_code=state["status"],
-                    error_category=_classify_exception(exc),
+                    status_code=upstream_status,
+                    error_category=error_category,
                     latency_ms=(time.perf_counter() - start) * 1000.0,
                     ttft_ms=state["ttft_ms"],
-                    response_raw=(bytes(state["body"]).decode("utf-8", errors="replace") if state["body"] else None),
+                    response_raw=upstream_body.decode("utf-8", errors="replace") if upstream_body else None,
                 )
             except Exception:
                 logger.warning("Model-call capture finalization failed.", exc_info=True)
diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py
index c401da9b9d..0e1775ff3f 100644
--- a/tests/unit_tests/test_base_responses_api_model.py
+++ b/tests/unit_tests/test_base_responses_api_model.py
@@ -402,6 +402,29 @@ async def _boom(body: dict = Body()) -> dict:
     assert calls[0].latency_ttft_ms is None  # nothing streamed before the raise
 
 
+def test_raised_upstream_error_preserves_status_and_body(tmp_path):
+    class UpstreamError(RuntimeError):
+        response = SimpleNamespace(status_code=403, content=b'{"error":{"message":"denied"}}')
+
+    app = FastAPI()
+
+    @app.post("/v1/responses")
+    async def _boom(body: dict = Body()) -> dict:
+        raise UpstreamError
+
+    _install_capture(app, tmp_path)
+    response = TestClient(app, raise_server_exceptions=False).post(
+        "/ng-rollout/r-upstream/v1/responses",
+        json={"input": "x"},
+    )
+
+    assert response.status_code == 500
+    [exchange] = CaptureStore(tmp_path).read("r-upstream")
+    assert exchange["status_code"] == 403
+    assert exchange["error_category"] == "auth"
+    assert json.loads(exchange["response_raw"]) == {"error": {"message": "denied"}}
+
+
 @pytest.mark.parametrize("request_bytes", [b"{not-json", b"[]"])
 def test_invalid_request_body_does_not_drop_capture(tmp_path, request_bytes):
     app = FastAPI()

From ffcb549f37ceb578fb9d419875f5398e03163ac2 Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Thu, 23 Jul 2026 15:12:41 +0200
Subject: [PATCH 6/8] Propagate rollout correlation to auxiliary model servers

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 resources_servers/gdpval/app.py               | 13 +++--
 resources_servers/gdpval/tests/test_app.py    |  7 ++-
 responses_api_agents/stirrup_agent/app.py     | 19 +++++--
 .../stirrup_agent/tests/test_app.py           | 55 ++++++++++++++++++-
 responses_api_agents/tau2/tests/test_app.py   | 42 ++++++++++++++
 5 files changed, 123 insertions(+), 13 deletions(-)

diff --git a/resources_servers/gdpval/app.py b/resources_servers/gdpval/app.py
index 47f36db923..111948c53c 100644
--- a/resources_servers/gdpval/app.py
+++ b/resources_servers/gdpval/app.py
@@ -41,7 +41,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from nemo_gym.base_resources_server import (
     BaseResourcesServerConfig,
@@ -50,7 +50,7 @@
     SimpleResourcesServer,
 )
 from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest, ModelServerRef
-from nemo_gym.server_utils import get_server_url
+from nemo_gym.server_utils import apply_rollout_prefix, get_server_url
 from resources_servers.gdpval.judge_panel import (
     ResolvedJudge,
     dir_contains_audio_video,
@@ -234,6 +234,7 @@ class GDPValResourcesServerConfig(BaseResourcesServerConfig):
 
 
 class GDPValVerifyRequest(BaseVerifyRequest):
+    rollout_id: Optional[str] = Field(default=None, exclude_if=lambda value: value is None)
     task_id: str
     sector: Optional[str] = None
     occupation: Optional[str] = None
@@ -324,7 +325,7 @@ def _effective_panel(self) -> List[JudgePanelMember]:
             JudgePanelMember(create_params_overrides=dict(self.config.judge_responses_create_params_overrides or {}))
         ]
 
-    def _resolve_judges(self) -> List[ResolvedJudge]:
+    def _resolve_judges(self, rollout_id: Optional[str] = None) -> List[ResolvedJudge]:
         """Resolve the (always non-empty) panel to concrete upstream coordinates.
 
         Every judge — including the single-judge special case (see
@@ -337,7 +338,7 @@ def _resolve_judges(self) -> List[ResolvedJudge]:
         legacy_overrides = dict(self.config.judge_responses_create_params_overrides or {})
 
         def _url(server: ModelServerRef) -> str:
-            return get_server_url(server.name) + "/v1"
+            return apply_rollout_prefix(get_server_url(server.name), rollout_id) + "/v1"
 
         judges: List[ResolvedJudge] = []
         for i, member in enumerate(self._effective_panel()):
@@ -373,7 +374,7 @@ async def _verify_rubric(self, body: GDPValVerifyRequest) -> GDPValVerifyRespons
                 invalid_judge_response=True,
             )
 
-        judges = self._resolve_judges()
+        judges = self._resolve_judges(body.rollout_id)
         # Route tasks with audio/video deliverables to the AV-capable judge(s) —
         # most judges can't read those modalities natively.
         if dir_contains_audio_video(body.deliverables_dir):
@@ -534,7 +535,7 @@ async def _verify_comparison(self, body: GDPValVerifyRequest) -> GDPValVerifyRes
         # Build the judge panel. Members may share a single proxy server (so we
         # reuse one OpenAI client per distinct upstream) and differ only by model
         # + reasoning settings. run_trials samples one member per trial.
-        resolved_judges = self._resolve_judges()
+        resolved_judges = self._resolve_judges(body.rollout_id)
         client_cache: Dict[tuple, Any] = {}
 
         def _client_for(judge: ResolvedJudge) -> Any:
diff --git a/resources_servers/gdpval/tests/test_app.py b/resources_servers/gdpval/tests/test_app.py
index dfbaa08475..ffdc6a4471 100644
--- a/resources_servers/gdpval/tests/test_app.py
+++ b/resources_servers/gdpval/tests/test_app.py
@@ -103,6 +103,9 @@ def test_missing_dir_returns_empty(self, tmp_path) -> None:
 
 
 class TestApp:
+    def test_rollout_id_is_absent_when_correlation_is_disabled(self) -> None:
+        assert "rollout_id" not in _verify_request().model_dump()
+
     def test_sanity_rubric(self) -> None:
         _server(reward_mode="rubric")
 
@@ -241,7 +244,7 @@ async def fake_score_with_rubric(**kwargs):
             captured.update(kwargs)
             return 0.5, {"overall_score": 0.5}
 
-        body = _verify_request(rubric_json=[{"criterion": "clarity", "score": 1}])
+        body = _verify_request(rubric_json=[{"criterion": "clarity", "score": 1}], rollout_id="7-3")
 
         with (
             patch("resources_servers.gdpval.scoring.score_with_rubric", side_effect=fake_score_with_rubric),
@@ -259,7 +262,7 @@ async def fake_score_with_rubric(**kwargs):
         assert judges[0].create_overrides == {"reasoning_effort": "medium"}
         assert judges[2].weight == 2.0
         # All share the single proxy base_url.
-        assert {j.base_url for j in judges} == {"http://localhost:9999/v1"}
+        assert {j.base_url for j in judges} == {"http://localhost:9999/ng-rollout/7-3/v1"}
         # A seeded rng is threaded through for reproducible sampling.
         assert captured["rng"] is not None
 
diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py
index 6d08dfaa86..8ca8aa632f 100644
--- a/responses_api_agents/stirrup_agent/app.py
+++ b/responses_api_agents/stirrup_agent/app.py
@@ -49,7 +49,7 @@
     NeMoGymResponseOutputMessage,
     NeMoGymResponseOutputText,
 )
-from nemo_gym.server_utils import get_response_json, raise_for_status
+from nemo_gym.server_utils import apply_rollout_prefix, get_response_json, raise_for_status
 from responses_api_agents.stirrup_agent.task_strategy import TaskSampleSkipError, TaskStrategy
 
 
@@ -936,7 +936,9 @@ class StirrupRunRequest(BaseRunRequest):
     "rubric_json",
     "rubric_pretty",
     "instance_id",
+    "_ng_task_index",
     "_ng_rollout_index",
+    "_ng_attempt_index",
 )
 
 
@@ -1036,20 +1038,21 @@ def model_post_init(self, __context: Any) -> None:
 
     # -- helpers ----------------------------------------------------------
 
-    def _get_model_base_url(self) -> str:
+    def _get_model_base_url(self, rollout_id: Optional[str] = None) -> str:
         from nemo_gym.global_config import get_first_server_config_dict
         from nemo_gym.server_utils import ServerClient
 
         global_config_dict = ServerClient.load_from_global_config().global_config_dict
         model_server_config = get_first_server_config_dict(global_config_dict, self.config.model_server.name)
-        return f"http://{model_server_config['host']}:{model_server_config['port']}/v1"
+        base_url = f"http://{model_server_config['host']}:{model_server_config['port']}"
+        return f"{apply_rollout_prefix(base_url, rollout_id)}/v1"
 
     # -- /v1/responses ----------------------------------------------------
 
     async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse:
         task_info = self.task_strategy.extract_task_info(body.metadata)
 
-        model_base_url = self._get_model_base_url()
+        model_base_url = self._get_model_base_url(self.rollout_id_from_run(body.metadata))
 
         if self.config.task == "gdpval":
             system_prompt = None
@@ -1180,10 +1183,16 @@ async def run(self, request: Request, body: StirrupRunRequest):
             for key in _TASK_METADATA_FIELDS:
                 top_value = body_dict.get(key)
                 meta_value = existing_metadata.get(key)
+                if key in ("_ng_task_index", "_ng_attempt_index") and not self._model_call_capture_enabled():
+                    continue
+                if key.startswith("_ng_") and top_value is not None:
+                    existing_metadata[key] = str(top_value)
+                    continue
                 if top_value is not None and meta_value is None:
                     existing_metadata[key] = top_value
                 elif meta_value is not None and top_value is None:
                     body_dict[key] = meta_value
+            rollout_id = self.rollout_id_from_run(body_dict)
             update: Dict[str, Any] = {"metadata": existing_metadata}
             if fixed_params.tool_choice is None:
                 update["tool_choice"] = "auto"
@@ -1388,6 +1397,8 @@ async def run(self, request: Request, body: StirrupRunRequest):
 
             verify_request_body = dict(body_dict)
             verify_request_body["response"] = response_clean.model_dump(mode="json")
+            if rollout_id is not None:
+                verify_request_body["rollout_id"] = rollout_id
             if deliverables_dir is not None:
                 verify_request_body["deliverables_dir"] = deliverables_dir
             # Surface the agent's runtime metadata for downstream logging.
diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py
index ad8ff6913a..d66c239720 100644
--- a/responses_api_agents/stirrup_agent/tests/test_app.py
+++ b/responses_api_agents/stirrup_agent/tests/test_app.py
@@ -125,6 +125,51 @@ def test_sanity(self) -> None:
         )
         StirrupAgentWrapper(config=config, server_client=MagicMock(spec=ServerClient))
 
+    def test_model_base_url_accepts_rollout_correlation(self) -> None:
+        wrapper = StirrupAgentWrapper(config=_make_config(), server_client=MagicMock(spec=ServerClient))
+        loaded = MagicMock(global_config_dict={"policy_model": {}})
+
+        with (
+            patch("nemo_gym.server_utils.ServerClient.load_from_global_config", return_value=loaded),
+            patch(
+                "nemo_gym.global_config.get_first_server_config_dict",
+                return_value={"host": "model-host", "port": 8000},
+            ),
+        ):
+            assert wrapper._get_model_base_url("7-3") == "http://model-host:8000/ng-rollout/7-3/v1"
+            assert wrapper._get_model_base_url() == "http://model-host:8000/v1"
+
+    @pytest.mark.asyncio
+    async def test_run_correlates_policy_and_judge_calls(self) -> None:
+        server_client = MagicMock(spec=ServerClient)
+        server_client.global_config_dict = {"observability_enabled": True}
+        server_client.post = AsyncMock(return_value=MagicMock())
+        wrapper = StirrupAgentWrapper(config=_make_config(), server_client=server_client)
+        body = StirrupRunRequest(
+            responses_create_params=NeMoGymResponseCreateParamsNonStreaming(
+                input="ignored",
+                metadata={"task_id": "task-1", "prompt": "do the thing", "_ng_rollout_index": "99"},
+            ),
+            task_id="task-1",
+            prompt="do the thing",
+            _ng_task_index=7,
+            _ng_rollout_index=3,
+        )
+        request = MagicMock(cookies={})
+        responses_mock = AsyncMock(return_value=_fake_response())
+
+        with (
+            patch.object(StirrupAgentWrapper, "responses", responses_mock),
+            patch("responses_api_agents.stirrup_agent.app.raise_for_status", AsyncMock()),
+            patch("responses_api_agents.stirrup_agent.app.get_response_json", AsyncMock(return_value={"reward": 1.0})),
+        ):
+            await wrapper.run(request, body)
+
+        policy_params = responses_mock.await_args.args[0]
+        assert wrapper.rollout_id_from_run(policy_params.metadata) == "7-3"
+        verify_calls = [call for call in server_client.post.await_args_list if call.kwargs["url_path"] == "/verify"]
+        assert verify_calls[0].kwargs["json"]["rollout_id"] == "7-3"
+
     def test_output_history_preserves_nemo_user_tool_results(self) -> None:
         """Run-history export should keep NeMo user-role tool results as tool outputs."""
         history = [
@@ -229,6 +274,7 @@ async def test_run_judge_only_scores_cached_deliverables(self, tmp_path) -> None
 
         config = _make_config(judge_only=True, persist_deliverables_dir=str(tmp_path))
         server_client = MagicMock(spec=ServerClient)
+        server_client.global_config_dict = {"observability_enabled": True}
         server_client.post = AsyncMock(return_value=MagicMock())
         wrapper = StirrupAgentWrapper(config=config, server_client=server_client)
 
@@ -236,7 +282,13 @@ async def test_run_judge_only_scores_cached_deliverables(self, tmp_path) -> None
             input="ignored",
             metadata={"task_id": "task-1", "prompt": "do the thing", "_ng_rollout_index": "0"},
         )
-        body = StirrupRunRequest(responses_create_params=params, task_id="task-1", prompt="do the thing")
+        body = StirrupRunRequest(
+            responses_create_params=params,
+            task_id="task-1",
+            prompt="do the thing",
+            _ng_task_index=7,
+            _ng_rollout_index=0,
+        )
         request = MagicMock()
         request.cookies = {}
 
@@ -258,6 +310,7 @@ async def test_run_judge_only_scores_cached_deliverables(self, tmp_path) -> None
         assert len(verify_calls) == 1
         verify_json = verify_calls[0].kwargs["json"]
         assert verify_json["deliverables_dir"].endswith(str(Path("task_task-1") / "repeat_0"))
+        assert verify_json["rollout_id"] == "7-0"
         assert result == {"reward": 0.9, "judge_response": "ok"}
 
     @pytest.mark.asyncio
diff --git a/responses_api_agents/tau2/tests/test_app.py b/responses_api_agents/tau2/tests/test_app.py
index 53dfd9afb0..be12e96b0f 100644
--- a/responses_api_agents/tau2/tests/test_app.py
+++ b/responses_api_agents/tau2/tests/test_app.py
@@ -17,7 +17,9 @@
 from typing import Tuple
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import pytest
 from fastapi.testclient import TestClient
+from tau2.data_model.simulation import RewardInfo, SimulationRun, TerminationReason
 
 from nemo_gym.base_responses_api_agent import AggregateMetricsRequest
 from nemo_gym.server_utils import ServerClient
@@ -114,6 +116,46 @@ def _clean(d):
 
         assert _clean(expected_response_dict) == _clean(actual_response_dict)
 
+    @pytest.mark.parametrize(
+        ("observability_enabled", "url_suffix"),
+        [(True, "/ng-rollout/7-2/v1"), (False, "/v1")],
+    )
+    def test_policy_and_user_model_calls_share_rollout_correlation(
+        self, observability_enabled: bool, url_suffix: str
+    ) -> None:
+        example_jsonl = Path(__file__).parent.parent / "data" / "example.jsonl"
+        request_body = json.loads(example_jsonl.read_text().splitlines()[0])
+        request_body |= {"_ng_task_index": 7, "_ng_rollout_index": 2}
+
+        config, server = self._dummy_server()
+        config.model_server.name = "policy"
+        config.user_model_server.name = "user"
+        server.server_client.global_config_dict = {"observability_enabled": observability_enabled}
+        with patch("responses_api_agents.tau2.app.ensure_tau2_data_dir"):
+            client = TestClient(server.setup_webserver())
+
+        result = SimulationRun(
+            id="run-1",
+            task_id="task-1",
+            start_time="2026-07-22T00:00:00Z",
+            end_time="2026-07-22T00:00:00Z",
+            duration=0,
+            termination_reason=TerminationReason.AGENT_STOP,
+            reward_info=RewardInfo(reward=1),
+            messages=[],
+        )
+        model_urls = {"policy": "http://policy:8000", "user": "http://user:8001"}
+        with (
+            patch("responses_api_agents.tau2.app.get_server_url", side_effect=model_urls.__getitem__),
+            patch("responses_api_agents.tau2.app.run_single_task", AsyncMock(return_value=result)),
+        ):
+            response = client.post("/run", json=request_body)
+
+        assert response.status_code == 200
+        response_config = response.json()["config"]
+        assert response_config["llm_args_agent"]["api_base"] == model_urls["policy"] + url_suffix
+        assert response_config["llm_args_user"]["api_base"] == model_urls["user"] + url_suffix
+
     async def test_compute_metrics(self) -> None:
         example_rollouts_fpath = Path(__file__).parent.parent / "data" / "example_rollouts.jsonl"
         with example_rollouts_fpath.open() as f:

From 72799cef371e977a4c084aa0d7e984afa0a31b9a Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Fri, 24 Jul 2026 18:32:14 +0200
Subject: [PATCH 7/8] Complete model call capture coverage

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 nemo_gym/base_responses_api_model.py          | 177 +++++++++++++++---
 resources_servers/gdpval/app.py               |  12 +-
 resources_servers/gdpval/tests/test_app.py    |   7 +-
 responses_api_agents/stirrup_agent/app.py     |  17 +-
 .../stirrup_agent/tests/test_app.py           |  36 +---
 .../test_responses_api_model_messages.py      | 168 ++++++++++++++++-
 6 files changed, 333 insertions(+), 84 deletions(-)

diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py
index 377de77daa..33ce6d528f 100644
--- a/nemo_gym/base_responses_api_model.py
+++ b/nemo_gym/base_responses_api_model.py
@@ -36,7 +36,7 @@
 import time
 from abc import abstractmethod
 from pathlib import Path
-from typing import Any, Mapping, Optional
+from typing import Any, Iterator, Mapping, Optional
 from uuid import uuid4
 
 import orjson
@@ -73,6 +73,67 @@
 
 # Stateless; shared by every model server's default /v1/messages handler.
 _ANTHROPIC_CONVERTER = AnthropicConverter()
+_NORMALIZED_CAPTURE_RESPONSE_SCOPE_KEY = "nemo_gym.normalized_capture_response"
+
+
+class _ChatCompletionsWireParams(NeMoGymChatCompletionCreateParamsNonStreaming):
+    """Wire shape accepted by the shared Chat Completions route."""
+
+    stream: Optional[bool] = None
+
+
+def _chat_completion_to_sse(response: NeMoGymChatCompletion, *, include_usage: bool) -> Iterator[str]:
+    """Emit a completed Chat Completions response as protocol-compatible SSE chunks."""
+    data = response.model_dump(mode="json", exclude_none=True)
+    chunk_base = {
+        "id": data["id"],
+        "object": "chat.completion.chunk",
+        "created": data["created"],
+        "model": data["model"],
+    }
+    for key in ("service_tier", "system_fingerprint"):
+        if key in data:
+            chunk_base[key] = data[key]
+
+    def _event(**fields: Any) -> str:
+        return f"data: {orjson.dumps({**chunk_base, **fields}).decode()}\n\n"
+
+    for choice in data["choices"]:
+        message = choice["message"]
+        delta = {"role": message["role"]}
+        for key in ("content", "refusal", "reasoning_content", "reasoning"):
+            if key in message:
+                delta[key] = message[key]
+        if message.get("tool_calls"):
+            delta["tool_calls"] = [
+                {**tool_call, "index": index} for index, tool_call in enumerate(message["tool_calls"])
+            ]
+
+        index = choice["index"]
+        yield _event(
+            choices=[
+                {
+                    "index": index,
+                    "delta": delta,
+                    "finish_reason": None,
+                    "logprobs": choice.get("logprobs"),
+                }
+            ]
+        )
+        yield _event(
+            choices=[
+                {
+                    "index": index,
+                    "delta": {},
+                    "finish_reason": choice["finish_reason"],
+                    "logprobs": None,
+                }
+            ]
+        )
+
+    if include_usage:
+        yield _event(choices=[], usage=data.get("usage"))
+    yield "data: [DONE]\n\n"
 
 
 class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig):
@@ -91,7 +152,7 @@ def setup_webserver(self) -> FastAPI:
         capture_config = ModelCallCaptureConfig.model_validate(self.server_client.global_config_dict)
         install_model_call_capture(app, capture_config, model_server_name=self.config.name)
 
-        app.post("/v1/chat/completions")(self.chat_completions)
+        app.post("/v1/chat/completions")(self._chat_completions_wire)
 
         app.post("/v1/responses")(self.responses_dispatch)
 
@@ -183,6 +244,33 @@ async def _invoke_responses(
             return await self.responses(request=request, body=params)
         return await self.responses(body=params)
 
+    async def _chat_completions_wire(self, request: Request, body: _ChatCompletionsWireParams = Body()) -> Any:
+        payload = body.model_dump(exclude_unset=True)
+        streaming = body.stream is True
+        if streaming:
+            payload["stream"] = False
+            payload.pop("stream_options", None)
+        params = NeMoGymChatCompletionCreateParamsNonStreaming.model_validate(payload)
+        response = await self._invoke_chat_completions(request, params)
+        if not streaming:
+            return response
+
+        # The synthetic stream may omit usage at the caller's request. Keep the complete normalized
+        # response in this request's ASGI scope so capture does not lose data that the model returned.
+        request.scope[_NORMALIZED_CAPTURE_RESPONSE_SCOPE_KEY] = response.model_dump(mode="json", exclude_none=True)
+        stream_options = body.stream_options or {}
+        return StreamingResponse(
+            _chat_completion_to_sse(response, include_usage=bool(stream_options.get("include_usage"))),
+            media_type="text/event-stream",
+        )
+
+    async def _invoke_chat_completions(
+        self, request: Request, params: NeMoGymChatCompletionCreateParamsNonStreaming
+    ) -> NeMoGymChatCompletion:
+        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)
+
 
 def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming:
     """Validate a /v1/responses body dict, surfacing failures as FastAPI's standard 422."""
@@ -787,34 +875,49 @@ def _reconstruct_anthropic_sse(events: list[dict[str, Any]]) -> Optional[dict[st
 
 def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, Any]]:
     """Rebuild a Chat Completions response from streamed chunks."""
-    content_parts: list[str] = []
-    reasoning_parts: list[str] = []
-    tool_calls: dict[int, dict[str, Any]] = {}
+    choices: dict[int, dict[str, Any]] = {}
     usage: Optional[dict[str, Any]] = None
     model: Optional[str] = None
     response_id: Optional[str] = None
-    role = "assistant"
-    finish_reason: Optional[str] = None
-    saw_choice = False
+    created: Optional[int] = None
+    service_tier: Optional[str] = None
+    system_fingerprint: Optional[str] = None
     for chunk in events:
         model = chunk.get("model") or model
         if isinstance(chunk.get("id"), str):
             response_id = chunk["id"]
+        if isinstance(chunk.get("created"), int):
+            created = chunk["created"]
+        service_tier = chunk.get("service_tier") or service_tier
+        system_fingerprint = chunk.get("system_fingerprint") or system_fingerprint
         if chunk.get("usage"):
             usage = chunk["usage"]
         for choice in chunk.get("choices") or []:
             if not isinstance(choice, dict):
                 continue
-            saw_choice = True
+            index = choice.get("index", 0)
+            state = choices.setdefault(
+                index,
+                {
+                    "role": "assistant",
+                    "content": [],
+                    "reasoning_content": [],
+                    "tool_calls": {},
+                    "finish_reason": None,
+                    "logprobs": None,
+                },
+            )
             delta = choice.get("delta") or {}
-            role = delta.get("role") or role
-            if delta.get("content"):
-                content_parts.append(delta["content"])
+            state["role"] = delta.get("role") or state["role"]
+            if isinstance(delta.get("content"), str):
+                state["content"].append(delta["content"])
             reasoning = delta.get("reasoning_content") or delta.get("reasoning")
-            if reasoning:
-                reasoning_parts.append(reasoning)
+            if isinstance(reasoning, str):
+                state["reasoning_content"].append(reasoning)
+            if "refusal" in delta:
+                state["refusal"] = delta["refusal"]
             for tc in delta.get("tool_calls") or []:
-                slot = tool_calls.setdefault(
+                slot = state["tool_calls"].setdefault(
                     tc.get("index", 0), {"id": None, "type": "function", "function": {"name": "", "arguments": ""}}
                 )
                 if tc.get("id"):
@@ -825,21 +928,44 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An
                 if fn.get("arguments"):
                     slot["function"]["arguments"] += fn["arguments"]
             if choice.get("finish_reason"):
-                finish_reason = choice["finish_reason"]
-    if not saw_choice:
+                state["finish_reason"] = choice["finish_reason"]
+            if choice.get("logprobs") is not None:
+                state["logprobs"] = choice["logprobs"]
+    if not choices:
         return None
-    message: dict[str, Any] = {"role": role, "content": "".join(content_parts) or None}
-    if reasoning_parts:
-        message["reasoning_content"] = "".join(reasoning_parts)
-    if tool_calls:
-        message["tool_calls"] = [tool_calls[i] for i in sorted(tool_calls)]
+
+    reconstructed_choices = []
+    for index in sorted(choices):
+        state = choices[index]
+        message: dict[str, Any] = {"role": state["role"], "content": "".join(state["content"]) or None}
+        if state["reasoning_content"]:
+            message["reasoning_content"] = "".join(state["reasoning_content"])
+        if "refusal" in state:
+            message["refusal"] = state["refusal"]
+        if state["tool_calls"]:
+            message["tool_calls"] = [state["tool_calls"][i] for i in sorted(state["tool_calls"])]
+        reconstructed_choices.append(
+            {
+                "index": index,
+                "message": message,
+                "finish_reason": state["finish_reason"],
+                "logprobs": state["logprobs"],
+            }
+        )
+
     result: dict[str, Any] = {
         "object": "chat.completion",
         "model": model,
-        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
+        "choices": reconstructed_choices,
     }
     if response_id is not None:
         result["id"] = response_id
+    if created is not None:
+        result["created"] = created
+    if service_tier is not None:
+        result["service_tier"] = service_tier
+    if system_fingerprint is not None:
+        result["system_fingerprint"] = system_fingerprint
     if usage:
         result["usage"] = usage
     return result
@@ -1071,8 +1197,9 @@ async def _flush_deferred_response() -> None:
         def _parse_and_record() -> None:
             # Off the event loop: body parse + SSE reassembly is best-effort and fully guarded, so a
             # malformed body can never surface as an ASGI error after the response was already sent.
-            response_body = None
-            if body_bytes:
+            normalized_response = scope.get(_NORMALIZED_CAPTURE_RESPONSE_SCOPE_KEY)
+            response_body = normalized_response if isinstance(normalized_response, dict) else None
+            if response_body is None and body_bytes:
                 try:
                     response_body = (
                         _reconstruct_streamed_response(body_bytes, dialect) if streaming else json.loads(body_bytes)
diff --git a/resources_servers/gdpval/app.py b/resources_servers/gdpval/app.py
index 111948c53c..65fdf53ae7 100644
--- a/resources_servers/gdpval/app.py
+++ b/resources_servers/gdpval/app.py
@@ -41,7 +41,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 from nemo_gym.base_resources_server import (
     BaseResourcesServerConfig,
@@ -50,6 +50,7 @@
     SimpleResourcesServer,
 )
 from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest, ModelServerRef
+from nemo_gym.rollout_correlation import current_rollout_id
 from nemo_gym.server_utils import apply_rollout_prefix, get_server_url
 from resources_servers.gdpval.judge_panel import (
     ResolvedJudge,
@@ -234,7 +235,6 @@ class GDPValResourcesServerConfig(BaseResourcesServerConfig):
 
 
 class GDPValVerifyRequest(BaseVerifyRequest):
-    rollout_id: Optional[str] = Field(default=None, exclude_if=lambda value: value is None)
     task_id: str
     sector: Optional[str] = None
     occupation: Optional[str] = None
@@ -325,7 +325,7 @@ def _effective_panel(self) -> List[JudgePanelMember]:
             JudgePanelMember(create_params_overrides=dict(self.config.judge_responses_create_params_overrides or {}))
         ]
 
-    def _resolve_judges(self, rollout_id: Optional[str] = None) -> List[ResolvedJudge]:
+    def _resolve_judges(self) -> List[ResolvedJudge]:
         """Resolve the (always non-empty) panel to concrete upstream coordinates.
 
         Every judge — including the single-judge special case (see
@@ -338,7 +338,7 @@ def _resolve_judges(self, rollout_id: Optional[str] = None) -> List[ResolvedJudg
         legacy_overrides = dict(self.config.judge_responses_create_params_overrides or {})
 
         def _url(server: ModelServerRef) -> str:
-            return apply_rollout_prefix(get_server_url(server.name), rollout_id) + "/v1"
+            return apply_rollout_prefix(get_server_url(server.name), current_rollout_id()) + "/v1"
 
         judges: List[ResolvedJudge] = []
         for i, member in enumerate(self._effective_panel()):
@@ -374,7 +374,7 @@ async def _verify_rubric(self, body: GDPValVerifyRequest) -> GDPValVerifyRespons
                 invalid_judge_response=True,
             )
 
-        judges = self._resolve_judges(body.rollout_id)
+        judges = self._resolve_judges()
         # Route tasks with audio/video deliverables to the AV-capable judge(s) —
         # most judges can't read those modalities natively.
         if dir_contains_audio_video(body.deliverables_dir):
@@ -535,7 +535,7 @@ async def _verify_comparison(self, body: GDPValVerifyRequest) -> GDPValVerifyRes
         # Build the judge panel. Members may share a single proxy server (so we
         # reuse one OpenAI client per distinct upstream) and differ only by model
         # + reasoning settings. run_trials samples one member per trial.
-        resolved_judges = self._resolve_judges(body.rollout_id)
+        resolved_judges = self._resolve_judges()
         client_cache: Dict[tuple, Any] = {}
 
         def _client_for(judge: ResolvedJudge) -> Any:
diff --git a/resources_servers/gdpval/tests/test_app.py b/resources_servers/gdpval/tests/test_app.py
index ffdc6a4471..a97bc99d48 100644
--- a/resources_servers/gdpval/tests/test_app.py
+++ b/resources_servers/gdpval/tests/test_app.py
@@ -22,6 +22,7 @@
     NeMoGymResponseOutputMessage,
     NeMoGymResponseOutputText,
 )
+from nemo_gym.rollout_correlation import rollout_context
 from nemo_gym.server_utils import ServerClient
 from resources_servers.gdpval.app import (
     GDPValResourcesServer,
@@ -103,9 +104,6 @@ def test_missing_dir_returns_empty(self, tmp_path) -> None:
 
 
 class TestApp:
-    def test_rollout_id_is_absent_when_correlation_is_disabled(self) -> None:
-        assert "rollout_id" not in _verify_request().model_dump()
-
     def test_sanity_rubric(self) -> None:
         _server(reward_mode="rubric")
 
@@ -244,9 +242,10 @@ async def fake_score_with_rubric(**kwargs):
             captured.update(kwargs)
             return 0.5, {"overall_score": 0.5}
 
-        body = _verify_request(rubric_json=[{"criterion": "clarity", "score": 1}], rollout_id="7-3")
+        body = _verify_request(rubric_json=[{"criterion": "clarity", "score": 1}])
 
         with (
+            rollout_context("7-3"),
             patch("resources_servers.gdpval.scoring.score_with_rubric", side_effect=fake_score_with_rubric),
             patch("resources_servers.gdpval.app.get_server_url", return_value="http://localhost:9999"),
         ):
diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py
index 8ca8aa632f..e3b959deae 100644
--- a/responses_api_agents/stirrup_agent/app.py
+++ b/responses_api_agents/stirrup_agent/app.py
@@ -49,6 +49,7 @@
     NeMoGymResponseOutputMessage,
     NeMoGymResponseOutputText,
 )
+from nemo_gym.rollout_correlation import current_rollout_id
 from nemo_gym.server_utils import apply_rollout_prefix, get_response_json, raise_for_status
 from responses_api_agents.stirrup_agent.task_strategy import TaskSampleSkipError, TaskStrategy
 
@@ -936,9 +937,7 @@ class StirrupRunRequest(BaseRunRequest):
     "rubric_json",
     "rubric_pretty",
     "instance_id",
-    "_ng_task_index",
     "_ng_rollout_index",
-    "_ng_attempt_index",
 )
 
 
@@ -1038,21 +1037,21 @@ def model_post_init(self, __context: Any) -> None:
 
     # -- helpers ----------------------------------------------------------
 
-    def _get_model_base_url(self, rollout_id: Optional[str] = None) -> str:
+    def _get_model_base_url(self) -> str:
         from nemo_gym.global_config import get_first_server_config_dict
         from nemo_gym.server_utils import ServerClient
 
         global_config_dict = ServerClient.load_from_global_config().global_config_dict
         model_server_config = get_first_server_config_dict(global_config_dict, self.config.model_server.name)
         base_url = f"http://{model_server_config['host']}:{model_server_config['port']}"
-        return f"{apply_rollout_prefix(base_url, rollout_id)}/v1"
+        return f"{apply_rollout_prefix(base_url, current_rollout_id())}/v1"
 
     # -- /v1/responses ----------------------------------------------------
 
     async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse:
         task_info = self.task_strategy.extract_task_info(body.metadata)
 
-        model_base_url = self._get_model_base_url(self.rollout_id_from_run(body.metadata))
+        model_base_url = self._get_model_base_url()
 
         if self.config.task == "gdpval":
             system_prompt = None
@@ -1183,16 +1182,10 @@ async def run(self, request: Request, body: StirrupRunRequest):
             for key in _TASK_METADATA_FIELDS:
                 top_value = body_dict.get(key)
                 meta_value = existing_metadata.get(key)
-                if key in ("_ng_task_index", "_ng_attempt_index") and not self._model_call_capture_enabled():
-                    continue
-                if key.startswith("_ng_") and top_value is not None:
-                    existing_metadata[key] = str(top_value)
-                    continue
                 if top_value is not None and meta_value is None:
                     existing_metadata[key] = top_value
                 elif meta_value is not None and top_value is None:
                     body_dict[key] = meta_value
-            rollout_id = self.rollout_id_from_run(body_dict)
             update: Dict[str, Any] = {"metadata": existing_metadata}
             if fixed_params.tool_choice is None:
                 update["tool_choice"] = "auto"
@@ -1397,8 +1390,6 @@ async def run(self, request: Request, body: StirrupRunRequest):
 
             verify_request_body = dict(body_dict)
             verify_request_body["response"] = response_clean.model_dump(mode="json")
-            if rollout_id is not None:
-                verify_request_body["rollout_id"] = rollout_id
             if deliverables_dir is not None:
                 verify_request_body["deliverables_dir"] = deliverables_dir
             # Surface the agent's runtime metadata for downstream logging.
diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py
index d66c239720..4283fb9bd5 100644
--- a/responses_api_agents/stirrup_agent/tests/test_app.py
+++ b/responses_api_agents/stirrup_agent/tests/test_app.py
@@ -25,6 +25,7 @@
     NeMoGymResponseOutputMessage,
     NeMoGymResponseOutputText,
 )
+from nemo_gym.rollout_correlation import rollout_context
 from nemo_gym.server_utils import ServerClient
 from responses_api_agents.stirrup_agent.app import (
     NG_FAILURE_CLASS_KEY,
@@ -136,40 +137,10 @@ def test_model_base_url_accepts_rollout_correlation(self) -> None:
                 return_value={"host": "model-host", "port": 8000},
             ),
         ):
-            assert wrapper._get_model_base_url("7-3") == "http://model-host:8000/ng-rollout/7-3/v1"
+            with rollout_context("7-3"):
+                assert wrapper._get_model_base_url() == "http://model-host:8000/ng-rollout/7-3/v1"
             assert wrapper._get_model_base_url() == "http://model-host:8000/v1"
 
-    @pytest.mark.asyncio
-    async def test_run_correlates_policy_and_judge_calls(self) -> None:
-        server_client = MagicMock(spec=ServerClient)
-        server_client.global_config_dict = {"observability_enabled": True}
-        server_client.post = AsyncMock(return_value=MagicMock())
-        wrapper = StirrupAgentWrapper(config=_make_config(), server_client=server_client)
-        body = StirrupRunRequest(
-            responses_create_params=NeMoGymResponseCreateParamsNonStreaming(
-                input="ignored",
-                metadata={"task_id": "task-1", "prompt": "do the thing", "_ng_rollout_index": "99"},
-            ),
-            task_id="task-1",
-            prompt="do the thing",
-            _ng_task_index=7,
-            _ng_rollout_index=3,
-        )
-        request = MagicMock(cookies={})
-        responses_mock = AsyncMock(return_value=_fake_response())
-
-        with (
-            patch.object(StirrupAgentWrapper, "responses", responses_mock),
-            patch("responses_api_agents.stirrup_agent.app.raise_for_status", AsyncMock()),
-            patch("responses_api_agents.stirrup_agent.app.get_response_json", AsyncMock(return_value={"reward": 1.0})),
-        ):
-            await wrapper.run(request, body)
-
-        policy_params = responses_mock.await_args.args[0]
-        assert wrapper.rollout_id_from_run(policy_params.metadata) == "7-3"
-        verify_calls = [call for call in server_client.post.await_args_list if call.kwargs["url_path"] == "/verify"]
-        assert verify_calls[0].kwargs["json"]["rollout_id"] == "7-3"
-
     def test_output_history_preserves_nemo_user_tool_results(self) -> None:
         """Run-history export should keep NeMo user-role tool results as tool outputs."""
         history = [
@@ -310,7 +281,6 @@ async def test_run_judge_only_scores_cached_deliverables(self, tmp_path) -> None
         assert len(verify_calls) == 1
         verify_json = verify_calls[0].kwargs["json"]
         assert verify_json["deliverables_dir"].endswith(str(Path("task_task-1") / "repeat_0"))
-        assert verify_json["rollout_id"] == "7-0"
         assert result == {"reward": 0.9, "judge_response": "ok"}
 
     @pytest.mark.asyncio
diff --git a/tests/unit_tests/test_responses_api_model_messages.py b/tests/unit_tests/test_responses_api_model_messages.py
index 0b7304f495..7e6129ae0c 100644
--- a/tests/unit_tests/test_responses_api_model_messages.py
+++ b/tests/unit_tests/test_responses_api_model_messages.py
@@ -19,14 +19,22 @@
 default mapping for both ``responses()`` signatures (with and without a leading ``request``).
 """
 
+import json
 from time import time
 from unittest.mock import MagicMock
 from uuid import uuid4
 
+import pytest
+from aiohttp import ClientResponseError
 from fastapi import Body, Request
 from fastapi.testclient import TestClient
 
-from nemo_gym.base_responses_api_model import BaseResponsesAPIModelConfig, SimpleResponsesAPIModel
+from nemo_gym.base_responses_api_model import (
+    BaseResponsesAPIModelConfig,
+    CaptureStore,
+    SimpleResponsesAPIModel,
+    _reconstruct_chat_sse,
+)
 from nemo_gym.openai_utils import (
     NeMoGymChatCompletion,
     NeMoGymChatCompletionCreateParamsNonStreaming,
@@ -57,11 +65,49 @@ def _build_response(text: str, model: str = "downstream-model") -> NeMoGymRespon
     )
 
 
+def _build_chat_completion() -> NeMoGymChatCompletion:
+    return NeMoGymChatCompletion.model_validate(
+        {
+            "id": "chatcmpl_test",
+            "created": 123,
+            "model": "downstream-model",
+            "object": "chat.completion",
+            "service_tier": "default",
+            "system_fingerprint": "fp_test",
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "tool_calls",
+                    "message": {
+                        "role": "assistant",
+                        "content": "first",
+                        "reasoning_content": "thinking",
+                        "tool_calls": [
+                            {
+                                "id": "call_1",
+                                "type": "function",
+                                "function": {"name": "lookup", "arguments": '{"id":1}'},
+                            }
+                        ],
+                    },
+                },
+                {
+                    "index": 1,
+                    "finish_reason": "stop",
+                    "message": {"role": "assistant", "content": "second"},
+                },
+            ],
+            "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18},
+        }
+    )
+
+
 class _BodyOnlyModel(SimpleResponsesAPIModel):
     """A server whose responses() takes only `body` (like openai_model)."""
 
     config: BaseResponsesAPIModelConfig
     last_params: object = None
+    last_chat_params: object = None
     model_config = {"arbitrary_types_allowed": True}
 
     async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse:
@@ -71,7 +117,8 @@ async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()
     async def chat_completions(
         self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
     ) -> NeMoGymChatCompletion:
-        raise NotImplementedError
+        object.__setattr__(self, "last_chat_params", body)
+        return _build_chat_completion()
 
 
 class _RequestAwareModel(SimpleResponsesAPIModel):
@@ -79,6 +126,8 @@ class _RequestAwareModel(SimpleResponsesAPIModel):
 
     config: BaseResponsesAPIModelConfig
     saw_request: bool = False
+    saw_chat_request: bool = False
+    last_chat_params: object = None
     model_config = {"arbitrary_types_allowed": True}
 
     async def responses(
@@ -87,10 +136,26 @@ async def responses(
         object.__setattr__(self, "saw_request", isinstance(request, Request))
         return _build_response("hi from request-aware")
 
+    async def chat_completions(
+        self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
+    ) -> NeMoGymChatCompletion:
+        object.__setattr__(self, "saw_chat_request", isinstance(request, Request))
+        object.__setattr__(self, "last_chat_params", body)
+        return _build_chat_completion()
+
+
+class _FailingModel(_BodyOnlyModel):
     async def chat_completions(
         self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
     ) -> NeMoGymChatCompletion:
-        raise NotImplementedError
+        error = ClientResponseError(
+            request_info=MagicMock(real_url="https://example.invalid/v1/chat/completions"),
+            history=(),
+            status=403,
+            message="Forbidden",
+        )
+        error.response_content = b'{"error":{"message":"denied"}}'
+        raise error
 
 
 def _config() -> BaseResponsesAPIModelConfig:
@@ -102,6 +167,14 @@ def _client(model_cls) -> TestClient:
     return TestClient(server.setup_webserver()), server
 
 
+def _sse_payloads(body: str) -> list[dict]:
+    return [
+        json.loads(line.removeprefix("data: "))
+        for line in body.splitlines()
+        if line.startswith("data: ") and line != "data: [DONE]"
+    ]
+
+
 class TestDefaultMessagesRoute:
     def test_messages_route_registered_alongside_openai_routes(self) -> None:
         server = _BodyOnlyModel(config=_config(), server_client=MagicMock(spec=ServerClient, global_config_dict={}))
@@ -150,3 +223,92 @@ def test_streaming_returns_anthropic_sse(self) -> None:
         assert "event: message_start" in body
         assert "event: content_block_delta" in body
         assert "event: message_stop" in body
+
+
+@pytest.mark.parametrize("model_cls", [_BodyOnlyModel, _RequestAwareModel])
+def test_chat_completions_streaming_wraps_existing_nonstreaming_handler(model_cls) -> None:
+    client, server = _client(model_cls)
+    response = client.post(
+        "/v1/chat/completions",
+        json={
+            "model": "requested-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("text/event-stream")
+    assert response.text.endswith("data: [DONE]\n\n")
+    assert server.last_chat_params.stream is False
+    assert server.last_chat_params.stream_options is None
+    if model_cls is _RequestAwareModel:
+        assert server.saw_chat_request is True
+
+    payloads = _sse_payloads(response.text)
+    reconstructed = NeMoGymChatCompletion.model_validate(_reconstruct_chat_sse(payloads))
+    assert reconstructed.model_dump(mode="json") == _build_chat_completion().model_dump(mode="json")
+
+
+def test_chat_completions_nonstreaming_is_unchanged() -> None:
+    client, server = _client(_BodyOnlyModel)
+    response = client.post(
+        "/v1/chat/completions",
+        json={"messages": [{"role": "user", "content": "hello"}]},
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("application/json")
+    assert response.json()["object"] == "chat.completion"
+    assert server.last_chat_params.stream is None
+
+
+def test_streaming_chat_completion_capture_keeps_wire_request_and_reconstructs_response(tmp_path) -> None:
+    server_client = MagicMock(spec=ServerClient)
+    server_client.global_config_dict = {
+        "observability_enabled": True,
+        "model_call_capture_dir": str(tmp_path),
+    }
+    server = _BodyOnlyModel(config=_config(), server_client=server_client)
+    client = TestClient(server.setup_webserver())
+
+    response = client.post(
+        "/ng-rollout/rollout-1/v1/chat/completions",
+        json={
+            "messages": [{"role": "user", "content": "hello"}],
+            "stream": True,
+        },
+    )
+
+    assert response.status_code == 200
+    exchange = CaptureStore(tmp_path).read("rollout-1")[0]
+    assert exchange["status_code"] == 200
+    assert exchange["error_category"] is None
+    assert exchange["request"]["stream"] is True
+    assert "stream_options" not in exchange["request"]
+    reconstructed = NeMoGymChatCompletion.model_validate(exchange["response"])
+    assert reconstructed.model_dump(mode="json") == _build_chat_completion().model_dump(mode="json")
+    assert all("usage" not in payload for payload in _sse_payloads(response.text))
+    assert exchange["response_raw"].endswith("data: [DONE]\n\n")
+
+
+def test_capture_preserves_upstream_error_status_and_body(tmp_path) -> None:
+    server_client = MagicMock(spec=ServerClient)
+    server_client.global_config_dict = {
+        "observability_enabled": True,
+        "model_call_capture_dir": str(tmp_path),
+    }
+    server = _FailingModel(config=_config(), server_client=server_client)
+    client = TestClient(server.setup_webserver(), raise_server_exceptions=False)
+
+    response = client.post(
+        "/ng-rollout/rollout-1/v1/chat/completions",
+        json={"messages": [{"role": "user", "content": "hello"}]},
+    )
+
+    assert response.status_code == 500
+    exchange = CaptureStore(tmp_path).read("rollout-1")[0]
+    assert exchange["status_code"] == 403
+    assert exchange["error_category"] == "auth"
+    assert json.loads(exchange["response_raw"]) == {"error": {"message": "denied"}}

From 0ca9488d79a2ea3a420c7b3752c651a18c22a141 Mon Sep 17 00:00:00 2001
From: Michal Bien <mbien@nvidia.com>
Date: Mon, 27 Jul 2026 18:34:11 +0200
Subject: [PATCH 8/8] Tighten auxiliary model correlation

Signed-off-by: Michal Bien <mbien@nvidia.com>
---
 nemo_gym/base_responses_api_model.py          | 177 +++---------------
 responses_api_agents/stirrup_agent/app.py     |  10 +-
 .../stirrup_agent/tests/test_app.py           |  20 +-
 responses_api_agents/tau2/tests/test_app.py   |  32 ++--
 .../test_responses_api_model_messages.py      | 168 +----------------
 5 files changed, 51 insertions(+), 356 deletions(-)

diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py
index 33ce6d528f..377de77daa 100644
--- a/nemo_gym/base_responses_api_model.py
+++ b/nemo_gym/base_responses_api_model.py
@@ -36,7 +36,7 @@
 import time
 from abc import abstractmethod
 from pathlib import Path
-from typing import Any, Iterator, Mapping, Optional
+from typing import Any, Mapping, Optional
 from uuid import uuid4
 
 import orjson
@@ -73,67 +73,6 @@
 
 # Stateless; shared by every model server's default /v1/messages handler.
 _ANTHROPIC_CONVERTER = AnthropicConverter()
-_NORMALIZED_CAPTURE_RESPONSE_SCOPE_KEY = "nemo_gym.normalized_capture_response"
-
-
-class _ChatCompletionsWireParams(NeMoGymChatCompletionCreateParamsNonStreaming):
-    """Wire shape accepted by the shared Chat Completions route."""
-
-    stream: Optional[bool] = None
-
-
-def _chat_completion_to_sse(response: NeMoGymChatCompletion, *, include_usage: bool) -> Iterator[str]:
-    """Emit a completed Chat Completions response as protocol-compatible SSE chunks."""
-    data = response.model_dump(mode="json", exclude_none=True)
-    chunk_base = {
-        "id": data["id"],
-        "object": "chat.completion.chunk",
-        "created": data["created"],
-        "model": data["model"],
-    }
-    for key in ("service_tier", "system_fingerprint"):
-        if key in data:
-            chunk_base[key] = data[key]
-
-    def _event(**fields: Any) -> str:
-        return f"data: {orjson.dumps({**chunk_base, **fields}).decode()}\n\n"
-
-    for choice in data["choices"]:
-        message = choice["message"]
-        delta = {"role": message["role"]}
-        for key in ("content", "refusal", "reasoning_content", "reasoning"):
-            if key in message:
-                delta[key] = message[key]
-        if message.get("tool_calls"):
-            delta["tool_calls"] = [
-                {**tool_call, "index": index} for index, tool_call in enumerate(message["tool_calls"])
-            ]
-
-        index = choice["index"]
-        yield _event(
-            choices=[
-                {
-                    "index": index,
-                    "delta": delta,
-                    "finish_reason": None,
-                    "logprobs": choice.get("logprobs"),
-                }
-            ]
-        )
-        yield _event(
-            choices=[
-                {
-                    "index": index,
-                    "delta": {},
-                    "finish_reason": choice["finish_reason"],
-                    "logprobs": None,
-                }
-            ]
-        )
-
-    if include_usage:
-        yield _event(choices=[], usage=data.get("usage"))
-    yield "data: [DONE]\n\n"
 
 
 class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig):
@@ -152,7 +91,7 @@ def setup_webserver(self) -> FastAPI:
         capture_config = ModelCallCaptureConfig.model_validate(self.server_client.global_config_dict)
         install_model_call_capture(app, capture_config, model_server_name=self.config.name)
 
-        app.post("/v1/chat/completions")(self._chat_completions_wire)
+        app.post("/v1/chat/completions")(self.chat_completions)
 
         app.post("/v1/responses")(self.responses_dispatch)
 
@@ -244,33 +183,6 @@ async def _invoke_responses(
             return await self.responses(request=request, body=params)
         return await self.responses(body=params)
 
-    async def _chat_completions_wire(self, request: Request, body: _ChatCompletionsWireParams = Body()) -> Any:
-        payload = body.model_dump(exclude_unset=True)
-        streaming = body.stream is True
-        if streaming:
-            payload["stream"] = False
-            payload.pop("stream_options", None)
-        params = NeMoGymChatCompletionCreateParamsNonStreaming.model_validate(payload)
-        response = await self._invoke_chat_completions(request, params)
-        if not streaming:
-            return response
-
-        # The synthetic stream may omit usage at the caller's request. Keep the complete normalized
-        # response in this request's ASGI scope so capture does not lose data that the model returned.
-        request.scope[_NORMALIZED_CAPTURE_RESPONSE_SCOPE_KEY] = response.model_dump(mode="json", exclude_none=True)
-        stream_options = body.stream_options or {}
-        return StreamingResponse(
-            _chat_completion_to_sse(response, include_usage=bool(stream_options.get("include_usage"))),
-            media_type="text/event-stream",
-        )
-
-    async def _invoke_chat_completions(
-        self, request: Request, params: NeMoGymChatCompletionCreateParamsNonStreaming
-    ) -> NeMoGymChatCompletion:
-        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)
-
 
 def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming:
     """Validate a /v1/responses body dict, surfacing failures as FastAPI's standard 422."""
@@ -875,49 +787,34 @@ def _reconstruct_anthropic_sse(events: list[dict[str, Any]]) -> Optional[dict[st
 
 def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, Any]]:
     """Rebuild a Chat Completions response from streamed chunks."""
-    choices: dict[int, dict[str, Any]] = {}
+    content_parts: list[str] = []
+    reasoning_parts: list[str] = []
+    tool_calls: dict[int, dict[str, Any]] = {}
     usage: Optional[dict[str, Any]] = None
     model: Optional[str] = None
     response_id: Optional[str] = None
-    created: Optional[int] = None
-    service_tier: Optional[str] = None
-    system_fingerprint: Optional[str] = None
+    role = "assistant"
+    finish_reason: Optional[str] = None
+    saw_choice = False
     for chunk in events:
         model = chunk.get("model") or model
         if isinstance(chunk.get("id"), str):
             response_id = chunk["id"]
-        if isinstance(chunk.get("created"), int):
-            created = chunk["created"]
-        service_tier = chunk.get("service_tier") or service_tier
-        system_fingerprint = chunk.get("system_fingerprint") or system_fingerprint
         if chunk.get("usage"):
             usage = chunk["usage"]
         for choice in chunk.get("choices") or []:
             if not isinstance(choice, dict):
                 continue
-            index = choice.get("index", 0)
-            state = choices.setdefault(
-                index,
-                {
-                    "role": "assistant",
-                    "content": [],
-                    "reasoning_content": [],
-                    "tool_calls": {},
-                    "finish_reason": None,
-                    "logprobs": None,
-                },
-            )
+            saw_choice = True
             delta = choice.get("delta") or {}
-            state["role"] = delta.get("role") or state["role"]
-            if isinstance(delta.get("content"), str):
-                state["content"].append(delta["content"])
+            role = delta.get("role") or role
+            if delta.get("content"):
+                content_parts.append(delta["content"])
             reasoning = delta.get("reasoning_content") or delta.get("reasoning")
-            if isinstance(reasoning, str):
-                state["reasoning_content"].append(reasoning)
-            if "refusal" in delta:
-                state["refusal"] = delta["refusal"]
+            if reasoning:
+                reasoning_parts.append(reasoning)
             for tc in delta.get("tool_calls") or []:
-                slot = state["tool_calls"].setdefault(
+                slot = tool_calls.setdefault(
                     tc.get("index", 0), {"id": None, "type": "function", "function": {"name": "", "arguments": ""}}
                 )
                 if tc.get("id"):
@@ -928,44 +825,21 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An
                 if fn.get("arguments"):
                     slot["function"]["arguments"] += fn["arguments"]
             if choice.get("finish_reason"):
-                state["finish_reason"] = choice["finish_reason"]
-            if choice.get("logprobs") is not None:
-                state["logprobs"] = choice["logprobs"]
-    if not choices:
+                finish_reason = choice["finish_reason"]
+    if not saw_choice:
         return None
-
-    reconstructed_choices = []
-    for index in sorted(choices):
-        state = choices[index]
-        message: dict[str, Any] = {"role": state["role"], "content": "".join(state["content"]) or None}
-        if state["reasoning_content"]:
-            message["reasoning_content"] = "".join(state["reasoning_content"])
-        if "refusal" in state:
-            message["refusal"] = state["refusal"]
-        if state["tool_calls"]:
-            message["tool_calls"] = [state["tool_calls"][i] for i in sorted(state["tool_calls"])]
-        reconstructed_choices.append(
-            {
-                "index": index,
-                "message": message,
-                "finish_reason": state["finish_reason"],
-                "logprobs": state["logprobs"],
-            }
-        )
-
+    message: dict[str, Any] = {"role": role, "content": "".join(content_parts) or None}
+    if reasoning_parts:
+        message["reasoning_content"] = "".join(reasoning_parts)
+    if tool_calls:
+        message["tool_calls"] = [tool_calls[i] for i in sorted(tool_calls)]
     result: dict[str, Any] = {
         "object": "chat.completion",
         "model": model,
-        "choices": reconstructed_choices,
+        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
     }
     if response_id is not None:
         result["id"] = response_id
-    if created is not None:
-        result["created"] = created
-    if service_tier is not None:
-        result["service_tier"] = service_tier
-    if system_fingerprint is not None:
-        result["system_fingerprint"] = system_fingerprint
     if usage:
         result["usage"] = usage
     return result
@@ -1197,9 +1071,8 @@ async def _flush_deferred_response() -> None:
         def _parse_and_record() -> None:
             # Off the event loop: body parse + SSE reassembly is best-effort and fully guarded, so a
             # malformed body can never surface as an ASGI error after the response was already sent.
-            normalized_response = scope.get(_NORMALIZED_CAPTURE_RESPONSE_SCOPE_KEY)
-            response_body = normalized_response if isinstance(normalized_response, dict) else None
-            if response_body is None and body_bytes:
+            response_body = None
+            if body_bytes:
                 try:
                     response_body = (
                         _reconstruct_streamed_response(body_bytes, dialect) if streaming else json.loads(body_bytes)
diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py
index e3b959deae..1f40bd0450 100644
--- a/responses_api_agents/stirrup_agent/app.py
+++ b/responses_api_agents/stirrup_agent/app.py
@@ -50,7 +50,7 @@
     NeMoGymResponseOutputText,
 )
 from nemo_gym.rollout_correlation import current_rollout_id
-from nemo_gym.server_utils import apply_rollout_prefix, get_response_json, raise_for_status
+from nemo_gym.server_utils import get_response_json, raise_for_status
 from responses_api_agents.stirrup_agent.task_strategy import TaskSampleSkipError, TaskStrategy
 
 
@@ -1038,13 +1038,7 @@ def model_post_init(self, __context: Any) -> None:
     # -- helpers ----------------------------------------------------------
 
     def _get_model_base_url(self) -> str:
-        from nemo_gym.global_config import get_first_server_config_dict
-        from nemo_gym.server_utils import ServerClient
-
-        global_config_dict = ServerClient.load_from_global_config().global_config_dict
-        model_server_config = get_first_server_config_dict(global_config_dict, self.config.model_server.name)
-        base_url = f"http://{model_server_config['host']}:{model_server_config['port']}"
-        return f"{apply_rollout_prefix(base_url, current_rollout_id())}/v1"
+        return self.resolve_model_base_url(self.config.model_server.name, current_rollout_id())
 
     # -- /v1/responses ----------------------------------------------------
 
diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py
index 4283fb9bd5..02441098bd 100644
--- a/responses_api_agents/stirrup_agent/tests/test_app.py
+++ b/responses_api_agents/stirrup_agent/tests/test_app.py
@@ -128,13 +128,12 @@ def test_sanity(self) -> None:
 
     def test_model_base_url_accepts_rollout_correlation(self) -> None:
         wrapper = StirrupAgentWrapper(config=_make_config(), server_client=MagicMock(spec=ServerClient))
-        loaded = MagicMock(global_config_dict={"policy_model": {}})
 
-        with (
-            patch("nemo_gym.server_utils.ServerClient.load_from_global_config", return_value=loaded),
-            patch(
-                "nemo_gym.global_config.get_first_server_config_dict",
-                return_value={"host": "model-host", "port": 8000},
+        with patch.object(
+            StirrupAgentWrapper,
+            "resolve_model_base_url",
+            side_effect=lambda _, rollout_id: (
+                f"http://model-host:8000/ng-rollout/{rollout_id}/v1" if rollout_id else "http://model-host:8000/v1"
             ),
         ):
             with rollout_context("7-3"):
@@ -245,7 +244,6 @@ async def test_run_judge_only_scores_cached_deliverables(self, tmp_path) -> None
 
         config = _make_config(judge_only=True, persist_deliverables_dir=str(tmp_path))
         server_client = MagicMock(spec=ServerClient)
-        server_client.global_config_dict = {"observability_enabled": True}
         server_client.post = AsyncMock(return_value=MagicMock())
         wrapper = StirrupAgentWrapper(config=config, server_client=server_client)
 
@@ -253,13 +251,7 @@ async def test_run_judge_only_scores_cached_deliverables(self, tmp_path) -> None
             input="ignored",
             metadata={"task_id": "task-1", "prompt": "do the thing", "_ng_rollout_index": "0"},
         )
-        body = StirrupRunRequest(
-            responses_create_params=params,
-            task_id="task-1",
-            prompt="do the thing",
-            _ng_task_index=7,
-            _ng_rollout_index=0,
-        )
+        body = StirrupRunRequest(responses_create_params=params, task_id="task-1", prompt="do the thing")
         request = MagicMock()
         request.cookies = {}
 
diff --git a/responses_api_agents/tau2/tests/test_app.py b/responses_api_agents/tau2/tests/test_app.py
index be12e96b0f..620b83c29f 100644
--- a/responses_api_agents/tau2/tests/test_app.py
+++ b/responses_api_agents/tau2/tests/test_app.py
@@ -19,7 +19,6 @@
 
 import pytest
 from fastapi.testclient import TestClient
-from tau2.data_model.simulation import RewardInfo, SimulationRun, TerminationReason
 
 from nemo_gym.base_responses_api_agent import AggregateMetricsRequest
 from nemo_gym.server_utils import ServerClient
@@ -134,27 +133,26 @@ def test_policy_and_user_model_calls_share_rollout_correlation(
         with patch("responses_api_agents.tau2.app.ensure_tau2_data_dir"):
             client = TestClient(server.setup_webserver())
 
-        result = SimulationRun(
-            id="run-1",
-            task_id="task-1",
-            start_time="2026-07-22T00:00:00Z",
-            end_time="2026-07-22T00:00:00Z",
-            duration=0,
-            termination_reason=TerminationReason.AGENT_STOP,
-            reward_info=RewardInfo(reward=1),
-            messages=[],
-        )
+        class StopRun(Exception):
+            pass
+
+        captured = {}
+
+        async def stop_after_config(**kwargs):
+            captured.update(kwargs)
+            raise StopRun
+
         model_urls = {"policy": "http://policy:8000", "user": "http://user:8001"}
         with (
             patch("responses_api_agents.tau2.app.get_server_url", side_effect=model_urls.__getitem__),
-            patch("responses_api_agents.tau2.app.run_single_task", AsyncMock(return_value=result)),
+            patch("responses_api_agents.tau2.app.run_single_task", stop_after_config),
+            pytest.raises(StopRun),
         ):
-            response = client.post("/run", json=request_body)
+            client.post("/run", json=request_body)
 
-        assert response.status_code == 200
-        response_config = response.json()["config"]
-        assert response_config["llm_args_agent"]["api_base"] == model_urls["policy"] + url_suffix
-        assert response_config["llm_args_user"]["api_base"] == model_urls["user"] + url_suffix
+        config = captured["config"]
+        assert config.llm_args_agent["api_base"] == model_urls["policy"] + url_suffix
+        assert config.llm_args_user["api_base"] == model_urls["user"] + url_suffix
 
     async def test_compute_metrics(self) -> None:
         example_rollouts_fpath = Path(__file__).parent.parent / "data" / "example_rollouts.jsonl"
diff --git a/tests/unit_tests/test_responses_api_model_messages.py b/tests/unit_tests/test_responses_api_model_messages.py
index 7e6129ae0c..0b7304f495 100644
--- a/tests/unit_tests/test_responses_api_model_messages.py
+++ b/tests/unit_tests/test_responses_api_model_messages.py
@@ -19,22 +19,14 @@
 default mapping for both ``responses()`` signatures (with and without a leading ``request``).
 """
 
-import json
 from time import time
 from unittest.mock import MagicMock
 from uuid import uuid4
 
-import pytest
-from aiohttp import ClientResponseError
 from fastapi import Body, Request
 from fastapi.testclient import TestClient
 
-from nemo_gym.base_responses_api_model import (
-    BaseResponsesAPIModelConfig,
-    CaptureStore,
-    SimpleResponsesAPIModel,
-    _reconstruct_chat_sse,
-)
+from nemo_gym.base_responses_api_model import BaseResponsesAPIModelConfig, SimpleResponsesAPIModel
 from nemo_gym.openai_utils import (
     NeMoGymChatCompletion,
     NeMoGymChatCompletionCreateParamsNonStreaming,
@@ -65,49 +57,11 @@ def _build_response(text: str, model: str = "downstream-model") -> NeMoGymRespon
     )
 
 
-def _build_chat_completion() -> NeMoGymChatCompletion:
-    return NeMoGymChatCompletion.model_validate(
-        {
-            "id": "chatcmpl_test",
-            "created": 123,
-            "model": "downstream-model",
-            "object": "chat.completion",
-            "service_tier": "default",
-            "system_fingerprint": "fp_test",
-            "choices": [
-                {
-                    "index": 0,
-                    "finish_reason": "tool_calls",
-                    "message": {
-                        "role": "assistant",
-                        "content": "first",
-                        "reasoning_content": "thinking",
-                        "tool_calls": [
-                            {
-                                "id": "call_1",
-                                "type": "function",
-                                "function": {"name": "lookup", "arguments": '{"id":1}'},
-                            }
-                        ],
-                    },
-                },
-                {
-                    "index": 1,
-                    "finish_reason": "stop",
-                    "message": {"role": "assistant", "content": "second"},
-                },
-            ],
-            "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18},
-        }
-    )
-
-
 class _BodyOnlyModel(SimpleResponsesAPIModel):
     """A server whose responses() takes only `body` (like openai_model)."""
 
     config: BaseResponsesAPIModelConfig
     last_params: object = None
-    last_chat_params: object = None
     model_config = {"arbitrary_types_allowed": True}
 
     async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse:
@@ -117,8 +71,7 @@ async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()
     async def chat_completions(
         self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
     ) -> NeMoGymChatCompletion:
-        object.__setattr__(self, "last_chat_params", body)
-        return _build_chat_completion()
+        raise NotImplementedError
 
 
 class _RequestAwareModel(SimpleResponsesAPIModel):
@@ -126,8 +79,6 @@ class _RequestAwareModel(SimpleResponsesAPIModel):
 
     config: BaseResponsesAPIModelConfig
     saw_request: bool = False
-    saw_chat_request: bool = False
-    last_chat_params: object = None
     model_config = {"arbitrary_types_allowed": True}
 
     async def responses(
@@ -136,26 +87,10 @@ async def responses(
         object.__setattr__(self, "saw_request", isinstance(request, Request))
         return _build_response("hi from request-aware")
 
-    async def chat_completions(
-        self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
-    ) -> NeMoGymChatCompletion:
-        object.__setattr__(self, "saw_chat_request", isinstance(request, Request))
-        object.__setattr__(self, "last_chat_params", body)
-        return _build_chat_completion()
-
-
-class _FailingModel(_BodyOnlyModel):
     async def chat_completions(
         self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()
     ) -> NeMoGymChatCompletion:
-        error = ClientResponseError(
-            request_info=MagicMock(real_url="https://example.invalid/v1/chat/completions"),
-            history=(),
-            status=403,
-            message="Forbidden",
-        )
-        error.response_content = b'{"error":{"message":"denied"}}'
-        raise error
+        raise NotImplementedError
 
 
 def _config() -> BaseResponsesAPIModelConfig:
@@ -167,14 +102,6 @@ def _client(model_cls) -> TestClient:
     return TestClient(server.setup_webserver()), server
 
 
-def _sse_payloads(body: str) -> list[dict]:
-    return [
-        json.loads(line.removeprefix("data: "))
-        for line in body.splitlines()
-        if line.startswith("data: ") and line != "data: [DONE]"
-    ]
-
-
 class TestDefaultMessagesRoute:
     def test_messages_route_registered_alongside_openai_routes(self) -> None:
         server = _BodyOnlyModel(config=_config(), server_client=MagicMock(spec=ServerClient, global_config_dict={}))
@@ -223,92 +150,3 @@ def test_streaming_returns_anthropic_sse(self) -> None:
         assert "event: message_start" in body
         assert "event: content_block_delta" in body
         assert "event: message_stop" in body
-
-
-@pytest.mark.parametrize("model_cls", [_BodyOnlyModel, _RequestAwareModel])
-def test_chat_completions_streaming_wraps_existing_nonstreaming_handler(model_cls) -> None:
-    client, server = _client(model_cls)
-    response = client.post(
-        "/v1/chat/completions",
-        json={
-            "model": "requested-model",
-            "messages": [{"role": "user", "content": "hello"}],
-            "stream": True,
-            "stream_options": {"include_usage": True},
-        },
-    )
-
-    assert response.status_code == 200
-    assert response.headers["content-type"].startswith("text/event-stream")
-    assert response.text.endswith("data: [DONE]\n\n")
-    assert server.last_chat_params.stream is False
-    assert server.last_chat_params.stream_options is None
-    if model_cls is _RequestAwareModel:
-        assert server.saw_chat_request is True
-
-    payloads = _sse_payloads(response.text)
-    reconstructed = NeMoGymChatCompletion.model_validate(_reconstruct_chat_sse(payloads))
-    assert reconstructed.model_dump(mode="json") == _build_chat_completion().model_dump(mode="json")
-
-
-def test_chat_completions_nonstreaming_is_unchanged() -> None:
-    client, server = _client(_BodyOnlyModel)
-    response = client.post(
-        "/v1/chat/completions",
-        json={"messages": [{"role": "user", "content": "hello"}]},
-    )
-
-    assert response.status_code == 200
-    assert response.headers["content-type"].startswith("application/json")
-    assert response.json()["object"] == "chat.completion"
-    assert server.last_chat_params.stream is None
-
-
-def test_streaming_chat_completion_capture_keeps_wire_request_and_reconstructs_response(tmp_path) -> None:
-    server_client = MagicMock(spec=ServerClient)
-    server_client.global_config_dict = {
-        "observability_enabled": True,
-        "model_call_capture_dir": str(tmp_path),
-    }
-    server = _BodyOnlyModel(config=_config(), server_client=server_client)
-    client = TestClient(server.setup_webserver())
-
-    response = client.post(
-        "/ng-rollout/rollout-1/v1/chat/completions",
-        json={
-            "messages": [{"role": "user", "content": "hello"}],
-            "stream": True,
-        },
-    )
-
-    assert response.status_code == 200
-    exchange = CaptureStore(tmp_path).read("rollout-1")[0]
-    assert exchange["status_code"] == 200
-    assert exchange["error_category"] is None
-    assert exchange["request"]["stream"] is True
-    assert "stream_options" not in exchange["request"]
-    reconstructed = NeMoGymChatCompletion.model_validate(exchange["response"])
-    assert reconstructed.model_dump(mode="json") == _build_chat_completion().model_dump(mode="json")
-    assert all("usage" not in payload for payload in _sse_payloads(response.text))
-    assert exchange["response_raw"].endswith("data: [DONE]\n\n")
-
-
-def test_capture_preserves_upstream_error_status_and_body(tmp_path) -> None:
-    server_client = MagicMock(spec=ServerClient)
-    server_client.global_config_dict = {
-        "observability_enabled": True,
-        "model_call_capture_dir": str(tmp_path),
-    }
-    server = _FailingModel(config=_config(), server_client=server_client)
-    client = TestClient(server.setup_webserver(), raise_server_exceptions=False)
-
-    response = client.post(
-        "/ng-rollout/rollout-1/v1/chat/completions",
-        json={"messages": [{"role": "user", "content": "hello"}]},
-    )
-
-    assert response.status_code == 500
-    exchange = CaptureStore(tmp_path).read("rollout-1")[0]
-    assert exchange["status_code"] == 403
-    assert exchange["error_category"] == "auth"
-    assert json.loads(exchange["response_raw"]) == {"error": {"message": "denied"}}