[shy982]

Problem

current_rollout_id() is Gym's public API for a server handler to find out which rollout the current request belongs to. It works from resources servers and agent servers. It does not work from a model server (responses_api_models/*) — it always returns None there, even though every correlated request a model server receives already carries the rollout id in its URL (/ng-rollout/<id>/...), and Gym's own capture middleware extracts that same id internally on every request for its own bookkeeping.

This blocks anything a model server might want to do with its own rollout id — propagate it to a downstream provider for correlation, key per-rollout state, log/debug per-rollout behavior — independent of which model server or downstream provider is involved.

Repro: any SimpleResponsesAPIModel subclass, called with a rollout-correlated URL —

  class EchoRolloutId(SimpleResponsesAPIModel):
      async def chat_completions(self, body=Body()):
          return {"rollout_id_seen": current_rollout_id()}  # always None

  POST /ng-rollout/4-2/v1/chat/completions returns
  {"rollout_id_seen": null}, even though 4-2 is right there in the path
  this same request arrived on.

Root cause

Gym has one general-purpose mechanism for exposing the rollout id to a handler: RolloutContextMiddleware, which matches the /ng-rollout/<id>/... prefix, publishes the id via a contextvar (current_rollout_id()), and strips the prefix before routing. It's installed by base_resources_server.py (resources servers) and base_responses_api_agent.py (agent servers), but never by base_responses_api_model.py (model servers).

server_utils.py reads current_rollout_id() on the sending side to stamp the prefix onto whatever URL it calls next — including a model server's — so the id is always sent; nothing on the receiving side ever reads it back into the same public API.

That the id does arrive intact is easy to prove: _CaptureMiddleware (installed by every SimpleResponsesAPIModel.setup_webserver()) extracts the exact same /ng-rollout/<id> prefix itself, via its own private regex, purely for internal training/eval-capture bookkeeping.

The information reaches the model server; it's just never exposed through the contextvar every other server type uses. Reusing RolloutContextMiddleware as-is on model servers isn't safe: it strips the prefix, but _CaptureMiddleware needs to see it intact — it does its own extraction and stripping, including a longer training-capture-specific prefix.

Installing the existing middleware naively would race _CaptureMiddleware over who strips the prefix first.

Fix

Add RolloutContextPeekMiddleware (nemo_gym/rollout_correlation.py): matches the same /ng-rollout/<id>/... prefix and publishes the id via rollout_context()/current_rollout_id() for the duration of the request, but never rewrites scope["path"]/scope["raw_path"] — it only reads, never touches routing, so _CaptureMiddleware downstream is unaffected. SimpleResponsesAPIModel.setup_webserver() installs it after install_model_call_capture() (which installs _CaptureMiddleware), so per Starlette's most-recently-added-runs-first order it's outermost and sees the untouched, prefixed path before _CaptureMiddleware strips it. Net effect: current_rollout_id() now works identically across resources, agent, and model servers, with no behavior change for anything that doesn't call it.

Tests

• test_rollout_context_peek_middleware_exposes_id_without_stripping_path — ASGI-level unit test: the middleware publishes the id during the call, leaves scope["path"]/scope["raw_path"] untouched, and the contextvar doesn't leak past the request.
• test_model_server_exposes_current_rollout_id_to_its_own_handler — integration test through a real SimpleResponsesAPIModel.setup_webserver(): a rollout-correlated request sees its own id from inside the handler; an uncorrelated request sees None.
• Existing capture regression test (test_verify_correlates_policy_and_judge_calls_and_preserves_raw_capture) passes unchanged, confirming _CaptureMiddleware's own prefix handling is unaffected.
• pytest tests/unit_tests/: 4456 passed, 0 new failures (two pre-existing failures reproduce identically without this change: test_e2b_provider.py::test_runtime_loader_routes_e2b_httpx_through_global_aiohttp, missing optional httpx_aiohttp; and test_enroot_provider.py::test_create_start_timeout_cleans_up, requires /proc, Linux-only).
• ruff check, ruff format --check, pre-commit run all pass on the changed files.

Context

Found while building a Gym inference_provider-backed integration that needed to propagate the rollout id to a downstream training service for per-rollout correlation, and had to work around this gap locally with a one-off, non-destructive middleware embedded in that one server. This PR generalizes that workaround into the base model-server class so any model server — current or future — gets the same capability Gym already gives resources and agent servers.

Checklist

• I have read the contributing guidelines (https://docs.nvidia.com/nemo/gym/latest/contribute/development-setup).
• The change is focused; unrelated "drive-by" edits are tracked as separate issues/PRs.
• Tests added or updated and pass locally, or N/A for docs-only / non-code changes (so CI unit/server checks pass when
  applicable).
• Pre-commit checks pass locally (pre-commit run --all-files) (so CI lint/format/copyright pass).
• All commits have DCO sign-off (git commit -s) (so the DCO check passes).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test c280916

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/36612179554