From f8de485d29c1476cfd8ec758c369aaa396555fbc Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 17:26:57 +0300 Subject: [PATCH 1/7] feat(mcp): carry session id + client identity across stateless pods Stateless / multi-pod MCP servers issue no session id, so $session_id fragments across pods and the client identity ("harness") sent only at initialize is lost on pods that never processed it. Mint a self-encoded token onto the Mcp-Session-Id response header at initialize (via a one-line ASGI middleware) and decode the replayed token on every request, recovering the same $session_id + client name/version on any pod with no shared state. Wire-compatible with the @posthog/mcp TypeScript SDK. - session_token.py: encode/decode codec (base64url JSON; sid/cn/cv/pv) - asgi.py: PostHogMcpStatelessSessionMiddleware + get_mcp_session - session.py / _internal.py: "token" session source, used verbatim + sticky - instrument() adapters: decode replayed header, backfill client identity Verified against a 2-pod stateless cluster behind round-robin nginx: all events share one $session_id and keep the harness (the published package fragments both). Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- examples/mcp_stateless_fastapi.py | 97 ++++++ posthog/mcp/__init__.py | 18 ++ posthog/mcp/_instrument_fastmcp.py | 9 + posthog/mcp/_instrument_lowlevel.py | 9 + posthog/mcp/_instrumentation.py | 26 +- posthog/mcp/_internal.py | 8 +- posthog/mcp/asgi.py | 203 +++++++++++++ posthog/mcp/session.py | 36 ++- posthog/mcp/session_token.py | 136 +++++++++ posthog/mcp/version.py | 2 +- posthog/test/mcp/test_session_token.py | 391 +++++++++++++++++++++++++ 11 files changed, 930 insertions(+), 5 deletions(-) create mode 100644 examples/mcp_stateless_fastapi.py create mode 100644 posthog/mcp/asgi.py create mode 100644 posthog/mcp/session_token.py create mode 100644 posthog/test/mcp/test_session_token.py diff --git a/examples/mcp_stateless_fastapi.py b/examples/mcp_stateless_fastapi.py new file mode 100644 index 00000000..84bd735c --- /dev/null +++ b/examples/mcp_stateless_fastapi.py @@ -0,0 +1,97 @@ +"""Stateless / multi-pod MCP analytics with the ``PostHogMCP`` custom dispatcher. + +A stateless MCP server issues no session id, so across pods (or per-request +transports) ``$session_id`` fragments and the client identity (the "harness", +e.g. Claude Code / Cursor) -- sent only at ``initialize`` -- is lost on any pod +that never processed the handshake. + +The fix is a self-encoded session token minted onto the ``Mcp-Session-Id`` +response header at ``initialize`` and replayed by the client on every request. +You do NOT set the header by hand: add +:class:`~posthog.mcp.PostHogMcpStatelessSessionMiddleware` once, then read the +recovered session with :func:`~posthog.mcp.get_mcp_session` and pass it into the +capture calls. + +Usage:: + + POSTHOG_PROJECT_API_KEY=phc_xxx uvicorn examples.mcp_stateless_fastapi:app + +The same one-line middleware also works in front of a mounted FastMCP app -- see +the note at the bottom. +""" + +import os + +from fastapi import FastAPI, Request + +from posthog.mcp import ( + PostHogMCP, + PostHogMcpStatelessSessionMiddleware, + get_mcp_session, +) + +posthog = PostHogMCP( + os.environ.get("POSTHOG_PROJECT_API_KEY", "phc_xxx"), + host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"), +) + +app = FastAPI() + +# One line. The middleware mints the session token onto the `Mcp-Session-Id` +# response header at `initialize` (when the client sent none) and decodes the +# replayed token on every later request. No manual header handling anywhere. +app.add_middleware(PostHogMcpStatelessSessionMiddleware) + + +@app.post("/mcp") +async def mcp_endpoint(request: Request): + body = await request.json() + method = body.get("method") + + # Recovered by the middleware from the replayed token. On the very first + # `initialize` it reflects the token just minted; on every later request + # (any pod) it carries the same session id + harness. + sess = get_mcp_session(request) + session_id = sess.session_id if sess else None + client_name = sess.client_name if sess else None + client_version = sess.client_version if sess else None + + if method == "initialize": + posthog.capture_initialize( + session_id=session_id, + client_name=client_name, + client_version=client_version, + parameters=body.get("params"), + ) + # ... return your real InitializeResult here ... + return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}} + + if method == "tools/call": + name = body["params"]["name"] + prepared = posthog.prepare_tool_call(name, body["params"].get("arguments")) + # ... dispatch prepared.args to your tool, then: ... + posthog.capture_tool_call( + tool_name=name, + session_id=session_id, + client_name=client_name, + client_version=client_version, + intent=prepared.intent, + intent_source=prepared.intent_source, + ) + return {"jsonrpc": "2.0", "id": body.get("id"), "result": {"content": []}} + + return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}} + + +# Mounted-FastMCP variant: the exact same middleware works in front of a FastMCP +# streamable-HTTP app, and instrument() reads the replayed token automatically -- +# no per-request code needed there: +# +# from mcp.server.fastmcp import FastMCP +# from posthog import Posthog +# from posthog.mcp import instrument, PostHogMcpStatelessSessionMiddleware +# +# server = FastMCP("my-server", stateless_http=True, json_response=True) +# instrument(server, Posthog("phc_xxx")) +# app = server.streamable_http_app() +# app.add_middleware(PostHogMcpStatelessSessionMiddleware) diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index b3cb0420..b9b7b1d4 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -44,6 +44,14 @@ from .logger import log, set_logger from .posthog_mcp import PostHogMCP from .session import derive_session_id_from_mcp_session, new_session_id +from .session_token import ( + MCP_SESSION_HEADER, + SessionTokenPayload, + decode_session_id, + encode_session_id, + read_mcp_session_header, +) +from .asgi import PostHogMcpStatelessSessionMiddleware, get_mcp_session from ._sink import McpEventSink from .tools import get_more_tools_result from .types import ( @@ -66,6 +74,16 @@ "PreparedToolCall", "get_more_tools_result", "derive_session_id_from_mcp_session", + # Self-encoded session tokens for stateless / multi-pod servers. Minted onto + # the `Mcp-Session-Id` response header by PostHogMcpStatelessSessionMiddleware + # and decoded on every request; codec is exported for custom HTTP layers. + "PostHogMcpStatelessSessionMiddleware", + "get_mcp_session", + "encode_session_id", + "decode_session_id", + "read_mcp_session_header", + "SessionTokenPayload", + "MCP_SESSION_HEADER", "set_logger", "POSTHOG_MCP_ANALYTICS_SOURCE", "PostHogMCPAnalyticsEvent", diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index c0fbf025..6d728fd5 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -47,6 +47,7 @@ record_tool_call, record_tools_list, request_to_dict, + resolve_session_and_client, ) from ._internal import MCPAnalyticsData from .logger import log @@ -91,6 +92,9 @@ async def wrapped( ) -> Any: client_name, client_version = _client_info(context) mcp_session_id = _mcp_session_id(context) + token, client_name, client_version = resolve_session_and_client( + mcp_session_id, client_name, client_version + ) request = build_tool_call_request(name, arguments) extra: Dict[str, Any] = {"session_id": mcp_session_id} @@ -101,6 +105,7 @@ async def wrapped( client_version=client_version, request=request, extra=extra, + token=token, ) missing_name = resolve_missing_capability_tool_name(data.options) @@ -211,6 +216,9 @@ async def list_handler(req: Any) -> Any: client_name, client_version = _low_level_client_info(server) mcp_session_id = _low_level_session_id(server) + token, client_name, client_version = resolve_session_and_client( + mcp_session_id, client_name, client_version + ) request = request_to_dict(req) extra: Dict[str, Any] = {"session_id": mcp_session_id} # Resolve session, emit $mcp_initialize (once per session) and identify here @@ -222,6 +230,7 @@ async def list_handler(req: Any) -> Any: client_version=client_version, request=request, extra=extra, + token=token, ) start = time.monotonic() diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 8f64f2b9..03372744 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -41,6 +41,7 @@ record_tool_call, record_tools_list, request_to_dict, + resolve_session_and_client, ) from ._internal import MCPAnalyticsData from .logger import log @@ -94,6 +95,9 @@ async def handler(req: Any) -> Any: arguments = dict(req.params.arguments or {}) client_name, client_version = _client_info(server) mcp_session_id = _mcp_session_id(server) + token, client_name, client_version = resolve_session_and_client( + mcp_session_id, client_name, client_version + ) request = build_tool_call_request(name, arguments) extra = {"session_id": mcp_session_id} @@ -104,6 +108,7 @@ async def handler(req: Any) -> Any: client_version=client_version, request=request, extra=extra, + token=token, ) missing_name = resolve_missing_capability_tool_name(data.options) @@ -223,6 +228,9 @@ async def handler(req: Any) -> Any: client_name, client_version = _client_info(server) mcp_session_id = _mcp_session_id(server) + token, client_name, client_version = resolve_session_and_client( + mcp_session_id, client_name, client_version + ) request = request_to_dict(req) extra = {"session_id": mcp_session_id} # Resolve session, emit $mcp_initialize (once per session) and identify here @@ -234,6 +242,7 @@ async def handler(req: Any) -> Any: client_version=client_version, request=request, extra=extra, + token=token, ) start = time.monotonic() diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index f7b2b6be..f3408f3c 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -23,6 +23,7 @@ from .logger import log from ._sanitization import build_captured_mcp_parameters from .session import resolve_session_id +from .session_token import SessionTokenPayload, decode_session_id # Keep strong refs to in-flight capture tasks/futures so they aren't GC'd mid-flight, # and so the asyncio ones can be awaited via drain_pending() before shutdown. Holds @@ -183,6 +184,24 @@ async def _apply_event_properties( event["properties"] = props +def resolve_session_and_client( + raw_session_id: Optional[str], + client_name: Optional[str], + client_version: Optional[str], +) -> tuple[Optional[SessionTokenPayload], Optional[str], Optional[str]]: + """Decode a replayed ``Mcp-Session-Id`` value as a self-encoded session token, + and backfill the client name/version from it when the live transport supplied + none (the stateless-pod case, where ``initialize`` was never seen here). + + Returns ``(token, client_name, client_version)``; ``token`` is ``None`` when the + header isn't one of our tokens (a plain transport UUID, JWT, or nothing).""" + token = decode_session_id(raw_session_id) + if token is not None: + client_name = client_name or token.client_name + client_version = client_version or token.client_version + return token, client_name, client_version + + async def prepare_request( data: MCPAnalyticsData, *, @@ -191,16 +210,21 @@ async def prepare_request( client_version: Optional[str], request: Dict[str, Any], extra: Optional[Dict[str, Any]], + token: Optional[SessionTokenPayload] = None, ) -> str: """Resolve the session id, run identify, then lazily emit initialize. Returns the session id to stamp on the event for this request. + ``token`` is the decoded self-encoded session token (see ``session_token.py``); + when present it takes precedence over ``mcp_session_id`` and carries the client + identity across stateless pods. + 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``.)""" - session_id = await resolve_session_id(data, mcp_session_id) + session_id = await resolve_session_id(data, mcp_session_id, token=token) identify_event = await handle_identify(data, session_id, request, extra) if identify_event: fire_and_forget(capture_event(data, identify_event)) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 0abe8cd6..39727e79 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -60,8 +60,14 @@ class MCPAnalyticsData: options: MCPAnalyticsOptions sink: Optional[McpEventSink] = None session_id: str = "" - session_source: str = "generated" # "generated" | "mcp" + session_source: str = "generated" # "generated" | "mcp" | "token" last_mcp_session_id: Optional[str] = None + # Client identity recovered from a self-encoded session token (see + # session_token.py). On a stateless pod that never processed `initialize`, + # the live `client_params` is empty, so these are the only harness source. + token_client_name: Optional[str] = None + token_client_version: Optional[str] = None + token_protocol_version: Optional[str] = None last_activity: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py new file mode 100644 index 00000000..cf98b175 --- /dev/null +++ b/posthog/mcp/asgi.py @@ -0,0 +1,203 @@ +"""Frictionless session-token minting for stateless / multi-pod MCP servers. + +A stateless MCP server issues no session id, so ``$session_id`` fragments across +pods and the client identity (the "harness") sent only at ``initialize`` is lost. +The fix is a self-encoded token (see :mod:`.session_token`) minted onto the +``Mcp-Session-Id`` response header at ``initialize``; clients replay it on every +request, so any pod recovers session + harness from the header alone. + +The MCP Python SDK owns ``initialize`` in its runner/session layer and forbids +overriding it via ``request_handlers``, so -- unlike the TS SDK, which wraps the +initialize handler -- there is no in-SDK seam to mint from. We mint at the HTTP +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. + +Add it once:: + + app.add_middleware(PostHogMcpStatelessSessionMiddleware) + +then read the recovered session on any request:: + + sess = get_mcp_session(request) + posthog.capture_tool_call( + tool_name=name, + session_id=sess.session_id if sess else None, + client_name=sess.client_name if sess else None, + client_version=sess.client_version if sess else None, + ) + +This module has no hard dependency on Starlette/FastAPI -- it speaks raw ASGI. +""" + +from __future__ import annotations + +import json +from typing import Any, Optional + +from .logger import log +from .session import new_session_id +from .session_token import ( + MCP_SESSION_HEADER, + SessionTokenPayload, + decode_session_id, + encode_session_id, + read_mcp_session_header, +) + +# Scope key the middleware stashes the decoded token payload under. +_SCOPE_KEY = "posthog_mcp_session" + +# 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 + + +def get_mcp_session(request_or_scope: Any) -> Optional[SessionTokenPayload]: + """Return the ``SessionTokenPayload`` the middleware recovered for this request, + or ``None``. Accepts a Starlette ``Request`` or a raw ASGI ``scope`` dict.""" + scope = getattr(request_or_scope, "scope", request_or_scope) + if not isinstance(scope, dict): + return None + value = scope.get(_SCOPE_KEY) + return value if isinstance(value, SessionTokenPayload) else None + + +class PostHogMcpStatelessSessionMiddleware: + """ASGI middleware that mints a session token onto the ``Mcp-Session-Id`` + response header at ``initialize`` (when the client sent none) and decodes the + replayed token on every request, exposing it via :func:`get_mcp_session`. + + 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: + self.app = app + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + if scope.get("type") != "http": + await self.app(scope, receive, send) + return + + try: + incoming_header = read_mcp_session_header(_headers_dict(scope)) + except Exception as error: # noqa: BLE001 + log(f"PostHog MCP session middleware: header read failed - {error}") + await self.app(scope, receive, send) + return + + # Decode a replayed token (if any) so the app can read it via get_mcp_session. + decoded = decode_session_id(incoming_header) + if decoded is not None: + scope[_SCOPE_KEY] = decoded + + # 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: + await self.app(scope, receive, send) + return + + body, receive = await _buffer_body(receive) + token = _mint_token_if_initialize(body) + if token 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)) + + +def _headers_dict(scope: Any) -> dict[str, str]: + """ASGI raw headers (list of (bytes, bytes)) -> a lowercased str dict.""" + result: dict[str, str] = {} + for key, value in scope.get("headers") or []: + try: + result[key.decode("latin-1").lower()] = value.decode("latin-1") + except Exception: # noqa: BLE001 + continue + return result + + +async def _buffer_body(receive: Any) -> tuple[bytes, Any]: + """Read the full request body, then return it plus a ``receive`` that replays + it downstream (so the app still sees an unconsumed stream).""" + chunks: list[bytes] = [] + more = True + while more: + message = await receive() + if message.get("type") != "http.request": + # A non-body message (e.g. http.disconnect); stop and replay what we have. + break + chunks.append(message.get("body", b"") or b"") + more = message.get("more_body", False) + + # Always buffer the whole body so replay is byte-faithful (Starlette's own + # Request.body() buffers fully too). This holds the request in memory, which + # is fine for tiny JSON-RPC POSTs on an MCP endpoint. + buffered = b"".join(chunks) + replayed = False + + async def replay() -> dict[str, Any]: + # Hand the whole buffered body back in one message, then defer to the + # original transport for anything after it (disconnect, etc.). + nonlocal replayed + if not replayed: + replayed = True + return {"type": "http.request", "body": buffered, "more_body": False} + return await receive() + + # Only *sniff* (parse to detect initialize) small bodies -- a giant POST is + # never our tiny initialize handshake, so skip minting but still replay it whole. + sniff = buffered if len(buffered) <= _MAX_SNIFF_BODY else b"" + 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).""" + if not body: + return None + try: + message = json.loads(body) + except (ValueError, TypeError): + 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 {} + try: + return encode_session_id( + SessionTokenPayload( + session_id=new_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(params.get("protocolVersion")), + ) + ) + except Exception as error: # noqa: BLE001 + log(f"PostHog MCP session middleware: mint failed - {error}") + return None + + +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 + stateful transport's own session id).""" + header = MCP_SESSION_HEADER.encode("latin-1") + token_bytes = token.encode("latin-1") + + async def wrapped(message: dict[str, Any]) -> None: + if message.get("type") == "http.response.start": + headers = list(message.get("headers") or []) + if not any(k.lower() == header for k, _ in headers): + headers.append((header, token_bytes)) + message = {**message, "headers": headers} + await send(message) + + return wrapped + + +def _str_or_none(value: Any) -> Optional[str]: + return value if isinstance(value, str) and value else None diff --git a/posthog/mcp/session.py b/posthog/mcp/session.py index 08531ef3..0bc19005 100644 --- a/posthog/mcp/session.py +++ b/posthog/mcp/session.py @@ -14,6 +14,7 @@ from .constants import INACTIVITY_TIMEOUT_IN_MINUTES from ._ids import deterministic_prefixed_id, new_prefixed_id from ._internal import MCPAnalyticsData +from .session_token import SessionTokenPayload __all__ = ["derive_session_id_from_mcp_session"] @@ -29,13 +30,42 @@ def derive_session_id_from_mcp_session(mcp_session_id: str) -> str: async def resolve_session_id( - data: MCPAnalyticsData, mcp_session_id: Optional[str] + data: MCPAnalyticsData, + mcp_session_id: Optional[str], + *, + token: Optional[SessionTokenPayload] = None, ) -> str: """Resolve the session id for a request. Mutates per-server state under a lock - so concurrent async requests can't race on session rotation.""" + so concurrent async requests can't race on session rotation. + + ``token`` is our self-encoded session token (see :mod:`.session_token`), + decoded from the replayed ``Mcp-Session-Id`` header. It is the only source + that survives a stateless / multi-pod deployment, so it takes precedence. + """ async with data.session_lock: now = datetime.now(timezone.utc) + if token is not None: + # A token we minted at `initialize`. Its session id is already a + # `ses_...` id, so use it verbatim -- do NOT re-hash. The client + # name/version/protocol version ride along for pods that never saw + # `initialize`; stash them so adapters can fall back to them when the + # live `client_params` is absent (the stateless-pod case). + data.session_id = token.session_id + data.session_source = "token" + data.token_client_name = token.client_name + data.token_client_version = token.client_version + data.token_protocol_version = token.protocol_version + data.last_activity = now + return data.session_id + + # A token session, like an MCP-derived one, lives as long as the client + # replays it -- keep the id even on a request that arrives without it, so + # the session doesn't fragment. + if data.session_source == "token": + data.last_activity = now + return data.session_id + if mcp_session_id: data.session_id = derive_session_id_from_mcp_session(mcp_session_id) data.last_mcp_session_id = mcp_session_id @@ -49,6 +79,8 @@ async def resolve_session_id( data.last_activity = now return data.session_id + # Only generated sessions roll over on inactivity -- token/MCP ids live + # as long as the client replays them, and regenerating would split them. timeout_seconds = INACTIVITY_TIMEOUT_IN_MINUTES * 60 if (now - data.last_activity).total_seconds() > timeout_seconds: data.session_id = new_session_id() diff --git a/posthog/mcp/session_token.py b/posthog/mcp/session_token.py new file mode 100644 index 00000000..228c0a59 --- /dev/null +++ b/posthog/mcp/session_token.py @@ -0,0 +1,136 @@ +"""Self-encoded session tokens for stateless / multi-pod MCP servers. + +A stateless server keeps nothing between requests, so every request starts a +new session and the client name/version (only sent at ``initialize``) is lost. +The one value clients replay on every request is the ``Mcp-Session-Id`` header. +So at ``initialize`` we mint that header as a token carrying the session id and +client identity -- any pod can read them back from the header alone. + +The token is unsigned: it holds only what the client already self-reports. It is +wire-compatible with the TypeScript SDK (same short keys ``sid``/``cn``/``cv``/``pv``), +so a token minted by one SDK decodes in the other. +""" + +from __future__ import annotations + +import base64 +import json +import re +from dataclasses import dataclass +from typing import Any, Optional + +MCP_SESSION_HEADER = "mcp-session-id" + +# On the wire the token is base64url(JSON) with shortened keys to keep the +# header small: sid = session_id, cn = client_name, cv = client_version, +# pv = protocol_version. +_MAX_TOKEN_LENGTH = 4096 +_MAX_SESSION_ID_LENGTH = 128 +_MAX_CLIENT_FIELD_LENGTH = 200 + +_BASE64URL_PATTERN = re.compile(r"^[A-Za-z0-9_-]+={0,2}$") + + +@dataclass +class SessionTokenPayload: + """What a session token carries.""" + + # PostHog session id (``ses_...``) -> ``$session_id``. + session_id: str + # MCP client name -> ``$mcp_client_name``. + client_name: Optional[str] = None + # MCP client version -> ``$mcp_client_version``. + client_version: Optional[str] = None + # MCP protocol (spec) version -> ``$mcp_protocol_version``. The client's + # *requested* version -- the only one known when the token is minted (before + # the initialize handshake negotiates). Lets pods that never saw ``initialize`` + # still stamp the spec version on their events. + protocol_version: Optional[str] = None + + +def encode_session_id(payload: SessionTokenPayload) -> str: + """Encode a session token for the ``Mcp-Session-Id`` response header. + + Raises ``ValueError`` for a missing/empty ``session_id`` (use ``new_session_id()``). + """ + if not isinstance(payload.session_id, str) or not payload.session_id: + raise ValueError( + "encode_session_id requires a non-empty `session_id` (use new_session_id())" + ) + wire: dict[str, str] = {"sid": payload.session_id} + if isinstance(payload.client_name, str) and payload.client_name: + wire["cn"] = payload.client_name[:_MAX_CLIENT_FIELD_LENGTH] + if isinstance(payload.client_version, str) and payload.client_version: + wire["cv"] = payload.client_version[:_MAX_CLIENT_FIELD_LENGTH] + if isinstance(payload.protocol_version, str) and payload.protocol_version: + wire["pv"] = payload.protocol_version[:_MAX_CLIENT_FIELD_LENGTH] + raw = json.dumps(wire, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") + + +def decode_session_id(value: Any) -> Optional[SessionTokenPayload]: + """Decode an ``Mcp-Session-Id`` value into a token payload. + + Returns ``None`` for anything that isn't one of our tokens (transport UUIDs, + JWTs, garbage) and never raises. + """ + if not isinstance(value, str) or not value or len(value) > _MAX_TOKEN_LENGTH: + return None + # JWTs carry dots; UUIDs pass this check but fail JSON parsing below. + if not _BASE64URL_PATTERN.match(value): + return None + try: + parsed = json.loads(_base64url_to_bytes(value)) + except (ValueError, TypeError): + return None + if not isinstance(parsed, dict): + return None + sid = parsed.get("sid") + if not isinstance(sid, str) or not sid or len(sid) > _MAX_SESSION_ID_LENGTH: + return None + payload = SessionTokenPayload(session_id=sid) + # A bad cn/cv/pv just means no client info -- it does not reject the token. + cn = parsed.get("cn") + if isinstance(cn, str) and cn: + payload.client_name = cn[:_MAX_CLIENT_FIELD_LENGTH] + cv = parsed.get("cv") + if isinstance(cv, str) and cv: + payload.client_version = cv[:_MAX_CLIENT_FIELD_LENGTH] + pv = parsed.get("pv") + if isinstance(pv, str) and pv: + payload.protocol_version = pv[:_MAX_CLIENT_FIELD_LENGTH] + return payload + + +def read_mcp_session_header(headers: Any) -> Optional[str]: + """Read the ``mcp-session-id`` value off a headers mapping. + + Handles case-insensitive keys (transports lowercase them, but hand-built + mappings may not) and list-valued headers, and trims whitespace. + """ + if headers is None: + return None + value: Any = None + # Mapping-like: prefer a direct get, else scan case-insensitively. + get = getattr(headers, "get", None) + if callable(get): + value = get(MCP_SESSION_HEADER) + if value is None: + try: + items = headers.items() + except AttributeError: + return None + for key, candidate in items: + if isinstance(key, str) and key.lower() == MCP_SESSION_HEADER: + value = candidate + break + first = value[0] if isinstance(value, (list, tuple)) and value else value + if not isinstance(first, str): + return None + trimmed = first.strip() + return trimmed or None + + +def _base64url_to_bytes(value: str) -> bytes: + padding = "=" * (-len(value) % 4) + return base64.urlsafe_b64decode(value + padding) diff --git a/posthog/mcp/version.py b/posthog/mcp/version.py index 77f0d164..c91ca805 100644 --- a/posthog/mcp/version.py +++ b/posthog/mcp/version.py @@ -4,4 +4,4 @@ # Version of the PostHog MCP analytics SDK surface. Stamped as ``sdk_version`` # on every captured event. The package itself ships inside posthog. -__version__ = "0.1.0" +__version__ = "0.2.0" diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py new file mode 100644 index 00000000..f193955d --- /dev/null +++ b/posthog/test/mcp/test_session_token.py @@ -0,0 +1,391 @@ +"""Tests for the self-encoded session token (stateless / multi-pod fix): the +codec, its use in session resolution, and the ASGI minting middleware.""" + +from __future__ import annotations + +import json + +import pytest + +from posthog.mcp._internal import MCPAnalyticsData +from posthog.mcp.asgi import ( + PostHogMcpStatelessSessionMiddleware, + get_mcp_session, +) +from posthog.mcp.session import new_session_id, resolve_session_id +from posthog.mcp.session_token import ( + MCP_SESSION_HEADER, + SessionTokenPayload, + decode_session_id, + encode_session_id, + read_mcp_session_header, +) +from posthog.mcp.types import MCPAnalyticsOptions + + +# --- codec ------------------------------------------------------------------- + + +def test_encode_decode_round_trip(): + token = encode_session_id( + SessionTokenPayload( + session_id="ses_abc", + client_name="Claude Code", + client_version="1.2.3", + protocol_version="2025-06-18", + ) + ) + payload = decode_session_id(token) + assert payload is not None + assert payload.session_id == "ses_abc" + assert payload.client_name == "Claude Code" + assert payload.client_version == "1.2.3" + assert payload.protocol_version == "2025-06-18" + + +def test_encode_decode_survives_non_ascii_client_name(): + token = encode_session_id( + SessionTokenPayload(session_id="ses_1", client_name="クロード🤖") + ) + payload = decode_session_id(token) + assert payload is not None and payload.client_name == "クロード🤖" + + +def test_token_matches_mcp_visible_ascii_session_pattern(): + # The MCP SDK validates a session id header against ^[\x21-\x7E]+$; a base64url + # token (no padding) must pass so the transport accepts it as the header. + token = encode_session_id( + SessionTokenPayload(session_id="ses_abc", client_name="x", client_version="y") + ) + assert all(0x21 <= ord(ch) <= 0x7E for ch in token) + + +def test_encode_requires_session_id(): + with pytest.raises(ValueError): + encode_session_id(SessionTokenPayload(session_id="")) + + +def test_decode_rejects_non_tokens_without_raising(): + # Transport UUID, JWT-ish (dots), empty, wrong types, and garbage. + assert decode_session_id("550e8400-e29b-41d4-a716-446655440000") is None + assert decode_session_id("aaa.bbb.ccc") is None + assert decode_session_id("") is None + assert decode_session_id(None) is None + assert decode_session_id(12345) is None # type: ignore[arg-type] + assert decode_session_id("!!!not base64!!!") is None + + +def test_decode_rejects_oversized_token(): + assert decode_session_id("A" * 5000) is None + + +def test_decode_rejects_payload_without_sid(): + import base64 + + raw = ( + base64.urlsafe_b64encode(json.dumps({"cn": "x"}).encode()).decode().rstrip("=") + ) + assert decode_session_id(raw) is None + + +def test_encode_truncates_long_client_fields(): + payload = decode_session_id( + encode_session_id( + SessionTokenPayload(session_id="ses_1", client_name="c" * 500) + ) + ) + assert payload is not None and len(payload.client_name or "") == 200 + + +def test_read_mcp_session_header_case_insensitive_list_and_trim(): + assert read_mcp_session_header({"Mcp-Session-Id": " tok "}) == "tok" + assert read_mcp_session_header({MCP_SESSION_HEADER: ["a", "b"]}) == "a" + assert read_mcp_session_header({"other": "x"}) is None + assert read_mcp_session_header({MCP_SESSION_HEADER: " "}) is None + assert read_mcp_session_header(None) is None + + +# --- session resolution ------------------------------------------------------ + + +def _data() -> MCPAnalyticsData: + data = MCPAnalyticsData(options=MCPAnalyticsOptions()) + data.session_id = new_session_id() + return data + + +async def test_resolve_session_id_uses_token_verbatim(): + data = _data() + token = decode_session_id( + encode_session_id( + SessionTokenPayload(session_id="ses_tok", client_name="Cursor") + ) + ) + sid = await resolve_session_id(data, "ignored-raw", token=token) + # Used verbatim -- NOT re-hashed through derive_session_id_from_mcp_session. + assert sid == "ses_tok" + assert data.session_source == "token" + assert data.token_client_name == "Cursor" + + +async def test_token_session_does_not_fragment_or_roll_over(): + from datetime import datetime, timedelta, timezone + + data = _data() + token = decode_session_id( + encode_session_id(SessionTokenPayload(session_id="ses_tok")) + ) + first = await resolve_session_id(data, None, token=token) + # A later request without the header keeps the id... + assert await resolve_session_id(data, None) == first + # ...and it never rolls over on inactivity (only generated sessions do). + data.last_activity = datetime.now(timezone.utc) - timedelta(minutes=31) + assert await resolve_session_id(data, None) == first + + +async def test_multi_pod_two_instances_resolve_same_session_and_harness(): + """The regression this fixes: independent pods (independent per-server state) + resolve the same replayed token to the same session id + harness.""" + token_str = encode_session_id( + SessionTokenPayload(session_id="ses_shared", client_name="Claude Code") + ) + results = [] + for _ in range(2): # two "pods" + data = _data() + token = decode_session_id(token_str) + sid = await resolve_session_id(data, token_str, token=token) + results.append((sid, data.session_source, data.token_client_name)) + assert results[0] == results[1] == ("ses_shared", "token", "Claude Code") + + +# --- ASGI minting middleware ------------------------------------------------- + + +def _init_body(client_name="Claude Code", client_version="1.0.0", pv="2025-06-18"): + return json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": pv, + "clientInfo": {"name": client_name, "version": client_version}, + }, + } + ).encode() + + +async def _run(app, scope, body: bytes): + """Drive an ASGI app once, returning the response-start headers as a dict.""" + sent = {} + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + async def send(message): + if message["type"] == "http.response.start": + sent["headers"] = { + k.decode().lower(): v.decode() for k, v in message.get("headers", []) + } + + await app(scope, receive, send) + return sent + + +def _scope(method="POST", headers=None): + raw = [(k.encode(), v.encode()) for k, v in (headers or {}).items()] + return {"type": "http", "method": method, "headers": raw} + + +async def test_middleware_mints_session_header_on_initialize(): + captured = {} + + async def app(scope, receive, send): + captured["session"] = get_mcp_session(scope) + # Read the body downstream to prove the middleware replayed it intact. + captured["body"] = (await receive())["body"] + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + mw = PostHogMcpStatelessSessionMiddleware(app) + scope = _scope() + body = _init_body() + sent = await _run(mw, scope, body) + + token = sent["headers"].get(MCP_SESSION_HEADER) + assert token is not None + payload = decode_session_id(token) + assert payload is not None + assert payload.client_name == "Claude Code" + assert payload.client_version == "1.0.0" + assert payload.protocol_version == "2025-06-18" + # The app saw the recovered session and the untouched body. + assert captured["session"].session_id == payload.session_id + assert captured["body"] == body + + +async def test_middleware_does_not_clobber_existing_response_header(): + async def app(scope, receive, send): + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(MCP_SESSION_HEADER.encode(), b"existing-uuid")], + } + ) + await send({"type": "http.response.body", "body": b"{}"}) + + mw = PostHogMcpStatelessSessionMiddleware(app) + sent = await _run(mw, _scope(), _init_body()) + assert sent["headers"][MCP_SESSION_HEADER] == "existing-uuid" + + +async def test_middleware_skips_when_client_replays_session_id(): + """A request already carrying a session id (ours or a stateful UUID) is left + alone; the replayed token is still decoded and exposed to the app.""" + token_str = encode_session_id( + SessionTokenPayload(session_id="ses_replayed", client_name="Cursor") + ) + seen = {} + + async def app(scope, receive, send): + seen["session"] = get_mcp_session(scope) + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + mw = PostHogMcpStatelessSessionMiddleware(app) + scope = _scope(headers={MCP_SESSION_HEADER: token_str}) + sent = await _run(mw, scope, _init_body()) + # No mint (client already has a session id) ... + assert MCP_SESSION_HEADER not in sent["headers"] + # ... but the replayed token is decoded for the app. + assert seen["session"].session_id == "ses_replayed" + assert seen["session"].client_name == "Cursor" + + +async def test_middleware_passes_through_non_initialize_post(): + async def app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + mw = PostHogMcpStatelessSessionMiddleware(app) + body = json.dumps({"jsonrpc": "2.0", "method": "tools/list", "id": 2}).encode() + sent = await _run(mw, _scope(), body) + assert MCP_SESSION_HEADER not in sent["headers"] + + +async def test_middleware_replays_body_split_across_chunks_and_over_sniff_cap(): + """The buffered body must be replayed byte-for-byte even when it arrives in + multiple ASGI messages and exceeds the sniff cap (a large POST is never our + initialize handshake, so we skip minting but must not corrupt the body).""" + from posthog.mcp import asgi as asgi_mod + + big = b'{"jsonrpc":"2.0","method":"tools/call","params":{"blob":"' + big += b"x" * (asgi_mod._MAX_SNIFF_BODY + 10) + big += b'"}}' + parts = [big[i : i + 4096] for i in range(0, len(big), 4096)] + + received = {} + + async def app(scope, receive, send): + buf = b"" + while True: + msg = await receive() + buf += msg.get("body", b"") + if not msg.get("more_body"): + break + received["body"] = buf + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"{}"}) + + async def receive(): + if parts: + chunk = parts.pop(0) + return {"type": "http.request", "body": chunk, "more_body": bool(parts)} + return {"type": "http.request", "body": b"", "more_body": False} + + sent = {} + + async def send(message): + if message["type"] == "http.response.start": + sent["headers"] = { + k.decode().lower(): v.decode() for k, v in message.get("headers", []) + } + + mw = PostHogMcpStatelessSessionMiddleware(app) + await mw(_scope(), receive, send) + + assert received["body"] == big # byte-faithful, not truncated at the cap + assert MCP_SESSION_HEADER not in sent["headers"] # too big to sniff -> no mint + + +# --- end-to-end against a real stateless FastMCP transport ------------------- + + +def test_middleware_end_to_end_with_stateless_fastmcp(): + """Cross-version sanity: mount a real stateless FastMCP streamable-HTTP app, + add the middleware, and confirm (1) the `initialize` response carries a + decodable minted token with the client's harness, and (2) a follow-up request + on a fresh stateless transport that replays the token is accepted (not + rejected) -- i.e. the same session survives across pods.""" + pytest.importorskip("starlette.testclient") + from mcp.server.fastmcp import FastMCP + from mcp.server.transport_security import TransportSecuritySettings + from starlette.testclient import TestClient + + srv = FastMCP( + "posthog-token-test", + stateless_http=True, + json_response=True, + # TestClient sends Host: testserver; allow it past DNS-rebinding protection. + transport_security=TransportSecuritySettings( + enable_dns_rebinding_protection=False + ), + ) + + @srv.tool() + def add(a: int, b: int) -> int: + return a + b + + app = srv.streamable_http_app() + app.add_middleware(PostHogMcpStatelessSessionMiddleware) + + hdrs = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + + def rpc(method, params=None, id=1, extra=None): + body = {"jsonrpc": "2.0", "id": id, "method": method} + if params is not None: + body["params"] = params + return {**hdrs, **(extra or {})}, json.dumps(body) + + with TestClient(app) as client: + h, b = rpc( + "initialize", + { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "Claude Code", "version": "9.9.9"}, + }, + ) + resp = client.post("/mcp", headers=h, content=b) + assert resp.status_code == 200, resp.text + + token = resp.headers.get(MCP_SESSION_HEADER) + payload = decode_session_id(token) + assert payload is not None, "initialize response did not carry our token" + assert payload.client_name == "Claude Code" + assert payload.client_version == "9.9.9" + assert payload.protocol_version == "2025-06-18" + + # A different pod (fresh stateless transport) replays the token: accepted. + h2, b2 = rpc( + "tools/list", + {}, + id=2, + extra={MCP_SESSION_HEADER: token, "mcp-protocol-version": "2025-06-18"}, + ) + resp2 = client.post("/mcp", headers=h2, content=b2) + assert resp2.status_code == 200, resp2.text From 2b05c276b5dd5ca89aaec2da359dca116057ad2c Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 17:50:48 +0300 Subject: [PATCH 2/7] feat(mcp): auto-wire the stateless mint from instrument() instrument() now wraps the FastMCP server's ASGI-app factories (streamable_http_app / sse_app, which mcp.run() also calls), so a stateless / multi-pod FastMCP keeps one $session_id + the client harness across pods with zero extra setup -- no app.add_middleware(). The middleware stays exported for the PostHogMCP custom-dispatcher path, where you own the ASGI app. Auto-wire is instance-scoped (no global / transport-internal patching), a no-op for stdio / low-level servers, and safe if the middleware is also added manually. Verified on a 2-pod stateless cluster behind round-robin nginx: with instrument() alone, all events across both pods share one $session_id. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- examples/mcp_stateless_fastapi.py | 13 +++--- posthog/mcp/__init__.py | 11 ++++- posthog/mcp/asgi.py | 47 +++++++++++++++++++- posthog/test/mcp/test_session_token.py | 59 ++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 8 deletions(-) diff --git a/examples/mcp_stateless_fastapi.py b/examples/mcp_stateless_fastapi.py index 84bd735c..69f41167 100644 --- a/examples/mcp_stateless_fastapi.py +++ b/examples/mcp_stateless_fastapi.py @@ -83,15 +83,16 @@ async def mcp_endpoint(request: Request): return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}} -# Mounted-FastMCP variant: the exact same middleware works in front of a FastMCP -# streamable-HTTP app, and instrument() reads the replayed token automatically -- -# no per-request code needed there: +# Mounted-FastMCP variant: nothing to add at all. instrument() auto-wires the +# stateless mint into the server's app factory (and mcp.run() uses the same +# factory), so a stateless FastMCP keeps one $session_id + the harness across pods +# with zero extra setup: # # from mcp.server.fastmcp import FastMCP # from posthog import Posthog -# from posthog.mcp import instrument, PostHogMcpStatelessSessionMiddleware +# from posthog.mcp import instrument # # server = FastMCP("my-server", stateless_http=True, json_response=True) # instrument(server, Posthog("phc_xxx")) -# app = server.streamable_http_app() -# app.add_middleware(PostHogMcpStatelessSessionMiddleware) +# app = server.streamable_http_app() # already mints/decodes the session token +# # (server.run(transport="streamable-http") works the same way) diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index b9b7b1d4..9ac09c4b 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -51,7 +51,11 @@ encode_session_id, read_mcp_session_header, ) -from .asgi import PostHogMcpStatelessSessionMiddleware, get_mcp_session +from .asgi import ( + PostHogMcpStatelessSessionMiddleware, + autowire_stateless_mint, + get_mcp_session, +) from ._sink import McpEventSink from .tools import get_more_tools_result from .types import ( @@ -245,6 +249,11 @@ def instrument( "fastmcp 2.0) or a low-level mcp.server.Server." ) + # Zero-config stateless minting: wrap the server's ASGI-app factories so a + # stateless/multi-pod deployment keeps one $session_id + the client harness + # across pods with no extra setup. No-op for stdio / low-level servers. + autowire_stateless_mint(server) + return McpAnalytics(key) except Exception as error: # noqa: BLE001 log(f"Warning: failed to instrument server - {error}") diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index cf98b175..342e5dd9 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -12,7 +12,12 @@ 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. -Add it once:: +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 +app it builds already carries the middleware -- nothing extra to add. + +For a custom ``PostHogMCP`` dispatcher (you own the ASGI app), add it once:: app.add_middleware(PostHogMcpStatelessSessionMiddleware) @@ -31,6 +36,7 @@ from __future__ import annotations +import functools import json from typing import Any, Optional @@ -201,3 +207,42 @@ async def wrapped(message: dict[str, Any]) -> None: def _str_or_none(value: Any) -> Optional[str]: return value if isinstance(value, str) and value else None + + +# Marker so we never double-wrap a factory (idempotent across repeat instrument()). +_AUTOWIRED = "__posthog_mcp_autowired__" + + +def autowire_stateless_mint(server: Any) -> None: + """Make stateless minting zero-config on the ``instrument()`` path. + + Wraps a FastMCP server's ASGI-app factories so the app they build already has + :class:`PostHogMcpStatelessSessionMiddleware` applied. Covers both + ``server.streamable_http_app()`` / ``sse_app()`` and ``mcp.run(transport=...)`` + (which calls those factories internally), so the user adds nothing. + + No-op for servers without app factories (stdio / low-level ``Server``), and safe + if the middleware is also added manually (it only mints when none is present).""" + for attr in ("streamable_http_app", "sse_app", "http_app"): + original = getattr(server, attr, None) + if not callable(original) or getattr(original, _AUTOWIRED, False): + continue + try: + setattr(server, attr, _wrap_app_factory(original)) + except Exception as error: # noqa: BLE001 - never let wiring break instrument() + log(f"PostHog MCP: could not auto-wire stateless mint on {attr} - {error}") + + +def _wrap_app_factory(original: Any) -> Any: + @functools.wraps(original) + def factory(*args: Any, **kwargs: Any) -> Any: + app = original(*args, **kwargs) + add_middleware = getattr(app, "add_middleware", None) + if callable(add_middleware): + add_middleware(PostHogMcpStatelessSessionMiddleware) + return app + # Not a Starlette app -- wrap as raw ASGI so minting still happens. + return PostHogMcpStatelessSessionMiddleware(app) + + setattr(factory, _AUTOWIRED, True) + return factory diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index f193955d..5ebc14dc 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -389,3 +389,62 @@ def rpc(method, params=None, id=1, extra=None): ) resp2 = client.post("/mcp", headers=h2, content=b2) assert resp2.status_code == 200, resp2.text + + +def test_instrument_autowires_stateless_mint_no_manual_middleware(): + """instrument() alone (no app.add_middleware) makes the FastMCP streamable-HTTP + app mint the session token -- the zero-config path. mcp.run() uses the same + factory internally, so it's covered too.""" + pytest.importorskip("starlette.testclient") + from mcp.server.fastmcp import FastMCP + from mcp.server.transport_security import TransportSecuritySettings + from starlette.testclient import TestClient + + from posthog.mcp import instrument + + class _Sink: + def capture(self, *_: object, **__: object) -> None: + pass + + srv = FastMCP( + "posthog-autowire-test", + stateless_http=True, + json_response=True, + transport_security=TransportSecuritySettings( + enable_dns_rebinding_protection=False + ), + ) + + @srv.tool() + def ping() -> str: + return "pong" + + # No app.add_middleware() anywhere -- instrument() wires the mint itself. + instrument(srv, _Sink()) + app = srv.streamable_http_app() + + with TestClient(app) as client: + resp = client.post( + "/mcp", + headers={ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + }, + content=json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "Cursor", "version": "0.42"}, + }, + } + ), + ) + assert resp.status_code == 200, resp.text + payload = decode_session_id(resp.headers.get(MCP_SESSION_HEADER)) + assert payload is not None, "instrument() did not auto-wire the mint" + assert payload.client_name == "Cursor" + assert payload.client_version == "0.42" From d6806f92969c2cb24ab371f889b36ee1559173bc Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 18:08:15 +0300 Subject: [PATCH 3/7] docs(mcp): lead the stateless example with the FastMCP + instrument() path Reword examples/mcp_stateless_fastapi.py to put the preferred, zero-config approach up top: FastMCP(stateless_http=True) + instrument() + serving via server.run(transport="streamable-http") or server.streamable_http_app(). The PostHogMCP custom-dispatcher path (which needs the middleware, since there's no server object to auto-wire) is now framed as the fallback. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- examples/mcp_stateless_fastapi.py | 60 +++++++++++++++---------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/examples/mcp_stateless_fastapi.py b/examples/mcp_stateless_fastapi.py index 69f41167..6a152df2 100644 --- a/examples/mcp_stateless_fastapi.py +++ b/examples/mcp_stateless_fastapi.py @@ -1,23 +1,37 @@ -"""Stateless / multi-pod MCP analytics with the ``PostHogMCP`` custom dispatcher. +"""Stateless / multi-pod MCP analytics -- two ways to set it up. A stateless MCP server issues no session id, so across pods (or per-request transports) ``$session_id`` fragments and the client identity (the "harness", e.g. Claude Code / Cursor) -- sent only at ``initialize`` -- is lost on any pod -that never processed the handshake. +that never processed the handshake. The fix is a self-encoded session token minted +onto the ``Mcp-Session-Id`` response header at ``initialize`` and replayed by the +client on every request. You never set that header by hand. -The fix is a self-encoded session token minted onto the ``Mcp-Session-Id`` -response header at ``initialize`` and replayed by the client on every request. -You do NOT set the header by hand: add -:class:`~posthog.mcp.PostHogMcpStatelessSessionMiddleware` once, then read the -recovered session with :func:`~posthog.mcp.get_mcp_session` and pass it into the -capture calls. +PREFERRED -- FastMCP + ``instrument()`` (no extra code):: -Usage:: + from mcp.server.fastmcp import FastMCP + from posthog import Posthog + from posthog.mcp import instrument - POSTHOG_PROJECT_API_KEY=phc_xxx uvicorn examples.mcp_stateless_fastapi:app + server = FastMCP("my-server", stateless_http=True) # stateless_http is the key bit + instrument(server, Posthog("phc_your_project_api_key")) + + # Serve via the FastMCP app factory -- instrument() already wired the mint into it: + server.run(transport="streamable-http") + # or: app = server.streamable_http_app() # then serve with uvicorn + +That's it -- no middleware, no header handling. ``instrument()`` auto-wires the +mint into the server's ``streamable_http_app()`` / ``sse_app()`` factories (which +``run()`` uses), so every pod keeps one ``$session_id`` + the harness. + +CUSTOM DISPATCHER -- ``PostHogMCP`` (the rest of this file): -The same one-line middleware also works in front of a mounted FastMCP app -- see -the note at the bottom. +Only when you hand-roll the MCP endpoint yourself, so there's no FastMCP server for +``instrument()`` to wire. Then add +:class:`~posthog.mcp.PostHogMcpStatelessSessionMiddleware` once and read the +recovered session with :func:`~posthog.mcp.get_mcp_session`. Run it with:: + + POSTHOG_PROJECT_API_KEY=phc_xxx uvicorn examples.mcp_stateless_fastapi:app """ import os @@ -37,9 +51,10 @@ app = FastAPI() -# One line. The middleware mints the session token onto the `Mcp-Session-Id` -# response header at `initialize` (when the client sent none) and decodes the -# replayed token on every later request. No manual header handling anywhere. +# Custom-dispatcher path: there's no FastMCP server for instrument() to auto-wire, +# so add the mint middleware yourself -- one line. It mints the session token onto +# the `Mcp-Session-Id` response header at `initialize` (when the client sent none) +# and decodes the replayed token on every later request. No manual header handling. app.add_middleware(PostHogMcpStatelessSessionMiddleware) @@ -81,18 +96,3 @@ async def mcp_endpoint(request: Request): return {"jsonrpc": "2.0", "id": body.get("id"), "result": {"content": []}} return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}} - - -# Mounted-FastMCP variant: nothing to add at all. instrument() auto-wires the -# stateless mint into the server's app factory (and mcp.run() uses the same -# factory), so a stateless FastMCP keeps one $session_id + the harness across pods -# with zero extra setup: -# -# from mcp.server.fastmcp import FastMCP -# from posthog import Posthog -# from posthog.mcp import instrument -# -# server = FastMCP("my-server", stateless_http=True, json_response=True) -# instrument(server, Posthog("phc_xxx")) -# app = server.streamable_http_app() # already mints/decodes the session token -# # (server.run(transport="streamable-http") works the same way) From f5e8d104cc7a1d3bae525d6f56d677779fe0295e Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 18:10:22 +0300 Subject: [PATCH 4/7] docs(mcp): tighten the stateless example comments Trim the docstring and inline comments; drop the "preferred" framing. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- examples/mcp_stateless_fastapi.py | 47 ++++++++----------------------- 1 file changed, 11 insertions(+), 36 deletions(-) diff --git a/examples/mcp_stateless_fastapi.py b/examples/mcp_stateless_fastapi.py index 6a152df2..7022ccb3 100644 --- a/examples/mcp_stateless_fastapi.py +++ b/examples/mcp_stateless_fastapi.py @@ -1,35 +1,15 @@ -"""Stateless / multi-pod MCP analytics -- two ways to set it up. +"""Stateless / multi-pod MCP analytics: keep one ``$session_id`` + the client +identity across pods via a session token in the ``Mcp-Session-Id`` header. -A stateless MCP server issues no session id, so across pods (or per-request -transports) ``$session_id`` fragments and the client identity (the "harness", -e.g. Claude Code / Cursor) -- sent only at ``initialize`` -- is lost on any pod -that never processed the handshake. The fix is a self-encoded session token minted -onto the ``Mcp-Session-Id`` response header at ``initialize`` and replayed by the -client on every request. You never set that header by hand. +FastMCP + ``instrument()`` needs no extra code -- ``instrument()`` wires the mint +into the server's ``streamable_http_app()`` / ``sse_app()`` factories:: -PREFERRED -- FastMCP + ``instrument()`` (no extra code):: + server = FastMCP("my-server", stateless_http=True) + instrument(server, Posthog("phc_...")) + server.run(transport="streamable-http") # or: app = server.streamable_http_app() - from mcp.server.fastmcp import FastMCP - from posthog import Posthog - from posthog.mcp import instrument - - server = FastMCP("my-server", stateless_http=True) # stateless_http is the key bit - instrument(server, Posthog("phc_your_project_api_key")) - - # Serve via the FastMCP app factory -- instrument() already wired the mint into it: - server.run(transport="streamable-http") - # or: app = server.streamable_http_app() # then serve with uvicorn - -That's it -- no middleware, no header handling. ``instrument()`` auto-wires the -mint into the server's ``streamable_http_app()`` / ``sse_app()`` factories (which -``run()`` uses), so every pod keeps one ``$session_id`` + the harness. - -CUSTOM DISPATCHER -- ``PostHogMCP`` (the rest of this file): - -Only when you hand-roll the MCP endpoint yourself, so there's no FastMCP server for -``instrument()`` to wire. Then add -:class:`~posthog.mcp.PostHogMcpStatelessSessionMiddleware` once and read the -recovered session with :func:`~posthog.mcp.get_mcp_session`. Run it with:: +This file shows the other case: a custom ``PostHogMCP`` dispatcher, where there's +no server for ``instrument()`` to wire, so you add the middleware yourself. Run:: POSTHOG_PROJECT_API_KEY=phc_xxx uvicorn examples.mcp_stateless_fastapi:app """ @@ -51,10 +31,7 @@ app = FastAPI() -# Custom-dispatcher path: there's no FastMCP server for instrument() to auto-wire, -# so add the mint middleware yourself -- one line. It mints the session token onto -# the `Mcp-Session-Id` response header at `initialize` (when the client sent none) -# and decodes the replayed token on every later request. No manual header handling. +# Mints the token at `initialize` and decodes the replayed one on every request. app.add_middleware(PostHogMcpStatelessSessionMiddleware) @@ -63,9 +40,7 @@ async def mcp_endpoint(request: Request): body = await request.json() method = body.get("method") - # Recovered by the middleware from the replayed token. On the very first - # `initialize` it reflects the token just minted; on every later request - # (any pod) it carries the same session id + harness. + # Session id + client identity recovered from the token, same across pods. sess = get_mcp_session(request) session_id = sess.session_id if sess else None client_name = sess.client_name if sess else None From c5771aa2817babde7d2d9315c9229165ac267cac Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 18:15:17 +0300 Subject: [PATCH 5/7] docs(mcp): make the stateless example a runnable FastMCP server Rename mcp_stateless_fastapi.py -> mcp_stateless.py and make the runnable body the FastMCP + instrument() + server.run() path (no FastAPI); the custom-dispatcher middleware approach is now a short comment at the bottom. Add a see-also pointer from mcp_analytics_demo.py. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- examples/mcp_analytics_demo.py | 2 + examples/mcp_stateless.py | 46 +++++++++++++++++++ examples/mcp_stateless_fastapi.py | 73 ------------------------------- 3 files changed, 48 insertions(+), 73 deletions(-) create mode 100644 examples/mcp_stateless.py delete mode 100644 examples/mcp_stateless_fastapi.py diff --git a/examples/mcp_analytics_demo.py b/examples/mcp_analytics_demo.py index 5cd784f3..5ce20981 100644 --- a/examples/mcp_analytics_demo.py +++ b/examples/mcp_analytics_demo.py @@ -11,6 +11,8 @@ This drives the instrumented server's seams directly (tools/list + tool calls) rather than spinning up a transport + client, so it's a self-contained way to generate events. + +Stateless / multi-pod servers: see ``examples/mcp_stateless.py``. """ import asyncio diff --git a/examples/mcp_stateless.py b/examples/mcp_stateless.py new file mode 100644 index 00000000..943c44c1 --- /dev/null +++ b/examples/mcp_stateless.py @@ -0,0 +1,46 @@ +"""Stateless / multi-pod MCP analytics: keep one ``$session_id`` + the client +identity across pods via a session token in the ``Mcp-Session-Id`` header. + +``instrument()`` wires the mint into FastMCP's ``streamable_http_app()`` / +``sse_app()`` factories, so a stateless server needs nothing extra. Run:: + + POSTHOG_PROJECT_API_KEY=phc_xxx python examples/mcp_stateless.py +""" + +import os + +from mcp.server.fastmcp import FastMCP + +from posthog import Posthog +from posthog.mcp import instrument + +posthog = Posthog( + os.environ.get("POSTHOG_PROJECT_API_KEY", "phc_xxx"), + host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"), +) + +server = FastMCP("my-server", stateless_http=True) + + +@server.tool() +def greet(name: str) -> str: + """Say hello.""" + return f"hi {name}" + + +instrument(server, posthog) + + +if __name__ == "__main__": + # instrument() already wired the session-token mint into this app, so every pod + # keeps one $session_id + the harness. (app = server.streamable_http_app() too.) + server.run(transport="streamable-http") + + +# No FastMCP server to wire (a custom dispatcher)? Add the middleware to your own +# ASGI app and read the recovered session per request: +# +# from posthog.mcp import PostHogMcpStatelessSessionMiddleware, get_mcp_session +# +# app.add_middleware(PostHogMcpStatelessSessionMiddleware) +# sess = get_mcp_session(request) # sess.session_id, sess.client_name, ... diff --git a/examples/mcp_stateless_fastapi.py b/examples/mcp_stateless_fastapi.py deleted file mode 100644 index 7022ccb3..00000000 --- a/examples/mcp_stateless_fastapi.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Stateless / multi-pod MCP analytics: keep one ``$session_id`` + the client -identity across pods via a session token in the ``Mcp-Session-Id`` header. - -FastMCP + ``instrument()`` needs no extra code -- ``instrument()`` wires the mint -into the server's ``streamable_http_app()`` / ``sse_app()`` factories:: - - server = FastMCP("my-server", stateless_http=True) - instrument(server, Posthog("phc_...")) - server.run(transport="streamable-http") # or: app = server.streamable_http_app() - -This file shows the other case: a custom ``PostHogMCP`` dispatcher, where there's -no server for ``instrument()`` to wire, so you add the middleware yourself. Run:: - - POSTHOG_PROJECT_API_KEY=phc_xxx uvicorn examples.mcp_stateless_fastapi:app -""" - -import os - -from fastapi import FastAPI, Request - -from posthog.mcp import ( - PostHogMCP, - PostHogMcpStatelessSessionMiddleware, - get_mcp_session, -) - -posthog = PostHogMCP( - os.environ.get("POSTHOG_PROJECT_API_KEY", "phc_xxx"), - host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"), -) - -app = FastAPI() - -# Mints the token at `initialize` and decodes the replayed one on every request. -app.add_middleware(PostHogMcpStatelessSessionMiddleware) - - -@app.post("/mcp") -async def mcp_endpoint(request: Request): - body = await request.json() - method = body.get("method") - - # Session id + client identity recovered from the token, same across pods. - sess = get_mcp_session(request) - session_id = sess.session_id if sess else None - client_name = sess.client_name if sess else None - client_version = sess.client_version if sess else None - - if method == "initialize": - posthog.capture_initialize( - session_id=session_id, - client_name=client_name, - client_version=client_version, - parameters=body.get("params"), - ) - # ... return your real InitializeResult here ... - return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}} - - if method == "tools/call": - name = body["params"]["name"] - prepared = posthog.prepare_tool_call(name, body["params"].get("arguments")) - # ... dispatch prepared.args to your tool, then: ... - posthog.capture_tool_call( - tool_name=name, - session_id=session_id, - client_name=client_name, - client_version=client_version, - intent=prepared.intent, - intent_source=prepared.intent_source, - ) - return {"jsonrpc": "2.0", "id": body.get("id"), "result": {"content": []}} - - return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}} From 8b161ea47befd2b035f53f367ae31bcf331f62ed Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 18:42:38 +0300 Subject: [PATCH 6/7] chore(mcp): update public API snapshot for stateless session-token exports New public surface: PostHogMcpStatelessSessionMiddleware, get_mcp_session, encode_session_id, decode_session_id, read_mcp_session_header, SessionTokenPayload, MCP_SESSION_HEADER (+ posthog.mcp.asgi / session_token modules), and mcp SDK-surface version 0.1.0 -> 0.2.0. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- references/public_api_snapshot.txt | 31 +++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 848b3075..daa423ac 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -322,15 +322,29 @@ alias posthog.integrations.django.contexts -> posthog.contexts alias posthog.mcp.CaptureEventData -> posthog.mcp.types.CaptureEventData alias posthog.mcp.MCPAnalyticsContextOptions -> posthog.mcp.types.MCPAnalyticsContextOptions alias posthog.mcp.MCPAnalyticsOptions -> posthog.mcp.types.MCPAnalyticsOptions +alias posthog.mcp.MCP_SESSION_HEADER -> posthog.mcp.session_token.MCP_SESSION_HEADER alias posthog.mcp.POSTHOG_MCP_ANALYTICS_SOURCE -> posthog.mcp.constants.POSTHOG_MCP_ANALYTICS_SOURCE alias posthog.mcp.PostHogMCP -> posthog.mcp.posthog_mcp.PostHogMCP alias posthog.mcp.PostHogMCPAnalyticsEvent -> posthog.mcp.constants.PostHogMCPAnalyticsEvent alias posthog.mcp.PostHogMCPAnalyticsProperty -> posthog.mcp.constants.PostHogMCPAnalyticsProperty +alias posthog.mcp.PostHogMcpStatelessSessionMiddleware -> posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware alias posthog.mcp.PreparedToolCall -> posthog.mcp.types.PreparedToolCall +alias posthog.mcp.SessionTokenPayload -> posthog.mcp.session_token.SessionTokenPayload alias posthog.mcp.UserIdentity -> posthog.mcp.types.UserIdentity alias posthog.mcp.__version__ -> posthog.mcp.version.__version__ +alias posthog.mcp.asgi.MCP_SESSION_HEADER -> posthog.mcp.session_token.MCP_SESSION_HEADER +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.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.decode_session_id -> posthog.mcp.session_token.decode_session_id alias posthog.mcp.derive_session_id_from_mcp_session -> posthog.mcp.session.derive_session_id_from_mcp_session +alias posthog.mcp.encode_session_id -> posthog.mcp.session_token.encode_session_id +alias posthog.mcp.get_mcp_session -> posthog.mcp.asgi.get_mcp_session alias posthog.mcp.get_more_tools_result -> posthog.mcp.tools.get_more_tools_result +alias posthog.mcp.read_mcp_session_header -> posthog.mcp.session_token.read_mcp_session_header alias posthog.mcp.set_logger -> posthog.mcp.logger.set_logger alias posthog.metrics_capture.VERSION -> posthog.version.VERSION alias posthog.metrics_capture.remove_trailing_slash -> posthog.utils.remove_trailing_slash @@ -657,6 +671,7 @@ attribute posthog.integrations.django.PosthogContextMiddleware.sync_capable = Tr attribute posthog.integrations.django.PosthogContextMiddleware.tag_map = cast('Optional[Callable[[Dict[str, Any]], Dict[str, Any]]]', settings.POSTHOG_MW_TAG_MAP) attribute posthog.is_server = True attribute posthog.log_captured_exceptions = False +attribute posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware.app = app attribute posthog.mcp.constants.POSTHOG_MCP_ANALYTICS_SOURCE = 'posthog_mcp_analytics' attribute posthog.mcp.constants.PostHogMCPAnalyticsEvent.CUSTOM = '$mcp_custom' attribute posthog.mcp.constants.PostHogMCPAnalyticsEvent.EXCEPTION = '$exception' @@ -687,6 +702,11 @@ attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.SOURCE = '$mcp_sourc attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.TOOL_CATEGORY = '$mcp_tool_category' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.TOOL_DESCRIPTION = '$mcp_tool_description' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.TOOL_NAME = '$mcp_tool_name' +attribute posthog.mcp.session_token.MCP_SESSION_HEADER = 'mcp-session-id' +attribute posthog.mcp.session_token.SessionTokenPayload.client_name: Optional[str] = None +attribute posthog.mcp.session_token.SessionTokenPayload.client_version: Optional[str] = None +attribute posthog.mcp.session_token.SessionTokenPayload.protocol_version: Optional[str] = None +attribute posthog.mcp.session_token.SessionTokenPayload.session_id: str attribute posthog.mcp.types.CaptureEventData.event: str attribute posthog.mcp.types.CaptureEventData.properties: Optional[JsonRecord] = None attribute posthog.mcp.types.MCPAnalyticsContextOptions.description: Optional[str] = None @@ -707,7 +727,7 @@ attribute posthog.mcp.types.PreparedToolCall.is_missing_capability: bool = False attribute posthog.mcp.types.UserIdentity.distinct_id: str attribute posthog.mcp.types.UserIdentity.groups: Optional[Dict[str, str]] = None attribute posthog.mcp.types.UserIdentity.properties: Optional[JsonRecord] = None -attribute posthog.mcp.version.__version__ = '0.1.0' +attribute posthog.mcp.version.__version__ = '0.2.0' attribute posthog.metrics = None attribute posthog.metrics_capture.DEFAULT_HISTOGRAM_BOUNDS = [0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000] attribute posthog.metrics_capture.MetricAttributeValue = Union[str, int, float, bool] @@ -883,9 +903,11 @@ class posthog.flag_definition_cache.FlagDefinitionCacheProvider class posthog.integrations.celery.PosthogCeleryIntegration(client: Optional[Client] = None, capture_exceptions: bool = True, capture_task_lifecycle_events: bool = True, propagate_context: bool = True, task_filter: Optional[Callable[[Optional[str], dict[str, Any]], bool]] = None) class posthog.integrations.django.PosthogContextMiddleware(get_response) class posthog.mcp.McpAnalytics(key: Any) +class posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware(app: Any) 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, **kwargs: Any) +class posthog.mcp.session_token.SessionTokenPayload(session_id: str, client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None) class posthog.mcp.types.CaptureEventData(event: str, properties: Optional[JsonRecord] = None) class posthog.mcp.types.MCPAnalyticsContextOptions(description: Optional[str] = None) class posthog.mcp.types.MCPAnalyticsOptions(logger: Optional[LoggerFn] = None, report_missing: bool = False, missing_capability_tool_name: Optional[str] = None, enable_conversation_id: bool = False, enable_exception_autocapture: bool = True, context: Union[bool, MCPAnalyticsContextOptions] = True, identify: Optional[Union[IdentifyFn, UserIdentity]] = None, intent_fallback: Optional[IntentFallbackFn] = None, before_send: Optional[BeforeSendFn] = None, event_properties: Optional[EventPropertiesFn] = None) @@ -1077,9 +1099,14 @@ function posthog.identify_context(distinct_id: str) function posthog.integrations.django.markcoroutinefunction(func) function posthog.join() -> None function posthog.load_feature_flags() +function posthog.mcp.asgi.autowire_stateless_mint(server: Any) -> None +function posthog.mcp.asgi.get_mcp_session(request_or_scope: Any) -> Optional[SessionTokenPayload] function posthog.mcp.instrument(server: Any, posthog_client: Optional[Client] = None, options: Optional[MCPAnalyticsOptions] = None) -> McpAnalytics function posthog.mcp.logger.set_logger(logger: Optional[LoggerFn]) -> None function posthog.mcp.session.derive_session_id_from_mcp_session(mcp_session_id: str) -> str +function posthog.mcp.session_token.decode_session_id(value: Any) -> Optional[SessionTokenPayload] +function posthog.mcp.session_token.encode_session_id(payload: SessionTokenPayload) -> str +function posthog.mcp.session_token.read_mcp_session_header(headers: Any) -> Optional[str] function posthog.mcp.tools.get_more_tools_result() -> Dict[str, Any] function posthog.new_context(fresh: bool = False, capture_exceptions: Optional[bool] = None, client: Optional[Client] = None) function posthog.request.batch_post(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, path: str = EVENTS_ENDPOINT, **kwargs) -> requests.Response @@ -1358,10 +1385,12 @@ module posthog.integrations module posthog.integrations.celery module posthog.integrations.django module posthog.mcp +module posthog.mcp.asgi module posthog.mcp.constants module posthog.mcp.logger module posthog.mcp.posthog_mcp module posthog.mcp.session +module posthog.mcp.session_token module posthog.mcp.tools module posthog.mcp.types module posthog.mcp.version From c7ab522da4068447235492c70eafb2d8ce910d20 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 24 Jul 2026 18:59:23 +0300 Subject: [PATCH 7/7] fix(mcp): address stateless review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - session.py: never reuse a token session for a request that didn't replay the token. `data` is shared by all clients on a server instance, so the old "sticky token" fallback merged unrelated clients under one $session_id. The token session is now resolved per request; the memory fallback only persists a genuine generated session (a leftover token/mcp session can't leak into it). - _internal.py: drop the token_client_name/version/protocol_version fields — they were written during token resolution but never read (client identity is recovered per request in the adapters), so they were dead shared state. - asgi.py: bound POST body buffering to _MAX_SNIFF_BODY. Once a session-less POST exceeds the cap it can't be an initialize handshake, so we stop buffering and stream the rest straight through instead of accumulating it — prevents an unauthenticated large/streamed body from exhausting memory. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2 --- posthog/mcp/_internal.py | 6 ---- posthog/mcp/asgi.py | 40 ++++++++++++++++---------- posthog/mcp/session.py | 32 ++++++++++----------- posthog/test/mcp/test_session_token.py | 33 +++++++++++++++------ 4 files changed, 65 insertions(+), 46 deletions(-) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 39727e79..5c529541 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -62,12 +62,6 @@ class MCPAnalyticsData: session_id: str = "" session_source: str = "generated" # "generated" | "mcp" | "token" last_mcp_session_id: Optional[str] = None - # Client identity recovered from a self-encoded session token (see - # session_token.py). On a stateless pod that never processed `initialize`, - # the live `client_params` is empty, so these are the only harness source. - token_client_name: Optional[str] = None - token_client_version: Optional[str] = None - token_protocol_version: Optional[str] = None last_activity: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index 342e5dd9..2c8c5346 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -124,36 +124,46 @@ def _headers_dict(scope: Any) -> dict[str, str]: async def _buffer_body(receive: Any) -> tuple[bytes, Any]: - """Read the full request body, then return it plus a ``receive`` that replays - it downstream (so the app still sees an unconsumed stream).""" + """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. + + 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 + straight from the original ``receive`` (never accumulated). This keeps an + unauthenticated large / streamed POST from exhausting memory.""" chunks: list[bytes] = [] - more = True - while more: + total = 0 + complete = False + overflow = False + while True: message = await receive() if message.get("type") != "http.request": - # A non-body message (e.g. http.disconnect); stop and replay what we have. + # A non-body message (e.g. http.disconnect); stop with what we have. break chunks.append(message.get("body", b"") or b"") - more = message.get("more_body", False) + total += len(chunks[-1]) + if not message.get("more_body", False): + complete = True + break + if total > _MAX_SNIFF_BODY: + overflow = True # too big to be initialize — stop buffering + break - # Always buffer the whole body so replay is byte-faithful (Starlette's own - # Request.body() buffers fully too). This holds the request in memory, which - # is fine for tiny JSON-RPC POSTs on an MCP endpoint. buffered = b"".join(chunks) replayed = False async def replay() -> dict[str, Any]: - # Hand the whole buffered body back in one message, then defer to the - # original transport for anything after it (disconnect, etc.). + # 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 if not replayed: replayed = True - return {"type": "http.request", "body": buffered, "more_body": False} + return {"type": "http.request", "body": buffered, "more_body": overflow} return await receive() - # Only *sniff* (parse to detect initialize) small bodies -- a giant POST is - # never our tiny initialize handshake, so skip minting but still replay it whole. - sniff = buffered if len(buffered) <= _MAX_SNIFF_BODY else b"" + # Only sniff (parse for `initialize`) when we captured the whole small body. + sniff = buffered if (complete and not overflow) else b"" return sniff, replay diff --git a/posthog/mcp/session.py b/posthog/mcp/session.py index 0bc19005..0b5527be 100644 --- a/posthog/mcp/session.py +++ b/posthog/mcp/session.py @@ -41,28 +41,23 @@ async def resolve_session_id( ``token`` is our self-encoded session token (see :mod:`.session_token`), decoded from the replayed ``Mcp-Session-Id`` header. It is the only source that survives a stateless / multi-pod deployment, so it takes precedence. + + The token session is resolved *per request*, never sticky: ``data`` is shared + by every client hitting this server instance, so reusing a stored token session + for a request that didn't replay the token would merge unrelated clients under + one ``$session_id``. A compliant client replays the header on every request, so + a genuine token session never needs the fallback. """ async with data.session_lock: now = datetime.now(timezone.utc) if token is not None: # A token we minted at `initialize`. Its session id is already a - # `ses_...` id, so use it verbatim -- do NOT re-hash. The client - # name/version/protocol version ride along for pods that never saw - # `initialize`; stash them so adapters can fall back to them when the - # live `client_params` is absent (the stateless-pod case). + # `ses_...` id, so use it verbatim -- do NOT re-hash. (Client + # name/version are recovered per request in the adapters, not stored + # on the shared `data`, for the same cross-client reason.) data.session_id = token.session_id data.session_source = "token" - data.token_client_name = token.client_name - data.token_client_version = token.client_version - data.token_protocol_version = token.protocol_version - data.last_activity = now - return data.session_id - - # A token session, like an MCP-derived one, lives as long as the client - # replays it -- keep the id even on a request that arrives without it, so - # the session doesn't fragment. - if data.session_source == "token": data.last_activity = now return data.session_id @@ -79,10 +74,13 @@ async def resolve_session_id( data.last_activity = now return data.session_id - # Only generated sessions roll over on inactivity -- token/MCP ids live - # as long as the client replays them, and regenerating would split them. + # Memory fallback (single-owner transports like stdio). A leftover token + # session must NOT leak to a credential-less request, so anything that + # isn't already a generated session starts fresh; generated sessions + # persist and roll over on inactivity. timeout_seconds = INACTIVITY_TIMEOUT_IN_MINUTES * 60 - if (now - data.last_activity).total_seconds() > timeout_seconds: + is_stale = (now - data.last_activity).total_seconds() > timeout_seconds + if data.session_source != "generated" or is_stale: data.session_id = new_session_id() data.session_source = "generated" data.last_activity = now diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index 5ebc14dc..f0a1a234 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -125,10 +125,9 @@ async def test_resolve_session_id_uses_token_verbatim(): # Used verbatim -- NOT re-hashed through derive_session_id_from_mcp_session. assert sid == "ses_tok" assert data.session_source == "token" - assert data.token_client_name == "Cursor" -async def test_token_session_does_not_fragment_or_roll_over(): +async def test_token_session_used_every_time_and_never_rolls_over(): from datetime import datetime, timedelta, timezone data = _data() @@ -136,16 +135,32 @@ async def test_token_session_does_not_fragment_or_roll_over(): encode_session_id(SessionTokenPayload(session_id="ses_tok")) ) first = await resolve_session_id(data, None, token=token) - # A later request without the header keeps the id... - assert await resolve_session_id(data, None) == first - # ...and it never rolls over on inactivity (only generated sessions do). + # Replayed on every request (what a compliant client does) -> same id, and it + # never rolls over on inactivity (unlike a generated session). data.last_activity = datetime.now(timezone.utc) - timedelta(minutes=31) - assert await resolve_session_id(data, None) == first + assert await resolve_session_id(data, None, token=token) == first + + +async def test_token_session_not_reused_for_tokenless_request(): + """`data` is shared across clients on one server, so a request that doesn't + replay the token must NOT inherit the previous client's token session.""" + data = _data() + token = decode_session_id( + encode_session_id(SessionTokenPayload(session_id="ses_a")) + ) + a = await resolve_session_id(data, None, token=token) + assert a == "ses_a" + # A different, tokenless client hits the same server -> fresh session, not ses_a. + b = await resolve_session_id(data, None) + assert b != "ses_a" and data.session_source == "generated" async def test_multi_pod_two_instances_resolve_same_session_and_harness(): """The regression this fixes: independent pods (independent per-server state) - resolve the same replayed token to the same session id + harness.""" + resolve the same replayed token to the same session id, and the harness is + recovered from the token by the adapter helper.""" + from posthog.mcp._instrumentation import resolve_session_and_client + token_str = encode_session_id( SessionTokenPayload(session_id="ses_shared", client_name="Claude Code") ) @@ -154,7 +169,9 @@ async def test_multi_pod_two_instances_resolve_same_session_and_harness(): data = _data() token = decode_session_id(token_str) sid = await resolve_session_id(data, token_str, token=token) - results.append((sid, data.session_source, data.token_client_name)) + # Harness comes back per request in the adapters (not from shared state). + _tok, name, _ver = resolve_session_and_client(token_str, None, None) + results.append((sid, data.session_source, name)) assert results[0] == results[1] == ("ses_shared", "token", "Claude Code")