From aafbd78ba326dc56d97d030ef523d2e74955e34e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 15:59:32 -0700 Subject: [PATCH 1/7] feat(protocol): ToolContext, an opt-in injected caller for environment tools A tool of an environment built on AgentEnvEnvironment can declare a parameter annotated ToolContext (any name) and receive, per call, who is calling: the role the gateway forwarded as MCP request _meta (the AgentEnv-Role header for direct callers), the caller's session, the tool name and its arguments. The parameter never reaches the wire: the registry sees a twin whose signature omits it, and tools without the slot are registered verbatim. One context is bound per call around ToolManager.call_tool, installed last at mount so it wraps any dispatch guard a subclass installed first. The handler's on_tool_call(context) may answer a call with a CallToolResult that reaches the wire verbatim, before output validation. ToolContext.current() hands the same object to any code running on behalf of the call. @extension routes fill the slot from the request headers. Co-Authored-By: Claude Fable 5.1 --- packages/agentenv-protocol/README.md | 29 ++ .../src/agentenv_protocol/__init__.py | 18 + .../agent_env_environment.py | 261 +++++++++- .../src/agentenv_protocol/testing.py | 18 + .../src/agentenv_protocol/tool_context.py | 211 ++++++++ .../tests/test_tool_context.py | 354 ++++++++++++++ .../tests/test_tool_context_injection.py | 459 ++++++++++++++++++ 7 files changed, 1335 insertions(+), 15 deletions(-) create mode 100644 packages/agentenv-protocol/src/agentenv_protocol/testing.py create mode 100644 packages/agentenv-protocol/src/agentenv_protocol/tool_context.py create mode 100644 packages/agentenv-protocol/tests/test_tool_context.py create mode 100644 packages/agentenv-protocol/tests/test_tool_context_injection.py diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index b2e7fb66..7de99c86 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -83,6 +83,35 @@ environment-prefixed tools without knowing the name at class-definition time; an FastMCP; since those apps have no MCP tool registry, constructing one with `@tool` methods raises. +### Who is calling + +A tool that needs the caller declares a parameter annotated `ToolContext` (any name) and receives +it on every call. The SDK fills it, never the client: it is absent from the advertised +`inputSchema` and from `tools/list`. Tools that declare no such parameter are registered unchanged. + +```python +from agentenv_protocol import ToolContext + + @tool(name="{environment_name}_list_inbox") + def list_inbox(self, ctx: ToolContext) -> list[dict]: + """Messages addressed to the caller.""" + return [m for m in self.messages if m["to"] == ctx.caller.role] +``` + +`ctx.caller.role` is the role the agent was deployed with, as the gateway forwarded it (`None` when +nothing was forwarded; the gateway's `default` role counts as nothing). `ctx.caller.session` tells +callers apart when the gateway shares one connection to the server. `ctx.tool`, `ctx.arguments`, +`ctx.call_id` and `ctx.transport` (`"mcp"` or `"rest"`) describe the call. Code that runs on behalf +of the call without a parameter of its own, such as a database method, reads the same object through +`ToolContext.current()`, which returns an empty context outside any call. The binding is per call +and follows the request's task, so concurrent callers never see each other's context. + +To answer a call before the tool runs, override `on_tool_call(self, context)`: return `None` to let +it proceed, or an `mcp.types.CallToolResult` (`isError=True` for a refusal) to send that as the +tool's reply. An `@extension` method may declare the same parameter; it is filled from the request's +`AgentEnv-Role` header. In tests, `agentenv_protocol.testing.tool_context(role=...)` binds a caller +around a direct call to a handler method. + ## Serving `serve()` builds the FastMCP app via `create_fastmcp_app()`, which encodes the agent-env deploy diff --git a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py index b7c4d85a..dbb9d474 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py @@ -12,9 +12,19 @@ environment_card, extension, get_data, + injecting, reset_data, tool, ) +from .tool_context import ( + DEFAULT_ROLE, + ROLE_HEADER, + ROLE_META_KEY, + SESSION_META_KEY, + Caller, + ToolContext, + normalize_role, +) from .manifest import ( CLI_COMMANDS, GET_INTERFACES_EXTENSION_URI, @@ -78,6 +88,14 @@ "get_data", "extension", "tool", + "injecting", + "ToolContext", + "Caller", + "normalize_role", + "ROLE_META_KEY", + "SESSION_META_KEY", + "ROLE_HEADER", + "DEFAULT_ROLE", "EnvironmentCapabilities", "EnvironmentCard", "EnvironmentExtension", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index 8f47a31a..9fc5e111 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -1,13 +1,15 @@ -"""Server SDK: decorate a handler's methods with @reset_data/@add_data/@get_data (JSON-RPC data plane at /agentenv; each optional, advertised as capabilities.operations), @extension (card-advertised REST routes), and @tool (MCP tools; FastMCP-backed apps only). Mount onto a FastMCP, Starlette, or FastAPI app via a per-framework AgentEnvApplication subclass — or subclass AgentEnvEnvironment, configure its card with @environment_card(...), and serve() it.""" +"""Server SDK: decorate a handler's methods with @reset_data/@add_data/@get_data (JSON-RPC data plane at /agentenv; each optional, advertised as capabilities.operations), @extension (card-advertised REST routes), and @tool (MCP tools; FastMCP-backed apps only; a parameter annotated ToolContext receives the caller and is never advertised). Mount onto a FastMCP, Starlette, or FastAPI app via a per-framework AgentEnvApplication subclass — or subclass AgentEnvEnvironment, configure its card with @environment_card(...), and serve() it.""" from __future__ import annotations import datetime import decimal import enum +import functools import inspect import json import logging import os +import sys import types import uuid from abc import ABC, abstractmethod @@ -19,6 +21,7 @@ from starlette.responses import Response from starlette.routing import Route +from .tool_context import EMPTY, ToolContext, bound, context_from_http, context_from_mcp from .transfers import WriteNamespaceGrant from .types import DATA_OBJECTS_EXTENSION_URI, MCP_TRANSPORT, METHOD_ADD, METHOD_GET, METHOD_RESET, RPC_PATH, WELL_KNOWN_PATH, AddDataRequest, AddDataResponse, EnvironmentCapabilities, EnvironmentCard, EnvironmentExtension, EnvironmentInterface, EnvironmentTool, GetDataResponse, ResetDataResponse, error_body @@ -29,6 +32,8 @@ _EXT_ROUTE_ATTR = "_aee_ext_route" _TOOL_ATTR = "_aee_tool" _CARD_CONFIG_ATTR = "_aee_card_config" +#: Set on the tool dispatch installed by ``_bind_tool_calls``; a guard downstream checks it at boot. +_TOOL_CONTEXT_BOUND_ATTR = "__agentenv_tool_context__" OP_RESET_DATA = "reset_data" OP_ADD_DATA = "add_data" @@ -111,6 +116,11 @@ def tool(name: str | None = None, *, description: str | None = None) -> Callable the card and the registration. Requires a FastMCP-backed application (non-FastMCP apps raise at construction). Duplicate resolved names raise; a card-declared name wins the advertisement while the method is still registered. + + A parameter annotated ``ToolContext`` (or ``Optional[ToolContext]``; any name) is filled by the + SDK with the caller of the current call and is omitted from the advertised ``inputSchema`` and + from tools/list; a FastMCP ``Context`` parameter is likewise omitted from the card while FastMCP + keeps injecting it. """ if callable(name): raise TypeError("use @tool(...) with parentheses, not bare @tool") @@ -180,6 +190,9 @@ def add_routes_to_app(self, app: Any, *, rpc_url: str = RPC_PATH, card_url: str self._add_route(app, path, [method], handler) for descriptor, fn in self._tools: self._register_tool(app, descriptor, fn) + # Last, so it is the outermost layer: a dispatch guard the handler installed before mounting + # then runs under the bound ToolContext instead of outside it. + self._bind_tool_calls(app) @abstractmethod def _add_route(self, app: Any, path: str, methods: list, handler: Callable) -> None: @@ -188,6 +201,10 @@ def _add_route(self, app: Any, path: str, methods: list, handler: Callable) -> N def _register_tool(self, app: Any, descriptor: EnvironmentTool, fn: Callable) -> None: raise NotImplementedError(f"{type(self).__name__} does not serve MCP tools") + def _bind_tool_calls(self, app: Any) -> None: + """Frameworks with a tool dispatch chokepoint bind a ToolContext around every call there; the others serve no tools.""" + return None + def _card_for(self, app: Any) -> EnvironmentCard: return self.environment_card @@ -199,7 +216,41 @@ def _add_route(self, app: Any, path: str, methods: list, handler: Callable) -> N app.custom_route(path, methods=methods)(handler) def _register_tool(self, app: Any, descriptor: EnvironmentTool, fn: Callable) -> None: - app.tool(name=descriptor.name, description=descriptor.description)(fn) + app.tool(name=descriptor.name, description=descriptor.description)(injecting(fn)) + + def _bind_tool_calls(self, app: Any) -> None: + """Wrap ``app._tool_manager.call_tool``, which every tools/call resolves at call time, so each call + of a registered tool runs under its ToolContext and the handler's ``on_tool_call`` may answer + it first. An unknown tool name goes straight to the dispatch, whose error answers it. + + The installed callable carries ``__agentenv_tool_context__ = True`` and ``__wrapped__`` (the + dispatch it wraps): the handles a downstream guard uses to assert the ordering contract. One + handler per app: a second binding would run its hook on the first handler's tools. + """ + manager = getattr(app, "_tool_manager", None) + inner = getattr(manager, "call_tool", None) + if not callable(inner): + return + if getattr(inner, _TOOL_CONTEXT_BOUND_ATTR, False): + raise RuntimeError("this FastMCP app already carries an AgentEnv tool dispatch; mount one handler per app") + hook = _tool_call_hook(self.handler) + get_tool = getattr(manager, "get_tool", None) + + async def call_tool(name: str, arguments: dict, *args: Any, **kwargs: Any) -> Any: + if get_tool is not None and get_tool(name) is None: + return await inner(name, arguments, *args, **kwargs) + context = kwargs.get("context", args[0] if args else None) + with bound(context_from_mcp(context, name, arguments)) as tool_context: + if hook is not None: + answer = await _maybe_await(hook(tool_context)) + if answer is not None: + _check_call_tool_result(answer) + return answer + return await inner(name, arguments, *args, **kwargs) + + functools.update_wrapper(call_tool, inner) + setattr(call_tool, _TOOL_CONTEXT_BOUND_ATTR, True) + manager.call_tool = call_tool def _card_for(self, app: Any) -> EnvironmentCard: """Declare the MCP endpoint only when the app says where it serves it; FastMCP-shaped targets may not.""" @@ -263,6 +314,16 @@ def serve(self, transport: str = "streamable-http") -> None: self.mount(self.mcp) self.mcp.run(transport=transport) + def on_tool_call(self, context: ToolContext) -> Any: + """Runs before every tools/call on this environment's app, under the bound ``context``. + + Return None to let the call proceed, or an ``mcp.types.CallToolResult`` (``isError=True`` for + a refusal) to answer the call without running the tool; it reaches the wire verbatim, before + output validation. Anything else raises TypeError. May be ``async``. Not invoked for + extension (REST) routes. + """ + return None + _SCALAR_TYPES = {str: "string", int: "integer", float: "number", bool: "boolean", list: "array", dict: "object"} _STRING_FORMATS = {datetime.datetime: "date-time", datetime.date: "date", datetime.time: "time", uuid.UUID: "uuid"} @@ -307,19 +368,179 @@ def _schema_for_annotation(annotation: Any) -> dict: return {"type": json_type} if json_type is not None else {} -def _schema_from_signature(fn: Callable) -> dict: - """Derive a JSON-Schema object for a handler's call params from its signature (Annotated[Field] descriptions included).""" - sig = inspect.signature(fn) +def _unwrap_annotation(annotation: Any) -> list: + """The concrete types an annotation names: Annotated metadata and None peeled off Optional/Union.""" + origin = get_origin(annotation) + if origin is Annotated: + return _unwrap_annotation(get_args(annotation)[0]) + if origin is Union or origin is getattr(types, "UnionType", None): + return [t for arg in get_args(annotation) if arg is not type(None) for t in _unwrap_annotation(arg)] + return [annotation] + + +def _split_top_level(text: str, sep: str) -> list: + parts: list = [] + depth = 0 + start = 0 + for i, ch in enumerate(text): + if ch in "[(": + depth += 1 + elif ch in "])": + depth -= 1 + elif ch == sep and depth == 0: + parts.append(text[start:i]) + start = i + 1 + parts.append(text[start:]) + return [p.strip() for p in parts] + + +def _string_annotation_names(text: str) -> list: + """The type names a string annotation spells out, unwrapped like ``_unwrap_annotation``: the fallback + under ``from __future__ import annotations`` when the function's hints cannot be resolved.""" + text = text.strip().strip("'\"") + parts = _split_top_level(text, "|") + if len(parts) > 1: + return [n for part in parts for n in _string_annotation_names(part)] + head, bracket, rest = text.partition("[") + if bracket and text.endswith("]"): + args = _split_top_level(rest[:-1], ",") + wrapper = head.rsplit(".", 1)[-1] + if wrapper == "Annotated": + return _string_annotation_names(args[0]) + if wrapper in ("Optional", "Union"): + return [n for arg in args for n in _string_annotation_names(arg)] + return [text] + return [] if text == "None" else [text.rsplit(".", 1)[-1]] + + +def _is_tool_context_annotation(annotation: Any) -> bool: + if isinstance(annotation, str): + return ToolContext.__name__ in _string_annotation_names(annotation) + return any(t is ToolContext for t in _unwrap_annotation(annotation)) + + +def _fastmcp_context_type() -> type | None: + """FastMCP's Context class when mcp is loaded, else None. mcp is an optional dependency, so it is + looked up rather than imported; a Context-annotated handler has necessarily imported it already.""" + return getattr(sys.modules.get("mcp.server.fastmcp.server"), "Context", None) + + +def _is_fastmcp_context_annotation(annotation: Any) -> bool: + context_type = _fastmcp_context_type() + if context_type is None: + return False + if isinstance(annotation, str): + return context_type.__name__ in _string_annotation_names(annotation) + return any(isinstance(t, type) and issubclass(t, context_type) for t in _unwrap_annotation(annotation)) + + +def _resolved_hints(fn: Callable) -> dict: try: - hints = get_type_hints(fn, include_extras=True) + return get_type_hints(fn, include_extras=True) except Exception: - hints = {} + return {} + + +def _callable_name(fn: Callable) -> str: + return getattr(fn, "__qualname__", None) or getattr(fn, "__name__", None) or repr(fn) + + +def _tool_context_slot(fn: Callable) -> str | None: + """Name of the parameter annotated ``ToolContext`` / ``Optional[ToolContext]`` (any name), or None. + The slot is filled by keyword, so a positional-only one is refused here rather than on every call.""" + hints = _resolved_hints(fn) + slots = [] + for name, param in inspect.signature(fn).parameters.items(): + if name in ("self", "cls") or param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + continue + if not _is_tool_context_annotation(hints.get(name, param.annotation)): + continue + if param.kind is inspect.Parameter.POSITIONAL_ONLY: + raise TypeError(f"{_callable_name(fn)}: ToolContext parameter {name!r} must not be positional-only") + slots.append(name) + if len(slots) > 1: + raise TypeError(f"{_callable_name(fn)} declares more than one ToolContext parameter: {slots}") + return slots[0] if slots else None + + +def _tool_call_hook(handler: Any) -> Callable | None: + """The handler's ``on_tool_call`` when it has one of its own (the base class's no-op does not count), + checked at mount to take the ToolContext as its one argument so a mismatch fails here, not per call.""" + hook = getattr(handler, "on_tool_call", None) + if hook is None or getattr(type(handler), "on_tool_call", None) is AgentEnvEnvironment.on_tool_call: + return None + try: + inspect.signature(hook).bind(EMPTY) + except TypeError as e: + raise TypeError(f"{_callable_name(hook)} must take the ToolContext as its only argument") from e + return hook + + +def _reduced_signature(fn: Callable, slot: str) -> tuple: + """``fn``'s signature and annotations without ``slot``, every annotation resolved: the registry + reads ``__signature__`` verbatim and cannot evaluate a string left in it.""" + sig = inspect.signature(fn) + hints = _resolved_hints(fn) + params = [p.replace(annotation=hints.get(p.name, p.annotation)) for p in sig.parameters.values() if p.name != slot] + returns = hints.get("return", sig.return_annotation) + unresolved = [p.name for p in params if isinstance(p.annotation, str)] + (["return"] if isinstance(returns, str) else []) + if unresolved: + raise TypeError( + f"{_callable_name(fn)}: cannot resolve the annotations of {unresolved}; " + "a function that declares a ToolContext parameter must have resolvable annotations" + ) + annotations = {p.name: p.annotation for p in params if p.annotation is not inspect.Parameter.empty} + if returns is not inspect.Signature.empty: + annotations["return"] = returns + return sig.replace(parameters=params, return_annotation=returns), annotations + + +def _inject_tool_context(fn: Callable, slot: str) -> Callable: + # Two bodies so the registry's sync/async detection sees the same kind as ``fn``. + if inspect.iscoroutinefunction(fn): + async def wrapper(*args: Any, **kwargs: Any) -> Any: + kwargs[slot] = ToolContext.current() + return await fn(*args, **kwargs) + else: + def wrapper(*args: Any, **kwargs: Any) -> Any: + kwargs[slot] = ToolContext.current() + return fn(*args, **kwargs) + functools.update_wrapper(wrapper, fn) + wrapper.__signature__, wrapper.__annotations__ = _reduced_signature(fn, slot) + return wrapper + + +def injecting(fn: Callable) -> Callable: + """``fn`` as an MCP tool registry should see it: unchanged when it declares no ``ToolContext`` + parameter, else a twin of the same name, doc and sync/async nature whose signature and annotations + omit that parameter and which fills it with ``ToolContext.current()`` on every call. ``@tool`` + methods get this at mount; use it for a slotted function registered any other way + (``app.tool()(fn)``, a separate ``ToolManager``), which otherwise fails at registration.""" + slot = _tool_context_slot(fn) + return _inject_tool_context(fn, slot) if slot else fn + + +def _check_call_tool_result(value: Any) -> None: + # The value goes to the wire as-is, so anything else would be emitted as a successful result. + result_type = getattr(sys.modules.get("mcp.types"), "CallToolResult", None) + if result_type is None or not isinstance(value, result_type): + raise TypeError(f"on_tool_call must return None or mcp.types.CallToolResult, got {type(value).__name__}") + + +def _schema_from_signature(fn: Callable) -> dict: + """Derive a JSON-Schema object for a handler's call params from its signature (Annotated[Field] descriptions + included). Injected parameters (``ToolContext``, FastMCP ``Context``) are not wire parameters and are left out.""" + sig = inspect.signature(fn) + hints = _resolved_hints(fn) properties: dict = {} required: list = [] for name, param in sig.parameters.items(): if name in ("self", "cls") or param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): continue - prop = _schema_for_annotation(hints.get(name, param.annotation)) + annotation = hints.get(name, param.annotation) + if _is_tool_context_annotation(annotation) or _is_fastmcp_context_annotation(annotation): + continue + prop = _schema_for_annotation(annotation) if param.default is inspect.Parameter.empty: required.append(name) else: @@ -601,7 +822,15 @@ def _coerce_params(fn: Callable, params: dict) -> dict: def _extension_handler(fn: Callable, http_method: str) -> Callable: - """REST handler for an @extension route: bind request params to the handler and return JSON.""" + """REST handler for an @extension route: bind request params to the handler and return JSON. A + ``ToolContext`` parameter is filled from the request headers (transport ``"rest"``) and never bound + from the params.""" + slot = _tool_context_slot(fn) + signature = inspect.signature(fn) + if slot: + signature = signature.replace(parameters=[p for p in signature.parameters.values() if p.name != slot]) + op = fn.__name__ or "invoke" + async def handler(request: Request) -> Response: try: if http_method == "GET": @@ -615,14 +844,16 @@ async def handler(request: Request) -> Response: except Exception as e: return _json_response(error_body("invalid_params", str(e)), 400) try: - inspect.signature(fn).bind(**params) + signature.bind(**params) except TypeError as e: return _json_response(error_body("invalid_params", str(e)), 400) - try: - result = await _maybe_await(fn(**params)) - except Exception as e: - logger.exception("extension invocation failed") - return _json_response(error_body("extension_failed", str(e)), 500) + with bound(context_from_http(request, op, params)) as tool_context: + call_kwargs = {**params, slot: tool_context} if slot else params + try: + result = await _maybe_await(fn(**call_kwargs)) + except Exception as e: + logger.exception("extension invocation failed") + return _json_response(error_body("extension_failed", str(e)), 500) return _json_response(result if result is not None else {}) return handler diff --git a/packages/agentenv-protocol/src/agentenv_protocol/testing.py b/packages/agentenv-protocol/src/agentenv_protocol/testing.py new file mode 100644 index 00000000..88caab48 --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/testing.py @@ -0,0 +1,18 @@ +"""Helpers for tests of environments built on this package.""" + +from __future__ import annotations + +import contextlib +from typing import Any, Iterator, Mapping, Optional + +from .tool_context import ToolContext, bound + + +@contextlib.contextmanager +def tool_context(*, role: Optional[str] = None, session: Optional[str] = None, tool: str = "test", + arguments: Optional[Mapping[str, Any]] = None) -> Iterator[ToolContext]: + """Bind a caller for the block, as the dispatch does for one tool call, so a handler or + database method called directly sees ``ToolContext.current()`` the way it would under the + gateway. Nested blocks shadow and restore.""" + with bound(ToolContext.for_test(role=role, session=session, tool=tool, arguments=arguments)) as context: + yield context diff --git a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py new file mode 100644 index 00000000..e33e6f7f --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py @@ -0,0 +1,211 @@ +"""Who is calling a tool, handed to the tool. + +A caller's role and session reach a server as MCP request ``_meta`` entries +(the reverse-DNS keys below) when a gateway in front of the server stamps +them on each proxied call, or as the ``AgentEnv-Role`` and ``mcp-session-id`` +headers when the caller reaches the server directly; a call that carries +neither has an empty caller. An environment built on +:class:`~agentenv_protocol.AgentEnvEnvironment` never reads either itself. +The protocol binds one :class:`ToolContext` per call, in the request's task, +before the handler runs, and hands it to the handler in two ways that always +agree: + +- a handler that declares a parameter annotated :class:`ToolContext` (any + name) receives it as an argument, and that parameter never appears in the + tool's advertised schema; +- any code running on behalf of the call (a database method, a shared helper) + can read the same object through :meth:`ToolContext.current`. + +The binding follows the request's task: work handed to ``asyncio.to_thread`` +or a new task inherits it, a raw ``threading.Thread`` does not and must be +given the context explicitly. Outside a call :meth:`ToolContext.current` +returns :data:`EMPTY`, never raises. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import dataclasses +import types +import uuid +from typing import Any, Iterator, Literal, Mapping, Optional + +from pydantic_core import core_schema + +__all__ = [ + "DEFAULT_ROLE", + "EMPTY", + "ROLE_HEADER", + "ROLE_META_KEY", + "SESSION_META_KEY", + "Caller", + "ToolContext", + "Transport", + "bound", + "caller_from_request", + "context_from_http", + "context_from_mcp", + "normalize_role", +] + +#: MCP request ``_meta`` key a gateway stamps the caller's role under. +ROLE_META_KEY = "agentenv.io/role" +#: MCP request ``_meta`` key a gateway stamps the caller's session under, so +#: per-session state keeps one owner per agent when the gateway shares one +#: connection to a server across agents. +SESSION_META_KEY = "agentenv.io/session" +#: Header carrying the role when no ``_meta`` does: non-MCP paths (a gateway +#: sets it on proxied REST requests) and direct MCP callers. +ROLE_HEADER = "AgentEnv-Role" +#: A gateway's unforwarded role. Same meaning as no role at all. +DEFAULT_ROLE = "default" + +#: How the call reached the server. ``"none"`` is :data:`EMPTY`, the context +#: outside any call; the other two name the dispatch that bound the context. +Transport = Literal["mcp", "rest", "none"] + + +def normalize_role(value: Any) -> Optional[str]: + """None for absent, blank and the gateway's unforwarded role; else the stripped role.""" + if not isinstance(value, str): + return None + value = value.strip() + return None if value in ("", DEFAULT_ROLE) else value + + +@dataclasses.dataclass(frozen=True) +class Caller: + """The identity a call arrived with. + + ``role`` is the role the agent was deployed with, normalised: None when + nothing was forwarded or the gateway's default role came through. + ``raw_role`` is what was received before normalisation, for logs. + ``session`` tells callers apart when the gateway shares one connection. + """ + + role: Optional[str] = None + session: Optional[str] = None + raw_role: Optional[str] = None + + @property + def is_forwarded(self) -> bool: + return self.role is not None + + +@dataclasses.dataclass(frozen=True) +class ToolContext: + """What the protocol knows about one tool call. + + Hashable, so per-call state can be keyed on the context itself: ``call_id`` + identifies the call, and ``arguments`` (a read-only mapping) and ``mcp`` + stay out of the hash. + """ + + caller: Caller = Caller() + tool: str = "" + arguments: Mapping[str, Any] = dataclasses.field(default_factory=lambda: types.MappingProxyType({}), hash=False) + call_id: str = "" + transport: Transport = "none" + #: FastMCP's own request context when the call came over MCP (progress, logging, session); else None. + mcp: Any = dataclasses.field(default=None, hash=False, compare=False) + + @property + def role(self) -> Optional[str]: + return self.caller.role + + @staticmethod + def current() -> "ToolContext": + """The context of the call this code runs on behalf of; :data:`EMPTY` outside any call.""" + return _current.get() + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: Any) -> Any: + # A pydantic model may carry a ToolContext (an audit record, an event): validate by instance. + return core_schema.is_instance_schema(cls) + + @classmethod + def __get_pydantic_json_schema__(cls, schema: Any, handler: Any) -> Any: + # Every field here is schema-able, so a slotted function registered on a tool registry + # without the injecting twin would silently advertise this type as a wire parameter; + # the registry renders its argument model to JSON Schema at registration, so failing + # here fails the mount. + raise TypeError( + "ToolContext is injected by the SDK and is not a wire parameter; " + "register the tool with @tool or wrap it with agentenv_protocol.injecting(fn)" + ) + + @classmethod + def for_test(cls, *, role: Optional[str] = None, session: Optional[str] = None, tool: str = "test", + arguments: Optional[Mapping[str, Any]] = None, transport: Transport = "mcp") -> "ToolContext": + """A context for a test that calls a handler directly or binds one with :func:`bound`.""" + return cls(caller=Caller(role=normalize_role(role), session=session, raw_role=role), tool=tool, + arguments=_frozen(arguments), call_id=uuid.uuid4().hex, transport=transport) + + +#: The context handed back outside any call. +EMPTY = ToolContext() + + +def _frozen(arguments: Optional[Mapping[str, Any]]) -> Mapping[str, Any]: + return types.MappingProxyType(dict(arguments or {})) + +_current: contextvars.ContextVar[ToolContext] = contextvars.ContextVar("agentenv_protocol.tool_context", default=EMPTY) + + +@contextlib.contextmanager +def bound(context: ToolContext) -> Iterator[ToolContext]: + """Bind ``context`` as the current one for the block, restoring the previous binding after.""" + token = _current.set(context) + try: + yield context + finally: + _current.reset(token) + + +def caller_from_request(meta_extra: Optional[Mapping[str, Any]], headers: Any) -> Caller: + """Build the caller from an MCP request's custom ``_meta`` entries and its HTTP headers. + + A ``_meta`` role entry, whatever its value, marks a call a gateway proxied from its own + view of the caller: the headers then describe the gateway's connection, not the caller, + and are not consulted, so a blank or ``default`` role stays "not forwarded" and the + session is the ``_meta`` session or none. Without that entry the caller reached the + server itself, and the role header and ``mcp-session-id`` are trusted exactly as the + caller's other requests are. + """ + meta = meta_extra or {} + forwarded = ROLE_META_KEY in meta + if forwarded: + raw_role = meta.get(ROLE_META_KEY) + else: + raw_role = headers.get(ROLE_HEADER) if headers is not None else None + session = meta.get(SESSION_META_KEY) + if not isinstance(session, str) or not session: + session = None if forwarded or headers is None else headers.get("mcp-session-id") + raw_role = raw_role if isinstance(raw_role, str) else None + return Caller(role=normalize_role(raw_role), session=session or None, raw_role=raw_role) + + +def context_from_mcp(context: Any, tool: str, arguments: Mapping[str, Any]) -> ToolContext: + """The context for one MCP tool call, from FastMCP's ``Context`` (or any object with a + ``request_context``). Tolerates every missing link: no context, no live request, no meta, + no HTTP request (stdio transports).""" + request_context = None + if context is not None: + try: + request_context = context.request_context + except (ValueError, AttributeError): + request_context = None + meta = getattr(request_context, "meta", None) + extra = getattr(meta, "model_extra", None) if meta is not None else None + request = getattr(request_context, "request", None) + headers = getattr(request, "headers", None) if request is not None else None + return ToolContext(caller=caller_from_request(extra, headers), tool=tool, arguments=_frozen(arguments), + call_id=uuid.uuid4().hex, transport="mcp", mcp=context) + + +def context_from_http(request: Any, tool: str, arguments: Mapping[str, Any]) -> ToolContext: + """The context for a REST or extension call: the caller from the request headers only.""" + headers = getattr(request, "headers", None) if request is not None else None + return ToolContext(caller=caller_from_request(None, headers), tool=tool, arguments=_frozen(arguments), + call_id=uuid.uuid4().hex, transport="rest") diff --git a/packages/agentenv-protocol/tests/test_tool_context.py b/packages/agentenv-protocol/tests/test_tool_context.py new file mode 100644 index 00000000..d476dbdd --- /dev/null +++ b/packages/agentenv-protocol/tests/test_tool_context.py @@ -0,0 +1,354 @@ +"""ToolContext without mcp: the value and its binding, how a call builds it, slot detection, the +registration twin, and the extension slot (FastMCP-shaped fake target, in-process).""" + +from __future__ import annotations + +import asyncio +import dataclasses +import inspect +import typing +from types import SimpleNamespace +from typing import Annotated, Optional, Union + +import httpx +import pytest +from pydantic import Field +from starlette.applications import Starlette +from starlette.routing import Route + +from agentenv_protocol import ( + DEFAULT_ROLE, + ROLE_HEADER, + ROLE_META_KEY, + SESSION_META_KEY, + AgentEnvFastMCPApplication, + Caller, + EnvironmentCard, + ToolContext, + extension, + injecting, + normalize_role, + tool, +) +from agentenv_protocol import testing +from agentenv_protocol.agent_env_environment import _is_tool_context_annotation, _schema_from_signature, _tool_context_slot +from agentenv_protocol.tool_context import EMPTY, bound, caller_from_request, context_from_http, context_from_mcp + + +class _FakeMCP: + """FastMCP-shaped target without a ``_tool_manager``: records routes and tool registrations.""" + + def __init__(self) -> None: + self.routes: list[Route] = [] + self.tools: dict[str, tuple] = {} + self.settings = SimpleNamespace(streamable_http_path="/mcp") + + def custom_route(self, path: str, methods: list[str]): + def deco(fn): + self.routes.append(Route(path, fn, methods=methods)) + return fn + return deco + + def tool(self, name: str | None = None, description: str | None = None): + def deco(fn): + self.tools[name or fn.__name__] = (description, fn) + return fn + return deco + + +def _mounted(handler) -> _FakeMCP: + mcp = _FakeMCP() + AgentEnvFastMCPApplication(environment_card=EnvironmentCard(name="items"), handler=handler).add_routes_to_app(mcp) + return mcp + + +def _client(handler) -> httpx.AsyncClient: + app = Starlette(routes=_mounted(handler).routes) + return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://t") + + +def test_normalize_role_absent_blank_default_and_non_string_are_none(): + for value in (None, "", " ", "default", " default ", 5): + assert normalize_role(value) is None + assert normalize_role(" alice ") == "alice" + assert ROLE_META_KEY == "agentenv.io/role" + assert SESSION_META_KEY == "agentenv.io/session" + assert ROLE_HEADER == "AgentEnv-Role" + assert DEFAULT_ROLE == "default" + + +def test_current_is_empty_outside_a_call_and_bound_restores(): + assert ToolContext.current() is EMPTY + assert EMPTY.transport == "none" and EMPTY.caller.role is None + with bound(ToolContext.for_test(role="a")) as outer: + assert ToolContext.current() is outer and outer.role == "a" + with bound(ToolContext.for_test(role="b")) as inner: + assert ToolContext.current() is inner + assert ToolContext.current() is outer + assert ToolContext.current() is EMPTY + with pytest.raises(dataclasses.FrozenInstanceError): + outer.tool = "x" + with pytest.raises(dataclasses.FrozenInstanceError): + outer.caller.role = "x" + + +def test_caller_from_request_precedence(): + headers = {ROLE_HEADER: "hdr", "mcp-session-id": "sid"} + assert caller_from_request({ROLE_META_KEY: "meta"}, headers).role == "meta" + assert caller_from_request({}, headers) == Caller("hdr", "sid", "hdr") + assert caller_from_request(None, headers) == Caller("hdr", "sid", "hdr") + assert caller_from_request({SESSION_META_KEY: "s-meta"}, headers).session == "s-meta" + assert caller_from_request({SESSION_META_KEY: ""}, headers).session == "sid" + assert caller_from_request({SESSION_META_KEY: 7}, headers).session == "sid" + assert caller_from_request({ROLE_META_KEY: "default"}, None) == Caller(None, None, "default") + assert caller_from_request(None, None) == Caller() + + +def test_context_from_mcp_tolerates_every_missing_link(): + empty = context_from_mcp(None, "t", {}) + assert empty.transport == "mcp" and empty.caller == Caller() and empty.mcp is None + + class _OutsideRequest: + @property + def request_context(self): + raise ValueError("Context is not available outside of a request") + + assert context_from_mcp(_OutsideRequest(), "t", {}).caller == Caller() + bare = SimpleNamespace(request_context=SimpleNamespace(meta=None, request=None)) + assert context_from_mcp(bare, "t", {}).caller == Caller() + + arguments = {"q": 1} + live = SimpleNamespace(request_context=SimpleNamespace( + meta=SimpleNamespace(model_extra={ROLE_META_KEY: "bob"}), + request=SimpleNamespace(headers={"mcp-session-id": "s1"}), + )) + ctx = context_from_mcp(live, "search", arguments) + assert ctx.caller == Caller("bob", None, "bob") and ctx.tool == "search" and ctx.mcp is live + assert ctx.call_id and ctx.arguments == {"q": 1} + arguments["q"] = 2 + assert ctx.arguments == {"q": 1} + + +def test_context_from_http_and_testing_helper(): + request = SimpleNamespace(headers={ROLE_HEADER: "carol"}) + ctx = context_from_http(request, "echo", {"m": 1}) + assert ctx.transport == "rest" and ctx.role == "carol" and ctx.tool == "echo" and ctx.arguments == {"m": 1} + + with testing.tool_context(role="alice", session="s") as bound_ctx: + assert ToolContext.current() is bound_ctx + assert bound_ctx.caller == Caller("alice", "s", "alice") + assert ToolContext.current() is EMPTY + with testing.tool_context(role="default") as default_ctx: + assert default_ctx.role is None and default_ctx.caller.raw_role == "default" + + +def test_tool_context_slot_detection_by_annotation(): + def plain(tc: ToolContext, q: str) -> str: ... + def optional(q: str, maybe: Optional[ToolContext] = None) -> str: ... + def pipe(q: str, u: ToolContext | None = None) -> str: ... + def union(q: str, un: Union[ToolContext, None] = None) -> str: ... + def annotated(a: Annotated[ToolContext, "x"], q: str) -> str: ... + def prefixed(t: typing.Optional[ToolContext], q: str) -> str: ... + def unresolvable_sibling(x: "Undefined", tc: ToolContext) -> str: ... # noqa: F821 + def none(q: str, n: int = 0) -> str: ... + def two(a: ToolContext, b: ToolContext) -> str: ... + + assert _tool_context_slot(plain) == "tc" + assert _tool_context_slot(optional) == "maybe" + assert _tool_context_slot(pipe) == "u" + assert _tool_context_slot(union) == "un" + assert _tool_context_slot(annotated) == "a" + assert _tool_context_slot(prefixed) == "t" + with pytest.raises(NameError): + typing.get_type_hints(unresolvable_sibling) # the string fallback is what finds the slot below + assert _tool_context_slot(unresolvable_sibling) == "tc" + assert _tool_context_slot(none) is None + with pytest.raises(TypeError, match="more than one ToolContext parameter"): + _tool_context_slot(two) + + +@pytest.mark.parametrize("text", [ + "ToolContext", + "pkg.ToolContext", + "Optional[ToolContext]", + "typing.Optional[ToolContext]", + "Union[ToolContext, None]", + "typing.Union[None, ToolContext]", + "ToolContext | None", + "None | agentenv_protocol.ToolContext", + "Annotated[ToolContext, Field(description='a, b')]", + "Annotated[Optional[ToolContext], 'x']", + "'ToolContext'", +]) +def test_string_annotation_forms_that_name_the_slot(text): + assert _is_tool_context_annotation(text) + + +@pytest.mark.parametrize("text", ["ToolContextual", "str", "Optional[str]", "list[ToolContext]", "dict[str, ToolContext]", "None"]) +def test_string_annotation_forms_that_do_not_name_the_slot(text): + assert not _is_tool_context_annotation(text) + + +class _SlotHandler: + def __init__(self) -> None: + self.seen: list = [] + + @tool(name="whoami", description="Who is calling.") + async def whoami(self, q: str, tc: ToolContext, limit: int = 3) -> dict: + self.seen.append(tc) + return {"q": q, "role": tc.role, "limit": limit} + + @tool(name="whoami_sync") + def whoami_sync(self, tc: Optional[ToolContext] = None) -> str: + self.seen.append(tc) + return tc.role or "" + + @tool(name="plain") + async def plain(self, item: Annotated[str, Field(description="The item.")], times: int = 1) -> dict: + """Add an item.""" + return {"item": item, "times": times} + + +@pytest.mark.asyncio +async def test_card_schema_omits_the_slot_and_registration_gets_the_twin(): + handler = _SlotHandler() + async with _client(handler) as client: + tools = {t["name"]: t for t in (await client.get("/.well-known/agent-env.json")).json()["capabilities"]["tools"]} + assert tools["whoami"]["inputSchema"] == { + "type": "object", + "properties": {"q": {"type": "string"}, "limit": {"type": "integer", "default": 3}}, + "required": ["q"], + } + assert tools["whoami"]["description"] == "Who is calling." + assert tools["whoami_sync"]["inputSchema"] == {"type": "object", "properties": {}} + + mcp = _mounted(handler) + twin = mcp.tools["whoami"][1] + assert "tc" not in inspect.signature(twin).parameters + assert list(inspect.signature(twin).parameters) == ["q", "limit"] + assert "tc" not in twin.__annotations__ and twin.__annotations__["return"] is dict + assert twin.__wrapped__ == handler.whoami and twin.__name__ == "whoami" + assert inspect.iscoroutinefunction(twin) and not inspect.iscoroutinefunction(mcp.tools["whoami_sync"][1]) + + with bound(ToolContext.for_test(role="alice")) as ctx: + assert await twin(q="x") == {"q": "x", "role": "alice", "limit": 3} + assert mcp.tools["whoami_sync"][1]() == "alice" + assert handler.seen == [ctx, ctx] + assert await twin(q="y") == {"q": "y", "role": None, "limit": 3} # EMPTY outside any binding + + +def test_slot_less_tool_is_registered_verbatim(): + handler = _SlotHandler() + plain = handler.plain + assert _mounted(handler).tools["plain"][1] == plain + assert injecting(plain) is plain + assert _schema_from_signature(plain) == { + "type": "object", + "properties": {"item": {"type": "string", "description": "The item."}, "times": {"type": "integer", "default": 1}}, + "required": ["item"], + } + + +def test_injecting_raises_when_a_slotted_function_has_unresolvable_annotations(): + def broken(x: "Undefined", tc: ToolContext) -> str: ... # noqa: F821 + + with pytest.raises(TypeError, match=r"broken: cannot resolve the annotations of \['x', 'return'\]"): + injecting(broken) + + +def test_fake_target_without_tool_manager_mounts_with_nothing_installed(): + mcp = _mounted(_SlotHandler()) + assert not hasattr(mcp, "_tool_manager") and not hasattr(mcp, "call_tool") + + +class _ExtHandler: + @extension(uri="urn:agentenv:whoami/v1") + async def whoami(self, tc: ToolContext, greeting: str = "hi") -> dict: + assert ToolContext.current() is tc + return {"role": tc.role, "transport": tc.transport, "tool": tc.tool, "greeting": greeting} + + @extension(uri="urn:agentenv:whoami-get/v1", method="GET") + def whoami_get(self, greeting: str, tc: Optional[ToolContext] = None) -> dict: + return {"role": tc.role, "transport": tc.transport, "greeting": greeting} + + +@pytest.mark.asyncio +async def test_extension_slot_binds_rest_caller_from_headers(): + async with _client(_ExtHandler()) as client: + card = (await client.get("/.well-known/agent-env.json")).json() + exts = {e["uri"]: e for e in card["capabilities"]["extensions"]} + assert exts["urn:agentenv:whoami/v1"]["params"]["methods"]["whoami"]["request"] == { + "type": "object", "properties": {"greeting": {"type": "string", "default": "hi"}}, + } + assert set(exts["urn:agentenv:whoami-get/v1"]["params"]["methods"]["whoami_get"]["request"]["properties"]) == {"greeting"} + + r = await client.post("/agentenv/ext/whoami", json={}, headers={ROLE_HEADER: "carol"}) + assert r.status_code == 200 + assert r.json() == {"role": "carol", "transport": "rest", "tool": "whoami", "greeting": "hi"} + r = await client.post("/agentenv/ext/whoami", json={"greeting": "yo"}) + assert r.status_code == 200 and r.json()["role"] is None and r.json()["greeting"] == "yo" + r = await client.post("/agentenv/ext/whoami", json={"greeting": "yo"}, headers={ROLE_HEADER: "default"}) + assert r.json()["role"] is None + + r = await client.post("/agentenv/ext/whoami", json={"tc": 1}) + assert r.status_code == 400 and r.json()["error"]["code"] == "invalid_params" + + r = await client.get("/agentenv/ext/whoami_get", params={"greeting": "hey"}, headers={ROLE_HEADER: "dave"}) + assert r.status_code == 200 and r.json() == {"role": "dave", "transport": "rest", "greeting": "hey"} + assert ToolContext.current() is EMPTY + + +def test_caller_from_request_consults_headers_only_when_no_meta_role_entry_exists(): + headers = {ROLE_HEADER: "hdr", "mcp-session-id": "conn"} + assert caller_from_request(None, headers) == Caller("hdr", "conn", "hdr") + assert caller_from_request({}, headers) == Caller("hdr", "conn", "hdr") + assert caller_from_request({SESSION_META_KEY: "s-only"}, headers) == Caller("hdr", "s-only", "hdr") + # a _meta role entry marks a proxied call: blank and default stay unforwarded, the connection id is not the caller's + assert caller_from_request({ROLE_META_KEY: ""}, headers) == Caller(None, None, "") + assert caller_from_request({ROLE_META_KEY: "default"}, headers) == Caller(None, None, "default") + assert caller_from_request({ROLE_META_KEY: "bob"}, headers) == Caller("bob", None, "bob") + assert caller_from_request({ROLE_META_KEY: "bob", SESSION_META_KEY: "s1"}, headers) == Caller("bob", "s1", "bob") + assert caller_from_request({ROLE_META_KEY: 7, SESSION_META_KEY: ""}, headers) == Caller(None, None, None) + + +def test_tool_context_is_hashable_and_its_arguments_are_read_only(): + ctx = ToolContext.for_test(role="a", arguments={"q": 1}) + assert hash(ctx) == hash(ctx) and {ctx: 1}[ctx] == 1 and ctx in {ctx} + assert hash(EMPTY) is not None and EMPTY != ctx + assert ctx.arguments == {"q": 1} and dict(ctx.arguments) == {"q": 1} + with pytest.raises(TypeError): + ctx.arguments["q"] = 2 # type: ignore[index] + with pytest.raises(TypeError): + EMPTY.arguments["k"] = 1 # type: ignore[index] + assert ToolContext.current().arguments == {} + + +def test_tool_context_embeds_in_a_pydantic_model_but_never_renders_as_a_wire_schema(): + from pydantic import BaseModel, TypeAdapter, ValidationError + + class Audit(BaseModel): + context: ToolContext + note: str = "" + + ctx = ToolContext.for_test(role="a") + assert Audit(context=ctx).context is ctx + with pytest.raises(ValidationError): + Audit(context={"caller": {}}) + assert TypeAdapter(Optional[ToolContext]).validate_python(None) is None + assert TypeAdapter(Optional[ToolContext]).validate_python(ctx) is ctx + with pytest.raises(TypeError, match=r"agentenv_protocol\.injecting\(fn\)"): + Audit.model_json_schema() + + +@pytest.mark.asyncio +async def test_binding_follows_the_task_into_to_thread_and_child_tasks(): + ctx = ToolContext.for_test(role="a") + with bound(ctx): + assert await asyncio.to_thread(ToolContext.current) is ctx + assert await asyncio.create_task(_current_later()) is ctx + assert await asyncio.to_thread(ToolContext.current) is EMPTY + + +async def _current_later() -> ToolContext: + await asyncio.sleep(0) + return ToolContext.current() diff --git a/packages/agentenv-protocol/tests/test_tool_context_injection.py b/packages/agentenv-protocol/tests/test_tool_context_injection.py new file mode 100644 index 00000000..36cf890f --- /dev/null +++ b/packages/agentenv-protocol/tests/test_tool_context_injection.py @@ -0,0 +1,459 @@ +"""ToolContext on a real FastMCP app (mcp installed, in-memory sessions): schema omission on the +wire, injection, per-call isolation, the on_tool_call chokepoint, and the ordering contract with a +pre-installed dispatch guard.""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace +from typing import Annotated, Any + +import anyio +import httpx +import pytest +from pydantic import Field + +pytest.importorskip("mcp.server.fastmcp") + +from mcp import ClientSession # noqa: E402 +from mcp.client import streamable_http as _streamable_http # noqa: E402 +from mcp.server.fastmcp import Context, FastMCP # noqa: E402 +from mcp.server.fastmcp.tools.base import Tool # noqa: E402 +from mcp.shared.memory import create_connected_server_and_client_session # noqa: E402 +from mcp.types import CallToolResult, TextContent # noqa: E402 + +from agentenv_protocol import ( # noqa: E402 + ROLE_HEADER, + ROLE_META_KEY, + SESSION_META_KEY, + AgentEnvEnvironment, + Caller, + EnvironmentCard, + ToolContext, + create_fastmcp_app, + environment_card, + injecting, + tool, +) +from agentenv_protocol.agent_env_environment import _is_fastmcp_context_annotation, _schema_from_signature # noqa: E402 +from agentenv_protocol.tool_context import EMPTY, bound # noqa: E402 + + +def _text(result: Any) -> str: + """The text of FastMCP.call_tool's result: a (content, structured) tuple for structured tools, else content.""" + content = result[0] if isinstance(result, tuple) else result + return content[0].text + + +def _request_context(meta_extra: dict | None = None, headers: dict | None = None) -> SimpleNamespace: + meta = SimpleNamespace(model_extra=meta_extra) if meta_extra is not None else None + request = SimpleNamespace(headers=headers) if headers is not None else None + return SimpleNamespace(request_context=SimpleNamespace(meta=meta, request=request)) + + +async def _card_tools(app: FastMCP) -> dict: + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app.streamable_http_app()), base_url="http://t") as client: + card = (await client.get("/.well-known/agent-env.json")).json() + return {t["name"]: t for t in card["capabilities"]["tools"]} + + +@environment_card(name="items") +class _Env(AgentEnvEnvironment): + def __init__(self) -> None: + self.seen: list = [] + self.hook_roles: list = [] + self.alice_in = asyncio.Event() + self.release = asyncio.Event() + + def on_tool_call(self, context: ToolContext) -> None: + self.hook_roles.append(context.caller.role) + + @tool(name="rendezvous") + async def rendezvous(self, who: str, tc: ToolContext) -> str: + """alice parks inside her call until bob has run his, so the two bodies provably overlap.""" + if who == "alice": + self.alice_in.set() + await self.release.wait() + else: + await self.alice_in.wait() + self.release.set() + return f"{tc.role}|{ToolContext.current().role}" + + @tool(name="whoami") + async def whoami(self, q: str, tc: ToolContext, limit: int = 3) -> dict[str, Any]: + self.seen.append((tc, ToolContext.current())) + return {"q": q, "role": tc.role, "limit": limit} + + @tool(name="whoami_sync") + def whoami_sync(self, tc: ToolContext) -> str: + self.seen.append((tc, ToolContext.current())) + return tc.role or "-" + + @tool(name="slow") + async def slow(self, delay: float, tc: ToolContext) -> str: + await asyncio.sleep(delay) + return f"{tc.role}|{ToolContext.current().role}" + + @tool(name="{environment_name}_add_item") + async def add_item( + self, + item: Annotated[str, Field(description="The item to add.")], + times: Annotated[int, Field(description="How many copies to add.")] = 1, + ) -> dict: + """Add an item to the store.""" + self.seen.append(ToolContext.current()) + return {"item": item, "times": times} + + +@pytest.mark.asyncio +async def test_card_and_fastmcp_tools_list_both_omit_the_slot(): + env = _Env() + app = env.create_app() + expected = { + "type": "object", + "properties": {"q": {"type": "string"}, "limit": {"type": "integer", "default": 3}}, + "required": ["q"], + } + card_tools = await _card_tools(app) + assert card_tools["whoami"]["inputSchema"] == expected + listed = {t.name: t for t in await app.list_tools()} + assert set(listed["whoami"].inputSchema["properties"]) == {"q", "limit"} + assert listed["whoami"].inputSchema["required"] == ["q"] + assert "$defs" not in listed["whoami"].inputSchema + assert listed["whoami_sync"].inputSchema["properties"] == {} + assert listed["whoami"].outputSchema is not None # dict[str, Any] stays structured through the twin + + # card and tools/list derive the schema independently from the same reduced signature + ours = _schema_from_signature(env.whoami) + theirs = Tool.from_function(injecting(env.whoami)).parameters + assert set(ours["properties"]) == set(theirs["properties"]) and ours["required"] == theirs["required"] + for prop, spec in ours["properties"].items(): + assert spec["type"] == theirs["properties"][prop]["type"] + + +@pytest.mark.asyncio +async def test_injection_sync_and_async_and_ambient_equals_injected(): + env = _Env() + app = env.create_app() + assert app._tool_manager.get_tool("whoami").is_async is True + assert app._tool_manager.get_tool("whoami_sync").is_async is False + + assert json.loads(_text(await app.call_tool("whoami", {"q": "x"}))) == {"q": "x", "role": None, "limit": 3} + assert _text(await app.call_tool("whoami_sync", {})) == "-" + assert len(env.seen) == 2 + for name, (injected, ambient) in zip(("whoami", "whoami_sync"), env.seen, strict=True): + assert isinstance(injected, ToolContext) and injected is ambient + assert injected.tool == name and injected.transport == "mcp" and injected.call_id + assert env.seen[0][0].arguments == {"q": "x"} + assert ToolContext.current() is EMPTY + + twin = app._tool_manager.get_tool("whoami").fn + with bound(ToolContext.for_test(role="x")) as ctx: + assert await twin(q="direct") == {"q": "direct", "role": "x", "limit": 3} + assert env.seen[-1][0] is ctx + + +@pytest.mark.asyncio +async def test_concurrent_calls_with_different_meta_roles_never_cross(): + env = _Env() + app = env.create_app() + results: dict[str, str] = {} + + async def call(role: str) -> None: + async with create_connected_server_and_client_session(app) as session: + result = await session.call_tool("rendezvous", {"who": role}, meta={ROLE_META_KEY: role}) + assert result.isError is False + results[role] = result.content[0].text + + with anyio.fail_after(10): + await asyncio.gather(call("alice"), call("bob")) + assert results == {"alice": "alice|alice", "bob": "bob|bob"} + assert set(env.hook_roles) == {"alice", "bob"} and len(env.hook_roles) == 2 + assert ToolContext.current() is EMPTY + + +@pytest.mark.asyncio +async def test_header_fallback_and_default_blank_roles(): + env = _Env() + app = env.create_app() + call = app._tool_manager.call_tool + headers = {ROLE_HEADER: "hdr", "mcp-session-id": "sid"} + + await call("whoami_sync", {}, context=_request_context(headers=headers)) + assert env.seen[-1][0].caller == Caller("hdr", "sid", "hdr") + await call("whoami_sync", {}, context=_request_context({ROLE_META_KEY: "meta", SESSION_META_KEY: "s-meta"}, headers)) + assert env.seen[-1][0].caller == Caller("meta", "s-meta", "meta") + await call("whoami_sync", {}, context=_request_context({ROLE_META_KEY: "default"}, headers)) + assert env.seen[-1][0].caller == Caller(None, None, "default") + await call("whoami_sync", {}, context=_request_context({ROLE_META_KEY: " "}, None)) + assert env.seen[-1][0].caller.role is None + + async with create_connected_server_and_client_session(app) as session: + result = await session.call_tool("whoami_sync", {}, meta={ROLE_META_KEY: "default"}) + assert result.content[0].text == "-" and env.seen[-1][0].caller.raw_role == "default" + + +def _refusal() -> CallToolResult: + return CallToolResult(isError=True, content=[TextContent(type="text", text="refused: intruder")]) + + +@environment_card(name="guarded") +class _GuardedEnv(AgentEnvEnvironment): + def __init__(self) -> None: + self.hook_calls: list = [] + self.ran: list = [] + + def on_tool_call(self, context: ToolContext) -> CallToolResult | None: + self.hook_calls.append((context.tool, dict(context.arguments), context.caller.role)) + return _refusal() if context.caller.role == "intruder" else None + + @tool(name="structured") + def structured(self, x: int, tc: ToolContext) -> dict[str, Any]: + self.ran.append("structured") + return {"x": x, "role": tc.role} + + @tool(name="plain") + def plain(self, x: int, tc: ToolContext) -> str: + self.ran.append("plain") + return f"{x}:{tc.role}" + + @tool(name="slotless") + def slotless(self, x: int) -> str: + self.ran.append("slotless") + return str(x) + + +class _AsyncGuardedEnv(_GuardedEnv): + async def on_tool_call(self, context: ToolContext) -> CallToolResult | None: + await asyncio.sleep(0) + return _GuardedEnv.on_tool_call(self, context) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("env_cls", [_GuardedEnv, _AsyncGuardedEnv]) +async def test_on_tool_call_refusal_reaches_the_wire_as_is_error_for_structured_and_unstructured_tools(env_cls): + env = env_cls() + app = env.create_app() + assert app._tool_manager.get_tool("structured").output_schema is not None + assert app._tool_manager.get_tool("plain").output_schema is not None + + async with create_connected_server_and_client_session(app) as session: + for name in ("structured", "plain", "slotless"): + result = await session.call_tool(name, {"x": 1}, meta={ROLE_META_KEY: "intruder"}) + assert result.isError is True, name + assert result.content[0].text == "refused: intruder", name + assert result.structuredContent is None + assert env.ran == [] + assert env.hook_calls == [(name, {"x": 1}, "intruder") for name in ("structured", "plain", "slotless")] + + result = await session.call_tool("structured", {"x": 1}, meta={ROLE_META_KEY: "alice"}) + assert result.isError is False and result.structuredContent == {"x": 1, "role": "alice"} + result = await session.call_tool("plain", {"x": 2}, meta={ROLE_META_KEY: "alice"}) + assert result.isError is False and result.content[0].text == "2:alice" + assert env.ran == ["structured", "plain"] + + +@pytest.mark.asyncio +async def test_on_tool_call_exceptions_and_bad_returns_surface_as_errors(): + class _Raising(_GuardedEnv): + def on_tool_call(self, context: ToolContext) -> None: + raise RuntimeError("boom") + + class _BadReturn(_GuardedEnv): + def on_tool_call(self, context: ToolContext) -> Any: + return {"not": "a CallToolResult"} + + async with create_connected_server_and_client_session(_Raising().create_app()) as session: + result = await session.call_tool("plain", {"x": 1}) + assert result.isError is True and "boom" in result.content[0].text + + bad = _BadReturn() + app = bad.create_app() + with pytest.raises(TypeError, match="on_tool_call must return None or mcp.types.CallToolResult, got dict"): + await app.call_tool("plain", {"x": 1}) + async with create_connected_server_and_client_session(app) as session: + result = await session.call_tool("plain", {"x": 1}) + assert result.isError is True and "on_tool_call must return None or mcp.types.CallToolResult" in result.content[0].text + assert bad.ran == [] + + +@environment_card(name="ctx") +class _ContextEnv(AgentEnvEnvironment): + @tool(name="t") + def t(self, x: int, ctx: Context) -> str: + return type(ctx).__name__ + + @tool(name="both") + def both(self, x: int, ctx: Context, tc: ToolContext) -> str: + return f"{type(ctx).__name__}:{tc.tool}" + + +@pytest.mark.asyncio +async def test_fastmcp_context_param_is_off_the_card_but_still_injected(): + app = _ContextEnv().create_app() + card_tools = await _card_tools(app) + assert set(card_tools["t"]["inputSchema"]["properties"]) == {"x"} + assert set(card_tools["both"]["inputSchema"]["properties"]) == {"x"} + listed = {t.name: t for t in await app.list_tools()} + assert set(listed["t"].inputSchema["properties"]) == {"x"} + assert set(listed["both"].inputSchema["properties"]) == {"x"} + assert app._tool_manager.get_tool("t").context_kwarg == "ctx" + assert app._tool_manager.get_tool("both").context_kwarg == "ctx" + assert _text(await app.call_tool("t", {"x": 1})) == "Context" + assert _text(await app.call_tool("both", {"x": 1})) == "Context:both" + + +def test_fastmcp_context_string_fallback_when_hints_do_not_resolve(): + for text in ("Context", "Optional[Context]", "mcp.server.fastmcp.Context", "Context | None"): + assert _is_fastmcp_context_annotation(text), text + assert not _is_fastmcp_context_annotation("ContextVar") + + def handler(x: "Undefined", ctx: Context, tc: ToolContext) -> str: ... # noqa: F821 + + assert set(_schema_from_signature(handler)["properties"]) == {"x"} + + +@pytest.mark.asyncio +async def test_tool_without_slot_is_registered_verbatim_with_byte_identical_schema(): + env = _Env() + app = env.create_app() + registered = app._tool_manager.get_tool("items_add_item") + assert registered.fn == env.add_item + assert registered.parameters == Tool.from_function(env.add_item, name="items_add_item").parameters + card_tools = await _card_tools(app) + assert json.dumps(card_tools["items_add_item"]["inputSchema"], sort_keys=True) == json.dumps({ + "type": "object", + "properties": { + "item": {"type": "string", "description": "The item to add."}, + "times": {"type": "integer", "description": "How many copies to add.", "default": 1}, + }, + "required": ["item"], + }, sort_keys=True) + + await app.call_tool("items_add_item", {"item": "x", "times": 2}) + ambient = env.seen[-1] + assert ambient.tool == "items_add_item" and ambient.arguments == {"item": "x", "times": 2} + assert ToolContext.current() is EMPTY + + +@pytest.mark.asyncio +async def test_pre_installed_guard_runs_inside_the_binding(): + app = FastMCP("items") + original = app._tool_manager.call_tool + guard_saw: list = [] + + async def guard(name: str, arguments: dict, context: Any = None, convert_result: bool = False) -> Any: + guard_saw.append((name, ToolContext.current().caller.role)) + return await original(name, arguments, context=context, convert_result=convert_result) + + app._tool_manager.call_tool = guard + env = _Env() + env.mount(app) + installed = app._tool_manager.call_tool + assert installed is not guard and installed.__wrapped__ is guard + assert installed.__agentenv_tool_context__ is True + + async with create_connected_server_and_client_session(app) as session: + result = await session.call_tool("slow", {"delay": 0}, meta={ROLE_META_KEY: "alice"}) + assert result.isError is False and result.content[0].text == "alice|alice" + assert guard_saw == [("slow", "alice")] + assert env.hook_roles == ["alice"] + + +def test_raw_slotted_function_cannot_be_registered_on_fastmcp(): + env = _Env() + with pytest.raises(TypeError, match=r"agentenv_protocol\.injecting\(fn\)"): + Tool.from_function(env.whoami) + app = FastMCP("raw") + with pytest.raises(TypeError, match="ToolContext is injected by the SDK and is not a wire parameter"): + app.tool(name="whoami")(env.whoami) + + def optional_slot(q: str, tc: ToolContext | None = None) -> str: + return q + + with pytest.raises(TypeError, match="injecting"): + Tool.from_function(optional_slot) + + app.tool(name="whoami")(injecting(env.whoami)) + assert set(app._tool_manager.get_tool("whoami").parameters["properties"]) == {"q", "limit"} + + +@pytest.mark.asyncio +async def test_create_fastmcp_app_binds_tool_calls_for_a_plain_handler_without_a_hook(): + class _Plain: + def __init__(self) -> None: + self.seen: list = [] + + @tool(name="echo") + def echo(self, q: str, tc: ToolContext) -> str: + self.seen.append(tc) + return q + + handler = _Plain() + app = create_fastmcp_app(handler, card=EnvironmentCard(name="plain")) + assert app._tool_manager.call_tool.__agentenv_tool_context__ is True + assert _text(await app.call_tool("echo", {"q": "hi"})) == "hi" + assert handler.seen[0].tool == "echo" and handler.seen[0].transport == "mcp" + + +@pytest.mark.asyncio +@pytest.mark.skipif(not hasattr(_streamable_http, "streamable_http_client"), reason="mcp without streamable_http_client") +async def test_direct_caller_over_streamable_http_gets_the_header_role_and_its_own_session(): + env = _Env() + app = FastMCP("direct", json_response=True) + env.mount(app) + base = "http://127.0.0.1:8000" + client = httpx.AsyncClient(transport=httpx.ASGITransport(app=app.streamable_http_app()), base_url=base, + headers={ROLE_HEADER: "grader"}) + + with anyio.fail_after(30): + async with app.session_manager.run(), client: + async with _streamable_http.streamable_http_client(f"{base}/mcp", http_client=client) as (read, write, session_id): + async with ClientSession(read, write) as session: + await session.initialize() + result = await session.call_tool("whoami_sync", {}) + forwarded = await session.call_tool("whoami_sync", {}, meta={ROLE_META_KEY: "bob"}) + sid = session_id() + assert sid and result.content[0].text == "grader" and forwarded.content[0].text == "bob" + direct, proxied = env.seen[-2][0], env.seen[-1][0] + assert direct.caller == Caller("grader", sid, "grader") and direct.mcp is not None + assert proxied.caller == Caller("bob", None, "bob") # a _meta role entry: the connection's id is not the caller's + + +def test_a_second_handler_on_the_same_app_is_refused_at_mount(): + app = FastMCP("shared") + _Env().mount(app) + with pytest.raises(RuntimeError, match="one handler per app"): + _Env().mount(app) + + +def test_positional_only_slot_and_mismatched_hook_fail_at_mount(): + class _PositionalOnly(AgentEnvEnvironment): + @tool(name="po") + def po(self, tc: ToolContext, /, x: int) -> int: + return x + + with pytest.raises(TypeError, match="must not be positional-only"): + _PositionalOnly().create_app() + + class _Arity: + def on_tool_call(self, role, name, arguments, result): + return None + + @tool(name="echo") + def echo(self, q: str) -> str: + return q + + with pytest.raises(TypeError, match="must take the ToolContext as its only argument"): + create_fastmcp_app(_Arity(), card=EnvironmentCard(name="arity")) + + +@pytest.mark.asyncio +async def test_unknown_tool_is_answered_by_the_dispatch_without_the_hook(): + env = _Env() + app = env.create_app() + async with create_connected_server_and_client_session(app) as session: + result = await session.call_tool("nope", {"x": 1}, meta={ROLE_META_KEY: "alice"}) + assert result.isError is True and "Unknown tool" in result.content[0].text + assert env.hook_roles == [] and env.seen == [] From 0fe9faa51cbcfa3997d976c94c0e0f92b1cc719d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 16:08:49 -0700 Subject: [PATCH 2/7] protocol: refuse a mount before touching the app; state what the role asserts The already-bound-dispatch and on_tool_call checks now run before any route or tool is added, so a refused mount leaves the app as it found it. The README says the role is the caller's side's assertion, not an authentication result, and the ToolContext docstring says which constructions snapshot arguments. Co-Authored-By: Claude Fable 5.1 --- packages/agentenv-protocol/README.md | 6 ++++- .../agent_env_environment.py | 22 +++++++++++++---- .../src/agentenv_protocol/tool_context.py | 6 +++-- .../tests/test_tool_context_injection.py | 24 +++++++++++++++++-- 4 files changed, 48 insertions(+), 10 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 7de99c86..f4f44b3d 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -99,7 +99,11 @@ from agentenv_protocol import ToolContext ``` `ctx.caller.role` is the role the agent was deployed with, as the gateway forwarded it (`None` when -nothing was forwarded; the gateway's `default` role counts as nothing). `ctx.caller.session` tells +nothing was forwarded; the gateway's `default` role counts as nothing). It is the caller's side's +assertion, trusted exactly as far as that side is: behind a gateway it is the gateway's view of the +agent it deployed; a caller that reaches the server directly chooses it in `AgentEnv-Role`. It says +whose data a call acts on, not that the caller was authenticated, so a server exposed to untrusted +direct callers keeps access control on its own boundary. `ctx.caller.session` tells callers apart when the gateway shares one connection to the server. `ctx.tool`, `ctx.arguments`, `ctx.call_id` and `ctx.transport` (`"mcp"` or `"rest"`) describe the call. Code that runs on behalf of the call without a parameter of its own, such as a database method, reads the same object through diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index 9fc5e111..0ce2627a 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -183,6 +183,7 @@ def __init__(self, environment_card: EnvironmentCard, handler: Any) -> None: self._tools = [(advertised.get(descriptor.name, descriptor), fn) for descriptor, fn in tool_handlers] def add_routes_to_app(self, app: Any, *, rpc_url: str = RPC_PATH, card_url: str = WELL_KNOWN_PATH) -> None: + self._check_tool_binding(app) self.environment_card = self._card_for(app) self._add_route(app, rpc_url, ["POST"], self._dispatch) self._add_route(app, card_url, ["GET"], _card_handler(self.environment_card)) @@ -201,6 +202,10 @@ def _add_route(self, app: Any, path: str, methods: list, handler: Callable) -> N def _register_tool(self, app: Any, descriptor: EnvironmentTool, fn: Callable) -> None: raise NotImplementedError(f"{type(self).__name__} does not serve MCP tools") + def _check_tool_binding(self, app: Any) -> None: + """Everything that can refuse the mount, before the app is touched; nothing to refuse without a tool dispatch.""" + return None + def _bind_tool_calls(self, app: Any) -> None: """Frameworks with a tool dispatch chokepoint bind a ToolContext around every call there; the others serve no tools.""" return None @@ -211,6 +216,7 @@ def _card_for(self, app: Any) -> EnvironmentCard: class AgentEnvFastMCPApplication(AgentEnvApplication): supports_tools: ClassVar[bool] = True + _hook: Callable | None = None def _add_route(self, app: Any, path: str, methods: list, handler: Callable) -> None: app.custom_route(path, methods=methods)(handler) @@ -218,22 +224,28 @@ def _add_route(self, app: Any, path: str, methods: list, handler: Callable) -> N def _register_tool(self, app: Any, descriptor: EnvironmentTool, fn: Callable) -> None: app.tool(name=descriptor.name, description=descriptor.description)(injecting(fn)) + def _check_tool_binding(self, app: Any) -> None: + """Refuse before any route or tool is added: a dispatch this protocol already wraps (one handler per + app, or the second hook would answer the first handler's tools) and an ``on_tool_call`` of the wrong + shape. The hook is kept for ``_bind_tool_calls``.""" + inner = getattr(getattr(app, "_tool_manager", None), "call_tool", None) + if callable(inner) and getattr(inner, _TOOL_CONTEXT_BOUND_ATTR, False): + raise RuntimeError("this FastMCP app already carries an AgentEnv tool dispatch; mount one handler per app") + self._hook = _tool_call_hook(self.handler) + def _bind_tool_calls(self, app: Any) -> None: """Wrap ``app._tool_manager.call_tool``, which every tools/call resolves at call time, so each call of a registered tool runs under its ToolContext and the handler's ``on_tool_call`` may answer it first. An unknown tool name goes straight to the dispatch, whose error answers it. The installed callable carries ``__agentenv_tool_context__ = True`` and ``__wrapped__`` (the - dispatch it wraps): the handles a downstream guard uses to assert the ordering contract. One - handler per app: a second binding would run its hook on the first handler's tools. + dispatch it wraps): the handles a downstream guard uses to assert the ordering contract. """ manager = getattr(app, "_tool_manager", None) inner = getattr(manager, "call_tool", None) if not callable(inner): return - if getattr(inner, _TOOL_CONTEXT_BOUND_ATTR, False): - raise RuntimeError("this FastMCP app already carries an AgentEnv tool dispatch; mount one handler per app") - hook = _tool_call_hook(self.handler) + hook = self._hook get_tool = getattr(manager, "get_tool", None) async def call_tool(name: str, arguments: dict, *args: Any, **kwargs: Any) -> Any: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py index e33e6f7f..a4fc410b 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py @@ -98,8 +98,10 @@ class ToolContext: """What the protocol knows about one tool call. Hashable, so per-call state can be keyed on the context itself: ``call_id`` - identifies the call, and ``arguments`` (a read-only mapping) and ``mcp`` - stay out of the hash. + identifies the call; ``arguments`` and ``mcp`` stay out of the hash. The + protocol's builders and :meth:`for_test` snapshot ``arguments`` into a + read-only mapping; a context constructed directly keeps the mapping it is + given. """ caller: Caller = Caller() diff --git a/packages/agentenv-protocol/tests/test_tool_context_injection.py b/packages/agentenv-protocol/tests/test_tool_context_injection.py index 36cf890f..9a2db4ae 100644 --- a/packages/agentenv-protocol/tests/test_tool_context_injection.py +++ b/packages/agentenv-protocol/tests/test_tool_context_injection.py @@ -28,6 +28,7 @@ ROLE_META_KEY, SESSION_META_KEY, AgentEnvEnvironment, + AgentEnvFastMCPApplication, Caller, EnvironmentCard, ToolContext, @@ -421,11 +422,24 @@ async def test_direct_caller_over_streamable_http_gets_the_header_role_and_its_o assert proxied.caller == Caller("bob", None, "bob") # a _meta role entry: the connection's id is not the caller's -def test_a_second_handler_on_the_same_app_is_refused_at_mount(): +def _app_shape(app: FastMCP) -> tuple: + return (sorted(t.name for t in app._tool_manager.list_tools()), len(app._custom_starlette_routes), app._tool_manager.call_tool) + + +def test_a_second_handler_on_the_same_app_is_refused_before_the_app_is_touched(): app = FastMCP("shared") _Env().mount(app) + before = _app_shape(app) + + @environment_card(name="other") + class _Other(AgentEnvEnvironment): + @tool(name="other_tool") + def other_tool(self, q: str) -> str: + return q + with pytest.raises(RuntimeError, match="one handler per app"): - _Env().mount(app) + _Other().mount(app) + assert _app_shape(app) == before def test_positional_only_slot_and_mismatched_hook_fail_at_mount(): @@ -448,6 +462,12 @@ def echo(self, q: str) -> str: with pytest.raises(TypeError, match="must take the ToolContext as its only argument"): create_fastmcp_app(_Arity(), card=EnvironmentCard(name="arity")) + app = FastMCP("untouched") + before = _app_shape(app) + with pytest.raises(TypeError, match="must take the ToolContext as its only argument"): + AgentEnvFastMCPApplication(EnvironmentCard(name="arity"), _Arity()).add_routes_to_app(app) + assert _app_shape(app) == before + @pytest.mark.asyncio async def test_unknown_tool_is_answered_by_the_dispatch_without_the_hook(): From 4e4433bd6cf7bf01defa4ee1e680eb5e93adcc13 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 14:47:00 -0700 Subject: [PATCH 3/7] protocol: refuse to mount on an app without a FastMCP tool dispatch The binding wraps the tool manager's call_tool because FastMCP in mcp 1.x has no public hook for it. A target without that manager used to mount with nothing installed, so every tool call ran unbound and nothing reading ToolContext.current() could tell that from an empty caller. The mount now refuses such an app before adding any route, the binding drops its silent fallback, and the test fakes carry the manager they stood in for. Co-Authored-By: Claude Fable 5.1 --- packages/agentenv-protocol/README.md | 4 ++- .../agent_env_environment.py | 33 +++++++++++-------- .../tests/test_agent_env_environment.py | 6 ++++ .../tests/test_data_objects.py | 13 ++++++++ .../tests/test_tool_context.py | 27 ++++++++++++--- 5 files changed, 65 insertions(+), 18 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index f4f44b3d..04001a77 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -108,7 +108,9 @@ callers apart when the gateway shares one connection to the server. `ctx.tool`, `ctx.call_id` and `ctx.transport` (`"mcp"` or `"rest"`) describe the call. Code that runs on behalf of the call without a parameter of its own, such as a database method, reads the same object through `ToolContext.current()`, which returns an empty context outside any call. The binding is per call -and follows the request's task, so concurrent callers never see each other's context. +and follows the request's task, so concurrent callers never see each other's context. It is made +at FastMCP's own tool dispatch, for which `mcp` 1.x offers no public hook, and the mount refuses an +app without one, so a server that mounted is a server whose tool calls are bound. To answer a call before the tool runs, override `on_tool_call(self, context)`: return `None` to let it proceed, or an `mcp.types.CallToolResult` (`isError=True` for a refusal) to send that as the diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index 0ce2627a..825c8067 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -225,31 +225,38 @@ def _register_tool(self, app: Any, descriptor: EnvironmentTool, fn: Callable) -> app.tool(name=descriptor.name, description=descriptor.description)(injecting(fn)) def _check_tool_binding(self, app: Any) -> None: - """Refuse before any route or tool is added: a dispatch this protocol already wraps (one handler per - app, or the second hook would answer the first handler's tools) and an ``on_tool_call`` of the wrong - shape. The hook is kept for ``_bind_tool_calls``.""" - inner = getattr(getattr(app, "_tool_manager", None), "call_tool", None) - if callable(inner) and getattr(inner, _TOOL_CONTEXT_BOUND_ATTR, False): + """Refuse before any route or tool is added: an app without the dispatch the binding wraps (its tool + calls would run unbound, which nothing reading ``ToolContext.current()`` can tell from an empty + caller), a dispatch this protocol already wraps (one handler per app, or the second hook would answer + the first handler's tools) and an ``on_tool_call`` of the wrong shape. The hook is kept for + ``_bind_tool_calls``.""" + manager = getattr(app, "_tool_manager", None) + if not (callable(getattr(manager, "call_tool", None)) and callable(getattr(manager, "get_tool", None))): + raise RuntimeError( + "this app has no FastMCP tool dispatch to bind a ToolContext around " + "(mcp>=1.25,<2 keeps it at FastMCP._tool_manager); " + "an app that serves no tools mounts through AgentEnvStarletteApplication" + ) + if getattr(manager.call_tool, _TOOL_CONTEXT_BOUND_ATTR, False): raise RuntimeError("this FastMCP app already carries an AgentEnv tool dispatch; mount one handler per app") self._hook = _tool_call_hook(self.handler) def _bind_tool_calls(self, app: Any) -> None: - """Wrap ``app._tool_manager.call_tool``, which every tools/call resolves at call time, so each call + """Wrap the tool manager's ``call_tool``, which every tools/call resolves at call time, so each call of a registered tool runs under its ToolContext and the handler's ``on_tool_call`` may answer - it first. An unknown tool name goes straight to the dispatch, whose error answers it. + it first. An unknown tool name goes straight to the dispatch, whose error answers it. FastMCP + has no public hook for this (its low-level handler is ``FastMCP.call_tool``, bound at + construction), so the manager is the one private seam, required by ``_check_tool_binding``. The installed callable carries ``__agentenv_tool_context__ = True`` and ``__wrapped__`` (the dispatch it wraps): the handles a downstream guard uses to assert the ordering contract. """ - manager = getattr(app, "_tool_manager", None) - inner = getattr(manager, "call_tool", None) - if not callable(inner): - return + manager = app._tool_manager + inner = manager.call_tool hook = self._hook - get_tool = getattr(manager, "get_tool", None) async def call_tool(name: str, arguments: dict, *args: Any, **kwargs: Any) -> Any: - if get_tool is not None and get_tool(name) is None: + if manager.get_tool(name) is None: return await inner(name, arguments, *args, **kwargs) context = kwargs.get("context", args[0] if args else None) with bound(context_from_mcp(context, name, arguments)) as tool_context: diff --git a/packages/agentenv-protocol/tests/test_agent_env_environment.py b/packages/agentenv-protocol/tests/test_agent_env_environment.py index 798d9d91..ed4ef77f 100644 --- a/packages/agentenv-protocol/tests/test_agent_env_environment.py +++ b/packages/agentenv-protocol/tests/test_agent_env_environment.py @@ -5,6 +5,7 @@ import datetime import decimal import enum +import inspect import json import sys import uuid @@ -47,6 +48,11 @@ def __init__(self) -> None: self.routes: list[Route] = [] self.tools: dict[str, tuple] = {} self.settings = SimpleNamespace(streamable_http_path="/mcp") + self._tool_manager = SimpleNamespace(get_tool=self.tools.get, call_tool=self._call_tool) + + async def _call_tool(self, name: str, arguments: dict, *_: object, **__: object): + result = self.tools[name][1](**arguments) + return await result if inspect.isawaitable(result) else result def custom_route(self, path: str, methods: list[str]): def deco(fn): diff --git a/packages/agentenv-protocol/tests/test_data_objects.py b/packages/agentenv-protocol/tests/test_data_objects.py index 6e42a271..dd235e69 100644 --- a/packages/agentenv-protocol/tests/test_data_objects.py +++ b/packages/agentenv-protocol/tests/test_data_objects.py @@ -3,6 +3,7 @@ from __future__ import annotations +import inspect import json from datetime import datetime, timedelta, timezone from types import SimpleNamespace @@ -51,7 +52,13 @@ def _grant(*, max_object_bytes: int = 1024) -> WriteNamespaceGrant: class _FakeMCP: def __init__(self) -> None: self.routes: list[Route] = [] + self.tools: dict[str, tuple] = {} self.settings = SimpleNamespace(streamable_http_path="/mcp") + self._tool_manager = SimpleNamespace(get_tool=self.tools.get, call_tool=self._call_tool) + + async def _call_tool(self, name: str, arguments: dict, *_: object, **__: object): + result = self.tools[name][1](**arguments) + return await result if inspect.isawaitable(result) else result def custom_route(self, path: str, methods: list[str]): def deco(fn): @@ -59,6 +66,12 @@ def deco(fn): return fn return deco + def tool(self, name: str | None = None, description: str | None = None): + def deco(fn): + self.tools[name or fn.__name__] = (description, fn) + return fn + return deco + class _InlineHandler: def __init__(self) -> None: diff --git a/packages/agentenv-protocol/tests/test_tool_context.py b/packages/agentenv-protocol/tests/test_tool_context.py index d476dbdd..f262dbb7 100644 --- a/packages/agentenv-protocol/tests/test_tool_context.py +++ b/packages/agentenv-protocol/tests/test_tool_context.py @@ -36,12 +36,17 @@ class _FakeMCP: - """FastMCP-shaped target without a ``_tool_manager``: records routes and tool registrations.""" + """FastMCP-shaped target: records routes and tool registrations; its tool manager is the dispatch the mount binds.""" def __init__(self) -> None: self.routes: list[Route] = [] self.tools: dict[str, tuple] = {} self.settings = SimpleNamespace(streamable_http_path="/mcp") + self._tool_manager = SimpleNamespace(get_tool=self.tools.get, call_tool=self._call_tool) + + async def _call_tool(self, name: str, arguments: dict, *_: object, **__: object): + result = self.tools[name][1](**arguments) + return await result if inspect.isawaitable(result) else result def custom_route(self, path: str, methods: list[str]): def deco(fn): @@ -256,9 +261,23 @@ def broken(x: "Undefined", tc: ToolContext) -> str: ... # noqa: F821 injecting(broken) -def test_fake_target_without_tool_manager_mounts_with_nothing_installed(): - mcp = _mounted(_SlotHandler()) - assert not hasattr(mcp, "_tool_manager") and not hasattr(mcp, "call_tool") +@pytest.mark.parametrize("strip", [lambda mcp: delattr(mcp, "_tool_manager"), lambda mcp: delattr(mcp._tool_manager, "call_tool")]) +def test_target_without_a_tool_dispatch_is_refused_before_any_route_is_added(strip): + mcp = _FakeMCP() + strip(mcp) + with pytest.raises(RuntimeError, match="no FastMCP tool dispatch"): + AgentEnvFastMCPApplication(environment_card=EnvironmentCard(name="items"), handler=_SlotHandler()).add_routes_to_app(mcp) + assert mcp.routes == [] and mcp.tools == {} + + +@pytest.mark.asyncio +async def test_mount_binds_the_target_dispatch(): + handler = _SlotHandler() + mcp = _mounted(handler) + assert mcp._tool_manager.call_tool.__agentenv_tool_context__ is True + assert await mcp._tool_manager.call_tool("whoami", {"q": "x"}, context=None) == {"q": "x", "role": None, "limit": 3} + assert handler.seen[-1].transport == "mcp" and handler.seen[-1].tool == "whoami" + assert ToolContext.current() is EMPTY class _ExtHandler: From ef4ed970d4241d3d639eb8d60b15a0b5b106a13b Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 15:17:17 -0700 Subject: [PATCH 4/7] protocol: the twin binds positional arguments against its own signature The reduced signature a slotted tool advertises omits the ToolContext parameter, but the twin forwarded positional arguments in the original order, so a positional call to a slot-first tool put the first value into the slot and then supplied the slot again by keyword. The twin now binds what it receives against the signature it advertises and rebinds onto the original one before adding the context, for sync and async tools alike. FastMCP always calls by keyword, so the wire path was unaffected. Co-Authored-By: Claude Fable 5.1 --- .../agent_env_environment.py | 26 +++++++++++++------ .../tests/test_tool_context.py | 22 ++++++++++++++++ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index 825c8067..e139f27f 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -466,7 +466,7 @@ def _callable_name(fn: Callable) -> str: def _tool_context_slot(fn: Callable) -> str | None: """Name of the parameter annotated ``ToolContext`` / ``Optional[ToolContext]`` (any name), or None. - The slot is filled by keyword, so a positional-only one is refused here rather than on every call.""" + The slot is filled by name, so a positional-only one is refused here rather than on every call.""" hints = _resolved_hints(fn) slots = [] for name, param in inspect.signature(fn).parameters.items(): @@ -515,25 +515,35 @@ def _reduced_signature(fn: Callable, slot: str) -> tuple: def _inject_tool_context(fn: Callable, slot: str) -> Callable: + signature, annotations = _reduced_signature(fn, slot) + full = inspect.signature(fn) + + def filled(args: tuple, kwargs: dict) -> inspect.BoundArguments: + # Positional arguments mean what the twin advertises, then land where ``fn`` declares them. + call = full.bind_partial() + call.arguments.update(signature.bind(*args, **kwargs).arguments) + call.arguments[slot] = ToolContext.current() + return call + # Two bodies so the registry's sync/async detection sees the same kind as ``fn``. if inspect.iscoroutinefunction(fn): async def wrapper(*args: Any, **kwargs: Any) -> Any: - kwargs[slot] = ToolContext.current() - return await fn(*args, **kwargs) + call = filled(args, kwargs) + return await fn(*call.args, **call.kwargs) else: def wrapper(*args: Any, **kwargs: Any) -> Any: - kwargs[slot] = ToolContext.current() - return fn(*args, **kwargs) + call = filled(args, kwargs) + return fn(*call.args, **call.kwargs) functools.update_wrapper(wrapper, fn) - wrapper.__signature__, wrapper.__annotations__ = _reduced_signature(fn, slot) + wrapper.__signature__, wrapper.__annotations__ = signature, annotations return wrapper def injecting(fn: Callable) -> Callable: """``fn`` as an MCP tool registry should see it: unchanged when it declares no ``ToolContext`` parameter, else a twin of the same name, doc and sync/async nature whose signature and annotations - omit that parameter and which fills it with ``ToolContext.current()`` on every call. ``@tool`` - methods get this at mount; use it for a slotted function registered any other way + omit that parameter and which fills it with ``ToolContext.current()`` on every call; positional + arguments follow the twin's signature, not ``fn``'s. ``@tool`` methods get this at mount; use it for a slotted function registered any other way (``app.tool()(fn)``, a separate ``ToolManager``), which otherwise fails at registration.""" slot = _tool_context_slot(fn) return _inject_tool_context(fn, slot) if slot else fn diff --git a/packages/agentenv-protocol/tests/test_tool_context.py b/packages/agentenv-protocol/tests/test_tool_context.py index f262dbb7..53c100c8 100644 --- a/packages/agentenv-protocol/tests/test_tool_context.py +++ b/packages/agentenv-protocol/tests/test_tool_context.py @@ -254,6 +254,28 @@ def test_slot_less_tool_is_registered_verbatim(): } +@pytest.mark.asyncio +async def test_twin_takes_positional_arguments_in_its_advertised_order(): + class _Handler: + async def first(self, tc: ToolContext, q: str, limit: int = 3) -> dict: + return {"q": q, "limit": limit, "role": tc.role} + + def middle(self, q: str, tc: ToolContext, limit: int = 3) -> dict: + return {"q": q, "limit": limit, "role": tc.role} + + handler = _Handler() + first, middle = injecting(handler.first), injecting(handler.middle) + assert list(inspect.signature(first).parameters) == ["q", "limit"] + with bound(ToolContext.for_test(role="alice")): + assert await first("x") == {"q": "x", "limit": 3, "role": "alice"} + assert await first("x", 5) == await first(q="x", limit=5) == {"q": "x", "limit": 5, "role": "alice"} + assert middle("y", 2) == middle(limit=2, q="y") == {"q": "y", "limit": 2, "role": "alice"} + with pytest.raises(TypeError, match="missing a required argument: 'q'"): + await first() + with pytest.raises(TypeError, match="too many positional arguments"): + await first("x", 5, 7) + + def test_injecting_raises_when_a_slotted_function_has_unresolvable_annotations(): def broken(x: "Undefined", tc: ToolContext) -> str: ... # noqa: F821 From e51399accb55f9968c345705c2b4d64e7c173822 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 15:42:00 -0700 Subject: [PATCH 5/7] protocol: drop the testing module; ToolContext.as_caller builds a caller's context Nothing test-only ships in the package. The context-manager helper in testing.py was bound(ToolContext.for_test(...)) under another name, and for_test itself was the same thing living in src. ToolContext.as_caller(role, session) replaces both: a normalised caller's context for code that acts as a caller outside any dispatch (a test calling a handler directly, a seeding script, an operator helper), bound with the public bound(), now exported at the package root. Its transport is "none", which now reads as "no dispatch bound this". Co-Authored-By: Claude Fable 5.1 --- packages/agentenv-protocol/README.md | 5 ++-- .../src/agentenv_protocol/__init__.py | 2 ++ .../src/agentenv_protocol/testing.py | 18 --------------- .../src/agentenv_protocol/tool_context.py | 14 ++++++----- .../tests/test_tool_context.py | 23 +++++++++---------- .../tests/test_tool_context_injection.py | 2 +- 6 files changed, 25 insertions(+), 39 deletions(-) delete mode 100644 packages/agentenv-protocol/src/agentenv_protocol/testing.py diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 04001a77..df40df90 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -115,8 +115,9 @@ app without one, so a server that mounted is a server whose tool calls are bound To answer a call before the tool runs, override `on_tool_call(self, context)`: return `None` to let it proceed, or an `mcp.types.CallToolResult` (`isError=True` for a refusal) to send that as the tool's reply. An `@extension` method may declare the same parameter; it is filled from the request's -`AgentEnv-Role` header. In tests, `agentenv_protocol.testing.tool_context(role=...)` binds a caller -around a direct call to a handler method. +`AgentEnv-Role` header. Code that acts as a caller outside any dispatch, such as a test calling a +handler directly or a seeding script, binds a context it builds: +`with bound(ToolContext.as_caller("alice")): ...`. ## Serving diff --git a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py index dbb9d474..1d84188e 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py @@ -23,6 +23,7 @@ SESSION_META_KEY, Caller, ToolContext, + bound, normalize_role, ) from .manifest import ( @@ -92,6 +93,7 @@ "ToolContext", "Caller", "normalize_role", + "bound", "ROLE_META_KEY", "SESSION_META_KEY", "ROLE_HEADER", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/testing.py b/packages/agentenv-protocol/src/agentenv_protocol/testing.py deleted file mode 100644 index 88caab48..00000000 --- a/packages/agentenv-protocol/src/agentenv_protocol/testing.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Helpers for tests of environments built on this package.""" - -from __future__ import annotations - -import contextlib -from typing import Any, Iterator, Mapping, Optional - -from .tool_context import ToolContext, bound - - -@contextlib.contextmanager -def tool_context(*, role: Optional[str] = None, session: Optional[str] = None, tool: str = "test", - arguments: Optional[Mapping[str, Any]] = None) -> Iterator[ToolContext]: - """Bind a caller for the block, as the dispatch does for one tool call, so a handler or - database method called directly sees ``ToolContext.current()`` the way it would under the - gateway. Nested blocks shadow and restore.""" - with bound(ToolContext.for_test(role=role, session=session, tool=tool, arguments=arguments)) as context: - yield context diff --git a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py index a4fc410b..63e2f132 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py @@ -61,8 +61,8 @@ #: A gateway's unforwarded role. Same meaning as no role at all. DEFAULT_ROLE = "default" -#: How the call reached the server. ``"none"`` is :data:`EMPTY`, the context -#: outside any call; the other two name the dispatch that bound the context. +#: How the call reached the server. ``"none"`` is a context no dispatch bound, :data:`EMPTY` +#: outside any call or one built with :meth:`ToolContext.as_caller`; the other two name the dispatch. Transport = Literal["mcp", "rest", "none"] @@ -99,7 +99,7 @@ class ToolContext: Hashable, so per-call state can be keyed on the context itself: ``call_id`` identifies the call; ``arguments`` and ``mcp`` stay out of the hash. The - protocol's builders and :meth:`for_test` snapshot ``arguments`` into a + protocol's builders and :meth:`as_caller` snapshot ``arguments`` into a read-only mapping; a context constructed directly keeps the mapping it is given. """ @@ -138,9 +138,11 @@ def __get_pydantic_json_schema__(cls, schema: Any, handler: Any) -> Any: ) @classmethod - def for_test(cls, *, role: Optional[str] = None, session: Optional[str] = None, tool: str = "test", - arguments: Optional[Mapping[str, Any]] = None, transport: Transport = "mcp") -> "ToolContext": - """A context for a test that calls a handler directly or binds one with :func:`bound`.""" + def as_caller(cls, role: Optional[str] = None, session: Optional[str] = None, *, tool: str = "", + arguments: Optional[Mapping[str, Any]] = None, transport: Transport = "none") -> "ToolContext": + """A context for code that acts as ``role`` outside any dispatch (a test calling a handler + directly, a seeding script, an operator helper), to bind with :func:`bound`. The role is + normalised as the builders normalise it.""" return cls(caller=Caller(role=normalize_role(role), session=session, raw_role=role), tool=tool, arguments=_frozen(arguments), call_id=uuid.uuid4().hex, transport=transport) diff --git a/packages/agentenv-protocol/tests/test_tool_context.py b/packages/agentenv-protocol/tests/test_tool_context.py index 53c100c8..1953dada 100644 --- a/packages/agentenv-protocol/tests/test_tool_context.py +++ b/packages/agentenv-protocol/tests/test_tool_context.py @@ -30,7 +30,6 @@ normalize_role, tool, ) -from agentenv_protocol import testing from agentenv_protocol.agent_env_environment import _is_tool_context_annotation, _schema_from_signature, _tool_context_slot from agentenv_protocol.tool_context import EMPTY, bound, caller_from_request, context_from_http, context_from_mcp @@ -85,9 +84,9 @@ def test_normalize_role_absent_blank_default_and_non_string_are_none(): def test_current_is_empty_outside_a_call_and_bound_restores(): assert ToolContext.current() is EMPTY assert EMPTY.transport == "none" and EMPTY.caller.role is None - with bound(ToolContext.for_test(role="a")) as outer: + with bound(ToolContext.as_caller("a")) as outer: assert ToolContext.current() is outer and outer.role == "a" - with bound(ToolContext.for_test(role="b")) as inner: + with bound(ToolContext.as_caller("b")) as inner: assert ToolContext.current() is inner assert ToolContext.current() is outer assert ToolContext.current() is EMPTY @@ -134,16 +133,16 @@ def request_context(self): assert ctx.arguments == {"q": 1} -def test_context_from_http_and_testing_helper(): +def test_context_from_http_and_as_caller(): request = SimpleNamespace(headers={ROLE_HEADER: "carol"}) ctx = context_from_http(request, "echo", {"m": 1}) assert ctx.transport == "rest" and ctx.role == "carol" and ctx.tool == "echo" and ctx.arguments == {"m": 1} - with testing.tool_context(role="alice", session="s") as bound_ctx: + with bound(ToolContext.as_caller("alice", session="s")) as bound_ctx: assert ToolContext.current() is bound_ctx - assert bound_ctx.caller == Caller("alice", "s", "alice") + assert bound_ctx.caller == Caller("alice", "s", "alice") and bound_ctx.transport == "none" assert ToolContext.current() is EMPTY - with testing.tool_context(role="default") as default_ctx: + with bound(ToolContext.as_caller("default")) as default_ctx: assert default_ctx.role is None and default_ctx.caller.raw_role == "default" @@ -235,7 +234,7 @@ async def test_card_schema_omits_the_slot_and_registration_gets_the_twin(): assert twin.__wrapped__ == handler.whoami and twin.__name__ == "whoami" assert inspect.iscoroutinefunction(twin) and not inspect.iscoroutinefunction(mcp.tools["whoami_sync"][1]) - with bound(ToolContext.for_test(role="alice")) as ctx: + with bound(ToolContext.as_caller("alice")) as ctx: assert await twin(q="x") == {"q": "x", "role": "alice", "limit": 3} assert mcp.tools["whoami_sync"][1]() == "alice" assert handler.seen == [ctx, ctx] @@ -266,7 +265,7 @@ def middle(self, q: str, tc: ToolContext, limit: int = 3) -> dict: handler = _Handler() first, middle = injecting(handler.first), injecting(handler.middle) assert list(inspect.signature(first).parameters) == ["q", "limit"] - with bound(ToolContext.for_test(role="alice")): + with bound(ToolContext.as_caller("alice")): assert await first("x") == {"q": "x", "limit": 3, "role": "alice"} assert await first("x", 5) == await first(q="x", limit=5) == {"q": "x", "limit": 5, "role": "alice"} assert middle("y", 2) == middle(limit=2, q="y") == {"q": "y", "limit": 2, "role": "alice"} @@ -353,7 +352,7 @@ def test_caller_from_request_consults_headers_only_when_no_meta_role_entry_exist def test_tool_context_is_hashable_and_its_arguments_are_read_only(): - ctx = ToolContext.for_test(role="a", arguments={"q": 1}) + ctx = ToolContext.as_caller("a", arguments={"q": 1}) assert hash(ctx) == hash(ctx) and {ctx: 1}[ctx] == 1 and ctx in {ctx} assert hash(EMPTY) is not None and EMPTY != ctx assert ctx.arguments == {"q": 1} and dict(ctx.arguments) == {"q": 1} @@ -371,7 +370,7 @@ class Audit(BaseModel): context: ToolContext note: str = "" - ctx = ToolContext.for_test(role="a") + ctx = ToolContext.as_caller("a") assert Audit(context=ctx).context is ctx with pytest.raises(ValidationError): Audit(context={"caller": {}}) @@ -383,7 +382,7 @@ class Audit(BaseModel): @pytest.mark.asyncio async def test_binding_follows_the_task_into_to_thread_and_child_tasks(): - ctx = ToolContext.for_test(role="a") + ctx = ToolContext.as_caller("a") with bound(ctx): assert await asyncio.to_thread(ToolContext.current) is ctx assert await asyncio.create_task(_current_later()) is ctx diff --git a/packages/agentenv-protocol/tests/test_tool_context_injection.py b/packages/agentenv-protocol/tests/test_tool_context_injection.py index 9a2db4ae..c1e8191f 100644 --- a/packages/agentenv-protocol/tests/test_tool_context_injection.py +++ b/packages/agentenv-protocol/tests/test_tool_context_injection.py @@ -150,7 +150,7 @@ async def test_injection_sync_and_async_and_ambient_equals_injected(): assert ToolContext.current() is EMPTY twin = app._tool_manager.get_tool("whoami").fn - with bound(ToolContext.for_test(role="x")) as ctx: + with bound(ToolContext.as_caller("x")) as ctx: assert await twin(q="direct") == {"q": "direct", "role": "x", "limit": 3} assert env.seen[-1][0] is ctx From 8fbfd5eb8527ddd4ce2f0c34dd871f3727396b17 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 20:46:31 -0700 Subject: [PATCH 6/7] refactor(protocol): drop ToolContext.transport and ToolContext.mcp A ToolContext now carries the caller, the tool name, the arguments and the call id. Neither removed field had a reader: how the call reached the server is not something a handler acts on, and FastMCP's own request context is already injected by FastMCP for the handlers that want it. Both can return with a reader; removing them keeps the first release of the type to what is used. Co-Authored-By: Claude Fable 5.1 --- packages/agentenv-protocol/README.md | 4 ++-- .../agent_env_environment.py | 2 +- .../src/agentenv_protocol/tool_context.py | 20 ++++++------------- .../tests/test_tool_context.py | 20 +++++++++---------- .../tests/test_tool_context_injection.py | 6 +++--- 5 files changed, 22 insertions(+), 30 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index df40df90..50edc99d 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -104,8 +104,8 @@ assertion, trusted exactly as far as that side is: behind a gateway it is the ga agent it deployed; a caller that reaches the server directly chooses it in `AgentEnv-Role`. It says whose data a call acts on, not that the caller was authenticated, so a server exposed to untrusted direct callers keeps access control on its own boundary. `ctx.caller.session` tells -callers apart when the gateway shares one connection to the server. `ctx.tool`, `ctx.arguments`, -`ctx.call_id` and `ctx.transport` (`"mcp"` or `"rest"`) describe the call. Code that runs on behalf +callers apart when the gateway shares one connection to the server. `ctx.tool`, `ctx.arguments` +and `ctx.call_id` describe the call. Code that runs on behalf of the call without a parameter of its own, such as a database method, reads the same object through `ToolContext.current()`, which returns an empty context outside any call. The binding is per call and follows the request's task, so concurrent callers never see each other's context. It is made diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index e139f27f..a0853a3d 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -852,7 +852,7 @@ def _coerce_params(fn: Callable, params: dict) -> dict: def _extension_handler(fn: Callable, http_method: str) -> Callable: """REST handler for an @extension route: bind request params to the handler and return JSON. A - ``ToolContext`` parameter is filled from the request headers (transport ``"rest"``) and never bound + ``ToolContext`` parameter is filled from the request headers and never bound from the params.""" slot = _tool_context_slot(fn) signature = inspect.signature(fn) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py index 63e2f132..6cc18ff8 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py @@ -29,7 +29,7 @@ import dataclasses import types import uuid -from typing import Any, Iterator, Literal, Mapping, Optional +from typing import Any, Iterator, Mapping, Optional from pydantic_core import core_schema @@ -41,7 +41,6 @@ "SESSION_META_KEY", "Caller", "ToolContext", - "Transport", "bound", "caller_from_request", "context_from_http", @@ -61,10 +60,6 @@ #: A gateway's unforwarded role. Same meaning as no role at all. DEFAULT_ROLE = "default" -#: How the call reached the server. ``"none"`` is a context no dispatch bound, :data:`EMPTY` -#: outside any call or one built with :meth:`ToolContext.as_caller`; the other two name the dispatch. -Transport = Literal["mcp", "rest", "none"] - def normalize_role(value: Any) -> Optional[str]: """None for absent, blank and the gateway's unforwarded role; else the stripped role.""" @@ -98,7 +93,7 @@ class ToolContext: """What the protocol knows about one tool call. Hashable, so per-call state can be keyed on the context itself: ``call_id`` - identifies the call; ``arguments`` and ``mcp`` stay out of the hash. The + identifies the call; ``arguments`` stays out of the hash. The protocol's builders and :meth:`as_caller` snapshot ``arguments`` into a read-only mapping; a context constructed directly keeps the mapping it is given. @@ -108,9 +103,6 @@ class ToolContext: tool: str = "" arguments: Mapping[str, Any] = dataclasses.field(default_factory=lambda: types.MappingProxyType({}), hash=False) call_id: str = "" - transport: Transport = "none" - #: FastMCP's own request context when the call came over MCP (progress, logging, session); else None. - mcp: Any = dataclasses.field(default=None, hash=False, compare=False) @property def role(self) -> Optional[str]: @@ -139,12 +131,12 @@ def __get_pydantic_json_schema__(cls, schema: Any, handler: Any) -> Any: @classmethod def as_caller(cls, role: Optional[str] = None, session: Optional[str] = None, *, tool: str = "", - arguments: Optional[Mapping[str, Any]] = None, transport: Transport = "none") -> "ToolContext": + arguments: Optional[Mapping[str, Any]] = None) -> "ToolContext": """A context for code that acts as ``role`` outside any dispatch (a test calling a handler directly, a seeding script, an operator helper), to bind with :func:`bound`. The role is normalised as the builders normalise it.""" return cls(caller=Caller(role=normalize_role(role), session=session, raw_role=role), tool=tool, - arguments=_frozen(arguments), call_id=uuid.uuid4().hex, transport=transport) + arguments=_frozen(arguments), call_id=uuid.uuid4().hex) #: The context handed back outside any call. @@ -205,11 +197,11 @@ def context_from_mcp(context: Any, tool: str, arguments: Mapping[str, Any]) -> T request = getattr(request_context, "request", None) headers = getattr(request, "headers", None) if request is not None else None return ToolContext(caller=caller_from_request(extra, headers), tool=tool, arguments=_frozen(arguments), - call_id=uuid.uuid4().hex, transport="mcp", mcp=context) + call_id=uuid.uuid4().hex) def context_from_http(request: Any, tool: str, arguments: Mapping[str, Any]) -> ToolContext: """The context for a REST or extension call: the caller from the request headers only.""" headers = getattr(request, "headers", None) if request is not None else None return ToolContext(caller=caller_from_request(None, headers), tool=tool, arguments=_frozen(arguments), - call_id=uuid.uuid4().hex, transport="rest") + call_id=uuid.uuid4().hex) diff --git a/packages/agentenv-protocol/tests/test_tool_context.py b/packages/agentenv-protocol/tests/test_tool_context.py index 1953dada..6670ccdf 100644 --- a/packages/agentenv-protocol/tests/test_tool_context.py +++ b/packages/agentenv-protocol/tests/test_tool_context.py @@ -83,7 +83,7 @@ def test_normalize_role_absent_blank_default_and_non_string_are_none(): def test_current_is_empty_outside_a_call_and_bound_restores(): assert ToolContext.current() is EMPTY - assert EMPTY.transport == "none" and EMPTY.caller.role is None + assert EMPTY.caller == Caller() and EMPTY.tool == "" and EMPTY.call_id == "" with bound(ToolContext.as_caller("a")) as outer: assert ToolContext.current() is outer and outer.role == "a" with bound(ToolContext.as_caller("b")) as inner: @@ -110,7 +110,7 @@ def test_caller_from_request_precedence(): def test_context_from_mcp_tolerates_every_missing_link(): empty = context_from_mcp(None, "t", {}) - assert empty.transport == "mcp" and empty.caller == Caller() and empty.mcp is None + assert empty.caller == Caller() and empty.tool == "t" and empty.call_id class _OutsideRequest: @property @@ -127,7 +127,7 @@ def request_context(self): request=SimpleNamespace(headers={"mcp-session-id": "s1"}), )) ctx = context_from_mcp(live, "search", arguments) - assert ctx.caller == Caller("bob", None, "bob") and ctx.tool == "search" and ctx.mcp is live + assert ctx.caller == Caller("bob", None, "bob") and ctx.tool == "search" assert ctx.call_id and ctx.arguments == {"q": 1} arguments["q"] = 2 assert ctx.arguments == {"q": 1} @@ -136,11 +136,11 @@ def request_context(self): def test_context_from_http_and_as_caller(): request = SimpleNamespace(headers={ROLE_HEADER: "carol"}) ctx = context_from_http(request, "echo", {"m": 1}) - assert ctx.transport == "rest" and ctx.role == "carol" and ctx.tool == "echo" and ctx.arguments == {"m": 1} + assert ctx.role == "carol" and ctx.tool == "echo" and ctx.arguments == {"m": 1} and ctx.call_id with bound(ToolContext.as_caller("alice", session="s")) as bound_ctx: assert ToolContext.current() is bound_ctx - assert bound_ctx.caller == Caller("alice", "s", "alice") and bound_ctx.transport == "none" + assert bound_ctx.caller == Caller("alice", "s", "alice") and bound_ctx.call_id assert ToolContext.current() is EMPTY with bound(ToolContext.as_caller("default")) as default_ctx: assert default_ctx.role is None and default_ctx.caller.raw_role == "default" @@ -297,7 +297,7 @@ async def test_mount_binds_the_target_dispatch(): mcp = _mounted(handler) assert mcp._tool_manager.call_tool.__agentenv_tool_context__ is True assert await mcp._tool_manager.call_tool("whoami", {"q": "x"}, context=None) == {"q": "x", "role": None, "limit": 3} - assert handler.seen[-1].transport == "mcp" and handler.seen[-1].tool == "whoami" + assert handler.seen[-1].tool == "whoami" assert ToolContext.current() is EMPTY @@ -305,11 +305,11 @@ class _ExtHandler: @extension(uri="urn:agentenv:whoami/v1") async def whoami(self, tc: ToolContext, greeting: str = "hi") -> dict: assert ToolContext.current() is tc - return {"role": tc.role, "transport": tc.transport, "tool": tc.tool, "greeting": greeting} + return {"role": tc.role, "tool": tc.tool, "greeting": greeting} @extension(uri="urn:agentenv:whoami-get/v1", method="GET") def whoami_get(self, greeting: str, tc: Optional[ToolContext] = None) -> dict: - return {"role": tc.role, "transport": tc.transport, "greeting": greeting} + return {"role": tc.role, "greeting": greeting} @pytest.mark.asyncio @@ -324,7 +324,7 @@ async def test_extension_slot_binds_rest_caller_from_headers(): r = await client.post("/agentenv/ext/whoami", json={}, headers={ROLE_HEADER: "carol"}) assert r.status_code == 200 - assert r.json() == {"role": "carol", "transport": "rest", "tool": "whoami", "greeting": "hi"} + assert r.json() == {"role": "carol", "tool": "whoami", "greeting": "hi"} r = await client.post("/agentenv/ext/whoami", json={"greeting": "yo"}) assert r.status_code == 200 and r.json()["role"] is None and r.json()["greeting"] == "yo" r = await client.post("/agentenv/ext/whoami", json={"greeting": "yo"}, headers={ROLE_HEADER: "default"}) @@ -334,7 +334,7 @@ async def test_extension_slot_binds_rest_caller_from_headers(): assert r.status_code == 400 and r.json()["error"]["code"] == "invalid_params" r = await client.get("/agentenv/ext/whoami_get", params={"greeting": "hey"}, headers={ROLE_HEADER: "dave"}) - assert r.status_code == 200 and r.json() == {"role": "dave", "transport": "rest", "greeting": "hey"} + assert r.status_code == 200 and r.json() == {"role": "dave", "greeting": "hey"} assert ToolContext.current() is EMPTY diff --git a/packages/agentenv-protocol/tests/test_tool_context_injection.py b/packages/agentenv-protocol/tests/test_tool_context_injection.py index c1e8191f..3c5aadd6 100644 --- a/packages/agentenv-protocol/tests/test_tool_context_injection.py +++ b/packages/agentenv-protocol/tests/test_tool_context_injection.py @@ -145,7 +145,7 @@ async def test_injection_sync_and_async_and_ambient_equals_injected(): assert len(env.seen) == 2 for name, (injected, ambient) in zip(("whoami", "whoami_sync"), env.seen, strict=True): assert isinstance(injected, ToolContext) and injected is ambient - assert injected.tool == name and injected.transport == "mcp" and injected.call_id + assert injected.tool == name and injected.call_id assert env.seen[0][0].arguments == {"q": "x"} assert ToolContext.current() is EMPTY @@ -395,7 +395,7 @@ def echo(self, q: str, tc: ToolContext) -> str: app = create_fastmcp_app(handler, card=EnvironmentCard(name="plain")) assert app._tool_manager.call_tool.__agentenv_tool_context__ is True assert _text(await app.call_tool("echo", {"q": "hi"})) == "hi" - assert handler.seen[0].tool == "echo" and handler.seen[0].transport == "mcp" + assert handler.seen[0].tool == "echo" @pytest.mark.asyncio @@ -418,7 +418,7 @@ async def test_direct_caller_over_streamable_http_gets_the_header_role_and_its_o sid = session_id() assert sid and result.content[0].text == "grader" and forwarded.content[0].text == "bob" direct, proxied = env.seen[-2][0], env.seen[-1][0] - assert direct.caller == Caller("grader", sid, "grader") and direct.mcp is not None + assert direct.caller == Caller("grader", sid, "grader") assert proxied.caller == Caller("bob", None, "bob") # a _meta role entry: the connection's id is not the caller's From a51cc74440258b2b3dae9f54ff98778569d808d0 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 20:52:38 -0700 Subject: [PATCH 7/7] refactor(protocol): Caller is role and session Drop raw_role and is_forwarded: nothing read either. The role stays normalised (absent, blank and the gateway's default all mean "not forwarded"); session stays, since a server keys per-session state on it when the gateway shares one connection across agents. Co-Authored-By: Claude Fable 5.1 --- .../src/agentenv_protocol/tool_context.py | 11 ++------ .../tests/test_tool_context.py | 28 +++++++++---------- .../tests/test_tool_context_injection.py | 12 ++++---- 3 files changed, 22 insertions(+), 29 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py index 6cc18ff8..a22aecc9 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py @@ -75,17 +75,11 @@ class Caller: ``role`` is the role the agent was deployed with, normalised: None when nothing was forwarded or the gateway's default role came through. - ``raw_role`` is what was received before normalisation, for logs. ``session`` tells callers apart when the gateway shares one connection. """ role: Optional[str] = None session: Optional[str] = None - raw_role: Optional[str] = None - - @property - def is_forwarded(self) -> bool: - return self.role is not None @dataclasses.dataclass(frozen=True) @@ -135,7 +129,7 @@ def as_caller(cls, role: Optional[str] = None, session: Optional[str] = None, *, """A context for code that acts as ``role`` outside any dispatch (a test calling a handler directly, a seeding script, an operator helper), to bind with :func:`bound`. The role is normalised as the builders normalise it.""" - return cls(caller=Caller(role=normalize_role(role), session=session, raw_role=role), tool=tool, + return cls(caller=Caller(role=normalize_role(role), session=session), tool=tool, arguments=_frozen(arguments), call_id=uuid.uuid4().hex) @@ -178,8 +172,7 @@ def caller_from_request(meta_extra: Optional[Mapping[str, Any]], headers: Any) - session = meta.get(SESSION_META_KEY) if not isinstance(session, str) or not session: session = None if forwarded or headers is None else headers.get("mcp-session-id") - raw_role = raw_role if isinstance(raw_role, str) else None - return Caller(role=normalize_role(raw_role), session=session or None, raw_role=raw_role) + return Caller(role=normalize_role(raw_role), session=session or None) def context_from_mcp(context: Any, tool: str, arguments: Mapping[str, Any]) -> ToolContext: diff --git a/packages/agentenv-protocol/tests/test_tool_context.py b/packages/agentenv-protocol/tests/test_tool_context.py index 6670ccdf..5f403aa1 100644 --- a/packages/agentenv-protocol/tests/test_tool_context.py +++ b/packages/agentenv-protocol/tests/test_tool_context.py @@ -99,12 +99,12 @@ def test_current_is_empty_outside_a_call_and_bound_restores(): def test_caller_from_request_precedence(): headers = {ROLE_HEADER: "hdr", "mcp-session-id": "sid"} assert caller_from_request({ROLE_META_KEY: "meta"}, headers).role == "meta" - assert caller_from_request({}, headers) == Caller("hdr", "sid", "hdr") - assert caller_from_request(None, headers) == Caller("hdr", "sid", "hdr") + assert caller_from_request({}, headers) == Caller("hdr", "sid") + assert caller_from_request(None, headers) == Caller("hdr", "sid") assert caller_from_request({SESSION_META_KEY: "s-meta"}, headers).session == "s-meta" assert caller_from_request({SESSION_META_KEY: ""}, headers).session == "sid" assert caller_from_request({SESSION_META_KEY: 7}, headers).session == "sid" - assert caller_from_request({ROLE_META_KEY: "default"}, None) == Caller(None, None, "default") + assert caller_from_request({ROLE_META_KEY: "default"}, None) == Caller() assert caller_from_request(None, None) == Caller() @@ -127,7 +127,7 @@ def request_context(self): request=SimpleNamespace(headers={"mcp-session-id": "s1"}), )) ctx = context_from_mcp(live, "search", arguments) - assert ctx.caller == Caller("bob", None, "bob") and ctx.tool == "search" + assert ctx.caller == Caller("bob") and ctx.tool == "search" assert ctx.call_id and ctx.arguments == {"q": 1} arguments["q"] = 2 assert ctx.arguments == {"q": 1} @@ -140,10 +140,10 @@ def test_context_from_http_and_as_caller(): with bound(ToolContext.as_caller("alice", session="s")) as bound_ctx: assert ToolContext.current() is bound_ctx - assert bound_ctx.caller == Caller("alice", "s", "alice") and bound_ctx.call_id + assert bound_ctx.caller == Caller("alice", "s") and bound_ctx.call_id assert ToolContext.current() is EMPTY with bound(ToolContext.as_caller("default")) as default_ctx: - assert default_ctx.role is None and default_ctx.caller.raw_role == "default" + assert default_ctx.role is None and default_ctx.caller == Caller() def test_tool_context_slot_detection_by_annotation(): @@ -340,15 +340,15 @@ async def test_extension_slot_binds_rest_caller_from_headers(): def test_caller_from_request_consults_headers_only_when_no_meta_role_entry_exists(): headers = {ROLE_HEADER: "hdr", "mcp-session-id": "conn"} - assert caller_from_request(None, headers) == Caller("hdr", "conn", "hdr") - assert caller_from_request({}, headers) == Caller("hdr", "conn", "hdr") - assert caller_from_request({SESSION_META_KEY: "s-only"}, headers) == Caller("hdr", "s-only", "hdr") + assert caller_from_request(None, headers) == Caller("hdr", "conn") + assert caller_from_request({}, headers) == Caller("hdr", "conn") + assert caller_from_request({SESSION_META_KEY: "s-only"}, headers) == Caller("hdr", "s-only") # a _meta role entry marks a proxied call: blank and default stay unforwarded, the connection id is not the caller's - assert caller_from_request({ROLE_META_KEY: ""}, headers) == Caller(None, None, "") - assert caller_from_request({ROLE_META_KEY: "default"}, headers) == Caller(None, None, "default") - assert caller_from_request({ROLE_META_KEY: "bob"}, headers) == Caller("bob", None, "bob") - assert caller_from_request({ROLE_META_KEY: "bob", SESSION_META_KEY: "s1"}, headers) == Caller("bob", "s1", "bob") - assert caller_from_request({ROLE_META_KEY: 7, SESSION_META_KEY: ""}, headers) == Caller(None, None, None) + assert caller_from_request({ROLE_META_KEY: ""}, headers) == Caller() + assert caller_from_request({ROLE_META_KEY: "default"}, headers) == Caller() + assert caller_from_request({ROLE_META_KEY: "bob"}, headers) == Caller("bob") + assert caller_from_request({ROLE_META_KEY: "bob", SESSION_META_KEY: "s1"}, headers) == Caller("bob", "s1") + assert caller_from_request({ROLE_META_KEY: 7, SESSION_META_KEY: ""}, headers) == Caller() def test_tool_context_is_hashable_and_its_arguments_are_read_only(): diff --git a/packages/agentenv-protocol/tests/test_tool_context_injection.py b/packages/agentenv-protocol/tests/test_tool_context_injection.py index 3c5aadd6..8ef927c4 100644 --- a/packages/agentenv-protocol/tests/test_tool_context_injection.py +++ b/packages/agentenv-protocol/tests/test_tool_context_injection.py @@ -182,17 +182,17 @@ async def test_header_fallback_and_default_blank_roles(): headers = {ROLE_HEADER: "hdr", "mcp-session-id": "sid"} await call("whoami_sync", {}, context=_request_context(headers=headers)) - assert env.seen[-1][0].caller == Caller("hdr", "sid", "hdr") + assert env.seen[-1][0].caller == Caller("hdr", "sid") await call("whoami_sync", {}, context=_request_context({ROLE_META_KEY: "meta", SESSION_META_KEY: "s-meta"}, headers)) - assert env.seen[-1][0].caller == Caller("meta", "s-meta", "meta") + assert env.seen[-1][0].caller == Caller("meta", "s-meta") await call("whoami_sync", {}, context=_request_context({ROLE_META_KEY: "default"}, headers)) - assert env.seen[-1][0].caller == Caller(None, None, "default") + assert env.seen[-1][0].caller == Caller() await call("whoami_sync", {}, context=_request_context({ROLE_META_KEY: " "}, None)) assert env.seen[-1][0].caller.role is None async with create_connected_server_and_client_session(app) as session: result = await session.call_tool("whoami_sync", {}, meta={ROLE_META_KEY: "default"}) - assert result.content[0].text == "-" and env.seen[-1][0].caller.raw_role == "default" + assert result.content[0].text == "-" and env.seen[-1][0].caller == Caller() def _refusal() -> CallToolResult: @@ -418,8 +418,8 @@ async def test_direct_caller_over_streamable_http_gets_the_header_role_and_its_o sid = session_id() assert sid and result.content[0].text == "grader" and forwarded.content[0].text == "bob" direct, proxied = env.seen[-2][0], env.seen[-1][0] - assert direct.caller == Caller("grader", sid, "grader") - assert proxied.caller == Caller("bob", None, "bob") # a _meta role entry: the connection's id is not the caller's + assert direct.caller == Caller("grader", sid) + assert proxied.caller == Caller("bob") # a _meta role entry: the connection's id is not the caller's def _app_shape(app: FastMCP) -> tuple: