feat(mcp): stateless and multi-pod server support#761
Conversation
posthog-python Compliance ReportDate: 2026-07-24 16:09:59 UTC ✅ All Tests Passed!111/111 tests passed Capture_V1 Tests✅ 94/94 tests passed View Details
Feature_Flags Tests✅ 17/17 tests passed View Details
|
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
435e698 to
f8de485
Compare
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
… 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
Trim the docstring and inline comments; drop the "preferred" framing. Generated-By: PostHog Code Task-Id: f44ec5e0-b836-4d13-98c4-c26887dd7ec2
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
…ports 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
gesh
changed the title
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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) |
There was a problem hiding this comment.
When an unauthenticated client sends a large or indefinitely streamed POST without an Mcp-Session-Id, this loop accumulates the entire body before applying _MAX_SNIFF_BODY, causing memory exhaustion or request starvation in the MCP worker.
Prompt To Fix With AI
This is a comment left during a code review.
Path: posthog/mcp/asgi.py
Line: 130-142
Comment:
**Unbounded POST body buffering**
When an unauthenticated client sends a large or indefinitely streamed POST without an `Mcp-Session-Id`, this loop accumulates the entire body before applying `_MAX_SNIFF_BODY`, causing memory exhaustion or request starvation in the MCP worker.
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Fixed in c7ab522: _buffer_body now stops at _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 remaining chunks straight through instead of accumulating them — memory is bounded to the cap.
| if data.session_source == "token": | ||
| data.last_activity = now | ||
| return data.session_id |
There was a problem hiding this comment.
When a token-bearing request is followed on the same server by another client without that token—or with a normal transport session ID—this branch reuses the previous client's shared token session before examining the new session ID, merging unrelated events under the same $session_id.
Knowledge Base Used: MCP Instrumentation (posthog/mcp)
Prompt To Fix With AI
This is a comment left during a code review.
Path: posthog/mcp/session.py
Line: 65-67
Comment:
**Token session leaks across clients**
When a token-bearing request is followed on the same server by another client without that token—or with a normal transport session ID—this branch reuses the previous client's shared token session before examining the new session ID, merging unrelated events under the same `$session_id`.
**Knowledge Base Used:** [MCP Instrumentation (posthog/mcp)](https://app.greptile.com/posthog-org-19734/-/custom-context/knowledge-base/posthog/posthog-python/-/docs/mcp-instrumentation.md)
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Fixed in c7ab522: removed the sticky-token fallback. data is shared by every client on a server instance, so a request that doesn't replay the token no longer inherits the previous client's token session — token sessions are resolved per request, and the memory fallback only persists a genuine generated session (a leftover token/mcp session can't leak into it).
| # 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 |
There was a problem hiding this comment.
These fields are written during token session resolution but never read by event enrichment; client identity is instead backfilled through local adapter variables, while the protocol version is dropped entirely. Remove the dead shared state or wire it into the intended metadata path so it does not imply a fallback that does not exist.
Knowledge Base Used: MCP Instrumentation (posthog/mcp)
Prompt To Fix With AI
This is a comment left during a code review.
Path: posthog/mcp/_internal.py
Line: 66-70
Comment:
**Token metadata fields are unused**
These fields are written during token session resolution but never read by event enrichment; client identity is instead backfilled through local adapter variables, while the protocol version is dropped entirely. Remove the dead shared state or wire it into the intended metadata path so it does not imply a fallback that does not exist.
**Knowledge Base Used:** [MCP Instrumentation (posthog/mcp)](https://app.greptile.com/posthog-org-19734/-/custom-context/knowledge-base/posthog/posthog-python/-/docs/mcp-instrumentation.md)
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Fixed in c7ab522: removed the token_client_name/version/protocol_version fields. Client identity is recovered per request in the adapters (resolve_session_and_client), not from shared state — so these were dead. (Storing them on the shared data would have the same cross-client leak as the session finding, so removing is the right call.)
🦔 ReviewHog reviewed this pull requestFound 1 must fix, 0 should fix, 0 consider. Published 1 finding (view the review). |
- 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
pauldambra
left a comment
pauldambra
left a comment
There was a problem hiding this comment.
approval to unblock
a bunch of robot comments to address though
thanks for sorting this so quickly!
|
i'm also intrigued to see what review hog manages here @gesh |
|
ReviewHog Alpha 🦔 If you find any issues helpful - please reply "valid", "invalid", etc., for evaluation purposes 🙏 |
![posthog[bot]](https://avatars.githubusercontent.com/in/1203049?s=60&v=4){kind=small}
There was a problem hiding this comment.
ReviewHog Report
Feature
Issues: 1 issue
Files (2)
posthog/mcp/session_token.pyposthog/mcp/asgi.py
What were the main changes
- New session_token.py: unsigned, self-encoded token codec (encode_session_id/decode_session_id) wire-compatible with the TS SDK — base64url(JSON{sid,cn,cv,pv})
- New asgi.py: PostHogMcpStatelessSessionMiddleware mints the Mcp-Session-Id response header at
initialize(when the client sent none) and decodes replayed tokens on every request; get_mcp_session() accessor for custom dispatchers - autowire_stateless_mint() wraps a FastMCP server's streamable_http_app/sse_app/http_app factories so instrument() gets zero-config minting
- Review flags unbounded POST body buffering in the sniff loop before the _MAX_SNIFF_BODY cap is applied, risking memory exhaustion on large/unauthenticated POSTs
| 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(), |
There was a problem hiding this comment.
Uncaught RecursionError from deeply-nested JSON body breaks the middleware's own "never raises" / fail-safe guarantee (unauthenticated DoS)
Why we think it's a valid issue
- Checked: the
json.loadsguard in_mint_token_if_initialize(asgi.py:166-169), the call site in__call__(asgi.py:105-112), the conditions gating that path (asgi.py:82-106), the sniff cap (asgi.py:156,_MAX_SNIFF_BODY), the stated fail-safe contract (docstring 76-77, 163), and the surrounding defensive guards (asgi.py:89, 121, 185, 232). - Found (verified programmatically):
RecursionErrorsubclassesRuntimeError, notValueError/TypeError, soexcept (ValueError, TypeError)at line 168 does not catch it. A 40,001-byte balanced-bracket payload (b'['*20000 + b'0' + b']'*20000), well under the 256 KB sniff cap, makesjson.loadsraiseRecursionError, and it escapes the narrow except. - Found: line 106 (
token = _mint_token_if_initialize(body)) has no try/except, so the exception propagates out of the middleware's__call__. The path is reachable by any http POST lacking anMcp-Session-Idheader with body ≤ 256 KB (asgi.py:101-106) — no prior handshake needed. On theinstrument()auto-wire path the middleware wraps the whole app (asgi.py:242), so every route is exposed. - Found: the module's own guards (
except Exceptionat 89/121/185/232) show fail-safe is the intended invariant; this single narrow clause is the inconsistency, and the docstrings (76-77, 163) explicitly promise it never breaks the host / never raises. - Impact: an unauthenticated caller can force the analytics middleware to convert a request the host would otherwise handle into a 500, violating the module's central 'analytics must never break the host' contract, on the default zero-config path. (Impact is a per-request 500, not a pod crash — ASGI isolates the exception — so the 'DoS' framing is slightly overstated, but the contract break is real, trivially triggerable, and fixed by one line: broaden the except to include
RecursionErroror useexcept Exception.)
Issue description
_mint_token_if_initialize parses the sniffed POST body with json.loads(body) and only catches (ValueError, TypeError). A deeply-nested JSON array/object (e.g. "[" * N + "0" + "]" * N) raises RecursionError, which is a RuntimeError subclass and is NOT caught by that except clause. I verified this empirically: a balanced-bracket payload with N=20000 (≈40KB, well under the _MAX_SNIFF_BODY = 256 * 1024 cap applied just before this call) reliably raises an uncaught RecursionError from json.loads. Because the call site in __call__ (line 106, token = _mint_token_if_initialize(body)) has no surrounding try/except either, this exception propagates all the way out of the ASGI __call__ coroutine for that request. This directly contradicts the module's own stated contract: the class docstring says "Fail-safe: any error while sniffing/minting is logged and the request passes through untouched -- analytics must never break the host," and the function's own docstring says "never raises." Any unauthenticated client can trigger this with a single ~40KB POST to any route mounted behind PostHogMcpStatelessSessionMiddleware that doesn't already carry an Mcp-Session-Id header (i.e. exactly the code path meant to handle the unauthenticated initialize handshake) -- no valid MCP session or handshake is required first. This is distinct from the already-flagged "unbounded POST body buffering" comment on this file: that issue is about memory use for bodies larger than the sniff cap; this one reproduces well within the cap and would still be exploitable even if that buffering issue were fixed.
Suggested fix
Broaden the except clause to catch the same failure classes the rest of this function already guards against, e.g. except (ValueError, TypeError, RecursionError):, or simply wrap the whole body of _mint_token_if_initialize in a single except Exception (as is already done a few lines later around the encode_session_id(...) call, and around the header-read in __call__). That keeps the function's documented "never raises" contract true for any malformed/adversarial input, not just the JSON-decode error types anticipated today.
Prompt to fix with AI (copy-paste)
## Context
@posthog/mcp/asgi.py#L160-179
@posthog/mcp/asgi.py#L105-107
<issue_description>
`_mint_token_if_initialize` parses the sniffed POST body with `json.loads(body)` and only catches `(ValueError, TypeError)`. A deeply-nested JSON array/object (e.g. `"[" * N + "0" + "]" * N`) raises `RecursionError`, which is a `RuntimeError` subclass and is NOT caught by that except clause. I verified this empirically: a balanced-bracket payload with N=20000 (≈40KB, well under the `_MAX_SNIFF_BODY = 256 * 1024` cap applied just before this call) reliably raises an uncaught `RecursionError` from `json.loads`. Because the call site in `__call__` (line 106, `token = _mint_token_if_initialize(body)`) has no surrounding try/except either, this exception propagates all the way out of the ASGI `__call__` coroutine for that request. This directly contradicts the module's own stated contract: the class docstring says "Fail-safe: any error while sniffing/minting is logged and the request passes through untouched -- analytics must never break the host," and the function's own docstring says "never raises." Any unauthenticated client can trigger this with a single ~40KB POST to any route mounted behind `PostHogMcpStatelessSessionMiddleware` that doesn't already carry an `Mcp-Session-Id` header (i.e. exactly the code path meant to handle the unauthenticated `initialize` handshake) -- no valid MCP session or handshake is required first. This is distinct from the already-flagged "unbounded POST body buffering" comment on this file: that issue is about memory use for bodies *larger* than the sniff cap; this one reproduces well *within* the cap and would still be exploitable even if that buffering issue were fixed.
</issue_description>
<issue_validation>
- **Checked:** the `json.loads` guard in `_mint_token_if_initialize` (asgi.py:166-169), the call site in `__call__` (asgi.py:105-112), the conditions gating that path (asgi.py:82-106), the sniff cap (asgi.py:156, `_MAX_SNIFF_BODY`), the stated fail-safe contract (docstring 76-77, 163), and the surrounding defensive guards (asgi.py:89, 121, 185, 232).
- **Found (verified programmatically):** `RecursionError` subclasses `RuntimeError`, not `ValueError`/`TypeError`, so `except (ValueError, TypeError)` at line 168 does not catch it. A 40,001-byte balanced-bracket payload (`b'['*20000 + b'0' + b']'*20000`), well under the 256 KB sniff cap, makes `json.loads` raise `RecursionError`, and it escapes the narrow except.
- **Found:** line 106 (`token = _mint_token_if_initialize(body)`) has no try/except, so the exception propagates out of the middleware's `__call__`. The path is reachable by any http POST lacking an `Mcp-Session-Id` header with body ≤ 256 KB (asgi.py:101-106) — no prior handshake needed. On the `instrument()` auto-wire path the middleware wraps the whole app (asgi.py:242), so every route is exposed.
- **Found:** the module's own guards (`except Exception` at 89/121/185/232) show fail-safe is the intended invariant; this single narrow clause is the inconsistency, and the docstrings (76-77, 163) explicitly promise it never breaks the host / never raises.
- **Impact:** an unauthenticated caller can force the analytics middleware to convert a request the host would otherwise handle into a 500, violating the module's central 'analytics must never break the host' contract, on the default zero-config path. (Impact is a per-request 500, not a pod crash — ASGI isolates the exception — so the 'DoS' framing is slightly overstated, but the contract break is real, trivially triggerable, and fixed by one line: broaden the except to include `RecursionError` or use `except Exception`.)
</issue_validation>
## Task
Investigate the issue and solve it
<potential_solution>
Broaden the except clause to catch the same failure classes the rest of this function already guards against, e.g. `except (ValueError, TypeError, RecursionError):`, or simply wrap the whole body of `_mint_token_if_initialize` in a single `except Exception` (as is already done a few lines later around the `encode_session_id(...)` call, and around the header-read in `__call__`). That keeps the function's documented "never raises" contract true for any malformed/adversarial input, not just the JSON-decode error types anticipated today.
</potential_solution>
There was a problem hiding this comment.
Pulled this down and ran the full mcp suite, 133 passing, including the stateless FastMCP end to end test and the zero config autowire path, and the public API snapshot check is clean.
Traced instrument() through the autowire into both adapters and the resolved session id is threaded as a local all the way to record_tool_call, so two requests sharing one pod can't stamp each other's session.
A few non blocking thoughts below, none of them need to hold the merge.
| 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``), |
There was a problem hiding this comment.
The whole point of this token is that it decodes in both this SDK and the TypeScript one (same sid/cn/cv/pv keys, base64url). But nothing in the suite actually protects that. If either side renames a key or changes the encoding, every test still passes and the two SDKs quietly stop reading each other's tokens. Consider one fixture test: a token string minted by the TS SDK, hardcoded here, asserted to decode into the expected fields, and ideally the reverse with a known Python token the TS side must read. That turns the compatibility claim into something CI defends.
| 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 |
There was a problem hiding this comment.
The docstring says any error while sniffing or minting is logged and the request passes through untouched, but only the header read is wrapped in try/except. If _buffer_body's receive() raises, or anything in the mint block throws, it propagates to the host instead of passing through. In practice a failing receive() is already a dead request so the blast radius is small, but the promise reads broader than what the code guarantees. Either tighten the wording or wrap the POST mint block in the same try/except so the failsafe matches what it claims.
| 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 ( |
There was a problem hiding this comment.
Seven new public names land here: the middleware, get_mcp_session, encode_session_id, decode_session_id, read_mcp_session_header, SessionTokenPayload, MCP_SESSION_HEADER. The middleware and get_mcp_session clearly need to be public for the custom dispatcher story. I'm less sure encode_session_id and read_mcp_session_header need to be public on day one versus staying internal until a real custom HTTP layer needs them. Not a blocker, just worth a second look since public surface is easy to add and hard to walk back.
|
|
||
| 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"): |
There was a problem hiding this comment.
On the bundled MCP SDK FastMCP this is fine: http_app doesn't exist and the other two build distinct apps. On fastmcp 2.x though, streamable_http_app and sse_app are often thin aliases over http_app, so wrapping all three could add the middleware twice to one app. It stays correct thanks to the no clobber check on the response header and decode being idempotent, but you'd mint two session ids and buffer the body twice per initialize. Worth confirming against fastmcp 2.x, or wrapping only the factory that actually builds the app.
|
The description points at examples/mcp_stateless_fastapi.py but the file that landed is examples/mcp_stateless.py. Worth fixing so nobody hunts for the wrong path. |
3 participants
Problem
Stateless / multi-pod MCP servers issue no session id. So
$session_idfragments (each pod generates its own) and the client identity — the "harness" (Claude Code, Cursor, …), sent only atinitialize— is lost on any pod that never processed the handshake. Works in single-pod staging, silently wrong in multi-pod production.Fix
Self-encoded session token (wire-compatible with the
@posthog/mcpTypeScript SDK). Atinitializewe mint theMcp-Session-Idresponse header asbase64url(JSON{sid,cn,cv,pv}); the client replays it on every request, so any pod recovers the same session id + client name/version with no shared state.Mint and decode are split because Python's MCP SDK owns
initializein its runner layer and forbids overriding it viarequest_handlers(unlike the TS SDK) — so there's no in-SDK seam to mint from:instrument()FastMCP path this is zero-config —instrument()auto-wires the mint into the server's ASGI-app factories (streamable_http_app/sse_app, whichmcp.run()also uses), so nothing to add. For a customPostHogMCPdispatcher (you own the ASGI app), addPostHogMcpStatelessSessionMiddlewareonce. Either way, no manual header handling.instrument()adapters decode the replayed header and backfill client identity;PostHogMCPusers readget_mcp_session(request).New:
posthog/mcp/session_token.py(codec),posthog/mcp/asgi.py(middleware,get_mcp_session, and theinstrument()auto-wire). Touched:session.py/_internal.py(a"token"session source, used verbatim + sticky, never rolls over) and the twoinstrument()adapters. Example:examples/mcp_stateless_fastapi.py.Validation
Unit + real stateless-FastMCP integration tests (
posthog/test/mcp/test_session_token.py), including the zero-configinstrument()auto-wire. Also verified against a 2-pod stateless cluster behind round-robin nginx: withinstrument()alone, all events across both pods share one$session_idand keep the harness — while the published package fragments both.Created with PostHog Code