fix(mcp): send one $mcp_initialize per initialize request

.sampo/changesets/mcp-initialize-per-handshake.md

@@ -0,0 +1,11 @@
+---
+pypi/posthog: patch
+---
+
+MCP analytics now sends one `$mcp_initialize` for each `initialize` request that a client sends.
+
+Before, the SDK sent the event on the first request of each session, once per process. A stateless fleet sent one event for each pod that a session reached. A conversation-anchored session sent another event. A client that connected and sent no other request sent no event.
+
+On MCP SDK 2.x, the event comes from `Server.middleware` on every transport. Requests on the 2026-07-28 spec, which has no `initialize`, no longer send one. On MCP SDK 1.x over HTTP, the event comes from `PostHogMcpStatelessSessionMiddleware`. The apps that `instrument()` autowires get this with no change. If you mount the middleware by hand, pass the server as `server=`. MCP SDK 1.x over stdio keeps the first request as the trigger.
+
+The event now also has the handshake request as `$mcp_parameters`. On MCP SDK 2.x, it also has the server reply as `$mcp_response`.

posthog/mcp/README.md

@@ -67,6 +67,31 @@ The value must contain 1 to 256 characters. The SDK records it as
 `$mcp_server_build` on all MCP events. `PostHogMCP.capture()` also adds it to
 custom events unless the event provides its own value.
 
+## When the SDK sends `$mcp_initialize`
+
+One `$mcp_initialize` per `initialize` request the client sends, in the session
+its later requests belong to. A client that reconnects five times reports five,
+however many pods its requests land on. A client that connects and leaves still
+reports one. Requests on the 2026-07-28 revision, which has no `initialize`,
+report none.
+
+The MCP SDK answers `initialize` itself, so the SDK observes it from outside:
+
+- **MCP SDK 2.x**: through `Server.middleware`, on every transport.
+- **MCP SDK 1.x over HTTP**: through `PostHogMcpStatelessSessionMiddleware`, which
+  `instrument()` adds to the apps it autowires. If you mount the middleware
+  yourself, pass the server: `app.add_middleware(PostHogMcpStatelessSessionMiddleware, server=mcp)`.
+  Without it, each pod reports the session's first request as its initialize.
+  The deprecated SSE transport also uses the session's first request.
+  The middleware reads only requests that carry no `Mcp-Session-Id`, so an
+  `initialize` that reuses an existing session id reports none.
+- **MCP SDK 1.x over stdio**: from the session's first request. A stdio server
+  serves one client, so this misses only a client that never sends a request.
+
+`identify` and `event_properties` run for the handshake too. On MCP SDK 1.x the
+SDK has no request context for it, so `extra["ctx"]` carries only `.request`.
+`get_request_headers(extra)` works the same as for any other request.
+
 ## Capture safe input field names
 
 Tool-call events include `$mcp_input_keys`. The SDK records names from the
@@ -385,11 +410,14 @@ a fragmented per-process id. Add the middleware to your app explicitly:
 from posthog.mcp import PostHogMcpStatelessSessionMiddleware, get_mcp_session
 
 app = mcp.streamable_http_app()
-app.add_middleware(PostHogMcpStatelessSessionMiddleware)
+app.add_middleware(PostHogMcpStatelessSessionMiddleware, server=mcp)
 ```
 
+`server=` lets the middleware capture `$mcp_initialize` from the handshake (see
+[When the SDK sends `$mcp_initialize`](#when-the-sdk-sends-mcp_initialize)).
+
 This is also the path for a custom `PostHogMCP` dispatcher (you own the ASGI app),
-where you then read the recovered session per request:
+where you leave `server=` out and read the recovered session per request:
 
 ```python
 sess = get_mcp_session(request)   # sess.session_id, sess.client_name, ...

posthog/mcp/__init__.py

@@ -51,6 +51,7 @@
     MCPAnalyticsData,
     get_server_tracking_data,
     set_server_tracking_data,
+    tracking_key,
 )
 from .logger import log, set_logger
 from .posthog_mcp import PostHogMCP
@@ -226,18 +227,6 @@ def _warn_if_unsupported_mcp_version() -> None:
         )
 
 
