variants: redirect inner server sending methods to fix dropping notifications 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 https://github.com/modelcontextprotocol/experimental-ext-variants/pull/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 commit is contained in:
parent
150c110128
commit
e5e8f117e8
1 changed files with 790 additions and 58 deletions
|
|
@ -7,6 +7,7 @@ package variants
|
|||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
|
||||
|
|
@ -31,10 +32,11 @@ type variantEntry struct {
|
|||
// (via [NewStreamableHTTPHandler] with Stateless option), a single set of
|
||||
// shared connections is created at construction and reused across all requests.
|
||||
type Server struct {
|
||||
impl *mcp.Implementation
|
||||
variants []variantEntry
|
||||
rankingFunc RankingFunc
|
||||
shared *sessionState // non-nil in stateless mode; cleaned up by Close
|
||||
impl *mcp.Implementation
|
||||
variants []variantEntry
|
||||
rankingFunc RankingFunc
|
||||
shared *sessionState // non-nil in stateless mode; cleaned up by Close
|
||||
frontSendingHandler mcp.MethodHandler // set by mcpServer(); used by sendingRedirectMiddleware
|
||||
}
|
||||
|
||||
// NewServer creates a new variant-aware server with no registered variants.
|
||||
|
|
@ -92,7 +94,7 @@ func (s *Server) addVariant(v ServerVariant, b backend, priority int) *Server {
|
|||
//
|
||||
// Variant IDs must be unique; registering a duplicate panics.
|
||||
func (s *Server) WithVariant(v ServerVariant, mcpServer *mcp.Server, priority int) *Server {
|
||||
return s.addVariant(v, &inMemoryBackend{server: mcpServer}, priority)
|
||||
return s.addVariant(v, newInMemoryBackend(mcpServer, v.ID, s), priority)
|
||||
}
|
||||
|
||||
// WithHTTPVariant registers a ServerVariant backed by an mcp.Server exposed
|
||||
|
|
@ -256,9 +258,45 @@ func (s *Server) mcpServer(stateless bool) (*mcp.Server, error) {
|
|||
|
||||
frontServer.AddReceivingMiddleware(s.sessionMiddleware(sessions, shared))
|
||||
|
||||
// Inject the front-facing session into the context so inner servers'
|
||||
// sending middleware can redirect notifications to the real client.
|
||||
frontServer.AddReceivingMiddleware(captureFrontSessionMiddleware)
|
||||
|
||||
s.frontSendingHandler, err = captureSendingMethodHandler(frontServer)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return frontServer, nil
|
||||
}
|
||||
|
||||
// captureFrontSessionMiddleware injects the front-facing session into the
|
||||
// context so inner servers' sending middleware can redirect notifications
|
||||
// and other method calls to the real client.
|
||||
func captureFrontSessionMiddleware(next mcp.MethodHandler) mcp.MethodHandler {
|
||||
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
|
||||
ctx = context.WithValue(ctx, frontSessionKeyType{}, req.GetSession().(*mcp.ServerSession))
|
||||
return next(ctx, method, req)
|
||||
}
|
||||
}
|
||||
|
||||
// captureSendingMethodHandler captures the server's sending method handler.
|
||||
// It uses AddSendingMiddleware to work around the lack of a public accessor
|
||||
// for mcp.Server.sendingMethodHandler_. Because the middleware is identity
|
||||
// (returns next unmodified), the handler chain is unchanged and no extra hop
|
||||
// is introduced even if called multiple times.
|
||||
func captureSendingMethodHandler(server *mcp.Server) (mcp.MethodHandler, error) {
|
||||
var handler mcp.MethodHandler
|
||||
server.AddSendingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler {
|
||||
handler = next
|
||||
return next
|
||||
})
|
||||
if handler == nil {
|
||||
return nil, fmt.Errorf("failed to capture sending method handler")
|
||||
}
|
||||
return handler, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Initialize helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Reference in a new issue