Preserve context.Context across variant dispatch by bypassing in-memory transport#6
Merged
StevenRChen merged 8 commits intomainfrom Mar 18, 2026
Merged
Preserve context.Context across variant dispatch by bypassing in-memory transport#6StevenRChen merged 8 commits intomainfrom
StevenRChen merged 8 commits intomainfrom
Conversation
…t and request extra directly to the backend handler bypassing the jsonrpc handling clientSession.
StevenRChen
changed the title
Degiorgio
reviewed
Mar 18, 2026
Collaborator
Degiorgio
left a comment
Degiorgio
left a comment
There was a problem hiding this comment.
Missing test for context propagation
The entire purpose of this PR is fixing context propagation, but there's no test proving it works?
Degiorgio
reviewed
Mar 18, 2026
Degiorgio
previously approved these changes
Mar 18, 2026
Degiorgio
approved these changes
Mar 18, 2026
Collaborator
Author
|
Will add integration tests in a separate PR. |
StevenRChen
added a commit
that referenced
this pull request
Mar 23, 2026
…ications in stateless mode (#7) Registers a sending middleware on each inner server that intercepts all outgoing messages and routes them through the front session using a session-swapped request wrapper. Replaces the previous proxy client handler approach which dropped log messages in stateless mode. ## Motivation and Context Tool handlers in variant servers use req.Session (the inner session) to send notifications. Previously, log messages were silently dropped in stateless mode because the inner session is not aware of the front session. This change injects the sendingMethodHandler from the frontSession into the context and adds a sending middleware to redirect all sending requests to the frontSession handler directly. A few considerations in this implementation: * Having frontSession in the context allows notifications to be sent to the correct frontSession, when the same backend session could be used by multiple front sessions. * This solution will work for synchronous notifications that happen within the same context (progress, log, sampling, elicitation) but not working for asynchronous notifications (tool list changed, resource updated etc), because asynchronous notifications are sent with a new context. * All the operations before the sending method will still happen in the backend. Most notable impact is in stateless mode, requests from one client could change the log level for other clients. For this reason, I default the logging level to debug for now. * This implementation relies on the MCP Go SDK sendingMethodHandler taking a generic `Request` not one that tied to specific function. Based on the current SDK implementation this is unlikely to change. * The same generic method forwarding mechanism could work for receiving methods too, but with some caveat. That can be addressed in separate PR. ## How Has This Been Tested? TestNotificationDelivery — verifies progress and log messages arrive at the front clients with multiple clients making concurrent request. New tests added for context preservation that was added in #6. Context preservation is required to make the notification forwarding work anyway. Since we use the context to carry frontSession handler to the backend session. All existing tests pass ## Breaking Changes None. Tool handlers continue to use req.Session normally. ## Types of changes Bug fix (non-breaking change which fixes an issue)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ClientSessionRPC calls to direct inner server handler invocation via a newbackendSessiontype, preservingcontext.Contextvalues (e.g.RequestContext) that were previously lost crossing the in-memory transport boundarybackendSessionabstraction insession.gothat mirrors themcp.ClientSessionAPI but calls the inner server's handler chain directly, forwardingRequestExtra(token info, HTTP headers, etc.)captureMCPMethodHandler) inbackend.go, reused across all connectionsinnerConnectionto hold*backendSession+ cleanup; in-memory transport/client kept alive only for notification forwarding (progress, logging)getSession→getConnectionreturning*innerConnectionto reduce indirection in dispatch codeinjectVariantMetacalls to after type assertion and nil checks in all dispatch handlers for consistency and nil-safetyTest plan
go test ./variants/...)RequestContext) are available in tool handlers when invoked through variant dispatchRequestExtrafields (token info, HTTP headers) are forwarded to inner server handlers