diff --git a/.sampo/changesets/mcp-virtual-tool-first-page.md b/.sampo/changesets/mcp-virtual-tool-first-page.md new file mode 100644 index 000000000..f2dbac2e4 --- /dev/null +++ b/.sampo/changesets/mcp-virtual-tool-first-page.md @@ -0,0 +1,11 @@ +--- +pypi/posthog: patch +--- + +Fix MCP analytics virtual tool injection on a paginated `tools/list`. `get_more_tools` was added to every page and `send_feedback` only to the last; both now go on the first page, matching `@posthog/mcp`. + +- A real tool of yours sharing a virtual tool's name now wins instead of being silently swallowed, and the warning names the option that renames PostHog's — `missing_capability_tool_name` or `collect_feedback`'s `tool_name`. Warnings also reach the `posthog.mcp` logger, so you see them without setting the `logger` option. +- If PostHog cannot tell whether a name is yours or its own, the call is delegated to your server rather than intercepted. +- A server that returns the same `tools/list` result object on every request no longer has PostHog's injected tool read back as a collision, which stopped both tools being advertised from the second listing on. + +On a paginated listing, the first page's `$mcp_tools_list.listed_tool_names` now carries the virtual tools even when a next page follows. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index ece6b5765..a8c8dd168 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -134,8 +134,8 @@ and adds no tools. The tool covers what `report_missing` covers (as feedback_type `missing_capability`), so new integrations should enable only one of the two. -If a real tool already uses the name, the SDK logs a warning, does not inject -the virtual tool, and never intercepts the real tool. +See [Virtual tools on a paginated `tools/list`](#virtual-tools-on-a-paginated-toolslist) +for name collisions and the page rule. Use the object form to rename the tool, declare host-specific fields, or route reports to a real backend: @@ -184,6 +184,70 @@ if call.is_feedback: `on_feedback` is ignored on this path — the dispatcher routes reports itself via `call.feedback_report`. +On this path you own the page rule, because you pass the switches per call. Pass +them for the first page only, and pass your own tool as `original_tool` so a real +tool by a virtual tool's name wins: + +```python +first_page = request.params.get("cursor") is None +tools = posthog.prepare_tool_list( + page_tools, report_missing=first_page, collect_feedback=first_page +) + +call = posthog.prepare_tool_call( + tool_name, raw_args, original_tool=my_tools.get(tool_name) +) +``` + +## Virtual tools on a paginated `tools/list` + +`instrument()` advertises up to two tools of its own: `get_more_tools` +(`report_missing=True`) and `send_feedback` (`collect_feedback=True`). + +A client concatenates every `tools/list` page into one list, so each virtual tool +is appended to the **first page only** — the page every client reads, including +clients that never follow `nextCursor`. "First page" means a request with no +cursor; an empty string is a valid opaque cursor, so `cursor: ""` is a +continuation page. + +### Tool name collisions + +Your tools win when the SDK can see them. Rename the SDK's to keep both: + +```python +MCPAnalyticsOptions( + report_missing=True, + missing_capability_tool_name="find_posthog_tools", + collect_feedback=CollectFeedbackOptions(tool_name="tell_posthog"), +) +``` + +- A real tool using the name **on the first page** wins: the SDK warns, injects + nothing, and never intercepts yours. +- A real tool that appears **only on a later page** is shadowed — page one cannot + see page two, so the SDK's tool is already advertised and calls to the name + reach it. The SDK warns when that page is served. `@posthog/mcp` is the same. +- Configuring **both** virtual tools with one name advertises only + `get_more_tools`, and warns. + +Warnings go to the `logger` option *and* the `posthog.mcp` standard-library +logger, so a default-configured host sees them on stderr. + +At call time the SDK checks ownership again, covering the process that serves a +call without having served a listing — the ordinary multi-pod case. FastMCP and +v2 `MCPServer` are asked via their tool registry; a raw low-level server has +none, so the SDK calls your own `tools/list` handler, once per call to a virtual +tool's name and never for ordinary traffic. If that check cannot answer — a +registry lookup or a listing handler raised, or you registered `tools/list` +after `instrument()` — the call is delegated to your server rather than +intercepted, and the reason is logged: guessing the other way would swallow a +real tool of yours silently. + +This check is the only ownership signal used at call time, so a server that +serves **different tool sets to different callers** from one instrumented +instance is handled correctly: nothing is carried over from whichever listing +happened to be served last. + ## Stateless / multi-pod servers A stateless MCP server issues no session id, so `$session_id` fragments across pods diff --git a/posthog/mcp/_conversation_id.py b/posthog/mcp/_conversation_id.py index 4602897ad..b1e024b3b 100644 --- a/posthog/mcp/_conversation_id.py +++ b/posthog/mcp/_conversation_id.py @@ -78,7 +78,7 @@ def resolve_conversation_id( enabled: bool, args: Any, tool_name: Optional[str], - missing_capability_tool_name: str, + missing_capability_tool_name: Optional[str], feedback_tool_name: Optional[str] = None, ) -> Tuple[Optional[str], bool]: """Return ``(conversation_id, minted)``. Disabled, get_more_tools, or @@ -86,12 +86,19 @@ def resolve_conversation_id( → ``(value, False)``; anything else (omitted, or a value the agent made up) → ``(new uuid, True)``. + Either virtual tool's name arrives as ``None`` when that tool is disabled, + so a real application tool by the same name mints and echoes a handle like + any other tool's. + Lowercased on the way in: the shape test is case-insensitive but the hash behind ``$session_id`` is not, so an uppercased echo (some hosts normalise uuids) would land in a different session than the call that minted it.""" if ( not enabled - or tool_name == missing_capability_tool_name + or ( + missing_capability_tool_name is not None + and tool_name == missing_capability_tool_name + ) or (feedback_tool_name is not None and tool_name == feedback_tool_name) ): return None, False diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 61d33c27a..d12846eaf 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -30,14 +30,13 @@ from ._instrument_lowlevel import _wrap_resource_requests from ._instrumentation import ( _to_jsonable, - append_get_more_tools, - append_send_feedback, + apply_virtual_tool_injection, collect_listed_tools, extract_tools, - listing_has_next_page, + is_first_listing_page, mutate_tool_schema, - refresh_feedback_shadow, request_to_dict, + resolve_virtual_tool_injection, resolve_session_and_client, start_tool_call_lifecycle, start_tools_list_lifecycle, @@ -50,7 +49,7 @@ ) 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 +from .tools import get_more_tools_result_text _WRAPPED_FLAG = "__posthog_mcp_wrapped__" @@ -115,15 +114,15 @@ async def wrapped( }, ) - if lifecycle.is_missing_capability: + if lifecycle.is_missing_capability and ( + _name_owned_by_real_tool(server, name) is False + ): await lifecycle.record_missing_capability() return [ mcp_types.TextContent(type="text", text=get_more_tools_result_text()) ] - if lifecycle.is_feedback and not _feedback_name_owned_by_real_tool( - server, name - ): + if lifecycle.is_feedback and (_name_owned_by_real_tool(server, name) is False): reply = await lifecycle.record_feedback() return [mcp_types.TextContent(type="text", text=reply)] @@ -206,6 +205,7 @@ def _inject_tool_schemas(server: Any, data: MCPAnalyticsData, tools: list) -> No schema_attribute="inputSchema", owns_context=_tool_owns_context(server, tool.name), context_required=True, + is_sdk_virtual_tool=False, ) @@ -229,10 +229,6 @@ async def list_handler(req: Any) -> Any: if req is None: result = await original(req) tools = extract_tools(result) - # Refresh the collision flag here too: this pass sees the real tool - # registry, so a real tool named like the feedback tool is detected - # before any client-facing listing. - refresh_feedback_shadow(data, tools) _inject_tool_schemas(server, data, tools) return result @@ -272,19 +268,17 @@ async def list_handler(req: Any) -> Any: tools = extract_tools(result) # Empty is computed before adding the virtual missing-capability tool. names, empty = collect_listed_tools(data, tools) - feedback_name = refresh_feedback_shadow(data, tools) + injection = resolve_virtual_tool_injection( + data, + tools, + is_first_page=is_first_listing_page(getattr(req, "params", None)), + ) _inject_tool_schemas(server, data, tools) - if data.options.report_missing: - missing_name = resolve_missing_capability_tool_name(data.options) - if not any(t.name == missing_name for t in tools): - append_get_more_tools(result, missing_name, data) - names.append(missing_name) - - if feedback_name is not None and not listing_has_next_page(result): - append_send_feedback(result, data) - names.append(feedback_name) + result = apply_virtual_tool_injection( + result, injection, names, data, schema_field="inputSchema" + ) await lifecycle.record_result( names=names, @@ -322,14 +316,21 @@ def _inject_prompt_back(result: Any, conversation_id: str) -> Any: return result -def _feedback_name_owned_by_real_tool(server: Any, name: str) -> bool: - """Live registry probe so a real tool by the feedback tool's name is never - shadowed even before the first listing refreshes the collision flag.""" - try: - tool_manager = getattr(server, "_tool_manager", None) - return tool_manager is not None and tool_manager.get_tool(name) is not None - except Exception: # noqa: BLE001 - unknown tool -> the name is not owned +def _name_owned_by_real_tool(server: Any, name: str) -> Optional[bool]: + """Live registry probe, so a real tool by a virtual tool's name is never + shadowed. Tri-state like its low-level twin: ``None`` when the lookup failed + rather than answered, and callers must not intercept on it.""" + tool_manager = getattr(server, "_tool_manager", None) + if tool_manager is None: return False + try: + return tool_manager.get_tool(name) is not None + except Exception as err: # noqa: BLE001 - analytics must not break the call + log( + f'Warning: could not determine whether "{name}" is a real tool of ' + f"yours; delegating the call to your server - {err}" + ) + return None def _tool_owns_param(server: Any, name: str, param: str) -> bool: diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 786597930..3f4cd0945 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -16,7 +16,8 @@ import inspect import time -from typing import Any, Optional, Tuple +from functools import lru_cache +from typing import Any, Optional, Set, Tuple import mcp.types as mcp_types @@ -25,26 +26,27 @@ from ._event_types import MCPAnalyticsEventType from ._instrumentation import ( _to_jsonable, - append_get_more_tools, - append_send_feedback, + advertised_tool_names, + apply_virtual_tool_injection, collect_listed_tools, extract_tools, - listing_has_next_page, + is_first_listing_page, mutate_tool_schema, prepare_request, + raw_listing_owns_tool_name, record_resource_request, - refresh_feedback_shadow, request_to_dict, resource_listing_response, resolve_session_and_client, + resolve_virtual_tool_injection, start_tool_call_lifecycle, start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData from ._model_parameters import request_meta_from_context 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 +from .logger import log, warn +from .tools import get_more_tools_result_text _WRAPPED_FLAG = "__posthog_mcp_wrapped__" @@ -206,7 +208,9 @@ async def handler(req: Any) -> Any: extra={"session_id": mcp_session_id, "ctx": _request_context(server)}, ) - if lifecycle.is_missing_capability: + if lifecycle.is_missing_capability and ( + await _name_owned_by_real_tool(high_level, data, name, server) is False + ): await lifecycle.record_missing_capability() return mcp_types.ServerResult( mcp_types.CallToolResult( @@ -219,8 +223,8 @@ async def handler(req: Any) -> Any: ) ) - if lifecycle.is_feedback and not await _feedback_name_owned_by_real_tool( - high_level, name + if lifecycle.is_feedback and ( + await _name_owned_by_real_tool(high_level, data, name, server) is False ): reply = await lifecycle.record_feedback() return mcp_types.ServerResult( @@ -333,6 +337,7 @@ def _inject_tool_schemas( schema_attribute="inputSchema", owns_context=schema_has_param(schema, "context"), context_required=context_required, + is_sdk_virtual_tool=False, ) @@ -344,6 +349,47 @@ def _wrap_list_tools( if original is None or getattr(original, _WRAPPED_FLAG, False): return + async def probe_raw_tool_names(_ctx: Any = None) -> Optional[Set[str]]: + """The names the host's own handler advertises on its first page. Calls + ``original``, not the wrapper, so the probe never recurses into + instrumentation or appends a virtual tool. Read by + ``_name_owned_by_real_tool`` on raw low-level servers, which have no + tool registry to ask instead.""" + # Late-bound on purpose: a host that replaces or removes tools/list + # *after* instrument() leaves `original` holding a catalogue the client + # never sees, and a confident answer from it would swallow a real tool. + # Report the question as unanswerable instead, so the call is delegated. + # `@posthog/mcp` re-captures the handler for the same reason. + current = handlers.get(mcp_types.ListToolsRequest) + if current is None or not getattr(current, _WRAPPED_FLAG, False): + # The virtual tools stay advertised -- a handler chained in front + # of ours still runs our injection -- but nothing is intercepted + # behind them, and a silent stop is invisible in the captured data. + if not data.warned_foreign_list_handler: + data.warned_foreign_list_handler = True + warn( + "Warning: your tools/list handler was replaced or removed " + "after instrument(), so PostHog can no longer tell whether a " + "tool " + "name is yours. Calls to PostHog's virtual tools are " + "delegated to your server, so no $mcp_missing_capability or " + "$mcp_feedback events are captured. Call instrument() after " + "registering your handlers." + ) + return None + result = await original(mcp_types.ListToolsRequest(method="tools/list")) + tools = extract_tools(result) + # `original` is usually the SDK's own list_tools decorator, which rebuilds + # `Server._tool_cache` from these un-injected schemas every time it runs. + # That cache is what the SDK validates real tool arguments against, so + # 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) + return advertised_tool_names(tools) + + data.raw_tool_names_probe = probe_raw_tool_names + async def handler(req: Any) -> Any: # The server calls the handler with None to populate its tool cache. # Skip analytics there — but still inject, because that cache is the @@ -355,10 +401,6 @@ async def handler(req: Any) -> Any: if req is None: result = await original(req) tools = extract_tools(result) - # Refresh the collision flag here too: this pass sees the real tool - # registry, so a real tool named like the feedback tool is detected - # before any client-facing listing. - refresh_feedback_shadow(data, tools) _inject_tool_schemas(data, tools, context_required=context_required) return result @@ -401,19 +443,17 @@ 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) - feedback_name = refresh_feedback_shadow(data, tools) + injection = resolve_virtual_tool_injection( + data, + tools, + is_first_page=is_first_listing_page(getattr(req, "params", None)), + ) _inject_tool_schemas(data, tools, context_required=context_required) - if data.options.report_missing: - missing_name = resolve_missing_capability_tool_name(data.options) - if not any(t.name == missing_name for t in tools): - append_get_more_tools(result, missing_name, data) - names.append(missing_name) - - if feedback_name is not None and not listing_has_next_page(result): - append_send_feedback(result, data) - names.append(feedback_name) + result = apply_virtual_tool_injection( + result, injection, names, data, schema_field="inputSchema" + ) await lifecycle.record_result( names=names, @@ -428,17 +468,67 @@ async def handler(req: Any) -> Any: handlers[mcp_types.ListToolsRequest] = handler -async def _feedback_name_owned_by_real_tool(high_level: Any, name: str) -> bool: - """Live registry probe on the standalone-fastmcp path, so a real tool by the - feedback tool's name is never shadowed even before the first listing refreshes - the collision flag. Raw low-level servers have no registry to probe; they rely - on the listing-derived flag alone.""" - if high_level is None: - return False +async def _name_owned_by_real_tool( + high_level: Any, data: MCPAnalyticsData, name: str, server: Any +) -> Optional[bool]: + """Whether a real application tool owns ``name``, so a virtual tool never + shadows it. Kind-agnostic: a lookup by name, shared by both virtual tools + rather than twin helpers that can drift. The other two adapters keep the + same tri-state contract against their own registries. + + On the standalone-fastmcp path the tool registry answers authoritatively. A + raw low-level server has no registry, so it asks the host's own tools/list + handler instead. + + Known limit, and *not* one a fallback can close: a ``FastMCP`` is an + ``AggregateProvider``, which gathers its providers with + ``return_exceptions=True`` and drops the failures, so an unreachable + mounted or proxied sub-server reads back as a plain ``None`` -- + indistinguishable from "no such tool" -- and we treat the name as free. + Consulting ``list_tools`` does not help: the same failure is dropped from + the listing too (``_collect_list_results``), and fastmcp 3.x exposes no + error strategy to opt out of. Narrower than it reads -- during the outage + the host's tool is absent from ``tools/list`` as well, so only a provider + that recovers between this check and dispatch loses a call that would have + worked. The remedy stays the documented one: rename PostHog's tool.""" + if high_level is not None: + try: + return await high_level.get_tool(name) is not None + except Exception as err: # noqa: BLE001 - analytics must not break the call + if isinstance(err, _tool_lookup_not_found_errors()): + return False + # The lookup failed rather than answered, so not "the name is free" + # -- guessing that would swallow a real tool of theirs. What reaches + # here is the visibility, transform and auth work layered on top of + # the providers; a provider failure never does (see the docstring). + log( + f'Warning: could not determine whether "{name}" is a real tool of ' + f"yours; delegating the call to your server - {err}" + ) + return None + # May be None: see `raw_listing_owns_tool_name`. Callers intercept only on a + # definite False. + return await raw_listing_owns_tool_name(data, name, server) + + +@lru_cache(maxsize=1) +def _tool_lookup_not_found_errors() -> Tuple[type, ...]: + """The fastmcp exceptions that mean "no live tool by that name" -- an answer. + Anything else out of ``get_tool`` is the lookup itself failing. Empty when + fastmcp is absent or has moved them, which makes every failure delegate: + the safe direction.""" try: - return await high_level.get_tool(name) is not None - except Exception: # noqa: BLE001 - unknown tool -> the name is not owned - return False + from fastmcp import exceptions + except Exception: # noqa: BLE001 - no fastmcp on this path + return () + return tuple( + err + for err in ( + getattr(exceptions, "NotFoundError", None), + getattr(exceptions, "DisabledError", None), + ) + if isinstance(err, type) and issubclass(err, BaseException) + ) async def _tool_owned_injected_keys(high_level: Any, name: str) -> set: diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 66a28e77b..21e22424b 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -32,7 +32,7 @@ import time from collections.abc import Mapping -from typing import Any, Dict, FrozenSet, Optional, Tuple +from typing import Any, Dict, FrozenSet, Optional, Set, Tuple import mcp.types as mcp_types @@ -41,19 +41,21 @@ from ._event_types import MCPAnalyticsEventType from ._instrumentation import ( _to_jsonable, + advertised_tool_names, + apply_virtual_tool_injection, collect_listed_tools, - listing_has_next_page, + is_first_listing_page, mutate_tool_schema, params_to_request_dict, prepare_request, + raw_listing_owns_tool_name, record_resource_request, - refresh_feedback_shadow, resource_listing_response, resolve_session_and_client, + resolve_virtual_tool_injection, start_tool_call_lifecycle, start_tools_list_lifecycle, ) -from .feedback import get_feedback_tool_descriptor, resolve_collect_feedback_options from ._internal import MCPAnalyticsData from ._model_parameters import ( can_inject_model_parameter, @@ -65,9 +67,7 @@ from .request_headers import get_request_headers from .session_token import read_mcp_session_header from .tools import ( - build_report_missing_descriptor, get_more_tools_result_text, - resolve_missing_capability_tool_name, ) _WRAPPED_FLAG = "__posthog_mcp_wrapped__" @@ -298,7 +298,9 @@ async def wrapped( extra={"session_id": mcp_session_id, "ctx": ctx}, ) - if lifecycle.is_missing_capability: + if lifecycle.is_missing_capability and ( + _name_owned_by_real_tool_v2(server, name) is False + ): await lifecycle.record_missing_capability() return mcp_types.CallToolResult( content=[ @@ -308,8 +310,8 @@ async def wrapped( ] ) - if lifecycle.is_feedback and not _feedback_name_owned_by_real_tool_v2( - server, name + if lifecycle.is_feedback and ( + _name_owned_by_real_tool_v2(server, name) is False ): reply = await lifecycle.record_feedback() return mcp_types.CallToolResult( @@ -518,7 +520,11 @@ async def handler(ctx: Any, params: Any) -> Any: extra={"session_id": mcp_session_id, "ctx": ctx}, ) - if lifecycle.is_missing_capability: + # No tool registry on a raw low-level server, so ownership is settled + # by asking the host's own tools/list handler. + if lifecycle.is_missing_capability and ( + await raw_listing_owns_tool_name(data, name, ctx) is False + ): await lifecycle.record_missing_capability() return mcp_types.CallToolResult( content=[ @@ -528,9 +534,9 @@ async def handler(ctx: Any, params: Any) -> Any: ] ) - # No registry to probe on a raw low-level server; interception relies on - # the listing-derived collision flag alone. - if lifecycle.is_feedback: + if lifecycle.is_feedback and ( + await raw_listing_owns_tool_name(data, name, ctx) is False + ): reply = await lifecycle.record_feedback() return mcp_types.CallToolResult( content=[mcp_types.TextContent(type="text", text=reply)] @@ -649,6 +655,14 @@ def _wrap_v2_list_tools( return original = entry.handler + async def probe_raw_tool_names(ctx: Any = None) -> Optional[Set[str]]: + """The host's own first-page tool names — ``original``, not the wrapper, + so the probe never recurses or appends a virtual tool.""" + result = await original(ctx, None) + return advertised_tool_names(list(getattr(result, "tools", []) or [])) + + data.raw_tool_names_probe = probe_raw_tool_names + async def handler(ctx: Any, params: Any) -> Any: token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) @@ -679,7 +693,9 @@ 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) - feedback_name = refresh_feedback_shadow(data, tools) + injection = resolve_virtual_tool_injection( + data, tools, is_first_page=is_first_listing_page(params) + ) for tool in tools: schema = getattr(tool, "input_schema", None) @@ -694,17 +710,12 @@ async def handler(ctx: Any, params: Any) -> Any: schema_attribute="input_schema", owns_context=owns_context, context_required=context_required, + is_sdk_virtual_tool=False, ) - if data.options.report_missing: - missing_name = resolve_missing_capability_tool_name(data.options) - if not any(t.name == missing_name for t in tools): - _append_get_more_tools_v2(result, missing_name, data) - names.append(missing_name) - - if feedback_name is not None and not listing_has_next_page(result): - _append_send_feedback_v2(result, data) - names.append(feedback_name) + result = apply_virtual_tool_injection( + result, injection, names, data, schema_field="input_schema" + ) await lifecycle.record_result( names=names, @@ -719,58 +730,15 @@ async def handler(ctx: Any, params: Any) -> Any: _replace_handler(server, _LIST_METHOD, handler, entry.params_type) -def _feedback_name_owned_by_real_tool_v2(high_level: Any, name: str) -> bool: - """Live registry probe so a real tool by the feedback tool's name is never - shadowed even before the first listing refreshes the collision flag.""" +def _name_owned_by_real_tool_v2(high_level: Any, name: str) -> Optional[bool]: + """Live registry probe, so a real tool by a virtual tool's name is never + shadowed. Tri-state like its low-level twin: ``None`` when the lookup failed + rather than answered, and callers must not intercept on it.""" try: return high_level._tool_manager.get_tool(name) is not None - except Exception: # noqa: BLE001 - unknown tool -> the name is not owned - return False - - -def _append_send_feedback_v2(result: Any, data: MCPAnalyticsData) -> None: - """Append the send_feedback virtual tool to a v2 ListToolsResult. Callers gate - on :func:`refresh_feedback_shadow` returning a name.""" - options = resolve_collect_feedback_options(data.options.collect_feedback) - if options is None: - return - descriptor = get_feedback_tool_descriptor(options) - tool = mcp_types.Tool( - name=descriptor["name"], - description=descriptor["description"], - input_schema=descriptor["inputSchema"], - annotations=descriptor["annotations"], - ) - # `owns_context=True`: the tool carries its intent in its own summary / - # details arguments, so no `context` parameter is injected — but the - # capture_model pass still runs, so it advertises `llm_model` too. - mutate_tool_schema( - data, - tool, - schema_attribute="input_schema", - owns_context=True, - context_required=True, - ) - tools_list = getattr(result, "tools", None) - if isinstance(tools_list, list): - tools_list.append(tool) - - -def _append_get_more_tools_v2(result: Any, name: str, data: MCPAnalyticsData) -> None: - descriptor = build_report_missing_descriptor(name) - tool = mcp_types.Tool( - name=descriptor["name"], - description=descriptor["description"], - input_schema=descriptor["inputSchema"], - annotations=descriptor["annotations"], - ) - mutate_tool_schema( - data, - tool, - schema_attribute="input_schema", - owns_context=True, - context_required=True, - ) - tools_list = getattr(result, "tools", None) - if isinstance(tools_list, list): - tools_list.append(tool) + except Exception as err: # noqa: BLE001 - analytics must not break the call + log( + f'Warning: could not determine whether "{name}" is a real tool of ' + f"yours; delegating the call to your server - {err}" + ) + return None diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index e657fd009..e3e459e6a 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -14,7 +14,7 @@ import threading from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set from ._capture import capture_event from ._context_parameters import ( @@ -50,9 +50,26 @@ from ._transport_identity import stamp_transport_identity from .session import resolve_session_id, resolve_session_id_with_source from .session_token import SessionTokenPayload, decode_session_id -from .tools import GET_MORE_TOOLS_NAME, resolve_missing_capability_tool_name +from .tools import resolve_missing_capability_tool_name from .types import CollectFeedbackOptions, FeedbackReport +# The virtual tools this SDK advertises into tools/list. Every piece of per-tool +# policy -- enable switch, configured name, warning text, warn-once +# bookkeeping, warning text -- is keyed by kind so the two can't drift apart. +VIRTUAL_TOOL_MISSING_CAPABILITY = "missing_capability" +VIRTUAL_TOOL_FEEDBACK = "feedback" + +# The option that renames each virtual tool, quoted verbatim in the collision +# warnings. +_VIRTUAL_TOOL_RENAME_OPTION = { + VIRTUAL_TOOL_MISSING_CAPABILITY: 'MCPAnalyticsOptions(missing_capability_tool_name="...")', + VIRTUAL_TOOL_FEEDBACK: 'MCPAnalyticsOptions(collect_feedback=CollectFeedbackOptions(tool_name="..."))', +} +_VIRTUAL_TOOL_EVENT = { + VIRTUAL_TOOL_MISSING_CAPABILITY: "$mcp_missing_capability", + VIRTUAL_TOOL_FEEDBACK: "$mcp_feedback", +} + # Keep strong refs to in-flight capture tasks/futures and their lifecycle owners so # they aren't GC'd mid-flight and lifecycle drains can select only their own work. _BACKGROUND_TASKS: Dict[Any, Any] = {} @@ -415,7 +432,9 @@ class ToolCallLifecycle: client_name: Optional[str] client_version: Optional[str] protocol_version: Optional[str] - missing_name: str + # ``None`` when the virtual tool is disabled, so every downstream decision + # treats a real tool by that name like any other tool's. + missing_name: Optional[str] feedback_options: Optional[CollectFeedbackOptions] feedback_name: Optional[str] conversation_id: Optional[str] @@ -423,17 +442,11 @@ class ToolCallLifecycle: @property def is_missing_capability(self) -> bool: - return self.data.options.report_missing and self.name == self.missing_name + return self.missing_name is not None and self.name == self.missing_name @property def is_feedback(self) -> bool: - # Never intercept a name a real application tool owns (fail-open): the - # listing pass records the collision on `feedback_tool_shadowed`. - return ( - self.feedback_name is not None - and self.name == self.feedback_name - and not self.data.feedback_tool_shadowed - ) + return self.feedback_name is not None and self.name == self.feedback_name async def prepare_session(self, conversation_id: Optional[str]) -> str: return await prepare_request( @@ -458,7 +471,7 @@ async def record_missing_capability(self) -> None: await record_missing_capability( self.data, session_id, - tool_name=self.missing_name, + tool_name=self.missing_name or self.name, context=(self.arguments or {}).get("context"), arguments=self.arguments, request_meta=self.request_meta, @@ -549,17 +562,12 @@ def start_tool_call_lifecycle( extra: Dict[str, Any], ) -> ToolCallLifecycle: """Resolve adapter-independent policy for a tool call without dispatching it.""" - missing_name = resolve_missing_capability_tool_name(data.options) + enabled = enabled_virtual_tool_names(data) + missing_name = enabled.get(VIRTUAL_TOOL_MISSING_CAPABILITY) + feedback_name = enabled.get(VIRTUAL_TOOL_FEEDBACK) + # Still needed whatever the name resolution says: parsing the report and + # running the host's `on_feedback` handler read the configured options. feedback_options = resolve_collect_feedback_options(data.options.collect_feedback) - feedback_name = ( - resolve_send_feedback_tool_name(feedback_options) - # Mirrors `ToolCallLifecycle.is_feedback`'s fail-open guard below: once a - # real application tool is known to own this name, conversation-id - # resolution must treat calls to it like any other tool too, not skip - # them as if they were the (shadowed) virtual feedback tool. - if feedback_options is not None and not data.feedback_tool_shadowed - else None - ) conversation_id, minted = resolve_conversation_id( data.options.enable_conversation_id, arguments, @@ -656,104 +664,266 @@ async def record_tool_call( def extract_tools(result: Any) -> list: - """Pull the tool list out of a ListTools ServerResult (a copy — to MUTATE the - real list use ``append_get_more_tools``).""" + """Pull the tool list out of a ListTools ServerResult, as a copy.""" root = getattr(result, "root", result) return list(getattr(root, "tools", []) or []) -def append_get_more_tools(result: Any, name: str, data: MCPAnalyticsData) -> None: - """Append the get_more_tools virtual tool to the real ListToolsResult.tools list.""" - import mcp.types as mcp_types +def append_virtual_tool(result: Any, tool: Any) -> Any: + """Return a copy of a ``tools/list`` result with ``tool`` added. - from .tools import build_report_missing_descriptor + A copy, not an in-place append: a host may return the *same* result object + from every ``tools/list``, and mutating it leaves PostHog's tool sitting in + what later reads back as the host's own catalogue. - descriptor = build_report_missing_descriptor(name) - tool = mcp_types.Tool( - name=descriptor["name"], - description=descriptor["description"], - inputSchema=descriptor["inputSchema"], - annotations=descriptor["annotations"], - ) + 1.x wraps ``ListToolsResult`` in a ``ServerResult`` root model, 2.x returns + it directly. The copy carries every other field, ``nextCursor`` included.""" root = getattr(result, "root", result) tools_list = getattr(root, "tools", None) - if isinstance(tools_list, list): - mutate_tool_schema( - data, - tool, - schema_attribute="inputSchema", - owns_context=True, - context_required=True, - ) - tools_list.append(tool) - - -def refresh_feedback_shadow(data: MCPAnalyticsData, tools: list) -> Optional[str]: - """Refresh the collision flag from this listing's tools. Returns the resolved - feedback tool name when the virtual tool may be appended, ``None`` when the - feature is off or a real application tool owns the name (fail-open: the real - tool is advertised and dispatched untouched). Run before the schema-injection - pass so it reads the fresh flag. - - Sticky for the instrumentation instance's lifetime: a paginated ``tools/list`` - delivers one page per request, so a collision seen on an earlier page must - survive a later page that doesn't list the real tool — otherwise that page - would re-arm interception and swallow the real tool's calls. The trade-off is - deliberate: un-shadowing after the host removes the real tool requires - re-instrumentation.""" - options = resolve_collect_feedback_options(data.options.collect_feedback) - if options is None: - return None - name = resolve_send_feedback_tool_name(options) - if any(getattr(tool, "name", None) == name for tool in tools): - if not data.feedback_tool_shadowed: - log( - f'Warning: Cannot inject agent-feedback tool "{name}" because a real tool ' - "already uses that name. The real tool will not be intercepted." - ) - data.feedback_tool_shadowed = True - return None if data.feedback_tool_shadowed else name + if not isinstance(tools_list, list): + return result + updated = root.model_copy(update={"tools": [*tools_list, tool]}) + return type(result)(updated) if hasattr(result, "root") else updated -def listing_has_next_page(result: Any) -> bool: - """Whether this ``tools/list`` result is a non-final page of a paginated - listing. The virtual feedback tool is only appended to the final page: an - earlier page could advertise it before a later page reveals a real tool by - the same name. Reads both SDK majors' cursor spelling (1.x models expose - ``nextCursor``, 2.x ``next_cursor``).""" - root = getattr(result, "root", result) - return bool(getattr(root, "nextCursor", None) or getattr(root, "next_cursor", None)) +def virtual_tool_descriptor( + data: MCPAnalyticsData, kind: str, name: str +) -> Dict[str, Any]: + """The advertised descriptor for a virtual tool, under its configured name.""" + from .tools import build_report_missing_descriptor + + if kind == VIRTUAL_TOOL_MISSING_CAPABILITY: + return build_report_missing_descriptor(name) + return get_feedback_tool_descriptor( + resolve_collect_feedback_options(data.options.collect_feedback) + ) -def append_send_feedback(result: Any, data: MCPAnalyticsData) -> None: - """Append the send_feedback virtual tool to the real ListToolsResult.tools - list. Callers gate on :func:`refresh_feedback_shadow` returning a name.""" +def append_virtual_tool_by_kind( + result: Any, kind: str, name: str, data: MCPAnalyticsData, *, schema_field: str +) -> Any: + """Add a virtual tool to a ``tools/list`` result and return the result to + serve. ``schema_field`` is the SDK major's spelling of the input schema + field: ``inputSchema`` on 1.x, ``input_schema`` on 2.x. + + ``name`` is passed in rather than re-resolved, so a rename can't drift + between :func:`resolve_virtual_tool_injection`'s decision and this append. + + ``owns_context=True``: a virtual tool carries its intent in its own + arguments, so no ``context`` is injected — but the capture_model pass still + runs, so it advertises ``llm_model``.""" import mcp.types as mcp_types - options = resolve_collect_feedback_options(data.options.collect_feedback) - if options is None: - return - descriptor = get_feedback_tool_descriptor(options) + descriptor = virtual_tool_descriptor(data, kind, name) tool = mcp_types.Tool( - name=descriptor["name"], + name=name, description=descriptor["description"], - inputSchema=descriptor["inputSchema"], annotations=descriptor["annotations"], + **{schema_field: descriptor["inputSchema"]}, ) - root = getattr(result, "root", result) - tools_list = getattr(root, "tools", None) - if isinstance(tools_list, list): - # `owns_context=True`: the tool carries its intent in its own summary / - # details arguments, so no `context` parameter is injected — but the - # capture_model pass still runs, so it advertises `llm_model` too. - mutate_tool_schema( - data, - tool, - schema_attribute="inputSchema", - owns_context=True, - context_required=True, + mutate_tool_schema( + data, + tool, + schema_attribute=schema_field, + owns_context=True, + context_required=True, + is_sdk_virtual_tool=True, + ) + return append_virtual_tool(result, tool) + + +def apply_virtual_tool_injection( + result: Any, + injection: Dict[str, str], + names: List[str], + data: MCPAnalyticsData, + *, + schema_field: str, +) -> Any: + """Append every virtual tool this page won, recording each name for the + ``$mcp_tools_list`` event. Missing-capability first: every call path tests + that kind first, which is what makes the ``duplicate`` rule in + :func:`resolve_virtual_tool_injection` resolve in its favour.""" + for kind in (VIRTUAL_TOOL_MISSING_CAPABILITY, VIRTUAL_TOOL_FEEDBACK): + name = injection.get(kind) + if name is None: + continue + result = append_virtual_tool_by_kind( + result, kind, name, data, schema_field=schema_field + ) + names.append(name) + return result + + +def enabled_virtual_tool_names(data: MCPAnalyticsData) -> Dict[str, str]: + """``{kind: configured name}`` for each virtual tool the options enable, + ignoring collisions. The only place the two enable switches and the two + rename options are read together, so a rename can't drift.""" + names: Dict[str, str] = {} + if data.options.report_missing: + names[VIRTUAL_TOOL_MISSING_CAPABILITY] = resolve_missing_capability_tool_name( + data.options + ) + feedback_options = resolve_collect_feedback_options(data.options.collect_feedback) + if feedback_options is not None: + names[VIRTUAL_TOOL_FEEDBACK] = resolve_send_feedback_tool_name(feedback_options) + return names + + +def is_first_listing_page(params: Any) -> bool: + """Whether a ``tools/list`` *request* is the first page of the listing. + + Per the MCP spec an absent ``cursor`` means "start of the listing". A present + one -- *including the empty string* -- is an opaque value from a previous + page, so it is a continuation. Takes the request params, so both handler + shapes share the rule.""" + return getattr(params, "cursor", None) is None + + +def advertised_tool_names(tools: list) -> Set[str]: + """The names a listing page advertises. The virtual tools are appended to a + *copy* of the host's result, so a page always reads back as the host wrote + it -- see :func:`append_virtual_tool`.""" + return { + name for tool in tools if isinstance(name := getattr(tool, "name", None), str) + } + + +def resolve_virtual_tool_injection( + data: MCPAnalyticsData, + tools: list, + *, + is_first_page: bool, +) -> Dict[str, str]: + """``{kind: name}`` for each virtual tool this listing page appends. A kind + is absent when the feature is off, a real tool owns the name, this is a + continuation page, or the other virtual tool already claimed the name. + + Run after ``collect_listed_tools``, so the virtual tools don't count towards + "this server advertises nothing". + + Stateless: this page's own tools are the whole input. Injection happens on a + first page only, so only a first page's view of the tool set decides + anything. + + * name free on the **first** page -> inject it. + * name taken on the **first** page -> warn, inject nothing. The call path + dispatches it normally and the host's tool wins. + * name taken on a **later** page -> the virtual tool is already advertised + from page one, so the host's tool is shadowed. Warn, naming the rename + option; page one cannot be taken back. + """ + enabled = enabled_virtual_tool_names(data) + if not enabled: + return {} + + listed = advertised_tool_names(tools) + + if not is_first_page: + for kind, name in enabled.items(): + # Only warn about a tool we actually shadowed. If this kind never + # made it onto the first page -- a real tool already held the name + # ("blocked"), or the other virtual tool won it ("duplicate") -- + # then nothing of ours is advertised under it and the host's tool + # runs untouched. Telling them otherwise sends them chasing a bug + # that isn't there. + if name in listed and not any( + (kind, name, variant) in data.warned_virtual_tool_collisions + for variant in ("blocked", "duplicate") + ): + _warn_virtual_tool_collision(data, kind, name, "shadowed") + return {} + + injectable: Dict[str, str] = {} + for kind, name in enabled.items(): + if name in listed: + _warn_virtual_tool_collision(data, kind, name, "blocked") + else: + injectable[kind] = name + + # Both virtual tools configured with one name would advertise it twice and + # dead-letter the feedback path, since every call path checks + # missing-capability first. Keep that precedence and say so. + missing_name = injectable.get(VIRTUAL_TOOL_MISSING_CAPABILITY) + if ( + missing_name is not None + and injectable.get(VIRTUAL_TOOL_FEEDBACK) == missing_name + ): + injectable.pop(VIRTUAL_TOOL_FEEDBACK) + _warn_virtual_tool_collision( + data, VIRTUAL_TOOL_FEEDBACK, missing_name, "duplicate" + ) + + return injectable + + +async def raw_listing_owns_tool_name( + data: MCPAnalyticsData, name: str, ctx: Any = None +) -> Optional[bool]: + """Whether the host's *own* ``tools/list`` handler advertises ``name`` on its + first page, asked at call time. For the raw low-level paths, 1.x and v2, + which have no tool registry to query instead. + + Tri-state. ``True``/``False`` are answers; ``None`` means the question could + not be asked, and callers must not intercept on it -- guessing "not owned" + would swallow a real tool of the host's. Runs once per call to a virtual + tool's name, never for ordinary traffic. + """ + probe = data.raw_tool_names_probe + if probe is None: + return None + try: + names = await probe(ctx) + except Exception as err: # noqa: BLE001 - analytics must not break the call + log( + f'Warning: could not determine whether "{name}" is a real tool of ' + f"yours; delegating the call to your server - {err}" + ) + return None + return None if names is None else name in names + + +def _warn_virtual_tool_collision( + data: MCPAnalyticsData, kind: str, name: str, variant: str +) -> None: + """Warn once per ``(kind, name, variant)`` for the life of the server's + tracking state, so a client that re-lists tools on every turn doesn't flood + the log with the same misconfiguration.""" + key = (kind, name, variant) + if key in data.warned_virtual_tool_collisions: + return + data.warned_virtual_tool_collisions.add(key) + warn(virtual_tool_collision_message(kind, name, variant)) + + +def virtual_tool_collision_message( + kind: str, name: str, variant: str, *, rename_option: Optional[str] = None +) -> str: + """The warning text for a virtual-tool name collision. Always names the option + that renames PostHog's tool -- a warning without its own remedy gets ignored. + ``rename_option`` overrides the ``instrument()`` spelling for hosts on the + ``PostHogMCP`` dispatcher path.""" + remedy = rename_option or _VIRTUAL_TOOL_RENAME_OPTION[kind] + if variant == "shadowed": + return ( + f'Warning: a later tools/list page advertises a real tool named "{name}", ' + "but PostHog already injected its own tool by that name on the first page. " + f'Calls to "{name}" are intercepted by PostHog and the real tool will not ' + f"run. Rename one of them; {remedy} renames PostHog's." + ) + event = _VIRTUAL_TOOL_EVENT[kind] + if variant == "duplicate": + return ( + "Warning: PostHog's missing-capability and agent-feedback tools are both " + f'configured to use the name "{name}". Only the missing-capability tool is ' + f"advertised and intercepted, so no {event} events will be captured. " + f"Rename one with {remedy}." ) - tools_list.append(tool) + return ( + f'Warning: Cannot inject PostHog\'s "{name}" tool because a real tool already ' + f"uses that name. PostHog will not intercept it and no {event} events will be " + f"captured. Rename PostHog's tool with {remedy}." + ) def read_tool_category(tool: Any) -> Optional[str]: @@ -779,19 +949,6 @@ def collect_listed_tools(data: MCPAnalyticsData, tools: list) -> tuple[List[str] return names, not tools -def _is_sdk_virtual_tool(data: MCPAnalyticsData, tool_name: Any) -> bool: - """Whether this name is one of the SDK's own virtual tools (``get_more_tools``, - ``send_feedback``) — those carry their intent in their own arguments, so they - never get ``context``/``conversation_id`` injected. A shadowed feedback name - belongs to a real application tool and keeps normal injection.""" - if tool_name == GET_MORE_TOOLS_NAME: - return True - options = resolve_collect_feedback_options(data.options.collect_feedback) - if options is None or data.feedback_tool_shadowed: - return False - return tool_name == resolve_send_feedback_tool_name(options) - - def mutate_tool_schema( data: MCPAnalyticsData, tool: Any, @@ -799,6 +956,7 @@ def mutate_tool_schema( schema_attribute: str, owns_context: bool, context_required: bool, + is_sdk_virtual_tool: bool, ) -> None: """Apply the common analytics schema pipeline and write it back in place. @@ -808,7 +966,6 @@ def mutate_tool_schema( """ schema = getattr(tool, schema_attribute, None) original_schema = schema - is_sdk_virtual_tool = _is_sdk_virtual_tool(data, tool.name) if ( not is_sdk_virtual_tool and is_context_enabled(data.options.context) diff --git a/posthog/mcp/_intent.py b/posthog/mcp/_intent.py index bac2e9146..076a3e1d8 100644 --- a/posthog/mcp/_intent.py +++ b/posthog/mcp/_intent.py @@ -56,14 +56,22 @@ async def resolve_tool_call_intent( request: Dict[str, Any], extra: Optional[Dict[str, Any]] = None, ) -> Optional[ResolvedIntent]: - from .tools import resolve_missing_capability_tool_name + from ._instrumentation import ( + VIRTUAL_TOOL_MISSING_CAPABILITY, + enabled_virtual_tool_names, + ) context_argument = _get_context_argument(request) name = (request.get("params") or {}).get("name") - missing_name = resolve_missing_capability_tool_name(data.options) + # The virtual tool carries its intent in its own `context` argument, which + # is captured as the event's own field rather than as `$mcp_intent`. Read + # through the enabled map, so with `report_missing` off a *real* application + # tool by that name keeps its `context` captured as intent like any other + # tool's. + missing_name = enabled_virtual_tool_names(data).get(VIRTUAL_TOOL_MISSING_CAPABILITY) if ( is_context_enabled(data.options.context) - and name != missing_name + and (missing_name is None or name != missing_name) and context_argument ): return (context_argument, "context_parameter") diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index f19840456..b10a9ef60 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -17,7 +17,7 @@ from collections import OrderedDict from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Any, Dict, Optional +from typing import Any, Awaitable, Callable, Dict, Optional, Set, Tuple from .logger import log from ._sink import McpEventSink @@ -66,11 +66,24 @@ class MCPAnalyticsData: # signature of a stateless server whose mint middleware never attached. Warned # a single time per server so the log isn't flooded on every request. warned_no_stateless_session: bool = False - # True when the last tools/list showed a real application tool using the - # feedback tool's name. Calls to that name then dispatch normally instead of - # being intercepted (fail-open), and the tool keeps its normal analytics - # schema injection. Refreshed at every listing pass. - feedback_tool_shadowed: bool = False + # True once the SDK has told the host that its tools/list handler was + # replaced after instrument(), so ownership can no longer be determined and + # the virtual tools are advertised but never intercepted. Warned a single + # time per server so the log isn't flooded on every call. + warned_foreign_list_handler: bool = False + # ``(kind, name, variant)`` collision warnings already emitted, so a client + # that re-lists tools on every turn logs each misconfiguration once. + warned_virtual_tool_collisions: Set[Tuple[str, str, str]] = field( + default_factory=set + ) + # Adapter-supplied probe returning the names the host's own (original, + # un-instrumented) tools/list handler advertises on its first page, or None + # when it can't be determined. Registered by the adapters with no tool + # registry to query — raw low-level servers. Takes the adapter's request + # context, which the 1.x handler shape ignores. + raw_tool_names_probe: Optional[Callable[[Any], Awaitable[Optional[Set[str]]]]] = ( + None + ) last_activity: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) diff --git a/posthog/mcp/logger.py b/posthog/mcp/logger.py index 33a41e1ea..6c116cc5f 100644 --- a/posthog/mcp/logger.py +++ b/posthog/mcp/logger.py @@ -41,16 +41,13 @@ def log(message: str) -> None: def warn(message: str) -> None: - """A misconfiguration the host almost certainly wants to know about, sent to - the ``logger`` option *and* to the ``posthog.mcp`` standard-library logger. - - Reserved for warnings that can only fire on an HTTP transport, where the - STDIO constraint above does not apply. A default-configured host still sees - these on stderr (logging's lastResort handler), which is the whole point: - the misconfigurations this is used for are invisible in the data, so a - warning nobody has opted in to receive is a warning nobody reads. Hosts that - do configure logging can route or silence them by name like any other - logger.""" + """A misconfiguration the host wants to know about, sent to the ``logger`` + option *and* to the ``posthog.mcp`` standard-library logger. + + For misconfigurations that are invisible in the captured data, so these go + out whether or not a ``logger`` option was passed. STDIO-safe: the + constraint above is on *stdout*, which carries the protocol stream, and the + MCP spec allows servers to log to stderr.""" log(message) try: _stdlib_logger.warning(message) diff --git a/posthog/mcp/posthog_mcp.py b/posthog/mcp/posthog_mcp.py index 614b96157..a2995dc07 100644 --- a/posthog/mcp/posthog_mcp.py +++ b/posthog/mcp/posthog_mcp.py @@ -13,7 +13,7 @@ import copy from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Set, Tuple, Union from posthog.client import Client @@ -24,9 +24,15 @@ ) from ._event_types import MCPAnalyticsEventType from ._exceptions import capture_exception -from ._instrumentation import drain_pending_sync, fire_and_forget +from ._instrumentation import ( + VIRTUAL_TOOL_FEEDBACK, + VIRTUAL_TOOL_MISSING_CAPABILITY, + drain_pending_sync, + fire_and_forget, + virtual_tool_collision_message, +) from ._lib_identity import apply_mcp_lib_identity -from .logger import log +from .logger import log, warn from ._model_parameters import ( add_model_parameter_to_schema, can_inject_model_parameter, @@ -106,6 +112,9 @@ def __init__( self._mcp_exception_autocapture = mcp_exception_autocapture self._capture_model = capture_model self._model_parameter_injected: Dict[str, bool] = {} + # (kind, name) collision warnings already emitted from prepare_tool_list, + # so a host that prepares a listing per request logs each once. + self._warned_virtual_tool_collisions: Set[Tuple[str, str]] = set() # --- lifecycle ----------------------------------------------------------- @@ -361,7 +370,19 @@ def prepare_tool_list( requires the constructor's ``collect_feedback`` option — the enable switch that gates detection in :meth:`prepare_tool_call`). Returns a new list; dict tools are copied, context injection mutates tool objects in place, - and model injection copies them to preserve field ownership.""" + and model injection copies them to preserve field ownership. + + **On a paginated listing, pass the two switches for the first page only** — + a client concatenates every page into one list:: + + first_page = request.params.get("cursor") is None + tools = posthog.prepare_tool_list( + page_tools, report_missing=first_page, collect_feedback=first_page + ) + + A real tool already using a virtual tool's name wins: it is left alone, + nothing is appended, and a warning names the option that renames + PostHog's tool.""" prepared = [] context_description = get_context_description(context) for tool in tools: @@ -372,21 +393,56 @@ def prepare_tool_list( ) prepared.append(current) - if report_missing and not any( - _tool_name(t) == self._missing_capability_tool_name for t in prepared - ): - prepared.append( - build_report_missing_descriptor(self._missing_capability_tool_name) - ) - if ( - collect_feedback - and self._collect_feedback is not None - and not any(_tool_name(t) == self._feedback_tool_name for t in prepared) - ): - prepared.append(get_feedback_tool_descriptor(self._collect_feedback)) + # A tool already using the name blocks injection — unless it is our own + # descriptor, which a host re-preparing an already-prepared list hands + # straight back. Warning about that would be warning about ourselves. + if report_missing: + name = self._missing_capability_tool_name + if self._real_tool_owns_name(prepared, name): + self._warn_virtual_tool_collision( + VIRTUAL_TOOL_MISSING_CAPABILITY, + name, + 'PostHogMCP(missing_capability_tool_name="...")', + ) + elif not any(_tool_name(t) == name for t in prepared): + prepared.append(build_report_missing_descriptor(name)) + if collect_feedback and self._collect_feedback is not None: + name = self._feedback_tool_name + if self._real_tool_owns_name(prepared, name): + self._warn_virtual_tool_collision( + VIRTUAL_TOOL_FEEDBACK, + name, + 'PostHogMCP(collect_feedback=CollectFeedbackOptions(tool_name="..."))', + ) + elif not any(_tool_name(t) == name for t in prepared): + prepared.append(get_feedback_tool_descriptor(self._collect_feedback)) prepared = self._inject_models(prepared) return prepared + def _real_tool_owns_name(self, prepared: List[Any], name: str) -> bool: + """Whether a *host* tool in this listing owns ``name``. Our own + descriptor doesn't count: a host may re-prepare an already-prepared + list, and that is not a collision to warn about.""" + return any( + _tool_name(tool) == name and not self._is_sdk_virtual_tool(tool) + for tool in prepared + ) + + def _warn_virtual_tool_collision( + self, kind: str, name: str, rename_option: str + ) -> None: + """Warn once per ``(kind, name)`` for this client's lifetime, so a host + that prepares a listing on every request doesn't flood the log.""" + key = (kind, name) + if key in self._warned_virtual_tool_collisions: + return + self._warned_virtual_tool_collisions.add(key) + warn( + virtual_tool_collision_message( + kind, name, "blocked", rename_option=rename_option + ) + ) + def prepare_tool_call( self, name: str, @@ -429,22 +485,28 @@ def prepare_tool_call( if analytics_owns_model: prepared_args = _strip_model(prepared_args) # A supplied `original_tool` is a real application tool by this name (it - # comes from the host's own list, which never holds the virtual tool), so - # the real tool wins — the stateless twin of instrument()'s listing-derived - # shadow flag. Without it the name match stands, and the documented remedy - # for a collision is configuring a non-colliding `tool_name`. + # comes from the host's own list, which never holds a virtual tool), so + # the real tool wins — the stateless twin of the ownership check + # instrument() runs. Without it the name match stands, and the + # documented remedy for a collision is renaming PostHog's tool. is_feedback = ( self._collect_feedback is not None and name == self._feedback_tool_name and original_tool is None ) + # Same guard for the missing-capability tool. Unlike feedback it has no + # constructor enable switch on this path (the name is always populated), + # so `original_tool` is the only ownership signal available here. + is_missing_capability = ( + name == self._missing_capability_tool_name and original_tool is None + ) return PreparedToolCall( args=prepared_args, intent=intent, intent_source="context_parameter" if intent else None, llm_model=llm_model, llm_model_source=llm_model_source, - is_missing_capability=name == self._missing_capability_tool_name, + is_missing_capability=is_missing_capability, is_feedback=is_feedback, feedback_report=( parse_feedback_report(args, self._collect_feedback) @@ -495,19 +557,29 @@ def _emit(self, event: Dict[str, Any]) -> None: # flush()/shutdown() able to drain without blocking their own event loop's tasks. fire_and_forget(self._mcp_sink.capture(event, options), self, background=True) - def _is_virtual_tool_name(self, name: Any) -> bool: - """The SDK's own virtual tools carry their intent in their own arguments, - so they never get the ``context`` parameter injected. The feedback name - only counts with the constructor opt-in — without it a real tool by that - name is an ordinary tool.""" + def _is_sdk_virtual_tool(self, tool: Any) -> bool: + """Whether this tool is one of the SDK's own descriptors, which carry + their intent in their own arguments and so never get ``context`` + injected. + + Matched on the description, not the name: a host may own a tool called + ``get_more_tools``, and a host re-preparing an already-prepared list + hands our descriptor straight back, so both arrive under the same name. + The description is ours and, unlike the schema, survives the + model-injection pass.""" + name = _tool_name(tool) if name == self._missing_capability_tool_name: - return True - return self._collect_feedback is not None and name == self._feedback_tool_name + expected = build_report_missing_descriptor(name) + elif self._collect_feedback is not None and name == self._feedback_tool_name: + expected = get_feedback_tool_descriptor(self._collect_feedback) + else: + return False + return _tool_description(tool) == expected["description"] def _inject_context(self, tool: Any, description: Optional[str]) -> Any: if isinstance(tool, dict): name = tool.get("name", "unknown") - if self._is_virtual_tool_name(name): + if self._is_sdk_virtual_tool(tool): return tool new_schema = add_context_parameter_to_schema( tool.get("inputSchema"), name, description @@ -515,7 +587,7 @@ def _inject_context(self, tool: Any, description: Optional[str]) -> Any: return {**tool, "inputSchema": new_schema} name = getattr(tool, "name", "unknown") - if self._is_virtual_tool_name(name): + if self._is_sdk_virtual_tool(tool): return tool new_schema = add_context_parameter_to_schema( getattr(tool, "inputSchema", None), name, description @@ -605,6 +677,13 @@ def _strip_model(args: Optional[JsonRecord]) -> Optional[JsonRecord]: return {k: v for k, v in args.items() if k != "llm_model"} +def _tool_description(tool: Any) -> Any: + """A tool's description, whether it is a dict or an SDK model.""" + if isinstance(tool, dict): + return tool.get("description") + return getattr(tool, "description", None) + + def _tool_name(tool: Any) -> Optional[str]: if isinstance(tool, dict): return tool.get("name") diff --git a/posthog/mcp/types.py b/posthog/mcp/types.py index d2f3afdd3..3c3c7be51 100644 --- a/posthog/mcp/types.py +++ b/posthog/mcp/types.py @@ -182,6 +182,9 @@ class MCPAnalyticsOptions: logger: Optional[LoggerFn] = None report_missing: bool = False + # Rename the `get_more_tools` virtual tool. Use it when a real tool of yours + # already owns the default name: the real tool wins, so without a rename no + # `$mcp_missing_capability` events are captured at all. missing_capability_tool_name: Optional[str] = None enable_conversation_id: bool = False enable_exception_autocapture: bool = True diff --git a/posthog/test/mcp/_helpers_lowlevel.py b/posthog/test/mcp/_helpers_lowlevel.py new file mode 100644 index 000000000..161313a05 --- /dev/null +++ b/posthog/test/mcp/_helpers_lowlevel.py @@ -0,0 +1,64 @@ +"""Shared fixtures for the raw low-level (mcp 1.x) virtual-tool tests. + +Kept out of ``_helpers`` because that module is imported under both SDK majors, +including by ``conftest``, while these seams are 1.x only. +""" + +import mcp.types as mcp_types +from mcp.server.lowlevel import Server + +ECHO_TOOL = mcp_types.Tool( + name="echo", + description="Echo", + inputSchema={"type": "object", "properties": {"msg": {"type": "string"}}}, +) + + +def make_paged_lowlevel(pages, name="paged"): + """A raw low-level server whose tools/list handler serves ``pages`` (a list of + tool lists) one page per request, chained by ``nextCursor``. The paged handler + is registered directly into ``request_handlers`` so the wire pagination shape + is exact; the real tool handler answers ``real tool ran``.""" + server = Server(name) + + @server.call_tool() + async def call_tool(name, arguments): + return [mcp_types.TextContent(type="text", text="real tool ran")] + + async def paged_list(req): + cursor = getattr(getattr(req, "params", None), "cursor", None) if req else None + index = int(cursor) if cursor else 0 + next_cursor = str(index + 1) if index + 1 < len(pages) else None + return mcp_types.ServerResult( + mcp_types.ListToolsResult(tools=list(pages[index]), nextCursor=next_cursor) + ) + + server.request_handlers[mcp_types.ListToolsRequest] = paged_list + return server + + +def list_page(server, cursor=None): + """Request one page. ``cursor=None`` is a first page; any string -- including + ``""``, a valid opaque cursor -- is a continuation, so the empty case must + not collapse to ``params=None``.""" + handler = server.request_handlers[mcp_types.ListToolsRequest] + params = ( + mcp_types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None + ) + return handler(mcp_types.ListToolsRequest(method="tools/list", params=params)) + + +def call_request(name, arguments): + return mcp_types.CallToolRequest( + method="tools/call", + params=mcp_types.CallToolRequestParams(name=name, arguments=arguments), + ) + + +async def call_tool(server, name, arguments): + handler = server.request_handlers[mcp_types.CallToolRequest] + return await handler(call_request(name, arguments)) + + +def tool_names(page): + return [tool.name for tool in page.root.tools] diff --git a/posthog/test/mcp/conftest.py b/posthog/test/mcp/conftest.py index 085263c66..e17970c23 100644 --- a/posthog/test/mcp/conftest.py +++ b/posthog/test/mcp/conftest.py @@ -19,6 +19,8 @@ "test_feedback.py", "test_lowlevel.py", "test_review_fixes.py", + # module-level `from mcp.server.fastmcp import ...` / v1 request_handlers seams + "test_virtual_tools.py", ] _V2_ONLY = [ diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index f98bca160..e15391745 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -298,3 +298,42 @@ async def test_public_call_tool_entrypoint_still_works_outside_a_request(): text_blocks = [c.text for c in result[0] if getattr(c, "type", None) == "text"] assert "5" in text_blocks + + +async def test_a_failed_registry_lookup_delegates_instead_of_swallowing(): + # A lookup that raises means "could not answer", not "the name is free". + # Answering False there swallows the host's own tool and returns PostHog's + # canned reply as a success. Same contract as the low-level adapter. + server = FastMCP("flaky-registry") + + @server.tool() + def get_more_tools(context: str) -> str: + return "real tool ran" + + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + original_get_tool = server._tool_manager.get_tool + failed = [] + + def flaky_get_tool(name, *args, **kwargs): + if name == "get_more_tools" and not failed: + failed.append(name) + raise ConnectionError("registry unreachable") + return original_get_tool(name, *args, **kwargs) + + server._tool_manager.get_tool = flaky_get_tool + + out = await server._tool_manager.call_tool( + "get_more_tools", {"context": "need csv export"} + ) + await _flush() + + assert "real tool ran" in str(out) + assert _events(client, "$mcp_missing_capability") == [] + assert any("delegating the call to your server" in m for m in messages) diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index 9ebc11459..61b6c2a91 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -151,3 +151,47 @@ def echo(msg: str) -> str: assert _events(client, "$mcp_tool_call")[0]["properties"]["$mcp_intent"] == ( "strict validation" ) + + +async def test_a_failed_registry_lookup_delegates_instead_of_swallowing(): + # A real tool of the host's owns the virtual tool's name, so the SDK must + # never answer that call itself. fastmcp resolves a tool through a provider + # chain that can reach a mounted or proxied upstream over the network, so + # `get_tool` raising means "could not look it up", not "the name is free" -- + # answering False there swallowed the host's tool and returned PostHog's + # canned reply as a success. + server = FastMCP("jlowin-flaky-registry") + + @server.tool + def get_more_tools(context: str) -> str: + return "real tool ran" + + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + await _list(server) + + original_get_tool = server.get_tool + failed = [] + + async def flaky_get_tool(name, *args, **kwargs): + # Transient, as a network blip is: the SDK's ownership lookup hits it, + # the host's own dispatch that follows does not. + if name == "get_more_tools" and not failed: + failed.append(name) + raise ConnectionError("upstream provider unreachable") + return await original_get_tool(name, *args, **kwargs) + + server.get_tool = flaky_get_tool + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert "real tool ran" in str(out.root.content[0].text) + assert _events(client, "$mcp_missing_capability") == [] + assert any("delegating the call to your server" in m for m in messages) diff --git a/posthog/test/mcp/test_feedback.py b/posthog/test/mcp/test_feedback.py index 0e9c02ae1..f3354856a 100644 --- a/posthog/test/mcp/test_feedback.py +++ b/posthog/test/mcp/test_feedback.py @@ -25,6 +25,12 @@ events_named as _events, flush_background as _flush, ) +from posthog.test.mcp._helpers_lowlevel import ( + ECHO_TOOL as _ECHO_TOOL, + call_request as _call_request, + list_page as _list_page, + make_paged_lowlevel, +) _REPORT_ARGS = { "feedback_type": "missing_capability", @@ -71,13 +77,6 @@ async def call_tool(name, arguments): return server -def _call_request(name, arguments): - return mcp_types.CallToolRequest( - method="tools/call", - params=mcp_types.CallToolRequestParams(name=name, arguments=arguments), - ) - - async def _list_tools_lowlevel(server): handler = server.request_handlers[mcp_types.ListToolsRequest] return await handler(mcp_types.ListToolsRequest(method="tools/list")) @@ -704,65 +703,56 @@ async def call_tool(name, arguments): assert _events(client, "$mcp_tool_call") -def _make_paged_lowlevel(pages): - """A raw low-level server whose tools/list handler serves ``pages`` (a list of - tool lists) one page per request, chained by ``nextCursor``. The paged handler - is registered directly into ``request_handlers`` so the wire pagination shape - is exact; the real ``send_feedback`` handler answers ``real tool ran``.""" - server = Server("feedback-lowlevel-paged") - - @server.call_tool() - async def call_tool(name, arguments): - return [mcp_types.TextContent(type="text", text="real tool ran")] - - async def paged_list(req): - cursor = getattr(getattr(req, "params", None), "cursor", None) if req else None - index = int(cursor) if cursor else 0 - next_cursor = str(index + 1) if index + 1 < len(pages) else None - return mcp_types.ServerResult( - mcp_types.ListToolsResult(tools=list(pages[index]), nextCursor=next_cursor) - ) - - server.request_handlers[mcp_types.ListToolsRequest] = paged_list - return server - - -def _list_page(server, cursor=None): - handler = server.request_handlers[mcp_types.ListToolsRequest] - params = mcp_types.PaginatedRequestParams(cursor=cursor) if cursor else None - return handler(mcp_types.ListToolsRequest(method="tools/list", params=params)) - - _REAL_SEND_FEEDBACK = mcp_types.Tool( name="send_feedback", description="A real application tool", inputSchema={"type": "object", "properties": {"note": {"type": "string"}}}, ) -_ECHO_TOOL = mcp_types.Tool( - name="echo", - description="Echo", - inputSchema={"type": "object", "properties": {"msg": {"type": "string"}}}, -) -@pytest.mark.parametrize("real_tool_page", [0, 1]) -async def test_paginated_listing_keeps_collision_across_pages(real_tool_page): - # A collision seen on any page must survive the other pages: recomputing the - # flag from one page alone would re-arm interception and swallow the real - # tool's calls (and an early page must not advertise the virtual tool before - # a later page reveals the real one). - pages = [[_ECHO_TOOL], [_ECHO_TOOL]] - pages[real_tool_page] = [_REAL_SEND_FEEDBACK] - server = _make_paged_lowlevel(pages) +async def test_paginated_listing_appends_virtual_tool_on_first_page_only(): + # A client concatenates every page into one list, so the virtual tool may + # appear on exactly one page -- the first, which every client reads. + server = make_paged_lowlevel([[_ECHO_TOOL], [_ECHO_TOOL]]) client = FakeClient() instrument(server, client, MCPAnalyticsOptions(collect_feedback=True)) + page_one = await _list_page(server) + page_two = await _list_page(server, cursor="1") + assert [t.name for t in page_one.root.tools] == ["echo", "send_feedback"] + assert [t.name for t in page_two.root.tools] == ["echo"] + + +async def test_empty_string_cursor_is_a_continuation_page(): + # `""` is a valid opaque cursor a client got from a previous page, not the + # absence of one, so it must not re-append the virtual tool. + server = make_paged_lowlevel([[_ECHO_TOOL]]) + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(collect_feedback=True)) + + assert [t.name for t in (await _list_page(server, cursor="")).root.tools] == [ + "echo" + ] + + +async def test_first_page_collision_lets_the_real_tool_win(): + # A real tool owning the name on the first page blocks injection outright, + # and its calls are dispatched, not intercepted. + server = make_paged_lowlevel([[_REAL_SEND_FEEDBACK], [_ECHO_TOOL]]) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(collect_feedback=True, logger=messages.append), + ) + page_one = await _list_page(server) page_two = await _list_page(server, cursor="1") listed = [t.name for t in page_one.root.tools] + [ t.name for t in page_two.root.tools ] - assert listed.count("send_feedback") == 1 # the real tool only, never appended + assert listed.count("send_feedback") == 1 # the real tool, never appended call_handler = server.request_handlers[mcp_types.CallToolRequest] out = await call_handler(_call_request("send_feedback", {"note": "hi"})) @@ -771,17 +761,56 @@ async def test_paginated_listing_keeps_collision_across_pages(real_tool_page): assert out.root.content[0].text == "real tool ran" assert _events(client, "$mcp_feedback") == [] assert _events(client, "$mcp_tool_call") + assert any("Cannot inject PostHog's" in m for m in messages) + # The warning has to name the way out, or nobody acts on it. + assert any("collect_feedback=CollectFeedbackOptions" in m for m in messages) -async def test_paginated_listing_appends_virtual_tool_once_on_final_page(): - server = _make_paged_lowlevel([[_ECHO_TOOL], [_ECHO_TOOL]]) +async def test_later_page_collision_shadows_the_real_tool(): + # The get_more_tools twin in test_virtual_tools.py carries the reasoning: + # page one cannot see page two, so the real tool is shadowed and the remedy + # is the rename option. Do not make this fail-open without also changing the + # JS SDK. + server = make_paged_lowlevel([[_ECHO_TOOL], [_REAL_SEND_FEEDBACK]]) client = FakeClient() - instrument(server, client, MCPAnalyticsOptions(collect_feedback=True)) + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(collect_feedback=True, logger=messages.append), + ) page_one = await _list_page(server) page_two = await _list_page(server, cursor="1") - assert [t.name for t in page_one.root.tools] == ["echo"] # non-final: no append - assert [t.name for t in page_two.root.tools] == ["echo", "send_feedback"] + listed = [t.name for t in page_one.root.tools] + [ + t.name for t in page_two.root.tools + ] + assert listed.count("send_feedback") == 2 # PostHog's, then the real one + + call_handler = server.request_handlers[mcp_types.CallToolRequest] + out = await call_handler(_call_request("send_feedback", {"summary": "shadowed"})) + await _flush() + + assert out.root.content[0].text != "real tool ran" + assert _events(client, "$mcp_feedback") + assert any("a later tools/list page advertises a real tool" in m for m in messages) + assert any("collect_feedback=CollectFeedbackOptions" in m for m in messages) + + +async def test_collision_warning_is_logged_once_across_repeated_listings(): + # A client that re-lists tools on every turn must not flood the log. + server = make_paged_lowlevel([[_REAL_SEND_FEEDBACK]]) + messages = [] + instrument( + server, + FakeClient(), + MCPAnalyticsOptions(collect_feedback=True, logger=messages.append), + ) + + for _ in range(3): + await _list_page(server) + + assert len([m for m in messages if "Cannot inject PostHog's" in m]) == 1 async def test_feedback_never_mints_conversation_id(): @@ -850,6 +879,49 @@ async def test_posthogmcp_prepare_tool_list_appends_descriptor(): collect_feedback=True, ) assert [t["name"] for t in collided] == ["send_feedback"] + # ...and that real tool still gets `context`, so its intent is captured. It + # is only PostHog's own descriptor — recognised by already declaring + # `context`, not by its name alone — that is left alone. + assert "context" in collided[0]["inputSchema"]["properties"] + + +async def test_posthogmcp_real_tool_named_get_more_tools_keeps_context(): + # A host tool named like a virtual tool used to be mistaken for PostHog's + # own descriptor and skipped, silently dropping its $mcp_intent. + client, _ = make_client() + prepared = client.prepare_tool_list( + [ + { + "name": "get_more_tools", + "inputSchema": {"type": "object", "properties": {}}, + } + ] + ) + assert "context" in prepared[0]["inputSchema"]["properties"] + + +async def test_posthogmcp_repreparing_a_prepared_list_is_not_a_collision(caplog): + # Hosts may re-prepare an already-prepared list. PostHog's own descriptors + # come back in it, and mistaking them for host tools would both warn about + # ourselves and duplicate the tools. + client, _ = make_client(collect_feedback=True) + tools = [{"name": "search", "inputSchema": {"type": "object", "properties": {}}}] + + once = client.prepare_tool_list(tools, report_missing=True, collect_feedback=True) + assert [t["name"] for t in once] == ["search", "get_more_tools", "send_feedback"] + + with caplog.at_level("WARNING", logger="posthog.mcp"): + twice = client.prepare_tool_list( + once, report_missing=True, collect_feedback=True + ) + + assert [t["name"] for t in twice] == [t["name"] for t in once] + assert not [r for r in caplog.records if r.name == "posthog.mcp"] + # The descriptors' own schemas are unchanged: get_more_tools keeps just its + # own `context`, and send_feedback — which states intent through `summary` + # and `details` instead — gains none. + assert list(twice[1]["inputSchema"]["properties"]) == ["context"] + assert "context" not in twice[2]["inputSchema"]["properties"] async def test_posthogmcp_prepare_tool_list_requires_constructor_option(): @@ -879,8 +951,8 @@ async def test_posthogmcp_prepare_tool_call_without_opt_in_never_flags(): async def test_posthogmcp_original_tool_wins_name_collision(): # A host whose own list holds a real `send_feedback` tool passes it as # `original_tool`; the call then dispatches as a real tool call instead of - # being swallowed as feedback — the stateless twin of instrument()'s - # listing-derived shadow flag. + # being swallowed as feedback — the stateless twin of the ownership check + # instrument() runs. client, _ = make_client(collect_feedback=True) real_tool = { "name": "send_feedback", diff --git a/posthog/test/mcp/test_units.py b/posthog/test/mcp/test_units.py index 998510df1..663f88cce 100644 --- a/posthog/test/mcp/test_units.py +++ b/posthog/test/mcp/test_units.py @@ -12,7 +12,14 @@ inject_prompt_back, resolve_conversation_id, ) -from posthog.mcp._instrumentation import mutate_tool_schema +from posthog.mcp._instrumentation import ( + VIRTUAL_TOOL_FEEDBACK, + VIRTUAL_TOOL_MISSING_CAPABILITY, + enabled_virtual_tool_names, + is_first_listing_page, + mutate_tool_schema, + resolve_virtual_tool_injection, +) from posthog.mcp._intent import _get_context_argument, resolve_tool_call_intent from posthog.mcp._internal import ( IdentityCache, @@ -25,7 +32,11 @@ new_session_id, resolve_session_id, ) -from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity +from posthog.mcp.types import ( + CollectFeedbackOptions, + MCPAnalyticsOptions, + UserIdentity, +) def _data(**opts): @@ -47,13 +58,26 @@ async def test_intent_from_context_argument(): async def test_intent_skips_context_for_missing_capability_tool(): - # a get_more_tools call's context is a capability report, not a tool-call intent + # a get_more_tools call's context is a capability report, not a tool-call + # intent — but only while the SDK actually owns that name out = await resolve_tool_call_intent( - _data(), _call(name="get_more_tools", args={"context": "need csv export"}) + _data(report_missing=True), + _call(name="get_more_tools", args={"context": "need csv export"}), ) assert out is None +async def test_intent_keeps_context_for_a_real_tool_named_get_more_tools(): + # With report_missing off the SDK advertises no such tool, so one by that + # name is the host's own and its context is an ordinary intent. Interception + # has always been gated on report_missing, so dropping the intent here left + # a real tool dispatched but unattributed. + out = await resolve_tool_call_intent( + _data(), _call(name="get_more_tools", args={"context": "need csv export"}) + ) + assert out == ("need csv export", "context_parameter") + + async def test_intent_fallback_sync(): data = _data(intent_fallback=lambda req, extra: "inferred it") assert await resolve_tool_call_intent(data, _call()) == ("inferred it", "inferred") @@ -87,6 +111,117 @@ def test_get_context_argument_ignores_non_string_and_blank(): assert _get_context_argument({}) is None +# --- virtual tools ----------------------------------------------------------- +# The kind-keyed resolver both virtual tools share. Duck-typed on `.name`, so it +# drives from plain namespaces and these run under both MCP SDK majors. + + +def _tool(name): + return SimpleNamespace(name=name) + + +def test_first_page_is_an_absent_cursor(): + assert is_first_listing_page(None) is True + assert is_first_listing_page(SimpleNamespace(cursor=None)) is True + + +def test_a_present_cursor_is_a_continuation_page(): + # `""` is a valid opaque cursor, not the absence of one. Reading it as + # falsy would re-append the virtual tools to that page. + assert is_first_listing_page(SimpleNamespace(cursor="")) is False + assert is_first_listing_page(SimpleNamespace(cursor="abc")) is False + + +def test_enabled_names_follow_the_switches_and_renames(): + assert enabled_virtual_tool_names(_data()) == {} + assert enabled_virtual_tool_names(_data(report_missing=True)) == { + VIRTUAL_TOOL_MISSING_CAPABILITY: "get_more_tools" + } + both = enabled_virtual_tool_names( + _data( + report_missing=True, + missing_capability_tool_name="find_tools", + collect_feedback=CollectFeedbackOptions(tool_name="tell_posthog"), + ) + ) + assert both == { + VIRTUAL_TOOL_MISSING_CAPABILITY: "find_tools", + VIRTUAL_TOOL_FEEDBACK: "tell_posthog", + } + + +def test_resolver_injects_on_a_first_page(): + data = _data(report_missing=True, collect_feedback=True) + injection = resolve_virtual_tool_injection( + data, [_tool("echo")], is_first_page=True + ) + assert injection == { + VIRTUAL_TOOL_MISSING_CAPABILITY: "get_more_tools", + VIRTUAL_TOOL_FEEDBACK: "send_feedback", + } + + +def test_resolver_injects_nothing_on_a_continuation_page(): + data = _data(report_missing=True, collect_feedback=True) + injection = resolve_virtual_tool_injection( + data, [_tool("echo")], is_first_page=False + ) + assert injection == {} + + +def test_resolver_skips_injection_on_a_first_page_collision(): + data = _data(report_missing=True) + injection = resolve_virtual_tool_injection( + data, [_tool("get_more_tools")], is_first_page=True + ) + assert injection == {} + + +def test_resolver_decides_each_page_on_its_own_tools(): + # Stateless: a collision on one first page does not shadow the virtual tool + # on the next, so dropping the colliding tool needs no re-instrumentation. + data = _data(report_missing=True) + blocked = resolve_virtual_tool_injection( + data, [_tool("get_more_tools")], is_first_page=True + ) + assert blocked == {} + + injection = resolve_virtual_tool_injection( + data, [_tool("echo")], is_first_page=True + ) + assert injection == {VIRTUAL_TOOL_MISSING_CAPABILITY: "get_more_tools"} + + +def test_resolver_warns_once_per_kind_name_and_variant(): + data = _data(report_missing=True) + for _ in range(3): + resolve_virtual_tool_injection( + data, [_tool("get_more_tools")], is_first_page=True + ) + assert data.warned_virtual_tool_collisions == { + (VIRTUAL_TOOL_MISSING_CAPABILITY, "get_more_tools", "blocked") + } + + +def test_resolver_keeps_missing_capability_when_both_share_a_name(): + # Every call path checks missing-capability first, so advertising both under + # one name would dead-letter the feedback path. + data = _data( + report_missing=True, + missing_capability_tool_name="ask_posthog", + collect_feedback=CollectFeedbackOptions(tool_name="ask_posthog"), + ) + injection = resolve_virtual_tool_injection( + data, [_tool("echo")], is_first_page=True + ) + assert injection == {VIRTUAL_TOOL_MISSING_CAPABILITY: "ask_posthog"} + assert ( + VIRTUAL_TOOL_FEEDBACK, + "ask_posthog", + "duplicate", + ) in data.warned_virtual_tool_collisions + + # --- conversation_id --------------------------------------------------------- @@ -114,6 +249,7 @@ def test_schema_pipeline_does_not_warn_for_owned_conversation_id(monkeypatch): schema_attribute="input_schema", owns_context=False, context_required=False, + is_sdk_virtual_tool=False, ) assert tool.input_schema is schema @@ -156,6 +292,14 @@ def test_resolve_conversation_id_skips_missing_capability_tool(): ) +def test_resolve_conversation_id_mints_for_a_shadowed_virtual_tool_name(): + # `None` means the virtual tool is disabled or a real application tool owns + # the name. Either way the call belongs to that real tool, so it mints and + # echoes a handle like any other tool's. + cid, minted = resolve_conversation_id(True, {}, "get_more_tools", None) + assert minted is True and cid + + def test_resolve_conversation_id_uses_supplied_when_mintable_shape(): # Only an echo of a handle we could have minted (a uuidv7) is accepted — # the handle becomes $session_id, so an invented value ("conv-1") must not diff --git a/posthog/test/mcp/test_v2_lowlevel.py b/posthog/test/mcp/test_v2_lowlevel.py index 4d53407ed..9b9142725 100644 --- a/posthog/test/mcp/test_v2_lowlevel.py +++ b/posthog/test/mcp/test_v2_lowlevel.py @@ -15,6 +15,7 @@ from mcp.server.lowlevel import Server from posthog.mcp import instrument +from posthog.mcp.tools import get_more_tools_result_text from posthog.mcp.types import MCPAnalyticsOptions from posthog.test.mcp._helpers import ( FakeClient, @@ -575,3 +576,142 @@ def identify(request, extra): await _flush() assert seen["headers"] == {"authorization": "Bearer t0ken", "user-agent": "probe/1"} + + +# --- virtual tools on a paginated listing ------------------------------------ + +_REAL_GET_MORE_TOOLS_V2 = mcp_types.Tool( + name="get_more_tools", + description="A real application tool that owns the name", + input_schema={"type": "object", "properties": {"context": {"type": "string"}}}, +) +_ECHO_TOOL_V2 = mcp_types.Tool( + name="echo", + description="Echo", + input_schema={"type": "object", "properties": {"msg": {"type": "string"}}}, +) + + +def make_paged_server_v2(pages): + """A raw v2 low-level server serving ``pages`` one page per request, chained + by ``next_cursor``. The tool handler answers ``real tool ran``.""" + + async def on_call_tool(ctx, params): + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type="text", text="real tool ran")] + ) + + async def on_list_tools(ctx, params): + cursor = getattr(params, "cursor", None) + index = int(cursor) if cursor else 0 + return mcp_types.ListToolsResult( + tools=list(pages[index]), + next_cursor=str(index + 1) if index + 1 < len(pages) else None, + ) + + return Server( + "test-paged-v2", + on_call_tool=on_call_tool, + on_list_tools=on_list_tools, + ) + + +async def _list_page_v2(server, cursor=None): + """Request one page. ``cursor=None`` is a first page; any string -- including + ``""``, a valid opaque cursor -- is a continuation.""" + entry = server.get_request_handler("tools/list") + params = ( + mcp_types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None + ) + return await entry.handler(fake_ctx(method="tools/list"), params) + + +async def test_v2_virtual_tools_appended_to_first_page_only(): + server = make_paged_server_v2([[_ECHO_TOOL_V2], [_ECHO_TOOL_V2]]) + instrument( + server, + FakeClient(), + MCPAnalyticsOptions(report_missing=True, collect_feedback=True), + ) + + first = await _list_page_v2(server) + second = await _list_page_v2(server, cursor="1") + assert [t.name for t in first.tools] == ["echo", "get_more_tools", "send_feedback"] + assert [t.name for t in second.tools] == ["echo"] + + +async def test_v2_empty_string_cursor_is_a_continuation_page(): + server = make_paged_server_v2([[_ECHO_TOOL_V2]]) + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + page = await _list_page_v2(server, cursor="") + assert [t.name for t in page.tools] == ["echo"] + + +async def test_v2_raw_list_probe_blocks_interception_before_any_listing(): + # A raw v2 low-level server has no tool registry, so ownership is settled by + # asking the host's own tools/list handler. + server = make_paged_server_v2([[_REAL_GET_MORE_TOOLS_V2]]) + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + result = await _call_tool(server, "get_more_tools", {"context": "need csv"}) + await _flush() + + assert result.content[0].text == "real tool ran" + assert _events(client, "$mcp_missing_capability") == [] + + +async def test_v2_first_page_collision_lets_the_real_tool_win(): + server = make_paged_server_v2([[_REAL_GET_MORE_TOOLS_V2], [_ECHO_TOOL_V2]]) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + first = await _list_page_v2(server) + second = await _list_page_v2(server, cursor="1") + listed = [t.name for t in first.tools] + [t.name for t in second.tools] + assert listed.count("get_more_tools") == 1 # the real tool, never appended + assert any("Cannot inject PostHog's" in m for m in messages) + assert any("missing_capability_tool_name" in m for m in messages) + + +async def test_v2_cached_result_object_does_not_collide_with_itself(): + # v2 returns ListToolsResult directly rather than wrapped in a root model, + # so it exercises the other branch of the non-mutating append. A host is + # free to hand back the same object every time; PostHog must not read its + # own injected tool back out of it as a real one. + cached = mcp_types.ListToolsResult(tools=[_ECHO_TOOL_V2]) + + async def on_call_tool(ctx, params): + raise ValueError(f"Unknown tool: {params.name}") + + async def on_list_tools(ctx, params): + return cached + + server = Server( + "test-cached-v2", on_call_tool=on_call_tool, on_list_tools=on_list_tools + ) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + first = [t.name for t in (await _list_page_v2(server)).tools] + second = [t.name for t in (await _list_page_v2(server)).tools] + assert first == ["echo", "get_more_tools"] + assert second == first + assert [t.name for t in cached.tools] == ["echo"] # the host's object is untouched + + result = await _call_tool(server, "get_more_tools", {"context": "need csv"}) + await _flush() + assert result.content[0].text == get_more_tools_result_text() + assert _events(client, "$mcp_missing_capability") + assert not [m for m in messages if "Cannot inject PostHog's" in m] diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index 80ae82ce6..d39c6162a 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -379,6 +379,28 @@ async def test_report_missing_accepts_omitted_arguments(): assert "$mcp_intent" not in missing[0]["properties"] +async def test_real_get_more_tools_is_not_intercepted(): + # The registry probe is the only ownership signal before any tools/list has + # run -- the multi-pod case. Without it the SDK swallows the host's tool and + # answers with its own canned reply. + server = MCPServer("test-server-v2") + + @server.tool() + def get_more_tools(context: str) -> str: + """A real application tool that owns the name.""" + return "real tool ran" + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + result = await _call_tool(server, "get_more_tools", {"context": "need csv"}) + await _flush() + + assert "real tool ran" in str(result.content) + assert _events(client, "$mcp_missing_capability") == [] + assert _events(client, "$mcp_tool_call") + + async def test_collect_feedback_advertises_and_captures(): server = make_server() client = FakeClient() @@ -429,33 +451,6 @@ def send_feedback(note: str) -> str: assert _events(client, "$mcp_tool_call") -async def test_collect_feedback_collision_keeps_conversation_id(): - # The name collision must fail open for every feature keyed off the - # feedback tool name, not just dispatch - conversation-id resolution used - # to keep skipping the real tool because it checked the configured name - # alone, ignoring the listing-derived shadow flag. - server = make_server() - - @server.tool() - def send_feedback(note: str) -> str: - return f"real tool got {note}" - - client = FakeClient() - instrument( - server, - client, - MCPAnalyticsOptions(collect_feedback=True, enable_conversation_id=True), - ) - - await _list_tools(server) - await _call_tool(server, "send_feedback", {"note": "hi", "context": "real tool"}) - await _flush() - - calls = _events(client, "$mcp_tool_call") - assert len(calls) == 1 - assert calls[0]["properties"].get("$mcp_conversation_id") - - async def test_instrument_is_idempotent(): server = make_server() client = FakeClient() @@ -511,3 +506,40 @@ async def test_anonymous_events_do_not_create_person_profiles(): calls = _events(client, "$mcp_tool_call") assert calls[0]["properties"]["$process_person_profile"] is False + + +async def test_a_failed_registry_lookup_delegates_instead_of_swallowing(): + # A lookup that raises means "could not answer", not "the name is free". + # Answering False there swallows the host's own tool and returns PostHog's + # canned reply as a success. Same contract as the low-level adapter. + server = MCPServer("flaky-registry-v2") + + @server.tool() + def get_more_tools(context: str) -> str: + return "real tool ran" + + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + original_get_tool = server._tool_manager.get_tool + failed = [] + + def flaky_get_tool(name, *args, **kwargs): + if name == "get_more_tools" and not failed: + failed.append(name) + raise ConnectionError("registry unreachable") + return original_get_tool(name, *args, **kwargs) + + server._tool_manager.get_tool = flaky_get_tool + + out = await _call_tool(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert "real tool ran" in str(out.content[0].text) + assert _events(client, "$mcp_missing_capability") == [] + assert any("delegating the call to your server" in m for m in messages) diff --git a/posthog/test/mcp/test_virtual_tools.py b/posthog/test/mcp/test_virtual_tools.py new file mode 100644 index 000000000..84481f9ca --- /dev/null +++ b/posthog/test/mcp/test_virtual_tools.py @@ -0,0 +1,576 @@ +"""Tests for the ``get_more_tools`` virtual tool (the ``report_missing`` option): +first-page-only injection on a paginated listing, and name collisions with a real +application tool. + +``send_feedback``'s equivalents live in ``test_feedback.py``. Both tools share one +kind-keyed resolver in ``_instrumentation``, so a change that breaks one of these +files should break the other. +""" + +import mcp.types as mcp_types +import pytest +from mcp.server.fastmcp import FastMCP +from mcp.server.lowlevel import Server + +from posthog.mcp import CollectFeedbackOptions, instrument +from posthog.mcp.tools import get_more_tools_result_text +from posthog.mcp.types import MCPAnalyticsOptions +from posthog.test.mcp._helpers import ( + FakeClient, + events_named as _events, + flush_background as _flush, +) +from posthog.test.mcp._helpers_lowlevel import ( + ECHO_TOOL as _ECHO_TOOL, + call_tool as _call, + list_page as _list_page, + make_paged_lowlevel, + tool_names as _names, +) + +_REAL_GET_MORE_TOOLS = mcp_types.Tool( + name="get_more_tools", + description="A real application tool that owns the name", + inputSchema={"type": "object", "properties": {"context": {"type": "string"}}}, +) + + +def _static_list(tools): + """A ``tools/list`` handler serving one unpaginated page, registered directly + into ``request_handlers`` the way a raw low-level host does.""" + + async def handler(req): + return mcp_types.ServerResult(mcp_types.ListToolsResult(tools=list(tools))) + + return handler + + +# --- pagination ---------------------------------------------------------------- + + +async def test_appended_to_first_page_only(): + # The regression test for the original bug: get_more_tools had no page gate + # at all, so a cursor-following client saw it once per page. + server = make_paged_lowlevel([[_ECHO_TOOL], [_ECHO_TOOL]]) + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + assert _names(await _list_page(server)) == ["echo", "get_more_tools"] + assert _names(await _list_page(server, cursor="1")) == ["echo"] + + +async def test_single_unpaginated_listing_still_gets_the_tool(): + server = make_paged_lowlevel([[_ECHO_TOOL]]) + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + assert _names(await _list_page(server)) == ["echo", "get_more_tools"] + + +async def test_empty_string_cursor_is_a_continuation_page(): + server = make_paged_lowlevel([[_ECHO_TOOL]]) + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + assert _names(await _list_page(server, cursor="")) == ["echo"] + + +async def test_both_virtual_tools_are_appended_to_the_first_page_only(): + # The two tools share one resolver, so they must agree on the page rule. + server = make_paged_lowlevel([[_ECHO_TOOL], [_ECHO_TOOL]]) + instrument( + server, + FakeClient(), + MCPAnalyticsOptions(report_missing=True, collect_feedback=True), + ) + + assert _names(await _list_page(server)) == [ + "echo", + "get_more_tools", + "send_feedback", + ] + assert _names(await _list_page(server, cursor="1")) == ["echo"] + + +# --- name collisions ----------------------------------------------------------- + + +async def test_first_page_collision_lets_the_real_tool_win(): + # Before this, get_more_tools had no collision handling at all: the host's + # real tool was silently swallowed and the agent got PostHog's canned reply. + server = make_paged_lowlevel([[_REAL_GET_MORE_TOOLS], [_ECHO_TOOL]]) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + listed = _names(await _list_page(server)) + _names( + await _list_page(server, cursor="1") + ) + assert listed.count("get_more_tools") == 1 # the real tool, never appended + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert out.root.content[0].text == "real tool ran" + assert _events(client, "$mcp_missing_capability") == [] + assert _events(client, "$mcp_tool_call") + assert any("Cannot inject PostHog's" in m for m in messages) + # The warning has to name the way out, or nobody acts on it. + assert any("missing_capability_tool_name" in m for m in messages) + + +async def test_later_page_collision_shadows_the_real_tool(): + # Page one cannot see page two, so the virtual tool is already advertised by + # the time the real one shows up. PostHog keeps intercepting the name and the + # real tool is shadowed -- the host's remedy is the rename option, which the + # warning names. @posthog/mcp behaves the same way; do not make this + # fail-open without also changing the JS SDK. + server = make_paged_lowlevel([[_ECHO_TOOL], [_REAL_GET_MORE_TOOLS]]) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + listed = _names(await _list_page(server)) + _names( + await _list_page(server, cursor="1") + ) + assert listed.count("get_more_tools") == 2 # PostHog's, then the real one + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert out.root.content[0].text == get_more_tools_result_text() + assert _events(client, "$mcp_missing_capability") + assert any("a later tools/list page advertises a real tool" in m for m in messages) + assert any("missing_capability_tool_name" in m for m in messages) + + +async def test_collision_un_shadows_once_the_real_tool_is_dropped(): + # Every first page is judged on its own tools, so a host that removes the + # colliding tool gets the virtual one back without re-instrumenting. + pages = [[_REAL_GET_MORE_TOOLS]] + server = make_paged_lowlevel(pages) + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + assert _names(await _list_page(server)) == ["get_more_tools"] + pages[0] = [_ECHO_TOOL] + assert _names(await _list_page(server)) == ["echo", "get_more_tools"] + + +async def test_registry_probe_blocks_interception_before_any_listing(): + # The multi-pod case: a call reaching a process that never served a + # tools/list has nothing to go on but the registry. + server = FastMCP("virtual-tools-fastmcp") + + @server.tool() + def get_more_tools(context: str) -> str: + return "real tool ran" + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + out = await server._tool_manager.call_tool( + "get_more_tools", {"context": "need csv export"} + ) + await _flush() + + assert "real tool ran" in str(out) + assert _events(client, "$mcp_missing_capability") == [] + + +async def test_raw_list_probe_blocks_interception_before_any_listing(): + # A raw low-level server has no tool registry, so ownership is settled by + # asking the host's own tools/list handler. + server = make_paged_lowlevel([[_REAL_GET_MORE_TOOLS]]) + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert out.root.content[0].text == "real tool ran" + assert _events(client, "$mcp_missing_capability") == [] + + +async def test_both_virtual_tools_configured_with_the_same_name(): + # Two tools by one name would be advertised twice and dead-letter the + # feedback path, since every call path checks missing-capability first. + server = make_paged_lowlevel([[_ECHO_TOOL]]) + messages = [] + instrument( + server, + FakeClient(), + MCPAnalyticsOptions( + report_missing=True, + missing_capability_tool_name="ask_posthog", + collect_feedback=CollectFeedbackOptions(tool_name="ask_posthog"), + logger=messages.append, + ), + ) + + assert _names(await _list_page(server)) == ["echo", "ask_posthog"] + assert any("both" in m and "ask_posthog" in m for m in messages) + + +# --- custom name --------------------------------------------------------------- + + +@pytest.mark.parametrize("enable_conversation_id", [False, True]) +async def test_renamed_tool_carries_its_own_intent(enable_conversation_id): + # A virtual tool states its intent in its own `context` argument, so it gets + # neither an injected `context` nor a `conversation_id` -- renamed or not. + server = make_paged_lowlevel([[_ECHO_TOOL]]) + instrument( + server, + FakeClient(), + MCPAnalyticsOptions( + report_missing=True, + missing_capability_tool_name="find_tools", + enable_conversation_id=enable_conversation_id, + ), + ) + + page = await _list_page(server) + virtual = [tool for tool in page.root.tools if tool.name == "find_tools"][0] + assert list(virtual.inputSchema["properties"]) == ["context"] + assert virtual.inputSchema["required"] == ["context"] + + +async def test_renamed_tool_is_intercepted_and_the_default_name_is_not(): + server = make_paged_lowlevel([[_ECHO_TOOL]]) + client = FakeClient() + instrument( + server, + client, + MCPAnalyticsOptions( + report_missing=True, missing_capability_tool_name="find_tools" + ), + ) + + await _list_page(server) + renamed = await _call(server, "find_tools", {"context": "need csv export"}) + default = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert renamed.root.content[0].text == get_more_tools_result_text() + assert default.root.content[0].text == "real tool ran" + assert len(_events(client, "$mcp_missing_capability")) == 1 + + +async def test_real_tool_named_get_more_tools_keeps_normal_injection(): + # With report_missing off the SDK advertises no such tool, so one by that + # name is an ordinary application tool: it gets `context` injected and its + # value captured as $mcp_intent, like any other tool's. + server = make_paged_lowlevel([[_REAL_GET_MORE_TOOLS]]) + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=False, context=True)) + + await _list_page(server) + out = await _call(server, "get_more_tools", {"context": "delete a cohort"}) + await _flush() + + assert out.root.content[0].text == "real tool ran" + calls = _events(client, "$mcp_tool_call") + assert calls + assert calls[0]["properties"]["$mcp_intent"] == "delete a cohort" + + +# --- a host that reuses one result object -------------------------------------- + + +def _make_cached_lowlevel(tools): + """A raw low-level server that returns the SAME ``ServerResult`` object from + every tools/list -- a module-level constant or the host's own cache. The + appends mutate that object in place, so the SDK must not later read its own + injected tool back out of it as if the host owned the name.""" + cached = mcp_types.ServerResult(mcp_types.ListToolsResult(tools=list(tools))) + server = Server("virtual-tools-cached") + + @server.call_tool() + async def call_tool(name, arguments): + if name != "echo": + raise ValueError(f"Unknown tool: {name}") + return [mcp_types.TextContent(type="text", text="real tool ran")] + + async def list_tools(req): + return cached + + server.request_handlers[mcp_types.ListToolsRequest] = list_tools + return server + + +async def test_cached_result_object_does_not_collide_with_itself(): + # Regression: the SDK saw its own injected tool in the host's reused result + # object, reported a collision against itself, stopped intercepting, and the + # agent got an unknown-tool error instead of its feedback being recorded. + server = _make_cached_lowlevel([_ECHO_TOOL]) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions( + report_missing=True, collect_feedback=True, logger=messages.append + ), + ) + + await _list_page(server) + feedback = await _call(server, "send_feedback", {"summary": "hi"}) + missing = await _call(server, "get_more_tools", {"context": "need csv"}) + await _flush() + + assert feedback.root.isError is not True + assert missing.root.isError is not True + assert _events(client, "$mcp_feedback") + assert _events(client, "$mcp_missing_capability") + assert not [m for m in messages if "Cannot inject PostHog's" in m] + + +async def test_cached_result_object_is_not_appended_to_twice(): + # The same contamination would also read as a collision on the second + # listing, so the virtual tools would silently stop being advertised. + server = _make_cached_lowlevel([_ECHO_TOOL]) + instrument( + server, + FakeClient(), + MCPAnalyticsOptions(report_missing=True, collect_feedback=True), + ) + + first = _names(await _list_page(server)) + second = _names(await _list_page(server)) + assert first.count("get_more_tools") == 1 + assert first.count("send_feedback") == 1 + assert second == first + + +async def test_the_probe_asks_the_host_once_per_virtual_tool_call(): + # The probe asks per call rather than trusting a listing, so a server with + # per-caller catalogues can't answer one caller from another's listing. + # Pinning the count keeps that cost visible. Counted inside the host's own + # handler, which the probe calls directly. + calls = [] + server = Server("virtual-tools-counted") + + @server.call_tool() + async def call_tool(name, arguments): + return [mcp_types.TextContent(type="text", text="real tool ran")] + + async def counted_list(req): + # `req is None` is the MCP SDK repopulating its own validation cache, + # which it does on a call to an unlisted name. Not ours, so not counted. + if req is not None: + calls.append(req) + return mcp_types.ServerResult(mcp_types.ListToolsResult(tools=[_ECHO_TOOL])) + + server.request_handlers[mcp_types.ListToolsRequest] = counted_list + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + await _list_page(server) + assert len(calls) == 1 + + await _call(server, "get_more_tools", {"context": "need csv"}) + await _call(server, "get_more_tools", {"context": "need csv"}) + assert len(calls) == 3 # one probe per virtual-tool call + + # Ordinary tool traffic never reaches the probe. + await _call(server, "echo", {"msg": "hi"}) + assert len(calls) == 3 + + +async def test_unanswerable_ownership_check_delegates_to_the_host(): + # When the SDK cannot tell whose tool a name is -- here the host's tools/list + # handler is failing -- it hands the call to the host. Guessing the other way + # swallows a real tool of theirs silently, for as long as the handler stays + # unwell. + server = Server("virtual-tools-unanswerable") + + @server.call_tool() + async def call_tool(name, arguments): + return [mcp_types.TextContent(type="text", text=f"host dispatched {name}")] + + async def failing_list(req): + # Fails only for real listing requests, so the MCP SDK's own `req is + # None` validation-cache pass still works and only the SDK's ownership + # check is affected. + if req is not None: + raise RuntimeError("catalogue backend unavailable") + return mcp_types.ServerResult(mcp_types.ListToolsResult(tools=[_ECHO_TOOL])) + + server.request_handlers[mcp_types.ListToolsRequest] = failing_list + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions( + report_missing=True, collect_feedback=True, logger=messages.append + ), + ) + + missing = await _call(server, "get_more_tools", {"context": "need csv"}) + feedback = await _call(server, "send_feedback", {"summary": "s"}) + await _flush() + + assert missing.root.content[0].text == "host dispatched get_more_tools" + assert feedback.root.content[0].text == "host dispatched send_feedback" + assert _events(client, "$mcp_missing_capability") == [] + assert _events(client, "$mcp_feedback") == [] + # The host is told why their call was not instrumented. + assert any("delegating the call to your server" in m for m in messages) + + +# --- the probe must not corrupt the SDK's validation cache --------------------- + + +def _make_decorated_lowlevel(): + """A raw low-level server that registers ``tools/list`` through the SDK's own + decorator and rebuilds its ``Tool`` objects on every call, the way a host + reading from a database would. The decorator refreshes ``Server._tool_cache`` + from whatever it returns, and that cache is the schema real tool calls are + validated against -- so anything that runs this handler must re-inject.""" + server = Server("virtual-tools-decorated") + + @server.list_tools() + async def list_tools(): + return [ + mcp_types.Tool( + name="echo", + description="Echo", + inputSchema={ + "type": "object", + "properties": {"msg": {"type": "string"}}, + "additionalProperties": False, + }, + ) + ] + + @server.call_tool() + async def call_tool(name, arguments): + return [mcp_types.TextContent(type="text", text="real tool ran")] + + return server + + +async def test_a_virtual_tool_call_leaves_real_tools_callable(): + # Regression: the ownership check ran the host's own list_tools, which made + # the SDK decorator rebuild `_tool_cache` from un-injected schemas. The next + # real call carrying the `context` we advertised was then rejected with + # "Additional properties are not allowed". + server = _make_decorated_lowlevel() + instrument( + server, FakeClient(), MCPAnalyticsOptions(report_missing=True, context=True) + ) + + await _list_page(server) + before = await _call(server, "echo", {"msg": "hi", "context": "say hi"}) + assert before.root.isError is not True + + await _call(server, "get_more_tools", {"context": "need csv export"}) + + after = await _call(server, "echo", {"msg": "hi", "context": "say hi again"}) + await _flush() + assert after.root.isError is not True, after.root.content[0].text + + +async def test_a_listing_handler_registered_after_instrument_is_not_swallowed(): + # Regression: the probe closed over the handler captured at instrument time, + # so a host registering `tools/list` afterwards had ownership answered from a + # catalogue no client ever saw -- a confident wrong answer that swallowed + # their real tool. Unknown must delegate instead. + server = Server("virtual-tools-late") + + @server.call_tool() + async def call_tool(name, arguments): + return [mcp_types.TextContent(type="text", text="real tool ran")] + + server.request_handlers[mcp_types.ListToolsRequest] = _static_list([_ECHO_TOOL]) + instrument(server, FakeClient(), MCPAnalyticsOptions(report_missing=True)) + + # The host replaces its listing after instrumentation, now owning the name. + server.request_handlers[mcp_types.ListToolsRequest] = _static_list( + [_REAL_GET_MORE_TOOLS] + ) + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + assert out.root.content[0].text == "real tool ran" + + +async def test_a_blocked_name_is_not_also_reported_as_shadowed(): + # The host owns the name on page one *and* a later page. Page one blocked + # injection, so nothing of PostHog's is advertised: warning that the real + # tool will not run would send them chasing a bug that isn't there. + server = make_paged_lowlevel([[_REAL_GET_MORE_TOOLS], [_REAL_GET_MORE_TOOLS]]) + messages = [] + instrument( + server, + FakeClient(), + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + await _list_page(server) + await _list_page(server, cursor="1") + + assert any("Cannot inject" in m for m in messages) + assert not any("already injected" in m for m in messages) + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + assert out.root.content[0].text == "real tool ran" + + +async def test_a_duplicate_dropped_name_is_not_reported_as_shadowed(): + # Both virtual tools configured to one name: missing-capability wins the + # first page, feedback is dropped. On a later page only the tool we actually + # injected shadows the host's -- the dropped one cannot. + server = make_paged_lowlevel([[_ECHO_TOOL], [_REAL_GET_MORE_TOOLS]]) + messages = [] + instrument( + server, + FakeClient(), + MCPAnalyticsOptions( + report_missing=True, + collect_feedback=CollectFeedbackOptions(tool_name="get_more_tools"), + logger=messages.append, + ), + ) + + await _list_page(server) + await _list_page(server, cursor="1") + + shadowed = [m for m in messages if "already injected" in m] + assert len(shadowed) == 1 + assert "missing_capability_tool_name" in shadowed[0] + + +async def test_a_removed_listing_handler_is_not_swallowed(): + # The sibling of a replaced handler: removed outright. Ownership is then + # unanswerable, so the call must be delegated rather than intercepted, and + # the host told why -- a silent stop is invisible in the captured data. + server = Server("virtual-tools-removed") + + @server.call_tool() + async def call_tool(name, arguments): + return [mcp_types.TextContent(type="text", text="real tool ran")] + + server.request_handlers[mcp_types.ListToolsRequest] = _static_list([_ECHO_TOOL]) + client = FakeClient() + messages = [] + instrument( + server, + client, + MCPAnalyticsOptions(report_missing=True, logger=messages.append), + ) + + del server.request_handlers[mcp_types.ListToolsRequest] + + out = await _call(server, "get_more_tools", {"context": "need csv export"}) + await _flush() + + assert out.root.content[0].text == "real tool ran" + assert _events(client, "$mcp_missing_capability") == [] + assert any("replaced or removed" in m for m in messages)