From e526f6c3b0569dfc5638ddf36b139b84bd6e663f Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Fri, 11 Sep 2026 15:48:51 -0300 Subject: [PATCH 1/7] chore(mcp): enable model capture and conversation correlation by default Enable both existing instrumentation features and model capture on PostHogMCP, with explicit opt-outs. Missing-capability reporting and feedback stay disabled. Resolve model argument ownership from bounded raw catalog lookups on fresh low-level instances. Do not emit synthetic tools/list events, alter application model arguments, or block dispatch when catalog resolution fails. Validation: 486 MCP v1 tests passed (1 skipped), 433 MCP v2 tests passed (21 skipped); Ruff lint/format, mypy baseline for 237 files, public API snapshot, warning-as-error import and wheel build passed. CodeScene gate passed. One feedback cursor deprecation warning remains outside the requested scope. Include a Sampo minor changeset and document changed defaults and opt-outs. --- .sampo/changesets/mcp-analytics-defaults.md | 5 ++ posthog/mcp/README.md | 27 ++++++- posthog/mcp/_instrument_lowlevel.py | 37 ++++++++- posthog/mcp/_instrument_v2.py | 55 +++++++++---- posthog/mcp/_tool_schema.py | 68 ++++++++++++++++ posthog/mcp/posthog_mcp.py | 2 +- posthog/mcp/types.py | 6 +- posthog/test/mcp/test_defaults.py | 90 +++++++++++++++++++++ posthog/test/mcp/test_lowlevel.py | 2 +- posthog/test/mcp/test_posthog_mcp.py | 2 +- posthog/test/mcp/test_session_token.py | 4 +- posthog/test/mcp/test_tool_schema.py | 71 ++++++++++++++++ posthog/test/mcp/test_units.py | 2 +- posthog/test/mcp/test_v2_lowlevel.py | 2 +- posthog/test/mcp/test_v2_mcpserver.py | 2 +- posthog/test/mcp/test_v2_wire_dual_era.py | 8 +- references/public_api_snapshot.txt | 8 +- 17 files changed, 350 insertions(+), 41 deletions(-) create mode 100644 .sampo/changesets/mcp-analytics-defaults.md create mode 100644 posthog/mcp/_tool_schema.py create mode 100644 posthog/test/mcp/test_defaults.py create mode 100644 posthog/test/mcp/test_tool_schema.py diff --git a/.sampo/changesets/mcp-analytics-defaults.md b/.sampo/changesets/mcp-analytics-defaults.md new file mode 100644 index 00000000..5bb909bf --- /dev/null +++ b/.sampo/changesets/mcp-analytics-defaults.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +Enable MCP model capture and conversation correlation by default, including on fresh low-level servers. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index ece6b576..aae31036 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -36,10 +36,24 @@ Request headers use the same identity and package version, so SDK Health can com Because `$lib` is a client-level identity, `instrument()` relabels every event sent by the client passed to it. Use a client dedicated to MCP analytics if the application also captures unrelated events. +## Defaults and opt-outs + +Intent, model capture, conversation correlation, and MCP exception capture are on by default. +Missing-capability reporting and feedback collection remain off. + +```python +instrument(server, posthog, MCPAnalyticsOptions(capture_model=False, enable_conversation_id=False)) +``` + +Conversation correlation adds an optional `conversation_id` argument and returns a handle in +eligible tool results. Clients must echo it to group later calls; calls without it mint new handles. +Set `enable_conversation_id=False` to retain transport-based session grouping and unchanged +response content. Custom `PostHogMCP` dispatchers enable model capture by default but still +supply their own session IDs. + ## Capture the calling model -Model capture is off by default. Enable it for an instrumented MCP Python SDK 1.x or -2.x server: +Model capture is on by default for instrumented MCP Python SDK 1.x and 2.x servers: ```python from posthog.mcp import MCPAnalyticsOptions, instrument @@ -47,7 +61,6 @@ from posthog.mcp import MCPAnalyticsOptions, instrument analytics = instrument( server, posthog, - MCPAnalyticsOptions(capture_model=True), ) ``` @@ -75,6 +88,12 @@ If a tool already declares `llm_model`, or uses a root `$ref`, `oneOf`, `allOf`, `anyOf` schema, PostHog leaves the schema and argument untouched. Client metadata can still be captured in those cases. +On fresh low-level instances, model argument ownership is resolved from the original tool +listing before dispatch. This internal lookup emits no discovery event and stops after 16 pages +or 250 ms. If the listing fails or omits the tool, its arguments remain unchanged and self-reported +model capture stays empty; recognized client metadata can still supply the model. Existing +listings and high-level registries continue to supply ownership directly. + For a custom dispatcher, use the same option on `PostHogMCP` and pass request metadata through explicitly: @@ -229,7 +248,7 @@ sess = get_mcp_session(request) # sess.session_id, sess.client_name, ... ### Or skip the middleware entirely: conversation ids -`MCPAnalyticsOptions(enable_conversation_id=True)` derives `$session_id` from the +`enable_conversation_id` is on by default and derives `$session_id` from the agent's conversation handle, deterministically and identically on every pod. That needs no middleware and no ordering discipline, and it is the only thing that correlates a session under the 2026-07-28 revision's per-request server instances. diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 78659793..18a84169 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -42,6 +42,7 @@ ) from ._internal import MCPAnalyticsData from ._model_parameters import request_meta_from_context +from ._tool_schema import resolve_model_ownership from ._output_instructions import mirror_instructions_into_structured_content from .logger import log from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name @@ -171,6 +172,31 @@ async def handler(req: Any) -> Any: handlers[request_type] = handler +async def _prepare_model_arguments( + server: Any, data: MCPAnalyticsData, req: Any, strip_injected: bool +) -> Tuple[Any, bool]: + async def list_page(cursor: Optional[str]) -> Any: + listing = server.request_handlers.get(mcp_types.ListToolsRequest) + raw = getattr(listing, "__posthog_mcp_original__", listing) + return await raw( + mcp_types.ListToolsRequest( + method="tools/list", + params=mcp_types.PaginatedRequestParams( + cursor=cursor, _meta=req.params.meta + ), + ) + ) + + owns_model = await resolve_model_ownership(data, req.params.name, list_page) + if strip_injected or not owns_model: + return req, owns_model + arguments = dict(req.params.arguments or {}) + arguments.pop("llm_model", None) + return req.model_copy( + update={"params": req.params.model_copy(update={"arguments": arguments})} + ), owns_model + + def _wrap_call_tool( server: Any, data: MCPAnalyticsData, *, strip_injected: bool, high_level: Any = None ) -> None: @@ -182,6 +208,10 @@ def _wrap_call_tool( async def handler(req: Any) -> Any: name = req.params.name arguments = dict(req.params.arguments or {}) + + req, analytics_owns_model = await _prepare_model_arguments( + server, data, req, strip_injected + ) client_name, client_version = _client_info(server) protocol_version = _protocol_version(server) mcp_session_id = _mcp_session_id(server) @@ -195,9 +225,7 @@ async def handler(req: Any) -> Any: name=name, arguments=arguments, request_meta=request_meta_from_context(_request_context(server)), - 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, @@ -240,7 +268,7 @@ async def handler(req: Any) -> Any: if strip_injected and req.params.arguments: owned = await _tool_owned_injected_keys(high_level, name) injected_keys = ["context", "conversation_id"] - if data.tool_model_parameter_injected.get(name, False): + if analytics_owns_model: injected_keys.append("llm_model") for key in injected_keys: if key not in owned: @@ -424,6 +452,7 @@ async def handler(req: Any) -> Any: return result + setattr(handler, "__posthog_mcp_original__", original) setattr(handler, _WRAPPED_FLAG, True) handlers[mcp_types.ListToolsRequest] = handler diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 66a28e77..dc788594 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -61,6 +61,7 @@ request_meta_from_context, ) from ._output_instructions import mirror_instructions_into_structured_content +from ._tool_schema import resolve_model_ownership from .logger import log from .request_headers import get_request_headers from .session_token import read_mcp_session_header @@ -477,6 +478,40 @@ async def _standalone_injected_parameters( return frozenset(key for key in injected if not schema_has_param(schema, key)) +async def _prepare_raw_v2_arguments( + server: Any, data: MCPAnalyticsData, ctx: Any, params: Any +) -> Tuple[Any, bool]: + async def list_page(cursor: Optional[str]) -> Any: + listing = server.get_request_handler(_LIST_METHOD) + raw = getattr(listing.handler, "__posthog_mcp_original__", listing.handler) + return await raw(ctx, mcp_types.PaginatedRequestParams(cursor=cursor)) + + owns_model = await resolve_model_ownership(data, params.name, list_page) + if not owns_model: + return params, False + arguments = dict(params.arguments or {}) + arguments.pop("llm_model", None) + return params.model_copy(update={"arguments": arguments}), True + + +async def _prepare_v2_arguments( + server: Any, data: MCPAnalyticsData, ctx: Any, params: Any +) -> Tuple[Any, bool]: + standalone = data.standalone_fastmcp() if data.standalone_fastmcp else None + if standalone is None: + return await _prepare_raw_v2_arguments(server, data, ctx, params) + version = _requested_tool_version(ctx) + injected = await _standalone_injected_parameters( + standalone, data, params.name, version + ) + if injected is None: + return params, False + arguments = dict(params.arguments or {}) + for key in injected: + arguments.pop(key, None) + return params.model_copy(update={"arguments": arguments}), "llm_model" in injected + + 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): @@ -486,21 +521,10 @@ 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) - 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( - standalone, data, name, version - ) - analytics_owns_model = injected is not None and "llm_model" in injected - if injected is not None: - call_arguments = { - key: value - for key, value in arguments.items() - if key not in injected - } - params = params.model_copy(update={"arguments": call_arguments}) + + params, analytics_owns_model = await _prepare_v2_arguments( + server, data, ctx, params + ) token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) ) @@ -715,6 +739,7 @@ async def handler(ctx: Any, params: Any) -> Any: return result + setattr(handler, "__posthog_mcp_original__", original) setattr(handler, _WRAPPED_FLAG, True) _replace_handler(server, _LIST_METHOD, handler, entry.params_type) diff --git a/posthog/mcp/_tool_schema.py b/posthog/mcp/_tool_schema.py new file mode 100644 index 00000000..1c7caac9 --- /dev/null +++ b/posthog/mcp/_tool_schema.py @@ -0,0 +1,68 @@ +"""Resolve model argument ownership on low-level servers without a tool registry.""" + +from __future__ import annotations + +import asyncio +from typing import Any, Awaitable, Callable, Optional + +from ._internal import MCPAnalyticsData +from ._model_parameters import can_inject_model_parameter, is_capture_model_enabled +from .logger import log + + +async def resolve_model_ownership( + data: MCPAnalyticsData, + name: str, + list_page: Callable[[Optional[str]], Awaitable[Any]], +) -> bool: + if not is_capture_model_enabled(data.options.capture_model): + return False + if name in data.tool_model_parameter_injected: + return data.tool_model_parameter_injected[name] + try: + return await asyncio.wait_for( + _find_model_ownership(name, list_page), timeout=0.25 + ) + except Exception: # noqa: BLE001 - discovery must not prevent tool dispatch + log( + "Warning: Could not resolve model argument ownership; leaving tool arguments unchanged." + ) + return False + + +async def _find_model_ownership( + name: str, list_page: Callable[[Optional[str]], Awaitable[Any]] +) -> bool: + cursor = None + seen = set() + for _ in range(16): + response = await list_page(cursor) + result = getattr(response, "root", response) + ownership = _model_ownership(result, name) + if ownership is not None: + return ownership + cursor = _next_cursor(result) + if cursor is None or cursor in seen: + return False + seen.add(cursor) + return False + + +def _model_ownership(result: Any, name: str) -> Optional[bool]: + for tool in getattr(result, "tools", []): + if getattr(tool, "name", None) == name: + schema = getattr(tool, "input_schema", None) + if schema is None: + schema = getattr(tool, "inputSchema", None) + return can_inject_model_parameter(schema) + return None + + +def _next_cursor(result: Any) -> Optional[str]: + if hasattr(result, "next_cursor"): + cursor = result.next_cursor + else: + cursor = getattr(result, "nextCursor", None) + if not isinstance(cursor, str): + return None + return cursor or None diff --git a/posthog/mcp/posthog_mcp.py b/posthog/mcp/posthog_mcp.py index 614b9615..a461e046 100644 --- a/posthog/mcp/posthog_mcp.py +++ b/posthog/mcp/posthog_mcp.py @@ -71,7 +71,7 @@ def __init__( api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, - capture_model: Union[bool, MCPAnalyticsModelOptions] = False, + capture_model: Union[bool, MCPAnalyticsModelOptions] = True, collect_feedback: Union[bool, CollectFeedbackOptions] = False, **kwargs: Any, ) -> None: diff --git a/posthog/mcp/types.py b/posthog/mcp/types.py index d2f3afdd..fb37a73e 100644 --- a/posthog/mcp/types.py +++ b/posthog/mcp/types.py @@ -183,7 +183,7 @@ class MCPAnalyticsOptions: logger: Optional[LoggerFn] = None report_missing: bool = False missing_capability_tool_name: Optional[str] = None - enable_conversation_id: bool = False + enable_conversation_id: bool = True enable_exception_autocapture: bool = True # Inject a required `context` parameter on every tool to capture user intent. context: Union[bool, MCPAnalyticsContextOptions] = True @@ -197,8 +197,8 @@ class MCPAnalyticsOptions: # Extra properties merged onto every auto-captured event. event_properties: Optional[EventPropertiesFn] = None # Capture the model from recognized client metadata, falling back to an - # SDK-injected llm_model argument. Off by default. - capture_model: Union[bool, MCPAnalyticsModelOptions] = False + # SDK-injected llm_model argument. On by default; False disables capture. + capture_model: Union[bool, MCPAnalyticsModelOptions] = True # Inject the `send_feedback` virtual tool so agents can send feedback about # this server to its developers — a missing capability (the priority # category), a tool that failed or confused them, or praise. Calls to it emit diff --git a/posthog/test/mcp/test_defaults.py b/posthog/test/mcp/test_defaults.py new file mode 100644 index 00000000..52341801 --- /dev/null +++ b/posthog/test/mcp/test_defaults.py @@ -0,0 +1,90 @@ +import pytest + +from posthog.mcp import PostHogMCP +from posthog.mcp.types import MCPAnalyticsOptions +from posthog.test.mcp._helpers import ( + MCP_MAJOR, + FakeClient, + events_named, + flush_background, +) + + +def test_model_capture_is_enabled_for_custom_dispatchers(): + for kwargs, enabled in [({}, True), ({"capture_model": False}, False)]: + client = PostHogMCP("test", disabled=True, **kwargs) + tools = client.prepare_tool_list( + [{"name": "echo", "inputSchema": {"type": "object", "properties": {}}}] + ) + assert ("llm_model" in tools[0]["inputSchema"]["properties"]) is enabled + client.shutdown() + + +@pytest.mark.skipif(MCP_MAJOR < 2, reason="v2 low-level handler API") +async def test_default_capture_on_fresh_lowlevel_instances(): + from posthog.mcp import instrument + from posthog.test.mcp.test_v2_lowlevel import make_server, _call_tool, _list_tools + + client = FakeClient() + + def fresh(): + server = make_server() + instrument(server, client) + return server + + listing = await _list_tools(fresh()) + assert [tool.name for tool in listing.tools] == ["add"] + assert {"context", "llm_model", "conversation_id"} <= set( + listing.tools[0].input_schema["properties"] + ) + await _call_tool( + fresh(), "add", {"a": 1, "b": 2, "context": "intent", "llm_model": "model-a"} + ) + await flush_background() + first = events_named(client, "$mcp_tool_call")[0]["properties"] + await _call_tool( + fresh(), + "add", + { + "a": 2, + "b": 3, + "context": "intent", + "llm_model": "model-a", + "conversation_id": first["$mcp_conversation_id"], + }, + ) + await flush_background() + calls = events_named(client, "$mcp_tool_call") + assert len(calls) == 2 + assert first["$mcp_llm_model"] == "model-a" + assert first["$session_id"] == calls[1]["properties"]["$session_id"] + assert len(events_named(client, "$mcp_tools_list")) == 1 + + +@pytest.mark.skipif(MCP_MAJOR >= 2, reason="v1 low-level handler API") +@pytest.mark.parametrize("enabled", [True, False]) +async def test_v1_cold_capture_and_opt_out(enabled): + import mcp.types as types + from posthog.mcp import instrument + from posthog.test.mcp.test_lowlevel import make_server, _call_request + + client = FakeClient() + server = make_server() + options = ( + None + if enabled + else MCPAnalyticsOptions(capture_model=False, enable_conversation_id=False) + ) + instrument(server, client, options) + request = _call_request("echo", {"msg": "ok", "llm_model": "model-a"}) + result = await server.request_handlers[types.CallToolRequest](request) + await flush_background() + calls = events_named(client, "$mcp_tool_call") + assert len(calls) == 1 + assert not calls[0]["properties"]["$mcp_is_error"] + assert calls[0]["properties"].get("$mcp_llm_model") == ( + "model-a" if enabled else None + ) + assert len(result.root.content) == (2 if enabled else 1) + assert len(events_named(client, "$mcp_tools_list")) == 0 + assert request.params.arguments == {"msg": "ok", "llm_model": "model-a"} diff --git a/posthog/test/mcp/test_lowlevel.py b/posthog/test/mcp/test_lowlevel.py index 992c558a..ac84541c 100644 --- a/posthog/test/mcp/test_lowlevel.py +++ b/posthog/test/mcp/test_lowlevel.py @@ -377,7 +377,7 @@ async def test_tool_call_error_captured_from_is_error_result(): async def test_initialize_emitted_once(): server = make_server() client = FakeClient() - instrument(server, client) + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=False)) handler = server.request_handlers[mcp_types.CallToolRequest] await handler(_call_request("echo", {"msg": "a", "context": "first call"})) diff --git a/posthog/test/mcp/test_posthog_mcp.py b/posthog/test/mcp/test_posthog_mcp.py index 7b9b89fc..9355e810 100644 --- a/posthog/test/mcp/test_posthog_mcp.py +++ b/posthog/test/mcp/test_posthog_mcp.py @@ -222,7 +222,7 @@ def test_prepare_tool_list_can_be_disabled(): ) async def test_prepare_and_capture_model(options: dict[str, bool]) -> None: client, captured = make_client(**options) - enabled = options.get("capture_model", False) + enabled = options.get("capture_model", True) tools = [ { "name": "search", diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index 10d79cba..4cb0ac81 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -645,7 +645,7 @@ def test_runtime_warns_when_app_was_built_before_instrument(): app = srv.streamable_http_app() # BEFORE instrument() -- the trap with _captured_logs() as logs: - instrument(srv, _Sink()) + instrument(srv, _Sink(), MCPAnalyticsOptions(enable_conversation_id=False)) with TestClient(app) as client: resp = _call_ping(client) assert resp.status_code == 200, resp.text @@ -716,7 +716,7 @@ def test_warnings_are_visible_without_configuring_a_logger(caplog): set_logger(None) # explicitly no `logger` option anywhere with caplog.at_level("WARNING", logger="posthog.mcp"): - instrument(srv, _Sink()) + instrument(srv, _Sink(), MCPAnalyticsOptions(enable_conversation_id=False)) with TestClient(app) as client: _call_ping(client) diff --git a/posthog/test/mcp/test_tool_schema.py b/posthog/test/mcp/test_tool_schema.py new file mode 100644 index 00000000..0c6df196 --- /dev/null +++ b/posthog/test/mcp/test_tool_schema.py @@ -0,0 +1,71 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from posthog.mcp._internal import MCPAnalyticsData +from posthog.mcp._tool_schema import resolve_model_ownership +from posthog.mcp.types import MCPAnalyticsOptions + + +def data(): + return MCPAnalyticsData(options=MCPAnalyticsOptions()) + + +@pytest.mark.parametrize("owned", [True, False]) +async def test_paginated_catalog_preserves_application_model(owned): + schema = { + "type": "object", + "properties": {"llm_model": {"type": "string"}} if owned else {}, + } + tool = SimpleNamespace(name="echo", input_schema=schema) + listing = AsyncMock( + side_effect=[ + SimpleNamespace(tools=[], next_cursor="next"), + SimpleNamespace(tools=[tool]), + ] + ) + assert await resolve_model_ownership(data(), "echo", listing) is not owned + assert [call.args for call in listing.call_args_list] == [(None,), ("next",)] + + +@pytest.mark.parametrize( + "mode,expected_calls", + [("cycle", 2), ("endless", 16), ("malformed", 1), ("error", 1)], +) +async def test_bounded_catalog_failures(mode, expected_calls): + calls = [] + + async def listing(cursor): + calls.append(cursor) + if mode == "error": + raise ValueError("unavailable") + if mode == "malformed": + return None + return SimpleNamespace( + tools=[], next_cursor="same" if mode == "cycle" else str(len(calls)) + ) + + assert await resolve_model_ownership(data(), "echo", listing) is False + assert len(calls) == expected_calls + + +async def test_slow_listing_does_not_block_dispatch(): + cancelled = asyncio.Event() + + async def listing(cursor): + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + assert await resolve_model_ownership(data(), "echo", listing) is False + assert cancelled.is_set() + + +async def test_opt_out_does_not_invoke_catalog(): + state = MCPAnalyticsData(options=MCPAnalyticsOptions(capture_model=False)) + listing = AsyncMock() + assert await resolve_model_ownership(state, "echo", listing) is False + listing.assert_not_called() diff --git a/posthog/test/mcp/test_units.py b/posthog/test/mcp/test_units.py index 998510df..e8312123 100644 --- a/posthog/test/mcp/test_units.py +++ b/posthog/test/mcp/test_units.py @@ -109,7 +109,7 @@ def test_schema_pipeline_does_not_warn_for_owned_conversation_id(monkeypatch): tool = SimpleNamespace(name="t", input_schema=schema) mutate_tool_schema( - _data(context=False, enable_conversation_id=True), + _data(context=False, capture_model=False, enable_conversation_id=True), tool, schema_attribute="input_schema", owns_context=False, diff --git a/posthog/test/mcp/test_v2_lowlevel.py b/posthog/test/mcp/test_v2_lowlevel.py index 4d53407e..e3ad65c7 100644 --- a/posthog/test/mcp/test_v2_lowlevel.py +++ b/posthog/test/mcp/test_v2_lowlevel.py @@ -453,7 +453,7 @@ async def late_read_resource(ctx, params): async def test_initialize_and_session_reuse_across_calls(): server = make_server() client = FakeClient() - instrument(server, client) + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=False)) await _call_tool(server, "add", {"a": 1, "b": 1, "context": "first"}) await _call_tool(server, "add", {"a": 2, "b": 2, "context": "second"}) diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index 80ae82ce..890e342c 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -175,7 +175,7 @@ async def test_tool_call_error_is_captured_and_converted(): async def test_initialize_emitted_once_per_session(): server = make_server() client = FakeClient() - instrument(server, client) + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=False)) await _call_tool( server, "add", {"a": 1, "b": 1, "context": "first call to warm up"} diff --git a/posthog/test/mcp/test_v2_wire_dual_era.py b/posthog/test/mcp/test_v2_wire_dual_era.py index 4b7c11b3..0e89b69b 100644 --- a/posthog/test/mcp/test_v2_wire_dual_era.py +++ b/posthog/test/mcp/test_v2_wire_dual_era.py @@ -227,7 +227,9 @@ async def test_modern_result_shape_survives_instrumentation(): bare_response = await modern_call(http, "add", {"a": 4, "b": 5}) instrumented = make_server() - instrument(instrumented, FakeClient()) + instrument( + instrumented, FakeClient(), MCPAnalyticsOptions(enable_conversation_id=False) + ) async with wire(instrumented) as http: response = await modern_call(http, "add", {"a": 4, "b": 5, "context": "alive"}) await _flush() @@ -295,7 +297,7 @@ async def test_legacy_stateless_token_survives_across_instances(): # Pod A mints the token on initialize. server_a = make_server() - instrument(server_a, client) + instrument(server_a, client, MCPAnalyticsOptions(enable_conversation_id=False)) async with wire(server_a) as http: init = rpc( "initialize", @@ -316,7 +318,7 @@ async def test_legacy_stateless_token_survives_across_instances(): # A compliant legacy client replays both the session header and the # negotiated MCP-Protocol-Version on every subsequent request. server_b = make_server() - instrument(server_b, client) + instrument(server_b, client, MCPAnalyticsOptions(enable_conversation_id=False)) async with wire(server_b) as http: headers = { **legacy_headers(), diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index bf7211fb..843003ab 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -818,10 +818,10 @@ attribute posthog.mcp.types.MCPAnalyticsContextOptions.description: Optional[str attribute posthog.mcp.types.MCPAnalyticsModelOptions.description: Optional[str] = None attribute posthog.mcp.types.MCPAnalyticsModelSource = Literal['client_metadata', 'self_reported'] attribute posthog.mcp.types.MCPAnalyticsOptions.before_send: Optional[BeforeSendFn] = None -attribute posthog.mcp.types.MCPAnalyticsOptions.capture_model: Union[bool, MCPAnalyticsModelOptions] = False +attribute posthog.mcp.types.MCPAnalyticsOptions.capture_model: Union[bool, MCPAnalyticsModelOptions] = True attribute posthog.mcp.types.MCPAnalyticsOptions.collect_feedback: Union[bool, CollectFeedbackOptions] = False attribute posthog.mcp.types.MCPAnalyticsOptions.context: Union[bool, MCPAnalyticsContextOptions] = True -attribute posthog.mcp.types.MCPAnalyticsOptions.enable_conversation_id: bool = False +attribute posthog.mcp.types.MCPAnalyticsOptions.enable_conversation_id: bool = True attribute posthog.mcp.types.MCPAnalyticsOptions.enable_exception_autocapture: bool = True attribute posthog.mcp.types.MCPAnalyticsOptions.event_properties: Optional[EventPropertiesFn] = None attribute posthog.mcp.types.MCPAnalyticsOptions.identify: Optional[Union[IdentifyFn, UserIdentity]] = None @@ -1021,14 +1021,14 @@ class posthog.mcp.McpAnalytics(key: Any) class posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware(app: Any) class posthog.mcp.constants.PostHogMCPAnalyticsEvent class posthog.mcp.constants.PostHogMCPAnalyticsProperty -class posthog.mcp.posthog_mcp.PostHogMCP(api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, capture_model: Union[bool, MCPAnalyticsModelOptions] = False, collect_feedback: Union[bool, CollectFeedbackOptions] = False, **kwargs: Any) +class posthog.mcp.posthog_mcp.PostHogMCP(api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, capture_model: Union[bool, MCPAnalyticsModelOptions] = True, collect_feedback: Union[bool, CollectFeedbackOptions] = False, **kwargs: Any) class posthog.mcp.session_token.SessionTokenPayload(session_id: str, client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None) class posthog.mcp.types.CaptureEventData(event: str, properties: Optional[JsonRecord] = None) class posthog.mcp.types.CollectFeedbackOptions(tool_name: Optional[str] = None, description: Optional[str] = None, extra_properties: Optional[Dict[str, Dict[str, Any]]] = None, extra_required: Optional[List[str]] = None, on_feedback: Optional[OnFeedbackFn] = None) class posthog.mcp.types.FeedbackReport(feedback_type: str = 'other', summary: str = '', sentiment: Optional[str] = None, friction_points: Optional[str] = None, suggested_improvement: Optional[str] = None, details: Optional[str] = None, tool_name: Optional[str] = None, task_completed: Optional[bool] = None, extras: JsonRecord = dict(), raw: JsonRecord = dict()) class posthog.mcp.types.MCPAnalyticsContextOptions(description: Optional[str] = None) class posthog.mcp.types.MCPAnalyticsModelOptions(description: Optional[str] = None) -class posthog.mcp.types.MCPAnalyticsOptions(logger: Optional[LoggerFn] = None, report_missing: bool = False, missing_capability_tool_name: Optional[str] = None, enable_conversation_id: bool = False, enable_exception_autocapture: bool = True, context: Union[bool, MCPAnalyticsContextOptions] = True, identify: Optional[Union[IdentifyFn, UserIdentity]] = None, intent_fallback: Optional[IntentFallbackFn] = None, before_send: Optional[BeforeSendFn] = None, event_properties: Optional[EventPropertiesFn] = None, capture_model: Union[bool, MCPAnalyticsModelOptions] = False, collect_feedback: Union[bool, CollectFeedbackOptions] = False) +class posthog.mcp.types.MCPAnalyticsOptions(logger: Optional[LoggerFn] = None, report_missing: bool = False, missing_capability_tool_name: Optional[str] = None, enable_conversation_id: bool = True, enable_exception_autocapture: bool = True, context: Union[bool, MCPAnalyticsContextOptions] = True, identify: Optional[Union[IdentifyFn, UserIdentity]] = None, intent_fallback: Optional[IntentFallbackFn] = None, before_send: Optional[BeforeSendFn] = None, event_properties: Optional[EventPropertiesFn] = None, capture_model: Union[bool, MCPAnalyticsModelOptions] = True, collect_feedback: Union[bool, CollectFeedbackOptions] = False) class posthog.mcp.types.PreparedToolCall(args: Optional[JsonRecord] = None, intent: Optional[str] = None, intent_source: Optional[str] = None, is_missing_capability: bool = False, llm_model: Optional[str] = None, llm_model_source: Optional[MCPAnalyticsModelSource] = None, is_feedback: bool = False, feedback_report: Optional[FeedbackReport] = None) class posthog.mcp.types.UserIdentity(distinct_id: str, properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None) class posthog.metrics_capture.PostHogMetrics(client, config: Optional[dict] = None) From ec5b6de979f4b8b00c9c08957f1c2dc776000038 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Fri, 11 Sep 2026 15:56:15 -0300 Subject: [PATCH 2/7] test(mcp): reuse custom-dispatcher default coverage Remove the duplicate custom-dispatcher default and opt-out test. The existing parameterized prepare-and-capture test already covers both. Keep the fresh-instance regressions and bounded catalog lookup coverage. Validation: MCP v1: 485 passed, 1 skipped. MCP v2: 432 passed, 21 skipped. Ruff lint/format and CodeScene pre-commit checks passed. The existing feedback nextCursor deprecation warning remains outside this change. Runtime behavior is unchanged. --- posthog/test/mcp/test_defaults.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/posthog/test/mcp/test_defaults.py b/posthog/test/mcp/test_defaults.py index 52341801..24270b46 100644 --- a/posthog/test/mcp/test_defaults.py +++ b/posthog/test/mcp/test_defaults.py @@ -1,6 +1,5 @@ import pytest -from posthog.mcp import PostHogMCP from posthog.mcp.types import MCPAnalyticsOptions from posthog.test.mcp._helpers import ( MCP_MAJOR, @@ -10,16 +9,6 @@ ) -def test_model_capture_is_enabled_for_custom_dispatchers(): - for kwargs, enabled in [({}, True), ({"capture_model": False}, False)]: - client = PostHogMCP("test", disabled=True, **kwargs) - tools = client.prepare_tool_list( - [{"name": "echo", "inputSchema": {"type": "object", "properties": {}}}] - ) - assert ("llm_model" in tools[0]["inputSchema"]["properties"]) is enabled - client.shutdown() - - @pytest.mark.skipif(MCP_MAJOR < 2, reason="v2 low-level handler API") async def test_default_capture_on_fresh_lowlevel_instances(): from posthog.mcp import instrument From 80a8dd27991e48e65212f589cf16d7f0b16d4d06 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Mon, 14 Sep 2026 13:51:38 -0300 Subject: [PATCH 3/7] fix(mcp): cache confirmed low-level model ownership Store both confirmed ownership outcomes in the existing model cache. Return None internally for unresolved catalogs so missing tools, page limits, errors and timeouts do not become permanent negative entries. Extended existing pagination/failure tests and added late-tool recovery. Caching regressions failed before the fix. Validation: MCP v1 486 passed, 1 skipped; MCP v2 433 passed, 21 skipped. Ruff, mypy baseline (237 files), and CodeScene pre-commit checks passed. The feedback nextCursor warning remains outside this change. Addresses review discussion_r3992430558 on #944. --- posthog/mcp/_tool_schema.py | 11 +++++++---- posthog/test/mcp/test_tool_schema.py | 26 ++++++++++++++++++++++---- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/posthog/mcp/_tool_schema.py b/posthog/mcp/_tool_schema.py index 1c7caac9..b6aa046c 100644 --- a/posthog/mcp/_tool_schema.py +++ b/posthog/mcp/_tool_schema.py @@ -20,9 +20,12 @@ async def resolve_model_ownership( if name in data.tool_model_parameter_injected: return data.tool_model_parameter_injected[name] try: - return await asyncio.wait_for( + ownership = await asyncio.wait_for( _find_model_ownership(name, list_page), timeout=0.25 ) + if ownership is not None: + data.tool_model_parameter_injected[name] = ownership + return ownership if ownership is not None else False except Exception: # noqa: BLE001 - discovery must not prevent tool dispatch log( "Warning: Could not resolve model argument ownership; leaving tool arguments unchanged." @@ -32,7 +35,7 @@ async def resolve_model_ownership( async def _find_model_ownership( name: str, list_page: Callable[[Optional[str]], Awaitable[Any]] -) -> bool: +) -> Optional[bool]: cursor = None seen = set() for _ in range(16): @@ -43,9 +46,9 @@ async def _find_model_ownership( return ownership cursor = _next_cursor(result) if cursor is None or cursor in seen: - return False + return None seen.add(cursor) - return False + return None def _model_ownership(result: Any, name: str) -> Optional[bool]: diff --git a/posthog/test/mcp/test_tool_schema.py b/posthog/test/mcp/test_tool_schema.py index 0c6df196..f45831cf 100644 --- a/posthog/test/mcp/test_tool_schema.py +++ b/posthog/test/mcp/test_tool_schema.py @@ -14,7 +14,7 @@ def data(): @pytest.mark.parametrize("owned", [True, False]) -async def test_paginated_catalog_preserves_application_model(owned): +async def test_paginated_catalog_caches_confirmed_model_ownership(owned): schema = { "type": "object", "properties": {"llm_model": {"type": "string"}} if owned else {}, @@ -26,10 +26,24 @@ async def test_paginated_catalog_preserves_application_model(owned): SimpleNamespace(tools=[tool]), ] ) - assert await resolve_model_ownership(data(), "echo", listing) is not owned + state = data() + for _ in range(2): + assert await resolve_model_ownership(state, "echo", listing) is not owned assert [call.args for call in listing.call_args_list] == [(None,), ("next",)] +async def test_missing_tool_is_retried_when_it_appears(): + tool = SimpleNamespace(name="echo", input_schema={"type": "object"}) + listing = AsyncMock( + side_effect=[SimpleNamespace(tools=[]), SimpleNamespace(tools=[tool])] + ) + state = data() + assert await resolve_model_ownership(state, "echo", listing) is False + assert "echo" not in state.tool_model_parameter_injected + assert await resolve_model_ownership(state, "echo", listing) is True + assert listing.call_count == 2 + + @pytest.mark.parametrize( "mode,expected_calls", [("cycle", 2), ("endless", 16), ("malformed", 1), ("error", 1)], @@ -47,7 +61,9 @@ async def listing(cursor): tools=[], next_cursor="same" if mode == "cycle" else str(len(calls)) ) - assert await resolve_model_ownership(data(), "echo", listing) is False + state = data() + assert await resolve_model_ownership(state, "echo", listing) is False + assert "echo" not in state.tool_model_parameter_injected assert len(calls) == expected_calls @@ -60,7 +76,9 @@ async def listing(cursor): finally: cancelled.set() - assert await resolve_model_ownership(data(), "echo", listing) is False + state = data() + assert await resolve_model_ownership(state, "echo", listing) is False + assert "echo" not in state.tool_model_parameter_injected assert cancelled.is_set() From ea63cf1d298fcfbe38e47f8b4f77ea12b850f095 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 15 Sep 2026 11:31:24 -0300 Subject: [PATCH 4/7] fix(mcp): read llm_model on cold instances under the ADR-0011 rule Replaces the raw-catalog replay added earlier on this branch with the ownership rule posthog-js ADR-0011 applies to `context`: a fresh raw low-level instance reads the self-reported model and strips nothing, and the SDK never replays the host's tools/list handler on the call path. Standalone FastMCP reads ownership of all three injected keys from the tool signature, so it strips llm_model without a prior listing. Tested: .venv pytest posthog/test/mcp (478 passed, 1 skipped), .venv-mcp-v2 (424 passed, 21 skipped), ruff check/format, mypy baseline (no issues), public API snapshot up to date. Reviewer notes: raw low-level servers keep main's behaviour of never stripping llm_model; posthog-js strips on positive ownership. That pre-existing difference is unchanged here. Claude-Session: https://claude.ai/code/session_012YBJXRzzCizpHFGn2mZaEg --- .sampo/changesets/mcp-analytics-defaults.md | 2 +- posthog/mcp/README.md | 34 ++++---- posthog/mcp/_instrument_lowlevel.py | 72 +++++++---------- posthog/mcp/_instrument_v2.py | 58 +++++--------- posthog/mcp/_tool_schema.py | 71 ---------------- posthog/test/mcp/test_defaults.py | 12 ++- posthog/test/mcp/test_fastmcp_v2.py | 18 +++++ posthog/test/mcp/test_tool_schema.py | 89 --------------------- 8 files changed, 90 insertions(+), 266 deletions(-) delete mode 100644 posthog/mcp/_tool_schema.py delete mode 100644 posthog/test/mcp/test_tool_schema.py diff --git a/.sampo/changesets/mcp-analytics-defaults.md b/.sampo/changesets/mcp-analytics-defaults.md index 5bb909bf..7dbb9b5d 100644 --- a/.sampo/changesets/mcp-analytics-defaults.md +++ b/.sampo/changesets/mcp-analytics-defaults.md @@ -2,4 +2,4 @@ pypi/posthog: minor --- -Enable MCP model capture and conversation correlation by default, including on fresh low-level servers. +Enable MCP model capture and conversation correlation by default. Advertised tool schemas gain an `llm_model` argument (never enforced at dispatch) and eligible tool results gain a conversation handle; `MCPAnalyticsOptions(capture_model=False, enable_conversation_id=False)` restores the previous shape. Fresh low-level instances now read the self-reported model instead of staying silent. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index aae31036..8352bbcb 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -42,28 +42,23 @@ Intent, model capture, conversation correlation, and MCP exception capture are o Missing-capability reporting and feedback collection remain off. ```python +from posthog.mcp import MCPAnalyticsOptions, instrument + instrument(server, posthog, MCPAnalyticsOptions(capture_model=False, enable_conversation_id=False)) ``` +Model capture adds an `llm_model` argument to compatible tool schemas, required on the official +high-level adapters and optional elsewhere. Dispatch never enforces it, so servers keep working; +strict-schema clients see the new field. Set `capture_model=False` to leave schemas untouched. Conversation correlation adds an optional `conversation_id` argument and returns a handle in eligible tool results. Clients must echo it to group later calls; calls without it mint new handles. Set `enable_conversation_id=False` to retain transport-based session grouping and unchanged response content. Custom `PostHogMCP` dispatchers enable model capture by default but still -supply their own session IDs. +supply their own session IDs. The reasoning is recorded in posthog-js `docs/adr/0013`. ## Capture the calling model -Model capture is on by default for instrumented MCP Python SDK 1.x and 2.x servers: - -```python -from posthog.mcp import MCPAnalyticsOptions, instrument - -analytics = instrument( - server, - posthog, -) -``` - +Model capture is on by default for instrumented MCP Python SDK 1.x and 2.x servers. The SDK records the best model identifier visible to the server as `$mcp_llm_model`. Recognized client metadata wins and sets `$mcp_llm_model_source` to `client_metadata`. The SDK also adds an `llm_model` @@ -88,19 +83,20 @@ If a tool already declares `llm_model`, or uses a root `$ref`, `oneOf`, `allOf`, `anyOf` schema, PostHog leaves the schema and argument untouched. Client metadata can still be captured in those cases. -On fresh low-level instances, model argument ownership is resolved from the original tool -listing before dispatch. This internal lookup emits no discovery event and stops after 16 pages -or 250 ms. If the listing fails or omits the tool, its arguments remain unchanged and self-reported -model capture stays empty; recognized client metadata can still supply the model. Existing -listings and high-level registries continue to supply ownership directly. +A raw low-level server learns ownership while serving `tools/list`. A fresh instance that never +served one has no answer, so it records `llm_model` as the self-reported model and strips +nothing (posthog-js ADR-0011: reads fail open, strips fail closed). A tool that declares its own +`llm_model` on such an instance is therefore recorded under `$mcp_llm_model` until a listing says +otherwise; `capture_model=False` or `before_send` are the escapes. High-level adapters and +standalone `fastmcp.FastMCP` read ownership from their live registry, so they are unaffected. -For a custom dispatcher, use the same option on `PostHogMCP` and pass request +For a custom dispatcher, `PostHogMCP` enables the same option by default; pass request metadata through explicitly: ```python from posthog.mcp import PostHogMCP -posthog = PostHogMCP("phc_...", capture_model=True) +posthog = PostHogMCP("phc_...") tools = posthog.prepare_tool_list(server_tools) original_tool = next(tool for tool in server_tools if tool["name"] == tool_name) call = posthog.prepare_tool_call( diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 18a84169..85fc796b 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -42,7 +42,6 @@ ) from ._internal import MCPAnalyticsData from ._model_parameters import request_meta_from_context -from ._tool_schema import resolve_model_ownership from ._output_instructions import mirror_instructions_into_structured_content from .logger import log from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name @@ -172,31 +171,6 @@ async def handler(req: Any) -> Any: handlers[request_type] = handler -async def _prepare_model_arguments( - server: Any, data: MCPAnalyticsData, req: Any, strip_injected: bool -) -> Tuple[Any, bool]: - async def list_page(cursor: Optional[str]) -> Any: - listing = server.request_handlers.get(mcp_types.ListToolsRequest) - raw = getattr(listing, "__posthog_mcp_original__", listing) - return await raw( - mcp_types.ListToolsRequest( - method="tools/list", - params=mcp_types.PaginatedRequestParams( - cursor=cursor, _meta=req.params.meta - ), - ) - ) - - owns_model = await resolve_model_ownership(data, req.params.name, list_page) - if strip_injected or not owns_model: - return req, owns_model - arguments = dict(req.params.arguments or {}) - arguments.pop("llm_model", None) - return req.model_copy( - update={"params": req.params.model_copy(update={"arguments": arguments})} - ), owns_model - - def _wrap_call_tool( server: Any, data: MCPAnalyticsData, *, strip_injected: bool, high_level: Any = None ) -> None: @@ -208,10 +182,7 @@ def _wrap_call_tool( async def handler(req: Any) -> Any: name = req.params.name arguments = dict(req.params.arguments or {}) - - req, analytics_owns_model = await _prepare_model_arguments( - server, data, req, strip_injected - ) + owned = await _tool_owned_injected_keys(high_level, name) client_name, client_version = _client_info(server) protocol_version = _protocol_version(server) mcp_session_id = _mcp_session_id(server) @@ -225,7 +196,7 @@ async def handler(req: Any) -> Any: name=name, arguments=arguments, request_meta=request_meta_from_context(_request_context(server)), - allow_self_reported_model=analytics_owns_model, + allow_self_reported_model=_analytics_reads_model(data, name, owned), mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -266,12 +237,8 @@ async def handler(req: Any) -> Any: # is read from the tool's own signature, so it holds with or without a prior # tools/list and across stateless per-request server instances. if strip_injected and req.params.arguments: - owned = await _tool_owned_injected_keys(high_level, name) - injected_keys = ["context", "conversation_id"] - if analytics_owns_model: - injected_keys.append("llm_model") - for key in injected_keys: - if key not in owned: + for key in _INJECTED_KEYS: + if key not in (owned or set()): req.params.arguments.pop(key, None) # Settle the shared session before the tool body runs, so an in-tool @@ -452,7 +419,6 @@ async def handler(req: Any) -> Any: return result - setattr(handler, "__posthog_mcp_original__", original) setattr(handler, _WRAPPED_FLAG, True) handlers[mcp_types.ListToolsRequest] = handler @@ -470,22 +436,38 @@ async def _feedback_name_owned_by_real_tool(high_level: Any, name: str) -> bool: return False -async def _tool_owned_injected_keys(high_level: Any, name: str) -> set: - """Which of (``context``, ``conversation_id``) the jlowin FastMCP tool declares - itself, read from its function signature. These are real tool arguments we must - not strip. On any lookup failure, return empty (strip both) — same as the prior - unconditional behaviour, so a flaky introspection never leaks an injected key.""" +_INJECTED_KEYS = ("context", "conversation_id", "llm_model") + + +async def _tool_owned_injected_keys(high_level: Any, name: str) -> Optional[set]: + """Which of ``_INJECTED_KEYS`` the jlowin FastMCP tool declares itself, read + from its function signature. These are real tool arguments we must not strip. + On any lookup failure, return empty (strip all) — same as the prior + unconditional behaviour, so a flaky introspection never leaks an injected key. + ``None`` means there is no registry to ask: raw low-level servers.""" if high_level is None: - return set() + return None try: tool = await high_level.get_tool(name) fn = getattr(tool, "fn", None) params = set(inspect.signature(fn).parameters) if fn is not None else set() - return {k for k in ("context", "conversation_id") if k in params} + return {k for k in _INJECTED_KEYS if k in params} except Exception: # noqa: BLE001 - introspection is best-effort return set() +def _analytics_reads_model( + data: MCPAnalyticsData, name: str, owned: Optional[set] +) -> bool: + """Whether ``llm_model`` in the arguments is read as the self-reported model. + The registry answers for standalone FastMCP. Raw servers only learn ownership + while serving ``tools/list``; a cold per-request instance has no answer and + reads anyway. Reading fails open, stripping fails closed (posthog-js ADR-0011).""" + if owned is not None: + return "llm_model" not in owned + return data.tool_model_parameter_injected.get(name) is not False + + def _request_context(server: Any) -> Any: try: return server.request_context diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index dc788594..844ca974 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -61,7 +61,6 @@ request_meta_from_context, ) from ._output_instructions import mirror_instructions_into_structured_content -from ._tool_schema import resolve_model_ownership from .logger import log from .request_headers import get_request_headers from .session_token import read_mcp_session_header @@ -478,40 +477,6 @@ async def _standalone_injected_parameters( return frozenset(key for key in injected if not schema_has_param(schema, key)) -async def _prepare_raw_v2_arguments( - server: Any, data: MCPAnalyticsData, ctx: Any, params: Any -) -> Tuple[Any, bool]: - async def list_page(cursor: Optional[str]) -> Any: - listing = server.get_request_handler(_LIST_METHOD) - raw = getattr(listing.handler, "__posthog_mcp_original__", listing.handler) - return await raw(ctx, mcp_types.PaginatedRequestParams(cursor=cursor)) - - owns_model = await resolve_model_ownership(data, params.name, list_page) - if not owns_model: - return params, False - arguments = dict(params.arguments or {}) - arguments.pop("llm_model", None) - return params.model_copy(update={"arguments": arguments}), True - - -async def _prepare_v2_arguments( - server: Any, data: MCPAnalyticsData, ctx: Any, params: Any -) -> Tuple[Any, bool]: - standalone = data.standalone_fastmcp() if data.standalone_fastmcp else None - if standalone is None: - return await _prepare_raw_v2_arguments(server, data, ctx, params) - version = _requested_tool_version(ctx) - injected = await _standalone_injected_parameters( - standalone, data, params.name, version - ) - if injected is None: - return params, False - arguments = dict(params.arguments or {}) - for key in injected: - arguments.pop(key, None) - return params.model_copy(update={"arguments": arguments}), "llm_model" in injected - - 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): @@ -521,10 +486,24 @@ 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 {}) - - params, analytics_owns_model = await _prepare_v2_arguments( - server, data, ctx, params - ) + # A raw instance that never served a listing has no ownership answer and + # 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 + 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( + 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) ) @@ -739,7 +718,6 @@ async def handler(ctx: Any, params: Any) -> Any: return result - setattr(handler, "__posthog_mcp_original__", original) setattr(handler, _WRAPPED_FLAG, True) _replace_handler(server, _LIST_METHOD, handler, entry.params_type) diff --git a/posthog/mcp/_tool_schema.py b/posthog/mcp/_tool_schema.py deleted file mode 100644 index b6aa046c..00000000 --- a/posthog/mcp/_tool_schema.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Resolve model argument ownership on low-level servers without a tool registry.""" - -from __future__ import annotations - -import asyncio -from typing import Any, Awaitable, Callable, Optional - -from ._internal import MCPAnalyticsData -from ._model_parameters import can_inject_model_parameter, is_capture_model_enabled -from .logger import log - - -async def resolve_model_ownership( - data: MCPAnalyticsData, - name: str, - list_page: Callable[[Optional[str]], Awaitable[Any]], -) -> bool: - if not is_capture_model_enabled(data.options.capture_model): - return False - if name in data.tool_model_parameter_injected: - return data.tool_model_parameter_injected[name] - try: - ownership = await asyncio.wait_for( - _find_model_ownership(name, list_page), timeout=0.25 - ) - if ownership is not None: - data.tool_model_parameter_injected[name] = ownership - return ownership if ownership is not None else False - except Exception: # noqa: BLE001 - discovery must not prevent tool dispatch - log( - "Warning: Could not resolve model argument ownership; leaving tool arguments unchanged." - ) - return False - - -async def _find_model_ownership( - name: str, list_page: Callable[[Optional[str]], Awaitable[Any]] -) -> Optional[bool]: - cursor = None - seen = set() - for _ in range(16): - response = await list_page(cursor) - result = getattr(response, "root", response) - ownership = _model_ownership(result, name) - if ownership is not None: - return ownership - cursor = _next_cursor(result) - if cursor is None or cursor in seen: - return None - seen.add(cursor) - return None - - -def _model_ownership(result: Any, name: str) -> Optional[bool]: - for tool in getattr(result, "tools", []): - if getattr(tool, "name", None) == name: - schema = getattr(tool, "input_schema", None) - if schema is None: - schema = getattr(tool, "inputSchema", None) - return can_inject_model_parameter(schema) - return None - - -def _next_cursor(result: Any) -> Optional[str]: - if hasattr(result, "next_cursor"): - cursor = result.next_cursor - else: - cursor = getattr(result, "nextCursor", None) - if not isinstance(cursor, str): - return None - return cursor or None diff --git a/posthog/test/mcp/test_defaults.py b/posthog/test/mcp/test_defaults.py index 24270b46..8a7a70a8 100644 --- a/posthog/test/mcp/test_defaults.py +++ b/posthog/test/mcp/test_defaults.py @@ -64,6 +64,14 @@ async def test_v1_cold_capture_and_opt_out(enabled): if enabled else MCPAnalyticsOptions(capture_model=False, enable_conversation_id=False) ) + received = [] + original = server.request_handlers[types.CallToolRequest] + + async def spy(req): + received.append(dict(req.params.arguments or {})) + return await original(req) + + server.request_handlers[types.CallToolRequest] = spy instrument(server, client, options) request = _call_request("echo", {"msg": "ok", "llm_model": "model-a"}) result = await server.request_handlers[types.CallToolRequest](request) @@ -76,4 +84,6 @@ async def test_v1_cold_capture_and_opt_out(enabled): ) assert len(result.root.content) == (2 if enabled else 1) assert len(events_named(client, "$mcp_tools_list")) == 0 - assert request.params.arguments == {"msg": "ok", "llm_model": "model-a"} + # A cold raw instance reads the self-reported model but strips nothing: it + # cannot prove it owns the argument, and raw handlers ignore extra keys. + assert received == [{"msg": "ok", "llm_model": "model-a"}] diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index 9ebc1145..5cf7fe66 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -151,3 +151,21 @@ def echo(msg: str) -> str: assert _events(client, "$mcp_tool_call")[0]["properties"]["$mcp_intent"] == ( "strict validation" ) + + +async def test_jlowin_cold_call_strips_llm_model_and_records_it(): + # No prior tools/list: ownership comes from the tool signature, so the + # injected llm_model is stripped before jlowin validates and still recorded. + server = make_server() + client = FakeClient() + instrument(server, client) + + out = await _call( + server, "add", {"a": 2, "b": 3, "context": "sum", "llm_model": "model-a"} + ) + await _flush() + + assert out.root.isError is False + calls = _events(client, "$mcp_tool_call") + assert calls[0]["properties"]["$mcp_llm_model"] == "model-a" + assert not _events(client, "$mcp_tools_list") diff --git a/posthog/test/mcp/test_tool_schema.py b/posthog/test/mcp/test_tool_schema.py deleted file mode 100644 index f45831cf..00000000 --- a/posthog/test/mcp/test_tool_schema.py +++ /dev/null @@ -1,89 +0,0 @@ -import asyncio -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest - -from posthog.mcp._internal import MCPAnalyticsData -from posthog.mcp._tool_schema import resolve_model_ownership -from posthog.mcp.types import MCPAnalyticsOptions - - -def data(): - return MCPAnalyticsData(options=MCPAnalyticsOptions()) - - -@pytest.mark.parametrize("owned", [True, False]) -async def test_paginated_catalog_caches_confirmed_model_ownership(owned): - schema = { - "type": "object", - "properties": {"llm_model": {"type": "string"}} if owned else {}, - } - tool = SimpleNamespace(name="echo", input_schema=schema) - listing = AsyncMock( - side_effect=[ - SimpleNamespace(tools=[], next_cursor="next"), - SimpleNamespace(tools=[tool]), - ] - ) - state = data() - for _ in range(2): - assert await resolve_model_ownership(state, "echo", listing) is not owned - assert [call.args for call in listing.call_args_list] == [(None,), ("next",)] - - -async def test_missing_tool_is_retried_when_it_appears(): - tool = SimpleNamespace(name="echo", input_schema={"type": "object"}) - listing = AsyncMock( - side_effect=[SimpleNamespace(tools=[]), SimpleNamespace(tools=[tool])] - ) - state = data() - assert await resolve_model_ownership(state, "echo", listing) is False - assert "echo" not in state.tool_model_parameter_injected - assert await resolve_model_ownership(state, "echo", listing) is True - assert listing.call_count == 2 - - -@pytest.mark.parametrize( - "mode,expected_calls", - [("cycle", 2), ("endless", 16), ("malformed", 1), ("error", 1)], -) -async def test_bounded_catalog_failures(mode, expected_calls): - calls = [] - - async def listing(cursor): - calls.append(cursor) - if mode == "error": - raise ValueError("unavailable") - if mode == "malformed": - return None - return SimpleNamespace( - tools=[], next_cursor="same" if mode == "cycle" else str(len(calls)) - ) - - state = data() - assert await resolve_model_ownership(state, "echo", listing) is False - assert "echo" not in state.tool_model_parameter_injected - assert len(calls) == expected_calls - - -async def test_slow_listing_does_not_block_dispatch(): - cancelled = asyncio.Event() - - async def listing(cursor): - try: - await asyncio.Event().wait() - finally: - cancelled.set() - - state = data() - assert await resolve_model_ownership(state, "echo", listing) is False - assert "echo" not in state.tool_model_parameter_injected - assert cancelled.is_set() - - -async def test_opt_out_does_not_invoke_catalog(): - state = MCPAnalyticsData(options=MCPAnalyticsOptions(capture_model=False)) - listing = AsyncMock() - assert await resolve_model_ownership(state, "echo", listing) is False - listing.assert_not_called() From 4d426f1155c884bc2fb8b4af605c70d4f701c2d5 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 15 Sep 2026 11:35:21 -0300 Subject: [PATCH 5/7] fix(mcp): read standalone FastMCP ownership from the advertised schema A jlowin `Tool` subclass declares its arguments in `parameters` and may have no `fn`, so the signature-based check returned nothing owned and the strip loop deleted an application-declared `llm_model`. Ownership now comes from the advertised schema first, matching the v2 standalone path, with the signature as fallback. Tested: .venv pytest posthog/test/mcp (480 passed, 1 skipped), .venv-mcp-v2 (424 passed, 21 skipped), ruff, mypy baseline. Found by `codex review --base main` on the previous commit. Claude-Session: https://claude.ai/code/session_012YBJXRzzCizpHFGn2mZaEg --- posthog/mcp/_instrument_lowlevel.py | 13 +++++++--- posthog/test/mcp/test_fastmcp_v2.py | 39 +++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 85fc796b..e37ca465 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -441,14 +441,19 @@ async def _feedback_name_owned_by_real_tool(high_level: Any, name: str) -> bool: async def _tool_owned_injected_keys(high_level: Any, name: str) -> Optional[set]: """Which of ``_INJECTED_KEYS`` the jlowin FastMCP tool declares itself, read - from its function signature. These are real tool arguments we must not strip. - On any lookup failure, return empty (strip all) — same as the prior - unconditional behaviour, so a flaky introspection never leaks an injected key. - ``None`` means there is no registry to ask: raw low-level servers.""" + from its advertised ``parameters`` schema (a ``Tool`` subclass may have no + function), else from its function signature. These are real tool arguments + we must not strip. On any lookup failure, return empty (strip all) — same as + the prior unconditional behaviour, so a flaky introspection never leaks an + injected key. ``None`` means there is no registry to ask: raw low-level + servers.""" if high_level is None: return None try: tool = await high_level.get_tool(name) + schema = getattr(tool, "parameters", None) + if isinstance(schema, dict): + return {k for k in _INJECTED_KEYS if schema_has_param(schema, k)} fn = getattr(tool, "fn", None) params = set(inspect.signature(fn).parameters) if fn is not None else set() return {k for k in _INJECTED_KEYS if k in params} diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index 5cf7fe66..b4b89268 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -169,3 +169,42 @@ async def test_jlowin_cold_call_strips_llm_model_and_records_it(): calls = _events(client, "$mcp_tool_call") assert calls[0]["properties"]["$mcp_llm_model"] == "model-a" assert not _events(client, "$mcp_tools_list") + + +@pytest.mark.parametrize("capture_model", [True, False]) +async def test_jlowin_schema_declared_llm_model_is_kept_and_not_misreported( + capture_model, +): + # A Tool subclass declares its arguments in `parameters` and has no `fn`. + # Its own llm_model must reach the tool and must not be read as the model. + from fastmcp.tools import Tool + from fastmcp.tools.tool import ToolResult + + class Router(Tool): + async def run(self, arguments): + return ToolResult( + content=[ + mcp_types.TextContent(type="text", text=arguments["llm_model"]) + ] + ) + + server = FastMCP("jlowin-schema-owner") + server.add_tool( + Router( + name="route", + parameters={ + "type": "object", + "properties": {"llm_model": {"type": "string"}}, + "required": ["llm_model"], + }, + ) + ) + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(capture_model=capture_model)) + + out = await _call(server, "route", {"llm_model": "gpt-5", "context": "routing"}) + await _flush() + + assert out.root.isError is False + assert out.root.content[0].text == "gpt-5" + assert "$mcp_llm_model" not in _events(client, "$mcp_tool_call")[0]["properties"] From cf885624326298536b32107cebdc7bac537bb248 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 15 Sep 2026 11:40:02 -0300 Subject: [PATCH 6/7] fix(mcp): share standalone FastMCP ownership across SDK majors The v1 and v2 standalone adapters resolved injected-argument ownership two different ways, and the v1 one stripped an application-declared `llm_model` inside a composed schema (`allOf`, `$ref`). Lift the v2 helper into `_standalone.py`, add the per-key composed-schema guard the listing-time injection already applies, and use it from both paths. Lookup failure now strips nothing on v1 too (fail closed), matching v2. Tested: .venv pytest posthog/test/mcp (482 passed, 1 skipped), .venv-mcp-v2 (424 passed, 21 skipped), ruff, mypy baseline, public API snapshot. Found by `codex review --base main` on the previous commit. Claude-Session: https://claude.ai/code/session_012YBJXRzzCizpHFGn2mZaEg --- posthog/mcp/_instrument_lowlevel.py | 50 +++++++------------------ posthog/mcp/_instrument_v2.py | 47 ++---------------------- posthog/mcp/_standalone.py | 57 +++++++++++++++++++++++++++++ posthog/test/mcp/test_fastmcp_v2.py | 31 ++++++++++------ 4 files changed, 94 insertions(+), 91 deletions(-) create mode 100644 posthog/mcp/_standalone.py diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index e37ca465..6d2ef99e 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -14,9 +14,8 @@ from __future__ import annotations -import inspect import time -from typing import Any, Optional, Tuple +from typing import Any, FrozenSet, Optional, Tuple import mcp.types as mcp_types @@ -43,6 +42,7 @@ from ._internal import MCPAnalyticsData from ._model_parameters import request_meta_from_context from ._output_instructions import mirror_instructions_into_structured_content +from ._standalone import standalone_injected_parameters from .logger import log from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name @@ -182,7 +182,11 @@ def _wrap_call_tool( async def handler(req: Any) -> Any: name = req.params.name arguments = dict(req.params.arguments or {}) - owned = await _tool_owned_injected_keys(high_level, name) + injected = ( + await standalone_injected_parameters(high_level, data, name, None) + if high_level is not None + else None + ) client_name, client_version = _client_info(server) protocol_version = _protocol_version(server) mcp_session_id = _mcp_session_id(server) @@ -196,7 +200,7 @@ async def handler(req: Any) -> Any: name=name, arguments=arguments, request_meta=request_meta_from_context(_request_context(server)), - allow_self_reported_model=_analytics_reads_model(data, name, owned), + allow_self_reported_model=_analytics_reads_model(data, name, injected), mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -236,10 +240,9 @@ async def handler(req: Any) -> Any: # but NOT a key the tool declares itself (that's a real argument). Ownership # is read from the tool's own signature, so it holds with or without a prior # tools/list and across stateless per-request server instances. - if strip_injected and req.params.arguments: - for key in _INJECTED_KEYS: - if key not in (owned or set()): - req.params.arguments.pop(key, None) + if strip_injected and injected and req.params.arguments: + for key in injected: + req.params.arguments.pop(key, None) # Settle the shared session before the tool body runs, so an in-tool # `analytics.capture()` is attributed to this caller and not the last one. @@ -436,40 +439,15 @@ async def _feedback_name_owned_by_real_tool(high_level: Any, name: str) -> bool: return False -_INJECTED_KEYS = ("context", "conversation_id", "llm_model") - - -async def _tool_owned_injected_keys(high_level: Any, name: str) -> Optional[set]: - """Which of ``_INJECTED_KEYS`` the jlowin FastMCP tool declares itself, read - from its advertised ``parameters`` schema (a ``Tool`` subclass may have no - function), else from its function signature. These are real tool arguments - we must not strip. On any lookup failure, return empty (strip all) — same as - the prior unconditional behaviour, so a flaky introspection never leaks an - injected key. ``None`` means there is no registry to ask: raw low-level - servers.""" - if high_level is None: - return None - try: - tool = await high_level.get_tool(name) - schema = getattr(tool, "parameters", None) - if isinstance(schema, dict): - return {k for k in _INJECTED_KEYS if schema_has_param(schema, k)} - fn = getattr(tool, "fn", None) - params = set(inspect.signature(fn).parameters) if fn is not None else set() - return {k for k in _INJECTED_KEYS if k in params} - except Exception: # noqa: BLE001 - introspection is best-effort - return set() - - def _analytics_reads_model( - data: MCPAnalyticsData, name: str, owned: Optional[set] + data: MCPAnalyticsData, name: str, injected: Optional[FrozenSet[str]] ) -> bool: """Whether ``llm_model`` in the arguments is read as the self-reported model. The registry answers for standalone FastMCP. Raw servers only learn ownership while serving ``tools/list``; a cold per-request instance has no answer and reads anyway. Reading fails open, stripping fails closed (posthog-js ADR-0011).""" - if owned is not None: - return "llm_model" not in owned + if injected is not None: + return "llm_model" in injected return data.tool_model_parameter_injected.get(name) is not False diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 844ca974..446d6be7 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -32,11 +32,11 @@ import time from collections.abc import Mapping -from typing import Any, Dict, FrozenSet, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import mcp.types as mcp_types -from ._context_parameters import is_context_enabled, schema_has_param +from ._context_parameters import schema_has_param from ._conversation_id import build_prompt_back from ._event_types import MCPAnalyticsEventType from ._instrumentation import ( @@ -61,6 +61,7 @@ request_meta_from_context, ) from ._output_instructions import mirror_instructions_into_structured_content +from ._standalone import standalone_injected_parameters from .logger import log from .request_headers import get_request_headers from .session_token import read_mcp_session_header @@ -437,46 +438,6 @@ def _requested_tool_version(ctx: Any) -> Optional[str]: return None -async def _standalone_injected_parameters( - server: Any, data: MCPAnalyticsData, name: str, version: Optional[str] -) -> Optional[FrozenSet[str]]: - """Resolve ownership in the current request, including middleware and versions. - - 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 = 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 ( - 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: entry = server.get_request_handler(_CALL_METHOD) if entry is None or getattr(entry.handler, _WRAPPED_FLAG, False): @@ -493,7 +454,7 @@ async def handler(ctx: Any, params: Any) -> Any: 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 = await standalone_injected_parameters( standalone, data, name, version ) if injected is not None: diff --git a/posthog/mcp/_standalone.py b/posthog/mcp/_standalone.py new file mode 100644 index 00000000..70de464a --- /dev/null +++ b/posthog/mcp/_standalone.py @@ -0,0 +1,57 @@ +"""Argument ownership for jlowin's standalone ``fastmcp.FastMCP`` (2.0+), shared by +the MCP SDK v1 and v2 adapters. It keeps a live registry, so ownership is read +from the tool's advertised schema per request instead of a prior listing.""" + +from __future__ import annotations + +from typing import Any, FrozenSet, Optional + +from ._context_parameters import is_context_enabled, schema_has_param +from ._internal import MCPAnalyticsData +from ._model_parameters import can_inject_model_parameter, is_capture_model_enabled +from .logger import log + + +async def standalone_injected_parameters( + server: Any, data: MCPAnalyticsData, name: str, version: Optional[str] +) -> Optional[FrozenSet[str]]: + """The analytics arguments the SDK injected into this tool's schema, which + are the only ones safe to strip before FastMCP validates the call. Mirrors + the listing-time injection rules: nothing is injected into a composed + schema, and ``llm_model`` follows ``can_inject_model_parameter``. + + Listings from other requests can have different application-owned + parameters (middleware, versions), so this resolves in the current request. + ``None`` means the schema could not be read; callers then strip nothing. + """ + 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 = set() + if not any(schema.get(key) for key in ("oneOf", "allOf", "anyOf")): + 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 ( + can_inject_model_parameter(schema) + ): + injected.add("llm_model") + return frozenset(key for key in injected if not schema_has_param(schema, key)) diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index b4b89268..6bcb175b 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -171,12 +171,28 @@ async def test_jlowin_cold_call_strips_llm_model_and_records_it(): assert not _events(client, "$mcp_tools_list") +_OWN_MODEL = {"llm_model": {"type": "string"}} + + @pytest.mark.parametrize("capture_model", [True, False]) +@pytest.mark.parametrize( + "parameters", + [ + {"type": "object", "properties": _OWN_MODEL, "required": ["llm_model"]}, + { + "allOf": [ + {"type": "object", "properties": _OWN_MODEL, "required": ["llm_model"]} + ] + }, + ], + ids=["properties", "allOf"], +) async def test_jlowin_schema_declared_llm_model_is_kept_and_not_misreported( - capture_model, + capture_model, parameters ): # A Tool subclass declares its arguments in `parameters` and has no `fn`. - # Its own llm_model must reach the tool and must not be read as the model. + # Its own llm_model must reach the tool and must not be read as the model, + # whether it sits in top-level properties or inside a composed schema. from fastmcp.tools import Tool from fastmcp.tools.tool import ToolResult @@ -189,16 +205,7 @@ async def run(self, arguments): ) server = FastMCP("jlowin-schema-owner") - server.add_tool( - Router( - name="route", - parameters={ - "type": "object", - "properties": {"llm_model": {"type": "string"}}, - "required": ["llm_model"], - }, - ) - ) + server.add_tool(Router(name="route", parameters=parameters)) client = FakeClient() instrument(server, client, MCPAnalyticsOptions(capture_model=capture_model)) From 5f20b75badb03b07939ceb2e52f89307b6bff005 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 15 Sep 2026 15:16:45 -0300 Subject: [PATCH 7/7] fix(mcp): strip every undeclared analytics key on standalone FastMCP from its schema With model capture on by default, a standalone FastMCP replica that never served the listing which advertised `llm_model` forwarded it to FastMCP's validator and the call failed. FastMCP rejects any undeclared argument, so stripping an analytics key the tool does not declare can never hurt, while stripping a declared one always does. The v1 standalone path now strips `llm_model` like `context` and `conversation_id`: unless the registered schema (or, without one, the function signature) declares it, with nothing stripped from a composed schema because nothing was injected into one. The registry is read directly, never through middleware, so rate limiters are not charged and no listing is needed. A local root `$ref` is dereferenced first, as FastMCP's built-in middleware does before the client sees the listing, and a pinned `_meta.fastmcp.version` is honoured only where FastMCP's own dispatch honours it, so ownership always follows the version that runs. Only keys the SDK injects under the current options are ever stripped, so a disabled feature leaves its key to the application, and sibling properties beside a root `$ref` count as declared, as does every node along a reference chain. For `llm_model` the effective listing is the first witness, because middleware can provide or shadow the tool the registry knows; the registry is second; with neither the argument stays and is still read. The registry is not trusted for `llm_model` while application middleware can change the listing or reroute dispatch, the server's own dereferencing setting decides whether a root `$ref` would have been injected into, and a listing that advertises two tools under one name (FastMCP 2.x with a shadowing middleware) marks the model argument as the application's. Application subclasses of FastMCP's built-in middleware count as the application's, and `$ref` segments decode JSON Pointer escapes. The model is read exactly when it was stripped; raw low-level servers keep reading it fail-open on unknown ownership (posthog-js ADR-0011). This replaces the listing-based resolver tried earlier on this branch, which five Codex rounds showed diverging from the advertised listing across FastMCP 2.5-4.0 and middleware combinations. Tested: .venv pytest posthog/test/mcp (502 passed, 1 skipped), .venv-mcp-v2 (424 passed, 21 skipped), a throwaway venv with fastmcp==2.14.5 (standalone and defaults tests: 28 passed, 5 skipped), and one with fastmcp==2.8.1, which has no middleware module (9 passed, 23 skipped; the one failure is a pre-existing test passing a constructor kwarg that release lacks), ruff, mypy baseline, public API snapshot. Claude-Session: https://claude.ai/code/session_012YBJXRzzCizpHFGn2mZaEg --- posthog/mcp/README.md | 3 +- posthog/mcp/_instrument_lowlevel.py | 234 ++++++++++++++++-- posthog/mcp/_instrument_v2.py | 47 +++- posthog/mcp/_standalone.py | 57 ----- posthog/test/mcp/test_fastmcp_v2.py | 364 +++++++++++++++++++++++++++- 5 files changed, 614 insertions(+), 91 deletions(-) delete mode 100644 posthog/mcp/_standalone.py diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 8352bbcb..bc3c68ef 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -88,7 +88,8 @@ served one has no answer, so it records `llm_model` as the self-reported model a nothing (posthog-js ADR-0011: reads fail open, strips fail closed). A tool that declares its own `llm_model` on such an instance is therefore recorded under `$mcp_llm_model` until a listing says otherwise; `capture_model=False` or `before_send` are the escapes. High-level adapters and -standalone `fastmcp.FastMCP` read ownership from their live registry, so they are unaffected. +standalone `fastmcp.FastMCP` read ownership from the registered tool schema, so they are +unaffected and need no prior listing. For a custom dispatcher, `PostHogMCP` enables the same option by default; pass request metadata through explicitly: diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 6d2ef99e..0a050046 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -14,12 +14,14 @@ from __future__ import annotations +import functools +import inspect import time -from typing import Any, FrozenSet, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple 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 ._event_types import MCPAnalyticsEventType from ._instrumentation import ( @@ -41,8 +43,8 @@ ) from ._internal import MCPAnalyticsData from ._model_parameters import request_meta_from_context +from ._model_parameters import can_inject_model_parameter, is_capture_model_enabled from ._output_instructions import mirror_instructions_into_structured_content -from ._standalone import standalone_injected_parameters from .logger import log from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name @@ -182,10 +184,10 @@ def _wrap_call_tool( async def handler(req: Any) -> Any: name = req.params.name arguments = dict(req.params.arguments or {}) - injected = ( - await standalone_injected_parameters(high_level, data, name, None) - if high_level is not None - else None + strip, model_ours = ( + await _standalone_ownership(data, high_level, name, req.params.meta) + if strip_injected + else (set(), data.tool_model_parameter_injected.get(name)) ) client_name, client_version = _client_info(server) protocol_version = _protocol_version(server) @@ -200,7 +202,8 @@ async def handler(req: Any) -> Any: name=name, arguments=arguments, request_meta=request_meta_from_context(_request_context(server)), - allow_self_reported_model=_analytics_reads_model(data, name, injected), + # Reads fail open on unknown ownership (posthog-js ADR-0011). + allow_self_reported_model=model_ours is not False, mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -240,8 +243,8 @@ async def handler(req: Any) -> Any: # but NOT a key the tool declares itself (that's a real argument). Ownership # is read from the tool's own signature, so it holds with or without a prior # tools/list and across stateless per-request server instances. - if strip_injected and injected and req.params.arguments: - for key in injected: + if strip and req.params.arguments: + for key in strip: req.params.arguments.pop(key, None) # Settle the shared session before the tool body runs, so an in-tool @@ -323,6 +326,7 @@ def _inject_tool_schemas( population pass, so the schema the SDK validates against always matches the one we advertised — see the note in ``handler``. """ + verdicts: Dict[str, bool] = {} for tool in tools: schema = getattr(tool, "inputSchema", None) mutate_tool_schema( @@ -332,6 +336,14 @@ def _inject_tool_schemas( owns_context=schema_has_param(schema, "context"), context_required=context_required, ) + verdict = data.tool_model_parameter_injected.get(tool.name) + if verdict is None: + continue + if verdicts.setdefault(tool.name, verdict) != verdict: + # Two advertised tools share this name and disagree (FastMCP 2.x + # lists a middleware tool beside the registered one it shadows). + # Which one dispatches is unknown, so the strip fails closed. + data.tool_model_parameter_injected[tool.name] = False def _wrap_list_tools( @@ -439,16 +451,198 @@ async def _feedback_name_owned_by_real_tool(high_level: Any, name: str) -> bool: return False -def _analytics_reads_model( - data: MCPAnalyticsData, name: str, injected: Optional[FrozenSet[str]] -) -> bool: - """Whether ``llm_model`` in the arguments is read as the self-reported model. - The registry answers for standalone FastMCP. Raw servers only learn ownership - while serving ``tools/list``; a cold per-request instance has no answer and - reads anyway. Reading fails open, stripping fails closed (posthog-js ADR-0011).""" - if injected is not None: - return "llm_model" in injected - return data.tool_model_parameter_injected.get(name) is not False +_INJECTED_KEYS = ("context", "conversation_id", "llm_model") + + +async def _standalone_ownership( + data: MCPAnalyticsData, high_level: Any, name: str, meta: Any +) -> Tuple[set, Optional[bool]]: + """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). + + Only keys injected under the current options are candidates. ``context`` + and ``conversation_id`` are stripped unless the registered schema (or, + without one, the function signature) declares them; a failed lookup strips + both — the prior behaviour, so a flaky introspection never leaks an injected + key into validation. ``llm_model`` is judged by the effective listing first, + because middleware can provide or shadow the tool the registry knows, then + by the registry; with neither witness it stays and is still read — strips + fail closed, reads fail open (posthog-js ADR-0011). + """ + try: + declared, model_injectable = await _registry_view(high_level, name, meta) + registry_trusted = model_injectable is not None and not _dispatch_can_differ( + high_level + ) + except Exception: # noqa: BLE001 - ownership inference must never prevent dispatch + declared, model_injectable, registry_trusted = None, None, False + listed = data.tool_model_parameter_injected.get(name) + if listed is not None: + model_ours: Optional[bool] = listed + elif registry_trusted: + model_ours = model_injectable + else: + model_ours = 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 + + +def _injected_keys(data: MCPAnalyticsData) -> set: + """The analytics arguments the SDK injects under the current options — the + only ones it may strip. A disabled feature injects nothing, so its key is + the application's even when the schema does not declare it.""" + keys = set() + if is_context_enabled(data.options.context): + keys.add("context") + if data.options.enable_conversation_id: + keys.add("conversation_id") + if is_capture_model_enabled(data.options.capture_model): + keys.add("llm_model") + return keys + + +async def _registry_view( + high_level: Any, name: str, meta: Any +) -> Tuple[Optional[set], Optional[bool]]: + """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.""" + try: + tool = await _registered_tool(high_level, name, meta) + except Exception: # noqa: BLE001 - introspection is best-effort + 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 + + +async def _registered_tool(high_level: Any, name: str, meta: Any) -> Any: + """The tool version the request pinned in ``_meta.fastmcp.version``, else + the highest one — the same choice FastMCP makes when dispatching.""" + version = _requested_tool_version(meta) + if version is None: + return await high_level.get_tool(name) + from fastmcp.utilities.versions import VersionSpec + + return await high_level.get_tool(name, version=VersionSpec(eq=version)) + + +def _requested_tool_version(meta: Any) -> Optional[str]: + """Ownership must follow dispatch. Only a FastMCP that exposes the + ``_meta`` version extractor its own dispatch uses honours a pinned + version; every earlier release calls the highest version regardless.""" + try: + from fastmcp.server.dependencies import extract_version_spec + except ImportError: + return None + dump = getattr(meta, "model_dump", None) + if callable(dump): + meta = dump(by_alias=True) + return extract_version_spec(meta) if isinstance(meta, dict) else None + + +def _schema_view(schema: Dict[str, Any], *, dereferenced: bool) -> Tuple[set, bool]: + """The injected keys a schema declares as its own, and whether a listing + would inject ``llm_model`` into it. With dereferencing on, every node along + a root ``$ref`` chain counts, as the client sees the merged schema; with it + off the client sees the reference itself, which the listing never injects + into. Nothing is injected into a composed or unresolvable schema either, so + all keys count as declared there and nothing is stripped.""" + nodes = _reference_chain(schema) if dereferenced else [schema] + if nodes is None or any( + node.get(key) for node in nodes for key in ("oneOf", "allOf", "anyOf") + ): + return set(_INJECTED_KEYS), False + declared = { + key + for node in nodes + if isinstance(node.get("properties"), dict) + for key in node["properties"] + if key in _INJECTED_KEYS + } + injectable = can_inject_model_parameter(nodes[-1]) and "llm_model" not in declared + return declared, injectable + + +def _server_dereferences(server: Any) -> bool: + """Whether this FastMCP dereferences schemas before advertising them (its + built-in middleware, on by default from 3.x; absent on 2.x).""" + return any( + type(middleware).__module__.endswith(".dereference") + for middleware in getattr(server, "middleware", ()) + ) + + +_DISPATCH_HOOKS = ("on_message", "on_request", "on_list_tools", "on_call_tool") + + +def _dispatch_can_differ(server: Any) -> bool: + """Whether application middleware can provide, shadow, or reroute a tool, + making the registry an unreliable witness for what actually runs. Any + override of a listing or dispatch hook can; FastMCP's own built-ins are + excluded. Without a trustworthy registry a cold instance keeps ``llm_model`` + (strips fail closed), which is exactly what it did before this path.""" + try: + from fastmcp.server.middleware import Middleware + except ImportError: # FastMCP before middleware existed: nothing can differ + return False + + return any( + type(middleware) not in _builtin_middleware_types() # subclasses are the app's + and any( + getattr(type(middleware), hook) is not getattr(Middleware, hook) + for hook in _DISPATCH_HOOKS + ) + for middleware in getattr(server, "middleware", ()) + ) + + +@functools.lru_cache(maxsize=1) +def _builtin_middleware_types() -> tuple: + """The middleware a bare ``FastMCP()`` installs on its own, probed rather + than named so a new built-in in a later release is still recognised.""" + from fastmcp import FastMCP + + return tuple(type(middleware) for middleware in FastMCP("posthog-probe").middleware) + + +def _reference_chain(schema: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: + """The schema and each local root ``$ref`` target in turn; ``None`` for a + reference that is external, dangling, or cyclic.""" + nodes = [schema] + while len(nodes) <= 8: + ref = nodes[-1].get("$ref") + if ref is None: + return nodes + if not isinstance(ref, str) or not ref.startswith("#/"): + return None + target: Any = schema + for part in ref[2:].split("/"): + key = part.replace("~1", "/").replace("~0", "~") # JSON Pointer escapes + target = target.get(key) if isinstance(target, dict) else None + if not isinstance(target, dict) or any(target is node for node in nodes): + return None + nodes.append(target) + return None def _request_context(server: Any) -> Any: diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 446d6be7..844ca974 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -32,11 +32,11 @@ 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 -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 ._event_types import MCPAnalyticsEventType from ._instrumentation import ( @@ -61,7 +61,6 @@ request_meta_from_context, ) from ._output_instructions import mirror_instructions_into_structured_content -from ._standalone import standalone_injected_parameters from .logger import log from .request_headers import get_request_headers from .session_token import read_mcp_session_header @@ -438,6 +437,46 @@ def _requested_tool_version(ctx: Any) -> Optional[str]: return None +async def _standalone_injected_parameters( + server: Any, data: MCPAnalyticsData, name: str, version: Optional[str] +) -> Optional[FrozenSet[str]]: + """Resolve ownership in the current request, including middleware and versions. + + 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 = 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 ( + 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: entry = server.get_request_handler(_CALL_METHOD) if entry is None or getattr(entry.handler, _WRAPPED_FLAG, False): @@ -454,7 +493,7 @@ async def handler(ctx: Any, params: Any) -> Any: 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 = await _standalone_injected_parameters( standalone, data, name, version ) if injected is not None: diff --git a/posthog/mcp/_standalone.py b/posthog/mcp/_standalone.py deleted file mode 100644 index 70de464a..00000000 --- a/posthog/mcp/_standalone.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Argument ownership for jlowin's standalone ``fastmcp.FastMCP`` (2.0+), shared by -the MCP SDK v1 and v2 adapters. It keeps a live registry, so ownership is read -from the tool's advertised schema per request instead of a prior listing.""" - -from __future__ import annotations - -from typing import Any, FrozenSet, Optional - -from ._context_parameters import is_context_enabled, schema_has_param -from ._internal import MCPAnalyticsData -from ._model_parameters import can_inject_model_parameter, is_capture_model_enabled -from .logger import log - - -async def standalone_injected_parameters( - server: Any, data: MCPAnalyticsData, name: str, version: Optional[str] -) -> Optional[FrozenSet[str]]: - """The analytics arguments the SDK injected into this tool's schema, which - are the only ones safe to strip before FastMCP validates the call. Mirrors - the listing-time injection rules: nothing is injected into a composed - schema, and ``llm_model`` follows ``can_inject_model_parameter``. - - Listings from other requests can have different application-owned - parameters (middleware, versions), so this resolves in the current request. - ``None`` means the schema could not be read; callers then strip nothing. - """ - 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 = set() - if not any(schema.get(key) for key in ("oneOf", "allOf", "anyOf")): - 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 ( - can_inject_model_parameter(schema) - ): - injected.add("llm_model") - return frozenset(key for key in injected if not schema_has_param(schema, key)) diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index 6bcb175b..f8ed5307 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -1,6 +1,8 @@ """Tests for jlowin's standalone FastMCP 2.0 (the `fastmcp` package), distinct from the official SDK's mcp.server.fastmcp.FastMCP.""" +import inspect + import pytest pytest.importorskip("fastmcp") @@ -27,6 +29,14 @@ def add(a: int, b: int) -> int: return server +def _tool_result_type(): + try: + from fastmcp.tools.tool import ToolResult + except ImportError: + pytest.skip("this FastMCP predates ToolResult") + return ToolResult + + async def _list(server): handler = server._mcp_server.request_handlers[mcp_types.ListToolsRequest] return await handler(mcp_types.ListToolsRequest(method="tools/list")) @@ -153,13 +163,17 @@ def echo(msg: str) -> str: ) -async def test_jlowin_cold_call_strips_llm_model_and_records_it(): - # No prior tools/list: ownership comes from the tool signature, so the - # injected llm_model is stripped before jlowin validates and still recorded. +@pytest.mark.parametrize("listed", [True, False], ids=["listed", "cold"]) +async def test_jlowin_call_strips_llm_model_and_records_it(listed): + # Ownership comes from the registered schema, so a cold instance (one that + # never served the listing that advertised llm_model, as in a multi-replica + # deployment) strips it before jlowin validates and still records it. server = make_server() client = FakeClient() instrument(server, client) + if listed: + await _list(server) out = await _call( server, "add", {"a": 2, "b": 3, "context": "sum", "llm_model": "model-a"} ) @@ -168,12 +182,12 @@ async def test_jlowin_cold_call_strips_llm_model_and_records_it(): assert out.root.isError is False calls = _events(client, "$mcp_tool_call") assert calls[0]["properties"]["$mcp_llm_model"] == "model-a" - assert not _events(client, "$mcp_tools_list") _OWN_MODEL = {"llm_model": {"type": "string"}} +@pytest.mark.parametrize("listed", [True, False], ids=["listed", "cold"]) @pytest.mark.parametrize("capture_model", [True, False]) @pytest.mark.parametrize( "parameters", @@ -187,14 +201,15 @@ async def test_jlowin_cold_call_strips_llm_model_and_records_it(): ], ids=["properties", "allOf"], ) -async def test_jlowin_schema_declared_llm_model_is_kept_and_not_misreported( - capture_model, parameters +async def test_jlowin_schema_declared_llm_model_is_kept( + capture_model, parameters, listed ): # A Tool subclass declares its arguments in `parameters` and has no `fn`. - # Its own llm_model must reach the tool and must not be read as the model, - # whether it sits in top-level properties or inside a composed schema. + # Its own llm_model must reach the tool and must never be read as the + # model, with or without a prior listing. from fastmcp.tools import Tool - from fastmcp.tools.tool import ToolResult + + ToolResult = _tool_result_type() class Router(Tool): async def run(self, arguments): @@ -209,9 +224,340 @@ async def run(self, arguments): client = FakeClient() instrument(server, client, MCPAnalyticsOptions(capture_model=capture_model)) + if listed: + await _list(server) out = await _call(server, "route", {"llm_model": "gpt-5", "context": "routing"}) await _flush() assert out.root.isError is False assert out.root.content[0].text == "gpt-5" assert "$mcp_llm_model" not in _events(client, "$mcp_tool_call")[0]["properties"] + + +@pytest.mark.parametrize( + "declares_model", + ["referenced", "sibling", "chained", "escaped", None], + ids=["referenced", "sibling", "chained", "escaped", "plain"], +) +async def test_jlowin_root_ref_schema_is_dereferenced_for_ownership(declares_model): + # FastMCP dereferences a root `$ref` before the client sees the listing, so + # ownership must be read from every node along the reference chain and from + # the root's own sibling properties: a key declared in any of them stays, + # and an undeclared llm_model is stripped so validation passes. + from fastmcp.tools import Tool + + ToolResult = _tool_result_type() + + properties = {"a": {"type": "integer"}} + extra = {} + if declares_model == "referenced": + properties["llm_model"] = {"type": "string"} + if declares_model == "sibling": + extra = {"properties": {"llm_model": {"type": "string"}}} + if declares_model == "escaped": + # A JSON Pointer escapes "/" as "~1"; FastMCP resolves it, so must we. + extra = { + "$ref": "#/$defs/Args~1Input", + "$defs": { + "Args/Input": { + "type": "object", + "properties": {**properties, "llm_model": {"type": "string"}}, + } + }, + } + if declares_model == "chained": + extra = { + "$defs": { + "Args": {"$ref": "#/$defs/Real"}, + "Real": { + "type": "object", + "properties": {**properties, "llm_model": {"type": "string"}}, + }, + } + } + + class Echo(Tool): + async def run(self, arguments): + return ToolResult( + content=[ + mcp_types.TextContent(type="text", text=",".join(sorted(arguments))) + ] + ) + + options = {} + if declares_model == "chained": + # FastMCP's own dereferencer recurses forever on a chained reference, + # instrumented or not; the chain case only exists with it turned off. + if "dereference_schemas" not in inspect.signature(FastMCP.__init__).parameters: + pytest.skip("this FastMCP cannot disable schema dereferencing") + options = {"dereference_schemas": False} + server = FastMCP("jlowin-root-ref", **options) + server.add_tool( + Echo( + name="echo", + parameters={ + "$ref": "#/$defs/Args", + "$defs": {"Args": {"type": "object", "properties": properties}}, + **extra, + }, + ) + ) + client = FakeClient() + instrument(server, client) + + out = await _call(server, "echo", {"a": 1, "context": "c", "llm_model": "gpt-5"}) + await _flush() + + assert out.root.isError is False + # A FastMCP that never dereferences advertises the `$ref` itself, which the + # listing never injects into, so the plain case keeps the argument there. + dereferences = ( + "dereference_schemas" in inspect.signature(FastMCP.__init__).parameters + ) + stripped = declares_model is None and dereferences + assert out.root.content[0].text == ("a" if stripped else "a,llm_model") + recorded = _events(client, "$mcp_tool_call")[0]["properties"].get("$mcp_llm_model") + assert recorded == ("gpt-5" if stripped else None) + + +async def test_jlowin_ownership_follows_the_dispatched_tool_version(): + # A request may pin a version in `_meta`; only a FastMCP that exposes + # `extract_version_spec` honours it when dispatching, every other release + # calls the highest version. Ownership must strip for the version that + # actually runs, or the call fails validation. + pytest.importorskip("fastmcp.utilities.versions") + try: + from fastmcp.server.dependencies import extract_version_spec # noqa: F401 + + expected = "v1:gpt-5" + except ImportError: + expected = "v2" + server = FastMCP("jlowin-versions") + + @server.tool(name="route", version="1") + def route_v1(prompt: str, llm_model: str) -> str: + return f"v1:{llm_model}" + + @server.tool(name="route", version="2") + def route_v2(prompt: str) -> str: + return "v2" + + client = FakeClient() + instrument(server, client) + handler = server._mcp_server.request_handlers[mcp_types.CallToolRequest] + + out = await handler( + mcp_types.CallToolRequest( + method="tools/call", + params=mcp_types.CallToolRequestParams( + name="route", + arguments={"prompt": "p", "context": "c", "llm_model": "gpt-5"}, + **{"_meta": {"fastmcp": {"version": "1"}}}, + ), + ) + ) + await _flush() + + assert out.root.isError is False, out.root.content + assert out.root.content[0].text == expected + recorded = _events(client, "$mcp_tool_call")[0]["properties"].get("$mcp_llm_model") + assert recorded == (None if expected.startswith("v1") else "gpt-5") + + +async def test_jlowin_disabled_capture_leaves_llm_model_for_permissive_tools(): + # With capture off nothing injected llm_model, so it is the application's + # even when the schema does not declare it: a Tool accepting arbitrary keys + # must still receive it. + from fastmcp.tools import Tool + + ToolResult = _tool_result_type() + + class Bag(Tool): + async def run(self, arguments): + return ToolResult( + content=[ + mcp_types.TextContent(type="text", text=",".join(sorted(arguments))) + ] + ) + + server = FastMCP("jlowin-permissive") + server.add_tool( + Bag( + name="bag", + parameters={"type": "object", "additionalProperties": {"type": "string"}}, + ) + ) + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(capture_model=False)) + + out = await _call(server, "bag", {"llm_model": "mine", "context": "c"}) + await _flush() + + assert out.root.isError is False + assert out.root.content[0].text == "llm_model" + assert "$mcp_llm_model" not in _events(client, "$mcp_tool_call")[0]["properties"] + + +@pytest.mark.parametrize("listed", [True, False], ids=["listed", "cold"]) +async def test_jlowin_middleware_provided_tool_keeps_its_llm_model(listed): + # The registry does not know a middleware-provided tool, so the effective + # listing is the witness for llm_model. Cold, with no witness at all, the + # argument stays (strips fail closed) and is read (reads fail open). + pytest.importorskip("fastmcp.server.middleware") + from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware + from fastmcp.tools import Tool + + def route(prompt: str, llm_model: str) -> str: + return llm_model + + server = FastMCP("jlowin-middleware-tool") + server.add_middleware(ToolInjectionMiddleware(tools=[Tool.from_function(route)])) + client = FakeClient() + instrument(server, client) + + if listed: + await _list(server) + out = await _call( + server, "route", {"prompt": "p", "context": "c", "llm_model": "own"} + ) + await _flush() + + assert out.root.isError is False, out.root.content + assert out.root.content[0].text == "own" + recorded = _events(client, "$mcp_tool_call")[0]["properties"].get("$mcp_llm_model") + assert recorded == (None if listed else "own") + + +@pytest.mark.parametrize( + "shadow, listed", + [ + ("listing", True), + ("listing", False), + ("dispatch", False), + ("builtin-subclass", False), + ], + ids=["listing-listed", "listing-cold", "dispatch-cold", "builtin-subclass-cold"], +) +async def test_jlowin_middleware_shadowed_tool_keeps_its_llm_model(shadow, listed): + # A registered tool without llm_model is shadowed by middleware serving one + # that declares it, either by also advertising it or only at dispatch. The + # registry would answer for the wrong tool, so with such middleware present + # it is not trusted: the argument stays and, cold, is read fail-open; a + # listing that advertised the shadowing tool settles it as the application's. + pytest.importorskip("fastmcp.server.middleware") + from fastmcp.server.middleware import Middleware + from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware + from fastmcp.tools import Tool + + server = FastMCP("jlowin-shadow") + + @server.tool(name="route") + def registered(prompt: str) -> str: + return "registered" + + def route(prompt: str, llm_model: str) -> str: + return llm_model + + shadowing = Tool.from_function(route) + + class DispatchShadow(Middleware): + async def on_call_tool(self, context, call_next): + if context.message.name == "route": + return await shadowing.run(context.message.arguments or {}) + return await call_next(context) + + if shadow == "builtin-subclass": + # An application subclass of one of FastMCP's own middleware is still + # the application's: its overrides count. + builtins = [type(m) for m in FastMCP("probe").middleware] + if not builtins: + pytest.skip("this FastMCP installs no built-in middleware") + builtin = builtins[0] + + class DispatchShadow(builtin): # type: ignore[no-redef,misc] + async def on_call_tool(self, context, call_next): + if context.message.name == "route": + return await shadowing.run(context.message.arguments or {}) + return await call_next(context) + + server.add_middleware( + ToolInjectionMiddleware(tools=[shadowing]) + if shadow == "listing" + else DispatchShadow() + ) + client = FakeClient() + instrument(server, client) + + if listed: + await _list(server) + out = await _call( + server, "route", {"prompt": "p", "context": "c", "llm_model": "own"} + ) + await _flush() + + assert out.root.isError is False, out.root.content + assert out.root.content[0].text == "own" + recorded = _events(client, "$mcp_tool_call")[0]["properties"].get("$mcp_llm_model") + assert recorded == (None if listed else "own") + + +async def test_jlowin_without_dereferencing_a_root_ref_is_never_injected_into(): + # With dereferencing off the client sees the `$ref` itself, which the + # listing never injects llm_model into, so a cold call must not strip it + # from a permissive tool either. + from fastmcp.tools import Tool + + ToolResult = _tool_result_type() + + if "dereference_schemas" not in inspect.signature(FastMCP.__init__).parameters: + pytest.skip("this FastMCP cannot disable schema dereferencing") + + class Bag(Tool): + async def run(self, arguments): + return ToolResult( + content=[ + mcp_types.TextContent(type="text", text=",".join(sorted(arguments))) + ] + ) + + server = FastMCP("jlowin-raw-ref", dereference_schemas=False) + server.add_tool( + Bag( + name="bag", + parameters={ + "$ref": "#/$defs/Args", + "$defs": {"Args": {"type": "object", "additionalProperties": True}}, + }, + ) + ) + client = FakeClient() + instrument(server, client) + + out = await _call(server, "bag", {"llm_model": "mine", "context": "c"}) + await _flush() + + assert out.root.isError is False + assert out.root.content[0].text == "llm_model" + + +async def test_jlowin_cold_call_survives_a_fastmcp_without_middleware(monkeypatch): + # FastMCP releases before middleware existed have no + # `fastmcp.server.middleware`; ownership inference must not import its way + # into breaking dispatch, and the registry is then simply trusted. + import sys + + monkeypatch.setitem(sys.modules, "fastmcp.server.middleware", None) + server = make_server() + client = FakeClient() + instrument(server, client) + + out = await _call( + server, "add", {"a": 2, "b": 3, "context": "sum", "llm_model": "model-a"} + ) + await _flush() + + assert out.root.isError is False, out.root.content + assert ( + _events(client, "$mcp_tool_call")[0]["properties"]["$mcp_llm_model"] + == "model-a" + )