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:
Steven Chen
GitHub
parent
150c110128
commit
e5e8f117e8
3 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.
|
||||
|
||||
**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
|
||||
|
||||
|
|
|
|||
|
|
@ -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. " +
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue