From d5bd5a82b4e3373048d443187f496f409904ff6c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 23:07:02 -0700 Subject: [PATCH 1/2] test(gateway): a real gateway in front of a real protocol child, over loopback Two loopback servers: a protocol environment and a Gateway with it as its only backing server. Two MCP clients with different AgentEnv-Role headers reach the child's tools as two callers, each with the gateway's per-session key rather than its connection id, stable across calls on one session; a tool that never asked for the caller keeps its schema and answers as before; a refusal returned by the child's on_tool_call reaches the client through the gateway as isError with its payload intact; a /step call carries the role and no session; a client on the child directly gets its header role and its own connection id. Co-Authored-By: Claude Fable 5.1 --- .../env/gateway/role_forwarding_local_test.py | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 tst/integration/env/gateway/role_forwarding_local_test.py diff --git a/tst/integration/env/gateway/role_forwarding_local_test.py b/tst/integration/env/gateway/role_forwarding_local_test.py new file mode 100644 index 00000000..a71cad3b --- /dev/null +++ b/tst/integration/env/gateway/role_forwarding_local_test.py @@ -0,0 +1,139 @@ +"""A real gateway in front of a real protocol child, over loopback: the role an MCP client connects with reaches the +child's tools as the caller the protocol binds; a tool that never asked for the caller is served exactly as before; and +a refusal the child's ``on_tool_call`` returns reaches the client through the gateway untouched.""" +from __future__ import annotations + +import json +import socket +from contextlib import asynccontextmanager + +import httpx +import pytest +import uvicorn +from mcp import ClientSession +from mcp.client.streamable_http import streamable_http_client +from mcp.types import CallToolResult, TextContent + +from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, InternalMCPServer +from agent_env.env.gateway.gateway import Gateway +from agentenv_protocol import ( + ROLE_HEADER, + ROLE_META_KEY, + SESSION_META_KEY, + AgentEnvEnvironment, + ToolContext, + environment_card, + tool, +) + +REFUSED = "nobody@example.com" + + +@environment_card(name="items") +class _Items(AgentEnvEnvironment): + def __init__(self) -> None: + self.items: list[dict] = [] + self.seen: list[ToolContext] = [] + + def on_tool_call(self, context: ToolContext) -> CallToolResult | None: + if context.caller.role == REFUSED: + payload = {"error": "unknown_identity", "role": context.caller.role, "tool": context.tool} + return CallToolResult(content=[TextContent(type="text", text=json.dumps(payload))], isError=True) + return None + + @tool(name="{environment_name}_add") + def add(self, item: str, times: int = 1) -> dict: + """Add an item.""" + self.items.append({"item": item, "times": times, "by": ToolContext.current().caller.role}) + return {"count": len(self.items)} + + @tool(name="{environment_name}_whoami") + def whoami(self, ctx: ToolContext) -> dict: + self.seen.append(ctx) + return {"role": ctx.caller.role, "session": ctx.caller.session} + + +async def _serve(app) -> tuple[uvicorn.Server, str]: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + config = uvicorn.Config(app, log_level="critical", access_log=False, timeout_graceful_shutdown=3) + config.load() + server = uvicorn.Server(config) + server.lifespan = config.lifespan_class(config) + await server.startup(sockets=[sock]) + return server, f"http://127.0.0.1:{sock.getsockname()[1]}" + + +@asynccontextmanager +async def _stack(): + """The child on one loopback server, the gateway on another with the child as its only backing server.""" + env = _Items() + child_server, child_url = await _serve(env.create_app().streamable_http_app()) + gw = Gateway(host="127.0.0.1", port=0, server_name="t", + internal_mcp_servers=[InternalMCPServer("items", f"{child_url}/mcp")]) + gw_server, gw_url = await _serve(gw._asgi_app()) + try: + yield env, gw_url, child_url + finally: + await gw_server.shutdown() + await child_server.shutdown() + + +@asynccontextmanager +async def _client(url: str, role: str | None = None): + headers = {AGENT_ENV_ROLE_HEADER: role} if role else {} + async with ( + httpx.AsyncClient(headers=headers, timeout=httpx.Timeout(10, read=60)) as http, + streamable_http_client(f"{url}/mcp", http_client=http) as (read, write, session_id), + ClientSession(read, write) as session, + ): + await session.initialize() + yield session, session_id + + +def _body(result: CallToolResult) -> dict: + return json.loads(result.content[0].text) + + +def test_gateway_and_protocol_agree_on_the_wire_names(): + assert (AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, AGENT_ENV_ROLE_HEADER) == (ROLE_META_KEY, SESSION_META_KEY, ROLE_HEADER) + + +@pytest.mark.asyncio +async def test_roles_reach_the_child_through_the_gateway_and_a_slotless_tool_is_served_as_before(): + async with _stack() as (env, gw_url, child_url): + async with _client(gw_url, "alice@example.com") as (alice, _), _client(gw_url, "bob@example.com") as (bob, _): + listed = {t.name: t for t in (await alice.list_tools()).tools} + assert set(listed["items_add"].inputSchema["properties"]) == {"item", "times"} + assert listed["items_whoami"].inputSchema.get("properties", {}) == {} # the injected parameter never reaches the wire + + a, b = _body(await alice.call_tool("items_whoami", {})), _body(await bob.call_tool("items_whoami", {})) + assert (a["role"], b["role"]) == ("alice@example.com", "bob@example.com") + assert a["session"] != b["session"] and len(a["session"]) == 32 # the gateway's per-session key, not its connection + assert a == _body(await alice.call_tool("items_whoami", {})) # stable across calls on one session + + added = await alice.call_tool("items_add", {"item": "pen", "times": 2}) + assert added.isError is False and _body(added) == {"count": 1} + assert env.items == [{"item": "pen", "times": 2, "by": "alice@example.com"}] + + async with _client(gw_url) as (anon, _): + assert _body(await anon.call_tool("items_whoami", {}))["role"] is None + assert _body(await anon.call_tool("items_add", {"item": "cup"})) == {"count": 2} + assert env.items[-1]["by"] is None + + async with _client(gw_url, REFUSED) as (nobody, _): + refused = await nobody.call_tool("items_add", {"item": "x"}) + assert refused.isError is True + assert _body(refused) == {"error": "unknown_identity", "role": REFUSED, "tool": "items_add"} + assert len(env.items) == 2 + + async with httpx.AsyncClient() as http: + step = await http.post(f"{gw_url}/step", headers={AGENT_ENV_ROLE_HEADER: "grader"}, + json={"action": "call_tool", "tool_name": "items_whoami", "arguments": {}}) + assert step.status_code == 200, step.text + assert (env.seen[-1].caller.role, env.seen[-1].caller.session) == ("grader", None) + + async with _client(child_url, "dave@example.com") as (direct, session_id): + d = _body(await direct.call_tool("items_whoami", {})) + assert d["role"] == "dave@example.com" and d["session"] == session_id() # no gateway: the header and its own connection From 3665debafcf7d879566cc6430f6518bb84ac1e81 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 10 Oct 2026 23:07:02 -0700 Subject: [PATCH 2/2] docs(protocol): say what the dispatch marker is for and when the _meta session is read The marker on the installed dispatch is read by _check_tool_binding to refuse a second handler on one app; no downstream guard reads it. caller_from_request takes the _meta session whenever one is present, with or without a _meta role entry, and falls back to the connection id only for a direct caller. Co-Authored-By: Claude Fable 5.1 --- .../src/agentenv_protocol/agent_env_environment.py | 4 ++-- .../src/agentenv_protocol/tool_context.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 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 a0853a3d..702eb434 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -248,8 +248,8 @@ def _bind_tool_calls(self, app: Any) -> None: 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. + The installed callable carries ``__agentenv_tool_context__ = True``, the marker + ``_check_tool_binding`` reads to refuse a second handler on the same app. """ manager = app._tool_manager inner = manager.call_tool diff --git a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py index a22aecc9..f943ac18 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/tool_context.py @@ -158,10 +158,11 @@ def caller_from_request(meta_extra: Optional[Mapping[str, Any]], headers: Any) - 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. + and the role header is not consulted, so a blank or ``default`` role stays "not + forwarded". Without that entry the caller reached the server itself and the role header + is trusted exactly as the caller's other requests are. The session is the ``_meta`` + session whenever one is present; otherwise it is the connection's ``mcp-session-id`` + for a direct caller and none for a proxied one, whose connection is the gateway's. """ meta = meta_extra or {} forwarded = ROLE_META_KEY in meta