fix(mcp): preserve tool arguments across FastMCP middleware

Skip model argument injection on standalone MCP SDK v1 servers when application middleware can alter listing or dispatch. Keep client metadata capture enabled, and invalidate earlier model ownership so a replacement tool retains its own llm_model argument.

Document the middleware limitation in the README and changeset. Extend the existing shadowing test with listed, reordered, and late middleware cases, and verify a fresh replica executes advertised calls while still capturing client metadata.

Validation: MCP v1 553 passed, 1 skipped; MCP v2 444 passed, 21 skipped. Ruff lint/format, mypy (239 files), public API snapshot, and CodeScene safeguard pass. CodeScene improves both runtime files; existing large-module and test-complexity warnings remain.

Author: lucasheriques

.sampo/changesets/mcp-analytics-defaults.md

@@ -3,3 +3,5 @@ pypi/posthog: minor
 ---
 
 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.
+
+Standalone FastMCP on MCP SDK 1.x skips `llm_model` injection when application middleware can change tool listing or dispatch. Model metadata capture remains enabled; this prevents cold replicas from rejecting injected arguments and preserves replacement tools' own arguments.

posthog/mcp/README.md

@@ -92,7 +92,15 @@ nothing (posthog-js ADR-0011: reads fail open, strips fail closed). A tool that
 `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 the registered tool schema, so they are
-unaffected and need no prior listing.
+unaffected and need no prior listing when dispatch uses that registry.
+
+On MCP SDK 1.x, standalone FastMCP application middleware can replace or reroute a tool.
+When it overrides a listing or dispatch hook, PostHog does not inject `llm_model`: a fresh
+replica cannot safely distinguish that field from a replacement tool's own argument.
+Client metadata capture remains enabled. This also applies to pass-through logging or
+authorization middleware with those hooks; argument-based model capture is skipped so an
+application-owned value cannot be mistaken for analytics. No additional catalog lookup runs
+during a tool call.
 
 For a custom dispatcher, `PostHogMCP` enables the same option by default; pass request
 metadata through explicitly:

posthog/mcp/_instrument_lowlevel.py

@@ -85,7 +85,7 @@ def instrument_fastmcp_v2(server: Any, data: MCPAnalyticsData) -> None:
     # schema that requires `context` contradicts the arguments the SDK actually
     # sees: under `FastMCP(strict_input_validation=True)` every call fails with
     # "'context' is a required property".
-    _wrap_list_tools(low_level, data, context_required=False)
+    _wrap_list_tools(low_level, data, context_required=False, high_level=server)
     _wrap_resource_requests(low_level, data)
 
 
@@ -327,14 +327,21 @@ async def handler(req: Any) -> Any:
 
 
 def _inject_tool_schemas(
-    data: MCPAnalyticsData, tools: list, *, context_required: bool
+    data: MCPAnalyticsData,
+    tools: list,
+    *,
+    context_required: bool,
+    high_level: Any = None,
 ) -> None:
     """Advertise the analytics parameters on a listing's tools, in place.
 
     Runs on both the client-facing listing and the SDK's internal cache-
     population pass, so the schema the SDK validates against always matches the
     one we advertised — see the note in ``handler``.
     """
+    # Middleware can replace the registered tool on another replica. Without
+    # proof of ownership there, advertising a field could break its validation.
+    inject_model = high_level is None or not _dispatch_can_differ(high_level)
     verdicts: Dict[str, bool] = {}
     for tool in tools:
         schema = getattr(tool, "inputSchema", None)
@@ -345,6 +352,7 @@ def _inject_tool_schemas(
             owns_context=schema_has_param(schema, "context"),
             context_required=context_required,
             is_sdk_virtual_tool=False,
+            inject_model=inject_model,
         )
         verdict = data.tool_model_parameter_injected.get(tool.name)
         if verdict is None:
@@ -357,7 +365,11 @@ def _inject_tool_schemas(
 
 
 def _wrap_list_tools(
-    server: Any, data: MCPAnalyticsData, *, context_required: bool
+    server: Any,
+    data: MCPAnalyticsData,
+    *,
+    context_required: bool,
+    high_level: Any = None,
 ) -> None:
     handlers = server.request_handlers
     original = handlers.get(mcp_types.ListToolsRequest)
@@ -400,7 +412,9 @@ async def probe_raw_tool_names(_ctx: Any = None) -> Optional[Set[str]]:
         # without re-injecting here the next real call is rejected for sending the
         # `context` we advertised. Same reason the `req is None` branch below
         # injects.
-        _inject_tool_schemas(data, tools, context_required=context_required)
+        _inject_tool_schemas(
+            data, tools, context_required=context_required, high_level=high_level
+        )
         return advertised_tool_names(tools)
 
     data.raw_tool_names_probe = probe_raw_tool_names
@@ -416,7 +430,9 @@ async def handler(req: Any) -> Any:
         if req is None:
             result = await original(req)
             tools = extract_tools(result)
-            _inject_tool_schemas(data, tools, context_required=context_required)
+            _inject_tool_schemas(
+                data, tools, context_required=context_required, high_level=high_level
+            )
             return result
 
         client_name, client_version = _client_info(server)
@@ -464,7 +480,9 @@ async def handler(req: Any) -> Any:
             is_first_page=is_first_listing_page(getattr(req, "params", None)),
         )
 
-        _inject_tool_schemas(data, tools, context_required=context_required)
+        _inject_tool_schemas(
+            data, tools, context_required=context_required, high_level=high_level
+        )
 
         result = apply_virtual_tool_injection(
             result, injection, names, data, schema_field="inputSchema"
@@ -557,25 +575,18 @@ async def _standalone_ownership(
     and ``conversation_id`` are stripped unless the registered schema (or,
     without one, the function signature) declares them. A registry failure
     protects all keys; a missing registry entry retains the middleware fallback
-    for ``context`` and ``conversation_id``. ``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).
+    for ``context`` and ``conversation_id``. Middleware that can replace a tool
+    disables model injection and invalidates earlier model ownership. Otherwise
+    the listing or registry decides; with neither witness the model argument
+    stays and is still read (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
-        )
+        model_ours = data.tool_model_parameter_injected.get(name, model_injectable)
+        if _dispatch_can_differ(high_level):
+            model_ours = False
     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
+        declared, model_ours = None, None
     candidates = _injected_keys(data)
     strip = {k for k in candidates - {"llm_model"} if k not in (declared or set())}
     if "llm_model" in candidates and model_ours:

posthog/mcp/_instrumentation.py

@@ -994,6 +994,7 @@ def mutate_tool_schema(
     owns_context: bool,
     context_required: bool,
     is_sdk_virtual_tool: bool,
+    inject_model: bool = True,
 ) -> None:
     """Apply the common analytics schema pipeline and write it back in place.
 
@@ -1014,21 +1015,14 @@ def mutate_tool_schema(
             get_context_description(data.options.context),
             required=context_required,
         )
-    if is_capture_model_enabled(data.options.capture_model):
-        model_was_injected = data.tool_model_parameter_injected.get(tool.name, False)
-        app_owns_model = (
-            schema_has_param(schema, "llm_model") and not model_was_injected
-        )
-        if not app_owns_model and not schema_has_param(schema, "llm_model"):
-            schema = add_model_parameter_to_schema(
-                schema,
-                tool.name,
-                get_model_description(data.options.capture_model),
-                required=context_required,
-            )
-        data.tool_model_parameter_injected[tool.name] = (
-            not app_owns_model and schema_has_param(schema, "llm_model")
+    if inject_model:
+        schema = _mutate_model_schema(
+            data, tool.name, schema, required=context_required
         )
+    else:
+        # Discard any earlier registered-tool verdict: middleware may now
+        # advertise or dispatch an application-owned field with the same name.
+        data.tool_model_parameter_injected[tool.name] = False
     if data.options.enable_conversation_id and not schema_has_param(
         schema, "conversation_id"
     ):
@@ -1044,6 +1038,26 @@ def mutate_tool_schema(
         )
 
 
+def _mutate_model_schema(
+    data: MCPAnalyticsData, name: str, schema: Any, *, required: bool
+) -> Any:
+    if not is_capture_model_enabled(data.options.capture_model):
+        return schema
+    model_was_injected = data.tool_model_parameter_injected.get(name, False)
+    app_owns_model = schema_has_param(schema, "llm_model") and not model_was_injected
+    if not app_owns_model and not schema_has_param(schema, "llm_model"):
+        schema = add_model_parameter_to_schema(
+            schema,
+            name,
+            get_model_description(data.options.capture_model),
+            required=required,
+        )
+    data.tool_model_parameter_injected[name] = not app_owns_model and schema_has_param(
+        schema, "llm_model"
+    )
+    return schema
+
+
 def request_to_dict(req: Any) -> Dict[str, Any]:
     """Shape a request object into the JSON-RPC-ish dict the sanitizer expects."""
     method = getattr(req, "method", None) or "tools/list"

posthog/test/mcp/test_fastmcp_v2.py

@@ -184,6 +184,44 @@ async def test_jlowin_call_strips_llm_model_and_records_it(listed):
     assert calls[0]["properties"]["$mcp_llm_model"] == "model-a"
 
 
+async def test_jlowin_middleware_defaults_work_across_fresh_instances(monkeypatch):
+    from types import SimpleNamespace
+
+    from fastmcp.server.middleware import Middleware
+
+    from posthog.mcp import _instrument_lowlevel
+
+    monkeypatch.setattr(
+        _instrument_lowlevel,
+        "_request_context",
+        lambda _: SimpleNamespace(meta={"x-codex-turn-metadata": {"model": "gpt-5"}}),
+    )
+    client = FakeClient()
+
+    class PassThrough(Middleware):
+        async def on_call_tool(self, context, call_next):
+            return await call_next(context)
+
+    def fresh():
+        server = make_server()
+        server.add_middleware(PassThrough())
+        instrument(server, client)
+        return server
+
+    listing = await _list(fresh())
+    schema = listing.root.tools[0].inputSchema
+    # A cold replica cannot distinguish pass-through from tool-replacing
+    # middleware, so discovery must not request an argument it cannot strip.
+    assert "llm_model" not in schema["properties"]
+    result = await _call(fresh(), "add", {"a": 2, "b": 3, "context": "sum"})
+    assert result.root.isError is False
+    assert result.root.content[0].text == "5"
+    await _flush()
+    event = _events(client, "$mcp_tool_call")[0]["properties"]
+    assert event["$mcp_llm_model"] == "gpt-5"
+    assert event["$mcp_llm_model_source"] == "client_metadata"
+
+
 _OWN_MODEL = {"llm_model": {"type": "string"}}
 
 
@@ -401,8 +439,8 @@ async def run(self, arguments):
 @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).
+    # listing may differ from dispatch. Model injection stays off and the
+    # application argument stays intact, even without a prior listing.
     pytest.importorskip("fastmcp.server.middleware")
     from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware
     from fastmcp.tools import Tool
@@ -425,7 +463,7 @@ def route(prompt: str, llm_model: str) -> str:
     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")
+    assert recorded is None
 
 
 @pytest.mark.parametrize(
@@ -434,16 +472,27 @@ def route(prompt: str, llm_model: str) -> str:
         ("listing", True),
         ("listing", False),
         ("dispatch", False),
+        ("dispatch", True),
+        ("listing-last", True),
+        ("late-dispatch", False),
         ("builtin-subclass", False),
     ],
-    ids=["listing-listed", "listing-cold", "dispatch-cold", "builtin-subclass-cold"],
+    ids=[
+        "listing-listed",
+        "listing-cold",
+        "dispatch-cold",
+        "dispatch-listed",
+        "listing-last",
+        "late-dispatch",
+        "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.
+    # it is not trusted: the argument stays and is not read as a self-report,
+    # regardless of listing order or a stale registered-tool verdict.
     pytest.importorskip("fastmcp.server.middleware")
     from fastmcp.server.middleware import Middleware
     from fastmcp.server.middleware.tool_injection import ToolInjectionMiddleware
@@ -461,6 +510,10 @@ def route(prompt: str, llm_model: str) -> str:
     shadowing = Tool.from_function(route)
 
     class DispatchShadow(Middleware):
+        async def on_list_tools(self, context, call_next):
+            tools = await call_next(context)
+            return [*tools, shadowing] if shadow == "listing-last" else tools
+
         async def on_call_tool(self, context, call_next):
             if context.message.name == "route":
                 return await shadowing.run(context.message.arguments or {})
@@ -480,14 +533,15 @@ async def on_call_tool(self, context, call_next):
                     return await shadowing.run(context.message.arguments or {})
                 return await call_next(context)
 
+    client = FakeClient()
+    instrument(server, client)
+    if shadow == "late-dispatch":
+        await _list(server)
     server.add_middleware(
         ToolInjectionMiddleware(tools=[shadowing])
         if shadow == "listing"
         else DispatchShadow()
     )
-    client = FakeClient()
-    instrument(server, client)
-
     if listed:
         await _list(server)
     out = await _call(
@@ -498,7 +552,7 @@ async def on_call_tool(self, context, call_next):
     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")
+    assert recorded is None
 
 
 async def test_jlowin_without_dereferencing_a_root_ref_is_never_injected_into():