fix(mcp): anchor virtual tools to conversations

Generated-By: PostHog Desktop
Task-Id: e71933e1-34f5-4e2f-b1dc-da92ab68ff26

Author: gesh

.sampo/changesets/mcp-virtual-tool-conversation-id.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+MCP virtual tools now use conversation IDs when `enable_conversation_id` is enabled.

posthog/mcp/README.md

@@ -299,6 +299,9 @@ needs no middleware and no ordering discipline, and it is the only thing that
 correlates a session under the 2026-07-28 revision's per-request server instances.
 Prefer it if you're on a recent client.
 
+The `get_more_tools` and `send_feedback` virtual tools also use the conversation
+handle when this option is enabled.
+
 ### How the SDK tells you it's misconfigured
 
 The failure used to be silent. It now surfaces two ways:

posthog/mcp/_conversation_id.py

@@ -74,33 +74,13 @@ def extract_conversation_id(args: Any) -> Optional[str]:
     return trimmed or None
 
 
-def resolve_conversation_id(
-    enabled: bool,
-    args: Any,
-    tool_name: Optional[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
-    send_feedback → ``(None, False)``; agent echoed a handle we could have minted
-    → ``(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.
+def resolve_conversation_id(enabled: bool, args: Any) -> Tuple[Optional[str], bool]:
+    """Return the conversation id and whether the SDK minted it.
 
     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 (
-            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)
-    ):
+    if not enabled:
         return None, False
     supplied = extract_conversation_id(args)
     if supplied and _MINTED_CONVERSATION_ID.match(supplied):

posthog/mcp/_instrument_fastmcp.py

@@ -118,14 +118,24 @@ async def wrapped(
         if lifecycle.is_missing_capability and (
             _name_owned_by_real_tool(server, name) is False
         ):
-            await lifecycle.record_missing_capability()
-            return [
+            virtual_content = [
                 mcp_types.TextContent(type="text", text=get_more_tools_result_text())
             ]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            await lifecycle.record_missing_capability(conversation_id_delivered=True)
+            return virtual_content
 
         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)]
+            reply = await lifecycle.record_feedback(conversation_id_delivered=True)
+            virtual_content = [mcp_types.TextContent(type="text", text=reply)]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            return virtual_content
 
         # Strip each injected key independently. A tool can declare its own
         # `context` (kept) while `conversation_id` is still SDK-injected (stripped),

posthog/mcp/_instrument_lowlevel.py

@@ -212,25 +212,33 @@ async def handler(req: Any) -> Any:
         if lifecycle.is_missing_capability and (
             await _name_owned_by_real_tool(high_level, data, name, server) is False
         ):
-            await lifecycle.record_missing_capability()
+            virtual_content = [
+                mcp_types.TextContent(type="text", text=get_more_tools_result_text())
+            ]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            await lifecycle.record_missing_capability(conversation_id_delivered=True)
             return mcp_types.ServerResult(
                 mcp_types.CallToolResult(
-                    content=[
-                        mcp_types.TextContent(
-                            type="text", text=get_more_tools_result_text()
-                        )
-                    ],
+                    content=virtual_content,
                     isError=False,
                 )
             )
 
         if lifecycle.is_feedback and (
             await _name_owned_by_real_tool(high_level, data, name, server) is False
         ):
-            reply = await lifecycle.record_feedback()
+            reply = await lifecycle.record_feedback(conversation_id_delivered=True)
+            virtual_content = [mcp_types.TextContent(type="text", text=reply)]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
             return mcp_types.ServerResult(
                 mcp_types.CallToolResult(
-                    content=[mcp_types.TextContent(type="text", text=reply)],
+                    content=virtual_content,
                     isError=False,
                 )
             )

posthog/mcp/_instrument_v2.py

@@ -302,22 +302,26 @@ async def wrapped(
         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=[
-                    mcp_types.TextContent(
-                        type="text", text=get_more_tools_result_text()
-                    )
-                ]
-            )
+            virtual_content = [
+                mcp_types.TextContent(type="text", text=get_more_tools_result_text())
+            ]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            await lifecycle.record_missing_capability(conversation_id_delivered=True)
+            return mcp_types.CallToolResult(content=virtual_content)
 
         if lifecycle.is_feedback and (
             _name_owned_by_real_tool_v2(server, name) is False
         ):
-            reply = await lifecycle.record_feedback()
-            return mcp_types.CallToolResult(
-                content=[mcp_types.TextContent(type="text", text=reply)]
-            )
+            reply = await lifecycle.record_feedback(conversation_id_delivered=True)
+            virtual_content = [mcp_types.TextContent(type="text", text=reply)]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            return mcp_types.CallToolResult(content=virtual_content)
 
         # v2 validates against the function signature and rejects unexpected
         # keys, so injected parameters are stripped before dispatch — but never
@@ -526,22 +530,26 @@ async def handler(ctx: Any, params: Any) -> Any:
         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=[
-                    mcp_types.TextContent(
-                        type="text", text=get_more_tools_result_text()
-                    )
-                ]
-            )
+            virtual_content = [
+                mcp_types.TextContent(type="text", text=get_more_tools_result_text())
+            ]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            await lifecycle.record_missing_capability(conversation_id_delivered=True)
+            return mcp_types.CallToolResult(content=virtual_content)
 
         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)]
-            )
+            reply = await lifecycle.record_feedback(conversation_id_delivered=True)
+            virtual_content = [mcp_types.TextContent(type="text", text=reply)]
+            if prompt_back := lifecycle.prompt_back_text():
+                virtual_content.append(
+                    mcp_types.TextContent(type="text", text=prompt_back)
+                )
+            return mcp_types.CallToolResult(content=virtual_content)
 
         # Settle the shared session before the tool body runs, so an in-tool
         # `analytics.capture()` is attributed to this caller and not the last one.

posthog/mcp/_instrumentation.py

@@ -23,7 +23,11 @@
     is_context_enabled,
     schema_has_param,
 )
-from ._conversation_id import add_conversation_id_to_schema, resolve_conversation_id
+from ._conversation_id import (
+    add_conversation_id_to_schema,
+    build_prompt_back,
+    resolve_conversation_id,
+)
 from ._event_types import MCPAnalyticsEventType
 from ._exceptions import capture_exception
 from .feedback import (
@@ -470,11 +474,25 @@ async def prime_session(self) -> None:
             self.data, mcp_session_id=self.mcp_session_id, token=self.token
         )
 
-    async def record_missing_capability(self) -> None:
-        session_id = await self.prepare_session(None)
+    def prompt_back_text(self) -> Optional[str]:
+        if not self.conversation_id or not self.minted_conversation_id:
+            return None
+        return build_prompt_back(self.conversation_id)["text"]
+
+    def _anchored_conversation_id(self, delivered: bool) -> Optional[str]:
+        if self.minted_conversation_id and not delivered:
+            return None
+        return self.conversation_id
+
+    async def record_missing_capability(
+        self, *, conversation_id_delivered: bool = False
+    ) -> None:
+        conversation_id = self._anchored_conversation_id(conversation_id_delivered)
+        session_id = await self.prepare_session(conversation_id)
         await record_missing_capability(
             self.data,
             session_id,
+            conversation_id=conversation_id,
             tool_name=self.missing_name or self.name,
             context=(self.arguments or {}).get("context"),
             arguments=self.arguments,
@@ -486,15 +504,17 @@ async def record_missing_capability(self) -> None:
             extra=self.extra,
         )
 
-    async def record_feedback(self) -> str:
+    async def record_feedback(self, *, conversation_id_delivered: bool = False) -> str:
         """Capture the ``$mcp_feedback`` event, then run the host's ``on_feedback``
         handler and return the reply text for the agent. The event is captured
         whether or not the handler raises."""
         report = parse_feedback_report(self.arguments, self.feedback_options)
-        session_id = await self.prepare_session(None)
+        conversation_id = self._anchored_conversation_id(conversation_id_delivered)
+        session_id = await self.prepare_session(conversation_id)
         await record_feedback(
             self.data,
             session_id,
+            conversation_id=conversation_id,
             report=report,
             tool_name=self.feedback_name or self.name,
             arguments=self.arguments,
@@ -509,7 +529,7 @@ async def record_feedback(self) -> str:
     async def record_error(self, error: Any, duration_ms: float) -> None:
         # A freshly minted handle cannot anchor or be captured when dispatch
         # raised: no adapter had an opportunity to deliver it to the agent.
-        conversation_id = None if self.minted_conversation_id else self.conversation_id
+        conversation_id = self._anchored_conversation_id(False)
         session_id = await self.prepare_session(conversation_id)
         await record_tool_call(
             self.data,
@@ -530,9 +550,7 @@ async def record_error(self, error: Any, duration_ms: float) -> None:
     async def record_result(
         self, result: Any, duration_ms: float, *, conversation_id_delivered: bool
     ) -> None:
-        conversation_id = self.conversation_id
-        if self.minted_conversation_id and not conversation_id_delivered:
-            conversation_id = None
+        conversation_id = self._anchored_conversation_id(conversation_id_delivered)
         session_id = await self.prepare_session(conversation_id)
         await record_tool_call(
             self.data,
@@ -573,11 +591,7 @@ def start_tool_call_lifecycle(
     # running the host's `on_feedback` handler read the configured options.
     feedback_options = resolve_collect_feedback_options(data.options.collect_feedback)
     conversation_id, minted = resolve_conversation_id(
-        data.options.enable_conversation_id,
-        arguments,
-        name,
-        missing_name,
-        feedback_name,
+        data.options.enable_conversation_id, arguments
     )
     return ToolCallLifecycle(
         data=data,
@@ -1010,10 +1024,8 @@ def mutate_tool_schema(
         data.tool_model_parameter_injected[tool.name] = (
             not app_owns_model and schema_has_param(schema, "llm_model")
         )
-    if (
-        not is_sdk_virtual_tool
-        and data.options.enable_conversation_id
-        and not schema_has_param(schema, "conversation_id")
+    if data.options.enable_conversation_id and not schema_has_param(
+        schema, "conversation_id"
     ):
         schema = add_conversation_id_to_schema(schema, tool.name)
     if schema is not original_schema:
@@ -1138,6 +1150,7 @@ async def record_missing_capability(
     data: MCPAnalyticsData,
     session_id: str,
     *,
+    conversation_id: Optional[str] = None,
     tool_name: str,
     context: Optional[str],
     arguments: Optional[Dict[str, Any]],
@@ -1155,6 +1168,7 @@ async def record_missing_capability(
         event: Dict[str, Any] = {
             "event_type": MCPAnalyticsEventType.MCP_MISSING_CAPABILITY,
             "session_id": session_id,
+            "conversation_id": conversation_id,
             "resource_name": tool_name,
             "parameters": build_captured_mcp_parameters(
                 request, strip_llm_model=allow_self_reported_model
@@ -1186,6 +1200,7 @@ async def record_feedback(
     data: MCPAnalyticsData,
     session_id: str,
     *,
+    conversation_id: Optional[str] = None,
     report: FeedbackReport,
     tool_name: str,
     arguments: Optional[Dict[str, Any]],
@@ -1206,6 +1221,7 @@ async def record_feedback(
         event: Dict[str, Any] = {
             "event_type": MCPAnalyticsEventType.MCP_FEEDBACK,
             "session_id": session_id,
+            "conversation_id": conversation_id,
             "resource_name": tool_name,
             "client_name": client_name,
             "client_version": client_version,

posthog/test/mcp/test_conversation_session.py

@@ -62,7 +62,7 @@ def test_derivation_is_deterministic_and_distinct():
 
 def test_echo_of_a_mintable_handle_is_accepted():
     cid, minted = resolve_conversation_id(
-        True, {"conversation_id": MINTED_SHAPE_HANDLE}, "t", "get_more_tools"
+        True, {"conversation_id": MINTED_SHAPE_HANDLE}
     )
     assert minted is False
     assert cid == MINTED_SHAPE_HANDLE
@@ -73,7 +73,7 @@ def test_uppercased_echo_is_lowercased_before_hashing():
     # case-sensitive, so the echo must be folded back or it lands in a
     # different session than the call that minted it.
     cid, minted = resolve_conversation_id(
-        True, {"conversation_id": MINTED_SHAPE_HANDLE.upper()}, "t", "get_more_tools"
+        True, {"conversation_id": MINTED_SHAPE_HANDLE.upper()}
     )
     assert minted is False
     assert cid == MINTED_SHAPE_HANDLE
@@ -82,9 +82,7 @@ def test_uppercased_echo_is_lowercased_before_hashing():
 def test_invented_handle_is_not_anchored():
     # Two unrelated users both sending "conv-1" must NOT share a session, so a
     # value we could not have minted is replaced with a fresh handle.
-    cid, minted = resolve_conversation_id(
-        True, {"conversation_id": "conv-1"}, "t", "get_more_tools"
-    )
+    cid, minted = resolve_conversation_id(True, {"conversation_id": "conv-1"})
     assert minted is True
     assert cid != "conv-1"
 

posthog/test/mcp/test_fastmcp.py

@@ -85,6 +85,28 @@ async def test_list_tools_injects_model_into_real_and_virtual_tools():
         assert "llm_model" in tools[name].inputSchema["required"]
 
 
+async def test_virtual_tool_uses_conversation_id():
+    server = make_server()
+    client = FakeClient()
+    instrument(
+        server,
+        client,
+        MCPAnalyticsOptions(report_missing=True, enable_conversation_id=True),
+    )
+
+    listed = await _list_tools(server)
+    tool = next(t for t in listed.root.tools if t.name == "get_more_tools")
+    assert "conversation_id" in tool.inputSchema["properties"]
+
+    result = await server._tool_manager.call_tool("get_more_tools", {"context": "csv"})
+    await _flush()
+
+    handle = _events(client, "$mcp_missing_capability")[0]["properties"][
+        "$mcp_conversation_id"
+    ]
+    assert any(handle in item.text for item in result)
+
+
 # --- tools/call --------------------------------------------------------------