fix(core): use sessionId for MCP transport correlation#19172
fix(core): use sessionId for MCP transport correlation#19172
Conversation
…n spans This introduces `recordInputs` and `recordOutputs` options to the MCP server wrapper and related functions, allowing for more granular control over the data captured in spans.
…to avoid breaking changes
…t input/output handlin
Wrapper transports (like NodeStreamableHTTPServerTransport wrapping WebStandardStreamableHTTPServerTransport) proxy onmessage/send via getters/setters, causing different 'this' values in each callback. This changes correlation from WeakMap<transport> to Map<sessionId> to correctly correlate requests with responses regardless of which transport object reference is used. Reported-by: @gotenxds
Add mock wrapper transport utility and tests verifying that instrumentation correctly handles transports that proxy onmessage/send via getters/setters (common pattern with StreamableHTTP transports).
Add E2E tests for MCP StreamableHTTP transport across node-express, node-express-v5, and tsx-express test applications. Tests verify that transactions are recorded correctly for the wrapper transport pattern used by NodeStreamableHTTPServerTransport.
Codecov Results 📊Generated by Codecov Action |
Codecov Results 📊✅ 23 passed | ⏭️ 7 skipped | Total: 30 | Pass Rate: 76.67% | Execution Time: 10.04s All tests are passing successfully. Generated by Codecov Action |
node-overhead report 🧳Note: This is a synthetic benchmark with a minimal express app and does not necessarily reflect the real-world performance impact in an application.
|
…er-transport-correlation
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.
| * different transport objects share the same logical session | ||
| */ | ||
| const transportToSpanMap = new WeakMap<MCPTransport, Map<RequestId, RequestSpanMapValue>>(); | ||
| const sessionToSpanMap = new Map<string, Map<RequestId, RequestSpanMapValue>>(); |
There was a problem hiding this comment.
Module-level Maps leak memory without transport close
Medium Severity
sessionToSpanMap and sessionToSessionData are module-level Map<string, ...> instances whose entries are only removed via explicit delete calls in cleanupPendingSpansForTransport / cleanupSessionDataForTransport, which are solely triggered by the transport's onclose handler. If a transport is abandoned without onclose firing (e.g., unclean disconnect), entries accumulate indefinitely. The previous WeakMap<MCPTransport, ...> approach provided automatic GC-based cleanup as a safety net, which this change loses for stateful transports.
Additional Locations (1)
There was a problem hiding this comment.
acknowledged, intentional.
So this would happen if somehow the mcp server skips onClose when closing a connection, and it shouldn't happen.
- the only way to skip
oncloseis a process crash (where memory is freed anyway) or a genuine bug in user code where they abandon a transport without closing it - the MCP
server.close()callsclose()on all connected transports, which triggersonclose - HTTP connection drops trigger
onclosein the transport implementation - SSE disconnects trigger
onclose
This is a normal flow as per MCP docs:
1. Transport created (sessionId assigned)
2. server.connect(transport) → we wrap onmessage/send/onclose
3. Messages flow:
- onmessage: stores span in sessionToSpanMap[sessionId][requestId]
- send: looks up span in sessionToSpanMap[sessionId][requestId], ends it, deletes entry
4. Transport closes:
- onclose fires
- cleanupPendingSpansForTransport(): ends any orphan spans, calls sessionToSpanMap.delete(sessionId)
- cleanupSessionDataForTransport(): calls sessionToSessionData.delete(sessionId)
5. Maps are now clean for that sessionIdSo everything is cleaned up.
2 participants
Summary
Fixes MCP server instrumentation not recording events for wrapper transport patterns (like
NodeStreamableHTTPServerTransportwhich wrapsWebStandardStreamableHTTPServerTransport).The root cause: wrapper transports proxy
onmessage/sendvia getters/setters, causing differentthisvalues. The previousWeakMap<transport, ...>correlation couldn't match requests to responses. This changes toMap<sessionId, ...>correlation which works regardless of transport object reference.Changes
correlation.tsto use sessionId-based Maps instead of transport-based WeakMapssessionManagement.tsto use sessionId-based session lookupTest Plan
cd packages/core && yarn test(2010 tests pass)yarn lint) & (yarn test).Closes #19233
Thanks @gotenxds for reporting and debugging this issue!