From 57d8dc6221388b299fc38d8f74453f3587ea6a0a Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 15:57:46 +0300 Subject: [PATCH 01/13] fix(mcp): send one $mcp_initialize per initialize request The Python MCP SDK answers initialize itself, so posthog.mcp sent the event on the first request of each session, once per process. A stateless fleet sent one event per pod. A client that sent only initialize sent no event. Now MCP SDK 2.x reports the event from Server.middleware, on every transport. MCP SDK 1.x over HTTP reports it from PostHogMcpStatelessSessionMiddleware, which takes a new optional server= argument. MCP SDK 1.x over stdio keeps the first-request fallback. Co-Authored-By: Claude Opus 5.5 --- .../mcp-initialize-per-handshake.md | 11 + posthog/mcp/README.md | 29 +- posthog/mcp/__init__.py | 15 +- posthog/mcp/_instrument_v2.py | 95 ++- posthog/mcp/_instrumentation.py | 80 +- posthog/mcp/_internal.py | 19 + posthog/mcp/asgi.py | 240 +++++- posthog/test/mcp/test_initialize_handshake.py | 738 ++++++++++++++++++ posthog/test/mcp/test_v2_lowlevel.py | 6 +- posthog/test/mcp/test_v2_mcpserver.py | 6 +- references/public_api_snapshot.txt | 3 +- 11 files changed, 1193 insertions(+), 49 deletions(-) create mode 100644 .sampo/changesets/mcp-initialize-per-handshake.md create mode 100644 posthog/test/mcp/test_initialize_handshake.py diff --git a/.sampo/changesets/mcp-initialize-per-handshake.md b/.sampo/changesets/mcp-initialize-per-handshake.md new file mode 100644 index 000000000..16afe513f --- /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 and the server reply as `$mcp_parameters` and `$mcp_response`. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 9db9bc31e..7b682d81b 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -67,6 +67,28 @@ 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. +- **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 +407,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..e5a8b7a59 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, + _to_jsonable, copy_tools_list_result, + fire_and_forget, 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 resolve_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,92 @@ 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 event is scheduled rather than awaited: + the runner parks the connection's read loop until the ``initialize`` chain + returns, so a slow ``identify`` callback must not hold up the handshake.""" + 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) + fire_and_forget(_record_v2_initialize(data, ctx, result, started), data) + 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, and the session lives in + 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.""" + 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 + 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: + params = dict(ctx.params) if isinstance(ctx.params, Mapping) else {} + client_info = params.get("clientInfo") + client_info = client_info if isinstance(client_info, Mapping) else {} + response = _to_jsonable(result) + negotiated = ( + response.get("protocolVersion") if isinstance(response, dict) else None + ) + mcp_session_id, token = _initialize_session(ctx) + session_id = await resolve_session_id(data, mcp_session_id, token=token) + await record_initialize( + data, + session_id=session_id, + client_name=_str_or_none(client_info.get("name")), + client_version=_str_or_none(client_info.get("version")), + protocol_version=_str_or_none(negotiated) + or _str_or_none(params.get("protocolVersion")), + request={"method": "initialize", "params": params}, + response=response, + 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}") + + +def _str_or_none(value: Any) -> Optional[str]: + return value if isinstance(value, str) and value else None # --- registry plumbing --------------------------------------------------------- diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 03b66392c..0f54d3581 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -240,17 +240,74 @@ def _wrap_response(result: Any) -> Any: return serialized +async def record_initialize( + data: MCPAnalyticsData, + *, + session_id: str, + client_name: Optional[str], + client_version: Optional[str], + protocol_version: Optional[str], + request: Dict[str, Any], + response: 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 once the handshake has succeeded. 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.""" + 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": client_name, + "client_version": client_version, + "protocol_version": protocol_version, + "parameters": request, + "response": response, + "timestamp": timestamp, + } + await _apply_event_properties(data, event, request, extra) + stamp_transport_identity(event, extra) + fire_and_forget(capture_event(data, event), data) + + +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 mcp 1.x over HTTP + without a bound ASGI middleware. 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.""" + if data.handshake_observed_in_process: + return False + if session_source == "conversation": + return False + return not (data.handshake_observed_at_asgi and get_request(extra) is not None) + + 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 +430,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 +455,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 +472,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..09599d249 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -107,6 +107,13 @@ 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) + # Where this server sees the client's real `initialize`. mcp 2.x hands it to + # `Server.middleware` on every transport; on mcp 1.x only an ASGI middleware + # bound to this server sees it, and only over HTTP. Whatever neither covers + # (mcp 1.x stdio, or HTTP without a bound middleware) falls back to the + # first request of the session. + handshake_observed_in_process: bool = False + handshake_observed_at_asgi: 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 +137,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..ccba8c1e2 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,14 @@ import functools import json -from typing import Any, Optional +import weakref +from datetime import datetime, timezone +from typing import Any, Dict, Optional +from ._instrumentation import fire_and_forget, 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 new_session_id, resolve_session_id from .session_token import ( MCP_SESSION_HEADER, SessionTokenPayload, @@ -53,6 +62,9 @@ # Scope key the middleware stashes the decoded token payload under. _SCOPE_KEY = "posthog_mcp_session" +# Scope key set once a middleware has taken an ``initialize`` to capture. +_HANDSHAKE_KEY = "posthog_mcp_initialize" + # JSON-RPC request bodies are tiny; cap what we buffer so a stray large POST on # the same app can't be read into memory in full before minting. _MAX_SNIFF_BODY = 256 * 1024 @@ -73,16 +85,37 @@ 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 = weakref.ref(tracking_key(server)) if server is not None else None + self._analytics() + + def _analytics(self) -> Optional[MCPAnalyticsData]: + """The bound server's tracking state, when this middleware is the one that + observes its handshakes: on mcp 2.x ``Server.middleware`` already does.""" + server = self._server() if self._server is not None else None + data = get_server_tracking_data(server) if server is not None else None + if data is None or data.handshake_observed_in_process: + return None + data.handshake_observed_at_asgi = True + return data async def __call__(self, scope: Any, receive: Any, send: Any) -> None: if scope.get("type") != "http": await self.app(scope, receive, send) return + # Bound late if instrument() ran after this middleware was built. + self._analytics() try: incoming_header = read_mcp_session_header(_headers_dict(scope)) @@ -105,24 +138,46 @@ 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) + token = _mint_token(initialize) if initialize is not None else None 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)) + # Two stacked middlewares (autowired plus a manual one) see the same + # handshake; the outer one captures it. + data = None if scope.get(_HANDSHAKE_KEY) else self._analytics() + recorder = None + if data is not None: + scope[_HANDSHAKE_KEY] = True + recorder = _InitializeResponse(send) + send = recorder + if token is not None: + scope[_SCOPE_KEY] = decode_session_id(token) + send = _sending_session_header(send, token) + + started = datetime.now(timezone.utc) + await self.app(scope, receive, send) + if data is not None and recorder is not None: + fire_and_forget( + _record_initialize(data, initialize, recorder, scope, started), data + ) 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 @@ -174,10 +229,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 +242,19 @@ 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_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).""" + message = _parse_initialize(body) + return _mint_token(message) if message is not None else None + + +def _mint_token(message: Dict[str, Any]) -> Optional[str]: + params, client_info = _initialize_params(message) try: return encode_session_id( SessionTokenPayload( @@ -206,6 +269,143 @@ 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 passes the ``initialize`` response through and keeps what + the event needs: the status, the ``Mcp-Session-Id`` the client will replay + (whichever of the transport or the mint set it), and the body, capped like the + request sniff.""" + + def __init__(self, send: Any) -> None: + self._send = send + self.status: Optional[int] = None + self.session_header: Optional[str] = None + self.content_type = "" + self._chunks: list[bytes] = [] + self._size = 0 + + async def __call__(self, message: Dict[str, Any]) -> None: + try: + self._observe(message) + except Exception as error: # noqa: BLE001 + log(f"PostHog MCP session middleware: response read failed - {error}") + await self._send(message) + + def _observe(self, message: Dict[str, Any]) -> None: + if message.get("type") == "http.response.start": + self.status = message.get("status") + headers = _decode_headers(message.get("headers") or []) + self.session_header = headers.get(MCP_SESSION_HEADER.lower()) or None + self.content_type = headers.get("content-type", "") + elif message.get("type") == "http.response.body": + chunk = message.get("body", b"") or b"" + if self._size + len(chunk) <= _MAX_SNIFF_BODY: + self._chunks.append(chunk) + self._size += len(chunk) + + def reply(self, request_id: Any) -> Optional[Dict[str, Any]]: + """The JSON-RPC reply to ``request_id``, from a JSON body or an SSE stream. + ``None`` when there is none to read: the deprecated SSE transport answers + the POST with 202 and delivers the reply on its event stream.""" + if self._size > _MAX_SNIFF_BODY: + return None + body = b"".join(self._chunks).decode("utf-8", errors="replace") + if self.content_type.startswith("text/event-stream"): + candidates = [ + line[len("data:") :] + for line in body.splitlines() + if line.startswith("data:") + ] + else: + candidates = [body] + for candidate in candidates: + try: + message = json.loads(candidate) + except Exception: # noqa: BLE001 + continue + if isinstance(message, dict) and message.get("id") == request_id: + return message + return None + + +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 + + +def _http_request(scope: Any) -> Any: + try: + from starlette.requests import Request + + return Request(scope) + except Exception: # noqa: BLE001 - only headers are read; never fail the capture + return _Headers(_headers_dict(scope)) + + +class _Headers: + def __init__(self, headers: Dict[str, str]) -> None: + self.headers = headers + + +async def _record_initialize( + data: MCPAnalyticsData, + message: Dict[str, Any], + response: _InitializeResponse, + scope: Any, + started: datetime, +) -> None: + """Emit ``$mcp_initialize`` for a handshake that succeeded, in the session the + client's later requests resolve to: the ``Mcp-Session-Id`` they replay (a + stateful transport's id or our token), or this process's own session when the + response carried none.""" + try: + if response.status is None or not 200 <= response.status < 300: + return + reply = response.reply(message.get("id")) + if reply is not None and "error" in reply: + return + result = reply.get("result") if reply is not None else None + params, client_info = _initialize_params(message) + header = response.session_header + session_id = await resolve_session_id( + data, header, token=decode_session_id(header) + ) + negotiated = result.get("protocolVersion") if isinstance(result, dict) else None + await record_initialize( + data, + session_id=session_id, + client_name=_str_or_none(client_info.get("name")), + client_version=_str_or_none(client_info.get("version")), + protocol_version=_str_or_none(negotiated) + or _str_or_none(params.get("protocolVersion")), + request={"method": "initialize", "params": params}, + response=result, + extra={ + "session_id": header, + "ctx": _InitializeContext(_http_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 +451,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 +512,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 +520,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/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py new file mode 100644 index 000000000..41ca43691 --- /dev/null +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -0,0 +1,738 @@ +"""``$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, + json_response: bool = True, + options: Optional[MCPAnalyticsOptions] = None, + bind: bool = True, +): + """One instrumented server of ``kind``, as ``(asgi_app, lifespan)``. + + ``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, json_response) + if kind == "lowlevel": + return _lowlevel_pod(client, options, stateless, json_response, 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=json_response, + 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=json_response, stateless_http=stateless + ) + return app, server.session_manager.run() + + +def _fastmcp_pod(client, options, stateless, json_response): + 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=json_response, stateless_http=stateless) + return app, app.router.lifespan_context(app) + + +def _lowlevel_pod(client, options, stateless, json_response, 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=json_response, stateless_http=stateless + ) + return app, server.session_manager.run() + + 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=json_response + ) + app = ( + PostHogMcpStatelessSessionMiddleware(manager.handle_request, server=server) + if bind + else PostHogMcpStatelessSessionMiddleware(manager.handle_request) + ) + return app, manager.run() + + +@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 result_of(response: httpx.Response) -> Dict[str, Any]: + """The JSON-RPC message in a response, whether it came back as JSON or as SSE.""" + if response.headers.get("content-type", "").startswith("text/event-stream"): + for line in response.text.splitlines(): + if line.startswith("data:"): + return json.loads(line[len("data:") :]) + raise AssertionError(f"no SSE data frame in {response.text!r}") + return response.json() + + +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", + { + "protocolVersion": PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": client_name, "version": "0.50.0"}, + }, + ), + headers, + ) + assert response.status_code == 200, response.text + negotiated = result_of(response)["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): + response = await lb.post( + rpc( + "tools/call", + {"name": "add", "arguments": {"a": 2, "b": 3, "context": "adding"}}, + 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 + + +@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_sse_response_mode(): + # json_response=False answers each POST as an SSE stream, so the negotiated + # version has to come out of an SSE frame. + client = FakeClient() + async with fleet(client, pods=2, json_response=False) as lb: + headers, negotiated = 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 + assert inits[0]["properties"]["$mcp_protocol_version"] == negotiated + + +async def test_initialize_carries_the_handshake_payload(): + client = FakeClient() + async with fleet(client) as lb: + await handshake(lb) + await _flush() + + (init,) = _events(client, "$mcp_initialize") + request = init["properties"]["$mcp_parameters"] + assert request["method"] == "initialize" + assert request["params"]["clientInfo"]["name"] == "codex-mcp-client" + response = init["properties"]["$mcp_response"] + assert response["serverInfo"]["name"] == "handshake" + assert init["properties"]["$mcp_duration_ms"] >= 0 + + +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"] + + +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 result_of(response) + 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): + response = await lb.post( + rpc( + "tools/call", + { + "name": "add", + "arguments": { + "a": 1, + "b": 1, + "context": "again", + "conversation_id": conversation_id, + }, + }, + request_id, + ), + headers, + ) + assert response.status_code == 200, response.text + 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() + async with fleet(client, pods=2) as lb: + for request_id in (1, 2, 3): + response = await lb.post( + rpc( + "tools/call", + { + "name": "add", + "arguments": {"a": 1, "b": 2, "context": "modern"}, + "_meta": modern_meta(), + }, + request_id, + ), + {**BASE_HEADERS, **modern_headers("tools/call", "add")}, + ) + assert response.status_code == 200, response.text + 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) + + def events(self, name): + if not self.events_file.exists(): + return [] + lines = self.events_file.read_text().splitlines() + return [e for e in map(json.loads, lines) if e["event"] == name] + + async def wait_for(self, name, count, timeout=5.0): + """Captures are fire-and-forget in the child, so give them a moment to land.""" + for _ in range(int(timeout / 0.05)): + if len(self.events(name)) >= count: + break + await asyncio.sleep(0.05) + await asyncio.sleep(0.2) # and anything that would make it one too many + return self.events(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() + + +def _initialize_params(): + return { + "protocolVersion": PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "claude-desktop", "version": "1.2.0"}, + } + + +async def test_stdio_handshake_then_calls(tmp_path): + async with stdio_server(tmp_path) as server: + reply = await server.send("initialize", _initialize_params(), 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.wait_for("$mcp_tool_call", 2) + inits = await server.wait_for("$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(), 1) + assert "result" in reply, reply + inits = await server.wait_for("$mcp_initialize", 1) + + assert len(inits) == 1 + assert ( + inits[0]["properties"]["$mcp_protocol_version"] + == reply["result"]["protocolVersion"] + ) + + +@pytest.mark.skipif(MCP_MAJOR >= 2, reason="mcp 2.x observes the stdio handshake") +async def test_stdio_bare_handshake_is_missed_on_mcp_1(tmp_path): + # The documented limit of the first-request fallback: a stdio client that + # handshakes and never sends a request reports nothing on mcp 1.x. + async with stdio_server(tmp_path) as server: + reply = await server.send("initialize", _initialize_params(), 1) + assert "result" in reply, reply + inits = await server.wait_for("$mcp_initialize", 1, timeout=1.0) + + assert inits == [] + + +# --- 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 posthog.mcp import get_mcp_session + + async def app(scope, receive, send): + body = b"" + while True: + message = await receive() + body += message.get("body", b"") + if not message.get("more_body"): + break + request = json.loads(body) + session = get_mcp_session(scope) + if "id" not in request: # a notification: accepted, nothing to answer + await send({"type": "http.response.start", "status": 202, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + return + if request["method"] == "initialize": + client_info = request["params"]["clientInfo"] + posthog_mcp.capture_initialize( + client_name=client_info["name"], + client_version=client_info["version"], + session_id=session.session_id if session else None, + ) + result = { + "protocolVersion": request["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "dispatcher", "version": "1.0.0"}, + } + else: + result = {"tools": []} + payload = json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send({"type": "http.response.body", "body": payload.encode()}) + + return app + + +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 == [] 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..c4e31d94e 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -368,6 +368,7 @@ alias posthog.mcp.asgi.encode_session_id -> posthog.mcp.session_token.encode_ses 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 +alias posthog.mcp.asgi.resolve_session_id -> posthog.mcp.session.resolve_session_id alias posthog.mcp.asgi.warn -> posthog.mcp.logger.warn alias posthog.mcp.decode_session_id -> posthog.mcp.session_token.decode_session_id alias posthog.mcp.derive_session_id_from_conversation -> posthog.mcp.session.derive_session_id_from_conversation @@ -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) From 8a66fd6f695a7ca82f8e01e6abfec78982f1d8f4 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 16:09:36 +0300 Subject: [PATCH 02/13] refactor(mcp): read only the initialize response head on mcp 1.x The mcp 1.x hook now keeps only the status and the Mcp-Session-Id header of the initialize response. It no longer parses the JSON or SSE body. The event uses the protocol version that the client asked for. record_initialize now reads the client name, version, and protocol from the initialize params and reply, so both hooks pass them through. The new tests share their request helpers and drop the SSE body case. Co-Authored-By: Claude Opus 5.5 --- .../mcp-initialize-per-handshake.md | 2 +- posthog/mcp/_instrument_v2.py | 22 +- posthog/mcp/_instrumentation.py | 26 +- posthog/mcp/asgi.py | 97 ++----- posthog/test/mcp/test_initialize_handshake.py | 237 +++++------------- 5 files changed, 104 insertions(+), 280 deletions(-) diff --git a/.sampo/changesets/mcp-initialize-per-handshake.md b/.sampo/changesets/mcp-initialize-per-handshake.md index 16afe513f..33d926d09 100644 --- a/.sampo/changesets/mcp-initialize-per-handshake.md +++ b/.sampo/changesets/mcp-initialize-per-handshake.md @@ -8,4 +8,4 @@ Before, the SDK sent the event on the first request of each session, once per pr 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 and the server reply as `$mcp_parameters` and `$mcp_response`. +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/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index e5a8b7a59..36141f301 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -192,24 +192,12 @@ async def _record_v2_initialize( data: MCPAnalyticsData, ctx: Any, result: Any, started: datetime ) -> None: try: - params = dict(ctx.params) if isinstance(ctx.params, Mapping) else {} - client_info = params.get("clientInfo") - client_info = client_info if isinstance(client_info, Mapping) else {} - response = _to_jsonable(result) - negotiated = ( - response.get("protocolVersion") if isinstance(response, dict) else None - ) mcp_session_id, token = _initialize_session(ctx) - session_id = await resolve_session_id(data, mcp_session_id, token=token) await record_initialize( data, - session_id=session_id, - client_name=_str_or_none(client_info.get("name")), - client_version=_str_or_none(client_info.get("version")), - protocol_version=_str_or_none(negotiated) - or _str_or_none(params.get("protocolVersion")), - request={"method": "initialize", "params": params}, - response=response, + session_id=await resolve_session_id(data, mcp_session_id, token=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, ) @@ -217,10 +205,6 @@ async def _record_v2_initialize( log(f"Error: failed to capture $mcp_initialize - {error}") -def _str_or_none(value: Any) -> Optional[str]: - return value if isinstance(value, str) and value else None - - # --- registry plumbing --------------------------------------------------------- diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 0f54d3581..69a5c7f0d 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -244,11 +244,8 @@ async def record_initialize( data: MCPAnalyticsData, *, session_id: str, - client_name: Optional[str], - client_version: Optional[str], - protocol_version: Optional[str], - request: Dict[str, Any], - response: Any, + params: Dict[str, Any], + result: Any, extra: Optional[Dict[str, Any]], timestamp: datetime, ) -> None: @@ -257,19 +254,28 @@ async def record_initialize( 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 once the handshake has succeeded. Identify runs first so the event is + this once the handshake 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.""" + + 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": client_name, - "client_version": client_version, - "protocol_version": protocol_version, + "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": response, + "response": result, "timestamp": timestamp, } await _apply_event_properties(data, event, request, extra) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index ccba8c1e2..945c3e994 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -279,61 +279,21 @@ def _initialize_params( class _InitializeResponse: - """A ``send`` that passes the ``initialize`` response through and keeps what - the event needs: the status, the ``Mcp-Session-Id`` the client will replay - (whichever of the transport or the mint set it), and the body, capped like the - request sniff.""" + """A ``send`` that passes the ``initialize`` response through and keeps its + 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) -> None: self._send = send self.status: Optional[int] = None self.session_header: Optional[str] = None - self.content_type = "" - self._chunks: list[bytes] = [] - self._size = 0 async def __call__(self, message: Dict[str, Any]) -> None: - try: - self._observe(message) - except Exception as error: # noqa: BLE001 - log(f"PostHog MCP session middleware: response read failed - {error}") - await self._send(message) - - def _observe(self, message: Dict[str, Any]) -> None: if message.get("type") == "http.response.start": self.status = message.get("status") headers = _decode_headers(message.get("headers") or []) self.session_header = headers.get(MCP_SESSION_HEADER.lower()) or None - self.content_type = headers.get("content-type", "") - elif message.get("type") == "http.response.body": - chunk = message.get("body", b"") or b"" - if self._size + len(chunk) <= _MAX_SNIFF_BODY: - self._chunks.append(chunk) - self._size += len(chunk) - - def reply(self, request_id: Any) -> Optional[Dict[str, Any]]: - """The JSON-RPC reply to ``request_id``, from a JSON body or an SSE stream. - ``None`` when there is none to read: the deprecated SSE transport answers - the POST with 202 and delivers the reply on its event stream.""" - if self._size > _MAX_SNIFF_BODY: - return None - body = b"".join(self._chunks).decode("utf-8", errors="replace") - if self.content_type.startswith("text/event-stream"): - candidates = [ - line[len("data:") :] - for line in body.splitlines() - if line.startswith("data:") - ] - else: - candidates = [body] - for candidate in candidates: - try: - message = json.loads(candidate) - except Exception: # noqa: BLE001 - continue - if isinstance(message, dict) and message.get("id") == request_id: - return message - return None + await self._send(message) class _InitializeContext: @@ -347,20 +307,6 @@ def __init__(self, request: Any) -> None: self.request = request -def _http_request(scope: Any) -> Any: - try: - from starlette.requests import Request - - return Request(scope) - except Exception: # noqa: BLE001 - only headers are read; never fail the capture - return _Headers(_headers_dict(scope)) - - -class _Headers: - def __init__(self, headers: Dict[str, str]) -> None: - self.headers = headers - - async def _record_initialize( data: MCPAnalyticsData, message: Dict[str, Any], @@ -368,36 +314,25 @@ async def _record_initialize( scope: Any, started: datetime, ) -> None: - """Emit ``$mcp_initialize`` for a handshake that succeeded, in the session the - client's later requests resolve to: the ``Mcp-Session-Id`` they replay (a + """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.""" + response carried none. The body is not read, so the event carries the protocol + version the client asked for.""" try: if response.status is None or not 200 <= response.status < 300: return - reply = response.reply(message.get("id")) - if reply is not None and "error" in reply: - return - result = reply.get("result") if reply is not None else None - params, client_info = _initialize_params(message) + from starlette.requests import Request + header = response.session_header - session_id = await resolve_session_id( - data, header, token=decode_session_id(header) - ) - negotiated = result.get("protocolVersion") if isinstance(result, dict) else None await record_initialize( data, - session_id=session_id, - client_name=_str_or_none(client_info.get("name")), - client_version=_str_or_none(client_info.get("version")), - protocol_version=_str_or_none(negotiated) - or _str_or_none(params.get("protocolVersion")), - request={"method": "initialize", "params": params}, - response=result, - extra={ - "session_id": header, - "ctx": _InitializeContext(_http_request(scope)), - }, + session_id=await resolve_session_id( + data, header, token=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 diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index 41ca43691..0ea88a91a 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -58,7 +58,6 @@ def make_pod( *, kind: str = "official", stateless: bool = True, - json_response: bool = True, options: Optional[MCPAnalyticsOptions] = None, bind: bool = True, ): @@ -68,9 +67,9 @@ def make_pod( 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, json_response) + return _fastmcp_pod(client, options, stateless) if kind == "lowlevel": - return _lowlevel_pod(client, options, stateless, json_response, bind) + return _lowlevel_pod(client, options, stateless, bind) if MCP_MAJOR < 2: from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings @@ -78,7 +77,7 @@ def make_pod( server = FastMCP( "handshake", stateless_http=stateless, - json_response=json_response, + json_response=True, transport_security=TransportSecuritySettings( enable_dns_rebinding_protection=False ), @@ -96,13 +95,11 @@ def add(a: int, b: int) -> int: if MCP_MAJOR < 2: app = server.streamable_http_app() else: - app = server.streamable_http_app( - json_response=json_response, stateless_http=stateless - ) + app = server.streamable_http_app(json_response=True, stateless_http=stateless) return app, server.session_manager.run() -def _fastmcp_pod(client, options, stateless, json_response): +def _fastmcp_pod(client, options, stateless): fastmcp = pytest.importorskip("fastmcp") server = fastmcp.FastMCP("handshake") @@ -112,11 +109,11 @@ def add(a: int, b: int) -> int: return a + b instrument(server, client, options) - app = server.http_app(json_response=json_response, stateless_http=stateless) + app = server.http_app(json_response=True, stateless_http=stateless) return app, app.router.lifespan_context(app) -def _lowlevel_pod(client, options, stateless, json_response, bind): +def _lowlevel_pod(client, options, stateless, bind): import mcp.types as mcp_types from mcp.server.lowlevel import Server @@ -146,9 +143,7 @@ async def on_call_tool(ctx, params): "handshake", on_list_tools=on_list_tools, on_call_tool=on_call_tool ) instrument(server, client, options) - app = server.streamable_http_app( - json_response=json_response, stateless_http=stateless - ) + app = server.streamable_http_app(json_response=True, stateless_http=stateless) return app, server.session_manager.run() from mcp.server.streamable_http_manager import StreamableHTTPSessionManager @@ -167,7 +162,7 @@ async def call_tool_handler(name, arguments): instrument(server, client, options) manager = StreamableHTTPSessionManager( - app=server, stateless=stateless, json_response=json_response + app=server, stateless=stateless, json_response=True ) app = ( PostHogMcpStatelessSessionMiddleware(manager.handle_request, server=server) @@ -220,14 +215,12 @@ def rpc(method: str, params: Dict[str, Any], request_id: Optional[int] = 1): return body -def result_of(response: httpx.Response) -> Dict[str, Any]: - """The JSON-RPC message in a response, whether it came back as JSON or as SSE.""" - if response.headers.get("content-type", "").startswith("text/event-stream"): - for line in response.text.splitlines(): - if line.startswith("data:"): - return json.loads(line[len("data:") :]) - raise AssertionError(f"no SSE data frame in {response.text!r}") - return response.json() +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( @@ -241,19 +234,9 @@ async def handshake( 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", - { - "protocolVersion": PROTOCOL_VERSION, - "capabilities": {}, - "clientInfo": {"name": client_name, "version": "0.50.0"}, - }, - ), - headers, - ) + response = await lb.post(rpc("initialize", initialize_params(client_name)), headers) assert response.status_code == 200, response.text - negotiated = result_of(response)["result"]["protocolVersion"] + negotiated = response.json()["result"]["protocolVersion"] replay = {**headers, "mcp-protocol-version": negotiated} session_header = response.headers.get("mcp-session-id") if session_header: @@ -263,13 +246,14 @@ async def handshake( return replay, negotiated -async def call_add(lb: RoundRobin, headers: Dict[str, str], request_id: int = 3): +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": {"a": 2, "b": 3, "context": "adding"}}, - request_id, - ), + rpc("tools/call", {"name": "add", "arguments": arguments, **meta}, request_id), headers, ) assert response.status_code == 200, response.text @@ -308,6 +292,10 @@ async def test_one_initialize_per_handshake_across_stateless_pods(kind): 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) @@ -359,36 +347,6 @@ async def test_initialize_session_matches_a_stateful_transport(kind): } -async def test_sse_response_mode(): - # json_response=False answers each POST as an SSE stream, so the negotiated - # version has to come out of an SSE frame. - client = FakeClient() - async with fleet(client, pods=2, json_response=False) as lb: - headers, negotiated = 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 - assert inits[0]["properties"]["$mcp_protocol_version"] == negotiated - - -async def test_initialize_carries_the_handshake_payload(): - client = FakeClient() - async with fleet(client) as lb: - await handshake(lb) - await _flush() - - (init,) = _events(client, "$mcp_initialize") - request = init["properties"]["$mcp_parameters"] - assert request["method"] == "initialize" - assert request["params"]["clientInfo"]["name"] == "codex-mcp-client" - response = init["properties"]["$mcp_response"] - assert response["serverInfo"]["name"] == "handshake" - assert init["properties"]["$mcp_duration_ms"] >= 0 - - async def test_initialize_is_attributed_on_the_pod_that_served_it(): def identify(request, extra): headers = get_request_headers(extra) or {} @@ -411,13 +369,17 @@ def identify(request, extra): 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 result_of(response) + assert "error" in response.json() assert _events(client, "$mcp_initialize") == [] @@ -432,23 +394,7 @@ async def test_an_echoed_conversation_does_not_initialize_again(): async with fleet(client, pods=2, options=options) as lb: headers, _ = await handshake(lb) for request_id in (3, 4, 5): - response = await lb.post( - rpc( - "tools/call", - { - "name": "add", - "arguments": { - "a": 1, - "b": 1, - "context": "again", - "conversation_id": conversation_id, - }, - }, - request_id, - ), - headers, - ) - assert response.status_code == 200, response.text + await call_add(lb, headers, request_id, conversation_id=conversation_id) await _flush() calls = _events(client, "$mcp_tool_call") @@ -465,21 +411,10 @@ 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): - response = await lb.post( - rpc( - "tools/call", - { - "name": "add", - "arguments": {"a": 1, "b": 2, "context": "modern"}, - "_meta": modern_meta(), - }, - request_id, - ), - {**BASE_HEADERS, **modern_headers("tools/call", "add")}, - ) - assert response.status_code == 200, response.text + await call_add(lb, headers, request_id, _meta=modern_meta()) await _flush() assert len(_events(client, "$mcp_tool_call")) == 3 @@ -557,21 +492,20 @@ async def send(self, method, params, request_id=None): line = await asyncio.wait_for(self.process.stdout.readline(), timeout=20) return json.loads(line) - def events(self, name): - if not self.events_file.exists(): - return [] + 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] - async def wait_for(self, name, count, timeout=5.0): - """Captures are fire-and-forget in the child, so give them a moment to land.""" - for _ in range(int(timeout / 0.05)): - if len(self.events(name)) >= count: - break - await asyncio.sleep(0.05) - await asyncio.sleep(0.2) # and anything that would make it one too many - return self.events(name) - @asynccontextmanager async def stdio_server(tmp_path): @@ -599,17 +533,9 @@ async def stdio_server(tmp_path): await process.wait() -def _initialize_params(): - return { - "protocolVersion": PROTOCOL_VERSION, - "capabilities": {}, - "clientInfo": {"name": "claude-desktop", "version": "1.2.0"}, - } - - async def test_stdio_handshake_then_calls(tmp_path): async with stdio_server(tmp_path) as server: - reply = await server.send("initialize", _initialize_params(), 1) + 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) @@ -620,8 +546,8 @@ async def test_stdio_handshake_then_calls(tmp_path): request_id, ) assert reply["result"]["isError"] is False, reply - calls = await server.wait_for("$mcp_tool_call", 2) - inits = await server.wait_for("$mcp_initialize", 1) + calls = await server.events("$mcp_tool_call", 2) + inits = await server.events("$mcp_initialize", 1) assert len(calls) == 2 assert len(inits) == 1 @@ -633,9 +559,9 @@ async def test_stdio_handshake_then_calls(tmp_path): @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(), 1) + reply = await server.send("initialize", initialize_params("claude-desktop"), 1) assert "result" in reply, reply - inits = await server.wait_for("$mcp_initialize", 1) + inits = await server.events("$mcp_initialize", 1) assert len(inits) == 1 assert ( @@ -644,64 +570,37 @@ async def test_stdio_bare_handshake_is_counted(tmp_path): ) -@pytest.mark.skipif(MCP_MAJOR >= 2, reason="mcp 2.x observes the stdio handshake") -async def test_stdio_bare_handshake_is_missed_on_mcp_1(tmp_path): - # The documented limit of the first-request fallback: a stdio client that - # handshakes and never sends a request reports nothing on mcp 1.x. - async with stdio_server(tmp_path) as server: - reply = await server.send("initialize", _initialize_params(), 1) - assert "result" in reply, reply - inits = await server.wait_for("$mcp_initialize", 1, timeout=1.0) - - assert inits == [] - - # --- 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 app(scope, receive, send): - body = b"" - while True: - message = await receive() - body += message.get("body", b"") - if not message.get("more_body"): - break - request = json.loads(body) - session = get_mcp_session(scope) - if "id" not in request: # a notification: accepted, nothing to answer - await send({"type": "http.response.start", "status": 202, "headers": []}) - await send({"type": "http.response.body", "body": b""}) - return - if request["method"] == "initialize": - client_info = request["params"]["clientInfo"] + 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=client_info["name"], - client_version=client_info["version"], + client_name=message["params"]["clientInfo"]["name"], session_id=session.session_id if session else None, ) result = { - "protocolVersion": request["params"]["protocolVersion"], - "capabilities": {"tools": {}}, + "protocolVersion": message["params"]["protocolVersion"], + "capabilities": {}, "serverInfo": {"name": "dispatcher", "version": "1.0.0"}, } - else: - result = {"tools": []} - payload = json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}) - await send( - { - "type": "http.response.start", - "status": 200, - "headers": [(b"content-type", b"application/json")], - } - ) - await send({"type": "http.response.body", "body": payload.encode()}) + return JSONResponse({"jsonrpc": "2.0", "id": message["id"], "result": result}) - return app + return Starlette(routes=[Route("/mcp", endpoint, methods=["POST"])]) async def test_a_posthogmcp_dispatcher_reports_only_its_own_initialize(): From 1b0fe8f9555393f8569d4afa8a24626658a7de45 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 17:26:41 +0300 Subject: [PATCH 03/13] fix(mcp): address review on the initialize hooks - SSE transport: the initialize event now lands in the session that the SSE calls use. On SDK 1.x the middleware leaves SSE requests to the first-request fallback. On SDK 2.x the hook uses the process session. - The bound middleware marks each request that it passes on. The fallback now stops only for marked requests, not for every HTTP request. - The hooks compute the handshake session without a write to the shared session that McpAnalytics.capture() reads. - A server= value that cannot be bound no longer breaks requests. - Remove the unused _mint_token_if_initialize. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/README.md | 1 + posthog/mcp/_instrument_v2.py | 9 +- posthog/mcp/_instrumentation.py | 18 +- posthog/mcp/_internal.py | 10 +- posthog/mcp/asgi.py | 77 ++++---- posthog/mcp/session.py | 19 ++ posthog/test/mcp/test_initialize_handshake.py | 168 +++++++++++++++++- references/public_api_snapshot.txt | 2 +- 8 files changed, 252 insertions(+), 52 deletions(-) diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 7b682d81b..6ca3e7324 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -82,6 +82,7 @@ The MCP SDK answers `initialize` itself, so the SDK observes it from outside: `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. - **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. diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 36141f301..439bb97c3 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -49,6 +49,7 @@ advertised_tool_names, apply_virtual_tool_injection, collect_listed_tools, + _is_sse_request, _to_jsonable, copy_tools_list_result, fire_and_forget, @@ -77,7 +78,7 @@ from .asgi import get_mcp_session from .logger import log from .request_headers import get_request_headers -from .session import resolve_session_id +from .session import handshake_session_id from .session_token import SessionTokenPayload, read_mcp_session_header from .tools import ( get_more_tools_result_text, @@ -180,6 +181,10 @@ def _initialize_session( exposes no public accessor. A stateless one has none, and the session lives in 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: @@ -195,7 +200,7 @@ async def _record_v2_initialize( mcp_session_id, token = _initialize_session(ctx) await record_initialize( data, - session_id=await resolve_session_id(data, mcp_session_id, token=token), + 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}, diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 69a5c7f0d..0429d3504 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -283,21 +283,29 @@ def text(value: Any) -> Optional[str]: 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 mcp 1.x over HTTP - without a bound ASGI middleware. 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.""" + 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 - return not (data.handshake_observed_at_asgi and get_request(extra) is not None) + 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( diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 09599d249..44a1eeb48 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -107,13 +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) - # Where this server sees the client's real `initialize`. mcp 2.x hands it to - # `Server.middleware` on every transport; on mcp 1.x only an ASGI middleware - # bound to this server sees it, and only over HTTP. Whatever neither covers - # (mcp 1.x stdio, or HTTP without a bound middleware) falls back to the - # first request of the session. + # 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 - handshake_observed_at_asgi: 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. diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 945c3e994..c50937a83 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -46,11 +46,16 @@ import weakref from datetime import datetime, timezone from typing import Any, Dict, Optional +from urllib.parse import parse_qs -from ._instrumentation import fire_and_forget, record_initialize +from ._instrumentation import ( + HANDSHAKE_OBSERVER_KEY, + fire_and_forget, + record_initialize, +) from ._internal import MCPAnalyticsData, get_server_tracking_data, tracking_key from .logger import log, warn -from .session import new_session_id, resolve_session_id +from .session import handshake_session_id, new_session_id from .session_token import ( MCP_SESSION_HEADER, SessionTokenPayload, @@ -62,9 +67,6 @@ # Scope key the middleware stashes the decoded token payload under. _SCOPE_KEY = "posthog_mcp_session" -# Scope key set once a middleware has taken an ``initialize`` to capture. -_HANDSHAKE_KEY = "posthog_mcp_initialize" - # JSON-RPC request bodies are tiny; cap what we buffer so a stray large POST on # the same app can't be read into memory in full before minting. _MAX_SNIFF_BODY = 256 * 1024 @@ -97,25 +99,36 @@ class PostHogMcpStatelessSessionMiddleware: def __init__(self, app: Any, server: Any = None) -> None: self.app = app - self._server = weakref.ref(tracking_key(server)) if server is not None else None - self._analytics() + 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 - observes its handshakes: on mcp 2.x ``Server.middleware`` already does.""" - server = self._server() if self._server is not None else None - data = get_server_tracking_data(server) if server is not None else None - if data is None or data.handshake_observed_in_process: + 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 - data.handshake_observed_at_asgi = True - return data + 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 - # Bound late if instrument() ran after this middleware was built. - self._analytics() + # 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)) @@ -150,12 +163,8 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: await self.app(scope, receive, send) return - # Two stacked middlewares (autowired plus a manual one) see the same - # handshake; the outer one captures it. - data = None if scope.get(_HANDSHAKE_KEY) else self._analytics() recorder = None - if data is not None: - scope[_HANDSHAKE_KEY] = True + if observer is not None: recorder = _InitializeResponse(send) send = recorder if token is not None: @@ -164,12 +173,22 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: started = datetime.now(timezone.utc) await self.app(scope, receive, send) - if data is not None and recorder is not None: + if observer is not None and recorder is not None: fire_and_forget( - _record_initialize(data, initialize, recorder, scope, started), data + _record_initialize(observer, initialize, recorder, scope, started), + observer, ) +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 []) @@ -245,15 +264,9 @@ def _parse_initialize(body: bytes) -> Optional[Dict[str, Any]]: return message -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).""" - message = _parse_initialize(body) - return _mint_token(message) if message is not None else None - - 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( @@ -327,8 +340,8 @@ async def _record_initialize( header = response.session_header await record_initialize( data, - session_id=await resolve_session_id( - data, header, token=decode_session_id(header) + session_id=await handshake_session_id( + data, header, decode_session_id(header) ), params=_initialize_params(message)[0], result=None, 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 index 0ea88a91a..b0bf5100a 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -61,7 +61,7 @@ def make_pod( options: Optional[MCPAnalyticsOptions] = None, bind: bool = True, ): - """One instrumented server of ``kind``, as ``(asgi_app, lifespan)``. + """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=``.""" @@ -96,7 +96,7 @@ def add(a: int, b: int) -> int: app = server.streamable_http_app() else: app = server.streamable_http_app(json_response=True, stateless_http=stateless) - return app, server.session_manager.run() + return app, server.session_manager.run(), server def _fastmcp_pod(client, options, stateless): @@ -110,7 +110,7 @@ def add(a: int, b: int) -> int: instrument(server, client, options) app = server.http_app(json_response=True, stateless_http=stateless) - return app, app.router.lifespan_context(app) + return app, app.router.lifespan_context(app), server def _lowlevel_pod(client, options, stateless, bind): @@ -144,7 +144,7 @@ async def on_call_tool(ctx, params): ) instrument(server, client, options) app = server.streamable_http_app(json_response=True, stateless_http=stateless) - return app, server.session_manager.run() + return app, server.session_manager.run(), server from mcp.server.streamable_http_manager import StreamableHTTPSessionManager @@ -169,7 +169,7 @@ async def call_tool_handler(name, arguments): if bind else PostHogMcpStatelessSessionMiddleware(manager.handle_request) ) - return app, manager.run() + return app, manager.run(), server @asynccontextmanager @@ -180,7 +180,7 @@ async def fleet(client: FakeClient, pods: int = 1, **pod_options: Any): async with AsyncExitStack() as stack: surfaces: List[httpx.AsyncClient] = [] for _ in range(pods): - app, lifespan = make_pod(client, **pod_options) + app, lifespan, _ = make_pod(client, **pod_options) await stack.enter_async_context(lifespan) surfaces.append( await stack.enter_async_context( @@ -635,3 +635,159 @@ def keep(event): 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 diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index c4e31d94e..518dd936b 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -365,10 +365,10 @@ 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 -alias posthog.mcp.asgi.resolve_session_id -> posthog.mcp.session.resolve_session_id alias posthog.mcp.asgi.warn -> posthog.mcp.logger.warn alias posthog.mcp.decode_session_id -> posthog.mcp.session_token.decode_session_id alias posthog.mcp.derive_session_id_from_conversation -> posthog.mcp.session.derive_session_id_from_conversation From da925d50cdd2eda37e1c767dbe92c6feda8be75b Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 18:29:55 +0300 Subject: [PATCH 04/13] fix(mcp): wait for the initialize capture inside its handshake task Both handshake hooks run record_initialize in a tracked background task. It started the capture as a second task, after flush() took its snapshot of pending tasks. With a slow before_send, flush() returned early and shutdown dropped $mcp_initialize. The handshake task now awaits the capture itself. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/_instrumentation.py | 6 +++- posthog/test/mcp/test_initialize_handshake.py | 32 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 0429d3504..776aea74c 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -280,7 +280,11 @@ def text(value: Any) -> Optional[str]: } await _apply_event_properties(data, event, request, extra) stamp_transport_identity(event, extra) - fire_and_forget(capture_event(data, event), data) + # Both seams already run this in a tracked background task. Await the capture + # inside it: a second task would start after flush() took its snapshot. + capture = capture_event(data, event) + if capture is not None: + await capture # ASGI scope key a bound ``PostHogMcpStatelessSessionMiddleware`` sets to the diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index b0bf5100a..9a9b18289 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -791,3 +791,35 @@ async def echo(scope, receive, send): headers=BASE_HEADERS, ) assert response.status_code == 200 + + +@pytest.mark.parametrize("kind", KINDS) +async def test_flush_waits_for_the_handshake_capture(kind): + # The hooks record the handshake in a background task. flush() must also wait + # for the capture that task starts, or shutdown drops $mcp_initialize when + # 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) # holds the recorder open past flush()'s snapshot + 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 From becff6a564f115a61a745e542d1b17dab165d4ed Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 18:31:03 +0300 Subject: [PATCH 05/13] fix(mcp): resolve the v2 handshake session from an incoming header On MCP SDK 2.x, a stateless client can send initialize with an existing Mcp-Session-Id and keep sending it. Its later calls resolve their session from that header, but the handshake hook read only the connection id and the minted token. The event landed in a generated session. The hook now falls back to the incoming header. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/_instrument_v2.py | 12 ++++++++-- posthog/test/mcp/test_initialize_handshake.py | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 439bb97c3..090130390 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -79,7 +79,11 @@ from .logger import log from .request_headers import get_request_headers from .session import handshake_session_id -from .session_token import SessionTokenPayload, read_mcp_session_header +from .session_token import ( + SessionTokenPayload, + decode_session_id, + read_mcp_session_header, +) from .tools import ( get_more_tools_result_text, ) @@ -178,7 +182,8 @@ def _initialize_session( 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, and the session lives in + exposes no public accessor. A stateless one has none: the session comes from + the ``Mcp-Session-Id`` the client sent, which it keeps sending, 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}): @@ -189,6 +194,9 @@ def _initialize_session( transport_session_id = getattr(connection, "session_id", None) if isinstance(transport_session_id, str) and transport_session_id: return transport_session_id, None + incoming = _ctx_mcp_session_id(ctx) + if incoming: + return incoming, decode_session_id(incoming) request = getattr(ctx, "request", None) return None, get_mcp_session(request) if request is not None else None diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index 9a9b18289..afe71b993 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -823,3 +823,25 @@ async def slow_before_send(event): 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: see the next fix") +@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"] + } From 0a7a6a45b9fe02b4f7f58c7aa9e7c0f85b7d6e37 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 18:32:26 +0300 Subject: [PATCH 06/13] fix(mcp): capture a v1 initialize that carries a session header On MCP SDK 1.x, the middleware read the body only of a POST with no Mcp-Session-Id, because only those need a token. A client that sent initialize with an existing session id got no event: the middleware marked the request, so the first-request fallback did not run either. A middleware bound to a server now reads those bodies too, behind a cheap check for "initialize". It still mints only without a header. The event uses the header that the client keeps sending. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/asgi.py | 25 +++++++++++++------ posthog/test/mcp/test_initialize_handshake.py | 1 - 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index c50937a83..093fb24ab 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -144,7 +144,11 @@ 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. - if scope.get("method") != "POST" or incoming_header is not None: + # A client can also send initialize with a session id it already holds, so + # a middleware that captures handshakes reads those bodies too. + if scope.get("method") != "POST" or ( + incoming_header is not None and observer is None + ): await self.app(scope, receive, send) return @@ -155,7 +159,8 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: try: body, receive = await _buffer_body(receive) initialize = _parse_initialize(body) - token = _mint_token(initialize) if initialize is not None else None + if initialize is not None and incoming_header is None: + token = _mint_token(initialize) except Exception as error: # noqa: BLE001 log(f"PostHog MCP session middleware: mint failed - {error}") @@ -175,7 +180,9 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: await self.app(scope, receive, send) if observer is not None and recorder is not None: fire_and_forget( - _record_initialize(observer, initialize, recorder, scope, started), + _record_initialize( + observer, initialize, recorder, incoming_header, scope, started + ), observer, ) @@ -251,7 +258,8 @@ async def replay() -> dict[str, Any]: 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: + # Cheap first: with a bound middleware this runs on every tool call too. + if not body or b'"initialize"' not in body: return None try: # RecursionError (deeply-nested JSON) is a RuntimeError, not a @@ -324,20 +332,21 @@ async def _record_initialize( data: MCPAnalyticsData, message: Dict[str, Any], response: _InitializeResponse, + incoming_header: Optional[str], scope: Any, started: datetime, ) -> 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.""" + stateful transport's id, our token, or the one the client already sent), or + this process's own session when there is none. The body is not read, so the + event carries the protocol version the client asked for.""" try: if response.status is None or not 200 <= response.status < 300: return from starlette.requests import Request - header = response.session_header + header = response.session_header or incoming_header await record_initialize( data, session_id=await handshake_session_id( diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index afe71b993..b6920ec30 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -825,7 +825,6 @@ async def slow_before_send(event): assert len(_events(client, "$mcp_initialize")) == 1 -@pytest.mark.skipif(MCP_MAJOR < 2, reason="mcp 1.x: see the next fix") @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 From ef0bb56d9e362463759a9d9070226fea5827431d Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 18:49:52 +0300 Subject: [PATCH 07/13] fix(mcp): pass a disconnect mid-body through the session middleware When the client disconnected while it sent the body, the middleware stopped its read and replayed the partial body as a complete request. The app parsed half a request and never saw the disconnect. The middleware now replays the partial body with more_body set, then hands the app the disconnect. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/asgi.py | 14 ++++++++---- posthog/test/mcp/test_session_token.py | 30 ++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 093fb24ab..e2a5723cb 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -224,10 +224,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]) @@ -243,11 +246,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. 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 ------------------- From 11e8b6e6f518d95aa4d892623f3b01e1901bd76b Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 18:50:59 +0300 Subject: [PATCH 08/13] fix(mcp): stop buffering tool calls once they are not an initialize A bound middleware reads the body of a POST that carries a session id, to catch an initialize that reuses one. That held every tool call in memory, up to 256 KB, before the app saw it. The middleware now stops after 4 KB when "initialize" is not in it, and streams the rest. JSON-RPC clients put "method" before "params", so a handshake shows up in that window. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/asgi.py | 20 +++++++++- posthog/test/mcp/test_session_token.py | 51 ++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index e2a5723cb..1dfce5731 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -71,6 +71,11 @@ # the same app can't be read into memory in full before minting. _MAX_SNIFF_BODY = 256 * 1024 +# A request that already carries a session id is almost always a tool call, which +# can be large. JSON-RPC clients put "method" before "params", so once this much +# of the body arrived with no "initialize" in it, stop buffering and stream it. +_PEEK_BODY = 4 * 1024 + def get_mcp_session(request_or_scope: Any) -> Optional[SessionTokenPayload]: """Return the ``SessionTokenPayload`` the middleware recovered for this request, @@ -157,7 +162,9 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: token = None initialize = None try: - body, receive = await _buffer_body(receive) + body, receive = await _buffer_body( + receive, peek=incoming_header is not None + ) initialize = _parse_initialize(body) if initialize is not None and incoming_header is None: token = _mint_token(initialize) @@ -211,10 +218,13 @@ def _decode_headers(raw: Any) -> dict[str, str]: return result -async def _buffer_body(receive: Any) -> tuple[bytes, Any]: +async def _buffer_body(receive: Any, *, peek: bool = False) -> tuple[bytes, Any]: """Read up to ``_MAX_SNIFF_BODY`` of the request body for sniffing, then return that prefix plus a ``receive`` that replays it and streams the rest. + With ``peek``, also stop once ``_PEEK_BODY`` bytes arrived and none of them is + ``"initialize"``: the request is not a handshake, so it is not held in memory. + Memory is bounded to the cap: an `initialize` handshake is tiny, so once we exceed the cap the request cannot be one and we stop buffering — the prefix is replayed with ``more_body: True`` and any remaining chunks are forwarded @@ -224,6 +234,7 @@ async def _buffer_body(receive: Any) -> tuple[bytes, Any]: total = 0 complete = False overflow = False + peeked = not peek # 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 @@ -240,6 +251,11 @@ async def _buffer_body(receive: Any) -> tuple[bytes, Any]: if total > _MAX_SNIFF_BODY: overflow = True # too big to be initialize — stop buffering break + if not peeked and total >= _PEEK_BODY: + peeked = True + if b'"initialize"' not in b"".join(chunks): + overflow = True # not a handshake — stop buffering + break buffered = b"".join(chunks) replayed = False diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index 9db5aca7e..3c0096015 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -888,3 +888,54 @@ def test_app_built_probe_fires_on_the_installed_sdk_major(): f"probe blind to a built app on MCP {MCP_MAJOR}.x -- the ordering warning " "cannot fire for these servers" ) + + +async def test_peek_streams_a_large_request_that_is_not_initialize(): + """With a session header, a tool call is read only until the peek window shows + it is not a handshake. The rest streams to the app, byte for byte.""" + from posthog.mcp import asgi as asgi_mod + + body = b'{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"blob":"' + body += b"x" * (16 * asgi_mod._PEEK_BODY) + b'"}}' + parts = [body[i : i + 1024] for i in range(0, len(body), 1024)] + reads = [] + + async def receive(): + chunk = parts.pop(0) + reads.append(chunk) + return {"type": "http.request", "body": chunk, "more_body": bool(parts)} + + sniff, replay = await asgi_mod._buffer_body(receive, peek=True) + + assert sniff == b"" + assert sum(map(len, reads)) == asgi_mod._PEEK_BODY # stopped at the window + replayed = b"" + while True: + message = await replay() + replayed += message["body"] + if not message["more_body"]: + break + assert replayed == body + + +async def test_peek_still_reads_an_initialize_split_into_small_chunks(): + from posthog.mcp import asgi as asgi_mod + + body = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}}, + } + ).encode() + parts = [body[i : i + 7] for i in range(0, len(body), 7)] + + async def receive(): + chunk = parts.pop(0) + return {"type": "http.request", "body": chunk, "more_body": bool(parts)} + + sniff, _ = await asgi_mod._buffer_body(receive, peek=True) + + assert sniff == body + assert asgi_mod._parse_initialize(sniff)["method"] == "initialize" From bb21c6a7075596dd23dbd17283674101f2974df0 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 22 Oct 2026 18:51:50 +0300 Subject: [PATCH 09/13] refactor(mcp): read the v2 handshake header through _resolve_ctx The v2 handshake hook read the Mcp-Session-Id header and decoded the token itself. Every later request gets the same pair from _resolve_ctx. The hook now calls it, so the handshake and later requests follow one rule and cannot drift apart into two sessions. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/_instrument_v2.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 090130390..7ce465167 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -79,11 +79,7 @@ from .logger import log from .request_headers import get_request_headers from .session import handshake_session_id -from .session_token import ( - SessionTokenPayload, - decode_session_id, - read_mcp_session_header, -) +from .session_token import SessionTokenPayload, read_mcp_session_header from .tools import ( get_more_tools_result_text, ) @@ -183,9 +179,10 @@ def _initialize_session( 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, 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.""" + 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. @@ -194,9 +191,9 @@ def _initialize_session( transport_session_id = getattr(connection, "session_id", None) if isinstance(transport_session_id, str) and transport_session_id: return transport_session_id, None - incoming = _ctx_mcp_session_id(ctx) + token, _, _, _, incoming = _resolve_ctx(ctx) if incoming: - return incoming, decode_session_id(incoming) + return incoming, token request = getattr(ctx, "request", None) return None, get_mcp_session(request) if request is not None else None From 20c1243f95a333dab99c0b65ffaa84a993ea2a08 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 23 Oct 2026 11:30:09 +0300 Subject: [PATCH 10/13] fix(mcp): run the initialize callbacks inline in the handshake request Both handshake hooks ran identify and event_properties in a detached background task. If a callback never returned, each handshake left one more task pending, with no bound, and flush() waited for all of them. The spec requires pending state to be bounded by count. The hooks now run the callbacks inline, the same as a tool call does. A hung callback holds that request, where the host sees it, and leaves no task behind. The capture still runs in the background after the callbacks, as for every other event. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/_instrument_v2.py | 11 ++-- posthog/mcp/_instrumentation.py | 15 ++--- posthog/mcp/asgi.py | 15 ++--- posthog/test/mcp/test_initialize_handshake.py | 55 +++++++++++++++++-- 4 files changed, 68 insertions(+), 28 deletions(-) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 7ce465167..58d45f778 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -52,7 +52,6 @@ _is_sse_request, _to_jsonable, copy_tools_list_result, - fire_and_forget, is_first_listing_page, mutate_tool_schema, params_to_request_dict, @@ -145,9 +144,11 @@ def _observe_initialize(server: Any, data: MCPAnalyticsData) -> None: 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 event is scheduled rather than awaited: - the runner parks the connection's read loop until the ``initialize`` chain - returns, so a slow ``identify`` callback must not hold up the handshake.""" + 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( @@ -162,7 +163,7 @@ async def observe_initialize(ctx: Any, call_next: Any) -> Any: return await call_next(ctx) started = datetime.now(timezone.utc) result = await call_next(ctx) - fire_and_forget(_record_v2_initialize(data, ctx, result, started), data) + await _record_v2_initialize(data, ctx, result, started) return result setattr(observe_initialize, _WRAPPED_FLAG, True) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 776aea74c..4e72e1de1 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -254,10 +254,11 @@ async def record_initialize( 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 once the handshake 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.""" + 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 @@ -280,11 +281,7 @@ def text(value: Any) -> Optional[str]: } await _apply_event_properties(data, event, request, extra) stamp_transport_identity(event, extra) - # Both seams already run this in a tracked background task. Await the capture - # inside it: a second task would start after flush() took its snapshot. - capture = capture_event(data, event) - if capture is not None: - await capture + fire_and_forget(capture_event(data, event), data) # ASGI scope key a bound ``PostHogMcpStatelessSessionMiddleware`` sets to the diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 1dfce5731..2b08becf1 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -48,11 +48,7 @@ from typing import Any, Dict, Optional from urllib.parse import parse_qs -from ._instrumentation import ( - HANDSHAKE_OBSERVER_KEY, - fire_and_forget, - record_initialize, -) +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 handshake_session_id, new_session_id @@ -186,11 +182,10 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: started = datetime.now(timezone.utc) await self.app(scope, receive, send) if observer is not None and recorder is not None: - fire_and_forget( - _record_initialize( - observer, initialize, recorder, incoming_header, scope, started - ), - observer, + # Inline, as on a tool call: the response is already out, and a hung + # identify holds this request rather than a task nobody tracks. + await _record_initialize( + observer, initialize, recorder, incoming_header, scope, started ) diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index b6920ec30..d18fe9537 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -795,14 +795,13 @@ async def echo(scope, receive, send): @pytest.mark.parametrize("kind", KINDS) async def test_flush_waits_for_the_handshake_capture(kind): - # The hooks record the handshake in a background task. flush() must also wait - # for the capture that task starts, or shutdown drops $mcp_initialize when - # before_send is slow. + # 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) # holds the recorder open past flush()'s snapshot + await asyncio.sleep(0.05) return UserIdentity(distinct_id="user_slow") async def slow_before_send(event): @@ -844,3 +843,51 @@ async def test_an_initialize_with_a_session_header_keeps_that_session(kind): 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 From 132eaef4cd9d9c2b3f02d3114a4c124d43eb6d65 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 23 Oct 2026 11:41:48 +0300 Subject: [PATCH 11/13] fix(mcp): record the v1 handshake inside the request's auth context The session middleware sits outside FastMCP's AuthContextMiddleware. It ran identify after the app returned, when that middleware had already reset the authenticated user. An identify that calls get_access_token() got None, so $mcp_initialize stayed anonymous under OAuth. The middleware now records the handshake when the response starts, while the app still handles the request. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/asgi.py | 50 +++++++++++-------- posthog/test/mcp/test_initialize_handshake.py | 49 ++++++++++++++++++ 2 files changed, 77 insertions(+), 22 deletions(-) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 2b08becf1..014e5bdd0 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -171,22 +171,24 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: await self.app(scope, receive, send) return - recorder = None + started = datetime.now(timezone.utc) if observer is not None: - recorder = _InitializeResponse(send) - send = recorder + send = _InitializeResponse( + send, + functools.partial( + _record_initialize, + observer, + initialize, + incoming_header, + scope, + started, + ), + ) if token is not None: scope[_SCOPE_KEY] = decode_session_id(token) send = _sending_session_header(send, token) - started = datetime.now(timezone.utc) await self.app(scope, receive, send) - if observer is not None and recorder is not None: - # Inline, as on a tool call: the response is already out, and a hung - # identify holds this request rather than a task nobody tracks. - await _record_initialize( - observer, initialize, recorder, incoming_header, scope, started - ) def _is_sse(scope: Any) -> bool: @@ -317,20 +319,23 @@ def _initialize_params( class _InitializeResponse: - """A ``send`` that passes the ``initialize`` response through and keeps its - 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.""" + """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) -> None: + def __init__(self, send: Any, record: Any) -> None: self._send = send - self.status: Optional[int] = None - self.session_header: Optional[str] = None + self._record = record async def __call__(self, message: Dict[str, Any]) -> None: - if message.get("type") == "http.response.start": - self.status = message.get("status") + 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 []) - self.session_header = headers.get(MCP_SESSION_HEADER.lower()) or None + await record( + message.get("status"), headers.get(MCP_SESSION_HEADER.lower()) or None + ) await self._send(message) @@ -348,10 +353,11 @@ def __init__(self, request: Any) -> None: async def _record_initialize( data: MCPAnalyticsData, message: Dict[str, Any], - response: _InitializeResponse, incoming_header: Optional[str], 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 @@ -359,11 +365,11 @@ async def _record_initialize( this process's own session when there is none. The body is not read, so the event carries the protocol version the client asked for.""" try: - if response.status is None or not 200 <= response.status < 300: + if status is None or not 200 <= status < 300: return from starlette.requests import Request - header = response.session_header or incoming_header + header = session_header or incoming_header await record_initialize( data, session_id=await handshake_session_id( diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index d18fe9537..a917fda86 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -891,3 +891,52 @@ def pending(): 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" From 23a0677bdd0fe261f6a1d756733acca29959c687 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 23 Oct 2026 11:42:40 +0300 Subject: [PATCH 12/13] fix(mcp): stop the body peek only on a method that is not initialize The 4 KB peek stopped when "initialize" was not in the first 4 KB. JSON key order is arbitrary, so an initialize with large params before its method got no event. The peek now stops only when the prefix shows a "method" and no "initialize". Without a method yet, it reads up to the full cap. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/asgi.py | 13 ++++++++----- posthog/test/mcp/test_session_token.py | 20 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 014e5bdd0..b1369b6d9 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -68,8 +68,9 @@ _MAX_SNIFF_BODY = 256 * 1024 # A request that already carries a session id is almost always a tool call, which -# can be large. JSON-RPC clients put "method" before "params", so once this much -# of the body arrived with no "initialize" in it, stop buffering and stream it. +# can be large. Once this much of the body arrived and shows a "method" that is not +# "initialize", stop buffering and stream it. Key order is not assumed: without a +# "method" yet, keep reading up to the full cap. _PEEK_BODY = 4 * 1024 @@ -219,8 +220,9 @@ async def _buffer_body(receive: Any, *, peek: bool = False) -> tuple[bytes, Any] """Read up to ``_MAX_SNIFF_BODY`` of the request body for sniffing, then return that prefix plus a ``receive`` that replays it and streams the rest. - With ``peek``, also stop once ``_PEEK_BODY`` bytes arrived and none of them is - ``"initialize"``: the request is not a handshake, so it is not held in memory. + With ``peek``, also stop once ``_PEEK_BODY`` bytes arrived with a ``"method"`` + but no ``"initialize"`` in them: the request is not a handshake, so it is not + held in memory. Memory is bounded to the cap: an `initialize` handshake is tiny, so once we exceed the cap the request cannot be one and we stop buffering — the prefix is @@ -250,7 +252,8 @@ async def _buffer_body(receive: Any, *, peek: bool = False) -> tuple[bytes, Any] break if not peeked and total >= _PEEK_BODY: peeked = True - if b'"initialize"' not in b"".join(chunks): + prefix = b"".join(chunks) + if b'"method"' in prefix and b'"initialize"' not in prefix: overflow = True # not a handshake — stop buffering break diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index 3c0096015..e7d967b05 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -939,3 +939,23 @@ async def receive(): assert sniff == body assert asgi_mod._parse_initialize(sniff)["method"] == "initialize" + + +async def test_peek_reads_an_initialize_whose_params_come_first(): + """JSON key order is arbitrary. An initialize with large params before its + "method" must still be read whole and recognised.""" + from posthog.mcp import asgi as asgi_mod + + params = {"capabilities": {"x": "y" * (2 * asgi_mod._PEEK_BODY)}} + body = json.dumps( + {"params": params, "jsonrpc": "2.0", "id": 1, "method": "initialize"} + ).encode() + parts = [body[i : i + 1024] for i in range(0, len(body), 1024)] + + async def receive(): + chunk = parts.pop(0) + return {"type": "http.request", "body": chunk, "more_body": bool(parts)} + + sniff, _ = await asgi_mod._buffer_body(receive, peek=True) + + assert asgi_mod._parse_initialize(sniff)["method"] == "initialize" From 0773e3289780987fe933011046ff24b53c665d59 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 23 Oct 2026 14:57:32 +0300 Subject: [PATCH 13/13] fix(mcp): drop the body byte checks from the session middleware The middleware guessed from raw bytes whether a request was initialize. A nested "method" key in params stopped the 4 KB peek too early, and a JSON-escaped "initialize" failed the quick byte check. That second case also stopped the token mint, which worked before this PR. The middleware now parses the body as JSON, as main does, and reads only requests with no Mcp-Session-Id. On MCP SDK 1.x, an initialize that reuses an existing session id is not captured. The README lists this limit. MCP SDK 2.x still captures it. Co-Authored-By: Claude Opus 5.5 --- posthog/mcp/README.md | 2 + posthog/mcp/asgi.py | 50 +++---------- posthog/test/mcp/test_initialize_handshake.py | 48 +++++++++++++ posthog/test/mcp/test_session_token.py | 71 ------------------- 4 files changed, 61 insertions(+), 110 deletions(-) diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index 6ca3e7324..7f3133f2f 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -83,6 +83,8 @@ The MCP SDK answers `initialize` itself, so the SDK observes it from outside: 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. diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index b1369b6d9..a66b41113 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -67,12 +67,6 @@ # the same app can't be read into memory in full before minting. _MAX_SNIFF_BODY = 256 * 1024 -# A request that already carries a session id is almost always a tool call, which -# can be large. Once this much of the body arrived and shows a "method" that is not -# "initialize", stop buffering and stream it. Key order is not assumed: without a -# "method" yet, keep reading up to the full cap. -_PEEK_BODY = 4 * 1024 - def get_mcp_session(request_or_scope: Any) -> Optional[SessionTokenPayload]: """Return the ``SessionTokenPayload`` the middleware recovered for this request, @@ -146,11 +140,9 @@ 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. - # A client can also send initialize with a session id it already holds, so - # a middleware that captures handshakes reads those bodies too. - if scope.get("method") != "POST" or ( - incoming_header is not None and observer is None - ): + # 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 @@ -159,11 +151,9 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: token = None initialize = None try: - body, receive = await _buffer_body( - receive, peek=incoming_header is not None - ) + body, receive = await _buffer_body(receive) initialize = _parse_initialize(body) - if initialize is not None and incoming_header is None: + if initialize is not None: token = _mint_token(initialize) except Exception as error: # noqa: BLE001 log(f"PostHog MCP session middleware: mint failed - {error}") @@ -177,12 +167,7 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: send = _InitializeResponse( send, functools.partial( - _record_initialize, - observer, - initialize, - incoming_header, - scope, - started, + _record_initialize, observer, initialize, scope, started ), ) if token is not None: @@ -216,14 +201,10 @@ def _decode_headers(raw: Any) -> dict[str, str]: return result -async def _buffer_body(receive: Any, *, peek: bool = False) -> tuple[bytes, Any]: +async def _buffer_body(receive: Any) -> tuple[bytes, Any]: """Read up to ``_MAX_SNIFF_BODY`` of the request body for sniffing, then return that prefix plus a ``receive`` that replays it and streams the rest. - With ``peek``, also stop once ``_PEEK_BODY`` bytes arrived with a ``"method"`` - but no ``"initialize"`` in them: the request is not a handshake, so it is not - held in memory. - Memory is bounded to the cap: an `initialize` handshake is tiny, so once we exceed the cap the request cannot be one and we stop buffering — the prefix is replayed with ``more_body: True`` and any remaining chunks are forwarded @@ -233,7 +214,6 @@ async def _buffer_body(receive: Any, *, peek: bool = False) -> tuple[bytes, Any] total = 0 complete = False overflow = False - peeked = not peek # 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 @@ -250,12 +230,6 @@ async def _buffer_body(receive: Any, *, peek: bool = False) -> tuple[bytes, Any] if total > _MAX_SNIFF_BODY: overflow = True # too big to be initialize — stop buffering break - if not peeked and total >= _PEEK_BODY: - peeked = True - prefix = b"".join(chunks) - if b'"method"' in prefix and b'"initialize"' not in prefix: - overflow = True # not a handshake — stop buffering - break buffered = b"".join(chunks) replayed = False @@ -280,8 +254,7 @@ async def replay() -> dict[str, Any]: 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).""" - # Cheap first: with a bound middleware this runs on every tool call too. - if not body or b'"initialize"' not in body: + if not body: return None try: # RecursionError (deeply-nested JSON) is a RuntimeError, not a @@ -356,7 +329,6 @@ def __init__(self, request: Any) -> None: async def _record_initialize( data: MCPAnalyticsData, message: Dict[str, Any], - incoming_header: Optional[str], scope: Any, started: datetime, status: Optional[int], @@ -364,15 +336,15 @@ async def _record_initialize( ) -> 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, our token, or the one the client already sent), or - this process's own session when there is none. The body is not read, so the + 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 or incoming_header + header = session_header await record_initialize( data, session_id=await handshake_session_id( diff --git a/posthog/test/mcp/test_initialize_handshake.py b/posthog/test/mcp/test_initialize_handshake.py index a917fda86..022235fa1 100644 --- a/posthog/test/mcp/test_initialize_handshake.py +++ b/posthog/test/mcp/test_initialize_handshake.py @@ -824,6 +824,10 @@ async def slow_before_send(event): 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 @@ -940,3 +944,47 @@ def identify(request, extra): (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 e7d967b05..9db5aca7e 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -888,74 +888,3 @@ def test_app_built_probe_fires_on_the_installed_sdk_major(): f"probe blind to a built app on MCP {MCP_MAJOR}.x -- the ordering warning " "cannot fire for these servers" ) - - -async def test_peek_streams_a_large_request_that_is_not_initialize(): - """With a session header, a tool call is read only until the peek window shows - it is not a handshake. The rest streams to the app, byte for byte.""" - from posthog.mcp import asgi as asgi_mod - - body = b'{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"blob":"' - body += b"x" * (16 * asgi_mod._PEEK_BODY) + b'"}}' - parts = [body[i : i + 1024] for i in range(0, len(body), 1024)] - reads = [] - - async def receive(): - chunk = parts.pop(0) - reads.append(chunk) - return {"type": "http.request", "body": chunk, "more_body": bool(parts)} - - sniff, replay = await asgi_mod._buffer_body(receive, peek=True) - - assert sniff == b"" - assert sum(map(len, reads)) == asgi_mod._PEEK_BODY # stopped at the window - replayed = b"" - while True: - message = await replay() - replayed += message["body"] - if not message["more_body"]: - break - assert replayed == body - - -async def test_peek_still_reads_an_initialize_split_into_small_chunks(): - from posthog.mcp import asgi as asgi_mod - - body = json.dumps( - { - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": {"protocolVersion": "2025-06-18", "capabilities": {}}, - } - ).encode() - parts = [body[i : i + 7] for i in range(0, len(body), 7)] - - async def receive(): - chunk = parts.pop(0) - return {"type": "http.request", "body": chunk, "more_body": bool(parts)} - - sniff, _ = await asgi_mod._buffer_body(receive, peek=True) - - assert sniff == body - assert asgi_mod._parse_initialize(sniff)["method"] == "initialize" - - -async def test_peek_reads_an_initialize_whose_params_come_first(): - """JSON key order is arbitrary. An initialize with large params before its - "method" must still be read whole and recognised.""" - from posthog.mcp import asgi as asgi_mod - - params = {"capabilities": {"x": "y" * (2 * asgi_mod._PEEK_BODY)}} - body = json.dumps( - {"params": params, "jsonrpc": "2.0", "id": 1, "method": "initialize"} - ).encode() - parts = [body[i : i + 1024] for i in range(0, len(body), 1024)] - - async def receive(): - chunk = parts.pop(0) - return {"type": "http.request", "body": chunk, "more_body": bool(parts)} - - sniff, _ = await asgi_mod._buffer_body(receive, peek=True) - - assert asgi_mod._parse_initialize(sniff)["method"] == "initialize"