From 142caf021b741a3e27b2ae267b91e5a2df27b9d0 Mon Sep 17 00:00:00 2001 From: Georgis Andonis <6430745+gesh@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:41:45 +0300 Subject: [PATCH 1/2] feat(mcp): record safe tool input names Generated-By: PostHog Desktop Task-Id: dda38523-5b8d-47b6-a11a-64cd59a812fe --- .sampo/changesets/mcp-input-names.md | 5 ++ posthog/mcp/README.md | 29 +++++++ posthog/mcp/__init__.py | 8 ++ posthog/mcp/_instrument_fastmcp.py | 2 +- posthog/mcp/_instrument_lowlevel.py | 2 +- posthog/mcp/_instrument_v2.py | 2 +- posthog/mcp/_instrumentation.py | 38 ++++++++- posthog/mcp/_internal.py | 5 ++ posthog/mcp/_tool_input.py | 114 +++++++++++++++++++++++++++ posthog/mcp/constants.py | 2 + posthog/mcp/types.py | 26 ++++++ posthog/test/mcp/test_features_m4.py | 31 ++++++++ posthog/test/mcp/test_tool_input.py | 90 +++++++++++++++++++++ references/public_api_snapshot.txt | 14 +++- 14 files changed, 360 insertions(+), 8 deletions(-) create mode 100644 .sampo/changesets/mcp-input-names.md create mode 100644 posthog/mcp/_tool_input.py create mode 100644 posthog/test/mcp/test_tool_input.py diff --git a/.sampo/changesets/mcp-input-names.md b/.sampo/changesets/mcp-input-names.md new file mode 100644 index 000000000..60015885c --- /dev/null +++ b/.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. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 287484dc8..798baccb6 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -50,6 +50,35 @@ from posthog.mcp import MCPAnalyticsOptions, instrument instrument(server, posthog, MCPAnalyticsOptions(capture_model=False, enable_conversation_id=False)) ``` +## 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. + +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. + +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. diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 298de241c..9080beea3 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -45,6 +45,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, @@ -81,11 +82,14 @@ CaptureEventData, CollectFeedbackOptions, FeedbackReport, + InputAliasMap, MCPAnalyticsContextOptions, MCPAnalyticsModelOptions, MCPAnalyticsModelSource, MCPAnalyticsOptions, PreparedToolCall, + ShouldRecordInputKeyFn, + ToolInputOptions, UserIdentity, ) from .version import __version__ @@ -102,7 +106,11 @@ "CaptureEventData", "CollectFeedbackOptions", "FeedbackReport", + "InputAliasMap", "PreparedToolCall", + "ShouldRecordInputKeyFn", + "ToolInputOptions", + "get_tool_input_properties", "get_more_tools_result", "send_feedback_result", "SEND_FEEDBACK_TOOL_NAME", diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 7f9af6f15..28ed09736 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -273,7 +273,7 @@ async def list_handler(req: Any) -> Any: duration_ms = (time.monotonic() - start) * 1000 tools = extract_tools(result) # Empty is computed before adding the virtual missing-capability tool. - names, empty = collect_listed_tools(data, tools) + names, empty = collect_listed_tools(data, tools, lifecycle.session_id) injection = resolve_virtual_tool_injection( data, tools, diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index ee24748aa..aca28949f 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -473,7 +473,7 @@ async def handler(req: Any) -> Any: # Zero advertised tools is treated as an errored tools/list before the # virtual missing-capability tool is appended. - names, empty = collect_listed_tools(data, tools) + names, empty = collect_listed_tools(data, tools, lifecycle.session_id) injection = resolve_virtual_tool_injection( data, tools, diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index b6d7ecd99..85ef69597 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -696,7 +696,7 @@ async def handler(ctx: Any, params: Any) -> Any: tools = list(getattr(result, "tools", []) or []) # Empty is computed before adding the virtual missing-capability tool. - names, empty = collect_listed_tools(data, tools) + names, empty = collect_listed_tools(data, tools, lifecycle.session_id) injection = resolve_virtual_tool_injection( data, tools, is_first_page=is_first_listing_page(params) ) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 2355bcdc6..6d254bb81 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -41,6 +41,7 @@ ) from ._intent import resolve_tool_call_intent, set_event_intent from ._internal import MCPAnalyticsData, handle_identify, resolve_event_properties +from ._tool_input import get_tool_input_properties from ._model_parameters import ( add_model_parameter_to_schema, get_model_description, @@ -55,7 +56,7 @@ from .session import resolve_session_id, resolve_session_id_with_source from .session_token import SessionTokenPayload, decode_session_id from .tools import resolve_missing_capability_tool_name -from .types import CollectFeedbackOptions, FeedbackReport +from .types import CollectFeedbackOptions, FeedbackReport, ToolInputOptions # The virtual tools this SDK advertises into tools/list. Every piece of per-tool # policy -- enable switch, configured name, warning text, warn-once @@ -677,8 +678,24 @@ async def record_tool_call( event["error"] = capture_exception(result) props = await resolve_event_properties(data, request, extra) - if props is not None: - event["properties"] = props + input_aliases = None + if data.options.resolve_input_aliases is not None: + try: + input_aliases = data.options.resolve_input_aliases(name) + except Exception as err: # noqa: BLE001 - analytics callbacks are isolated + log(f"Warning: resolve_input_aliases failed for tool {name}: {err}") + schema = data.tool_input_schemas.get(session_id, {}).get(name) + event["properties"] = { + **(props or {}), + **get_tool_input_properties( + arguments or {}, + schema, + ToolInputOptions( + should_record_input_key=data.options.should_record_input_key, + input_aliases=input_aliases, + ), + ), + } stamp_transport_identity(event, extra) fire_and_forget(capture_event(data, event), data) @@ -973,9 +990,12 @@ def read_tool_category(tool: Any) -> Optional[str]: return None -def collect_listed_tools(data: MCPAnalyticsData, tools: list) -> tuple[List[str], bool]: +def collect_listed_tools( + data: MCPAnalyticsData, tools: list, session_id: Optional[str] = None +) -> tuple[List[str], bool]: """Cache common tool metadata and return the pre-injection listing summary.""" names = [] + schemas = data.tool_input_schemas.get(session_id, {}) if session_id else None for tool in tools: names.append(tool.name) if getattr(tool, "description", None): @@ -983,6 +1003,16 @@ def collect_listed_tools(data: MCPAnalyticsData, tools: list) -> tuple[List[str] category = read_tool_category(tool) if category: data.tool_categories[tool.name] = category + if schemas is not None: + schema = getattr(tool, "inputSchema", None) + if schema is None: + schema = getattr(tool, "input_schema", None) + schemas[tool.name] = schema + if session_id and schemas is not None: + data.tool_input_schemas[session_id] = schemas + data.tool_input_schemas.move_to_end(session_id) + while len(data.tool_input_schemas) > 1000: + data.tool_input_schemas.popitem(last=False) return names, not tools diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index b10a9ef60..61ea145ad 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -88,6 +88,11 @@ class MCPAnalyticsData: identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) tool_descriptions: Dict[str, str] = field(default_factory=dict) + # Original tool input schemas by PostHog session. This keeps field-name + # privacy decisions isolated when a server advertises user-specific tools. + tool_input_schemas: "OrderedDict[str, Dict[str, Any]]" = field( + default_factory=OrderedDict + ) # True only when PostHog added llm_model to this tool's advertised schema. # Missing/False fails closed so an application-owned field is never read or stripped. tool_model_parameter_injected: Dict[str, bool] = field(default_factory=dict) diff --git a/posthog/mcp/_tool_input.py b/posthog/mcp/_tool_input.py new file mode 100644 index 000000000..0596000f4 --- /dev/null +++ b/posthog/mcp/_tool_input.py @@ -0,0 +1,114 @@ +"""Describe tool arguments by field name without reading their values.""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from .constants import PostHogMCPAnalyticsProperty +from .types import InputAliasMap, JsonRecord, ToolInputOptions + +_MAX_INPUT_KEYS = 20 +_MAX_KEY_LENGTH = 64 +_ANALYTICS_KEYS = {"context", "llm_model", "conversation_id"} + + +def _declared_properties(schema: Any) -> Dict[str, Any]: + if not isinstance(schema, dict): + return {} + properties = schema.get("properties") + return properties if isinstance(properties, dict) else {} + + +def _alias_names(aliases: Optional[InputAliasMap]) -> List[str]: + if not isinstance(aliases, dict): + return [] + return [ + name + for names in aliases.values() + if isinstance(names, (list, tuple)) + for name in names + if isinstance(name, str) + ] + + +def _aliases_used( + aliases: Optional[InputAliasMap], input_value: Dict[str, Any] +) -> List[str]: + if not isinstance(aliases, dict): + return [] + used = [] + for canonical, names in aliases.items(): + if not isinstance(canonical, str) or canonical in input_value: + continue + names_value: Any = names + if not isinstance(names_value, (list, tuple)): + continue + alias = next( + ( + name + for name in names_value + if isinstance(name, str) and name in input_value + ), + None, + ) + if ( + alias + and len(alias) <= _MAX_KEY_LENGTH + and len(canonical) <= _MAX_KEY_LENGTH + ): + used.append(f"{alias}:{canonical}") + return sorted(used)[:_MAX_INPUT_KEYS] + + +def get_tool_input_properties( + input_value: Any, + input_schema: Any = None, + options: Optional[ToolInputOptions] = None, +) -> JsonRecord: + """Describe tool arguments without reading their values. + + The schema and aliases must come from the server. Unknown names become one + ``[redacted]`` entry because an argument name can contain private data. + """ + try: + if type(input_value) is not dict: + return {} + resolved = options or ToolInputOptions() + properties = _declared_properties(input_schema) + known = set(properties) | set(_alias_names(resolved.input_aliases)) + keys = [ + key + for key in input_value + if isinstance(key, str) and (key in known or key not in _ANALYTICS_KEYS) + ] + declared: List[str] = [] + undeclared: List[str] = [] + has_redacted = False + for key in keys: + is_declared = key in known + record = is_declared + if resolved.should_record_input_key is not None: + try: + record = ( + resolved.should_record_input_key(key, {"declared": is_declared}) + is True + ) + except Exception: # noqa: BLE001 - analytics callbacks fail closed + record = False + if len(key) <= _MAX_KEY_LENGTH and record: + (declared if is_declared else undeclared).append(key) + else: + has_redacted = True + + visible = [*sorted(declared), *sorted(undeclared)][:_MAX_INPUT_KEYS] + if has_redacted and len(visible) < _MAX_INPUT_KEYS: + visible.append("[redacted]") + aliases_used = _aliases_used(resolved.input_aliases, input_value) + result: JsonRecord = { + PostHogMCPAnalyticsProperty.INPUT_KEYS: visible, + } + if aliases_used: + result[PostHogMCPAnalyticsProperty.INPUT_ALIASES_USED] = aliases_used + return result + except Exception: # noqa: BLE001 - analytics must not change tool dispatch + return {} diff --git a/posthog/mcp/constants.py b/posthog/mcp/constants.py index 85f2fed63..95af83a2e 100644 --- a/posthog/mcp/constants.py +++ b/posthog/mcp/constants.py @@ -87,6 +87,8 @@ class PostHogMCPAnalyticsProperty: IS_ERROR = "$mcp_is_error" INTENT = "$mcp_intent" INTENT_SOURCE = "$mcp_intent_source" + INPUT_ALIASES_USED = "$mcp_input_aliases_used" + INPUT_KEYS = "$mcp_input_keys" LLM_MODEL = "$mcp_llm_model" LLM_MODEL_SOURCE = "$mcp_llm_model_source" LISTED_TOOL_NAMES = "$mcp_listed_tool_names" diff --git a/posthog/mcp/types.py b/posthog/mcp/types.py index 6b8bfc444..fb5ce09fc 100644 --- a/posthog/mcp/types.py +++ b/posthog/mcp/types.py @@ -43,7 +43,10 @@ "FeedbackReport", "FeedbackSentiment", "FeedbackType", + "InputAliasMap", "PreparedToolCall", + "ShouldRecordInputKeyFn", + "ToolInputOptions", ] JsonRecord = Dict[str, Any] @@ -174,6 +177,23 @@ class FeedbackReport: ] # (request, extra) -> Optional[UserIdentity] | awaitable IntentFallbackFn = Callable[..., Any] # (request, extra) -> Optional[str] | awaitable EventPropertiesFn = Callable[..., Any] # (request, extra) -> Optional[dict] | awaitable +InputAliasMap = Dict[str, List[str]] + + +class ToolInputKeyDetails(TypedDict): + declared: bool + + +ShouldRecordInputKeyFn = Callable[[str, ToolInputKeyDetails], bool] +ResolveInputAliasesFn = Callable[[str], Optional[InputAliasMap]] + + +@dataclass +class ToolInputOptions: + """Configure safe tool input field-name capture.""" + + should_record_input_key: Optional[ShouldRecordInputKeyFn] = None + input_aliases: Optional[InputAliasMap] = None @dataclass @@ -212,6 +232,12 @@ class MCPAnalyticsOptions: # so new integrations should enable only one of the two. New field appended # last: positional construction of the earlier fields must keep working. collect_feedback: Union[bool, CollectFeedbackOptions] = False + # Decide which top-level argument names `$mcp_input_keys` records. The + # default records only names that the server's input schema declares. + should_record_input_key: Optional[ShouldRecordInputKeyFn] = None + # Return the alternative names that one tool accepts. The SDK records alias + # use but does not change tool arguments. + resolve_input_aliases: Optional[ResolveInputAliasesFn] = None @dataclass diff --git a/posthog/test/mcp/test_features_m4.py b/posthog/test/mcp/test_features_m4.py index 6f0b4ce9c..998703217 100644 --- a/posthog/test/mcp/test_features_m4.py +++ b/posthog/test/mcp/test_features_m4.py @@ -3,6 +3,7 @@ import json import mcp.types as mcp_types +import pytest from mcp.server.fastmcp import FastMCP from mcp.server.lowlevel import Server @@ -124,6 +125,36 @@ async def test_fastmcp_conversation_id_captured(): assert calls and calls[0]["properties"].get("$mcp_conversation_id") # minted +async def test_fastmcp_captures_safe_input_names_and_aliases(): + server = make_fastmcp() + client = FakeClient() + instrument( + server, + client, + MCPAnalyticsOptions( + resolve_input_aliases=lambda name: {"a": ["first"]} + if name == "add" + else None + ), + ) + + list_handler = server._mcp_server.request_handlers[mcp_types.ListToolsRequest] + await list_handler(mcp_types.ListToolsRequest(method="tools/list")) + # Alias telemetry does not normalize arguments. This server does not accept + # the alias, so the call fails after the SDK records the original names. + with pytest.raises(Exception): + await server._tool_manager.call_tool( + "add", + {"first": 1, "b": 2, "unknown@example.com": 3}, + convert_result=True, + ) + await _flush() + + properties = _events(client, "$mcp_tool_call")[0]["properties"] + assert properties["$mcp_input_keys"] == ["b", "first", "[redacted]"] + assert properties["$mcp_input_aliases_used"] == ["first:a"] + + async def test_lowlevel_conversation_id_captured_and_prompt_back(): server = make_lowlevel() client = FakeClient() diff --git a/posthog/test/mcp/test_tool_input.py b/posthog/test/mcp/test_tool_input.py new file mode 100644 index 000000000..7f6ce4984 --- /dev/null +++ b/posthog/test/mcp/test_tool_input.py @@ -0,0 +1,90 @@ +from posthog.mcp import ToolInputOptions, get_tool_input_properties + + +def test_keeps_declared_names_and_redacts_unknown_names(): + input_value = { + "id": "private-value", + "context": "analytics-value", + "llm_model": "example-model", + "conversation_id": "example-conversation", + "person@example.com": True, + } + schema = {"type": "object", "properties": {"id": {}, "context": {}}} + + assert get_tool_input_properties(input_value, schema) == { + "$mcp_input_keys": ["context", "id", "[redacted]"] + } + assert input_value["id"] == "private-value" + + +def test_custom_rule_replaces_the_default_and_keeps_declared_names_first(): + seen = [] + + def record(key, details): + seen.append((key, details["declared"])) + return key.replace("_", "").isalnum() + + result = get_tool_input_properties( + {"id": 1, "experiment_id": 1, "person@example.com": 1}, + {"properties": {"id": {}}}, + ToolInputOptions(should_record_input_key=record), + ) + + assert result == {"$mcp_input_keys": ["id", "experiment_id", "[redacted]"]} + assert ("id", True) in seen + assert ("experiment_id", False) in seen + + +def test_records_alias_names_and_the_first_alias_used(): + result = get_tool_input_properties( + {"experimentId": 1, "experiment_id": 2, "flagKey": "k", "other": 3}, + {"properties": {"id": {}, "key": {}}}, + ToolInputOptions( + input_aliases={ + "id": ["experimentId", "experiment_id"], + "key": ["flagKey"], + } + ), + ) + + assert result == { + "$mcp_input_keys": [ + "experimentId", + "experiment_id", + "flagKey", + "[redacted]", + ], + "$mcp_input_aliases_used": ["experimentId:id", "flagKey:key"], + } + + +def test_canonical_name_prevents_alias_use_record(): + result = get_tool_input_properties( + {"id": 1, "experiment_id": 2}, + {"properties": {"id": {}}}, + ToolInputOptions(input_aliases={"id": ["experiment_id"]}), + ) + + assert result == {"$mcp_input_keys": ["experiment_id", "id"]} + + +def test_bounds_names_and_fails_closed(): + properties = {f"key{i}": {} for i in range(30)} + result = get_tool_input_properties(properties, {"properties": properties}) + assert len(result["$mcp_input_keys"]) == 20 + + assert get_tool_input_properties(None) == {} + assert get_tool_input_properties({"x" * 65: 1}, {"properties": {"x" * 65: {}}}) == { + "$mcp_input_keys": ["[redacted]"] + } + + +def test_callback_error_redacts_the_name(): + def fail(_key, _details): + raise RuntimeError("boom") + + assert get_tool_input_properties( + {"id": 1}, + {"properties": {"id": {}}}, + ToolInputOptions(should_record_input_key=fail), + ) == {"$mcp_input_keys": ["[redacted]"]} diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 95783ef4c..47ea056a8 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -339,6 +339,7 @@ alias posthog.integrations.django.contexts -> posthog.contexts alias posthog.mcp.CaptureEventData -> posthog.mcp.types.CaptureEventData alias posthog.mcp.CollectFeedbackOptions -> posthog.mcp.types.CollectFeedbackOptions alias posthog.mcp.FeedbackReport -> posthog.mcp.types.FeedbackReport +alias posthog.mcp.InputAliasMap -> posthog.mcp.types.InputAliasMap alias posthog.mcp.MCPAnalyticsContextOptions -> posthog.mcp.types.MCPAnalyticsContextOptions alias posthog.mcp.MCPAnalyticsModelOptions -> posthog.mcp.types.MCPAnalyticsModelOptions alias posthog.mcp.MCPAnalyticsModelSource -> posthog.mcp.types.MCPAnalyticsModelSource @@ -352,6 +353,8 @@ alias posthog.mcp.PostHogMcpStatelessSessionMiddleware -> posthog.mcp.asgi.PostH alias posthog.mcp.PreparedToolCall -> posthog.mcp.types.PreparedToolCall alias posthog.mcp.SEND_FEEDBACK_TOOL_NAME -> posthog.mcp.feedback.SEND_FEEDBACK_TOOL_NAME alias posthog.mcp.SessionTokenPayload -> posthog.mcp.session_token.SessionTokenPayload +alias posthog.mcp.ShouldRecordInputKeyFn -> posthog.mcp.types.ShouldRecordInputKeyFn +alias posthog.mcp.ToolInputOptions -> posthog.mcp.types.ToolInputOptions alias posthog.mcp.UserIdentity -> posthog.mcp.types.UserIdentity alias posthog.mcp.__version__ -> posthog.mcp.version.__version__ alias posthog.mcp.asgi.MCP_SESSION_HEADER -> posthog.mcp.session_token.MCP_SESSION_HEADER @@ -785,6 +788,8 @@ attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.FEEDBACK_SUMMARY = ' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.FEEDBACK_TASK_COMPLETED = '$mcp_feedback_task_completed' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.FEEDBACK_TOOL = '$mcp_feedback_tool' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.FEEDBACK_TYPE = '$mcp_feedback_type' +attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.INPUT_ALIASES_USED = '$mcp_input_aliases_used' +attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.INPUT_KEYS = '$mcp_input_keys' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.INTENT = '$mcp_intent' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.INTENT_SOURCE = '$mcp_intent_source' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.IS_ERROR = '$mcp_is_error' @@ -829,6 +834,7 @@ attribute posthog.mcp.types.FeedbackReport.task_completed: Optional[bool] = None attribute posthog.mcp.types.FeedbackReport.tool_name: Optional[str] = None attribute posthog.mcp.types.FeedbackSentiment = Literal['positive', 'neutral', 'negative', 'mixed'] attribute posthog.mcp.types.FeedbackType = Literal['missing_capability', 'issue', 'praise', 'other'] +attribute posthog.mcp.types.InputAliasMap = Dict[str, List[str]] attribute posthog.mcp.types.MCPAnalyticsContextOptions.description: Optional[str] = None attribute posthog.mcp.types.MCPAnalyticsModelOptions.description: Optional[str] = None attribute posthog.mcp.types.MCPAnalyticsModelSource = Literal['client_metadata', 'self_reported'] @@ -844,6 +850,8 @@ attribute posthog.mcp.types.MCPAnalyticsOptions.intent_fallback: Optional[Intent attribute posthog.mcp.types.MCPAnalyticsOptions.logger: Optional[LoggerFn] = None attribute posthog.mcp.types.MCPAnalyticsOptions.missing_capability_tool_name: Optional[str] = None attribute posthog.mcp.types.MCPAnalyticsOptions.report_missing: bool = False +attribute posthog.mcp.types.MCPAnalyticsOptions.resolve_input_aliases: Optional[ResolveInputAliasesFn] = None +attribute posthog.mcp.types.MCPAnalyticsOptions.should_record_input_key: Optional[ShouldRecordInputKeyFn] = None attribute posthog.mcp.types.PreparedToolCall.args: Optional[JsonRecord] = None attribute posthog.mcp.types.PreparedToolCall.feedback_report: Optional[FeedbackReport] = None attribute posthog.mcp.types.PreparedToolCall.intent: Optional[str] = None @@ -852,6 +860,9 @@ attribute posthog.mcp.types.PreparedToolCall.is_feedback: bool = False attribute posthog.mcp.types.PreparedToolCall.is_missing_capability: bool = False attribute posthog.mcp.types.PreparedToolCall.llm_model: Optional[str] = None attribute posthog.mcp.types.PreparedToolCall.llm_model_source: Optional[MCPAnalyticsModelSource] = None +attribute posthog.mcp.types.ShouldRecordInputKeyFn = Callable[[str, ToolInputKeyDetails], bool] +attribute posthog.mcp.types.ToolInputOptions.input_aliases: Optional[InputAliasMap] = None +attribute posthog.mcp.types.ToolInputOptions.should_record_input_key: Optional[ShouldRecordInputKeyFn] = None attribute posthog.mcp.types.UserIdentity.distinct_id: str attribute posthog.mcp.types.UserIdentity.groups: Optional[Dict[str, str]] = None attribute posthog.mcp.types.UserIdentity.properties: Optional[JsonRecord] = None @@ -1051,8 +1062,9 @@ class posthog.mcp.types.CollectFeedbackOptions(tool_name: Optional[str] = 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 = 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.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, should_record_input_key: Optional[ShouldRecordInputKeyFn] = None, resolve_input_aliases: Optional[ResolveInputAliasesFn] = None) 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.ToolInputOptions(should_record_input_key: Optional[ShouldRecordInputKeyFn] = None, input_aliases: Optional[InputAliasMap] = 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) class posthog.poller.Poller(interval, execute, *args, **kwargs) From 2f602c906d7e50a71eb24366059ecd3e59c8a1da Mon Sep 17 00:00:00 2001 From: Georgis Andonis <6430745+gesh@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:06:37 +0300 Subject: [PATCH 2/2] fix(mcp): scope input schemas to tool calls Generated-By: PostHog Desktop Task-Id: dda38523-5b8d-47b6-a11a-64cd59a812fe --- posthog/mcp/_instrument_fastmcp.py | 14 +++++- posthog/mcp/_instrument_lowlevel.py | 65 +++++++++++++++++----------- posthog/mcp/_instrument_v2.py | 35 +++++++++++---- posthog/mcp/_instrumentation.py | 24 ++++------ posthog/mcp/_internal.py | 5 --- posthog/mcp/_tool_input.py | 43 ++++++++++++------ posthog/test/mcp/test_features_m4.py | 34 ++++++++++++++- posthog/test/mcp/test_tool_input.py | 27 ++++++++++++ 8 files changed, 177 insertions(+), 70 deletions(-) diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 28ed09736..f50b818d4 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/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 @@ -132,6 +133,8 @@ async def wrapped( for text in lifecycle.virtual_result_texts(reply) ] + lifecycle = replace(lifecycle, input_schema=_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. @@ -273,7 +276,7 @@ async def list_handler(req: Any) -> Any: duration_ms = (time.monotonic() - start) * 1000 tools = extract_tools(result) # Empty is computed before adding the virtual missing-capability tool. - names, empty = collect_listed_tools(data, tools, lifecycle.session_id) + names, empty = collect_listed_tools(data, tools) injection = resolve_virtual_tool_injection( data, tools, @@ -336,6 +339,15 @@ def _name_owned_by_real_tool(server: Any, name: str) -> Optional[bool]: return None +def _tool_input_schema(server: Any, name: str) -> Optional[Dict[str, Any]]: + """Return the schema from the tool that this server will call.""" + try: + schema = server._tool_manager.get_tool(name).parameters + except Exception: # noqa: BLE001 - analytics must not break the call + return None + return schema if isinstance(schema, dict) else None + + def _tool_owns_param(server: Any, name: str, param: str) -> bool: """True when the tool's own function declares ``param`` — then it's a real tool argument we must neither inject nor strip (the agent's value belongs to the tool).""" diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index aca28949f..a4424e77f 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -187,10 +187,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) @@ -213,6 +213,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 ( @@ -473,7 +474,7 @@ async def handler(req: Any) -> Any: # Zero advertised tools is treated as an errored tools/list before the # virtual missing-capability tool is appended. - names, empty = collect_listed_tools(data, tools, lifecycle.session_id) + names, empty = collect_listed_tools(data, tools) injection = resolve_virtual_tool_injection( data, tools, @@ -566,10 +567,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, @@ -581,17 +583,41 @@ 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 = None 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): + declared, injectable = _schema_view( + schema, dereferenced=_server_dereferences(high_level) + ) + 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: @@ -610,7 +636,7 @@ 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, @@ -618,27 +644,16 @@ async def _registry_view( 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: diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 85ef69597..f31dc1539 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/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 @@ -231,8 +232,7 @@ 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 = (_tool_input_schema_v2(high_level, name) or {}).get("properties") except Exception: # noqa: BLE001 return {} # Fail closed on a malformed schema: the caller does `param in ` in the @@ -241,6 +241,16 @@ def _tool_own_properties_v2(high_level: Any, name: str) -> Dict[str, Any]: return properties if isinstance(properties, dict) else {} +def _tool_input_schema_v2( + high_level: Any, name: str +) -> Optional[Dict[str, Any]]: + try: + schema = high_level._tool_manager.get_tool(name).parameters + except Exception: # noqa: BLE001 - analytics must not break the call + return None + return schema if isinstance(schema, dict) else None + + def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool: """Whether the tool's own JSON schema declares ``param`` — then it's a real tool argument we must neither inject over nor strip. Read from the tool's @@ -319,6 +329,10 @@ async def wrapped( ] return mcp_types.CallToolResult(content=virtual_content) + lifecycle = replace( + lifecycle, input_schema=_tool_input_schema_v2(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). @@ -442,7 +456,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. @@ -465,9 +479,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") @@ -477,7 +491,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: @@ -493,10 +510,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: @@ -522,6 +540,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 @@ -696,7 +715,7 @@ async def handler(ctx: Any, params: Any) -> Any: tools = list(getattr(result, "tools", []) or []) # Empty is computed before adding the virtual missing-capability tool. - names, empty = collect_listed_tools(data, tools, lifecycle.session_id) + names, empty = collect_listed_tools(data, tools) injection = resolve_virtual_tool_injection( data, tools, is_first_page=is_first_listing_page(params) ) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 6d254bb81..8c02139dd 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -448,6 +448,7 @@ class ToolCallLifecycle: feedback_name: Optional[str] conversation_id: Optional[str] minted_conversation_id: bool + input_schema: Any @property def is_missing_capability(self) -> bool: @@ -547,6 +548,7 @@ async def record_error(self, error: Any, duration_ms: float) -> None: protocol_version=self.protocol_version, conversation_id=conversation_id, extra=self.extra, + input_schema=self.input_schema, ) async def record_result( @@ -568,6 +570,7 @@ async def record_result( protocol_version=self.protocol_version, conversation_id=conversation_id, extra=self.extra, + input_schema=self.input_schema, ) @@ -584,6 +587,7 @@ def start_tool_call_lifecycle( client_version: Optional[str], protocol_version: Optional[str], extra: Dict[str, Any], + input_schema: Any = None, ) -> ToolCallLifecycle: """Resolve adapter-independent policy for a tool call without dispatching it.""" enabled = enabled_virtual_tool_names(data) @@ -617,6 +621,7 @@ def start_tool_call_lifecycle( feedback_name=feedback_name, conversation_id=conversation_id, minted_conversation_id=minted, + input_schema=input_schema, ) @@ -636,6 +641,7 @@ async def record_tool_call( protocol_version: Optional[str] = None, conversation_id: Optional[str] = None, extra: Optional[Dict[str, Any]] = None, + input_schema: Any = None, ) -> None: # Analytics must never change what the tool returns or raises: any failure # building/publishing the event is logged and swallowed here. @@ -684,12 +690,11 @@ async def record_tool_call( input_aliases = data.options.resolve_input_aliases(name) except Exception as err: # noqa: BLE001 - analytics callbacks are isolated log(f"Warning: resolve_input_aliases failed for tool {name}: {err}") - schema = data.tool_input_schemas.get(session_id, {}).get(name) event["properties"] = { **(props or {}), **get_tool_input_properties( arguments or {}, - schema, + input_schema, ToolInputOptions( should_record_input_key=data.options.should_record_input_key, input_aliases=input_aliases, @@ -990,12 +995,9 @@ def read_tool_category(tool: Any) -> Optional[str]: return None -def collect_listed_tools( - data: MCPAnalyticsData, tools: list, session_id: Optional[str] = None -) -> tuple[List[str], bool]: +def collect_listed_tools(data: MCPAnalyticsData, tools: list) -> tuple[List[str], bool]: """Cache common tool metadata and return the pre-injection listing summary.""" names = [] - schemas = data.tool_input_schemas.get(session_id, {}) if session_id else None for tool in tools: names.append(tool.name) if getattr(tool, "description", None): @@ -1003,16 +1005,6 @@ def collect_listed_tools( category = read_tool_category(tool) if category: data.tool_categories[tool.name] = category - if schemas is not None: - schema = getattr(tool, "inputSchema", None) - if schema is None: - schema = getattr(tool, "input_schema", None) - schemas[tool.name] = schema - if session_id and schemas is not None: - data.tool_input_schemas[session_id] = schemas - data.tool_input_schemas.move_to_end(session_id) - while len(data.tool_input_schemas) > 1000: - data.tool_input_schemas.popitem(last=False) return names, not tools diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 61ea145ad..b10a9ef60 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -88,11 +88,6 @@ class MCPAnalyticsData: identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) tool_descriptions: Dict[str, str] = field(default_factory=dict) - # Original tool input schemas by PostHog session. This keeps field-name - # privacy decisions isolated when a server advertises user-specific tools. - tool_input_schemas: "OrderedDict[str, Dict[str, Any]]" = field( - default_factory=OrderedDict - ) # True only when PostHog added llm_model to this tool's advertised schema. # Missing/False fails closed so an application-owned field is never read or stripped. tool_model_parameter_injected: Dict[str, bool] = field(default_factory=dict) diff --git a/posthog/mcp/_tool_input.py b/posthog/mcp/_tool_input.py index 0596000f4..1095c57a0 100644 --- a/posthog/mcp/_tool_input.py +++ b/posthog/mcp/_tool_input.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional from .constants import PostHogMCPAnalyticsProperty from .types import InputAliasMap, JsonRecord, ToolInputOptions @@ -32,13 +32,20 @@ def _alias_names(aliases: Optional[InputAliasMap]) -> List[str]: def _aliases_used( - aliases: Optional[InputAliasMap], input_value: Dict[str, Any] + aliases: Optional[InputAliasMap], + input_value: Dict[str, Any], + can_record: Callable[[str], bool], ) -> List[str]: if not isinstance(aliases, dict): return [] used = [] for canonical, names in aliases.items(): - if not isinstance(canonical, str) or canonical in input_value: + if ( + not isinstance(canonical, str) + or canonical in input_value + or len(canonical) > _MAX_KEY_LENGTH + or not can_record(canonical) + ): continue names_value: Any = names if not isinstance(names_value, (list, tuple)): @@ -47,15 +54,14 @@ def _aliases_used( ( name for name in names_value - if isinstance(name, str) and name in input_value + if isinstance(name, str) + and name in input_value + and len(name) <= _MAX_KEY_LENGTH + and can_record(name) ), None, ) - if ( - alias - and len(alias) <= _MAX_KEY_LENGTH - and len(canonical) <= _MAX_KEY_LENGTH - ): + if alias: used.append(f"{alias}:{canonical}") return sorted(used)[:_MAX_INPUT_KEYS] @@ -84,8 +90,12 @@ def get_tool_input_properties( declared: List[str] = [] undeclared: List[str] = [] has_redacted = False - for key in keys: - is_declared = key in known + recording_decisions: Dict[tuple[str, bool], bool] = {} + + def can_record(key: str, is_declared: bool) -> bool: + cache_key = (key, is_declared) + if cache_key in recording_decisions: + return recording_decisions[cache_key] record = is_declared if resolved.should_record_input_key is not None: try: @@ -95,7 +105,12 @@ def get_tool_input_properties( ) except Exception: # noqa: BLE001 - analytics callbacks fail closed record = False - if len(key) <= _MAX_KEY_LENGTH and record: + recording_decisions[cache_key] = record + return record + + for key in keys: + is_declared = key in known + if len(key) <= _MAX_KEY_LENGTH and can_record(key, is_declared): (declared if is_declared else undeclared).append(key) else: has_redacted = True @@ -103,7 +118,9 @@ def get_tool_input_properties( visible = [*sorted(declared), *sorted(undeclared)][:_MAX_INPUT_KEYS] if has_redacted and len(visible) < _MAX_INPUT_KEYS: visible.append("[redacted]") - aliases_used = _aliases_used(resolved.input_aliases, input_value) + aliases_used = _aliases_used( + resolved.input_aliases, input_value, lambda key: can_record(key, True) + ) result: JsonRecord = { PostHogMCPAnalyticsProperty.INPUT_KEYS: visible, } diff --git a/posthog/test/mcp/test_features_m4.py b/posthog/test/mcp/test_features_m4.py index 998703217..b8e4da5e3 100644 --- a/posthog/test/mcp/test_features_m4.py +++ b/posthog/test/mcp/test_features_m4.py @@ -138,8 +138,6 @@ async def test_fastmcp_captures_safe_input_names_and_aliases(): ), ) - list_handler = server._mcp_server.request_handlers[mcp_types.ListToolsRequest] - await list_handler(mcp_types.ListToolsRequest(method="tools/list")) # Alias telemetry does not normalize arguments. This server does not accept # the alias, so the call fails after the SDK records the original names. with pytest.raises(Exception): @@ -155,6 +153,38 @@ async def test_fastmcp_captures_safe_input_names_and_aliases(): assert properties["$mcp_input_aliases_used"] == ["first:a"] +async def test_fastmcp_keeps_input_names_after_conversation_anchoring(): + server = make_fastmcp() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + handle = "0198d3a7-1111-7222-8333-444455556666" + + await server._tool_manager.call_tool( + "add", + {"a": 1, "b": 2, "conversation_id": handle}, + convert_result=True, + ) + await _flush() + + properties = _events(client, "$mcp_tool_call")[0]["properties"] + assert properties["$mcp_input_keys"] == ["a", "b"] + + +async def test_lowlevel_does_not_reuse_a_listed_schema_for_input_names(): + server = make_lowlevel() + client = FakeClient() + instrument(server, client) + + list_handler = server.request_handlers[mcp_types.ListToolsRequest] + await list_handler(mcp_types.ListToolsRequest(method="tools/list")) + call_handler = server.request_handlers[mcp_types.CallToolRequest] + await call_handler(_call_request("echo", {"msg": "hi"})) + await _flush() + + properties = _events(client, "$mcp_tool_call")[0]["properties"] + assert properties["$mcp_input_keys"] == ["[redacted]"] + + async def test_lowlevel_conversation_id_captured_and_prompt_back(): server = make_lowlevel() client = FakeClient() diff --git a/posthog/test/mcp/test_tool_input.py b/posthog/test/mcp/test_tool_input.py index 7f6ce4984..89a637287 100644 --- a/posthog/test/mcp/test_tool_input.py +++ b/posthog/test/mcp/test_tool_input.py @@ -68,6 +68,33 @@ def test_canonical_name_prevents_alias_use_record(): assert result == {"$mcp_input_keys": ["experiment_id", "id"]} +def test_alias_telemetry_uses_the_recording_rule(): + result = get_tool_input_properties( + {"privateAlias": 1}, + {"properties": {"id": {}}}, + ToolInputOptions( + input_aliases={"id": ["privateAlias"]}, + should_record_input_key=lambda key, _details: key != "privateAlias", + ), + ) + + assert result == {"$mcp_input_keys": ["[redacted]"]} + + +def test_long_alias_does_not_hide_a_later_valid_alias(): + long_alias = "x" * 65 + result = get_tool_input_properties( + {long_alias: 1, "validAlias": 2}, + {"properties": {"id": {}}}, + ToolInputOptions(input_aliases={"id": [long_alias, "validAlias"]}), + ) + + assert result == { + "$mcp_input_keys": ["validAlias", "[redacted]"], + "$mcp_input_aliases_used": ["validAlias:id"], + } + + def test_bounds_names_and_fails_closed(): properties = {f"key{i}": {} for i in range(30)} result = get_tool_input_properties(properties, {"properties": properties})