fix(mcp): scope input schemas to tool calls

Author: gesh

posthog/mcp/_instrument_fastmcp.py

@@ -22,6 +22,7 @@
 
 import inspect
 import time
+from dataclasses import replace
 from typing import Any, Dict, Optional, Tuple
 
 import mcp.types as mcp_types
@@ -131,6 +132,8 @@ async def wrapped(
                 for text in lifecycle.virtual_result_texts(reply)
             ]
 
+        lifecycle = replace(lifecycle, input_schema=_tool_input_schema(server, name))
+
         # Strip each injected key independently. A tool can declare its own
         # `context` (kept) while `conversation_id` is still SDK-injected (stripped),
         # so coupling both to context-ownership leaked conversation_id into the tool.
@@ -272,7 +275,7 @@ async def list_handler(req: Any) -> Any:
         duration_ms = (time.monotonic() - start) * 1000
         tools = extract_tools(result)
         # Empty is computed before adding the virtual missing-capability tool.
-        names, empty = collect_listed_tools(data, tools, lifecycle.session_id)
+        names, empty = collect_listed_tools(data, tools)
         injection = resolve_virtual_tool_injection(
             data,
             tools,
@@ -335,6 +338,15 @@ def _name_owned_by_real_tool(server: Any, name: str) -> Optional[bool]:
         return None
 
 
+def _tool_input_schema(server: Any, name: str) -> Optional[Dict[str, Any]]:
+    """Return the schema from the tool that this server will call."""
+    try:
+        schema = server._tool_manager.get_tool(name).parameters
+    except Exception:  # noqa: BLE001 - analytics must not break the call
+        return None
+    return schema if isinstance(schema, dict) else None
+
+
 def _tool_owns_param(server: Any, name: str, param: str) -> bool:
     """True when the tool's own function declares ``param`` — then it's a real tool
     argument we must neither inject nor strip (the agent's value belongs to the tool)."""

posthog/mcp/_instrument_lowlevel.py

@@ -186,10 +186,10 @@ def _wrap_call_tool(
     async def handler(req: Any) -> Any:
         name = req.params.name
         arguments = dict(req.params.arguments or {})
-        strip, model_ours = (
+        strip, model_ours, input_schema = (
             await _standalone_ownership(data, high_level, name, req.params.meta)
             if strip_injected
-            else (set(), data.tool_model_parameter_injected.get(name))
+            else (set(), data.tool_model_parameter_injected.get(name), None)
         )
         client_name, client_version = _client_info(server)
         protocol_version = _protocol_version(server)
@@ -212,6 +212,7 @@ async def handler(req: Any) -> Any:
             client_version=client_version,
             protocol_version=protocol_version,
             extra={"session_id": mcp_session_id, "ctx": _request_context(server)},
+            input_schema=input_schema,
         )
 
         if lifecycle.is_missing_capability and (
@@ -472,7 +473,7 @@ async def handler(req: Any) -> Any:
 
         # Zero advertised tools is treated as an errored tools/list before the
         # virtual missing-capability tool is appended.
-        names, empty = collect_listed_tools(data, tools, lifecycle.session_id)
+        names, empty = collect_listed_tools(data, tools)
         injection = resolve_virtual_tool_injection(
             data,
             tools,
@@ -565,10 +566,11 @@ def _tool_lookup_not_found_errors() -> Tuple[type, ...]:
 
 async def _standalone_ownership(
     data: MCPAnalyticsData, high_level: Any, name: str, meta: Any
-) -> Tuple[set, Optional[bool]]:
+) -> Tuple[set, Optional[bool], Optional[Dict[str, Any]]]:
     """Ownership of the injected arguments on jlowin's standalone FastMCP: the
     keys to strip before it validates the call, and whether ``llm_model`` is
-    ours (``None`` when nothing can say).
+    ours (``None`` when nothing can say). The third item is the trusted schema
+    for input-name analytics, or ``None`` when middleware can change dispatch.
 
     Only keys injected under the current options are candidates. ``context``
     and ``conversation_id`` are stripped unless the registered schema (or,
@@ -580,17 +582,41 @@ async def _standalone_ownership(
     stays and is still read (posthog-js ADR-0011).
     """
     try:
-        declared, model_injectable = await _registry_view(high_level, name, meta)
+        declared, model_injectable, input_schema = await _registry_view(
+            high_level, name, meta
+        )
         model_ours = data.tool_model_parameter_injected.get(name, model_injectable)
         if _dispatch_can_differ(high_level):
             model_ours = False
+            input_schema = None
     except Exception:  # noqa: BLE001 - ownership inference must never prevent dispatch
-        declared, model_ours = None, None
+        declared, model_ours, input_schema = None, None, None
     candidates = _injected_keys(data)
     strip = {k for k in candidates - {"llm_model"} if k not in (declared or set())}
     if "llm_model" in candidates and model_ours:
         strip.add("llm_model")
-    return strip, model_ours
+    return strip, model_ours, input_schema
+
+
+def _tool_schema_view(
+    high_level: Any, tool: Any
+) -> Tuple[Optional[set], Optional[bool], Optional[Dict[str, Any]]]:
+    if tool is None:
+        return None, None, None
+    schema = getattr(tool, "parameters", None)
+    if isinstance(schema, dict):
+        declared, injectable = _schema_view(
+            schema, dereferenced=_server_dereferences(high_level)
+        )
+        return declared, injectable, schema
+    fn = getattr(tool, "fn", None)
+    if fn is None:
+        return set(), True, None
+    try:
+        declared = {k for k in _INJECTED_KEYS if k in inspect.signature(fn).parameters}
+    except Exception:  # noqa: BLE001 - introspection is best-effort
+        return set(), True, None
+    return declared, "llm_model" not in declared, None
 
 
 def _injected_keys(data: MCPAnalyticsData) -> set:
@@ -609,35 +635,24 @@ def _injected_keys(data: MCPAnalyticsData) -> set:
 
 async def _registry_view(
     high_level: Any, name: str, meta: Any
-) -> Tuple[Optional[set], Optional[bool]]:
+) -> Tuple[Optional[set], Optional[bool], Optional[Dict[str, Any]]]:
     """What the registered tool says about the injected keys: which of
     ``_INJECTED_KEYS`` it declares itself, and whether a listing would have
     injected ``llm_model`` into its schema (the same test the listing applies,
     on the schema as the client would see it). Read from the schema (a ``Tool``
     subclass may have no function) else the signature. The registry is read
     directly, never through middleware, so a cold instance answers without a
     listing and rate limiters are not charged. ``(None, None)`` when the
-    registry has no such tool or cannot be read."""
+    registry has no such tool or cannot be read. The third item is the tool's
+    input schema when the registry supplies one."""
     try:
         tool = await _registered_tool(high_level, name, meta)
     except Exception as error:  # noqa: BLE001 - introspection is best-effort
         if not isinstance(error, _tool_lookup_not_found_errors()):
             warn_ownership_lookup_failed(name, error)
-            return set(_INJECTED_KEYS), None
-        return None, None
-    if tool is None:
-        return None, None
-    schema = getattr(tool, "parameters", None)
-    if isinstance(schema, dict):
-        return _schema_view(schema, dereferenced=_server_dereferences(high_level))
-    fn = getattr(tool, "fn", None)
-    if fn is None:
-        return set(), True
-    try:
-        declared = {k for k in _INJECTED_KEYS if k in inspect.signature(fn).parameters}
-    except Exception:  # noqa: BLE001 - introspection is best-effort
-        return set(), True
-    return declared, "llm_model" not in declared
+            return set(_INJECTED_KEYS), None, None
+        return None, None, None
+    return _tool_schema_view(high_level, tool)
 
 
 async def _registered_tool(high_level: Any, name: str, meta: Any) -> Any:

posthog/mcp/_instrument_v2.py

@@ -32,6 +32,7 @@
 
 import time
 from collections.abc import Mapping
+from dataclasses import replace
 from typing import Any, Dict, FrozenSet, Optional, Set, Tuple
 
 import mcp.types as mcp_types
@@ -230,8 +231,7 @@ def _tool_own_properties_v2(high_level: Any, name: str) -> Dict[str, Any]:
     site so checking ownership of both ``context`` and ``conversation_id``
     doesn't look the tool up from the manager twice."""
     try:
-        tool = high_level._tool_manager.get_tool(name)
-        properties = (getattr(tool, "parameters", None) or {}).get("properties")
+        properties = (_tool_input_schema_v2(high_level, name) or {}).get("properties")
     except Exception:  # noqa: BLE001
         return {}
     # Fail closed on a malformed schema: the caller does `param in <this>` in the
@@ -240,6 +240,14 @@ def _tool_own_properties_v2(high_level: Any, name: str) -> Dict[str, Any]:
     return properties if isinstance(properties, dict) else {}
 
 
+def _tool_input_schema_v2(high_level: Any, name: str) -> Optional[Dict[str, Any]]:
+    try:
+        schema = high_level._tool_manager.get_tool(name).parameters
+    except Exception:  # noqa: BLE001 - analytics must not break the call
+        return None
+    return schema if isinstance(schema, dict) else None
+
+
 def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool:
     """Whether the tool's own JSON schema declares ``param`` — then it's a real
     tool argument we must neither inject over nor strip. Read from the tool's
@@ -318,6 +326,8 @@ async def wrapped(
             ]
             return mcp_types.CallToolResult(content=virtual_content)
 
+        lifecycle = replace(lifecycle, input_schema=_tool_input_schema_v2(server, name))
+
         # v2 validates against the function signature and rejects unexpected
         # keys, so injected parameters are stripped before dispatch — but never
         # one the tool's own schema declares (that's a real argument).
@@ -441,7 +451,7 @@ def _requested_tool_version(ctx: Any) -> Optional[str]:
 
 async def _standalone_injected_parameters(
     server: Any, data: MCPAnalyticsData, name: str, version: Optional[str]
-) -> Optional[FrozenSet[str]]:
+) -> Tuple[Optional[FrozenSet[str]], Optional[Dict[str, Any]]]:
     """Resolve ownership in the current request, including middleware and versions.
 
     Listings from other requests can have different application-owned parameters.
@@ -464,9 +474,9 @@ async def _standalone_injected_parameters(
         schema = getattr(tool, "parameters", None)
     except Exception as error:  # noqa: BLE001 - schema lookup must not prevent dispatch
         log(f"PostHog MCP: could not resolve schema for tool {name!r} - {error}")
-        return None
+        return None, None
     if not isinstance(schema, dict):
-        return None
+        return None, None
     injected = set()
     if is_context_enabled(data.options.context):
         injected.add("context")
@@ -476,7 +486,10 @@ async def _standalone_injected_parameters(
         can_inject_model_parameter(schema)
     ):
         injected.add("llm_model")
-    return frozenset(key for key in injected if not schema_has_param(schema, key))
+    return (
+        frozenset(key for key in injected if not schema_has_param(schema, key)),
+        schema,
+    )
 
 
 def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None:
@@ -492,10 +505,11 @@ async def handler(ctx: Any, params: Any) -> Any:
         # reads the self-reported model anyway; only a listing that proved the
         # application owns `llm_model` stops it (posthog-js ADR-0011).
         analytics_owns_model = data.tool_model_parameter_injected.get(name) is not False
+        input_schema = None
         standalone = data.standalone_fastmcp() if data.standalone_fastmcp else None
         if standalone is not None:
             version = _requested_tool_version(ctx)
-            injected = await _standalone_injected_parameters(
+            injected, input_schema = await _standalone_injected_parameters(
                 standalone, data, name, version
             )
             if injected is not None:
@@ -521,6 +535,7 @@ async def handler(ctx: Any, params: Any) -> Any:
             client_version=client_version,
             protocol_version=protocol_version,
             extra={"session_id": mcp_session_id, "ctx": ctx},
+            input_schema=input_schema,
         )
 
         # No tool registry on a raw low-level server, so ownership is settled
@@ -695,7 +710,7 @@ async def handler(ctx: Any, params: Any) -> Any:
 
         tools = list(getattr(result, "tools", []) or [])
         # Empty is computed before adding the virtual missing-capability tool.
-        names, empty = collect_listed_tools(data, tools, lifecycle.session_id)
+        names, empty = collect_listed_tools(data, tools)
         injection = resolve_virtual_tool_injection(
             data, tools, is_first_page=is_first_listing_page(params)
         )

posthog/mcp/_instrumentation.py

@@ -448,6 +448,7 @@ class ToolCallLifecycle:
     feedback_name: Optional[str]
     conversation_id: Optional[str]
     minted_conversation_id: bool
+    input_schema: Any
 
     @property
     def is_missing_capability(self) -> bool:
@@ -547,6 +548,7 @@ async def record_error(self, error: Any, duration_ms: float) -> None:
             protocol_version=self.protocol_version,
             conversation_id=conversation_id,
             extra=self.extra,
+            input_schema=self.input_schema,
         )
 
     async def record_result(
@@ -568,6 +570,7 @@ async def record_result(
             protocol_version=self.protocol_version,
             conversation_id=conversation_id,
             extra=self.extra,
+            input_schema=self.input_schema,
         )
 
 
@@ -584,6 +587,7 @@ def start_tool_call_lifecycle(
     client_version: Optional[str],
     protocol_version: Optional[str],
     extra: Dict[str, Any],
+    input_schema: Any = None,
 ) -> ToolCallLifecycle:
     """Resolve adapter-independent policy for a tool call without dispatching it."""
     enabled = enabled_virtual_tool_names(data)
@@ -617,6 +621,7 @@ def start_tool_call_lifecycle(
         feedback_name=feedback_name,
         conversation_id=conversation_id,
         minted_conversation_id=minted,
+        input_schema=input_schema,
     )
 
 
@@ -636,6 +641,7 @@ async def record_tool_call(
     protocol_version: Optional[str] = None,
     conversation_id: Optional[str] = None,
     extra: Optional[Dict[str, Any]] = None,
+    input_schema: Any = None,
 ) -> None:
     # Analytics must never change what the tool returns or raises: any failure
     # building/publishing the event is logged and swallowed here.
@@ -684,12 +690,11 @@ async def record_tool_call(
                 input_aliases = data.options.resolve_input_aliases(name)
             except Exception as err:  # noqa: BLE001 - analytics callbacks are isolated
                 log(f"Warning: resolve_input_aliases failed for tool {name}: {err}")
-        schema = data.tool_input_schemas.get(session_id, {}).get(name)
         event["properties"] = {
             **(props or {}),
             **get_tool_input_properties(
                 arguments or {},
-                schema,
+                input_schema,
                 ToolInputOptions(
                     should_record_input_key=data.options.should_record_input_key,
                     input_aliases=input_aliases,
@@ -1008,29 +1013,16 @@ def read_tool_category(tool: Any) -> Optional[str]:
     return None
 
 
-def collect_listed_tools(
-    data: MCPAnalyticsData, tools: list, session_id: Optional[str] = None
-) -> tuple[List[str], bool]:
+def collect_listed_tools(data: MCPAnalyticsData, tools: list) -> tuple[List[str], bool]:
     """Cache common tool metadata and return the pre-injection listing summary."""
     names = []
-    schemas = data.tool_input_schemas.get(session_id, {}) if session_id else None
     for tool in tools:
         names.append(tool.name)
         if getattr(tool, "description", None):
             data.tool_descriptions[tool.name] = tool.description
         category = read_tool_category(tool)
         if category:
             data.tool_categories[tool.name] = category
-        if schemas is not None:
-            schema = getattr(tool, "inputSchema", None)
-            if schema is None:
-                schema = getattr(tool, "input_schema", None)
-            schemas[tool.name] = schema
-    if session_id and schemas is not None:
-        data.tool_input_schemas[session_id] = schemas
-        data.tool_input_schemas.move_to_end(session_id)
-        while len(data.tool_input_schemas) > 1000:
-            data.tool_input_schemas.popitem(last=False)
     return names, not tools
 
 

posthog/mcp/_internal.py

@@ -88,11 +88,6 @@ class MCPAnalyticsData:
     identified_sessions: IdentityCache = field(default_factory=IdentityCache)
     tool_categories: Dict[str, str] = field(default_factory=dict)
     tool_descriptions: Dict[str, str] = field(default_factory=dict)
-    # Original tool input schemas by PostHog session. This keeps field-name
-    # privacy decisions isolated when a server advertises user-specific tools.
-    tool_input_schemas: "OrderedDict[str, Dict[str, Any]]" = field(
-        default_factory=OrderedDict
-    )
     # True only when PostHog added llm_model to this tool's advertised schema.
     # Missing/False fails closed so an application-owned field is never read or stripped.
     tool_model_parameter_injected: Dict[str, bool] = field(default_factory=dict)