Skip to content

Commit f8de485

Browse files
committed
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
1 parent c2604d7 commit f8de485

11 files changed

Lines changed: 930 additions & 5 deletions

examples/mcp_stateless_fastapi.py

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
"""Stateless / multi-pod MCP analytics with the ``PostHogMCP`` custom dispatcher.
2+
3+
A stateless MCP server issues no session id, so across pods (or per-request
4+
transports) ``$session_id`` fragments and the client identity (the "harness",
5+
e.g. Claude Code / Cursor) -- sent only at ``initialize`` -- is lost on any pod
6+
that never processed the handshake.
7+
8+
The fix is a self-encoded session token minted onto the ``Mcp-Session-Id``
9+
response header at ``initialize`` and replayed by the client on every request.
10+
You do NOT set the header by hand: add
11+
:class:`~posthog.mcp.PostHogMcpStatelessSessionMiddleware` once, then read the
12+
recovered session with :func:`~posthog.mcp.get_mcp_session` and pass it into the
13+
capture calls.
14+
15+
Usage::
16+
17+
POSTHOG_PROJECT_API_KEY=phc_xxx uvicorn examples.mcp_stateless_fastapi:app
18+
19+
The same one-line middleware also works in front of a mounted FastMCP app -- see
20+
the note at the bottom.
21+
"""
22+
23+
import os
24+
25+
from fastapi import FastAPI, Request
26+
27+
from posthog.mcp import (
28+
PostHogMCP,
29+
PostHogMcpStatelessSessionMiddleware,
30+
get_mcp_session,
31+
)
32+
33+
posthog = PostHogMCP(
34+
os.environ.get("POSTHOG_PROJECT_API_KEY", "phc_xxx"),
35+
host=os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com"),
36+
)
37+
38+
app = FastAPI()
39+
40+
# One line. The middleware mints the session token onto the `Mcp-Session-Id`
41+
# response header at `initialize` (when the client sent none) and decodes the
42+
# replayed token on every later request. No manual header handling anywhere.
43+
app.add_middleware(PostHogMcpStatelessSessionMiddleware)
44+
45+
46+
@app.post("/mcp")
47+
async def mcp_endpoint(request: Request):
48+
body = await request.json()
49+
method = body.get("method")
50+
51+
# Recovered by the middleware from the replayed token. On the very first
52+
# `initialize` it reflects the token just minted; on every later request
53+
# (any pod) it carries the same session id + harness.
54+
sess = get_mcp_session(request)
55+
session_id = sess.session_id if sess else None
56+
client_name = sess.client_name if sess else None
57+
client_version = sess.client_version if sess else None
58+
59+
if method == "initialize":
60+
posthog.capture_initialize(
61+
session_id=session_id,
62+
client_name=client_name,
63+
client_version=client_version,
64+
parameters=body.get("params"),
65+
)
66+
# ... return your real InitializeResult here ...
67+
return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}}
68+
69+
if method == "tools/call":
70+
name = body["params"]["name"]
71+
prepared = posthog.prepare_tool_call(name, body["params"].get("arguments"))
72+
# ... dispatch prepared.args to your tool, then: ...
73+
posthog.capture_tool_call(
74+
tool_name=name,
75+
session_id=session_id,
76+
client_name=client_name,
77+
client_version=client_version,
78+
intent=prepared.intent,
79+
intent_source=prepared.intent_source,
80+
)
81+
return {"jsonrpc": "2.0", "id": body.get("id"), "result": {"content": []}}
82+
83+
return {"jsonrpc": "2.0", "id": body.get("id"), "result": {}}
84+
85+
86+
# Mounted-FastMCP variant: the exact same middleware works in front of a FastMCP
87+
# streamable-HTTP app, and instrument() reads the replayed token automatically --
88+
# no per-request code needed there:
89+
#
90+
# from mcp.server.fastmcp import FastMCP
91+
# from posthog import Posthog
92+
# from posthog.mcp import instrument, PostHogMcpStatelessSessionMiddleware
93+
#
94+
# server = FastMCP("my-server", stateless_http=True, json_response=True)
95+
# instrument(server, Posthog("phc_xxx"))
96+
# app = server.streamable_http_app()
97+
# app.add_middleware(PostHogMcpStatelessSessionMiddleware)

