[ffrujeri]

Stacked on #5349 (MCP). The stack: #5346 (episode tokens) → #5347 (the base class) → #5348 (shutdown sweep) → #5349 (MCP) → this PR. Review those first; this PR's diff against #5349 is the routing.

What changed

With several uvicorn workers, a tool call for an episode can reach a worker that does not hold the episode. BaseEnvironmentServerWithTools now routes each call to the worker that runs the episode, with the session routing that Resources and Agent Servers already use, and the interim num_workers > 1 refusal from #5347 is gone.

agent ──POST /add, episode token──▶ worker 2 ──forward (claim names worker 1)──▶ worker 1: runs the episode, checks the nonce, calls the tool

• Claim. serve_tools() claims the episode, by its capture key, on the worker that runs it. The episode's cleanup releases the claim when it stops serving the tools. A retried attempt (-a<N>) has its own key, so its own claim. Serving the same attempt on another worker is refused.
• Routing reads the token. The router reads which episode a call names from the signed episode token in its header (X-NeMo-Gym-Episode-Token). It does not need the nonce; only the episode's worker checks it, so a forged token is forwarded and then refused there with 403.
• MCP is routed the same way. feat(environment): serve BaseEnvironmentServerWithTools episode tools over MCP #5349's MCP grant carries the same header on its connection, so every MCP request (initialize, tools/list, tools/call) is routed by it. /mcp is stateless Streamable HTTP, so no MCP session lives on a worker; each request only has to reach the episode's worker, which the e2e test checks.
• nemo_gym/server_runtime/request_handling/session_routing.py, generic and opt-in:
  • session_header and session_from_header: a header that names the session, with the server's reader of its value, checked before the cookie and MCP token.
  • SessionHolder: a server that claims and releases each session itself, as an Agent Server releases its own.
  • SessionKind's cookie key, body field, and seed and close routes may be None, for a kind named only in a header.
  • Resources and Agent Servers pass none of these, so their routing is unchanged.
  • Main's rollout-stop path claim (SessionKind.path_claim, feat(agent): let a caller stop a running rollout and still grade it #5145) is unchanged. The episode kind sets none, so a tool call never takes that path.
• CI. The session-routing e2e job also runs when nemo_gym/server_runtime/base_environment_server.py changes.

Tool call	Answer
Reaches another worker	Forwarded to the episode's worker
The episode's worker exited	410 from the router ("no longer running")
The episode ended (claim released)	410 from the worker it reached ("not running on this server")
Token that names no episode	401 where it lands

Why

Linear GYMINT-645. Single-worker is not an option at Gym's scale, so an Environment Server that serves tools needs the same per-worker routing as a Resources Server. This removes the limit #5347 had to add.

Validation

Command	Result
pytest -n 12 tests/unit_tests/server_runtime/ tests/unit_tests/test_public_api.py tests/unit_tests/test_import_linter_coverage.py tests/unit_tests/test_moved_paths.py (after the rebase onto main at dda609f86)	1337 passed
NEMO_GYM_MULTIWORKER_E2E=1 pytest tests/e2e/session_routing (after the rebase onto main at dda609f86)	8 passed. One earlier run, right after the full unit suite, failed test_every_mcp_request_runs_on_the_worker_that_runs_its_episode and main's test_session_of_an_exited_worker_gets_an_explicit_error; both passed when rerun
NEMO_GYM_MULTIWORKER_STRESS=1 NEMO_GYM_STRESS_EPISODES=400 NEMO_GYM_STRESS_SESSIONS=100 pytest tests/e2e/session_routing/test_session_routing_stress.py (after the rebase onto main at dda609f86)	8 passed
pytest tests/unit_tests/server_runtime/test_episode_tools_session_routing.py	8 passed. Two in-process workers share a real socket directory: a call that reaches the other worker runs on the episode's worker; a forged nonce is forwarded and refused there (403); a token that names no episode gets 401 where it lands; each attempt has its own claim; one attempt cannot be served on two workers; the claim ends with the episode (then 410); an exited worker's episode gets 410; one worker serves calls without claims.
Same, with the claim in serve_tools disabled	5 of 8 fail
NEMO_GYM_MULTIWORKER_E2E=1 pytest tests/e2e/session_routing/test_episode_tools_routing_e2e.py	3 passed. A test-only BaseEnvironmentServerWithTools with 4 real workers and tool_transports: [direct_http, mcp]: 200 episodes each make 8 direct HTTP tool calls through the public port on new connections; all 1,600 return 200 and ran on their episode's worker, and the episodes ran on several workers. 50 episodes each act as an MCP client (the MCP SDK's, without keep-alive, so each request lands on any worker): initialize, tools/list and 4 tools/call; every episode listed exactly its tool and all 200 calls ran on its worker. Killing the worker that runs an episode makes its next tool call 410.
Same, with the claim disabled (before the rebase, direct HTTP tests)	Both fail: 1,150 of 1,600 calls got 410, and the dead worker's call got the handler's 410 instead of the router's
Same, with the claim disabled, MCP test	Fails: tools/list requests that reached other workers listed no tools
NEMO_GYM_MULTIWORKER_E2E=1 pytest tests/e2e/session_routing	8 passed (stress tests skipped)
NEMO_GYM_MULTIWORKER_STRESS=1 NEMO_GYM_STRESS_EPISODES=400 NEMO_GYM_STRESS_SESSIONS=100 pytest tests/e2e/session_routing/test_session_routing_stress.py (before the rebase)	8 passed (reduced from the default 4,000 episodes and 1,000 sessions)
pytest tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py tests/unit_tests/server_runtime/test_episode_tools_session_routing.py tests/unit_tests/server_runtime/test_session_routing.py tests/unit_tests/server_runtime/test_episode_tools.py tests/unit_tests/server_runtime/test_episode_tokens.py tests/unit_tests/server_runtime/test_mcp_auto_exposure.py tests/unit_tests/server_runtime/test_checkpoint_resources_mcp.py tests/unit_tests/server_runtime/test_environment_server.py tests/unit_tests/server_runtime/test_base_server.py (after the rebase)	511 passed
pytest tests/unit_tests/server_runtime/test_base_resources_server.py tests/unit_tests/server_runtime/test_checkpoint_participant_workers.py tests/unit_tests/server_runtime/test_base_server.py tests/unit_tests/server_runtime/test_mcp_auto_exposure.py (before the rebase)	445 passed
pre-commit run on the changed files	passed

Rollout evidence

N/A: infrastructure for a base class that no environment uses yet. The e2e test runs real multi-worker server processes; BIRD-Interact's move onto the base class (Linear GYMINT-646) will run rollouts.

Compatibility

• Resources, Agent and other Environment Servers route as before; the new install_session_routing arguments are optional and they pass none.
• SessionKind fields are now Optional. The two existing kinds set all of them.
• A BaseEnvironmentServerWithTools now accepts num_workers > 1.

🤖 Generated with Claude Code

[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.

[ffrujeri]

/ok to test efac265

[ffrujeri]

/ok to test d565468

[ffrujeri]

/ok to test dcd7c47

[ffrujeri]

/ok to test 34deb08

[ffrujeri]

/ok to test e9a058d

[ffrujeri]

/ok to test 3893869

[ffrujeri]

/ok to test 5daf3e5