agentenv-framework-protocol 0.1.331__tar.gz → 0.1.332__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/PKG-INFO +43 -2
  2. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/README.md +42 -1
  3. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/pyproject.toml +1 -1
  4. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/__init__.py +4 -0
  5. agentenv_framework_protocol-0.1.332/src/agentenv_protocol/tool_client.py +205 -0
  6. agentenv_framework_protocol-0.1.332/tests/test_tool_client.py +385 -0
  7. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/.gitignore +0 -0
  8. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/LICENSE +0 -0
  9. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/NOTICE +0 -0
  10. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/THIRD_PARTY_NOTICES.md +0 -0
  11. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/__init__.py +0 -0
  12. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/advanced_agent.py +0 -0
  13. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/basic_agent.py +0 -0
  14. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/custom_extensions_agent.py +0 -0
  15. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/multimodal_agent.py +0 -0
  16. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/reference_model.py +0 -0
  17. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/examples/streaming_agent.py +0 -0
  18. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/__init__.py +0 -0
  19. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/_trajectory_log.py +0 -0
  20. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/_triggers.py +0 -0
  21. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/extensions.py +0 -0
  22. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/framework.py +0 -0
  23. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/registry.py +0 -0
  24. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/staging.py +0 -0
  25. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/tasks/__init__.py +0 -0
  26. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/a2a_agent/tasks/v1.py +0 -0
  27. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/agent_env_environment.py +0 -0
  28. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/client.py +0 -0
  29. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/manifest.py +0 -0
  30. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/preflight.py +0 -0
  31. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/tool_context.py +0 -0
  32. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/transfers.py +0 -0
  33. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/src/agentenv_protocol/types.py +0 -0
  34. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_a2a_agent.py +0 -0
  35. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_a2a_agent_examples.py +0 -0
  36. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_a2a_object_transfers.py +0 -0
  37. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_a2a_staging.py +0 -0
  38. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_a2a_trajectory_log.py +0 -0
  39. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_agent_env_environment.py +0 -0
  40. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_client.py +0 -0
  41. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_data_objects.py +0 -0
  42. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_intake.py +0 -0
  43. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_manifest.py +0 -0
  44. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_preflight.py +0 -0
  45. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_tool_context.py +0 -0
  46. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_tool_context_fastmcp.py +0 -0
  47. {agentenv_framework_protocol-0.1.331 → agentenv_framework_protocol-0.1.332}/tests/test_tool_context_injection.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: agentenv-framework-protocol
3
- Version: 0.1.331
3
+ Version: 0.1.332
4
4
  Summary: Open data-plane protocol and server SDK for agent environments.
5
5
  Project-URL: Homepage, https://www.agentenvframework.com
6
6
  Project-URL: Documentation, https://www.agentenvframework.com/docs
@@ -245,7 +245,48 @@ class SlackEnv(AgentEnvEnvironment):
245
245
  return [DataPart(data=self.state())]
246
246
  ```
247
247
 
248
- 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.
248
+ 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.
249
+
250
+ ## Calling an env's tools
251
+
252
+ A loop of your own, such as a trainer's rollout or an eval harness, reaches an env's tools through `ToolSession`. It
253
+ is an MCP client for the env's streamable-HTTP endpoint, written on `httpx`, so it works wherever this package
254
+ installs, whichever `mcp` version the loop's own environment pins.
255
+
256
+ ```python
257
+ from agentenv_protocol import ToolSession
258
+
259
+ async with ToolSession(mcp_url=mcp_url, role="alice@example.com") as session:
260
+ tools = await session.list_tools() # EnvironmentTool: name, description, inputSchema
261
+ result = await session.call_tool("slack_send_message", {"channel": "general", "text": "hi"})
262
+ print(result.text, result.isError)
263
+ ```
264
+
265
+ - **The URL:** `mcp_url` is `DeployedEnv.mcp_url` in agent-env, or the address a card was fetched from joined with
266
+ `client.mcp_path(card)`.
267
+ - **Sessions:** a `ToolSession` is one MCP session, and calls on it may run concurrently. `async with` opens and ends
268
+ it; `open()` and `close()` do the same for a session held across calls. A tool tells sessions apart by
269
+ `ctx.caller.session`: the session's id when the server is reached directly, or, behind a gateway, a key the gateway
270
+ gives each session. A session belongs to the event loop that opened it, and using or closing it on another raises
271
+ `RuntimeError`, so a caller with a synchronous API runs all of a session's calls on one loop of its own.
272
+ - **Roles:** `role` is sent as `AgentEnv-Role` (`ROLE_HEADER`) and says who the caller is. A gateway lists only the
273
+ tools that role may use, answers a call to a hidden one with an error result, and forwards the role to the servers
274
+ behind it; a server reached directly reads the header itself. Either way a tool reads it as `ctx.caller.role`
275
+ ([Who is calling](#who-is-calling)), where `default`, like no role, is `None`. `headers` adds any others.
276
+ - **Errors:**
277
+ - A tool that fails, or a call the server refuses (a tool hidden from the role, or an `on_tool_call` refusal),
278
+ returns a result with `isError` set and the server's text.
279
+ - `ToolSessionError` means the request failed at the MCP level: a JSON-RPC error, an HTTP error status, a reply
280
+ that can't be read, or a session the server no longer knows. After a restart, the env's state is usually gone too.
281
+ Only the server the session talks to reports this: a gateway reconnects to a child that restarted and retries the
282
+ call, and a stateless server has no session to lose, so neither restart raises anything.
283
+ - `httpx.TransportError` means the connection failed, and `TimeoutError` that a request outlived `timeout`
284
+ (600 s by default, the gateway's own limit on a tool call).
285
+ - `ToolSession` retries nothing. A call that raised may still have run on the server, or may yet run: a request
286
+ stalled on a bad link can arrive after the client has given up. So after any failed call the env's state is
287
+ unknown; read it back (`data/get`) or end the episode.
288
+ - **Tools:** `list_tools()` returns the tools the role may use as `EnvironmentTool`, the type a card's
289
+ `capabilities.tools` holds. Each `inputSchema` is the JSON Schema a model API takes as the tool's parameters.
249
290
 
250
291
  ## A2A agent framework
251
292
 
@@ -200,7 +200,48 @@ class SlackEnv(AgentEnvEnvironment):
200
200
  return [DataPart(data=self.state())]
201
201
  ```