-def _canonical_server(server: Any) -> Any:
-    """The underlying low-level server for high-level wrappers (SDK 1.x FastMCP and
-    jlowin's fastmcp expose ``_mcp_server``; SDK 2.x MCPServer renamed it
-    ``_lowlevel_server``), else the server itself. Used as the tracking key so
-    instrumenting a wrapper and its underlying server resolve to one state instead
-    of two divergent ones (matching the TS SDK)."""
-    low_level = getattr(server, "_mcp_server", None) or getattr(
-        server, "_lowlevel_server", None
-    )
-    return low_level if low_level is not None else server
-
-
 def instrument(
     server: Any,
     posthog_client: Optional[Client] = None,
@@ -286,7 +275,7 @@ def instrument(
     if feedback_options is not None:
         get_feedback_tool_descriptor(feedback_options)
 
-    key = _canonical_server(server)
+    key = tracking_key(server)
 
     try:
         # MCP is an optional peer: load adapters only when instrumentation is

posthog/mcp/_instrument_v2.py

@@ -33,6 +33,7 @@
 import time
 from collections.abc import Mapping
 from dataclasses import replace
+from datetime import datetime, timezone
 from typing import Any, Dict, FrozenSet, Optional, Set, Tuple
 
 import mcp.types as mcp_types
@@ -48,12 +49,15 @@
     advertised_tool_names,
     apply_virtual_tool_injection,
     collect_listed_tools,
+    _is_sse_request,
+    _to_jsonable,
     copy_tools_list_result,
     is_first_listing_page,
     mutate_tool_schema,
     params_to_request_dict,
     prepare_request,
     raw_listing_owns_tool_name,
+    record_initialize,
     record_resource_request,
     resource_listing_response,
     resolve_session_and_client,
@@ -70,9 +74,11 @@
 )
 from ._output_instructions import mirror_instructions_into_structured_content
 from ._tool_input import get_registered_tool_input_schema
+from .asgi import get_mcp_session
 from .logger import log
 from .request_headers import get_request_headers
-from .session_token import read_mcp_session_header
+from .session import handshake_session_id
+from .session_token import SessionTokenPayload, read_mcp_session_header
 from .tools import (
     get_more_tools_result_text,
 )
@@ -111,6 +117,7 @@ def instrument_mcpserver_v2(server: Any, data: MCPAnalyticsData) -> None:
     for method, event_type in _RESOURCE_METHODS.items():
         _wrap_v2_resource_request(low_level, data, method, event_type)
     _patch_add_request_handler(low_level, data, wrap_call=False, high_level=server)
+    _observe_initialize(low_level, data)
 
 
 def instrument_lowlevel_v2(server: Any, data: MCPAnalyticsData) -> None:
@@ -126,6 +133,87 @@ def instrument_lowlevel_v2(server: Any, data: MCPAnalyticsData) -> None:
     for method, event_type in _RESOURCE_METHODS.items():
         _wrap_v2_resource_request(server, data, method, event_type)
     _patch_add_request_handler(server, data, wrap_call=True)
+    _observe_initialize(server, data)
+
+
+# --- initialize ------------------------------------------------------------------
+
+
+def _observe_initialize(server: Any, data: MCPAnalyticsData) -> None:
+    """Emit ``$mcp_initialize`` from the real handshake via ``Server.middleware``.
+
+    mcp 2.x refuses a request handler for ``initialize`` and points at
+    ``Server.middleware`` instead, which wraps every inbound request on every
+    transport, ``initialize`` included. The host callbacks run inline, as they do
+    for a tool call: a hung ``identify`` hangs this handshake, where the client
+    sees it, and leaves no task behind. The runner parks the connection's read
+    loop until the chain returns, but a client sends nothing else before the
+    ``initialize`` reply."""
+    middleware = getattr(server, "middleware", None)
+    if not isinstance(middleware, list):
+        log(
+            "Warning: this MCP server has no Server.middleware list; "
+            "$mcp_initialize falls back to the first request of each session."
+        )
+        return
+    if not any(getattr(entry, _WRAPPED_FLAG, False) for entry in middleware):
+
+        async def observe_initialize(ctx: Any, call_next: Any) -> Any:
+            if getattr(ctx, "method", None) != "initialize":
+                return await call_next(ctx)
+            started = datetime.now(timezone.utc)
+            result = await call_next(ctx)
+            await _record_v2_initialize(data, ctx, result, started)
+            return result
+
+        setattr(observe_initialize, _WRAPPED_FLAG, True)
+        middleware.append(observe_initialize)
+    data.handshake_observed_in_process = True
+
+
+def _initialize_session(
+    ctx: Any,
+) -> Tuple[Optional[str], Optional[SessionTokenPayload]]:
+    """The ``(mcp_session_id, token)`` every later request of this connection will
+    resolve its session from.
+
+    A stateful transport has already assigned the id it returns in the
+    ``Mcp-Session-Id`` header; ``ServerSession`` keeps it on its connection and
+    exposes no public accessor. A stateless one has none: the session comes from
+    the ``Mcp-Session-Id`` the client sent, which it keeps sending, read the way
+    :func:`_resolve_ctx` reads it for every later request, or else from the token
+    our ASGI middleware minted for this response. stdio has neither and resolves
+    to this process's own session, as its later requests do."""
+    if _is_sse_request({"ctx": ctx}):
+        # The SSE transport keys its session on a query parameter that later
+        # requests do not resolve from: they land in this process's own session.
+        return None, None
+    connection = getattr(getattr(ctx, "session", None), "_connection", None)
+    transport_session_id = getattr(connection, "session_id", None)
+    if isinstance(transport_session_id, str) and transport_session_id:
+        return transport_session_id, None
+    token, _, _, _, incoming = _resolve_ctx(ctx)
+    if incoming:
+        return incoming, token
+    request = getattr(ctx, "request", None)
+    return None, get_mcp_session(request) if request is not None else None
+
+
+async def _record_v2_initialize(
+    data: MCPAnalyticsData, ctx: Any, result: Any, started: datetime
+) -> None:
+    try:
+        mcp_session_id, token = _initialize_session(ctx)
+        await record_initialize(
+            data,
+            session_id=await handshake_session_id(data, mcp_session_id, token),
+            params=dict(ctx.params) if isinstance(ctx.params, Mapping) else {},
+            result=_to_jsonable(result),
+            extra={"session_id": mcp_session_id, "ctx": ctx},
+            timestamp=started,
+        )
+    except Exception as error:  # noqa: BLE001 - analytics must never break the handshake
+        log(f"Error: failed to capture $mcp_initialize - {error}")
 
 
 # --- registry plumbing ---------------------------------------------------------

posthog/mcp/_instrumentation.py

@@ -240,17 +240,89 @@ def _wrap_response(result: Any) -> Any:
     return serialized
 
 
+async def record_initialize(
+    data: MCPAnalyticsData,
+    *,
+    session_id: str,
+    params: Dict[str, Any],
+    result: Any,
+    extra: Optional[Dict[str, Any]],
+    timestamp: datetime,
+) -> None:
+    """Emit ``$mcp_initialize`` for an ``initialize`` the client sent.
+
+    The Python MCP SDK answers ``initialize`` itself (mcp 1.x inside
+    ``ServerSession``; mcp 2.x refuses a handler for it), so the seams that observe
+    it — ``Server.middleware`` on 2.x, the ASGI middleware on 1.x over HTTP — call
+    this inline in the handshake request once it has succeeded. ``result`` is the
+    server's reply where the seam can read it, else ``None``. Identify runs first
+    so the event is attributed on the pod that served the handshake, which on a
+    stateless fleet may never see this client again. The capture itself runs after,
+    in the background, as for every other event."""
+
+    def text(value: Any) -> Optional[str]:
+        return value if isinstance(value, str) and value else None
+
+    client_info = params.get("clientInfo")
+    client_info = client_info if isinstance(client_info, dict) else {}
+    agreed = result.get("protocolVersion") if isinstance(result, dict) else None
+    request = {"method": "initialize", "params": params}
+    await handle_identify(data, session_id, request, extra)
+    data.mark_session_initialized(session_id)
+    event: Dict[str, Any] = {
+        "event_type": MCPAnalyticsEventType.MCP_INITIALIZE,
+        "session_id": session_id,
+        "client_name": text(client_info.get("name")),
+        "client_version": text(client_info.get("version")),
+        "protocol_version": text(agreed) or text(params.get("protocolVersion")),
+        "parameters": request,
+        "response": result,
+        "timestamp": timestamp,
+    }
+    await _apply_event_properties(data, event, request, extra)
+    stamp_transport_identity(event, extra)
+    fire_and_forget(capture_event(data, event), data)
+
+
+# ASGI scope key a bound ``PostHogMcpStatelessSessionMiddleware`` sets to the
+# tracking state whose handshakes it captures, on each request it passes on.
+HANDSHAKE_OBSERVER_KEY = "posthog_mcp_handshake_observer"
+
+
+def _first_request_stands_in_for_initialize(
+    data: MCPAnalyticsData, session_source: str, extra: Optional[Dict[str, Any]]
+) -> bool:
+    """Whether this request has to stand in for an ``initialize`` nothing observed.
+
+    Only where no seam sees the real handshake: mcp 1.x over stdio (one client per
+    process, so its first request is as good as the handshake), and an mcp 1.x HTTP
+    request that did not pass a middleware bound to this server. Never for a
+    conversation-anchored session: it belongs to a chat, not to a handshake, and
+    synthesizing one there is how a single connection used to report two. So a
+    client whose first request already echoes a conversation reports no initialize
+    here."""
+    if data.handshake_observed_in_process:
+        return False
+    if session_source == "conversation":
+        return False
+    scope = getattr(get_request(extra), "scope", None)
+    return not (isinstance(scope, dict) and scope.get(HANDSHAKE_OBSERVER_KEY) is data)
+
+
 async def _maybe_emit_initialize(
     data: MCPAnalyticsData,
     session_id: str,
+    session_source: str,
     client_name: Optional[str],
     client_version: Optional[str],
     extra: Optional[Dict[str, Any]],
     protocol_version: Optional[str] = None,
 ) -> None:
-    """Lazily emit ``$mcp_initialize`` once per session. The Python MCP SDK handles
-    ``InitializeRequest`` inside the session layer (not ``request_handlers``), so we
-    synthesize the event from the first instrumented request that carries client info."""
+    """Emit ``$mcp_initialize`` once per session from its first request, where no
+    seam observes the real ``initialize`` (see
+    :func:`_first_request_stands_in_for_initialize`)."""
+    if not _first_request_stands_in_for_initialize(data, session_source, extra):
+        return
     if session_id in data.initialized_sessions:
         return
     data.mark_session_initialized(session_id)
@@ -373,8 +445,9 @@ async def prepare_request(
     protocol_version: Optional[str] = None,
     conversation_id: Optional[str] = None,
 ) -> str:
-    """Resolve the session id, run identify, then lazily emit initialize. Returns
-    the session id to stamp on the event for this request.
+    """Resolve the session id, run identify, then emit initialize if this request
+    has to stand in for an unobserved handshake. Returns the session id to stamp on
+    the event for this request.
 
     ``conversation_id`` is the agent's handle for this request, and when present
     it anchors the session (ADR-0004) so every event of the request — initialize
@@ -397,8 +470,8 @@ async def prepare_request(
     Identify runs *before* initialize so the resolved identity is already in the cache
     when ``capture_event`` builds the initialize event — otherwise the first
     ``$mcp_initialize`` is anonymous even when identify resolves on the same request.
-    (Still not byte-parity with the TS SDK, which wraps the real initialize handler;
-    the Python SDK handles initialize in the session layer, not ``request_handlers``.)
+    Where a seam observes the real handshake (:func:`record_initialize`), nothing
+    is synthesized here.
 
     A request that reached us over HTTP yet still resolved to this process's memory
     has nothing correlating it across pods, which on a stateless server means the
@@ -414,7 +487,13 @@ async def prepare_request(
         _warn_stateless_session_not_wired(data)
     await handle_identify(data, session_id, request, extra)
     await _maybe_emit_initialize(
-        data, session_id, client_name, client_version, extra, protocol_version
+        data,
+        session_id,
+        session_source,
+        client_name,
+        client_version,
+        extra,
+        protocol_version,
     )
     return session_id
 

posthog/mcp/_internal.py

@@ -107,6 +107,11 @@ class MCPAnalyticsData:
     # Bounded FIFO of sessions we've emitted $mcp_initialize for, so a long-lived
     # server can't accumulate one entry per session forever.
     initialized_sessions: "OrderedDict[str, None]" = field(default_factory=OrderedDict)
+    # True once mcp 2.x `Server.middleware` observes this server's real
+    # `initialize`, on every transport. On mcp 1.x an ASGI middleware bound to
+    # this server observes it over HTTP and marks each request it passes on;
+    # whatever neither covers falls back to the first request of the session.
+    handshake_observed_in_process: bool = False
     server_name: Optional[str] = None
     # Validated once during setup. Keep it separate from the caller-owned
     # options object so later option changes cannot alter event attribution.
@@ -130,6 +135,18 @@ def mark_session_initialized(self, session_id: str) -> None:
 )
 
 
+def tracking_key(server: Any) -> Any:
+    """The underlying low-level server for high-level wrappers (SDK 1.x FastMCP and
+    jlowin's fastmcp expose ``_mcp_server``; SDK 2.x MCPServer renamed it
+    ``_lowlevel_server``), else the server itself. Used as the tracking key so
+    instrumenting a wrapper and its underlying server resolve to one state instead
+    of two divergent ones (matching the TS SDK)."""
+    low_level = getattr(server, "_mcp_server", None) or getattr(
+        server, "_lowlevel_server", None
+    )
+    return low_level if low_level is not None else server
+
+
 def get_server_tracking_data(server: Any) -> Optional[MCPAnalyticsData]:
     return _server_tracking.get(server)