fix(mcp): preserve conversation result isolation

Generated-By: PostHog Desktop
Task-Id: 6fa508f7-af04-48ef-b7ea-31732d3ca42e

Author: gesh

.sampo/changesets/mcp-dispatcher-conversations.md

@@ -2,4 +2,4 @@
 pypi/posthog: minor
 ---
 
-Add conversation and session correlation to custom `PostHogMCP` dispatchers, matching `@posthog/mcp`. `prepare_tool_list()` adds an optional `conversation_id` argument and a compatible `_mcp_instructions` output field, `prepare_tool_call()` accepts a carried `session_id` and returns the resolved `session_id` and `conversation_id`, and the new `prepare_tool_result()` delivers a minted handle without changing the original result. The capture methods accept `conversation_id`. Set `PostHogMCP(enable_conversation_id=False)` to keep the previous behavior.
+Add conversation and session correlation to custom `PostHogMCP` dispatchers, matching `@posthog/mcp`. `prepare_tool_list()` adds an optional `conversation_id` field to each compatible tool input schema and a compatible `_mcp_instructions` output field, `prepare_tool_call()` accepts a carried `session_id` and returns the resolved `session_id` and `conversation_id`, and the new `prepare_tool_result()` delivers a minted handle without changing the original result. The capture methods accept `conversation_id`. Set `PostHogMCP(enable_conversation_id=False)` to keep the previous behavior.

posthog/mcp/_output_instructions.py

