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
9 changed files with 790 additions and 58 deletions
|
|
@ -2,7 +2,9 @@
|
||||||
|
|
||||||
A variant-aware MCP server that manages context budget by providing the same research tools (`search_papers`, `get_paper`, `summarize`) at different verbosity levels.
|
A variant-aware MCP server that manages context budget by providing the same research tools (`search_papers`, `get_paper`, `summarize`) at different verbosity levels.
|
||||||
|
|
||||||
**Pattern demonstrated:** Context budget management with description verbosity control.
|
Each tool sends progress notifications and log messages as it works, simulating a deep-research workflow where the client can observe each step (e.g. "Searching arXiv", "Resolving references", "Generating summary").
|
||||||
|
|
||||||
|
**Patterns demonstrated:** Context budget management with description verbosity control, notification streaming.
|
||||||
|
|
||||||
## Variants
|
## Variants
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,12 @@
|
||||||
// at different verbosity levels. Demonstrates how variants can optimize tool
|
// at different verbosity levels. Demonstrates how variants can optimize tool
|
||||||
// descriptions for different context window sizes and use cases.
|
// descriptions for different context window sizes and use cases.
|
||||||
//
|
//
|
||||||
// Capability demonstrated: Context budget management, description verbosity.
|
// Each tool sends progress notifications and log messages as it works,
|
||||||
|
// simulating a deep-research workflow where the client can observe each
|
||||||
|
// step (e.g. "Searching arXiv", "Resolving references").
|
||||||
|
//
|
||||||
|
// Capability demonstrated: Context budget management, description verbosity,
|
||||||
|
// notification streaming.
|
||||||
//
|
//
|
||||||
// Variants:
|
// Variants:
|
||||||
// - deep-research: Verbose multi-paragraph descriptions with usage examples
|
// - deep-research: Verbose multi-paragraph descriptions with usage examples
|
||||||
|
|
@ -31,7 +36,9 @@ import (
|
||||||
func main() {
|
func main() {
|
||||||
// Deep-research variant: verbose descriptions with usage examples
|
// Deep-research variant: verbose descriptions with usage examples
|
||||||
// for thorough research workflows with large context windows.
|
// for thorough research workflows with large context windows.
|
||||||
deepServer := mcp.NewServer(&mcp.Implementation{Name: "research-assistant", Version: "v1.0.0"}, nil)
|
deepServer := mcp.NewServer(&mcp.Implementation{Name: "research-assistant", Version: "v1.0.0"}, &mcp.ServerOptions{
|
||||||
|
Capabilities: &mcp.ServerCapabilities{Logging: &mcp.LoggingCapabilities{}},
|
||||||
|
})
|
||||||
mcp.AddTool(deepServer, &mcp.Tool{
|
mcp.AddTool(deepServer, &mcp.Tool{
|
||||||
Name: "search_papers",
|
Name: "search_papers",
|
||||||
Description: "Search academic papers across multiple databases including arXiv, Semantic Scholar, and PubMed. " +
|
Description: "Search academic papers across multiple databases including arXiv, Semantic Scholar, and PubMed. " +
|
||||||
|
|
@ -63,7 +70,9 @@ func main() {
|
||||||
|
|
||||||
// Quick-lookup variant: concise descriptions for fast Q&A sessions
|
// Quick-lookup variant: concise descriptions for fast Q&A sessions
|
||||||
// where context budget is limited.
|
// where context budget is limited.
|
||||||
quickServer := mcp.NewServer(&mcp.Implementation{Name: "research-assistant", Version: "v1.0.0"}, nil)
|
quickServer := mcp.NewServer(&mcp.Implementation{Name: "research-assistant", Version: "v1.0.0"}, &mcp.ServerOptions{
|
||||||
|
Capabilities: &mcp.ServerCapabilities{Logging: &mcp.LoggingCapabilities{}},
|
||||||
|
})
|
||||||
mcp.AddTool(quickServer, &mcp.Tool{
|
mcp.AddTool(quickServer, &mcp.Tool{
|
||||||
Name: "search_papers",
|
Name: "search_papers",
|
||||||
Description: "Search academic papers by query, with optional year and field filters.",
|
Description: "Search academic papers by query, with optional year and field filters.",
|
||||||
|
|
@ -79,7 +88,9 @@ func main() {
|
||||||
|
|
||||||
// Synthesis variant: balanced descriptions for report generation
|
// Synthesis variant: balanced descriptions for report generation
|
||||||
// workflows that need moderate detail.
|
// workflows that need moderate detail.
|
||||||
synthesisServer := mcp.NewServer(&mcp.Implementation{Name: "research-assistant", Version: "v1.0.0"}, nil)
|
synthesisServer := mcp.NewServer(&mcp.Implementation{Name: "research-assistant", Version: "v1.0.0"}, &mcp.ServerOptions{
|
||||||
|
Capabilities: &mcp.ServerCapabilities{Logging: &mcp.LoggingCapabilities{}},
|
||||||
|
})
|
||||||
mcp.AddTool(synthesisServer, &mcp.Tool{
|
mcp.AddTool(synthesisServer, &mcp.Tool{
|
||||||
Name: "search_papers",
|
Name: "search_papers",
|
||||||
Description: "Search academic papers across arXiv, Semantic Scholar, and PubMed. " +
|
Description: "Search academic papers across arXiv, Semantic Scholar, and PubMed. " +
|
||||||
|
|
|
||||||
|
|
@ -2,6 +2,8 @@ package main
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
)
|
)
|
||||||
|
|
@ -27,7 +29,33 @@ type SearchPapersOutput struct {
|
||||||
TotalFound int `json:"totalFound"`
|
TotalFound int `json:"totalFound"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func searchPapers(_ context.Context, _ *mcp.CallToolRequest, in SearchPapersInput) (*mcp.CallToolResult, SearchPapersOutput, error) {
|
var databases = []string{"arXiv", "Semantic Scholar", "PubMed"}
|
||||||
|
|
||||||
|
func searchPapers(ctx context.Context, req *mcp.CallToolRequest, in SearchPapersInput) (*mcp.CallToolResult, SearchPapersOutput, error) {
|
||||||
|
token := req.Params.GetProgressToken()
|
||||||
|
total := len(databases)
|
||||||
|
|
||||||
|
for i, db := range databases {
|
||||||
|
req.Session.Log(ctx, &mcp.LoggingMessageParams{
|
||||||
|
Level: "info",
|
||||||
|
Logger: "search_papers",
|
||||||
|
Data: fmt.Sprintf("Searching %s for %q", db, in.Query),
|
||||||
|
})
|
||||||
|
req.Session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{
|
||||||
|
ProgressToken: token,
|
||||||
|
Progress: float64(i + 1),
|
||||||
|
Total: float64(total),
|
||||||
|
Message: fmt.Sprintf("Searching %s (%d/%d)", db, i+1, total),
|
||||||
|
})
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
}
|
||||||
|
|
||||||
|
req.Session.Log(ctx, &mcp.LoggingMessageParams{
|
||||||
|
Level: "info",
|
||||||
|
Logger: "search_papers",
|
||||||
|
Data: fmt.Sprintf("Found 1247 results for %q, returning top 2", in.Query),
|
||||||
|
})
|
||||||
|
|
||||||
return nil, SearchPapersOutput{
|
return nil, SearchPapersOutput{
|
||||||
Papers: []Paper{
|
Papers: []Paper{
|
||||||
{
|
{
|
||||||
|
|
@ -71,7 +99,25 @@ type GetPaperOutput struct {
|
||||||
Paper PaperDetail `json:"paper"`
|
Paper PaperDetail `json:"paper"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func getPaper(_ context.Context, _ *mcp.CallToolRequest, in GetPaperInput) (*mcp.CallToolResult, GetPaperOutput, error) {
|
func getPaper(ctx context.Context, req *mcp.CallToolRequest, in GetPaperInput) (*mcp.CallToolResult, GetPaperOutput, error) {
|
||||||
|
token := req.Params.GetProgressToken()
|
||||||
|
steps := []string{"Fetching metadata", "Resolving references", "Loading citations"}
|
||||||
|
|
||||||
|
for i, step := range steps {
|
||||||
|
req.Session.Log(ctx, &mcp.LoggingMessageParams{
|
||||||
|
Level: "info",
|
||||||
|
Logger: "get_paper",
|
||||||
|
Data: fmt.Sprintf("%s for %s", step, in.ID),
|
||||||
|
})
|
||||||
|
req.Session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{
|
||||||
|
ProgressToken: token,
|
||||||
|
Progress: float64(i + 1),
|
||||||
|
Total: float64(len(steps)),
|
||||||
|
Message: fmt.Sprintf("%s (%d/%d)", step, i+1, len(steps)),
|
||||||
|
})
|
||||||
|
time.Sleep(80 * time.Millisecond)
|
||||||
|
}
|
||||||
|
|
||||||
return nil, GetPaperOutput{
|
return nil, GetPaperOutput{
|
||||||
Paper: PaperDetail{
|
Paper: PaperDetail{
|
||||||
ID: in.ID,
|
ID: in.ID,
|
||||||
|
|
@ -98,11 +144,30 @@ type SummarizeOutput struct {
|
||||||
Style string `json:"style"`
|
Style string `json:"style"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func summarize(_ context.Context, _ *mcp.CallToolRequest, in SummarizeInput) (*mcp.CallToolResult, SummarizeOutput, error) {
|
func summarize(ctx context.Context, req *mcp.CallToolRequest, in SummarizeInput) (*mcp.CallToolResult, SummarizeOutput, error) {
|
||||||
style := in.Style
|
style := in.Style
|
||||||
if style == "" {
|
if style == "" {
|
||||||
style = "abstract"
|
style = "abstract"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
token := req.Params.GetProgressToken()
|
||||||
|
steps := []string{"Analyzing text", "Generating summary"}
|
||||||
|
|
||||||
|
for i, step := range steps {
|
||||||
|
req.Session.Log(ctx, &mcp.LoggingMessageParams{
|
||||||
|
Level: "info",
|
||||||
|
Logger: "summarize",
|
||||||
|
Data: fmt.Sprintf("%s (%s style)", step, style),
|
||||||
|
})
|
||||||
|
req.Session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{
|
||||||
|
ProgressToken: token,
|
||||||
|
Progress: float64(i + 1),
|
||||||
|
Total: float64(len(steps)),
|
||||||
|
Message: fmt.Sprintf("%s (%d/%d)", step, i+1, len(steps)),
|
||||||
|
})
|
||||||
|
time.Sleep(80 * time.Millisecond)
|
||||||
|
}
|
||||||
|
|
||||||
return nil, SummarizeOutput{
|
return nil, SummarizeOutput{
|
||||||
Summary: "This paper presents key findings in the field, demonstrating significant improvements over prior work through novel methodology and comprehensive evaluation across multiple benchmarks.",
|
Summary: "This paper presents key findings in the field, demonstrating significant improvements over prior work through novel methodology and comprehensive evaluation across multiple benchmarks.",
|
||||||
Style: style,
|
Style: style,
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,6 @@ package variants
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
|
||||||
|
|
||||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
)
|
)
|
||||||
|
|
@ -31,30 +30,75 @@ type backend interface {
|
||||||
// inMemoryBackend connects to a co-located *mcp.Server via in-memory
|
// inMemoryBackend connects to a co-located *mcp.Server via in-memory
|
||||||
// transports.
|
// transports.
|
||||||
//
|
//
|
||||||
// Limitation: list-changed and resource-updated notifications from
|
// The sending redirect middleware intercepts all outgoing messages from
|
||||||
// inner servers are silently dropped. The mcp.ServerSession API only
|
// the inner server and forwards them through the front session. This
|
||||||
// exposes NotifyProgress and Log; there is no public method to send
|
// covers notifications (progress, log) and server-to-client requests
|
||||||
// tool/resource/prompt list-changed or resource-updated notifications
|
// (Elicit, CreateMessage, ListRoots) during request handling.
|
||||||
// to the front client. This means that if an inner server dynamically
|
|
||||||
// adds or removes tools/resources/prompts at runtime, the front
|
|
||||||
// client will NOT be notified.
|
|
||||||
//
|
//
|
||||||
|
// Limitation: async notifications (e.g. tool/resource/prompt
|
||||||
|
// list-changed) triggered outside request handling use
|
||||||
|
// context.Background() and lack the front session in the context, so
|
||||||
|
// the middleware falls through and these notifications are dropped.
|
||||||
// In practice this is acceptable because inner servers are typically
|
// In practice this is acceptable because inner servers are typically
|
||||||
// statically configured (tools registered at startup). If dynamic
|
// statically configured (tools registered at startup).
|
||||||
// capability changes are needed, this will require the Go MCP SDK to
|
|
||||||
// expose generic notification sending on ServerSession.
|
|
||||||
type inMemoryBackend struct {
|
type inMemoryBackend struct {
|
||||||
server *mcp.Server
|
variantID string
|
||||||
|
server *mcp.Server
|
||||||
|
mcpMethodHandler mcp.MethodHandler
|
||||||
}
|
}
|
||||||
|
|
||||||
// captureMCPMethodHandler captures and returns a reference to the inner
|
// sessionSwappedRequest wraps an existing mcp.Request but returns a
|
||||||
|
// different session from GetSession(). The embedded mcp.Request satisfies
|
||||||
|
// the unexported isRequest() method required by the interface.
|
||||||
|
type sessionSwappedRequest struct {
|
||||||
|
mcp.Request
|
||||||
|
session mcp.Session
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *sessionSwappedRequest) GetSession() mcp.Session { return r.session }
|
||||||
|
|
||||||
|
// sendingRedirectMiddleware returns a sending middleware that intercepts all
|
||||||
|
// outgoing messages from the inner server (notifications and server-to-client
|
||||||
|
// requests) and redirects them through the front server's sending handler with
|
||||||
|
// the front session swapped in. By swapping the session we route messages
|
||||||
|
// through the front connection to the real client.
|
||||||
|
//
|
||||||
|
// The front session (per-request) is read from context; the sending handler
|
||||||
|
// (constant) is read from the Server struct.
|
||||||
|
func sendingRedirectMiddleware(variantID string, vs *Server) mcp.Middleware {
|
||||||
|
return func(next mcp.MethodHandler) mcp.MethodHandler {
|
||||||
|
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
|
||||||
|
frontSession, _ := ctx.Value(frontSessionKeyType{}).(*mcp.ServerSession)
|
||||||
|
if frontSession == nil || vs.frontSendingHandler == nil {
|
||||||
|
return next(ctx, method, req)
|
||||||
|
}
|
||||||
|
return vs.frontSendingHandler(ctx, method, &sessionSwappedRequest{
|
||||||
|
Request: req,
|
||||||
|
session: frontSession,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// newInMemoryBackend creates an inMemoryBackend, registers the sending
|
||||||
|
// redirect middleware on the inner server, and captures the inner
|
||||||
|
// server's handler chain for direct dispatch.
|
||||||
|
func newInMemoryBackend(server *mcp.Server, variantID string, vs *Server) *inMemoryBackend {
|
||||||
|
server.AddSendingMiddleware(sendingRedirectMiddleware(variantID, vs))
|
||||||
|
return &inMemoryBackend{
|
||||||
|
variantID: variantID,
|
||||||
|
server: server,
|
||||||
|
mcpMethodHandler: captureReceivingMethodHandler(server),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// captureReceivingMethodHandler captures and returns a reference to the inner
|
||||||
// server's handler chain. This is a workaround using AddReceivingMiddleware
|
// server's handler chain. This is a workaround using AddReceivingMiddleware
|
||||||
// to gain a reference to mcp.Server.receivingMethodHandler_, since the SDK
|
// to gain a reference to mcp.Server.receivingMethodHandler_, since the SDK
|
||||||
// does not expose a public accessor for it. This can be replaced once the
|
// does not expose a public accessor for it. This can be replaced once the
|
||||||
// SDK exposes a public accessor for the receiving handler chain.
|
// SDK exposes a public accessor for the receiving handler chain.
|
||||||
func captureMCPMethodHandler(server *mcp.Server) (mcp.MethodHandler, error) {
|
func captureReceivingMethodHandler(server *mcp.Server) mcp.MethodHandler {
|
||||||
var handler mcp.MethodHandler
|
var handler mcp.MethodHandler
|
||||||
|
|
||||||
// The middleware is identity (returns next unmodified), so the handler
|
// The middleware is identity (returns next unmodified), so the handler
|
||||||
// chain is unchanged, no extra hop introduced even if called multiple times.
|
// chain is unchanged, no extra hop introduced even if called multiple times.
|
||||||
server.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler {
|
server.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler {
|
||||||
|
|
@ -62,22 +106,15 @@ func captureMCPMethodHandler(server *mcp.Server) (mcp.MethodHandler, error) {
|
||||||
return next
|
return next
|
||||||
})
|
})
|
||||||
|
|
||||||
if handler == nil {
|
return handler
|
||||||
return nil, fmt.Errorf("failed to capture backend MCP method handler")
|
|
||||||
}
|
|
||||||
return handler, nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// connect creates an in-memory transport pair and connects the inner server.
|
// connect creates an in-memory transport pair and connects the inner server.
|
||||||
// Requests bypass the transport via serverHandler to preserve context values.
|
// Requests bypass the transport via mcpMethodHandler to preserve context
|
||||||
// The transport is kept alive only for notification forwarding (progress,
|
// values. Notifications are redirected by the sending middleware registered
|
||||||
// logging) from the inner server to the front client.
|
// in newInMemoryBackend; the proxy client is kept only for the initialize
|
||||||
|
// handshake and to set the inner session's log level.
|
||||||
func (b *inMemoryBackend) connect(ctx context.Context, variant ServerVariant, frontSession *mcp.ServerSession) (*innerConnection, error) {
|
func (b *inMemoryBackend) connect(ctx context.Context, variant ServerVariant, frontSession *mcp.ServerSession) (*innerConnection, error) {
|
||||||
mcpMethodHandler, err := captureMCPMethodHandler(b.server)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
serverTransport, clientSideTransport := mcp.NewInMemoryTransports()
|
serverTransport, clientSideTransport := mcp.NewInMemoryTransports()
|
||||||
|
|
||||||
serverSession, err := b.server.Connect(ctx, serverTransport, nil)
|
serverSession, err := b.server.Connect(ctx, serverTransport, nil)
|
||||||
|
|
@ -85,26 +122,26 @@ func (b *inMemoryBackend) connect(ctx context.Context, variant ServerVariant, fr
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// The proxy client completes the initialize handshake and sets the
|
||||||
|
// inner session's log level. It must remain open for the lifetime of
|
||||||
|
// the serverSession — closing it tears down the in-memory transport,
|
||||||
|
// causing the serverSession to shut down. The nop handlers advertise
|
||||||
|
// capabilities so the inner ServerSession doesn't short-circuit
|
||||||
|
// methods like Elicit with "client does not support X". They are
|
||||||
|
// never called because the sending redirect middleware intercepts
|
||||||
|
// all outgoing messages before they reach the transport.
|
||||||
|
nopElicit := func(context.Context, *mcp.ElicitRequest) (*mcp.ElicitResult, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
nopSampling := func(context.Context, *mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
client := mcp.NewClient(&mcp.Implementation{
|
client := mcp.NewClient(&mcp.Implementation{
|
||||||
Name: "variant-proxy-client",
|
Name: "variant-proxy-client",
|
||||||
Version: "1.0.0",
|
Version: "1.0.0",
|
||||||
}, &mcp.ClientOptions{
|
}, &mcp.ClientOptions{
|
||||||
ProgressNotificationHandler: func(ctx context.Context, req *mcp.ProgressNotificationClientRequest) {
|
ElicitationHandler: nopElicit,
|
||||||
if frontSession != nil {
|
CreateMessageHandler: nopSampling,
|
||||||
injectVariantMeta(req.Params, variant.ID)
|
|
||||||
_ = frontSession.NotifyProgress(ctx, req.Params)
|
|
||||||
}
|
|
||||||
},
|
|
||||||
LoggingMessageHandler: func(ctx context.Context, req *mcp.LoggingMessageRequest) {
|
|
||||||
if frontSession != nil {
|
|
||||||
injectVariantMeta(req.Params, variant.ID)
|
|
||||||
_ = frontSession.Log(ctx, req.Params)
|
|
||||||
}
|
|
||||||
},
|
|
||||||
ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) {},
|
|
||||||
ResourceListChangedHandler: func(context.Context, *mcp.ResourceListChangedRequest) {},
|
|
||||||
PromptListChangedHandler: func(context.Context, *mcp.PromptListChangedRequest) {},
|
|
||||||
ResourceUpdatedHandler: func(context.Context, *mcp.ResourceUpdatedNotificationRequest) {},
|
|
||||||
})
|
})
|
||||||
|
|
||||||
clientSession, err := client.Connect(ctx, clientSideTransport, nil)
|
clientSession, err := client.Connect(ctx, clientSideTransport, nil)
|
||||||
|
|
@ -113,11 +150,19 @@ func (b *inMemoryBackend) connect(ctx context.Context, variant ServerVariant, fr
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Set the inner session's log level to the lowest so that
|
||||||
|
// ServerSession.Log() does not short-circuit before reaching the
|
||||||
|
// sending middleware. The actual log-level filtering is performed by
|
||||||
|
// the front-facing session when the middleware redirects the
|
||||||
|
// notification. Errors are ignored: if the inner server does not
|
||||||
|
// advertise the Logging capability the call simply fails harmlessly.
|
||||||
|
_ = clientSession.SetLoggingLevel(ctx, &mcp.SetLoggingLevelParams{Level: "debug"})
|
||||||
|
|
||||||
return &innerConnection{
|
return &innerConnection{
|
||||||
backendSession: &backendSession{
|
backendSession: &backendSession{
|
||||||
variantID: variant.ID,
|
variantID: b.variantID,
|
||||||
serverSession: serverSession,
|
serverSession: serverSession,
|
||||||
mcpMethodHandler: mcpMethodHandler,
|
mcpMethodHandler: b.mcpMethodHandler,
|
||||||
},
|
},
|
||||||
cleanupFn: func() {
|
cleanupFn: func() {
|
||||||
clientSession.Close()
|
clientSession.Close()
|
||||||
|
|
|
||||||
149
go/sdk/variants/context_test.go
Normal file
149
go/sdk/variants/context_test.go
Normal file
|
|
@ -0,0 +1,149 @@
|
||||||
|
package variants
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
type contextKey string
|
||||||
|
|
||||||
|
const testCtxKey contextKey = "test-value"
|
||||||
|
|
||||||
|
type checkContextOutput struct {
|
||||||
|
ContextValue string `json:"contextValue"`
|
||||||
|
VariantID string `json:"variantId"`
|
||||||
|
HasExtra bool `json:"hasExtra"`
|
||||||
|
CustomHeader string `json:"customHeader"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type emptyInput struct{}
|
||||||
|
|
||||||
|
// headerRoundTripper injects custom headers into every outgoing HTTP request.
|
||||||
|
type headerRoundTripper struct {
|
||||||
|
base http.RoundTripper
|
||||||
|
header http.Header
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
req = req.Clone(req.Context())
|
||||||
|
for k, v := range h.header {
|
||||||
|
req.Header[k] = v
|
||||||
|
}
|
||||||
|
return h.base.RoundTrip(req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newContextTestInnerServer() *mcp.Server {
|
||||||
|
inner := mcp.NewServer(&mcp.Implementation{Name: "ctx-test", Version: "v1.0.0"}, nil)
|
||||||
|
|
||||||
|
mcp.AddTool(inner, &mcp.Tool{
|
||||||
|
Name: "check_context",
|
||||||
|
Description: "Reports context and request metadata",
|
||||||
|
}, func(ctx context.Context, req *mcp.CallToolRequest, _ emptyInput) (*mcp.CallToolResult, checkContextOutput, error) {
|
||||||
|
out := checkContextOutput{}
|
||||||
|
|
||||||
|
if v, ok := ctx.Value(testCtxKey).(string); ok {
|
||||||
|
out.ContextValue = v
|
||||||
|
}
|
||||||
|
|
||||||
|
if req.Extra != nil {
|
||||||
|
out.HasExtra = true
|
||||||
|
if req.Extra.Header != nil {
|
||||||
|
out.CustomHeader = req.Extra.Header.Get("X-Test-Header")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if req.Params != nil {
|
||||||
|
meta := req.Params.GetMeta()
|
||||||
|
if v, ok := meta[metaKeyVariant].(string); ok {
|
||||||
|
out.VariantID = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil, out, nil
|
||||||
|
})
|
||||||
|
|
||||||
|
return inner
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseCheckContextOutput(t *testing.T, result *mcp.CallToolResult) checkContextOutput {
|
||||||
|
t.Helper()
|
||||||
|
require.NotEmpty(t, result.Content)
|
||||||
|
tc, ok := result.Content[0].(*mcp.TextContent)
|
||||||
|
require.True(t, ok, "expected TextContent, got %T", result.Content[0])
|
||||||
|
|
||||||
|
var out checkContextOutput
|
||||||
|
require.NoError(t, json.Unmarshal([]byte(tc.Text), &out))
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestContextPreservation verifies that context values set by middleware on
|
||||||
|
// the front (outer) mcp.Server and HTTP request headers are preserved through
|
||||||
|
// the variant dispatch layer and reach the inner server's tool handler.
|
||||||
|
func TestContextPreservation(t *testing.T) {
|
||||||
|
vs := NewServer(&mcp.Implementation{Name: "ctx-test", Version: "v1.0.0"}).
|
||||||
|
WithVariant(ServerVariant{
|
||||||
|
ID: "test-variant",
|
||||||
|
Description: "Context test variant",
|
||||||
|
Status: Stable,
|
||||||
|
}, newContextTestInnerServer(), 0)
|
||||||
|
|
||||||
|
frontServer, err := vs.mcpServer(false)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
frontServer.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler {
|
||||||
|
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
|
||||||
|
ctx = context.WithValue(ctx, testCtxKey, "from-middleware")
|
||||||
|
return next(ctx, method, req)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
handler := mcp.NewStreamableHTTPHandler(
|
||||||
|
func(*http.Request) *mcp.Server { return frontServer },
|
||||||
|
nil,
|
||||||
|
)
|
||||||
|
httpSrv := httptest.NewServer(handler)
|
||||||
|
t.Cleanup(httpSrv.Close)
|
||||||
|
|
||||||
|
client := mcp.NewClient(
|
||||||
|
&mcp.Implementation{Name: "ctx-test-client", Version: "v0.0.1"},
|
||||||
|
nil,
|
||||||
|
)
|
||||||
|
|
||||||
|
session, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{
|
||||||
|
Endpoint: httpSrv.URL,
|
||||||
|
HTTPClient: &http.Client{
|
||||||
|
Transport: &headerRoundTripper{
|
||||||
|
base: http.DefaultTransport,
|
||||||
|
header: http.Header{
|
||||||
|
"X-Test-Header": []string{"test-header-value"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}, nil)
|
||||||
|
require.NoError(t, err)
|
||||||
|
t.Cleanup(func() { session.Close() })
|
||||||
|
|
||||||
|
result, err := session.CallTool(context.Background(), &mcp.CallToolParams{
|
||||||
|
Name: "check_context",
|
||||||
|
Arguments: map[string]json.RawMessage{},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
out := parseCheckContextOutput(t, result)
|
||||||
|
|
||||||
|
assert.Equal(t, "from-middleware", out.ContextValue,
|
||||||
|
"context value from front server middleware should be preserved through dispatch")
|
||||||
|
assert.Equal(t, "test-variant", out.VariantID,
|
||||||
|
"variant ID should be injected into request _meta")
|
||||||
|
assert.True(t, out.HasExtra,
|
||||||
|
"RequestExtra should be forwarded to inner tool handler")
|
||||||
|
assert.Equal(t, "test-header-value", out.CustomHeader,
|
||||||
|
"HTTP headers should be captured in RequestExtra and forwarded through dispatch")
|
||||||
|
}
|
||||||
215
go/sdk/variants/notification_test.go
Normal file
215
go/sdk/variants/notification_test.go
Normal file
|
|
@ -0,0 +1,215 @@
|
||||||
|
// Copyright 2025 The MCP Variants Authors. All rights reserved.
|
||||||
|
// Use of this source code is governed by a Apache-2.0
|
||||||
|
// license that can be found in the LICENSE file.
|
||||||
|
|
||||||
|
package variants
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http/httptest"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
const notifyCount = 3
|
||||||
|
|
||||||
|
type notifyInput struct {
|
||||||
|
ClientID string `json:"client_id"`
|
||||||
|
Count int `json:"count"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func notifyHandler(ctx context.Context, req *mcp.CallToolRequest, input notifyInput) (*mcp.CallToolResult, any, error) {
|
||||||
|
count := input.Count
|
||||||
|
if count <= 0 {
|
||||||
|
count = notifyCount
|
||||||
|
}
|
||||||
|
|
||||||
|
progressToken := req.Params.GetProgressToken()
|
||||||
|
|
||||||
|
for i := 1; i <= count; i++ {
|
||||||
|
req.Session.NotifyProgress(ctx, &mcp.ProgressNotificationParams{
|
||||||
|
ProgressToken: progressToken,
|
||||||
|
Progress: float64(i),
|
||||||
|
Total: float64(count),
|
||||||
|
Message: fmt.Sprintf("[%s] progress %d/%d", input.ClientID, i, count),
|
||||||
|
})
|
||||||
|
req.Session.Log(ctx, &mcp.LoggingMessageParams{
|
||||||
|
Level: "info",
|
||||||
|
Logger: "notify-test",
|
||||||
|
Data: fmt.Sprintf("[%s] log %d/%d", input.ClientID, i, count),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil, map[string]string{"status": "done", "client_id": input.ClientID}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func newNotifyVariantServer() *Server {
|
||||||
|
inner := mcp.NewServer(
|
||||||
|
&mcp.Implementation{Name: "notify-test", Version: "v1.0.0"},
|
||||||
|
&mcp.ServerOptions{
|
||||||
|
Capabilities: &mcp.ServerCapabilities{
|
||||||
|
Logging: &mcp.LoggingCapabilities{},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
mcp.AddTool(inner, &mcp.Tool{
|
||||||
|
Name: "notify",
|
||||||
|
Description: "Sends progress and log notifications via req.Session",
|
||||||
|
}, notifyHandler)
|
||||||
|
|
||||||
|
return NewServer(&mcp.Implementation{Name: "notify-test", Version: "v1.0.0"}).
|
||||||
|
WithVariant(ServerVariant{
|
||||||
|
ID: "default",
|
||||||
|
Description: "Notification test variant",
|
||||||
|
Status: Stable,
|
||||||
|
}, inner, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
type notificationCollector struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
progress []*mcp.ProgressNotificationParams
|
||||||
|
logs []*mcp.LoggingMessageParams
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *notificationCollector) clientOptions() *mcp.ClientOptions {
|
||||||
|
return &mcp.ClientOptions{
|
||||||
|
ProgressNotificationHandler: func(_ context.Context, req *mcp.ProgressNotificationClientRequest) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
c.progress = append(c.progress, req.Params)
|
||||||
|
},
|
||||||
|
LoggingMessageHandler: func(_ context.Context, req *mcp.LoggingMessageRequest) {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
c.logs = append(c.logs, req.Params)
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *notificationCollector) progressCount() int {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
return len(c.progress)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *notificationCollector) logCount() int {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
return len(c.logs)
|
||||||
|
}
|
||||||
|
|
||||||
|
const numClients = 3
|
||||||
|
|
||||||
|
func TestNotificationDelivery(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
stateless bool
|
||||||
|
}{
|
||||||
|
{name: "Stateful", stateless: false},
|
||||||
|
{name: "Stateless", stateless: true},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
vs := newNotifyVariantServer()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
handler := NewStreamableHTTPHandler(vs, &mcp.StreamableHTTPOptions{Stateless: tt.stateless})
|
||||||
|
httpSrv := httptest.NewServer(handler)
|
||||||
|
t.Cleanup(httpSrv.Close)
|
||||||
|
|
||||||
|
// Set up N clients, each with its own collector.
|
||||||
|
type clientState struct {
|
||||||
|
id string
|
||||||
|
session *mcp.ClientSession
|
||||||
|
collector *notificationCollector
|
||||||
|
}
|
||||||
|
clients := make([]clientState, numClients)
|
||||||
|
for i := range clients {
|
||||||
|
cs := &clients[i]
|
||||||
|
cs.id = fmt.Sprintf("client-%d", i)
|
||||||
|
cs.collector = ¬ificationCollector{}
|
||||||
|
|
||||||
|
client := mcp.NewClient(
|
||||||
|
&mcp.Implementation{Name: cs.id, Version: "v0.0.1"},
|
||||||
|
cs.collector.clientOptions(),
|
||||||
|
)
|
||||||
|
var err error
|
||||||
|
cs.session, err = client.Connect(ctx, &mcp.StreamableClientTransport{
|
||||||
|
Endpoint: httpSrv.URL,
|
||||||
|
}, nil)
|
||||||
|
require.NoError(t, err)
|
||||||
|
t.Cleanup(func() { cs.session.Close() })
|
||||||
|
|
||||||
|
err = cs.session.SetLoggingLevel(ctx, &mcp.SetLoggingLevelParams{Level: "debug"})
|
||||||
|
require.NoError(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// All clients call the tool concurrently.
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
errs := make([]error, numClients)
|
||||||
|
for i := range clients {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(idx int) {
|
||||||
|
defer wg.Done()
|
||||||
|
cs := &clients[idx]
|
||||||
|
args := map[string]json.RawMessage{
|
||||||
|
"client_id": json.RawMessage(fmt.Sprintf("%q", cs.id)),
|
||||||
|
"count": json.RawMessage(fmt.Sprintf("%d", notifyCount)),
|
||||||
|
}
|
||||||
|
result, err := cs.session.CallTool(ctx, &mcp.CallToolParams{
|
||||||
|
Name: "notify",
|
||||||
|
Arguments: args,
|
||||||
|
})
|
||||||
|
if err == nil && result == nil {
|
||||||
|
err = fmt.Errorf("nil result for %s", cs.id)
|
||||||
|
}
|
||||||
|
errs[idx] = err
|
||||||
|
}(i)
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
for i, err := range errs {
|
||||||
|
require.NoError(t, err, "client %d tool call failed", i)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for all notifications to arrive at each collector.
|
||||||
|
for i := range clients {
|
||||||
|
cs := &clients[i]
|
||||||
|
require.Eventually(t, func() bool {
|
||||||
|
return cs.collector.progressCount() == notifyCount && cs.collector.logCount() == notifyCount
|
||||||
|
}, 2*time.Second, 10*time.Millisecond,
|
||||||
|
"client %s: expected %d progress and %d log notifications", cs.id, notifyCount, notifyCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify each client received only its own notifications.
|
||||||
|
for i := range clients {
|
||||||
|
cs := &clients[i]
|
||||||
|
cs.collector.mu.Lock()
|
||||||
|
|
||||||
|
for j := 0; j < notifyCount; j++ {
|
||||||
|
p := cs.collector.progress[j]
|
||||||
|
assert.Equal(t, float64(j+1), p.Progress)
|
||||||
|
assert.Equal(t, float64(notifyCount), p.Total)
|
||||||
|
assert.Equal(t, fmt.Sprintf("[%s] progress %d/%d", cs.id, j+1, notifyCount), p.Message,
|
||||||
|
"client %s got a progress notification belonging to another client", cs.id)
|
||||||
|
|
||||||
|
l := cs.collector.logs[j]
|
||||||
|
assert.Equal(t, "info", string(l.Level))
|
||||||
|
assert.Equal(t, "notify-test", l.Logger)
|
||||||
|
assert.Equal(t, fmt.Sprintf("[%s] log %d/%d", cs.id, j+1, notifyCount), l.Data,
|
||||||
|
"client %s got a log notification belonging to another client", cs.id)
|
||||||
|
}
|
||||||
|
|
||||||
|
cs.collector.mu.Unlock()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
204
go/sdk/variants/sending_test.go
Normal file
204
go/sdk/variants/sending_test.go
Normal file
|
|
@ -0,0 +1,204 @@
|
||||||
|
// Copyright 2025 The MCP Variants Authors. All rights reserved.
|
||||||
|
// Use of this source code is governed by a Apache-2.0
|
||||||
|
// license that can be found in the LICENSE file.
|
||||||
|
|
||||||
|
package variants
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
// sendingInput is the tool input for the sending-methods test tool.
|
||||||
|
type sendingInput struct {
|
||||||
|
// Which server-to-client methods to call: "elicit", "createMessage", "listRoots".
|
||||||
|
Methods []string `json:"methods"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendingResult is returned by the tool with results from each method call.
|
||||||
|
type sendingResult struct {
|
||||||
|
ElicitAction string `json:"elicitAction,omitempty"`
|
||||||
|
ElicitContent map[string]any `json:"elicitContent,omitempty"`
|
||||||
|
SamplingModel string `json:"samplingModel,omitempty"`
|
||||||
|
SamplingText string `json:"samplingText,omitempty"`
|
||||||
|
RootURIs []string `json:"rootURIs,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendingHandler calls server-to-client methods via req.Session and returns
|
||||||
|
// the results so the test can verify they were forwarded through the variant layer.
|
||||||
|
func sendingHandler(ctx context.Context, req *mcp.CallToolRequest, input sendingInput) (*mcp.CallToolResult, sendingResult, error) {
|
||||||
|
ss := req.Session
|
||||||
|
var out sendingResult
|
||||||
|
|
||||||
|
for _, m := range input.Methods {
|
||||||
|
switch m {
|
||||||
|
case "elicit":
|
||||||
|
result, err := ss.Elicit(ctx, &mcp.ElicitParams{
|
||||||
|
Message: "Please confirm",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, out, err
|
||||||
|
}
|
||||||
|
out.ElicitAction = result.Action
|
||||||
|
out.ElicitContent = result.Content
|
||||||
|
|
||||||
|
case "createMessage":
|
||||||
|
result, err := ss.CreateMessage(ctx, &mcp.CreateMessageParams{
|
||||||
|
MaxTokens: 100,
|
||||||
|
Messages: []*mcp.SamplingMessage{
|
||||||
|
{Role: "user", Content: &mcp.TextContent{Text: "Hello"}},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, out, err
|
||||||
|
}
|
||||||
|
out.SamplingModel = result.Model
|
||||||
|
if tc, ok := result.Content.(*mcp.TextContent); ok {
|
||||||
|
out.SamplingText = tc.Text
|
||||||
|
}
|
||||||
|
|
||||||
|
case "listRoots":
|
||||||
|
result, err := ss.ListRoots(ctx, &mcp.ListRootsParams{})
|
||||||
|
if err != nil {
|
||||||
|
return nil, out, err
|
||||||
|
}
|
||||||
|
for _, r := range result.Roots {
|
||||||
|
out.RootURIs = append(out.RootURIs, r.URI)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil, out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func newSendingVariantServer() *Server {
|
||||||
|
inner := mcp.NewServer(
|
||||||
|
&mcp.Implementation{Name: "sending-test", Version: "v1.0.0"},
|
||||||
|
nil,
|
||||||
|
)
|
||||||
|
mcp.AddTool(inner, &mcp.Tool{
|
||||||
|
Name: "send",
|
||||||
|
Description: "Calls server-to-client methods and returns results",
|
||||||
|
}, sendingHandler)
|
||||||
|
|
||||||
|
return NewServer(&mcp.Implementation{Name: "sending-test", Version: "v1.0.0"}).
|
||||||
|
WithVariant(ServerVariant{
|
||||||
|
ID: "default",
|
||||||
|
Description: "Sending methods test variant",
|
||||||
|
Status: Stable,
|
||||||
|
}, inner, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestSendingMethodsForwarding verifies that server-to-client request methods
|
||||||
|
// (Elicit, CreateMessage, ListRoots) are correctly forwarded through the
|
||||||
|
// variant layer's sending redirect middleware to the real client.
|
||||||
|
//
|
||||||
|
// Stateless mode is excluded: the StreamableHTTP transport rejects
|
||||||
|
// server-to-client requests in stateless mode ("stateless servers cannot
|
||||||
|
// make requests").
|
||||||
|
func TestSendingMethodsForwarding(t *testing.T) {
|
||||||
|
vs := newSendingVariantServer()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
handler := NewStreamableHTTPHandler(vs, nil)
|
||||||
|
httpSrv := httptest.NewServer(handler)
|
||||||
|
t.Cleanup(httpSrv.Close)
|
||||||
|
|
||||||
|
client := mcp.NewClient(
|
||||||
|
&mcp.Implementation{Name: "test-client", Version: "v0.0.1"},
|
||||||
|
&mcp.ClientOptions{
|
||||||
|
CreateMessageHandler: func(_ context.Context, req *mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
|
||||||
|
return &mcp.CreateMessageResult{
|
||||||
|
Model: "test-model",
|
||||||
|
Role: "assistant",
|
||||||
|
Content: &mcp.TextContent{Text: "mock response"},
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
ElicitationHandler: func(_ context.Context, req *mcp.ElicitRequest) (*mcp.ElicitResult, error) {
|
||||||
|
return &mcp.ElicitResult{
|
||||||
|
Action: "accept",
|
||||||
|
Content: map[string]any{"confirmed": true},
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
client.AddRoots(
|
||||||
|
&mcp.Root{URI: "file:///workspace/project-a", Name: "Project A"},
|
||||||
|
&mcp.Root{URI: "file:///workspace/project-b", Name: "Project B"},
|
||||||
|
)
|
||||||
|
|
||||||
|
session, err := client.Connect(ctx, &mcp.StreamableClientTransport{
|
||||||
|
Endpoint: httpSrv.URL,
|
||||||
|
}, nil)
|
||||||
|
require.NoError(t, err)
|
||||||
|
t.Cleanup(func() { session.Close() })
|
||||||
|
|
||||||
|
t.Run("Elicit", func(t *testing.T) {
|
||||||
|
result := callSend(t, ctx, session, "elicit")
|
||||||
|
assert.Equal(t, "accept", result.ElicitAction)
|
||||||
|
assert.Equal(t, true, result.ElicitContent["confirmed"])
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateMessage", func(t *testing.T) {
|
||||||
|
result := callSend(t, ctx, session, "createMessage")
|
||||||
|
assert.Equal(t, "test-model", result.SamplingModel)
|
||||||
|
assert.Equal(t, "mock response", result.SamplingText)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("ListRoots", func(t *testing.T) {
|
||||||
|
result := callSend(t, ctx, session, "listRoots")
|
||||||
|
assert.Equal(t, []string{
|
||||||
|
"file:///workspace/project-a",
|
||||||
|
"file:///workspace/project-b",
|
||||||
|
}, result.RootURIs)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("AllMethods", func(t *testing.T) {
|
||||||
|
result := callSend(t, ctx, session, "elicit", "createMessage", "listRoots")
|
||||||
|
assert.Equal(t, "accept", result.ElicitAction)
|
||||||
|
assert.Equal(t, "test-model", result.SamplingModel)
|
||||||
|
assert.Equal(t, "mock response", result.SamplingText)
|
||||||
|
assert.Len(t, result.RootURIs, 2)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// callSend invokes the "send" tool with the given methods and returns
|
||||||
|
// the parsed result.
|
||||||
|
func callSend(t *testing.T, ctx context.Context, session *mcp.ClientSession, methods ...string) sendingResult {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
methodsJSON, err := json.Marshal(methods)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
result, err := session.CallTool(ctx, &mcp.CallToolParams{
|
||||||
|
Name: "send",
|
||||||
|
Arguments: map[string]json.RawMessage{
|
||||||
|
"methods": json.RawMessage(methodsJSON),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NotNil(t, result)
|
||||||
|
|
||||||
|
var out sendingResult
|
||||||
|
require.NoError(t, unmarshalToolResult(result, &out))
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// unmarshalToolResult extracts the JSON content from a CallToolResult.
|
||||||
|
func unmarshalToolResult(result *mcp.CallToolResult, v any) error {
|
||||||
|
if len(result.Content) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
tc, ok := result.Content[0].(*mcp.TextContent)
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return json.Unmarshal([]byte(tc.Text), v)
|
||||||
|
}
|
||||||
|
|
@ -7,6 +7,7 @@ package variants
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
|
|
@ -31,10 +32,11 @@ type variantEntry struct {
|
||||||
// (via [NewStreamableHTTPHandler] with Stateless option), a single set of
|
// (via [NewStreamableHTTPHandler] with Stateless option), a single set of
|
||||||
// shared connections is created at construction and reused across all requests.
|
// shared connections is created at construction and reused across all requests.
|
||||||
type Server struct {
|
type Server struct {
|
||||||
impl *mcp.Implementation
|
impl *mcp.Implementation
|
||||||
variants []variantEntry
|
variants []variantEntry
|
||||||
rankingFunc RankingFunc
|
rankingFunc RankingFunc
|
||||||
shared *sessionState // non-nil in stateless mode; cleaned up by Close
|
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.
|
// 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.
|
// Variant IDs must be unique; registering a duplicate panics.
|
||||||
func (s *Server) WithVariant(v ServerVariant, mcpServer *mcp.Server, priority int) *Server {
|
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
|
// 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))
|
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
|
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
|
// Initialize helpers
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
|
||||||
|
|
@ -160,6 +160,9 @@ func (ss *sessionState) close() {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// frontSessionKeyType is the context key for the front-facing ServerSession.
|
||||||
|
type frontSessionKeyType struct{}
|
||||||
|
|
||||||
// injectVariantMeta sets the variant ID in a Params' _meta map,
|
// injectVariantMeta sets the variant ID in a Params' _meta map,
|
||||||
// preserving any existing metadata.
|
// preserving any existing metadata.
|
||||||
func injectVariantMeta(p mcp.Params, variantID string) {
|
func injectVariantMeta(p mcp.Params, variantID string) {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue