feat(protocol): add ToolSession to list and call an env's MCP tools

[earakely-scale]

What and why

A training framework that runs its own loop, such as TRL's GRPOTrainer, NeMo Gym, verifiers, verl, SkyRL or Tinker, has to list an env's tools, declare them to its model and call them. The protocol package covers an env's card, data plane and extensions, but not its tools. The only tool-calling code is core's internal _tool_names and each agent's own, all on the MCP SDK, and installing the SDK clashes with environments that pin mcp==2.0.0 (verifiers) or fastmcp>=3 (OpenEnv).

agentenv_protocol gains ToolSession, an MCP client for an env's streamable-HTTP endpoint, written on httpx, which is already a base dependency:

from agentenv_protocol import ToolSession

async with ToolSession(mcp_url=deployed.mcp_url, role="alice@example.com") as session:
    tools = await session.list_tools()                  # list[EnvironmentTool], each with its inputSchema
    result = await session.call_tool("items_add_item", {"item": "crate"})
    result.text, result.isError, result.content         # ToolResult

• Constructor: ToolSession(*, mcp_url, role=None, headers=None, timeout=600.0).
  • mcp_url is keyword-only. A deployed env has several URLs (environment_url, gateway_url, the card's, mcp_url), and passing the wrong one positionally reads fine and fails at runtime. A later transport, such as the gateway's /step, can then sit beside it as its own keyword without MCP being the privileged positional one. None is added here.
  • headers is there for sandbox providers' REQUEST_HEADERS, the headers a platform needs on requests to its URLs (deploy_agent passes them to agents). timeout is there because 600 s is a ceiling, not a deadline a training loop wants.
  • No verify: nothing passes one, and httpx reads SSL_CERT_FILE and SSL_CERT_DIR for a private CA.
• No new dependency: mcp stays out of the protocol package, as its README promises, so the client imports wherever the package does.
• Wire behaviour:
  • initialize negotiates the protocol version, then notifications/initialized follows.
  • Mcp-Session-Id and MCP-Protocol-Version are echoed on every later request.
  • A reply is read as JSON or as an event stream; the stream parser skips comments, priming events, notifications and other ids.
  • tools/list follows nextCursor, and redirects are followed.
  • close() sends DELETE and never raises when the server is already gone.
  • It works against the SDK server's three modes: event stream (the default), json_response and stateless_http.
• Results: list_tools() returns the existing EnvironmentTool (name, description, inputSchema), the type a card's capabilities.tools holds. Each inputSchema is the JSON Schema a model API takes as the tool's parameters. Converting to a model API's tool format is left to the loop: TRL builds schemas from the adapter's own methods, and elsewhere it's a one-liner. ToolResult carries MCP's content and isError, and text joins its text blocks. It leaves out structuredContent, because servers also send that JSON as text; adding it later breaks nobody.
• Errors:
  • A tool that fails, or a tool the server doesn't know, comes back as a result with isError set; it does not raise.
  • ToolSessionError is an MCP-level failure: a JSON-RPC error, an HTTP error status, an unreadable reply, or a session the server no longer knows. It says so when a restart has probably taken the env's state.
  • Only the server the session talks to reports a lost session. Behind a gateway, a child that restarted is reconnected and the call retried (_call_child_tool), and a stateless server has no session to lose, so neither restart raises. The README and the error's docstring say so; telling that an env was reset under an episode is a job for the episode API that comes next.
  • httpx.TransportError means the connection failed.
  • TimeoutError means a request outlived timeout. That is one overall deadline per request, 600 s by default (the gateway's TOOL_CALL_TIMEOUT_S). It also bounds an event stream that trickles.
  • The client retries nothing: a lost session usually means lost env state, which the caller has to handle.
• Role and caller: this is the client side of the ToolContext that feat(protocol): ToolContext, an opt-in injected caller for environment tools #153 hands to tools.
  • role= sends AgentEnv-Role, the protocol's ROLE_HEADER from tool_context. Since feat(gateway): forward the caller's role to child servers and share one child session generation per process #145 it is the caller's identity: a gateway filters tools/list and tools/call by it and forwards it to the servers behind it, and a server reached directly reads the header itself. A tool reads it as ctx.caller.role, and default, like no role, reads as None.
  • Each ToolSession is its own caller session, ctx.caller.session: the session's id when the server is reached directly, or a key the gateway gives that session, stable across its calls.
  • A call the server refuses, whether a tool hidden from the role or an on_tool_call refusal, comes back as an isError result, like a tool that fails.
• One event loop per session: a session's connections belong to the loop that opened it. Used or closed on another loop, it raises RuntimeError("the ToolSession was opened on another event loop; ..."). Without that check, the first such call failed with RuntimeError: Event loop is closed and later ones went through. TRL calls an environment's reset synchronously, so its adapter will run every session on one loop thread of its own.
• Python 3.10: the protocol package supports 3.10, so deadlines use asyncio.wait_for and catch both asyncio.TimeoutError and the builtin.
• Docs: the protocol README gains "Calling an env's tools", and its dependency line now names httpx, which was already a dependency.
• Exports: ToolSession, ToolResult and ToolSessionError. The MCP protocol version the client asks for stays private; session.protocol_version shows the agreed one.
• Additive: nothing is renamed or removed, so the title has no !.

How it was tested

• packages/agentenv-protocol/tests/test_tool_client.py (new, 34 tests, in process, no sockets):
  • A real server: an @tool env served by create_fastmcp_app in all three SDK modes, covering:
    • listing, with descriptions and schemas from Annotated fields;
    • calls, and failing or unknown tools as isError;
    • 25 concurrent calls on one session, each getting its own reply;
    • a deadline that raises TimeoutError at 0.3 s while the session carries on;
    • a tool that declares a ToolContext reading what three sessions send: alice@example.com and her own session id, while no role and default both read as None. The ToolContext parameter stays out of the listed schema.
  • A scripted server:
    • role and extra headers on every request, the session and version echoed, DELETE on close;
    • re-opening a closed session as a new one, and opening an open one ending the session it held;
    • a session used or closed on another event loop raising, then closing normally on its own;
    • three pages of tools/list;
    • a reply found in a stream full of comments, priming events, notifications, other ids and multi-line data:;
    • a session the server dropped;
    • twelve broken replies, each a ToolSessionError with its own message, including text blocks without string text;
    • a tools/list result that can't be read (a nameless tool, or tools that is not a list, null or missing), a failed initialize, closing after the server is gone, and closing within the session's timeout when DELETE stalls.
  • Without mcp: only the 7 tests that serve FastMCP need mcp; without it the other 27 pass.
• Integration, fast tier: tst/integration/env/envs/server_env_local_test.py gains test_a_deployed_envs_tools_are_listed_and_called_from_a_loop_of_ones_own. On a deployed items env, it checks that:
  • deployed.mcp_url lists its tools, with their parameter descriptions in each inputSchema;
  • calls agree with the env's own data/get;
  • error injection through urn:agentenv:set-errors/v1 comes back as isError;
  • 8 concurrent sessions, each with its own role, get their own count, and the items env's new items_whoami tool reads each one's role and session id;
  • after a docker restart, the old session raises "no longer knows session" and a new one sees the emptied state.
• Integration, slow tier:
  • test_gateway_with_single_mcp_server's role section now also checks the gateway through ToolSession. Role default doesn't list list_emails, and calling it returns isError with "is disabled for role 'default'"; role cli lists it and reads the seeded emails.
  • test_gateway_with_agentenv_environment: through a deployed gateway, the role a ToolSession sends reaches the child's items_whoami as its caller. One session's key is the same across its calls, and a second session gets another. The card lists items_whoami without the ToolContext parameter, so this test's and test_renamed_registration_serves_its_registered_name's card tool lists gain it.
  • The gateway-to-protocol constant parity test is test(gateway): a real gateway in front of a real protocol child, over loopback #160's (AgentEnv-Role, agentenv.io/role, agentenv.io/session), so this PR no longer adds its own.
• Runs, on this head:
  • macOS, Python 3.14: make unit-test gives 7490 passed, 15 skipped and 1 failed. The failure, tst/unit/bundle/parse_test.py::TestFiles::test_deep_nesting[json], is pre-existing: origin/main fails it the same way on this machine, it passes when run alone, and main's CI is green.
  • server_env_local_test.py: 3 passed.
  • The three slow gateway tests this PR touches pass (3 min 15 s with cached images): test_gateway_with_single_mcp_server, test_gateway_with_agentenv_environment and test_renamed_registration_serves_its_registered_name.
  • test_tool_client.py on Python 3.10: 34 pass with mcp 1.30.0 (the version this repo's CodeBuild job resolves); without mcp, 27 pass and 7 skip.
• Runs before the rebase: Linux (Ubuntu 24.04 VM, Python 3.12): make unit-test gave 7175 passed and 5 skipped; server_env_local_test.py 3 passed; test_gateway_with_single_mcp_server passed in 3 min 40 s.

Chaos testing

Chaos ran on the client before the rebase onto main. Since then the client has changed in these ways:

• where it imports ROLE_HEADER;
• a stricter check on a tools/list result;
• the constructor (keyword-only mcp_url, no verify);
• the event-loop check;
• httpx no longer has timeouts of its own;
• notifications/initialized goes through the same streamed POST as every request.

The requests on the wire are the same as before. The deadline chaos measured was always the wait_for one, which is now the only one. Chaos wasn't rerun.

On an Ubuntu 24.04 VM (8 vCPU, Docker 29), with four harness processes in parallel:

• Envs: 32, deployed with MCPServerEnv.deploy(sandbox_type="local") from an @tool test env built with this checkout's protocol package. One in four runs json_response.
• Clients: 48 workers, plus a short-deadline prober and a session churner.

Workers:

• echo of awkward strings: unicode, RTL, NUL, text that looks like data: lines, 5 KB;
• 3–20 concurrent echos on one session;
• add then list, checked against an exact model of the env (owned envs) or for containment (shared envs);
• deterministic blobs up to 4 MB, checked by hash;
• a failing tool, and a tool that sends up to 60 log notifications before it replies;
• tools/list, cancelling a call mid-flight, and reopening the session.

Faults: on 1–3 envs every 2–10 s, each held 2–25 s:

• docker pause, restart, and kill then start;
• tc netem on both ends of the link: 20–400 ms delay with jitter, up to 30% loss, duplication, reordering, corruption;
• iptables drops of the published port, or of the container;
• ss -K connection resets.

Runs:

	Run 1	Run 2	Run 3
Commit, before the rebase	ed50222..b2fa205	7764ebb	2397f3d
Length	15 min	15 min	10 min
Faults	852	840	575
Timed requests	375,391	372,044	260,089
Checked answers	292,140	289,722	204,478
Failures, all ConnectError / ReadError / RemoteProtocolError / TimeoutError	7,775	6,958	4,921
Violations	1 (below)	0	0
Longest single request at a 15 s deadline	15.14 s	15.21 s	15.21 s
Harness event-loop lag, max	0.27 s	0.30 s	0.29 s
Churned sessions; live sessions left on the env	33,897; 1	34,374; 1	24,157; 1

Every run:

• No exception outside the documented types, and close() never raised.
• The short-deadline probes (0.5–2 s) got TimeoutError at most 0.26 s late.
• Client fds stayed between 17 and 87 (they follow how many requests are in flight), and RSS between 158 and 214 MB, with no trend.
• Every env answered once faults cleared.
• open() stayed within three timeouts: two requests, plus a cleanup DELETE bounded the same way when the second fails. The longest was 5.6 s at a 2 s deadline. close() stayed within 5 s.

What it found:

• A call that raised can still run on the server. In run 1, an add timed out behind a 22 s port blackhole. Once the port reopened, the kernel's retransmission delivered it after the worker's next two adds, on a session the dropped DELETE had never ended. Any HTTP client behaves this way. The README and the docstring now say that after a failed call the env's state is unknown, and runs 2 and 3 checked state with that allowance. Of 301 adds that raised, none landed late in those runs.
• Opening could overrun its deadline. A session with a 0.5 s deadline took 2.4 s to fail opening, because open()'s cleanup DELETE was capped at a fixed 5 s. close() now waits at most the session's timeout when that is shorter, with a test.
• A test was too tight for CI. CodeBuild's runner took more than 0.3 s for an in-process initialize, which tripped the deadline test that opened its session at 0.3 s. The test now shortens the deadline only for the call that must time out.

Notes for review

• Since the rebase onto main (feat(gateway): forward the caller's role to child servers and share one child session generation per process #145, feat(deploy_agent): deliver the role to the agent or fail the step #146, feat(protocol): ToolContext, an opt-in injected caller for environment tools #153, fix(providers): draw local sandbox host ports below the ephemeral range, in turn #161 and later): ROLE_HEADER is imported from tool_context rather than defined a second time in types. The docs say what a role and a session mean to a tool, and the tests check that a tool reads both. The constant parity test gave way to test(gateway): a real gateway in front of a real protocol child, over loopback #160's, which is still open.
• From review:
  • keyword-only mcp_url, no verify, the event-loop check, the private protocol-version constant, and the docs on where a lost session is reported;
  • then a smaller client: tool_definitions removed until a loop needs it, one timer per request instead of httpx's plus wait_for, one POST path (no separate _notify), and docstrings cut to the two rules a caller must follow, with the README carrying the rest.
• Considered and left as is:
  • Redirects are followed, and httpx strips only Authorization when one leaves the host, so the role and any extra headers go to whatever host a redirect names. A 301, 302 or 303 also turns the POST into a GET, which then waits out its deadline. The only redirect our servers send is a trailing-slash 307 on the same host, and the env is already the party being trusted.
  • list_tools() follows nextCursor without a loop guard, so a server that keeps returning the same cursor keeps it paging, each page inside its own deadline.
• What's not in it: answering server-to-client requests (sampling and elicitation; the client declares neither capability), a GET stream for server-sent notifications, and resuming a broken stream with Last-Event-ID. A server that sends a request on the reply stream gets no answer and the call runs to its deadline. No AgentEnv server does that today.
• Errors left as httpx's own: a transport failure is httpx's own TransportError subclass (ConnectError, ReadError, RemoteProtocolError, ...), not wrapped. DecodingError and TooManyRedirects are httpx errors outside TransportError. Neither can happen against these servers, which send no content-encoding and no redirect chains, and neither showed up in chaos, so they aren't wrapped either. Say if you'd rather ToolSessionError cover them.
• open() and the deadline: timeout bounds each request. open() is two requests, initialize and notifications/initialized, plus a cleanup DELETE when the second fails, so it can take up to three timeouts. Say if you'd rather have one deadline over all of open().
• Follow-ups this enables: core's _tool_names and the agents' own clients could move onto ToolSession; that isn't done here. A converter to OpenAI or Anthropic tool formats should arrive with the first loop that needs one, such as PortSim's agent or verifiers.

🤖 Generated with Claude Code

Confidence Score: 5/5

The PR appears safe to merge; no new actionable issue remains.

Summary

Adds ToolSession to list and call an environment's MCP tools without installing the MCP SDK.

• Training loops can list and call environment tools through ToolSession.

_{Reviews (6) · Last reviewed commit: "refactor(protocol): a smaller ToolSessio..." · Reviewed by Greptile}