feat: Add Gym-owned MCP resources server support - #1682
Merged
Conversation
adil-a
changed the title
Signed-off-by: Codex <codex@openai.com>
Signed-off-by: Codex <codex@openai.com>
The MCP SDK enables DNS-rebinding protection by default, which only accepts loopback Host headers and returns HTTP 421 otherwise. That breaks multi-node / use_absolute_ip deployments, where the agent reaches the resources server by a routable host: every MCP tool call 421s, no call is recorded, and verify() returns reward=0. Disable it in MCPResourcesServer (the endpoint is gated by the per-rollout X-NeMo-Gym-Session-Token instead) and add a regression test that drives /mcp with a non-loopback Host header. Also add an 'MCP Resources Server' environment tutorial (covering both the Gym-owned spin-up flow and the external/existing-server flow) and expand the example_mcp_weather README with runnable instructions. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
Authoring an MCP tool previously meant overriding register_mcp_tools, wrapping each tool in @mcp.tool(), and calling require_mcp_session_id() by hand. Add a @gym_tool decorator so a tool is just a typed method: MCPResourcesServer.register_mcp_tools now has a concrete default that discovers @gym_tool methods and registers each as an MCP tool, deriving the input schema from the method's parameters. Declare a 'session_id: str' parameter to receive the per-rollout Gym session; it is injected from the hidden token and omitted from the MCP input schema (the model only sees the real args). Stateless tools omit it. Both sync and async methods are supported. Constraints are enforced at registration: no 'request' parameter (there is no FastAPI Request on the MCP path) and no reserved names (verify/seed_session/ aggregate_metrics/mcp). Overriding register_mcp_tools (and calling super()) still works for hand-written tools. Rewrite example_mcp_weather to the decorator form, add base-class tests (auto-register, schema hides session_id, stateless tool, reserved-name + request-param rejection), and update the MCP Resources Server tutorial. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
verify() matched the city case-insensitively but compared the weather sentence exactly, so a correct get_weather call that used different casing (e.g. 'PARIS') scored 0. Match on the case-folded city alone (the sentence is derived from it) and compare the final answer case-insensitively. Add a regression test. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
adil-a
force-pushed
the
codex/gym-mcp-claude-code
branch
from
June 25, 2026 08:24
b54ba7e to
8e014bc
Compare
adil-a
marked this pull request as ready for review
June 25, 2026 08:24
The tutorial still said you implement register_mcp_tools (pre-@gym_tool) and its verify() snippet showed the old exact weather-sentence match (which the case-insensitivity fix replaced) plus an undefined variable. Update both to match the current code. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
Resolve the only conflict (uv.lock) by regenerating it from the auto-merged pyproject.toml (which keeps the mcp>=1.27,<2 dependency alongside main's changes). Verified: uv lock --check passes, 46 unit tests pass, and the live MCP endpoint round-trip works under the merged deps. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
Contributor
|
🌿 Preview your docs: https://nvidia-preview-codex-gym-mcp-claude-code.docs.buildwithfern.com/nemo/gym Here are the markdown pages you've updated: |
ruff-format (pre-commit, ruff v0.9.9) collapses four wrapped one-line statements in the TestRolloutMCPConfig tests that fit within the 119-col limit. These lines predate this work; the Lint check flagged them on --all-files. Verified clean against ruff v0.9.9. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
Contributor
Author
|
/ok to test 8aa4e91 |
ananthsub
reviewed
Jun 25, 2026
Comment on lines
+251
to
+257
| else: | ||
|
|
||
| @functools.wraps(method) | ||
| def wrapper(**kwargs: Any) -> Any: | ||
| if inject_session: | ||
| kwargs["session_id"] = self.require_mcp_session_id() | ||
| return method(**kwargs) |
Contributor
There was a problem hiding this comment.
it looks like gym_tool methods are dispatched by FastMCP directly on the event loop without any threadpool offload. this means a blocking sync tool will stall every concurrent rollout in the worker. can these sync tools be offloaded?
| def require_mcp_session_id(self) -> str: | ||
| token = _MCP_SESSION_TOKEN.get() | ||
| if not token: | ||
| raise HTTPException( |
Contributor
There was a problem hiding this comment.
not sure about this exception handling. it looks like HTTPException -> is caught and re-raised as a ToolError, where it's returned as isError:true over HTTP 200? so the client doesn't see the 401 error. should this raise another type of error instead?
| request.session[SESSION_ID_KEY] = session_id | ||
|
|
||
| token = uuid4().hex | ||
| self._mcp_session_id_by_token[token] = session_id |
Contributor
There was a problem hiding this comment.
- should we clean up the session token after? these are currently never removed
- also, if the server is configured to use
num_workers > 1, then we don't have the tokens from worker 1 visible in worker 2. we'd need either a stateless signed token like the existig session cookie or enforcingnum_workers == 1forMCPResourcesServer
| return {"mcpServers": {}} | ||
|
|
||
| config_path = Path(self.config.mcp_config).expanduser() | ||
| config = json.loads(config_path.read_text()) |
Contributor
There was a problem hiding this comment.
since the mcp config is static, we shouldn't need to read this from disk on every rollout. this can go into the model_post_init instead
| entry["headers"] = {str(key): str(value) for key, value in headers.items()} | ||
|
|
||
| config = self._load_static_mcp_config() | ||
| config.setdefault("mcpServers", {})[str(server_name)] = entry |
Contributor
There was a problem hiding this comment.
what about a risk of name collision here? do we need to document that the dynamic entry overwrites a static entry with the same server name?
| "url": url, | ||
| } | ||
| headers = metadata.get("headers") | ||
| if isinstance(headers, dict) and headers: |
Contributor
There was a problem hiding this comment.
do we need to warn/error if the MCP server returns missing headers?
| @@ -0,0 +1,5 @@ | |||
| {"responses_create_params":{"input":"Use the get_weather MCP tool for Paris, then answer with exactly the weather sentence returned by the tool.","tools":[]},"expected_city":"Paris"} | |||
Contributor
There was a problem hiding this comment.
expected_city should be under the verifier_metadata instead right?
also, input here is a plain string vs the message list format used by the other examples
Comment on lines
+417
to
+420
| url = metadata.get("url") | ||
| if not url: | ||
| url_path = str(metadata.get("url_path") or "/mcp") | ||
| url = f"{self._resources_server_base_url().rstrip('/')}/{url_path.lstrip('/')}" |
Contributor
There was a problem hiding this comment.
where is url coming from in metadata?
…ool error Addresses review feedback on MCPResourcesServer: - (#3) Replace the per-process token->session dict with a stateless signed token (itsdangerous URLSafeSerializer keyed by the deterministic session-middleware secret). Any worker can verify a token another worker minted, so this works with num_workers > 1 and there is nothing to evict. - (#1) Offload blocking sync @gym_tool methods to a threadpool so they don't stall the event loop (and every concurrent rollout in the worker). - (#2) Raise a plain MCPSessionError with a clean message instead of HTTPException(401). MCP runs over JSON-RPC (HTTP 200), so the status code never reaches the client; FastMCP surfaces this as a tool error (isError: true). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
…audit) - (#4) Cache the static mcp_config (read once, reuse) instead of reading it from disk every rollout. - (NVIDIA-NeMo#5) Document that the per-rollout Gym entry overwrites a same-named static mcp_config server. - (NVIDIA-NeMo#6) Warn when seed MCP metadata has no headers (tool calls would otherwise be unauthenticated). - (NVIDIA-NeMo#7) Add an agent test asserting the session cookie is threaded seed_session -> verify. - (#8a) Move the example task ground truth (expected_city) under verifier_metadata, per convention. - (#8b) Use the message-list input format in the example dataset, like the other examples. - (NVIDIA-NeMo#9) Remove the dead 'url' branch in the rollout mcp_config (MCPServerMetadata has no 'url'); audited that every metadata field accessed (server_name/url_path/transport/headers) actually exists. - (#2v) Add a test that a token-less MCP tool call surfaces as a clean isError (HTTP 200), not a 401. Also patch only ensure_claude_code (not all of model_post_init) in the claude_code_agent test helper, so the model's private attributes initialize. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
Signed-off-by: Codex <codex@openai.com> # Conflicts: # fern/versions/latest/pages/evaluation/adding-a-benchmark.mdx # fern/versions/latest/pages/evaluation/aggregate-metrics.mdx
Contributor
Author
|
/ok to test 6d45d25 |
- Teach the verifier_metadata request shape (tutorial snippet, Episode Flow, README round-trip) instead of the old top-level expected_city, which the server now silently ignores (extra=allow -> 'Paris' default). - Migrate example commands from the deprecated ng_run/ng_test legacy CLI to gym env start --config / gym env test --resources-server, with the Anthropic key in a repo-root env.yaml (matches the quickstart). - Fix the tutorial 'raises 401' claim -> MCPSessionError surfaced as a tool error (isError) over HTTP 200. - Reword the MCPResourcesServer class docstring for the @gym_tool default + signed session token. - Correct the --bare description in claude_code_agent (skills/MCP servers are not what it disables; align with 'claude --help') and drop a garbled comment fragment. - Make environment-tutorials section positions contiguous. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Codex <codex@openai.com>
Contributor
Author
|
/ok to test 3126260 |
ananthsub
approved these changes
Jun 25, 2026
Comment on lines
+166
to
+170
| Subclasses decorate tool methods with ``@gym_tool`` (the default ``register_mcp_tools`` | ||
| auto-registers them; override only for manual control) and call ``build_mcp_session_metadata`` | ||
| from ``seed_session`` to hand the agent a per-rollout token. A ``@gym_tool`` method receives the | ||
| Gym session by declaring a ``session_id`` parameter, which the base resolves from that token (a | ||
| stateless signed value) so tool calls share the session id used by /seed_session and /verify. |
Contributor
There was a problem hiding this comment.
for next steps, what's the process for supporting other resources servers in gym to be mcp compatible?
Contributor
There was a problem hiding this comment.
would be good to add a small section to docs: Converting an existing resources server to be mcp compatible
Contributor
Author
There was a problem hiding this comment.
I think best foot forward might be to modify the gym_tool decorator to register a given tool as both an HTTP POST endpoint and an MCP tool. The resources servers that expose a tool can then go ahead and inherit from MCPResourcesServer
kajalj22
approved these changes
Jun 26, 2026
adil-a
enabled auto-merge (squash)
June 26, 2026 01:05
cmunley1
approved these changes
Jun 26, 2026
This was referenced Jun 26, 2026
This was referenced Jun 26, 2026
This was referenced Jul 12, 2026
adil-a
added a commit
that referenced
this pull request
Jul 20, 2026
…CP mechanism PR #2059 supersedes the decorator-based MCP API from #1682: tools are the plain POST routes authors already write, exposed over MCP by the expose_tools_over_mcp config flag, with no decorators and no signature changes. Removed from the user-facing API: - @gym_tool, MCPResourcesServer, MCPSessionError, the token contextvar, and the header middleware (base_resources_server.py, 357 -> 151 lines) - their unit tests (test_base_resources_server.py) Kept: the wire contract the engine reuses (MCPServerMetadata, the X-NeMo-Gym-Session-Token header, the token salt, RESERVED_MCP_TOOL_NAMES). example_mcp_weather is rewritten as a plain SimpleResourcesServer with a typed POST /get_weather route and expose_tools_over_mcp: true in its yaml - the in-tree example of the new mechanism. Its tests now cover both doors: plain HTTP with the session cookie, and MCP tools/call through the real engine mount (plus a missing-token rejection test). Docs: the tutorial page now teaches auto-exposure as the Gym-owned flow (two integration shapes, not three); the external-server section and the workplace_assistant catch-all example are unchanged. Signed-off-by: Codex <adasif@nvidia.com>
bxyu-nvidia
pushed a commit
that referenced
this pull request
Jul 22, 2026
> **Draft / RFC.** Consolidates the exploration in #2002 and #2053 into one tracked module + a small framework hook, **removes #1682's decorator API** (`@gym_tool` / `MCPResourcesServer`) so there is exactly one user-facing MCP mechanism, then hardened by a full review pass (details below). Branched off `main`. ## What this is Serve a resources server's **existing** FastAPI tool routes over MCP with **zero handler changes** and **one opt-in config field** — no code change at all for a typed-route server: ```yaml resources_servers: my_server: entrypoint: app.py expose_tools_over_mcp: true # the only addition — set per instance in the config ``` `run_webserver` installs the `/mcp` endpoint automatically after the app is built (and imports the MCP SDK only for servers that opt in). Handlers keep their `request: Request` param and `request.session[...]` reads exactly as written. The only in-tree code a server ever needs is for shapes the config can't express, and it is one override: `mcp_tools(harvested, catchall)` — return the harvested typed routes (the default), filter one out to exclude a harness-only route (e.g. `/end_session`), or append catch-all-backed tools (`harvested + [catchall.tool(name, input_schema, description)]`) for dispatcher servers. A dropped route is never advertised, never callable over MCP, never shape-checked, and unchanged over HTTP. This PR also **removes the decorator-based MCP API** that #1682 added (`@gym_tool`, `MCPResourcesServer`, `MCPSessionError`, the token contextvar and header middleware): tools no longer need decorators or MCP-specific signatures, so the old mechanism is superseded rather than parallel. The wire contract survives (`MCPServerMetadata`, the `X-NeMo-Gym-Session-Token` header, the token salt, the reserved-name set) — `claude_code_agent` on `main` consumes it unchanged. `example_mcp_weather` is rewritten onto auto-exposure: a plain typed `POST /get_weather` route plus the yaml flag, with tests driving the MCP door (and the wrapped `/seed_session` + `/verify`) through the real engine mount, and the plain-HTTP tool route on the stock app. ## How to review this (~30 min reading order) 1. `nemo_gym/mcp_auto_exposure.py` module docstring — the contract in one screen. 2. `bind_route` — the detector. The accept/refuse table below is what it enforces. 3. `call_direct` — dispatch parity with the plain HTTP route (threadpool for sync handlers, body bytes, session cookie, `response_model` filtering, error-text parity). 4. `harvest_tools` → `install_auto_exposure` — tool map, startup guards, session token, `/seed_session` wrap, `/mcp` mount. 5. `base_resources_server.py` diff — the `expose_tools_over_mcp` config field, the two overridables (`mcp_tools(harvested, catchall)` — what to expose; `mcp_allowed_tools_for_session(seed_body)` — per-rollout narrowing), and the removal of the #1682 decorator API (the file drops from 315 lines on `main` to 150; it had grown to 356 on this branch before the removal commit). 6. `server_utils.py` — the 6-line activation point (plus a 4-line session-middleware idempotency guard). 7. `tests/unit_tests/test_mcp_auto_exposure.py` — 45 tests double as the spec; each guard and parity shape has a named test. ## What the detector accepts vs refuses **Accepts (dispatches identically to the plain HTTP route):** async and sync handlers (sync runs via `run_in_threadpool`, exactly like FastAPI); `body: Model`; `body: dict` / `dict[str, Any]`; raw-body catch-alls reading `await request.json()`; handlers taking both a body model *and* a `request` they read raw; `str` path params; `response_model=` filtering (decorator kwarg or return annotation). **Refuses loudly at startup (`ValueError` naming the route and reason):** non-Gym middleware; `Depends`/`Security`; `*args/**kwargs`; multiple body models; union/optional body params; defaulted query params; unresolvable annotations; tool names outside `^[A-Za-z0-9_-]+$` (e.g. nested routes); tool names (from `mcp_tools()`) colliding with `verify`/`seed_session`/`aggregate_metrics`/`mcp`; duplicate tool names; multiple parameterized catch-all routes (auto-exposure cannot tell which backs the tools); a server that already serves `/mcp` (a hand-rolled MCP mount would be shadowed); a missing `/seed_session` or `/verify`. These rules apply to route-handler signatures only — inner tool functions may use `**kwargs` freely, and no in-tree server trips any refusal. No silent fallback exists: a shape either dispatches provably like HTTP or the server does not start. One soft case: a dispatcher whose `mcp_tools()` override ignores its catch-all logs a warning (the catch-all-backed tools stay HTTP-only). ## Session & token `/seed_session` gains an additive `"mcp"` key (`MCPServerMetadata` — the same shape `claude_code_agent` on `main` already consumes): the `/mcp` URL plus a signed per-rollout token in `X-NeMo-Gym-Session-Token`. The token payload is always `{"sid", "tools"}` — the session id plus this rollout's allow-list from `mcp_allowed_tools_for_session(seed_body)` (`null` = unrestricted) — signed with `URLSafeSerializer` (untimed; expiry was cut as consumer-less scope). The header/salt/metadata constants live in `base_resources_server` — the wire contract retained from #1682's scheme after its decorator API was removed. The token is verified on every call (no cache). `tools/list` works tokenless and advertises the full exposed set; a token carrying an allow-list narrows both `tools/list` and `tools/call` for that session. Known limitation, deliberately out of scope: the signing secret derivation (`class___config-name`) is pre-existing `main` behavior (the session-middleware secret) — hardening it is a separate change. ## Verify-time tool-name normalization (install-time, gated) MCP-native agents record trajectory tool calls namespaced (`mcp__<server>__<tool>`); verifiers know bare names — found live when MCP rollouts scored 0.0 on perfect trajectories. For **flag-on servers only**, installation wraps the `/verify` route's current endpoint (whatever handler it holds — servers that re-register `/verify` are covered by construction): names are normalized for scoring on a deep copy, then the model's original names are restored in the echoed response (matched by `call_id`). Flag-off servers — every existing benchmark — keep verify byte-for-byte, so no baselines move; persisted rollout artifacts always keep the names the model actually emitted. ## Diff shape (+2,136 / −631, 11 files) | Area | Change | |---|---| | `nemo_gym/mcp_auto_exposure.py` | **new, one file, 841 lines** — detector, direct dispatcher, harvest + guards, token, `/seed_session` + `/verify` wraps, `/mcp` mount | | `nemo_gym/base_resources_server.py` | +48/−207: adds the `expose_tools_over_mcp` config field, the `normalize_tool_name` helper, and the `mcp_tools()` / `mcp_allowed_tools_for_session()` overridables; **removes** `@gym_tool`, `MCPResourcesServer`, `MCPSessionError`, and the header middleware (their unit tests go too: `test_base_resources_server.py` −277) | | `nemo_gym/server_utils.py` | +10: `run_webserver` activates auto-exposure for opted-in servers (lazy MCP import); `setup_session_middleware` is now idempotent | | `tests/unit_tests/test_mcp_auto_exposure.py` | **new — 45 self-contained tests** (TestClient, synthetic servers, ~3 s); +37 in `test_server_utils.py`; docs section in `fern/.../mcp-resources-server.mdx` | | `resources_servers/` | `example_mcp_weather` rewritten onto auto-exposure (plain typed route + yaml flag; 6 tests: the MCP door — plus the wrapped `/seed_session` and `/verify` — through the real engine mount, and the plain-HTTP tool route on the stock app). **Every other server: byte-identical to `main`.** Enablement stays per-server follow-up: pure config for typed-route servers; a small `mcp_tools()` override for catch-all dispatchers. | MCP engine: the official SDK's **public low-level** `Server` — no private-attr access. ## Review provenance This branch already went through a 7-reviewer + devil's-advocate review pass (59 confirmed findings) followed by a fix pass with adversarial re-verification: all 55 code findings are fixed here with tests (silently-wrong detector shapes, refuse-loudly gaps, an unbounded token cache — the token is now verified on every call with no cache, and expiry was later cut as consumer-less scope — event-loop blocking for sync handlers, provenance-preserving verify normalization), and a docstring-accuracy audit (77 reviewed, 3 corrected). The remaining follow-ups are listed under "For discussion". ## Verified - **Reproducible from this PR:** `pytest tests/unit_tests/test_mcp_auto_exposure.py` → **45/45** (~3 s; covers every refusal, parity shape, token path, exclusion semantics, and the normalization gating/provenance — including a re-registered-`/verify` regression test). - **Live e2e at this HEAD (`3891dd65`):** `workplace_assistant` MCP door (`claude_code_agent`) — **5/5 `reward: 1.0`**, all 8 distinct tools namespaced; HTTP door (`simple_agent`) — **5/5** on the second run, 4/5 on the first (the failing rollout omitted the `status: "Lead"` search filter — a trajectory-verified policy flake, not an engine failure; the HTTP door runs the stock route code). `example_mcp_weather` (`claude_code_agent`) — **5/5 `reward: 1.0`** at the removal commit `4d6d569c`; the one commit since is docs plus a comment reword. Earlier demo (below): both doors 5/5 tasks `reward: 1.0`, ~30 s wall for 5 concurrent rollouts, MCP-door trajectories persist the namespaced names the model emitted. n=5 single runs of a stochastic policy — a smoke demonstration that both transports work end to end, not a statistical parity claim. - **Development-time only (local, gitignored — not reproducible from this PR):** a 44-check live acceptance harness exercising six representative in-tree servers (finance, workplace, aviary, newton_bench, openenv, ns_tools) against pristine `origin/main` handlers, covering the factory-`__signature__`, `dict`-body, and raw-body/`PlainTextResponse` shapes; 44/44 as of 2026-07-17 against a pre-refactor revision of the engine (development-time evidence; the harness predates the `mcp_tools()` API). Other in-tree servers were covered by a static route sweep, not this live harness. ## Runnable e2e — both doors through `gym eval run` The full 5-task `workplace_assistant` example dataset, driven end to end by Gym. Because `resources_servers/` ships unchanged, the demo uses a small user-side entrypoint (`app_mcp.py`, shown in the docs added by this PR) that subclasses the stock server and overrides `mcp_tools()` to append the 27 catch-all-backed tools; the flag itself is plain yaml: Compose config (`workplace_2059.yaml`) — MCP door via `claude_code_agent` (already on `main`; it reads the `mcp` key `/seed_session` returns and wires the Claude Code binary itself), HTTP door via `simple_agent`: ```yaml workplace_assistant: resources_servers: workplace_assistant: entrypoint: app_mcp.py # thin subclass overriding mcp_tools(); see the docs page domain: agent expose_tools_over_mcp: true workplace_claude: responses_api_agents: claude_code_agent: entrypoint: app.py resources_server: {type: resources_servers, name: workplace_assistant} model: ${anthropic_model_name} anthropic_api_key: ${anthropic_api_key} # env.yaml; any Anthropic-format endpoint anthropic_base_url: ${anthropic_base_url} # optional gateway; CLI appends /v1/messages policy_model: responses_api_models: inference_provider: entrypoint: app.py base_url: ${policy_base_url} # any OpenAI-format endpoint api_key: ${policy_api_key} model: ${policy_model_name} workplace_simple: responses_api_agents: simple_agent: entrypoint: app.py max_steps: 12 resources_server: {type: resources_servers, name: workplace_assistant} model_server: {type: responses_api_models, name: policy_model} ``` ```bash gym env start --config workplace_2059.yaml # (a) MCP door — Claude Code discovers and calls the tools over /mcp with the per-rollout token gym eval run --no-serve --agent workplace_claude \ --input resources_servers/workplace_assistant/data/example.jsonl \ --output results/workplace_claude_rollouts.jsonl # (b) HTTP door — the same tools as plain POST /<tool_name> gym eval run --no-serve --agent workplace_simple \ --input resources_servers/workplace_assistant/data/example.jsonl \ --output results/workplace_simple_rollouts.jsonl ``` Tip: if your provider's OpenAI *Responses-API* emulation is incomplete (a common gateway gap on multi-turn tool exchanges), use `inference_provider` (chat-completions + Gym's converter) rather than `openai_model`. ## Enablement (deliberately none in this PR) This PR ships the engine plus the one in-tree example (`example_mcp_weather`) — every other server under `resources_servers/` is byte-identical to `main` (we cannot properly test each env here, per review). Enabling a server is per-server follow-up work: set `expose_tools_over_mcp: true` in its instance config (typed-route servers need no code at all), plus a small `mcp_tools()` override for catch-all dispatcher servers — the docs added in this PR walk through workplace_assistant end to end. During development, a fleet exercise ran `claude_code_agent` over `/mcp` against 12 servers enabled this way (workplace_assistant 27 tools, math_advanced_calculations, finance_sec_search, indirect_prompt_injection with per-task narrowing, the examples, and others), all dispatching correctly; those results are development-time evidence, not shipped code. **Deliberately not exposed even as follow-up (each needs design work first):** aviary (its `/step`/`/close` bodies carry an `env_id` that can address other rollouts' environments — needs session-derived lookup), genrm_compare (its only route is a harness-side batch API), and the gymnasium family (reward accrues from `/reset`/`/step`, which a model could drive to reroll its own episode). The 74 verify-only benchmarks have no tool routes to expose. ## For discussion - Flag is **opt-in (default off)** — auto-exposing every route isn't always wanted (aviary's `/step` carries `env_id`; exposing it lets a model address other rollouts' envs). - Follow-ups deliberately not in this PR: session-binding aviary's `env_id` so it can be exposed safely; an MCP-safe reward design for the gymnasium family; and hardening the pre-existing token-signing secret derivation shared with #1682. - **#1682's decorator API is removed in this PR** rather than left as a parallel mechanism: auto-exposure reads the routes authors already wrote, so `@gym_tool` had no remaining role, and keeping both would have forced every author to choose. `example_mcp_weather` (the only in-tree consumer) is migrated here; #2002's dual-registration approach is superseded. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Codex <adasif@nvidia.com> Signed-off-by: Codex <codex@openai.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
OlegSudakov
pushed a commit
to OlegSudakov/Gym
that referenced
this pull request
Aug 7, 2026
> **Draft / RFC.** Consolidates the exploration in NVIDIA-NeMo#2002 and NVIDIA-NeMo#2053 into one tracked module + a small framework hook, **removes NVIDIA-NeMo#1682's decorator API** (`@gym_tool` / `MCPResourcesServer`) so there is exactly one user-facing MCP mechanism, then hardened by a full review pass (details below). Branched off `main`. ## What this is Serve a resources server's **existing** FastAPI tool routes over MCP with **zero handler changes** and **one opt-in config field** — no code change at all for a typed-route server: ```yaml resources_servers: my_server: entrypoint: app.py expose_tools_over_mcp: true # the only addition — set per instance in the config ``` `run_webserver` installs the `/mcp` endpoint automatically after the app is built (and imports the MCP SDK only for servers that opt in). Handlers keep their `request: Request` param and `request.session[...]` reads exactly as written. The only in-tree code a server ever needs is for shapes the config can't express, and it is one override: `mcp_tools(harvested, catchall)` — return the harvested typed routes (the default), filter one out to exclude a harness-only route (e.g. `/end_session`), or append catch-all-backed tools (`harvested + [catchall.tool(name, input_schema, description)]`) for dispatcher servers. A dropped route is never advertised, never callable over MCP, never shape-checked, and unchanged over HTTP. This PR also **removes the decorator-based MCP API** that NVIDIA-NeMo#1682 added (`@gym_tool`, `MCPResourcesServer`, `MCPSessionError`, the token contextvar and header middleware): tools no longer need decorators or MCP-specific signatures, so the old mechanism is superseded rather than parallel. The wire contract survives (`MCPServerMetadata`, the `X-NeMo-Gym-Session-Token` header, the token salt, the reserved-name set) — `claude_code_agent` on `main` consumes it unchanged. `example_mcp_weather` is rewritten onto auto-exposure: a plain typed `POST /get_weather` route plus the yaml flag, with tests driving the MCP door (and the wrapped `/seed_session` + `/verify`) through the real engine mount, and the plain-HTTP tool route on the stock app. ## How to review this (~30 min reading order) 1. `nemo_gym/mcp_auto_exposure.py` module docstring — the contract in one screen. 2. `bind_route` — the detector. The accept/refuse table below is what it enforces. 3. `call_direct` — dispatch parity with the plain HTTP route (threadpool for sync handlers, body bytes, session cookie, `response_model` filtering, error-text parity). 4. `harvest_tools` → `install_auto_exposure` — tool map, startup guards, session token, `/seed_session` wrap, `/mcp` mount. 5. `base_resources_server.py` diff — the `expose_tools_over_mcp` config field, the two overridables (`mcp_tools(harvested, catchall)` — what to expose; `mcp_allowed_tools_for_session(seed_body)` — per-rollout narrowing), and the removal of the NVIDIA-NeMo#1682 decorator API (the file drops from 315 lines on `main` to 150; it had grown to 356 on this branch before the removal commit). 6. `server_utils.py` — the 6-line activation point (plus a 4-line session-middleware idempotency guard). 7. `tests/unit_tests/test_mcp_auto_exposure.py` — 45 tests double as the spec; each guard and parity shape has a named test. ## What the detector accepts vs refuses **Accepts (dispatches identically to the plain HTTP route):** async and sync handlers (sync runs via `run_in_threadpool`, exactly like FastAPI); `body: Model`; `body: dict` / `dict[str, Any]`; raw-body catch-alls reading `await request.json()`; handlers taking both a body model *and* a `request` they read raw; `str` path params; `response_model=` filtering (decorator kwarg or return annotation). **Refuses loudly at startup (`ValueError` naming the route and reason):** non-Gym middleware; `Depends`/`Security`; `*args/**kwargs`; multiple body models; union/optional body params; defaulted query params; unresolvable annotations; tool names outside `^[A-Za-z0-9_-]+$` (e.g. nested routes); tool names (from `mcp_tools()`) colliding with `verify`/`seed_session`/`aggregate_metrics`/`mcp`; duplicate tool names; multiple parameterized catch-all routes (auto-exposure cannot tell which backs the tools); a server that already serves `/mcp` (a hand-rolled MCP mount would be shadowed); a missing `/seed_session` or `/verify`. These rules apply to route-handler signatures only — inner tool functions may use `**kwargs` freely, and no in-tree server trips any refusal. No silent fallback exists: a shape either dispatches provably like HTTP or the server does not start. One soft case: a dispatcher whose `mcp_tools()` override ignores its catch-all logs a warning (the catch-all-backed tools stay HTTP-only). ## Session & token `/seed_session` gains an additive `"mcp"` key (`MCPServerMetadata` — the same shape `claude_code_agent` on `main` already consumes): the `/mcp` URL plus a signed per-rollout token in `X-NeMo-Gym-Session-Token`. The token payload is always `{"sid", "tools"}` — the session id plus this rollout's allow-list from `mcp_allowed_tools_for_session(seed_body)` (`null` = unrestricted) — signed with `URLSafeSerializer` (untimed; expiry was cut as consumer-less scope). The header/salt/metadata constants live in `base_resources_server` — the wire contract retained from NVIDIA-NeMo#1682's scheme after its decorator API was removed. The token is verified on every call (no cache). `tools/list` works tokenless and advertises the full exposed set; a token carrying an allow-list narrows both `tools/list` and `tools/call` for that session. Known limitation, deliberately out of scope: the signing secret derivation (`class___config-name`) is pre-existing `main` behavior (the session-middleware secret) — hardening it is a separate change. ## Verify-time tool-name normalization (install-time, gated) MCP-native agents record trajectory tool calls namespaced (`mcp__<server>__<tool>`); verifiers know bare names — found live when MCP rollouts scored 0.0 on perfect trajectories. For **flag-on servers only**, installation wraps the `/verify` route's current endpoint (whatever handler it holds — servers that re-register `/verify` are covered by construction): names are normalized for scoring on a deep copy, then the model's original names are restored in the echoed response (matched by `call_id`). Flag-off servers — every existing benchmark — keep verify byte-for-byte, so no baselines move; persisted rollout artifacts always keep the names the model actually emitted. ## Diff shape (+2,136 / −631, 11 files) | Area | Change | |---|---| | `nemo_gym/mcp_auto_exposure.py` | **new, one file, 841 lines** — detector, direct dispatcher, harvest + guards, token, `/seed_session` + `/verify` wraps, `/mcp` mount | | `nemo_gym/base_resources_server.py` | +48/−207: adds the `expose_tools_over_mcp` config field, the `normalize_tool_name` helper, and the `mcp_tools()` / `mcp_allowed_tools_for_session()` overridables; **removes** `@gym_tool`, `MCPResourcesServer`, `MCPSessionError`, and the header middleware (their unit tests go too: `test_base_resources_server.py` −277) | | `nemo_gym/server_utils.py` | +10: `run_webserver` activates auto-exposure for opted-in servers (lazy MCP import); `setup_session_middleware` is now idempotent | | `tests/unit_tests/test_mcp_auto_exposure.py` | **new — 45 self-contained tests** (TestClient, synthetic servers, ~3 s); +37 in `test_server_utils.py`; docs section in `fern/.../mcp-resources-server.mdx` | | `resources_servers/` | `example_mcp_weather` rewritten onto auto-exposure (plain typed route + yaml flag; 6 tests: the MCP door — plus the wrapped `/seed_session` and `/verify` — through the real engine mount, and the plain-HTTP tool route on the stock app). **Every other server: byte-identical to `main`.** Enablement stays per-server follow-up: pure config for typed-route servers; a small `mcp_tools()` override for catch-all dispatchers. | MCP engine: the official SDK's **public low-level** `Server` — no private-attr access. ## Review provenance This branch already went through a 7-reviewer + devil's-advocate review pass (59 confirmed findings) followed by a fix pass with adversarial re-verification: all 55 code findings are fixed here with tests (silently-wrong detector shapes, refuse-loudly gaps, an unbounded token cache — the token is now verified on every call with no cache, and expiry was later cut as consumer-less scope — event-loop blocking for sync handlers, provenance-preserving verify normalization), and a docstring-accuracy audit (77 reviewed, 3 corrected). The remaining follow-ups are listed under "For discussion". ## Verified - **Reproducible from this PR:** `pytest tests/unit_tests/test_mcp_auto_exposure.py` → **45/45** (~3 s; covers every refusal, parity shape, token path, exclusion semantics, and the normalization gating/provenance — including a re-registered-`/verify` regression test). - **Live e2e at this HEAD (`3891dd65`):** `workplace_assistant` MCP door (`claude_code_agent`) — **5/5 `reward: 1.0`**, all 8 distinct tools namespaced; HTTP door (`simple_agent`) — **5/5** on the second run, 4/5 on the first (the failing rollout omitted the `status: "Lead"` search filter — a trajectory-verified policy flake, not an engine failure; the HTTP door runs the stock route code). `example_mcp_weather` (`claude_code_agent`) — **5/5 `reward: 1.0`** at the removal commit `4d6d569c`; the one commit since is docs plus a comment reword. Earlier demo (below): both doors 5/5 tasks `reward: 1.0`, ~30 s wall for 5 concurrent rollouts, MCP-door trajectories persist the namespaced names the model emitted. n=5 single runs of a stochastic policy — a smoke demonstration that both transports work end to end, not a statistical parity claim. - **Development-time only (local, gitignored — not reproducible from this PR):** a 44-check live acceptance harness exercising six representative in-tree servers (finance, workplace, aviary, newton_bench, openenv, ns_tools) against pristine `origin/main` handlers, covering the factory-`__signature__`, `dict`-body, and raw-body/`PlainTextResponse` shapes; 44/44 as of 2026-07-17 against a pre-refactor revision of the engine (development-time evidence; the harness predates the `mcp_tools()` API). Other in-tree servers were covered by a static route sweep, not this live harness. ## Runnable e2e — both doors through `gym eval run` The full 5-task `workplace_assistant` example dataset, driven end to end by Gym. Because `resources_servers/` ships unchanged, the demo uses a small user-side entrypoint (`app_mcp.py`, shown in the docs added by this PR) that subclasses the stock server and overrides `mcp_tools()` to append the 27 catch-all-backed tools; the flag itself is plain yaml: Compose config (`workplace_2059.yaml`) — MCP door via `claude_code_agent` (already on `main`; it reads the `mcp` key `/seed_session` returns and wires the Claude Code binary itself), HTTP door via `simple_agent`: ```yaml workplace_assistant: resources_servers: workplace_assistant: entrypoint: app_mcp.py # thin subclass overriding mcp_tools(); see the docs page domain: agent expose_tools_over_mcp: true workplace_claude: responses_api_agents: claude_code_agent: entrypoint: app.py resources_server: {type: resources_servers, name: workplace_assistant} model: ${anthropic_model_name} anthropic_api_key: ${anthropic_api_key} # env.yaml; any Anthropic-format endpoint anthropic_base_url: ${anthropic_base_url} # optional gateway; CLI appends /v1/messages policy_model: responses_api_models: inference_provider: entrypoint: app.py base_url: ${policy_base_url} # any OpenAI-format endpoint api_key: ${policy_api_key} model: ${policy_model_name} workplace_simple: responses_api_agents: simple_agent: entrypoint: app.py max_steps: 12 resources_server: {type: resources_servers, name: workplace_assistant} model_server: {type: responses_api_models, name: policy_model} ``` ```bash gym env start --config workplace_2059.yaml # (a) MCP door — Claude Code discovers and calls the tools over /mcp with the per-rollout token gym eval run --no-serve --agent workplace_claude \ --input resources_servers/workplace_assistant/data/example.jsonl \ --output results/workplace_claude_rollouts.jsonl # (b) HTTP door — the same tools as plain POST /<tool_name> gym eval run --no-serve --agent workplace_simple \ --input resources_servers/workplace_assistant/data/example.jsonl \ --output results/workplace_simple_rollouts.jsonl ``` Tip: if your provider's OpenAI *Responses-API* emulation is incomplete (a common gateway gap on multi-turn tool exchanges), use `inference_provider` (chat-completions + Gym's converter) rather than `openai_model`. ## Enablement (deliberately none in this PR) This PR ships the engine plus the one in-tree example (`example_mcp_weather`) — every other server under `resources_servers/` is byte-identical to `main` (we cannot properly test each env here, per review). Enabling a server is per-server follow-up work: set `expose_tools_over_mcp: true` in its instance config (typed-route servers need no code at all), plus a small `mcp_tools()` override for catch-all dispatcher servers — the docs added in this PR walk through workplace_assistant end to end. During development, a fleet exercise ran `claude_code_agent` over `/mcp` against 12 servers enabled this way (workplace_assistant 27 tools, math_advanced_calculations, finance_sec_search, indirect_prompt_injection with per-task narrowing, the examples, and others), all dispatching correctly; those results are development-time evidence, not shipped code. **Deliberately not exposed even as follow-up (each needs design work first):** aviary (its `/step`/`/close` bodies carry an `env_id` that can address other rollouts' environments — needs session-derived lookup), genrm_compare (its only route is a harness-side batch API), and the gymnasium family (reward accrues from `/reset`/`/step`, which a model could drive to reroll its own episode). The 74 verify-only benchmarks have no tool routes to expose. ## For discussion - Flag is **opt-in (default off)** — auto-exposing every route isn't always wanted (aviary's `/step` carries `env_id`; exposing it lets a model address other rollouts' envs). - Follow-ups deliberately not in this PR: session-binding aviary's `env_id` so it can be exposed safely; an MCP-safe reward design for the gymnasium family; and hardening the pre-existing token-signing secret derivation shared with NVIDIA-NeMo#1682. - **NVIDIA-NeMo#1682's decorator API is removed in this PR** rather than left as a parallel mechanism: auto-exposure reads the routes authors already wrote, so `@gym_tool` had no remaining role, and keeping both would have forced every author to choose. `example_mcp_weather` (the only in-tree consumer) is migrated here; NVIDIA-NeMo#2002's dual-registration approach is superseded. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Codex <adasif@nvidia.com> Signed-off-by: Codex <codex@openai.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a Gym-owned MCP integration layer so an MCP-native agent (Claude Code) can use tools that a
Resources Server both serves and verifies — i.e. MCP tool implementations +
verify()= aResources Server. Closes #212.
MCPResourcesServer(nemo_gym/base_resources_server.py): aSimpleResourcesServerthat alsomounts a Streamable-HTTP MCP endpoint at
/mcpon the same app as/seed_sessionand/verify.Per-rollout isolation via a hidden
X-NeMo-Gym-Session-Tokenthat maps a tool call back to the sameGym
session_idused by/seed_session//verify. DNS-rebinding protection is disabled so calls workoff-loopback (multi-node /
use_absolute_ip=True); the token is the auth instead.@gym_toolauthoring API: an MCP tool is just a decorated, typed method — noregister_mcp_toolsboilerplate. The base class auto-registers each
@gym_toolmethod as an MCP tool (input schema derivedfrom its parameters). Declare a
session_id: strparameter to receive the per-rollout session; it isinjected from the token and hidden from the tool's input schema. Constraints are enforced at
registration (no
requestparameter — there is no FastAPIRequeston the MCP path; no reserved namesverify/seed_session/aggregate_metrics/mcp). Overrideregister_mcp_toolsfor full manual control.example_mcp_weatherresources server: aget_weather@gym_toolwith same-session verification.claude_code_agent: generates a per-rollout--mcp-configfrom the/seed_sessionMCP metadata,while preserving static
mcp_configbehavior (the two are merged).fern/and an expanded example README.How to launch the Claude Code example test
Prereqs
Full agent path (
ng_run) — runs Claude Code against the Gym MCP server end to end. Requires anAnthropic API key (the agent runs
claudewith an injectedANTHROPIC_API_KEY):A correct rollout shows Claude Code calling
mcp__example_mcp_weather__get_weatherandreward = 1.0.Unit tests
Quick manual MCP check (no LLM) — start the server, then drive
/seed_session → /mcp tools/call → /verify(arequests.Sessionpreserves the session cookie). See the example README for a copy-pastesnippet. This also confirms the endpoint is reachable from a non-loopback host.
Authoring a tool
The model sees a
get_weathertool taking onlycity;session_idis injected from the per-rollouttoken so the tool can record state that
/verifylater reads.Scope note on issue #212's "two flows"
MCPResourcesServer(this PR). This is the flowthat needs new infrastructure, because mounting MCP inside the Resources Server is what lets a tool call
be bound to the rollout's session and thus verified.
directly via a static
mcp_config, and a plainSimpleResourcesServer.verify()scores the trajectory.Gym's session cookie only flows agent↔resources-server, so it does not interfere with the external MCP
connection; verification for external servers is trajectory-based (Gym can't observe their calls). Both
flows — and their combination — are documented in the new tutorial.
Verification
pyteston the new/changed test files: all pass — incl. a regression test driving/mcpwith anon-loopback
Hostheader (asserts no HTTP 421),@gym_toolauto-registration (input schema hidessession_id), reserved-name /request-param rejection, and a differently-cased-city verify case.seed_session → MCP tools/call → verifyreturnsreward=1.0; a different session returns0.0.claudeCLI with a generated--mcp-config: Claude invokes theget_weatherMCP tool and theserver records it in-session.
ruff check/ruff formatclean;fern checkpasses.