feat(mcp): stateless and multi-pod server support

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

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, ...

posthog/mcp/__init__.py

@@ -44,6 +44,18 @@
 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,
+    autowire_stateless_mint,
+    get_mcp_session,
+)
 from ._sink import McpEventSink
 from .tools import get_more_tools_result
 from .types import (
@@ -66,6 +78,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",
@@ -227,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}")

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()

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()

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))

posthog/mcp/_internal.py

@@ -60,7 +60,7 @@ 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
     last_activity: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
     identified_sessions: IdentityCache = field(default_factory=IdentityCache)