fix(mcp): resolve referenced input schemas

Generated-By: PostHog Desktop
Task-Id: dda38523-5b8d-47b6-a11a-64cd59a812fe

Author: gesh

posthog/mcp/_instrument_lowlevel.py

@@ -605,10 +605,13 @@ def _tool_schema_view(
         return None, None, None
     schema = getattr(tool, "parameters", None)
     if isinstance(schema, dict):
-        declared, injectable = _schema_view(
-            schema, dereferenced=_server_dereferences(high_level)
+        dereferenced = _server_dereferences(high_level)
+        declared, injectable = _schema_view(schema, dereferenced=dereferenced)
+        return (
+            declared,
+            injectable,
+            _input_schema_view(schema, dereferenced=dereferenced),
         )
-        return declared, injectable, schema
     fn = getattr(tool, "fn", None)
     if fn is None:
         return set(), True, None
@@ -703,6 +706,26 @@ def _schema_view(schema: Dict[str, Any], *, dereferenced: bool) -> Tuple[set, bo
     return declared, injectable
 
 
+def _input_schema_view(
+    schema: Dict[str, Any], *, dereferenced: bool
+) -> Optional[Dict[str, Any]]:
+    """Return the declared fields that the client sees for input analytics."""
+    if not dereferenced:
+        return schema
+    nodes = _reference_chain(schema)
+    if nodes is None:
+        return None
+    if len(nodes) == 1:
+        return schema
+    properties = {
+        key: value
+        for node in nodes
+        if isinstance(node.get("properties"), dict)
+        for key, value in node["properties"].items()
+    }
+    return {"type": "object", "properties": properties}
+
+
 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)."""

posthog/test/mcp/test_fastmcp_v2.py

@@ -358,6 +358,43 @@ async def run(self, arguments):
     assert recorded == ("gpt-5" if stripped else None)
 
 
+async def test_jlowin_root_ref_schema_is_dereferenced_for_input_names():
+    from fastmcp.tools import Tool
+
+    ToolResult = _tool_result_type()
+
+    class Echo(Tool):
+        async def run(self, arguments):
+            return ToolResult(
+                content=[mcp_types.TextContent(type="text", text=str(arguments["a"]))]
+            )
+
+    server = FastMCP("jlowin-input-ref")
+    server.add_tool(
+        Echo(
+            name="echo",
+            parameters={
+                "$ref": "#/$defs/Args",
+                "$defs": {
+                    "Args": {
+                        "type": "object",
+                        "properties": {"a": {"type": "integer"}},
+                    }
+                },
+            },
+        )
+    )
+    client = FakeClient()
+    instrument(server, client)
+
+    out = await _call(server, "echo", {"a": 1})
+    await _flush()
+
+    assert out.root.isError is False
+    properties = _events(client, "$mcp_tool_call")[0]["properties"]
+    assert properties["$mcp_input_keys"] == ["a"]
+
+
 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