fix(mcp): preserve FastMCP schemas across setup order and versions

Keep a weak reference to the standalone FastMCP wrapper in shared tracking
state so a later wrapper install enriches existing low-level instrumentation
without wrapping handlers twice. Resolve schemas for the client-requested
tool version before stripping analytics-owned arguments.

Consolidate adapter imports at the optional MCP dependency boundary instead
of repeating imports in dispatch branches. Keep FastMCP-specific imports
lazy so the official SDK remains usable without that optional package.

Validation: reproduced both setup-order failures and the versioned-tool
argument failure before fixing them. MCP v1 suite 307 passed; MCP v2 suite
296 passed, 13 skipped. Ruff format/lint and filtered mypy passed.
Fresh-process checks passed with MCP absent and with FastMCP absent.

Author: lucasheriques

posthog/mcp/README.md

@@ -166,6 +166,8 @@ Standalone `fastmcp` 4 uses the MCP SDK v2 handler registry. `instrument()` dete
 that registry automatically and captures tool calls over stdio and streamable HTTP,
 including the stateless protocol. Mounted tools retain their own arguments; analytics
 parameters are removed before dispatch only when the tool does not declare them.
+Instrumenting both the wrapper and its underlying server works in either order.
+For versioned tools, argument ownership follows the version requested by the client.
 The same installation code continues to support standalone FastMCP 2.x/3.x on MCP SDK v1.
 
 Two gaps worth knowing: jlowin's `fastmcp` 2.x/3.x doesn't expose the attribute the

posthog/mcp/__init__.py

@@ -31,6 +31,7 @@
 
 from __future__ import annotations
 
+import weakref
 from datetime import datetime, timezone
 from typing import Any, Optional
 