@@ -218,17 +218,35 @@ def mirror_instructions_into_structured_content(
             new_target = copy_model(update={attr: updated})
         except Exception:  # noqa: BLE001 - never let delivery break the tool path
             return result, False
-        if target is result:
-            return new_target, True
-        rewrap = getattr(result, "model_copy", None)
-        if callable(rewrap):
-            try:
-                return rewrap(update={"root": new_target}), True
-            except Exception:  # noqa: BLE001
-                return result, False
+    else:
+        new_target = _copy_with_attr(target, attr, updated)
+        if new_target is None:
+            return result, False
+    if target is result:
+        return new_target, True
+    rewrap = getattr(result, "model_copy", None)
+    if callable(rewrap):
+        try:
+            return rewrap(update={"root": new_target}), True
+        except Exception:  # noqa: BLE001
+            return result, False
+    new_result = _copy_with_attr(result, "root", new_target)
+    if new_result is None:
         return result, False
+    return new_result, True
+
+
+def _copy_with_attr(value: Any, attr: str, updated: Any) -> Optional[Any]:
+    """Return a shallow copy with one changed attribute, or ``None`` when the
+    object cannot be copied safely."""
     try:
-        setattr(target, attr, updated)
-    except Exception:  # noqa: BLE001 - never let delivery break the tool path
-        return result, False
-    return result, True
+        copied = copy.copy(value)
+    except Exception:  # noqa: BLE001 - analytics must not break tool results
+        return None
+    if copied is value:
+        return None
+    try:
+        setattr(copied, attr, updated)
+    except Exception:  # noqa: BLE001 - read-only result objects fail closed
+        return None
+    return copied

posthog/mcp/posthog_mcp.py

@@ -918,22 +918,27 @@ def _collect_conversation_ownership(
 
 
 def _inject_prompt_back(result: Any, conversation_id: str) -> Any:
-    """Append the handle to a dict or ``CallToolResult`` result's ``content``,
-    including a ``CallToolResult`` inside an MCP SDK 1.x ``ServerResult``.
-    Returns ``result`` itself when there is no content list to append to."""
+    """Append the handle to a tuple, list, dict, or ``CallToolResult`` result's
+    ``content``, including a ``CallToolResult`` inside an MCP SDK 1.x
+    ``ServerResult``. Returns ``result`` itself when there is no content list to
+    append to."""
     if isinstance(result, dict):
         return inject_prompt_back(result, conversation_id)
-    target = getattr(result, "root", result)
-    content = getattr(target, "content", None)
-    copy_model = getattr(target, "model_copy", None)
-    if not isinstance(content, list) or not callable(copy_model):
-        return result
     # A model result means the MCP SDK is installed; it stays a peer dependency.
     import mcp.types as mcp_types  # noqa: PLC0415
 
     block = mcp_types.TextContent(
         type="text", text=build_prompt_back(conversation_id)["text"]
     )
+    if isinstance(result, tuple) and len(result) == 2 and isinstance(result[0], list):
+        return ([*result[0], block], result[1])
+    if isinstance(result, list):
+        return [*result, block]
+    target = getattr(result, "root", result)
+    content = getattr(target, "content", None)
+    copy_model = getattr(target, "model_copy", None)
+    if not isinstance(content, list) or not callable(copy_model):
+        return result
     try:
         updated = copy_model(update={"content": [*content, block]})
         if target is result:

posthog/test/mcp/test_output_instructions.py

@@ -118,10 +118,13 @@ def test_mirrors_into_a_model_result_both_attr_shapes():
     for attr in ("structured_content", "structuredContent"):
         result = SimpleNamespace(**{attr: {"total": 7}})
 
-        _, delivered = mirror_instructions_into_structured_content(result, "conv-2")
+        prepared, delivered = mirror_instructions_into_structured_content(
+            result, "conv-2"
+        )
 
         assert delivered is True
-        assert getattr(result, attr)[MCP_INSTRUCTIONS_KEY] == {
+        assert getattr(result, attr) == {"total": 7}
+        assert getattr(prepared, attr)[MCP_INSTRUCTIONS_KEY] == {
             "conversation_id": "conv-2"
         }
 
@@ -130,10 +133,11 @@ def test_mirrors_through_a_serverresult_wrapper():
     inner = SimpleNamespace(structuredContent={"total": 1})
     result = SimpleNamespace(root=inner)
 
-    _, delivered = mirror_instructions_into_structured_content(result, "conv-3")
+    prepared, delivered = mirror_instructions_into_structured_content(result, "conv-3")
 
     assert delivered is True
-    assert inner.structuredContent[MCP_INSTRUCTIONS_KEY] == {
+    assert inner.structuredContent == {"total": 1}
+    assert prepared.root.structuredContent[MCP_INSTRUCTIONS_KEY] == {
         "conversation_id": "conv-3"
     }
 

posthog/test/mcp/test_posthog_mcp.py

@@ -530,6 +530,42 @@ def test_prepare_tool_result_delivers_into_call_tool_result_models(wrapped):
     assert prepared.conversation_id == call.conversation_id
 
 
+def test_prepare_tool_result_delivers_minted_handle_into_tuple_content():
+    client, _ = make_client()
+    client.prepare_tool_list([_sql_tool()])
+    call = client.prepare_tool_call("execute-sql", {})
+    content = [TextContent(type="text", text="done")]
+    tool_result = (content, None)
+
+    prepared = client.prepare_tool_result(tool_result, call)
+
+    assert len(content) == 1
+    assert prepared.result[0][-1].text == _handle_block(call.conversation_id)["text"]
+    assert prepared.conversation_id == call.conversation_id
+
+
+def test_prepare_tool_result_copies_shared_custom_result_objects():
+    client, _ = make_client()
+    client.prepare_tool_list([_sql_tool()])
+    shared_result = SimpleNamespace(content=[], structuredContent={"rows": []})
+
+    first_call = client.prepare_tool_call("execute-sql", {})
+    first = client.prepare_tool_result(shared_result, first_call)
+    second_call = client.prepare_tool_call("execute-sql", {})
+    second = client.prepare_tool_result(shared_result, second_call)
+
+    assert shared_result.structuredContent == {"rows": []}
+    assert first.result is not shared_result
+    assert second.result is not shared_result
+    assert first.result.structuredContent[MCP_INSTRUCTIONS_KEY] == {
+        "conversation_id": first_call.conversation_id
+    }
+    assert second.result.structuredContent[MCP_INSTRUCTIONS_KEY] == {
+        "conversation_id": second_call.conversation_id
+    }
+    assert first_call.conversation_id != second_call.conversation_id
+
+
 def test_prepare_tool_result_omits_conversation_without_delivery_state():
     client, _ = make_client()
     tool_result = {"content": []}