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

packages/agentenv-protocol/README.md

@@ -200,7 +200,48 @@ class SlackEnv(AgentEnvEnvironment):
         return [DataPart(data=self.state())]
 ```
 
-Dependencies are intentionally light (`pydantic`, `starlette`) so the package can be added to environment server images without pulling a heavier framework — `mcp` is imported lazily inside `create_fastmcp_app()` and is deliberately not a dependency.
+Dependencies are intentionally light (`pydantic`, `starlette`, `httpx`) so the package can be added to environment server images without pulling a heavier framework — `mcp` is imported lazily inside `create_fastmcp_app()` and is deliberately not a dependency.
+
+## Calling an env's tools
+
+A loop of your own, such as a trainer's rollout or an eval harness, reaches an env's tools through `ToolSession`. It
+is an MCP client for the env's streamable-HTTP endpoint, written on `httpx`, so it works wherever this package
+installs, whichever `mcp` version the loop's own environment pins.
+
+```python
+from agentenv_protocol import ToolSession
+
+async with ToolSession(mcp_url=mcp_url, role="alice@example.com") as session:
+    tools = await session.list_tools()                 # EnvironmentTool: name, description, inputSchema
+    result = await session.call_tool("slack_send_message", {"channel": "general", "text": "hi"})
+    print(result.text, result.isError)
+```
+
+- **The URL:** `mcp_url` is `DeployedEnv.mcp_url` in agent-env, or the address a card was fetched from joined with
+  `client.mcp_path(card)`.
+- **Sessions:** a `ToolSession` is one MCP session, and calls on it may run concurrently. `async with` opens and ends
+  it; `open()` and `close()` do the same for a session held across calls. A tool tells sessions apart by
+  `ctx.caller.session`: the session's id when the server is reached directly, or, behind a gateway, a key the gateway
+  gives each session. A session belongs to the event loop that opened it, and using or closing it on another raises
+  `RuntimeError`, so a caller with a synchronous API runs all of a session's calls on one loop of its own.
+- **Roles:** `role` is sent as `AgentEnv-Role` (`ROLE_HEADER`) and says who the caller is. A gateway lists only the
+  tools that role may use, answers a call to a hidden one with an error result, and forwards the role to the servers
+  behind it; a server reached directly reads the header itself. Either way a tool reads it as `ctx.caller.role`
+  ([Who is calling](#who-is-calling)), where `default`, like no role, is `None`. `headers` adds any others.
+- **Errors:**
+  - A tool that fails, or a call the server refuses (a tool hidden from the role, or an `on_tool_call` refusal),
+    returns a result with `isError` set and the server's text.
+  - `ToolSessionError` means the request failed at the MCP level: a JSON-RPC error, an HTTP error status, a reply
+    that can't be read, or a session the server no longer knows. After a restart, the env's state is usually gone too.
+    Only the server the session talks to reports this: a gateway reconnects to a child that restarted and retries the
+    call, and a stateless server has no session to lose, so neither restart raises anything.
+  - `httpx.TransportError` means the connection failed, and `TimeoutError` that a request outlived `timeout`
+    (600 s by default, the gateway's own limit on a tool call).
+  - `ToolSession` retries nothing. A call that raised may still have run on the server, or may yet run: a request
+    stalled on a bad link can arrive after the client has given up. So after any failed call the env's state is
+    unknown; read it back (`data/get`) or end the episode.
+- **Tools:** `list_tools()` returns the tools the role may use as `EnvironmentTool`, the type a card's
+  `capabilities.tools` holds. Each `inputSchema` is the JSON Schema a model API takes as the tool's parameters.
 
 ## A2A agent framework
 

packages/agentenv-protocol/src/agentenv_protocol/__init__.py

@@ -2,6 +2,7 @@
 
 from . import client
 from .preflight import intake_fit_check
+from .tool_client import ToolResult, ToolSession, ToolSessionError
 from .agent_env_environment import (
     AgentEnvApplication,
     AgentEnvEnvironment,
@@ -116,6 +117,9 @@
     "uploaded_object_path",
     "MCP_PATH",
     "MCP_TRANSPORT",
+    "ToolSession",
+    "ToolResult",
+    "ToolSessionError",
     "MANIFEST_VERSION",
     "INTERFACE_MANIFEST_PATH",
     "GET_INTERFACES_EXTENSION_URI",

packages/agentenv-protocol/src/agentenv_protocol/tool_client.py

@@ -0,0 +1,205 @@
+"""MCP client for an environment's tools, over the streamable-HTTP transport on httpx alone: a trainer's or eval
+harness's own loop lists and calls an env's tools with it, whatever `mcp` version that loop's environment pins."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import itertools
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+from pydantic import BaseModel, Field, field_validator
+
+from .tool_context import ROLE_HEADER
+from .types import PROTOCOL_VERSION, EnvironmentTool
+
+_MCP_PROTOCOL_VERSION = "2025-11-25"
+_SESSION_HEADER = "Mcp-Session-Id"
+_VERSION_HEADER = "MCP-Protocol-Version"
+_ACCEPT = "application/json, text/event-stream"
+_CLOSE_TIMEOUT_S = 5.0
+
+
+class ToolSessionError(RuntimeError):
+    """An MCP-level failure: a JSON-RPC error, an HTTP error status, a reply that can't be read, or a session the server
+    no longer knows. A tool that fails is not one; its result carries ``isError``."""
+
+
+class ToolResult(BaseModel):
+    content: list[dict] = Field(default_factory=list)
+    isError: bool = False
+
+    @field_validator("content")
+    @classmethod
+    def _text_blocks_carry_text(cls, content: list[dict]) -> list[dict]:
+        if any(block.get("type") == "text" and not isinstance(block.get("text"), str) for block in content):
+            raise ValueError("a text block has no string text")
+        return content
+
+    @property
+    def text(self) -> str:
+        return "\n".join(block["text"] for block in self.content if block.get("type") == "text")
+
+
+class ToolSession:
+    """One MCP session on an env's ``mcp_url``. Use it on the event loop that opened it, and after any failed call treat
+    the env's state as unknown: nothing is retried, and a call that raised may still run on the server."""
+
+    def __init__(self, *, mcp_url: str, role: str | None = None, headers: dict[str, str] | None = None,
+                 timeout: float = 600.0) -> None:
+        self.mcp_url = mcp_url
+        self.timeout = timeout
+        self.session_id: str | None = None
+        self.protocol_version: str | None = None
+        self._headers = {**(headers or {}), **({ROLE_HEADER: role} if role else {}), "Accept": _ACCEPT}
+        self._client: httpx.AsyncClient | None = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._ids = itertools.count(1)
+
+    async def __aenter__(self) -> ToolSession:
+        return await self.open()
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.close()
+
+    async def open(self) -> ToolSession:
+        """Start a new MCP session, ending the one this object holds, if any."""
+        await self.close()
+        self.session_id = self.protocol_version = None
+        self._loop = asyncio.get_running_loop()
+        self._client = httpx.AsyncClient(headers=self._headers, timeout=None, follow_redirects=True)
+        try:
+            result = await self._request("initialize", {
+                "protocolVersion": _MCP_PROTOCOL_VERSION,
+                "capabilities": {},
+                "clientInfo": {"name": "agentenv-protocol", "version": PROTOCOL_VERSION},
+            })
+            self.protocol_version = result.get("protocolVersion") or _MCP_PROTOCOL_VERSION
+            await self._deadline("notifications/initialized",
+                                 self._exchange({"jsonrpc": "2.0", "method": "notifications/initialized"}))
+        except BaseException:
+            await self.close()
+            raise
+        return self
+
+    async def close(self) -> None:
+        """End the session on the server, waiting at most 5 s (or ``timeout``, if shorter), then close the connection; a
+        server that is already gone is not an error."""
+        if self._client is None:
+            return
+        client, self._client = self._open_client(), None
+        try:
+            if self.session_id:
+                with contextlib.suppress(httpx.HTTPError, asyncio.TimeoutError):
+                    await asyncio.wait_for(client.delete(self.mcp_url, headers=self._session_headers()),
+                                           min(self.timeout, _CLOSE_TIMEOUT_S))
+        finally:
+            await client.aclose()
+
+    async def list_tools(self) -> list[EnvironmentTool]:
+        tools: list[EnvironmentTool] = []
+        params: dict[str, Any] = {}
+        while True:
+            result = await self._request("tools/list", params)
+            if not isinstance(result.get("tools"), list):
+                raise ToolSessionError("tools/list: unreadable result: tools is not a list")
+            tools += [_validated(EnvironmentTool, tool, "tools/list") for tool in result["tools"]]
+            if not result.get("nextCursor"):
+                return tools
+            params = {"cursor": result["nextCursor"]}
+
+    async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> ToolResult:
+        result = await self._request("tools/call", {"name": name, "arguments": arguments or {}})
+        return _validated(ToolResult, result, "tools/call")
+
+    async def _request(self, method: str, params: dict[str, Any]) -> dict:
+        message = {"jsonrpc": "2.0", "id": next(self._ids), "method": method, "params": params}
+        reply = await self._deadline(method, self._exchange(message))
+        if "error" in reply:
+            error = reply["error"] if isinstance(reply["error"], dict) else {"message": reply["error"]}
+            raise ToolSessionError(f"{method} failed: {error.get('message')} (code {error.get('code')})")
+        if not isinstance(reply.get("result"), dict):
+            raise ToolSessionError(f"{method}: the reply has no result object")
+        return reply["result"]
+
+    async def _deadline(self, method: str, exchange: Any) -> Any:
+        try:
+            return await asyncio.wait_for(exchange, self.timeout)
+        except asyncio.TimeoutError as e:
+            raise TimeoutError(f"{method} did not finish within {self.timeout}s") from e
+
+    async def _exchange(self, message: dict) -> dict:
+        async with self._open_client().stream("POST", self.mcp_url, json=message,
+                                              headers=self._session_headers()) as response:
+            await self._raise_for_status(response, message["method"])
+            if "id" not in message:
+                return {}
+            if message["method"] == "initialize":
+                self.session_id = response.headers.get(_SESSION_HEADER)
+            content_type = response.headers.get("content-type", "")
+            try:
+                if content_type.startswith("text/event-stream"):
+                    async for data in _sse_data(response):
+                        reply = json.loads(data)
+                        if isinstance(reply, dict) and reply.get("id") == message["id"] and (
+                                "result" in reply or "error" in reply):
+                            return reply
+                    raise ToolSessionError(f"{message['method']}: the event stream ended before the reply")
+                if content_type.startswith("application/json"):
+                    reply = json.loads(await response.aread())
+                    if isinstance(reply, dict):
+                        return reply
+                    raise ToolSessionError(f"{message['method']}: the reply is not a JSON-RPC object")
+            except ValueError as e:
+                raise ToolSessionError(f"{message['method']}: unreadable reply: {e}") from e
+            raise ToolSessionError(f"{message['method']}: unexpected content type {content_type!r}")
+
+    async def _raise_for_status(self, response: httpx.Response, method: str) -> None:
+        if response.status_code == 404 and self.session_id:
+            raise ToolSessionError(f"{method}: the server no longer knows session {self.session_id}; it restarted "
+                                   "or ended the session, and the env's state may be gone with it")
+        if response.status_code >= 400:
+            body = (await response.aread()).decode(errors="replace")[:1000]
+            raise ToolSessionError(f"{method}: HTTP {response.status_code} from {self.mcp_url}: {body}")
+
+    def _session_headers(self) -> dict[str, str]:
+        headers = {_SESSION_HEADER: self.session_id} if self.session_id else {}
+        if self.protocol_version:
+            headers[_VERSION_HEADER] = self.protocol_version
+        return headers
+
+    def _open_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            raise RuntimeError("the ToolSession is not open: use `async with ToolSession(...)` or `await open()`")
+        if asyncio.get_running_loop() is not self._loop:
+            raise RuntimeError("the ToolSession was opened on another event loop; use and close it on that one")
+        return self._client
+
+
+def _validated(model: type[BaseModel], data: Any, method: str) -> Any:
+    try:
+        return model.model_validate(data)
+    except ValueError as e:
+        raise ToolSessionError(f"{method}: unreadable result: {e}") from e
+
+
+async def _sse_data(response: httpx.Response) -> AsyncIterator[str]:
+    """The data of each ``message`` event in a server-sent event stream; an event the stream ends inside is dropped,
+    as the SSE spec has it, and so are events with no data, such as a resumable server's priming event."""
+    event, data = "message", []
+    async for line in response.aiter_lines():
+        if not line:
+            payload = "\n".join(data)
+            if payload.strip() and event == "message":
+                yield payload
+            event, data = "message", []
+            continue
+        field, _, value = line.partition(":")
+        value = value[1:] if value.startswith(" ") else value
+        if field == "data":
+            data.append(value)
+        elif field == "event":
+            event = value