posthog/mcp/__init__.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,14 @@
4444
from .logger import log, set_logger
4545
from .posthog_mcp import PostHogMCP
4646
from .session import derive_session_id_from_mcp_session, new_session_id
47+
from .session_token import (
48+
MCP_SESSION_HEADER,
49+
SessionTokenPayload,
50+
decode_session_id,
51+
encode_session_id,
52+
read_mcp_session_header,
53+
)
54+
from .asgi import PostHogMcpStatelessSessionMiddleware, get_mcp_session
4755
from ._sink import McpEventSink
4856
from .tools import get_more_tools_result
4957
from .types import (
@@ -66,6 +74,16 @@
6674
"PreparedToolCall",
6775
"get_more_tools_result",
6876
"derive_session_id_from_mcp_session",
77+
# Self-encoded session tokens for stateless / multi-pod servers. Minted onto
78+
# the `Mcp-Session-Id` response header by PostHogMcpStatelessSessionMiddleware
79+
# and decoded on every request; codec is exported for custom HTTP layers.
80+
"PostHogMcpStatelessSessionMiddleware",
81+
"get_mcp_session",
82+
"encode_session_id",
83+
"decode_session_id",
84+
"read_mcp_session_header",
85+
"SessionTokenPayload",
86+
"MCP_SESSION_HEADER",
6987
"set_logger",
7088
"POSTHOG_MCP_ANALYTICS_SOURCE",
7189
"PostHogMCPAnalyticsEvent",

posthog/mcp/_instrument_fastmcp.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
record_tool_call,
4848
record_tools_list,
4949
request_to_dict,
50+
resolve_session_and_client,
5051
)
5152
from ._internal import MCPAnalyticsData
5253
from .logger import log
@@ -91,6 +92,9 @@ async def wrapped(
9192
) -> Any:
9293
client_name, client_version = _client_info(context)
9394
mcp_session_id = _mcp_session_id(context)
95+
token, client_name, client_version = resolve_session_and_client(
96+
mcp_session_id, client_name, client_version
97+
)
9498
request = build_tool_call_request(name, arguments)
9599
extra: Dict[str, Any] = {"session_id": mcp_session_id}
96100

@@ -101,6 +105,7 @@ async def wrapped(
101105
client_version=client_version,
102106
request=request,
103107
extra=extra,
108+
token=token,
104109
)
105110

106111
missing_name = resolve_missing_capability_tool_name(data.options)
@@ -211,6 +216,9 @@ async def list_handler(req: Any) -> Any:
211216

212217
client_name, client_version = _low_level_client_info(server)
213218
mcp_session_id = _low_level_session_id(server)
219+
token, client_name, client_version = resolve_session_and_client(
220+
mcp_session_id, client_name, client_version
221+
)
214222
request = request_to_dict(req)
215223
extra: Dict[str, Any] = {"session_id": mcp_session_id}
216224
# Resolve session, emit $mcp_initialize (once per session) and identify here
@@ -222,6 +230,7 @@ async def list_handler(req: Any) -> Any:
222230
client_version=client_version,
223231
request=request,
224232
extra=extra,
233+
token=token,
225234
)
226235

227236
start = time.monotonic()

posthog/mcp/_instrument_lowlevel.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
record_tool_call,
4242
record_tools_list,
4343
request_to_dict,
44+
resolve_session_and_client,
4445
)
4546
from ._internal import MCPAnalyticsData
4647
from .logger import log
@@ -94,6 +95,9 @@ async def handler(req: Any) -> Any:
9495
arguments = dict(req.params.arguments or {})
9596
client_name, client_version = _client_info(server)
9697
mcp_session_id = _mcp_session_id(server)
98+
token, client_name, client_version = resolve_session_and_client(
99+
mcp_session_id, client_name, client_version
100+
)
97101
request = build_tool_call_request(name, arguments)
98102
extra = {"session_id": mcp_session_id}
99103

@@ -104,6 +108,7 @@ async def handler(req: Any) -> Any:
104108
client_version=client_version,
105109
request=request,
106110
extra=extra,
111+
token=token,
107112
)
108113

109114
missing_name = resolve_missing_capability_tool_name(data.options)
@@ -223,6 +228,9 @@ async def handler(req: Any) -> Any:
223228

