Repository navigation
fix(mcp): send one $mcp_initialize per initialize request - #1103
Conversation
The Python MCP SDK answers initialize itself, so posthog.mcp sent the event on the first request of each session, once per process. A stateless fleet sent one event per pod. A client that sent only initialize sent no event. Now MCP SDK 2.x reports the event from Server.middleware, on every transport. MCP SDK 1.x over HTTP reports it from PostHogMcpStatelessSessionMiddleware, which takes a new optional server= argument. MCP SDK 1.x over stdio keeps the first-request fallback. Co-Authored-By: Claude Opus 5.5 <[email protected]>
posthog-python Compliance ReportDate: 2026-10-23T12:02:36.488961+00:00 ✅ All Tests Passed!121/121 tests passed Capture_V1 Tests✅ 95/95 tests passed View Details
Capture_Ai Tests✅ 5/5 tests passed View Details
Feature_Flags Tests✅ 17/17 tests passed View Details
Feature_Flags_Local_Evaluation Tests✅ 4/4 tests passed View Details
|
![posthog[bot]](https://avatars.githubusercontent.com/in/1203049?s=60&v=4){kind=small}
There was a problem hiding this comment.
PostHog Review · Standard · reviewed 57d8dc6: 3 Should fix.
This is the only automatic review while this pull request is a draft. Comment @posthog review to review new changes sooner. Otherwise, a review of the new changes runs when you mark it ready for review.
The mcp 1.x hook now keeps only the status and the Mcp-Session-Id header of the initialize response. It no longer parses the JSON or SSE body. The event uses the protocol version that the client asked for. record_initialize now reads the client name, version, and protocol from the initialize params and reply, so both hooks pass them through. The new tests share their request helpers and drop the SSE body case. Co-Authored-By: Claude Opus 5.5 <[email protected]>
- SSE transport: the initialize event now lands in the session that the SSE calls use. On SDK 1.x the middleware leaves SSE requests to the first-request fallback. On SDK 2.x the hook uses the process session. - The bound middleware marks each request that it passes on. The fallback now stops only for marked requests, not for every HTTP request. - The hooks compute the handshake session without a write to the shared session that McpAnalytics.capture() reads. - A server= value that cannot be bound no longer breaks requests. - Remove the unused _mint_token_if_initialize. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Both handshake hooks run record_initialize in a tracked background task. It started the capture as a second task, after flush() took its snapshot of pending tasks. With a slow before_send, flush() returned early and shutdown dropped $mcp_initialize. The handshake task now awaits the capture itself. Co-Authored-By: Claude Opus 5.5 <[email protected]>
On MCP SDK 2.x, a stateless client can send initialize with an existing Mcp-Session-Id and keep sending it. Its later calls resolve their session from that header, but the handshake hook read only the connection id and the minted token. The event landed in a generated session. The hook now falls back to the incoming header. Co-Authored-By: Claude Opus 5.5 <[email protected]>
On MCP SDK 1.x, the middleware read the body only of a POST with no Mcp-Session-Id, because only those need a token. A client that sent initialize with an existing session id got no event: the middleware marked the request, so the first-request fallback did not run either. A middleware bound to a server now reads those bodies too, behind a cheap check for "initialize". It still mints only without a header. The event uses the header that the client keeps sending. Co-Authored-By: Claude Opus 5.5 <[email protected]>
PostHog Review · Standard · couldn't review
|
PostHog Review · Standard · couldn't review
|
When the client disconnected while it sent the body, the middleware stopped its read and replayed the partial body as a complete request. The app parsed half a request and never saw the disconnect. The middleware now replays the partial body with more_body set, then hands the app the disconnect. Co-Authored-By: Claude Opus 5.5 <[email protected]>
A bound middleware reads the body of a POST that carries a session id, to catch an initialize that reuses one. That held every tool call in memory, up to 256 KB, before the app saw it. The middleware now stops after 4 KB when "initialize" is not in it, and streams the rest. JSON-RPC clients put "method" before "params", so a handshake shows up in that window. Co-Authored-By: Claude Opus 5.5 <[email protected]>
The v2 handshake hook read the Mcp-Session-Id header and decoded the token itself. Every later request gets the same pair from _resolve_ctx. The hook now calls it, so the handshake and later requests follow one rule and cannot drift apart into two sessions. Co-Authored-By: Claude Opus 5.5 <[email protected]>
dustinbyrne
left a comment
dustinbyrne
left a comment
There was a problem hiding this comment.
The handshake-count fixes are well covered and reproduce correctly against the base. One optional resource-safety follow-up is noted inline.
AI-assisted review.
| return await call_next(ctx) | ||
| started = datetime.now(timezone.utc) | ||
| result = await call_next(ctx) | ||
| fire_and_forget(_record_v2_initialize(data, ctx, result, started), data) |
There was a problem hiding this comment.
[P2] Optional follow-up: bound detached handshake recorder admission
Consider adding a finite, nonblocking per-server admission cap before creating handshake recorder tasks, including the equivalent v1 HTTP path in asgi.py:186–194.
This matters when a configured async identify or event_properties callback stays pending—for example, an identity lookup using httpx.AsyncClient(timeout=None) against a service that accepts the connection but never replies. Each successful handshake can then leave another analytics task waiting after the request finishes and the client disconnects.
With the same gated-callback test on base e17a8be and head bb21c6a, both MCP majors retained 0 tasks on base versus 16, 32, and 1,056 on head after that many completed handshakes. Head flush waited for callback release; releasing the gate drained everything to zero. The regression used an illustrative 1,024-task threshold, not a specification-prescribed number; unconditional scheduling in source establishes the missing bound.
The MCP observer-only contract explicitly calls for count-bounded pending state. Sensible caller-side timeouts and cached lookups substantially mitigate this, though, so I’m treating it as optional hardening rather than a merge requirement. This is not an ordinary-operation leak or a demonstrated production OOM.
A saturation regression should exceed the chosen cap, verify successful unchanged handshake responses and bounded work after disconnect, then verify recording resumes once callbacks finish. An unlimited semaphore wait queue would not bound retained tasks. An admission cap also would not, by itself, guarantee flush termination for a callback that never returns; callback-deadline policy can remain separate.
There was a problem hiding this comment.
kind of a stretch, fwiw, given this is only truly unbounded if the caller allows an infinite timeout period - i left this in considering the reference to sdk-specs, but it seems pretty benign to me overall
There was a problem hiding this comment.
Thanks, good catch. The detached recorder was the real problem, so I removed it instead of capping it.
20c1243 runs identify and event_properties inline in the handshake request on both hooks, the same as on a tool call. A hung callback now holds that one request, where the host sees it, and leaves no analytics task behind. The capture still goes out in the background afterwards.
test_a_hung_identify_leaves_no_task_behind runs your scenario on all three server types and both SDK majors. It sends 5 handshakes against an identify that never returns, then asserts 0 pending tasks and that flush() returns. After the callback recovers, the SDK records the next handshake. On bb21c6a the test fails on the pending recorder tasks.
The trade-off: on both SDK majors the handshake reply now waits for identify, as a tool call already does. Since 132eaef the v1 middleware records it when the response starts, inside the request context.
Both handshake hooks ran identify and event_properties in a detached background task. If a callback never returned, each handshake left one more task pending, with no bound, and flush() waited for all of them. The spec requires pending state to be bounded by count. The hooks now run the callbacks inline, the same as a tool call does. A hung callback holds that request, where the host sees it, and leaves no task behind. The capture still runs in the background after the callbacks, as for every other event. Co-Authored-By: Claude Opus 5.5 <[email protected]>
There was a problem hiding this comment.
PostHog Review · Standard · reviewed 20c1243: 2 Should fix.
This is review round 3 on this pull request. When findings keep coming back in the same code, the approach is usually the problem. Before the next fix, check whether a different design would remove the whole class of findings.
The session middleware sits outside FastMCP's AuthContextMiddleware. It ran identify after the app returned, when that middleware had already reset the authenticated user. An identify that calls get_access_token() got None, so $mcp_initialize stayed anonymous under OAuth. The middleware now records the handshake when the response starts, while the app still handles the request. Co-Authored-By: Claude Opus 5.5 <[email protected]>
The 4 KB peek stopped when "initialize" was not in the first 4 KB. JSON key order is arbitrary, so an initialize with large params before its method got no event. The peek now stops only when the prefix shows a "method" and no "initialize". Without a method yet, it reads up to the full cap. Co-Authored-By: Claude Opus 5.5 <[email protected]>
|
already approved so wont approve again, left 2 new comments tho, address them if they make sense |
marandaneto
left a comment
marandaneto
left a comment
There was a problem hiding this comment.
Two reproduced blocking findings from the review. The complete patch was revalidated before posting.
| if not peeked and total >= _PEEK_BODY: | ||
| peeked = True | ||
| prefix = b"".join(chunks) | ||
| if b'"method"' in prefix and b'"initialize"' not in prefix: |
There was a problem hiding this comment.
blocking: A nested method can still hide the actual initialize method. The latest cutoff handles a missing method in the prefix, but it also matches nested keys: an initialize with an existing session header and params first can contain capabilities.experimental.vendor = {"method": "custom", "padding": ...} before its top-level method. With over 4 KB of parameters, buffering stops at that nested key, so SDK 1.x successfully answers the handshake but emits no $mcp_initialize; the bound middleware also suppresses the fallback. Only the top-level RPC method can establish that the request is not initialize. Reproduction: reproduced — test_nested_method_before_top_level_initialize_is_counted received HTTP 200 with a successful result but zero initialize events on the reviewed head; disabling the premature cutoff made it pass.
There was a problem hiding this comment.
Good catch, thanks. Both of your findings came from the same byte checks, so 0773e32 removes them instead of patching them again.
On SDK 1.x the middleware no longer reads requests that carry an Mcp-Session-Id. Nothing guesses at the method from raw bytes now. The cost is that on SDK 1.x an initialize that reuses an existing session id gets no event. The README lists that limit. SDK 2.x still counts it, because Server.middleware gets the parsed method from the SDK.
Your case is now test_a_nested_method_before_the_top_level_initialize_is_counted. It runs on SDK 2.x for all three server types.
| """The JSON-RPC message in ``body`` if it is an ``initialize`` request, else | ||
| ``None`` (never raises).""" | ||
| # Cheap first: with a bound middleware this runs on every tool call too. | ||
| if not body or b'"initialize"' not in body: |
There was a problem hiding this comment.
blocking: Preserve JSON-escaped initialize methods. The valid JSON method string "\u0069nitialize" decodes to initialize, but this new raw-byte filter rejects it before JSON parsing. This regresses session-token minting even for unbound middleware, losing cross-pod session correlation, and prevents bound SDK 1.x middleware from observing the handshake. Reproduction: reproduced — test_escaped_initialize_still_mints_a_session returned HTTP 200 without Mcp-Session-Id on the reviewed head under both SDK majors, while the identical test passed on the merge base.
There was a problem hiding this comment.
Agreed, that one broke behavior that worked on main. Fixed in 0773e32: the middleware parses the body as JSON again before it checks the method, as main does, so "initialize" gets its token and its event.
Your case is now test_an_escaped_initialize_still_mints_a_session, on both SDK majors for all three server types. It fails on 23a0677.
The middleware guessed from raw bytes whether a request was initialize. A nested "method" key in params stopped the 4 KB peek too early, and a JSON-escaped "initialize" failed the quick byte check. That second case also stopped the token mint, which worked before this PR. The middleware now parses the body as JSON, as main does, and reads only requests with no Mcp-Session-Id. On MCP SDK 1.x, an initialize that reuses an existing session id is not captured. The README lists this limit. MCP SDK 2.x still captures it. Co-Authored-By: Claude Opus 5.5 <[email protected]>
💡 Motivation and Context
initializeitself, soposthog.mcpcannot wrap it. Until now, it sent$mcp_initializeon the first request of each session, once per process.initialize.$mcp_initializeonly for aninitializethat the client sent.Changes
The SDK now sends one
$mcp_initializefor eachinitializerequest, in the session that the client's later requests use.Server.middleware. The SDK hands it the parsedinitialize.PostHogMcpStatelessSessionMiddleware. It parses each request with noMcp-Session-Idas JSON, as it already did to mint the token. It records the event when the response starts.instrument()binds the middleware on the apps that it autowires. If you mount the middleware by hand, pass the new optionalserver=argument. Without it, each pod still sends the event on its first request.identifyandevent_propertiesrun inline in the handshake request, the same as on a tool call. The capture runs in the background, andanalytics.flush()waits for it.$mcp_parameterswith the handshake request. On MCP SDK 2.x, it also has$mcp_responsewith the server reply.conversation_idanchors no longer gets its own$mcp_initialize.Check in review
ServerSession._connection. A guard handles it, and the stateful tests fail if the SDK renames it.Mcp-Session-Id. So the event has the protocol version that the client asked for, and aninitializethat fails with a JSON-RPC error still counts.extra["ctx"]for the handshake callbacks has only.request.get_request_headers(extra)works as before.$mcp_initializecount goes down for stateless Python fleets. The default MCP dashboard counts users and sessions from this event.Known limits
server=keeps the per-pod event.initializethat reuses an existingMcp-Session-Idgets no event. The middleware does not read requests that carry one, so tool calls stay unread.initializeand nothing else gets no event.conversation_idgets no event.💚 How did you test it?
test_initialize_handshake.pydrives real apps with raw JSON-RPC and counts the events. Two pods with round-robin requests stand in for a fleet.Serveron both majors. It also runs real stdio and SSE servers, and aPostHogMCPdispatcher.identify, a hungidentify,flush(), an escaped or nested method, and the 2026-07-28 spec.main, 23 of the new tests fail on SDK 1.x and 28 fail on SDK 2.x.mcp1.30.0 withfastmcp3.x, and SDK 2.x runsmcp2.3.0 withfastmcp4.1.0.0773e32with real uvicorn servers and a client that reconnects 3 times. The table shows the$mcp_initializecount. Themaincolumn is posthog 7.75.4 on SDK 1.x.tools/listtools/listtools/list, and 2 tool callsOn this PR, the counts are the same for stateful and stateless servers, on both SDK majors.
📝 Checklist
If releasing new changes
sampo addto generate a changeset file🤖 Agent context
Autonomy: Human-driven (agent-assisted)
claude-opus-5-5).initialize. SDK 1.x has no such hook, so the ASGI middleware observes the handshake over HTTP, and stdio keeps the first-request fallback.🤖 Generated with Claude Code