Merge branch 'posthog-code/mcp-sdk-v2' into posthog-code/mcp-error-properties

Author: gesh

posthog/mcp/_instrument_fastmcp.py

@@ -42,6 +42,7 @@
     build_tool_call_request,
     extract_tools,
     prepare_request,
+    prime_session,
     read_tool_category,
     record_missing_capability,
     record_tool_call,
@@ -171,6 +172,10 @@ async def _session(anchor: Optional[str]) -> str:
                     k: v for k, v in arguments.items() if k not in strip_keys
                 }
 
+        # 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.
+        await prime_session(data, mcp_session_id=mcp_session_id, token=token)
+
         start = time.monotonic()
         try:
             result = await original(

posthog/mcp/_instrument_lowlevel.py

@@ -37,6 +37,7 @@
     build_tool_call_request,
     extract_tools,
     prepare_request,
+    prime_session,
     read_tool_category,
     record_missing_capability,
     record_tool_call,
@@ -178,6 +179,10 @@ async def _session(anchor: Optional[str]) -> str:
                 if key not in owned:
                     req.params.arguments.pop(key, None)
 
+        # 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.
+        await prime_session(data, mcp_session_id=mcp_session_id, token=token)
+
         start = time.monotonic()
         try:
             result = await original(req)

posthog/mcp/_instrument_v2.py

@@ -51,6 +51,7 @@
     build_tool_call_request,
     params_to_request_dict,
     prepare_request,
+    prime_session,
     read_tool_category,
     record_missing_capability,
     record_tool_call,
@@ -310,6 +311,10 @@ async def _session(anchor: Optional[str]) -> str:
                     k: v for k, v in arguments.items() if k not in strip_keys
                 }
 
+        # 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.
+        await prime_session(data, mcp_session_id=mcp_session_id, token=token)
+
         start = time.monotonic()
         try:
             result = await original(
@@ -474,6 +479,10 @@ async def _session(anchor: Optional[str]) -> str:
                 ]
             )
 
+        # 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.
+        await prime_session(data, mcp_session_id=mcp_session_id, token=token)
+
         start = time.monotonic()
         try:
             result = await original(ctx, params)

posthog/mcp/_instrumentation.py

@@ -248,6 +248,24 @@ def resolve_session_and_client(
     return token, client_name, client_version, protocol_version
 
 
+async def prime_session(
+    data: MCPAnalyticsData,
+    *,
+    mcp_session_id: Optional[str],
+    token: Optional[SessionTokenPayload] = None,
+) -> None:
+    """Point the shared per-server session at *this* request before the tool body runs.
+
+    ``McpAnalytics.capture()`` reads ``data.session_id`` for custom in-tool
+    events. The conversation anchor can only be resolved after the call (we
+    don't know until then whether the agent received the handle), so without
+    this the tool body would read whatever the *previous* request left behind
+    and attribute a custom event to the wrong caller. Emits nothing — it only
+    settles the transport/memory session an in-tool event should belong to.
+    """
+    await resolve_session_id(data, mcp_session_id, token=token)
+
+
 async def prepare_request(
     data: MCPAnalyticsData,
     *,

posthog/test/mcp/test_conversation_session.py

@@ -295,3 +295,60 @@ def boom() -> str:
     )
     # and nothing claims a conversation the agent never received
     assert all("$mcp_conversation_id" not in e["properties"] for e in client.events)
+
+
+@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server")
+async def test_in_tool_events_are_not_attributed_to_the_previous_caller():
+    """A custom event captured *inside* a tool body reads the shared
+    ``data.session_id``. The conversation anchor can only be resolved after the
+    call, so unless that field is settled first, caller B's in-tool event is
+    attributed to caller A's session — and, through the identity cache, to
+    caller A's person."""
+    from types import SimpleNamespace
+
+    from mcp.server.fastmcp import FastMCP
+
+    from posthog.mcp import instrument
+    from posthog.mcp._internal import get_server_tracking_data
+    from posthog.mcp.session import derive_session_id_from_mcp_session
+
+    def caller(session_header):
+        return SimpleNamespace(
+            request_context=SimpleNamespace(
+                request=SimpleNamespace(headers={"mcp-session-id": session_header}),
+                session=SimpleNamespace(client_params=None),
+            )
+        )
+
+    server = FastMCP("in-tool")
+    seen = {}
+
+    @server.tool()
+    def echo(msg: str) -> str:
+        return msg
+
+    client = FakeClient()
+    handle = instrument(server, client, MCPAnalyticsOptions())
+
+    @server.tool()
+    def emits(msg: str) -> str:
+        data = get_server_tracking_data(handle._key)
+        seen["session_during_body"] = data.session_id if data else None
+        return msg
+
+    # Caller A runs first and leaves its session behind on the shared state.
+    await server._tool_manager.call_tool(
+        "echo", {"msg": "a", "context": "caller A"}, context=caller("session-A")
+    )
+    # Caller B's tool body must see *its own* session, not A's.
+    await server._tool_manager.call_tool(
+        "emits", {"msg": "b", "context": "caller B"}, context=caller("session-B")
+    )
+    await _flush()
+
+    assert seen["session_during_body"] == derive_session_id_from_mcp_session(
+        "session-B"
+    )
+    assert seen["session_during_body"] != derive_session_id_from_mcp_session(
+        "session-A"
+    )