202
202
 
203
- 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.
203
+ 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.
204
+
205
+ ## Calling an env's tools
206
+
207
+ A loop of your own, such as a trainer's rollout or an eval harness, reaches an env's tools through `ToolSession`. It
208
+ is an MCP client for the env's streamable-HTTP endpoint, written on `httpx`, so it works wherever this package
209
+ installs, whichever `mcp` version the loop's own environment pins.
210
+
211
+ ```python
212
+ from agentenv_protocol import ToolSession
213
+
214
+ async with ToolSession(mcp_url=mcp_url, role="alice@example.com") as session:
215
+ tools = await session.list_tools() # EnvironmentTool: name, description, inputSchema
216
+ result = await session.call_tool("slack_send_message", {"channel": "general", "text": "hi"})
217
+ print(result.text, result.isError)
218
+ ```
219
+
220
+ - **The URL:** `mcp_url` is `DeployedEnv.mcp_url` in agent-env, or the address a card was fetched from joined with
221
+ `client.mcp_path(card)`.
222
+ - **Sessions:** a `ToolSession` is one MCP session, and calls on it may run concurrently. `async with` opens and ends
223
+ it; `open()` and `close()` do the same for a session held across calls. A tool tells sessions apart by
224
+ `ctx.caller.session`: the session's id when the server is reached directly, or, behind a gateway, a key the gateway
225
+ gives each session. A session belongs to the event loop that opened it, and using or closing it on another raises
226
+ `RuntimeError`, so a caller with a synchronous API runs all of a session's calls on one loop of its own.
227
+ - **Roles:** `role` is sent as `AgentEnv-Role` (`ROLE_HEADER`) and says who the caller is. A gateway lists only the
228
+ tools that role may use, answers a call to a hidden one with an error result, and forwards the role to the servers
229
+ behind it; a server reached directly reads the header itself. Either way a tool reads it as `ctx.caller.role`
230
+ ([Who is calling](#who-is-calling)), where `default`, like no role, is `None`. `headers` adds any others.
231
+ - **Errors:**
232
+ - A tool that fails, or a call the server refuses (a tool hidden from the role, or an `on_tool_call` refusal),
233
+ returns a result with `isError` set and the server's text.
234
+ - `ToolSessionError` means the request failed at the MCP level: a JSON-RPC error, an HTTP error status, a reply
235
+ that can't be read, or a session the server no longer knows. After a restart, the env's state is usually gone too.
236
+ Only the server the session talks to reports this: a gateway reconnects to a child that restarted and retries the
237
+ call, and a stateless server has no session to lose, so neither restart raises anything.
238
+ - `httpx.TransportError` means the connection failed, and `TimeoutError` that a request outlived `timeout`
239
+ (600 s by default, the gateway's own limit on a tool call).
240
+ - `ToolSession` retries nothing. A call that raised may still have run on the server, or may yet run: a request
241
+ stalled on a bad link can arrive after the client has given up. So after any failed call the env's state is
242
+ unknown; read it back (`data/get`) or end the episode.
243
+ - **Tools:** `list_tools()` returns the tools the role may use as `EnvironmentTool`, the type a card's
244
+ `capabilities.tools` holds. Each `inputSchema` is the JSON Schema a model API takes as the tool's parameters.
204
245
 
205
246
  ## A2A agent framework
206
247
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "agentenv-framework-protocol"
3
- version = "0.1.331"
3
+ version = "0.1.332"
4
4
  description = "Open data-plane protocol and server SDK for agent environments."
5
5
  readme = "README.md"
6
6
  license = "Apache-2.0"
@@ -2,6 +2,7 @@
2
2
 
3
3
  from . import client
4
4
  from .preflight import intake_fit_check
5
+ from .tool_client import ToolResult, ToolSession, ToolSessionError
5
6
  from .agent_env_environment import (
6
7
  AgentEnvApplication,
7
8
  AgentEnvEnvironment,
@@ -116,6 +117,9 @@ __all__ = [
116
117
  "uploaded_object_path",
117
118
  "MCP_PATH",
118
119
  "MCP_TRANSPORT",
120
+ "ToolSession",
121
+ "ToolResult",
122
+ "ToolSessionError",
119
123
  "MANIFEST_VERSION",
120
124
  "INTERFACE_MANIFEST_PATH",
121
125
  "GET_INTERFACES_EXTENSION_URI",
@@ -0,0 +1,205 @@
1
+ """MCP client for an environment's tools, over the streamable-HTTP transport on httpx alone: a trainer's or eval
2
+ harness's own loop lists and calls an env's tools with it, whatever `mcp` version that loop's environment pins."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import asyncio
7
+ import contextlib
8
+ import itertools
9
+ import json
10
+ from collections.abc import AsyncIterator
11
+ from typing import Any
12
+
13
+ import httpx
14
+ from pydantic import BaseModel, Field, field_validator
15
+
16
+ from .tool_context import ROLE_HEADER
17
+ from .types import PROTOCOL_VERSION, EnvironmentTool
18
+
19
+ _MCP_PROTOCOL_VERSION = "2025-11-25"
20
+ _SESSION_HEADER = "Mcp-Session-Id"
21
+ _VERSION_HEADER = "MCP-Protocol-Version"
22
+ _ACCEPT = "application/json, text/event-stream"
23
+ _CLOSE_TIMEOUT_S = 5.0
24
+
25
+
26
+ class ToolSessionError(RuntimeError):
27
+ """An MCP-level failure: a JSON-RPC error, an HTTP error status, a reply that can't be read, or a session the server
28
+ no longer knows. A tool that fails is not one; its result carries ``isError``."""
29
+
30
+
31
+ class ToolResult(BaseModel):
32
+ content: list[dict] = Field(default_factory=list)
33
+ isError: bool = False
34
+
35
+ @field_validator("content")
36
+ @classmethod
37
+ def _text_blocks_carry_text(cls, content: list[dict]) -> list[dict]:
38
+ if any(block.get("type") == "text" and not isinstance(block.get("text"), str) for block in content):
39
+ raise ValueError("a text block has no string text")
40
+ return content
41
+
42
+ @property
43
+ def text(self) -> str:
44
+ return "\n".join(block["text"] for block in self.content if block.get("type") == "text")
45
+
46
+
47
+ class ToolSession:
48
+ """One MCP session on an env's ``mcp_url``. Use it on the event loop that opened it, and after any failed call treat
49
+ the env's state as unknown: nothing is retried, and a call that raised may still run on the server."""
50
+
51
+ def __init__(self, *, mcp_url: str, role: str | None = None, headers: dict[str, str] | None = None,
52
+ timeout: float = 600.0) -> None:
53
+ self.mcp_url = mcp_url
54
+ self.timeout = timeout
55
+ self.session_id: str | None = None
56
+ self.protocol_version: str | None = None
57
+ self._headers = {**(headers or {}), **({ROLE_HEADER: role} if role else {}), "Accept": _ACCEPT}
58
+ self._client: httpx.AsyncClient | None = None
59
+ self._loop: asyncio.AbstractEventLoop | None = None
60
+ self._ids = itertools.count(1)
61
+
62
+ async def __aenter__(self) -> ToolSession:
63
+ return await self.open()
64
+
65
+ async def __aexit__(self, *exc_info: object) -> None:
66
+ await self.close()
67
+
68
+ async def open(self) -> ToolSession:
69
+ """Start a new MCP session, ending the one this object holds, if any."""
70
+ await self.close()
71
+ self.session_id = self.protocol_version = None
72
+ self._loop = asyncio.get_running_loop()
73
+ self._client = httpx.AsyncClient(headers=self._headers, timeout=None, follow_redirects=True)
74
+ try:
75
+ result = await self._request("initialize", {
76
+ "protocolVersion": _MCP_PROTOCOL_VERSION,
77
+ "capabilities": {},
78
+ "clientInfo": {"name": "agentenv-protocol", "version": PROTOCOL_VERSION},
79
+ })
80
+ self.protocol_version = result.get("protocolVersion") or _MCP_PROTOCOL_VERSION
81
+ await self._deadline("notifications/initialized",
82
+ self._exchange({"jsonrpc": "2.0", "method": "notifications/initialized"}))
83
+ except BaseException:
84
+ await self.close()
85
+ raise
86
+ return self
87
+
88
+ async def close(self) -> None:
89
+ """End the session on the server, waiting at most 5 s (or ``timeout``, if shorter), then close the connection; a
90
+ server that is already gone is not an error."""
91
+ if self._client is None:
92
+ return
93
+ client, self._client = self._open_client(), None
94
+ try:
95
+ if self.session_id:
96
+ with contextlib.suppress(httpx.HTTPError, asyncio.TimeoutError):
97
+ await asyncio.wait_for(client.delete(self.mcp_url, headers=self._session_headers()),
98
+ min(self.timeout, _CLOSE_TIMEOUT_S))
99
+ finally:
100
+ await client.aclose()
101
+
102
+ async def list_tools(self) -> list[EnvironmentTool]:
103
+ tools: list[EnvironmentTool] = []
104
+ params: dict[str, Any] = {}
105
+ while True:
106
+ result = await self._request("tools/list", params)
107
+ if not isinstance(result.get("tools"), list):
108
+ raise ToolSessionError("tools/list: unreadable result: tools is not a list")
109
+ tools += [_validated(EnvironmentTool, tool, "tools/list") for tool in result["tools"]]
110
+ if not result.get("nextCursor"):
111
+ return tools
112
+ params = {"cursor": result["nextCursor"]}
113
+
114
+ async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> ToolResult:
115
+ result = await self._request("tools/call", {"name": name, "arguments": arguments or {}})
116
+ return _validated(ToolResult, result, "tools/call")
117
+
118
+ async def _request(self, method: str, params: dict[str, Any]) -> dict:
119
+ message = {"jsonrpc": "2.0", "id": next(self._ids), "method": method, "params": params}
120
+ reply = await self._deadline(method, self._exchange(message))
121
+ if "error" in reply:
122
+ error = reply["error"] if isinstance(reply["error"], dict) else {"message": reply["error"]}
123
+ raise ToolSessionError(f"{method} failed: {error.get('message')} (code {error.get('code')})")
124
+ if not isinstance(reply.get("result"), dict):
125
+ raise ToolSessionError(f"{method}: the reply has no result object")
126
+ return reply["result"]
127
+
128
+ async def _deadline(self, method: str, exchange: Any) -> Any:
129
+ try:
130
+ return await asyncio.wait_for(exchange, self.timeout)
131
+ except asyncio.TimeoutError as e:
132
+ raise TimeoutError(f"{method} did not finish within {self.timeout}s") from e
133
+
134
+ async def _exchange(self, message: dict) -> dict:
135
+ async with self._open_client().stream("POST", self.mcp_url, json=message,
136
+ headers=self._session_headers()) as response:
137
+ await self._raise_for_status(response, message["method"])
138
+ if "id" not in message:
139
+ return {}
140
+ if message["method"] == "initialize":
141
+ self.session_id = response.headers.get(_SESSION_HEADER)
142
+ content_type = response.headers.get("content-type", "")
143
+ try:
144
+ if content_type.startswith("text/event-stream"):
145
+ async for data in _sse_data(response):
146
+ reply = json.loads(data)
147
+ if isinstance(reply, dict) and reply.get("id") == message["id"] and (
148
+ "result" in reply or "error" in reply):
149
+ return reply
150
+ raise ToolSessionError(f"{message['method']}: the event stream ended before the reply")
151
+ if content_type.startswith("application/json"):
152
+ reply = json.loads(await response.aread())
153
+ if isinstance(reply, dict):
154
+ return reply
155
+ raise ToolSessionError(f"{message['method']}: the reply is not a JSON-RPC object")
156
+ except ValueError as e:
157
+ raise ToolSessionError(f"{message['method']}: unreadable reply: {e}") from e
158
+ raise ToolSessionError(f"{message['method']}: unexpected content type {content_type!r}")
159
+
160
+ async def _raise_for_status(self, response: httpx.Response, method: str) -> None:
161
+ if response.status_code == 404 and self.session_id:
162
+ raise ToolSessionError(f"{method}: the server no longer knows session {self.session_id}; it restarted "
163
+ "or ended the session, and the env's state may be gone with it")
164
+ if response.status_code >= 400:
165
+ body = (await response.aread()).decode(errors="replace")[:1000]
166
+ raise ToolSessionError(f"{method}: HTTP {response.status_code} from {self.mcp_url}: {body}")
167
+
168
+ def _session_headers(self) -> dict[str, str]:
169
+ headers = {_SESSION_HEADER: self.session_id} if self.session_id else {}
170
+ if self.protocol_version:
171
+ headers[_VERSION_HEADER] = self.protocol_version
172
+ return headers
173
+
174
+ def _open_client(self) -> httpx.AsyncClient:
175
+ if self._client is None:
176
+ raise RuntimeError("the ToolSession is not open: use `async with ToolSession(...)` or `await open()`")
177
+ if asyncio.get_running_loop() is not self._loop:
178
+ raise RuntimeError("the ToolSession was opened on another event loop; use and close it on that one")
179
+ return self._client
180
+
181
+
182
+ def _validated(model: type[BaseModel], data: Any, method: str) -> Any:
183
+ try:
184
+ return model.model_validate(data)
185
+ except ValueError as e:
186
+ raise ToolSessionError(f"{method}: unreadable result: {e}") from e
187
+
188
+
189
+ async def _sse_data(response: httpx.Response) -> AsyncIterator[str]:
190
+ """The data of each ``message`` event in a server-sent event stream; an event the stream ends inside is dropped,
191
+ as the SSE spec has it, and so are events with no data, such as a resumable server's priming event."""
192
+ event, data = "message", []
193
+ async for line in response.aiter_lines():
194
+ if not line:
195
+ payload = "\n".join(data)
196
+ if payload.strip() and event == "message":
197
+ yield payload
198
+ event, data = "message", []
199
+ continue
200
+ field, _, value = line.partition(":")
201
+ value = value[1:] if value.startswith(" ") else value
202
+ if field == "data":
203
+ data.append(value)
204
+ elif field == "event":
205
+ event = value
@@ -0,0 +1,385 @@
1
+ """ToolSession against a real AgentEnv FastMCP server in each streamable-HTTP mode, and against a scripted server for
2
+ the wire cases a real one rarely produces: paged lists, busy event streams, lost sessions and broken replies."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import asyncio
7
+ import json
8
+ import time
9
+ from typing import Annotated
10
+
11
+ import httpx
12
+ import pytest
13
+ from pydantic import Field
14
+
15
+ from agentenv_protocol import (
16
+ ROLE_HEADER,
17
+ EnvironmentCard,
18
+ ToolContext,
19
+ ToolResult,
20
+ ToolSession,
21
+ ToolSessionError,
22
+ create_fastmcp_app,
23
+ tool,
24
+ )
25
+
26
+ _URL = "http://env/mcp"
27
+
28
+
29
+ class _Items:
30
+ def __init__(self) -> None:
31
+ self.items: list[str] = []
32
+
33
+ @tool(name="{environment_name}_add_item")
34
+ async def add_item(self, item: Annotated[str, Field(description="The item to add.")], times: int = 1) -> str:
35
+ """Add an item to the store."""
36
+ self.items.extend([item] * times)
37
+ return json.dumps({"count": len(self.items)})
38
+
39
+ @tool()
40
+ async def list_items(self) -> str:
41
+ """List the stored items."""
42
+ return json.dumps({"items": self.items})
43
+
44
+ @tool()
45
+ async def echo_after(self, value: str, seconds: float) -> str:
46
+ """Answer with value after a pause."""
47
+ await asyncio.sleep(seconds)
48
+ return value
49
+
50
+ @tool()
51
+ def explode(self) -> str:
52
+ """Always fails."""
53
+ raise RuntimeError("the tool blew up")
54
+
55
+ @tool()
56
+ def whoami(self, ctx: ToolContext) -> str:
57
+ """Who is calling."""
58
+ return json.dumps({"role": ctx.caller.role, "session": ctx.caller.session})
59
+
60
+
61
+ def _route_clients_to(monkeypatch, transport: httpx.AsyncBaseTransport) -> None:
62
+ real = httpx.AsyncClient
63
+ monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: real(transport=transport, **kwargs))
64
+
65
+
66
+ @pytest.fixture
67
+ def serve_items(monkeypatch):
68
+ """Serve an _Items env with create_fastmcp_app in the given mode, in process; yields the handler."""
69
+ pytest.importorskip("mcp.server.fastmcp")
70
+
71
+ async def serve(**fastmcp_kwargs):
72
+ handler = _Items()
73
+ app = create_fastmcp_app(handler, card=EnvironmentCard(name="items"), **fastmcp_kwargs).streamable_http_app()
74
+ _route_clients_to(monkeypatch, httpx.ASGITransport(app=app))
75
+ return handler, app.router.lifespan_context(app)
76
+ return serve
77
+
78
+
79
+ class _Scripted:
80
+ """A streamable-HTTP MCP server whose answer to each method a test sets, recording every request it receives."""
81
+
82
+ def __init__(self) -> None:
83
+ self.requests: list[httpx.Request] = []
84
+ self.sessions = {"s-1"}
85
+ self.answers: dict = {"tools/list": {"tools": []}}
86
+
87
+ async def __call__(self, request: httpx.Request) -> httpx.Response:
88
+ self.requests.append(request)
89
+ session = request.headers.get("mcp-session-id")
90
+ if request.method == "DELETE":
91
+ self.sessions.discard(session)
92
+ return httpx.Response(200)
93
+ message = json.loads(request.content)
94
+ if session is not None and session not in self.sessions:
95
+ return httpx.Response(404, json={"jsonrpc": "2.0", "id": "server-error",
96
+ "error": {"code": -32600, "message": "Session not found"}})
97
+ if "id" not in message:
98
+ return httpx.Response(202)
99
+ answer = self.answers.get(message["method"])
100
+ if callable(answer):
101
+ return answer(message)
102
+ if message["method"] == "initialize":
103
+ return httpx.Response(200, headers={"mcp-session-id": "s-1"}, json=_reply(message, {
104
+ "protocolVersion": "2025-06-18", "capabilities": {"tools": {}}, "serverInfo": {"name": "s", "version": "1"}}))
105
+ return httpx.Response(200, json=_reply(message, answer if answer is not None else {}))
106
+
107
+ def bodies(self, method: str) -> list[dict]:
108
+ return [json.loads(r.content) for r in self.requests if r.content and json.loads(r.content)["method"] == method]
109
+
110
+
111
+ def _reply(message: dict, result: dict) -> dict:
112
+ return {"jsonrpc": "2.0", "id": message["id"], "result": result}
113
+
114
+
115
+ def _events(*blocks: str) -> httpx.Response:
116
+ return httpx.Response(200, headers={"content-type": "text/event-stream"}, content="".join(blocks).encode())
117
+
118
+
119
+ @pytest.fixture
120
+ def scripted(monkeypatch) -> _Scripted:
121
+ server = _Scripted()
122
+ _route_clients_to(monkeypatch, httpx.MockTransport(server))
123
+ return server
124
+
125
+
126
+ @pytest.mark.asyncio
127
+ @pytest.mark.parametrize("mode", [{}, {"json_response": True}, {"stateless_http": True}],
128
+ ids=["event-stream", "json", "stateless"])
129
+ async def test_an_agentenv_servers_tools_are_listed_and_called_in_each_streamable_http_mode(serve_items, mode):
130
+ handler, lifespan = await serve_items(**mode)
131
+ async with lifespan, ToolSession(mcp_url=_URL) as session:
132
+ tools = {t.name: t for t in await session.list_tools()}
133
+ added = await session.call_tool("items_add_item", {"item": "crate", "times": 2})
134
+ listed = await session.call_tool("list_items")
135
+
136
+ assert {"items_add_item", "list_items", "echo_after", "explode"} <= set(tools)
137
+ assert tools["items_add_item"].description == "Add an item to the store."
138
+ assert tools["items_add_item"].inputSchema["properties"]["item"]["description"] == "The item to add."
139
+ assert (added.isError, json.loads(added.text)) == (False, {"count": 2})
140
+ assert json.loads(listed.text) == {"items": ["crate", "crate"]} and handler.items == ["crate", "crate"]
141
+ assert (session.session_id is None) == ("stateless_http" in mode)
142
+
143
+
144
+ @pytest.mark.asyncio
145
+ async def test_a_failing_or_unknown_tool_is_an_error_result_not_an_exception(serve_items):
146
+ _, lifespan = await serve_items()
147
+ async with lifespan, ToolSession(mcp_url=_URL) as session:
148
+ failed = await session.call_tool("explode")
149
+ unknown = await session.call_tool("no_such_tool")
150
+ after = await session.call_tool("list_items")
151
+
152
+ assert failed.isError and "the tool blew up" in failed.text
153
+ assert unknown.isError and "no_such_tool" in unknown.text
154
+ assert not after.isError
155
+
156
+
157
+ @pytest.mark.asyncio
158
+ async def test_a_tool_reads_the_role_and_session_a_tool_session_sends_as_its_caller(serve_items):
159
+ _, lifespan = await serve_items()
160
+ async with (lifespan, ToolSession(mcp_url=_URL, role="alice@example.com") as alice,
161
+ ToolSession(mcp_url=_URL) as anonymous, ToolSession(mcp_url=_URL, role="default") as default):
162
+ whoami = {t.name: t for t in await alice.list_tools()}["whoami"]
163
+ seen = [json.loads((await session.call_tool("whoami")).text) for session in (alice, anonymous, default)]
164
+
165
+ assert not whoami.inputSchema.get("properties")
166
+ assert seen == [{"role": "alice@example.com", "session": alice.session_id},
167
+ {"role": None, "session": anonymous.session_id},
168
+ {"role": None, "session": default.session_id}]
169
+ assert len({alice.session_id, anonymous.session_id, default.session_id}) == 3
170
+
171
+
172
+ @pytest.mark.asyncio
173
+ async def test_concurrent_calls_on_one_session_each_get_their_own_reply(serve_items):
174
+ _, lifespan = await serve_items()
175
+ async with lifespan, ToolSession(mcp_url=_URL) as session:
176
+ results = await asyncio.gather(*(session.call_tool("echo_after", {"value": f"v{i}", "seconds": (i % 5) / 50})
177
+ for i in range(25)))
178
+
179
+ assert [r.text for r in results] == [f"v{i}" for i in range(25)]
180
+
181
+
182
+ @pytest.mark.asyncio
183
+ async def test_a_request_that_outlives_the_timeout_raises_timeout_error_and_the_session_carries_on(serve_items):
184
+ _, lifespan = await serve_items()
185
+ async with lifespan, ToolSession(mcp_url=_URL) as session:
186
+ session.timeout = 0.3
187
+ started = time.monotonic()
188
+ with pytest.raises(TimeoutError, match="tools/call did not finish within 0.3s"):
189
+ await session.call_tool("echo_after", {"value": "late", "seconds": 5})
190
+ waited = time.monotonic() - started
191
+ session.timeout = 30
192
+ assert (await session.call_tool("echo_after", {"value": "on time", "seconds": 0})).text == "on time"
193
+
194
+ assert waited < 2
195
+
196
+
197
+ def test_text_joins_the_text_blocks_and_skips_the_rest():
198
+ result = ToolResult(content=[{"type": "text", "text": "first"}, {"type": "image", "data": "...", "mimeType": "image/png"},
199
+ {"type": "text", "text": "second"}])
200
+ assert result.text == "first\nsecond" and ToolResult().text == ""
201
+
202
+
203
+ @pytest.mark.asyncio
204
+ async def test_the_session_sends_its_role_and_headers_echoes_the_session_and_ends_it_on_close(scripted):
205
+ async with ToolSession(mcp_url=_URL, role="cli", headers={"X-Trace": "t-1"}) as session:
206
+ await session.list_tools()
207
+
208
+ initialize, initialized, listing, delete = scripted.requests
209
+ assert [r.method for r in scripted.requests] == ["POST", "POST", "POST", "DELETE"]
210
+ assert all(r.headers[ROLE_HEADER] == "cli" and r.headers["x-trace"] == "t-1" for r in scripted.requests)
211
+ assert initialize.headers["accept"] == "application/json, text/event-stream"
212
+ assert "mcp-session-id" not in initialize.headers and "mcp-protocol-version" not in initialize.headers
213
+ assert json.loads(initialize.content)["params"]["protocolVersion"] == "2025-11-25"
214
+ assert json.loads(initialized.content) == {"jsonrpc": "2.0", "method": "notifications/initialized"}
215
+ for later in (initialized, listing, delete):
216
+ assert (later.headers["mcp-session-id"], later.headers["mcp-protocol-version"]) == ("s-1", "2025-06-18")
217
+ assert scripted.sessions == set()
218
+
219
+
220
+ @pytest.mark.asyncio
221
+ async def test_a_closed_session_opens_again_as_a_new_session(scripted):
222
+ session = ToolSession(mcp_url=_URL)
223
+ async with session:
224
+ pass
225
+ scripted.sessions.add("s-1")
226
+ async with session:
227
+ await session.list_tools()
228
+
229
+ second_initialize = [r for r in scripted.requests if r.content and json.loads(r.content)["method"] == "initialize"][1]
230
+ assert "mcp-session-id" not in second_initialize.headers and session.session_id == "s-1"
231
+
232
+
233
+ @pytest.mark.asyncio
234
+ async def test_opening_an_open_session_ends_the_one_it_held_first(scripted):
235
+ session = await ToolSession(mcp_url=_URL).open()
236
+ scripted.answers["initialize"] = lambda m: httpx.Response(200, headers={"mcp-session-id": "s-2"}, json=_reply(m, {
237
+ "protocolVersion": "2025-06-18"}))
238
+ scripted.sessions.add("s-2")
239
+ await session.open()
240
+ await session.close()
241
+
242
+ deletes = [r.headers["mcp-session-id"] for r in scripted.requests if r.method == "DELETE"]
243
+ assert deletes == ["s-1", "s-2"] and scripted.sessions == set()
244
+
245
+
246
+ def test_a_session_is_used_and_closed_only_on_the_event_loop_that_opened_it(scripted):
247
+ session, opener, other = ToolSession(mcp_url=_URL), asyncio.new_event_loop(), asyncio.new_event_loop()
248
+ try:
249
+ opener.run_until_complete(session.open())
250
+ with pytest.raises(RuntimeError, match="opened on another event loop"):
251
+ other.run_until_complete(session.call_tool("t"))
252
+ with pytest.raises(RuntimeError, match="opened on another event loop"):
253
+ other.run_until_complete(session.close())
254
+ assert not opener.run_until_complete(session.call_tool("t")).isError
255
+ opener.run_until_complete(session.close())
256
+ finally:
257
+ opener.close()
258
+ other.close()
259
+
260
+ assert [r.method for r in scripted.requests].count("DELETE") == 1 and scripted.sessions == set()
261
+
262
+
263
+ @pytest.mark.asyncio
264
+ async def test_tools_list_follows_cursors_until_the_last_page(scripted):
265
+ pages = {None: (["a", "b"], "p2"), "p2": (["c"], "p3"), "p3": (["d"], None)}
266
+
267
+ def page(message):
268
+ names, cursor = pages[message["params"].get("cursor")]
269
+ return httpx.Response(200, json=_reply(message, {"tools": [{"name": n} for n in names],
270
+ **({"nextCursor": cursor} if cursor else {})}))
271
+ scripted.answers["tools/list"] = page
272
+
273
+ async with ToolSession(mcp_url=_URL) as session:
274
+ assert [t.name for t in await session.list_tools()] == ["a", "b", "c", "d"]
275
+ assert [b["params"] for b in scripted.bodies("tools/list")] == [{}, {"cursor": "p2"}, {"cursor": "p3"}]
276
+
277
+
278
+ @pytest.mark.asyncio
279
+ async def test_an_event_stream_reply_is_found_among_comments_priming_events_notifications_and_other_ids(scripted):
280
+ def stream(message):
281
+ reply = json.dumps(_reply(message, {"content": [{"type": "text", "text": "found"}]}), indent=1)
282
+ return _events(
283
+ ": keep-alive\n\n",
284
+ "id: prime-1\ndata:\n\n",
285
+ 'event: message\ndata: {"jsonrpc": "2.0", "method": "notifications/progress", "params": {}}\n\n',
286
+ f'data: {json.dumps(_reply({"id": 999}, {"content": []}))}\n\n',
287
+ "event: ping\ndata: {}\n\n",
288
+ "".join(f"data: {line}\n" for line in reply.splitlines()) + "\n",
289
+ )
290
+ scripted.answers["tools/call"] = stream
291
+
292
+ async with ToolSession(mcp_url=_URL) as session:
293
+ assert (await session.call_tool("t")).text == "found"
294
+
295
+
296
+ @pytest.mark.asyncio
297
+ async def test_a_session_the_server_dropped_raises_and_says_the_state_may_be_gone(scripted):
298
+ async with ToolSession(mcp_url=_URL) as session:
299
+ scripted.sessions.clear()
300
+ with pytest.raises(ToolSessionError, match="no longer knows session s-1.*state may be gone"):
301
+ await session.call_tool("t")
302
+
303
+
304
+ @pytest.mark.asyncio
305
+ @pytest.mark.parametrize("answer, match", [
306
+ (lambda m: httpx.Response(200, json={"jsonrpc": "2.0", "id": m["id"], "error": {"code": -32602, "message": "bad args"}}),
307
+ r"tools/call failed: bad args \(code -32602\)"),
308
+ (lambda m: httpx.Response(500, text="upstream exploded"), "tools/call: HTTP 500 from http://env/mcp: upstream exploded"),
309
+ (lambda m: httpx.Response(200, headers={"content-type": "application/json"}, content=b"{not json"),
310
+ "tools/call: unreadable reply"),
311
+ (lambda m: _events("data: {broken\n\n"), "tools/call: unreadable reply"),
312
+ (lambda m: _events('data: {"jsonrpc": "2.0", "method": "notifications/message"}\n\n'),
313
+ "tools/call: the event stream ended before the reply"),
314
+ (lambda m: _events(f"data: {json.dumps(_reply(m, {}))}"), "tools/call: the event stream ended before the reply"),
315
+ (lambda m: httpx.Response(200, headers={"content-type": "text/plain"}, content=b"hi"),
316
+ "tools/call: unexpected content type 'text/plain'"),
317
+ (lambda m: httpx.Response(200, json=[_reply(m, {})]), "tools/call: the reply is not a JSON-RPC object"),
318
+ (lambda m: httpx.Response(200, json={"jsonrpc": "2.0", "id": m["id"], "result": "ok"}),
319
+ "tools/call: the reply has no result object"),
320
+ (lambda m: httpx.Response(200, json=_reply(m, {"content": "not a list"})), "tools/call: unreadable result"),
321
+ (lambda m: httpx.Response(200, json=_reply(m, {"content": [{"type": "text"}]})), "tools/call: unreadable result"),
322
+ (lambda m: httpx.Response(200, json=_reply(m, {"content": [{"type": "text", "text": 7}]})),
323
+ "tools/call: unreadable result"),
324
+ ], ids=["jsonrpc-error", "http-500", "bad-json", "bad-event-json", "no-reply-in-stream", "unterminated-event",
325
+ "wrong-content-type", "batch-reply", "result-not-object", "result-wrong-shape", "text-block-without-text",
326
+ "text-block-with-a-number"])
327
+ async def test_a_failure_at_the_mcp_level_raises_a_tool_session_error(scripted, answer, match):
328
+ scripted.answers["tools/call"] = answer
329
+ async with ToolSession(mcp_url=_URL) as session:
330
+ with pytest.raises(ToolSessionError, match=match):
331
+ await session.call_tool("t")
332
+
333
+
334
+ @pytest.mark.asyncio
335
+ @pytest.mark.parametrize("result", [{"tools": [{"description": "nameless"}]}, {"tools": 7}, {"tools": None}, {}],
336
+ ids=["nameless-tool", "tools-not-a-list", "tools-null", "tools-missing"])
337
+ async def test_a_tools_list_that_cannot_be_read_raises_a_tool_session_error(scripted, result):
338
+ scripted.answers["tools/list"] = result
339
+ async with ToolSession(mcp_url=_URL) as session:
340
+ with pytest.raises(ToolSessionError, match="tools/list: unreadable result"):
341
+ await session.list_tools()
342
+
343
+
344
+ @pytest.mark.asyncio
345
+ async def test_a_failed_initialize_raises_and_leaves_the_session_closed(scripted):
346
+ scripted.answers["initialize"] = lambda m: httpx.Response(503, text="starting up")
347
+ session = ToolSession(mcp_url=_URL)
348
+ with pytest.raises(ToolSessionError, match="initialize: HTTP 503"):
349
+ await session.open()
350
+ with pytest.raises(RuntimeError, match="the ToolSession is not open"):
351
+ await session.list_tools()
352
+
353
+
354
+ @pytest.mark.asyncio
355
+ async def test_closing_after_the_server_is_gone_does_not_raise(monkeypatch):
356
+ async def gone_on_delete(request: httpx.Request) -> httpx.Response:
357
+ if request.method == "DELETE":
358
+ raise httpx.ConnectError("connection refused", request=request)
359
+ message = json.loads(request.content)
360
+ if "id" not in message:
361
+ return httpx.Response(202)
362
+ return httpx.Response(200, headers={"mcp-session-id": "s-9"}, json=_reply(message, {"protocolVersion": "2025-11-25"}))
363
+ _route_clients_to(monkeypatch, httpx.MockTransport(gone_on_delete))
364
+
365
+ async with ToolSession(mcp_url=_URL) as session:
366
+ assert session.session_id == "s-9"
367
+ await session.close()
368
+
369
+
370
+ @pytest.mark.asyncio
371
+ async def test_closing_waits_for_the_server_no_longer_than_the_sessions_timeout(monkeypatch):
372
+ async def stalls_on_delete(request: httpx.Request) -> httpx.Response:
373
+ if request.method == "DELETE":
374
+ await asyncio.sleep(30)
375
+ message = json.loads(request.content or b"{}")
376
+ if "id" not in message:
377
+ return httpx.Response(202)
378
+ return httpx.Response(200, headers={"mcp-session-id": "s-2"}, json=_reply(message, {"protocolVersion": "2025-11-25"}))
379
+ _route_clients_to(monkeypatch, httpx.MockTransport(stalls_on_delete))
380
+
381
+ session = await ToolSession(mcp_url=_URL).open()
382
+ session.timeout = 0.3
383
+ started = time.monotonic()
384
+ await session.close()
385
+ assert time.monotonic() - started < 1