diff --git a/go/sdk/examples/server/research/README.md b/go/sdk/examples/server/research/README.md index 9c55be3..37f7585 100644 --- a/go/sdk/examples/server/research/README.md +++ b/go/sdk/examples/server/research/README.md @@ -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. -**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 diff --git a/go/sdk/examples/server/research/main.go b/go/sdk/examples/server/research/main.go index 51edf6c..0e5b6d1 100644 --- a/go/sdk/examples/server/research/main.go +++ b/go/sdk/examples/server/research/main.go @@ -2,7 +2,12 @@ // at different verbosity levels. Demonstrates how variants can optimize tool // 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: // - deep-research: Verbose multi-paragraph descriptions with usage examples @@ -31,7 +36,9 @@ import ( func main() { // Deep-research variant: verbose descriptions with usage examples // 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{ Name: "search_papers", 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 // 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{ Name: "search_papers", 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 // 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{ Name: "search_papers", Description: "Search academic papers across arXiv, Semantic Scholar, and PubMed. " + diff --git a/go/sdk/examples/server/research/tools.go b/go/sdk/examples/server/research/tools.go index 46cf68d..0ab8a2a 100644 --- a/go/sdk/examples/server/research/tools.go +++ b/go/sdk/examples/server/research/tools.go @@ -2,6 +2,8 @@ package main import ( "context" + "fmt" + "time" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -27,7 +29,33 @@ type SearchPapersOutput struct { 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{ Papers: []Paper{ { @@ -71,7 +99,25 @@ type GetPaperOutput struct { 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{ Paper: PaperDetail{ ID: in.ID, @@ -98,11 +144,30 @@ type SummarizeOutput struct { 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 if style == "" { 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{ 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, diff --git a/go/sdk/variants/backend.go b/go/sdk/variants/backend.go index 2d3ffbd..615788d 100644 --- a/go/sdk/variants/backend.go +++ b/go/sdk/variants/backend.go @@ -6,7 +6,6 @@ package variants import ( "context" - "fmt" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -31,30 +30,75 @@ type backend interface { // inMemoryBackend connects to a co-located *mcp.Server via in-memory // transports. // -// Limitation: list-changed and resource-updated notifications from -// inner servers are silently dropped. The mcp.ServerSession API only -// exposes NotifyProgress and Log; there is no public method to send -// tool/resource/prompt list-changed or resource-updated notifications -// 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. +// The sending redirect middleware intercepts all outgoing messages from +// the inner server and forwards them through the front session. This +// covers notifications (progress, log) and server-to-client requests +// (Elicit, CreateMessage, ListRoots) during request handling. // +// 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 -// statically configured (tools registered at startup). If dynamic -// capability changes are needed, this will require the Go MCP SDK to -// expose generic notification sending on ServerSession. +// statically configured (tools registered at startup). 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 // 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 // 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 - // The middleware is identity (returns next unmodified), so the handler // chain is unchanged, no extra hop introduced even if called multiple times. server.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { @@ -62,22 +106,15 @@ func captureMCPMethodHandler(server *mcp.Server) (mcp.MethodHandler, error) { return next }) - if handler == nil { - return nil, fmt.Errorf("failed to capture backend MCP method handler") - } - return handler, nil + return handler } // connect creates an in-memory transport pair and connects the inner server. -// Requests bypass the transport via serverHandler to preserve context values. -// The transport is kept alive only for notification forwarding (progress, -// logging) from the inner server to the front client. +// Requests bypass the transport via mcpMethodHandler to preserve context +// values. Notifications are redirected by the sending middleware registered +// 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) { - mcpMethodHandler, err := captureMCPMethodHandler(b.server) - if err != nil { - return nil, err - } - serverTransport, clientSideTransport := mcp.NewInMemoryTransports() 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 } + // 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{ Name: "variant-proxy-client", Version: "1.0.0", }, &mcp.ClientOptions{ - ProgressNotificationHandler: func(ctx context.Context, req *mcp.ProgressNotificationClientRequest) { - if frontSession != nil { - 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) {}, + ElicitationHandler: nopElicit, + CreateMessageHandler: nopSampling, }) clientSession, err := client.Connect(ctx, clientSideTransport, nil) @@ -113,11 +150,19 @@ func (b *inMemoryBackend) connect(ctx context.Context, variant ServerVariant, fr 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{ backendSession: &backendSession{ - variantID: variant.ID, + variantID: b.variantID, serverSession: serverSession, - mcpMethodHandler: mcpMethodHandler, + mcpMethodHandler: b.mcpMethodHandler, }, cleanupFn: func() { clientSession.Close() diff --git a/go/sdk/variants/context_test.go b/go/sdk/variants/context_test.go new file mode 100644 index 0000000..d4de6ec --- /dev/null +++ b/go/sdk/variants/context_test.go @@ -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") +} diff --git a/go/sdk/variants/notification_test.go b/go/sdk/variants/notification_test.go new file mode 100644 index 0000000..6deab61 --- /dev/null +++ b/go/sdk/variants/notification_test.go @@ -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() + } + }) + } +} diff --git a/go/sdk/variants/sending_test.go b/go/sdk/variants/sending_test.go new file mode 100644 index 0000000..540f963 --- /dev/null +++ b/go/sdk/variants/sending_test.go @@ -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) +} diff --git a/go/sdk/variants/server.go b/go/sdk/variants/server.go index 9052e87..7405a97 100644 --- a/go/sdk/variants/server.go +++ b/go/sdk/variants/server.go @@ -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 // --------------------------------------------------------------------------- diff --git a/go/sdk/variants/session.go b/go/sdk/variants/session.go index 12b869f..add4563 100644 --- a/go/sdk/variants/session.go +++ b/go/sdk/variants/session.go @@ -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, // preserving any existing metadata. func injectVariantMeta(p mcp.Params, variantID string) {