feat(protocol): ToolContext, an opt-in injected caller for environment tools

packages/agentenv-protocol/README.md

@@ -83,6 +83,42 @@ 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). 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`
+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
+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
+tool's reply. An `@extension` method may declare the same parameter; it is filled from the request's
+`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
 
 `serve()` builds the FastMCP app via `create_fastmcp_app()`, which encodes the agent-env deploy

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

@@ -12,9 +12,20 @@
     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,
+    bound,
+    normalize_role,
+)
 from .manifest import (
     CLI_COMMANDS,
     GET_INTERFACES_EXTENSION_URI,
@@ -78,6 +89,15 @@
     "get_data",
     "extension",
     "tool",
+    "injecting",
+    "ToolContext",
+    "Caller",
+    "normalize_role",
+    "bound",
+    "ROLE_META_KEY",
+    "SESSION_META_KEY",
+    "ROLE_HEADER",
+    "DEFAULT_ROLE",
     "EnvironmentCapabilities",
     "EnvironmentCard",
     "EnvironmentExtension",

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

@@ -0,0 +1,200 @@
+"""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, Mapping, Optional
+
+from pydantic_core import core_schema
+
+__all__ = [
+    "DEFAULT_ROLE",
+    "EMPTY",
+    "ROLE_HEADER",
+    "ROLE_META_KEY",
+    "SESSION_META_KEY",
+    "Caller",
+    "ToolContext",
+    "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"
+
+
+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.
+    ``session`` tells callers apart when the gateway shares one connection.
+    """
+
+    role: Optional[str] = None
+    session: Optional[str] = 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; ``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.
+    """
+
+    caller: Caller = Caller()
+    tool: str = ""
+    arguments: Mapping[str, Any] = dataclasses.field(default_factory=lambda: types.MappingProxyType({}), hash=False)
+    call_id: str = ""
+
+    @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 as_caller(cls, role: Optional[str] = None, session: Optional[str] = None, *, tool: str = "",
+                  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), tool=tool,
+                   arguments=_frozen(arguments), call_id=uuid.uuid4().hex)
+
+
+#: 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")
+    return Caller(role=normalize_role(raw_role), session=session or None)
+
+
+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)
+
+
+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)

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):

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,14 +52,26 @@ 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):
             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
+
 
 class _InlineHandler:
     def __init__(self) -> None: