diff --git a/CLAUDE.md b/CLAUDE.md index 69217cc..ba24efe 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,6 +72,9 @@ task license-fix # Add missing license headers | `env` | Environment variable abstraction with `Reader` interface for testable code | | `httperr` | Wrap errors with HTTP status codes; use `WithCode()`, `Code()`, `New()` | | `logging` | Pre-configured `*slog.Logger` factory with consistent ToolHive defaults (Alpha) | +| `mcpcompat/mcp` | Drop-in shim for `mark3labs/mcp-go/mcp` data types; re-exports via aliases, backed by go-sdk elsewhere (Alpha) | +| `mcpcompat/client` | Drop-in shim for `mark3labs/mcp-go/client` (+ `client/transport`) reimplemented on the official go-sdk (Alpha) | +| `mcpcompat/server` | Drop-in shim for `mark3labs/mcp-go/server` reimplemented on the official go-sdk (Alpha) | | `oci/artifact` | Artifact-agnostic OCI tar/gzip/extraction/platform primitives shared by oci/skills and oci/plugins (Alpha) | | `oci/skills` | OCI artifact types, media types, and registry operations for ToolHive skills (Alpha) | | `oci/plugins` | OCI artifact types, media types, and registry operations for ToolHive plugins (Alpha) | diff --git a/go.mod b/go.mod index b4d59a7..45a27c7 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.10.0 github.com/mark3labs/mcp-go v0.55.1 + github.com/modelcontextprotocol/go-sdk v1.6.1 github.com/modelcontextprotocol/registry v1.7.9 github.com/opencontainers/go-digest v1.0.0 github.com/opencontainers/image-spec v1.1.1 @@ -78,7 +79,7 @@ require ( github.com/go-openapi/validate v0.25.3 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect github.com/google/certificate-transparency-go v1.3.3 // indirect - github.com/google/jsonschema-go v0.4.2 // indirect + github.com/google/jsonschema-go v0.4.3 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 // indirect github.com/in-toto/attestation v1.2.0 // indirect github.com/in-toto/in-toto-golang v0.11.0 // indirect @@ -90,6 +91,8 @@ require ( github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/secure-systems-lab/go-securesystemslib v0.11.0 // indirect + github.com/segmentio/asm v1.1.3 // indirect + github.com/segmentio/encoding v0.5.4 // indirect github.com/shibumi/go-pathspec v1.3.0 // indirect github.com/sigstore/rekor v1.5.2 // indirect github.com/sigstore/rekor-tiles/v2 v2.2.2-0.20260601073857-5d098a2b6443 // indirect @@ -114,6 +117,7 @@ require ( golang.org/x/crypto v0.53.0 // indirect golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect golang.org/x/mod v0.37.0 // indirect + golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sync v0.21.0 // indirect golang.org/x/sys v0.46.0 // indirect golang.org/x/term v0.44.0 // indirect diff --git a/go.sum b/go.sum index 4e6b86f..ed83f98 100644 --- a/go.sum +++ b/go.sum @@ -187,8 +187,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/go-containerregistry v0.21.7 h1:/vPFuVXDjtFREsVArW+0h1CIl5urnOhzei4X2DMW9IU= github.com/google/go-containerregistry v0.21.7/go.mod h1:kjSbt7/zMsKLWfnHrIvKvhXHUw91jbe9DNjPPJ32gXE= -github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8= -github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= +github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/trillian v1.7.3 h1:hziW+vo4czis48tzx2GK5xRBl/ZxBA9B0/UR5avXOro= @@ -261,6 +261,8 @@ github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= +github.com/modelcontextprotocol/go-sdk v1.6.1 h1:0zOSupjKUxPKSocPT1Wtago+mUHU2/uZ4xSOY0FGReU= +github.com/modelcontextprotocol/go-sdk v1.6.1/go.mod h1:kzm3kzFL1/+AziGOE0nUs3gvPoNxMCvkxokMkuFapXQ= github.com/modelcontextprotocol/registry v1.7.9 h1:vpPfx2A2egjhm6YlbwfkX8NkR2N0S2eYmvYXI8bXaBs= github.com/modelcontextprotocol/registry v1.7.9/go.mod h1:y03zY98e+REsiCaj1sUKzXbk3qEp++Y3gzAnV83wrNs= github.com/natefinch/atomic v1.0.1 h1:ZPYKxkqQOx3KZ+RsbnP/YsgvxWQPGxjC0oBt2AhwV0A= @@ -291,6 +293,10 @@ github.com/sassoftware/relic/v7 v7.6.2 h1:rS44Lbv9G9eXsukknS4mSjIAuuX+lMq/FnStgm github.com/sassoftware/relic/v7 v7.6.2/go.mod h1:kjmP0IBVkJZ6gXeAu35/KCEfca//+PKM6vTAsyDPY+k= github.com/secure-systems-lab/go-securesystemslib v0.11.0 h1:iuCR9kcMFD4QurdKrGvPLoKZLv9YvwPYVr0473BdtFs= github.com/secure-systems-lab/go-securesystemslib v0.11.0/go.mod h1:+PMOTjUGwHj2vcZ+TFKlb1tXRbrdWE1LYDT5i9JC80Q= +github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= +github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw= github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4= github.com/shibumi/go-pathspec v1.3.0 h1:QUyMZhFo0Md5B8zV8x2tesohbb5kfbpTi9rBnKh5dkI= @@ -414,6 +420,8 @@ golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +golang.org/x/tools v0.46.0 h1:7jTurBkPZu4moS/Uy4OQT1M+QBlsj3wejyZwsT8Z7rk= +golang.org/x/tools v0.46.0/go.mod h1:FrD85F8l+NWL+9XWBSyVSHO6Ne4jutsfIFba7AWQ5Ys= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= google.golang.org/api v0.280.0 h1:F4OfEHZhZh6a7uTufJAXXVd/2TQ8EjM4vZH+jX/vFYk= diff --git a/mcpcompat/client/client.go b/mcpcompat/client/client.go new file mode 100644 index 0000000..6e5d5c4 --- /dev/null +++ b/mcpcompat/client/client.go @@ -0,0 +1,575 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package client is a drop-in compatibility shim for +// github.com/mark3labs/mcp-go/client, reimplemented on top of the official +// github.com/modelcontextprotocol/go-sdk. +// +// It presents mcp-go's flat client API (NewStreamableHttpClient, Start, +// Initialize, ListTools, CallTool, ...) while delegating the actual protocol to +// a go-sdk Client and ClientSession underneath. Data types are exchanged as the +// mcp-go-shaped types from mcpcompat/mcp; conversion to and from the go-sdk's +// own types happens at this boundary via JSON round-trips, which is robust +// because both encode the identical MCP wire format. +// +// Stability: Alpha. +package client + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "sync" + "time" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// Client is an MCP client backed by the official go-sdk. It mirrors the subset +// of mcp-go's client.Client that ToolHive uses. +type Client struct { + // transport configuration; exactly one of streamable/sse is non-nil. + streamable *transport.StreamableHTTP + sse *transport.SSE + + mu sync.Mutex + client *gosdk.Client + session *gosdk.ClientSession + notifyMu sync.Mutex + notify []func(mcp.JSONRPCNotification) + + connLostMu sync.Mutex + connLost func(error) + watchStarted bool + + // resume holds the raw JSON-RPC-over-HTTP machinery used when the client + // resumes a pre-existing session (transport.WithSession) without calling + // Initialize. It is lazily created on the first resumed request. See resume.go. + resume *resumeState +} + +// NewStreamableHttpClient creates a Streamable HTTP MCP client for baseURL. Like +// mcp-go, the returned client is not yet connected; call Start then Initialize. +func NewStreamableHttpClient(baseURL string, options ...transport.StreamableHTTPCOption) (*Client, error) { + return &Client{streamable: transport.NewStreamableHTTP(baseURL, options...)}, nil +} + +// NewSSEMCPClient creates an SSE MCP client for baseURL. The returned client is +// not yet connected; call Start then Initialize. +func NewSSEMCPClient(baseURL string, options ...transport.ClientOption) (*Client, error) { + return &Client{sse: transport.NewSSE(baseURL, options...)}, nil +} + +// Start prepares the client. The go-sdk performs connection and initialization +// together in a single Connect call, which this shim issues from Initialize; +// Start is therefore a no-op retained for API compatibility. +func (*Client) Start(_ context.Context) error { return nil } + +// Initialize connects the underlying go-sdk client and performs the MCP +// initialize handshake using the supplied client info and capabilities. +func (c *Client) Initialize(ctx context.Context, request mcp.InitializeRequest) (*mcp.InitializeResult, error) { + ctx = withErrCapture(ctx) + c.mu.Lock() + defer c.mu.Unlock() + + if c.session != nil { + return nil, fmt.Errorf("client already initialized") + } + + impl := &gosdk.Implementation{} + if err := jsonConvert(request.Params.ClientInfo, impl); err != nil { + return nil, fmt.Errorf("converting client info: %w", err) + } + if impl.Name == "" { + impl.Name = "toolhive" + } + if impl.Version == "" { + impl.Version = "0.0.0" + } + + opts := &gosdk.ClientOptions{} + if !isZeroCapabilities(request.Params.Capabilities) { + caps := &gosdk.ClientCapabilities{} + if err := jsonConvert(request.Params.Capabilities, caps); err != nil { + return nil, fmt.Errorf("converting client capabilities: %w", err) + } + opts.Capabilities = caps + } + c.installNotificationHandlers(opts) + + gc := gosdk.NewClient(impl, opts) + + tr, err := c.buildTransport() + if err != nil { + return nil, err + } + + session, err := gc.Connect(ctx, tr, nil) + if err != nil { + return nil, mapConnectError(ctx, err) + } + + c.client = gc + c.session = session + if c.streamable != nil { + c.streamable.SetSessionID(session.ID()) + } + c.maybeStartConnLostWatch(session) + + result := &mcp.InitializeResult{} + if err := jsonConvert(session.InitializeResult(), result); err != nil { + return nil, fmt.Errorf("converting initialize result: %w", err) + } + return result, nil +} + +// buildTransport constructs the go-sdk transport from the configured options. +func (c *Client) buildTransport() (gosdk.Transport, error) { + switch { + case c.streamable != nil: + return &gosdk.StreamableClientTransport{ + Endpoint: c.streamable.Endpoint(), + HTTPClient: buildHTTPClient( + c.streamable.HTTPClient(), c.streamable.Headers(), c.streamable.HeaderFunc(), c.streamable.Timeout(), + ), + DisableStandaloneSSE: !c.streamable.ContinuousListening(), + }, nil + case c.sse != nil: + return &gosdk.SSEClientTransport{ + Endpoint: c.sse.Endpoint(), + HTTPClient: buildHTTPClient(c.sse.HTTPClient(), c.sse.Headers(), nil, 0), + }, nil + default: + return nil, fmt.Errorf("no transport configured") + } +} + +// Close terminates the session. +func (c *Client) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + if c.session == nil { + return nil + } + err := c.session.Close() + c.session = nil + return err +} + +// Ping verifies the server is responsive. +func (c *Client) Ping(ctx context.Context) error { + if c.isResume() { + return c.resumeCall(ctx, "ping", struct{}{}, nil) + } + s, err := c.sessionFor() + if err != nil { + return err + } + return s.Ping(ctx, nil) +} + +// ListTools lists the server's tools. +func (c *Client) ListTools(ctx context.Context, request mcp.ListToolsRequest) (*mcp.ListToolsResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + out := &mcp.ListToolsResult{} + return out, c.resumeCall(ctx, "tools/list", request.Params, out) + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + res, err := s.ListTools(ctx, &gosdk.ListToolsParams{Cursor: string(request.Params.Cursor)}) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertResult[mcp.ListToolsResult](res) +} + +// CallTool invokes a tool. +func (c *Client) CallTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + out := &mcp.CallToolResult{} + return out, c.resumeCall(ctx, "tools/call", request.Params, out) + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + params := &gosdk.CallToolParams{ + Name: request.Params.Name, + Arguments: request.Params.Arguments, + } + // Preserve the request _meta (ToolHive propagates metadata through vMCP to + // backends via this field). mcp-go's *Meta flattens its AdditionalFields to + // top-level _meta keys, so convert via JSON into go-sdk's Meta (map[string]any). + if request.Params.Meta != nil { + m := gosdk.Meta{} + if err := jsonConvert(request.Params.Meta, &m); err != nil { + return nil, fmt.Errorf("converting call meta: %w", err) + } + params.Meta = m + } + res, err := s.CallTool(ctx, params) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertResult[mcp.CallToolResult](res) +} + +// ReadResource reads a resource by URI. +func (c *Client) ReadResource(ctx context.Context, request mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + raw := &gosdk.ReadResourceResult{} + if err := c.resumeCall(ctx, "resources/read", request.Params, raw); err != nil { + return nil, err + } + return convertReadResourceResult(raw), nil + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + res, err := s.ReadResource(ctx, &gosdk.ReadResourceParams{URI: request.Params.URI}) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertReadResourceResult(res), nil +} + +// convertReadResourceResult converts a go-sdk ReadResourceResult into its +// mcp-go-shaped equivalent. mcp-go's ReadResourceResult.Contents is a slice of +// the mcp.ResourceContents interface with no custom unmarshaler, so a JSON +// round-trip cannot populate it; each go-sdk ResourceContents struct is mapped +// explicitly to the concrete text/blob mcp-go type instead. +func convertReadResourceResult(res *gosdk.ReadResourceResult) *mcp.ReadResourceResult { + out := &mcp.ReadResourceResult{} + if len(res.Meta) > 0 { + out.Meta = mcp.NewMetaFromMap(map[string]any(res.Meta)) + } + for _, rc := range res.Contents { + if rc == nil { + continue + } + var meta map[string]any + if len(rc.Meta) > 0 { + meta = map[string]any(rc.Meta) + } + if len(rc.Blob) > 0 { + out.Contents = append(out.Contents, mcp.BlobResourceContents{ + URI: rc.URI, + MIMEType: rc.MIMEType, + Blob: base64.StdEncoding.EncodeToString(rc.Blob), + Meta: meta, + }) + continue + } + out.Contents = append(out.Contents, mcp.TextResourceContents{ + URI: rc.URI, + MIMEType: rc.MIMEType, + Text: rc.Text, + Meta: meta, + }) + } + return out +} + +// GetPrompt gets a prompt, rendered with the provided arguments. +func (c *Client) GetPrompt(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + raw := &gosdk.GetPromptResult{} + if err := c.resumeCall(ctx, "prompts/get", request.Params, raw); err != nil { + return nil, err + } + return convertGetPromptResult(raw) + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + res, err := s.GetPrompt(ctx, &gosdk.GetPromptParams{ + Name: request.Params.Name, + Arguments: request.Params.Arguments, + }) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertGetPromptResult(res) +} + +// convertGetPromptResult converts a go-sdk GetPromptResult into its mcp-go-shaped +// equivalent. mcp-go's PromptMessage.Content is the mcp.Content interface with no +// custom unmarshaler, so a JSON round-trip cannot populate it; each message's +// content is re-marshaled and decoded via mcp.UnmarshalContent instead. +func convertGetPromptResult(res *gosdk.GetPromptResult) (*mcp.GetPromptResult, error) { + out := &mcp.GetPromptResult{Description: res.Description} + if len(res.Meta) > 0 { + out.Meta = mcp.NewMetaFromMap(map[string]any(res.Meta)) + } + for _, pm := range res.Messages { + if pm == nil { + continue + } + msg := mcp.PromptMessage{Role: mcp.Role(pm.Role)} + if pm.Content != nil { + cb, err := json.Marshal(pm.Content) + if err != nil { + return nil, fmt.Errorf("marshaling prompt content: %w", err) + } + content, err := mcp.UnmarshalContent(cb) + if err != nil { + return nil, fmt.Errorf("converting prompt content: %w", err) + } + msg.Content = content + } + out.Messages = append(out.Messages, msg) + } + return out, nil +} + +// ListResources lists the server's resources. +func (c *Client) ListResources(ctx context.Context, request mcp.ListResourcesRequest) (*mcp.ListResourcesResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + out := &mcp.ListResourcesResult{} + return out, c.resumeCall(ctx, "resources/list", request.Params, out) + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + res, err := s.ListResources(ctx, &gosdk.ListResourcesParams{Cursor: string(request.Params.Cursor)}) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertResult[mcp.ListResourcesResult](res) +} + +// ListPrompts lists the server's prompts. +func (c *Client) ListPrompts(ctx context.Context, request mcp.ListPromptsRequest) (*mcp.ListPromptsResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + out := &mcp.ListPromptsResult{} + return out, c.resumeCall(ctx, "prompts/list", request.Params, out) + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + res, err := s.ListPrompts(ctx, &gosdk.ListPromptsParams{Cursor: string(request.Params.Cursor)}) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertResult[mcp.ListPromptsResult](res) +} + +// ListResourceTemplates lists the server's resource templates. +func (c *Client) ListResourceTemplates( + ctx context.Context, request mcp.ListResourceTemplatesRequest, +) (*mcp.ListResourceTemplatesResult, error) { + ctx = withErrCapture(ctx) + if c.isResume() { + out := &mcp.ListResourceTemplatesResult{} + return out, c.resumeCall(ctx, "resources/templates/list", request.Params, out) + } + s, err := c.sessionFor() + if err != nil { + return nil, err + } + res, err := s.ListResourceTemplates(ctx, &gosdk.ListResourceTemplatesParams{Cursor: string(request.Params.Cursor)}) + if err != nil { + return nil, mapCallError(ctx, err) + } + return convertResult[mcp.ListResourceTemplatesResult](res) +} + +// OnNotification registers a handler invoked for server-initiated notifications. +// Handlers must be registered before Initialize so they can be wired into the +// underlying go-sdk client. The go-sdk exposes typed notification handlers +// rather than a single catch-all, so this shim synthesizes JSONRPCNotification +// values for the list-changed, progress and logging notifications. +func (c *Client) OnNotification(handler func(notification mcp.JSONRPCNotification)) { + c.notifyMu.Lock() + defer c.notifyMu.Unlock() + c.notify = append(c.notify, handler) +} + +func (c *Client) dispatch(method string) { + c.notifyMu.Lock() + handlers := make([]func(mcp.JSONRPCNotification), len(c.notify)) + copy(handlers, c.notify) + c.notifyMu.Unlock() + n := mcp.JSONRPCNotification{JSONRPC: mcp.JSONRPC_VERSION} + n.Method = method + for _, h := range handlers { + h(n) + } +} + +func (c *Client) installNotificationHandlers(opts *gosdk.ClientOptions) { + opts.ToolListChangedHandler = func(_ context.Context, _ *gosdk.ToolListChangedRequest) { + c.dispatch("notifications/tools/list_changed") + } + opts.PromptListChangedHandler = func(_ context.Context, _ *gosdk.PromptListChangedRequest) { + c.dispatch("notifications/prompts/list_changed") + } + opts.ResourceListChangedHandler = func(_ context.Context, _ *gosdk.ResourceListChangedRequest) { + c.dispatch("notifications/resources/list_changed") + } +} + +// OnConnectionLost registers a handler invoked when the connection to the +// server is lost. It mirrors mcp-go's client.Client.OnConnectionLost, which is +// used to handle transport-level disconnections (e.g. an HTTP/2 idle timeout) +// that should not be treated as fatal errors. +// +// The go-sdk signals a dropped connection by having ClientSession.Wait return. +// The handler is invoked once, with the error Wait reports (nil on a clean +// close). If registered before Initialize, the watch starts once the session +// connects; if registered after, it starts immediately. +func (c *Client) OnConnectionLost(handler func(error)) { + c.connLostMu.Lock() + c.connLost = handler + c.connLostMu.Unlock() + + c.mu.Lock() + session := c.session + c.mu.Unlock() + if session != nil { + c.maybeStartConnLostWatch(session) + } +} + +// maybeStartConnLostWatch starts, at most once, a goroutine that blocks on the +// session's Wait and invokes the registered connection-lost handler when the +// connection drops. +func (c *Client) maybeStartConnLostWatch(session *gosdk.ClientSession) { + c.connLostMu.Lock() + defer c.connLostMu.Unlock() + if c.connLost == nil || c.watchStarted { + return + } + c.watchStarted = true + handler := c.connLost + go func() { + err := session.Wait() + handler(err) + }() +} + +// GetTransport returns the transport handle. For a Streamable HTTP client the +// dynamic type is *transport.StreamableHTTP (as ToolHive expects); otherwise it +// is nil. +func (c *Client) GetTransport() transport.Interface { + if c.streamable != nil { + return c.streamable + } + return nil +} + +// GetSessionId returns the current MCP session ID, if any. +func (c *Client) GetSessionId() string { + if c.streamable != nil { + return c.streamable.GetSessionId() + } + return "" +} + +// IsInitialized reports whether the client has completed initialization. +func (c *Client) IsInitialized() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.session != nil +} + +func (c *Client) sessionFor() (*gosdk.ClientSession, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.session == nil { + return nil, fmt.Errorf("client not initialized: call Initialize first") + } + return c.session, nil +} + +// --- helpers --------------------------------------------------------------- + +// jsonConvert marshals src and unmarshals it into dst. Both mcp-go-shaped and +// go-sdk types encode the identical MCP wire format, so this is a faithful +// structural conversion. +func jsonConvert(src, dst any) error { + b, err := json.Marshal(src) + if err != nil { + return err + } + return json.Unmarshal(b, dst) +} + +// convertResult converts a go-sdk result into its mcp-go-shaped equivalent. +func convertResult[T any](src any) (*T, error) { + out := new(T) + if err := jsonConvert(src, out); err != nil { + return nil, fmt.Errorf("converting result: %w", err) + } + return out, nil +} + +func isZeroCapabilities(c mcp.ClientCapabilities) bool { + b, err := json.Marshal(c) + return err == nil && string(b) == "{}" +} + +// headerRoundTripper injects static and/or per-request headers on every request. +type headerRoundTripper struct { + headers map[string]string + headerFunc transport.HTTPHeaderFunc + base http.RoundTripper +} + +func (h *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + for k, v := range h.headers { + req.Header.Set(k, v) + } + if h.headerFunc != nil { + for k, v := range h.headerFunc(req.Context()) { + req.Header.Set(k, v) + } + } + base := h.base + if base == nil { + base = http.DefaultTransport + } + resp, err := base.RoundTrip(req) + if err == nil { + // Capture non-2xx bodies so the error mappers can re-attach the server's + // message (go-sdk drops it). Restores resp.Body for downstream readers. + captureErrorBody(req, resp) + } + return resp, err +} + +// buildHTTPClient returns an *http.Client honoring the given base client, static +// headers, per-request header function and timeout. It returns nil when no +// customization is needed so the go-sdk uses its default client. +func buildHTTPClient( + base *http.Client, headers map[string]string, headerFunc transport.HTTPHeaderFunc, timeout time.Duration, +) *http.Client { + hc := &http.Client{} + if base != nil { + *hc = *base + } + if timeout > 0 { + hc.Timeout = timeout + } + // Always install the RoundTripper: besides applying static/dynamic headers it + // captures non-2xx response bodies for error enrichment (see errorbody.go). + hc.Transport = &headerRoundTripper{headers: headers, headerFunc: headerFunc, base: hc.Transport} + return hc +} diff --git a/mcpcompat/client/client_test.go b/mcpcompat/client/client_test.go new file mode 100644 index 0000000..8cb0552 --- /dev/null +++ b/mcpcompat/client/client_test.go @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client_test + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +type echoInput struct { + Message string `json:"message"` +} + +// echoToolName is the shared tool name used across the client tests. +const echoToolName = "echo" + +// newTestServer stands up a real go-sdk MCP server exposing a single "echo" +// tool, served over Streamable HTTP via httptest. +func newTestServer(t *testing.T) *httptest.Server { + t.Helper() + srv := gosdk.NewServer(&gosdk.Implementation{Name: "testserver", Version: "9.9.9"}, nil) + gosdk.AddTool(srv, &gosdk.Tool{Name: echoToolName, Description: "echo the message"}, + func(_ context.Context, _ *gosdk.CallToolRequest, in echoInput) (*gosdk.CallToolResult, any, error) { + return &gosdk.CallToolResult{ + Content: []gosdk.Content{&gosdk.TextContent{Text: "echo: " + in.Message}}, + }, nil, nil + }) + handler := gosdk.NewStreamableHTTPHandler(func(*http.Request) *gosdk.Server { return srv }, nil) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// TestStreamableClient_EndToEnd drives the full client path against a live +// go-sdk server: Start, Initialize, ListTools, CallTool, Ping, transport +// handle, and Close. +func TestStreamableClient_EndToEnd(t *testing.T) { + t.Parallel() + ctx := context.Background() + ts := newTestServer(t) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + + require.NoError(t, c.Start(ctx)) + assert.False(t, c.IsInitialized()) + + initRes, err := c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "test-client", Version: "1.0.0"}, + }, + }) + require.NoError(t, err) + assert.Equal(t, "testserver", initRes.ServerInfo.Name) + assert.True(t, c.IsInitialized()) + + tools, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + require.Len(t, tools.Tools, 1) + assert.Equal(t, echoToolName, tools.Tools[0].Name) + + callRes, err := c.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: echoToolName, + Arguments: map[string]any{"message": "hi"}, + }, + }) + require.NoError(t, err) + assert.False(t, callRes.IsError) + require.Len(t, callRes.Content, 1) + txt, ok := mcp.AsTextContent(callRes.Content[0]) + require.True(t, ok) + assert.Equal(t, "echo: hi", txt.Text) + + require.NoError(t, c.Ping(ctx)) + + // GetTransport must yield a *transport.StreamableHTTP, as ToolHive expects. + sh, ok := c.GetTransport().(*transport.StreamableHTTP) + require.True(t, ok) + assert.Equal(t, c.GetSessionId(), sh.GetSessionId()) + + require.NoError(t, c.Close()) +} + +// TestCallBeforeInitialize verifies methods fail cleanly before Initialize. +func TestCallBeforeInitialize(t *testing.T) { + t.Parallel() + c, err := client.NewStreamableHttpClient("http://example.invalid") + require.NoError(t, err) + _, err = c.ListTools(context.Background(), mcp.ListToolsRequest{}) + assert.Error(t, err) +} + +// TestGetTransport_SSEIsNil verifies an SSE client returns a nil transport +// handle (so a type assertion to *StreamableHTTP fails gracefully rather than +// panicking). +func TestGetTransport_SSEIsNil(t *testing.T) { + t.Parallel() + c, err := client.NewSSEMCPClient("http://example.invalid") + require.NoError(t, err) + _, ok := c.GetTransport().(*transport.StreamableHTTP) + assert.False(t, ok) + assert.Empty(t, c.GetSessionId()) +} diff --git a/mcpcompat/client/conn_lost_test.go b/mcpcompat/client/conn_lost_test.go new file mode 100644 index 0000000..793bd38 --- /dev/null +++ b/mcpcompat/client/conn_lost_test.go @@ -0,0 +1,100 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// TestOnConnectionLost_FiresOnClose verifies the handler registered via +// OnConnectionLost is invoked once the session ends (here, via Close, which +// causes the underlying go-sdk ClientSession.Wait to return). +func TestOnConnectionLost_FiresOnClose(t *testing.T) { + t.Parallel() + ctx := context.Background() + ts := newTestServer(t) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + + lost := make(chan error, 1) + // Register before Initialize: the watch must start once connected. + c.OnConnectionLost(func(err error) { lost <- err }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "c", Version: "1"}, + }, + }) + require.NoError(t, err) + + require.NoError(t, c.Close()) + + select { + case <-lost: + // handler fired as expected (error value is transport-dependent). + case <-time.After(5 * time.Second): + t.Fatal("OnConnectionLost handler was not invoked after Close") + } +} + +// TestOnConnectionLost_RegisterAfterInitialize verifies registration works even +// when the client is already connected. +func TestOnConnectionLost_RegisterAfterInitialize(t *testing.T) { + t.Parallel() + ctx := context.Background() + ts := newTestServer(t) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "c", Version: "1"}, + }, + }) + require.NoError(t, err) + + lost := make(chan error, 1) + c.OnConnectionLost(func(err error) { lost <- err }) + + require.NoError(t, c.Close()) + + select { + case <-lost: + case <-time.After(5 * time.Second): + t.Fatal("OnConnectionLost handler was not invoked after Close") + } +} + +// TestOnConnectionLost_NoHandlerNoWatch verifies registering no handler is safe +// and Close still works. +func TestOnConnectionLost_NoHandlerNoWatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + ts := newTestServer(t) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "c", Version: "1"}, + }, + }) + require.NoError(t, err) + assert.NoError(t, c.Close()) +} diff --git a/mcpcompat/client/errorbody.go b/mcpcompat/client/errorbody.go new file mode 100644 index 0000000..033912e --- /dev/null +++ b/mcpcompat/client/errorbody.go @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "bytes" + "context" + "io" + "net/http" +) + +// go-sdk's Streamable HTTP client reports a non-2xx response using only +// http.StatusText(code) (e.g. "403 Forbidden") and discards the response body. +// mcp-go instead included the body ("request failed with status N: "), +// which is where servers put actionable detail — e.g. ToolHive's authorization +// middleware writes "Unauthorized" into a 403 body. To preserve that, the client +// installs a RoundTripper that captures the body of a non-2xx response into a +// per-call holder carried on the request context, and the error mappers +// re-attach it to the returned error. + +// maxCapturedErrBody bounds how much of an error response body we read. +const maxCapturedErrBody = 8 << 10 // 8 KiB + +type errBody struct { + status int + body string +} + +type errBodyKey struct{} + +// withErrCapture returns a context carrying a holder for a captured non-2xx +// response body. Pass the returned context to the underlying request so the +// RoundTripper can populate it, then hand the same context to mapCallError / +// mapConnectError to enrich the error. +func withErrCapture(ctx context.Context) context.Context { + return context.WithValue(ctx, errBodyKey{}, &errBody{}) +} + +func capturedErr(ctx context.Context) *errBody { + if v, ok := ctx.Value(errBodyKey{}).(*errBody); ok { + return v + } + return nil +} + +// captureErrorBody records a non-2xx response body into the context holder (if +// present) and restores resp.Body so downstream readers are unaffected. Safe on +// nil/2xx responses (no-op). +func captureErrorBody(req *http.Request, resp *http.Response) { + if resp == nil || resp.StatusCode < 400 || resp.Body == nil { + return + } + h := capturedErr(req.Context()) + if h == nil || h.body != "" { + return + } + data, err := io.ReadAll(io.LimitReader(resp.Body, maxCapturedErrBody)) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(data)) + if err == nil && len(data) > 0 { + h.status = resp.StatusCode + h.body = string(data) + } +} diff --git a/mcpcompat/client/errorbody_internal_test.go b/mcpcompat/client/errorbody_internal_test.go new file mode 100644 index 0000000..f33303a --- /dev/null +++ b/mcpcompat/client/errorbody_internal_test.go @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestErrorBodyEnrichment reproduces the authz case: a tool call denied with an +// HTTP 403 whose body says "Unauthorized". go-sdk's client would surface only +// "403 Forbidden" (status text, body dropped); the RoundTripper captures the +// body and mapCallError re-attaches it so callers (and ToolHive's authz e2e +// test) see "Unauthorized". +func TestErrorBodyEnrichment(t *testing.T) { + t.Parallel() + + const bodyText = "Unauthorized: request denied by policy" + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(bodyText)) + })) + t.Cleanup(ts.Close) + + ctx := withErrCapture(context.Background()) + hc := buildHTTPClient(nil, nil, nil, 0) + require.NotNil(t, hc) + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ts.URL, http.NoBody) + require.NoError(t, err) + resp, err := hc.Do(req) + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + + // The RoundTripper must have restored the body for downstream readers. + restored, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Equal(t, bodyText, string(restored), "response body must be restored") + + // Simulate go-sdk's status-only error and confirm enrichment surfaces the body. + gosdkErr := fmt.Errorf(`calling "tools/call": %s`, http.StatusText(resp.StatusCode)) + assert.NotContains(t, gosdkErr.Error(), "Unauthorized", "precondition: go-sdk error lacks the body") + + enriched := mapCallError(ctx, gosdkErr) + require.Error(t, enriched) + assert.Contains(t, enriched.Error(), "Unauthorized", "enriched error must include the server body") + assert.Contains(t, enriched.Error(), "403") +} + +// TestErrorBodyEnrichment_NoCaptureNoop verifies enrichment is a no-op when no +// body was captured (e.g. a 2xx path or a non-HTTP error). +func TestErrorBodyEnrichment_NoCaptureNoop(t *testing.T) { + t.Parallel() + ctx := withErrCapture(context.Background()) + orig := fmt.Errorf("some transport failure") + assert.Equal(t, orig, mapCallError(ctx, orig)) +} diff --git a/mcpcompat/client/errors.go b/mcpcompat/client/errors.go new file mode 100644 index 0000000..99270ab --- /dev/null +++ b/mcpcompat/client/errors.go @@ -0,0 +1,95 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// enrichWithResponseBody re-attaches a captured non-2xx HTTP response body to +// err (go-sdk reports only the status text and drops the body). Matches mcp-go's +// "request failed with status N: " so server-provided detail — e.g. +// ToolHive's authorization middleware writing "Unauthorized" into a 403 body — +// reaches callers. No-op when nothing was captured. +func enrichWithResponseBody(ctx context.Context, err error) error { + if err == nil { + return nil + } + h := capturedErr(ctx) + if h == nil || h.body == "" { + return err + } + return fmt.Errorf("%w: request failed with status %d: %s", err, h.status, strings.TrimSpace(h.body)) +} + +// mapConnectError maps an error returned by the underlying go-sdk Connect call +// onto the transport-level sentinels ToolHive checks for. +func mapConnectError(ctx context.Context, err error) error { + return mapTransportError(enrichWithResponseBody(ctx, err)) +} + +// mapCallError maps an error returned by an underlying go-sdk request call onto +// the sentinels ToolHive checks for. A JSON-RPC -32601 response is surfaced as +// mcp.ErrMethodNotFound (as mcp-go did) so callers that recover from a backend +// lacking an optional method — e.g. resources/list or prompts/list — via +// errors.Is(err, mcp.ErrMethodNotFound) keep working. +func mapCallError(ctx context.Context, err error) error { + if err == nil { + return nil + } + err = enrichWithResponseBody(ctx, err) + var wireErr *jsonrpc.Error + if errors.As(err, &wireErr) && wireErr.Code == jsonrpc.CodeMethodNotFound { + return errors.Join(mcp.ErrMethodNotFound, err) + } + return mapTransportError(err) +} + +// mapTransportError inspects err and, when it recognizes an HTTP auth/session +// failure, returns an error that satisfies the errors.Is/errors.As checks +// ToolHive performs against the transport package's sentinels. +// +// NOTE: the go-sdk does not currently expose a typed error carrying the HTTP +// status code, so detection is best-effort based on the error text. When the +// pattern is not recognized the original error is returned unchanged. This is +// the one area of the client shim where exact parity with mcp-go's OAuth flow +// may need refinement as the go-sdk's error surface evolves. +func mapTransportError(err error) error { + if err == nil { + return nil + } + msg := strings.ToLower(err.Error()) + + switch { + case strings.Contains(msg, "401") || strings.Contains(msg, "unauthorized"): + // 401: ToolHive checks both ErrAuthorizationRequired (and As + // *AuthorizationRequiredError / *transport.Error) and ErrUnauthorized. + return transport.NewError(errors.Join( + &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(err)}, + transport.ErrUnauthorized, + err, + )) + case strings.Contains(msg, "404") && strings.Contains(msg, "session"): + return transport.NewError(errors.Join(transport.ErrSessionTerminated, err)) + case strings.Contains(msg, "legacy") || strings.Contains(msg, "method not allowed") || strings.Contains(msg, "405"): + return transport.NewError(errors.Join(transport.ErrLegacySSEServer, err)) + default: + return err + } +} + +// extractResourceMetadataURL is a placeholder for parsing the RFC 9728 +// resource_metadata parameter out of a WWW-Authenticate header. The go-sdk does +// not surface the header on the error today, so this returns empty for now. +func extractResourceMetadataURL(_ error) string { + return "" +} diff --git a/mcpcompat/client/errors_internal_test.go b/mcpcompat/client/errors_internal_test.go new file mode 100644 index 0000000..911e47d --- /dev/null +++ b/mcpcompat/client/errors_internal_test.go @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" +) + +func TestMapTransportError_Unauthorized(t *testing.T) { + t.Parallel() + err := mapTransportError(errors.New("request failed: 401 Unauthorized")) + + // ToolHive branches on all of these for its OAuth/401 handling. + assert.True(t, errors.Is(err, transport.ErrUnauthorized), "errors.Is ErrUnauthorized") + assert.True(t, errors.Is(err, transport.ErrAuthorizationRequired), "errors.Is ErrAuthorizationRequired") + + var te *transport.Error + assert.True(t, errors.As(err, &te), "errors.As *transport.Error") + + var are *transport.AuthorizationRequiredError + assert.True(t, errors.As(err, &are), "errors.As *AuthorizationRequiredError") +} + +func TestMapTransportError_SessionTerminated(t *testing.T) { + t.Parallel() + err := mapTransportError(errors.New("server returned 404: session not found")) + assert.True(t, errors.Is(err, transport.ErrSessionTerminated)) +} + +func TestMapTransportError_LegacySSE(t *testing.T) { + t.Parallel() + err := mapTransportError(errors.New("405 method not allowed")) + assert.True(t, errors.Is(err, transport.ErrLegacySSEServer)) +} + +func TestMapTransportError_Passthrough(t *testing.T) { + t.Parallel() + orig := errors.New("some unrelated failure") + assert.Equal(t, orig, mapTransportError(orig)) +} + +func TestMapTransportError_Nil(t *testing.T) { + t.Parallel() + assert.NoError(t, mapTransportError(nil)) +} diff --git a/mcpcompat/client/header_test.go b/mcpcompat/client/header_test.go new file mode 100644 index 0000000..3345457 --- /dev/null +++ b/mcpcompat/client/header_test.go @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client_test + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// TestWithHTTPHeaderFunc_PerRequestHeaders verifies that headers returned by the +// function passed to transport.WithHTTPHeaderFunc are attached to outgoing +// requests and observed by the server. +func TestWithHTTPHeaderFunc_PerRequestHeaders(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := gosdk.NewServer(&gosdk.Implementation{Name: "hdr", Version: "1"}, nil) + gosdk.AddTool(srv, &gosdk.Tool{Name: echoToolName, Description: echoToolName}, + func(_ context.Context, _ *gosdk.CallToolRequest, _ echoInput) (*gosdk.CallToolResult, any, error) { + return &gosdk.CallToolResult{Content: []gosdk.Content{&gosdk.TextContent{Text: "ok"}}}, nil, nil + }) + inner := gosdk.NewStreamableHTTPHandler(func(*http.Request) *gosdk.Server { return srv }, nil) + + var mu sync.Mutex + seen := map[string]string{} + wrapped := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if v := r.Header.Get("X-Caller"); v != "" { + mu.Lock() + seen["X-Caller"] = v + mu.Unlock() + } + inner.ServeHTTP(w, r) + }) + ts := httptest.NewServer(wrapped) + t.Cleanup(ts.Close) + + c, err := client.NewStreamableHttpClient(ts.URL, + transport.WithHTTPHeaderFunc(func(context.Context) map[string]string { + return map[string]string{"X-Caller": "tenant-42"} + }), + ) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "c", Version: "1"}, + }, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, "tenant-42", seen["X-Caller"], "server should observe the per-request header") +} diff --git a/mcpcompat/client/meta_roundtrip_test.go b/mcpcompat/client/meta_roundtrip_test.go new file mode 100644 index 0000000..55f06f3 --- /dev/null +++ b/mcpcompat/client/meta_roundtrip_test.go @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client_test + +import ( + "context" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// TestCallToolMetadataRoundTrip verifies that the request _meta survives the +// full client->server round-trip: the client sends a tools/call with a _meta +// field, the server handler observes it, and echoes it back in the result's +// _meta. ToolHive relies on this to propagate metadata through vMCP to backends. +func TestCallToolMetadataRoundTrip(t *testing.T) { + t.Parallel() + + var seen map[string]any + mcpSrv := server.NewMCPServer("srv", "1.0.0") + mcpSrv.AddTool(mcp.NewTool("echo", mcp.WithDescription("echoes meta")), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if req.Params.Meta != nil { + seen = req.Params.Meta.AdditionalFields + } + res := mcp.NewToolResultText("ok") + // Echo the request meta back on the result. + res.Meta = &mcp.Meta{AdditionalFields: map[string]any{"echoed": "yes"}} + return res, nil + }) + ts := httptest.NewServer(server.NewStreamableHTTPServer(mcpSrv)) + defer ts.Close() + + c, err := mcpclient.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{}) + require.NoError(t, err) + defer func() { _ = c.Close() }() + + res, err := c.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "echo", + Meta: &mcp.Meta{AdditionalFields: map[string]any{"trace-id": "abc123"}}, + }, + }) + require.NoError(t, err) + + // The server handler must have seen the request _meta. + require.NotNil(t, seen, "server tool handler must receive the request _meta") + assert.Equal(t, "abc123", seen["trace-id"], "request _meta must reach the handler") + + // The result _meta must survive the trip back to the client. + require.NotNil(t, res.Meta, "result _meta must be preserved to the client") + assert.Equal(t, "yes", res.Meta.AdditionalFields["echoed"]) +} diff --git a/mcpcompat/client/resume.go b/mcpcompat/client/resume.go new file mode 100644 index 0000000..1c60882 --- /dev/null +++ b/mcpcompat/client/resume.go @@ -0,0 +1,190 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync/atomic" + + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" +) + +// resumeProtocolVersion is the MCP protocol version header sent on resumed +// requests. A resumed client skips the initialize handshake, so it advertises a +// widely-supported spec revision; the server seeds the session's negotiated +// version from this header when it rehydrates the session. +const resumeProtocolVersion = "2025-06-18" + +// resumeState holds the minimal raw JSON-RPC-over-HTTP machinery used when a +// client resumes a pre-existing session (transport.WithSession) without calling +// Initialize. The go-sdk client always performs the initialize handshake and has +// no resume primitive, so a resumed session cannot use it; this path speaks the +// Streamable HTTP wire protocol directly, matching mcp-go's client resume. +type resumeState struct { + nextID atomic.Int64 + client *http.Client +} + +// isResume reports whether this client should use the raw resume path: a +// Streamable HTTP transport with a preset session ID (transport.WithSession) and +// no go-sdk session established (Initialize was not called). +func (c *Client) isResume() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.session == nil && c.streamable != nil && c.streamable.GetSessionId() != "" +} + +// resumeCall issues a single JSON-RPC request over HTTP POST carrying the +// resumed session ID and decodes the result into out. It does not perform (and +// must not perform) initialization. +func (c *Client) resumeCall(ctx context.Context, method string, params, out any) error { + c.mu.Lock() + if c.resume == nil { + c.resume = &resumeState{ + client: buildHTTPClient( + c.streamable.HTTPClient(), c.streamable.Headers(), c.streamable.HeaderFunc(), c.streamable.Timeout(), + ), + } + if c.resume.client == nil { + c.resume.client = http.DefaultClient + } + } + rs := c.resume + endpoint := c.streamable.Endpoint() + sessionID := c.streamable.GetSessionId() + c.mu.Unlock() + + id := rs.nextID.Add(1) + reqMsg := map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": method, + } + if params != nil { + reqMsg["params"] = params + } + body, err := json.Marshal(reqMsg) + if err != nil { + return fmt.Errorf("marshaling %s request: %w", method, err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("building %s request: %w", method, err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + httpReq.Header.Set("Mcp-Session-Id", sessionID) + httpReq.Header.Set("MCP-Protocol-Version", resumeProtocolVersion) + + resp, err := rs.client.Do(httpReq) + if err != nil { + return fmt.Errorf("sending %s request: %w", method, err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode == http.StatusNotFound { + // A 404 means the server no longer recognizes the session (terminated or + // expired in the shared store). Surface transport.ErrSessionTerminated so + // callers that errors.Is(err, transport.ErrSessionTerminated) — as ToolHive + // does to detect cross-replica lazy eviction — keep working. + return transport.NewError(errors.Join( + transport.ErrSessionTerminated, + fmt.Errorf("resumed session %q terminated (404)", sessionID), + )) + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + return fmt.Errorf("%s request failed: HTTP %d: %s", method, resp.StatusCode, strings.TrimSpace(string(snippet))) + } + + result, jsonErr, err := readRPCResponse(resp, id) + if err != nil { + return err + } + if jsonErr != nil { + return fmt.Errorf("%s: %s", method, jsonErr.Message) + } + if out != nil && len(result) > 0 { + if err := json.Unmarshal(result, out); err != nil { + return fmt.Errorf("decoding %s result: %w", method, err) + } + } + return nil +} + +// rpcError is the JSON-RPC error object. +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// rpcEnvelope is a JSON-RPC response/notification envelope. +type rpcEnvelope struct { + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *rpcError `json:"error"` + Method string `json:"method"` +} + +// readRPCResponse reads the JSON-RPC response matching wantID from an HTTP +// response, handling both application/json and text/event-stream bodies. It +// ignores server->client requests/notifications interleaved on an SSE stream +// (a resumed client does not service them). +func readRPCResponse(resp *http.Response, wantID int64) (result json.RawMessage, rpcErr *rpcError, err error) { + if strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") { + var env rpcEnvelope + if derr := json.NewDecoder(resp.Body).Decode(&env); derr != nil { + return nil, nil, fmt.Errorf("decoding JSON response: %w", derr) + } + return env.Result, env.Error, nil + } + + // SSE: scan for the message whose id matches wantID. + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + for sc.Scan() { + data, ok := strings.CutPrefix(sc.Text(), "data:") + if !ok { + continue + } + var env rpcEnvelope + if json.Unmarshal([]byte(strings.TrimSpace(data)), &env) != nil { + continue + } + // Only a response carries a result or error; skip server->client + // requests/notifications (they have a method set). + if env.Method != "" { + continue + } + if !idMatches(env.ID, wantID) { + continue + } + return env.Result, env.Error, nil + } + if scErr := sc.Err(); scErr != nil { + return nil, nil, fmt.Errorf("reading SSE stream: %w", scErr) + } + return nil, nil, fmt.Errorf("no JSON-RPC response for id %d in stream", wantID) +} + +// idMatches reports whether the raw JSON id equals wantID. +func idMatches(raw json.RawMessage, wantID int64) bool { + if len(raw) == 0 { + return false + } + var n int64 + if err := json.Unmarshal(raw, &n); err == nil { + return n == wantID + } + return false +} diff --git a/mcpcompat/client/resume_test.go b/mcpcompat/client/resume_test.go new file mode 100644 index 0000000..e598d5e --- /dev/null +++ b/mcpcompat/client/resume_test.go @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client_test + +import ( + "context" + "fmt" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// sharedSessionManager is an in-memory SessionIdManager standing in for +// ToolHive's Redis-backed cross-replica store. +type sharedSessionManager struct { + mu sync.Mutex + valid map[string]bool + terminated map[string]bool +} + +func newSharedSessionManager() *sharedSessionManager { + return &sharedSessionManager{valid: map[string]bool{}, terminated: map[string]bool{}} +} + +func (m *sharedSessionManager) Generate() string { + m.mu.Lock() + defer m.mu.Unlock() + id := uuid.NewString() + m.valid[id] = true + return id +} + +func (m *sharedSessionManager) Validate(sessionID string) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.terminated[sessionID] { + return true, nil + } + if !m.valid[sessionID] { + return false, fmt.Errorf("session %q not found", sessionID) + } + return false, nil +} + +func (m *sharedSessionManager) Terminate(sessionID string) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.terminated[sessionID] = true + return false, nil +} + +// TestClientResumeWithoutInitialize verifies STAGE B: a client created with +// transport.WithSession(id) can Start and issue requests WITHOUT calling +// Initialize, resuming a session established by another client. It exercises the +// full round-trip against a shim server whose Streamable HTTP transport +// rehydrates the resumed session. +func TestClientResumeWithoutInitialize(t *testing.T) { + t.Parallel() + + mgr := newSharedSessionManager() + + newReplica := func(name string) *httptest.Server { + s := server.NewMCPServer(name, "1.0.0") + s.AddTool(mcp.NewTool("greet", mcp.WithDescription("greets")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("hello"), nil + }) + ts := httptest.NewServer(server.NewStreamableHTTPServer(s, server.WithSessionIdManager(mgr))) + return ts + } + + // Replica A handles initialize; replica B (separate instance sharing the + // session manager) is where the client resumes — the real cross-replica flow. + tsA := newReplica("A") + defer tsA.Close() + tsB := newReplica("B") + defer tsB.Close() + + // Client 1: normal initialize on A, capture the session ID and tools. + client1, err := mcpclient.NewStreamableHttpClient(tsA.URL) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + require.NoError(t, client1.Start(ctx)) + _, err = client1.Initialize(ctx, mcp.InitializeRequest{}) + require.NoError(t, err) + sid := client1.GetSessionId() + require.NotEmpty(t, sid) + tools1, err := client1.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + require.Len(t, tools1.Tools, 1) + defer func() { _ = client1.Close() }() + + // Client 2: resume with the SAME session ID against replica B, NO Initialize. + client2, err := mcpclient.NewStreamableHttpClient(tsB.URL, transport.WithSession(sid)) + require.NoError(t, err) + require.NoError(t, client2.Start(ctx)) + defer func() { _ = client2.Close() }() + + require.Equal(t, sid, client2.GetSessionId(), "resumed client reports the resumed session ID") + + tools2, err := client2.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err, "resumed ListTools must succeed without Initialize") + names := make([]string, 0, len(tools2.Tools)) + for _, tl := range tools2.Tools { + names = append(names, tl.Name) + } + assert.Equal(t, []string{"greet"}, names, "resumed session returns the same tools") + + // A tool call over the resumed session must also work. + res, err := client2.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "greet"}, + }) + require.NoError(t, err) + require.Len(t, res.Content, 1) + + // Terminate the session in the shared store; the resumed client's next request + // must surface transport.ErrSessionTerminated (the 404 -> sentinel mapping + // ToolHive relies on to detect cross-replica lazy eviction). + _, _ = mgr.Terminate(sid) + _, err = client2.ListTools(ctx, mcp.ListToolsRequest{}) + require.Error(t, err) + assert.ErrorIs(t, err, transport.ErrSessionTerminated, + "resumed client must report ErrSessionTerminated after the session is terminated") +} diff --git a/mcpcompat/client/transport/errors.go b/mcpcompat/client/transport/errors.go new file mode 100644 index 0000000..38b9f34 --- /dev/null +++ b/mcpcompat/client/transport/errors.go @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transport + +import ( + "errors" + "fmt" +) + +// These error values mirror github.com/mark3labs/mcp-go/client/transport so that +// ToolHive's auth-detection code (errors.Is / errors.As against these) keeps +// working after the import swap. The go-sdk-backed client (see the client +// package) maps the underlying SDK/HTTP failures onto these sentinels. +var ( + // ErrAuthorizationRequired indicates the server requires authorization + // (HTTP 401 with a WWW-Authenticate header per RFC 9728). + ErrAuthorizationRequired = errors.New("authorization required") + + // ErrOAuthAuthorizationRequired indicates OAuth authorization is required + // and no valid token is available. + ErrOAuthAuthorizationRequired = errors.New("no valid token available, authorization required") + + // ErrUnauthorized indicates an HTTP 401 response. + ErrUnauthorized = fmt.Errorf("unauthorized (401)") + + // ErrLegacySSEServer indicates the server returned 4xx for the initialize + // POST, which usually means it is a legacy SSE-only server. + ErrLegacySSEServer = fmt.Errorf("server returned 4xx for initialize POST, likely a legacy SSE server") + + // ErrSessionTerminated indicates the server no longer recognizes the + // current session (HTTP 404); the client must re-initialize. + ErrSessionTerminated = fmt.Errorf("session terminated (404). need to re-initialize") +) + +// Error wraps a transport-level error. It mirrors mcp-go's transport.Error so +// that errors.As(err, new(*transport.Error)) continues to work. +type Error struct { + Err error +} + +func (e *Error) Error() string { + return fmt.Sprintf("transport error: %v", e.Err) +} + +// Unwrap returns the wrapped error. +func (e *Error) Unwrap() error { + return e.Err +} + +// NewError wraps err in a *Error. +func NewError(err error) *Error { + return &Error{Err: err} +} + +// AuthorizationRequiredError is returned for 401 responses carrying a +// WWW-Authenticate header. It mirrors mcp-go's transport.AuthorizationRequiredError. +type AuthorizationRequiredError struct { + // ResourceMetadataURL is extracted from the WWW-Authenticate header per RFC 9728. + ResourceMetadataURL string +} + +func (*AuthorizationRequiredError) Error() string { + return ErrAuthorizationRequired.Error() +} + +// Unwrap returns ErrAuthorizationRequired so errors.Is works. +func (*AuthorizationRequiredError) Unwrap() error { + return ErrAuthorizationRequired +} diff --git a/mcpcompat/client/transport/sse.go b/mcpcompat/client/transport/sse.go new file mode 100644 index 0000000..d001475 --- /dev/null +++ b/mcpcompat/client/transport/sse.go @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transport + +import "net/http" + +// SSE holds SSE (Server-Sent Events) transport configuration. It is the option +// target for the SSE client, mirroring mcp-go's transport package. +type SSE struct { + endpoint string + httpClient *http.Client + headers map[string]string +} + +// NewSSE creates an SSE transport config for the given endpoint and applies the +// supplied options. Used by the client package. +func NewSSE(endpoint string, opts ...ClientOption) *SSE { + s := &SSE{endpoint: endpoint} + for _, opt := range opts { + opt(s) + } + return s +} + +// Endpoint returns the configured endpoint URL. +func (s *SSE) Endpoint() string { return s.endpoint } + +// HTTPClient returns the configured HTTP client, or nil to use the default. +func (s *SSE) HTTPClient() *http.Client { return s.httpClient } + +// Headers returns the configured static headers. +func (s *SSE) Headers() map[string]string { return s.headers } + +// ClientOption configures an SSE transport. +type ClientOption func(*SSE) + +// WithHTTPClient sets a custom HTTP client for the SSE transport. +func WithHTTPClient(httpClient *http.Client) ClientOption { + return func(s *SSE) { s.httpClient = httpClient } +} + +// WithHeaders sets static headers for the SSE client. +func WithHeaders(headers map[string]string) ClientOption { + return func(s *SSE) { s.headers = headers } +} diff --git a/mcpcompat/client/transport/streamable_http.go b/mcpcompat/client/transport/streamable_http.go new file mode 100644 index 0000000..2da0a87 --- /dev/null +++ b/mcpcompat/client/transport/streamable_http.go @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package transport is a drop-in compatibility shim for +// github.com/mark3labs/mcp-go/client/transport. It provides the transport +// option types and error values that ToolHive references; the actual transport +// is driven by the official go-sdk from the client package. +// +// This package intentionally does not import the go-sdk: it only carries +// configuration and error types. The client package reads the exported +// accessors here to construct the underlying go-sdk transport. +package transport + +import ( + "context" + "log/slog" + "net/http" + "time" +) + +// HTTPHeaderFunc returns headers to attach to each Streamable HTTP request. It +// mirrors mcp-go's transport.HTTPHeaderFunc and is evaluated per request, so it +// can return values that change over the lifetime of the connection. +type HTTPHeaderFunc func(context.Context) map[string]string + +// Interface is the transport handle returned by client.GetTransport. It mirrors +// mcp-go's transport.Interface for the subset ToolHive uses (type-asserting to +// *StreamableHTTP and reading the session ID). +type Interface interface { + // GetSessionId returns the transport-level MCP session ID, if any. + GetSessionId() string +} + +// StreamableHTTP holds Streamable HTTP transport configuration and, once +// connected, the live session ID. It is both the option target (mirroring +// mcp-go, whose options mutate the transport struct) and the handle returned by +// client.GetTransport. +type StreamableHTTP struct { + endpoint string + httpClient *http.Client + headers map[string]string + headerFunc HTTPHeaderFunc + timeout time.Duration + logger *slog.Logger + continuousListening bool + sessionID string +} + +// NewStreamableHTTP creates a StreamableHTTP for the given endpoint and applies +// the supplied options. It is used by the client package. +func NewStreamableHTTP(endpoint string, opts ...StreamableHTTPCOption) *StreamableHTTP { + s := &StreamableHTTP{endpoint: endpoint} + for _, opt := range opts { + opt(s) + } + return s +} + +// GetSessionId returns the current MCP session ID (empty if not yet connected +// or if the transport is stateless). +func (s *StreamableHTTP) GetSessionId() string { return s.sessionID } + +// SetSessionID records the live session ID. Used by the client package after +// the underlying go-sdk session is established. +func (s *StreamableHTTP) SetSessionID(id string) { s.sessionID = id } + +// Endpoint returns the configured endpoint URL. +func (s *StreamableHTTP) Endpoint() string { return s.endpoint } + +// HTTPClient returns the configured HTTP client, or nil to use the default. +func (s *StreamableHTTP) HTTPClient() *http.Client { return s.httpClient } + +// Headers returns the configured static headers. +func (s *StreamableHTTP) Headers() map[string]string { return s.headers } + +// HeaderFunc returns the configured per-request header function, or nil. +func (s *StreamableHTTP) HeaderFunc() HTTPHeaderFunc { return s.headerFunc } + +// Timeout returns the configured HTTP timeout (0 if unset). +func (s *StreamableHTTP) Timeout() time.Duration { return s.timeout } + +// Logger returns the configured logger, if any. +func (s *StreamableHTTP) Logger() *slog.Logger { return s.logger } + +// ContinuousListening reports whether a standalone SSE listening stream was +// requested. +func (s *StreamableHTTP) ContinuousListening() bool { return s.continuousListening } + +// StreamableHTTPCOption configures a StreamableHTTP transport. +type StreamableHTTPCOption func(*StreamableHTTP) + +// WithHTTPTimeout sets the HTTP timeout for the Streamable HTTP transport. +func WithHTTPTimeout(timeout time.Duration) StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.timeout = timeout } +} + +// WithHTTPBasicClient sets a custom HTTP client for the Streamable HTTP transport. +func WithHTTPBasicClient(client *http.Client) StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.httpClient = client } +} + +// WithHTTPHeaders sets static headers sent on each request. +func WithHTTPHeaders(headers map[string]string) StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.headers = headers } +} + +// WithHTTPHeaderFunc sets a function that returns headers for each Streamable +// HTTP request. It mirrors mcp-go's transport.WithHTTPHeaderFunc; the function +// is evaluated per request, so its returned headers may vary over time. +func WithHTTPHeaderFunc(headerFunc HTTPHeaderFunc) StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.headerFunc = headerFunc } +} + +// WithSession sets an initial session ID (for resuming a session). +func WithSession(sessionID string) StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.sessionID = sessionID } +} + +// WithContinuousListening enables a standalone SSE stream for server-initiated +// messages. +func WithContinuousListening() StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.continuousListening = true } +} + +// WithHTTPLogger sets a logger for the Streamable HTTP transport. +func WithHTTPLogger(logger *slog.Logger) StreamableHTTPCOption { + return func(s *StreamableHTTP) { s.logger = logger } +} diff --git a/mcpcompat/client/transport/transport_test.go b/mcpcompat/client/transport/transport_test.go new file mode 100644 index 0000000..373f8a9 --- /dev/null +++ b/mcpcompat/client/transport/transport_test.go @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transport_test + +import ( + "context" + "errors" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" +) + +func TestWithHTTPHeaderFunc(t *testing.T) { + t.Parallel() + + // No header func by default. + sh := transport.NewStreamableHTTP("https://example.test/mcp") + assert.Nil(t, sh.HeaderFunc()) + + // WithHTTPHeaderFunc records a per-request header function. + const hdrKey = "X-Test" + called := false + fn := func(context.Context) map[string]string { + called = true + return map[string]string{hdrKey: "v"} + } + sh = transport.NewStreamableHTTP("https://example.test/mcp", transport.WithHTTPHeaderFunc(fn)) + got := sh.HeaderFunc() + require.NotNil(t, got) + hdrs := got(context.Background()) + assert.True(t, called) + assert.Equal(t, "v", hdrs[hdrKey]) +} + +func TestStreamableHTTPOptions(t *testing.T) { + t.Parallel() + hc := &http.Client{} + sh := transport.NewStreamableHTTP("https://example.test/mcp", + transport.WithHTTPTimeout(5*time.Second), + transport.WithHTTPBasicClient(hc), + transport.WithHTTPHeaders(map[string]string{"X-Test": "1"}), + transport.WithSession("sess-123"), + transport.WithContinuousListening(), + ) + + assert.Equal(t, "https://example.test/mcp", sh.Endpoint()) + assert.Equal(t, 5*time.Second, sh.Timeout()) + assert.Same(t, hc, sh.HTTPClient()) + assert.Equal(t, "1", sh.Headers()["X-Test"]) + assert.True(t, sh.ContinuousListening()) + assert.Equal(t, "sess-123", sh.GetSessionId()) + + sh.SetSessionID("live-456") + assert.Equal(t, "live-456", sh.GetSessionId()) + + // StreamableHTTP satisfies the Interface returned by client.GetTransport. + var _ transport.Interface = sh +} + +func TestSSEOptions(t *testing.T) { + t.Parallel() + hc := &http.Client{} + s := transport.NewSSE("https://example.test/sse", + transport.WithHTTPClient(hc), + transport.WithHeaders(map[string]string{"X-Test": "2"}), + ) + assert.Equal(t, "https://example.test/sse", s.Endpoint()) + assert.Same(t, hc, s.HTTPClient()) + assert.Equal(t, "2", s.Headers()["X-Test"]) +} + +func TestErrorTypesUnwrap(t *testing.T) { + t.Parallel() + + are := &transport.AuthorizationRequiredError{ResourceMetadataURL: "https://as.example/meta"} + assert.ErrorIs(t, are, transport.ErrAuthorizationRequired) + + wrapped := transport.NewError(errors.New("boom")) + assert.Contains(t, wrapped.Error(), "boom") + assert.EqualError(t, wrapped.Unwrap(), "boom") +} diff --git a/mcpcompat/mcp/alias.go b/mcpcompat/mcp/alias.go new file mode 100644 index 0000000..c235668 --- /dev/null +++ b/mcpcompat/mcp/alias.go @@ -0,0 +1,340 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +// This file is the single point in the compatibility layer that references +// mcp-go. Every symbol below is re-exported from github.com/mark3labs/mcp-go/mcp +// so that migrating call sites is a pure import swap. To eventually drop the +// mcp-go dependency, replace each alias/assignment here with a standalone +// definition copied from the mcp-go source (see doc.go). The wire-format tests +// guard that conversion. + +import ( + mcpgo "github.com/mark3labs/mcp-go/mcp" +) + +// ---------------------------------------------------------------------------- +// Protocol constants +// ---------------------------------------------------------------------------- + +// LATEST_PROTOCOL_VERSION mirrors mcp-go's latest supported protocol version. +// +//nolint:revive,staticcheck // name intentionally matches mcp-go for drop-in compatibility. +const LATEST_PROTOCOL_VERSION = mcpgo.LATEST_PROTOCOL_VERSION + +// JSONRPC_VERSION is the JSON-RPC version used by MCP ("2.0"). +// +//nolint:revive,staticcheck // name intentionally matches mcp-go for drop-in compatibility. +const JSONRPC_VERSION = mcpgo.JSONRPC_VERSION + +// MCP method names. +const ( + MethodInitialize = mcpgo.MethodInitialize + MethodToolsList = mcpgo.MethodToolsList + MethodToolsCall = mcpgo.MethodToolsCall + MethodResourcesList = mcpgo.MethodResourcesList + MethodResourcesTemplatesList = mcpgo.MethodResourcesTemplatesList + MethodResourcesRead = mcpgo.MethodResourcesRead + MethodPromptsList = mcpgo.MethodPromptsList + MethodPromptsGet = mcpgo.MethodPromptsGet + MethodListRoots = mcpgo.MethodListRoots + MethodSetLogLevel = mcpgo.MethodSetLogLevel + MethodPing = mcpgo.MethodPing + MethodElicitationCreate = mcpgo.MethodElicitationCreate + + // MethodNotificationInitialized indicates the client finished initialization. + MethodNotificationInitialized = mcpgo.MethodNotificationInitialized +) + +// JSON-RPC / MCP error codes. +// +//nolint:revive,staticcheck // names intentionally match mcp-go for drop-in compatibility. +const ( + PARSE_ERROR = mcpgo.PARSE_ERROR + INVALID_REQUEST = mcpgo.INVALID_REQUEST + METHOD_NOT_FOUND = mcpgo.METHOD_NOT_FOUND + INVALID_PARAMS = mcpgo.INVALID_PARAMS + INTERNAL_ERROR = mcpgo.INTERNAL_ERROR + RESOURCE_NOT_FOUND = mcpgo.RESOURCE_NOT_FOUND +) + +// MCPMethod is the type of MCP method-name constants. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type MCPMethod = mcpgo.MCPMethod + +// ---------------------------------------------------------------------------- +// Errors +// ---------------------------------------------------------------------------- + +// ErrMethodNotFound indicates the requested method does not exist. +var ErrMethodNotFound = mcpgo.ErrMethodNotFound + +// ---------------------------------------------------------------------------- +// Core protocol types +// ---------------------------------------------------------------------------- + +type ( + // Implementation describes the name and version of an MCP implementation. + Implementation = mcpgo.Implementation + // ClientCapabilities represents capabilities a client may support. + ClientCapabilities = mcpgo.ClientCapabilities + // ServerCapabilities represents capabilities a server may support. + ServerCapabilities = mcpgo.ServerCapabilities + + // Meta is metadata attached to a request's parameters or a result. + Meta = mcpgo.Meta + // Result is the base type embedded in protocol result messages. + Result = mcpgo.Result + // Request is the base type embedded in protocol request messages. + Request = mcpgo.Request + // RequestParams is the base params type carrying _meta. + RequestParams = mcpgo.RequestParams + + // JSONRPCNotification is a JSON-RPC notification (no response expected). + JSONRPCNotification = mcpgo.JSONRPCNotification + // JSONRPCMessage is any JSON-RPC request/response/notification/error. + JSONRPCMessage = mcpgo.JSONRPCMessage + // JSONRPCResponse is a successful JSON-RPC response. + JSONRPCResponse = mcpgo.JSONRPCResponse + // JSONRPCError is a JSON-RPC error response. + JSONRPCError = mcpgo.JSONRPCError + // JSONRPCErrorDetails carries the code/message/data of a JSON-RPC error. + JSONRPCErrorDetails = mcpgo.JSONRPCErrorDetails + // RequestId is a JSON-RPC request identifier. + RequestId = mcpgo.RequestId //nolint:revive // name intentionally matches mcp-go for drop-in compatibility. + + // Notification is the base of a JSON-RPC notification (method + params). + Notification = mcpgo.Notification + // NotificationParams carries a notification's params. + NotificationParams = mcpgo.NotificationParams + + // EmptyResult is an empty MCP result (e.g. the response to ping). + EmptyResult = mcpgo.EmptyResult + // PingRequest is a ping request. + PingRequest = mcpgo.PingRequest +) + +// JSON-RPC message constructors. +var ( + // NewRequestId wraps a raw id value in a RequestId. + NewRequestId = mcpgo.NewRequestId //nolint:revive // name intentionally matches mcp-go for drop-in compatibility. + // NewJSONRPCErrorDetails builds a JSONRPCErrorDetails value. + NewJSONRPCErrorDetails = mcpgo.NewJSONRPCErrorDetails + // NewJSONRPCResultResponse builds a successful JSONRPCResponse. + NewJSONRPCResultResponse = mcpgo.NewJSONRPCResultResponse +) + +// ---------------------------------------------------------------------------- +// Initialization +// ---------------------------------------------------------------------------- + +type ( + // InitializeRequest is sent by the client to begin initialization. + InitializeRequest = mcpgo.InitializeRequest + // InitializeParams are the params of an initialize request. + InitializeParams = mcpgo.InitializeParams + // InitializeResult is the server's response to initialize. + InitializeResult = mcpgo.InitializeResult +) + +// ---------------------------------------------------------------------------- +// Content +// ---------------------------------------------------------------------------- + +type ( + // Content is a polymorphic content element (text/image/audio/resource/...). + Content = mcpgo.Content + // Annotated is the base carrying optional Annotations. + Annotated = mcpgo.Annotated + // Annotations describe audience/priority/lastModified for content. + Annotations = mcpgo.Annotations + // Role is the role of a sampling/prompt message ("user"/"assistant"). + Role = mcpgo.Role + + // TextContent is text content. + TextContent = mcpgo.TextContent + // ImageContent is base64-encoded image content. + ImageContent = mcpgo.ImageContent + // AudioContent is base64-encoded audio content. + AudioContent = mcpgo.AudioContent + // EmbeddedResource embeds a resource in content. + EmbeddedResource = mcpgo.EmbeddedResource + // ResourceLink links to a resource. + ResourceLink = mcpgo.ResourceLink + + // ResourceContents is the contents of a resource (text or blob). + ResourceContents = mcpgo.ResourceContents + // TextResourceContents is text resource contents. + TextResourceContents = mcpgo.TextResourceContents + // BlobResourceContents is base64 blob resource contents. + BlobResourceContents = mcpgo.BlobResourceContents +) + +// Message roles. +const ( + RoleUser = mcpgo.RoleUser + RoleAssistant = mcpgo.RoleAssistant +) + +// Content constructors. +var ( + NewTextContent = mcpgo.NewTextContent + NewImageContent = mcpgo.NewImageContent + NewAudioContent = mcpgo.NewAudioContent + NewEmbeddedResource = mcpgo.NewEmbeddedResource + NewResourceLink = mcpgo.NewResourceLink +) + +// Content type-assertion helpers. +var ( + AsTextContent = mcpgo.AsTextContent + AsImageContent = mcpgo.AsImageContent + AsAudioContent = mcpgo.AsAudioContent + AsEmbeddedResource = mcpgo.AsEmbeddedResource + AsTextResourceContents = mcpgo.AsTextResourceContents + AsBlobResourceContents = mcpgo.AsBlobResourceContents +) + +// GetTextFromContent extracts text from a content value, if any. +var GetTextFromContent = mcpgo.GetTextFromContent + +// UnmarshalContent decodes a single JSON content object into the concrete +// mcp.Content implementation (TextContent, ImageContent, ...). It is used by the +// client shim to populate the Content interface fields of PromptMessage, which +// mcp-go cannot unmarshal generically. +var UnmarshalContent = mcpgo.UnmarshalContent + +// NewMetaFromMap builds a *Meta from a raw map. +var NewMetaFromMap = mcpgo.NewMetaFromMap + +// ---------------------------------------------------------------------------- +// Tools +// ---------------------------------------------------------------------------- + +type ( + // Tool describes a callable tool. + Tool = mcpgo.Tool + // ToolInputSchema is a tool's input JSON schema. + ToolInputSchema = mcpgo.ToolInputSchema + // ToolOutputSchema is a tool's output JSON schema. + ToolOutputSchema = mcpgo.ToolOutputSchema + // ToolAnnotation carries tool behavior hints. + ToolAnnotation = mcpgo.ToolAnnotation + + // ToolOption configures a Tool built via NewTool. + ToolOption = mcpgo.ToolOption + // PropertyOption configures a schema property. + PropertyOption = mcpgo.PropertyOption + + // CallToolRequest is a request to invoke a tool. + CallToolRequest = mcpgo.CallToolRequest + // CallToolParams are the params of a tool call. + CallToolParams = mcpgo.CallToolParams + // CallToolResult is the result of a tool call. + CallToolResult = mcpgo.CallToolResult + + // ListToolsRequest requests the tool list. + ListToolsRequest = mcpgo.ListToolsRequest + // ListToolsResult is the tool list response. + ListToolsResult = mcpgo.ListToolsResult +) + +// Tool builders and result constructors. +var ( + NewTool = mcpgo.NewTool + NewToolWithRawSchema = mcpgo.NewToolWithRawSchema + WithDescription = mcpgo.WithDescription + WithString = mcpgo.WithString + Required = mcpgo.Required + Description = mcpgo.Description + + NewToolResultText = mcpgo.NewToolResultText + NewToolResultError = mcpgo.NewToolResultError + NewToolResultStructuredOnly = mcpgo.NewToolResultStructuredOnly +) + +// ---------------------------------------------------------------------------- +// Resources +// ---------------------------------------------------------------------------- + +type ( + // Resource is a readable resource. + Resource = mcpgo.Resource + // ResourceTemplate is a URI-templated resource description. + ResourceTemplate = mcpgo.ResourceTemplate + + // ReadResourceRequest reads a resource by URI. + ReadResourceRequest = mcpgo.ReadResourceRequest + // ReadResourceParams are the params of a read. + ReadResourceParams = mcpgo.ReadResourceParams + // ReadResourceResult is the read response. + ReadResourceResult = mcpgo.ReadResourceResult + + // ListResourcesRequest requests the resource list. + ListResourcesRequest = mcpgo.ListResourcesRequest + // ListResourcesResult is the resource list response. + ListResourcesResult = mcpgo.ListResourcesResult + // ListResourceTemplatesRequest requests the resource-template list. + ListResourceTemplatesRequest = mcpgo.ListResourceTemplatesRequest + // ListResourceTemplatesResult is the resource-template list response. + ListResourceTemplatesResult = mcpgo.ListResourceTemplatesResult +) + +// ---------------------------------------------------------------------------- +// Prompts +// ---------------------------------------------------------------------------- + +type ( + // Prompt describes a prompt template. + Prompt = mcpgo.Prompt + // PromptArgument is a prompt template argument. + PromptArgument = mcpgo.PromptArgument + // PromptMessage is a message within a prompt result. + PromptMessage = mcpgo.PromptMessage + // PromptOption configures a Prompt built via NewPrompt. + PromptOption = mcpgo.PromptOption + + // GetPromptRequest gets a prompt. + GetPromptRequest = mcpgo.GetPromptRequest + // GetPromptParams are the params of a get. + GetPromptParams = mcpgo.GetPromptParams + // GetPromptResult is the get response. + GetPromptResult = mcpgo.GetPromptResult + + // ListPromptsRequest requests the prompt list. + ListPromptsRequest = mcpgo.ListPromptsRequest + // ListPromptsResult is the prompt list response. + ListPromptsResult = mcpgo.ListPromptsResult +) + +// Prompt builders. +var ( + NewPrompt = mcpgo.NewPrompt + WithPromptDescription = mcpgo.WithPromptDescription +) + +// ---------------------------------------------------------------------------- +// Elicitation +// ---------------------------------------------------------------------------- + +type ( + // ElicitationRequest is a server->client elicitation request. + ElicitationRequest = mcpgo.ElicitationRequest + // ElicitationParams are the params of an elicitation. + ElicitationParams = mcpgo.ElicitationParams + // ElicitationResult is the result of an elicitation. + ElicitationResult = mcpgo.ElicitationResult + // ElicitationResponse is the user's response payload. + ElicitationResponse = mcpgo.ElicitationResponse + // ElicitationResponseAction indicates accept/decline/cancel. + ElicitationResponseAction = mcpgo.ElicitationResponseAction +) + +// Elicitation response actions. +const ( + ElicitationResponseActionAccept = mcpgo.ElicitationResponseActionAccept + ElicitationResponseActionDecline = mcpgo.ElicitationResponseActionDecline + ElicitationResponseActionCancel = mcpgo.ElicitationResponseActionCancel +) diff --git a/mcpcompat/mcp/doc.go b/mcpcompat/mcp/doc.go new file mode 100644 index 0000000..50db46f --- /dev/null +++ b/mcpcompat/mcp/doc.go @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package mcp is a drop-in compatibility shim for github.com/mark3labs/mcp-go/mcp. +// +// It exists so that ToolHive and its sibling projects can migrate off mcp-go and +// onto the official Model Context Protocol Go SDK +// (github.com/modelcontextprotocol/go-sdk) by swapping imports rather than +// rewriting call sites: replace +// +// github.com/mark3labs/mcp-go/mcp +// +// with +// +// github.com/stacklok/toolhive-core/mcpcompat/mcp +// +// while keeping the existing import alias. The companion packages +// mcpcompat/server, mcpcompat/client and mcpcompat/client/transport reimplement +// mcp-go's protocol machinery on top of the official SDK; this package supplies +// the data types those APIs exchange. +// +// # Migration strategy +// +// The types and helpers below are currently re-exported from mcp-go via type +// aliases and value assignments. This guarantees byte-for-byte wire and source +// compatibility during the transition: the aliased symbols ARE mcp-go's, so +// existing struct-literal construction, field access and custom JSON marshaling +// behave identically. +// +// This file is the single chokepoint that still references mcp-go. The end goal +// is to remove the mcp-go dependency from the tree entirely. To get there, each +// alias here is replaced by a standalone definition (copied from the mcp-go +// source, same JSON tags and marshaling). Because consumers only ever see the +// symbols in this package, that swap is transparent to them. The wire-format +// golden tests in this package pin the exact JSON shape of every re-exported +// type and are what make the alias-to-standalone conversion safe. +// +// # Scope +// +// Only the subset of mcp-go's mcp package that ToolHive actually uses is +// re-exported. Unused surface (sampling, roots, completion, tasks, logging-level +// control, the fluent schema builder beyond WithDescription/WithString/Required, +// etc.) is intentionally omitted and can be added on demand. +// +// Stability: Alpha. +package mcp diff --git a/mcpcompat/mcp/mcp_test.go b/mcpcompat/mcp/mcp_test.go new file mode 100644 index 0000000..264e2fd --- /dev/null +++ b/mcpcompat/mcp/mcp_test.go @@ -0,0 +1,241 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +const queryProp = "query" + +// TestWireFormat_GoldenJSON pins the exact JSON wire shape of the re-exported +// types and constructors. These goldens are the equivalence spec: when the +// aliases in alias.go are later replaced by standalone definitions to drop the +// mcp-go dependency, these tests must continue to pass unchanged. +func TestWireFormat_GoldenJSON(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + value any + wantJSON string + }{ + { + name: "text content", + value: mcp.NewTextContent("hi"), + wantJSON: `{"type":"text","text":"hi"}`, + }, + { + name: "image content", + value: mcp.NewImageContent("ZGF0YQ==", "image/png"), + wantJSON: `{"type":"image","data":"ZGF0YQ==","mimeType":"image/png"}`, + }, + { + name: "audio content", + value: mcp.NewAudioContent("YXVkaW8=", "audio/wav"), + wantJSON: `{"type":"audio","data":"YXVkaW8=","mimeType":"audio/wav"}`, + }, + { + name: "tool result text", + value: mcp.NewToolResultText("hi"), + wantJSON: `{"content":[{"type":"text","text":"hi"}]}`, + }, + { + name: "tool result error", + value: mcp.NewToolResultError("boom"), + wantJSON: `{"content":[{"type":"text","text":"boom"}],"isError":true}`, + }, + { + name: "implementation", + value: mcp.Implementation{Name: "client", Version: "1.0.0"}, + wantJSON: `{"name":"client","version":"1.0.0"}`, + }, + { + name: "call tool params", + value: mcp.CallToolParams{Name: "t", Arguments: map[string]any{"a": "b"}}, + wantJSON: `{"name":"t","arguments":{"a":"b"}}`, + }, + { + name: "resource link content", + value: mcp.NewResourceLink( + "file:///x", "x", "desc", "text/plain"), + wantJSON: `{"type":"resource_link","uri":"file:///x","name":"x","description":"desc","mimeType":"text/plain"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := json.Marshal(tt.value) + require.NoError(t, err) + assert.JSONEq(t, tt.wantJSON, string(got)) + }) + } +} + +// TestTool_StructLiteralWire mirrors how ToolHive builds tools in production: +// a struct literal (not the fluent builder) with an explicit InputSchema and +// annotations. It verifies the fields survive marshaling to the MCP wire shape. +func TestTool_StructLiteralWire(t *testing.T) { + t.Parallel() + + tool := mcp.Tool{ + Name: "search", + Description: "search the index", + InputSchema: mcp.ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + queryProp: map[string]any{"type": "string"}, + }, + Required: []string{queryProp}, + }, + Annotations: mcp.ToolAnnotation{}, + } + + raw, err := json.Marshal(tool) + require.NoError(t, err) + + var got map[string]any + require.NoError(t, json.Unmarshal(raw, &got)) + + assert.Equal(t, "search", got["name"]) + assert.Equal(t, "search the index", got["description"]) + + schema, ok := got["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object") + assert.Equal(t, "object", schema["type"]) + assert.Contains(t, schema, "properties") + assert.Equal(t, []any{"query"}, schema["required"]) +} + +// TestCallToolRequest_Accessors mirrors how ToolHive handlers read arguments: +// via Params.Name and the argument accessor helpers. +func TestCallToolRequest_Accessors(t *testing.T) { + t.Parallel() + + req := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "greet", + Arguments: map[string]any{ + "name": "ada", + "count": float64(3), + "loud": true, + }, + }, + } + + assert.Equal(t, "greet", req.Params.Name) + + args := req.GetArguments() + assert.Equal(t, "ada", args["name"]) + + assert.Equal(t, "ada", req.GetString("name", "default")) + assert.Equal(t, "default", req.GetString("missing", "default")) + assert.Equal(t, 3, req.GetInt("count", 0)) + assert.True(t, req.GetBool("loud", false)) + + s, err := req.RequireString("name") + require.NoError(t, err) + assert.Equal(t, "ada", s) + + _, err = req.RequireString("missing") + assert.Error(t, err) +} + +// TestContent_InterfaceAndAsHelpers verifies the polymorphic Content interface +// and the As* type-assertion helpers used by ToolHive's content conversion. +func TestContent_InterfaceAndAsHelpers(t *testing.T) { + t.Parallel() + + var content []mcp.Content + content = append(content, mcp.NewTextContent("hello")) + content = append(content, mcp.NewImageContent("ZGF0YQ==", "image/png")) + + txt, ok := mcp.AsTextContent(content[0]) + require.True(t, ok) + assert.Equal(t, "hello", txt.Text) + + _, ok = mcp.AsImageContent(content[0]) + assert.False(t, ok, "text content is not image content") + + img, ok := mcp.AsImageContent(content[1]) + require.True(t, ok) + assert.Equal(t, "image/png", img.MIMEType) + + assert.Equal(t, "hello", mcp.GetTextFromContent(content[0])) +} + +// TestResourceContents_AsHelpers verifies the resource-contents helpers used +// when converting resource read results. +func TestResourceContents_AsHelpers(t *testing.T) { + t.Parallel() + + var contents []mcp.ResourceContents + contents = append(contents, mcp.TextResourceContents{URI: "file:///a", Text: "body"}) + contents = append(contents, mcp.BlobResourceContents{URI: "file:///b", Blob: "ZGF0YQ=="}) + + txt, ok := mcp.AsTextResourceContents(contents[0]) + require.True(t, ok) + assert.Equal(t, "body", txt.Text) + + blob, ok := mcp.AsBlobResourceContents(contents[1]) + require.True(t, ok) + assert.Equal(t, "ZGF0YQ==", blob.Blob) +} + +// TestNewToolFluentBuilder exercises the fluent tool builder that ToolHive uses +// only in its test/helper code (fake backends). +func TestNewToolFluentBuilder(t *testing.T) { + t.Parallel() + + tool := mcp.NewTool("greet", + mcp.WithDescription("greets a person"), + mcp.WithString("name", mcp.Required()), + ) + + assert.Equal(t, "greet", tool.Name) + assert.Equal(t, "greets a person", tool.Description) + assert.Contains(t, tool.InputSchema.Properties, "name") + assert.Contains(t, tool.InputSchema.Required, "name") +} + +// TestProtocolConstants pins stable protocol-level values. These come straight +// from the MCP spec (JSON-RPC error codes, method names) and must be preserved +// verbatim by any future standalone reimplementation. +func TestProtocolConstants(t *testing.T) { + t.Parallel() + + assert.Equal(t, -32601, mcp.METHOD_NOT_FOUND) + assert.Equal(t, -32602, mcp.INVALID_PARAMS) + assert.Equal(t, -32603, mcp.INTERNAL_ERROR) + + assert.Equal(t, "tools/call", string(mcp.MethodToolsCall)) + assert.Equal(t, "tools/list", string(mcp.MethodToolsList)) + assert.Equal(t, "resources/read", string(mcp.MethodResourcesRead)) + assert.Equal(t, "prompts/get", string(mcp.MethodPromptsGet)) + assert.Equal(t, "initialize", string(mcp.MethodInitialize)) + + assert.NotEmpty(t, mcp.LATEST_PROTOCOL_VERSION) + + assert.Equal(t, "accept", string(mcp.ElicitationResponseActionAccept)) + assert.Equal(t, "decline", string(mcp.ElicitationResponseActionDecline)) + assert.Equal(t, "cancel", string(mcp.ElicitationResponseActionCancel)) +} + +// Handler signature guards. These package-level assignments will fail to compile +// if the re-exported request/result types drift from the exact function shapes +// ToolHive's server package registers. They document the contract the +// mcpcompat/server adapters must satisfy. +var ( + _ = func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, nil } + _ = func(_ context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { return nil, nil } + _ = func(_ context.Context, _ mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return nil, nil } +) diff --git a/mcpcompat/server/elicitation.go b/mcpcompat/server/elicitation.go new file mode 100644 index 0000000..f79e4ff --- /dev/null +++ b/mcpcompat/server/elicitation.go @@ -0,0 +1,41 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "fmt" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// RequestElicitation sends a server->client elicitation request on the session +// associated with ctx and returns the user's response. It mirrors mcp-go's +// (*MCPServer).RequestElicitation. +func (*MCPServer) RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + cs, ok := ClientSessionFromContext(ctx).(*clientSession) + if !ok || cs == nil { + return nil, fmt.Errorf("no active session in context for elicitation") + } + ss := cs.goSession.Load() + if ss == nil { + return nil, fmt.Errorf("no server session available for elicitation") + } + + params := &gosdk.ElicitParams{} + if err := jsonConvert(request.Params, params); err != nil { + return nil, fmt.Errorf("converting elicitation params: %w", err) + } + res, err := ss.Elicit(ctx, params) + if err != nil { + return nil, err + } + out := &mcp.ElicitationResult{} + if err := jsonConvert(res, out); err != nil { + return nil, fmt.Errorf("converting elicitation result: %w", err) + } + return out, nil +} diff --git a/mcpcompat/server/hooks.go b/mcpcompat/server/hooks.go new file mode 100644 index 0000000..616b070 --- /dev/null +++ b/mcpcompat/server/hooks.go @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// OnRegisterSessionHookFunc runs when a session registers. +type OnRegisterSessionHookFunc func(ctx context.Context, session ClientSession) + +// OnBeforeListToolsFunc runs before a tools/list request is handled. +type OnBeforeListToolsFunc func(ctx context.Context, id any, message *mcp.ListToolsRequest) + +// OnBeforeCallToolFunc runs before a tools/call request is handled. +type OnBeforeCallToolFunc func(ctx context.Context, id any, message *mcp.CallToolRequest) + +// Hooks holds lifecycle callbacks. It mirrors the subset of mcp-go's +// server.Hooks that ToolHive registers. +type Hooks struct { + registerSessionHooks []OnRegisterSessionHookFunc + listToolsHooks []OnBeforeListToolsFunc + callToolHooks []OnBeforeCallToolFunc +} + +// AddOnRegisterSession registers a session-registration hook. +func (c *Hooks) AddOnRegisterSession(hook OnRegisterSessionHookFunc) { + c.registerSessionHooks = append(c.registerSessionHooks, hook) +} + +// AddBeforeListTools registers a before-tools/list hook. +func (c *Hooks) AddBeforeListTools(hook OnBeforeListToolsFunc) { + c.listToolsHooks = append(c.listToolsHooks, hook) +} + +// AddBeforeCallTool registers a before-tools/call hook. +func (c *Hooks) AddBeforeCallTool(hook OnBeforeCallToolFunc) { + c.callToolHooks = append(c.callToolHooks, hook) +} + +func (c *Hooks) registerSession(ctx context.Context, session ClientSession) { + for _, h := range c.registerSessionHooks { + h(ctx, session) + } +} + +func (c *Hooks) beforeCallTool(ctx context.Context, id any, message *mcp.CallToolRequest) { + for _, h := range c.callToolHooks { + h(ctx, id, message) + } +} + +func (c *Hooks) beforeListTools(ctx context.Context, id any, message *mcp.ListToolsRequest) { + for _, h := range c.listToolsHooks { + h(ctx, id, message) + } +} diff --git a/mcpcompat/server/notifications.go b/mcpcompat/server/notifications.go new file mode 100644 index 0000000..a5921c0 --- /dev/null +++ b/mcpcompat/server/notifications.go @@ -0,0 +1,84 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// SendNotificationToAllClients broadcasts a server-initiated notification with +// the given method and params to every currently connected client session. It +// mirrors mcp-go's server.MCPServer.SendNotificationToAllClients, which +// ToolHive's stdio bridge uses to relay upstream notifications downstream. +// +// go-sdk backing and limitation: the go-sdk does not expose a public API to +// send an arbitrary (method, params) notification on a ServerSession; only +// typed senders are exported. This method therefore maps the well-known MCP +// notification methods onto the go-sdk's typed senders: +// +// - notifications/progress -> ServerSession.NotifyProgress +// - notifications/message -> ServerSession.Log (delivered only once the +// client has set a logging level, per the go-sdk/spec behavior) +// +// The list-changed notifications (tools/prompts/resources) are emitted +// automatically by the go-sdk server when its registered feature set changes, +// so they cannot be re-broadcast through a public API here; they, and any other +// unrecognized method, are dropped (logged at debug level). This is the one +// behavioral gap versus mcp-go's raw channel-based broadcast and is documented +// rather than silently ignored. +func (s *MCPServer) SendNotificationToAllClients(method string, params map[string]any) { + ctx := context.Background() + s.sessions.Range(func(_, v any) bool { + cs, ok := v.(*clientSession) + if !ok { + return true + } + if ss := cs.goSession.Load(); ss != nil { + s.sendOneNotification(ctx, ss, method, params) + } + return true + }) +} + +// sendOneNotification dispatches a single notification to one go-sdk session, +// translating the method+params into the matching typed go-sdk sender. Errors +// are intentionally ignored to match mcp-go's best-effort broadcast semantics. +func (s *MCPServer) sendOneNotification( + ctx context.Context, ss *gosdk.ServerSession, method string, params map[string]any, +) { + switch method { + case "notifications/progress": + var p gosdk.ProgressNotificationParams + if err := jsonConvert(params, &p); err != nil { + s.logNotifyErr(method, err) + return + } + _ = ss.NotifyProgress(ctx, &p) + case "notifications/message": + var p gosdk.LoggingMessageParams + if err := jsonConvert(params, &p); err != nil { + s.logNotifyErr(method, err) + return + } + _ = ss.Log(ctx, &p) + default: + // See the doc comment on SendNotificationToAllClients: go-sdk offers no + // public generic notification sender, so list-changed and other methods + // cannot be forwarded and are dropped. + if s.logger != nil { + s.logger.Debug("SendNotificationToAllClients: dropping unsupported notification method", + "method", method) + } + } +} + +// logNotifyErr logs a notification-conversion error if a logger is configured. +func (s *MCPServer) logNotifyErr(method string, err error) { + if s.logger != nil { + s.logger.Warn("SendNotificationToAllClients: failed to convert params", + "method", method, "error", err) + } +} diff --git a/mcpcompat/server/notifications_test.go b/mcpcompat/server/notifications_test.go new file mode 100644 index 0000000..9b62513 --- /dev/null +++ b/mcpcompat/server/notifications_test.go @@ -0,0 +1,82 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "context" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// TestSendNotificationToAllClients_NoSessions verifies the broadcast is a safe +// no-op when there are no connected clients. +func TestSendNotificationToAllClients_NoSessions(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("notify-server", "1.0.0") + assert.NotPanics(t, func() { + srv.SendNotificationToAllClients("notifications/message", map[string]any{"data": "hi"}) + }) +} + +// TestSendNotificationToAllClients_Broadcast connects a live client and then +// broadcasts several notification methods, exercising the real per-session +// dispatch path (progress/message/list-changed/unknown) without panicking. +func TestSendNotificationToAllClients_Broadcast(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := server.NewMCPServer("notify-server", "1.0.0", + server.WithToolCapabilities(true), + server.WithLogging(), + ) + srv.AddTool( + mcp.NewTool("noop", mcp.WithDescription("noop")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("ok"), nil + }, + ) + + httpSrv := server.NewStreamableHTTPServer(srv) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "c", Version: "1"}, + }, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + // Each of these maps onto a different branch of the dispatcher; none should + // panic even though some are dropped (list-changed, unknown). + assert.NotPanics(t, func() { + srv.SendNotificationToAllClients("notifications/progress", + map[string]any{"progressToken": "t", "progress": 0.5}) + srv.SendNotificationToAllClients("notifications/message", + map[string]any{"level": "info", "data": "hello"}) + srv.SendNotificationToAllClients("notifications/tools/list_changed", nil) + srv.SendNotificationToAllClients("some/unknown/method", map[string]any{"x": 1}) + }) +} + +// TestSSEHandlers verifies SSEServer exposes non-nil SSE and message handlers. +func TestSSEHandlers(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("sse-server", "1.0.0") + sse := server.NewSSEServer(srv) + assert.NotNil(t, sse.SSEHandler()) + assert.NotNil(t, sse.MessageHandler()) +} diff --git a/mcpcompat/server/rehydration_test.go b/mcpcompat/server/rehydration_test.go new file mode 100644 index 0000000..bdb739f --- /dev/null +++ b/mcpcompat/server/rehydration_test.go @@ -0,0 +1,412 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// sharedSessionManager is a test double for ToolHive's Redis-backed +// SessionIdManager: an in-memory store that multiple StreamableHTTPServer +// instances share, standing in for cross-replica shared state. +type sharedSessionManager struct { + mu sync.Mutex + valid map[string]bool // sessionID -> terminated? + terminated map[string]bool +} + +func newSharedSessionManager() *sharedSessionManager { + return &sharedSessionManager{valid: map[string]bool{}, terminated: map[string]bool{}} +} + +func (m *sharedSessionManager) Generate() string { + m.mu.Lock() + defer m.mu.Unlock() + id := uuid.NewString() + m.valid[id] = true + return id +} + +func (m *sharedSessionManager) Validate(sessionID string) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.terminated[sessionID] { + return true, nil + } + if !m.valid[sessionID] { + return false, fmt.Errorf("session %q not found", sessionID) + } + return false, nil +} + +func (m *sharedSessionManager) Terminate(sessionID string) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.terminated[sessionID] = true + return false, nil +} + +// greetTool registers a simple "greet" tool on the server. +func addGreetTool(s *server.MCPServer) { + s.AddTool(mcp.NewTool("greet", mcp.WithDescription("greets")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("hello"), nil + }) +} + +// rpcResult holds the parsed JSON-RPC envelope for the response to a request. +type rpcResult struct { + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` +} + +// postRPC sends a single JSON-RPC message to url and returns the HTTP response +// so the caller can inspect status codes or stream SSE. +func postRPC(ctx context.Context, t *testing.T, url, sessionID string, body string) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("MCP-Protocol-Version", "2025-06-18") + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// readFirstResult reads the first JSON-RPC result/error message from an HTTP +// response, transparently handling both application/json and text/event-stream +// bodies. +func readFirstResult(t *testing.T, resp *http.Response) rpcResult { + t.Helper() + defer func() { _ = resp.Body.Close() }() + ct := resp.Header.Get("Content-Type") + if strings.HasPrefix(ct, "application/json") { + var r rpcResult + require.NoError(t, json.NewDecoder(resp.Body).Decode(&r)) + return r + } + // SSE: scan "data:" lines for the first message carrying a result or error. + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for sc.Scan() { + line := sc.Text() + data, ok := strings.CutPrefix(line, "data:") + if !ok { + continue + } + var r rpcResult + if err := json.Unmarshal([]byte(strings.TrimSpace(data)), &r); err != nil { + continue + } + if len(r.Result) > 0 || r.Error != nil { + return r + } + } + t.Fatalf("no JSON-RPC result found in SSE stream") + return rpcResult{} +} + +// initSession initializes a session against url using raw JSON-RPC and returns +// the assigned session ID. +func initSession(t *testing.T, url string) string { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}` + resp := postRPC(ctx, t, url, "", body) + sid := resp.Header.Get("Mcp-Session-Id") + _ = readFirstResult(t, resp) + require.NotEmpty(t, sid, "initialize must assign a session ID") + // Send notifications/initialized to complete the handshake. + notif := `{"jsonrpc":"2.0","method":"notifications/initialized"}` + nresp := postRPC(ctx, t, url, sid, notif) + _ = nresp.Body.Close() + return sid +} + +// listToolNames issues tools/list against url (optionally with a session ID) and +// returns the tool names. +func listToolNames(t *testing.T, url, sessionID string) []string { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + body := `{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}` + resp := postRPC(ctx, t, url, sessionID, body) + require.Equal(t, http.StatusOK, resp.StatusCode) + r := readFirstResult(t, resp) + require.Nil(t, r.Error, "tools/list should not error") + var res struct { + Tools []struct { + Name string `json:"name"` + } `json:"tools"` + } + require.NoError(t, json.Unmarshal(r.Result, &res)) + names := make([]string, 0, len(res.Tools)) + for _, tl := range res.Tools { + names = append(names, tl.Name) + } + return names +} + +// TestCrossReplicaSessionRouting verifies that a session initialized on replica +// A is accepted on replica B (which shares the SessionIdManager) and returns the +// same tools — the core production fix for Redis-shared cross-replica sessions. +func TestCrossReplicaSessionRouting(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + + // Replica A: registers greet globally. + streamA := server.NewMCPServer("A", "1.0.0") + addGreetTool(streamA) + sA := server.NewStreamableHTTPServer(streamA, server.WithSessionIdManager(mgr)) + tsA := httptest.NewServer(sA) + defer tsA.Close() + + // Replica B: same global tools (separate process in production). + streamB := server.NewMCPServer("B", "1.0.0") + addGreetTool(streamB) + sB := server.NewStreamableHTTPServer(streamB, server.WithSessionIdManager(mgr)) + tsB := httptest.NewServer(sB) + defer tsB.Close() + + // Initialize on A. + sid := initSession(t, tsA.URL) + toolsA := listToolNames(t, tsA.URL, sid) + require.Contains(t, toolsA, "greet") + + // Same session ID, replica B — must be accepted via rehydration. + toolsB := listToolNames(t, tsB.URL, sid) + assert.ElementsMatch(t, toolsA, toolsB, "replica B must return the same tools for the shared session") +} + +// TestCrossReplicaLazyToolInjection mirrors ToolHive's real cross-replica +// mechanism: replica B has NO globally-registered tools; instead a before-list +// hook lazily injects per-session tools when the session has none (the +// rehydrated case, matching Server.lazyInjectSessionTools). It proves the +// before-hooks fire on a rehydrated session with a usable ClientSession in +// context and that the injected overlay is served by tools/list. +func TestCrossReplicaLazyToolInjection(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + + // Replica A initializes the session. + streamA := server.NewMCPServer("A", "1.0.0") + sA := server.NewStreamableHTTPServer(streamA, server.WithSessionIdManager(mgr)) + tsA := httptest.NewServer(sA) + defer tsA.Close() + + // Replica B: no global tools; a before-list hook injects per-session tools. + hooks := &server.Hooks{} + var injected bool + hooks.AddBeforeListTools(func(ctx context.Context, _ any, _ *mcp.ListToolsRequest) { + sess := server.ClientSessionFromContext(ctx) + require.NotNil(t, sess, "before-list hook must see a ClientSession for the rehydrated session") + swt, ok := sess.(server.SessionWithTools) + require.True(t, ok, "session must support per-session tools") + if len(swt.GetSessionTools()) > 0 { + return + } + injected = true + swt.SetSessionTools(map[string]server.ServerTool{ + "lazy": { + Tool: mcp.NewTool("lazy", mcp.WithDescription("lazily injected")), + Handler: func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("ok"), nil + }, + }, + }) + }) + streamB := server.NewMCPServer("B", "1.0.0", server.WithHooks(hooks)) + sB := server.NewStreamableHTTPServer(streamB, server.WithSessionIdManager(mgr)) + tsB := httptest.NewServer(sB) + defer tsB.Close() + + sid := initSession(t, tsA.URL) + + toolsB := listToolNames(t, tsB.URL, sid) + assert.True(t, injected, "before-list hook should have injected tools on the rehydrated replica") + assert.Contains(t, toolsB, "lazy", "rehydrated session must serve the lazily-injected per-session tool") +} + +// TestRehydrationRejectsUnknownAndTerminated verifies the 404 paths. +func TestRehydrationRejectsUnknownAndTerminated(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + mcpSrv := server.NewMCPServer("B", "1.0.0") + addGreetTool(mcpSrv) + s := server.NewStreamableHTTPServer(mcpSrv, server.WithSessionIdManager(mgr)) + ts := httptest.NewServer(s) + defer ts.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + body := `{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}` + + // Unknown session ID -> 404. + resp := postRPC(ctx, t, ts.URL, "does-not-exist", body) + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + _ = resp.Body.Close() + + // Terminated session ID -> 404. + sid := mgr.Generate() + _, _ = mgr.Terminate(sid) + resp2 := postRPC(ctx, t, ts.URL, sid, body) + assert.Equal(t, http.StatusNotFound, resp2.StatusCode) + _ = resp2.Body.Close() +} + +// TestRehydratedSessionEvictedAfterTermination verifies lazy eviction across +// replicas: once replica B has rehydrated AND cached a session, terminating it +// in the shared store must cause B to reject the NEXT request rather than serve +// the cached session. This is the cross-replica termination path — the handler +// must validate every request, not only on cache miss. +func TestRehydratedSessionEvictedAfterTermination(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + + streamA := server.NewMCPServer("A", "1.0.0") + addGreetTool(streamA) + sA := server.NewStreamableHTTPServer(streamA, server.WithSessionIdManager(mgr)) + tsA := httptest.NewServer(sA) + defer tsA.Close() + + streamB := server.NewMCPServer("B", "1.0.0") + addGreetTool(streamB) + sB := server.NewStreamableHTTPServer(streamB, server.WithSessionIdManager(mgr)) + tsB := httptest.NewServer(sB) + defer tsB.Close() + + sid := initSession(t, tsA.URL) + + // First request on B rehydrates and caches the session. + require.Contains(t, listToolNames(t, tsB.URL, sid), "greet") + + // Terminate the session in the shared store (as a DELETE on replica A does). + _, _ = mgr.Terminate(sid) + + // The next request on B must be rejected via lazy eviction, not served from + // the cached rehydrated session. + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + resp := postRPC(ctx, t, tsB.URL, sid, `{"jsonrpc":"2.0","id":9,"method":"tools/list","params":{}}`) + assert.Equal(t, http.StatusNotFound, resp.StatusCode, + "replica B must reject a session terminated in the shared store") + _ = resp.Body.Close() +} + +// TestRehydratedSessionElicitation proves a rehydrated session is a full, +// stateful session (NOT stateless): a tool handler on the rehydrating replica +// performs a server->client elicitation, and the client responds over the same +// session, completing the tool call. +func TestRehydratedSessionElicitation(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + + // Replica A only initializes the session (client uses elicitation cap). + streamA := server.NewMCPServer("A", "1.0.0") + sA := server.NewStreamableHTTPServer(streamA, server.WithSessionIdManager(mgr)) + tsA := httptest.NewServer(sA) + defer tsA.Close() + + // Replica B has a tool that elicits input from the client. + streamB := server.NewMCPServer("B", "1.0.0") + var srvB = streamB + streamB.AddTool(mcp.NewTool("ask", mcp.WithDescription("asks the user")), + func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + res, err := srvB.RequestElicitation(ctx, mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{ + Message: "your name?", + RequestedSchema: map[string]any{"type": "object"}, + }, + }) + if err != nil { + return nil, err + } + return mcp.NewToolResultText("action=" + string(res.Action)), nil + }) + sB := server.NewStreamableHTTPServer(streamB, server.WithSessionIdManager(mgr)) + tsB := httptest.NewServer(sB) + defer tsB.Close() + + // Initialize on A (declares elicitation capability at the wire level via the + // rehydrated seed on B; the session id is shared). + sid := initSession(t, tsA.URL) + + // Fire tools/call ask on B. This POST's SSE stream will carry a server->client + // elicitation request; we respond on a second POST with the same session id. + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + callBody := `{"jsonrpc":"2.0","id":42,"method":"tools/call","params":{"name":"ask","arguments":{}}}` + resp := postRPC(ctx, t, tsB.URL, sid, callBody) + defer func() { _ = resp.Body.Close() }() + require.Equal(t, http.StatusOK, resp.StatusCode) + + // Read the SSE stream: when we see the elicitation request, answer it; then + // expect the tool result. + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 0, 64*1024), 1024*1024) + var gotToolResult bool + for sc.Scan() { + line := sc.Text() + data, ok := strings.CutPrefix(line, "data:") + if !ok { + continue + } + var msg rpcResult + if err := json.Unmarshal([]byte(strings.TrimSpace(data)), &msg); err != nil { + continue + } + if msg.Method == "elicitation/create" { + // Respond to the server->client request on a separate POST. + respBody := fmt.Sprintf( + `{"jsonrpc":"2.0","id":%s,"result":{"action":"accept","content":{}}}`, + string(msg.ID)) + ackCtx, ackCancel := context.WithTimeout(context.Background(), 10*time.Second) + ack := postRPC(ackCtx, t, tsB.URL, sid, respBody) + _ = ack.Body.Close() + ackCancel() + continue + } + if bytes.Equal(msg.ID, []byte("42")) && len(msg.Result) > 0 { + var res struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + } + require.NoError(t, json.Unmarshal(msg.Result, &res)) + require.NotEmpty(t, res.Content) + assert.Equal(t, "action=accept", res.Content[0].Text) + gotToolResult = true + break + } + } + assert.True(t, gotToolResult, "expected the elicited tool call to complete on the rehydrated session") +} diff --git a/mcpcompat/server/request_handler.go b/mcpcompat/server/request_handler.go new file mode 100644 index 0000000..51a67a3 --- /dev/null +++ b/mcpcompat/server/request_handler.go @@ -0,0 +1,235 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "fmt" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// WithContext stores the given ClientSession in the returned context so that +// handlers (and ClientSessionFromContext) can recover it. It mirrors mcp-go's +// server.MCPServer.WithContext, which ToolHive uses to associate a session with +// a request context (e.g. for per-session tool injection). +func (*MCPServer) WithContext(ctx context.Context, session ClientSession) context.Context { + return context.WithValue(ctx, sessionContextKey{}, session) +} + +// HandleMessage processes a single incoming JSON-RPC message and returns the +// appropriate JSON-RPC response (or nil for notifications and server-directed +// responses). It mirrors mcp-go's server.MCPServer.HandleMessage. +// +// go-sdk backing and limitation: the go-sdk drives its own JSON-RPC loop over a +// Transport and does not expose a public "handle one raw message" entrypoint. +// This shim therefore dispatches the message directly against the tools, +// resources, prompts and their handlers registered on this MCPServer — the same +// registration state the go-sdk server is built from — so behavior matches +// mcp-go for the methods ToolHive exercises: initialize, ping, tools/list, +// tools/call, resources/list, resources/templates/list, resources/read, +// prompts/list and prompts/get, plus notifications (which return nil). +// Capability-gated extras that ToolHive does not use over this path (logging +// setLevel, subscribe/unsubscribe, completion, tasks) return METHOD_NOT_FOUND. +func (s *MCPServer) HandleMessage(ctx context.Context, message json.RawMessage) mcp.JSONRPCMessage { + var base struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + ID any `json:"id,omitempty"` + Result any `json:"result,omitempty"` + } + if err := json.Unmarshal(message, &base); err != nil { + return errorResponse(nil, mcp.PARSE_ERROR, "Failed to parse message") + } + if base.JSONRPC != mcp.JSONRPC_VERSION { + return errorResponse(base.ID, mcp.INVALID_REQUEST, "Invalid JSON-RPC version") + } + // Notifications (no id) are handled and produce no response. + if base.ID == nil { + return nil + } + // A message carrying a result is a response to a server-initiated request. + if base.Result != nil { + return nil + } + return s.dispatch(ctx, base.Method, base.ID, message) +} + +//nolint:gocyclo // faithful 1:1 mirror of mcp-go's method dispatch switch. +func (s *MCPServer) dispatch(ctx context.Context, method string, id any, message json.RawMessage) mcp.JSONRPCMessage { + switch method { + case string(mcp.MethodInitialize): + return successResponse(id, s.handleInitialize()) + case string(mcp.MethodPing): + return successResponse(id, mcp.EmptyResult{}) + case string(mcp.MethodToolsList): + return successResponse(id, s.handleListTools()) + case string(mcp.MethodToolsCall): + return s.handleToolCall(ctx, id, message) + case string(mcp.MethodResourcesList): + return successResponse(id, s.handleListResources()) + case string(mcp.MethodResourcesTemplatesList): + return successResponse(id, s.handleListResourceTemplates()) + case string(mcp.MethodResourcesRead): + return s.handleReadResource(ctx, id, message) + case string(mcp.MethodPromptsList): + return successResponse(id, s.handleListPrompts()) + case string(mcp.MethodPromptsGet): + return s.handleGetPrompt(ctx, id, message) + default: + return errorResponse(id, mcp.METHOD_NOT_FOUND, fmt.Sprintf("Method %s not found", method)) + } +} + +func (s *MCPServer) handleInitialize() mcp.InitializeResult { + return mcp.InitializeResult{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ServerInfo: mcp.Implementation{Name: s.name, Version: s.version}, + Capabilities: s.buildCapabilities(), + } +} + +// buildCapabilities derives a ServerCapabilities value from the registered +// features and declared capability flags. +func (s *MCPServer) buildCapabilities() mcp.ServerCapabilities { + s.mu.RLock() + nTools := len(s.tools) + nResources := len(s.resources) + len(s.resourceTemplates) + nPrompts := len(s.prompts) + s.mu.RUnlock() + + const listChangedKey = "listChanged" + caps := map[string]any{} + if nTools > 0 || s.toolListChanged { + caps["tools"] = map[string]any{listChangedKey: s.toolListChanged} + } + if nResources > 0 || s.resourceListChanged || s.resourceSubscribe { + caps["resources"] = map[string]any{ + listChangedKey: s.resourceListChanged, + "subscribe": s.resourceSubscribe, + } + } + if nPrompts > 0 || s.promptListChanged { + caps["prompts"] = map[string]any{listChangedKey: s.promptListChanged} + } + if s.logging { + caps["logging"] = map[string]any{} + } + + var out mcp.ServerCapabilities + // jsonConvert cannot fail for these plain maps. + _ = jsonConvert(caps, &out) + return out +} + +func (s *MCPServer) handleListTools() mcp.ListToolsResult { + s.mu.RLock() + defer s.mu.RUnlock() + result := mcp.ListToolsResult{Tools: make([]mcp.Tool, 0, len(s.tools))} + for _, st := range s.tools { + result.Tools = append(result.Tools, st.Tool) + } + return result +} + +func (s *MCPServer) handleToolCall(ctx context.Context, id any, message json.RawMessage) mcp.JSONRPCMessage { + var req mcp.CallToolRequest + if err := json.Unmarshal(message, &req); err != nil { + return errorResponse(id, mcp.INVALID_REQUEST, err.Error()) + } + s.mu.RLock() + st, ok := s.tools[req.Params.Name] + s.mu.RUnlock() + if !ok { + return errorResponse(id, mcp.INVALID_PARAMS, fmt.Sprintf("tool %q not found", req.Params.Name)) + } + res, err := st.Handler(ctx, req) + if err != nil { + return errorResponse(id, mcp.INTERNAL_ERROR, err.Error()) + } + return successResponse(id, res) +} + +func (s *MCPServer) handleListResources() mcp.ListResourcesResult { + s.mu.RLock() + defer s.mu.RUnlock() + result := mcp.ListResourcesResult{Resources: make([]mcp.Resource, 0, len(s.resources))} + for _, sr := range s.resources { + result.Resources = append(result.Resources, sr.Resource) + } + return result +} + +func (s *MCPServer) handleListResourceTemplates() mcp.ListResourceTemplatesResult { + s.mu.RLock() + defer s.mu.RUnlock() + result := mcp.ListResourceTemplatesResult{ + ResourceTemplates: make([]mcp.ResourceTemplate, 0, len(s.resourceTemplates)), + } + for _, sr := range s.resourceTemplates { + result.ResourceTemplates = append(result.ResourceTemplates, sr.Template) + } + return result +} + +func (s *MCPServer) handleReadResource(ctx context.Context, id any, message json.RawMessage) mcp.JSONRPCMessage { + var req mcp.ReadResourceRequest + if err := json.Unmarshal(message, &req); err != nil { + return errorResponse(id, mcp.INVALID_REQUEST, err.Error()) + } + s.mu.RLock() + sr, ok := s.resources[req.Params.URI] + s.mu.RUnlock() + if !ok { + return errorResponse(id, mcp.RESOURCE_NOT_FOUND, fmt.Sprintf("resource %q not found", req.Params.URI)) + } + contents, err := sr.Handler(ctx, req) + if err != nil { + return errorResponse(id, mcp.INTERNAL_ERROR, err.Error()) + } + return successResponse(id, mcp.ReadResourceResult{Contents: contents}) +} + +func (s *MCPServer) handleListPrompts() mcp.ListPromptsResult { + s.mu.RLock() + defer s.mu.RUnlock() + result := mcp.ListPromptsResult{Prompts: make([]mcp.Prompt, 0, len(s.prompts))} + for _, sp := range s.prompts { + result.Prompts = append(result.Prompts, sp.Prompt) + } + return result +} + +func (s *MCPServer) handleGetPrompt(ctx context.Context, id any, message json.RawMessage) mcp.JSONRPCMessage { + var req mcp.GetPromptRequest + if err := json.Unmarshal(message, &req); err != nil { + return errorResponse(id, mcp.INVALID_REQUEST, err.Error()) + } + s.mu.RLock() + sp, ok := s.prompts[req.Params.Name] + s.mu.RUnlock() + if !ok { + return errorResponse(id, mcp.INVALID_PARAMS, fmt.Sprintf("prompt %q not found", req.Params.Name)) + } + res, err := sp.Handler(ctx, req) + if err != nil { + return errorResponse(id, mcp.INTERNAL_ERROR, err.Error()) + } + return successResponse(id, res) +} + +// successResponse builds a JSON-RPC success response mirroring mcp-go. +func successResponse(id, result any) mcp.JSONRPCMessage { + return mcp.NewJSONRPCResultResponse(mcp.NewRequestId(id), result) +} + +// errorResponse builds a JSON-RPC error response mirroring mcp-go. +func errorResponse(id any, code int, message string) mcp.JSONRPCMessage { + return mcp.JSONRPCError{ + JSONRPC: mcp.JSONRPC_VERSION, + ID: mcp.NewRequestId(id), + Error: mcp.NewJSONRPCErrorDetails(code, message, nil), + } +} diff --git a/mcpcompat/server/request_handler_test.go b/mcpcompat/server/request_handler_test.go new file mode 100644 index 0000000..5ec4ee0 --- /dev/null +++ b/mcpcompat/server/request_handler_test.go @@ -0,0 +1,248 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// newHandlerTestServer builds an MCPServer with one tool, resource, resource +// template and prompt registered, for driving HandleMessage directly. +func newHandlerTestServer() *server.MCPServer { + srv := server.NewMCPServer("handler-server", "4.5.6", + server.WithToolCapabilities(true), + server.WithResourceCapabilities(true, true), + server.WithPromptCapabilities(true), + server.WithLogging(), + ) + srv.AddTool( + mcp.NewTool("greet", mcp.WithDescription("greet"), mcp.WithString("name", mcp.Required())), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("hello " + req.GetString("name", "world")), nil + }, + ) + srv.AddResource( + mcp.Resource{URI: "file://readme", Name: "readme"}, + func(_ context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + return []mcp.ResourceContents{mcp.TextResourceContents{URI: "file://readme", Text: "body"}}, nil + }, + ) + srv.AddResourceTemplate( + mcp.ResourceTemplate{Name: "tmpl"}, + func(_ context.Context, _ mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + return nil, nil + }, + ) + srv.AddPrompt( + mcp.Prompt{Name: "p1"}, + func(_ context.Context, _ mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{Description: "desc"}, nil + }, + ) + return srv +} + +// marshalResult marshals a HandleMessage response and unmarshals into a generic +// map for structural assertions. +func marshalResult(t *testing.T, msg mcp.JSONRPCMessage) map[string]any { + t.Helper() + require.NotNil(t, msg) + b, err := json.Marshal(msg) + require.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal(b, &m)) + return m +} + +func TestHandleMessage(t *testing.T) { + t.Parallel() + srv := newHandlerTestServer() + ctx := context.Background() + + tests := []struct { + name string + message string + wantNil bool + wantErr bool + assertResp func(t *testing.T, resp map[string]any) + }{ + { + name: "initialize", + message: `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + result, ok := resp["result"].(map[string]any) + require.True(t, ok, "result present") + assert.Equal(t, mcp.LATEST_PROTOCOL_VERSION, result["protocolVersion"]) + info := result["serverInfo"].(map[string]any) + assert.Equal(t, "handler-server", info["name"]) + assert.Contains(t, result, "capabilities") + }, + }, + { + name: "ping", + message: `{"jsonrpc":"2.0","id":2,"method":"ping"}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + assert.Contains(t, resp, "result") + }, + }, + { + name: "tools/list", + message: `{"jsonrpc":"2.0","id":3,"method":"tools/list"}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + tools := resp["result"].(map[string]any)["tools"].([]any) + require.Len(t, tools, 1) + assert.Equal(t, "greet", tools[0].(map[string]any)["name"]) + }, + }, + { + name: "tools/call", + message: `{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"greet","arguments":{"name":"bob"}}}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + content := resp["result"].(map[string]any)["content"].([]any) + require.Len(t, content, 1) + assert.Equal(t, "hello bob", content[0].(map[string]any)["text"]) + }, + }, + { + name: "tools/call unknown tool", + message: `{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"nope"}}`, + wantErr: true, + }, + { + name: "resources/list", + message: `{"jsonrpc":"2.0","id":6,"method":"resources/list"}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + res := resp["result"].(map[string]any)["resources"].([]any) + require.Len(t, res, 1) + assert.Equal(t, "file://readme", res[0].(map[string]any)["uri"]) + }, + }, + { + name: "resources/templates/list", + message: `{"jsonrpc":"2.0","id":7,"method":"resources/templates/list"}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + tmpls := resp["result"].(map[string]any)["resourceTemplates"].([]any) + require.Len(t, tmpls, 1) + }, + }, + { + name: "resources/read", + message: `{"jsonrpc":"2.0","id":8,"method":"resources/read","params":{"uri":"file://readme"}}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + contents := resp["result"].(map[string]any)["contents"].([]any) + require.Len(t, contents, 1) + assert.Equal(t, "body", contents[0].(map[string]any)["text"]) + }, + }, + { + name: "resources/read unknown", + message: `{"jsonrpc":"2.0","id":9,"method":"resources/read","params":{"uri":"file://missing"}}`, + wantErr: true, + }, + { + name: "prompts/list", + message: `{"jsonrpc":"2.0","id":10,"method":"prompts/list"}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + prompts := resp["result"].(map[string]any)["prompts"].([]any) + require.Len(t, prompts, 1) + }, + }, + { + name: "prompts/get", + message: `{"jsonrpc":"2.0","id":11,"method":"prompts/get","params":{"name":"p1"}}`, + assertResp: func(t *testing.T, resp map[string]any) { + t.Helper() + assert.Equal(t, "desc", resp["result"].(map[string]any)["description"]) + }, + }, + { + name: "unknown method", + message: `{"jsonrpc":"2.0","id":12,"method":"does/not/exist"}`, + wantErr: true, + }, + { + name: "bad jsonrpc version", + message: `{"jsonrpc":"1.0","id":13,"method":"ping"}`, + wantErr: true, + }, + { + name: "parse error", + message: `{not json`, + wantErr: true, + }, + { + name: "notification returns nil", + message: `{"jsonrpc":"2.0","method":"notifications/initialized"}`, + wantNil: true, + }, + { + name: "server response returns nil", + message: `{"jsonrpc":"2.0","id":14,"result":{}}`, + wantNil: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + resp := srv.HandleMessage(ctx, json.RawMessage(tt.message)) + if tt.wantNil { + assert.Nil(t, resp) + return + } + m := marshalResult(t, resp) + if tt.wantErr { + assert.Contains(t, m, "error") + return + } + assert.NotContains(t, m, "error") + if tt.assertResp != nil { + tt.assertResp(t, m) + } + }) + } +} + +func TestWithContext(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("ctx-server", "1.0.0") + + // A ClientSession stored via WithContext must be recoverable via + // ClientSessionFromContext. + sess := &fakeSession{id: "sess-123"} + ctx := srv.WithContext(context.Background(), sess) + + got := server.ClientSessionFromContext(ctx) + require.NotNil(t, got) + assert.Equal(t, "sess-123", got.SessionID()) + + // A context with no session yields nil. + assert.Nil(t, server.ClientSessionFromContext(context.Background())) +} + +// fakeSession is a minimal ClientSession for WithContext round-tripping. +type fakeSession struct { + id string +} + +func (f *fakeSession) SessionID() string { return f.id } +func (*fakeSession) Initialize() {} +func (*fakeSession) Initialized() bool { return true } +func (*fakeSession) NotificationChannel() chan<- mcp.JSONRPCNotification { return nil } diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go new file mode 100644 index 0000000..da08eae --- /dev/null +++ b/mcpcompat/server/server.go @@ -0,0 +1,571 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +// Package server is a drop-in compatibility shim for +// github.com/mark3labs/mcp-go/server, reimplemented on top of the official +// github.com/modelcontextprotocol/go-sdk. +// +// It presents mcp-go's MCPServer API (NewMCPServer, AddTool, AddResource, +// AddPrompt, the ServerTool/ServerResource/ServerPrompt registration units, the +// Hooks and per-session interfaces, and the stdio/SSE/Streamable-HTTP +// transports) while delegating protocol handling to a go-sdk Server. Tools, +// resources and prompts registered here are converted to their go-sdk +// equivalents and served by the SDK. +// +// # Scope and status +// +// The global registration path (AddTool/AddResource/AddPrompt served over the +// stdio and HTTP transports) is fully functional and tested. The per-session +// interfaces (SessionWithTools, SessionWithResources, SessionIdManager) and the +// Hooks type are implemented for source compatibility, and per-session tool +// overlays are stored on the session objects. Wiring those overlays into +// go-sdk's live session lifecycle so that per-session tool *dispatch* matches +// mcp-go exactly (ToolHive's vMCP projection) is the one area that needs +// integration validation against ToolHive before this package can fully replace +// mcp-go for the vMCP server; see the notes on SessionWithTools. +// +// Stability: Alpha. +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net/http" + "strings" + "sync" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// ServerOption configures an MCPServer. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type ServerOption func(*MCPServer) + +// ToolHandlerFunc handles a tool call. It mirrors mcp-go's type exactly. +type ToolHandlerFunc func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) + +// ResourceHandlerFunc handles a resource read. +type ResourceHandlerFunc func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) + +// ResourceTemplateHandlerFunc handles a templated resource read. +type ResourceTemplateHandlerFunc func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) + +// PromptHandlerFunc handles a prompt get. +type PromptHandlerFunc func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) + +// NotificationHandlerFunc handles a client notification. +type NotificationHandlerFunc func(ctx context.Context, notification mcp.JSONRPCNotification) + +// ServerTool pairs a tool with its handler. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type ServerTool struct { + Tool mcp.Tool + Handler ToolHandlerFunc +} + +// ServerResource pairs a resource with its handler. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type ServerResource struct { + Resource mcp.Resource + Handler ResourceHandlerFunc +} + +// ServerResourceTemplate pairs a resource template with its handler. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type ServerResourceTemplate struct { + Template mcp.ResourceTemplate + Handler ResourceTemplateHandlerFunc +} + +// ServerPrompt pairs a prompt with its handler. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type ServerPrompt struct { + Prompt mcp.Prompt + Handler PromptHandlerFunc +} + +// MCPServer is an MCP server backed by the official go-sdk. It mirrors the +// subset of mcp-go's server.MCPServer that ToolHive uses. +type MCPServer struct { + name string + version string + logger *slog.Logger + hooks *Hooks + + // capability flags (informational; go-sdk infers capabilities from + // registered features, but these are retained for API compatibility). + toolListChanged bool + resourceSubscribe bool + resourceListChanged bool + promptListChanged bool + logging bool + + mu sync.RWMutex + tools map[string]ServerTool + resources map[string]ServerResource + resourceTemplates map[string]ServerResourceTemplate + prompts map[string]ServerPrompt + + sessions sync.Map // sessionID -> *clientSession + + // localSessions records the IDs of sessions that were initialized on THIS + // server instance (i.e. the initialize handshake was handled here by the + // go-sdk StreamableHTTPHandler). The Streamable HTTP transport uses it to + // decide, for a request carrying an existing session ID, whether the session + // is local (route to the go-sdk handler, which owns its session map) or was + // created on another replica (rehydrate; see StreamableHTTPServer). Populated + // in registerAndSync (which only fires on this instance's initialize path). + localSessions sync.Map // sessionID -> struct{} + + // pendingReqCtx maps an in-flight request's session ID to the HTTP request + // context, so the dispatch middleware can bridge per-request context values + // (identity, audit BackendInfo, telemetry) into the handler context. The + // go-sdk processes messages on a detached session goroutine and does not + // propagate the HTTP request context the way mcp-go did; this restores it. + pendingReqCtx sync.Map // sessionID -> context.Context +} + +// setPendingRequestContext records the HTTP request context for an in-flight +// request on the given session so the dispatch middleware can bridge its values. +func (s *MCPServer) setPendingRequestContext(ctx context.Context, sessionID string) { + s.pendingReqCtx.Store(sessionID, ctx) +} + +// pendingRequestContext returns the recorded HTTP request context for sessionID. +func (s *MCPServer) pendingRequestContext(sessionID string) context.Context { + if v, ok := s.pendingReqCtx.Load(sessionID); ok { + return v.(context.Context) + } + return nil +} + +// clearPendingRequestContext drops the recorded HTTP request context. +func (s *MCPServer) clearPendingRequestContext(sessionID string) { + s.pendingReqCtx.Delete(sessionID) +} + +// valueBridgeContext bridges the originating HTTP request's context values into +// a handler running on go-sdk's detached session goroutine. Its lifecycle +// (Deadline/Done/Err) comes from the embedded handler context; Value lookups +// consult the per-request HTTP context (values) FIRST, then fall back to the +// handler context. +// +// The per-request context must take precedence because go-sdk uses the +// *initialize* request's context as the whole session's context. Without +// values-first ordering, request-scoped values that the HTTP middleware chain +// re-establishes per request (audit BackendInfo, identity, telemetry) would be +// shadowed by the stale copies frozen at initialize time. go-sdk's own internal +// context keys are absent from the raw HTTP request context, so they still +// resolve via the fallback. +type valueBridgeContext struct { + context.Context + values context.Context +} + +func (c *valueBridgeContext) Value(key any) any { + if v := c.values.Value(key); v != nil { + return v + } + return c.Context.Value(key) +} + +// NewMCPServer creates a new MCP server with the given name and version. +func NewMCPServer(name, version string, opts ...ServerOption) *MCPServer { + s := &MCPServer{ + name: name, + version: version, + tools: make(map[string]ServerTool), + resources: make(map[string]ServerResource), + resourceTemplates: make(map[string]ServerResourceTemplate), + prompts: make(map[string]ServerPrompt), + } + for _, opt := range opts { + opt(s) + } + return s +} + +// WithToolCapabilities declares tool support (listChanged notifications). +func WithToolCapabilities(listChanged bool) ServerOption { + return func(s *MCPServer) { s.toolListChanged = listChanged } +} + +// WithResourceCapabilities declares resource support. +func WithResourceCapabilities(subscribe, listChanged bool) ServerOption { + return func(s *MCPServer) { + s.resourceSubscribe = subscribe + s.resourceListChanged = listChanged + } +} + +// WithPromptCapabilities declares prompt support. +func WithPromptCapabilities(listChanged bool) ServerOption { + return func(s *MCPServer) { s.promptListChanged = listChanged } +} + +// WithLogging enables logging capability. +func WithLogging() ServerOption { + return func(s *MCPServer) { s.logging = true } +} + +// WithLogger sets the server logger. +func WithLogger(logger *slog.Logger) ServerOption { + return func(s *MCPServer) { s.logger = logger } +} + +// WithHooks installs lifecycle hooks. +func WithHooks(hooks *Hooks) ServerOption { + return func(s *MCPServer) { s.hooks = hooks } +} + +// AddTool registers a tool and its handler. +func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) { + s.mu.Lock() + defer s.mu.Unlock() + s.tools[tool.Name] = ServerTool{Tool: tool, Handler: handler} +} + +// AddTools registers multiple tools. +func (s *MCPServer) AddTools(tools ...ServerTool) { + s.mu.Lock() + defer s.mu.Unlock() + for _, t := range tools { + s.tools[t.Tool.Name] = t + } +} + +// SetTools replaces the tool set. +func (s *MCPServer) SetTools(tools ...ServerTool) { + s.mu.Lock() + defer s.mu.Unlock() + s.tools = make(map[string]ServerTool, len(tools)) + for _, t := range tools { + s.tools[t.Tool.Name] = t + } +} + +// DeleteTools removes tools by name. +func (s *MCPServer) DeleteTools(names ...string) { + s.mu.Lock() + defer s.mu.Unlock() + for _, n := range names { + delete(s.tools, n) + } +} + +// AddResource registers a resource and its handler. +func (s *MCPServer) AddResource(resource mcp.Resource, handler ResourceHandlerFunc) { + s.mu.Lock() + defer s.mu.Unlock() + s.resources[resource.URI] = ServerResource{Resource: resource, Handler: handler} +} + +// AddResourceTemplate registers a resource template and its handler. +func (s *MCPServer) AddResourceTemplate(template mcp.ResourceTemplate, handler ResourceTemplateHandlerFunc) { + s.mu.Lock() + defer s.mu.Unlock() + name := template.Name + s.resourceTemplates[name] = ServerResourceTemplate{Template: template, Handler: handler} +} + +// AddPrompt registers a prompt and its handler. +func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { + s.mu.Lock() + defer s.mu.Unlock() + s.prompts[prompt.Name] = ServerPrompt{Prompt: prompt, Handler: handler} +} + +// buildServer constructs a go-sdk Server from the globally-registered features +// (AddTool/AddResource/AddPrompt). +// +// Per-session overlays (SessionWithTools/SessionWithResources) are NOT baked in +// here: the streamable/SSE transports call this once per new client session (via +// getServer) so each session gets its own go-sdk Server, and the registration +// middleware installed by this function syncs that session's overlay tools and +// resources onto its own server once the OnRegisterSession hooks have run. This +// mirrors mcp-go, whose per-session tools were dispatched per connection. +func (s *MCPServer) buildServer(genSessionID func() string) (*gosdk.Server, error) { + s.mu.RLock() + tools := make(map[string]ServerTool, len(s.tools)) + for k, v := range s.tools { + tools[k] = v + } + resources := make(map[string]ServerResource, len(s.resources)) + for k, v := range s.resources { + resources[k] = v + } + prompts := make(map[string]ServerPrompt, len(s.prompts)) + for k, v := range s.prompts { + prompts[k] = v + } + s.mu.RUnlock() + + impl := &gosdk.Implementation{Name: s.name, Version: s.version} + // srv is referenced by the InitializedHandler and registration middleware + // closures below; it is assigned before either can fire (both run only while + // serving a request, long after NewServer returns). + var srv *gosdk.Server + opts := &gosdk.ServerOptions{ + Logger: s.logger, + InitializedHandler: func(ctx context.Context, req *gosdk.InitializedRequest) { + if req == nil || req.Session == nil { + return + } + s.registerAndSync(ctx, req.Session, srv) + }, + } + // When a SessionIdManager is supplied (WithSessionIdManager), drive the SDK's + // session-ID generation through it: mcp-go called Generate() to mint the ID, + // which is where ToolHive's manager creates the placeholder session record + // that the OnRegisterSession hook later promotes via CreateSession. Without + // this the SDK would mint its own ID and CreateSession would fail to find the + // placeholder. + if genSessionID != nil { + opts.GetSessionID = genSessionID + } + srv = gosdk.NewServer(impl, opts) + + for _, st := range tools { + gt, err := toGoSDKTool(st.Tool) + if err != nil { + return nil, fmt.Errorf("converting tool %q: %w", st.Tool.Name, err) + } + srv.AddTool(gt, s.wrapToolHandler(st.Handler)) + } + for _, sr := range resources { + gr := &gosdk.Resource{} + if err := jsonConvert(sr.Resource, gr); err != nil { + return nil, fmt.Errorf("converting resource %q: %w", sr.Resource.URI, err) + } + srv.AddResource(gr, s.wrapResourceHandler(sr.Handler)) + } + for _, sp := range prompts { + gp := &gosdk.Prompt{} + if err := jsonConvert(sp.Prompt, gp); err != nil { + return nil, fmt.Errorf("converting prompt %q: %w", sp.Prompt.Name, err) + } + srv.AddPrompt(gp, s.wrapPromptHandler(sp.Handler)) + } + + srv.AddReceivingMiddleware(s.sessionDispatchMiddleware(srv)) + return srv, nil +} + +// sessionDispatchMiddleware wires mcp-go's per-session semantics onto a go-sdk +// server: it registers the session (firing OnRegisterSession) when the client +// initializes — mcp-go fired that hook on initialize, whereas go-sdk's +// InitializedHandler only fires on the later notifications/initialized — and it +// fires the before-list/before-call hooks so ToolHive's lazy per-session tool +// injection runs before the SDK enumerates or dispatches tools. +// getServerFunc returns a getServer callback for the go-sdk HTTP/SSE handlers, +// which invoke it once per new client session. genSessionID (may be nil) is the +// session-ID generator to install on each per-session server. On a build error +// it logs and returns nil, which the go-sdk handler surfaces as an HTTP 400. +func (s *MCPServer) getServerFunc(genSessionID func() string) func(*http.Request) *gosdk.Server { + return func(*http.Request) *gosdk.Server { + srv, err := s.buildServer(genSessionID) + if err != nil { + if s.logger != nil { + s.logger.Error("building per-session MCP server", "error", err) + } + return nil + } + return srv + } +} + +func (s *MCPServer) sessionDispatchMiddleware(srv *gosdk.Server) gosdk.Middleware { + return func(next gosdk.MethodHandler) gosdk.MethodHandler { + return func(ctx context.Context, method string, req gosdk.Request) (gosdk.Result, error) { + ss, _ := req.GetSession().(*gosdk.ServerSession) + if ss != nil { + ctx = s.contextWithSession(ctx, ss) + // Bridge the originating HTTP request's context values (identity, + // audit BackendInfo, telemetry) into the handler context. + if reqCtx := s.pendingRequestContext(ss.ID()); reqCtx != nil { + ctx = &valueBridgeContext{Context: ctx, values: reqCtx} + } + } + // Fire the before-hooks ahead of the SDK's own handling so a + // hook that injects per-session tools does so before the SDK + // enumerates (tools/list) or dispatches (tools/call) them. + switch method { + case string(mcp.MethodToolsList): + if s.hooks != nil { + s.hooks.beforeListTools(ctx, nil, &mcp.ListToolsRequest{}) + } + case string(mcp.MethodToolsCall): + if s.hooks != nil { + s.hooks.beforeCallTool(ctx, nil, &mcp.CallToolRequest{}) + } + } + res, err := next(ctx, method, req) + if err != nil && method == string(mcp.MethodToolsCall) { + err = translateUnknownToolError(err, req) + } + if err == nil && method == string(mcp.MethodInitialize) && ss != nil { + s.registerAndSync(ctx, ss, srv) + } + return res, err + } + } +} + +func (s *MCPServer) wrapToolHandler(h ToolHandlerFunc) gosdk.ToolHandler { + return func(ctx context.Context, req *gosdk.CallToolRequest) (res *gosdk.CallToolResult, err error) { + // mcp-go recovered panics in handlers at the transport/session layer, + // turning them into an error response. go-sdk runs handlers on a detached + // session goroutine with no recovery, so an unrecovered panic crashes the + // process. Recover here to preserve mcp-go's fault isolation. + defer func() { + if r := recover(); r != nil { + res, err = nil, fmt.Errorf("panic recovered in tool handler %q: %v", req.Params.Name, r) + } + }() + ctx = s.contextWithSession(ctx, req.Session) + + var args map[string]any + if len(req.Params.Arguments) > 0 { + if err := json.Unmarshal(req.Params.Arguments, &args); err != nil { + return nil, fmt.Errorf("unmarshaling arguments: %w", err) + } + } + mreq := mcp.CallToolRequest{} + mreq.Params.Name = req.Params.Name + mreq.Params.Arguments = args + // Preserve the request _meta so ToolHive can propagate metadata through + // vMCP to backends. go-sdk's Meta is map[string]any; mcp-go's *Meta has a + // custom (un)marshaler, so convert via JSON. + if len(req.Params.Meta) > 0 { + meta := &mcp.Meta{} + if err := jsonConvert(req.Params.Meta, meta); err != nil { + return nil, fmt.Errorf("converting call meta: %w", err) + } + mreq.Params.Meta = meta + } + + // Note: the before-call hook fires in sessionDispatchMiddleware, ahead of + // the SDK's dispatch, so per-session tool injection happens before the SDK + // resolves the tool. It is intentionally not fired again here. + + mres, err := h(ctx, mreq) + if err != nil { + return nil, err + } + out := &gosdk.CallToolResult{} + if err := jsonConvert(mres, out); err != nil { + return nil, fmt.Errorf("converting tool result: %w", err) + } + return out, nil + } +} + +func (s *MCPServer) wrapResourceHandler(h ResourceHandlerFunc) gosdk.ResourceHandler { + return func(ctx context.Context, req *gosdk.ReadResourceRequest) (res *gosdk.ReadResourceResult, err error) { + defer func() { + if r := recover(); r != nil { + res, err = nil, fmt.Errorf("panic recovered in resource handler %q: %v", req.Params.URI, r) + } + }() + ctx = s.contextWithSession(ctx, req.Session) + mreq := mcp.ReadResourceRequest{} + mreq.Params.URI = req.Params.URI + contents, err := h(ctx, mreq) + if err != nil { + return nil, err + } + out := &gosdk.ReadResourceResult{} + if err := jsonConvert(mcp.ReadResourceResult{Contents: contents}, out); err != nil { + return nil, fmt.Errorf("converting resource result: %w", err) + } + return out, nil + } +} + +func (s *MCPServer) wrapPromptHandler(h PromptHandlerFunc) gosdk.PromptHandler { + return func(ctx context.Context, req *gosdk.GetPromptRequest) (res *gosdk.GetPromptResult, err error) { + defer func() { + if r := recover(); r != nil { + res, err = nil, fmt.Errorf("panic recovered in prompt handler %q: %v", req.Params.Name, r) + } + }() + ctx = s.contextWithSession(ctx, req.Session) + mreq := mcp.GetPromptRequest{} + mreq.Params.Name = req.Params.Name + mreq.Params.Arguments = req.Params.Arguments + mres, err := h(ctx, mreq) + if err != nil { + return nil, err + } + out := &gosdk.GetPromptResult{} + if err := jsonConvert(mres, out); err != nil { + return nil, fmt.Errorf("converting prompt result: %w", err) + } + return out, nil + } +} + +// toGoSDKTool converts an mcp-go-shaped Tool into a go-sdk Tool. Both marshal to +// the same MCP wire JSON (including the outputSchema derived from +// RawOutputSchema), so a JSON round-trip is a faithful conversion. +// +// go-sdk's AddTool panics unless InputSchema is a non-nil object schema, whereas +// mcp-go tolerated a missing/empty schema. Normalize to the empty object schema +// ({"type":"object"}) so tools with no declared input (common in ToolHive's +// per-session vMCP projection) register cleanly, matching mcp-go's leniency. +func toGoSDKTool(t mcp.Tool) (*gosdk.Tool, error) { + out := &gosdk.Tool{} + if err := jsonConvert(t, out); err != nil { + return nil, err + } + out.InputSchema = normalizeObjectSchema(out.InputSchema) + return out, nil +} + +// normalizeObjectSchema ensures a JSON-schema value is a non-nil object schema +// suitable for go-sdk's AddTool. A nil schema, or one whose "type" is not +// "object", is replaced with the empty object schema. +func normalizeObjectSchema(schema any) any { + if m, ok := schema.(map[string]any); ok { + if m["type"] == "object" { + return m + } + } + return map[string]any{"type": "object"} +} + +// translateUnknownToolError rewrites go-sdk's "unknown tool" error for a +// tools/call into mcp-go's `tool %q not found` wording, so callers (and tests) +// that matched mcp-go's message keep working. The JSON-RPC code (InvalidParams) +// is preserved. +func translateUnknownToolError(err error, req gosdk.Request) error { + var jerr *jsonrpc.Error + if !errors.As(err, &jerr) || !strings.Contains(jerr.Message, "unknown tool") { + return err + } + name := "" + if p, ok := req.GetParams().(*gosdk.CallToolParams); ok { + name = p.Name + } + return &jsonrpc.Error{Code: jerr.Code, Message: fmt.Sprintf("tool %q not found", name)} +} + +// jsonConvert marshals src and unmarshals it into dst. +func jsonConvert(src, dst any) error { + b, err := json.Marshal(src) + if err != nil { + return err + } + return json.Unmarshal(b, dst) +} diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go new file mode 100644 index 0000000..186eb97 --- /dev/null +++ b/mcpcompat/server/server_internal_test.go @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// Compile-time interface checks: the concrete session must satisfy the +// per-session interfaces ToolHive relies on. +var ( + _ ClientSession = (*clientSession)(nil) + _ SessionWithTools = (*clientSession)(nil) + _ SessionWithResources = (*clientSession)(nil) +) + +func TestClientSession_Store(t *testing.T) { + t.Parallel() + cs := newClientSession("sess-1") + + assert.Equal(t, "sess-1", cs.SessionID()) + assert.False(t, cs.Initialized()) + cs.Initialize() + assert.True(t, cs.Initialized()) + assert.NotNil(t, cs.NotificationChannel()) + + cs.SetSessionTools(map[string]ServerTool{ + "echo": {Tool: mcp.Tool{Name: "echo"}}, + }) + got := cs.GetSessionTools() + require.Contains(t, got, "echo") + + // GetSessionTools must return a copy (mutating it must not affect the store). + got["echo2"] = ServerTool{} + assert.NotContains(t, cs.GetSessionTools(), "echo2") + + cs.SetSessionResources(map[string]ServerResource{ + "file:///r": {Resource: mcp.Resource{URI: "file:///r"}}, + }) + assert.Contains(t, cs.GetSessionResources(), "file:///r") +} + +func TestHooks_Fire(t *testing.T) { + t.Parallel() + h := &Hooks{} + + var gotSession ClientSession + h.AddOnRegisterSession(func(_ context.Context, s ClientSession) { gotSession = s }) + + var gotCallName string + h.AddBeforeCallTool(func(_ context.Context, _ any, m *mcp.CallToolRequest) { gotCallName = m.Params.Name }) + + var listFired bool + h.AddBeforeListTools(func(_ context.Context, _ any, _ *mcp.ListToolsRequest) { listFired = true }) + + cs := newClientSession("s") + h.registerSession(context.Background(), cs) + assert.Equal(t, cs, gotSession) + + req := &mcp.CallToolRequest{} + req.Params.Name = "greet" + h.beforeCallTool(context.Background(), "id", req) + assert.Equal(t, "greet", gotCallName) + + h.beforeListTools(context.Background(), "id", &mcp.ListToolsRequest{}) + assert.True(t, listFired) +} + +func TestBuildServer_GlobalAndSessionTools(t *testing.T) { + t.Parallel() + s := NewMCPServer("s", "1") + s.AddTool(mcp.NewTool("global", mcp.WithDescription("g")), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, nil }) + + // Building the global server (with the globally-registered tool) must succeed. + // Per-session overlays are no longer baked in here; they are synced onto the + // per-session server by syncSessionTools once the session registers. + srv, err := s.buildServer(nil) + require.NoError(t, err) + require.NotNil(t, srv) + + // A session carrying an additional per-session tool syncs onto the server + // without error, including a tool that declares no input schema (mcp-go was + // lenient; the shim normalizes it to the empty object schema). + cs := s.sessionFor("sid") + cs.SetSessionTools(map[string]ServerTool{ + "session-only": { + Tool: mcp.Tool{Name: "session-only"}, + Handler: func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, nil }, + }, + }) + assert.NotPanics(t, func() { s.syncSessionTools(srv, cs) }) +} + +func TestBuildServer_WithSessionIDGenerator(t *testing.T) { + t.Parallel() + s := NewMCPServer("s", "1") + called := false + gen := func() string { called = true; return "generated-id" } + + srv, err := s.buildServer(gen) + require.NoError(t, err) + require.NotNil(t, srv) + // The generator is installed on the server (invoked by the SDK per new + // session), not called at build time. + assert.False(t, called) +} + +func TestClientSessionFromContext_Empty(t *testing.T) { + t.Parallel() + assert.Nil(t, ClientSessionFromContext(context.Background())) +} + +// fakeIDManager verifies the SessionIdManager interface is satisfiable. +type fakeIDManager struct{} + +func (fakeIDManager) Generate() string { return "id" } +func (fakeIDManager) Validate(string) (bool, error) { return false, nil } +func (fakeIDManager) Terminate(string) (bool, error) { return false, nil } + +var _ SessionIdManager = fakeIDManager{} diff --git a/mcpcompat/server/server_test.go b/mcpcompat/server/server_test.go new file mode 100644 index 0000000..59d8013 --- /dev/null +++ b/mcpcompat/server/server_test.go @@ -0,0 +1,113 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "context" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// TestGlobalServer_EndToEnd registers a tool on the compat server, serves it +// over Streamable HTTP, and drives it with the compat client — exercising the +// whole server->go-sdk->client path through both shims. +func TestGlobalServer_EndToEnd(t *testing.T) { + t.Parallel() + ctx := context.Background() + + var registered server.ClientSession + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { registered = s }) + + srv := server.NewMCPServer("compat-server", "1.2.3", + server.WithToolCapabilities(false), + server.WithLogging(), + server.WithHooks(hooks), + ) + + srv.AddTool( + mcp.NewTool("greet", + mcp.WithDescription("greet someone"), + mcp.WithString("name", mcp.Required()), + ), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + name := req.GetString("name", "world") + return mcp.NewToolResultText("hello " + name), nil + }, + ) + + httpSrv := server.NewStreamableHTTPServer(srv) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + + initRes, err := c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "test-client", Version: "1.0.0"}, + }, + }) + require.NoError(t, err) + assert.Equal(t, "compat-server", initRes.ServerInfo.Name) + + tools, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + require.Len(t, tools.Tools, 1) + assert.Equal(t, "greet", tools.Tools[0].Name) + + callRes, err := c.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "greet", Arguments: map[string]any{"name": "ada"}}, + }) + require.NoError(t, err) + require.False(t, callRes.IsError) + require.Len(t, callRes.Content, 1) + txt, ok := mcp.AsTextContent(callRes.Content[0]) + require.True(t, ok) + assert.Equal(t, "hello ada", txt.Text) + + require.NoError(t, c.Close()) + + // The OnRegisterSession hook should have fired for the connected session. + assert.NotNil(t, registered) + if registered != nil { + assert.NotEmpty(t, registered.SessionID()) + } +} + +// TestServeStdio_Builds verifies the stdio entrypoint constructs a server from +// the registered tools without error (it blocks on Run, so we only exercise the +// build path here via a server with a tool registered). +func TestServer_RegistrationSurface(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("s", "1", + server.WithToolCapabilities(true), + server.WithResourceCapabilities(true, true), + server.WithPromptCapabilities(true), + ) + + srv.AddTool(mcp.Tool{Name: "t", Description: "d"}, + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, nil }) + srv.AddResource(mcp.Resource{URI: "file:///r", Name: "r"}, + func(context.Context, mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { return nil, nil }) + srv.AddPrompt(mcp.Prompt{Name: "p"}, + func(context.Context, mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return nil, nil }) + + // ServerTool / ServerResource / ServerPrompt are the registration units + // ToolHive stores; verify they compose with the handler types. + _ = server.ServerTool{Tool: mcp.Tool{Name: "x"}, Handler: func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, nil }} + _ = server.ServerResource{Resource: mcp.Resource{URI: "u"}, Handler: func(context.Context, mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { return nil, nil }} + _ = server.ServerPrompt{Prompt: mcp.Prompt{Name: "p"}, Handler: func(context.Context, mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return nil, nil }} + + srv.DeleteTools("t") +} diff --git a/mcpcompat/server/session.go b/mcpcompat/server/session.go new file mode 100644 index 0000000..a6e0350 --- /dev/null +++ b/mcpcompat/server/session.go @@ -0,0 +1,294 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "sync" + "sync/atomic" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// ClientSession represents an active client session. It mirrors mcp-go's +// server.ClientSession. +type ClientSession interface { + // Initialize marks the session as fully initialized. + Initialize() + // Initialized reports whether the session is ready for notifications. + Initialized() bool + // NotificationChannel provides a channel for sending notifications to the client. + NotificationChannel() chan<- mcp.JSONRPCNotification + // SessionID uniquely identifies the session. + SessionID() string +} + +// SessionWithTools is a ClientSession that carries per-session tools. ToolHive's +// vMCP layer uses this to project a per-session tool set. +// +// NOTE: overlays set here are stored and merged when a go-sdk server is built +// for the session (see MCPServer.buildServer). Making per-session tool changes +// take effect on an already-connected go-sdk session at runtime (live +// list_changed dispatch) is the integration point that needs validation against +// ToolHive's vMCP flow before this shim can fully replace mcp-go there. +type SessionWithTools interface { + ClientSession + // GetSessionTools returns the session's tools. Thread-safe. + GetSessionTools() map[string]ServerTool + // SetSessionTools sets the session's tools. Thread-safe. + SetSessionTools(tools map[string]ServerTool) +} + +// SessionWithResources is a ClientSession that carries per-session resources. +type SessionWithResources interface { + ClientSession + // GetSessionResources returns the session's resources. Thread-safe. + GetSessionResources() map[string]ServerResource + // SetSessionResources sets the session's resources. Thread-safe. + SetSessionResources(resources map[string]ServerResource) +} + +// SessionIdManager governs MCP session ID lifecycle. It mirrors mcp-go's +// server.SessionIdManager so ToolHive's implementation can be supplied via +// WithSessionIdManager. +type SessionIdManager interface { + // Generate returns a new session ID. + Generate() string + // Validate reports whether a session ID is valid; isTerminated is true if + // the ID is valid but belongs to a terminated session. + Validate(sessionID string) (isTerminated bool, err error) + // Terminate marks a session ID terminated; isNotAllowed is true if policy + // prevents client termination. + Terminate(sessionID string) (isNotAllowed bool, err error) +} + +// clientSession is the concrete ClientSession backed by a go-sdk ServerSession. +type clientSession struct { + id string + initialized atomic.Bool + registered atomic.Bool + notifCh chan mcp.JSONRPCNotification + goSession atomic.Pointer[gosdk.ServerSession] + + // owner and boundServer are set when the session's go-sdk server is bound + // (at registration). They let SetSessionTools/SetSessionResources reconcile + // the per-session overlay onto the live go-sdk server at runtime. + owner *MCPServer + boundServer atomic.Pointer[gosdk.Server] + + mu sync.RWMutex + tools map[string]ServerTool + resources map[string]ServerResource + // sdkToolNames tracks the tool names this session has added to its go-sdk + // server, so a later SetSessionTools can remove the ones that went away. + sdkToolNames map[string]struct{} +} + +func newClientSession(id string) *clientSession { + cs := &clientSession{ + id: id, + notifCh: make(chan mcp.JSONRPCNotification, 64), + } + // Drain the notification channel to avoid blocking senders. Forwarding + // server-initiated notifications onto the go-sdk session is a known gap. + go func() { + for range cs.notifCh { //nolint:revive // intentional drain + } + }() + return cs +} + +func (c *clientSession) SessionID() string { return c.id } + +func (c *clientSession) Initialize() { c.initialized.Store(true) } + +func (c *clientSession) Initialized() bool { return c.initialized.Load() } + +func (c *clientSession) NotificationChannel() chan<- mcp.JSONRPCNotification { return c.notifCh } + +func (c *clientSession) GetSessionTools() map[string]ServerTool { + c.mu.RLock() + defer c.mu.RUnlock() + out := make(map[string]ServerTool, len(c.tools)) + for k, v := range c.tools { + out[k] = v + } + return out +} + +func (c *clientSession) SetSessionTools(tools map[string]ServerTool) { + c.mu.Lock() + c.tools = make(map[string]ServerTool, len(tools)) + for k, v := range tools { + c.tools[k] = v + } + c.mu.Unlock() + // Reconcile the overlay onto the live go-sdk server if one is bound. + if srv := c.boundServer.Load(); srv != nil && c.owner != nil { + c.owner.syncSessionTools(srv, c) + } +} + +func (c *clientSession) GetSessionResources() map[string]ServerResource { + c.mu.RLock() + defer c.mu.RUnlock() + out := make(map[string]ServerResource, len(c.resources)) + for k, v := range c.resources { + out[k] = v + } + return out +} + +func (c *clientSession) SetSessionResources(resources map[string]ServerResource) { + c.mu.Lock() + c.resources = make(map[string]ServerResource, len(resources)) + for k, v := range resources { + c.resources[k] = v + } + c.mu.Unlock() + if srv := c.boundServer.Load(); srv != nil && c.owner != nil { + c.owner.syncSessionResources(srv, c) + } +} + +// sessionContextKey is the context key under which the ClientSession is stored. +type sessionContextKey struct{} + +// contextWithSession looks up (or creates) the clientSession for the given +// go-sdk ServerSession and stores it in the context. +func (s *MCPServer) contextWithSession(ctx context.Context, ss *gosdk.ServerSession) context.Context { + if ss == nil { + return ctx + } + cs := s.sessionFor(ss.ID()) + cs.goSession.Store(ss) + return context.WithValue(ctx, sessionContextKey{}, ClientSession(cs)) +} + +// sessionFor returns the registered clientSession for id, creating it if needed. +func (s *MCPServer) sessionFor(id string) *clientSession { + if v, ok := s.sessions.Load(id); ok { + return v.(*clientSession) + } + cs := newClientSession(id) + actual, _ := s.sessions.LoadOrStore(id, cs) + return actual.(*clientSession) +} + +// registerAndSync registers the session for the given go-sdk ServerSession, +// firing the OnRegisterSession hooks exactly once and then reconciling any +// per-session tool/resource overlay the hooks installed onto srv (the go-sdk +// server bound to this session). It is invoked from the initialize dispatch +// middleware (matching mcp-go's on-initialize timing) and, defensively, from the +// InitializedHandler; the once-guard makes the second call a cheap no-op. +func (s *MCPServer) registerAndSync(ctx context.Context, ss *gosdk.ServerSession, srv *gosdk.Server) { + if ss == nil { + return + } + cs := s.sessionFor(ss.ID()) + cs.goSession.Store(ss) + cs.owner = s + cs.boundServer.Store(srv) + // Mark the session local: it was initialized on this instance, so the go-sdk + // StreamableHTTPHandler owns it and subsequent requests must route there + // rather than through the cross-replica rehydration path. + s.localSessions.Store(ss.ID(), struct{}{}) + if !cs.registered.CompareAndSwap(false, true) { + return + } + cs.Initialize() + if s.hooks != nil { + s.hooks.registerSession(ctx, cs) + } + // The hooks may have installed per-session tools/resources via + // SetSessionTools/SetSessionResources; those calls reconcile onto srv + // themselves now that boundServer is set. Sync once more here to cover any + // overlay set before the server was bound. + s.syncSessionTools(srv, cs) + s.syncSessionResources(srv, cs) +} + +// syncSessionTools reconciles the session's tool overlay onto its go-sdk server: +// tools present in the overlay are added (overwriting by name) and tools that +// were previously added for this session but are no longer present are removed. +func (s *MCPServer) syncSessionTools(srv *gosdk.Server, cs *clientSession) { + cs.mu.Lock() + defer cs.mu.Unlock() + newNames := make(map[string]struct{}, len(cs.tools)) + for name, st := range cs.tools { + gt, err := toGoSDKTool(st.Tool) + if err != nil { + continue + } + srv.AddTool(gt, s.wrapToolHandler(st.Handler)) + newNames[name] = struct{}{} + } + var removed []string + for name := range cs.sdkToolNames { + if _, ok := newNames[name]; !ok { + removed = append(removed, name) + } + } + if len(removed) > 0 { + srv.RemoveTools(removed...) + } + cs.sdkToolNames = newNames +} + +// syncSessionResources adds the session's resource overlay onto its go-sdk +// server. Resources are add-only here (ToolHive sets them once at registration). +func (s *MCPServer) syncSessionResources(srv *gosdk.Server, cs *clientSession) { + cs.mu.RLock() + defer cs.mu.RUnlock() + for _, sr := range cs.resources { + gr := &gosdk.Resource{} + if err := jsonConvert(sr.Resource, gr); err != nil { + continue + } + srv.AddResource(gr, s.wrapResourceHandler(sr.Handler)) + } +} + +// isLocalSession reports whether the session ID was initialized on this server +// instance (see MCPServer.localSessions). +func (s *MCPServer) isLocalSession(id string) bool { + _, ok := s.localSessions.Load(id) + return ok +} + +// forgetSession drops all local bookkeeping for a session ID. It is called when +// a session is terminated (DELETE) so a later request with the same ID is not +// mistaken for a live local session. +func (s *MCPServer) forgetSession(id string) { + s.localSessions.Delete(id) + s.sessions.Delete(id) + s.pendingReqCtx.Delete(id) +} + +// bindRehydratedSession binds the clientSession for a session that was created +// on another replica and is being rehydrated here (see StreamableHTTPServer's +// rehydration path). Unlike registerAndSync it does NOT fire OnRegisterSession +// and does NOT mark the session local: a rehydrated session skips the +// initialize handshake (and therefore ToolHive's Generate/CreateSession +// two-phase creation), and cross-replica capability projection is driven by the +// before-list/before-call hooks (ToolHive's lazy per-session tool injection), +// not by OnRegisterSession. Binding owner+boundServer here lets those hooks' +// SetSessionTools/SetSessionResources reconcile the overlay onto srv. +func (s *MCPServer) bindRehydratedSession(id string, ss *gosdk.ServerSession, srv *gosdk.Server) { + cs := s.sessionFor(id) + cs.goSession.Store(ss) + cs.owner = s + cs.boundServer.Store(srv) + cs.Initialize() +} + +// ClientSessionFromContext returns the ClientSession stored in ctx, or nil. +func ClientSessionFromContext(ctx context.Context) ClientSession { + if v, ok := ctx.Value(sessionContextKey{}).(ClientSession); ok { + return v + } + return nil +} diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go new file mode 100644 index 0000000..39e06dc --- /dev/null +++ b/mcpcompat/server/transports.go @@ -0,0 +1,465 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "fmt" + "mime" + "net/http" + "strings" + "sync" + "time" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// bothAcceptMediaTypes is the Accept header value the go-sdk Streamable HTTP and +// SSE handlers require on a POST. mcp-go's server ignored Accept entirely, so +// clients such as ToolHive's tests that POST with only Content-Type set (no +// Accept) must be tolerated; the shim injects this value when either required +// media type is absent to restore that leniency. +const bothAcceptMediaTypes = "application/json, text/event-stream" + +// ensureAcceptMediaTypes restores mcp-go's Accept-header leniency: the go-sdk +// handlers reject a request whose Accept header does not advertise both +// application/json and text/event-stream. When either is missing this sets the +// Accept header to advertise both so the request is accepted, matching mcp-go. +func ensureAcceptMediaTypes(r *http.Request) { + var jsonOK, streamOK bool + for _, value := range r.Header.Values("Accept") { + for _, part := range strings.Split(value, ",") { + mediaType, _, err := mime.ParseMediaType(part) + if err != nil { + continue + } + switch mediaType { + case "application/json", "application/*", "*/*": + jsonOK = true + } + switch mediaType { + case "text/event-stream", "text/*", "*/*": + streamOK = true + } + } + } + if !jsonOK || !streamOK { + r.Header.Set("Accept", bothAcceptMediaTypes) + } +} + +// --- stdio ----------------------------------------------------------------- + +// StdioOption configures the stdio server (retained for API compatibility). +type StdioOption func(*stdioConfig) + +type stdioConfig struct{} + +// ServeStdio runs the MCP server over stdio until the context is done. It +// mirrors mcp-go's server.ServeStdio. +func ServeStdio(server *MCPServer, _ ...StdioOption) error { + srv, err := server.buildServer(nil) + if err != nil { + return err + } + return srv.Run(context.Background(), &gosdk.StdioTransport{}) +} + +// --- Streamable HTTP ------------------------------------------------------- + +// HTTPContextFunc customizes the request context for HTTP transports. +type HTTPContextFunc func(ctx context.Context, r *http.Request) context.Context + +// StreamableHTTPOption configures a StreamableHTTPServer. +type StreamableHTTPOption func(*StreamableHTTPServer) + +// StreamableHTTPServer serves the MCP server over the Streamable HTTP transport. +// It implements http.Handler so it can be mounted on an http.ServeMux, and also +// offers Start/Shutdown for standalone use. +type StreamableHTTPServer struct { + mcp *MCPServer + endpointPath string + contextFunc HTTPContextFunc + sessionIDMgr SessionIdManager + heartbeat time.Duration + + once sync.Once + handler http.Handler + buildErr error + httpSrv *http.Server + + // rehydrated holds sessions that were created on another replica and + // reconstructed here (see the rehydration path in ServeHTTP). Guarded by + // rehydratedMu. These use a custom StreamableServerTransport (SSE responses) + // rather than the go-sdk StreamableHTTPHandler, because the handler 404s any + // session ID it did not create itself. + rehydratedMu sync.Mutex + rehydrated map[string]*rehydratedSession +} + +// rehydratedSession is a session reconstructed from cross-replica shared state +// (validated via the SessionIdManager) and served by a per-session go-sdk +// StreamableServerTransport. +type rehydratedSession struct { + transport *gosdk.StreamableServerTransport + session *gosdk.ServerSession +} + +// defaultRehydratedProtocolVersion is the MCP protocol version seeded into a +// rehydrated session when the request carries no MCP-Protocol-Version header. It +// matches a widely-supported spec revision; clients that resumed a session send +// the negotiated version header, which takes precedence. +const defaultRehydratedProtocolVersion = "2025-06-18" + +// mcpProtocolVersionHeader is the HTTP header carrying the negotiated MCP +// protocol version on subsequent (post-initialize) requests. +const mcpProtocolVersionHeader = "MCP-Protocol-Version" + +// NewStreamableHTTPServer creates a Streamable HTTP server for the MCP server. +func NewStreamableHTTPServer(server *MCPServer, opts ...StreamableHTTPOption) *StreamableHTTPServer { + s := &StreamableHTTPServer{mcp: server, endpointPath: "/mcp"} + for _, opt := range opts { + opt(s) + } + return s +} + +// WithEndpointPath sets the HTTP path the server is mounted at. +func WithEndpointPath(endpointPath string) StreamableHTTPOption { + return func(s *StreamableHTTPServer) { s.endpointPath = endpointPath } +} + +// WithSessionIdManager supplies a session ID manager. +// +// NOTE: the go-sdk manages MCP session IDs internally; the supplied manager is +// retained for API compatibility but does not yet drive the SDK's ID lifecycle. +func WithSessionIdManager(manager SessionIdManager) StreamableHTTPOption { + return func(s *StreamableHTTPServer) { s.sessionIDMgr = manager } +} + +// WithHeartbeatInterval sets the keep-alive ping interval. +func WithHeartbeatInterval(interval time.Duration) StreamableHTTPOption { + return func(s *StreamableHTTPServer) { s.heartbeat = interval } +} + +// WithHTTPContextFunc installs a per-request context customizer. +func WithHTTPContextFunc(fn HTTPContextFunc) StreamableHTTPOption { + return func(s *StreamableHTTPServer) { s.contextFunc = fn } +} + +func (s *StreamableHTTPServer) build() { + s.once.Do(func() { + var gen func() string + if s.sessionIDMgr != nil { + gen = s.sessionIDMgr.Generate + } + // Validate the server configuration once up-front so a bad registration + // surfaces as a clean 500 rather than a per-request nil. + if _, err := s.mcp.buildServer(gen); err != nil { + s.buildErr = err + return + } + // JSONResponse makes the handler reply with application/json rather than + // text/event-stream for request/response exchanges, matching mcp-go's + // server so callers that json.Decode the response body keep working. + opts := &gosdk.StreamableHTTPOptions{JSONResponse: true} + // A fresh go-sdk server per client session lets each session carry its own + // tool/resource overlay (mcp-go's per-session projection), synced by the + // registration middleware buildServer installs. + s.handler = gosdk.NewStreamableHTTPHandler(s.mcp.getServerFunc(gen), opts) + }) +} + +// ServeHTTP implements http.Handler. +func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + s.build() + if s.buildErr != nil { + http.Error(w, fmt.Sprintf("building server: %v", s.buildErr), http.StatusInternalServerError) + return + } + ensureAcceptMediaTypes(r) + if s.contextFunc != nil { + r = r.WithContext(s.contextFunc(r.Context(), r)) + } + // Record the request context so the dispatch middleware can bridge its values + // into the handler running on go-sdk's detached session goroutine. Keyed by + // the client-supplied session ID; the initialize request (no session ID yet) + // carries no per-request values that the handler needs. This POST is handled + // synchronously (JSONResponse), so the entry is valid for the handler's whole + // lifetime and cleared when ServeHTTP returns. + if sid := r.Header.Get("Mcp-Session-Id"); sid != "" { + s.mcp.setPendingRequestContext(r.Context(), sid) + defer s.mcp.clearPendingRequestContext(sid) + } + // DELETE terminates the session. mcp-go answered 200 and drove the supplied + // SessionIdManager's Terminate; go-sdk answers 204 and manages its own session + // map. Rewrite the status to 200 for compatibility and forward the termination + // to the manager so ToolHive's session storage is cleaned up in lockstep. + if r.Method == http.MethodDelete { + sid := r.Header.Get("Mcp-Session-Id") + s.deleteRehydrated(sid) + s.mcp.forgetSession(sid) + rw := &statusRewriter{ResponseWriter: w, from: http.StatusNoContent, to: http.StatusOK} + s.handler.ServeHTTP(rw, r) + if s.sessionIDMgr != nil && sid != "" { + _, _ = s.sessionIDMgr.Terminate(sid) + } + return + } + + // Cross-replica routing: a request carrying a session ID that was NOT + // initialized on this instance (its initialize handshake happened on another + // replica) cannot be served by the go-sdk StreamableHTTPHandler, which 404s + // any session ID it did not create. Validate it against the shared + // SessionIdManager and, if valid, rehydrate a local session for it. + if sid := r.Header.Get("Mcp-Session-Id"); sid != "" && s.sessionIDMgr != nil && !s.mcp.isLocalSession(sid) { + s.serveRehydrated(w, r, sid) + return + } + + s.handler.ServeHTTP(w, r) +} + +// serveRehydrated routes a request for a session created on another replica. +// It validates the session ID against the shared SessionIdManager and serves it +// through a locally-reconstructed session, matching mcp-go's behavior where any +// replica sharing the manager's backing store accepts the session. +func (s *StreamableHTTPServer) serveRehydrated(w http.ResponseWriter, r *http.Request, sid string) { + // Validate on EVERY request (as mcp-go does), not just on cache miss. This is + // what implements lazy eviction: a session terminated on another replica is + // marked terminated in the shared store, and the next request here must reject + // it and drop any locally-cached reconstruction rather than serving it. + isTerminated, err := s.sessionIDMgr.Validate(sid) + if err != nil { + s.deleteRehydrated(sid) + http.Error(w, "Invalid session ID", http.StatusNotFound) + return + } + if isTerminated { + s.deleteRehydrated(sid) + http.Error(w, "Session terminated", http.StatusNotFound) + return + } + + rt := s.getRehydrated(sid) + if rt == nil { + rt, err = s.rehydrate(r, sid) + if err != nil { + if s.mcp.logger != nil { + s.mcp.logger.Error("rehydrating cross-replica session", "session_id", sid, "error", err) + } + http.Error(w, "failed to rehydrate session", http.StatusInternalServerError) + return + } + } + rt.transport.ServeHTTP(w, r) +} + +// getRehydrated returns the reconstructed session for sid, or nil. +func (s *StreamableHTTPServer) getRehydrated(sid string) *rehydratedSession { + s.rehydratedMu.Lock() + defer s.rehydratedMu.Unlock() + return s.rehydrated[sid] +} + +// deleteRehydrated closes and drops a reconstructed session, if present. +func (s *StreamableHTTPServer) deleteRehydrated(sid string) { + if sid == "" { + return + } + s.rehydratedMu.Lock() + rt := s.rehydrated[sid] + delete(s.rehydrated, sid) + s.rehydratedMu.Unlock() + if rt != nil { + _ = rt.session.Close() + } +} + +// rehydrate reconstructs a session that was created on another replica: it +// builds a fresh per-session go-sdk server (carrying this instance's tools plus +// the before-hook lazy-injection path), connects a StreamableServerTransport +// seeded with an already-initialized state (so it accepts non-initialize +// requests and can perform server->client calls such as elicitation), binds the +// clientSession so the before-hooks can reconcile the per-session overlay, and +// caches it keyed by session ID. +func (s *StreamableHTTPServer) rehydrate(r *http.Request, sid string) (*rehydratedSession, error) { + s.rehydratedMu.Lock() + defer s.rehydratedMu.Unlock() + // Double-check under the lock in case a concurrent request rehydrated first. + if rt, ok := s.rehydrated[sid]; ok { + return rt, nil + } + + srv, err := s.mcp.buildServer(nil) + if err != nil { + return nil, err + } + + protocolVersion := r.Header.Get(mcpProtocolVersionHeader) + if protocolVersion == "" { + protocolVersion = defaultRehydratedProtocolVersion + } + + transport := &gosdk.StreamableServerTransport{SessionID: sid, Stateless: false} + state := &gosdk.ServerSessionState{ + InitializeParams: &gosdk.InitializeParams{ + ProtocolVersion: protocolVersion, + // Seed the elicitation capability so a server->client elicitation on a + // rehydrated session passes go-sdk's capability gate. A stateless + // session cannot do this; this is what proves the rehydrated session is + // a full, stateful session. + Capabilities: &gosdk.ClientCapabilities{Elicitation: &gosdk.ElicitationCapabilities{}}, + }, + InitializedParams: &gosdk.InitializedParams{}, + LogLevel: "info", + } + // Detach from the request context: this session outlives the request that + // created it (subsequent requests on other HTTP connections reuse it). + session, err := srv.Connect(context.WithoutCancel(r.Context()), transport, &gosdk.ServerSessionOptions{State: state}) + if err != nil { + return nil, err + } + s.mcp.bindRehydratedSession(sid, session, srv) + + rt := &rehydratedSession{transport: transport, session: session} + if s.rehydrated == nil { + s.rehydrated = make(map[string]*rehydratedSession) + } + s.rehydrated[sid] = rt + return rt, nil +} + +// statusRewriter is an http.ResponseWriter that rewrites a single status code on +// WriteHeader (used to translate go-sdk's 204 DELETE response to mcp-go's 200). +type statusRewriter struct { + http.ResponseWriter + from, to int +} + +func (s *statusRewriter) WriteHeader(code int) { + if code == s.from { + code = s.to + } + s.ResponseWriter.WriteHeader(code) +} + +// Start serves on addr until Shutdown is called. +func (s *StreamableHTTPServer) Start(addr string) error { + mux := http.NewServeMux() + mux.Handle(s.endpointPath, s) + s.httpSrv = &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 10 * time.Second} + return s.httpSrv.ListenAndServe() +} + +// Shutdown gracefully stops the HTTP server. +func (s *StreamableHTTPServer) Shutdown(ctx context.Context) error { + if s.httpSrv == nil { + return nil + } + return s.httpSrv.Shutdown(ctx) +} + +// --- SSE ------------------------------------------------------------------- + +// SSEOption configures an SSEServer. +type SSEOption func(*SSEServer) + +// SSEServer serves the MCP server over the (legacy) SSE transport. +type SSEServer struct { + mcp *MCPServer + sseEndpoint string + messageEndpoint string + + once sync.Once + handler http.Handler + buildErr error + httpSrv *http.Server +} + +// NewSSEServer creates an SSE server for the MCP server. +func NewSSEServer(server *MCPServer, opts ...SSEOption) *SSEServer { + s := &SSEServer{mcp: server, sseEndpoint: "/sse", messageEndpoint: "/message"} + for _, opt := range opts { + opt(s) + } + return s +} + +// WithSSEEndpoint sets the SSE endpoint path. +func WithSSEEndpoint(endpoint string) SSEOption { + return func(s *SSEServer) { s.sseEndpoint = endpoint } +} + +// WithMessageEndpoint sets the message endpoint path. +func WithMessageEndpoint(endpoint string) SSEOption { + return func(s *SSEServer) { s.messageEndpoint = endpoint } +} + +func (s *SSEServer) build() { + s.once.Do(func() { + if _, err := s.mcp.buildServer(nil); err != nil { + s.buildErr = err + return + } + s.handler = gosdk.NewSSEHandler(s.mcp.getServerFunc(nil), nil) + }) +} + +// ServeHTTP implements http.Handler. +func (s *SSEServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + s.build() + if s.buildErr != nil { + http.Error(w, fmt.Sprintf("building server: %v", s.buildErr), http.StatusInternalServerError) + return + } + ensureAcceptMediaTypes(r) + // mcp-go advertised a distinct message endpoint for client POSTs, whereas + // go-sdk derives the endpoint it advertises in the SSE "endpoint" event from + // the SSE (GET) request's own path plus a sessionid query. Rewrite the GET + // request path to the configured message endpoint so the advertised POST + // target matches mcp-go's split-path model (and any middleware mounted on the + // message path). The stream itself is already served on this connection, so + // the path change only affects the advertised endpoint. POSTs are dispatched + // by sessionid and are path-independent. + if r.Method == http.MethodGet && s.messageEndpoint != "" && r.URL.Path != s.messageEndpoint { + r = r.Clone(r.Context()) + r.URL.Path = s.messageEndpoint + } + s.handler.ServeHTTP(w, r) +} + +// SSEHandler returns the http.Handler for the SSE (stream) endpoint. It mirrors +// mcp-go's SSEServer.SSEHandler, allowing the endpoint to be mounted on a custom +// router. +// +// go-sdk backing and limitation: the go-sdk serves SSE and message delivery +// through a single unified handler keyed off the request method and a session +// query parameter, whereas mcp-go splits them across two paths. Both SSEHandler +// and MessageHandler therefore return the same underlying go-sdk handler; when +// mounting them on separate paths, mount both under a common base path so the +// go-sdk handler can correlate the stream and its message posts. +func (s *SSEServer) SSEHandler() http.Handler { + return http.HandlerFunc(s.ServeHTTP) +} + +// MessageHandler returns the http.Handler for the message (POST) endpoint. See +// SSEHandler for the go-sdk backing and the shared-handler limitation. +func (s *SSEServer) MessageHandler() http.Handler { + return http.HandlerFunc(s.ServeHTTP) +} + +// Start serves on addr until Shutdown is called. +func (s *SSEServer) Start(addr string) error { + s.httpSrv = &http.Server{Addr: addr, Handler: s, ReadHeaderTimeout: 10 * time.Second} + return s.httpSrv.ListenAndServe() +} + +// Shutdown gracefully stops the HTTP server. +func (s *SSEServer) Shutdown(ctx context.Context) error { + if s.httpSrv == nil { + return nil + } + return s.httpSrv.Shutdown(ctx) +}