224229
client_name, client_version = _client_info(server)
225230
mcp_session_id = _mcp_session_id(server)
231+
token, client_name, client_version = resolve_session_and_client(
232+
mcp_session_id, client_name, client_version
233+
)
226234
request = request_to_dict(req)
227235
extra = {"session_id": mcp_session_id}
228236
# Resolve session, emit $mcp_initialize (once per session) and identify here
@@ -234,6 +242,7 @@ async def handler(req: Any) -> Any:
234242
client_version=client_version,
235243
request=request,
236244
extra=extra,
245+
token=token,
237246
)
238247

239248
start = time.monotonic()

posthog/mcp/_instrumentation.py

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from .logger import log
2424
from ._sanitization import build_captured_mcp_parameters
2525
from .session import resolve_session_id
26+
from .session_token import SessionTokenPayload, decode_session_id
2627

2728
# Keep strong refs to in-flight capture tasks/futures so they aren't GC'd mid-flight,
2829
# and so the asyncio ones can be awaited via drain_pending() before shutdown. Holds
@@ -183,6 +184,24 @@ async def _apply_event_properties(
183184
event["properties"] = props
184185

185186

187+
def resolve_session_and_client(
188+
raw_session_id: Optional[str],
189+
client_name: Optional[str],
190+
client_version: Optional[str],
191+
) -> tuple[Optional[SessionTokenPayload], Optional[str], Optional[str]]:
192+
"""Decode a replayed ``Mcp-Session-Id`` value as a self-encoded session token,
193+
and backfill the client name/version from it when the live transport supplied
194+
none (the stateless-pod case, where ``initialize`` was never seen here).
195+
196+
Returns ``(token, client_name, client_version)``; ``token`` is ``None`` when the
197+
header isn't one of our tokens (a plain transport UUID, JWT, or nothing)."""
198+
token = decode_session_id(raw_session_id)
199+
if token is not None:
200+
client_name = client_name or token.client_name
201+
client_version = client_version or token.client_version
202+
return token, client_name, client_version
203+
204+
186205
async def prepare_request(
187206
data: MCPAnalyticsData,
188207
*,
@@ -191,16 +210,21 @@ async def prepare_request(
191210
client_version: Optional[str],
192211
request: Dict[str, Any],
193212
extra: Optional[Dict[str, Any]],
213+
token: Optional[SessionTokenPayload] = None,
194214
) -> str:
195215
"""Resolve the session id, run identify, then lazily emit initialize. Returns
196216
the session id to stamp on the event for this request.
197217
218+
``token`` is the decoded self-encoded session token (see ``session_token.py``);
219+
when present it takes precedence over ``mcp_session_id`` and carries the client
220+
identity across stateless pods.
221+
198222
Identify runs *before* initialize so the resolved identity is already in the cache
199223
when ``capture_event`` builds the initialize event — otherwise the first
200224
``$mcp_initialize`` is anonymous even when identify resolves on the same request.
201225
(Still not byte-parity with the TS SDK, which wraps the real initialize handler;
202226
the Python SDK handles initialize in the session layer, not ``request_handlers``.)"""
203-
session_id = await resolve_session_id(data, mcp_session_id)
227+
session_id = await resolve_session_id(data, mcp_session_id, token=token)
204228
identify_event = await handle_identify(data, session_id, request, extra)
205229
if identify_event:
206230
fire_and_forget(capture_event(data, identify_event))

posthog/mcp/_internal.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,14 @@ class MCPAnalyticsData:
6060
options: MCPAnalyticsOptions
6161
sink: Optional[McpEventSink] = None
6262
session_id: str = ""
63-
session_source: str = "generated" # "generated" | "mcp"
63+
session_source: str = "generated" # "generated" | "mcp" | "token"
6464
last_mcp_session_id: Optional[str] = None
65+
# Client identity recovered from a self-encoded session token (see
66+
# session_token.py). On a stateless pod that never processed `initialize`,
67+
# the live `client_params` is empty, so these are the only harness source.
68+
token_client_name: Optional[str] = None
69+
token_client_version: Optional[str] = None
70+
token_protocol_version: Optional[str] = None
6571
last_activity: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
6672
identified_sessions: IdentityCache = field(default_factory=IdentityCache)
6773
tool_categories: Dict[str, str] = field(default_factory=dict)

0 commit comments

Comments
 (0)