Implement variant-aware MCP server (#3)
Summary
Builds on the types and architecture from the previous PR to deliver a
working variant-aware MCP
server that routes requests to inner mcp.Server instances based on
per-request _meta variant selection,
enabling a single MCP endpoint to serve multiple tool/resource/prompt
variants simultaneously.
```go
// Create standard mcp.Server instances — one per variant
codeReviewServer := mcp.NewServer(&mcp.Implementation{Name:
"devplatform", Version: "v1.0.0"}, nil)
mcp.AddTool(codeReviewServer, &mcp.Tool{Name: "list_pull_requests",
...}, handler)
mcp.AddTool(codeReviewServer, &mcp.Tool{Name: "get_diff", ...}, handler)
pmServer := mcp.NewServer(&mcp.Implementation{Name: "devplatform",
Version: "v1.0.0"}, nil)
mcp.AddTool(pmServer, &mcp.Tool{Name: "list_issues", ...}, handler)
mcp.AddTool(pmServer, &mcp.Tool{Name: "create_issue", ...}, handler)
// Register them as variants on a single server
vs := variants.NewServer(&mcp.Implementation{Name: "devplatform",
Version: "v1.0.0"}).
WithVariant(variants.ServerVariant{
ID: "code-review",
Description: "PR and code review operations",
Status: variants.Stable,
}, codeReviewServer, 0).
WithVariant(variants.ServerVariant{
ID: "project-management",
Description: "Issue and project tracking operations",
Status: variants.Stable,
}, pmServer, 1)
// Future: register variants backed by remote MCP servers
// WithHTTPVariant(v, mcpServer, priority) — variant backed by an
mcp.Server exposed over HTTP
// WithRemoteVariant(v, "https://...", priority) — variant backed by a
remote MCP endpoint
// Serve over HTTP for multiple concurrent clients
handler := variants.NewStreamableHTTPHandler(vs, nil)
http.ListenAndServe(":8080", handler)
// Or serve over stdio for single-client use:
// vs.Run(ctx, &mcp.StdioTransport{})
```
Key additions
- Backend abstraction (inMemoryBackend) — creates in-memory transport pairs between the variant server and
inner mcp.Server instances, with notification forwarding for progress and logging events.
- Request dispatcher (dispatch.go) — handles routing for tools/list, tools/call, resources/list,
resources/read, resources/subscribe, prompts/list, prompts/get, and completion/complete methods.
Implements variant-scoped cursor wrapping to correctly paginate across multiple inner servers.
- Session middleware (session.go) — manages per-session inner connections in stateful mode and shared
connections in stateless HTTP mode, ensuring proper lifecycle management.
- Initialize enrichment — injects availableVariants into the server's capabilities response so clients can
discover which variants are available.
- Capability discovery — merges inner server capabilities via union to present a unified capability set to
clients.
- NewStreamableHTTPHandler — provides an HTTP handler for multi-client serving over Streamable HTTP
transport.
- Integration tests (server_test.go, 458 lines) — covers end-to-end variant selection, backward
compatibility (no variant specified), HTTP transport, and stateless modes.
- Examples — stdio-based (examples/server/variants/main.go) and HTTP-based
(examples/server/variants-http/main.go) example servers demonstrating the API.
- README — rewritten with API documentation and architecture diagrams.
TODO (subsequent PRs):
- Unit Testing
- More integration tests
Signed-off-by: Kurt Degiorgio <kdegiorgio@bloomberg.net>
This commit is contained in:
Kurt Degiorgio
GitHub
parent
dbf99bebf0
commit
c01fd188d2
15 changed files with 2217 additions and 245 deletions
258
go/sdk/README.md
258
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).
|
||||
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.
|
||||
|
|
|
|||
254
go/sdk/examples/server/exampletools/tools.go
Normal file
254
go/sdk/examples/server/exampletools/tools.go
Normal file
|
|
@ -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
|
||||
}
|
||||
111
go/sdk/examples/server/variants-http/main.go
Normal file
111
go/sdk/examples/server/variants-http/main.go
Normal file
|
|
@ -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))
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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=
|
||||
|
|
|
|||
126
go/sdk/variants/backend.go
Normal file
126
go/sdk/variants/backend.go
Normal file
|
|
@ -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
|
||||
}
|
||||
448
go/sdk/variants/dispatch.go
Normal file
448
go/sdk/variants/dispatch.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
138
go/sdk/variants/session.go
Normal file
138
go/sdk/variants/session.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Reference in a new issue