feat(mcp): support conversations in custom dispatchers

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

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+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`. The new `prepare_tool_result()` delivers a minted handle without changing the original result. Tool and report capture methods accept `conversation_id`. Existing dispatchers must call `prepare_tool_result()` to deliver new handles. Set `PostHogMCP(enable_conversation_id=False)` to keep the previous behavior.

posthog/mcp/README.md

@@ -73,8 +73,8 @@ strict-schema clients see the new field. Set `capture_model=False` to leave sche
 Conversation correlation adds an optional `conversation_id` argument and returns a handle in
 eligible tool results. Clients must echo it to group later calls; calls without it mint new handles.
 Set `enable_conversation_id=False` to retain transport-based session grouping and unchanged
-response content. Custom `PostHogMCP` dispatchers enable model capture by default but still
-supply their own session IDs. The reasoning is recorded in posthog-js `docs/adr/0013`.
+response content. Custom `PostHogMCP` dispatchers also enable model capture and conversation
+correlation by default. The reasoning is recorded in posthog-js `docs/adr/0013`.
 
 ## Capture the calling model
 
@@ -119,8 +119,10 @@ authorization middleware with those hooks; argument-based model capture is skipp
 application-owned value cannot be mistaken for analytics. No additional catalog lookup runs
 during a tool call.
 
-For a custom dispatcher, `PostHogMCP` enables the same option by default; pass request
-metadata through explicitly:
+For a custom dispatcher, `PostHogMCP` enables model capture and conversation correlation
+by default. `prepare_tool_list()` injects the analytics fields and records ownership by tool
+name. `prepare_tool_call()` removes SDK-owned arguments and resolves the conversation and
+session. `prepare_tool_result()` returns the result to send and the final values to capture:
 
 ```python
 from posthog.mcp import PostHogMCP
@@ -133,22 +135,34 @@ call = posthog.prepare_tool_call(
     raw_args,
     request_meta=request.get("params", {}).get("_meta"),
     original_tool=original_tool,
+    session_id=transport_session_id,
 )
-result = dispatch(tool_name, call.args)
+prepared = posthog.prepare_tool_result(dispatch(tool_name, call.args), call)
 posthog.capture_tool_call(
     tool_name,
     llm_model=call.llm_model,
     llm_model_source=call.llm_model_source,
+    session_id=prepared.session_id,
+    conversation_id=prepared.conversation_id,
 )
+return prepared.result
 ```
 
 Passing `original_tool` keeps ownership accurate when `tools/list` and
 `tools/call` reach different server replicas. A persistent single-process
 dispatcher can omit it after calling `prepare_tool_list()`.
-Model injection copies tool objects instead of changing their original schemas.
+Model and conversation injection copy tool objects instead of changing their original schemas.
 Always advertise the returned list and pass the original application tool to
 `prepare_tool_call()`. Repeatedly preparing the original list preserves ownership.
 
+Pass an existing transport or request session as `session_id`. A valid echoed
+`conversation_id` takes precedence. Otherwise the existing session stays, and no new handle is
+minted. `prepare_tool_result()` appends a new handle to the result's `content` and mirrors it
+into `structuredContent` when the tool declares an output schema. If the result has no channel
+that can carry a new handle, `conversation_id` is `None` and the derived `session_id` is kept.
+Set `enable_conversation_id=False` on `PostHogMCP` to keep the previous custom-dispatcher
+behavior.
+
 ## Collect agent feedback
 
 Feedback collection is off by default. Enable it to advertise a `send_feedback`
@@ -221,8 +235,14 @@ tools = posthog.prepare_tool_list(server_tools, collect_feedback=True)
 # tools/call dispatcher
 call = posthog.prepare_tool_call(tool_name, raw_args)
 if call.is_feedback:
-    posthog.capture_feedback(report=call.feedback_report)  # emits $mcp_feedback
-    return send_feedback_result()  # replies to the agent and stops dispatch
+    # Replies to the agent and stops dispatch.
+    prepared = posthog.prepare_tool_result(send_feedback_result(), call)
+    posthog.capture_feedback(  # emits $mcp_feedback
+        report=call.feedback_report,
+        session_id=prepared.session_id,
+        conversation_id=prepared.conversation_id,
+    )
+    return prepared.result
 ```
 
 `on_feedback` is ignored on this path — the dispatcher routes reports itself via

posthog/mcp/__init__.py

@@ -87,6 +87,7 @@
     MCPAnalyticsModelSource,
     MCPAnalyticsOptions,
     PreparedToolCall,
+    PreparedToolResult,
     UserIdentity,
 )
 from .version import __version__
@@ -104,6 +105,7 @@
     "CollectFeedbackOptions",
     "FeedbackReport",
     "PreparedToolCall",
+    "PreparedToolResult",
     "get_more_tools_result",
     "send_feedback_result",
     "SEND_FEEDBACK_TOOL_NAME",

posthog/mcp/_conversation_id.py

@@ -42,6 +42,8 @@ def add_conversation_id_to_schema(
         and isinstance(schema.get("properties"), dict)
         and CONVERSATION_ID_PARAM_NAME in schema["properties"]
     ):
