feat(mcp): record safe tool input names (#989)

* feat(mcp): record safe tool input names

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

* fix(mcp): scope input schemas to tool calls

* fix(mcp): resolve referenced input schemas

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

* fix(mcp): address input name review

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

Author: gesh

.sampo/changesets/mcp-input-names.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+Record safe tool input field names on MCP tool-call events. Add server-owned input alias maps for automatic instrumentation and a public helper for custom dispatchers. The SDK records field names and alias use without reading argument values or changing tool calls.

posthog/mcp/README.md

@@ -67,6 +67,46 @@ The value must contain 1 to 256 characters. The SDK records it as
 `$mcp_server_build` on all MCP events. `PostHogMCP.capture()` also adds it to
 custom events unless the event provides its own value.
 
+## Capture safe input field names
+
+Tool-call events include `$mcp_input_keys`. The SDK records names from the
+server's input schema. It replaces unknown names with one `[redacted]` entry.
+Argument values do not affect this property.
+
+The SDK records at most 20 names. Each name can contain at most 64 characters.
+Use `should_record_input_key` to control names in both input properties. The
+callback receives the name and `{"declared": bool}`. Return `True` only for
+names that the SDK can record.
+
+Raw low-level servers have no trusted tool registry. The SDK redacts names when
+it cannot resolve the current schema for a call.
+
+Use `resolve_input_aliases` when a tool accepts alternative names. The map uses
+each canonical name as a key. Its value lists accepted aliases in server order.
+
+```python
+instrument(
+    server,
+    posthog,
+    MCPAnalyticsOptions(
+        resolve_input_aliases=lambda tool_name: (
+            {"location": ["city", "place"]}
+            if tool_name == "weather-current"
+            else None
+        )
+    ),
+)
+```
+
+The SDK records `city` in `$mcp_input_keys`. It also records
+`city:location` in `$mcp_input_aliases_used`. The SDK does not change the call.
+
+This safe-name rule does not change `$mcp_parameters`. That property still
+contains the sanitized tool arguments and their original names.
+
+Custom dispatchers can call `get_tool_input_properties()` and add its result to
+the `properties` argument of `capture_tool_call()`.
+
 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.

posthog/mcp/__init__.py

@@ -46,6 +46,7 @@
 )
 from ._event_types import MCPAnalyticsEventType
 from ._instrumentation import drain_pending
+from ._tool_input import get_tool_input_properties
 from ._internal import (
     MCPAnalyticsData,
     get_server_tracking_data,
@@ -82,12 +83,15 @@
     CaptureEventData,
     CollectFeedbackOptions,
     FeedbackReport,
+    InputAliasMap,
     MCPAnalyticsContextOptions,
     MCPAnalyticsModelOptions,
     MCPAnalyticsModelSource,
     MCPAnalyticsOptions,
     PreparedToolCall,
     PreparedToolResult,
+    ShouldRecordInputKeyFn,
+    ToolInputOptions,
     UserIdentity,
 )
 from .version import __version__
@@ -104,8 +108,12 @@
     "CaptureEventData",
     "CollectFeedbackOptions",
     "FeedbackReport",
+    "InputAliasMap",
     "PreparedToolCall",
     "PreparedToolResult",
+    "ShouldRecordInputKeyFn",
+    "ToolInputOptions",
+    "get_tool_input_properties",
     "get_more_tools_result",
     "send_feedback_result",
     "SEND_FEEDBACK_TOOL_NAME",

posthog/mcp/_instrument_fastmcp.py

@@ -22,6 +22,7 @@
 
 import inspect
 import time
+from dataclasses import replace
 from typing import Any, Dict, Optional, Tuple
 
 import mcp.types as mcp_types
@@ -48,6 +49,7 @@
     request_meta_from_context,
 )
 from ._output_instructions import mirror_instructions_into_structured_content
+from ._tool_input import get_registered_tool_input_schema
 from .logger import log
 from .tools import get_more_tools_result_text
 
@@ -131,6 +133,11 @@ async def wrapped(
                 for text in lifecycle.virtual_result_texts(reply)
             ]
 
+        lifecycle = replace(
+            lifecycle,
+            input_schema=get_registered_tool_input_schema(server, name),
+        )
+
         # Strip each injected key independently. A tool can declare its own
         # `context` (kept) while `conversation_id` is still SDK-injected (stripped),
         # so coupling both to context-ownership leaked conversation_id into the tool.

