feat(mcp): resolve fresh low-level tool ownership (#1009)

* feat(mcp): resolve fresh low-level tool ownership

Add resolve_original_tool for raw MCP SDK 1.x and 2.x servers. Prefer ownership learned from a served listing and keep unknown ownership fail-open.

Generated-By: PostHog Desktop
Task-Id: db1a8a08-fe03-43b6-9e49-705e8b1f375e

* fix(mcp): preserve listed argument ownership

Generated-By: PostHog Desktop
Task-Id: db1a8a08-fe03-43b6-9e49-705e8b1f375e

* fix(mcp): preserve arguments for unknown schemas

Generated-By: PostHog Desktop
Task-Id: db1a8a08-fe03-43b6-9e49-705e8b1f375e

* fix(mcp): preserve low-level tool schemas

Generated-By: PostHog Desktop
Task-Id: db1a8a08-fe03-43b6-9e49-705e8b1f375e
(cherry picked from commit 127009d)

Author: gesh

.sampo/changesets/mcp-lowlevel-tool-ownership.md

@@ -0,0 +1,7 @@
+---
+pypi/posthog: minor
+---
+
+Add `resolve_original_tool` for low-level MCP servers. Fresh server instances can now remove PostHog-owned arguments before strict tool validation.
+
+Raw low-level servers also remove PostHog-owned arguments after `tools/list`. A tool-owned `context` remains tool data.

posthog/mcp/_argument_ownership.py

@@ -0,0 +1,98 @@
+"""Resolve ownership of arguments that PostHog adds to MCP tool schemas."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Mapping
+from typing import Any, Dict, FrozenSet, Optional, Tuple
+
+from ._context_parameters import is_context_enabled, schema_has_param
+from ._model_parameters import is_capture_model_enabled
+from .logger import log
+
+_COMPLEX_SCHEMA_KEYS = ("$ref", "oneOf", "allOf", "anyOf")
+
+
+def analytics_owned_parameters(
+    options: Any,
+    input_schema: Any,
+) -> FrozenSet[str]:
+    """Return the enabled arguments that PostHog can add to this schema."""
+    enabled = set()
+    if is_context_enabled(options.context):
+        enabled.add("context")
+    if options.enable_conversation_id:
+        enabled.add("conversation_id")
+    if is_capture_model_enabled(options.capture_model):
+        enabled.add("llm_model")
+
+    if isinstance(input_schema, dict) and any(
+        input_schema.get(key) for key in _COMPLEX_SCHEMA_KEYS
+    ):
+        return frozenset()
+    return frozenset(
+        name for name in enabled if not schema_has_param(input_schema, name)
+    )
+
+
+def cache_listed_tool_ownership(
+    data: Any, tool: Any, *, schema_attribute: str
+) -> FrozenSet[str]:
+    """Cache ownership from the host schema before PostHog changes it."""
+    name = getattr(tool, "name", None)
+    if not isinstance(name, str):
+        return frozenset()
+    schema = getattr(tool, schema_attribute, None)
+    ownership = analytics_owned_parameters(data.options, schema)
+    data.tool_analytics_parameter_ownership[name] = ownership
+    if isinstance(schema, dict):
+        data.tool_input_schemas[name] = schema
+    else:
+        data.tool_input_schemas.pop(name, None)
+    return ownership
+
+
+async def resolve_lowlevel_tool_ownership(
+    data: Any, name: str
+) -> Tuple[Optional[FrozenSet[str]], Optional[Dict[str, Any]]]:
+    """Resolve one raw tool. A served listing on this instance has priority."""
+    if name in data.tool_analytics_parameter_ownership:
+        return (
+            data.tool_analytics_parameter_ownership[name],
+            data.tool_input_schemas.get(name),
+        )
+
+    resolver = data.options.resolve_original_tool
+    if resolver is None:
+        return None, None
+
+    try:
+        descriptor = resolver(name)
+        if inspect.isawaitable(descriptor):
+            descriptor = await descriptor
+        if descriptor is None:
+            return None, None
+        schema = _descriptor_input_schema(descriptor)
+        if not isinstance(schema, dict):
+            log(
+                f"Warning: resolve_original_tool failed for tool {name!r}: "
+                "resolver returned no usable input schema"
+            )
+            return None, None
+        return analytics_owned_parameters(data.options, schema), schema
+    except Exception as error:  # noqa: BLE001 - analytics must not break dispatch
+        log(f"Warning: resolve_original_tool failed for tool {name!r}: {error}")
+        return None, None
+
+
+def _descriptor_input_schema(descriptor: Any) -> Any:
+    """Read MCP 1.x, MCP 2.x, and dictionary tool descriptors."""
+    if isinstance(descriptor, Mapping):
+        if "inputSchema" in descriptor:
+            return descriptor["inputSchema"]
+        return descriptor.get("input_schema")
+
+    schema = getattr(descriptor, "input_schema", None)
+    if schema is not None:
+        return schema
+    return getattr(descriptor, "inputSchema", None)

posthog/mcp/_context_parameters.py

@@ -48,10 +48,11 @@ def add_context_parameter_to_schema(
     """Return a new JSON Schema dict with a ``context`` string property added.
 
     Returns the input unchanged (logging a warning) for schemas that already
-    define ``context`` or use ``oneOf``/``allOf``/``anyOf``. ``required`` controls
-    whether ``context`` is added to the schema's ``required`` list — pass ``False``
-    where the advertised schema is also used to validate inbound calls (the
-    low-level server), so a call omitting ``context`` is not rejected."""
+    define ``context`` or use ``$ref``/``oneOf``/``allOf``/``anyOf``. ``required``
+    controls whether ``context`` is added to the schema's ``required`` list —
+    pass ``False`` where the advertised schema is also used to validate inbound
+    calls (the low-level server), so a call omitting ``context`` is not rejected.
+    """
     schema = input_schema
 
     if (
@@ -64,9 +65,10 @@ def add_context_parameter_to_schema(
         )
         return schema
 
-    if schema and (schema.get("oneOf") or schema.get("allOf") or schema.get("anyOf")):
+    if schema and any(schema.get(key) for key in ("$ref", "oneOf", "allOf", "anyOf")):
         log(
-            f'WARN: Tool "{tool_name}" has complex schema (oneOf/allOf/anyOf). Skipping context injection.'
+            f'WARN: Tool "{tool_name}" has complex schema '
+            "($ref/oneOf/allOf/anyOf). Skipping context injection."
         )
         return schema
 

posthog/mcp/_conversation_id.py

@@ -35,7 +35,9 @@ def add_conversation_id_to_schema(
     input_schema: Optional[Dict[str, Any]], tool_name: str = "unknown"
 ) -> Optional[Dict[str, Any]]:
     """Return a new JSON Schema with an optional ``conversation_id`` string property.
-    Skips schemas that already define it or use ``oneOf``/``allOf``/``anyOf``."""
+    Skips schemas that already define it or use
+    ``$ref``/``oneOf``/``allOf``/``anyOf``.
+    """
     schema = input_schema
     if (
         schema
@@ -48,7 +50,7 @@ def add_conversation_id_to_schema(
             f"WARN: Tool \"{tool_name}\" already has '{CONVERSATION_ID_PARAM_NAME}'. Skipping injection."
         )
         return schema
-    if schema and (schema.get("oneOf") or schema.get("allOf") or schema.get("anyOf")):
+    if schema and any(schema.get(key) for key in ("$ref", "oneOf", "allOf", "anyOf")):
         log(
             f'WARN: Tool "{tool_name}" has complex schema. Skipping conversation_id injection.'
         )

posthog/mcp/_instrument_lowlevel.py

@@ -22,13 +22,18 @@
 
 import mcp.types as mcp_types
 
+from ._argument_ownership import (
+    cache_listed_tool_ownership,
+    resolve_lowlevel_tool_ownership,
+)
 from ._context_parameters import is_context_enabled, schema_has_param
 from ._conversation_id import build_prompt_back
 from ._event_types import MCPAnalyticsEventType
 from ._instrumentation import (
     advertised_tool_names,
     apply_virtual_tool_injection,
     collect_listed_tools,
+    copy_tools_list_result,
     extract_tools,
     is_first_listing_page,
     mutate_tool_schema,
@@ -54,9 +59,11 @@
 
 
 def instrument_low_level(server: Any, data: MCPAnalyticsData) -> None:
-    """Instrument a raw ``mcp.server.Server``. ``context`` is injected as an
-    optional schema property and NOT stripped — that schema is also the call's
-    validation schema, and a typical ``(name, arguments)`` handler ignores extra keys."""
+    """Instrument a raw ``mcp.server.Server``.
+
+    The adapter removes arguments that a listing or resolver proves PostHog
+    owns. Unknown arguments pass through unchanged.
+    """
     data.server_name = getattr(server, "name", None)
     data.server_version = getattr(server, "version", None)
     _wrap_call_tool(server, data, strip_injected=False)
@@ -186,11 +193,21 @@ def _wrap_call_tool(
     async def handler(req: Any) -> Any:
         name = req.params.name
         arguments = dict(req.params.arguments or {})
-        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), None)
-        )
+        parameter_ownership = None
+        if strip_injected:
+            strip, model_ours, input_schema = await _standalone_ownership(
+                data, high_level, name, req.params.meta
+            )
+        else:
+            parameter_ownership, input_schema = await resolve_lowlevel_tool_ownership(
+                data, name
+            )
+            strip = set(parameter_ownership or ())
+            model_ours = (
+                "llm_model" in parameter_ownership
+                if parameter_ownership is not None
+                else data.tool_model_parameter_injected.get(name)
+            )
         client_name, client_version = _client_info(server)
         protocol_version = _protocol_version(server)
         mcp_session_id = _mcp_session_id(server)
@@ -213,6 +230,7 @@ async def handler(req: Any) -> Any:
             protocol_version=protocol_version,
             extra={"session_id": mcp_session_id, "ctx": _request_context(server)},
             input_schema=input_schema,
+            analytics_owned_parameters=parameter_ownership,
         )
 
         if lifecycle.is_missing_capability and (
@@ -327,6 +345,7 @@ async def handler(req: Any) -> Any:
 
 
 def _inject_tool_schemas(
+    server: Any,
     data: MCPAnalyticsData,
     tools: list,
     *,
@@ -345,6 +364,7 @@ def _inject_tool_schemas(
     verdicts: Dict[str, bool] = {}
     for tool in tools:
         schema = getattr(tool, "inputSchema", None)
+        cache_listed_tool_ownership(data, tool, schema_attribute="inputSchema")
         mutate_tool_schema(
             data,
             tool,
@@ -363,6 +383,11 @@ def _inject_tool_schemas(
             # Which one dispatches is unknown, so the strip fails closed.
             data.tool_model_parameter_injected[tool.name] = False
 
+    cache = getattr(server, "_tool_cache", None)
+    if isinstance(cache, dict):
+        for tool in tools:
+            cache[tool.name] = tool
+
 
 def _wrap_list_tools(
     server: Any,
@@ -404,34 +429,36 @@ async def probe_raw_tool_names(_ctx: Any = None) -> Optional[Set[str]]:
                     "registering your handlers."
                 )
             return None
-        result = await original(mcp_types.ListToolsRequest(method="tools/list"))
+        result = copy_tools_list_result(
+            await original(mcp_types.ListToolsRequest(method="tools/list"))
+        )
         tools = extract_tools(result)
-        # `original` is usually the SDK's own list_tools decorator, which rebuilds
-        # `Server._tool_cache` from these un-injected schemas every time it runs.
-        # That cache is what the SDK validates real tool arguments against, so
-        # without re-injecting here the next real call is rejected for sending the
-        # `context` we advertised. Same reason the `req is None` branch below
-        # injects.
+        # Resolve ownership from the host listing. Apply injection only to the
+        # copy so a shared host descriptor remains original.
         _inject_tool_schemas(
-            data, tools, context_required=context_required, high_level=high_level
+            server,
+            data,
+            tools,
+            context_required=context_required,
+            high_level=high_level,
         )
         return advertised_tool_names(tools)
 
     data.raw_tool_names_probe = probe_raw_tool_names
 
     async def handler(req: Any) -> Any:
         # The server calls the handler with None to populate its tool cache.
-        # Skip analytics there — but still inject, because that cache is the
-        # schema the SDK validates calls against. This adapter advertises
-        # `context`/`conversation_id` without stripping them, so a cache built
-        # from un-injected schemas rejects the very arguments we told the agent
-        # to send ("Additional properties are not allowed") on any tool with
-        # `additionalProperties: false`.
+        # Skip analytics there, but return the same injected schema as a client
+        # listing.
         if req is None:
-            result = await original(req)
+            result = copy_tools_list_result(await original(req))
             tools = extract_tools(result)
             _inject_tool_schemas(
-                data, tools, context_required=context_required, high_level=high_level
+                server,
+                data,
+                tools,
+                context_required=context_required,
+                high_level=high_level,
             )
             return result
 
@@ -464,7 +491,7 @@ async def handler(req: Any) -> Any:
 
         start = time.monotonic()
         try:
-            result = await original(req)
+            result = copy_tools_list_result(await original(req))
         except Exception as error:
             await lifecycle.record_error(error, (time.monotonic() - start) * 1000)
             raise
@@ -481,7 +508,11 @@ async def handler(req: Any) -> Any:
         )
 
         _inject_tool_schemas(
-            data, tools, context_required=context_required, high_level=high_level
+            server,
+            data,
+            tools,
+            context_required=context_required,
+            high_level=high_level,
         )
 
         result = apply_virtual_tool_injection(