+        if _is_our_declaration(schema["properties"][CONVERSATION_ID_PARAM_NAME]):
+            return schema
         log(
             f"WARN: Tool \"{tool_name}\" already has '{CONVERSATION_ID_PARAM_NAME}'. Skipping injection."
         )
@@ -64,6 +66,26 @@ def add_conversation_id_to_schema(
     return schema
 
 
+def can_inject_conversation_id(input_schema: Any) -> bool:
+    """Whether the SDK can own ``conversation_id`` on this input schema. An
+    application-declared field or a composed schema stays the application's,
+    so its value is never read as a handle or stripped before dispatch."""
+    if not isinstance(input_schema, dict):
+        return True
+    properties = input_schema.get("properties")
+    if isinstance(properties, dict) and CONVERSATION_ID_PARAM_NAME in properties:
+        return _is_our_declaration(properties[CONVERSATION_ID_PARAM_NAME])
+    return not any(input_schema.get(key) for key in ("$ref", "oneOf", "allOf", "anyOf"))
+
+
+def _is_our_declaration(declaration: Any) -> bool:
+    return (
+        isinstance(declaration, dict)
+        and declaration.get("type") == "string"
+        and declaration.get("description") == DEFAULT_CONVERSATION_ID_DESCRIPTION
+    )
+
+
 def extract_conversation_id(args: Any) -> Optional[str]:
     if not isinstance(args, dict):
         return None
@@ -116,9 +138,60 @@ def build_prompt_back(conversation_id: str) -> Dict[str, Any]:
 
 
 def inject_prompt_back(result: Any, conversation_id: str) -> Any:
-    if not can_inject_prompt_back(result):
+    """Append a handle block to a result copy when it has content."""
+    block: Any = build_prompt_back(conversation_id)
+    if isinstance(result, dict):
+        if not isinstance(result.get("content"), list):
+            return result
+        return {**result, "content": [*result["content"], block]}
+    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)
+    if not isinstance(content, list):
         return result
-    return {
-        **result,
-        "content": [*result["content"], build_prompt_back(conversation_id)],
-    }
+    try:
+        import mcp.types as mcp_types  # noqa: PLC0415
+
+        block = mcp_types.TextContent(type="text", text=block["text"])
+    except ImportError:
+        pass
+
+    copy_model = getattr(target, "model_copy", None)
+    if callable(copy_model):
+        try:
+            updated = copy_model(update={"content": [*content, block]})
+        except Exception:  # noqa: BLE001 - delivery must not break a tool call
+            return result
+    else:
+        updated = _copy_with_attr(target, "content", [*content, block])
+        if updated is None:
+            return result
+    if target is result:
+        return updated
+
+    rewrap = getattr(result, "model_copy", None)
+    if callable(rewrap):
+        try:
+            return rewrap(update={"root": updated})
+        except Exception:  # noqa: BLE001 - delivery must not break a tool call
+            return result
+    wrapped = _copy_with_attr(result, "root", updated)
+    return result if wrapped is None else wrapped
+
+
+def _copy_with_attr(value: Any, attr: str, updated: Any) -> Optional[Any]:
+    try:
+        copied = copy.copy(value)
+    except Exception:  # noqa: BLE001 - delivery must not break a tool call
+        return None
+    if copied is value:
+        return None
+    try:
+        setattr(copied, attr, updated)
+    except Exception:  # noqa: BLE001 - read-only objects fail closed
+        return None
+    return copied

posthog/mcp/_output_instructions.py

@@ -123,8 +123,19 @@ def add_instructions_to_output_schema(tool: Any) -> bool:
             )
         return False
 
+    try:
+        setattr(tool, attr, declare_output_instructions(original))
+    except Exception:  # noqa: BLE001 - some schema attrs may be read-only
+        log(f"WARN: could not set {attr} on tool {name}")
+        return False
+    return True
+
+
+def declare_output_instructions(output_schema: Dict[str, Any]) -> Dict[str, Any]:
+    """A copy of ``output_schema`` with the optional :data:`MCP_INSTRUCTIONS_KEY`
+    declared. Callers check :func:`can_declare_output_instructions` first."""
     # Deep copy: the server may reuse or freeze the schema object it handed us.
-    schema = copy.deepcopy(original)
+    schema = copy.deepcopy(output_schema)
     if not isinstance(schema.get("properties"), dict):
         schema["properties"] = {}
     schema["properties"][MCP_INSTRUCTIONS_KEY] = {
@@ -137,12 +148,14 @@ def add_instructions_to_output_schema(tool: Any) -> bool:
             }
         },
     }
-    try:
-        setattr(tool, attr, schema)
-    except Exception:  # noqa: BLE001 - some schema attrs may be read-only
-        log(f"WARN: could not set {attr} on tool {name}")
-        return False
-    return True
+    return schema
+
+
+def tool_output_schema(tool: Any) -> Any:
+    """A tool's advertised output schema, whether it is a dict or an SDK model."""
+    if isinstance(tool, dict):
+        return tool.get("outputSchema")
+    return _read_attr(tool, _OUTPUT_SCHEMA_ATTRS)[1]
 
 
 def build_conversation_instructions(conversation_id: str) -> Dict[str, Any]:
@@ -205,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