posthog/mcp/_instrument_lowlevel.py

@@ -186,10 +186,10 @@ def _wrap_call_tool(
     async def handler(req: Any) -> Any:
         name = req.params.name
         arguments = dict(req.params.arguments or {})
-        strip, model_ours = (
+        strip, model_ours, input_schema = (
             await _standalone_ownership(data, high_level, name, req.params.meta)
             if strip_injected
-            else (set(), data.tool_model_parameter_injected.get(name))
+            else (set(), data.tool_model_parameter_injected.get(name), None)
         )
         client_name, client_version = _client_info(server)
         protocol_version = _protocol_version(server)
@@ -212,6 +212,7 @@ async def handler(req: Any) -> Any:
             client_version=client_version,
             protocol_version=protocol_version,
             extra={"session_id": mcp_session_id, "ctx": _request_context(server)},
+            input_schema=input_schema,
         )
 
         if lifecycle.is_missing_capability and (
@@ -565,10 +566,11 @@ def _tool_lookup_not_found_errors() -> Tuple[type, ...]:
 
 async def _standalone_ownership(
     data: MCPAnalyticsData, high_level: Any, name: str, meta: Any
-) -> Tuple[set, Optional[bool]]:
+) -> Tuple[set, Optional[bool], Optional[Dict[str, Any]]]:
     """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).
+    ours (``None`` when nothing can say). The third item is the trusted schema
+    for input-name analytics, or ``None`` when middleware can change dispatch.
 
     Only keys injected under the current options are candidates. ``context``
     and ``conversation_id`` are stripped unless the registered schema (or,
@@ -580,17 +582,44 @@ async def _standalone_ownership(
     stays and is still read (posthog-js ADR-0011).
     """
     try:
-        declared, model_injectable = await _registry_view(high_level, name, meta)
+        declared, model_injectable, input_schema = await _registry_view(
+            high_level, name, meta
+        )
         model_ours = data.tool_model_parameter_injected.get(name, model_injectable)
         if _dispatch_can_differ(high_level):
             model_ours = False
+            input_schema = await _listed_input_schema(high_level, name, meta)
     except Exception:  # noqa: BLE001 - ownership inference must never prevent dispatch
-        declared, model_ours = None, None
+        declared, model_ours, input_schema = None, 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:
         strip.add("llm_model")
-    return strip, model_ours
+    return strip, model_ours, input_schema
+
+
+def _tool_schema_view(
+    high_level: Any, tool: Any
+) -> Tuple[Optional[set], Optional[bool], Optional[Dict[str, Any]]]:
+    if tool is None:
+        return None, None, None
+    schema = getattr(tool, "parameters", None)
+    if isinstance(schema, dict):
+        dereferenced = _server_dereferences(high_level)
+        declared, injectable = _schema_view(schema, dereferenced=dereferenced)
+        return (
+            declared,
+            injectable,
+            schema,
+        )
+    fn = getattr(tool, "fn", None)
+    if fn is None:
+        return set(), True, None
+    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, None
+    return declared, "llm_model" not in declared, None
 
 
 def _injected_keys(data: MCPAnalyticsData) -> set:
@@ -609,35 +638,24 @@ def _injected_keys(data: MCPAnalyticsData) -> set:
 
 async def _registry_view(
     high_level: Any, name: str, meta: Any
-) -> Tuple[Optional[set], Optional[bool]]:
+) -> Tuple[Optional[set], Optional[bool], Optional[Dict[str, Any]]]:
     """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."""
