[ffrujeri]

Stacked on #5346 (episode tokens). Review that one first; this PR's diff against it is the base class and the toolset.

What changed

BaseEnvironmentServerWithTools (nemo_gym/server_runtime/base_environment_server.py, next to BaseEnvironmentServer) is a base class for an Environment Server that serves each episode's tools itself, so it needs no Resources Server. Other Environment Servers grant the agent a Resources Server's tools; this one owns the tools, the episode's state and verification. It does not impose turns: a subclass runs its episodes as it likes, and the user-simulation layer adds the turn loop on top. BaseEnvironmentServer does not change.

class MyEnvironment(BaseEnvironmentServerWithTools[MyRequest, MyResponse]):
    async def run(self, request, cleanup):
        state = MyEpisodeState(request.task)                       # the episode's state lives here
        served = await self.serve_tools(request, TypedToolset(state), cleanup)
        ...                                                        # seed the agent with served.tool_accesses
        return grade(state, served.calls)                          # every call, in order

• EpisodeToolset (nemo_gym/server_runtime/episode_tools.py): the tools of one episode, as specs() and call(name, arguments). TypedToolset builds one from @episode_tool methods that take a Pydantic model of arguments; the docstring is the description and the model's JSON schema the parameters. A toolset can also come from the task, as nemo_user_sim's probe tools will.
• serve_tools(request, toolset, cleanup) grants the toolset until the episode's cleanup runs. It returns a DirectHTTPToolAccess that points at this server, with a signed episode token (feat(environment): sign episode tokens for Environment Server tool grants #5346) in its headers, and the list that records each call.
• Tool calls go to POST /<tool name> at the server's root, as on a Resources Server, so simple_agent calls them unchanged, also under a /ng-rollout/<id>/ prefix. A tool cannot take the name of one of the server's routes, such as run.
• Public API. nemo_gym.server_runtime exports BaseEnvironmentServerWithTools, ServedTools and the toolset names (EpisodeToolset, TypedToolset, episode_tool, ToolSpec, ToolResult, ToolCallRecord, ToolArgumentsError), next to BaseEnvironmentServer, because a server outside nemo_gym/ imports from the package, not from its submodules.
• Calls to one episode run one at a time, in arrival order, and each is recorded (ToolCallRecord: name, arguments, output, error flag, duration).

Call	Answer
Ran (a tool can return ToolResult(is_error=True) for an error the agent should handle)	200 with the tool's output
No token, or one this server did not sign	401
Signed for this episode, but not the nonce of its grant	403
The episode is not running here (it ended, or another episode's token)	410
A tool the episode does not have	404
Arguments the tool rejects, or a body that is not a JSON object	422, recorded
The tool raised	500, recorded and logged

Why

This is the core abstraction for environments that drop their Resources Server: BIRD-Interact first, then nemo_user_sim (Linear GYMINT-640; the user-simulation RFC's layer 1, BaseEnvironmentServerWithTools). Per review, it is a new base class rather than a change to BaseEnvironmentServer, so existing Environment Servers are untouched. It was briefly named BaseSelfContainedEnvironmentServer, then BaseTurnBasedEnvironmentServer; it is now named for what it adds, since it does not impose turns.

Interim limit: until session routing covers these servers (the next PR in the stack), a BaseEnvironmentServerWithTools refuses num_workers > 1. With several workers, a tool call can reach a worker that does not hold the episode. Single-worker is not the target; the routing PR removes the limit.

Validation

Command	Result
pytest tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py tests/unit_tests/server_runtime/test_episode_tools.py tests/unit_tests/server_runtime/test_environment_server.py tests/unit_tests/test_public_api.py tests/unit_tests/test_import_linter_coverage.py (after the rebase onto main at dda609f86)	169 passed
pytest tests/unit_tests/server_runtime/test_base_environment_server_with_tools.py tests/unit_tests/server_runtime/test_episode_tools.py	29 passed. The server tests run real HTTP tool calls against the app (httpx.ASGITransport) while run() is in progress: calls change the episode's state and are recorded in order; a call after the episode ends gets 410; five token cases (missing, garbage, another server's, a guessed nonce, another episode's); unknown tool, rejected arguments, a non-object body and a raising tool; a rollout-prefixed call; two concurrent episodes with separate tokens; conflicting toolsets fail the episode; one grant per episode; num_workers=2 refused.
Same test with the per-episode lock removed	test_calls_to_one_episode_run_one_at_a_time fails (three calls ran at once), so the test checks the lock
pytest tests/unit_tests/server_runtime/test_environment_server.py tests/unit_tests/server_runtime/test_episode_tokens.py	57 passed
pre-commit run on the changed files	passed

Rollout evidence

N/A: a base class that no environment uses yet. BIRD-Interact's move onto it (Linear GYMINT-646) will run rollouts with simple_agent and Codex.

Compatibility

New modules only. BaseEnvironmentServer and every existing server are unchanged.

Next in the stack

MCP exposure of an episode's toolset, multi-worker session routing on the episode token, and running open episodes' cleanups at shutdown.

🤖 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 7db0586

[ffrujeri]

/ok to test fb8001f

[ffrujeri]

/ok to test cabd187

[ffrujeri]

/ok to test 24cd72d

[ffrujeri]

/ok to test ae441ee

[ffrujeri]

/ok to test 991166f

[ffrujeri]

/ok to test 72c6981