diff --git a/.sampo/changesets/mcp-initialize-per-handshake.md b/.sampo/changesets/mcp-initialize-per-handshake.md new file mode 100644 index 000000000..33d926d09 --- /dev/null +++ b/.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`. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 9db9bc31e..7f3133f2f 100644 --- a/posthog/mcp/README.md +++ b/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, ... diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 70d37eb48..98a73bede 100644 --- a/posthog/mcp/__init__.py +++ b/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 diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 843f2d032..58d45f778 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/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 --------------------------------------------------------- diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 03b66392c..4e72e1de1 100644 --- a/posthog/mcp/_instrumentation.py +++ b/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 diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 8da3cd8d0..44a1eeb48 100644 --- a/posthog/mcp/_internal.py +++ b/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) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 4e742a143..a66b41113 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -12,6 +12,11 @@ layer instead, with a pure-ASGI middleware that works for both a mounted FastMCP app and a custom ``PostHogMCP`` dispatcher, across every SDK routing path. +The same blind spot hides the handshake from analytics on MCP SDK 1.x, so a +middleware bound to an instrumented server also captures ``$mcp_initialize`` here, +once per ``initialize``, in the session the response hands the client. (MCP SDK +2.x exposes ``Server.middleware``, which ``instrument()`` uses instead.) + On the ``instrument()`` path this is **wired up automatically**: ``instrument(server, ...)`` wraps the FastMCP server's app factories (``streamable_http_app()`` / ``sse_app()``, which ``mcp.run()`` also calls), so the @@ -38,10 +43,15 @@ import functools import json -from typing import Any, Optional +import weakref +from datetime import datetime, timezone +from typing import Any, Dict, Optional +from urllib.parse import parse_qs +from ._instrumentation import HANDSHAKE_OBSERVER_KEY, record_initialize +from ._internal import MCPAnalyticsData, get_server_tracking_data, tracking_key from .logger import log, warn -from .session import new_session_id +from .session import handshake_session_id, new_session_id from .session_token import ( MCP_SESSION_HEADER, SessionTokenPayload, @@ -73,16 +83,48 @@ class PostHogMcpStatelessSessionMiddleware: response header at ``initialize`` (when the client sent none) and decodes the replayed token on every request, exposing it via :func:`get_mcp_session`. + Pass the instrumented ``server`` to also capture ``$mcp_initialize`` from the + handshake itself on MCP SDK 1.x, which answers ``initialize`` outside anything + ``instrument()`` can wrap. ``instrument()`` does this for the apps it autowires; + add it yourself when you mount the middleware by hand on an instrumented server. + Leave it out for a custom ``PostHogMCP`` dispatcher, which captures its own + events. + Fail-safe: any error while sniffing/minting is logged and the request passes through untouched -- analytics must never break the host.""" - def __init__(self, app: Any) -> None: + def __init__(self, app: Any, server: Any = None) -> None: self.app = app + self._server: Optional[weakref.ReferenceType[Any]] = None + if server is not None: + try: + self._server = weakref.ref(tracking_key(server)) + except TypeError as error: + log(f"PostHog MCP session middleware: cannot bind server= - {error}") + + def _analytics(self) -> Optional[MCPAnalyticsData]: + """The bound server's tracking state, when this middleware is the one that + captures its handshakes: on mcp 2.x ``Server.middleware`` already does.""" + try: + server = self._server() if self._server is not None else None + data = get_server_tracking_data(server) if server is not None else None + except TypeError: # an unhashable server= can never be instrumented + return None + return None if data is None or data.handshake_observed_in_process else data async def __call__(self, scope: Any, receive: Any, send: Any) -> None: if scope.get("type") != "http": await self.app(scope, receive, send) return + # Tell the in-process fallback that this request's handshake is captured + # here. Two stacked middlewares see the same request; the outer one owns it. + # The SSE transport stays unmarked: its session is in a query parameter, so + # its first request captures the handshake, in the session its calls use. + observer = None + if HANDSHAKE_OBSERVER_KEY not in scope and not _is_sse(scope): + observer = self._analytics() + if observer is not None: + scope[HANDSHAKE_OBSERVER_KEY] = observer try: incoming_header = read_mcp_session_header(_headers_dict(scope)) @@ -98,6 +140,8 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: # Only POSTs carry JSON-RPC. Mint only when the client replayed no session # id at all -- if one is present (ours or a stateful transport's), leave it. + # Such a request is not read, so an initialize that reuses a session id is + # not captured here. Tool calls, which all carry one, stay unread. if scope.get("method") != "POST" or incoming_header is not None: await self.app(scope, receive, send) return @@ -105,24 +149,51 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: # Guard the whole sniff+mint path: on any failure (a raising receive() or # adversarial body) the request passes through untouched, per the contract. token = None + initialize = None try: body, receive = await _buffer_body(receive) - token = _mint_token_if_initialize(body) + initialize = _parse_initialize(body) + if initialize is not None: + token = _mint_token(initialize) except Exception as error: # noqa: BLE001 log(f"PostHog MCP session middleware: mint failed - {error}") - if token is None: + if initialize is None: await self.app(scope, receive, send) return - scope[_SCOPE_KEY] = decode_session_id(token) - await self.app(scope, receive, _sending_session_header(send, token)) + started = datetime.now(timezone.utc) + if observer is not None: + send = _InitializeResponse( + send, + functools.partial( + _record_initialize, observer, initialize, scope, started + ), + ) + if token is not None: + scope[_SCOPE_KEY] = decode_session_id(token) + send = _sending_session_header(send, token) + + await self.app(scope, receive, send) + + +def _is_sse(scope: Any) -> bool: + """The deprecated SSE transport, which keys its session on ``?session_id=``.""" + try: + query = parse_qs((scope.get("query_string") or b"").decode("latin-1")) + except Exception: # noqa: BLE001 + return False + return "session_id" in query def _headers_dict(scope: Any) -> dict[str, str]: """ASGI raw headers (list of (bytes, bytes)) -> a lowercased str dict.""" + return _decode_headers(scope.get("headers") or []) + + +def _decode_headers(raw: Any) -> dict[str, str]: result: dict[str, str] = {} - for key, value in scope.get("headers") or []: + for key, value in raw: try: result[key.decode("latin-1").lower()] = value.decode("latin-1") except Exception: # noqa: BLE001 @@ -143,10 +214,13 @@ async def _buffer_body(receive: Any) -> tuple[bytes, Any]: total = 0 complete = False overflow = False + # A non-body message (e.g. http.disconnect) ends the read early. It is kept + # and handed to the app after the prefix, so the app sees the disconnect. + held: Optional[dict[str, Any]] = None while True: message = await receive() if message.get("type") != "http.request": - # A non-body message (e.g. http.disconnect); stop with what we have. + held = message break chunks.append(message.get("body", b"") or b"") total += len(chunks[-1]) @@ -162,11 +236,14 @@ async def _buffer_body(receive: Any) -> tuple[bytes, Any]: async def replay() -> dict[str, Any]: # Replay the buffered prefix once; if we stopped early, flag more_body so - # the app keeps pulling the rest straight from the original transport. - nonlocal replayed + # the app keeps reading: the held message, then the original transport. + nonlocal replayed, held if not replayed: replayed = True - return {"type": "http.request", "body": buffered, "more_body": overflow} + return {"type": "http.request", "body": buffered, "more_body": not complete} + if held is not None: + message, held = held, None + return message return await receive() # Only sniff (parse for `initialize`) when we captured the whole small body. @@ -174,10 +251,9 @@ async def replay() -> dict[str, Any]: return sniff, replay -def _mint_token_if_initialize(body: bytes) -> Optional[str]: - """If ``body`` is an ``initialize`` JSON-RPC request, mint and return a token - string carrying a fresh session id + the client's self-reported identity. - Returns ``None`` for anything else (never raises).""" +def _parse_initialize(body: bytes) -> Optional[Dict[str, Any]]: + """The JSON-RPC message in ``body`` if it is an ``initialize`` request, else + ``None`` (never raises).""" if not body: return None try: @@ -188,10 +264,13 @@ def _mint_token_if_initialize(body: bytes) -> Optional[str]: return None if not isinstance(message, dict) or message.get("method") != "initialize": return None - params = message.get("params") - params = params if isinstance(params, dict) else {} - client_info = params.get("clientInfo") - client_info = client_info if isinstance(client_info, dict) else {} + return message + + +def _mint_token(message: Dict[str, Any]) -> Optional[str]: + """Mint a token carrying a fresh session id + the client's self-reported + identity from an ``initialize`` message. Returns ``None`` on failure.""" + params, client_info = _initialize_params(message) try: return encode_session_id( SessionTokenPayload( @@ -206,6 +285,82 @@ def _mint_token_if_initialize(body: bytes) -> Optional[str]: return None +def _initialize_params( + message: Dict[str, Any], +) -> tuple[Dict[str, Any], Dict[str, Any]]: + params = message.get("params") + params = params if isinstance(params, dict) else {} + client_info = params.get("clientInfo") + return params, client_info if isinstance(client_info, dict) else {} + + +class _InitializeResponse: + """A ``send`` that records the handshake when the ``initialize`` response + starts, then passes the response on. The app calls it from inside its own + request context, so an auth middleware's user is still set for ``identify``. + The start carries the status and the ``Mcp-Session-Id`` the client will send + back (whichever of the transport or the mint set it). The body is not read.""" + + def __init__(self, send: Any, record: Any) -> None: + self._send = send + self._record = record + + async def __call__(self, message: Dict[str, Any]) -> None: + if message.get("type") == "http.response.start" and self._record is not None: + record, self._record = self._record, None + headers = _decode_headers(message.get("headers") or []) + await record( + message.get("status"), headers.get(MCP_SESSION_HEADER.lower()) or None + ) + await self._send(message) + + +class _InitializeContext: + """``extra["ctx"]`` for callbacks run on an ``initialize`` observed here. + + MCP SDK 1.x answers ``initialize`` inside its session layer, outside any request + context, so there is no SDK context object to hand over. This carries only the + HTTP request, which is what :func:`~posthog.mcp.get_request_headers` reads.""" + + def __init__(self, request: Any) -> None: + self.request = request + + +async def _record_initialize( + data: MCPAnalyticsData, + message: Dict[str, Any], + scope: Any, + started: datetime, + status: Optional[int], + session_header: Optional[str], +) -> None: + """Emit ``$mcp_initialize`` for a handshake the app answered, in the session the + client's later requests resolve to: the ``Mcp-Session-Id`` they send back (a + stateful transport's id or our token), or this process's own session when the + response carried none. The body is not read, so the + event carries the protocol version the client asked for.""" + try: + if status is None or not 200 <= status < 300: + return + from starlette.requests import Request + + header = session_header + await record_initialize( + data, + session_id=await handshake_session_id( + data, header, decode_session_id(header) + ), + params=_initialize_params(message)[0], + result=None, + extra={"session_id": header, "ctx": _InitializeContext(Request(scope))}, + timestamp=started, + ) + except Exception as error: # noqa: BLE001 - analytics must never break the host + log( + f"PostHog MCP session middleware: failed to capture $mcp_initialize - {error}" + ) + + def _sending_session_header(send: Any, token: str) -> Any: """Wrap ``send`` so the outgoing response start carries ``Mcp-Session-Id: token`` -- but only if the app/transport didn't already set one (never clobber a @@ -251,7 +406,7 @@ def autowire_stateless_mint(server: Any) -> None: if not callable(original) or getattr(original, _AUTOWIRED, False): continue try: - setattr(server, attr, _wrap_app_factory(original)) + setattr(server, attr, _wrap_app_factory(original, server)) except Exception as error: # noqa: BLE001 - never let wiring break instrument() log(f"PostHog MCP: could not auto-wire stateless mint on {attr} - {error}") @@ -312,7 +467,7 @@ def _app_already_wrapped(app: Any) -> bool: return False -def _wrap_app_factory(original: Any) -> Any: +def _wrap_app_factory(original: Any, server: Any = None) -> Any: @functools.wraps(original) def factory(*args: Any, **kwargs: Any) -> Any: app = original(*args, **kwargs) @@ -320,10 +475,10 @@ def factory(*args: Any, **kwargs: Any) -> Any: return app add_middleware = getattr(app, "add_middleware", None) if callable(add_middleware): - add_middleware(PostHogMcpStatelessSessionMiddleware) + add_middleware(PostHogMcpStatelessSessionMiddleware, server=server) return app # Not a Starlette app -- wrap as raw ASGI so minting still happens. - return PostHogMcpStatelessSessionMiddleware(app) + return PostHogMcpStatelessSessionMiddleware(app, server=server) setattr(factory, _AUTOWIRED, True) return factory diff --git a/posthog/mcp/session.py b/posthog/mcp/session.py index dc59f51fb..4da6a6273 100644 --- a/posthog/mcp/session.py +++ b/posthog/mcp/session.py @@ -60,6 +60,25 @@ async def resolve_session_id( return session_id +async def handshake_session_id( + data: MCPAnalyticsData, + mcp_session_id: Optional[str], + token: Optional[SessionTokenPayload], +) -> str: + """The session an ``initialize`` belongs to: the one the client's later requests + resolve to through :func:`resolve_session_id_with_source`. + + A token or a transport session id gives it directly, so this leaves the server's + shared session alone. One server serves many clients, and a handshake must not + move the session that ``McpAnalytics.capture()`` reads. Without either, the + session is this process's own, which only the resolver can hand out.""" + if token is not None: + return token.session_id + if mcp_session_id: + return derive_session_id_from_mcp_session(mcp_session_id) + return await resolve_session_id(data, None) + + async def resolve_session_id_with_source( data: MCPAnalyticsData, mcp_session_id: Optional[str], diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py new file mode 100644 index 000000000..022235fa1 --- /dev/null +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -0,0 +1,990 @@ +"""``$mcp_initialize`` is one event per ``initialize`` the client actually sent. + +The spec (sdk-specs ``mcp-analytics``) allows ``$mcp_initialize`` only for a legacy +``initialize`` request, and never without one. Python cannot wrap ``initialize`` +the way posthog-js does — mcp 1.x answers it inside ``ServerSession`` and mcp 2.x +refuses a handler for it — so these tests drive a real streamable-HTTP app over +raw JSON-RPC and count what comes out. Written once, collected under both SDK +majors; the version-specific code is :func:`make_pod`, which builds each server +shape a host might run. + +The shapes are the ones agents produce in the field: a client that reconnects +before every few calls (Codex, hosted connectors), a probe that handshakes and +leaves, and a stateless fleet behind a load balancer with no stickiness, where +one handshake's requests land on several pods. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from contextlib import AsyncExitStack, asynccontextmanager +from typing import Any, Dict, List, Optional + +import httpx +import pytest + +from posthog.mcp import ( + MCPAnalyticsOptions, + UserIdentity, + decode_session_id, + get_request_headers, + instrument, +) +from posthog.test.mcp._helpers import ( + MCP_MAJOR, + FakeClient, + events_named as _events, + flush_background as _flush, +) + +PROTOCOL_VERSION = "2025-06-18" +BASE_HEADERS = { + "content-type": "application/json", + "accept": "application/json, text/event-stream", +} + + +#: The server shapes a host builds, each served the way that shape is served: +#: the official SDK's high-level server, standalone ``fastmcp`` (jlowin's 2.x/3.x +#: on mcp 1.x, 4.x on mcp 2.x), and a low-level ``Server``. +KINDS = ["official", "fastmcp", "lowlevel"] + + +def make_pod( + client: FakeClient, + *, + kind: str = "official", + stateless: bool = True, + options: Optional[MCPAnalyticsOptions] = None, + bind: bool = True, +): + """One instrumented server of ``kind``, as ``(asgi_app, lifespan, server)``. + + ``bind`` only matters for an mcp 1.x low-level server: nothing autowires it, + so the middleware is mounted by hand, with or without ``server=``.""" + options = options or MCPAnalyticsOptions(enable_conversation_id=False) + if kind == "fastmcp": + return _fastmcp_pod(client, options, stateless) + if kind == "lowlevel": + return _lowlevel_pod(client, options, stateless, bind) + if MCP_MAJOR < 2: + from mcp.server.fastmcp import FastMCP + from mcp.server.transport_security import TransportSecuritySettings + + server = FastMCP( + "handshake", + stateless_http=stateless, + json_response=True, + transport_security=TransportSecuritySettings( + enable_dns_rebinding_protection=False + ), + ) + else: + from mcp.server.mcpserver import MCPServer + + server = MCPServer("handshake") + + @server.tool() + def add(a: int, b: int) -> int: + return a + b + + instrument(server, client, options) + if MCP_MAJOR < 2: + app = server.streamable_http_app() + else: + app = server.streamable_http_app(json_response=True, stateless_http=stateless) + return app, server.session_manager.run(), server + + +def _fastmcp_pod(client, options, stateless): + fastmcp = pytest.importorskip("fastmcp") + + server = fastmcp.FastMCP("handshake") + + @server.tool() + def add(a: int, b: int) -> int: + return a + b + + instrument(server, client, options) + app = server.http_app(json_response=True, stateless_http=stateless) + return app, app.router.lifespan_context(app), server + + +def _lowlevel_pod(client, options, stateless, bind): + import mcp.types as mcp_types + from mcp.server.lowlevel import Server + + schema = { + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + } + + def total(arguments): + return str((arguments or {}).get("a", 0) + (arguments or {}).get("b", 0)) + + if MCP_MAJOR >= 2: + + async def on_list_tools(ctx, params): + return mcp_types.ListToolsResult( + tools=[mcp_types.Tool(name="add", input_schema=schema)] + ) + + async def on_call_tool(ctx, params): + return mcp_types.CallToolResult( + content=[ + mcp_types.TextContent(type="text", text=total(params.arguments)) + ] + ) + + server = Server( + "handshake", on_list_tools=on_list_tools, on_call_tool=on_call_tool + ) + instrument(server, client, options) + app = server.streamable_http_app(json_response=True, stateless_http=stateless) + return app, server.session_manager.run(), server + + from mcp.server.streamable_http_manager import StreamableHTTPSessionManager + + from posthog.mcp import PostHogMcpStatelessSessionMiddleware + + server = Server("handshake") + + @server.list_tools() + async def list_tools_handler(): + return [mcp_types.Tool(name="add", inputSchema=schema)] + + @server.call_tool() + async def call_tool_handler(name, arguments): + return [mcp_types.TextContent(type="text", text=total(arguments))] + + instrument(server, client, options) + manager = StreamableHTTPSessionManager( + app=server, stateless=stateless, json_response=True + ) + app = ( + PostHogMcpStatelessSessionMiddleware(manager.handle_request, server=server) + if bind + else PostHogMcpStatelessSessionMiddleware(manager.handle_request) + ) + return app, manager.run(), server + + +@asynccontextmanager +async def fleet(client: FakeClient, pods: int = 1, **pod_options: Any): + """``pods`` independently instrumented servers, each a live HTTP surface. Every + pod has its own tracking state, exactly like separate processes behind a + load balancer.""" + async with AsyncExitStack() as stack: + surfaces: List[httpx.AsyncClient] = [] + for _ in range(pods): + app, lifespan, _ = make_pod(client, **pod_options) + await stack.enter_async_context(lifespan) + surfaces.append( + await stack.enter_async_context( + httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://127.0.0.1:8000", + ) + ) + ) + yield RoundRobin(surfaces) + + +class RoundRobin: + """A load balancer with no stickiness: each request goes to the next pod.""" + + def __init__(self, pods: List[httpx.AsyncClient]): + self.pods = pods + self.turn = 0 + + async def post( + self, body: Dict[str, Any], headers: Dict[str, str] + ) -> httpx.Response: + pod = self.pods[self.turn % len(self.pods)] + self.turn += 1 + return await pod.post("/mcp", json=body, headers=headers) + + +def rpc(method: str, params: Dict[str, Any], request_id: Optional[int] = 1): + body: Dict[str, Any] = {"jsonrpc": "2.0", "method": method, "params": params} + if request_id is not None: + body["id"] = request_id + return body + + +def initialize_params(client_name: str) -> Dict[str, Any]: + return { + "protocolVersion": PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": client_name, "version": "0.50.0"}, + } + + +async def handshake( + lb: RoundRobin, + *, + client_name: str = "codex-mcp-client", + extra_headers: Optional[Dict[str, str]] = None, + send_initialized: bool = True, +): + """``initialize`` then ``notifications/initialized``, as every MCP client + connects. Returns the headers the client replays afterwards and the + negotiated protocol version.""" + headers = {**BASE_HEADERS, **(extra_headers or {})} + response = await lb.post(rpc("initialize", initialize_params(client_name)), headers) + assert response.status_code == 200, response.text + negotiated = response.json()["result"]["protocolVersion"] + replay = {**headers, "mcp-protocol-version": negotiated} + session_header = response.headers.get("mcp-session-id") + if session_header: + replay["mcp-session-id"] = session_header + if send_initialized: + await lb.post(rpc("notifications/initialized", {}, None), replay) + return replay, negotiated + + +async def call_add( + lb: RoundRobin, headers: Dict[str, str], request_id: int = 3, **extra: Any +): + """``tools/call add``. ``extra`` adds arguments or params (``_meta``).""" + meta = {"_meta": extra.pop("_meta")} if "_meta" in extra else {} + arguments = {"a": 2, "b": 3, "context": "adding", **extra} + response = await lb.post( + rpc("tools/call", {"name": "add", "arguments": arguments, **meta}, request_id), + headers, + ) + assert response.status_code == 200, response.text + return response + + +async def list_tools(lb: RoundRobin, headers: Dict[str, str]): + response = await lb.post(rpc("tools/list", {}, 2), headers) + assert response.status_code == 200, response.text + return response + + +# --- one event per handshake ------------------------------------------------------ + + +@pytest.mark.parametrize("kind", KINDS) +async def test_one_initialize_per_handshake_across_stateless_pods(kind): + # The overcount this file exists for: dedupe used to be per process, so every + # pod a session touched emitted its own $mcp_initialize. + client = FakeClient() + async with fleet(client, pods=2, kind=kind) as lb: + headers, negotiated = await handshake(lb) + await list_tools(lb, headers) + await call_add(lb, headers, 3) + await call_add(lb, headers, 4) + await call_add(lb, headers, 5) + await _flush() + + inits = _events(client, "$mcp_initialize") + assert len(inits) == 1 + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 3 + sessions = {c["properties"]["$session_id"] for c in calls} + assert sessions == {inits[0]["properties"]["$session_id"]} + props = inits[0]["properties"] + assert props["$mcp_client_name"] == "codex-mcp-client" + assert props["$mcp_client_version"] == "0.50.0" + assert props["$mcp_protocol_version"] == negotiated + assert ( + props["$mcp_parameters"]["params"]["clientInfo"]["name"] == "codex-mcp-client" + ) + assert props["$mcp_duration_ms"] >= 0 + + +@pytest.mark.parametrize("kind", KINDS) +async def test_each_reconnect_is_one_initialize(kind): + client = FakeClient() + async with fleet(client, pods=2, kind=kind) as lb: + for _ in range(3): + headers, _ = await handshake(lb) + await list_tools(lb, headers) + await call_add(lb, headers) + await call_add(lb, headers, 4) + await _flush() + + inits = _events(client, "$mcp_initialize") + assert len(inits) == 3 + assert len({e["properties"]["$session_id"] for e in inits}) == 3 + assert len(_events(client, "$mcp_tool_call")) == 6 + + +@pytest.mark.parametrize("kind", KINDS) +async def test_a_bare_handshake_is_counted(kind): + # Probes and health checkers connect and leave. There is no later request to + # hang a synthesized event on, which used to make them invisible. + client = FakeClient() + async with fleet(client, kind=kind) as lb: + await handshake(lb, send_initialized=False) + await handshake(lb) + await _flush() + + assert len(_events(client, "$mcp_initialize")) == 2 + + +@pytest.mark.parametrize("kind", KINDS) +async def test_initialize_session_matches_a_stateful_transport(kind): + # On a stateful server the transport assigns Mcp-Session-Id in the initialize + # response; the event must land in the session every later request resolves to. + client = FakeClient() + async with fleet(client, stateless=False, kind=kind) as lb: + headers, _ = await handshake(lb) + await call_add(lb, headers) + await call_add(lb, headers, 4) + await _flush() + + inits = _events(client, "$mcp_initialize") + assert len(inits) == 1 + calls = _events(client, "$mcp_tool_call") + assert {c["properties"]["$session_id"] for c in calls} == { + inits[0]["properties"]["$session_id"] + } + + +async def test_initialize_is_attributed_on_the_pod_that_served_it(): + def identify(request, extra): + headers = get_request_headers(extra) or {} + token = headers.get("authorization", "") + return ( + UserIdentity(distinct_id=token.removeprefix("Bearer ")) if token else None + ) + + client = FakeClient() + options = MCPAnalyticsOptions(enable_conversation_id=False, identify=identify) + async with fleet(client, pods=2, options=options) as lb: + headers, _ = await handshake( + lb, extra_headers={"authorization": "Bearer user_alice"} + ) + await call_add(lb, headers) + await _flush() + + (init,) = _events(client, "$mcp_initialize") + assert init["distinct_id"] == "user_alice" + assert init["properties"]["$mcp_client_user_agent"] + + +@pytest.mark.skipif( + MCP_MAJOR < 2, + reason="mcp 1.x reads only the response head, and a JSON-RPC error is an HTTP 200", +) +async def test_a_failed_initialize_is_not_captured(): + client = FakeClient() + async with fleet(client) as lb: + response = await lb.post(rpc("initialize", {"capabilities": {}}), BASE_HEADERS) + await _flush() + + assert "error" in response.json() + assert _events(client, "$mcp_initialize") == [] + + +async def test_an_echoed_conversation_does_not_initialize_again(): + # ADR-0004 anchors echoed calls to the conversation's session. That session + # never sent an initialize, so it gets no $mcp_initialize of its own. + from posthog.mcp._conversation_id import _uuid7 + + conversation_id = _uuid7() + client = FakeClient() + options = MCPAnalyticsOptions(enable_conversation_id=True) + async with fleet(client, pods=2, options=options) as lb: + headers, _ = await handshake(lb) + for request_id in (3, 4, 5): + await call_add(lb, headers, request_id, conversation_id=conversation_id) + await _flush() + + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 3 + assert {c["properties"]["$mcp_conversation_id"] for c in calls} == {conversation_id} + assert len(_events(client, "$mcp_initialize")) == 1 + + +# --- the 2026-07-28 revision has no initialize ------------------------------------- + + +@pytest.mark.skipif(MCP_MAJOR < 2, reason="mcp 1.x cannot serve 2026-07-28") +async def test_modern_era_requests_emit_no_initialize(): + from posthog.test.mcp._helpers_v2 import modern_headers, modern_meta + + client = FakeClient() + headers = {**BASE_HEADERS, **modern_headers("tools/call", "add")} + async with fleet(client, pods=2) as lb: + for request_id in (1, 2, 3): + await call_add(lb, headers, request_id, _meta=modern_meta()) + await _flush() + + assert len(_events(client, "$mcp_tool_call")) == 3 + assert _events(client, "$mcp_initialize") == [] + + +# --- mcp 1.x low-level server, middleware mounted by hand -------------------------- + + +@pytest.mark.skipif(MCP_MAJOR >= 2, reason="mcp 2.x observes initialize in-process") +async def test_an_unbound_middleware_keeps_the_first_request_fallback(): + # Without server= the middleware cannot reach the tracking state, so the first + # request of the session on each pod still stands in for the handshake: the + # pre-existing behaviour, kept so a hand-mounted middleware never goes dark. + client = FakeClient() + async with fleet(client, pods=2, kind="lowlevel", bind=False) as lb: + headers, _ = await handshake(lb) + await call_add(lb, headers) + await call_add(lb, headers, 4) + await _flush() + + assert len(_events(client, "$mcp_initialize")) == 2 + + +# --- stdio: a real server process, driven over its pipes --------------------------- + +_STDIO_SERVER = """ +import importlib.metadata +import json +import os + +from posthog.mcp import MCPAnalyticsOptions, instrument + + +class FileClient: + # Appends each capture to a file the test reads, one JSON line per event. + def capture(self, event, distinct_id=None, properties=None, **kwargs): + with open(os.environ["EVENTS_FILE"], "a") as sink: + sink.write(json.dumps({"event": event, "properties": properties or {}}, default=str) + "\\n") + + +if int(importlib.metadata.version("mcp").split(".")[0]) < 2: + from mcp.server.fastmcp import FastMCP as Server +else: + from mcp.server.mcpserver import MCPServer as Server + +server = Server("stdio-handshake") + + +@server.tool() +def add(a: int, b: int) -> int: + return a + b + + +instrument(server, FileClient(), MCPAnalyticsOptions(enable_conversation_id=False)) +server.run() +""" + + +class StdioServer: + """The instrumented server as a child process on stdio, the way a desktop client + launches it: newline-delimited JSON-RPC on stdin/stdout.""" + + def __init__(self, process, events_file): + self.process = process + self.events_file = events_file + + async def send(self, method, params, request_id=None): + self.process.stdin.write( + (json.dumps(rpc(method, params, request_id)) + "\n").encode() + ) + await self.process.stdin.drain() + if request_id is None: + return None + line = await asyncio.wait_for(self.process.stdout.readline(), timeout=20) + return json.loads(line) + + async def events(self, name, count): + """Captures are fire-and-forget in the child: wait for ``count`` of them, + then a moment more for any that would make it one too many.""" + for _ in range(100): + if self.events_file.exists(): + lines = self.events_file.read_text().splitlines() + found = [e for e in map(json.loads, lines) if e["event"] == name] + if len(found) >= count: + break + await asyncio.sleep(0.05) + await asyncio.sleep(0.2) + lines = self.events_file.read_text().splitlines() + return [e for e in map(json.loads, lines) if e["event"] == name] + + +@asynccontextmanager +async def stdio_server(tmp_path): + script = tmp_path / "stdio_server.py" + script.write_text(_STDIO_SERVER) + events_file = tmp_path / "events.jsonl" + process = await asyncio.create_subprocess_exec( + sys.executable, + str(script), + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + env={**os.environ, "EVENTS_FILE": str(events_file)}, + cwd=str(tmp_path), + ) + try: + yield StdioServer(process, events_file) + finally: + if process.stdin and not process.stdin.is_closing(): + process.stdin.close() + try: + await asyncio.wait_for(process.wait(), timeout=10) + except asyncio.TimeoutError: + process.kill() + await process.wait() + + +async def test_stdio_handshake_then_calls(tmp_path): + async with stdio_server(tmp_path) as server: + reply = await server.send("initialize", initialize_params("claude-desktop"), 1) + assert "result" in reply, reply + await server.send("notifications/initialized", {}) + await server.send("tools/list", {}, 2) + for request_id in (3, 4): + reply = await server.send( + "tools/call", + {"name": "add", "arguments": {"a": 1, "b": 2, "context": "stdio"}}, + request_id, + ) + assert reply["result"]["isError"] is False, reply + calls = await server.events("$mcp_tool_call", 2) + inits = await server.events("$mcp_initialize", 1) + + assert len(calls) == 2 + assert len(inits) == 1 + props = inits[0]["properties"] + assert props["$mcp_client_name"] == "claude-desktop" + assert {c["properties"]["$session_id"] for c in calls} == {props["$session_id"]} + + +@pytest.mark.skipif(MCP_MAJOR < 2, reason="mcp 1.x stdio has no handshake seam") +async def test_stdio_bare_handshake_is_counted(tmp_path): + async with stdio_server(tmp_path) as server: + reply = await server.send("initialize", initialize_params("claude-desktop"), 1) + assert "result" in reply, reply + inits = await server.events("$mcp_initialize", 1) + + assert len(inits) == 1 + assert ( + inits[0]["properties"]["$mcp_protocol_version"] + == reply["result"]["protocolVersion"] + ) + + +# --- custom dispatcher: PostHogMCP behind an unbound middleware --------------------- + + +def make_dispatcher(posthog_mcp): + """A hand-rolled JSON-RPC endpoint, the shape ``PostHogMCP`` exists for: the host + answers ``initialize`` itself and captures it through the manual API.""" + from starlette.applications import Starlette + from starlette.responses import JSONResponse, Response + from starlette.routing import Route + + from posthog.mcp import get_mcp_session + + async def endpoint(request): + message = await request.json() + if "id" not in message: # a notification: accepted, nothing to answer + return Response(status_code=202) + result: Dict[str, Any] = {"tools": []} + if message["method"] == "initialize": + session = get_mcp_session(request) + posthog_mcp.capture_initialize( + client_name=message["params"]["clientInfo"]["name"], + session_id=session.session_id if session else None, + ) + result = { + "protocolVersion": message["params"]["protocolVersion"], + "capabilities": {}, + "serverInfo": {"name": "dispatcher", "version": "1.0.0"}, + } + return JSONResponse({"jsonrpc": "2.0", "id": message["id"], "result": result}) + + return Starlette(routes=[Route("/mcp", endpoint, methods=["POST"])]) + + +async def test_a_posthogmcp_dispatcher_reports_only_its_own_initialize(): + from posthog.mcp import PostHogMCP, PostHogMcpStatelessSessionMiddleware + + captured: List[Dict[str, Any]] = [] + + def keep(event): + captured.append(event) + return event + + posthog_mcp = PostHogMCP("phc_test", send=False, before_send=keep) + # An instrumented server in the same process: an unbound middleware must never + # attach the dispatcher's handshakes to it. + bystander = FakeClient() + make_pod(bystander) + + app = PostHogMcpStatelessSessionMiddleware(make_dispatcher(posthog_mcp)) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://127.0.0.1:8000" + ) as http: + lb = RoundRobin([http]) + headers, _ = await handshake(lb, client_name="dispatcher-client") + await list_tools(lb, headers) + await _flush() + + inits = [e for e in captured if e["event"] == "$mcp_initialize"] + assert len(inits) == 1 + assert inits[0]["properties"]["$mcp_client_name"] == "dispatcher-client" + # The session the host read back from the minted token, not a second one. + token = decode_session_id(headers["mcp-session-id"]) + assert token is not None + assert inits[0]["properties"]["$session_id"] == token.session_id + assert bystander.events == [] + + +# --- the deprecated SSE transport: session in a query parameter --------------------- + + +@asynccontextmanager +async def sse_server(client: FakeClient): + """An instrumented server's ``sse_app()`` on a real port. The SSE transport keeps + its session in the ``session_id`` query parameter of ``/messages/``, not in a + header, and answers each POST with 202 and the reply on the event stream.""" + import uvicorn + + if MCP_MAJOR < 2: + from mcp.server.fastmcp import FastMCP + from mcp.server.transport_security import TransportSecuritySettings + + server = FastMCP( + "handshake", + transport_security=TransportSecuritySettings( + enable_dns_rebinding_protection=False + ), + ) + else: + from mcp.server.mcpserver import MCPServer + + server = MCPServer("handshake") + + @server.tool() + def add(a: int, b: int) -> int: + return a + b + + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=False)) + config = uvicorn.Config(server.sse_app(), port=0, log_level="error", ws="none") + http_server = uvicorn.Server(config) + serving = asyncio.create_task(http_server.serve()) + while not http_server.started: + await asyncio.sleep(0.01) + port = http_server.servers[0].sockets[0].getsockname()[1] + try: + yield f"http://127.0.0.1:{port}" + finally: + http_server.should_exit = True + await asyncio.wait_for(serving, timeout=10) + + +async def test_sse_transport_handshake_then_calls(): + client = FakeClient() + async with sse_server(client) as url: + async with httpx.AsyncClient(base_url=url, timeout=20) as http: + async with http.stream("GET", "/sse") as stream: + lines = stream.aiter_lines() + + async def next_data() -> str: + async for line in lines: + if line.startswith("data:"): + return line[len("data:") :].strip() + raise AssertionError("the event stream closed") + + endpoint = await next_data() + + async def send(method, params, request_id=None): + response = await http.post( + endpoint, json=rpc(method, params, request_id) + ) + assert response.status_code == 202, response.text + return json.loads(await next_data()) if request_id else None + + await send("initialize", initialize_params("sse-client"), 1) + await send("notifications/initialized", {}) + await send("tools/list", {}, 2) + for request_id in (3, 4): + arguments = {"a": 1, "b": 2, "context": "sse"} + reply = await send( + "tools/call", + {"name": "add", "arguments": arguments}, + request_id, + ) + assert reply["result"]["isError"] is False, reply + await _flush() + + (init,) = _events(client, "$mcp_initialize") + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 2 + assert {c["properties"]["$session_id"] for c in calls} == { + init["properties"]["$session_id"] + } + + +# --- the middleware is per request and fail-safe ----------------------------------- + + +@pytest.mark.skipif(MCP_MAJOR >= 2, reason="mcp 2.x has no first-request fallback") +async def test_a_route_without_the_middleware_keeps_the_fallback(): + # One process, one server, two routes: the autowired app with the bound + # middleware, and the session manager mounted by hand without it. Each route + # reports its own handshakes, once. + client = FakeClient() + app, lifespan, server = make_pod(client) + async with lifespan: + for surface in (app, server.session_manager.handle_request): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=surface), + base_url="http://127.0.0.1:8000", + ) as http: + lb = RoundRobin([http]) + headers, _ = await handshake(lb) + await call_add(lb, headers) + await call_add(lb, headers, 4) + await _flush() + + inits = _events(client, "$mcp_initialize") + assert len(inits) == 2 + calls = _events(client, "$mcp_tool_call") + assert {c["properties"]["$session_id"] for c in calls} == { + e["properties"]["$session_id"] for e in inits + } + + +async def test_a_handshake_leaves_the_shared_session_alone(): + # McpAnalytics.capture() reads the server's shared session. A handshake that + # carries its own session (a token here) must not move it. + from posthog.mcp._internal import get_server_tracking_data, tracking_key + + client = FakeClient() + app, lifespan, server = make_pod(client) + data = get_server_tracking_data(tracking_key(server)) + before = (data.session_id, data.session_source) + async with lifespan: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://127.0.0.1:8000" + ) as http: + await handshake(RoundRobin([http])) + await _flush() + + assert len(_events(client, "$mcp_initialize")) == 1 + assert (data.session_id, data.session_source) == before + + +async def test_a_server_that_cannot_be_bound_does_not_break_requests(): + from posthog.mcp import PostHogMcpStatelessSessionMiddleware + + async def echo(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + for unbindable in (object(), {"not": "hashable"}): + app = PostHogMcpStatelessSessionMiddleware(echo, server=unbindable) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://127.0.0.1:8000" + ) as http: + response = await http.post( + "/mcp", + json=rpc("initialize", initialize_params("x")), + headers=BASE_HEADERS, + ) + assert response.status_code == 200 + + +@pytest.mark.parametrize("kind", KINDS) +async def test_flush_waits_for_the_handshake_capture(kind): + # The handshake capture starts during the request. flush() must wait for it, + # or shutdown drops $mcp_initialize when identify or before_send is slow. + from posthog.mcp import McpAnalytics + from posthog.mcp._internal import tracking_key + + async def slow_identify(request, extra): + await asyncio.sleep(0.05) + return UserIdentity(distinct_id="user_slow") + + async def slow_before_send(event): + await asyncio.sleep(0.05) + return event + + client = FakeClient() + options = MCPAnalyticsOptions( + enable_conversation_id=False, + identify=slow_identify, + before_send=slow_before_send, + ) + app, lifespan, server = make_pod(client, kind=kind, options=options) + async with lifespan: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://127.0.0.1:8000" + ) as http: + await handshake(RoundRobin([http]), send_initialized=False) + await McpAnalytics(tracking_key(server)).flush() + assert len(_events(client, "$mcp_initialize")) == 1 + + +@pytest.mark.skipif( + MCP_MAJOR < 2, + reason="mcp 1.x does not read a request that carries a session id (README)", +) +@pytest.mark.parametrize("kind", KINDS) +async def test_an_initialize_with_a_session_header_keeps_that_session(kind): + # A stateless client that sends initialize with an existing Mcp-Session-Id + # keeps sending it, and its later requests resolve their session from it. + client = FakeClient() + async with fleet(client, pods=2, kind=kind) as lb: + headers, _ = await handshake( + lb, extra_headers={"mcp-session-id": "an-existing-session"} + ) + assert headers["mcp-session-id"] == "an-existing-session" + await call_add(lb, headers) + await call_add(lb, headers, 4) + await _flush() + + (init,) = _events(client, "$mcp_initialize") + calls = _events(client, "$mcp_tool_call") + assert {c["properties"]["$session_id"] for c in calls} == { + init["properties"]["$session_id"] + } + + +@pytest.mark.parametrize("kind", KINDS) +async def test_a_hung_identify_leaves_no_task_behind(kind): + # identify runs inline in the handshake, as it does on a tool call. A callback + # that never returns holds those requests, and leaves no analytics task behind + # for each connection. Once the callback recovers, recording resumes. + import contextlib + + import posthog.mcp._instrumentation as instr + from posthog.mcp import McpAnalytics + from posthog.mcp._internal import get_server_tracking_data, tracking_key + + recovered = asyncio.Event() + + async def identify(request, extra): + await recovered.wait() + return UserIdentity(distinct_id="user_1") + + client = FakeClient() + options = MCPAnalyticsOptions(enable_conversation_id=False, identify=identify) + app, lifespan, server = make_pod(client, kind=kind, options=options) + data = get_server_tracking_data(tracking_key(server)) + + def pending(): + with instr._tasks_lock: + tasks = list(instr._BACKGROUND_TASKS.items()) + return [task for task, owner in tasks if owner is data and not task.done()] + + async with lifespan: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://127.0.0.1:8000" + ) as http: + lb = RoundRobin([http]) + for _ in range(5): # each client gives up, as a timeout would + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(handshake(lb, send_initialized=False), 0.2) + assert pending() == [] + await asyncio.wait_for(McpAnalytics(tracking_key(server)).flush(), 1) + + recovered.set() + await handshake(lb, client_name="after-recovery", send_initialized=False) + await _flush() + + names = [ + e["properties"]["$mcp_client_name"] for e in _events(client, "$mcp_initialize") + ] + assert "after-recovery" in names + + +@pytest.mark.skipif(MCP_MAJOR >= 2, reason="the mcp 1.x auth context middleware") +async def test_identify_sees_the_auth_context_of_the_handshake(): + # FastMCP's AuthContextMiddleware sets the authenticated user for the length of + # the request and resets it on the way out. Our middleware sits outside it, so + # the handshake must be recorded while the app still handles the request. + from mcp.server.auth.middleware.auth_context import ( + auth_context_var, + get_access_token, + ) + from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser + from mcp.server.auth.provider import AccessToken + + from posthog.mcp import PostHogMcpStatelessSessionMiddleware + + def with_user(app): + user = AuthenticatedUser( + AccessToken(token="t", client_id="alice", scopes=[], expires_at=None) + ) + + async def authenticated(scope, receive, send): + reset = auth_context_var.set(user) + try: + await app(scope, receive, send) + finally: + auth_context_var.reset(reset) + + return authenticated + + def identify(request, extra): + token = get_access_token() + return UserIdentity(distinct_id=token.client_id) if token else None + + client = FakeClient() + options = MCPAnalyticsOptions(enable_conversation_id=False, identify=identify) + _, lifespan, server = make_pod(client, options=options) + app = PostHogMcpStatelessSessionMiddleware( + with_user(server.session_manager.handle_request), server=server + ) + async with lifespan: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://127.0.0.1:8000" + ) as http: + await handshake(RoundRobin([http]), send_initialized=False) + await _flush() + + (init,) = _events(client, "$mcp_initialize") + assert init["distinct_id"] == "alice" + + +@pytest.mark.parametrize("kind", KINDS) +async def test_an_escaped_initialize_still_mints_a_session(kind): + # "\u0069nitialize" is valid JSON for "initialize". The middleware parses the + # body as JSON, so it mints the token and captures the handshake. + raw = ( + b'{"jsonrpc":"2.0","id":1,"method":"\\u0069nitialize","params":' + + json.dumps(initialize_params("escaped-client")).encode() + + b"}" + ) + client = FakeClient() + async with fleet(client, kind=kind) as lb: + response = await lb.pods[0].post("/mcp", content=raw, headers=BASE_HEADERS) + await _flush() + + assert response.status_code == 200, response.text + assert decode_session_id(response.headers.get("mcp-session-id")) is not None + (init,) = _events(client, "$mcp_initialize") + assert init["properties"]["$mcp_client_name"] == "escaped-client" + + +@pytest.mark.skipif( + MCP_MAJOR < 2, + reason="mcp 1.x does not read a request that carries a session id (README)", +) +@pytest.mark.parametrize("kind", KINDS) +async def test_a_nested_method_before_the_top_level_initialize_is_counted(kind): + # JSON key order is arbitrary: params can come first, and can hold their own + # "method" key. Only the top-level method says whether this is a handshake. + params = initialize_params("nested-client") + params["capabilities"] = { + "experimental": {"vendor": {"method": "custom", "padding": "x" * 8192}} + } + body = {"params": params, "jsonrpc": "2.0", "id": 1, "method": "initialize"} + headers = {**BASE_HEADERS, "mcp-session-id": "an-existing-session"} + client = FakeClient() + async with fleet(client, kind=kind) as lb: + response = await lb.post(body, headers) + await _flush() + + assert response.status_code == 200, response.text + (init,) = _events(client, "$mcp_initialize") + assert init["properties"]["$mcp_client_name"] == "nested-client" diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index 4cb0ac810..9db5aca7e 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -418,6 +418,36 @@ async def send(message): assert MCP_SESSION_HEADER not in sent["headers"] # too big to sniff -> no mint +async def test_middleware_passes_a_disconnect_mid_body_to_the_app(): + """A client that disconnects while it sends the body must reach the app as a + disconnect, not as a complete request made of the part that arrived.""" + messages = [ + {"type": "http.request", "body": b'{"jsonrpc":"2.0","met', "more_body": True}, + {"type": "http.disconnect"}, + ] + seen = [] + + async def app(scope, receive, send): + while True: + message = await receive() + seen.append(message) + if message["type"] != "http.request" or not message.get("more_body"): + return + + async def receive(): + return messages.pop(0) if messages else {"type": "http.disconnect"} + + async def send(message): + pass + + await PostHogMcpStatelessSessionMiddleware(app)(_scope(), receive, send) + + assert seen == [ + {"type": "http.request", "body": b'{"jsonrpc":"2.0","met', "more_body": True}, + {"type": "http.disconnect"}, + ] + + # --- end-to-end against a real stateless FastMCP transport ------------------- diff --git a/posthog/test/mcp/test_v2_lowlevel.py b/posthog/test/mcp/test_v2_lowlevel.py index 905292b1d..319cbc3eb 100644 --- a/posthog/test/mcp/test_v2_lowlevel.py +++ b/posthog/test/mcp/test_v2_lowlevel.py @@ -481,7 +481,7 @@ async def late_read_resource(ctx, params): assert reads[0]["properties"]["$mcp_resource_name"] == "file:///late.txt" -async def test_initialize_and_session_reuse_across_calls(): +async def test_session_reuse_across_calls_without_a_handshake(): server = make_server() client = FakeClient() instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=False)) @@ -490,7 +490,9 @@ async def test_initialize_and_session_reuse_across_calls(): await _call_tool(server, "add", {"a": 2, "b": 2, "context": "second"}) await _flush() - assert len(_events(client, "$mcp_initialize")) == 1 + # No initialize was sent, so none is reported: Server.middleware sees every + # real handshake (test_initialize_handshake.py). + assert _events(client, "$mcp_initialize") == [] calls = _events(client, "$mcp_tool_call") session_ids = {c["properties"]["$session_id"] for c in calls} assert len(session_ids) == 1 diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index cdc991ac1..4e7a5d2c1 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -195,7 +195,7 @@ async def test_tool_call_error_is_captured_and_converted(): assert exception_list[-1]["value"] == "explode" -async def test_initialize_emitted_once_per_session(): +async def test_tool_calls_alone_emit_no_initialize(): server = make_server() client = FakeClient() instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=False)) @@ -208,7 +208,9 @@ async def test_initialize_emitted_once_per_session(): ) await _flush() - assert len(_events(client, "$mcp_initialize")) == 1 + # No initialize was sent, so none is reported: Server.middleware sees every + # real handshake (test_initialize_handshake.py). + assert _events(client, "$mcp_initialize") == [] assert len(_events(client, "$mcp_tool_call")) == 2 diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index dc2bbe85f..518dd936b 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -365,6 +365,7 @@ alias posthog.mcp.asgi.MCP_SESSION_HEADER -> posthog.mcp.session_token.MCP_SESSI alias posthog.mcp.asgi.SessionTokenPayload -> posthog.mcp.session_token.SessionTokenPayload alias posthog.mcp.asgi.decode_session_id -> posthog.mcp.session_token.decode_session_id alias posthog.mcp.asgi.encode_session_id -> posthog.mcp.session_token.encode_session_id +alias posthog.mcp.asgi.handshake_session_id -> posthog.mcp.session.handshake_session_id alias posthog.mcp.asgi.log -> posthog.mcp.logger.log alias posthog.mcp.asgi.new_session_id -> posthog.mcp.session.new_session_id alias posthog.mcp.asgi.read_mcp_session_header -> posthog.mcp.session_token.read_mcp_session_header @@ -1295,7 +1296,7 @@ class posthog.integrations.django.PosthogContextMiddleware(get_response) class posthog.integrations.flask.PosthogFlaskIntegration(app: Optional[Flask] = None, *, client: Optional[Client] = None, capture_exceptions: Optional[bool] = None, request_filter: Optional[Callable[[Request], bool]] = None, extra_properties: Optional[Callable[[Request], Mapping[str, Any]]] = None, trust_tracing_headers: Optional[bool] = None, capture_tracing_headers: Optional[bool] = None) class posthog.integrations.otel.PostHogContextSpanProcessor class posthog.mcp.McpAnalytics(key: Any) -class posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware(app: Any) +class posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware(app: Any, server: Any = None) class posthog.mcp.constants.PostHogMCPAnalyticsEvent class posthog.mcp.constants.PostHogMCPAnalyticsProperty class posthog.mcp.posthog_mcp.PostHogMCP(api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, capture_model: Union[bool, MCPAnalyticsModelOptions] = True, collect_feedback: Union[bool, CollectFeedbackOptions] = False, server_build: Optional[str] = None, enable_conversation_id: bool = True, **kwargs: Any)