diff --git a/go/sdk/README.md b/go/sdk/README.md index 56d7d0c..ab7a783 100644 --- a/go/sdk/README.md +++ b/go/sdk/README.md @@ -1,3 +1,259 @@ # MCP Variants - Go Implementation -This will contain the Go implementation of the MCP Variants based on [SEP-2053](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2053). \ No newline at end of file +Go implementation of [SEP-2053: Server Variants](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2053) — a server-level variant mechanism for MCP that lets servers expose multiple tool/resource/prompt sets simultaneously, selectable per-request via `_meta`. + +## Quick Start + +```go +package main + +import ( + "log" + "net/http" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/experimental-ext-variants/go/sdk/variants" +) + +func main() { + // Code review variant: PR-focused tools + codeReviewServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "list_pull_requests", + Description: "List open pull requests, optionally filtered by author", + }, listPullRequests) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "get_diff", + Description: "Get the diff for a pull request", + }, getDiff) + + // Project management variant: issue-focused tools + pmServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "list_issues", + Description: "List issues, optionally filtered by state and labels", + }, listIssues) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "create_issue", + Description: "Create a new issue with title, body, and optional labels", + }, createIssue) + + vs := variants.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}). + WithVariant(variants.ServerVariant{ + ID: "code-review", + Description: "Pull request and code review operations. Includes diff viewing, review comments, approval workflows, and merge controls.", + Hints: map[string]string{"domain": "code-review", "accessLevel": "read-write"}, + Status: variants.Stable, + }, codeReviewServer, 0). + WithVariant(variants.ServerVariant{ + ID: "project-management", + Description: "Issue and project tracking operations. Includes issue CRUD, labels, milestones, assignments, and project boards.", + Hints: map[string]string{"domain": "project-management", "accessLevel": "read-write"}, + Status: variants.Stable, + }, pmServer, 1) + + // Serve over HTTP (or use vs.Run(ctx, &mcp.StdioTransport{}) for stdio) + handler := variants.NewStreamableHTTPHandler(vs, nil) + log.Fatal(http.ListenAndServe(":8080", handler)) +} +``` + +Clients select a variant per-request via `_meta`: + +```json +{ + "method": "tools/list", + "params": { + "_meta": { + "io.modelcontextprotocol/server-variant": "code-review" + } + } +} +``` + +Clients that don't know about variants get the default (first-ranked) variant automatically. + +## How It Works + +During `initialize`, the server responds with ranked `availableVariants`: + +```json +{ + "capabilities": { + "experimental": { + "io.modelcontextprotocol/server-variants": { + "availableVariants": [ + { + "id": "code-review", + "description": "Pull request and code review operations...", + "hints": { "domain": "code-review", "accessLevel": "read-write" }, + "status": "stable" + }, + { + "id": "project-management", + "description": "Issue and project tracking operations...", + "hints": { "domain": "project-management", "accessLevel": "read-write" }, + "status": "stable" + } + ], + "moreVariantsAvailable": false + } + } + } +} +``` + +Each subsequent request can target a specific variant via `_meta`. The server routes the request to the appropriate backing `mcp.Server`: + +``` +Client ── transport ──▸ variants.Server ──▸ sessionMiddleware + │ + ┌────────────┼────────────┐ + │ │ │ + initialize route by pass through + (create _meta (ping, etc.) + per-session variant + connections) │ + ▼ + dispatcher + │ + ┌────────────┬────────┴────────┬────────────┐ + ▼ ▼ ▼ ▼ + code-review project-mgmt security-readonly ci-automation + (mcp.Server) (mcp.Server) (mcp.Server) (mcp.Server) +``` + +In **stateful mode** (default, stdio and HTTP), per-session inner connections are created during `initialize` and scoped to the client session's lifetime. In **stateless mode** (via `NewStreamableHTTPHandler` with `Stateless: true`), a single set of shared connections is created at construction and reused across all requests. + +## Features + +- **Variant isolation**: each variant is a full `mcp.Server` with its own tools, resources, and prompts +- **Per-request selection**: variant chosen via `_meta` field, no session state needed +- **Default fallback**: clients without variant support get the first-ranked variant +- **Custom ranking**: provide a `RankingFunc` to rank variants based on client hints +- **Cursor scoping**: pagination cursors are variant-scoped and cannot be reused across variants (per SEP-2053) +- **Namespace scoping**: tool names, prompt names, and resource URIs resolve within the active variant's namespace; errors include `activeVariant` in error data +- **Notification forwarding**: progress and logging notifications from inner servers are forwarded to the front client with variant metadata injected +- **HTTP and stdio**: works with both `StdioTransport` and `StreamableHTTPHandler` + +## Examples + +See [`examples/server/`](examples/server/) for runnable examples: + +- [`variants/`](examples/server/variants/) — stdio transport (single client) +- [`variants-http/`](examples/server/variants-http/) — HTTP transport (multiple concurrent clients) + +## API + +### Server + +#### `variants.NewServer(impl *mcp.Implementation) *Server` + +Creates a new variant-aware server with no registered variants. `impl` must not be nil. + +#### `(*Server).WithVariant(v ServerVariant, mcpServer *mcp.Server, priority int) *Server` + +Registers a variant backed by an in-memory `mcp.Server`. `priority` determines the default ordering when no `RankingFunc` is set — lower values rank higher (0 = highest priority). Panics on duplicate variant IDs. Returns the receiver for chaining. + +#### `(*Server).WithRanking(fn RankingFunc) *Server` + +Sets a custom ranking function used to order variants based on client hints during initialization. If nil, variants are ordered by priority value. + +#### `(*Server).Variants() []ServerVariant` + +Returns a copy of all registered variants in registration order. + +#### `(*Server).RankedVariants(ctx context.Context, hints VariantHints) []ServerVariant` + +Returns registered variants ranked by the configured `RankingFunc` (or default priority-based ranking). + +#### `(*Server).Run(ctx context.Context, t mcp.Transport) error` + +Starts the server on the given transport (e.g., `&mcp.StdioTransport{}`). For multi-client HTTP support, use `NewStreamableHTTPHandler` instead. + +#### `(*Server).Close() error` + +Releases resources held by all registered backends. + +#### `variants.NewStreamableHTTPHandler(vs *Server, opts *mcp.StreamableHTTPOptions) *mcp.StreamableHTTPHandler` + +Returns an `http.Handler` for serving multiple concurrent clients over HTTP. Pass `&mcp.StreamableHTTPOptions{Stateless: true}` for stateless mode. + +### Types + +#### `ServerVariant` + +Describes a selectable variant: + +```go +type ServerVariant struct { + ID string `json:"id"` + Description string `json:"description"` + Hints map[string]string `json:"hints,omitempty"` + Status VariantStatus `json:"status,omitempty"` + DeprecationInfo *DeprecationInfo `json:"deprecationInfo,omitempty"` +} +``` + +`Priority() int` returns the priority value set during registration. + +#### `VariantStatus` + +```go +const ( + Stable VariantStatus = "stable" + Experimental VariantStatus = "experimental" + Deprecated VariantStatus = "deprecated" +) +``` + +#### `DeprecationInfo` + +Migration guidance for deprecated variants: + +```go +type DeprecationInfo struct { + Message string `json:"message"` + Replacement string `json:"replacement,omitempty"` + RemovalDate string `json:"removalDate,omitempty"` +} +``` + +#### `VariantHints` + +Client-provided hints for variant selection, sent during `initialize`: + +```go +type VariantHints struct { + Description string `json:"description,omitempty"` + Hints map[string]any `json:"hints,omitempty"` +} +``` + +#### `HintValue[T any](h VariantHints, key string) (T, bool)` + +Generic helper to extract a typed value from a `VariantHints` map. + +#### `RankingFunc` + +```go +type RankingFunc func(ctx context.Context, hints VariantHints, variants []ServerVariant) []ServerVariant +``` + +Called during initialization to rank variants based on client hints. Must return variants sorted by relevance, most appropriate first. + +#### Well-known hint keys + +| Constant | Key | Example values | +|---|---|---| +| `HintModelFamily` | `"modelFamily"` | `"anthropic"`, `"openai"`, `"local"`, `"any"` | +| `HintUseCase` | `"useCase"` | `"autonomous-agent"`, `"ide"`, `"chat"` | +| `HintContextSize` | `"contextSize"` | `"compact"`, `"standard"`, `"verbose"` | +| `HintRenderingCapabilities` | `"renderingCapabilities"` | `"rich"`, `"markdown"`, `"text-only"` | +| `HintLanguageOptimization` | `"languageOptimization"` | `"en"`, `"multilingual"`, `"code-focused"` | + +## Known Limitations + +- **Default variant resolution**: When a client omits `_meta` variant selection, the server re-ranks variants with empty hints to determine the default. This may differ from the ranking returned during `initialize` (where client hints were used). Per SEP-2053, the default should be the first variant from the `initialize` response. To fix this, the per-session ranked order needs to be stored during `initialize` and reused for subsequent requests. +- **List-changed notifications**: Dynamic capability changes from inner servers (tool/resource/prompt list changes) are not forwarded to front clients. The Go MCP SDK does not expose generic notification sending on `ServerSession`. In practice this is acceptable because inner servers are typically statically configured. +- **HTTP and remote backends**: `WithHTTPVariant` and `WithRemoteVariant` are not yet implemented. diff --git a/go/sdk/examples/.gitkeep b/go/sdk/examples/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/go/sdk/examples/server/exampletools/tools.go b/go/sdk/examples/server/exampletools/tools.go new file mode 100644 index 0000000..85561ec --- /dev/null +++ b/go/sdk/examples/server/exampletools/tools.go @@ -0,0 +1,254 @@ +// Package exampletools defines the shared tool types and handlers used by the +// variant-aware MCP server examples. The tools simulate a developer +// productivity platform (à la GitHub/GitLab) with four variants: +// +// - code-review: PR operations, diffs, and review comments +// - project-management: issue tracking, labels, and assignments +// - security-readonly: security scanning alerts and advisories (read-only) +// - ci-automation: CI/CD workflow management and dispatch +package exampletools + +import ( + "context" + "fmt" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// -- Code Review variant tools ------------------------------------------------ + +type ListPullRequestsInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + State string `json:"state,omitempty" jsonschema:"filter by state: open, closed, all (default: open)"` + Author string `json:"author,omitempty" jsonschema:"filter by PR author username"` +} + +type PullRequest struct { + Number int `json:"number"` + Title string `json:"title"` + Author string `json:"author"` + State string `json:"state"` +} + +type ListPullRequestsOutput struct { + PullRequests []PullRequest `json:"pullRequests"` +} + +func ListPullRequests(_ context.Context, _ *mcp.CallToolRequest, in ListPullRequestsInput) (*mcp.CallToolResult, ListPullRequestsOutput, error) { + return nil, ListPullRequestsOutput{ + PullRequests: []PullRequest{ + {Number: 42, Title: "Add retry logic to API client", Author: "alice", State: "open"}, + {Number: 38, Title: "Fix race condition in cache layer", Author: "bob", State: "open"}, + {Number: 35, Title: "Update dependencies to latest", Author: "alice", State: "open"}, + }, + }, nil +} + +type GetDiffInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Number int `json:"number" jsonschema:"pull request number"` +} + +type GetDiffOutput struct { + Diff string `json:"diff"` + Files []string `json:"filesChanged"` + Additions int `json:"additions"` + Deletions int `json:"deletions"` +} + +func GetDiff(_ context.Context, _ *mcp.CallToolRequest, in GetDiffInput) (*mcp.CallToolResult, GetDiffOutput, error) { + return nil, GetDiffOutput{ + Diff: "--- a/client.go\n+++ b/client.go\n@@ -45,6 +45,12 @@\n+ for attempt := 0; attempt < maxRetries; attempt++ {\n+ resp, err := c.do(req)\n+ if err == nil { return resp, nil }\n+ time.Sleep(backoff(attempt))\n+ }", + Files: []string{"client.go", "client_test.go"}, + Additions: 12, + Deletions: 3, + }, nil +} + +type AddReviewCommentInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Number int `json:"number" jsonschema:"pull request number"` + Path string `json:"path" jsonschema:"file path to comment on"` + Line int `json:"line" jsonschema:"line number to comment on"` + Body string `json:"body" jsonschema:"review comment body (markdown)"` +} + +type AddReviewCommentOutput struct { + CommentID int `json:"commentId"` + URL string `json:"url"` +} + +func AddReviewComment(_ context.Context, _ *mcp.CallToolRequest, in AddReviewCommentInput) (*mcp.CallToolResult, AddReviewCommentOutput, error) { + return nil, AddReviewCommentOutput{ + CommentID: 1001, + URL: fmt.Sprintf("https://github.com/%s/pull/%d#discussion_r1001", in.Repo, in.Number), + }, nil +} + +// -- Project Management variant tools ----------------------------------------- + +type ListIssuesInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + State string `json:"state,omitempty" jsonschema:"filter by state: open, closed, all (default: open)"` + Labels []string `json:"labels,omitempty" jsonschema:"filter by label names"` +} + +type Issue struct { + Number int `json:"number"` + Title string `json:"title"` + State string `json:"state"` + Labels []string `json:"labels"` + Assignee string `json:"assignee,omitempty"` +} + +type ListIssuesOutput struct { + Issues []Issue `json:"issues"` +} + +func ListIssues(_ context.Context, _ *mcp.CallToolRequest, in ListIssuesInput) (*mcp.CallToolResult, ListIssuesOutput, error) { + return nil, ListIssuesOutput{ + Issues: []Issue{ + {Number: 101, Title: "API rate limiting returns wrong status code", State: "open", Labels: []string{"bug", "api"}, Assignee: "alice"}, + {Number: 98, Title: "Add support for webhook retries", State: "open", Labels: []string{"enhancement"}, Assignee: ""}, + {Number: 95, Title: "Document authentication flow", State: "open", Labels: []string{"documentation"}, Assignee: "carol"}, + }, + }, nil +} + +type CreateIssueInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Title string `json:"title" jsonschema:"issue title"` + Body string `json:"body" jsonschema:"issue body (markdown)"` + Labels []string `json:"labels,omitempty" jsonschema:"labels to apply"` +} + +type CreateIssueOutput struct { + Number int `json:"number"` + URL string `json:"url"` +} + +func CreateIssue(_ context.Context, _ *mcp.CallToolRequest, in CreateIssueInput) (*mcp.CallToolResult, CreateIssueOutput, error) { + return nil, CreateIssueOutput{ + Number: 102, + URL: fmt.Sprintf("https://github.com/%s/issues/102", in.Repo), + }, nil +} + +type AddLabelInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Number int `json:"number" jsonschema:"issue number"` + Labels []string `json:"labels" jsonschema:"labels to add"` +} + +type AddLabelOutput struct { + Labels []string `json:"currentLabels"` +} + +func AddLabel(_ context.Context, _ *mcp.CallToolRequest, in AddLabelInput) (*mcp.CallToolResult, AddLabelOutput, error) { + return nil, AddLabelOutput{ + Labels: append([]string{"bug", "api"}, in.Labels...), + }, nil +} + +// -- Security variant tools (read-only) --------------------------------------- + +type ListSecurityAlertsInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Severity string `json:"severity,omitempty" jsonschema:"filter by severity: critical, high, medium, low (default: all)"` + State string `json:"state,omitempty" jsonschema:"filter by state: open, dismissed, fixed (default: open)"` +} + +type SecurityAlert struct { + Number int `json:"number"` + Rule string `json:"rule"` + Severity string `json:"severity"` + State string `json:"state"` + Path string `json:"path"` +} + +type ListSecurityAlertsOutput struct { + Alerts []SecurityAlert `json:"alerts"` +} + +func ListSecurityAlerts(_ context.Context, _ *mcp.CallToolRequest, in ListSecurityAlertsInput) (*mcp.CallToolResult, ListSecurityAlertsOutput, error) { + return nil, ListSecurityAlertsOutput{ + Alerts: []SecurityAlert{ + {Number: 1, Rule: "sql-injection", Severity: "critical", State: "open", Path: "src/db/query.go"}, + {Number: 2, Rule: "hardcoded-secret", Severity: "high", State: "open", Path: "config/settings.go"}, + {Number: 3, Rule: "insecure-hash", Severity: "medium", State: "open", Path: "pkg/auth/hash.go"}, + }, + }, nil +} + +type GetAdvisoryInput struct { + AdvisoryID string `json:"advisoryId" jsonschema:"security advisory identifier (e.g. GHSA-xxxx-xxxx-xxxx)"` +} + +type GetAdvisoryOutput struct { + ID string `json:"id"` + Summary string `json:"summary"` + Severity string `json:"severity"` + CVEs []string `json:"cves"` + AffectedPkg string `json:"affectedPackage"` + PatchedIn string `json:"patchedIn"` +} + +func GetAdvisory(_ context.Context, _ *mcp.CallToolRequest, in GetAdvisoryInput) (*mcp.CallToolResult, GetAdvisoryOutput, error) { + return nil, GetAdvisoryOutput{ + ID: in.AdvisoryID, + Summary: "Remote code execution via crafted request payload", + Severity: "critical", + CVEs: []string{"CVE-2024-12345"}, + AffectedPkg: "example.com/server", + PatchedIn: "v1.2.4", + }, nil +} + +// -- CI/CD variant tools ------------------------------------------------------ + +type ListWorkflowRunsInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Workflow string `json:"workflow,omitempty" jsonschema:"filter by workflow name or filename"` + Status string `json:"status,omitempty" jsonschema:"filter by status: completed, in_progress, queued (default: all)"` +} + +type WorkflowRun struct { + ID int `json:"id"` + Workflow string `json:"workflow"` + Status string `json:"status"` + Conclusion string `json:"conclusion,omitempty"` + Branch string `json:"branch"` +} + +type ListWorkflowRunsOutput struct { + Runs []WorkflowRun `json:"runs"` +} + +func ListWorkflowRuns(_ context.Context, _ *mcp.CallToolRequest, in ListWorkflowRunsInput) (*mcp.CallToolResult, ListWorkflowRunsOutput, error) { + return nil, ListWorkflowRunsOutput{ + Runs: []WorkflowRun{ + {ID: 5001, Workflow: "ci.yml", Status: "completed", Conclusion: "success", Branch: "main"}, + {ID: 5002, Workflow: "ci.yml", Status: "in_progress", Branch: "feat/new-api"}, + {ID: 5003, Workflow: "deploy.yml", Status: "completed", Conclusion: "failure", Branch: "main"}, + }, + }, nil +} + +type TriggerWorkflowInput struct { + Repo string `json:"repo" jsonschema:"repository in owner/name format"` + Workflow string `json:"workflow" jsonschema:"workflow filename (e.g. ci.yml)"` + Ref string `json:"ref" jsonschema:"git ref to run against (branch, tag, or SHA)"` + Inputs map[string]string `json:"inputs,omitempty" jsonschema:"workflow input parameters"` +} + +type TriggerWorkflowOutput struct { + RunID int `json:"runId"` + URL string `json:"url"` +} + +func TriggerWorkflow(_ context.Context, _ *mcp.CallToolRequest, in TriggerWorkflowInput) (*mcp.CallToolResult, TriggerWorkflowOutput, error) { + return nil, TriggerWorkflowOutput{ + RunID: 5004, + URL: fmt.Sprintf("https://github.com/%s/actions/runs/5004", in.Repo), + }, nil +} diff --git a/go/sdk/examples/server/variants-http/main.go b/go/sdk/examples/server/variants-http/main.go new file mode 100644 index 0000000..6e7a80e --- /dev/null +++ b/go/sdk/examples/server/variants-http/main.go @@ -0,0 +1,111 @@ +// Example: A developer productivity platform (à la GitHub) that exposes +// different tool sets for different agent types using MCP server variants, +// served over HTTP for multiple concurrent clients. +// +// Variants: +// - code-review: PR operations, diffs, and review comments +// - project-management: issue tracking, labels, and assignments +// - security-readonly: security scanning alerts and advisories (read-only) +// - ci-automation: CI/CD workflow management and dispatch +// +// Run: +// +// go run . +// +// Then connect any MCP client to http://localhost:8080. +package main + +import ( + "log" + "net/http" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/modelcontextprotocol/experimental-ext-variants/go/sdk/examples/server/exampletools" + "github.com/modelcontextprotocol/experimental-ext-variants/go/sdk/variants" +) + +func main() { + // Code review variant: PR-focused tools for code review agents + codeReviewServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "list_pull_requests", + Description: "List open pull requests, optionally filtered by author", + }, exampletools.ListPullRequests) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "get_diff", + Description: "Get the diff for a pull request, including changed files and line counts", + }, exampletools.GetDiff) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "add_review_comment", + Description: "Post a review comment on a specific line of a pull request", + }, exampletools.AddReviewComment) + + // Project management variant: issue-focused tools for PM agents + pmServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "list_issues", + Description: "List issues, optionally filtered by state and labels", + }, exampletools.ListIssues) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "create_issue", + Description: "Create a new issue with title, body, and optional labels", + }, exampletools.CreateIssue) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "add_label", + Description: "Add labels to an existing issue", + }, exampletools.AddLabel) + + // Security variant: read-only security scanning tools + securityServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(securityServer, &mcp.Tool{ + Name: "list_security_alerts", + Description: "List code scanning alerts for a repository", + }, exampletools.ListSecurityAlerts) + mcp.AddTool(securityServer, &mcp.Tool{ + Name: "get_advisory", + Description: "Get details of a security advisory", + }, exampletools.GetAdvisory) + + // CI/CD variant: workflow management tools for automation agents + ciServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(ciServer, &mcp.Tool{ + Name: "list_workflow_runs", + Description: "List recent workflow runs for a repository", + }, exampletools.ListWorkflowRuns) + mcp.AddTool(ciServer, &mcp.Tool{ + Name: "trigger_workflow", + Description: "Trigger a workflow dispatch event", + }, exampletools.TriggerWorkflow) + + vs := variants.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}). + WithVariant(variants.ServerVariant{ + ID: "code-review", + Description: "Pull request and code review operations. Includes diff viewing, review comments, approval workflows, and merge controls.", + Hints: map[string]string{"domain": "code-review", "accessLevel": "read-write"}, + Status: variants.Stable, + }, codeReviewServer, 0). + WithVariant(variants.ServerVariant{ + ID: "project-management", + Description: "Issue and project tracking operations. Includes issue CRUD, labels, milestones, assignments, and project boards.", + Hints: map[string]string{"domain": "project-management", "accessLevel": "read-write"}, + Status: variants.Stable, + }, pmServer, 1). + WithVariant(variants.ServerVariant{ + ID: "security-readonly", + Description: "Security scanning and vulnerability management. Read-only access to code scanning alerts, secret detection, and security advisories.", + Hints: map[string]string{"domain": "security", "accessLevel": "readonly"}, + Status: variants.Stable, + }, securityServer, 2). + WithVariant(variants.ServerVariant{ + ID: "ci-automation", + Description: "CI/CD workflow management. Trigger runs, monitor jobs, manage deployments designed for automation agents.", + Hints: map[string]string{"domain": "ci-cd", "accessLevel": "automation"}, + Status: variants.Stable, + }, ciServer, 3) + + handler := variants.NewStreamableHTTPHandler(vs, nil) + + log.Println("Listening on :8080") + log.Fatal(http.ListenAndServe(":8080", handler)) +} diff --git a/go/sdk/examples/server/variants/main.go b/go/sdk/examples/server/variants/main.go index 5232db3..812bf32 100644 --- a/go/sdk/examples/server/variants/main.go +++ b/go/sdk/examples/server/variants/main.go @@ -1,111 +1,108 @@ -// Example: Setting up a variant-aware MCP server with multiple variants. +// Example: A developer productivity platform (à la GitHub) that exposes +// different tool sets for different agent types using MCP server variants. +// +// Variants: +// - code-review: PR operations, diffs, and review comments +// - project-management: issue tracking, labels, and assignments +// - security-readonly: security scanning alerts and advisories (read-only) +// - ci-automation: CI/CD workflow management and dispatch +// +// Run over stdio: +// +// go run . package main import ( "context" "log" - mcp "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/experimental-ext-variants/go/sdk/examples/server/exampletools" "github.com/modelcontextprotocol/experimental-ext-variants/go/sdk/variants" ) -// -- Coding variant tool types ------------------------------------------------ - -type AnalyzeCodeInput struct { - Code string `json:"code" jsonschema:"source code to analyze"` - Language string `json:"language" jsonschema:"programming language"` -} - -type AnalyzeCodeOutput struct { - Issues []string `json:"issues"` - Suggestions []string `json:"suggestions"` -} - -func AnalyzeCode(_ context.Context, _ *mcp.CallToolRequest, in AnalyzeCodeInput) (*mcp.CallToolResult, AnalyzeCodeOutput, error) { - return nil, AnalyzeCodeOutput{ - Issues: []string{"unused variable on line 3"}, - Suggestions: []string{"consider using a switch statement"}, - }, nil -} - -type RefactorInput struct { - Code string `json:"code" jsonschema:"source code to refactor"` - Action string `json:"action" jsonschema:"refactoring action, e.g. extract-function"` -} - -type RefactorOutput struct { - Refactored string `json:"refactored"` -} - -func Refactor(_ context.Context, _ *mcp.CallToolRequest, in RefactorInput) (*mcp.CallToolResult, RefactorOutput, error) { - return nil, RefactorOutput{Refactored: "// refactored\n" + in.Code}, nil -} - -// -- Compact variant tool types ----------------------------------------------- - -type SummarizeInput struct { - Text string `json:"text" jsonschema:"text to summarize"` -} - -type SummarizeOutput struct { - Summary string `json:"summary"` -} - -func Summarize(_ context.Context, _ *mcp.CallToolRequest, in SummarizeInput) (*mcp.CallToolResult, SummarizeOutput, error) { - return nil, SummarizeOutput{Summary: in.Text[:min(len(in.Text), 50)]}, nil -} - -type LookupInput struct { - Query string `json:"query" jsonschema:"lookup query"` -} - -type LookupOutput struct { - Result string `json:"result"` -} - -func Lookup(_ context.Context, _ *mcp.CallToolRequest, in LookupInput) (*mcp.CallToolResult, LookupOutput, error) { - return nil, LookupOutput{Result: "result for: " + in.Query}, nil -} - -// ----------------------------------------------------------------------------- - func main() { - // Server 1: coding-focused tools - server1 := mcp.NewServer(&mcp.Implementation{Name: "my-server", Version: "v1.0.0"}, nil) - mcp.AddTool(server1, &mcp.Tool{Name: "analyze_code", Description: "Perform static analysis on source code"}, AnalyzeCode) - mcp.AddTool(server1, &mcp.Tool{Name: "refactor", Description: "Apply a refactoring action to source code"}, Refactor) + // Code review variant: PR-focused tools for code review agents + codeReviewServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "list_pull_requests", + Description: "List open pull requests, optionally filtered by author", + }, exampletools.ListPullRequests) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "get_diff", + Description: "Get the diff for a pull request, including changed files and line counts", + }, exampletools.GetDiff) + mcp.AddTool(codeReviewServer, &mcp.Tool{ + Name: "add_review_comment", + Description: "Post a review comment on a specific line of a pull request", + }, exampletools.AddReviewComment) - // Server 2: compact / minimal tools - server2 := mcp.NewServer(&mcp.Implementation{Name: "my-server", Version: "v1.0.0"}, nil) - mcp.AddTool(server2, &mcp.Tool{Name: "summarize", Description: "Summarize text"}, Summarize) - mcp.AddTool(server2, &mcp.Tool{Name: "lookup", Description: "Quick fact lookup"}, Lookup) + // Project management variant: issue-focused tools for PM agents + pmServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "list_issues", + Description: "List issues, optionally filtered by state and labels", + }, exampletools.ListIssues) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "create_issue", + Description: "Create a new issue with title, body, and optional labels", + }, exampletools.CreateIssue) + mcp.AddTool(pmServer, &mcp.Tool{ + Name: "add_label", + Description: "Add labels to an existing issue", + }, exampletools.AddLabel) - coding := variants.ServerVariant{ - ID: "coding-assistant", - Description: "Optimized for coding workflows", - Hints: map[string]string{"useCase": "coding", "contextSize": "standard"}, - Status: variants.Stable, - } + // Security variant: read-only security scanning tools + securityServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(securityServer, &mcp.Tool{ + Name: "list_security_alerts", + Description: "List code scanning alerts for a repository", + }, exampletools.ListSecurityAlerts) + mcp.AddTool(securityServer, &mcp.Tool{ + Name: "get_advisory", + Description: "Get details of a security advisory", + }, exampletools.GetAdvisory) - compact := variants.ServerVariant{ - ID: "compact", - Description: "Minimal token usage", - Hints: map[string]string{"contextSize": "small"}, - Status: variants.Experimental, - } + // CI/CD variant: workflow management tools for automation agents + ciServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}, nil) + mcp.AddTool(ciServer, &mcp.Tool{ + Name: "list_workflow_runs", + Description: "List recent workflow runs for a repository", + }, exampletools.ListWorkflowRuns) + mcp.AddTool(ciServer, &mcp.Tool{ + Name: "trigger_workflow", + Description: "Trigger a workflow dispatch event", + }, exampletools.TriggerWorkflow) - vs := variants.NewServer(). - WithVariant(coding, server1, 0). - WithVariant(compact, server2, 1). - WithRanking(func(_ context.Context, _ variants.VariantHints, vs []variants.ServerVariant) []variants.ServerVariant { - // rank based on client hints, return sorted - // note: this is optional, defults to ranking by priority. - return vs - }) + vs := variants.NewServer(&mcp.Implementation{Name: "devplatform", Version: "v1.0.0"}). + WithVariant(variants.ServerVariant{ + ID: "code-review", + Description: "Pull request and code review operations. Includes diff viewing, review comments, approval workflows, and merge controls.", + Hints: map[string]string{"domain": "code-review", "accessLevel": "read-write"}, + Status: variants.Stable, + }, codeReviewServer, 0). + WithVariant(variants.ServerVariant{ + ID: "project-management", + Description: "Issue and project tracking operations. Includes issue CRUD, labels, milestones, assignments, and project boards.", + Hints: map[string]string{"domain": "project-management", "accessLevel": "read-write"}, + Status: variants.Stable, + }, pmServer, 1). + WithVariant(variants.ServerVariant{ + ID: "security-readonly", + Description: "Security scanning and vulnerability management. Read-only access to code scanning alerts, secret detection, and security advisories.", + Hints: map[string]string{"domain": "security", "accessLevel": "readonly"}, + Status: variants.Stable, + }, securityServer, 2). + WithVariant(variants.ServerVariant{ + ID: "ci-automation", + Description: "CI/CD workflow management. Trigger runs, monitor jobs, manage deployments designed for automation agents.", + Hints: map[string]string{"domain": "ci-cd", "accessLevel": "automation"}, + Status: variants.Stable, + }, ciServer, 3) ctx := context.Background() - if err := vs.Run(ctx, nil); err != nil { + if err := vs.Run(ctx, &mcp.StdioTransport{}); err != nil { log.Fatal(err) } } diff --git a/go/sdk/go.mod b/go/sdk/go.mod index 12f14d6..e4a7abe 100644 --- a/go/sdk/go.mod +++ b/go/sdk/go.mod @@ -4,10 +4,16 @@ go 1.23.0 toolchain go1.24.3 -require github.com/modelcontextprotocol/go-sdk v1.2.0 +require ( + github.com/modelcontextprotocol/go-sdk v1.2.0 + github.com/stretchr/testify v1.11.1 +) require ( + github.com/davecgh/go-spew v1.1.1 // indirect github.com/google/jsonschema-go v0.3.0 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect golang.org/x/oauth2 v0.30.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go/sdk/go.sum b/go/sdk/go.sum index 3bc80f3..2c0fd3d 100644 --- a/go/sdk/go.sum +++ b/go/sdk/go.sum @@ -1,3 +1,5 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= @@ -6,9 +8,17 @@ github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIy github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/modelcontextprotocol/go-sdk v1.2.0 h1:Y23co09300CEk8iZ/tMxIX1dVmKZkzoSBZOpJwUnc/s= github.com/modelcontextprotocol/go-sdk v1.2.0/go.mod h1:6fM3LCm3yV7pAs8isnKLn07oKtB0MP9LHd3DfAcKw10= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/go/sdk/internal/.gitkeep b/go/sdk/internal/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/go/sdk/variants/backend.go b/go/sdk/variants/backend.go new file mode 100644 index 0000000..bcc805b --- /dev/null +++ b/go/sdk/variants/backend.go @@ -0,0 +1,126 @@ +// 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" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// backend abstracts how a variant connects to its backing MCP server. +// The in-memory implementation lives in this file; future implementations +// (HTTP, remote) will satisfy the same interface. +type backend interface { + // connect creates a connection to the backing server and returns an + // innerConnection for dispatching requests. frontSession may be nil + // in stateless mode. + connect(ctx context.Context, variant ServerVariant, frontSession *mcp.ServerSession) (*innerConnection, error) + + // capabilities performs an ephemeral connect to discover the server's + // advertised capabilities, then tears down the probe connection. + capabilities(ctx context.Context) (*mcp.ServerCapabilities, error) + + // close releases any resources held by the backend. + close() error +} + +// 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. +// +// 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. +type inMemoryBackend struct { + server *mcp.Server +} + +// connect creates an in-memory transport pair, connects the inner server, +// and creates a proxy client to communicate with it. Notifications for +// progress and logging are forwarded directly to the front session with +// _meta variant ID injection. +func (b *inMemoryBackend) connect(ctx context.Context, variant ServerVariant, frontSession *mcp.ServerSession) (*innerConnection, error) { + serverTransport, clientSideTransport := mcp.NewInMemoryTransports() + + serverSession, err := b.server.Connect(ctx, serverTransport, nil) + if err != nil { + return nil, err + } + + 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) {}, + }) + + clientSession, err := client.Connect(ctx, clientSideTransport, nil) + if err != nil { + serverSession.Close() + return nil, err + } + + return &innerConnection{ + clientSession: clientSession, + cleanupFn: func() { + serverSession.Close() + }, + }, nil +} + +// capabilities performs an ephemeral in-memory connect to discover the +// server's advertised capabilities. +func (b *inMemoryBackend) capabilities(ctx context.Context) (*mcp.ServerCapabilities, error) { + st, ct := mcp.NewInMemoryTransports() + ss, err := b.server.Connect(ctx, st, nil) + if err != nil { + return nil, err + } + + c := mcp.NewClient(&mcp.Implementation{Name: "cap-probe", Version: "1.0.0"}, nil) + cs, err := c.Connect(ctx, ct, nil) + if err != nil { + ss.Close() + return nil, err + } + + var caps *mcp.ServerCapabilities + if ir := cs.InitializeResult(); ir != nil { + caps = ir.Capabilities + } + + cs.Close() + ss.Close() + return caps, nil +} + +// close is a no-op for in-memory backends. +func (b *inMemoryBackend) close() error { + return nil +} diff --git a/go/sdk/variants/dispatch.go b/go/sdk/variants/dispatch.go new file mode 100644 index 0000000..89c5a28 --- /dev/null +++ b/go/sdk/variants/dispatch.go @@ -0,0 +1,448 @@ +// 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/base64" + "encoding/json" + "errors" + "reflect" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// dispatcher routes requests to inner variant servers. In stateful mode one +// dispatcher exists per client session; in stateless mode a single dispatcher +// is shared across all requests. +type dispatcher struct { + server *Server + connections map[string]*innerConnection +} + +// handle dispatches a request to the appropriate inner variant server. +// Unknown methods are passed through to next. +func (d *dispatcher) handle(ctx context.Context, method string, req mcp.Request, next mcp.MethodHandler) (mcp.Result, error) { + switch method { + case "tools/list", "resources/list", "prompts/list", "resources/templates/list": + return d.handleList(ctx, method, req) + case "tools/call", "resources/read", "prompts/get": + return d.handleCall(ctx, method, req) + case "resources/subscribe": + return d.handleSubscribe(ctx, req) + case "resources/unsubscribe": + return d.handleUnsubscribe(ctx, req) + case "completion/complete": + return d.handleCompletion(ctx, req) + default: + return next(ctx, method, req) + } +} + +// createInvalidVariantError creates a spec-compliant error for invalid variant selection. +func (d *dispatcher) createInvalidVariantError(ctx context.Context, requestedVariant string) error { + ranked := d.server.RankedVariants(ctx, VariantHints{}) + availableIDs := make([]string, len(ranked)) + for i, v := range ranked { + availableIDs[i] = v.ID + } + + errorData := map[string]any{ + "requestedVariant": requestedVariant, + "availableVariants": availableIDs, + } + dataJSON, err := json.Marshal(errorData) + if err != nil { + dataJSON = []byte("{}") + } + + return &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "Invalid server variant", + Data: json.RawMessage(dataJSON), + } +} + +// variantIDFromMeta extracts the variant ID from the request's _meta field. +// Returns empty string if no variant is specified. Guards against typed-nil +// params (e.g. (*ListToolsParams)(nil) wrapped in the mcp.Params interface) +// which the SDK can produce for requests with no parameters. +func variantIDFromMeta(req mcp.Request) string { + params := req.GetParams() + if params == nil { + return "" + } + if v := reflect.ValueOf(params); v.Kind() == reflect.Ptr && v.IsNil() { + return "" + } + meta := params.GetMeta() + if meta == nil { + return "" + } + id, _ := meta[metaKeyVariant].(string) + return id +} + +// getSession extracts the variant ID from request _meta and returns the +// corresponding client session for dispatching. Falls back to the +// first-ranked variant when no variant is specified. +func (d *dispatcher) getSession(ctx context.Context, req mcp.Request) (string, *mcp.ClientSession, error) { + variantID := variantIDFromMeta(req) + + // If no variant specified, use first-ranked (default). + // + // BUG: This re-ranks with empty hints, which may differ from the + // ranking returned during initialize (where client hints were used). + // Per SEP-2053, the default should be the first variant from the + // initialize response. To fix this properly, the per-session ranked + // order should be stored during initialize and reused here. + if variantID == "" { + ranked := d.server.RankedVariants(ctx, VariantHints{}) + if len(ranked) == 0 { + return "", nil, errors.New("no variants available") + } + variantID = ranked[0].ID + } + + conn, ok := d.connections[variantID] + if !ok { + return "", nil, d.createInvalidVariantError(ctx, variantID) + } + + return variantID, conn.clientSession, nil +} + +// enrichError adds activeVariant to a jsonrpc.Error's Data field for +// variant-scoped resolution failures. Per SEP-2053 Implementation Notes: +// "Servers SHOULD include activeVariant in error data for variant-scoped +// resolution failures (unknown tool/prompt/resource, invalid cursor, invalid +// subscription context)." +// +// Only errors with codes -32602 (InvalidParams) or -32601 (MethodNotFound) +// are enriched; business-logic errors from tool execution are passed through +// unmodified. +func enrichError(err error, variantID string) error { + var jErr *jsonrpc.Error + if !errors.As(err, &jErr) { + return err + } + + // Only enrich resolution-class errors. + switch jErr.Code { + case jsonrpc.CodeInvalidParams, jsonrpc.CodeMethodNotFound: + // fall through to enrich + default: + return err + } + + // Parse existing data (if any) and inject activeVariant. + // Copy the error to avoid mutating the original + data := make(map[string]any) + if len(jErr.Data) > 0 { + _ = json.Unmarshal(jErr.Data, &data) + } + data["activeVariant"] = variantID + + enriched := &jsonrpc.Error{ + Code: jErr.Code, + Message: jErr.Message, + } + if encoded, mErr := json.Marshal(data); mErr == nil { + enriched.Data = json.RawMessage(encoded) + } + return enriched +} + +// --------------------------------------------------------------------------- +// List methods +// --------------------------------------------------------------------------- + +// handleList handles list methods by forwarding to the appropriate variant. +// Implements cursor scoping per SEP-2053: unwraps incoming cursors and wraps outgoing cursors. +func (d *dispatcher) handleList(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + variantID, session, err := d.getSession(ctx, req) + if err != nil { + return nil, err + } + + params := req.GetParams() + + switch method { + case "tools/list": + p, _ := params.(*mcp.ListToolsParams) + if p != nil && p.Cursor != "" { + innerCursor, err := unwrapCursor(p.Cursor, variantID) + if err != nil { + return nil, err + } + p.Cursor = innerCursor + } + result, err := session.ListTools(ctx, p) + if err != nil { + return nil, enrichError(err, variantID) + } + if result != nil && result.NextCursor != "" { + result.NextCursor = wrapCursor(result.NextCursor, variantID) + } + return result, nil + + case "resources/list": + p, _ := params.(*mcp.ListResourcesParams) + if p != nil && p.Cursor != "" { + innerCursor, err := unwrapCursor(p.Cursor, variantID) + if err != nil { + return nil, err + } + p.Cursor = innerCursor + } + result, err := session.ListResources(ctx, p) + if err != nil { + return nil, enrichError(err, variantID) + } + if result != nil && result.NextCursor != "" { + result.NextCursor = wrapCursor(result.NextCursor, variantID) + } + return result, nil + + case "prompts/list": + p, _ := params.(*mcp.ListPromptsParams) + if p != nil && p.Cursor != "" { + innerCursor, err := unwrapCursor(p.Cursor, variantID) + if err != nil { + return nil, err + } + p.Cursor = innerCursor + } + result, err := session.ListPrompts(ctx, p) + if err != nil { + return nil, enrichError(err, variantID) + } + if result != nil && result.NextCursor != "" { + result.NextCursor = wrapCursor(result.NextCursor, variantID) + } + return result, nil + + case "resources/templates/list": + p, _ := params.(*mcp.ListResourceTemplatesParams) + if p != nil && p.Cursor != "" { + innerCursor, err := unwrapCursor(p.Cursor, variantID) + if err != nil { + return nil, err + } + p.Cursor = innerCursor + } + result, err := session.ListResourceTemplates(ctx, p) + if err != nil { + return nil, enrichError(err, variantID) + } + if result != nil && result.NextCursor != "" { + result.NextCursor = wrapCursor(result.NextCursor, variantID) + } + return result, nil + + default: + return nil, errors.New("unsupported list method: " + method) + } +} + +// --------------------------------------------------------------------------- +// Call methods +// --------------------------------------------------------------------------- + +// handleCall handles call methods (tools/call, resources/read, prompts/get). +func (d *dispatcher) handleCall(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + variantID, session, err := d.getSession(ctx, req) + if err != nil { + return nil, err + } + + params := req.GetParams() + var result mcp.Result + switch method { + case "tools/call": + // The server SDK unmarshals tools/call params as *CallToolParamsRaw + // (with json.RawMessage arguments). Convert to *CallToolParams for + // the client-side call. + raw, _ := params.(*mcp.CallToolParamsRaw) + if raw == nil { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "missing or invalid tools/call params", + } + } + result, err = session.CallTool(ctx, &mcp.CallToolParams{ + Meta: raw.Meta, + Name: raw.Name, + Arguments: raw.Arguments, + }) + case "resources/read": + p, _ := params.(*mcp.ReadResourceParams) + if p == nil { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "missing or invalid resources/read params", + } + } + result, err = session.ReadResource(ctx, p) + case "prompts/get": + p, _ := params.(*mcp.GetPromptParams) + if p == nil { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "missing or invalid prompts/get params", + } + } + result, err = session.GetPrompt(ctx, p) + default: + return nil, errors.New("unsupported call method: " + method) + } + + if err != nil { + return nil, enrichError(err, variantID) + } + return result, nil +} + +// --------------------------------------------------------------------------- +// Subscription methods +// --------------------------------------------------------------------------- + +// handleSubscribe handles resources/subscribe. +func (d *dispatcher) handleSubscribe(ctx context.Context, req mcp.Request) (mcp.Result, error) { + variantID, session, err := d.getSession(ctx, req) + if err != nil { + return nil, err + } + + params, _ := req.GetParams().(*mcp.SubscribeParams) + if params == nil { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "missing or invalid resources/subscribe params", + } + } + if err := session.Subscribe(ctx, params); err != nil { + return nil, enrichError(err, variantID) + } + return nil, nil +} + +// handleUnsubscribe handles resources/unsubscribe. +// Per SEP-2053: "Servers MUST continue to accept resources/unsubscribe for +// existing subscription ids even if the underlying resource is no longer available." +func (d *dispatcher) handleUnsubscribe(ctx context.Context, req mcp.Request) (mcp.Result, error) { + variantID, session, err := d.getSession(ctx, req) + if err != nil { + return nil, err + } + + params, _ := req.GetParams().(*mcp.UnsubscribeParams) + if params == nil { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "missing or invalid resources/unsubscribe params", + } + } + if err := session.Unsubscribe(ctx, params); err != nil { + return nil, enrichError(err, variantID) + } + return nil, nil +} + +// --------------------------------------------------------------------------- +// Completion +// --------------------------------------------------------------------------- + +// handleCompletion handles completion/complete. +func (d *dispatcher) handleCompletion(ctx context.Context, req mcp.Request) (mcp.Result, error) { + variantID, session, err := d.getSession(ctx, req) + if err != nil { + return nil, err + } + + params, _ := req.GetParams().(*mcp.CompleteParams) + if params == nil { + return nil, &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "missing or invalid completion/complete params", + } + } + result, err := session.Complete(ctx, params) + if err != nil { + return nil, enrichError(err, variantID) + } + return result, nil +} + +// --------------------------------------------------------------------------- +// Cursor scoping +// --------------------------------------------------------------------------- + +// variantCursor wraps pagination cursors with variant ID for scoping. +// Per SEP-2053: "Cursors MUST be treated as opaque and variant-scoped" +type variantCursor struct { + VariantID string `json:"v"` + InnerCursor string `json:"c"` +} + +// wrapCursor wraps a cursor from an inner server with the variant ID. +// Returns empty string if the inner cursor is empty. +func wrapCursor(cursor string, variantID string) string { + if cursor == "" { + return "" + } + wrapped := variantCursor{ + VariantID: variantID, + InnerCursor: cursor, + } + data, err := json.Marshal(wrapped) + if err != nil { + // Should never happen with simple struct + return "" + } + return base64.StdEncoding.EncodeToString(data) +} + +// unwrapCursor validates and unwraps a cursor for the expected variant. +// Returns the inner cursor if valid, or an error if the cursor is invalid +// or belongs to a different variant. +func unwrapCursor(cursor string, expectedVariant string) (string, error) { + if cursor == "" { + return "", nil + } + + data, err := base64.StdEncoding.DecodeString(cursor) + if err != nil { + return "", &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "Invalid cursor format", + } + } + + var wrapped variantCursor + if err := json.Unmarshal(data, &wrapped); err != nil { + return "", &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "Invalid cursor format", + } + } + + if wrapped.VariantID != expectedVariant { + errorData := map[string]any{ + "cursorVariant": wrapped.VariantID, + "requestedVariant": expectedVariant, + } + dataJSON, _ := json.Marshal(errorData) + + return "", &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "Cursor invalid for requested variant", + Data: json.RawMessage(dataJSON), + } + } + + return wrapped.InnerCursor, nil +} diff --git a/go/sdk/variants/ranking.go b/go/sdk/variants/ranking.go index 9012af3..b548351 100644 --- a/go/sdk/variants/ranking.go +++ b/go/sdk/variants/ranking.go @@ -12,15 +12,17 @@ import ( // defaultRankingFunc is the built-in ranking function used when no custom // RankingFunc is provided. It sorts variants by priority (lowest first), // using stable-before-experimental-before-deprecated as a tiebreaker. +// +// The input slice is already a copy (from RankedVariants), so sorting +// in place is safe. func defaultRankingFunc(_ context.Context, _ VariantHints, vs []ServerVariant) []ServerVariant { - ranked := slices.Clone(vs) - slices.SortStableFunc(ranked, func(a, b ServerVariant) int { + slices.SortStableFunc(vs, func(a, b ServerVariant) int { if a.Priority() != b.Priority() { return a.Priority() - b.Priority() } return statusWeight(a.Status) - statusWeight(b.Status) }) - return ranked + return vs } // statusWeight returns a sort weight for a VariantStatus. diff --git a/go/sdk/variants/server.go b/go/sdk/variants/server.go index aad7b56..9052e87 100644 --- a/go/sdk/variants/server.go +++ b/go/sdk/variants/server.go @@ -7,14 +7,16 @@ package variants import ( "context" "errors" + "net/http" + "sync" "github.com/modelcontextprotocol/go-sdk/mcp" ) -// variantEntry binds a ServerVariant to its backing mcp.Server. +// variantEntry binds a ServerVariant to its backing backend. type variantEntry struct { variant ServerVariant - server *mcp.Server + backend backend } // Server is a variant-aware MCP server that multiplexes across multiple @@ -22,17 +24,63 @@ type variantEntry struct { // initialization, clients provide VariantHints and the server responds with a // ranked list of available variants. Per-request variant selection is carried // in the _meta field. +// +// Server holds only configuration and is safe for concurrent use. In stateful +// mode (the default), per-session inner connections are created during +// initialize and scoped to the front session's lifetime. In stateless mode +// (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 - variantMap map[string]*mcp.Server // ID -> server for O(1) lookup rankingFunc RankingFunc + shared *sessionState // non-nil in stateless mode; cleaned up by Close } // NewServer creates a new variant-aware server with no registered variants. -func NewServer() *Server { - return &Server{ - variantMap: make(map[string]*mcp.Server), +// +// The first argument must not be nil. +func NewServer(impl *mcp.Implementation) *Server { + if impl == nil { + panic("variants: nil Implementation") } + return &Server{ + impl: impl, + } +} + +// NewStreamableHTTPHandler returns a new [mcp.StreamableHTTPHandler] for +// serving multiple concurrent clients over HTTP. It mirrors +// [mcp.NewStreamableHTTPHandler]. +// +// handler := variants.NewStreamableHTTPHandler(vs, nil) +// http.ListenAndServe(":8080", handler) +func NewStreamableHTTPHandler(vs *Server, opts *mcp.StreamableHTTPOptions) *mcp.StreamableHTTPHandler { + if vs == nil { + panic("variants: nil Server") + } + stateless := opts != nil && opts.Stateless + srv, err := vs.mcpServer(stateless) + if err != nil { + panic("variants: " + err.Error()) + } + return mcp.NewStreamableHTTPHandler( + func(r *http.Request) *mcp.Server { return srv }, + opts, + ) +} + +// addVariant is the shared registration logic for all With* methods. +// It checks for duplicates, sets priority, and appends the entry. +func (s *Server) addVariant(v ServerVariant, b backend, priority int) *Server { + for _, e := range s.variants { + if e.variant.ID == v.ID { + panic("variants: duplicate variant ID: " + v.ID) + } + } + v.priority = priority + s.variants = append(s.variants, variantEntry{variant: v, backend: b}) + return s } // WithVariant registers a ServerVariant backed by the given mcp.Server. @@ -42,21 +90,21 @@ func NewServer() *Server { // in the list and serve as the recommended default for clients. This // behavior can be overridden by providing a custom RankingFunc. // -// Per the SEP, the first variant in the ranked list MUST have status -// "stable" (unless the client explicitly requested experimental variants -// via variantHints). Priority ordering is respected within that -// constraint: variants are sorted by priority, but the first stable -// variant is always promoted to the front of the list. -// // Variant IDs must be unique; registering a duplicate panics. func (s *Server) WithVariant(v ServerVariant, mcpServer *mcp.Server, priority int) *Server { - if _, exists := s.variantMap[v.ID]; exists { - panic("variants: duplicate variant ID: " + v.ID) - } - v.priority = priority - s.variants = append(s.variants, variantEntry{variant: v, server: mcpServer}) - s.variantMap[v.ID] = mcpServer - return s + return s.addVariant(v, &inMemoryBackend{server: mcpServer}, priority) +} + +// WithHTTPVariant registers a ServerVariant backed by an mcp.Server exposed +// over HTTP. Not yet implemented. +func (s *Server) WithHTTPVariant(v ServerVariant, mcpServer *mcp.Server, priority int) *Server { + panic("variants: WithHTTPVariant not yet implemented") +} + +// WithRemoteVariant registers a ServerVariant backed by a remote MCP server +// at the given endpoint URL. Not yet implemented. +func (s *Server) WithRemoteVariant(v ServerVariant, endpoint string, priority int) *Server { + panic("variants: WithRemoteVariant not yet implemented") } // WithRanking sets a custom ranking function used to order variants based @@ -64,13 +112,6 @@ func (s *Server) WithVariant(v ServerVariant, mcpServer *mcp.Server, priority in // sorted by relevance, with the most appropriate variant first. If nil, // variants are ordered by their priority value (lowest first). // -// The ranking function does not need to enforce the SEP's stable-first -// invariant. After the function returns, the SDK will automatically promote -// the first stable variant to the front of the list (per the SEP: "The first -// element MUST have status stable unless the client explicitly requested -// experimental variants via variantHints"). The rest of the ranking order -// is preserved. -// // Returns the receiver for chaining. func (s *Server) WithRanking(fn RankingFunc) *Server { s.rankingFunc = fn @@ -102,122 +143,253 @@ func (s *Server) RankedVariants(ctx context.Context, hints VariantHints) []Serve return rankFn(ctx, hints, all) } -// ServerFor returns the mcp.Server associated with the given variant ID, or -// nil if no variant with that ID is registered. -func (s *Server) ServerFor(id string) *mcp.Server { - return s.variantMap[id] +// Close releases resources held by all registered backends and, in stateless +// mode, tears down the shared inner connections. +func (s *Server) Close() error { + if s.shared != nil { + s.shared.close() + s.shared = nil + } + var firstErr error + for _, entry := range s.variants { + if err := entry.backend.close(); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr } -// Run starts the variant server. +// Run starts the variant server on the given transport (e.g., stdio). +// For multi-client HTTP support, use [NewStreamableHTTPHandler] instead. func (s *Server) Run(ctx context.Context, t mcp.Transport) error { - // TODO: Implement Server.Run using a method-dispatching proxy pattern. - // - // Architecture: the variant Server acts as a "front" MCP server that owns - // the real transport, handles lifecycle/discovery methods itself, and - // dispatches tool/resource/prompt calls to the correct inner *mcp.Server - // based on _meta["io.modelcontextprotocol/server-variant"]. - // - // --------------------------------------------------------------------------- - // Constants - // --------------------------------------------------------------------------- - // - // Extension ID (plural, for capability negotiation): - // "io.modelcontextprotocol/server-variants" - // - // Per-request _meta key (singular, for variant selection): - // "io.modelcontextprotocol/server-variant" - // - // --------------------------------------------------------------------------- - // Assumptions (v1) - // --------------------------------------------------------------------------- - // - // We need to make sure the SEP's session stability invariant is satisfied: - // - // "The server MUST ensure that the same variant is recommended and selected for the entire session, unless the client explicitly requests otherwise via variantHints." - // - // For the first version we will relax this constraint to avoid maintaining state. - // - // --------------------------------------------------------------------------- - // Steps - // --------------------------------------------------------------------------- - // - // 1. Create a "front" *mcp.Server inside Run. Configure it with - // ServerOptions.Capabilities set to the union of all variants' - // capability flags. For example, if any variant has tools, set - // Capabilities.Tools = &ToolCapabilities{ListChanged: true}. This - // causes the SDK to advertise the capability during initialize - // without registering any actual tools, resources, or prompts. - // The inner *mcp.Server instances are never connected to a - // transport — they are only used as handler registries. - // - // 2. Add a receiving middleware via AddReceivingMiddleware on the - // front server. This single middleware handles all variant-aware - // dispatch: - // - // a. "initialize": read variantHints from the client's extension - // payload at capabilities.extensions[ - // "io.modelcontextprotocol/server-variants"]. Call - // RankedVariants() and call next() to let the SDK produce - // the base InitializeResult, then inject availableVariants - // (including id, description, hints, status, and - // deprecationInfo where applicable) into the response at - // capabilities.extensions[ - // "io.modelcontextprotocol/server-variants"]. Include - // moreVariantsAvailable if only a subset is returned. - // - // b. "tools/list", "resources/list", "prompts/list", - // "resources/templates/list", "completion/complete": - // read _meta["io.modelcontextprotocol/server-variant"] to - // determine the active variant (default to first-ranked if - // absent). Look up the corresponding inner *mcp.Server and - // forward the request to its handler. Return only that - // variant's items. Cursors returned in paginated responses - // MUST be variant-scoped — embed the variant ID so that - // cross-variant cursor reuse is detectable. - // - // c. "tools/call", "resources/read", "prompts/get": - // read _meta["io.modelcontextprotocol/server-variant"] from - // the request. Validate the variant ID: - // - If the ID is not in the session's availableVariants, - // return JSON-RPC error -32602 with message - // "Invalid server variant" and data containing - // { requestedVariant, availableVariants } (scoped to - // the authenticated principal's visible variants). - // - If absent, default to the first-ranked variant. - // Resolve the inner *mcp.Server via ServerFor(id) and - // forward the request. - // - // d. "resources/subscribe": read the variant from _meta, - // validate as in (c). Resource notifications for this subscription - // MUST carry _meta["io.modelcontextprotocol/server-variant"] - // matching the bound variant. If the resource disappears - // from the variant, send a resource_updated notification - // followed by list_changed. - // - // e. "notifications/tools/list_changed", - // "notifications/prompts/list_changed", - // "notifications/resources/list_changed": when an inner - // server emits a list_changed notification, the middleware - // re-emits it on the front server with - // _meta["io.modelcontextprotocol/server-variant"] set to - // the originating variant's ID. - // - // f. All other methods ("ping", etc.): - // pass through to the front server's default handler. - // - // 3. Call frontServer.Run(ctx, t) to own the transport and manage - // the session lifecycle. The front server handles the JSON-RPC - // connection, SSE streaming, progress notifications, and - // session teardown. - // - // --------------------------------------------------------------------------- - // Backward compatibility - // --------------------------------------------------------------------------- - // - // If the client does not send the extension payload during - // initialize, the server omits availableVariants from the response - // and uses the first-ranked variant for all requests. The server - // behaves identically to a non-variant-aware server. - // - return errors.New("variants: Run not implemented") + defer s.Close() + srv, err := s.mcpServer(false) + if err != nil { + return err + } + return srv.Run(ctx, t) +} + +// discoverCapabilities probes each backend to determine its advertised +// capabilities. The results are merged into a single set for the front +// proxy server. +func (s *Server) discoverCapabilities() (*mcp.ServerCapabilities, error) { + ctx := context.Background() + var allCaps []*mcp.ServerCapabilities + + for _, entry := range s.variants { + caps, err := entry.backend.capabilities(ctx) + if err != nil { + return nil, err + } + if caps != nil { + allCaps = append(allCaps, caps) + } + } + + return unionCapabilities(allCaps), nil +} + +// mcpServer returns a configured *mcp.Server that routes requests to the +// appropriate inner variant server based on the _meta variant field. +// +// The returned server is a thin proxy: it handles initialization (injecting +// variant metadata into the response) and delegates all other requests to the +// backing inner servers via dispatchers. +// +// Request flow (stateful mode): +// +// Client ── transport ──▸ frontServer ──▸ sessionMiddleware +// │ +// ┌────────────┼────────────┐ +// │ │ │ +// initialize route by pass through +// (create _meta (ping, etc.) +// per-session variant +// connections) │ +// ▼ +// dispatcher +// │ +// ┌───────────┴───────────┐ +// ▼ ▼ +// innerConnection innerConnection +// (variant "A") (variant "B") +// │ │ +// backend backend +// (in-memory, (in-memory, +// HTTP, or HTTP, or +// Remote) remote) +// +// Each variant's [backend] determines how the inner connection is +// established. See [inMemoryBackend] for the co-located case. +// +// In stateless mode (see [NewStreamableHTTPHandler]), the inner connections +// are created once and shared across all requests instead of per-session. +func (s *Server) mcpServer(stateless bool) (*mcp.Server, error) { + if len(s.variants) == 0 { + return nil, errors.New("variants: no variants registered") + } + + caps, err := s.discoverCapabilities() + if err != nil { + return nil, err + } + + // Per-session state, keyed by *mcp.ServerSession pointer identity. + sessions := &sync.Map{} + + // In stateless mode, create shared connections once and reuse them + // across all requests (no per-session state). The shared state is + // stored on the Server so Close() can release it. + var shared *sessionState + if stateless { + shared, err = s.createSessionState(context.Background(), nil) + if err != nil { + return nil, err + } + s.shared = shared + } + + frontServer := mcp.NewServer(s.impl, &mcp.ServerOptions{ + Capabilities: caps, + }) + + frontServer.AddReceivingMiddleware(s.sessionMiddleware(sessions, shared)) + + return frontServer, nil +} + +// --------------------------------------------------------------------------- +// Initialize helpers +// --------------------------------------------------------------------------- + +// extractVariantHints extracts client-provided variant hints from the +// initialize request's extension payload. Per SEP-2053, the client sends: +// +// experimental["io.modelcontextprotocol/server-variants"]["variantHints"] +func extractVariantHints(req mcp.Request) VariantHints { + params, _ := req.GetParams().(*mcp.InitializeParams) + if params == nil || params.Capabilities == nil || params.Capabilities.Experimental == nil { + return VariantHints{} + } + ext, ok := params.Capabilities.Experimental[extensionID] + if !ok { + return VariantHints{} + } + extMap, _ := ext.(map[string]any) + vh, _ := extMap["variantHints"].(map[string]any) + if vh == nil { + return VariantHints{} + } + var hints VariantHints + hints.Description, _ = vh["description"].(string) + if h, ok := vh["hints"].(map[string]any); ok { + hints.Hints = h + } + return hints +} + +// enrichInitResult injects variant information into the initialize response. +func (s *Server) enrichInitResult(ctx context.Context, result mcp.Result, req mcp.Request) (mcp.Result, error) { + initResult, ok := result.(*mcp.InitializeResult) + if !ok { + return result, nil + } + + ranked := s.RankedVariants(ctx, extractVariantHints(req)) + + // Build availableVariants payload + availableVariants := make([]map[string]any, len(ranked)) + for i, v := range ranked { + variant := map[string]any{ + "id": v.ID, + "description": v.Description, + } + if v.Hints != nil { + variant["hints"] = v.Hints + } + if v.Status != "" { + variant["status"] = v.Status + } + if v.DeprecationInfo != nil { + variant["deprecationInfo"] = v.DeprecationInfo + } + availableVariants[i] = variant + } + + if initResult.Capabilities == nil { + initResult.Capabilities = &mcp.ServerCapabilities{} + } + if initResult.Capabilities.Experimental == nil { + initResult.Capabilities.Experimental = make(map[string]any) + } + initResult.Capabilities.Experimental[extensionID] = map[string]any{ + "availableVariants": availableVariants, + "moreVariantsAvailable": len(ranked) < len(s.variants), + } + + return initResult, nil +} + +// unionCapabilities merges multiple ServerCapabilities into a single set +// that the front proxy server advertises to clients. The merge strategy is: +// +// - Tools, Resources, Prompts: the capability is advertised if any variant +// advertises it. Boolean sub-fields (ListChanged, Subscribe) are OR-ed +// across all variants so the front server enables the feature if at least +// one inner server supports it. +// - Completions, Logging: marker capabilities (empty structs). Advertised +// if any variant advertises them; the first non-nil value is kept. +// - Experimental: keys are merged into a single map. The first variant to +// register a given key wins; later duplicates are ignored. +func unionCapabilities(allCaps []*mcp.ServerCapabilities) *mcp.ServerCapabilities { + union := &mcp.ServerCapabilities{} + + for _, caps := range allCaps { + if caps == nil { + continue + } + + if caps.Tools != nil { + if union.Tools == nil { + union.Tools = &mcp.ToolCapabilities{} + } + union.Tools.ListChanged = union.Tools.ListChanged || caps.Tools.ListChanged + } + + if caps.Resources != nil { + if union.Resources == nil { + union.Resources = &mcp.ResourceCapabilities{} + } + union.Resources.Subscribe = union.Resources.Subscribe || caps.Resources.Subscribe + union.Resources.ListChanged = union.Resources.ListChanged || caps.Resources.ListChanged + } + + if caps.Prompts != nil { + if union.Prompts == nil { + union.Prompts = &mcp.PromptCapabilities{} + } + union.Prompts.ListChanged = union.Prompts.ListChanged || caps.Prompts.ListChanged + } + + if caps.Completions != nil && union.Completions == nil { + union.Completions = caps.Completions + } + if caps.Logging != nil && union.Logging == nil { + union.Logging = caps.Logging + } + + if caps.Experimental != nil { + if union.Experimental == nil { + union.Experimental = make(map[string]any) + } + for k, v := range caps.Experimental { + if _, exists := union.Experimental[k]; !exists { + union.Experimental[k] = v + } + } + } + } + + return union } diff --git a/go/sdk/variants/server_test.go b/go/sdk/variants/server_test.go index 4962133..3110bc6 100644 --- a/go/sdk/variants/server_test.go +++ b/go/sdk/variants/server_test.go @@ -4,11 +4,455 @@ package variants -import "testing" +import ( + "context" + "encoding/json" + "net/http/httptest" + "sync" + "testing" -func TestPlaceholder(t *testing.T) { - // Placeholder test to ensure go test runs successfully - if false { - t.Error("This should not happen") + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestServers creates two mcp.Servers with distinct tools: +// - "coding": analyze_code, refactor +// - "compact": summarize, lookup +func newTestServers() (coding *mcp.Server, compact *mcp.Server) { + coding = mcp.NewServer(&mcp.Implementation{Name: "coding-server", Version: "v1.0.0"}, nil) + mcp.AddTool(coding, &mcp.Tool{Name: "analyze_code", Description: "Static analysis"}, analyzeCode) + mcp.AddTool(coding, &mcp.Tool{Name: "refactor", Description: "Refactor code"}, refactor) + + compact = mcp.NewServer(&mcp.Implementation{Name: "compact-server", Version: "v1.0.0"}, nil) + mcp.AddTool(compact, &mcp.Tool{Name: "summarize", Description: "Summarize text"}, summarize) + mcp.AddTool(compact, &mcp.Tool{Name: "lookup", Description: "Quick lookup"}, lookup) + + return coding, compact +} + +// newTestVariantServer creates a variant server with two variants. +func newTestVariantServer() *Server { + codingServer, compactServer := newTestServers() + + return NewServer(&mcp.Implementation{Name: "test-server", Version: "1.0.0"}). + WithVariant(ServerVariant{ + ID: "coding", + Description: "Optimized for coding workflows", + Status: Stable, + }, codingServer, 0). + WithVariant(ServerVariant{ + ID: "compact", + Description: "Minimal token usage", + Status: Experimental, + }, compactServer, 1) +} + +// connectTestClient starts the variant server and connects a test client. +// Returns the client session; cleanup is handled via t.Cleanup. +func connectTestClient(t *testing.T, vs *Server, clientOpts *mcp.ClientOptions) *mcp.ClientSession { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + + serverTransport, clientTransport := mcp.NewInMemoryTransports() + + errCh := make(chan error, 1) + go func() { + errCh <- vs.Run(ctx, serverTransport) + }() + + client := mcp.NewClient( + &mcp.Implementation{Name: "test-client", Version: "v0.0.1"}, + clientOpts, + ) + + session, err := client.Connect(ctx, clientTransport, nil) + if err != nil { + cancel() + t.Fatalf("client.Connect: %v", err) } + + t.Cleanup(func() { + session.Close() + cancel() + <-errCh + }) + + return session +} + +// TestIntegration_EndToEnd verifies the full variant server lifecycle with a +// variant-aware client: the client advertises variant support and sends hints +// during initialization, then selects variants per-request via _meta. +func TestIntegration_EndToEnd(t *testing.T) { + vs := newTestVariantServer() + session := connectTestClient(t, vs, nil) + ctx := context.Background() + + // --- 1. Verify initialize returned variant information --- + initResult := session.InitializeResult() + require.NotNil(t, initResult) + require.NotNil(t, initResult.Capabilities) + + ext, ok := initResult.Capabilities.Experimental[extensionID] + require.True(t, ok, "expected experimental capability with extension ID") + + extJSON, err := json.Marshal(ext) + require.NoError(t, err) + + var extData struct { + AvailableVariants []struct { + ID string `json:"id"` + Description string `json:"description"` + Status string `json:"status"` + } `json:"availableVariants"` + MoreVariantsAvailable bool `json:"moreVariantsAvailable"` + } + require.NoError(t, json.Unmarshal(extJSON, &extData)) + require.Len(t, extData.AvailableVariants, 2) + + // Default ranking: coding (priority 0, stable) first + assert.Equal(t, "coding", extData.AvailableVariants[0].ID) + assert.Equal(t, "compact", extData.AvailableVariants[1].ID) + assert.False(t, extData.MoreVariantsAvailable) + + // --- 2. List tools for the "coding" variant --- + codingTools, err := session.ListTools(ctx, &mcp.ListToolsParams{ + Meta: mcp.Meta{metaKeyVariant: "coding"}, + }) + require.NoError(t, err) + codingNames := toolNames(codingTools.Tools) + assert.Contains(t, codingNames, "analyze_code") + assert.Contains(t, codingNames, "refactor") + assert.NotContains(t, codingNames, "summarize") + + // --- 3. List tools for the "compact" variant --- + compactTools, err := session.ListTools(ctx, &mcp.ListToolsParams{ + Meta: mcp.Meta{metaKeyVariant: "compact"}, + }) + require.NoError(t, err) + compactNames := toolNames(compactTools.Tools) + assert.Contains(t, compactNames, "summarize") + assert.Contains(t, compactNames, "lookup") + assert.NotContains(t, compactNames, "analyze_code") + + // --- 4. Call a tool on the "coding" variant --- + callResult, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "analyze_code", + Meta: mcp.Meta{metaKeyVariant: "coding"}, + Arguments: map[string]json.RawMessage{ + "code": json.RawMessage(`"fmt.Println(x)"`), + "language": json.RawMessage(`"go"`), + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, callResult.Content) + + // --- 5. Call a tool on the "compact" variant --- + callResult2, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "summarize", + Meta: mcp.Meta{metaKeyVariant: "compact"}, + Arguments: map[string]json.RawMessage{ + "text": json.RawMessage(`"This is a long text that should be summarized into something shorter."`), + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, callResult2.Content) + + // --- 6. Cross-variant tool call: coding tool fails on compact variant --- + _, err = session.CallTool(ctx, &mcp.CallToolParams{ + Name: "analyze_code", + Meta: mcp.Meta{metaKeyVariant: "compact"}, + Arguments: map[string]json.RawMessage{ + "code": json.RawMessage(`"x := 1"`), + "language": json.RawMessage(`"go"`), + }, + }) + assert.Error(t, err, "coding tool should not be reachable via compact variant") + + // --- 7. Invalid variant should return an error --- + _, err = session.ListTools(ctx, &mcp.ListToolsParams{ + Meta: mcp.Meta{metaKeyVariant: "nonexistent"}, + }) + assert.Error(t, err) +} + +// TestIntegration_ClientWithoutVariantSupport verifies that a client that +// doesn't know about variants can still use the server normally. All requests +// go to the first-ranked (default) variant without any _meta selection. +func TestIntegration_ClientWithoutVariantSupport(t *testing.T) { + vs := newTestVariantServer() + session := connectTestClient(t, vs, nil) + ctx := context.Background() + + // --- 1. Initialize still succeeds and advertises variants (server always does) --- + initResult := session.InitializeResult() + require.NotNil(t, initResult) + require.NotNil(t, initResult.Capabilities) + assert.Contains(t, initResult.Capabilities.Experimental, extensionID, + "server should advertise variants even for unaware clients") + + // --- 2. List tools without _meta — should get default variant ("coding") tools --- + tools, err := session.ListTools(ctx, nil) + require.NoError(t, err) + names := toolNames(tools.Tools) + assert.Contains(t, names, "analyze_code") + assert.Contains(t, names, "refactor") + assert.NotContains(t, names, "summarize") + assert.NotContains(t, names, "lookup") + + // --- 3. Call a tool without _meta — should route to default variant --- + result, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "analyze_code", + Arguments: map[string]json.RawMessage{ + "code": json.RawMessage(`"x := 1"`), + "language": json.RawMessage(`"go"`), + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, result.Content) + + // --- 4. Calling a tool that only exists on a non-default variant should fail --- + _, err = session.CallTool(ctx, &mcp.CallToolParams{ + Name: "summarize", + Arguments: map[string]json.RawMessage{ + "text": json.RawMessage(`"hello"`), + }, + }) + assert.Error(t, err, "tool from non-default variant should not be reachable without _meta") +} + +// connectHTTPTestClient connects a client to an existing httptest server via +// StreamableClientTransport. Returns the client session; cleanup is handled +// via t.Cleanup. +func connectHTTPTestClient(t *testing.T, httpSrv *httptest.Server) *mcp.ClientSession { + t.Helper() + ctx := context.Background() + + client := mcp.NewClient( + &mcp.Implementation{Name: "test-http-client", Version: "v0.0.1"}, + nil, + ) + + session, err := client.Connect(ctx, &mcp.StreamableClientTransport{ + Endpoint: httpSrv.URL, + }, nil) + require.NoError(t, err) + + t.Cleanup(func() { + session.Close() + }) + + return session +} + +// TestIntegration_HTTP verifies the variant server works over HTTP with +// multiple concurrent clients. Each client gets its own session with +// independent variant routing. +func TestIntegration_HTTP(t *testing.T) { + vs := newTestVariantServer() + + handler := NewStreamableHTTPHandler(vs, nil) + + httpSrv := httptest.NewServer(handler) + t.Cleanup(httpSrv.Close) + + // --- Connect two clients concurrently --- + session1 := connectHTTPTestClient(t, httpSrv) + session2 := connectHTTPTestClient(t, httpSrv) + + ctx := context.Background() + + // --- 1. Both clients see variant metadata --- + for _, session := range []*mcp.ClientSession{session1, session2} { + initResult := session.InitializeResult() + require.NotNil(t, initResult) + require.NotNil(t, initResult.Capabilities) + assert.Contains(t, initResult.Capabilities.Experimental, extensionID) + } + + // --- 2. Client 1 uses "coding" variant, Client 2 uses "compact" --- + var wg sync.WaitGroup + wg.Add(2) + + // Client 1: coding variant + go func() { + defer wg.Done() + + tools, err := session1.ListTools(ctx, &mcp.ListToolsParams{ + Meta: mcp.Meta{metaKeyVariant: "coding"}, + }) + if !assert.NoError(t, err) { + return + } + names := toolNames(tools.Tools) + assert.Contains(t, names, "analyze_code") + assert.Contains(t, names, "refactor") + assert.NotContains(t, names, "summarize") + + result, err := session1.CallTool(ctx, &mcp.CallToolParams{ + Name: "analyze_code", + Meta: mcp.Meta{metaKeyVariant: "coding"}, + Arguments: map[string]json.RawMessage{ + "code": json.RawMessage(`"x := 1"`), + "language": json.RawMessage(`"go"`), + }, + }) + if !assert.NoError(t, err) { + return + } + assert.NotEmpty(t, result.Content) + }() + + // Client 2: compact variant + go func() { + defer wg.Done() + + tools, err := session2.ListTools(ctx, &mcp.ListToolsParams{ + Meta: mcp.Meta{metaKeyVariant: "compact"}, + }) + if !assert.NoError(t, err) { + return + } + names := toolNames(tools.Tools) + assert.Contains(t, names, "summarize") + assert.Contains(t, names, "lookup") + assert.NotContains(t, names, "analyze_code") + + result, err := session2.CallTool(ctx, &mcp.CallToolParams{ + Name: "summarize", + Meta: mcp.Meta{metaKeyVariant: "compact"}, + Arguments: map[string]json.RawMessage{ + "text": json.RawMessage(`"A long text to summarize"`), + }, + }) + if !assert.NoError(t, err) { + return + } + assert.NotEmpty(t, result.Content) + }() + + wg.Wait() + + // --- 3. Cross-variant isolation: coding tool fails on compact --- + _, err := session1.CallTool(ctx, &mcp.CallToolParams{ + Name: "analyze_code", + Meta: mcp.Meta{metaKeyVariant: "compact"}, + Arguments: map[string]json.RawMessage{ + "code": json.RawMessage(`"x := 1"`), + "language": json.RawMessage(`"go"`), + }, + }) + assert.Error(t, err, "coding tool should not be reachable via compact variant") +} + +// TestIntegration_HTTP_Stateless verifies the variant server works in +// stateless HTTP mode where shared connections are reused across requests. +func TestIntegration_HTTP_Stateless(t *testing.T) { + vs := newTestVariantServer() + + handler := NewStreamableHTTPHandler(vs, + &mcp.StreamableHTTPOptions{Stateless: true}, + ) + + httpSrv := httptest.NewServer(handler) + t.Cleanup(httpSrv.Close) + + session := connectHTTPTestClient(t, httpSrv) + ctx := context.Background() + + // Variant metadata is still returned + initResult := session.InitializeResult() + require.NotNil(t, initResult) + assert.Contains(t, initResult.Capabilities.Experimental, extensionID) + + // List tools for coding variant + codingTools, err := session.ListTools(ctx, &mcp.ListToolsParams{ + Meta: mcp.Meta{metaKeyVariant: "coding"}, + }) + require.NoError(t, err) + names := toolNames(codingTools.Tools) + assert.Contains(t, names, "analyze_code") + assert.NotContains(t, names, "summarize") + + // Call tool on compact variant + result, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "summarize", + Meta: mcp.Meta{metaKeyVariant: "compact"}, + Arguments: map[string]json.RawMessage{ + "text": json.RawMessage(`"stateless test"`), + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, result.Content) +} + +// --------------------------------------------------------------------------- +// Test tool handlers +// --------------------------------------------------------------------------- + +type analyzeCodeInput struct { + Code string `json:"code"` + Language string `json:"language"` +} + +type analyzeCodeOutput struct { + Issues []string `json:"issues"` +} + +func analyzeCode(_ context.Context, _ *mcp.CallToolRequest, in analyzeCodeInput) (*mcp.CallToolResult, analyzeCodeOutput, error) { + return nil, analyzeCodeOutput{Issues: []string{"unused variable"}}, nil +} + +type refactorInput struct { + Code string `json:"code"` + Action string `json:"action"` +} + +type refactorOutput struct { + Refactored string `json:"refactored"` +} + +func refactor(_ context.Context, _ *mcp.CallToolRequest, in refactorInput) (*mcp.CallToolResult, refactorOutput, error) { + return nil, refactorOutput{Refactored: "// refactored\n" + in.Code}, nil +} + +type summarizeInput struct { + Text string `json:"text"` +} + +type summarizeOutput struct { + Summary string `json:"summary"` +} + +func summarize(_ context.Context, _ *mcp.CallToolRequest, in summarizeInput) (*mcp.CallToolResult, summarizeOutput, error) { + n := len(in.Text) + if n > 50 { + n = 50 + } + return nil, summarizeOutput{Summary: in.Text[:n]}, nil +} + +type lookupInput struct { + Query string `json:"query"` +} + +type lookupOutput struct { + Result string `json:"result"` +} + +func lookup(_ context.Context, _ *mcp.CallToolRequest, in lookupInput) (*mcp.CallToolResult, lookupOutput, error) { + return nil, lookupOutput{Result: "result for: " + in.Query}, nil +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +func toolNames(tools []*mcp.Tool) []string { + names := make([]string, len(tools)) + for i, t := range tools { + names[i] = t.Name + } + return names } diff --git a/go/sdk/variants/session.go b/go/sdk/variants/session.go new file mode 100644 index 0000000..c3e4d22 --- /dev/null +++ b/go/sdk/variants/session.go @@ -0,0 +1,138 @@ +// 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" + "sync" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// --------------------------------------------------------------------------- +// Per-session types +// --------------------------------------------------------------------------- + +// innerConnection holds the resources for communicating with one inner +// server. In stateful mode one is created per client session; in +// stateless mode a single instance is shared across all requests. +type innerConnection struct { + clientSession *mcp.ClientSession + cleanupFn func() +} + +// close tears down the client session and invokes the backend-specific +// cleanup function. +func (c *innerConnection) close() { + c.clientSession.Close() + if c.cleanupFn != nil { + c.cleanupFn() + } +} + +// sessionState holds all per-session state for one front client. +type sessionState struct { + dispatcher *dispatcher +} + +// close tears down all inner connections for this session. +func (ss *sessionState) close() { + for _, c := range ss.dispatcher.connections { + c.close() + } +} + +// injectVariantMeta sets the variant ID in a Params' _meta map, +// preserving any existing metadata. +func injectVariantMeta(p mcp.Params, variantID string) { + meta := p.GetMeta() + if meta == nil { + meta = map[string]any{} + p.SetMeta(meta) + } + meta[metaKeyVariant] = variantID +} + +// --------------------------------------------------------------------------- +// Session lifecycle +// --------------------------------------------------------------------------- + +// createSessionState sets up inner connections for all variants and returns +// the per-session state. +func (s *Server) createSessionState(ctx context.Context, frontSession *mcp.ServerSession) (*sessionState, error) { + connections := make(map[string]*innerConnection, len(s.variants)) + + for _, entry := range s.variants { + conn, err := entry.backend.connect(ctx, entry.variant, frontSession) + if err != nil { + for _, c := range connections { + c.close() + } + return nil, err + } + connections[entry.variant.ID] = conn + } + + return &sessionState{ + dispatcher: &dispatcher{ + server: s, + connections: connections, + }, + }, nil +} + +// sessionMiddleware builds the receiving middleware that manages per-session +// state and delegates to the variant dispatcher. +// +// When shared is non-nil (stateless mode), all requests use the shared +// connections instead of creating per-session state. +func (s *Server) sessionMiddleware(sessions *sync.Map, shared *sessionState) mcp.Middleware { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + ss := req.GetSession().(*mcp.ServerSession) + + if method == "initialize" { + // Let the SDK handle init first (capability negotiation etc.) + result, err := next(ctx, method, req) + if err != nil { + return nil, err + } + + // In stateless mode, skip per-session connection creation; + // requests will use the shared connections. + if shared == nil { + state, err := s.createSessionState(ctx, ss) + if err != nil { + return nil, err + } + sessions.Store(ss, state) + + // Clean up when the front session closes. + go func() { + ss.Wait() + sessions.Delete(ss) + state.close() + }() + } + + // Enrich the init result with variant information + return s.enrichInitResult(ctx, result, req) + } + + // Try per-session state first + if v, ok := sessions.Load(ss); ok { + state := v.(*sessionState) + return state.dispatcher.handle(ctx, method, req, next) + } + + // Fall back to shared state (stateless mode) + if shared != nil { + return shared.dispatcher.handle(ctx, method, req, next) + } + + return next(ctx, method, req) + } + } +} diff --git a/go/sdk/variants/types.go b/go/sdk/variants/types.go index 120dbe3..12d5667 100644 --- a/go/sdk/variants/types.go +++ b/go/sdk/variants/types.go @@ -6,6 +6,14 @@ package variants import "context" +const ( + // Extension ID for variant capability negotiation (plural) + extensionID = "io.modelcontextprotocol/server-variants" + + // Per-request _meta key for variant selection (singular) + metaKeyVariant = "io.modelcontextprotocol/server-variant" +) + // --------------------------------------------------------------------------- // Variant status // ---------------------------------------------------------------------------