From 62bb16f5a316c2550f1ad8c7200fcfa53aa821aa Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Wed, 9 Sep 2026 17:21:08 -0300 Subject: [PATCH 1/5] fix(mcp): capture standalone FastMCP on the v2 registry Select the adapter by handler registry so standalone FastMCP 4 captures tool calls without changing the customer's instrument() invocation. Preserve application-owned arguments, including mounted tools, and use native v2 model fields when FastMCP exposes deprecated v1 aliases. Add HTTP regression coverage for both protocol eras, late tool registration, repeat instrumentation, failed tools, unavailable capture, and mounted argument ownership. Include FastMCP 4 in the existing MCP v2 CI matrix, document adapter behavior, and add a Sampo patch changeset. Validation on current main: MCP v1 suite 307 passed; MCP v2 suite 291 passed, 13 skipped. Ruff lint/format, filtered mypy (231 source files), import warning check, and actionlint passed. Real stdio tools emitted HTTP capture batches to a local receiver under both protocol eras, two events per run. Authenticated wizard installation and hosted ingestion remain untested. --- .github/workflows/ci.yml | 9 +- .sampo/changesets/chivalrous-witch-goulven.md | 5 + posthog/mcp/README.md | 6 + posthog/mcp/__init__.py | 11 +- posthog/mcp/_exceptions.py | 2 +- posthog/mcp/_instrument_v2.py | 59 +++++- posthog/mcp/_instrumentation.py | 8 +- posthog/mcp/_output_instructions.py | 4 +- posthog/test/mcp/conftest.py | 1 + posthog/test/mcp/test_fastmcp_v4.py | 173 ++++++++++++++++++ posthog/test/mcp/test_output_instructions.py | 4 +- 11 files changed, 257 insertions(+), 25 deletions(-) create mode 100644 .sampo/changesets/chivalrous-witch-goulven.md create mode 100644 posthog/test/mcp/test_fastmcp_v4.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9057b3da3..7f2b83cf7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -155,10 +155,11 @@ jobs: # The MCP suite as a named gate per MCP Python SDK major. The v1 leg uses # the lockfile's mcp 1.x (also exercised incidentally by the `tests` # matrix — this leg exists as an explicit, named signal); the v2 leg - # (spec 2026-07-28) swaps in mcp>=2 and drops jlowin fastmcp, which pins - # mcp<2. posthog/test/mcp/conftest.py splits collection by major. + # (spec 2026-07-28) swaps in mcp>=2 and standalone FastMCP 4, which uses + # the v2 registry. posthog/test/mcp/conftest.py splits collection by major. name: MCP SDK ${{ matrix.mcp-major }} (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest + timeout-minutes: 20 strategy: matrix: python-version: ['3.10', '3.14'] @@ -189,8 +190,8 @@ jobs: if: matrix.mcp-major == 'v2' shell: bash run: | - uv pip uninstall --python $pythonLocation fastmcp - uv pip install --python $pythonLocation 'mcp>=2,<3' + uv pip uninstall --python "$pythonLocation" fastmcp + uv pip install --python "$pythonLocation" 'mcp>=2,<3' 'fastmcp>=4,<5' - name: Run MCP tests against SDK ${{ matrix.mcp-major }} run: | diff --git a/.sampo/changesets/chivalrous-witch-goulven.md b/.sampo/changesets/chivalrous-witch-goulven.md new file mode 100644 index 000000000..381f48c48 --- /dev/null +++ b/.sampo/changesets/chivalrous-witch-goulven.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Fix missing MCP analytics events with standalone FastMCP 4 while preserving tool arguments and compatibility with MCP SDK v1. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index a032629f7..ce487d9ee 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -162,6 +162,12 @@ Neither fires for stdio, for a correctly-wired server, or for a conversation-anc session. The instrument-time check can't see whether you added the middleware yourself (the app is already built by then), so ignore it if you did. +Standalone `fastmcp` 4 uses the MCP SDK v2 handler registry. `instrument()` detects +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. +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 instrument-time check reads, so those servers get the runtime warning only. And the deprecated SSE transport is excluded — it keys sessions off a query parameter, and the diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 633459300..d9b18ee57 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -288,9 +288,16 @@ def instrument( instrument_mcpserver_v2(server, data) elif is_fastmcp_v2(server): - from ._instrument_lowlevel import instrument_fastmcp_v2 + 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 + ) + else: + from ._instrument_lowlevel import instrument_fastmcp_v2 - instrument_fastmcp_v2(server, data) + instrument_fastmcp_v2(server, data) elif is_low_level_server(server): if uses_v2_handler_registry(server): from ._instrument_v2 import instrument_lowlevel_v2 diff --git a/posthog/mcp/_exceptions.py b/posthog/mcp/_exceptions.py index bf6e6f424..f34f30c7b 100644 --- a/posthog/mcp/_exceptions.py +++ b/posthog/mcp/_exceptions.py @@ -56,7 +56,7 @@ def _is_call_tool_result(value: Any) -> bool: dict or a pydantic model from the ``mcp`` SDK.""" if isinstance(value, dict): return "isError" in value and isinstance(value.get("content"), list) - return hasattr(value, "isError") and isinstance( + return (hasattr(value, "is_error") or hasattr(value, "isError")) and isinstance( getattr(value, "content", None), list ) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index c740ae273..6e1e65397 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -89,16 +89,21 @@ 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) -> None: +def instrument_lowlevel_v2( + server: Any, data: MCPAnalyticsData, *, strip_injected_for: Any = None +) -> 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.""" + extra argument keys. For standalone FastMCP, ``strip_injected_for`` 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) + _wrap_v2_call_tool(server, data, strip_injected_for=strip_injected_for) _wrap_v2_list_tools(server, data, context_required=False) - _patch_add_request_handler(server, data, wrap_call=True) + _patch_add_request_handler( + server, data, wrap_call=True, strip_injected_for=strip_injected_for + ) # --- registry plumbing --------------------------------------------------------- @@ -110,7 +115,12 @@ 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 + server: Any, + data: MCPAnalyticsData, + *, + wrap_call: bool, + high_level: Any = None, + strip_injected_for: Any = None, ) -> None: """Wrap ``add_request_handler`` so handlers registered *after* instrument() for the instrumented methods get wrapped too. Registrations for other @@ -124,7 +134,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) + _wrap_v2_call_tool(server, data, strip_injected_for=strip_injected_for) elif method == _LIST_METHOD: _wrap_v2_list_tools( server, @@ -394,7 +404,17 @@ def _deliver_conversation_id( # --- low-level: tools/call ------------------------------------------------------ -def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None: +async def _standalone_tool_schema(server: Any, name: str) -> Any: + try: + tool = await server.get_tool(name) + 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: entry = server.get_request_handler(_CALL_METHOD) if entry is None or getattr(entry.handler, _WRAPPED_FLAG, False): return @@ -403,6 +423,27 @@ def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None: 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) + analytics_owns_model = ( + isinstance(schema, dict) + and is_capture_model_enabled(data.options.capture_model) + and can_inject_model_parameter(schema) + ) + injected = {"context"} + if data.options.enable_conversation_id: + injected.add("conversation_id") + if analytics_owns_model: + injected.add("llm_model") + call_arguments = { + key: value + for key, value in arguments.items() + if not isinstance(schema, dict) + or key not in injected + or schema_has_param(schema, key) + } + params = params.model_copy(update={"arguments": call_arguments}) token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) ) @@ -411,9 +452,7 @@ async def handler(ctx: Any, params: Any) -> Any: name=name, arguments=arguments, request_meta=request_meta_from_context(ctx), - allow_self_reported_model=data.tool_model_parameter_injected.get( - name, False - ), + allow_self_reported_model=analytics_owns_model, mcp_session_id=mcp_session_id, token=token, client_name=client_name, diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 328fc597f..68b3e7bee 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -164,10 +164,10 @@ def is_tool_result_error(result: Any) -> bool: (wire JSON unchanged); check both shapes.""" if isinstance(result, dict): return result.get("isError") is True or result.get("is_error") is True - return ( - getattr(result, "isError", None) is True - or getattr(result, "is_error", None) is True - ) + is_error = getattr(result, "is_error", None) + if is_error is not None: + return is_error is True + return getattr(result, "isError", None) is True def build_tool_call_request( diff --git a/posthog/mcp/_output_instructions.py b/posthog/mcp/_output_instructions.py index e1fb084ae..d340f36ab 100644 --- a/posthog/mcp/_output_instructions.py +++ b/posthog/mcp/_output_instructions.py @@ -37,8 +37,8 @@ _CONVERSATION_ID_FIELD_DESCRIPTION = "The server-issued conversation identifier." # `outputSchema` on MCP SDK 1.x models, `output_schema` on 2.x (same wire field). -_OUTPUT_SCHEMA_ATTRS = ("outputSchema", "output_schema") -_STRUCTURED_CONTENT_ATTRS = ("structuredContent", "structured_content") +_OUTPUT_SCHEMA_ATTRS = ("output_schema", "outputSchema") +_STRUCTURED_CONTENT_ATTRS = ("structured_content", "structuredContent") def _read_attr(obj: Any, names: Tuple[str, ...]) -> Tuple[Optional[str], Any]: diff --git a/posthog/test/mcp/conftest.py b/posthog/test/mcp/conftest.py index 65c0e887d..d01802d00 100644 --- a/posthog/test/mcp/conftest.py +++ b/posthog/test/mcp/conftest.py @@ -25,6 +25,7 @@ "test_v2_mcpserver.py", "test_v2_lowlevel.py", "test_v2_wire_dual_era.py", + "test_fastmcp_v4.py", ] collect_ignore = _V2_ONLY if MCP_MAJOR < 2 else _V1_ONLY diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py new file mode 100644 index 000000000..bbbb88b19 --- /dev/null +++ b/posthog/test/mcp/test_fastmcp_v4.py @@ -0,0 +1,173 @@ +"""Standalone FastMCP on the MCP SDK v2 registry, exercised through HTTP.""" + +from contextlib import asynccontextmanager + +import httpx +import pytest + +pytest.importorskip("fastmcp", minversion="4") + +from fastmcp import FastMCP # noqa: E402 + +from posthog.mcp import MCPAnalyticsOptions, instrument # noqa: E402 +from posthog.test.mcp._helpers import ( # noqa: E402 + FakeClient, + events_named, + flush_background, +) +from posthog.test.mcp._helpers_v2 import ( # noqa: E402 + LEGACY_PROTOCOL_VERSION, + MODERN_PROTOCOL_VERSION, + legacy_headers, + modern_headers, + modern_meta, +) + + +@asynccontextmanager +async def wire(server): + app = server.http_app(json_response=True, stateless_http=True) + async with app.router.lifespan_context(app): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://localhost" + ) as http: + yield http + + +async def rpc(http, protocol, method, params): + params = dict(params) + headers = legacy_headers() + if protocol == MODERN_PROTOCOL_VERSION: + params["_meta"] = modern_meta() + headers = modern_headers(method, params.get("name")) + else: + headers["mcp-protocol-version"] = protocol + response = await http.post( + "/mcp", + headers=headers, + json={"jsonrpc": "2.0", "id": 1, "method": method, "params": params}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert "error" not in body, body + return body["result"] + + +async def initialize(http, protocol): + if protocol == LEGACY_PROTOCOL_VERSION: + await rpc( + http, + protocol, + "initialize", + { + "protocolVersion": protocol, + "capabilities": {}, + "clientInfo": {"name": "example-client", "version": "1.0"}, + }, + ) + + +@pytest.mark.parametrize("protocol", [LEGACY_PROTOCOL_VERSION, MODERN_PROTOCOL_VERSION]) +async def test_capture_success_failure_and_sink_outage(protocol): + server = FastMCP("example-server") + sink = FakeClient() + options = MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True) + instrument(server, sink, options) + instrument(server, sink, options) + received = [] + + @server.tool() + def add(a: int, b: int) -> str: + received.append({"a": a, "b": b}) + return str(a + b) + + @server.tool() + def fail() -> str: + raise ValueError("example failure") + + async with wire(server) as http: + await initialize(http, protocol) + listed = await rpc(http, protocol, "tools/list", {}) + schema = next(t for t in listed["tools"] if t["name"] == "add")["inputSchema"] + assert {"context", "conversation_id", "llm_model"} <= schema[ + "properties" + ].keys() + assert not {"context", "conversation_id", "llm_model"} & set( + schema.get("required", []) + ) + args = {"context": "example addition", "llm_model": "example-model"} + result = await rpc( + http, + protocol, + "tools/call", + {"name": "add", "arguments": {"a": 2, "b": 3, **args}}, + ) + assert not result.get("isError", False) + assert result["content"][0]["text"] == "5" + assert received == [{"a": 2, "b": 3}] + assert len(result["content"]) == 2 + failed = await rpc( + http, protocol, "tools/call", {"name": "fail", "arguments": args} + ) + assert failed["isError"] + await flush_background() + calls = events_named(sink, "$mcp_tool_call") + assert len(calls) == 2 + assert [c["properties"]["$mcp_is_error"] for c in calls] == [False, True] + assert all(c["properties"]["$mcp_protocol_version"] == protocol for c in calls) + assert calls[0]["properties"]["$mcp_intent"] == "example addition" + assert calls[0]["properties"]["$mcp_llm_model"] == "example-model" + + def unavailable(*args, **kwargs): + raise RuntimeError("example capture outage") + + sink.capture = unavailable + offline = await rpc( + http, protocol, "tools/call", {"name": "add", "arguments": {"a": 3, "b": 4}} + ) + assert not offline.get("isError", False) + assert offline["content"][0]["text"] == "7" + await flush_background() + + +@pytest.mark.parametrize("list_first", [False, True]) +@pytest.mark.parametrize("mounted", [False, True]) +async def test_preserve_application_parameters(list_first, mounted): + child = FastMCP("example-tools") + + @child.tool() + def echo(context: str, conversation_id: str, llm_model: str) -> str: + return f"{context}|{conversation_id}|{llm_model}" + + server = FastMCP("example-app") if mounted else child + if mounted: + server.mount(child, namespace="shared") + name = "shared_echo" if mounted else "echo" + sink = FakeClient() + instrument( + server, + sink, + MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True), + ) + async with wire(server) as http: + if list_first: + await rpc(http, MODERN_PROTOCOL_VERSION, "tools/list", {}) + result = await rpc( + http, + MODERN_PROTOCOL_VERSION, + "tools/call", + { + "name": name, + "arguments": { + "context": "own-context", + "conversation_id": "own-id", + "llm_model": "own-model", + }, + }, + ) + assert not result.get("isError", False) + assert result["content"][0]["text"] == "own-context|own-id|own-model" + await flush_background() + calls = events_named(sink, "$mcp_tool_call") + assert len(calls) == 1 + assert "$mcp_llm_model" not in calls[0]["properties"] diff --git a/posthog/test/mcp/test_output_instructions.py b/posthog/test/mcp/test_output_instructions.py index 0fd569606..656a82ec4 100644 --- a/posthog/test/mcp/test_output_instructions.py +++ b/posthog/test/mcp/test_output_instructions.py @@ -115,7 +115,7 @@ def test_mirrors_into_a_fastmcp_v1_tuple_result(): def test_mirrors_into_a_model_result_both_attr_shapes(): - for attr in ("structuredContent", "structured_content"): + for attr in ("structured_content", "structuredContent"): result = SimpleNamespace(**{attr: {"total": 7}}) _, delivered = mirror_instructions_into_structured_content(result, "conv-2") @@ -202,7 +202,7 @@ def test_mirror_does_not_mutate_a_shared_model_result(): # `structuredContent` on SDK 1.x, `structured_content` on 2.x (same wire field). def structured(result): - for attr in ("structuredContent", "structured_content"): + for attr in ("structured_content", "structuredContent"): if hasattr(result, attr): return getattr(result, attr) raise AssertionError("no structured content attribute") From 8865969b712fe36069cfaad2bdea533ffb3b833c Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Wed, 9 Sep 2026 18:19:51 -0300 Subject: [PATCH 2/5] 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. --- posthog/mcp/README.md | 2 ++ posthog/mcp/__init__.py | 40 +++++++++++++------------- posthog/mcp/_instrument_v2.py | 44 ++++++++++++++--------------- posthog/mcp/_internal.py | 2 ++ posthog/test/mcp/test_fastmcp_v4.py | 44 ++++++++++++++++++++++++++--- 5 files changed, 86 insertions(+), 46 deletions(-) diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index ce487d9ee..0046f4126 100644 --- a/posthog/mcp/README.md +++ b/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 diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index d9b18ee57..90fb94a83 100644 --- a/posthog/mcp/__init__.py +++ b/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,8 +258,8 @@ 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, @@ -266,46 +267,45 @@ def instrument( 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( diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 6e1e65397..7da82282d 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/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) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 34e547535..cb82aca33 100644 --- a/posthog/mcp/_internal.py +++ b/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: diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py index bbbb88b19..81a9e87b5 100644 --- a/posthog/test/mcp/test_fastmcp_v4.py +++ b/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 From 702797790da1b4e62305f2d8ddc50f46529cccb2 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Thu, 10 Sep 2026 12:39:04 -0300 Subject: [PATCH 3/5] fix(mcp): strip FastMCP analytics arguments from what was advertised Decide which analytics parameters to remove before dispatch from the tools/list the process served, recorded per tool in mutate_tool_schema, instead of asking FastMCP for the tool schema on every call. The live lookup remains only for tools never listed here or when the client pins a tool version, and now logs when it fails. Tools supplied by ToolInjectionMiddleware are listed but not resolvable via get_tool(), so the per-call lookup advertised `context` and then forwarded it, failing validation. Tested: posthog/test/mcp under mcp 2.2.0 + fastmcp 4.0.3 (297 passed, 13 skipped, new middleware test red before the fix) and under mcp 1.30.0 + fastmcp 3.2.0 (307 passed); ruff, mypy-baseline clean. Reviewer note: tool_model_parameter_injected still exists alongside the new tool_injected_parameters set because four adapters read it; folding it in is a follow-up. Claude-Session: https://claude.ai/code/session_015wZCNFdfPNK9k4utwJ5bMu --- posthog/mcp/__init__.py | 3 ++ posthog/mcp/_instrument_v2.py | 73 +++++++++++++++++++---------- posthog/mcp/_instrumentation.py | 6 +++ posthog/mcp/_internal.py | 7 ++- posthog/test/mcp/test_fastmcp_v4.py | 43 +++++++++++++++++ 5 files changed, 106 insertions(+), 26 deletions(-) diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 90fb94a83..79f6ae53c 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -286,6 +286,9 @@ def instrument( if is_fastmcp_v2(server) and uses_v2_handler_registry(key): data.standalone_fastmcp = weakref.ref(server) + # A standalone FastMCP wrapper and its low-level server share one tracking + # key, so instrumenting the second of the pair must still attach what only + # that object provides: the wrapper's schema lookup and ASGI app factories. if existing_data is not None: autowire_stateless_mint(server) log("instrument() - server already instrumented, skipping initialization") diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 7da82282d..170e44340 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -32,7 +32,7 @@ import time from collections.abc import Mapping -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict, FrozenSet, Optional, Tuple import mcp.types as mcp_types @@ -396,21 +396,46 @@ def _deliver_conversation_id( # --- low-level: tools/call ------------------------------------------------------ -async def _standalone_tool_schema(server: Any, name: str, ctx: Any) -> Any: +def _requested_tool_version(ctx: Any) -> Optional[str]: + """The FastMCP tool version a client pinned via request ``_meta``, if any.""" try: # 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) + return extract_version_spec(meta) + except Exception: # noqa: BLE001 - version parsing must not prevent dispatch + return None + + +async def _standalone_injected_parameters( + server: Any, data: MCPAnalyticsData, name: str, version: Optional[str] +) -> Optional[FrozenSet[str]]: + """Which analytics parameters to strip, derived from the tool's own schema, for + a tool this process never listed or a client-pinned version. ``None`` when the + tool cannot be resolved (for example, tools supplied by middleware), in which + case nothing is stripped.""" + try: + from fastmcp.utilities.versions import VersionSpec + 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 + 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 + if not isinstance(schema, dict): return None + injected = {"context"} + if data.options.enable_conversation_id: + injected.add("conversation_id") + if is_capture_model_enabled(data.options.capture_model) and ( + can_inject_model_parameter(schema) + ): + injected.add("llm_model") + return frozenset(key for key in injected if not schema_has_param(schema, key)) def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None: @@ -425,25 +450,23 @@ async def handler(ctx: Any, params: Any) -> Any: analytics_owns_model = data.tool_model_parameter_injected.get(name, False) 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) - and can_inject_model_parameter(schema) - ) - injected = {"context"} - if data.options.enable_conversation_id: - injected.add("conversation_id") - if analytics_owns_model: - injected.add("llm_model") - call_arguments = { - key: value - for key, value in arguments.items() - if not isinstance(schema, dict) - or key not in injected - or schema_has_param(schema, key) - } - params = params.model_copy(update={"arguments": call_arguments}) + # The listing this process served is the source of truth for what was + # advertised. A client-pinned version may differ from the listed one, + # so only then is the tool's own schema consulted. + version = _requested_tool_version(ctx) + injected = data.tool_injected_parameters.get(name) + if injected is None or version is not None: + injected = await _standalone_injected_parameters( + standalone, data, name, version + ) + if injected is not None: + analytics_owns_model = "llm_model" in injected + call_arguments = { + key: value + for key, value in arguments.items() + if key not in injected + } + params = params.model_copy(update={"arguments": call_arguments}) token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) ) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 68b3e7bee..78a19f42e 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -670,6 +670,7 @@ def mutate_tool_schema( """ schema = getattr(tool, schema_attribute, None) original_schema = schema + injected = set() if ( tool.name != GET_MORE_TOOLS_NAME and is_context_enabled(data.options.context) @@ -681,6 +682,7 @@ def mutate_tool_schema( get_context_description(data.options.context), required=context_required, ) + injected.add("context") if is_capture_model_enabled(data.options.capture_model): model_was_injected = data.tool_model_parameter_injected.get(tool.name, False) app_owns_model = ( @@ -696,12 +698,16 @@ def mutate_tool_schema( data.tool_model_parameter_injected[tool.name] = ( not app_owns_model and schema_has_param(schema, "llm_model") ) + if data.tool_model_parameter_injected[tool.name]: + injected.add("llm_model") if ( tool.name != GET_MORE_TOOLS_NAME and data.options.enable_conversation_id and not schema_has_param(schema, "conversation_id") ): schema = add_conversation_id_to_schema(schema, tool.name) + injected.add("conversation_id") + data.tool_injected_parameters[tool.name] = frozenset(injected) if schema is not original_schema: try: setattr(tool, schema_attribute, schema) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index cb82aca33..684dbd863 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -17,7 +17,7 @@ from collections import OrderedDict from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, FrozenSet, Optional from .logger import log from ._sink import McpEventSink @@ -73,6 +73,11 @@ class MCPAnalyticsData: # 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) + # Every analytics parameter PostHog added to a tool's advertised schema at + # tools/list. Standalone FastMCP validates arguments against the tool's own + # schema, so exactly these keys are stripped before dispatch. Absent means + # "never served a listing for this tool" and falls back to a live lookup. + tool_injected_parameters: Dict[str, FrozenSet[str]] = field(default_factory=dict) # Which tools got `_mcp_instructions` declared on their advertised output # schema at tools/list. Only those may be mirrored into on a call — writing # an undeclared key fails the customer's whole result under diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py index 81a9e87b5..145eed36d 100644 --- a/posthog/test/mcp/test_fastmcp_v4.py +++ b/posthog/test/mcp/test_fastmcp_v4.py @@ -207,3 +207,46 @@ def newer(text: str) -> str: assert result["content"][0]["text"] == "example|application-context" await flush_background() assert len(events_named(sink, "$mcp_tool_call")) == 1 + + +async def test_strip_analytics_parameters_from_middleware_tools(): + from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware + from fastmcp.tools import Tool + + def echo(text: str) -> str: + return text + + server = FastMCP("example-middleware-tools") + server.add_middleware(ToolInjectionMiddleware(tools=[Tool.from_function(echo)])) + sink = FakeClient() + instrument( + server, + sink, + MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True), + ) + async with wire(server) as http: + listed = await rpc(http, MODERN_PROTOCOL_VERSION, "tools/list", {}) + schema = next(t for t in listed["tools"] if t["name"] == "echo")["inputSchema"] + assert {"context", "conversation_id", "llm_model"} <= schema[ + "properties" + ].keys() + result = await rpc( + http, + MODERN_PROTOCOL_VERSION, + "tools/call", + { + "name": "echo", + "arguments": { + "text": "example", + "context": "example intent", + "llm_model": "example-model", + }, + }, + ) + assert not result.get("isError", False), result + assert result["content"][0]["text"] == "example" + await flush_background() + calls = events_named(sink, "$mcp_tool_call") + assert len(calls) == 1 + assert calls[0]["properties"]["$mcp_intent"] == "example intent" + assert calls[0]["properties"]["$mcp_llm_model"] == "example-model" From 2b91f144221c8c82e52c52bc5e3966cd7d729abb Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Fri, 11 Sep 2026 13:23:35 -0300 Subject: [PATCH 4/5] chore(mcp): verify middleware dispatch without schema lookup Extend FastMCP HTTP coverage to prove middleware tools dispatch when get_tool returns None. Exercise calls before listing, all advertised analytics arguments after listing, and application-owned arguments in both paths. Document why unknown arguments remain intact. Validation: MCP v1 suite 307 passed; MCP v2 suite 299 passed, 13 skipped. Ruff check and format checks passed for MCP source and tests. Repository-wide filtered mypy passed for 231 source files. --- posthog/mcp/_instrument_v2.py | 5 ++-- posthog/test/mcp/test_fastmcp_v4.py | 39 ++++++++++++++++++++--------- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 170e44340..80c866ac5 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -414,8 +414,9 @@ async def _standalone_injected_parameters( ) -> Optional[FrozenSet[str]]: """Which analytics parameters to strip, derived from the tool's own schema, for a tool this process never listed or a client-pinned version. ``None`` when the - tool cannot be resolved (for example, tools supplied by middleware), in which - case nothing is stripped.""" + tool cannot be resolved. Without a schema, stripping could delete application + arguments. Middleware tools normally use the recorded tools/list ownership + instead, since they can dispatch without resolving through get_tool().""" try: from fastmcp.utilities.versions import VersionSpec diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py index 145eed36d..1a3ed4bd9 100644 --- a/posthog/test/mcp/test_fastmcp_v4.py +++ b/posthog/test/mcp/test_fastmcp_v4.py @@ -8,6 +8,8 @@ pytest.importorskip("fastmcp", minversion="4") from fastmcp import FastMCP # noqa: E402 +from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware # noqa: E402 +from fastmcp.tools import Tool # noqa: E402 from posthog.mcp import MCPAnalyticsOptions, instrument # noqa: E402 from posthog.test.mcp._helpers import ( # noqa: E402 @@ -137,18 +139,23 @@ def unavailable(*args, **kwargs): @pytest.mark.parametrize("list_first", [False, True]) -@pytest.mark.parametrize("mounted", [False, True]) -async def test_preserve_application_parameters(list_first, mounted): +@pytest.mark.parametrize("source", ["direct", "mounted", "middleware"]) +async def test_preserve_application_parameters(list_first, source): child = FastMCP("example-tools") - @child.tool() def echo(context: str, conversation_id: str, llm_model: str) -> str: return f"{context}|{conversation_id}|{llm_model}" - server = FastMCP("example-app") if mounted else child - if mounted: + tool = Tool.from_function(echo) + if source == "middleware": + child.add_middleware(ToolInjectionMiddleware(tools=[tool])) + assert await child.get_tool("echo") is None + else: + child.add_tool(tool) + server = FastMCP("example-app") if source == "mounted" else child + if source == "mounted": server.mount(child, namespace="shared") - name = "shared_echo" if mounted else "echo" + name = "shared_echo" if source == "mounted" else "echo" sink = FakeClient() instrument( server, @@ -210,14 +217,12 @@ def newer(text: str) -> str: async def test_strip_analytics_parameters_from_middleware_tools(): - from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware - from fastmcp.tools import Tool - def echo(text: str) -> str: return text server = FastMCP("example-middleware-tools") server.add_middleware(ToolInjectionMiddleware(tools=[Tool.from_function(echo)])) + assert await server.get_tool("echo") is None sink = FakeClient() instrument( server, @@ -225,6 +230,14 @@ def echo(text: str) -> str: MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True), ) async with wire(server) as http: + unlisted = await rpc( + http, + MODERN_PROTOCOL_VERSION, + "tools/call", + {"name": "echo", "arguments": {"text": "before listing"}}, + ) + assert not unlisted.get("isError", False), unlisted + assert unlisted["content"][0]["text"] == "before listing" listed = await rpc(http, MODERN_PROTOCOL_VERSION, "tools/list", {}) schema = next(t for t in listed["tools"] if t["name"] == "echo")["inputSchema"] assert {"context", "conversation_id", "llm_model"} <= schema[ @@ -239,6 +252,7 @@ def echo(text: str) -> str: "arguments": { "text": "example", "context": "example intent", + "conversation_id": "example-conversation", "llm_model": "example-model", }, }, @@ -247,6 +261,7 @@ def echo(text: str) -> str: assert result["content"][0]["text"] == "example" await flush_background() calls = events_named(sink, "$mcp_tool_call") - assert len(calls) == 1 - assert calls[0]["properties"]["$mcp_intent"] == "example intent" - assert calls[0]["properties"]["$mcp_llm_model"] == "example-model" + assert len(calls) == 2 + assert all(not call["properties"]["$mcp_is_error"] for call in calls) + assert calls[1]["properties"]["$mcp_intent"] == "example intent" + assert calls[1]["properties"]["$mcp_llm_model"] == "example-model" From 43fa340bd51ff1ffb936ce58c5d23127e20f3549 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Fri, 11 Sep 2026 13:38:17 -0300 Subject: [PATCH 5/5] fix(mcp): resolve FastMCP argument ownership per request Remove the server-wide injected-parameter cache. Resolve effective schemas with the current request context, middleware, and requested tool version before stripping analytics arguments. Avoid using cached model ownership when schema resolution fails. Add interleaved-client regression cases for dynamic providers, middleware tools, and middleware overrides in both listing orders. The provider and middleware cases reproduced a missing application-owned llm_model argument before the fix. Document the per-call schema lookup. Validation: MCP v1 suite 307 passed; MCP v2 suite 305 passed, 13 skipped. Ruff check and format passed for MCP source and tests. Repository-wide filtered mypy passed for 231 files. --- posthog/mcp/README.md | 2 + posthog/mcp/_instrument_v2.py | 47 +++++++++-------- posthog/mcp/_instrumentation.py | 6 --- posthog/mcp/_internal.py | 7 +-- posthog/test/mcp/test_fastmcp_v4.py | 80 +++++++++++++++++++++++++++++ 5 files changed, 109 insertions(+), 33 deletions(-) diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 0046f4126..51c1a22bc 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -168,6 +168,8 @@ including the stateless protocol. Mounted tools retain their own arguments; anal 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. +Each tool call resolves the schema through FastMCP's tool listing in the current request context, including middleware and session transforms. +This adds a schema lookup per call so clients with different tool schemas cannot change how another client's arguments are handled. 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 diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 80c866ac5..6e9e6da43 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -36,7 +36,7 @@ import mcp.types as mcp_types -from ._context_parameters import schema_has_param +from ._context_parameters import is_context_enabled, schema_has_param from ._conversation_id import build_prompt_back from ._instrumentation import ( _to_jsonable, @@ -412,24 +412,34 @@ 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]]: - """Which analytics parameters to strip, derived from the tool's own schema, for - a tool this process never listed or a client-pinned version. ``None`` when the - tool cannot be resolved. Without a schema, stripping could delete application - arguments. Middleware tools normally use the recorded tools/list ownership - instead, since they can dispatch without resolving through get_tool().""" - try: - from fastmcp.utilities.versions import VersionSpec + """Resolve ownership in the current request, including middleware and versions. - tool = await server.get_tool( - name, version=VersionSpec(eq=version) if version else None - ) + Listings from other requests can have different application-owned parameters. + Without a schema, stripping could delete application arguments. + """ + try: + from fastmcp.utilities.versions import VersionSpec, version_sort_key + + version_spec = VersionSpec(eq=version) if version else None + # Middleware can shadow registered tools, so resolve the effective listing. + candidates = [ + tool + for tool in await server.list_tools() + if tool.name == name + and (version_spec is None or version_spec.matches(tool.version)) + ] + tool = max(candidates, key=version_sort_key, default=None) + if tool is None: + tool = await server.get_tool(name, version=version_spec) 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 if not isinstance(schema, dict): return None - injected = {"context"} + injected = set() + if is_context_enabled(data.options.context): + injected.add("context") if data.options.enable_conversation_id: injected.add("conversation_id") if is_capture_model_enabled(data.options.capture_model) and ( @@ -451,17 +461,12 @@ async def handler(ctx: Any, params: Any) -> Any: analytics_owns_model = data.tool_model_parameter_injected.get(name, False) standalone = data.standalone_fastmcp() if data.standalone_fastmcp else None if standalone is not None: - # The listing this process served is the source of truth for what was - # advertised. A client-pinned version may differ from the listed one, - # so only then is the tool's own schema consulted. version = _requested_tool_version(ctx) - injected = data.tool_injected_parameters.get(name) - if injected is None or version is not None: - injected = await _standalone_injected_parameters( - standalone, data, name, version - ) + injected = await _standalone_injected_parameters( + standalone, data, name, version + ) + analytics_owns_model = injected is not None and "llm_model" in injected if injected is not None: - analytics_owns_model = "llm_model" in injected call_arguments = { key: value for key, value in arguments.items() diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 78a19f42e..68b3e7bee 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -670,7 +670,6 @@ def mutate_tool_schema( """ schema = getattr(tool, schema_attribute, None) original_schema = schema - injected = set() if ( tool.name != GET_MORE_TOOLS_NAME and is_context_enabled(data.options.context) @@ -682,7 +681,6 @@ def mutate_tool_schema( get_context_description(data.options.context), required=context_required, ) - injected.add("context") if is_capture_model_enabled(data.options.capture_model): model_was_injected = data.tool_model_parameter_injected.get(tool.name, False) app_owns_model = ( @@ -698,16 +696,12 @@ def mutate_tool_schema( data.tool_model_parameter_injected[tool.name] = ( not app_owns_model and schema_has_param(schema, "llm_model") ) - if data.tool_model_parameter_injected[tool.name]: - injected.add("llm_model") if ( tool.name != GET_MORE_TOOLS_NAME and data.options.enable_conversation_id and not schema_has_param(schema, "conversation_id") ): schema = add_conversation_id_to_schema(schema, tool.name) - injected.add("conversation_id") - data.tool_injected_parameters[tool.name] = frozenset(injected) if schema is not original_schema: try: setattr(tool, schema_attribute, schema) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 684dbd863..cb82aca33 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -17,7 +17,7 @@ from collections import OrderedDict from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Any, Dict, FrozenSet, Optional +from typing import Any, Dict, Optional from .logger import log from ._sink import McpEventSink @@ -73,11 +73,6 @@ class MCPAnalyticsData: # 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) - # Every analytics parameter PostHog added to a tool's advertised schema at - # tools/list. Standalone FastMCP validates arguments against the tool's own - # schema, so exactly these keys are stripped before dispatch. Absent means - # "never served a listing for this tool" and falls back to a live lookup. - tool_injected_parameters: Dict[str, FrozenSet[str]] = field(default_factory=dict) # Which tools got `_mcp_instructions` declared on their advertised output # schema at tools/list. Only those may be mirrored into on a call — writing # an undeclared key fails the customer's whole result under diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py index 1a3ed4bd9..ac29c3a9f 100644 --- a/posthog/test/mcp/test_fastmcp_v4.py +++ b/posthog/test/mcp/test_fastmcp_v4.py @@ -8,7 +8,10 @@ pytest.importorskip("fastmcp", minversion="4") from fastmcp import FastMCP # noqa: E402 +from fastmcp.server.dependencies import get_http_headers # noqa: E402 +from fastmcp.server.middleware import Middleware # noqa: E402 from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware # noqa: E402 +from fastmcp.server.providers import Provider # noqa: E402 from fastmcp.tools import Tool # noqa: E402 from posthog.mcp import MCPAnalyticsOptions, instrument # noqa: E402 @@ -265,3 +268,80 @@ def echo(text: str) -> str: assert all(not call["properties"]["$mcp_is_error"] for call in calls) assert calls[1]["properties"]["$mcp_intent"] == "example intent" assert calls[1]["properties"]["$mcp_llm_model"] == "example-model" + + +@pytest.mark.parametrize("source", ["provider", "middleware", "middleware_shadow"]) +@pytest.mark.parametrize( + "listing_order", [("application", "analytics"), ("analytics", "application")] +) +async def test_tool_argument_ownership_isolated_between_clients(source, listing_order): + def application(context: str, conversation_id: str, llm_model: str) -> str: + return f"{context}|{conversation_id}|{llm_model}" + + def analytics() -> str: + return "analytics" + + tools = { + "application": Tool.from_function(application, name="echo"), + "analytics": Tool.from_function(analytics, name="echo"), + } + + def current_tool() -> Tool: + return tools[get_http_headers()["x-example-client"]] + + class ClientTools(Provider): + async def _list_tools(self): + return [current_tool()] + + async def _get_tool(self, name, version=None): + return current_tool() if name == "echo" else None + + class ClientMiddleware(Middleware): + async def on_list_tools(self, context, call_next): + return [current_tool()] + + async def on_call_tool(self, context, call_next): + return await current_tool().run(context.message.arguments or {}) + + server = FastMCP("example-client-tools") + if source == "provider": + server.add_provider(ClientTools()) + else: + if source == "middleware_shadow": + server.add_tool(tools["analytics"]) + server.add_middleware(ClientMiddleware()) + sink = FakeClient() + instrument( + server, + sink, + MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True), + ) + async with wire(server) as http: + for client in listing_order: + http.headers["x-example-client"] = client + await rpc(http, MODERN_PROTOCOL_VERSION, "tools/list", {}) + for client, expected in ( + ("application", "own-context|own-id|own-model"), + ("analytics", "analytics"), + ): + http.headers["x-example-client"] = client + result = await rpc( + http, + MODERN_PROTOCOL_VERSION, + "tools/call", + { + "name": "echo", + "arguments": { + "context": "own-context", + "conversation_id": "own-id", + "llm_model": "own-model", + }, + }, + ) + assert not result.get("isError", False), result + assert result["content"][0]["text"] == expected + await flush_background() + calls = events_named(sink, "$mcp_tool_call") + assert len(calls) == 2 + assert "$mcp_llm_model" not in calls[0]["properties"] + assert calls[1]["properties"]["$mcp_llm_model"] == "own-model"