@@ -257,55 +258,54 @@ def instrument(
     key = _canonical_server(server)
 
     try:
-        # Imported inside the try: the adapters touch major-specific modules, and
-        # an import error must degrade to the no-op handle, not crash the host.
+        # MCP is an optional peer: load adapters only when instrumentation is
+        # requested, inside the no-crash boundary. Class probes stay major-specific.
         from ._compatibility import (
             is_fastmcp,
             is_fastmcp_v2,
             is_low_level_server,
             is_mcpserver,
             uses_v2_handler_registry,
         )
+        from ._instrument_fastmcp import instrument_fastmcp
+        from ._instrument_lowlevel import instrument_fastmcp_v2, instrument_low_level
+        from ._instrument_v2 import instrument_lowlevel_v2, instrument_mcpserver_v2
 
         client = _resolve_client(posthog_client)
         if client is None:
             log("Warning: no PostHog client available; MCP events will not be sent.")
 
-        if get_server_tracking_data(key) is not None:
+        existing_data = get_server_tracking_data(key)
+        data = existing_data
+        if data is None:
+            sink = McpEventSink(client) if client is not None else None
+            data = MCPAnalyticsData(
+                options=opts, sink=sink, session_id=new_session_id()
+            )
+
+        if is_fastmcp_v2(server) and uses_v2_handler_registry(key):
+            data.standalone_fastmcp = weakref.ref(server)
+
+        if existing_data is not None:
+            autowire_stateless_mint(server)
             log("instrument() - server already instrumented, skipping initialization")
             return McpAnalytics(key)
 
-        sink = McpEventSink(client) if client is not None else None
-        data = MCPAnalyticsData(options=opts, sink=sink, session_id=new_session_id())
         set_server_tracking_data(key, data)
 
         if is_fastmcp(server):
-            from ._instrument_fastmcp import instrument_fastmcp
-
             instrument_fastmcp(server, data)
         elif is_mcpserver(server):
-            from ._instrument_v2 import instrument_mcpserver_v2
-
             instrument_mcpserver_v2(server, data)
         elif is_fastmcp_v2(server):
             if uses_v2_handler_registry(server._mcp_server):
-                from ._instrument_v2 import instrument_lowlevel_v2
-
-                instrument_lowlevel_v2(
-                    server._mcp_server, data, strip_injected_for=server
-                )
+                instrument_lowlevel_v2(server._mcp_server, data)
             else:
-                from ._instrument_lowlevel import instrument_fastmcp_v2
-
                 instrument_fastmcp_v2(server, data)
         elif is_low_level_server(server):
             if uses_v2_handler_registry(server):
-                from ._instrument_v2 import instrument_lowlevel_v2
-
                 instrument_lowlevel_v2(server, data)
             else:
-                from ._instrument_lowlevel import instrument_low_level
-
                 instrument_low_level(server, data)
         else:
             raise TypeError(

posthog/mcp/_instrument_v2.py

@@ -31,6 +31,7 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Mapping
 from typing import Any, Dict, Optional, Tuple
 
 import mcp.types as mcp_types
@@ -89,21 +90,17 @@ def instrument_mcpserver_v2(server: Any, data: MCPAnalyticsData) -> None:
     _patch_add_request_handler(low_level, data, wrap_call=False, high_level=server)
 
 
-def instrument_lowlevel_v2(
-    server: Any, data: MCPAnalyticsData, *, strip_injected_for: Any = None
-) -> None:
+def instrument_lowlevel_v2(server: Any, data: MCPAnalyticsData) -> None:
     """Instrument a raw v2 low-level ``Server``. ``context`` is injected as an
     *optional* schema property and NOT stripped — the schema doubles as the
     call's validation surface, and a typical ``(ctx, params)`` handler ignores
-    extra argument keys. For standalone FastMCP, ``strip_injected_for`` supplies
+    extra argument keys. For standalone FastMCP, the shared tracking state supplies
     the tool schemas so injected arguments are removed before validation."""
     data.server_name = getattr(server, "name", None)
     data.server_version = getattr(server, "version", None)
-    _wrap_v2_call_tool(server, data, strip_injected_for=strip_injected_for)
+    _wrap_v2_call_tool(server, data)
     _wrap_v2_list_tools(server, data, context_required=False)
-    _patch_add_request_handler(
-        server, data, wrap_call=True, strip_injected_for=strip_injected_for
-    )
+    _patch_add_request_handler(server, data, wrap_call=True)
 
 
 # --- registry plumbing ---------------------------------------------------------
@@ -115,12 +112,7 @@ def _replace_handler(server: Any, method: str, wrapped: Any, params_type: Any) -
 
 
 def _patch_add_request_handler(
-    server: Any,
-    data: MCPAnalyticsData,
-    *,
-    wrap_call: bool,
-    high_level: Any = None,
-    strip_injected_for: Any = None,
+    server: Any, data: MCPAnalyticsData, *, wrap_call: bool, high_level: Any = None
 ) -> None:
     """Wrap ``add_request_handler`` so handlers registered *after* instrument()
     for the instrumented methods get wrapped too. Registrations for other
@@ -134,7 +126,7 @@ def add_request_handler(method: str, params_type: Any, handler: Any) -> None:
         if getattr(handler, _WRAPPED_FLAG, False):
             return
         if method == _CALL_METHOD and wrap_call:
-            _wrap_v2_call_tool(server, data, strip_injected_for=strip_injected_for)
+            _wrap_v2_call_tool(server, data)
         elif method == _LIST_METHOD:
             _wrap_v2_list_tools(
                 server,
@@ -404,17 +396,24 @@ def _deliver_conversation_id(
 # --- low-level: tools/call ------------------------------------------------------
 
 
-async def _standalone_tool_schema(server: Any, name: str) -> Any:
+async def _standalone_tool_schema(server: Any, name: str, ctx: Any) -> Any:
     try:
-        tool = await server.get_tool(name)
+        # Standalone FastMCP is optional even when the official MCP SDK is installed.
+        from fastmcp.server.dependencies import extract_version_spec
+        from fastmcp.utilities.versions import VersionSpec
+
+        params = getattr(ctx, "params", None)
+        meta = params.get("_meta") if isinstance(params, Mapping) else None
+        version = extract_version_spec(meta)
+        tool = await server.get_tool(
+            name, version=VersionSpec(eq=version) if version else None
+        )
         return getattr(tool, "parameters", None)
     except Exception:  # noqa: BLE001 - schema lookup must not prevent dispatch
         return None
 
 
-def _wrap_v2_call_tool(
-    server: Any, data: MCPAnalyticsData, *, strip_injected_for: Any = None
-) -> None:
+def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None:
     entry = server.get_request_handler(_CALL_METHOD)
     if entry is None or getattr(entry.handler, _WRAPPED_FLAG, False):
         return
@@ -424,8 +423,9 @@ async def handler(ctx: Any, params: Any) -> Any:
         name = params.name
         arguments = dict(params.arguments or {})
         analytics_owns_model = data.tool_model_parameter_injected.get(name, False)
-        if strip_injected_for is not None:
-            schema = await _standalone_tool_schema(strip_injected_for, name)
+        standalone = data.standalone_fastmcp() if data.standalone_fastmcp else None
+        if standalone is not None:
+            schema = await _standalone_tool_schema(standalone, name, ctx)
             analytics_owns_model = (
                 isinstance(schema, dict)
                 and is_capture_model_enabled(data.options.capture_model)

posthog/mcp/_internal.py

@@ -84,6 +84,8 @@ class MCPAnalyticsData:
     initialized_sessions: "OrderedDict[str, None]" = field(default_factory=OrderedDict)
     server_name: Optional[str] = None
     server_version: Optional[str] = None
+    # A strong wrapper reference would retain the low-level WeakKeyDictionary key.
+    standalone_fastmcp: Optional["weakref.ReferenceType[Any]"] = None
     session_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
 
     def mark_session_initialized(self, session_id: str) -> None:

posthog/test/mcp/test_fastmcp_v4.py

@@ -38,7 +38,7 @@ async def rpc(http, protocol, method, params):
     params = dict(params)
     headers = legacy_headers()
     if protocol == MODERN_PROTOCOL_VERSION:
-        params["_meta"] = modern_meta()
+        params["_meta"] = {**modern_meta(), **params.get("_meta", {})}
         headers = modern_headers(method, params.get("name"))
     else:
         headers["mcp-protocol-version"] = protocol
@@ -68,12 +68,18 @@ async def initialize(http, protocol):
 
 
 @pytest.mark.parametrize("protocol", [LEGACY_PROTOCOL_VERSION, MODERN_PROTOCOL_VERSION])
-async def test_capture_success_failure_and_sink_outage(protocol):
+@pytest.mark.parametrize("order", ["wrapper_first", "lowlevel_first", "wrapper_twice"])
+async def test_capture_success_failure_and_sink_outage(protocol, order):
     server = FastMCP("example-server")
     sink = FakeClient()
     options = MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True)
-    instrument(server, sink, options)
-    instrument(server, sink, options)
+    targets = {
+        "wrapper_first": [server, server._mcp_server],
+        "lowlevel_first": [server._mcp_server, server],
+        "wrapper_twice": [server, server],
+    }
+    for target in targets[order]:
+        instrument(target, sink, options)
     received = []
 
     @server.tool()
@@ -171,3 +177,33 @@ def echo(context: str, conversation_id: str, llm_model: str) -> str:
     calls = events_named(sink, "$mcp_tool_call")
     assert len(calls) == 1
     assert "$mcp_llm_model" not in calls[0]["properties"]
+
+
+async def test_preserve_parameters_of_requested_tool_version():
+    server = FastMCP("example-versioned-tools")
+
+    @server.tool(name="echo", version="1")
+    def older(text: str, context: str) -> str:
+        return f"{text}|{context}"
+
+    @server.tool(name="echo", version="2")
+    def newer(text: str) -> str:
+        return text
+
+    sink = FakeClient()
+    instrument(server, sink)
+    async with wire(server) as http:
+        result = await rpc(
+            http,
+            MODERN_PROTOCOL_VERSION,
+            "tools/call",
+            {
+                "name": "echo",
+                "arguments": {"text": "example", "context": "application-context"},
+                "_meta": {"fastmcp": {"version": "1"}},
+            },
+        )
+        assert not result.get("isError", False), result
+        assert result["content"][0]["text"] == "example|application-context"
+        await flush_background()
+    assert len(events_named(sink, "$mcp_tool_call")) == 1