[rollout-observability][follow-up] Correlate auxiliary Model Server calls

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,11 +114,40 @@ 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
 
 The model-server boundary observes model HTTP requests and responses. It can record tool calls,
 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.

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)

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,14 +25,14 @@
     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 (
     NeMoGymResponse,
     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