diff --git a/ATTRIBUTIONS.md b/ATTRIBUTIONS.md index 0cb5f03315..d095e5ebf2 100644 --- a/ATTRIBUTIONS.md +++ b/ATTRIBUTIONS.md @@ -14,6 +14,7 @@ All third-party software is used as obtained, without modification, unless other | **Hydra-core** | Configuration management | MIT | Facebook Research | https://github.com/facebookresearch/hydra | | **MLflow** | Experiment tracking | Apache-2.0 | Databricks Inc. | https://github.com/mlflow/mlflow | | **Math-Verify** | Math reasoning verifier | Apache-2.0 | Hugging Face | https://github.com/huggingface/Math-Verify | +| **MCP Python SDK** | Model Context Protocol SDK | MIT | Anthropic, PBC | https://github.com/modelcontextprotocol/python-sdk | | **OmegaConf** | Config library | BSD-3-Clause | Omry Yadan | https://github.com/omry/omegaconf | | **OpenAI Python** | API client | Apache-2.0 | OpenAI | https://github.com/openai/openai-python | | **spaCy Model (en_core_web_sm)** | NLP embedding model | MIT | Explosion AI | https://github.com/explosion/spacy-models/releases/tag/en_core_web_sm-3.8.0 | diff --git a/fern/versions/latest/pages/environment-tutorials/integrate-external-environments.mdx b/fern/versions/latest/pages/environment-tutorials/integrate-external-environments.mdx index 2362d620b2..1a8a15a43e 100644 --- a/fern/versions/latest/pages/environment-tutorials/integrate-external-environments.mdx +++ b/fern/versions/latest/pages/environment-tutorials/integrate-external-environments.mdx @@ -1,7 +1,7 @@ --- title: "Integrate external libraries" description: "" -position: 6 +position: 7 --- Fundamentally, a training environment in NeMo Gym is some Python logic that performs a sequence or graph of model calls and tool calls. For native NeMo Gym environments, all of the model and tool calls are present within an agent server like [`simple_agent`](https://github.com/NVIDIA-NeMo/Gym/blob/main/responses_api_agents/simple_agent/app.py). For environments that we consider "external", the orchestration of model and tool calls is offloaded to a third-party library rather than implemented within NeMo Gym itself. diff --git a/fern/versions/latest/pages/environment-tutorials/mcp-resources-server.mdx b/fern/versions/latest/pages/environment-tutorials/mcp-resources-server.mdx new file mode 100644 index 0000000000..fc2bd2af9c --- /dev/null +++ b/fern/versions/latest/pages/environment-tutorials/mcp-resources-server.mdx @@ -0,0 +1,244 @@ +--- +title: "MCP Resources Server" +description: "Expose tools to an agent over the Model Context Protocol (MCP) and verify their use, with a runnable Claude Code example" +position: 5 +--- +import { NavButton } from "../../../../components/NavButton"; + +This tutorial shows how to expose environment tools over the **Model Context Protocol (MCP)** so that an MCP-native agent — such as Claude Code — can discover and call them, while the Resources Server still owns verification. The pattern is: **MCP tool implementations + a `verify()` function = a Resources Server.** + + + +--- + +## Two ways to combine MCP with a Resources Server + +There are two distinct integration shapes, and they need different amounts of plumbing: + +| Flow | When | What you build | +|---|---|---| +| **Gym-owned MCP server** | You want the tools *and* their verification to live in Gym, with per-rollout session isolation | Subclass `MCPResourcesServer` — Gym mounts a Streamable-HTTP MCP endpoint at `/mcp` on the same app as `/seed_session` and `/verify` | +| **Existing / external MCP server** | The MCP server already runs outside Gym (a third-party or shared service) | Point the agent at it directly with a static `mcp_config`; write a plain `SimpleResourcesServer.verify()` that scores the resulting trajectory | + +The rest of this page builds the **Gym-owned** flow (the one that needs new infrastructure) and then explains the **external** flow at the end. + + +**Why a Gym-owned MCP server at all?** Mounting the MCP endpoint inside the Resources Server lets a tool call be bound to the *same per-rollout session* as `/seed_session` and `/verify`. That is what makes "was this tool actually used in this episode?" a verifiable, isolated question. An external MCP server can't offer that — Gym can't observe its calls — so external-server verification has to work off the agent's trajectory instead. + + +--- + +## What You'll Build + +A weather environment with a single MCP tool, `get_weather(city)`. The agent must call the tool and then answer with exactly the sentence the tool returned. The Resources Server rewards the rollout only if the tool was called **in this session** and the final answer contains the returned sentence. + +### Episode Flow + +```text +Goal + - Learn MCP tool usage bound to a Gym session: call an MCP tool, then answer using its result. + +Inputs + - seed input: expected_city (e.g., "Paris") + +Flow (the MCP endpoint and /verify share one session_id) + 1) Agent -> ResourcesServer POST /seed_session {"verifier_metadata": {"expected_city": "Paris"}} + - returns hidden MCP metadata: a per-rollout X-NeMo-Gym-Session-Token bound to this session_id + 2) Agent writes a per-rollout mcp_config and launches Claude Code with --mcp-config + 3) Claude Code -> ResourcesServer POST /mcp (tools/call get_weather, carrying the token header) + - the tool resolves the token back to session_id and records the call + 4) Agent -> ResourcesServer POST /verify {"verifier_metadata": {"expected_city": "Paris"}, "response": ...} + - reward = 1.0 iff the tool was called in this session AND the answer contains the sentence +``` + +--- + +## Implementation + +The base class `MCPResourcesServer` (in `nemo_gym/base_resources_server.py`) mounts the MCP endpoint and manages the per-rollout token. You write a **`@gym_tool` method** (your tool), a `seed_session()` that returns the MCP metadata so the agent can connect, and a `verify()` that scores the rollout. + +**File ([`resources_servers/example_mcp_weather/app.py`](https://github.com/NVIDIA-NeMo/Gym/tree/main/resources_servers/example_mcp_weather/app.py)):** + +```python +# simplified +from typing import Any, Optional + +from fastapi import Request +from pydantic import ConfigDict, Field + +from nemo_gym.base_resources_server import ( + BaseResourcesServerConfig, + BaseSeedSessionRequest, + BaseSeedSessionResponse, + BaseVerifyRequest, + BaseVerifyResponse, + MCPResourcesServer, + MCPServerMetadata, + gym_tool, +) +from nemo_gym.server_utils import SESSION_ID_KEY + + +def _weather_sentence(city: str) -> str: + return f"The weather in {city} is sunny and 72 F." + + +class ExampleMCPWeatherResourcesServerConfig(BaseResourcesServerConfig): + pass + + +class ExampleMCPWeatherSeedSessionRequest(BaseSeedSessionRequest): + model_config = ConfigDict(extra="allow") + # Task ground truth travels in verifier_metadata, e.g. {"expected_city": "Paris"}. + verifier_metadata: Optional[dict[str, Any]] = None + + +# seed_session returns the MCP metadata under the `mcp` key +class ExampleMCPWeatherSeedSessionResponse(BaseSeedSessionResponse): + mcp: MCPServerMetadata + + +class ExampleMCPWeatherVerifyRequest(BaseVerifyRequest): + model_config = ConfigDict(extra="allow") + verifier_metadata: Optional[dict[str, Any]] = None + + +class ExampleMCPWeatherResourcesServer(MCPResourcesServer): + config: ExampleMCPWeatherResourcesServerConfig + session_id_to_state: dict[str, dict[str, Any]] = Field(default_factory=dict) + + async def seed_session( + self, request: Request, body: ExampleMCPWeatherSeedSessionRequest + ) -> ExampleMCPWeatherSeedSessionResponse: + session_id = request.session[SESSION_ID_KEY] + expected_city = (body.verifier_metadata or {}).get("expected_city", "Paris") + self.session_id_to_state[session_id] = {"expected_city": expected_city, "weather_calls": []} + # build_mcp_session_metadata() mints a per-rollout token bound to this session_id + return ExampleMCPWeatherSeedSessionResponse(mcp=self.build_mcp_session_metadata(request)) + + # Decorate a method with @gym_tool and it is auto-registered as an MCP tool named `get_weather`. + # Declare a `session_id: str` param to receive the Gym session; it is injected from the per-rollout + # token and hidden from the tool's input schema, so the model only sees `city`. + @gym_tool + def get_weather(self, session_id: str, city: str) -> str: + """Get a deterministic weather report for a city.""" + state = self.session_id_to_state.setdefault(session_id, {"weather_calls": []}) + weather = _weather_sentence(city) + state["weather_calls"].append({"city": city, "weather": weather}) + return weather + + async def verify( + self, request: Request, body: ExampleMCPWeatherVerifyRequest + ) -> BaseVerifyResponse: + session_id = request.session[SESSION_ID_KEY] + state = self.session_id_to_state.get(session_id, {"weather_calls": []}) + expected_city_value = (body.verifier_metadata or {}).get("expected_city", "Paris") + expected_city = expected_city_value.casefold() + expected = _weather_sentence(expected_city_value) + # reward iff the tool was called for this city in this session AND the final answer repeats it + # (match case-insensitively, so a correct call/answer that used different casing still counts) + tool_called = any(str(c.get("city", "")).casefold() == expected_city for c in state["weather_calls"]) + final_text = _extract_assistant_text(body) # join the assistant message text from body.response + reward = float(tool_called and expected.casefold() in final_text.casefold()) + return BaseVerifyResponse(**body.model_dump(), reward=reward) + + +if __name__ == "__main__": + ExampleMCPWeatherResourcesServer.run_webserver() +``` + +### Key Pattern + +Writing a tool is just decorating a method: + +1. **`@gym_tool`** — mark a method and the base class auto-registers it as an MCP tool (name = method name), mounted at `/mcp` (Streamable HTTP). The MCP input schema is derived from the method's typed parameters. To receive the Gym session, declare a **`session_id: str`** parameter — it is injected from the per-rollout token and **hidden** from the tool's input schema (the model only sees the real args). Omit it for a stateless tool. A missing/invalid token raises `MCPSessionError`, which — because MCP runs over JSON-RPC — FastMCP surfaces to the client as a tool error (`isError: true`) on an HTTP 200 response, not an HTTP status code. Both sync and async methods work. Tool names may not collide with reserved endpoints (`verify`, `seed_session`, `aggregate_metrics`, `mcp`), and a tool must **not** take a `request` parameter (there is no FastAPI `Request` on the MCP path — use `session_id`). +2. **`build_mcp_session_metadata(request)`** — call this from `seed_session` and return it under the response's `mcp` key. It mints the one-time `X-NeMo-Gym-Session-Token` bound to the current `session_id`. + +> Need full control (e.g. a hand-written `@mcp.tool()` with custom schema)? Override `register_mcp_tools(self, mcp)` — call `super().register_mcp_tools(mcp)` first to keep the auto-registered `@gym_tool` ones. + + +`MCPResourcesServer` disables the MCP SDK's default DNS-rebinding protection (`TransportSecuritySettings(enable_dns_rebinding_protection=False)`). That protection only accepts loopback `Host` headers and returns HTTP `421` otherwise — which would break multi-node / `use_absolute_ip=True` deployments where the agent reaches the server by a routable host. The endpoint is instead protected by the per-rollout session token. You don't need to set this yourself; the base class handles it. + + +--- + +## Wiring the agent (Claude Code) + +The [`claude_code_agent`](https://github.com/NVIDIA-NeMo/Gym/tree/main/responses_api_agents/claude_code_agent) reads the `mcp` metadata from `/seed_session`, writes a per-rollout `gym_mcp_config.json`, and launches Claude Code with `--mcp-config`. The generated config looks like: + +```json +{ + "mcpServers": { + "example_mcp_weather": { + "type": "http", + "url": "http://:/mcp", + "headers": { "X-NeMo-Gym-Session-Token": "" } + } + } +} +``` + +A minimal config (`resources_servers/example_mcp_weather/configs/example_mcp_weather.yaml`) wires the server and the agent together: + +```yaml +example_mcp_weather: + resources_servers: + example_mcp_weather: + entrypoint: app.py + domain: agent + +example_mcp_weather_claude_code_agent: + responses_api_agents: + claude_code_agent: + entrypoint: app.py + resources_server: { type: resources_servers, name: example_mcp_weather } + model: claude-sonnet-4-6 + anthropic_api_key: ${anthropic_api_key} + datasets: + - { name: example, type: example, jsonl_fpath: resources_servers/example_mcp_weather/data/example.jsonl } +``` + +### Run it + +Put your key in a repo-root `env.yaml` (the config above interpolates `${anthropic_api_key}`): + +```yaml +anthropic_api_key: sk-ant-... +``` + +Then start the servers: + +```bash +gym env start --config resources_servers/example_mcp_weather/configs/example_mcp_weather.yaml +``` + +Then collect rollouts against the `example` dataset and reward-profile as in the [quickstart](/get-started/quickstart). A correct rollout shows Claude Code calling `mcp__example_mcp_weather__get_weather` and a `reward` of `1.0`. + + +To watch the MCP round-trip without a full `gym env start`, start the Resources Server on its own and drive `/seed_session → /mcp tools/call → /verify` directly (a `requests.Session` preserves the session cookie). This is also the fastest way to confirm the endpoint is reachable from another host. + + +--- + +## Pointing at an existing / external MCP server + +If the MCP server already runs outside Gym, the agent talks to it **directly** — you do not need an `MCPResourcesServer`. Give the agent a static `mcp_config` pointing at the external server, and write a plain `SimpleResourcesServer.verify()` that scores the agent's trajectory: + +```yaml +my_external_mcp_agent: + responses_api_agents: + claude_code_agent: + entrypoint: app.py + resources_server: { type: resources_servers, name: my_verifier } # a SimpleResourcesServer with verify() + mcp_config: /abs/path/to/external_mcp_config.json # static config passed via --mcp-config +``` + +Things to know about this flow: + +- **No cookie/session entanglement.** Gym's session cookie flows only between the agent server and the Resources Server (`/seed_session` ↔ `/verify`). The agent-to-external-MCP connection is a separate channel with its own auth (whatever `headers` you put in the static config). They don't interfere. +- **Verify off the trajectory.** Gym can't observe the external server's calls, so `verify()` must score the `function_call` / `function_call_output` items in the agent's Responses-API output — not server-side session state. +- **Static + per-rollout compose.** When both are present, the agent merges your static `mcp_config` with the per-rollout Gym-owned entry, so a single rollout can use external tools *and* a Gym-owned MCP server at once. If a static server happens to share the **same name** as the Gym resources server, the per-rollout Gym entry takes precedence and overwrites it. + +--- + + diff --git a/fern/versions/latest/pages/environment-tutorials/real-world-environment/index.mdx b/fern/versions/latest/pages/environment-tutorials/real-world-environment/index.mdx index 8ca3f778f0..aada5ed279 100644 --- a/fern/versions/latest/pages/environment-tutorials/real-world-environment/index.mdx +++ b/fern/versions/latest/pages/environment-tutorials/real-world-environment/index.mdx @@ -1,13 +1,13 @@ --- title: "Real-World Environment" description: "" -position: 5 +position: 6 --- import { NavButton } from "../../../../../components/NavButton"; The Workplace Assistant environment simulates an office with email, calendar, analytics, project management, and CRM toolkits. It uses dynamic routing, per-session state, and state-based verification to grade outcomes. - + --- diff --git a/fern/versions/latest/pages/environment-tutorials/stateful-environment.mdx b/fern/versions/latest/pages/environment-tutorials/stateful-environment.mdx index 3503b1f264..fb449eac68 100644 --- a/fern/versions/latest/pages/environment-tutorials/stateful-environment.mdx +++ b/fern/versions/latest/pages/environment-tutorials/stateful-environment.mdx @@ -194,4 +194,4 @@ ResourcesServer: --- - + diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py index 83155df73b..4d519aca7b 100644 --- a/nemo_gym/base_resources_server.py +++ b/nemo_gym/base_resources_server.py @@ -12,10 +12,20 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import functools +import inspect from abc import abstractmethod +from contextlib import asynccontextmanager +from contextvars import ContextVar +from typing import Any, Optional, get_type_hints +from uuid import uuid4 -from fastapi import FastAPI +from fastapi import FastAPI, Request +from itsdangerous import BadSignature, URLSafeSerializer from pydantic import BaseModel +from starlette.concurrency import run_in_threadpool +from starlette.datastructures import Headers +from starlette.routing import Route from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest from nemo_gym.openai_utils import ( @@ -23,7 +33,43 @@ NeMoGymResponseCreateParamsNonStreaming, ) from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics -from nemo_gym.server_utils import BaseRunServerInstanceConfig, BaseServer, SimpleServer +from nemo_gym.server_utils import SESSION_ID_KEY, BaseRunServerInstanceConfig, BaseServer, SimpleServer + + +NEMO_GYM_MCP_SESSION_TOKEN_HEADER = "X-NeMo-Gym-Session-Token" +NEMO_GYM_MCP_METADATA_KEY = "mcp" +_MCP_SESSION_TOKEN: ContextVar[Optional[str]] = ContextVar("nemo_gym_mcp_session_token", default=None) +# Salt namespacing the signed MCP session token, so it can't be confused with another signer +# that happens to share the same session-middleware secret. +_MCP_TOKEN_SALT = "nemo-gym-mcp-session-token" + + +class MCPSessionError(Exception): + """A Gym MCP tool call lacked a valid per-rollout session token. + + Deliberately not an HTTP error: MCP runs over JSON-RPC, so FastMCP returns HTTP 200 and surfaces + this to the client as a tool error (``isError: true``). An HTTP status code raised here would + never reach the caller, so we raise a plain error with a clear message instead. + """ + + +# Names a @gym_tool method may not use, because they collide with the resources server's own +# endpoints (and would silently shadow them on HTTP while still registering as MCP tools). +RESERVED_MCP_TOOL_NAMES = frozenset({"verify", "seed_session", "aggregate_metrics", "mcp"}) + + +def gym_tool(fn): + """Mark a resources-server method as a tool to auto-expose over MCP. + + The method is registered as an MCP tool named after the method, and its MCP input schema is + derived from the method's typed parameters. Declare a ``session_id: str`` parameter to receive + the per-rollout Gym session id; it is injected automatically (from the hidden session token) and + hidden from the tool's input schema. The method must NOT take a ``request`` parameter — there is + no FastAPI ``Request`` on the MCP path; use ``session_id`` instead. Both sync and async methods + are supported. + """ + fn.__gym_tool__ = True + return fn class BaseResourcesServerConfig(BaseRunServerInstanceConfig): @@ -54,6 +100,32 @@ class BaseSeedSessionResponse(BaseModel): pass +class MCPServerMetadata(BaseModel): + """Metadata returned from /seed_session for per-rollout Gym MCP access.""" + + server_name: str + url_path: str = "/mcp" + transport: str = "http" + headers: dict[str, str] + + +class _MCPHeaderSessionMiddleware: + def __init__(self, app: Any): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + token = Headers(scope=scope).get(NEMO_GYM_MCP_SESSION_TOKEN_HEADER) + context_token = _MCP_SESSION_TOKEN.set(token) + try: + await self.app(scope, receive, send) + finally: + _MCP_SESSION_TOKEN.reset(context_token) + + class SimpleResourcesServer(BaseResourcesServer, AggregateMetricsMixin, SimpleServer): config: BaseResourcesServerConfig @@ -86,3 +158,158 @@ async def aggregate_metrics(self, body: AggregateMetricsRequest) -> AggregateMet compute_metrics_fn=self.compute_metrics, get_key_metrics_fn=self.get_key_metrics, ) + + +class MCPResourcesServer(SimpleResourcesServer): + """SimpleResourcesServer variant that also exposes Gym-owned MCP tools. + + 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. + """ + + mcp_url_path: str = "/mcp" + + def setup_webserver(self) -> FastAPI: + app = super().setup_webserver() + + try: + from mcp.server.fastmcp import FastMCP + from mcp.server.transport_security import TransportSecuritySettings + except ImportError as exc: # pragma: no cover - exercised only without the optional runtime dependency + raise RuntimeError( + "MCPResourcesServer requires the official MCP Python SDK. Install the 'mcp' package." + ) from exc + + mcp = FastMCP( + self.config.name or self.__class__.__name__, + stateless_http=True, + json_response=True, + streamable_http_path="/", + # The MCP SDK enables DNS-rebinding protection by default, which only accepts loopback + # Host headers and returns HTTP 421 for anything else. Gym mounts this endpoint for + # server-to-server access: the agent reaches it via the resources server's resolved host, + # which is a routable IP/hostname when use_absolute_ip=True (required for multi-node runs). + # The endpoint is already gated by the per-rollout X-NeMo-Gym-Session-Token, so we disable + # Host/Origin validation to keep MCP tool calls working off-loopback. + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + self.register_mcp_tools(mcp) + + main_app_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def lifespan_wrapper(app: FastAPI): + async with mcp.session_manager.run(): + async with main_app_lifespan(app) as maybe_state: + yield maybe_state + + app.router.lifespan_context = lifespan_wrapper + mcp_app = mcp.streamable_http_app() + streamable_http_route = next(route for route in mcp_app.routes if getattr(route, "path", None) == "/") + + # Mounting serves the slash-suffixed path; this exact route avoids relying on client redirects. + app.router.routes.append( + Route( + self.mcp_url_path, + _MCPHeaderSessionMiddleware(streamable_http_route.endpoint), + include_in_schema=False, + ) + ) + app.mount(self.mcp_url_path, _MCPHeaderSessionMiddleware(mcp_app)) + return app + + def register_mcp_tools(self, mcp: Any) -> None: + """Auto-register methods decorated with ``@gym_tool`` as MCP tools. + + Subclasses can either rely on this default (just decorate tool methods with ``@gym_tool``) or + override it for full manual control. To add manual ``@mcp.tool()`` functions on top of the + auto-registered ones, call ``super().register_mcp_tools(mcp)`` first. + """ + for name, func in inspect.getmembers(type(self), predicate=inspect.isfunction): + if getattr(func, "__gym_tool__", False): + self._register_gym_tool(mcp, name, getattr(self, name)) + + def _register_gym_tool(self, mcp: Any, name: str, method: Any) -> None: + """Register one bound ``@gym_tool`` method as an MCP tool. + + Builds a wrapper whose signature mirrors the method's parameters minus ``session_id`` (so the + session id stays out of the model-visible input schema) and injects the resolved Gym session id + at call time. Enforces the ``@gym_tool`` constraints. + """ + if name in RESERVED_MCP_TOOL_NAMES: + raise ValueError( + f"@gym_tool method {name!r} collides with a reserved endpoint name " + f"{sorted(RESERVED_MCP_TOOL_NAMES)}; rename the tool." + ) + + signature = inspect.signature(method) + hints = get_type_hints(method) + for param_name, param in signature.parameters.items(): + if param_name == "request" or hints.get(param_name, param.annotation) is Request: + raise ValueError( + f"@gym_tool method {name!r} must not take a 'request' parameter; there is no FastAPI " + "Request on the MCP path. Declare a 'session_id: str' parameter to access the Gym session." + ) + + inject_session = "session_id" in signature.parameters + + if inspect.iscoroutinefunction(method): + + @functools.wraps(method) + async def wrapper(**kwargs: Any) -> Any: + if inject_session: + kwargs["session_id"] = self.require_mcp_session_id() + return await method(**kwargs) + else: + + @functools.wraps(method) + async def wrapper(**kwargs: Any) -> Any: + if inject_session: + kwargs["session_id"] = self.require_mcp_session_id() + # Offload blocking sync tools to a thread so they don't stall the event loop + # (which would otherwise block every concurrent rollout in this worker). + return await run_in_threadpool(method, **kwargs) + + # Mirror the method's parameters (with resolved annotations) minus session_id, so FastMCP builds + # the tool's input schema from real types even under ``from __future__ import annotations``. + visible_params = [ + param.replace(annotation=hints.get(param_name, param.annotation)) + for param_name, param in signature.parameters.items() + if param_name != "session_id" + ] + wrapper.__signature__ = signature.replace( + parameters=visible_params, + return_annotation=hints.get("return", signature.return_annotation), + ) + wrapper.__annotations__ = {k: v for k, v in hints.items() if k != "session_id"} + mcp.add_tool(wrapper, name=name, description=(method.__doc__ or "").strip() or None) + + def build_mcp_session_metadata(self, request: Request) -> MCPServerMetadata: + session_id = request.session.get(SESSION_ID_KEY) + if not session_id: + session_id = str(uuid4()) + request.session[SESSION_ID_KEY] = session_id + + return MCPServerMetadata( + server_name=self.config.name or self.__class__.__name__, + url_path=self.mcp_url_path, + headers={NEMO_GYM_MCP_SESSION_TOKEN_HEADER: self._mcp_token_serializer().dumps(session_id)}, + ) + + def _mcp_token_serializer(self) -> URLSafeSerializer: + # Stateless signed token: the session-middleware secret is derived deterministically from the + # server class + config name, so any worker can verify a token another worker signed. This needs + # no per-worker token storage (it works with num_workers > 1, and there is nothing to evict). + return URLSafeSerializer(self.get_session_middleware_key(), salt=_MCP_TOKEN_SALT) + + def require_mcp_session_id(self) -> str: + token = _MCP_SESSION_TOKEN.get() + if not token: + raise MCPSessionError(f"Missing {NEMO_GYM_MCP_SESSION_TOKEN_HEADER} for Gym MCP tool call.") + try: + return self._mcp_token_serializer().loads(token) + except BadSignature as exc: + raise MCPSessionError("Invalid Gym MCP session token.") from exc diff --git a/pyproject.toml b/pyproject.toml index 0605e373ab..83caf14ff1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,11 @@ dependencies = [ # License: MIT https://github.com/fastapi/fastapi/blob/25a3697cedc6e7dfb84e93c8ff965801486f00f4/LICENSE "fastapi", + # MCP Python SDK: Used by Gym-owned Streamable HTTP MCP resources servers. + # Updated Tue Jun 23, 2026 with mcp>=1.27,<2 + # License: MIT https://github.com/modelcontextprotocol/python-sdk/blob/main/LICENSE + "mcp>=1.27,<2", + # itsdangerous: Requirement for Starlette SessionMiddleware # Updated Tue Feb 17, 2026 with itsdangerous==2.2.0 # License: BSD 3-Clause https://github.com/pallets/itsdangerous/blob/672971d66a2ef9f85151e53283113f33d642dabd/LICENSE.txt diff --git a/resources_servers/example_mcp_weather/README.md b/resources_servers/example_mcp_weather/README.md new file mode 100644 index 0000000000..f7702a6755 --- /dev/null +++ b/resources_servers/example_mcp_weather/README.md @@ -0,0 +1,59 @@ +# Example MCP Weather + +A minimal **Gym-owned MCP Resources Server**: it mounts a Streamable-HTTP MCP endpoint at `/mcp` on the +same FastAPI app as `/seed_session` and `/verify`. The `get_weather` MCP tool records its calls against the +per-rollout Gym session (resolved from a hidden `X-NeMo-Gym-Session-Token`), and `/verify` rewards a rollout +only if the tool was used **in that same session** and the final answer repeats the returned sentence. + +This is the runnable companion to the [MCP Resources Server tutorial](https://github.com/NVIDIA-NeMo/Gym/tree/main/fern/versions/latest/pages/environment-tutorials/mcp-resources-server.mdx). + +## Run with an agent (Claude Code) + +Put your key in a repo-root `env.yaml` (the config interpolates `${anthropic_api_key}`), then start the servers +— the `claude_code_agent` runs `claude` with that key injected: + +```bash +# env.yaml: anthropic_api_key: sk-ant-... +gym env start --config resources_servers/example_mcp_weather/configs/example_mcp_weather.yaml +``` + +Then collect rollouts against `data/example.jsonl` and reward-profile as in the +[quickstart](https://github.com/NVIDIA-NeMo/Gym/tree/main/fern/versions/latest/pages/get-started/quickstart.mdx). +A correct rollout shows Claude Code calling `mcp__example_mcp_weather__get_weather` and a `reward` of `1.0`. + +## Inspect the MCP round-trip without an LLM + +Start the server and drive the endpoint directly — a `requests.Session` preserves the session cookie so +`/verify` sees the same session as the tool call: + +```python +import requests + +s = requests.Session() +meta = s.post("http://127.0.0.1:/seed_session", json={"verifier_metadata": {"expected_city": "Paris"}}).json()["mcp"] +token = meta["headers"]["X-NeMo-Gym-Session-Token"] + +# call the MCP tool over the mounted /mcp route, carrying the per-rollout token +s.post( + f"http://127.0.0.1:{meta['url_path']}", + headers={"Accept": "application/json, text/event-stream", "X-NeMo-Gym-Session-Token": token}, + json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": {"name": "get_weather", "arguments": {"city": "Paris"}}}, +) + +# verify in the same session -> reward 1.0 +print(s.post("http://127.0.0.1:/verify", json={ + "responses_create_params": {"input": [{"role": "user", "content": "use the weather tool"}]}, + "verifier_metadata": {"expected_city": "Paris"}, + "response": {"id": "r", "created_at": 0, "model": "t", "object": "response", "output": [ + {"id": "m", "type": "message", "role": "assistant", "status": "completed", + "content": [{"type": "output_text", "text": "The weather in Paris is sunny and 72 F.", "annotations": []}]}], + "parallel_tool_calls": False, "tool_choice": "none", "tools": []}, +}).json()["reward"]) +``` + +## Tests + +```bash +gym env test --resources-server example_mcp_weather +``` diff --git a/resources_servers/example_mcp_weather/__init__.py b/resources_servers/example_mcp_weather/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/resources_servers/example_mcp_weather/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/resources_servers/example_mcp_weather/app.py b/resources_servers/example_mcp_weather/app.py new file mode 100644 index 0000000000..4361a96a3e --- /dev/null +++ b/resources_servers/example_mcp_weather/app.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Optional + +from fastapi import Request +from pydantic import ConfigDict, Field + +from nemo_gym.base_resources_server import ( + BaseResourcesServerConfig, + BaseSeedSessionRequest, + BaseSeedSessionResponse, + BaseVerifyRequest, + BaseVerifyResponse, + MCPResourcesServer, + MCPServerMetadata, + gym_tool, +) +from nemo_gym.server_utils import SESSION_ID_KEY + + +def _weather_sentence(city: str) -> str: + return f"The weather in {city} is sunny and 72 F." + + +def _extract_assistant_text(body: BaseVerifyRequest) -> str: + texts: list[str] = [] + for output_item in body.response.output: + if getattr(output_item, "type", None) != "message" or getattr(output_item, "role", None) != "assistant": + continue + content = getattr(output_item, "content", None) + if isinstance(content, list): + for part in content: + text = getattr(part, "text", None) + if isinstance(text, str): + texts.append(text) + elif isinstance(content, str): + texts.append(content) + return "\n".join(texts).strip() + + +class ExampleMCPWeatherResourcesServerConfig(BaseResourcesServerConfig): + pass + + +class ExampleMCPWeatherSeedSessionRequest(BaseSeedSessionRequest): + model_config = ConfigDict(extra="allow") + + # Task-specific ground truth travels in verifier_metadata (per Gym convention), e.g. + # {"expected_city": "Paris"}. + verifier_metadata: Optional[dict[str, Any]] = None + + +class ExampleMCPWeatherSeedSessionResponse(BaseSeedSessionResponse): + mcp: MCPServerMetadata + + +class ExampleMCPWeatherVerifyRequest(BaseVerifyRequest): + model_config = ConfigDict(extra="allow") + + verifier_metadata: Optional[dict[str, Any]] = None + + +class ExampleMCPWeatherVerifyResponse(BaseVerifyResponse): + model_config = ConfigDict(extra="allow") + + expected_weather: str + tool_call_seen: bool + final_response_mentions_weather: bool + + +class ExampleMCPWeatherResourcesServer(MCPResourcesServer): + config: ExampleMCPWeatherResourcesServerConfig + session_id_to_state: dict[str, dict[str, Any]] = Field(default_factory=dict) + + async def seed_session( + self, + request: Request, + body: ExampleMCPWeatherSeedSessionRequest, + ) -> ExampleMCPWeatherSeedSessionResponse: + session_id = request.session[SESSION_ID_KEY] + expected_city = (body.verifier_metadata or {}).get("expected_city", "Paris") + self.session_id_to_state[session_id] = { + "expected_city": expected_city, + "weather_calls": [], + } + return ExampleMCPWeatherSeedSessionResponse(mcp=self.build_mcp_session_metadata(request)) + + @gym_tool + def get_weather(self, session_id: str, city: str) -> str: + """Get a deterministic weather report for a city.""" + # session_id is injected by the base class (from the per-rollout MCP token); it is hidden from + # the tool's MCP input schema, so the model only sees `city`. + state = self.session_id_to_state.setdefault(session_id, {"weather_calls": []}) + weather = _weather_sentence(city) + state["weather_calls"].append({"city": city, "weather": weather}) + return weather + + async def verify( + self, + request: Request, + body: ExampleMCPWeatherVerifyRequest, + ) -> ExampleMCPWeatherVerifyResponse: + session_id = request.session[SESSION_ID_KEY] + state = self.session_id_to_state.get(session_id, {"weather_calls": []}) + expected_city_value = (body.verifier_metadata or {}).get("expected_city", "Paris") + expected_weather = _weather_sentence(expected_city_value) + expected_city = expected_city_value.casefold() + + # Match the city case-insensitively. The weather sentence is derived deterministically from the + # city, so a city match is sufficient; comparing the sentence exactly would spuriously reject a + # correct call that used different casing (e.g. get_weather("PARIS")). + tool_call_seen = any(str(call.get("city", "")).casefold() == expected_city for call in state["weather_calls"]) + final_text = _extract_assistant_text(body) + final_response_mentions_weather = expected_weather.casefold() in final_text.casefold() + reward = float(tool_call_seen and final_response_mentions_weather) + + return ExampleMCPWeatherVerifyResponse( + **body.model_dump(), + reward=reward, + expected_weather=expected_weather, + tool_call_seen=tool_call_seen, + final_response_mentions_weather=final_response_mentions_weather, + ) + + +if __name__ == "__main__": + ExampleMCPWeatherResourcesServer.run_webserver() diff --git a/resources_servers/example_mcp_weather/configs/example_mcp_weather.yaml b/resources_servers/example_mcp_weather/configs/example_mcp_weather.yaml new file mode 100644 index 0000000000..65d8a75568 --- /dev/null +++ b/resources_servers/example_mcp_weather/configs/example_mcp_weather.yaml @@ -0,0 +1,32 @@ +example_mcp_weather: + resources_servers: + example_mcp_weather: + entrypoint: app.py + domain: agent + verified: false + description: Claude Code MCP smoke test with a Gym-owned weather tool + +example_mcp_weather_claude_code_agent: + responses_api_agents: + claude_code_agent: + entrypoint: app.py + resources_server: + type: resources_servers + name: example_mcp_weather + concurrency: 1 + model: claude-sonnet-4-6 + anthropic_api_key: ${anthropic_api_key} + anthropic_base_url: null + max_turns: 10 + timeout: 300 + system_prompt: null + allowed_tools: null + disallowed_tools: null + claude_code_version: null + bare: true + mcp_config: null + settings: null + datasets: + - name: example + type: example + jsonl_fpath: resources_servers/example_mcp_weather/data/example.jsonl diff --git a/resources_servers/example_mcp_weather/data/example.jsonl b/resources_servers/example_mcp_weather/data/example.jsonl new file mode 100644 index 0000000000..d17ef923bc --- /dev/null +++ b/resources_servers/example_mcp_weather/data/example.jsonl @@ -0,0 +1,5 @@ +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Paris, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "verifier_metadata": {"expected_city": "Paris"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Tokyo, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "verifier_metadata": {"expected_city": "Tokyo"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Seattle, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "verifier_metadata": {"expected_city": "Seattle"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Nairobi, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "verifier_metadata": {"expected_city": "Nairobi"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Toronto, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "verifier_metadata": {"expected_city": "Toronto"}} diff --git a/resources_servers/example_mcp_weather/data/example_metrics.json b/resources_servers/example_mcp_weather/data/example_metrics.json new file mode 100644 index 0000000000..0f93825e33 --- /dev/null +++ b/resources_servers/example_mcp_weather/data/example_metrics.json @@ -0,0 +1,36 @@ +{ + "name": "example", + "type": "example", + "jsonl_fpath": "resources_servers/example_mcp_weather/data/example.jsonl", + "gitlab_identifier": null, + "license": null, + "Number of examples": 5, + "Number of tools": { + "Total # non-null values": 5, + "Average": 0.0, + "Min": 0.0, + "Max": 0.0, + "Standard deviation": 0.0 + }, + "Json-dumped number of words (proxy for token count)": { + "Total # non-null values": 5, + "Average": 27.0, + "Min": 27.0, + "Max": 27.0, + "Standard deviation": 0.0 + }, + "Number of turns": { + "Total # non-null values": 5, + "Average": 1.0, + "Min": 1.0, + "Max": 1.0, + "Standard deviation": 0.0 + }, + "Temperature": { + "Total # non-null values": 0, + "Average": 0.0, + "Min": 0.0, + "Max": 0.0, + "Standard deviation": 0.0 + } +} diff --git a/resources_servers/example_mcp_weather/data/example_rollouts.jsonl b/resources_servers/example_mcp_weather/data/example_rollouts.jsonl new file mode 100644 index 0000000000..ba1488e680 --- /dev/null +++ b/resources_servers/example_mcp_weather/data/example_rollouts.jsonl @@ -0,0 +1,5 @@ +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Paris, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "response": {"id": "resp_example_paris", "created_at": 0, "model": "claude-sonnet-4-6", "object": "response", "output": [{"id": "toolu_weather_paris", "arguments": "{\"city\":\"Paris\"}", "call_id": "toolu_weather_paris", "name": "get_weather", "type": "function_call", "status": "completed"}, {"type": "function_call_output", "call_id": "toolu_weather_paris", "output": "The weather in Paris is sunny and 72 F.", "status": "completed"}, {"id": "msg_weather_paris", "content": [{"text": "The weather in Paris is sunny and 72 F.", "type": "output_text", "annotations": []}], "role": "assistant", "status": "completed", "type": "message"}], "parallel_tool_calls": true, "tool_choice": "auto", "tools": []}, "reward": 1.0, "expected_weather": "The weather in Paris is sunny and 72 F.", "tool_call_seen": true, "final_response_mentions_weather": true, "verifier_metadata": {"expected_city": "Paris"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Tokyo, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "response": {"id": "resp_example_tokyo", "created_at": 0, "model": "claude-sonnet-4-6", "object": "response", "output": [{"id": "toolu_weather_tokyo", "arguments": "{\"city\":\"Tokyo\"}", "call_id": "toolu_weather_tokyo", "name": "get_weather", "type": "function_call", "status": "completed"}, {"type": "function_call_output", "call_id": "toolu_weather_tokyo", "output": "The weather in Tokyo is sunny and 72 F.", "status": "completed"}, {"id": "msg_weather_tokyo", "content": [{"text": "The weather in Tokyo is sunny and 72 F.", "type": "output_text", "annotations": []}], "role": "assistant", "status": "completed", "type": "message"}], "parallel_tool_calls": true, "tool_choice": "auto", "tools": []}, "reward": 1.0, "expected_weather": "The weather in Tokyo is sunny and 72 F.", "tool_call_seen": true, "final_response_mentions_weather": true, "verifier_metadata": {"expected_city": "Tokyo"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Seattle, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "response": {"id": "resp_example_seattle", "created_at": 0, "model": "claude-sonnet-4-6", "object": "response", "output": [{"id": "toolu_weather_seattle", "arguments": "{\"city\":\"Seattle\"}", "call_id": "toolu_weather_seattle", "name": "get_weather", "type": "function_call", "status": "completed"}, {"type": "function_call_output", "call_id": "toolu_weather_seattle", "output": "The weather in Seattle is sunny and 72 F.", "status": "completed"}, {"id": "msg_weather_seattle", "content": [{"text": "The weather in Seattle is sunny and 72 F.", "type": "output_text", "annotations": []}], "role": "assistant", "status": "completed", "type": "message"}], "parallel_tool_calls": true, "tool_choice": "auto", "tools": []}, "reward": 1.0, "expected_weather": "The weather in Seattle is sunny and 72 F.", "tool_call_seen": true, "final_response_mentions_weather": true, "verifier_metadata": {"expected_city": "Seattle"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Nairobi, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "response": {"id": "resp_example_nairobi", "created_at": 0, "model": "claude-sonnet-4-6", "object": "response", "output": [{"id": "toolu_weather_nairobi", "arguments": "{\"city\":\"Nairobi\"}", "call_id": "toolu_weather_nairobi", "name": "get_weather", "type": "function_call", "status": "completed"}, {"type": "function_call_output", "call_id": "toolu_weather_nairobi", "output": "The weather in Nairobi is sunny and 72 F.", "status": "completed"}, {"id": "msg_weather_nairobi", "content": [{"text": "The weather in Nairobi is sunny and 72 F.", "type": "output_text", "annotations": []}], "role": "assistant", "status": "completed", "type": "message"}], "parallel_tool_calls": true, "tool_choice": "auto", "tools": []}, "reward": 1.0, "expected_weather": "The weather in Nairobi is sunny and 72 F.", "tool_call_seen": true, "final_response_mentions_weather": true, "verifier_metadata": {"expected_city": "Nairobi"}} +{"responses_create_params": {"input": [{"role": "user", "content": "Use the get_weather MCP tool for Toronto, then answer with exactly the weather sentence returned by the tool."}], "tools": []}, "response": {"id": "resp_example_toronto", "created_at": 0, "model": "claude-sonnet-4-6", "object": "response", "output": [{"id": "toolu_weather_toronto", "arguments": "{\"city\":\"Toronto\"}", "call_id": "toolu_weather_toronto", "name": "get_weather", "type": "function_call", "status": "completed"}, {"type": "function_call_output", "call_id": "toolu_weather_toronto", "output": "The weather in Toronto is sunny and 72 F.", "status": "completed"}, {"id": "msg_weather_toronto", "content": [{"text": "The weather in Toronto is sunny and 72 F.", "type": "output_text", "annotations": []}], "role": "assistant", "status": "completed", "type": "message"}], "parallel_tool_calls": true, "tool_choice": "auto", "tools": []}, "reward": 1.0, "expected_weather": "The weather in Toronto is sunny and 72 F.", "tool_call_seen": true, "final_response_mentions_weather": true, "verifier_metadata": {"expected_city": "Toronto"}} diff --git a/resources_servers/example_mcp_weather/requirements.txt b/resources_servers/example_mcp_weather/requirements.txt new file mode 100644 index 0000000000..00ed83213e --- /dev/null +++ b/resources_servers/example_mcp_weather/requirements.txt @@ -0,0 +1 @@ +-e nemo-gym[dev] @ ../../ diff --git a/resources_servers/example_mcp_weather/tests/__init__.py b/resources_servers/example_mcp_weather/tests/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/resources_servers/example_mcp_weather/tests/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/resources_servers/example_mcp_weather/tests/test_app.py b/resources_servers/example_mcp_weather/tests/test_app.py new file mode 100644 index 0000000000..4e0be847a6 --- /dev/null +++ b/resources_servers/example_mcp_weather/tests/test_app.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from unittest.mock import MagicMock + +import pytest +from fastapi import Request +from fastapi.testclient import TestClient + +from nemo_gym.base_resources_server import MCPSessionError +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, + NeMoGymResponseOutputMessage, + NeMoGymResponseOutputText, +) +from nemo_gym.server_utils import SESSION_ID_KEY, ServerClient +from resources_servers.example_mcp_weather.app import ( + ExampleMCPWeatherResourcesServer, + ExampleMCPWeatherResourcesServerConfig, + ExampleMCPWeatherSeedSessionRequest, + ExampleMCPWeatherVerifyRequest, +) + + +class FakeMCP: + """Captures tools registered via the FastMCP-style ``add_tool`` API used by gym_tool.""" + + def __init__(self): + self.tools = {} + + def add_tool(self, fn, name=None, description=None): + self.tools[name or fn.__name__] = fn + + +def _server() -> ExampleMCPWeatherResourcesServer: + config = ExampleMCPWeatherResourcesServerConfig( + host="127.0.0.1", + port=12345, + entrypoint="app.py", + name="example_mcp_weather", + ) + return ExampleMCPWeatherResourcesServer(config=config, server_client=MagicMock(spec=ServerClient)) + + +def _request(session_id: str) -> Request: + request = MagicMock(spec=Request) + request.session = {SESSION_ID_KEY: session_id} + return request + + +def _verify_request(expected_city: str, final_text: str) -> ExampleMCPWeatherVerifyRequest: + return ExampleMCPWeatherVerifyRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming( + input=[NeMoGymEasyInputMessage(role="user", content="use the MCP weather tool")] + ), + response=NeMoGymResponse( + id="resp_1", + created_at=0, + model="test", + object="response", + output=[ + NeMoGymResponseOutputMessage( + id="msg_1", + content=[NeMoGymResponseOutputText(text=final_text, annotations=[])], + role="assistant", + status="completed", + type="message", + ) + ], + parallel_tool_calls=False, + tool_choice="none", + tools=[], + ), + verifier_metadata={"expected_city": expected_city}, + ) + + +@pytest.mark.asyncio +async def test_verify_rewards_tool_call_from_same_session() -> None: + server = _server() + seed = await server.seed_session( + _request("session-1"), ExampleMCPWeatherSeedSessionRequest(verifier_metadata={"expected_city": "Paris"}) + ) + token = seed.mcp.headers["X-NeMo-Gym-Session-Token"] + + fake_mcp = FakeMCP() + server.register_mcp_tools(fake_mcp) + + from nemo_gym.base_resources_server import _MCP_SESSION_TOKEN + + # The auto-registered MCP wrapper takes only `city` (session_id is injected from the token) and is + # async (sync tools are offloaded to a threadpool), so it must be awaited. + context_token = _MCP_SESSION_TOKEN.set(token) + try: + assert await fake_mcp.tools["get_weather"](city="Paris") == "The weather in Paris is sunny and 72 F." + finally: + _MCP_SESSION_TOKEN.reset(context_token) + + result = await server.verify( + _request("session-1"), + _verify_request("Paris", "The weather in Paris is sunny and 72 F."), + ) + + assert result.reward == 1.0 + assert result.tool_call_seen is True + assert result.final_response_mentions_weather is True + + +@pytest.mark.asyncio +async def test_verify_accepts_differently_cased_city() -> None: + # A correct tool call that used different casing than the seed city must still be rewarded. + server = _server() + await server.seed_session( + _request("session-1"), ExampleMCPWeatherSeedSessionRequest(verifier_metadata={"expected_city": "Paris"}) + ) + server.session_id_to_state["session-1"]["weather_calls"].append( + {"city": "PARIS", "weather": "The weather in PARIS is sunny and 72 F."} + ) + + result = await server.verify( + _request("session-1"), + _verify_request("Paris", "The weather in PARIS is sunny and 72 F."), + ) + + assert result.reward == 1.0 + assert result.tool_call_seen is True + assert result.final_response_mentions_weather is True + + +@pytest.mark.asyncio +async def test_verify_rejects_tool_call_from_different_session() -> None: + server = _server() + await server.seed_session( + _request("session-1"), ExampleMCPWeatherSeedSessionRequest(verifier_metadata={"expected_city": "Paris"}) + ) + server.session_id_to_state["session-2"] = { + "expected_city": "Paris", + "weather_calls": [{"city": "Paris", "weather": "The weather in Paris is sunny and 72 F."}], + } + + result = await server.verify( + _request("session-1"), + _verify_request("Paris", "The weather in Paris is sunny and 72 F."), + ) + + assert result.reward == 0.0 + assert result.tool_call_seen is False + + +@pytest.mark.asyncio +async def test_mcp_tool_requires_valid_session_token() -> None: + server = _server() + fake_mcp = FakeMCP() + server.register_mcp_tools(fake_mcp) + + from nemo_gym.base_resources_server import _MCP_SESSION_TOKEN + + context_token = _MCP_SESSION_TOKEN.set("invalid-token") + try: + with pytest.raises(MCPSessionError): + await fake_mcp.tools["get_weather"](city="Paris") + finally: + _MCP_SESSION_TOKEN.reset(context_token) + + +def test_streamable_http_mcp_endpoint_records_same_session() -> None: + pytest.importorskip("mcp") + server = _server() + app = server.setup_webserver() + rpc_headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + + with TestClient(app, base_url="http://127.0.0.1:8000") as client: + seed_response = client.post("/seed_session", json={"verifier_metadata": {"expected_city": "Paris"}}) + assert seed_response.status_code == 200 + token = seed_response.json()["mcp"]["headers"]["X-NeMo-Gym-Session-Token"] + + tool_response = client.post( + "/mcp", + headers={**rpc_headers, "X-NeMo-Gym-Session-Token": token}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "get_weather", "arguments": {"city": "Paris"}}, + }, + follow_redirects=False, + ) + + assert tool_response.status_code == 200 + assert tool_response.json()["result"]["structuredContent"]["result"] == ( + "The weather in Paris is sunny and 72 F." + ) + + verify_response = client.post( + "/verify", + json=_verify_request("Paris", "The weather in Paris is sunny and 72 F.").model_dump(mode="json"), + ) + assert verify_response.status_code == 200 + assert verify_response.json()["reward"] == 1.0 + assert verify_response.json()["tool_call_seen"] is True diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index 3acf5ca06a..93be891a33 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -14,11 +14,13 @@ # limitations under the License. import asyncio +import copy import json import logging import os import shutil import subprocess +import tempfile from asyncio import Semaphore from pathlib import Path from time import time @@ -26,9 +28,9 @@ from uuid import uuid4 from fastapi import Request -from pydantic import ConfigDict +from pydantic import ConfigDict, PrivateAttr -from nemo_gym.base_resources_server import BaseRunRequest, BaseVerifyResponse +from nemo_gym.base_resources_server import NEMO_GYM_MCP_METADATA_KEY, BaseRunRequest, BaseVerifyResponse from nemo_gym.base_responses_api_agent import ( BaseResponsesAPIAgentConfig, Body, @@ -213,7 +215,7 @@ def _extract_instruction(body_input) -> tuple[str, Optional[str]]: class ClaudeCodeAgentConfig(BaseResponsesAPIAgentConfig): resources_server: ResourcesServerRef # When model_server is set, ANTHROPIC_BASE_URL is resolved from the Gym model - # server's URL (requires the server to expose POST /v1/messages. None is pushed yet). + # server's URL (requires the server to expose POST /v1/messages). # When None, anthropic_base_url is used directly. model_server: Optional[ModelServerRef] = None concurrency: int = 32 @@ -228,9 +230,9 @@ class ClaudeCodeAgentConfig(BaseResponsesAPIAgentConfig): claude_code_version: Optional[str] = None thinking: Optional[str] = None max_thinking_tokens: Optional[int] = None - # Runtime capability knobs. The default (bare=True, no mcp_config/settings) - # reproduces the original isolated behavior: Claude Code skips auto-discovery - # of skills, hooks, plugins, MCP servers, auto memory, and CLAUDE.md. + # Runtime capability knobs. The default (bare=True, no mcp_config/settings) reproduces the original + # isolated behavior: Claude Code skips hooks, LSP, plugin sync, attribution, auto-memory, background + # prefetches, keychain reads, and CLAUDE.md auto-discovery (skills still resolve via /skill-name). bare: bool = True mcp_config: Optional[str] = None settings: Optional[str] = None @@ -249,6 +251,7 @@ class ClaudeCodeAgentVerifyResponse(BaseVerifyResponse): class ClaudeCodeAgent(SimpleResponsesAPIAgent): config: ClaudeCodeAgentConfig sem: Semaphore = None + _static_mcp_config: Optional[dict[str, Any]] = PrivateAttr(default=None) model_config = ConfigDict(arbitrary_types_allowed=True) def model_post_init(self, __context: Any) -> None: @@ -301,12 +304,19 @@ def _setup_config_dir(self) -> Path: (claude_config_dir / "settings.json").write_text(json.dumps(self._build_settings())) return claude_config_dir - def _build_command(self, model: str, instruction: str, system_prompt: Optional[str] = None) -> list[str]: + def _build_command( + self, + model: str, + instruction: str, + system_prompt: Optional[str] = None, + mcp_config: Optional[str] = None, + ) -> list[str]: """Construct the ``claude`` CLI argv from config. - ``--bare`` is only passed when ``config.bare`` is True; it disables auto-discovery of - skills, hooks, plugins, MCP servers, auto memory, and CLAUDE.md. Explicit capabilities - like ``--mcp-config`` are passed regardless of ``--bare`` since they are not auto-discovered. + ``--bare`` is only passed when ``config.bare`` is True; it skips hooks, LSP, plugin sync, + attribution, auto-memory, background prefetches, keychain reads, and CLAUDE.md auto-discovery + (skills still resolve via /skill-name). Explicit capabilities like ``--mcp-config`` are passed + regardless of ``--bare`` since they are not auto-discovered. """ cmd = [ "claude", @@ -319,8 +329,9 @@ def _build_command(self, model: str, instruction: str, system_prompt: Optional[s if self.config.bare: cmd.append("--bare") cmd += ["--max-turns", str(self.config.max_turns), "--model", model] - if self.config.mcp_config: - cmd += ["--mcp-config", self.config.mcp_config] + effective_mcp_config = mcp_config if mcp_config is not None else self.config.mcp_config + if effective_mcp_config: + cmd += ["--mcp-config", effective_mcp_config] if system_prompt: cmd += ["--append-system-prompt", system_prompt] if self.config.allowed_tools: @@ -334,7 +345,12 @@ def _build_command(self, model: str, instruction: str, system_prompt: Optional[s cmd += ["--", instruction] return cmd - async def _run_claude_code(self, instruction: str, system_prompt: Optional[str] = None) -> tuple[str, str]: + async def _run_claude_code( + self, + instruction: str, + system_prompt: Optional[str] = None, + mcp_config: Optional[str] = None, + ) -> tuple[str, str]: """Run claude -p --output-format=stream-json and return (stdout, model_name).""" base_url = self._resolve_base_url() # Keep full model name for local/custom endpoints; strip provider prefix for real Anthropic API. @@ -358,7 +374,7 @@ async def _run_claude_code(self, instruction: str, system_prompt: Optional[str] env["ANTHROPIC_BASE_URL"] = base_url env["ANTHROPIC_AUTH_TOKEN"] = api_key or "local" - cmd = self._build_command(model, instruction, system_prompt=system_prompt) + cmd = self._build_command(model, instruction, system_prompt=system_prompt, mcp_config=mcp_config) proc = await asyncio.create_subprocess_exec( *cmd, @@ -382,10 +398,70 @@ async def _run_claude_code(self, instruction: str, system_prompt: Optional[str] finally: shutil.rmtree(claude_config_dir, ignore_errors=True) - async def responses( + def _resources_server_base_url(self) -> str: + cfg = get_first_server_config_dict( + self.server_client.global_config_dict, + self.config.resources_server.name, + ) + return self.server_client._build_server_base_url(cfg) + + def _load_static_mcp_config(self) -> dict[str, Any]: + if not self.config.mcp_config: + return {"mcpServers": {}} + + config_path = Path(self.config.mcp_config).expanduser() + config = json.loads(config_path.read_text()) + if not isinstance(config, dict): + raise ValueError(f"Claude Code mcp_config must be a JSON object: {config_path}") + mcp_servers = config.setdefault("mcpServers", {}) + if not isinstance(mcp_servers, dict): + raise ValueError(f"Claude Code mcp_config has non-object mcpServers: {config_path}") + return config + + def _get_static_mcp_config(self) -> dict[str, Any]: + # The static mcp_config is immutable, so read it from disk at most once and reuse the cached + # copy for every rollout instead of re-reading the file each time. + if self._static_mcp_config is None: + self._static_mcp_config = self._load_static_mcp_config() + return self._static_mcp_config + + def _write_rollout_mcp_config(self, seed_response_json: dict[str, Any], output_dir: Path) -> Optional[str]: + metadata = seed_response_json.get(NEMO_GYM_MCP_METADATA_KEY) + if not isinstance(metadata, dict): + return None + + server_name = metadata.get("server_name") or self.config.resources_server.name + url_path = str(metadata.get("url_path") or "/mcp") + url = f"{self._resources_server_base_url().rstrip('/')}/{url_path.lstrip('/')}" + + entry: dict[str, Any] = { + "type": metadata.get("transport") or "http", + "url": url, + } + headers = metadata.get("headers") + if isinstance(headers, dict) and headers: + entry["headers"] = {str(key): str(value) for key, value in headers.items()} + else: + LOG.warning( + "MCP seed metadata for %r has no headers; the tool endpoint will be called without a " + "session token and will reject the calls.", + server_name, + ) + + # Start from a copy of the (cached) static config and add the per-rollout Gym entry. If a static + # mcp_config server already uses this name, the per-rollout Gym entry takes precedence over it. + config = copy.deepcopy(self._get_static_mcp_config()) + config.setdefault("mcpServers", {})[str(server_name)] = entry + + output_dir.mkdir(parents=True, exist_ok=True) + config_path = output_dir / "gym_mcp_config.json" + config_path.write_text(json.dumps(config, indent=2, sort_keys=True)) + return str(config_path) + + async def _create_response( self, - request: Request, - body: NeMoGymResponseCreateParamsNonStreaming = Body(), + body: NeMoGymResponseCreateParamsNonStreaming, + mcp_config: Optional[str] = None, ) -> NeMoGymResponse: body = body.model_copy(deep=True) if isinstance(body.input, str): @@ -395,7 +471,11 @@ async def responses( system_parts = [p for p in [self.config.system_prompt, input_system] if p] system_prompt = "\n\n".join(system_parts) if system_parts else None - stdout, model_name = await self._run_claude_code(user_message, system_prompt=system_prompt) + stdout, model_name = await self._run_claude_code( + user_message, + system_prompt=system_prompt, + mcp_config=mcp_config, + ) output_items, usage = parse_stream_json(stdout) if not any( @@ -434,6 +514,13 @@ async def responses( ), ) + async def responses( + self, + request: Request, + body: NeMoGymResponseCreateParamsNonStreaming = Body(), + ) -> NeMoGymResponse: + return await self._create_response(body) + async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> ClaudeCodeAgentVerifyResponse: async with self.sem: cookies = request.cookies @@ -446,16 +533,12 @@ async def run(self, request: Request, body: ClaudeCodeAgentRunRequest) -> Claude ) await raise_for_status(seed_resp) cookies = seed_resp.cookies + seed_resp_json = await get_response_json(seed_resp) - agent_resp = await self.server_client.post( - server_name=self.config.name, - url_path="/v1/responses", - json=body.responses_create_params, - cookies=cookies, - ) - await raise_for_status(agent_resp) - cookies = agent_resp.cookies - agent_resp_json = await get_response_json(agent_resp) + with tempfile.TemporaryDirectory(prefix="nemo_gym_claude_mcp_") as mcp_config_dir: + mcp_config = self._write_rollout_mcp_config(seed_resp_json, Path(mcp_config_dir)) + agent_resp = await self._create_response(body.responses_create_params, mcp_config=mcp_config) + agent_resp_json = agent_resp.model_dump(mode="json") verify_resp = await self.server_client.post( server_name=self.config.resources_server.name, diff --git a/responses_api_agents/claude_code_agent/tests/test_app.py b/responses_api_agents/claude_code_agent/tests/test_app.py index dc9f6caeb0..ae024d5fdc 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -19,10 +19,12 @@ from unittest.mock import MagicMock, patch import yaml +from fastapi import Request from nemo_gym.openai_utils import ( NeMoGymEasyInputMessage, NeMoGymFunctionCallOutput, + NeMoGymResponseCreateParamsNonStreaming, NeMoGymResponseFunctionToolCall, NeMoGymResponseOutputMessage, ) @@ -30,6 +32,7 @@ from responses_api_agents.claude_code_agent.app import ( ClaudeCodeAgent, ClaudeCodeAgentConfig, + ClaudeCodeAgentRunRequest, ResourcesServerRef, _extract_instruction, parse_stream_json, @@ -37,27 +40,38 @@ def _config(**kwargs) -> ClaudeCodeAgentConfig: + kwargs.setdefault("resources_server", ResourcesServerRef(type="resources_servers", name="")) return ClaudeCodeAgentConfig( host="0.0.0.0", port=8080, entrypoint="", name="", - resources_server=ResourcesServerRef(type="resources_servers", name=""), **kwargs, ) def _make_agent(**kwargs) -> ClaudeCodeAgent: - with patch("responses_api_agents.claude_code_agent.app.ClaudeCodeAgent.model_post_init"): - agent = ClaudeCodeAgent(config=_config(**kwargs), server_client=MagicMock(spec=ServerClient)) - agent.sem = asyncio.Semaphore(agent.config.concurrency) - return agent + # Patch only the external side effect (claude-code install/version check) so the real + # model_post_init still runs — it initializes the model's private attrs and the semaphore. + with patch("responses_api_agents.claude_code_agent.app.ensure_claude_code"): + return ClaudeCodeAgent(config=_config(**kwargs), server_client=MagicMock(spec=ServerClient)) def _event(type_: str, **kwargs) -> str: return json.dumps({"type": type_, **kwargs}) +class FakeAioHTTPResponse: + ok = True + + def __init__(self, payload: dict, cookies: dict | None = None): + self.payload = payload + self.cookies = cookies or {} + + async def read(self) -> bytes: + return json.dumps(self.payload).encode() + + class TestSanity: def test_config_defaults(self) -> None: cfg = _config() @@ -106,6 +120,11 @@ def test_mcp_config_passed_independently_of_bare(self) -> None: assert "--bare" in cmd assert cmd[cmd.index("--mcp-config") + 1] == "/path/to/mcp.json" + def test_dynamic_mcp_config_overrides_static_for_command(self) -> None: + agent = _make_agent(mcp_config="/path/to/static.json") + cmd = agent._build_command("m", "x", mcp_config="/tmp/dynamic.json") + assert cmd[cmd.index("--mcp-config") + 1] == "/tmp/dynamic.json" + def test_optional_flags_threaded_through(self) -> None: agent = _make_agent( allowed_tools="Bash,Read", @@ -235,6 +254,188 @@ async def fake_wait_for(coro, timeout): assert model == "claude-sonnet-4-6" +class TestRolloutMCPConfig: + def test_no_metadata_preserves_static_config(self, tmp_path: Path) -> None: + agent = _make_agent(mcp_config="/path/to/static.json") + assert agent._write_rollout_mcp_config({}, tmp_path) is None + + def test_writes_rollout_mcp_config_with_session_header(self, tmp_path: Path) -> None: + agent = _make_agent(resources_server=ResourcesServerRef(type="resources_servers", name="example_mcp_weather")) + agent.server_client.global_config_dict = { + "example_mcp_weather": { + "resources_servers": { + "example_mcp_weather": { + "host": "127.0.0.1", + "port": 8123, + } + } + } + } + agent.server_client._build_server_base_url.side_effect = lambda cfg: f"http://{cfg['host']}:{cfg['port']}" + + config_path = agent._write_rollout_mcp_config( + { + "mcp": { + "server_name": "example_mcp_weather", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "secret-token"}, + } + }, + tmp_path, + ) + + assert config_path is not None + config = json.loads(Path(config_path).read_text()) + server = config["mcpServers"]["example_mcp_weather"] + assert server["type"] == "http" + assert server["url"] == "http://127.0.0.1:8123/mcp" + assert server["headers"]["X-NeMo-Gym-Session-Token"] == "secret-token" + + def test_merges_static_mcp_config_when_metadata_present(self, tmp_path: Path) -> None: + static_config = tmp_path / "static_mcp.json" + static_config.write_text(json.dumps({"mcpServers": {"static": {"type": "stdio", "command": "server"}}})) + agent = _make_agent( + mcp_config=str(static_config), + resources_server=ResourcesServerRef(type="resources_servers", name="example_mcp_weather"), + ) + agent.server_client.global_config_dict = { + "example_mcp_weather": { + "resources_servers": { + "example_mcp_weather": { + "host": "127.0.0.1", + "port": 8123, + } + } + } + } + agent.server_client._build_server_base_url.side_effect = lambda cfg: f"http://{cfg['host']}:{cfg['port']}" + + config_path = agent._write_rollout_mcp_config( + { + "mcp": { + "server_name": "dynamic", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "tok"}, + } + }, + tmp_path / "run", + ) + + config = json.loads(Path(config_path).read_text()) + assert "static" in config["mcpServers"] + assert config["mcpServers"]["dynamic"]["headers"]["X-NeMo-Gym-Session-Token"] == "tok" + + def test_run_passes_generated_mcp_config(self, tmp_path: Path) -> None: + agent = _make_agent(resources_server=ResourcesServerRef(type="resources_servers", name="example_mcp_weather")) + agent.server_client.global_config_dict = { + "example_mcp_weather": { + "resources_servers": { + "example_mcp_weather": { + "host": "127.0.0.1", + "port": 8123, + } + } + } + } + agent.server_client._build_server_base_url.side_effect = lambda cfg: f"http://{cfg['host']}:{cfg['port']}" + + async def fake_post(server_name, url_path, json=None, cookies=None): + if url_path == "/seed_session": + return FakeAioHTTPResponse( + { + "mcp": { + "server_name": "example_mcp_weather", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "tok"}, + } + }, + cookies={"session": "abc"}, + ) + if url_path == "/verify": + return FakeAioHTTPResponse(json | {"reward": 1.0}) + raise AssertionError(f"unexpected post: {server_name} {url_path}") + + captured: dict = {} + + async def fake_run_claude_code(instruction, system_prompt=None, mcp_config=None): + captured["instruction"] = instruction + captured["mcp_config"] = mcp_config + captured["config_exists_during_run"] = Path(mcp_config).is_file() + captured["config"] = json.loads(Path(mcp_config).read_text()) + return _event( + "assistant", + message={"content": [{"type": "text", "text": "The weather in Paris is sunny and 72 F."}]}, + ), "claude-sonnet-4-6" + + agent.server_client.post.side_effect = fake_post + object.__setattr__(agent, "_run_claude_code", fake_run_claude_code) + request = MagicMock(spec=Request) + request.cookies = {} + body = ClaudeCodeAgentRunRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming(input="use the weather tool"), + expected_city="Paris", + ) + + result = asyncio.run(agent.run(request, body)) + + assert result.reward == 1.0 + assert captured["instruction"] == "use the weather tool" + assert captured["config_exists_during_run"] is True + server = captured["config"]["mcpServers"]["example_mcp_weather"] + assert server["url"] == "http://127.0.0.1:8123/mcp" + assert server["headers"]["X-NeMo-Gym-Session-Token"] == "tok" + assert not Path(captured["mcp_config"]).exists() + + def test_run_threads_session_cookie_seed_to_verify(self, tmp_path: Path) -> None: + agent = _make_agent(resources_server=ResourcesServerRef(type="resources_servers", name="example_mcp_weather")) + agent.server_client.global_config_dict = { + "example_mcp_weather": {"resources_servers": {"example_mcp_weather": {"host": "127.0.0.1", "port": 8123}}} + } + agent.server_client._build_server_base_url.side_effect = lambda cfg: f"http://{cfg['host']}:{cfg['port']}" + + captured: dict = {} + + async def fake_post(server_name, url_path, json=None, cookies=None): + if url_path == "/seed_session": + # the resources server sets a session cookie on the seed response + return FakeAioHTTPResponse( + { + "mcp": { + "server_name": "example_mcp_weather", + "url_path": "/mcp", + "headers": {"X-NeMo-Gym-Session-Token": "tok"}, + } + }, + cookies={"session": "sess-cookie"}, + ) + if url_path == "/verify": + captured["verify_cookies"] = cookies + return FakeAioHTTPResponse(json | {"reward": 1.0}) + raise AssertionError(f"unexpected post: {server_name} {url_path}") + + async def fake_run_claude_code(instruction, system_prompt=None, mcp_config=None): + captured["config_token"] = json.loads(Path(mcp_config).read_text())["mcpServers"]["example_mcp_weather"][ + "headers" + ]["X-NeMo-Gym-Session-Token"] + return _event("assistant", message={"content": [{"type": "text", "text": "ok"}]}), "claude-sonnet-4-6" + + agent.server_client.post.side_effect = fake_post + object.__setattr__(agent, "_run_claude_code", fake_run_claude_code) + request = MagicMock(spec=Request) + request.cookies = {} + body = ClaudeCodeAgentRunRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming(input="use the weather tool"), + verifier_metadata={"expected_city": "Paris"}, + ) + + asyncio.run(agent.run(request, body)) + + # the cookie set on /seed_session is threaded into the /verify call (same rollout session), + # and the per-rollout token from seed metadata reaches the generated MCP config. + assert captured["verify_cookies"] == {"session": "sess-cookie"} + assert captured["config_token"] == "tok" + + class TestExtractInstruction: def test_user_only(self) -> None: items = [NeMoGymEasyInputMessage(role="user", content="hello")] diff --git a/tests/unit_tests/test_base_resources_server.py b/tests/unit_tests/test_base_resources_server.py index 6dd00dc6b5..076e1fe979 100644 --- a/tests/unit_tests/test_base_resources_server.py +++ b/tests/unit_tests/test_base_resources_server.py @@ -14,11 +14,20 @@ # limitations under the License. from unittest.mock import MagicMock +import pytest +from fastapi import Request + from nemo_gym.base_resources_server import ( BaseResourcesServerConfig, + BaseSeedSessionRequest, + BaseSeedSessionResponse, + MCPResourcesServer, + MCPServerMetadata, + MCPSessionError, SimpleResourcesServer, + gym_tool, ) -from nemo_gym.server_utils import ServerClient +from nemo_gym.server_utils import SESSION_ID_KEY, ServerClient class TestBaseResourcesServer: @@ -31,3 +40,266 @@ async def verify(self, body): agent = TestSimpleResourcesServer(config=config, server_client=MagicMock(spec=ServerClient)) agent.setup_webserver() + + +class TestMCPResourcesServer: + def test_mounts_mcp_endpoint_with_normal_gym_endpoints(self) -> None: + pytest.importorskip("mcp") + config = BaseResourcesServerConfig(host="", port=0, entrypoint="", name="test_mcp_resources_server") + + class TestMCPServer(MCPResourcesServer): + def register_mcp_tools(self, mcp): + @mcp.tool() + def ping() -> str: + return "pong" + + async def verify(self, body): + pass + + server = TestMCPServer(config=config, server_client=MagicMock(spec=ServerClient)) + app = server.setup_webserver() + paths = {getattr(route, "path", None) for route in app.routes} + + assert "/seed_session" in paths + assert "/verify" in paths + assert "/aggregate_metrics" in paths + assert "/mcp" in paths + + def test_build_mcp_session_metadata_maps_token_to_session_id(self) -> None: + pytest.importorskip("mcp") + config = BaseResourcesServerConfig(host="", port=0, entrypoint="", name="test_mcp_resources_server") + + class TestMCPServer(MCPResourcesServer): + def register_mcp_tools(self, mcp): + pass + + async def verify(self, body): + pass + + server = TestMCPServer(config=config, server_client=MagicMock(spec=ServerClient)) + request = MagicMock(spec=Request) + request.session = {SESSION_ID_KEY: "gym-session-1"} + + metadata = server.build_mcp_session_metadata(request) + token = metadata.headers["X-NeMo-Gym-Session-Token"] + + assert metadata.server_name == "test_mcp_resources_server" + assert metadata.url_path == "/mcp" + # The signed token round-trips back to the session id (no server-side storage). + from nemo_gym.base_resources_server import _MCP_SESSION_TOKEN + + ctx = _MCP_SESSION_TOKEN.set(token) + try: + assert server.require_mcp_session_id() == "gym-session-1" + finally: + _MCP_SESSION_TOKEN.reset(ctx) + + def test_missing_mcp_session_token_raises(self) -> None: + pytest.importorskip("mcp") + config = BaseResourcesServerConfig(host="", port=0, entrypoint="", name="test_mcp_resources_server") + + class TestMCPServer(MCPResourcesServer): + def register_mcp_tools(self, mcp): + pass + + async def verify(self, body): + pass + + server = TestMCPServer(config=config, server_client=MagicMock(spec=ServerClient)) + + with pytest.raises(MCPSessionError): + server.require_mcp_session_id() + + def test_invalid_mcp_session_token_raises(self) -> None: + pytest.importorskip("mcp") + config = BaseResourcesServerConfig(host="", port=0, entrypoint="", name="test_mcp_resources_server") + + class TestMCPServer(MCPResourcesServer): + def register_mcp_tools(self, mcp): + pass + + async def verify(self, body): + pass + + server = TestMCPServer(config=config, server_client=MagicMock(spec=ServerClient)) + + from nemo_gym.base_resources_server import _MCP_SESSION_TOKEN + + context_token = _MCP_SESSION_TOKEN.set("bad-token") + try: + with pytest.raises(MCPSessionError): + server.require_mcp_session_id() + finally: + _MCP_SESSION_TOKEN.reset(context_token) + + def test_mcp_endpoint_accepts_non_loopback_host(self) -> None: + """Regression: the MCP SDK's default DNS-rebinding protection returns HTTP 421 for any + non-loopback Host header, which breaks multi-node/absolute-IP deployments. MCPResourcesServer + must disable it so server-to-server MCP calls keep working off-loopback.""" + pytest.importorskip("mcp") + from fastapi.testclient import TestClient + + config = BaseResourcesServerConfig(host="", port=0, entrypoint="", name="test_mcp_resources_server") + + class TestMCPServer(MCPResourcesServer): + def register_mcp_tools(self, mcp): + @mcp.tool() + def ping() -> str: + return "pong" + + async def verify(self, body): + pass + + server = TestMCPServer(config=config, server_client=MagicMock(spec=ServerClient)) + app = server.setup_webserver() + + with TestClient(app, base_url="http://127.0.0.1:8000") as client: + resp = client.post( + "/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + # A routable, non-loopback host as seen on a multi-node deployment. + "Host": "10.20.30.40:8000", + }, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "ping", "arguments": {}}, + }, + follow_redirects=False, + ) + + assert resp.status_code != 421, "MCP endpoint rejected a non-loopback Host (DNS-rebinding protection)" + assert resp.status_code == 200 + assert resp.json()["result"]["structuredContent"]["result"] == "pong" + + +class _GymToolSeedResponse(BaseSeedSessionResponse): + mcp: MCPServerMetadata + + +class _GymToolServer(MCPResourcesServer): + """Exercises the @gym_tool auto-registration: a session-bound tool and a stateless one.""" + + async def seed_session(self, request: Request, body: BaseSeedSessionRequest) -> _GymToolSeedResponse: + return _GymToolSeedResponse(mcp=self.build_mcp_session_metadata(request)) + + @gym_tool + def echo(self, session_id: str, text: str) -> str: + """Echo text tagged with the session id.""" + return f"{session_id}:{text}" + + @gym_tool + def add(self, a: int, b: int) -> int: + """Add two numbers (stateless — no session_id).""" + return a + b + + async def verify(self, body): + pass + + +def _gym_tool_server() -> _GymToolServer: + config = BaseResourcesServerConfig(host="", port=0, entrypoint="", name="test_gym_tool_server") + return _GymToolServer(config=config, server_client=MagicMock(spec=ServerClient)) + + +class TestGymToolAutoRegistration: + def test_decorated_methods_auto_register_over_mcp_with_session_hidden(self) -> None: + pytest.importorskip("mcp") + from fastapi.testclient import TestClient + + server = _gym_tool_server() + app = server.setup_webserver() + rpc_headers = {"Accept": "application/json, text/event-stream", "Content-Type": "application/json"} + + with TestClient(app, base_url="http://127.0.0.1:8000") as client: + token = client.post("/seed_session", json={}).json()["mcp"]["headers"]["X-NeMo-Gym-Session-Token"] + + listing = client.post( + "/mcp", + headers=rpc_headers, + json={"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}}, + follow_redirects=False, + ) + tools = {t["name"]: t for t in listing.json()["result"]["tools"]} + assert set(tools) == {"echo", "add"} + # session_id is injected, never surfaced to the model + assert set(tools["echo"]["inputSchema"]["properties"]) == {"text"} + assert set(tools["add"]["inputSchema"]["properties"]) == {"a", "b"} + assert tools["echo"]["description"] == "Echo text tagged with the session id." + + # the session-bound tool resolves session_id from the token + echoed = client.post( + "/mcp", + headers={**rpc_headers, "X-NeMo-Gym-Session-Token": token}, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"text": "hi"}}, + }, + follow_redirects=False, + ) + assert echoed.json()["result"]["structuredContent"]["result"].endswith(":hi") + + # the stateless tool needs no token + summed = client.post( + "/mcp", + headers={**rpc_headers, "X-NeMo-Gym-Session-Token": token}, + json={ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": {"name": "add", "arguments": {"a": 2, "b": 3}}, + }, + follow_redirects=False, + ) + assert summed.json()["result"]["structuredContent"]["result"] == 5 + + def test_missing_token_surfaces_as_clean_tool_error(self) -> None: + """A session-bound tool called without a token must come back as a clean MCP tool error + (HTTP 200, isError) — not an HTTP 401, and without leaking the raw status into the message.""" + pytest.importorskip("mcp") + from fastapi.testclient import TestClient + + server = _gym_tool_server() + app = server.setup_webserver() + rpc_headers = {"Accept": "application/json, text/event-stream", "Content-Type": "application/json"} + + with TestClient(app, base_url="http://127.0.0.1:8000") as client: + resp = client.post( + "/mcp", + headers=rpc_headers, # note: no X-NeMo-Gym-Session-Token + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": "echo", "arguments": {"text": "hi"}}, + }, + follow_redirects=False, + ) + + assert resp.status_code == 200 # MCP/JSON-RPC: transport succeeds; the failure is in the body + result = resp.json()["result"] + assert result["isError"] is True + text = result["content"][0]["text"] + assert "X-NeMo-Gym-Session-Token" in text # clean, specific message + assert "401" not in text # no leaked HTTP status code + + def test_rejects_reserved_tool_name(self) -> None: + pytest.importorskip("mcp") + server = _gym_tool_server() + with pytest.raises(ValueError, match="reserved endpoint name"): + server._register_gym_tool(MagicMock(), "aggregate_metrics", lambda **_: None) + + def test_rejects_request_parameter(self) -> None: + pytest.importorskip("mcp") + server = _gym_tool_server() + + def needs_request(request: Request, city: str) -> str: + return city + + with pytest.raises(ValueError, match="must not take a 'request' parameter"): + server._register_gym_tool(MagicMock(), "bad_tool", needs_request) diff --git a/uv.lock b/uv.lock index a8f5f76a10..71de4a036a 100644 --- a/uv.lock +++ b/uv.lock @@ -2,8 +2,12 @@ version = 1 revision = 3 requires-python = ">=3.12" resolution-markers = [ - "python_full_version >= '3.13' and sys_platform == 'linux'", - "python_full_version >= '3.13' and sys_platform != 'linux'", + "python_full_version >= '3.14' and sys_platform == 'linux'", + "python_full_version == '3.13.*' and sys_platform == 'linux'", + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform != 'linux' and sys_platform != 'win32'", "python_full_version < '3.13' and sys_platform == 'linux'", "python_full_version < '3.13' and sys_platform != 'linux'", ] @@ -900,6 +904,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, ] +[[package]] +name = "httpx-sse" +version = "0.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, +] + [[package]] name = "huggingface-hub" version = "0.34.4" @@ -1127,6 +1140,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4f/65/6079a46068dfceaeabb5dcad6d674f5f5c61a6fa5673746f42a9f4c233b3/MarkupSafe-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:e444a31f8db13eb18ada366ab3cf45fd4b31e4db1236a4448f68778c1d1a5a2f", size = 15739, upload-time = "2024-10-18T15:21:42.784Z" }, ] +[[package]] +name = "mcp" +version = "1.28.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "httpx" }, + { name = "httpx-sse" }, + { name = "jsonschema" }, + { name = "pydantic" }, + { name = "pydantic-settings" }, + { name = "pyjwt" }, + { name = "python-multipart" }, + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "sse-starlette" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, + { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c1/ee/94c6c50ffc5b5cf4737052275d11b57367f32d1a8516e31dcd60591b3916/mcp-1.28.0.tar.gz", hash = "sha256:559d3f9943674cafbe5744c5d3794f3237e8b47f9bbc58e20c0fad680d8487c2", size = 636040, upload-time = "2026-06-16T21:37:17.996Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/e1/4c1dc1fbb688641a712d34650c3d58bbbdcb314ddb75bc5817bbf33515a4/mcp-1.28.0-py3-none-any.whl", hash = "sha256:9c1e7cf3a9125557e418ecd4fed8e9adddce81b0dfdae4d6601d700f5beb71a4", size = 221959, upload-time = "2026-06-16T21:37:16.579Z" }, +] + [[package]] name = "mdit-py-plugins" version = "0.5.0" @@ -1412,6 +1450,7 @@ dependencies = [ { name = "gprof2dot" }, { name = "hydra-core" }, { name = "itsdangerous" }, + { name = "mcp" }, { name = "mlflow" }, { name = "mlflow-skinny" }, { name = "omegaconf" }, @@ -1476,6 +1515,7 @@ requires-dist = [ { name = "gprof2dot" }, { name = "hydra-core" }, { name = "itsdangerous" }, + { name = "mcp", specifier = ">=1.27,<2" }, { name = "mlflow", specifier = ">=3.14.0" }, { name = "mlflow-skinny", specifier = ">=3.14.0" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" }, @@ -2149,6 +2189,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, ] +[[package]] +name = "pydantic-settings" +version = "2.14.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5c/b5/8f48e906c3e0205276e8bd8cb7512217a87b2685304d64be27cad5b3019f/pydantic_settings-2.14.2.tar.gz", hash = "sha256:c19dd64b19097f1de80184f0cc7b0272a13ae6e170cbf240a3e27e381ed14a5f", size = 237700, upload-time = "2026-06-19T13:44:56.324Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/c1/6e422f34e569cf8e18df68d1939c81c099d2b61e4f7d9621c8a77560799c/pydantic_settings-2.14.2-py3-none-any.whl", hash = "sha256:a20c97b37910b6550d5ea50fbcc2d4187defe58cd57070b73863d069419c9440", size = 61715, upload-time = "2026-06-19T13:44:55.02Z" }, +] + [[package]] name = "pydata-sphinx-theme" version = "0.16.1" @@ -2188,6 +2242,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, ] +[[package]] +name = "pyjwt" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423", size = 107515, upload-time = "2026-05-21T19:54:36.618Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" }, +] + [[package]] name = "pyparsing" version = "3.3.2" @@ -2801,17 +2864,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a9/5c/bfd6bd0bf979426d405cc6e71eceb8701b148b16c21d2dc3c261efc61c7b/sqlparse-0.5.3-py3-none-any.whl", hash = "sha256:cf2196ed3418f3ba5de6af7e82c694a9fbdbfecccdfc72e281548517081f16ca", size = 44415, upload-time = "2024-12-10T12:05:27.824Z" }, ] +[[package]] +name = "sse-starlette" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "starlette" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d2/1b/bc9e3e7a72dcdad7dc7888758f5d00f56f8909ed5cfdff822bd72bb4c520/sse_starlette-3.4.5.tar.gz", hash = "sha256:83072538bc211a2f68b7b0422226c4af3e9b62e106e07034664b832ca019842a", size = 35249, upload-time = "2026-06-20T17:36:58.322Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/75/c88d3f5dafd59c791da1ce27650d30bf5b70cbf1cbf01cd00e5f9e360915/sse_starlette-3.4.5-py3-none-any.whl", hash = "sha256:e71bad53323f65573c3864a6c3bd0c1eb6e5f092b2e48082b0c35927d19ca296", size = 16518, upload-time = "2026-06-20T17:36:56.729Z" }, +] + [[package]] name = "starlette" -version = "0.47.3" +version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/15/b9/cc3017f9a9c9b6e27c5106cc10cc7904653c3eec0729793aec10479dd669/starlette-0.47.3.tar.gz", hash = "sha256:6bc94f839cc176c4858894f1f8908f0ab79dfec1a6b8402f6da9be26ebea52e9", size = 2584144, upload-time = "2025-08-24T13:36:42.122Z" } +sdist = { url = "https://files.pythonhosted.org/packages/eb/e3/7c1dc7381d9f8ab7d854328ebfa884e62cb3f3d8549ddfd37c7814f42afa/starlette-1.3.1.tar.gz", hash = "sha256:05d0213193f2fbaae60e2ecb593b4add4262ad4e46536b54abe36f11a71724e0", size = 2703240, upload-time = "2026-06-12T09:23:11.602Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ce/fd/901cfa59aaa5b30a99e16876f11abe38b59a1a2c51ffb3d7142bb6089069/starlette-0.47.3-py3-none-any.whl", hash = "sha256:89c0778ca62a76b826101e7c709e70680a1699ca7da6b44d38eb0a7e61fe4b51", size = 72991, upload-time = "2025-08-24T13:36:40.887Z" }, + { url = "https://files.pythonhosted.org/packages/ec/bb/2799cc2ede3ed41131f8975621e7213dfc7ef4acbbaadfa440f32500c370/starlette-1.3.1-py3-none-any.whl", hash = "sha256:c7372aae11c3c3f26a42df7bd626cec2f47d03483d261d369516a615a53714c6", size = 73632, upload-time = "2026-06-12T09:23:10.017Z" }, ] [[package]]