+    registry has no such tool or cannot be read. The third item is the tool's
+    input schema when the registry supplies one."""
     try:
         tool = await _registered_tool(high_level, name, meta)
     except Exception as error:  # noqa: BLE001 - introspection is best-effort
         if not isinstance(error, _tool_lookup_not_found_errors()):
             warn_ownership_lookup_failed(name, error)
-            return set(_INJECTED_KEYS), None
-        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
+            return set(_INJECTED_KEYS), None, None
+        return None, None, None
+    return _tool_schema_view(high_level, tool)
 
 
 async def _registered_tool(high_level: Any, name: str, meta: Any) -> Any:
@@ -651,6 +669,29 @@ async def _registered_tool(high_level: Any, name: str, meta: Any) -> Any:
     return await high_level.get_tool(name, version=VersionSpec(eq=version))
 
 
+async def _listed_input_schema(
+    high_level: Any, name: str, meta: Any
+) -> Optional[Dict[str, Any]]:
+    """Return the schema that middleware advertises for the current request."""
+    try:
+        version = _requested_tool_version(meta)
+        candidates = [
+            tool for tool in await high_level.list_tools() if tool.name == name
+        ]
+        if version is not None:
+            candidates = [
+                tool
+                for tool in candidates
+                if str(getattr(tool, "version", "")) == version
+            ]
+        tool = candidates[-1] if candidates else None
+        schema = getattr(tool, "parameters", None)
+    except Exception as error:  # noqa: BLE001 - analytics must not break dispatch
+        log(f"PostHog MCP: could not resolve schema for tool {name!r} - {error}")
+        return None
+    return schema if isinstance(schema, dict) else None
+
+
 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

posthog/mcp/_instrument_v2.py

@@ -32,6 +32,7 @@
 
 import time
 from collections.abc import Mapping
+from dataclasses import replace
 from typing import Any, Dict, FrozenSet, Optional, Set, Tuple
 
 import mcp.types as mcp_types
@@ -63,6 +64,7 @@
     request_meta_from_context,
 )
 from ._output_instructions import mirror_instructions_into_structured_content
+from ._tool_input import get_registered_tool_input_schema
 from .logger import log
 from .request_headers import get_request_headers
 from .session_token import read_mcp_session_header
@@ -230,8 +232,9 @@ def _tool_own_properties_v2(high_level: Any, name: str) -> Dict[str, Any]:
     site so checking ownership of both ``context`` and ``conversation_id``
     doesn't look the tool up from the manager twice."""
     try:
-        tool = high_level._tool_manager.get_tool(name)
-        properties = (getattr(tool, "parameters", None) or {}).get("properties")
+        properties = (get_registered_tool_input_schema(high_level, name) or {}).get(
+            "properties"
+        )
     except Exception:  # noqa: BLE001
         return {}
     # Fail closed on a malformed schema: the caller does `param in <this>` in the
@@ -318,6 +321,11 @@ async def wrapped(
             ]
             return mcp_types.CallToolResult(content=virtual_content)
 
+        lifecycle = replace(
+            lifecycle,
+            input_schema=get_registered_tool_input_schema(server, name),
+        )
+
         # v2 validates against the function signature and rejects unexpected
         # keys, so injected parameters are stripped before dispatch — but never
         # one the tool's own schema declares (that's a real argument).
@@ -441,7 +449,7 @@ def _requested_tool_version(ctx: Any) -> Optional[str]:
 
 async def _standalone_injected_parameters(
     server: Any, data: MCPAnalyticsData, name: str, version: Optional[str]
-) -> Optional[FrozenSet[str]]:
+) -> Tuple[Optional[FrozenSet[str]], Optional[Dict[str, Any]]]:
     """Resolve ownership in the current request, including middleware and versions.
 
     Listings from other requests can have different application-owned parameters.
@@ -464,9 +472,9 @@ async def _standalone_injected_parameters(
         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
+        return None, None
     if not isinstance(schema, dict):
-        return None
+        return None, None
     injected = set()
     if is_context_enabled(data.options.context):
         injected.add("context")
@@ -476,7 +484,10 @@ async def _standalone_injected_parameters(
         can_inject_model_parameter(schema)
     ):
         injected.add("llm_model")
-    return frozenset(key for key in injected if not schema_has_param(schema, key))
+    return (
+        frozenset(key for key in injected if not schema_has_param(schema, key)),
+        schema,
+    )
 
 
 def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None:
@@ -492,10 +503,11 @@ async def handler(ctx: Any, params: Any) -> Any:
         # 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
+        input_schema = None
         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, input_schema = await _standalone_injected_parameters(
                 standalone, data, name, version
             )
             if injected is not None:
@@ -521,6 +533,7 @@ async def handler(ctx: Any, params: Any) -> Any:
             client_version=client_version,
             protocol_version=protocol_version,
             extra={"session_id": mcp_session_id, "ctx": ctx},
+            input_schema=input_schema,
         )
 
         # No tool registry on a raw low-level server, so ownership is settled