From b22b8e162b15f43326894e176ac27f9b8ef260bd Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Wed, 8 Jul 2026 22:20:38 +0300 Subject: [PATCH 1/8] fix(mcpcompat): close capability, context-bridge, and schema gaps (issue #156 wave 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit U1 — Map capability flags to go-sdk ServerOptions.Capabilities so vMCP advertises tools/resources/prompts at initialize even with zero pre-registered features (merge blocker). U3 — Rekey the pending-request-context bridge from session ID to a per-POST crypto/rand nonce (X-MCP-Req-Nonce), eliminating the cross-request context bleed / identity-attribution race for concurrent POSTs on the same session (merge blocker, security). go-sdk does not propagate the per-POST HTTP context into session middleware, so the bridge is retained but keyed per-request. U6 — Narrow normalizeObjectSchema to replace only nil/empty/empty-type schemas with {"type":"object"}; pass $ref/oneOf/boolean/object- with-type-omitted through verbatim (matching mcp-go). Recover go-sdk AddTool panics on the per-session path so a bad overlay schema skips the tool instead of crashing the session. Tests: capability advertisement e2e, per-request context isolation under -race (verified to fail without the bridge), schema normalization table, non-object schema panic recovery, nil-schema tools/list e2e. Refs #156 --- mcpcompat/server/server.go | 225 ++++++++++++++++++----- mcpcompat/server/server_internal_test.go | 154 ++++++++++++++++ mcpcompat/server/server_test.go | 214 ++++++++++++++++++++- mcpcompat/server/session.go | 42 ++++- mcpcompat/server/transports.go | 33 +++- 5 files changed, 613 insertions(+), 55 deletions(-) diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index da08eae..2175756 100644 --- a/mcpcompat/server/server.go +++ b/mcpcompat/server/server.go @@ -43,6 +43,14 @@ import ( mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) +// schemaTypeObject is the empty object JSON schema ("type":"object") that +// go-sdk's AddTool requires a tool input schema to carry. schemaTypeKey is the +// JSON-schema "type" property name. +const ( + schemaTypeKey = "type" + schemaTypeObject = "object" +) + // ServerOption configures an MCPServer. // //nolint:revive // name intentionally matches mcp-go for drop-in compatibility. @@ -103,11 +111,21 @@ type MCPServer struct { logger *slog.Logger hooks *Hooks - // capability flags (informational; go-sdk infers capabilities from - // registered features, but these are retained for API compatibility). + // Capability flags set via WithToolCapabilities/WithResourceCapabilities/ + // WithPromptCapabilities. These are wired into the go-sdk ServerOptions in + // buildServer (see ServerOptions.Capabilities): go-sdk otherwise infers + // capabilities only from registered features, so a server that registers + // tools per-session AFTER initialize (ToolHive's vMCP projection) would + // advertise no tools capability at initialize time. The *Declared flags + // record that the corresponding With*Capabilities option was invoked (mcp-go + // advertises the capability whenever the option is used, regardless of the + // sub-flag value); the sub-flags carry the ListChanged/Subscribe settings. + toolsDeclared bool toolListChanged bool + resourcesDeclared bool resourceSubscribe bool resourceListChanged bool + promptsDeclared bool promptListChanged bool logging bool @@ -128,46 +146,60 @@ type MCPServer struct { // 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 -} + // pendingReqCtx bridges per-request context values (identity, audit + // BackendInfo, telemetry) into handlers running on go-sdk's session + // goroutine. go-sdk does NOT propagate the per-POST HTTP request's context + // into the receiving-middleware path for subsequent requests on an existing + // session: messages published from servePOST are handled on the connection + // goroutine whose context was captured at session-creation (initialize) + // time, so request-scoped values added via WithHTTPContextFunc would + // otherwise be lost. + // + // To avoid the per-session race where two concurrent POSTs on the same + // session clobber each other's context (issue #156, item U3), entries are + // keyed by a per-POST nonce rather than the session ID. ServeHTTP generates + // a nonce, stores the request context under it, and sets it as the + // X-MCP-Req-Nonce header; the dispatch middleware reads that header off the + // per-request RequestExtra (req.GetExtra().Header, which go-sdk populates + // from the POST's headers) to look up the correct context. Entries are + // cleared when ServeHTTP returns. + pendingReqCtx sync.Map // nonce -> context.Context +} + +// reqNonceHeader is the HTTP header carrying the per-POST nonce that +// correlates a request's context (stored by ServeHTTP) with the handler +// invocation on go-sdk's session goroutine. +const reqNonceHeader = "X-MCP-Req-Nonce" // 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) +// POST under the given nonce so the dispatch middleware can bridge its values. +func (s *MCPServer) setPendingRequestContext(ctx context.Context, nonce string) { + s.pendingReqCtx.Store(nonce, 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 { +// pendingRequestContext returns the recorded HTTP request context for nonce. +func (s *MCPServer) pendingRequestContext(nonce string) context.Context { + if v, ok := s.pendingReqCtx.Load(nonce); ok { return v.(context.Context) } return nil } // clearPendingRequestContext drops the recorded HTTP request context. -func (s *MCPServer) clearPendingRequestContext(sessionID string) { - s.pendingReqCtx.Delete(sessionID) +func (s *MCPServer) clearPendingRequestContext(nonce string) { + s.pendingReqCtx.Delete(nonce) } // valueBridgeContext bridges the originating HTTP request's context values into -// a handler running on go-sdk's detached session goroutine. Its lifecycle +// a handler running on go-sdk's 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. +// The per-request context takes precedence so request-scoped values that the +// HTTP middleware chain re-establishes per request (audit BackendInfo, +// identity, telemetry) are visible to the handler rather than shadowed by the +// copies frozen onto the session context at initialize time. type valueBridgeContext struct { context.Context values context.Context @@ -196,22 +228,64 @@ func NewMCPServer(name, version string, opts ...ServerOption) *MCPServer { return s } +// serverCapabilities translates the mcp-go capability flags declared via +// WithToolCapabilities/WithResourceCapabilities/WithPromptCapabilities (and +// WithLogging) into a go-sdk ServerCapabilities value for ServerOptions. A +// capability declared via its option is advertised with its sub-flags +// (ListChanged/Subscribe); an undeclared capability is left nil so go-sdk's +// inference from registered features applies. See buildServer for why this +// mapping is required (per-session tool registration after initialize). +func (s *MCPServer) serverCapabilities() *gosdk.ServerCapabilities { + caps := &gosdk.ServerCapabilities{} + if s.toolsDeclared { + caps.Tools = &gosdk.ToolCapabilities{ListChanged: s.toolListChanged} + } + if s.resourcesDeclared { + caps.Resources = &gosdk.ResourceCapabilities{ + Subscribe: s.resourceSubscribe, + ListChanged: s.resourceListChanged, + } + } + if s.promptsDeclared { + caps.Prompts = &gosdk.PromptCapabilities{ListChanged: s.promptListChanged} + } + if s.logging { + caps.Logging = &gosdk.LoggingCapabilities{} + } + return caps +} + // WithToolCapabilities declares tool support (listChanged notifications). +// +// Invoking this option advertises the tools capability in the initialize result +// regardless of whether any tools are registered at initialize time: vMCP +// registers tools per-session AFTER initialize, so without this the capability +// would be absent (go-sdk otherwise infers capabilities from registered +// features). func WithToolCapabilities(listChanged bool) ServerOption { - return func(s *MCPServer) { s.toolListChanged = listChanged } + return func(s *MCPServer) { s.toolsDeclared = true; s.toolListChanged = listChanged } } // WithResourceCapabilities declares resource support. +// +// Invoking this option advertises the resources capability in the initialize +// result regardless of whether any resources are registered at initialize time +// (see WithToolCapabilities for the per-session registration rationale). func WithResourceCapabilities(subscribe, listChanged bool) ServerOption { return func(s *MCPServer) { + s.resourcesDeclared = true s.resourceSubscribe = subscribe s.resourceListChanged = listChanged } } // WithPromptCapabilities declares prompt support. +// +// Invoking this option advertises the prompts capability in the initialize +// result regardless of whether any prompts are registered at initialize time +// (see WithToolCapabilities for the per-session registration rationale). func WithPromptCapabilities(listChanged bool) ServerOption { - return func(s *MCPServer) { s.promptListChanged = listChanged } + return func(s *MCPServer) { s.promptsDeclared = true; s.promptListChanged = listChanged } } // WithLogging enables logging capability. @@ -325,6 +399,19 @@ func (s *MCPServer) buildServer(genSessionID func() string) (*gosdk.Server, erro s.registerAndSync(ctx, req.Session, srv) }, } + // Map the mcp-go capability flags to go-sdk's ServerCapabilities. go-sdk + // otherwise infers capabilities solely from registered features (see + // (*Server).capabilities): a server that registers tools per-session AFTER + // initialize — ToolHive's vMCP projection, where go-sdk has zero tools at + // initialize time — would advertise no tools capability, and spec-compliant + // clients that gate tools/list on capabilities.tools would see no tools. + // WithToolCapabilities/WithResourceCapabilities/WithPromptCapabilities + // declare those capabilities up front, mirroring mcp-go. The non-deprecated + // ServerOptions.Capabilities field is used (HasTools/HasResources/HasPrompts + // exist but are deprecated). Setting a capability to a non-nil value forces + // it to be advertised regardless of registered features; a nil entry leaves + // go-sdk's inference (from registered features) in place. + opts.Capabilities = s.serverCapabilities() // 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 @@ -391,10 +478,21 @@ func (s *MCPServer) sessionDispatchMiddleware(srv *gosdk.Server) gosdk.Middlewar 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} + } + // Bridge the originating HTTP request's context values (identity, + // audit BackendInfo, telemetry) into the handler context. go-sdk + // does not propagate the per-POST request context into the handler + // for existing sessions, so ServeHTTP stored it keyed by a per-POST + // nonce (X-MCP-Req-Nonce). The nonce is read off the per-request + // RequestExtra.Header — which go-sdk populates from the POST's + // headers — so concurrent POSTs on the same session each resolve + // their OWN context (issue #156, item U3: per-request, not per + // session). + if re := req.GetExtra(); re != nil { + if nonce := re.Header.Get(reqNonceHeader); nonce != "" { + if reqCtx := s.pendingRequestContext(nonce); reqCtx != nil { + ctx = &valueBridgeContext{Context: ctx, values: reqCtx} + } } } // Fire the before-hooks ahead of the SDK's own handling so a @@ -520,10 +618,13 @@ func (s *MCPServer) wrapPromptHandler(h PromptHandlerFunc) gosdk.PromptHandler { // 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. +// mcp-go passed RawInputSchema through verbatim, tolerating a missing/empty +// schema. go-sdk's AddTool panics unless InputSchema is non-nil (see +// normalizeObjectSchema), so a nil/empty schema is normalized to the empty +// object schema ({"type":"object"}). All other schemas pass through verbatim, +// matching mcp-go's behavior; go-sdk's AddTool may still panic on a non-object +// schema, which the per-session registration path recovers from (see +// normalizeObjectSchema and MCPServer.addSessionTool). func toGoSDKTool(t mcp.Tool) (*gosdk.Tool, error) { out := &gosdk.Tool{} if err := jsonConvert(t, out); err != nil { @@ -533,16 +634,56 @@ func toGoSDKTool(t mcp.Tool) (*gosdk.Tool, error) { 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. +// normalizeObjectSchema ensures a tool input schema is suitable for go-sdk's +// AddTool, which panics on a nil InputSchema. A nil schema, an empty map, an +// empty string, or a map whose "type" is the empty string (the value mcp-go's +// ToolInputSchema marshals to when no type is declared) is replaced with the +// empty object schema ({"type":"object"}). Every other schema is returned +// verbatim, mirroring mcp-go, which passed RawInputSchema through unchanged — +// including object schemas with "type" omitted, $ref, oneOf/anyOf, and boolean +// schemas. +// +// TODO(upstream): go-sdk v1.6.1 AddTool additionally panics unless the +// schema's top-level "type" is literally "object", so a non-object schema that +// passes through here (e.g. $ref, oneOf, a boolean schema, or an object schema +// with type omitted) will panic at registration time. We pass such schemas +// through verbatim (rather than silently rewriting them) so callers surface the +// incompatibility rather than silently receiving a different schema than the one +// they declared; mcp-go performed no validation here either. The two registration +// paths handle the panic differently: the per-session path +// (MCPServer.addSessionTool, used by syncSessionTools) recovers and skips the +// offending tool so one bad overlay schema cannot crash a live session; the +// global path (buildServer) lets the panic surface as a construction-time error +// (a 500 at handler build). func normalizeObjectSchema(schema any) any { - if m, ok := schema.(map[string]any); ok { - if m["type"] == "object" { - return m + switch s := schema.(type) { + case nil: + return map[string]any{schemaTypeKey: schemaTypeObject} + case map[string]any: + if len(s) == 0 { + return map[string]any{schemaTypeKey: schemaTypeObject} + } + // mcp-go's ToolInputSchema always marshals a "type" field, even when + // unset (it serializes as ""). An empty type string is the sentinel for + // "no schema declared"; normalize it to "object" so AddTool accepts the + // tool (matching mcp-go's leniency for tools with no declared input). + if t, ok := s[schemaTypeKey].(string); ok && t == "" { + cpy := make(map[string]any, len(s)) + for k, v := range s { + cpy[k] = v + } + cpy[schemaTypeKey] = schemaTypeObject + return cpy + } + return s + case string: + if s == "" { + return map[string]any{schemaTypeKey: schemaTypeObject} } + return s + default: + return s } - return map[string]any{"type": "object"} } // translateUnknownToolError rewrites go-sdk's "unknown tool" error for a diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index 186eb97..b659ace 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/mcpcompat/server/server_internal_test.go @@ -5,6 +5,7 @@ package server import ( "context" + "encoding/json" "testing" "github.com/stretchr/testify/assert" @@ -47,6 +48,56 @@ func TestClientSession_Store(t *testing.T) { assert.Contains(t, cs.GetSessionResources(), "file:///r") } +// TestSetSessionTools_NonObjectSchemaDoesNotPanic verifies that a per-session +// overlay tool whose RawInputSchema is a non-object schema ($ref) does NOT crash +// the session when synced onto the go-sdk server. go-sdk v1.6.1 AddTool panics +// unless the input schema's top-level type is "object"; normalizeObjectSchema +// passes non-object schemas through verbatim (to mirror mcp-go), so the panic +// must be recovered by MCPServer.addSessionTool and the offending tool skipped +// while other tools in the same overlay still register. Without the recover in +// syncSessionTools, this test panics mid-call. +func TestSetSessionTools_NonObjectSchemaDoesNotPanic(t *testing.T) { + t.Parallel() + s := NewMCPServer("s", "1") + srv, err := s.buildServer(nil) + require.NoError(t, err) + + cs := s.sessionFor("sid-bad-schema") + cs.SetSessionTools(map[string]ServerTool{ + // A non-object schema (a $ref). normalizeObjectSchema returns this + // verbatim; go-sdk AddTool panics on it. addSessionTool must recover. + "bad-schema": { + Tool: mcp.Tool{ + Name: "bad-schema", + RawInputSchema: json.RawMessage(`{"$ref":"#"}`), + }, + Handler: func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, nil }, + }, + // A well-formed object-schema tool in the same overlay. It must still + // register despite the sibling's bad schema. + "good-schema": { + Tool: mcp.NewTool("good-schema", mcp.WithDescription("ok")), + Handler: func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("ok"), nil + }, + }, + }) + + // Must not panic (go-sdk AddTool would panic on the $ref schema without + // the recover in addSessionTool). + assert.NotPanics(t, func() { s.syncSessionTools(srv, cs) }) + + // The offending tool is skipped; the well-formed sibling is registered. + cs.mu.RLock() + registered := make(map[string]struct{}, len(cs.sdkToolNames)) + for n := range cs.sdkToolNames { + registered[n] = struct{}{} + } + cs.mu.RUnlock() + assert.NotContains(t, registered, "bad-schema", "the non-object-schema tool must be skipped, not registered") + assert.Contains(t, registered, "good-schema", "the well-formed sibling tool must still register") +} + func TestHooks_Fire(t *testing.T) { t.Parallel() h := &Hooks{} @@ -126,3 +177,106 @@ func (fakeIDManager) Validate(string) (bool, error) { return false, nil } func (fakeIDManager) Terminate(string) (bool, error) { return false, nil } var _ SessionIdManager = fakeIDManager{} + +// TestNormalizeObjectSchema verifies that only nil/empty input schemas are +// normalized to {"type":"object"}; all other schemas pass through verbatim, +// matching mcp-go (which passed RawInputSchema through unchanged). Previously +// any schema whose top-level "type" was not "object" was clobbered, stripping +// valid $ref/oneOf/boolean/object-with-type-omitted schemas. +// +//nolint:goconst // test fixtures legitimately repeat schema keys +func TestNormalizeObjectSchema(t *testing.T) { + t.Parallel() + emptyObject := map[string]any{schemaTypeKey: schemaTypeObject} + + tests := []struct { + name string + input any + want any + }{ + { + name: "nil", + input: nil, + want: emptyObject, + }, + { + name: "empty map", + input: map[string]any{}, + want: emptyObject, + }, + { + name: "type is the empty string (mcp-go wire sentinel)", + input: map[string]any{schemaTypeKey: ""}, + want: emptyObject, + }, + { + name: "type empty with properties (mcp-go wire sentinel)", + input: map[string]any{schemaTypeKey: "", "properties": map[string]any{}}, + want: map[string]any{ + schemaTypeKey: schemaTypeObject, + "properties": map[string]any{}, + }, + }, + { + name: "empty string", + input: "", + want: emptyObject, + }, + { + name: "object schema with type object", + input: map[string]any{ + schemaTypeKey: schemaTypeObject, + "properties": map[string]any{ + "name": map[string]any{schemaTypeKey: "string"}, + }, + }, + want: map[string]any{ + schemaTypeKey: schemaTypeObject, + "properties": map[string]any{ + "name": map[string]any{schemaTypeKey: "string"}, + }, + }, + }, + { + name: "object schema with type omitted", + input: map[string]any{ + "properties": map[string]any{ + "name": map[string]any{schemaTypeKey: "string"}, + }, + }, + want: map[string]any{ + "properties": map[string]any{ + "name": map[string]any{schemaTypeKey: "string"}, + }, + }, + }, + { + name: "$ref schema", + input: map[string]any{"$ref": "#"}, + want: map[string]any{"$ref": "#"}, + }, + { + name: "oneOf schema", + input: map[string]any{"oneOf": []any{map[string]any{schemaTypeKey: "string"}, map[string]any{schemaTypeKey: "number"}}}, + want: map[string]any{"oneOf": []any{map[string]any{schemaTypeKey: "string"}, map[string]any{schemaTypeKey: "number"}}}, + }, + { + name: "boolean true schema", + input: true, + want: true, + }, + { + name: "boolean false schema", + input: false, + want: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := normalizeObjectSchema(tc.input) + assert.Equal(t, tc.want, got) + }) + } +} diff --git a/mcpcompat/server/server_test.go b/mcpcompat/server/server_test.go index 59d8013..6e38484 100644 --- a/mcpcompat/server/server_test.go +++ b/mcpcompat/server/server_test.go @@ -5,17 +5,27 @@ package server_test import ( "context" + "net/http" "net/http/httptest" + "sync" "testing" + "time" "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" "github.com/stacklok/toolhive-core/mcpcompat/server" ) +// testClientVersion and testClientName are reused across the e2e tests below. +const ( + testClientName = "test-client" + testClientVersion = "1.0.0" +) + // 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. @@ -55,7 +65,7 @@ func TestGlobalServer_EndToEnd(t *testing.T) { 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"}, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, }, }) require.NoError(t, err) @@ -85,6 +95,50 @@ func TestGlobalServer_EndToEnd(t *testing.T) { } } +// TestGlobalServer_NilSchemaToolAdvertisesObject verifies that a globally +// registered tool with no declared input schema (mcp-go's leniency) is served +// over Streamable HTTP with its input schema normalized to {"type":"object"}, +// which is what go-sdk's AddTool requires. go-sdk would otherwise panic on a +// nil/empty schema; normalizeObjectSchema rewrites it so the tool registers +// and advertises the empty object schema on the wire. +func TestGlobalServer_NilSchemaToolAdvertisesObject(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := server.NewMCPServer("schema-server", testClientVersion, server.WithToolCapabilities(false)) + // A tool with neither InputSchema nor RawInputSchema set: mcp-go tolerated + // this; the shim must normalize it to {"type":"object"}. + srv.AddTool(mcp.Tool{Name: "no-schema"}, + 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)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + tools, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + require.Len(t, tools.Tools, 1) + assert.Equal(t, "no-schema", tools.Tools[0].Name) + // The advertised input schema must be the normalized empty object schema. + assert.Equal(t, "object", tools.Tools[0].InputSchema.Type) + assert.Empty(t, tools.Tools[0].InputSchema.Properties) +} + // 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). @@ -111,3 +165,161 @@ func TestServer_RegistrationSurface(t *testing.T) { srv.DeleteTools("t") } + +// TestAdvertisesToolsCapability_WithNoGlobalTools verifies that calling +// WithToolCapabilities forces the tools capability to be advertised in the +// initialize result even when NO global tools are registered at build time. +// This matters for ToolHive's vMCP projection, which registers tools +// per-session AFTER initialize: go-sdk otherwise infers capabilities only from +// registered features and would advertise no tools capability, so +// spec-compliant clients that gate tools/list on capabilities.tools would see +// no tools. Without the fix (mapping capability flags to ServerOptions), the +// Tools field is absent. +func TestAdvertisesToolsCapability_WithNoGlobalTools(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := server.NewMCPServer("cap-server", testClientVersion, + server.WithToolCapabilities(false), + server.WithResourceCapabilities(false, false), + server.WithPromptCapabilities(false), + ) + // Deliberately register NO tools/resources/prompts globally. + + 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: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + require.NotNil(t, initRes.Capabilities, "server must advertise capabilities") + assert.NotNil(t, initRes.Capabilities.Tools, "capabilities.tools must be advertised when WithToolCapabilities is set") + assert.NotNil(t, initRes.Capabilities.Resources, "capabilities.resources must be advertised when WithResourceCapabilities is set") + assert.NotNil(t, initRes.Capabilities.Prompts, "capabilities.prompts must be advertised when WithPromptCapabilities is set") + + require.NoError(t, c.Close()) +} + +// reqIDKey is a per-request context value used by +// TestPerRequestContext_NotClobberedConcurrently. The client sets it via a +// header (read by the server's WithHTTPContextFunc) so the value flows from the +// call's context, through an HTTP header, back into the handler context via +// go-sdk's request-context propagation. +type reqIDKey struct{} + +// TestPerRequestContext_NotClobberedConcurrently verifies that per-request +// context values are NOT keyed by session ID: two concurrent tools/call POSTs +// on the SAME session, each carrying a distinct request-scoped value, must each +// see their OWN value in the tool handler. Before the fix, the shim stored a +// single context per session (pendingReqCtx[sessionID]), so concurrent POSTs +// clobbered each other and one handler intermittently saw the other's value. +// +// The per-request value is injected by the server's WithHTTPContextFunc from an +// X-Request-Id header, and the client sets that header per call via its +// per-request HTTPHeaderFunc. go-sdk does NOT propagate the per-POST HTTP +// request's context into handlers for requests on an existing session (it +// handles them on the session's connection goroutine using the initialize-time +// context), so the value reaches the handler via the nonce bridge: ServeHTTP +// stores the request context keyed by a per-POST X-MCP-Req-Nonce header, the +// dispatch middleware reads that header off the per-request RequestExtra and +// bridges the stored context into the handler via valueBridgeContext. Without +// the bridge, handlers see empty/clobbered values (verified by short-circuiting +// the bridge — the test then fails with handlers recording ""). +func TestPerRequestContext_NotClobberedConcurrently(t *testing.T) { + t.Parallel() + + const ( + toolName = "echo-req" + headerName = "X-Request-Id" + ) + + var seen struct { + sync.Mutex + values []string + } + srv := server.NewMCPServer("ctx-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool( + mcp.NewTool(toolName, mcp.WithDescription("echo request id")), + func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + v, _ := ctx.Value(reqIDKey{}).(string) + seen.Lock() + seen.values = append(seen.values, v) + seen.Unlock() + // Block briefly to widen the concurrency window and make a race + // observable when the value is keyed by session. + time.Sleep(20 * time.Millisecond) + return mcp.NewToolResultText(v), nil + }, + ) + + httpSrv := server.NewStreamableHTTPServer(srv, + server.WithHTTPContextFunc(func(ctx context.Context, r *http.Request) context.Context { + if id := r.Header.Get(headerName); id != "" { + return context.WithValue(ctx, reqIDKey{}, id) + } + return ctx + }), + ) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + c, err := client.NewStreamableHttpClient(ts.URL, + transport.WithHTTPHeaderFunc(func(ctx context.Context) map[string]string { + if v, ok := ctx.Value(reqIDKey{}).(string); ok { + return map[string]string{headerName: v} + } + return nil + }), + ) + require.NoError(t, err) + require.NoError(t, c.Start(context.Background())) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(context.Background(), mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + const iterations = 100 + for i := 0; i < iterations; i++ { + var wg sync.WaitGroup + var a, b = "A", "B" + wg.Add(2) + go func() { + defer wg.Done() + ctx := context.WithValue(context.Background(), reqIDKey{}, a) + _, err := c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: toolName}}) + assert.NoError(t, err) + }() + go func() { + defer wg.Done() + ctx := context.WithValue(context.Background(), reqIDKey{}, b) + _, err := c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: toolName}}) + assert.NoError(t, err) + }() + wg.Wait() + + seen.Lock() + got := append([]string(nil), seen.values...) + seen.values = seen.values[:0] + seen.Unlock() + + require.Len(t, got, 2, "iteration %d: both handlers must record a value", i) + // Each handler must have seen its OWN value; if the per-request context + // were keyed by session, one value would appear twice and the other zero + // times. + assert.ElementsMatch(t, []string{a, b}, got, "iteration %d: handlers saw clobbered values %v", i, got) + } +} diff --git a/mcpcompat/server/session.go b/mcpcompat/server/session.go index a6e0350..73dfc48 100644 --- a/mcpcompat/server/session.go +++ b/mcpcompat/server/session.go @@ -214,6 +214,17 @@ func (s *MCPServer) registerAndSync(ctx context.Context, ss *gosdk.ServerSession // 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. +// +// go-sdk's AddTool panics unless the tool's input schema is non-nil with +// top-level type "object". normalizeObjectSchema only normalizes nil/empty +// schemas to {"type":"object"}; non-object schemas ($ref, oneOf, boolean, or an +// object schema with type omitted) pass through verbatim to mirror mcp-go, so +// AddTool can still panic here for a malformed overlay schema. This path runs on +// the live session goroutine (from before-hooks during a request), so an +// unrecovered panic would crash the session/process mid-request. Each tool is +// therefore registered via addSessionTool, which recovers such a panic, logs it, +// and skips the offending tool — one bad overlay schema does not take down the +// session, and the remaining tools in the same overlay still register. func (s *MCPServer) syncSessionTools(srv *gosdk.Server, cs *clientSession) { cs.mu.Lock() defer cs.mu.Unlock() @@ -221,9 +232,17 @@ func (s *MCPServer) syncSessionTools(srv *gosdk.Server, cs *clientSession) { for name, st := range cs.tools { gt, err := toGoSDKTool(st.Tool) if err != nil { + if s.logger != nil { + s.logger.Warn("skipping per-session tool: conversion failed", "tool", name, "error", err) + } + continue + } + if !s.addSessionTool(srv, gt, st.Handler) { + if s.logger != nil { + s.logger.Warn("skipping per-session tool: AddTool rejected its schema", "tool", name) + } continue } - srv.AddTool(gt, s.wrapToolHandler(st.Handler)) newNames[name] = struct{}{} } var removed []string @@ -238,6 +257,26 @@ func (s *MCPServer) syncSessionTools(srv *gosdk.Server, cs *clientSession) { cs.sdkToolNames = newNames } +// addSessionTool registers a tool on a per-session go-sdk server, recovering the +// panic go-sdk's AddTool raises when a tool's input schema is non-nil but not +// type:"object" (see normalizeObjectSchema). It returns false (and logs) when the +// tool could not be registered, so the caller skips it rather than crashing the +// session. This mirrors the fault-isolation recover pattern used in +// wrapToolHandler (server.go). +func (s *MCPServer) addSessionTool(srv *gosdk.Server, gt *gosdk.Tool, h ToolHandlerFunc) (ok bool) { + defer func() { + if r := recover(); r != nil { + if s.logger != nil { + s.logger.Error("per-session AddTool panicked; skipping tool", + "tool", gt.Name, "panic", r) + } + ok = false + } + }() + srv.AddTool(gt, s.wrapToolHandler(h)) + return true +} + // 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) { @@ -265,7 +304,6 @@ func (s *MCPServer) isLocalSession(id string) bool { 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 diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index 39e06dc..25ab92f 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -5,6 +5,7 @@ package server import ( "context" + crand "crypto/rand" "fmt" "mime" "net/http" @@ -144,6 +145,12 @@ func WithHeartbeatInterval(interval time.Duration) StreamableHTTPOption { } // WithHTTPContextFunc installs a per-request context customizer. +// +// Context values injected here are applied to ALL POSTs, including the +// initialize request (previously the nonce bridge only bridged non-initialize +// POSTs; the contextFunc itself now runs unconditionally on every request). +// This is harmless/desired but observable: an initialize handler that reads a +// context value populated by contextFunc will now see it. func WithHTTPContextFunc(fn HTTPContextFunc) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.contextFunc = fn } } @@ -182,16 +189,22 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) 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) - } + // Bridge per-request context values into the handler. go-sdk does not + // propagate the per-POST request context into handlers for existing + // sessions (it handles messages on the session's connection goroutine using + // the initialize-time context), so values added by contextFunc (identity, + // audit BackendInfo, telemetry) would be lost. Store the request context + // keyed by a per-POST nonce and stamp the nonce as a header; the dispatch + // middleware reads it back via req.GetExtra().Header to bridge the values. + // Keying by nonce (not session ID) avoids the race where concurrent POSTs + // on the same session clobber each other's context (issue #156, item U3). + // The initialize POST (no session yet) carries no per-request values the + // handler needs; only non-initialize POSTs need bridging, but bridging + // initialize is harmless. + nonce := crand.Text() + s.mcp.setPendingRequestContext(r.Context(), nonce) + r.Header.Set(reqNonceHeader, nonce) + defer s.mcp.clearPendingRequestContext(nonce) // 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 From a3020896decc03b4a221d10c8b5c3de536e9ea8b Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Wed, 8 Jul 2026 22:52:49 +0300 Subject: [PATCH 2/8] fix(mcpcompat): wire client notification handlers and carry params (issue #156 wave 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1a + U8 — Register go-sdk ProgressNotificationHandler and LoggingMessageHandler in the mcpcompat client; route to the compat OnNotification callback. dispatch() now carries notification params via toNotificationParams (JSON round-trip, _meta split), restoring progress and logging notifications end-to-end with their params. 1b — Confirmed go-sdk auto-emits tools/list_changed on SetSessionTools (no production code needed beyond U1); corrected the stale doc comments in notifications.go that claimed list_changed is dropped. Added SetLoggingLevel + mcp.LoggingLevel aliases (required for the logging-notification path; renamed/simplified vs mcp-go's SetLevel, documented). Tests: progress+logging e2e with param assertions, list_changed e2e, toNotificationParams unit table (meta split, error paths, nil guards), handler nil-param guards, method-string coverage for all three list_changed variants, broadcast delivery assertion, early-exit on list_changed. Refs #156 --- mcpcompat/client/client.go | 113 ++++++- mcpcompat/client/client_test.go | 142 +++++++- .../client/notifications_internal_test.go | 305 ++++++++++++++++++ mcpcompat/mcp/alias.go | 21 ++ mcpcompat/server/notifications.go | 20 +- mcpcompat/server/notifications_test.go | 157 ++++++++- 6 files changed, 742 insertions(+), 16 deletions(-) create mode 100644 mcpcompat/client/notifications_internal_test.go diff --git a/mcpcompat/client/client.go b/mcpcompat/client/client.go index 6e5d5c4..fd02d0f 100644 --- a/mcpcompat/client/client.go +++ b/mcpcompat/client/client.go @@ -173,6 +173,23 @@ func (c *Client) Ping(ctx context.Context) error { return s.Ping(ctx, nil) } +// SetLoggingLevel sets the server's logging level. This is a renamed and +// simplified counterpart to mcp-go's client.Client.SetLevel: rather than +// taking an mcp.SetLevelRequest, it accepts the mcp.LoggingLevel directly. +// The SetLevel, SetLevelRequest and SetLevelParams types from mcp-go are +// intentionally NOT re-exported by this shim. The server only delivers +// notifications/message notifications at or above the requested level, so this +// must be called before the OnNotification handler will see logging +// notifications. level is one of the MCP logging levels ("debug", "info", +// "notice", "warning", "error", "critical", "alert", "emergency"). +func (c *Client) SetLoggingLevel(ctx context.Context, level mcp.LoggingLevel) error { + s, err := c.sessionFor() + if err != nil { + return err + } + return s.SetLoggingLevel(ctx, &gosdk.SetLoggingLevelParams{Level: gosdk.LoggingLevel(level)}) +} + // ListTools lists the server's tools. func (c *Client) ListTools(ctx context.Context, request mcp.ListToolsRequest) (*mcp.ListToolsResult, error) { ctx = withErrCapture(ctx) @@ -395,34 +412,120 @@ func (c *Client) ListResourceTemplates( // 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. +// values (method + params) for the list-changed, progress and logging +// notifications, forwarding each to every registered handler. +// +// Server-initiated notifications (including the list_changed notifications, +// which the server emits outside any in-flight request) are only delivered if +// the client enabled continuous listening via +// transport.WithContinuousListening(). Without it the go-sdk streamable +// transport has no standalone SSE stream to carry such notifications and they +// are silently dropped; no callback fires. 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) { +// dispatch fans a synthesized JSONRPCNotification out to every registered +// OnNotification handler. params is converted into mcp-go's +// JSONRPCNotificationParams (Meta + AdditionalFields) when non-nil, so the +// notification carries its params as mcp-go would; nil leaves the params empty +// (used by the list_changed notifications, which carry no params on the wire). +func (c *Client) dispatch(method string, params any) { 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 + if params != nil { + n.Params = toNotificationParams(params) + } for _, h := range handlers { h(n) } } +// toNotificationParams converts a go-sdk notification params value into the +// mcp-go NotificationParams shape (Meta + AdditionalFields). Both encode to the +// same MCP wire JSON, so a JSON round-trip faithfully maps the go-sdk struct's +// fields onto mcp-go's AdditionalFields map. This mirrors the jsonConvert +// convention used elsewhere in the shim for cross-type conversion. +func toNotificationParams(src any) mcp.NotificationParams { + // Marshal the go-sdk params and unmarshal into a generic map, then split the + // reserved "_meta" key (if present) from the remaining fields, which become + // AdditionalFields. + b, err := json.Marshal(src) + if err != nil { + return mcp.NotificationParams{} + } + var m map[string]any + if err := json.Unmarshal(b, &m); err != nil { + return mcp.NotificationParams{} + } + out := mcp.NotificationParams{AdditionalFields: make(map[string]any, len(m))} + for k, v := range m { + if k == "_meta" { + if mm, ok := v.(map[string]any); ok { + out.Meta = mm + } + continue + } + out.AdditionalFields[k] = v + } + return out +} + func (c *Client) installNotificationHandlers(opts *gosdk.ClientOptions) { opts.ToolListChangedHandler = func(_ context.Context, _ *gosdk.ToolListChangedRequest) { - c.dispatch("notifications/tools/list_changed") + c.dispatch("notifications/tools/list_changed", nil) } opts.PromptListChangedHandler = func(_ context.Context, _ *gosdk.PromptListChangedRequest) { - c.dispatch("notifications/prompts/list_changed") + c.dispatch("notifications/prompts/list_changed", nil) } opts.ResourceListChangedHandler = func(_ context.Context, _ *gosdk.ResourceListChangedRequest) { - c.dispatch("notifications/resources/list_changed") + c.dispatch("notifications/resources/list_changed", nil) + } + // Forward server->client progress notifications. go-sdk hands the params via + // *ProgressNotificationClientRequest; convert them to mcp-go's notification + // params (progressToken, progress, total, message) and dispatch. + opts.ProgressNotificationHandler = newProgressNotificationHandler(c.dispatch) + // Forward server->client logging notifications. go-sdk hands the params via + // *LoggingMessageRequest; convert them to mcp-go's notification params + // (level, data, logger) and dispatch. + opts.LoggingMessageHandler = newLoggingMessageHandler(c.dispatch) +} + +// dispatchFunc is the signature of Client.dispatch, factored out so the +// notification handlers can be unit-tested in isolation without a live session. +type dispatchFunc func(method string, params any) + +// newProgressNotificationHandler builds the go-sdk progress-notification +// handler that forwards notifications/progress to the supplied dispatch func. +// If the go-sdk request or its params are nil, dispatch is called with nil +// params (an empty NotificationParams) rather than panicking. +func newProgressNotificationHandler(dispatch dispatchFunc) func(context.Context, *gosdk.ProgressNotificationClientRequest) { + return func(_ context.Context, req *gosdk.ProgressNotificationClientRequest) { + if req == nil || req.Params == nil { + dispatch("notifications/progress", nil) + return + } + dispatch("notifications/progress", req.Params) + } +} + +// newLoggingMessageHandler builds the go-sdk logging-notification handler that +// forwards notifications/message to the supplied dispatch func. If the go-sdk +// request or its params are nil, dispatch is called with nil params rather than +// panicking. +func newLoggingMessageHandler(dispatch dispatchFunc) func(context.Context, *gosdk.LoggingMessageRequest) { + return func(_ context.Context, req *gosdk.LoggingMessageRequest) { + if req == nil || req.Params == nil { + dispatch("notifications/message", nil) + return + } + dispatch("notifications/message", req.Params) } } diff --git a/mcpcompat/client/client_test.go b/mcpcompat/client/client_test.go index 8cb0552..120d3d0 100644 --- a/mcpcompat/client/client_test.go +++ b/mcpcompat/client/client_test.go @@ -7,7 +7,9 @@ import ( "context" "net/http" "net/http/httptest" + "sync" "testing" + "time" gosdk "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/assert" @@ -25,6 +27,9 @@ type echoInput struct { // echoToolName is the shared tool name used across the client tests. const echoToolName = "echo" +// testClientVersion is the shared client/server version string used in tests. +const testClientVersion = "1.0.0" + // 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 { @@ -59,7 +64,7 @@ func TestStreamableClient_EndToEnd(t *testing.T) { 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"}, + ClientInfo: mcp.Implementation{Name: "test-client", Version: testClientVersion}, }, }) require.NoError(t, err) @@ -114,3 +119,138 @@ func TestGetTransport_SSEIsNil(t *testing.T) { assert.False(t, ok) assert.Empty(t, c.GetSessionId()) } + +// newProgressLoggingServer stands up a real go-sdk MCP server whose "notify" +// tool handler emits both a progress notification and a log message on the +// calling session. The client must register the progress token it wants +// echoed back via the call's _meta; the server reads it off the request _meta. +func newProgressLoggingServer(t *testing.T) *httptest.Server { + t.Helper() + srv := gosdk.NewServer( + &gosdk.Implementation{Name: "notify-server", Version: testClientVersion}, + &gosdk.ServerOptions{Capabilities: &gosdk.ServerCapabilities{Logging: &gosdk.LoggingCapabilities{}}}, + ) + srv.AddTool(&gosdk.Tool{ + Name: "notify", + Description: "emit progress + log", + InputSchema: map[string]any{"type": "object"}, + }, + func(ctx context.Context, req *gosdk.CallToolRequest) (*gosdk.CallToolResult, error) { + // Pull the progress token the client sent in _meta and echo it back + // on a notifications/progress. If absent, use a fixed token. + token := any("fallback-token") + if len(req.Params.Meta) > 0 { + if v, ok := req.Params.Meta["progressToken"]; ok { + token = v + } + } + _ = req.Session.NotifyProgress(ctx, &gosdk.ProgressNotificationParams{ + ProgressToken: token, + Progress: 0.5, + Total: 1.0, + Message: "halfway", + }) + _ = req.Session.Log(ctx, &gosdk.LoggingMessageParams{ + Level: gosdk.LoggingLevel("info"), + Data: "hello from server", + }) + return &gosdk.CallToolResult{ + Content: []gosdk.Content{&gosdk.TextContent{Text: "notified"}}, + }, nil + }) + handler := gosdk.NewStreamableHTTPHandler(func(*http.Request) *gosdk.Server { return srv }, nil) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// TestOnNotification_ProgressAndLogging verifies the shim wires the go-sdk +// ProgressNotificationHandler and LoggingMessageHandler so server->client +// notifications/progress and notifications/message reach the registered +// OnNotification callback with their params (progressToken/progress and +// level/data). Without the fix, neither handler is registered and the callback +// never fires. +func TestOnNotification_ProgressAndLogging(t *testing.T) { + t.Parallel() + ctx := context.Background() + ts := newProgressLoggingServer(t) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + type recordedNotif struct { + method string + params mcp.NotificationParams + } + var ( + mu sync.Mutex + got []recordedNotif + progChan = make(chan recordedNotif, 1) + logChan = make(chan recordedNotif, 1) + ) + c.OnNotification(func(n mcp.JSONRPCNotification) { + mu.Lock() + got = append(got, recordedNotif{method: n.Method, params: n.Params}) + mu.Unlock() + switch n.Method { + case "notifications/progress": + select { + case progChan <- recordedNotif{method: n.Method, params: n.Params}: + default: + } + case "notifications/message": + select { + case logChan <- recordedNotif{method: n.Method, params: n.Params}: + default: + } + } + }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "test-client", Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + // The server only delivers notifications/message once the client has set a + // logging level, so raise it before invoking the tool. + require.NoError(t, c.SetLoggingLevel(ctx, mcp.LoggingLevelInfo)) + + const token = "tok-123" + _, err = c.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "notify", + Meta: mcp.NewMetaFromMap(map[string]any{"progressToken": token}), + }, + }) + require.NoError(t, err) + + // Both notifications should arrive shortly; they are delivered on the + // session's connection goroutine. + select { + case p := <-progChan: + assert.Equal(t, "notifications/progress", p.method) + assert.Equal(t, token, p.params.AdditionalFields["progressToken"]) + assert.InDelta(t, 0.5, p.params.AdditionalFields["progress"], 1e-9) + assert.InDelta(t, 1.0, p.params.AdditionalFields["total"], 1e-9) + assert.Equal(t, "halfway", p.params.AdditionalFields["message"]) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for notifications/progress") + } + select { + case l := <-logChan: + assert.Equal(t, "notifications/message", l.method) + assert.Equal(t, "info", l.params.AdditionalFields["level"]) + assert.Equal(t, "hello from server", l.params.AdditionalFields["data"]) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for notifications/message") + } + + mu.Lock() + defer mu.Unlock() + require.GreaterOrEqual(t, len(got), 2, "at least progress and logging notifications expected") +} diff --git a/mcpcompat/client/notifications_internal_test.go b/mcpcompat/client/notifications_internal_test.go new file mode 100644 index 0000000..3b389f1 --- /dev/null +++ b/mcpcompat/client/notifications_internal_test.go @@ -0,0 +1,305 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "sync" + "testing" + + gosdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// TestToNotificationParams covers the conversion helper in isolation: the JSON +// round-trip, the reserved "_meta" key extraction into Meta, and the silent +// dropping of a non-map _meta. It is otherwise only exercised by e2e tests. +func TestToNotificationParams(t *testing.T) { + t.Parallel() + + // fooKey is the JSON key used across the test structs and assertions. + const fooKey = "foo" + + type scalarsOnly struct { + // Only top-level scalar fields; no _meta. + Foo string `json:"foo"` + Bar int `json:"bar"` + } + + type withMeta struct { + Foo string `json:"foo"` + Meta map[string]any `json:"_meta"` + } + + type withNonMapMeta struct { + Foo string `json:"foo"` + Meta string `json:"_meta"` + } + + type unmarshalable struct { + Ch chan struct{} `json:"ch"` + } + + tests := []struct { + name string + src any + wantMeta map[string]any // nil means expect Meta unset (nil). + wantExtra map[string]any // expected AdditionalFields (subset checked). + }{ + { + name: "scalar fields land in AdditionalFields, Meta empty", + src: scalarsOnly{Foo: "x", Bar: 7}, + wantExtra: map[string]any{fooKey: "x", "bar": float64(7)}, + }, + { + name: "populated _meta lands in Meta, not AdditionalFields", + src: withMeta{ + Foo: "x", + Meta: map[string]any{"trace-id": "abc"}, + }, + wantMeta: map[string]any{"trace-id": "abc"}, + wantExtra: map[string]any{fooKey: "x"}, + }, + { + name: "non-map _meta silently dropped, other fields preserved", + src: withNonMapMeta{Foo: "x", Meta: "not-a-map"}, + wantExtra: map[string]any{fooKey: "x"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + out := toNotificationParams(tc.src) + + require.NotNil(t, out.AdditionalFields, "AdditionalFields map must be initialized") + for k, v := range tc.wantExtra { + assert.Equal(t, v, out.AdditionalFields[k], "AdditionalFields[%q]", k) + } + // _meta must never leak into AdditionalFields. + _, hasMetaKey := out.AdditionalFields["_meta"] + assert.False(t, hasMetaKey, "_meta must not appear in AdditionalFields") + + if tc.wantMeta == nil { + assert.Nil(t, out.Meta, "Meta should be unset/nil") + } else { + assert.Equal(t, tc.wantMeta, out.Meta, "Meta mismatch") + } + }) + } + + // Nil input: marshal of nil yields "null", unmarshal into a map leaves it + // nil, so the loop body never runs and we get an empty (but non-nil + // AdditionalFields) NotificationParams. Assert no panic and an empty result. + t.Run("nil input returns empty params without panic", func(t *testing.T) { + t.Parallel() + out := toNotificationParams(nil) + assert.NotNil(t, out.AdditionalFields, "AdditionalFields map still initialized") + assert.Empty(t, out.AdditionalFields) + assert.Nil(t, out.Meta) + }) + + // Unmarshalable input (channel) makes json.Marshal fail; the helper must + // swallow the error and return the zero-value NotificationParams (nil + // AdditionalFields, nil Meta), not panic. + t.Run("marshal error swallowed, empty params returned", func(t *testing.T) { + t.Parallel() + out := toNotificationParams(unmarshalable{Ch: make(chan struct{})}) + assert.Nil(t, out.AdditionalFields) + assert.Nil(t, out.Meta) + }) +} + +// TestProgressNotificationHandler_NilGuards exercises the nil-req / nil-params +// branches of the progress notification handler directly, asserting no panic +// and that dispatch is invoked with nil params (matching the guard behavior). +func TestProgressNotificationHandler_NilGuards(t *testing.T) { + t.Parallel() + + type dispatched struct { + method string + params any + } + + // newRecorder builds a dispatch func that appends to a fresh slice guarded + // by a mutex, returning a snapshot accessor. Each subtest gets its own + // recorder so the subtests are independent and can run in parallel. + newRecorder := func() (func(string, any), func() []dispatched) { + var ( + mu sync.Mutex + got []dispatched + ) + dispatch := func(method string, params any) { + mu.Lock() + got = append(got, dispatched{method: method, params: params}) + mu.Unlock() + } + snapshot := func() []dispatched { + mu.Lock() + defer mu.Unlock() + out := make([]dispatched, len(got)) + copy(out, got) + return out + } + return dispatch, snapshot + } + hOf := func(dispatch func(string, any)) func(context.Context, *gosdk.ProgressNotificationClientRequest) { + return newProgressNotificationHandler(dispatch) + } + + t.Run("nil request dispatches progress with nil params", func(t *testing.T) { + t.Parallel() + dispatch, snapshot := newRecorder() + h := hOf(dispatch) + assert.NotPanics(t, func() { h(context.Background(), nil) }) + got := snapshot() + require.Len(t, got, 1) + assert.Equal(t, "notifications/progress", got[0].method) + assert.Nil(t, got[0].params) + }) + + t.Run("nil params dispatches progress with nil params", func(t *testing.T) { + t.Parallel() + dispatch, snapshot := newRecorder() + h := hOf(dispatch) + assert.NotPanics(t, func() { + h(context.Background(), &gosdk.ProgressNotificationClientRequest{Params: nil}) + }) + got := snapshot() + require.Len(t, got, 1) + assert.Equal(t, "notifications/progress", got[0].method) + assert.Nil(t, got[0].params) + }) + + t.Run("populated params dispatched through", func(t *testing.T) { + t.Parallel() + dispatch, snapshot := newRecorder() + h := hOf(dispatch) + assert.NotPanics(t, func() { + h(context.Background(), &gosdk.ProgressNotificationClientRequest{ + Params: &gosdk.ProgressNotificationParams{ProgressToken: "t", Progress: 0.5}, + }) + }) + got := snapshot() + require.Len(t, got, 1) + assert.Equal(t, "notifications/progress", got[0].method) + assert.NotNil(t, got[0].params) + }) +} + +// TestLoggingMessageHandler_NilGuards exercises the nil-req / nil-params +// branches of the logging notification handler directly, asserting no panic +// and that dispatch is invoked with nil params. +func TestLoggingMessageHandler_NilGuards(t *testing.T) { + t.Parallel() + + type dispatched struct { + method string + params any + } + newRecorder := func() (func(string, any), func() []dispatched) { + var ( + mu sync.Mutex + got []dispatched + ) + dispatch := func(method string, params any) { + mu.Lock() + got = append(got, dispatched{method: method, params: params}) + mu.Unlock() + } + snapshot := func() []dispatched { + mu.Lock() + defer mu.Unlock() + out := make([]dispatched, len(got)) + copy(out, got) + return out + } + return dispatch, snapshot + } + hOf := func(dispatch func(string, any)) func(context.Context, *gosdk.LoggingMessageRequest) { + return newLoggingMessageHandler(dispatch) + } + + t.Run("nil request dispatches message with nil params", func(t *testing.T) { + t.Parallel() + dispatch, snapshot := newRecorder() + h := hOf(dispatch) + assert.NotPanics(t, func() { h(context.Background(), nil) }) + got := snapshot() + require.Len(t, got, 1) + assert.Equal(t, "notifications/message", got[0].method) + assert.Nil(t, got[0].params) + }) + + t.Run("nil params dispatches message with nil params", func(t *testing.T) { + t.Parallel() + dispatch, snapshot := newRecorder() + h := hOf(dispatch) + assert.NotPanics(t, func() { + h(context.Background(), &gosdk.LoggingMessageRequest{Params: nil}) + }) + got := snapshot() + require.Len(t, got, 1) + assert.Equal(t, "notifications/message", got[0].method) + assert.Nil(t, got[0].params) + }) + + t.Run("populated params dispatched through", func(t *testing.T) { + t.Parallel() + dispatch, snapshot := newRecorder() + h := hOf(dispatch) + assert.NotPanics(t, func() { + h(context.Background(), &gosdk.LoggingMessageRequest{ + Params: &gosdk.LoggingMessageParams{Level: gosdk.LoggingLevel("info"), Data: "hi"}, + }) + }) + got := snapshot() + require.Len(t, got, 1) + assert.Equal(t, "notifications/message", got[0].method) + assert.NotNil(t, got[0].params) + }) +} + +// TestListChangedHandlers_MethodStrings guards against a typo in the three +// list-changed method strings wired up by installNotificationHandlers. It +// installs the handlers onto a real (unconnected) Client, invokes each go-sdk +// handler directly, and asserts the dispatched method string matches the MCP +// spec name. A regression here would silently break consumers expecting +// notifications/{tools,prompts,resources}/list_changed. +func TestListChangedHandlers_MethodStrings(t *testing.T) { + t.Parallel() + + c := &Client{} + var ( + mu sync.Mutex + got []string + ) + c.OnNotification(func(n mcp.JSONRPCNotification) { + mu.Lock() + got = append(got, n.Method) + mu.Unlock() + }) + + opts := &gosdk.ClientOptions{} + c.installNotificationHandlers(opts) + + require.NotNil(t, opts.ToolListChangedHandler) + require.NotNil(t, opts.PromptListChangedHandler) + require.NotNil(t, opts.ResourceListChangedHandler) + + opts.ToolListChangedHandler(context.Background(), nil) + opts.PromptListChangedHandler(context.Background(), nil) + opts.ResourceListChangedHandler(context.Background(), nil) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{ + "notifications/tools/list_changed", + "notifications/prompts/list_changed", + "notifications/resources/list_changed", + }, got, "list-changed method strings must match the MCP spec names") +} diff --git a/mcpcompat/mcp/alias.go b/mcpcompat/mcp/alias.go index c235668..cffe931 100644 --- a/mcpcompat/mcp/alias.go +++ b/mcpcompat/mcp/alias.go @@ -209,6 +209,27 @@ var UnmarshalContent = mcpgo.UnmarshalContent // NewMetaFromMap builds a *Meta from a raw map. var NewMetaFromMap = mcpgo.NewMetaFromMap +// ---------------------------------------------------------------------------- +// Logging +// ---------------------------------------------------------------------------- + +// LoggingLevel is the severity of a server log message, mirrored from mcp-go. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type LoggingLevel = mcpgo.LoggingLevel + +// MCP logging levels, mirrored from mcp-go. +const ( + LoggingLevelDebug = mcpgo.LoggingLevelDebug + LoggingLevelInfo = mcpgo.LoggingLevelInfo + LoggingLevelNotice = mcpgo.LoggingLevelNotice + LoggingLevelWarning = mcpgo.LoggingLevelWarning + LoggingLevelError = mcpgo.LoggingLevelError + LoggingLevelCritical = mcpgo.LoggingLevelCritical + LoggingLevelAlert = mcpgo.LoggingLevelAlert + LoggingLevelEmergency = mcpgo.LoggingLevelEmergency +) + // ---------------------------------------------------------------------------- // Tools // ---------------------------------------------------------------------------- diff --git a/mcpcompat/server/notifications.go b/mcpcompat/server/notifications.go index a5921c0..a7cbc39 100644 --- a/mcpcompat/server/notifications.go +++ b/mcpcompat/server/notifications.go @@ -23,12 +23,14 @@ import ( // - 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. +// The list-changed notifications (tools/prompts/resources) are NOT re-broadcast +// through this method: the go-sdk server emits them automatically to connected +// clients whenever its registered feature set changes on a live session +// (per-session srv.AddTool / srv.RemoveTools, as driven by syncSessionTools +// when WithToolCapabilities(true) is set), so a manual broadcast would double +// them. Any list-changed method (and any other unrecognized method) is dropped +// here rather than forwarded. 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 { @@ -66,8 +68,10 @@ func (s *MCPServer) sendOneNotification( _ = 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. + // public generic notification sender. The list-changed notifications are + // emitted automatically by the go-sdk server on feature mutation, so + // they (and any other unrecognized method) are dropped here to avoid + // double-delivery. if s.logger != nil { s.logger.Debug("SendNotificationToAllClients: dropping unsupported notification method", "method", method) diff --git a/mcpcompat/server/notifications_test.go b/mcpcompat/server/notifications_test.go index 9b62513..c3ec195 100644 --- a/mcpcompat/server/notifications_test.go +++ b/mcpcompat/server/notifications_test.go @@ -6,12 +6,16 @@ package server_test import ( "context" "net/http/httptest" + "sync" + "sync/atomic" "testing" + "time" "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" "github.com/stacklok/toolhive-core/mcpcompat/server" ) @@ -28,7 +32,11 @@ func TestSendNotificationToAllClients_NoSessions(t *testing.T) { // 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. +// dispatch path (progress/message/list-changed/unknown). It registers an +// OnNotification handler and asserts that the progress and message +// notifications actually arrive with their expected params, so that a silent +// regression in sendOneNotification's jsonConvert or the go-sdk's +// NotifyProgress/Log is caught rather than masked by a NotPanics check. func TestSendNotificationToAllClients_Broadcast(t *testing.T) { t.Parallel() ctx := context.Background() @@ -48,7 +56,10 @@ func TestSendNotificationToAllClients_Broadcast(t *testing.T) { ts := httptest.NewServer(httpSrv) t.Cleanup(ts.Close) - c, err := client.NewStreamableHttpClient(ts.URL) + // WithContinuousListening opens the standalone SSE stream so + // server-initiated notifications (which arrive outside any in-flight + // request) are delivered to the client rather than dropped. + c, err := client.NewStreamableHttpClient(ts.URL, transport.WithContinuousListening()) require.NoError(t, err) require.NoError(t, c.Start(ctx)) _, err = c.Initialize(ctx, mcp.InitializeRequest{ @@ -60,6 +71,39 @@ func TestSendNotificationToAllClients_Broadcast(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = c.Close() }) + // The server only delivers notifications/message once the client has set a + // logging level, so raise it before broadcasting. + require.NoError(t, c.SetLoggingLevel(ctx, mcp.LoggingLevelInfo)) + + type recorded struct { + method string + params mcp.NotificationParams + } + var ( + mu sync.Mutex + got []recorded + progChan = make(chan recorded, 1) + logChan = make(chan recorded, 1) + ) + c.OnNotification(func(n mcp.JSONRPCNotification) { + r := recorded{method: n.Method, params: n.Params} + mu.Lock() + got = append(got, r) + mu.Unlock() + switch n.Method { + case "notifications/progress": + select { + case progChan <- r: + default: + } + case "notifications/message": + select { + case logChan <- r: + default: + } + } + }) + // 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() { @@ -70,6 +114,25 @@ func TestSendNotificationToAllClients_Broadcast(t *testing.T) { srv.SendNotificationToAllClients("notifications/tools/list_changed", nil) srv.SendNotificationToAllClients("some/unknown/method", map[string]any{"x": 1}) }) + + // Confirm the progress notification actually arrived with its params. + select { + case p := <-progChan: + assert.Equal(t, "notifications/progress", p.method) + assert.Equal(t, "t", p.params.AdditionalFields["progressToken"]) + assert.InDelta(t, 0.5, p.params.AdditionalFields["progress"], 1e-9) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for notifications/progress broadcast") + } + // Confirm the message notification actually arrived with its params. + select { + case l := <-logChan: + assert.Equal(t, "notifications/message", l.method) + assert.Equal(t, "info", l.params.AdditionalFields["level"]) + assert.Equal(t, "hello", l.params.AdditionalFields["data"]) + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for notifications/message broadcast") + } } // TestSSEHandlers verifies SSEServer exposes non-nil SSE and message handlers. @@ -80,3 +143,93 @@ func TestSSEHandlers(t *testing.T) { assert.NotNil(t, sse.SSEHandler()) assert.NotNil(t, sse.MessageHandler()) } + +// TestListChanged_EmittedOnSetSessionTools verifies that the go-sdk server +// auto-emits a notifications/tools/list_changed notification to the connected +// client when a per-session tool overlay is applied via SetSessionTools. This +// restores the mcp-go behavior where mutating a session's tool set notifies +// the client, and depends on WithToolCapabilities(true) so the server +// advertises (and emits) list_changed. Without the per-session sync onto the +// go-sdk server (syncSessionTools), the notification would never fire. +func TestListChanged_EmittedOnSetSessionTools(t *testing.T) { + t.Parallel() + ctx := context.Background() + + var ( + regSession server.ClientSession + regDone = make(chan struct{}) + ) + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { + regSession = s + close(regDone) + }) + + srv := server.NewMCPServer("lc-server", "1.0.0", + server.WithToolCapabilities(true), + server.WithHooks(hooks), + ) + + httpSrv := server.NewStreamableHTTPServer(srv) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // WithContinuousListening opens the standalone SSE stream so server-initiated + // notifications (the list_changed fired by SetSessionTools, which arrives + // outside any in-flight request) can be delivered. Without it the go-sdk + // streamable transport has no channel to carry such notifications and drops + // them (see streamable.go streamableServerConn.Write -> streams[""]). + c, err := client.NewStreamableHttpClient(ts.URL, transport.WithContinuousListening()) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + var ( + gotListChanged atomic.Bool + done = make(chan struct{}) + ) + c.OnNotification(func(n mcp.JSONRPCNotification) { + if n.Method == "notifications/tools/list_changed" && gotListChanged.CompareAndSwap(false, true) { + close(done) + } + }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "lc-client", Version: "1"}, + }, + }) + require.NoError(t, err) + + // Wait for the OnRegisterSession hook to fire (it runs during initialize) so + // we have a handle to the session to mutate. + select { + case <-regDone: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for OnRegisterSession") + } + + // SetSessionTools must be exposed on the session; assert it and mutate the + // overlay, which reconciles onto the live go-sdk server and triggers the + // auto-emitted list_changed notification. + swt, ok := regSession.(server.SessionWithTools) + require.True(t, ok, "session must implement SessionWithTools") + swt.SetSessionTools(map[string]server.ServerTool{ + "injected": { + Tool: mcp.NewTool("injected", mcp.WithDescription("injected at runtime")), + Handler: func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("ok"), nil + }, + }, + }) + + // go-sdk debounces list_changed by ~10ms. Exit as soon as the list_changed + // notification arrives (via done), falling back to a generous timeout. + select { + case <-done: + case <-time.After(1500 * time.Millisecond): + } + assert.True(t, gotListChanged.Load(), + "expected notifications/tools/list_changed after SetSessionTools") +} From 3c7ab695dc15f0df871b642546849050117b9c91 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Wed, 8 Jul 2026 23:13:11 +0300 Subject: [PATCH 3/8] fix(mcpcompat): wire heartbeat, localhost protection, page size, local-session validate (issue #156 wave 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 2 — Wire WithHeartbeatInterval to go-sdk ServerOptions.KeepAlive (previously a silent no-op); sessions now close on unanswered pings. 5 — Expose WithDisableLocalhostProtection(bool) so callers can preserve pre-migration behavior for local proxies with custom Host headers (default: protection ON, matching go-sdk). 6 — Add WithPageSize(n int) ServerOption so aggregating servers (vMCP) can raise the tools/list page size above go-sdk's default 1000. 3/U5 — Call SessionIdManager.Validate on local sessions too, not just rehydrated ones; restores per-request liveness and sliding TTL so a session terminated cross-pod or via auth-failure takes effect on the origin pod immediately. Tests: heartbeat eviction e2e, localhost-protection toggle e2e (raw non-localhost Host, both enabled and disabled), page-size pagination e2e with nextCursor, local-session validate-evicts-on-shared-terminate and served-without-manager. Refs #156 --- mcpcompat/server/rehydration_test.go | 64 ++++++++ mcpcompat/server/server.go | 54 +++++-- mcpcompat/server/server_internal_test.go | 6 +- mcpcompat/server/server_test.go | 186 +++++++++++++++++++++++ mcpcompat/server/transports.go | 72 +++++++-- 5 files changed, 359 insertions(+), 23 deletions(-) diff --git a/mcpcompat/server/rehydration_test.go b/mcpcompat/server/rehydration_test.go index bdb739f..01f4e8a 100644 --- a/mcpcompat/server/rehydration_test.go +++ b/mcpcompat/server/rehydration_test.go @@ -410,3 +410,67 @@ func TestRehydratedSessionElicitation(t *testing.T) { } assert.True(t, gotToolResult, "expected the elicited tool call to complete on the rehydrated session") } + +// TestLocalSession_ValidateEvictsWhenTerminatedInSharedStore verifies that a +// LOCAL session (initialized on this instance) is validated against the shared +// SessionIdManager on EVERY request (issue #156, item U5), so that a session +// terminated in the shared store (e.g. via a DELETE on another replica, or a +// direct manager.Terminate) is rejected here rather than served forever from +// the local go-sdk handler. mcp-go validated every request; before the fix the +// local fast-path trusted a local session indefinitely. +// +// Single server with a SessionIdManager: initialize locally, terminate the id +// in the shared store, then issue a request with the same sid → must 404, not +// 200. +func TestLocalSession_ValidateEvictsWhenTerminatedInSharedStore(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() + + // Initialize locally on replica A — this marks the session local and hands + // its lifecycle to the go-sdk StreamableHTTPHandler. + sid := initSession(t, tsA.URL) + require.Contains(t, listToolNames(t, tsA.URL, sid), "greet", + "local session must be served before termination") + + // Terminate the session in the shared store (as a DELETE on another replica + // would, or an explicit lifecycle operation). The local session is still + // cached in the go-sdk handler's map. + _, _ = mgr.Terminate(sid) + + // The next request on A with the same sid must be rejected via per-request + // validation of the local session, not served from the cached local handler. + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + resp := postRPC(ctx, t, tsA.URL, sid, `{"jsonrpc":"2.0","id":9,"method":"tools/list","params":{}}`) + assert.Equal(t, http.StatusNotFound, resp.StatusCode, + "local session terminated in the shared store must be rejected, not served") + _ = resp.Body.Close() +} + +// TestLocalSession_ServedWithoutManager verifies the single-replica fast-path is +// preserved: when NO SessionIdManager is configured, a local session is NOT +// validated on every request (there is nothing to validate against), so a +// request after the session exists is served normally (200, not 404). This +// guards against the fix over-eagerly rejecting local sessions in single-replica +// deployments. +func TestLocalSession_ServedWithoutManager(t *testing.T) { + t.Parallel() + + stream := server.NewMCPServer("solo", "1.0.0") + addGreetTool(stream) + // No WithSessionIdManager: single-replica, nothing to validate. + s := server.NewStreamableHTTPServer(stream) + ts := httptest.NewServer(s) + defer ts.Close() + + sid := initSession(t, ts.URL) + // A follow-up request must be served (200) from the local handler. + require.Contains(t, listToolNames(t, ts.URL, sid), "greet", + "local session must be served in single-replica mode without a manager") +} diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index 2175756..8910b33 100644 --- a/mcpcompat/server/server.go +++ b/mcpcompat/server/server.go @@ -36,6 +36,7 @@ import ( "net/http" "strings" "sync" + "time" "github.com/modelcontextprotocol/go-sdk/jsonrpc" gosdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -129,6 +130,13 @@ type MCPServer struct { promptListChanged bool logging bool + // pageSize is the maximum number of items returned in a single page for + // list methods (tools/list, resources/list, prompts/list). go-sdk defaults + // to DefaultPageSize (1000) when zero; mcp-go returned everything in one + // page. Setting it via WithPageSize lets aggregators with >1000 tools raise + // (or otherwise configure) the page size so tools/list is not paginated. + pageSize int + mu sync.RWMutex tools map[string]ServerTool resources map[string]ServerResource @@ -303,6 +311,20 @@ func WithHooks(hooks *Hooks) ServerOption { return func(s *MCPServer) { s.hooks = hooks } } +// WithPageSize configures the server's list pagination page size (the maximum +// number of items returned in a single tools/list, resources/list, or +// prompts/list response). A value of 0 leaves go-sdk's default +// (DefaultPageSize=1000) in place. mcp-go returned all items in one page; +// aggregators with more than 1000 tools must raise this to avoid pagination. +func WithPageSize(n int) ServerOption { + return func(s *MCPServer) { + if n < 0 { + n = 0 + } + s.pageSize = n + } +} + // AddTool registers a tool and its handler. func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) { s.mu.Lock() @@ -369,7 +391,7 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { // 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) { +func (s *MCPServer) buildServer(genSessionID func() string, keepalive time.Duration) (*gosdk.Server, error) { s.mu.RLock() tools := make(map[string]ServerTool, len(s.tools)) for k, v := range s.tools { @@ -392,6 +414,15 @@ func (s *MCPServer) buildServer(genSessionID func() string) (*gosdk.Server, erro var srv *gosdk.Server opts := &gosdk.ServerOptions{ Logger: s.logger, + // PageSize configures list pagination. A zero value leaves go-sdk's + // DefaultPageSize (1000) in place, preserving the pre-existing behavior + // for callers that do not set WithPageSize. + PageSize: s.pageSize, + // KeepAlive wires WithHeartbeatInterval (streamable HTTP transport) to + // go-sdk's per-session keep-alive ping: the SDK pings the peer at this + // interval and closes the session if a ping goes unanswered. stdio/SSE + // pass 0 (no keep-alive), matching mcp-go. + KeepAlive: keepalive, InitializedHandler: func(ctx context.Context, req *gosdk.InitializedRequest) { if req == nil || req.Session == nil { return @@ -449,19 +480,14 @@ func (s *MCPServer) buildServer(genSessionID func() string) (*gosdk.Server, erro 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 { +// session-ID generator to install on each per-session server; keepalive is the +// per-session keep-alive ping interval (0 disables). 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, keepalive time.Duration) func(*http.Request) *gosdk.Server { return func(*http.Request) *gosdk.Server { - srv, err := s.buildServer(genSessionID) + srv, err := s.buildServer(genSessionID, keepalive) if err != nil { if s.logger != nil { s.logger.Error("building per-session MCP server", "error", err) @@ -472,6 +498,12 @@ func (s *MCPServer) getServerFunc(genSessionID func() string) func(*http.Request } } +// 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. 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) { diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index b659ace..576549f 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/mcpcompat/server/server_internal_test.go @@ -59,7 +59,7 @@ func TestClientSession_Store(t *testing.T) { func TestSetSessionTools_NonObjectSchemaDoesNotPanic(t *testing.T) { t.Parallel() s := NewMCPServer("s", "1") - srv, err := s.buildServer(nil) + srv, err := s.buildServer(nil, 0) require.NoError(t, err) cs := s.sessionFor("sid-bad-schema") @@ -133,7 +133,7 @@ func TestBuildServer_GlobalAndSessionTools(t *testing.T) { // 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) + srv, err := s.buildServer(nil, 0) require.NoError(t, err) require.NotNil(t, srv) @@ -156,7 +156,7 @@ func TestBuildServer_WithSessionIDGenerator(t *testing.T) { called := false gen := func() string { called = true; return "generated-id" } - srv, err := s.buildServer(gen) + srv, err := s.buildServer(gen, 0) require.NoError(t, err) require.NotNil(t, srv) // The generator is installed on the server (invoked by the SDK per new diff --git a/mcpcompat/server/server_test.go b/mcpcompat/server/server_test.go index 6e38484..f779221 100644 --- a/mcpcompat/server/server_test.go +++ b/mcpcompat/server/server_test.go @@ -5,8 +5,11 @@ package server_test import ( "context" + "fmt" + "io" "net/http" "net/http/httptest" + "strings" "sync" "testing" "time" @@ -323,3 +326,186 @@ func TestPerRequestContext_NotClobberedConcurrently(t *testing.T) { assert.ElementsMatch(t, []string{a, b}, got, "iteration %d: handlers saw clobbered values %v", i, got) } } + +// TestPageSize_Configurable verifies that WithPageSize threads go-sdk's +// ServerOptions.PageSize so tools/list is paginated at the configured size +// (issue #156, item 6). With PageSize(2) and 3 tools, the first tools/list must +// return at most 2 tools and a nextCursor to fetch the remaining page. go-sdk's +// default (DefaultPageSize=1000) would otherwise return all tools in one page, +// breaking vMCP aggregators with >1000 tools. +func TestPageSize_Configurable(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := server.NewMCPServer("page-server", testClientVersion, + server.WithToolCapabilities(false), + server.WithPageSize(2), + ) + for i := 0; i < 3; i++ { + name := fmt.Sprintf("t%d", i) + srv.AddTool(mcp.NewTool(name, mcp.WithDescription("tool "+name)), + 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)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + // First page: at most 2 tools, plus a nextCursor signalling more remain. + first, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + assert.LessOrEqual(t, len(first.Tools), 2, "first page must respect the configured page size") + assert.NotEmpty(t, first.NextCursor, "a nextCursor must be returned when more tools remain") + + // Second page: fetch the rest with the cursor; together they cover all 3. + secondReq := mcp.ListToolsRequest{} + secondReq.Params.Cursor = first.NextCursor + second, err := c.ListTools(ctx, secondReq) + require.NoError(t, err) + total := len(first.Tools) + len(second.Tools) + assert.Equal(t, 3, total, "both pages together must return all registered tools") +} + +// TestDisableLocalhostProtection_PassedThrough verifies that +// WithDisableLocalhostProtection propagates to go-sdk's +// StreamableHTTPOptions.DisableLocalhostProtection (issue #156, item 5). go-sdk +// by default 403s requests on a loopback listener with a non-localhost Host +// header; with the option set, a request carrying a custom Host header must +// proceed rather than be rejected. +// +// This mirrors TestDisableLocalhostProtection_DefaultRejectsNonLocalhostHost: +// a raw http.Request with req.Host set to a non-localhost value is the only +// way to exercise go-sdk's DNS-rebinding check (the compat client sends a +// localhost Host header, which go-sdk never 403s). With the option disabled +// the same request that the default-rejects test expects to 403 must instead +// succeed (200), proving the option actually toggles the protection. +func TestDisableLocalhostProtection_PassedThrough(t *testing.T) { + t.Parallel() + + srv := server.NewMCPServer("host-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool(mcp.NewTool("ping", mcp.WithDescription("ping")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("pong"), nil + }) + + httpSrv := server.NewStreamableHTTPServer(srv, server.WithDisableLocalhostProtection(true)) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}` + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ts.URL, strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + // A non-localhost Host header would trigger go-sdk's DNS-rebinding + // protection by default; WithDisableLocalhostProtection(true) must let it + // through. + req.Host = "evil.example.com" + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusOK, resp.StatusCode, "with protection disabled a non-localhost Host must be accepted") +} + +// TestDisableLocalhostProtection_DefaultRejectsNonLocalhostHost verifies that +// WITHOUT WithDisableLocalhostProtection, go-sdk's default localhost protection +// is in effect: a request on a loopback listener carrying a non-localhost Host +// header is rejected with 403. This confirms the knob actually toggles the +// protection (the default is "protected"). +func TestDisableLocalhostProtection_DefaultRejectsNonLocalhostHost(t *testing.T) { + t.Parallel() + + srv := server.NewMCPServer("prot-server", testClientVersion, server.WithToolCapabilities(false)) + httpSrv := server.NewStreamableHTTPServer(srv) // no disable option + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}` + req, err := http.NewRequestWithContext(ctx, http.MethodPost, ts.URL, strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + // A non-localhost Host header triggers go-sdk's DNS-rebinding protection. + req.Host = "evil.example.com" + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusForbidden, resp.StatusCode, "default localhost protection must reject a non-localhost Host") +} + +// TestHeartbeat_WiredToKeepAlive verifies that WithHeartbeatInterval is wired to +// go-sdk's per-session keep-alive (issue #156, item 2) via an e2e check: a +// server built with a heartbeat disconnects a session whose transport stops +// responding to pings. go-sdk pings the peer at KeepAlive and closes the session +// when a ping goes unanswered, observable as the client's stream/sessions +// closing. Because ServerOptions.KeepAlive is not observable directly, this is +// an e2e test. +// +// We observe the wiring through the resulting server's session lifecycle: a +// raw JSON-RPC initialize against a server configured with a short heartbeat, +// after which the client stops reading, must lead to the session being closed +// within a bounded window. This proves the heartbeat is no longer a no-op. +func TestHeartbeat_WiredToKeepAlive(t *testing.T) { + t.Parallel() + + srv := server.NewMCPServer("hb-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool(mcp.NewTool("ping", mcp.WithDescription("ping")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("pong"), nil + }) + + httpSrv := server.NewStreamableHTTPServer(srv, server.WithHeartbeatInterval(150*time.Millisecond)) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // Initialize a session via raw JSON-RPC, then stop reading the response + // stream entirely (abandon the connection). With KeepAlive wired, go-sdk's + // per-session pinger will fail to get a pong back and must close the session + // within a bounded window. We assert the session is gone by issuing a + // follow-up request with the same session id: once the SDK has evicted the + // dead session, it 404s the id (the go-sdk handler does not recognize it). + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + sid := initSession(t, ts.URL) + + // Give the keep-alive pinger time to fire and evict the unresponsive session. + // The heartbeat is 150ms; allow a generous multiple for the ping timeout + + // eviction. + deadline := time.Now().Add(3 * time.Second) + var got404 bool + for time.Now().Before(deadline) { + resp := postRPC(ctx, t, ts.URL, sid, `{"jsonrpc":"2.0","id":7,"method":"tools/list","params":{}}`) + sc := resp.StatusCode + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + if sc == http.StatusNotFound { + got404 = true + break + } + // A 200 means the session is still alive; back off briefly and retry. + time.Sleep(100 * time.Millisecond) + } + assert.True(t, got404, "session must be evicted by keep-alive within the bounded window once the client stops responding to pings") +} diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index 25ab92f..63b10fd 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -60,7 +60,7 @@ 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) + srv, err := server.buildServer(nil, 0) if err != nil { return err } @@ -84,6 +84,11 @@ type StreamableHTTPServer struct { contextFunc HTTPContextFunc sessionIDMgr SessionIdManager heartbeat time.Duration + // disableLocalhostProtection turns off go-sdk's DNS-rebinding/localhost + // protection (which 403s requests on a loopback listener with a non-localhost + // Host header). mcp-go had no such protection; local proxies with custom Host + // headers need it disabled. See WithDisableLocalhostProtection. + disableLocalhostProtection bool once sync.Once handler http.Handler @@ -139,11 +144,27 @@ func WithSessionIdManager(manager SessionIdManager) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.sessionIDMgr = manager } } -// WithHeartbeatInterval sets the keep-alive ping interval. +// WithHeartbeatInterval sets the keep-alive ping interval. This option is now +// LIVE (it was previously a no-op): it is wired to go-sdk's +// ServerOptions.KeepAlive, so the SDK pings the peer at this interval and +// closes the session when a ping goes unanswered. Set a small interval only if +// you want unanswered pings to terminate the session. Applies to the Streamable +// HTTP transport only; stdio and SSE pass 0 (no keep-alive), matching mcp-go. func WithHeartbeatInterval(interval time.Duration) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.heartbeat = interval } } +// WithDisableLocalhostProtection disables go-sdk's DNS-rebinding/localhost +// protection on the Streamable HTTP transport. go-sdk by default rejects (403) +// requests that arrive on a loopback listener but carry a non-localhost Host +// header; mcp-go had no such protection, so local proxies that set a custom +// Host header must disable it. Pass true to disable the protection (restoring +// mcp-go's behavior); pass false (or omit) to leave go-sdk's default +// (protected) in place. +func WithDisableLocalhostProtection(disable bool) StreamableHTTPOption { + return func(s *StreamableHTTPServer) { s.disableLocalhostProtection = disable } +} + // WithHTTPContextFunc installs a per-request context customizer. // // Context values injected here are applied to ALL POSTs, including the @@ -163,18 +184,27 @@ func (s *StreamableHTTPServer) build() { } // 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 { + if _, err := s.mcp.buildServer(gen, s.heartbeat); 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} + opts := &gosdk.StreamableHTTPOptions{ + JSONResponse: true, + // DisableLocalhostProtection propagates the option to go-sdk. It + // defaults to false (protection enabled, matching go-sdk's safe + // default); set it via WithDisableLocalhostProtection to turn the + // protection off, restoring mcp-go's acceptance of local proxies + // with custom Host headers. + DisableLocalhostProtection: s.disableLocalhostProtection, + } // 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) + // registration middleware buildServer installs. The heartbeat interval is + // threaded into each per-session go-sdk server's KeepAlive. + s.handler = gosdk.NewStreamableHTTPHandler(s.mcp.getServerFunc(gen, s.heartbeat), opts) }) } @@ -231,6 +261,30 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) return } + // Local-session validation (issue #156, item U5): a session initialized on + // THIS instance is trusted by the local go-sdk handler, but mcp-go validated + // every request. When a shared SessionIdManager is configured, validate local + // sessions on every request too, so a session terminated in the shared store + // (e.g. via a DELETE on another replica) is rejected here and its local + // bookkeeping dropped, rather than being served forever from the local + // handler. The single-replica fast-path (no manager) is preserved: there is + // nothing to validate against, so local sessions route straight through. + if sid := r.Header.Get("Mcp-Session-Id"); sid != "" && s.sessionIDMgr != nil && s.mcp.isLocalSession(sid) { + isTerminated, err := s.sessionIDMgr.Validate(sid) + if err != nil { + s.mcp.forgetSession(sid) + s.deleteRehydrated(sid) + http.Error(w, "Invalid session ID", http.StatusNotFound) + return + } + if isTerminated { + s.mcp.forgetSession(sid) + s.deleteRehydrated(sid) + http.Error(w, "Session terminated", http.StatusNotFound) + return + } + } + s.handler.ServeHTTP(w, r) } @@ -305,7 +359,7 @@ func (s *StreamableHTTPServer) rehydrate(r *http.Request, sid string) (*rehydrat return rt, nil } - srv, err := s.mcp.buildServer(nil) + srv, err := s.mcp.buildServer(nil, 0) if err != nil { return nil, err } @@ -412,11 +466,11 @@ func WithMessageEndpoint(endpoint string) SSEOption { func (s *SSEServer) build() { s.once.Do(func() { - if _, err := s.mcp.buildServer(nil); err != nil { + if _, err := s.mcp.buildServer(nil, 0); err != nil { s.buildErr = err return } - s.handler = gosdk.NewSSEHandler(s.mcp.getServerFunc(nil), nil) + s.handler = gosdk.NewSSEHandler(s.mcp.getServerFunc(nil, 0), nil) }) } From 19c482ba460fea2194a4f92fe086452ea20bb8ea Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Wed, 8 Jul 2026 23:40:25 +0300 Subject: [PATCH 4/8] fix(mcpcompat): scope 401 string-matching to transport-level errors (issue #156 wave 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 4/U7 — Constrain 401/unauthorized/404 string-matching to transport-level errors only by checking the captured HTTP status (from errorbody.go) first and using errors.As(*jsonrpc.Error) to distinguish RPC-level from transport-level errors. A JSON-RPC tool error whose message contains "unauthorized" no longer triggers a false-positive OAuth refresh. Restore mcp-go's any-4xx-on-initialize → ErrLegacySSEServer classification (400/403/404/405 on connect → legacy SSE; 401 stays ErrUnauthorized). String-matching remains as a best-effort fallback for transport failures with no captured body. Tests: false-positive JSON-RPC error isolation, transport 401/404/403 classification by captured status, 4xx-on-initialize legacy SSE restoration, string-matching fallback. Refs #156 --- mcpcompat/client/errors.go | 71 +++++++- mcpcompat/client/errors_internal_test.go | 204 ++++++++++++++++++++++- 2 files changed, 264 insertions(+), 11 deletions(-) diff --git a/mcpcompat/client/errors.go b/mcpcompat/client/errors.go index 99270ab..3d3a980 100644 --- a/mcpcompat/client/errors.go +++ b/mcpcompat/client/errors.go @@ -7,6 +7,7 @@ import ( "context" "errors" "fmt" + "net/http" "strings" "github.com/modelcontextprotocol/go-sdk/jsonrpc" @@ -32,9 +33,29 @@ func enrichWithResponseBody(ctx context.Context, err error) error { } // mapConnectError maps an error returned by the underlying go-sdk Connect call -// onto the transport-level sentinels ToolHive checks for. +// (the initialize handshake) onto the transport-level sentinels ToolHive checks +// for. +// +// mcp-go classified any 4xx (except 401) on the initialize POST as a legacy +// SSE-only server. We restore that behavior here using the captured HTTP status +// (set by captureErrorBody for non-2xx responses), falling back to best-effort +// string matching when no body was captured (e.g. a failure before any +// response). 401 stays ErrUnauthorized so ToolHive's OAuth refresh flow still +// triggers. func mapConnectError(ctx context.Context, err error) error { - return mapTransportError(enrichWithResponseBody(ctx, err)) + err = enrichWithResponseBody(ctx, err) + h := capturedErr(ctx) + if h != nil && h.status >= 400 && h.status < 500 { + if h.status == http.StatusUnauthorized { + return transport.NewError(errors.Join( + &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(err)}, + transport.ErrUnauthorized, + err, + )) + } + return transport.NewError(errors.Join(transport.ErrLegacySSEServer, err)) + } + return mapTransportError(err, h) } // mapCallError maps an error returned by an underlying go-sdk request call onto @@ -42,31 +63,69 @@ func mapConnectError(ctx context.Context, err error) error { // 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. +// +// JSON-RPC-level errors (a *jsonrpc.Error returned inside a 2xx response) that +// carry no captured HTTP status are application errors, not transport failures: +// they must NOT be reclassified as transport auth/session failures based on +// their text — otherwise a tool error whose message contains "unauthorized" +// would wrongly trigger ToolHive's OAuth refresh flow. Only errors that look +// like transport-level HTTP failures (a captured non-2xx status, or a plain +// non-RPC error from the transport) are passed to mapTransportError. func mapCallError(ctx context.Context, err error) error { if err == nil { return nil } err = enrichWithResponseBody(ctx, err) + h := capturedErr(ctx) var wireErr *jsonrpc.Error if errors.As(err, &wireErr) && wireErr.Code == jsonrpc.CodeMethodNotFound { return errors.Join(mcp.ErrMethodNotFound, err) } - return mapTransportError(err) + if (h == nil || h.status == 0) && errors.As(err, &wireErr) { + // RPC-level error with no captured HTTP status: surface unchanged. + return err + } + return mapTransportError(err, h) } // 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. // +// When the caller captured a non-2xx HTTP status (h.status set by +// captureErrorBody), classification is driven by the status code, which avoids +// false positives from server-provided body text. When no status was captured +// (e.g. a transport failure before any response, or a response whose body +// capture was skipped), detection falls back to best-effort string matching. +// // 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 +// status code, so the string-matching fallback is inherently best-effort. 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 { +func mapTransportError(err error, h *errBody) error { if err == nil { return nil } + // Prefer the captured HTTP status over string matching: the status is + // authoritative, while body text (re-attached by enrichWithResponseBody) + // can legitimately contain words like "unauthorized" for non-401 responses. + if h != nil && h.status >= 400 && h.status < 500 { + switch h.status { + case http.StatusUnauthorized: + return transport.NewError(errors.Join( + &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(err)}, + transport.ErrUnauthorized, + err, + )) + case http.StatusNotFound: + return transport.NewError(errors.Join(transport.ErrSessionTerminated, err)) + default: + // Other 4xx on a regular call are not legacy-SSE (that classification + // is connect-time only, in mapConnectError); surface unchanged. + return err + } + } msg := strings.ToLower(err.Error()) switch { diff --git a/mcpcompat/client/errors_internal_test.go b/mcpcompat/client/errors_internal_test.go index 911e47d..7fe55ce 100644 --- a/mcpcompat/client/errors_internal_test.go +++ b/mcpcompat/client/errors_internal_test.go @@ -4,17 +4,22 @@ package client import ( + "context" "errors" + "fmt" + "net/http" "testing" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/stacklok/toolhive-core/mcpcompat/client/transport" ) func TestMapTransportError_Unauthorized(t *testing.T) { t.Parallel() - err := mapTransportError(errors.New("request failed: 401 Unauthorized")) + err := mapTransportError(errors.New("request failed: 401 Unauthorized"), nil) // ToolHive branches on all of these for its OAuth/401 handling. assert.True(t, errors.Is(err, transport.ErrUnauthorized), "errors.Is ErrUnauthorized") @@ -29,23 +34,212 @@ func TestMapTransportError_Unauthorized(t *testing.T) { func TestMapTransportError_SessionTerminated(t *testing.T) { t.Parallel() - err := mapTransportError(errors.New("server returned 404: session not found")) + err := mapTransportError(errors.New("server returned 404: session not found"), nil) assert.True(t, errors.Is(err, transport.ErrSessionTerminated)) } func TestMapTransportError_LegacySSE(t *testing.T) { t.Parallel() - err := mapTransportError(errors.New("405 method not allowed")) + err := mapTransportError(errors.New("405 method not allowed"), nil) 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)) + assert.Equal(t, orig, mapTransportError(orig, nil)) } func TestMapTransportError_Nil(t *testing.T) { t.Parallel() - assert.NoError(t, mapTransportError(nil)) + assert.NoError(t, mapTransportError(nil, nil)) +} + +// seedCapture returns a context carrying an errBody holder seeded with the +// given captured HTTP status and body, mirroring what captureErrorBody records +// for a non-2xx response. It is used to simulate a go-sdk transport-level HTTP +// failure without standing up a real server. +func seedCapture(status int, body string) context.Context { + ctx := withErrCapture(context.Background()) + h := capturedErr(ctx) + h.status = status + h.body = body + return ctx +} + +// TestMapCallError_DoesNotMisclassifyJSONRPCErrors is the false-positive fix: +// a JSON-RPC tool error whose message text contains "unauthorized" (e.g. a +// backend returning {"error":"unauthorized to access X"} inside a 2xx response) +// must NOT be reclassified as a transport 401, since that would wrongly trigger +// ToolHive's OAuth refresh flow. There is no captured HTTP status (it's an +// RPC-level error returned in a 2xx body), so the auth/session classification is +// skipped. +func TestMapCallError_DoesNotMisclassifyJSONRPCErrors(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + err error + }{ + { + name: "unauthorized in message", + err: &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "unauthorized to access resource X"}, + }, + { + name: "401 in message", + err: &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "got 401 from downstream"}, + }, + { + name: "method not allowed in message but non-method-not-found code", + err: &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "method not allowed here"}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + // No captured body: this is an RPC-level error. + ctx := withErrCapture(context.Background()) + mapped := mapCallError(ctx, tc.err) + require.Error(t, mapped) + assert.False(t, errors.Is(mapped, transport.ErrUnauthorized), + "RPC-level error must not satisfy ErrUnauthorized: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrAuthorizationRequired), + "RPC-level error must not satisfy ErrAuthorizationRequired: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrLegacySSEServer), + "RPC-level error must not satisfy ErrLegacySSEServer: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrSessionTerminated), + "RPC-level error must not satisfy ErrSessionTerminated: %v", mapped) + }) + } +} + +// TestMapCallError_Transport401 confirms a transport-level 401 (with a captured +// HTTP status) through mapCallError still maps to ErrUnauthorized. Scoping, not +// removal: real transport 401s must keep triggering ToolHive's OAuth flow. +func TestMapCallError_Transport401(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status int + body string + // wantErr is the sentinel the mapped error must satisfy. + wantErr error + // notErr is a sentinel it must NOT satisfy (sanity). + notErr error + }{ + { + name: "captured 401 body", + status: http.StatusUnauthorized, + body: "Unauthorized", + wantErr: transport.ErrUnauthorized, + notErr: transport.ErrLegacySSEServer, + }, + { + name: "captured 404 session", + status: http.StatusNotFound, + body: "session not found", + wantErr: transport.ErrSessionTerminated, + notErr: transport.ErrUnauthorized, + }, + { + name: "captured 403 on call is not unauthorized", + status: http.StatusForbidden, + body: "Unauthorized: denied by policy", + wantErr: nil, // surfaced unchanged + notErr: transport.ErrUnauthorized, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := seedCapture(tc.status, tc.body) + // A plain (non-RPC) transport error, as go-sdk surfaces for non-2xx. + src := fmt.Errorf(`calling "tools/call": %s`, http.StatusText(tc.status)) + mapped := mapCallError(ctx, src) + require.Error(t, mapped) + if tc.wantErr != nil { + assert.True(t, errors.Is(mapped, tc.wantErr), "want %v, got %v", tc.wantErr, mapped) + } + if tc.notErr != nil { + assert.False(t, errors.Is(mapped, tc.notErr), "must not satisfy %v, got %v", tc.notErr, mapped) + } + }) + } +} + +// TestMapConnectError_4xxOnInitialize_LegacySSE restores mcp-go's behavior: +// any 4xx except 401 on the initialize/connect POST is classified as a legacy +// SSE-only server. 401 on connect stays ErrUnauthorized (so OAuth refresh +// triggers) rather than legacy SSE. +func TestMapConnectError_4xxOnInitialize_LegacySSE(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status int + body string + wantErr error + notErr error + }{ + { + name: "405 method not allowed", + status: http.StatusMethodNotAllowed, + body: "Method Not Allowed", + wantErr: transport.ErrLegacySSEServer, + notErr: transport.ErrUnauthorized, + }, + { + name: "404 on initialize", + status: http.StatusNotFound, + body: "Not Found", + wantErr: transport.ErrLegacySSEServer, + notErr: transport.ErrSessionTerminated, + }, + { + name: "400 bad request", + status: http.StatusBadRequest, + body: "Bad Request", + wantErr: transport.ErrLegacySSEServer, + notErr: transport.ErrUnauthorized, + }, + { + name: "403 forbidden on initialize", + status: http.StatusForbidden, + body: "Forbidden", + wantErr: transport.ErrLegacySSEServer, + notErr: transport.ErrUnauthorized, + }, + { + name: "401 on connect stays unauthorized", + status: http.StatusUnauthorized, + body: "Unauthorized", + wantErr: transport.ErrUnauthorized, + notErr: transport.ErrLegacySSEServer, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := seedCapture(tc.status, tc.body) + src := fmt.Errorf("calling initialize: %s", http.StatusText(tc.status)) + mapped := mapConnectError(ctx, src) + require.Error(t, mapped) + assert.True(t, errors.Is(mapped, tc.wantErr), "want %v, got %v", tc.wantErr, mapped) + assert.False(t, errors.Is(mapped, tc.notErr), "must not satisfy %v, got %v", tc.notErr, mapped) + }) + } +} + +// TestMapConnectError_NoCaptureFallback verifies that when no body was captured +// (e.g. a transport failure before any response), mapConnectError falls back to +// best-effort string matching, preserving the prior behavior for go-sdk errors +// that don't surface a captured status. +func TestMapConnectError_NoCaptureFallback(t *testing.T) { + t.Parallel() + ctx := withErrCapture(context.Background()) // empty holder, no status + src := errors.New("request failed: 401 Unauthorized") + mapped := mapConnectError(ctx, src) + require.Error(t, mapped) + assert.True(t, errors.Is(mapped, transport.ErrUnauthorized), "fallback should still detect 401: %v", mapped) } From 884cfd345dc89a9f67af1356fb976d4b6f420547 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Thu, 9 Jul 2026 00:30:50 +0300 Subject: [PATCH 5/8] fix(mcpcompat): elicitation surface, before-hook params, preset session docs, stale comments (issue #156 wave 5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit U4 — Document elicitation delivery under JSONResponse mode (requires client continuous listening); add client-side ElicitationHandler interface, WithElicitationHandler option, and NewStreamableHttpClientWithOpts; rewrite server elicitation.go to mirror mcp-go (SessionWithElicitation, RequestElicitation, ErrNoActiveSession/ErrElicitationNotSupported). U9 — Populate before-hook request objects (mcp.CallToolRequest with tool name + args, mcp.ListToolsRequest with cursor) from go-sdk params; fix latent translateUnknownToolError type assertion (CallToolParamsRaw). U2 — Document the preset backend session ID limitation (upstream-only: go-sdk has no SessionID field on StreamableClientTransport; honored via resume path only, not Initialize). 7 — Fix remaining stale doc comments (WithSessionIdManager, package doc, SessionWithTools). Tests: elicitation e2e (works with continuous listening, fails without), handler error return, ErrNoActiveSession/ErrElicitationNotSupported guards, before-hook tool name+args+cursor, unknown-tool error contains tool name. Refs #156 --- mcpcompat/client/client.go | 126 +++++ mcpcompat/client/transport/streamable_http.go | 18 + mcpcompat/mcp/alias.go | 2 + mcpcompat/server/elicitation.go | 75 ++- mcpcompat/server/rehydration_test.go | 2 +- mcpcompat/server/server.go | 82 +++- mcpcompat/server/server_internal_test.go | 7 +- mcpcompat/server/server_test.go | 440 ++++++++++++++++++ mcpcompat/server/session.go | 12 +- mcpcompat/server/transports.go | 44 +- 10 files changed, 768 insertions(+), 40 deletions(-) diff --git a/mcpcompat/client/client.go b/mcpcompat/client/client.go index fd02d0f..40694fa 100644 --- a/mcpcompat/client/client.go +++ b/mcpcompat/client/client.go @@ -51,14 +51,89 @@ type Client struct { // resumes a pre-existing session (transport.WithSession) without calling // Initialize. It is lazily created on the first resumed request. See resume.go. resume *resumeState + + // elicitationHandler, when set, is wired into the go-sdk ClientOptions so + // the client declares the elicitation capability at initialize and answers + // server->client elicitation/create requests. It mirrors mcp-go's + // client.WithElicitationHandler. + elicitationHandler ElicitationHandler +} + +// ElicitationHandler handles server->client elicitation/create requests. It +// mirrors mcp-go's client.ElicitationHandler. +type ElicitationHandler interface { + Elicit(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) +} + +// ElicitationHandlerFunc is an adapter to allow the use of ordinary functions +// as elicitation handlers. It is a shim convenience adapter not present in +// mcp-go (mcp-go's client.WithElicitationHandler accepts a plain func), kept +// here so the shim's ElicitationHandler interface and the func form compose. +type ElicitationHandlerFunc func(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) + +// Elicit calls f(ctx, request). +func (f ElicitationHandlerFunc) Elicit(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + return f(ctx, request) +} + +// ClientOption configures a Client. It is applied at construction time via +// NewStreamableHttpClientWithOpts (the only constructor that accepts +// ClientOptions); NewStreamableHttpClient and NewSSEMCPClient take only +// transport-level options. +// +//nolint:revive // name intentionally matches mcp-go for drop-in compatibility. +type ClientOption func(*Client) + +// WithElicitationHandler installs a handler for server->client +// elicitation/create requests. The handler must be registered before Initialize +// so it is wired into the underlying go-sdk client, which (a) declares the +// elicitation capability at initialize time so the server's Elicit calls pass +// the go-sdk capability gate, and (b) dispatches incoming elicitation/create +// requests to the handler. +// +// Elicitation under the Streamable HTTP transport requires the client to hold +// an open standalone SSE stream: the shim server replies with application/json +// (JSONResponse is on by design), so the go-sdk routes a server->client +// elicitation request made during request handling to the standalone SSE +// stream rather than the POST response. Configure the client with +// transport.WithContinuousListening() to open that stream; without it +// (DisableStandaloneSSE defaults to true) elicitation cannot be delivered and +// the server's Elicit call is rejected. See the build() doc comment in +// mcpcompat/server/transports.go. +// +// It mirrors mcp-go's client.WithElicitationHandler. +func WithElicitationHandler(h ElicitationHandler) ClientOption { + return func(c *Client) { c.elicitationHandler = h } +} + +// applyClientOptions applies the client-level options to c. +func applyClientOptions(c *Client, opts []ClientOption) { + for _, opt := range opts { + opt(c) + } } // NewStreamableHttpClient creates a Streamable HTTP MCP client for baseURL. Like // mcp-go, the returned client is not yet connected; call Start then Initialize. +// Transport options configure the underlying Streamable HTTP transport; client +// options (client.WithElicitationHandler, ...) configure client-level behavior. func NewStreamableHttpClient(baseURL string, options ...transport.StreamableHTTPCOption) (*Client, error) { return &Client{streamable: transport.NewStreamableHTTP(baseURL, options...)}, nil } +// NewStreamableHttpClientWithOpts creates a Streamable HTTP MCP client for +// baseURL, applying transport-level options followed by client-level options. +// It is the entry point for options that are not transport options (e.g. +// WithElicitationHandler); NewStreamableHttpClient remains for callers that +// only need transport options. +func NewStreamableHttpClientWithOpts( + baseURL string, transportOpts []transport.StreamableHTTPCOption, clientOpts []ClientOption, +) (*Client, error) { + c := &Client{streamable: transport.NewStreamableHTTP(baseURL, transportOpts...)} + applyClientOptions(c, clientOpts) + return c, 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) { @@ -72,6 +147,16 @@ 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. +// +// LIMITATION (issue #156, item U2): if a preset session ID was supplied via +// transport.WithSession, Initialize does NOT honor it. Initialize always +// negotiates a fresh, server-assigned session ID: it builds a go-sdk +// StreamableClientTransport (buildTransport) which has no preset session-ID +// field (go-sdk has no such field in any version, verified through v1.6.1), +// performs the initialize handshake, and then adopts the ID the server returns +// (overwriting any value set via WithSession). To resume a pre-existing session +// by ID, skip Initialize and use the resume path instead (see resume.go). The +// resume path is exercised by resume_test.go. func (c *Client) Initialize(ctx context.Context, request mcp.InitializeRequest) (*mcp.InitializeResult, error) { ctx = withErrCapture(ctx) c.mu.Lock() @@ -101,6 +186,7 @@ func (c *Client) Initialize(ctx context.Context, request mcp.InitializeRequest) opts.Capabilities = caps } c.installNotificationHandlers(opts) + c.installElicitationHandler(opts) gc := gosdk.NewClient(impl, opts) @@ -497,6 +583,46 @@ func (c *Client) installNotificationHandlers(opts *gosdk.ClientOptions) { opts.LoggingMessageHandler = newLoggingMessageHandler(c.dispatch) } +// installElicitationHandler wires the user-supplied ElicitationHandler (set via +// WithElicitationHandler) onto the go-sdk ClientOptions. The go-sdk +// auto-declares the elicitation capability when ElicitationHandler is non-nil +// (see Client.capabilities), so the server's Elicit calls pass the capability +// gate; incoming elicitation/create requests are dispatched to the handler. +// +// The go-sdk handler receives a *gosdk.ElicitRequest and must return a +// *gosdk.ElicitResult. The shim converts between the go-sdk and mcp-go-shaped +// elicitation types via JSON round-trips (both encode the identical MCP wire +// format), so a handler registered against the mcp-go ElicitationRequest/ +// ElicitationResult types works unchanged. The go-sdk validates an accepted +// result against the requested schema itself (in (*Client).elicit), so the +// shim does not re-validate here. +func (c *Client) installElicitationHandler(opts *gosdk.ClientOptions) { + if c.elicitationHandler == nil { + return + } + opts.ElicitationHandler = func(ctx context.Context, req *gosdk.ElicitRequest) (*gosdk.ElicitResult, error) { + mreq := mcp.ElicitationRequest{} + if req != nil { + // Reconstruct the mcp-go-shaped request from the go-sdk params. The + // go-sdk ElicitRequest embeds its Params; a JSON round-trip maps them + // onto mcp-go's ElicitationParams (mode, message, requestedSchema, url, + // elicitationId, _meta). + if err := jsonConvert(req.Params, &mreq.Params); err != nil { + return nil, fmt.Errorf("converting elicitation request: %w", err) + } + } + res, err := c.elicitationHandler.Elicit(ctx, mreq) + if err != nil { + return nil, err + } + out := &gosdk.ElicitResult{} + if err := jsonConvert(res, out); err != nil { + return nil, fmt.Errorf("converting elicitation result: %w", err) + } + return out, nil + } +} + // dispatchFunc is the signature of Client.dispatch, factored out so the // notification handlers can be unit-tested in isolation without a live session. type dispatchFunc func(method string, params any) diff --git a/mcpcompat/client/transport/streamable_http.go b/mcpcompat/client/transport/streamable_http.go index 2da0a87..15d35ce 100644 --- a/mcpcompat/client/transport/streamable_http.go +++ b/mcpcompat/client/transport/streamable_http.go @@ -112,6 +112,24 @@ func WithHTTPHeaderFunc(headerFunc HTTPHeaderFunc) StreamableHTTPCOption { } // WithSession sets an initial session ID (for resuming a session). +// +// LIMITATION (issue #156, item U2): a preset session ID is honored ONLY via +// the resume path — i.e. when the client skips Initialize and issues requests +// through Client's raw JSON-RPC resume machinery (see resume.go, exercised by +// resume_test.go), which sends the preset ID as the Mcp-Session-Id header on +// each POST. It is NOT honored by Initialize: Initialize always builds a fresh +// go-sdk StreamableClientTransport (see Client.buildTransport) with no preset +// session-ID field and performs a full initialize handshake, so the server +// assigns (and the client adopts) a freshly-negotiated session ID that +// overwrites the value set here. +// +// This cannot be fixed in the shim alone: go-sdk's StreamableClientTransport +// has no SessionID/preset field in any version (verified through v1.6.1); the +// session ID is solely server-assigned at initialize time. Callers that need to +// pin a specific session ID must use the resume path (transport.WithSession + +// direct method calls, no Initialize), not Initialize. +// TODO(upstream): ask go-sdk to support a preset/client-requested session ID on +// StreamableClientTransport, or a Connect option that supplies one. func WithSession(sessionID string) StreamableHTTPCOption { return func(s *StreamableHTTP) { s.sessionID = sessionID } } diff --git a/mcpcompat/mcp/alias.go b/mcpcompat/mcp/alias.go index cffe931..2003b13 100644 --- a/mcpcompat/mcp/alias.go +++ b/mcpcompat/mcp/alias.go @@ -260,6 +260,8 @@ type ( ListToolsRequest = mcpgo.ListToolsRequest // ListToolsResult is the tool list response. ListToolsResult = mcpgo.ListToolsResult + // Cursor is an opaque pagination token. + Cursor = mcpgo.Cursor ) // Tool builders and result constructors. diff --git a/mcpcompat/server/elicitation.go b/mcpcompat/server/elicitation.go index f79e4ff..c76b5f1 100644 --- a/mcpcompat/server/elicitation.go +++ b/mcpcompat/server/elicitation.go @@ -5,6 +5,7 @@ package server import ( "context" + "errors" "fmt" gosdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -12,19 +13,75 @@ import ( 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. +// ErrNoActiveSession is returned when an elicitation is requested outside a +// session's request context (no ClientSession in ctx). It mirrors mcp-go's +// server.ErrNoActiveSession. +var ErrNoActiveSession = errors.New("no active session") + +// ErrElicitationNotSupported is returned when the session in scope does not +// support elicitation (it does not implement SessionWithElicitation, or its +// bound go-sdk ServerSession is unavailable). It mirrors mcp-go's +// server.ErrElicitationNotSupported. +var ErrElicitationNotSupported = errors.New("session does not support elicitation") + +// SessionWithElicitation is a ClientSession that can send elicitation requests +// to the client. It mirrors mcp-go's server.SessionWithElicitation so callers +// that type-assert against it (as mcp-go's RequestElicitation does) keep +// working; the shim's clientSession implements it. +type SessionWithElicitation interface { + ClientSession + RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) +} + +// RequestElicitation sends an elicitation/create request to the client bound to +// the session in ctx (the session a tool/prompt/resource handler is running +// under) and returns the user's response. The client must have declared the +// elicitation capability during initialization; otherwise the go-sdk rejects +// the call with "client does not support elicitation". +// +// Under the shim's Streamable HTTP transport the server replies with +// application/json (JSONResponse is on by design so ToolHive callers that +// json.Decode the body keep working). The go-sdk therefore routes a +// server->client elicitation request made during request handling to the +// standalone SSE stream (the "" stream) rather than the request's response. +// The client MUST hold an open standalone SSE stream for the elicitation to be +// delivered and answered: configure the shim client with +// transport.WithContinuousListening() (the default is DisableStandaloneSSE: +// true, under which elicitation cannot complete). See the build() doc comment +// in transports.go. +// +// It mirrors mcp-go's server.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") + session := ClientSessionFromContext(ctx) + if session == nil { + return nil, ErrNoActiveSession } - ss := cs.goSession.Load() - if ss == nil { - return nil, fmt.Errorf("no server session available for elicitation") + es, ok := session.(SessionWithElicitation) + if !ok { + return nil, ErrElicitationNotSupported } + if err := request.Params.Validate(); err != nil { + // mcp-go returns this validation error unwrapped; mirror that for parity. + return nil, err + } + return es.RequestElicitation(ctx, request) +} +// RequestElicitation sends an elicitation/create request to the client over the +// bound go-sdk ServerSession. It converts the mcp-go-shaped ElicitationRequest +// to the go-sdk's ElicitParams, calls ServerSession.Elicit, and converts the +// result back. If no go-sdk session is bound yet it returns +// ErrElicitationNotSupported. +// +// The go-sdk's Elicit enforces the capability gate (the client must have +// declared Elicitation at initialize), infers the mode from the params when +// unset, and validates an accepted result against the requested schema; those +// behaviors are delegated to the SDK rather than reimplemented here. +func (c *clientSession) RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + ss := c.goSession.Load() + if ss == nil { + return nil, ErrElicitationNotSupported + } params := &gosdk.ElicitParams{} if err := jsonConvert(request.Params, params); err != nil { return nil, fmt.Errorf("converting elicitation params: %w", err) diff --git a/mcpcompat/server/rehydration_test.go b/mcpcompat/server/rehydration_test.go index 01f4e8a..118a1ad 100644 --- a/mcpcompat/server/rehydration_test.go +++ b/mcpcompat/server/rehydration_test.go @@ -344,7 +344,7 @@ func TestRehydratedSessionElicitation(t *testing.T) { res, err := srvB.RequestElicitation(ctx, mcp.ElicitationRequest{ Params: mcp.ElicitationParams{ Message: "your name?", - RequestedSchema: map[string]any{"type": "object"}, + RequestedSchema: map[string]any{schemaType: objectType}, }, }) if err != nil { diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index 8910b33..042c4db 100644 --- a/mcpcompat/server/server.go +++ b/mcpcompat/server/server.go @@ -16,13 +16,15 @@ // // 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. +// interfaces (SessionWithTools, SessionWithResources, SessionWithElicitation, +// SessionIdManager) and the Hooks type are implemented and wired: per-session +// tool/resource overlays set via SetSessionTools/SetSessionResources are +// reconciled onto the session's live go-sdk server (syncSessionTools/ +// syncSessionResources), and the before-list/before-call hooks fire ahead of +// SDK dispatch so ToolHive's lazy per-session tool injection runs first. +// Cross-replica session rehydration (Validate-driven lazy eviction) and the +// Streamable HTTP transports are functional. See the notes on SessionWithTools +// for the live-overlay reconciliation details. // // Stability: Alpha. package server @@ -529,17 +531,10 @@ func (s *MCPServer) sessionDispatchMiddleware(srv *gosdk.Server) gosdk.Middlewar } // 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{}) - } - } + // enumerates (tools/list) or dispatches (tools/call) them. The + // hook's request object is populated from req.GetParams(); see + // fireBeforeHooks for the extraction and fallback behavior. + s.fireBeforeHooks(ctx, method, req) res, err := next(ctx, method, req) if err != nil && method == string(mcp.MethodToolsCall) { err = translateUnknownToolError(err, req) @@ -552,6 +547,55 @@ func (s *MCPServer) sessionDispatchMiddleware(srv *gosdk.Server) gosdk.Middlewar } } +// fireBeforeHooks fires the before-list-tools / before-call-tool 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. +// +// The hook's request object is populated from req.GetParams() so a hook that +// reads the tool name/args/cursor (any future hook) sees the actual request +// rather than an empty value, matching mcp-go (whose hooks fired with the parsed +// request). The go-sdk hands the params via req.GetParams(): a +// *CallToolParamsRaw for tools/call (carrying Name and raw Arguments) and a +// *ListToolsParams for tools/list (carrying Cursor). If extraction fails (batch, +// unexpected type), the hook receives the empty request and a Debug log is +// emitted — dispatch is never broken by a hook-parameter extraction failure. +func (s *MCPServer) fireBeforeHooks(ctx context.Context, method string, req gosdk.Request) { + if s.hooks == nil { + return + } + switch method { + case string(mcp.MethodToolsList): + listReq := &mcp.ListToolsRequest{} + if p, ok := req.GetParams().(*gosdk.ListToolsParams); ok && p != nil { + listReq.Params.Cursor = mcp.Cursor(p.Cursor) + } else if s.logger != nil { + s.logger.Debug("before-list-tools hook: params not *ListToolsParams; hook receives empty request", + "method", method) + } + s.hooks.beforeListTools(ctx, nil, listReq) + case string(mcp.MethodToolsCall): + callReq := &mcp.CallToolRequest{} + if p, ok := req.GetParams().(*gosdk.CallToolParamsRaw); ok && p != nil { + callReq.Params.Name = p.Name + if len(p.Arguments) > 0 { + var args map[string]any + if err := json.Unmarshal(p.Arguments, &args); err != nil { + if s.logger != nil { + s.logger.Debug("before-call-tool hook: unmarshaling arguments failed; hook receives name only", + "tool", p.Name, "error", err) + } + } else { + callReq.Params.Arguments = args + } + } + } else if s.logger != nil { + s.logger.Debug("before-call-tool hook: params not *CallToolParamsRaw; hook receives empty request", + "method", method) + } + s.hooks.beforeCallTool(ctx, nil, callReq) + } +} + 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, @@ -728,7 +772,7 @@ func translateUnknownToolError(err error, req gosdk.Request) error { return err } name := "" - if p, ok := req.GetParams().(*gosdk.CallToolParams); ok { + if p, ok := req.GetParams().(*gosdk.CallToolParamsRaw); ok && p != nil { name = p.Name } return &jsonrpc.Error{Code: jerr.Code, Message: fmt.Sprintf("tool %q not found", name)} diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index 576549f..45ba378 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/mcpcompat/server/server_internal_test.go @@ -17,9 +17,10 @@ import ( // 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) + _ ClientSession = (*clientSession)(nil) + _ SessionWithTools = (*clientSession)(nil) + _ SessionWithResources = (*clientSession)(nil) + _ SessionWithElicitation = (*clientSession)(nil) ) func TestClientSession_Store(t *testing.T) { diff --git a/mcpcompat/server/server_test.go b/mcpcompat/server/server_test.go index f779221..b9d4dcc 100644 --- a/mcpcompat/server/server_test.go +++ b/mcpcompat/server/server_test.go @@ -29,6 +29,21 @@ const ( testClientVersion = "1.0.0" ) +// schema fixture literals shared across the elicitation tests below; factored +// into consts so goconst does not flag them across the test file. +const ( + schemaType = "type" + objectType = "object" + stringType = "string" + nameKey = "name" + properties = "properties" + + // elicitMessage and elicitToolName are reused across the elicitation tests; + // factored into consts so goconst does not flag them. + elicitMessage = "What is your name?" + elicitToolName = "ask" +) + // 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. @@ -509,3 +524,428 @@ func TestHeartbeat_WiredToKeepAlive(t *testing.T) { } assert.True(t, got404, "session must be evicted by keep-alive within the bounded window once the client stops responding to pings") } + +// TestBeforeCallTool_ReceivesToolNameAndArgs verifies that the before-call hook +// fired by sessionDispatchMiddleware receives a populated mcp.CallToolRequest +// (tool name and arguments) rather than an empty one (issue #156, item U9). +// Before the fix the hook was fired with mcp.CallToolRequest{} (no name, no +// args): vMCP's hooks only use ctx today, but any future hook reading the +// request would break silently. The hook records the request it saw; after a +// tools/call the recorded name must match and the arguments must be present. +func TestBeforeCallTool_ReceivesToolNameAndArgs(t *testing.T) { + t.Parallel() + ctx := context.Background() + + var ( + hookMu sync.Mutex + hookName string + hookArgs any + ) + hooks := &server.Hooks{} + hooks.AddBeforeCallTool(func(_ context.Context, _ any, req *mcp.CallToolRequest) { + hookMu.Lock() + defer hookMu.Unlock() + hookName = req.Params.Name + hookArgs = req.Params.Arguments + }) + + srv := server.NewMCPServer("hook-server", testClientVersion, + server.WithToolCapabilities(false), + server.WithHooks(hooks), + ) + srv.AddTool( + mcp.NewTool("greet", mcp.WithDescription("greet"), mcp.WithString("name", mcp.Required())), + func(_ context.Context, r mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("hi " + r.GetString("name", "world")), 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)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + _, err = c.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "greet", Arguments: map[string]any{"name": "ada"}}, + }) + require.NoError(t, err) + + hookMu.Lock() + defer hookMu.Unlock() + assert.Equal(t, "greet", hookName, "before-call hook must receive the tool name") + args, _ := hookArgs.(map[string]any) + require.Contains(t, args, "name", "before-call hook must receive the arguments") + assert.Equal(t, "ada", args["name"], "before-call hook must receive the argument values") +} + +// TestBeforeListTools_ReceivesCursor verifies that the before-list-tools hook +// receives the pagination cursor from the request (issue #156, item U9). +func TestBeforeListTools_ReceivesCursor(t *testing.T) { + t.Parallel() + ctx := context.Background() + + var ( + hookMu sync.Mutex + hookCursor string + ) + hooks := &server.Hooks{} + hooks.AddBeforeListTools(func(_ context.Context, _ any, req *mcp.ListToolsRequest) { + hookMu.Lock() + defer hookMu.Unlock() + hookCursor = string(req.Params.Cursor) + }) + + srv := server.NewMCPServer("hook-server", testClientVersion, + server.WithToolCapabilities(false), + server.WithPageSize(1), + server.WithHooks(hooks), + ) + srv.AddTool(mcp.NewTool("t0", mcp.WithDescription("d")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("ok"), nil + }) + srv.AddTool(mcp.NewTool("t1", mcp.WithDescription("d")), + 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)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + first, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + require.NotEmpty(t, first.NextCursor, "need a nextCursor to exercise the cursor path") + + // Second list with the cursor: the before-list hook must see that cursor. + secondReq := mcp.ListToolsRequest{} + secondReq.Params.Cursor = first.NextCursor + _, err = c.ListTools(ctx, secondReq) + require.NoError(t, err) + + hookMu.Lock() + defer hookMu.Unlock() + assert.Equal(t, string(first.NextCursor), hookCursor, "before-list-tools hook must receive the request cursor") +} + +// TestElicitation_DefaultConfig verifies the elicitation delivery contract +// (issue #156, item U4): the shim server keeps JSONResponse on, so a +// server->client elicitation made during a tools/call is routed by go-sdk to +// the standalone SSE stream. With transport.WithContinuousListening() the +// client opens that stream and an installed elicitation handler answers, so +// elicitation completes and the tool handler receives the response. Without +// continuous listening the standalone SSE stream is absent and elicitation +// cannot be delivered (documented separately below). +func TestElicitation_DefaultConfig(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + srv := server.NewMCPServer("elicit-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool( + mcp.NewTool(elicitToolName, mcp.WithDescription("elicit then answer")), + func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + res, err := srv.RequestElicitation(ctx, mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{ + Message: elicitMessage, + RequestedSchema: map[string]any{ + schemaType: objectType, + properties: map[string]any{ + nameKey: map[string]any{schemaType: stringType}, + }, + }, + }, + }) + if err != nil { + return nil, err + } + name, _ := res.Content.(map[string]any)[nameKey].(string) + return mcp.NewToolResultText("hello " + name), nil + }, + ) + + httpSrv := server.NewStreamableHTTPServer(srv) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + c, err := client.NewStreamableHttpClientWithOpts( + ts.URL, + []transport.StreamableHTTPCOption{transport.WithContinuousListening()}, + []client.ClientOption{client.WithElicitationHandler(client.ElicitationHandlerFunc( + func(_ context.Context, _ mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + return &mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{ + Action: mcp.ElicitationResponseActionAccept, + Content: map[string]any{nameKey: "grace"}, + }, + }, nil + }, + ))}, + ) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + res, err := c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: elicitToolName}}) + require.NoError(t, err) + require.False(t, res.IsError, "tool call must succeed; elicitation must complete") + require.Len(t, res.Content, 1) + txt, ok := mcp.AsTextContent(res.Content[0]) + require.True(t, ok) + assert.Equal(t, "hello grace", txt.Text, "the tool handler must have received the elicited response") +} + +// TestElicitation_WithoutContinuousListening_Documented documents the +// limitation (issue #156, item U4): with the shim client's default config +// (DisableStandaloneSSE: true, i.e. NO transport.WithContinuousListening), the +// client never opens a standalone SSE stream. Under JSONResponse the go-sdk +// routes a server->client elicitation request to that (absent) stream, so the +// write is rejected and the server's Elicit call fails; the tool call surfaces +// an error rather than completing. This asserts that expected failure so the +// limitation is regression-tested: if a future go-sdk change made elicitation +// work without a standalone stream, this test would start failing and the +// decision docs would need updating. +func TestElicitation_WithoutContinuousListening_Documented(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + srv := server.NewMCPServer("elicit-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool( + mcp.NewTool(elicitToolName, mcp.WithDescription("elicit then answer")), + func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + _, err := srv.RequestElicitation(ctx, mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{ + Message: elicitMessage, + RequestedSchema: map[string]any{ + schemaType: objectType, + properties: map[string]any{ + nameKey: map[string]any{schemaType: stringType}, + }, + }, + }, + }) + if err != nil { + return nil, err + } + return mcp.NewToolResultText("ok"), nil + }, + ) + + httpSrv := server.NewStreamableHTTPServer(srv) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // Default config: no WithContinuousListening, so DisableStandaloneSSE is + // true and no standalone SSE stream is opened. An elicitation handler is + // still installed so the capability is declared (otherwise the failure + // would be "client does not support elicitation" rather than the + // stream-delivery failure this test documents). + c, err := client.NewStreamableHttpClientWithOpts( + ts.URL, + nil, + []client.ClientOption{client.WithElicitationHandler(client.ElicitationHandlerFunc( + func(_ context.Context, _ mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + return &mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{ + Action: mcp.ElicitationResponseActionAccept, + Content: map[string]any{nameKey: "grace"}, + }, + }, nil + }, + ))}, + ) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + _, err = c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: elicitToolName}}) + require.Error(t, err, "elicitation must fail without a standalone SSE stream (documented limitation)") + // Distinguish "elicitation was rejected" from "the test timed out": a + // context.DeadlineExceeded would mean elicitation hung rather than failed, + // and a future go-sdk change making elicitation silently succeed (no error) + // is already caught by the require.Error above. This guard catches the + // inverse regression where the failure masquerades as a timeout. + require.NotErrorIs(t, err, context.DeadlineExceeded, + "elicitation must be rejected, not surface as a context timeout") +} + +// TestElicitation_HandlerReturnsError verifies the error-return path of +// elicitation: when the client's elicitation handler returns an error, the +// server's RequestElicitation surfaces it and the surrounding tools/call fails +// wrapping that error. go-sdk carries the elicitation handler's error back to +// the server-side Elicit call, so the tool handler returns it and the +// CallTool result carries it. +func TestElicitation_HandlerReturnsError(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + srv := server.NewMCPServer("elicit-err-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool( + mcp.NewTool(elicitToolName, mcp.WithDescription("elicit then answer")), + func(ctx context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + _, err := srv.RequestElicitation(ctx, mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{ + Message: elicitMessage, + RequestedSchema: map[string]any{ + schemaType: objectType, + properties: map[string]any{ + nameKey: map[string]any{schemaType: stringType}, + }, + }, + }, + }) + if err != nil { + return nil, err + } + return mcp.NewToolResultText("ok"), nil + }, + ) + + httpSrv := server.NewStreamableHTTPServer(srv) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + c, err := client.NewStreamableHttpClientWithOpts( + ts.URL, + []transport.StreamableHTTPCOption{transport.WithContinuousListening()}, + []client.ClientOption{client.WithElicitationHandler(client.ElicitationHandlerFunc( + func(_ context.Context, _ mcp.ElicitationRequest) (*mcp.ElicitationResult, error) { + return nil, fmt.Errorf("declined") + }, + ))}, + ) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + _, err = c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: elicitToolName}}) + require.Error(t, err, "tool call must fail when the elicitation handler returns an error") + assert.Contains(t, err.Error(), "declined", "the CallTool error must wrap the handler's error") +} + +// TestCallUnknownTool_ErrorContainsToolName is a regression test for the +// translateUnknownToolError fix: calling a tool that does not exist over the +// Streamable HTTP transport must surface an error naming the missing tool +// (e.g. `tool "nope" not found`), NOT the empty-name `tool "" not found`. +// The fix changed the type assertion in translateUnknownToolError from +// *gosdk.CallToolParams to *gosdk.CallToolParamsRaw so the tool name is +// populated; the HandleMessage path (request_handler_test.go) does not exercise +// translateUnknownToolError, so this drives it end-to-end through the HTTP +// transport and the go-sdk dispatch middleware. +func TestCallUnknownTool_ErrorContainsToolName(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := server.NewMCPServer("unk-tool-server", testClientVersion, server.WithToolCapabilities(false)) + srv.AddTool( + mcp.NewTool("greet", mcp.WithDescription("greet")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("hi"), 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)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + _, err = c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "nope"}}) + require.Error(t, err, "calling an unknown tool must error") + // The translated error must carry the requested tool name, proving the + // *CallToolParamsRaw assertion populated it (rather than leaving it ""). + assert.Contains(t, err.Error(), `tool "nope" not found`, + "unknown-tool error must name the requested tool, not the empty name") + assert.NotContains(t, err.Error(), `tool "" not found`, + "unknown-tool error must not carry the empty tool name") +} + +// TestRequestElicitation_NoActiveSession is a fast unit-level test (no HTTP +// server) exercising the ErrNoActiveSession guard: calling RequestElicitation +// outside any session's request context returns ErrNoActiveSession. +func TestRequestElicitation_NoActiveSession(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("elicit-unit", testClientVersion) + _, err := srv.RequestElicitation(context.Background(), mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{Message: "hi"}, + }) + require.ErrorIs(t, err, server.ErrNoActiveSession) +} + +// TestRequestElicitation_SessionWithoutElicitationSupport is a fast unit-level +// test exercising the ErrElicitationNotSupported guard: a session bound via +// WithContext that does NOT implement SessionWithElicitation causes +// RequestElicitation to return ErrElicitationNotSupported (not a nil-panic or +// a silent success). +func TestRequestElicitation_SessionWithoutElicitationSupport(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("elicit-unit", testClientVersion) + // fakeSession implements ClientSession but NOT SessionWithElicitation. + ctx := srv.WithContext(context.Background(), &fakeSession{id: "sess-no-elicit"}) + _, err := srv.RequestElicitation(ctx, mcp.ElicitationRequest{ + Params: mcp.ElicitationParams{Message: "hi"}, + }) + require.ErrorIs(t, err, server.ErrElicitationNotSupported) +} diff --git a/mcpcompat/server/session.go b/mcpcompat/server/session.go index 73dfc48..9872a38 100644 --- a/mcpcompat/server/session.go +++ b/mcpcompat/server/session.go @@ -29,11 +29,13 @@ type ClientSession interface { // 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. +// Overlays set here are stored and reconciled onto the session's live go-sdk +// server (syncSessionTools) once a server is bound to the session (at +// registration, or immediately if already bound): added tools are registered +// via srv.AddTool and removed tools via srv.RemoveTools, which also drives +// go-sdk's automatic notifications/tools/list_changed emission when +// WithToolCapabilities(true) is set. So per-session tool changes take effect +// on an already-connected session at runtime. type SessionWithTools interface { ClientSession // GetSessionTools returns the session's tools. Thread-safe. diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index 63b10fd..606b993 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -136,10 +136,25 @@ func WithEndpointPath(endpointPath string) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.endpointPath = endpointPath } } -// WithSessionIdManager supplies a session ID manager. +// WithSessionIdManager supplies a session ID manager that drives cross-replica +// session lifecycle. // -// 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. +// The manager IS wired into the SDK's session lifecycle: +// - Generate is installed as the go-sdk Server's GetSessionID (buildServer), +// so the SDK mints session IDs through it — this is where ToolHive's manager +// creates the placeholder session record that OnRegisterSession later +// promotes. +// - Validate is called on EVERY request carrying a session ID: for local +// sessions (initialized on this instance) it is validated per-request so a +// session terminated in the shared store (e.g. via a DELETE on another +// replica) is rejected here and its local bookkeeping dropped; for +// cross-replica sessions it drives lazy eviction on each request. +// - Terminate is forwarded the session ID on a DELETE so ToolHive's shared +// session storage is cleaned up in lockstep with the local session. +// +// The single-replica fast path (no manager) is preserved: there is nothing to +// validate against, so local sessions route straight through to the go-sdk +// handler. func WithSessionIdManager(manager SessionIdManager) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.sessionIDMgr = manager } } @@ -176,6 +191,29 @@ func WithHTTPContextFunc(fn HTTPContextFunc) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.contextFunc = fn } } +// build constructs the go-sdk StreamableHTTPHandler. It is idempotent (once). +// +// Decision (issue #156, item U4): JSONResponse is ON by design and is NOT +// configurable. mcp-go's server replied to POSTs with application/json, and +// ToolHive callers json.Decode the response body, so the handler must keep +// replying with application/json rather than text/event-stream. +// +// Consequence for elicitation (and any server->client request made during +// request handling, e.g. sampling): under JSONResponse the go-sdk cannot +// inline a server->client request on the POST response stream, so it routes +// such messages to the standalone SSE stream (the "" stream) — see +// streamable.go's streamableServerConn.Write: when c.jsonResponse is set and +// the message is not a response, relatedRequest is reset to the empty id, +// which selects c.streams[""]. If no standalone SSE stream exists (the client +// never opened one), that stream is nil and the write is rejected +// (ErrRejected: "write to closed stream"), so the server's Elicit call fails. +// +// The shim client defaults DisableStandaloneSSE to +// !ContinuousListening(), i.e. the standalone SSE stream is OFF by default. +// For elicitation to work the client MUST open the standalone SSE stream via +// transport.WithContinuousListening() (and install an elicitation handler via +// client.WithElicitationHandler). See TestElicitation_DefaultConfig for the +// runtime verification of this contract. func (s *StreamableHTTPServer) build() { s.once.Do(func() { var gen func() string From d898410f2dd8c3768fec48c5761610bf432a3ef6 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Thu, 9 Jul 2026 09:52:10 +0300 Subject: [PATCH 6/8] fix(mcpcompat): address PR #157 review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CRITICAL — Revert WithHeartbeatInterval → ServerOptions.KeepAlive wiring: go-sdk's KeepAlive sends an active ping request that closes sessions without a connected standalone SSE stream, incompatible with JSONResponse mode + JSON-only clients. WithHeartbeatInterval is now a documented no-op (passive SSE-comment design is TODO). HIGH — Recover global AddTool panics in buildServer so a non-object schema tool doesn't poison sync.Once (was nil-panic on every subsequent request). Improve normalizeObjectSchema: type-omitted object schemas with properties/required get type:object added (callable, matching mcp-go); truly non-object schemas pass through and are caught by recover. MEDIUM — Only forgetSession on genuine termination (isTerminated), not on transient Validate errors (Redis blip), preventing split-brain where a retry builds a second go-sdk session with the same ID. MEDIUM — captureErrorBody now records h.status unconditionally (regardless of body length) so empty-body 4xx (common 401/404) are classified by status, not lost to the string fallback. LOW — Document unbounded session growth gap and bridge notification POST limitation. Fix pre-existing data race: clientSession.owner converted to atomic.Pointer[MCPServer] (matching boundServer/goSession pattern) — SetSessionTools and registerAndSync were racing on the plain pointer. Tests: healthy session survives heartbeat interval (no eviction), non-object global tool doesn't poison server, type-omitted schema callable, transient Validate error doesn't forget session, empty-body 4xx classification, 10x -race clean. Refs #156 --- mcpcompat/client/errorbody.go | 26 +++- mcpcompat/client/errors_internal_test.go | 38 +++++ mcpcompat/server/rehydration_test.go | 71 +++++++++ mcpcompat/server/server.go | 152 ++++++++++++++----- mcpcompat/server/server_internal_test.go | 37 ++++- mcpcompat/server/server_test.go | 177 +++++++++++++++++------ mcpcompat/server/session.go | 24 ++- mcpcompat/server/transports.go | 71 +++++++-- 8 files changed, 490 insertions(+), 106 deletions(-) diff --git a/mcpcompat/client/errorbody.go b/mcpcompat/client/errorbody.go index 033912e..8a13ddf 100644 --- a/mcpcompat/client/errorbody.go +++ b/mcpcompat/client/errorbody.go @@ -47,19 +47,39 @@ func capturedErr(ctx context.Context) *errBody { // 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). +// +// The HTTP status code is recorded regardless of whether the body is non-empty +// (issue #156, finding 4): a 4xx with an empty body — common for 401 and 404 — +// still carries an authoritative status code, and the status-driven +// classification in mapTransportError/mapConnectError relies on h.status being +// set. Previously h.status was only set when the body was non-empty, so an +// empty-body 4xx left h.status == 0 and classification fell to the +// best-effort string fallback, which does not map bare "400 Bad Request"/"404 +// Not Found" on initialize to ErrLegacySSEServer. 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 != "" { + if h == nil { + return + } + // Record the status code unconditionally: it is available from the response + // regardless of body length, and the status-driven classification depends on + // it. Do this before reading the body so a read error cannot prevent it. + h.status = resp.StatusCode + // Only capture the body once (the first non-2xx response on this context). + if h.body != "" { + // Body already captured; still restore resp.Body for downstream readers. + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(nil)) 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 + if err == nil { h.body = string(data) } } diff --git a/mcpcompat/client/errors_internal_test.go b/mcpcompat/client/errors_internal_test.go index 7fe55ce..a01e9eb 100644 --- a/mcpcompat/client/errors_internal_test.go +++ b/mcpcompat/client/errors_internal_test.go @@ -135,6 +135,16 @@ func TestMapCallError_Transport401(t *testing.T) { wantErr: transport.ErrUnauthorized, notErr: transport.ErrLegacySSEServer, }, + { + // issue #156, finding 4: a 401 with an EMPTY body must still + // classify as ErrUnauthorized via the status code (not fall back to + // string matching, which misses bare "401 Unauthorized"). + name: "empty-body 401", + status: http.StatusUnauthorized, + body: "", + wantErr: transport.ErrUnauthorized, + notErr: transport.ErrLegacySSEServer, + }, { name: "captured 404 session", status: http.StatusNotFound, @@ -196,6 +206,15 @@ func TestMapConnectError_4xxOnInitialize_LegacySSE(t *testing.T) { wantErr: transport.ErrLegacySSEServer, notErr: transport.ErrSessionTerminated, }, + { + // issue #156, finding 4: a 404 with an EMPTY body on initialize + // must still classify as ErrLegacySSEServer via the status code. + name: "empty-body 404 on initialize", + status: http.StatusNotFound, + body: "", + wantErr: transport.ErrLegacySSEServer, + notErr: transport.ErrSessionTerminated, + }, { name: "400 bad request", status: http.StatusBadRequest, @@ -203,6 +222,15 @@ func TestMapConnectError_4xxOnInitialize_LegacySSE(t *testing.T) { wantErr: transport.ErrLegacySSEServer, notErr: transport.ErrUnauthorized, }, + { + // issue #156, finding 4: a 400 with an EMPTY body on initialize + // must still classify as ErrLegacySSEServer via the status code. + name: "empty-body 400 on initialize", + status: http.StatusBadRequest, + body: "", + wantErr: transport.ErrLegacySSEServer, + notErr: transport.ErrUnauthorized, + }, { name: "403 forbidden on initialize", status: http.StatusForbidden, @@ -217,6 +245,16 @@ func TestMapConnectError_4xxOnInitialize_LegacySSE(t *testing.T) { wantErr: transport.ErrUnauthorized, notErr: transport.ErrLegacySSEServer, }, + { + // issue #156, finding 4: a 401 with an EMPTY body on connect must + // still classify as ErrUnauthorized via the status code (so OAuth + // refresh triggers), not fall back to string matching. + name: "empty-body 401 on connect", + status: http.StatusUnauthorized, + body: "", + wantErr: transport.ErrUnauthorized, + notErr: transport.ErrLegacySSEServer, + }, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { diff --git a/mcpcompat/server/rehydration_test.go b/mcpcompat/server/rehydration_test.go index 118a1ad..25e743e 100644 --- a/mcpcompat/server/rehydration_test.go +++ b/mcpcompat/server/rehydration_test.go @@ -9,6 +9,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" "strings" @@ -474,3 +475,73 @@ func TestLocalSession_ServedWithoutManager(t *testing.T) { require.Contains(t, listToolNames(t, ts.URL, sid), "greet", "local session must be served in single-replica mode without a manager") } + +// flakySessionManager is a sharedSessionManager whose Validate can be made to +// return a transient (non-terminated) error on demand, simulating a Redis blip. +type flakySessionManager struct { + sharedSessionManager + mu sync.Mutex + flakeErr error // when non-nil, Validate returns (false, flakeErr) +} + +func (m *flakySessionManager) Validate(sessionID string) (bool, error) { + m.mu.Lock() + err := m.flakeErr + m.flakeErr = nil // one-shot: a single flake, then recover + m.mu.Unlock() + if err != nil { + return false, err + } + return m.sharedSessionManager.Validate(sessionID) +} + +func (m *flakySessionManager) flake(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.flakeErr = err +} + +// TestLocalSession_TransientValidateError_DoesNotForgetSession verifies the fix +// for issue #156, finding 3: a transient Validate error (e.g. a Redis blip) +// must NOT drop local session bookkeeping. If it did, the live go-sdk session +// would remain in the handler's map while the local marker is gone, so the +// client's retry would take the !isLocalSession branch and build a SECOND +// go-sdk session with the same ID (split-brain). Instead the request is +// rejected (503) and the local session stays local; once Validate succeeds +// again the request is served normally through the SAME local go-sdk session. +func TestLocalSession_TransientValidateError_DoesNotForgetSession(t *testing.T) { + t.Parallel() + mgr := &flakySessionManager{sharedSessionManager: *newSharedSessionManager()} + + stream := server.NewMCPServer("flaky-A", "1.0.0") + addGreetTool(stream) + s := server.NewStreamableHTTPServer(stream, server.WithSessionIdManager(mgr)) + ts := httptest.NewServer(s) + defer ts.Close() + + // Initialize locally — this marks the session local and hands its lifecycle + // to the go-sdk StreamableHTTPHandler. + sid := initSession(t, ts.URL) + require.Contains(t, listToolNames(t, ts.URL, sid), "greet", + "local session must be served before the flake") + + // Make Validate return a transient (non-terminated) error once. + mgr.flake(fmt.Errorf("redis blip: connection refused")) + + // The request must be rejected (503), NOT 404, and NOT served. + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + resp := postRPC(ctx, t, ts.URL, sid, `{"jsonrpc":"2.0","id":9,"method":"tools/list","params":{}}`) + assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode, + "a transient Validate error must reject with 503, not 404 (which would imply local state was dropped)") + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + // The session must STILL be local: a subsequent request (Validate now + // succeeding) must be served normally through the same local go-sdk session, + // NOT routed through rehydration. If local state had been dropped, this + // would either 404 (unknown to the shared store's rehydration path) or build + // a second session; being served proves the local marker survived. + require.Contains(t, listToolNames(t, ts.URL, sid), "greet", + "local session must survive a transient Validate error and be served once it recovers") +} diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index 042c4db..5ec3baa 100644 --- a/mcpcompat/server/server.go +++ b/mcpcompat/server/server.go @@ -145,6 +145,19 @@ type MCPServer struct { resourceTemplates map[string]ServerResourceTemplate prompts map[string]ServerPrompt + // sessions holds every clientSession created on this instance, keyed by + // sessionID. It is the registry consulted by contextWithSession / sessionFor + // and SendNotificationToAllClients. + // + // KNOWN GAP (issue #156, finding 5): entries are only removed on the DELETE + // path and on a Validate-failure/termination path (see forgetSession). A + // session whose client vanishes, or that the go-sdk handler closes internally + // (e.g. on a transport error), never has forgetSession called, so its entry + // leaks for the process lifetime and SendNotificationToAllClients iterates + // corpses. A reaping mechanism (e.g. a periodic sweep that drops entries + // whose go-sdk ServerSession is closed, or a go-sdk close callback wired into + // forgetSession) is needed; this is lower priority and tracked separately. + // Do not over-engineer here without the upstream close hook. sessions sync.Map // sessionID -> *clientSession // localSessions records the IDs of sessions that were initialized on THIS @@ -154,6 +167,10 @@ type MCPServer struct { // 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). + // + // Shares the same unbounded-growth gap as sessions above (finding 5): + // entries are dropped only via forgetSession (DELETE / Validate-termination), + // not on a vanished client or an internal go-sdk close. localSessions sync.Map // sessionID -> struct{} // pendingReqCtx bridges per-request context values (identity, audit @@ -393,7 +410,22 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { // 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, keepalive time.Duration) (*gosdk.Server, error) { +func (s *MCPServer) buildServer(genSessionID func() string, _ time.Duration) (*gosdk.Server, error) { + // NOTE: the keepalive parameter is intentionally unused. Wave 3 wired + // WithHeartbeatInterval to go-sdk's ServerOptions.KeepAlive, but go-sdk's + // KeepAlive sends an active ping REQUEST (server→client) each interval and + // session.Close()s on failure. Under JSONResponse mode that ping routes to + // the standalone SSE stream; a JSON-only client (no GET stream, the shim + // client's own default) has that stream unconnected → ping rejected → + // session closed on the first tick. mcp-go's heartbeat was passive (pings + // written only to an existing GET stream, never awaited, never + // terminating). Wiring KeepAlive therefore evicts every healthy JSON-only + // client session at the heartbeat interval. The interval is stored by + // WithHeartbeatInterval (see transports.go) but NOT wired to KeepAlive. + // + // TODO(issue #156): implement a passive keep-alive (SSE comment injection + // via a ResponseWriter wrapper around the go-sdk handler) matching + // mcp-go's design, rather than the active ping go-sdk's KeepAlive performs. s.mu.RLock() tools := make(map[string]ServerTool, len(s.tools)) for k, v := range s.tools { @@ -420,11 +452,10 @@ func (s *MCPServer) buildServer(genSessionID func() string, keepalive time.Durat // DefaultPageSize (1000) in place, preserving the pre-existing behavior // for callers that do not set WithPageSize. PageSize: s.pageSize, - // KeepAlive wires WithHeartbeatInterval (streamable HTTP transport) to - // go-sdk's per-session keep-alive ping: the SDK pings the peer at this - // interval and closes the session if a ping goes unanswered. stdio/SSE - // pass 0 (no keep-alive), matching mcp-go. - KeepAlive: keepalive, + // KeepAlive is intentionally NOT set: see the note on buildServer's + // keepalive parameter above. go-sdk's KeepAlive sends an active ping + // request that closes sessions without a connected standalone SSE + // stream, incompatible with JSONResponse mode + JSON-only clients. InitializedHandler: func(ctx context.Context, req *gosdk.InitializedRequest) { if req == nil || req.Session == nil { return @@ -456,12 +487,22 @@ func (s *MCPServer) buildServer(genSessionID func() string, keepalive time.Durat } srv = gosdk.NewServer(impl, opts) + // Register global tools/resources/prompts. go-sdk's AddTool/AddResource/ + // AddPrompt panic on a malformed schema (e.g. a $ref or non-object type), + // and buildServer runs inside sync.Once.Do in the transports' build(): a + // panic here would mark Once done while buildErr stays nil and handler stays + // nil, so every subsequent request nil-panics forever. Recover such panics + // and convert them into an error so they flow into buildErr and properly + // poison the Once (issue #156, finding 2). This mirrors the per-session + // recover in addSessionTool (session.go). 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)) + if err := addGlobalTool(srv, gt, s.wrapToolHandler(st.Handler), st.Tool.Name); err != nil { + return nil, err + } } for _, sr := range resources { gr := &gosdk.Resource{} @@ -482,11 +523,29 @@ func (s *MCPServer) buildServer(genSessionID func() string, keepalive time.Durat return srv, nil } +// addGlobalTool registers a globally-declared tool on a go-sdk server, recovering +// the panic go-sdk's AddTool raises when a tool's input schema is non-nil but not +// top-level type:"object" (e.g. $ref, oneOf, boolean). It converts the panic +// into a returned error so buildServer surfaces a clean construction-time error +// (flowing into buildErr/sync.Once) rather than poisoning the server's once with +// a nil-handler-nil-error state. Mirrors addSessionTool (session.go), which +// recovers and skips for per-session overlays. +func addGlobalTool(srv *gosdk.Server, gt *gosdk.Tool, h gosdk.ToolHandler, name string) (err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("registering global tool %q: go-sdk AddTool rejected its input schema: %v", name, r) + } + }() + srv.AddTool(gt, h) + return nil +} + // 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; keepalive is the -// per-session keep-alive ping interval (0 disables). On a build error it logs -// and returns nil, which the go-sdk handler surfaces as an HTTP 400. +// session-ID generator to install on each per-session server. The keepalive +// parameter is retained for signature stability but is NOT wired to go-sdk's +// KeepAlive (see buildServer); it is effectively ignored. 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, keepalive time.Duration) func(*http.Request) *gosdk.Server { return func(*http.Request) *gosdk.Server { srv, err := s.buildServer(genSessionID, keepalive) @@ -711,26 +770,19 @@ func toGoSDKTool(t mcp.Tool) (*gosdk.Tool, error) { } // normalizeObjectSchema ensures a tool input schema is suitable for go-sdk's -// AddTool, which panics on a nil InputSchema. A nil schema, an empty map, an -// empty string, or a map whose "type" is the empty string (the value mcp-go's -// ToolInputSchema marshals to when no type is declared) is replaced with the -// empty object schema ({"type":"object"}). Every other schema is returned -// verbatim, mirroring mcp-go, which passed RawInputSchema through unchanged — -// including object schemas with "type" omitted, $ref, oneOf/anyOf, and boolean -// schemas. -// -// TODO(upstream): go-sdk v1.6.1 AddTool additionally panics unless the -// schema's top-level "type" is literally "object", so a non-object schema that -// passes through here (e.g. $ref, oneOf, a boolean schema, or an object schema -// with type omitted) will panic at registration time. We pass such schemas -// through verbatim (rather than silently rewriting them) so callers surface the -// incompatibility rather than silently receiving a different schema than the one -// they declared; mcp-go performed no validation here either. The two registration -// paths handle the panic differently: the per-session path -// (MCPServer.addSessionTool, used by syncSessionTools) recovers and skips the -// offending tool so one bad overlay schema cannot crash a live session; the -// global path (buildServer) lets the panic surface as a construction-time error -// (a 500 at handler build). +// AddTool, which panics on a nil InputSchema and (v1.6.1) unless the schema's +// top-level "type" is literally "object". The following are normalized to +// {"type":"object"} (preserving any existing fields): a nil schema, an empty +// map, an empty string, a map whose "type" is the empty string (the value +// mcp-go's ToolInputSchema marshals to when no type is declared), and a map +// with NO "type" key (or type:"") that DOES carry "properties" or "required" — +// a spec-loose but common shape ({"properties":...}) that mcp-go served +// verbatim and callable. Forcing type:"object" here makes such schemas callable +// under go-sdk, matching mcp-go. Truly non-object schemas ($ref, oneOf, a +// boolean schema, type:"string", type:["object","null"], etc.) pass through +// verbatim; go-sdk's AddTool will panic on them at registration time, which the +// registration paths recover from (addGlobalTool for globals surfaces a clean +// construction error; addSessionTool skips the offending per-session tool). func normalizeObjectSchema(schema any) any { switch s := schema.(type) { case nil: @@ -744,12 +796,30 @@ func normalizeObjectSchema(schema any) any { // "no schema declared"; normalize it to "object" so AddTool accepts the // tool (matching mcp-go's leniency for tools with no declared input). if t, ok := s[schemaTypeKey].(string); ok && t == "" { - cpy := make(map[string]any, len(s)) - for k, v := range s { - cpy[k] = v + // If the schema declares properties/required, this is an object + // schema whose type was simply omitted (spec-loose but common); + // normalize type to "object" preserving the other fields so the tool + // is callable under go-sdk. Otherwise (no properties/required) it's + // a truly empty-declared schema, normalized to the bare object. + if _, hasProps := s["properties"]; hasProps { + return withTypeObject(s) + } + if _, hasReq := s["required"]; hasReq { + return withTypeObject(s) + } + return map[string]any{schemaTypeKey: schemaTypeObject} + } + // No "type" key at all: if the schema declares properties or required, + // treat it as an object schema with type omitted (spec-loose but common) + // and add type:"object" so go-sdk's AddTool accepts it. This matches + // mcp-go, which served such schemas verbatim and callable. + if _, hasType := s[schemaTypeKey]; !hasType { + if _, hasProps := s["properties"]; hasProps { + return withTypeObject(s) + } + if _, hasReq := s["required"]; hasReq { + return withTypeObject(s) } - cpy[schemaTypeKey] = schemaTypeObject - return cpy } return s case string: @@ -762,6 +832,18 @@ func normalizeObjectSchema(schema any) any { } } +// withTypeObject returns a copy of s with the top-level "type" set to "object", +// preserving all other fields. Used for spec-loose object schemas whose type +// was omitted or empty. +func withTypeObject(s map[string]any) map[string]any { + cpy := make(map[string]any, len(s)+1) + for k, v := range s { + cpy[k] = v + } + cpy[schemaTypeKey] = schemaTypeObject + return cpy +} + // 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) diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index 45ba378..45aa5f8 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/mcpcompat/server/server_internal_test.go @@ -179,11 +179,15 @@ func (fakeIDManager) Terminate(string) (bool, error) { return false, nil } var _ SessionIdManager = fakeIDManager{} -// TestNormalizeObjectSchema verifies that only nil/empty input schemas are -// normalized to {"type":"object"}; all other schemas pass through verbatim, -// matching mcp-go (which passed RawInputSchema through unchanged). Previously -// any schema whose top-level "type" was not "object" was clobbered, stripping -// valid $ref/oneOf/boolean/object-with-type-omitted schemas. +// TestNormalizeObjectSchema verifies the normalization rules in +// normalizeObjectSchema: nil/empty/empty-type schemas become {"type":"object"}; +// a type-omitted (or empty-type) schema that carries "properties" or +// "required" is treated as a spec-loose object schema and gets type:"object" +// added (preserving the other fields), so it is callable under go-sdk (matching +// mcp-go's verbatim, callable behavior). Truly non-object schemas ($ref, oneOf, +// boolean, type:["object","null"], a map with no type and no +// properties/required) pass through verbatim; go-sdk's AddTool will panic on +// the non-callable ones, recovered by addGlobalTool/addSessionTool. // //nolint:goconst // test fixtures legitimately repeat schema keys func TestNormalizeObjectSchema(t *testing.T) { @@ -239,18 +243,39 @@ func TestNormalizeObjectSchema(t *testing.T) { }, }, { - name: "object schema with type omitted", + name: "object schema with type omitted gets type object", input: map[string]any{ "properties": map[string]any{ "name": map[string]any{schemaTypeKey: "string"}, }, }, want: map[string]any{ + schemaTypeKey: schemaTypeObject, "properties": map[string]any{ "name": map[string]any{schemaTypeKey: "string"}, }, }, }, + { + name: "object schema with type omitted and required gets type object", + input: map[string]any{ + "required": []any{"name"}, + }, + want: map[string]any{ + schemaTypeKey: schemaTypeObject, + "required": []any{"name"}, + }, + }, + { + name: "type omitted with no properties/required passes through", + input: map[string]any{"title": "t"}, + want: map[string]any{"title": "t"}, + }, + { + name: "type object null array passes through verbatim", + input: map[string]any{schemaTypeKey: []any{schemaTypeObject, "null"}}, + want: map[string]any{schemaTypeKey: []any{schemaTypeObject, "null"}}, + }, { name: "$ref schema", input: map[string]any{"$ref": "#"}, diff --git a/mcpcompat/server/server_test.go b/mcpcompat/server/server_test.go index b9d4dcc..97adf87 100644 --- a/mcpcompat/server/server_test.go +++ b/mcpcompat/server/server_test.go @@ -5,6 +5,7 @@ package server_test import ( "context" + "encoding/json" "fmt" "io" "net/http" @@ -469,20 +470,18 @@ func TestDisableLocalhostProtection_DefaultRejectsNonLocalhostHost(t *testing.T) assert.Equal(t, http.StatusForbidden, resp.StatusCode, "default localhost protection must reject a non-localhost Host") } -// TestHeartbeat_WiredToKeepAlive verifies that WithHeartbeatInterval is wired to -// go-sdk's per-session keep-alive (issue #156, item 2) via an e2e check: a -// server built with a heartbeat disconnects a session whose transport stops -// responding to pings. go-sdk pings the peer at KeepAlive and closes the session -// when a ping goes unanswered, observable as the client's stream/sessions -// closing. Because ServerOptions.KeepAlive is not observable directly, this is -// an e2e test. -// -// We observe the wiring through the resulting server's session lifecycle: a -// raw JSON-RPC initialize against a server configured with a short heartbeat, -// after which the client stops reading, must lead to the session being closed -// within a bounded window. This proves the heartbeat is no longer a no-op. -func TestHeartbeat_WiredToKeepAlive(t *testing.T) { +// TestHeartbeat_NoOpDoesNotEvictHealthySession documents the current state of +// WithHeartbeatInterval (issue #156): it stores the interval but does NOT wire +// it to go-sdk's ServerOptions.KeepAlive, because go-sdk's KeepAlive sends an +// active ping request that closes JSON-only client sessions (no standalone SSE +// stream) on the first tick. This test asserts the inverse of the old +// session-killing behavior: an idle, healthy client session configured with a +// short heartbeat (150ms) must still be alive well past the heartbeat interval +// (500ms), proving the interval is not wired to an evicting KeepAlive. If a +// future passive keep-alive is wired here, this test will need updating. +func TestHeartbeat_NoOpDoesNotEvictHealthySession(t *testing.T) { t.Parallel() + ctx := context.Background() srv := server.NewMCPServer("hb-server", testClientVersion, server.WithToolCapabilities(false)) srv.AddTool(mcp.NewTool("ping", mcp.WithDescription("ping")), @@ -494,35 +493,30 @@ func TestHeartbeat_WiredToKeepAlive(t *testing.T) { ts := httptest.NewServer(httpSrv) t.Cleanup(ts.Close) - // Initialize a session via raw JSON-RPC, then stop reading the response - // stream entirely (abandon the connection). With KeepAlive wired, go-sdk's - // per-session pinger will fail to get a pong back and must close the session - // within a bounded window. We assert the session is gone by issuing a - // follow-up request with the same session id: once the SDK has evicted the - // dead session, it 404s the id (the go-sdk handler does not recognize it). - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) - sid := initSession(t, ts.URL) - - // Give the keep-alive pinger time to fire and evict the unresponsive session. - // The heartbeat is 150ms; allow a generous multiple for the ping timeout + - // eviction. - deadline := time.Now().Add(3 * time.Second) - var got404 bool - for time.Now().Before(deadline) { - resp := postRPC(ctx, t, ts.URL, sid, `{"jsonrpc":"2.0","id":7,"method":"tools/list","params":{}}`) - sc := resp.StatusCode - _, _ = io.Copy(io.Discard, resp.Body) - _ = resp.Body.Close() - if sc == http.StatusNotFound { - got404 = true - break - } - // A 200 means the session is still alive; back off briefly and retry. - time.Sleep(100 * time.Millisecond) - } - assert.True(t, got404, "session must be evicted by keep-alive within the bounded window once the client stops responding to pings") + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + // Idle past the heartbeat interval (150ms). Wait 500ms — well over 3 ticks — + // so an evicting KeepAlive would have fired by now. + time.Sleep(500 * time.Millisecond) + + // The session must still be alive and serve the call. + res, err := c.CallTool(ctx, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "ping"}}) + require.NoError(t, err, "healthy idle session must survive past the heartbeat interval (no-op KeepAlive)") + require.False(t, res.IsError) + txt, ok := mcp.AsTextContent(res.Content[0]) + require.True(t, ok) + assert.Equal(t, "pong", txt.Text) } // TestBeforeCallTool_ReceivesToolNameAndArgs verifies that the before-call hook @@ -949,3 +943,104 @@ func TestRequestElicitation_SessionWithoutElicitationSupport(t *testing.T) { }) require.ErrorIs(t, err, server.ErrElicitationNotSupported) } + +// TestGlobalTool_NonObjectSchema_DoesNotPoisonServer is a regression test for +// issue #156, finding 2: a global tool whose RawInputSchema is a non-object +// schema (a $ref) makes go-sdk's AddTool panic at build time. Before the fix +// that panic ran inside sync.Once.Do, marking Once done while buildErr stayed +// nil and handler stayed nil — so every subsequent request nil-panicked +// forever, bricking the server. With the recover in addGlobalTool, buildServer +// returns a clean error, buildErr is set, and requests get a 500 (not a +// nil-pointer panic). A subsequent request (e.g. against a SEPARATE server with +// a valid tool) proves the server object itself is usable; this test asserts +// the poisoned server returns a clean 500 on every request rather than panicking. +func TestGlobalTool_NonObjectSchema_DoesNotPoisonServer(t *testing.T) { + t.Parallel() + + srv := server.NewMCPServer("bad-schema-server", testClientVersion, server.WithToolCapabilities(false)) + // A non-object schema ($ref). normalizeObjectSchema passes it through + // verbatim; go-sdk AddTool panics on it. addGlobalTool must recover and + // convert the panic into a build error. + srv.AddTool(mcp.NewToolWithRawSchema("bad", "bad schema", json.RawMessage(`{"$ref":"#"}`)), + 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) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + // The first request must get a clean 500 (build error), NOT a nil-pointer + // panic (which would surface as a connection reset / 502). + resp := postRPC(ctx, t, ts.URL, "", `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}`) + assert.Equal(t, http.StatusInternalServerError, resp.StatusCode, + "a bad global tool schema must surface as a clean 500, not a nil panic") + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + // A second request with the same poisoned server must ALSO get a clean 500 + // (sync.Once is properly poisoned with buildErr, not nil-handler-nil-error). + resp2 := postRPC(ctx, t, ts.URL, "", `{"jsonrpc":"2.0","id":2,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}`) + assert.Equal(t, http.StatusInternalServerError, resp2.StatusCode, + "sync.Once must remain poisoned with the build error, not nil-panic on retry") + _, _ = io.Copy(io.Discard, resp2.Body) + _ = resp2.Body.Close() +} + +// TestGlobalTool_TypeOmittedObjectSchema_Callable verifies the +// normalizeObjectSchema improvement (issue #156, finding 2): a spec-loose +// object schema with NO "type" key but WITH "properties" (a shape mcp-go +// served verbatim and callable) is normalized to include type:"object" so the +// tool registers under go-sdk and is callable end-to-end. Without the +// normalization, go-sdk's AddTool would panic (top-level type not "object") +// and the tool would fail to register. +func TestGlobalTool_TypeOmittedObjectSchema_Callable(t *testing.T) { + t.Parallel() + ctx := context.Background() + + srv := server.NewMCPServer("loose-schema-server", testClientVersion, server.WithToolCapabilities(false)) + // Object schema with type OMITTED but properties present. mcp-go served + // this verbatim and the tool was callable; normalizeObjectSchema now adds + // type:"object" so go-sdk's AddTool accepts it. + srv.AddTool(mcp.NewToolWithRawSchema("loose", "loose object schema", + json.RawMessage(`{"properties":{"name":{"type":"string"}},"required":["name"]}`)), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + name := req.GetString("name", "world") + return mcp.NewToolResultText("hi " + 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)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err, "server with a type-omitted object schema tool must build and initialize") + + // The tool must be advertised and callable. + tools, err := c.ListTools(ctx, mcp.ListToolsRequest{}) + require.NoError(t, err) + require.Len(t, tools.Tools, 1) + assert.Equal(t, "loose", tools.Tools[0].Name) + + res, err := c.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "loose", Arguments: map[string]any{nameKey: "loose-user"}}, + }) + require.NoError(t, err, "type-omitted object schema tool must be callable") + require.False(t, res.IsError) + txt, ok := mcp.AsTextContent(res.Content[0]) + require.True(t, ok) + assert.Equal(t, "hi loose-user", txt.Text) +} diff --git a/mcpcompat/server/session.go b/mcpcompat/server/session.go index 9872a38..5dbb283 100644 --- a/mcpcompat/server/session.go +++ b/mcpcompat/server/session.go @@ -78,7 +78,13 @@ type clientSession struct { // 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 + // + // owner is an atomic.Pointer (not a plain field) because SetSessionTools/ + // SetSessionResources may run from any goroutine (test code, before-hooks on + // a live request) concurrently with registerAndSync/bindRehydratedSession, + // which set owner on the session's connection goroutine. boundServer uses the + // same atomic pattern for the same reason. + owner atomic.Pointer[MCPServer] boundServer atomic.Pointer[gosdk.Server] mu sync.RWMutex @@ -129,8 +135,10 @@ func (c *clientSession) SetSessionTools(tools map[string]ServerTool) { } 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) + if srv := c.boundServer.Load(); srv != nil { + if owner := c.owner.Load(); owner != nil { + owner.syncSessionTools(srv, c) + } } } @@ -151,8 +159,10 @@ func (c *clientSession) SetSessionResources(resources map[string]ServerResource) c.resources[k] = v } c.mu.Unlock() - if srv := c.boundServer.Load(); srv != nil && c.owner != nil { - c.owner.syncSessionResources(srv, c) + if srv := c.boundServer.Load(); srv != nil { + if owner := c.owner.Load(); owner != nil { + owner.syncSessionResources(srv, c) + } } } @@ -192,7 +202,7 @@ func (s *MCPServer) registerAndSync(ctx context.Context, ss *gosdk.ServerSession } cs := s.sessionFor(ss.ID()) cs.goSession.Store(ss) - cs.owner = s + cs.owner.Store(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 @@ -320,7 +330,7 @@ func (s *MCPServer) forgetSession(id string) { func (s *MCPServer) bindRehydratedSession(id string, ss *gosdk.ServerSession, srv *gosdk.Server) { cs := s.sessionFor(id) cs.goSession.Store(ss) - cs.owner = s + cs.owner.Store(s) cs.boundServer.Store(srv) cs.Initialize() } diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index 606b993..a62a395 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -159,12 +159,20 @@ func WithSessionIdManager(manager SessionIdManager) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.sessionIDMgr = manager } } -// WithHeartbeatInterval sets the keep-alive ping interval. This option is now -// LIVE (it was previously a no-op): it is wired to go-sdk's -// ServerOptions.KeepAlive, so the SDK pings the peer at this interval and -// closes the session when a ping goes unanswered. Set a small interval only if -// you want unanswered pings to terminate the session. Applies to the Streamable -// HTTP transport only; stdio and SSE pass 0 (no keep-alive), matching mcp-go. +// WithHeartbeatInterval sets the keep-alive ping interval. This option is +// currently a NO-OP: it stores the interval but does NOT wire it to go-sdk's +// ServerOptions.KeepAlive. go-sdk's KeepAlive sends an active ping REQUEST +// (server→client) each interval and session.Close()s on failure; under +// JSONResponse mode that ping routes to the standalone SSE stream, and a +// JSON-only client (no GET stream, the shim client's own default) has that +// stream unconnected → ping rejected → session closed on the first tick. That +// would evict every healthy JSON-only client session at the interval. mcp-go's +// heartbeat was passive (pings written only to an existing GET stream, never +// awaited, never terminating). +// +// TODO(issue #156): implement a passive keep-alive (SSE comment injection via +// a ResponseWriter wrapper around the go-sdk handler) matching mcp-go's design +// and wire this option to it. Until then the value is stored but unused. func WithHeartbeatInterval(interval time.Duration) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.heartbeat = interval } } @@ -221,8 +229,9 @@ func (s *StreamableHTTPServer) build() { 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, s.heartbeat); err != nil { + // surfaces as a clean 500 rather than a per-request nil. Pass 0 for + // keepalive: WithHeartbeatInterval is a documented no-op (see its doc). + if _, err := s.mcp.buildServer(gen, 0); err != nil { s.buildErr = err return } @@ -240,9 +249,9 @@ func (s *StreamableHTTPServer) build() { } // 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. The heartbeat interval is - // threaded into each per-session go-sdk server's KeepAlive. - s.handler = gosdk.NewStreamableHTTPHandler(s.mcp.getServerFunc(gen, s.heartbeat), opts) + // registration middleware buildServer installs. Keepalive is 0 + // (WithHeartbeatInterval is a no-op; see its doc). + s.handler = gosdk.NewStreamableHTTPHandler(s.mcp.getServerFunc(gen, 0), opts) }) } @@ -272,6 +281,18 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) nonce := crand.Text() s.mcp.setPendingRequestContext(r.Context(), nonce) r.Header.Set(reqNonceHeader, nonce) + // KNOWN LIMITATION (issue #156, finding 6): this bridge covers request/ + // response calls (tools/call, tools/list, initialize), which block in + // servePOST until the response is ready, so the deferred clear runs only + // AFTER the session goroutine has dispatched the message and read the nonce. + // It does NOT cover notification-only POSTs (e.g. notifications/initialized, + // notifications/cancel): those are acknowledged with an immediate 202 before + // the session goroutine dispatches the message, so the deferred clear can + // run before the dispatch middleware looks up the nonce → per-request + // context values are silently absent for notification handling. There is no + // cross-request bleed (the entry is keyed by a unique nonce), but the bridge + // is best-effort for notifications. A proper fix awaits go-sdk per-request + // context propagation (upstream ask). defer s.mcp.clearPendingRequestContext(nonce) // DELETE terminates the session. mcp-go answered 200 and drove the supplied // SessionIdManager's Terminate; go-sdk answers 204 and manages its own session @@ -307,12 +328,34 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) // bookkeeping dropped, rather than being served forever from the local // handler. The single-replica fast-path (no manager) is preserved: there is // nothing to validate against, so local sessions route straight through. + // + // Transient Validate errors (issue #156, finding 3): a Validate error that is + // NOT a genuine termination (e.g. a Redis blip) must NOT drop local state. If + // we called forgetSession here, the live go-sdk session would remain in the + // handler's map while the local marker is gone → the client's retry takes the + // !isLocalSession branch → serveRehydrated builds a SECOND go-sdk session with + // the same ID → split-brain (two live sessions, one ID, one pod; the standalone + // GET stream stays on the old one). So on a non-terminated Validate error we + // only reject THIS request (503, retry-able) and leave the local session + // intact; if the shared store recovers, the next Validate succeeds and the + // request is served normally. Only a genuine termination (isTerminated) drops + // local bookkeeping. if sid := r.Header.Get("Mcp-Session-Id"); sid != "" && s.sessionIDMgr != nil && s.mcp.isLocalSession(sid) { isTerminated, err := s.sessionIDMgr.Validate(sid) if err != nil { - s.mcp.forgetSession(sid) - s.deleteRehydrated(sid) - http.Error(w, "Invalid session ID", http.StatusNotFound) + if isTerminated { + // Genuinely terminated: drop local bookkeeping so a later request + // with the same ID is not mistaken for a live local session. + s.mcp.forgetSession(sid) + s.deleteRehydrated(sid) + http.Error(w, "Session terminated", http.StatusNotFound) + return + } + // Transient error (e.g. shared-store blip): reject this request + // without dropping local state. The local go-sdk session stays + // alive; the client retries and, once Validate succeeds again, is + // served normally rather than routed through rehydration. + http.Error(w, "session validation unavailable", http.StatusServiceUnavailable) return } if isTerminated { From eb170641d00648ee28e803e4e4571d41edd27401 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Mon, 13 Jul 2026 12:43:47 +0300 Subject: [PATCH 7/8] fix(mcpcompat): address PR #157 review findings (wave 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses all unresolved review findings from jhrozek (2026-07-13): 1. 5xx false-positive OAuth refresh (MEDIUM): a captured 5xx status whose body contained "401"/"unauthorized" still fell through to the string fallback → false OAuth refresh. mapTransportError now surfaces 5xx unchanged (only string-fallbacks when no status was captured at all). Refactored into mapStatusError/mapStringError for clarity. 2. Rehydrated transient Validate error inconsistency (MEDIUM): serveRehydrated returned 404 "Invalid session ID" on a transient Validate error, while the local path returned 503. Now: a cached rehydrated session with a transient error gets 503 (preserving the cache); an unknown (never cached) session still gets 404. 3. extractResourceMetadataURL stub (MINOR): the shim's headerRoundTripper has direct access to resp.Header and now captures WWW-Authenticate values into errBody. extractResourceMetadataURL parses the RFC 9728 §5.1 resource_metadata parameter (ported from mcp-go's well-tested parser) and populates AuthorizationRequiredError.ResourceMetadataURL end-to-end. No longer blocked by the go-sdk. 4. normalizeObjectSchema branch unification (NIT): unified the empty-type and no-type-key branches into a single typeIsStr/typeStr check, preserving the semantic distinction (empty-string type → normalize; absent type without properties/required → pass through). 5. notifCh drain goroutine leak (pre-existing, MINOR): forgetSession now closes notifCh via sync.Once, letting the drain goroutine exit. sessionFor re-creates a fresh channel for the same ID after forgetSession. 6. rehydrate lock scope (LOW): documented the cross-sid serialization cost of holding rehydratedMu across buildServer+Connect (not worth restructuring now per reviewer). Tests: - TestMapTransportError_5xxWithUnauthorizedBodyNotMisclassified (3 cases) - TestMapCallError_5xxSurfacesUnchanged - TestExtractResourceMetadataURL_FromCapturedHeader (5 cases) - TestMapCallError_Transport401_PopulatesResourceMetadataURL (e2e) - TestMapConnectError_Transport401_PopulatesResourceMetadataURL (e2e) - TestCaptureErrorBody_CapturesWWWAuthenticate (e2e through RoundTripper) - TestRehydratedSession_TransientValidateError_Returns503 - TestForgetSession_ClosesNotifChannel Verification: task lint (0 issues), task test (-race, all green), task license-check (clean). --- mcpcompat/client/errorbody.go | 12 +- mcpcompat/client/errors.go | 95 +++++++--- mcpcompat/client/errors_internal_test.go | 213 ++++++++++++++++++++++- mcpcompat/client/rfc9728.go | 125 +++++++++++++ mcpcompat/server/rehydration_test.go | 47 +++++ mcpcompat/server/server.go | 23 +-- mcpcompat/server/server_internal_test.go | 38 ++++ mcpcompat/server/session.go | 17 +- mcpcompat/server/transports.go | 28 ++- 9 files changed, 547 insertions(+), 51 deletions(-) create mode 100644 mcpcompat/client/rfc9728.go diff --git a/mcpcompat/client/errorbody.go b/mcpcompat/client/errorbody.go index 8a13ddf..edcea35 100644 --- a/mcpcompat/client/errorbody.go +++ b/mcpcompat/client/errorbody.go @@ -23,8 +23,9 @@ import ( const maxCapturedErrBody = 8 << 10 // 8 KiB type errBody struct { - status int - body string + status int + body string + wwwAuthHdrs []string // captured WWW-Authenticate header values (RFC 9728) } type errBodyKey struct{} @@ -68,6 +69,13 @@ func captureErrorBody(req *http.Request, resp *http.Response) { // regardless of body length, and the status-driven classification depends on // it. Do this before reading the body so a read error cannot prevent it. h.status = resp.StatusCode + // Capture WWW-Authenticate header values (RFC 9728 §5.1) so the error + // mappers can populate AuthorizationRequiredError.ResourceMetadataURL. The + // go-sdk does not surface this header on its errors; the shim's own + // RoundTripper is the only place it is available. + if vals := resp.Header.Values("WWW-Authenticate"); len(vals) > 0 { + h.wwwAuthHdrs = vals + } // Only capture the body once (the first non-2xx response on this context). if h.body != "" { // Body already captured; still restore resp.Body for downstream readers. diff --git a/mcpcompat/client/errors.go b/mcpcompat/client/errors.go index 3d3a980..8bc26c9 100644 --- a/mcpcompat/client/errors.go +++ b/mcpcompat/client/errors.go @@ -48,14 +48,14 @@ func mapConnectError(ctx context.Context, err error) error { if h != nil && h.status >= 400 && h.status < 500 { if h.status == http.StatusUnauthorized { return transport.NewError(errors.Join( - &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(err)}, + &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(ctx)}, transport.ErrUnauthorized, err, )) } return transport.NewError(errors.Join(transport.ErrLegacySSEServer, err)) } - return mapTransportError(err, h) + return mapTransportError(ctx, err, h) } // mapCallError maps an error returned by an underlying go-sdk request call onto @@ -85,7 +85,7 @@ func mapCallError(ctx context.Context, err error) error { // RPC-level error with no captured HTTP status: surface unchanged. return err } - return mapTransportError(err, h) + return mapTransportError(ctx, err, h) } // mapTransportError inspects err and, when it recognizes an HTTP auth/session @@ -98,42 +98,76 @@ func mapCallError(ctx context.Context, err error) error { // (e.g. a transport failure before any response, or a response whose body // capture was skipped), detection falls back to best-effort string matching. // +// A captured 5xx status surfaces the error unchanged: a 5xx body containing the +// substring "401" or "unauthorized" must NOT trigger ToolHive's OAuth refresh +// flow (issue #156, wave 4 review finding). The string fallback only runs when +// NO status was captured at all. +// // NOTE: the go-sdk does not currently expose a typed error carrying the HTTP // status code, so the string-matching fallback is inherently best-effort. 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, h *errBody) error { +func mapTransportError(ctx context.Context, err error, h *errBody) error { if err == nil { return nil } // Prefer the captured HTTP status over string matching: the status is // authoritative, while body text (re-attached by enrichWithResponseBody) // can legitimately contain words like "unauthorized" for non-401 responses. - if h != nil && h.status >= 400 && h.status < 500 { - switch h.status { - case http.StatusUnauthorized: - return transport.NewError(errors.Join( - &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(err)}, - transport.ErrUnauthorized, - err, - )) - case http.StatusNotFound: - return transport.NewError(errors.Join(transport.ErrSessionTerminated, err)) - default: - // Other 4xx on a regular call are not legacy-SSE (that classification - // is connect-time only, in mapConnectError); surface unchanged. - return err + if h != nil && h.status >= 400 { + if ok, mapped := mapStatusError(ctx, err, h); ok { + return mapped } } - msg := strings.ToLower(err.Error()) + // No captured status (or a captured status we don't classify, e.g. 5xx): + // fall back to best-effort string matching. A captured 5xx is surfaced + // unchanged by mapStatusError (ok=true, mapped=err), so the string fallback + // only runs when no status was captured at all. + return mapStringError(ctx, err) +} +// mapStatusError classifies an error using the captured HTTP status code. It +// returns (ok, mapped): ok=true when the status was captured and classified +// (including the 5xx pass-through), or ok=false when no status was captured (h +// is nil or h.status == 0), signaling the caller to use the string fallback. +func mapStatusError(ctx context.Context, err error, h *errBody) (bool, error) { + if h == nil || h.status == 0 { + return false, nil + } + if h.status < 400 { + return false, nil + } + if h.status >= 500 { + // A captured 5xx must not fall through to string matching: a 5xx body + // containing "401"/"unauthorized" would falsely trigger OAuth refresh. + return true, err + } + switch h.status { + case http.StatusUnauthorized: + return true, transport.NewError(errors.Join( + &transport.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(ctx)}, + transport.ErrUnauthorized, + err, + )) + case http.StatusNotFound: + return true, transport.NewError(errors.Join(transport.ErrSessionTerminated, err)) + default: + // Other 4xx on a regular call are not legacy-SSE (that classification + // is connect-time only, in mapConnectError); surface unchanged. + return true, err + } +} + +// mapStringError classifies an error using best-effort string matching on its +// message. This is the fallback path used when no HTTP status was captured +// (e.g. a transport failure before any response). +func mapStringError(ctx context.Context, err error) error { + 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.AuthorizationRequiredError{ResourceMetadataURL: extractResourceMetadataURL(ctx)}, transport.ErrUnauthorized, err, )) @@ -146,9 +180,16 @@ func mapTransportError(err error, h *errBody) error { } } -// 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 "" +// extractResourceMetadataURL parses the RFC 9728 §5.1 resource_metadata +// parameter from the WWW-Authenticate header captured by the shim's own +// headerRoundTripper (see captureErrorBody). The go-sdk does not surface the +// header on its errors, but the shim's RoundTripper has direct access to +// resp.Header and captures it into the per-call errBody holder. Returns empty +// when no header was captured or no resource_metadata parameter is present. +func extractResourceMetadataURL(ctx context.Context) string { + h := capturedErr(ctx) + if h == nil { + return "" + } + return extractResourceMetadataURLFromHeaders(h.wwwAuthHdrs) } diff --git a/mcpcompat/client/errors_internal_test.go b/mcpcompat/client/errors_internal_test.go index a01e9eb..7f1ec3e 100644 --- a/mcpcompat/client/errors_internal_test.go +++ b/mcpcompat/client/errors_internal_test.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "net/http" + "net/http/httptest" "testing" "github.com/modelcontextprotocol/go-sdk/jsonrpc" @@ -17,9 +18,14 @@ import ( "github.com/stacklok/toolhive-core/mcpcompat/client/transport" ) +// unauthorizedBody is the canonical body text used across test cases for a 401 +// response. Extracted as a constant to satisfy goconst (4+ occurrences). +const unauthorizedBody = "Unauthorized" + func TestMapTransportError_Unauthorized(t *testing.T) { t.Parallel() - err := mapTransportError(errors.New("request failed: 401 Unauthorized"), nil) + ctx := withErrCapture(context.Background()) + err := mapTransportError(ctx, errors.New("request failed: 401 Unauthorized"), nil) // ToolHive branches on all of these for its OAuth/401 handling. assert.True(t, errors.Is(err, transport.ErrUnauthorized), "errors.Is ErrUnauthorized") @@ -34,25 +40,29 @@ func TestMapTransportError_Unauthorized(t *testing.T) { func TestMapTransportError_SessionTerminated(t *testing.T) { t.Parallel() - err := mapTransportError(errors.New("server returned 404: session not found"), nil) + ctx := withErrCapture(context.Background()) + err := mapTransportError(ctx, errors.New("server returned 404: session not found"), nil) assert.True(t, errors.Is(err, transport.ErrSessionTerminated)) } func TestMapTransportError_LegacySSE(t *testing.T) { t.Parallel() - err := mapTransportError(errors.New("405 method not allowed"), nil) + ctx := withErrCapture(context.Background()) + err := mapTransportError(ctx, errors.New("405 method not allowed"), nil) assert.True(t, errors.Is(err, transport.ErrLegacySSEServer)) } func TestMapTransportError_Passthrough(t *testing.T) { t.Parallel() + ctx := withErrCapture(context.Background()) orig := errors.New("some unrelated failure") - assert.Equal(t, orig, mapTransportError(orig, nil)) + assert.Equal(t, orig, mapTransportError(ctx, orig, nil)) } func TestMapTransportError_Nil(t *testing.T) { t.Parallel() - assert.NoError(t, mapTransportError(nil, nil)) + ctx := withErrCapture(context.Background()) + assert.NoError(t, mapTransportError(ctx, nil, nil)) } // seedCapture returns a context carrying an errBody holder seeded with the @@ -131,7 +141,7 @@ func TestMapCallError_Transport401(t *testing.T) { { name: "captured 401 body", status: http.StatusUnauthorized, - body: "Unauthorized", + body: unauthorizedBody, wantErr: transport.ErrUnauthorized, notErr: transport.ErrLegacySSEServer, }, @@ -241,7 +251,7 @@ func TestMapConnectError_4xxOnInitialize_LegacySSE(t *testing.T) { { name: "401 on connect stays unauthorized", status: http.StatusUnauthorized, - body: "Unauthorized", + body: unauthorizedBody, wantErr: transport.ErrUnauthorized, notErr: transport.ErrLegacySSEServer, }, @@ -281,3 +291,192 @@ func TestMapConnectError_NoCaptureFallback(t *testing.T) { require.Error(t, mapped) assert.True(t, errors.Is(mapped, transport.ErrUnauthorized), "fallback should still detect 401: %v", mapped) } + +// TestMapTransportError_5xxWithUnauthorizedBodyNotMisclassified is the residual +// edge case from wave 4's review: a captured 5xx status whose body contains the +// substring "401" or "unauthorized" must NOT fall through to the string +// fallback and falsely trigger ToolHive's OAuth refresh flow. The status-driven +// classification only covers 4xx (< 500), so without the 5xx guard a 500 body +// saying "unauthorized to access X" would be misclassified as ErrUnauthorized. +func TestMapTransportError_5xxWithUnauthorizedBodyNotMisclassified(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status int + body string + }{ + { + name: "500 with unauthorized in body", + status: http.StatusInternalServerError, + body: "unauthorized to access internal resource", + }, + { + name: "502 with 401 in body", + status: http.StatusBadGateway, + body: "upstream returned 401 Unauthorized", + }, + { + name: "503 with unauthorized in body", + status: http.StatusServiceUnavailable, + body: "service unauthorized to handle request", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := seedCapture(tc.status, tc.body) + src := fmt.Errorf(`calling "tools/call": %s`, http.StatusText(tc.status)) + mapped := mapCallError(ctx, src) + require.Error(t, mapped) + assert.False(t, errors.Is(mapped, transport.ErrUnauthorized), + "5xx must not satisfy ErrUnauthorized: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrAuthorizationRequired), + "5xx must not satisfy ErrAuthorizationRequired: %v", mapped) + }) + } +} + +// TestMapCallError_5xxSurfacesUnchanged confirms a captured 5xx status surfaces +// the error unchanged (not reclassified as any transport sentinel). +func TestMapCallError_5xxSurfacesUnchanged(t *testing.T) { + t.Parallel() + ctx := seedCapture(http.StatusInternalServerError, "internal server error") + src := fmt.Errorf(`calling "tools/call": %s`, http.StatusText(http.StatusInternalServerError)) + mapped := mapCallError(ctx, src) + require.Error(t, mapped) + assert.False(t, errors.Is(mapped, transport.ErrUnauthorized)) + assert.False(t, errors.Is(mapped, transport.ErrSessionTerminated)) + assert.False(t, errors.Is(mapped, transport.ErrLegacySSEServer)) + assert.False(t, errors.Is(mapped, transport.ErrAuthorizationRequired)) +} + +// TestExtractResourceMetadataURL_FromCapturedHeader verifies the shim's own +// headerRoundTripper captures the WWW-Authenticate header and the error mappers +// populate AuthorizationRequiredError.ResourceMetadataURL from it (RFC 9728 +// §5.1). The go-sdk does not surface this header on its errors; the shim's +// RoundTripper is the only place it is available. +func TestExtractResourceMetadataURL_FromCapturedHeader(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + header string // WWW-Authenticate header value + want string // expected ResourceMetadataURL + }{ + { + name: "bearer with resource_metadata", + header: `Bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource"`, + want: "https://example.com/.well-known/oauth-protected-resource", + }, + { + name: "bearer with realm and resource_metadata", + header: `Bearer realm="api", resource_metadata="https://auth.example.com/prm"`, + want: "https://auth.example.com/prm", + }, + { + name: "no resource_metadata parameter", + header: `Bearer realm="api"`, + want: "", + }, + { + name: "empty header", + header: "", + want: "", + }, + { + name: "token value form (unquoted, limited to token chars)", + header: `Bearer resource_metadata=token123`, + want: "token123", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := withErrCapture(context.Background()) + h := capturedErr(ctx) + h.status = http.StatusUnauthorized + if tc.header != "" { + h.wwwAuthHdrs = []string{tc.header} + } + got := extractResourceMetadataURL(ctx) + assert.Equal(t, tc.want, got) + }) + } +} + +// TestMapCallError_Transport401_PopulatesResourceMetadataURL verifies that a +// transport-level 401 with a WWW-Authenticate header populates the +// AuthorizationRequiredError.ResourceMetadataURL field end-to-end through +// mapCallError. +func TestMapCallError_Transport401_PopulatesResourceMetadataURL(t *testing.T) { + t.Parallel() + ctx := withErrCapture(context.Background()) + h := capturedErr(ctx) + h.status = http.StatusUnauthorized + h.body = unauthorizedBody + h.wwwAuthHdrs = []string{ + `Bearer resource_metadata="https://resource.example.com/.well-known/oauth-protected-resource"`, + } + src := fmt.Errorf(`calling "tools/call": %s`, http.StatusText(http.StatusUnauthorized)) + mapped := mapCallError(ctx, src) + require.Error(t, mapped) + assert.True(t, errors.Is(mapped, transport.ErrUnauthorized)) + + var are *transport.AuthorizationRequiredError + require.True(t, errors.As(mapped, &are)) + assert.Equal(t, "https://resource.example.com/.well-known/oauth-protected-resource", are.ResourceMetadataURL) +} + +// TestMapConnectError_Transport401_PopulatesResourceMetadataURL verifies the +// same end-to-end population through mapConnectError (the initialize path). +func TestMapConnectError_Transport401_PopulatesResourceMetadataURL(t *testing.T) { + t.Parallel() + ctx := withErrCapture(context.Background()) + h := capturedErr(ctx) + h.status = http.StatusUnauthorized + h.body = unauthorizedBody + h.wwwAuthHdrs = []string{ + `Bearer realm="api", resource_metadata="https://auth.example.com/prm"`, + } + src := fmt.Errorf("calling initialize: %s", http.StatusText(http.StatusUnauthorized)) + mapped := mapConnectError(ctx, src) + require.Error(t, mapped) + assert.True(t, errors.Is(mapped, transport.ErrUnauthorized)) + + var are *transport.AuthorizationRequiredError + require.True(t, errors.As(mapped, &are)) + assert.Equal(t, "https://auth.example.com/prm", are.ResourceMetadataURL) +} + +// TestCaptureErrorBody_CapturesWWWAuthenticate verifies the headerRoundTripper +// captures the WWW-Authenticate header into the errBody holder for a 401 +// response, so the error mappers can parse resource_metadata from it. +func TestCaptureErrorBody_CapturesWWWAuthenticate(t *testing.T) { + t.Parallel() + + const wwwAuth = `Bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource"` + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("WWW-Authenticate", wwwAuth) + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(unauthorizedBody)) + })) + 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() }) + + h := capturedErr(ctx) + require.NotNil(t, h) + assert.Equal(t, http.StatusUnauthorized, h.status) + assert.Contains(t, h.wwwAuthHdrs, wwwAuth) + assert.Equal(t, "https://example.com/.well-known/oauth-protected-resource", + extractResourceMetadataURL(ctx)) +} diff --git a/mcpcompat/client/rfc9728.go b/mcpcompat/client/rfc9728.go new file mode 100644 index 0000000..03fea6c --- /dev/null +++ b/mcpcompat/client/rfc9728.go @@ -0,0 +1,125 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import "strings" + +// This file implements RFC 9728 §5.1 WWW-Authenticate header parsing to extract +// the resource_metadata parameter, which carries the URL to the protected +// resource metadata document. The parsing logic mirrors mcp-go's +// extractResourceMetadataURL/extractResourceMetadataURLs so the shim's +// AuthorizationRequiredError.ResourceMetadataURL is populated identically. +// +// The go-sdk does not surface the WWW-Authenticate header on its errors, so the +// shim's own headerRoundTripper captures it (see captureErrorBody) and the error +// mappers call extractResourceMetadataURL to parse it. + +// extractResourceMetadataURLFromHeaders returns the first resource_metadata +// parameter value found across the given WWW-Authenticate header values, per +// RFC 9728 §5.1. Returns empty string when none is present. +func extractResourceMetadataURLFromHeaders(wwwAuthHeaders []string) string { + for _, header := range wwwAuthHeaders { + for _, u := range extractResourceMetadataURLs(header) { + if u != "" { + return u + } + } + } + return "" +} + +// extractResourceMetadataURLs returns every resource_metadata parameter value +// from a single WWW-Authenticate header value per RFC 9728 §5.1, in the order +// they appear. Returns an empty slice when the header is empty or no such +// parameters are present. Parameter names are matched case-insensitively per +// RFC 9110 §11.2; both quoted-string and token value forms are accepted. +// +//nolint:gocyclo // direct port of mcp-go's well-tested parser; restructuring harms readability +func extractResourceMetadataURLs(header string) []string { + const target = "resource_metadata" + var out []string + i := 0 + for i < len(header) { + // Advance to the next token start. + for i < len(header) && !isAuthTokenChar(header[i]) { + i++ + } + nameStart := i + for i < len(header) && isAuthTokenChar(header[i]) { + i++ + } + name := header[nameStart:i] + // Skip optional whitespace between the name and '='. + for i < len(header) && (header[i] == ' ' || header[i] == '\t') { + i++ + } + if i >= len(header) || header[i] != '=' { + // Name was a scheme token (e.g. "Bearer"), not a parameter. + continue + } + // Skip '=' and optional whitespace. + i++ + for i < len(header) && (header[i] == ' ' || header[i] == '\t') { + i++ + } + value, next, ok := parseAuthParamValue(header, i) + i = next + if !ok { + continue + } + if value != "" && strings.EqualFold(name, target) { + out = append(out, value) + } + } + return out +} + +// parseAuthParamValue reads a single WWW-Authenticate parameter value starting +// at offset i: a quoted-string (with backslash escapes) when the first byte is +// '"', otherwise a bare token. It returns the decoded value, the index of the +// first byte after it, and whether the value was well-formed. Truncated quoted +// strings (no closing '"') and lone trailing backslashes yield ok=false so +// malformed input is rejected rather than producing a partial value. +func parseAuthParamValue(s string, i int) (string, int, bool) { + if i >= len(s) { + return "", i, false + } + if s[i] == '"' { + i++ + var b strings.Builder + for i < len(s) { + c := s[i] + if c == '\\' { + if i+1 >= len(s) { + return "", i + 1, false + } + b.WriteByte(s[i+1]) + i += 2 + continue + } + if c == '"' { + return b.String(), i + 1, true + } + b.WriteByte(c) + i++ + } + return "", i, false + } + start := i + for i < len(s) && isAuthTokenChar(s[i]) { + i++ + } + return s[start:i], i, i > start +} + +// isAuthTokenChar reports whether c is a valid RFC 9110 §5.6.2 token character +// — the character class used for scheme and parameter names in +// WWW-Authenticate. +func isAuthTokenChar(c byte) bool { + switch { + case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9': + return true + } + return strings.IndexByte("!#$%&'*+-.^_`|~", c) >= 0 +} diff --git a/mcpcompat/server/rehydration_test.go b/mcpcompat/server/rehydration_test.go index 25e743e..81c5458 100644 --- a/mcpcompat/server/rehydration_test.go +++ b/mcpcompat/server/rehydration_test.go @@ -323,6 +323,53 @@ func TestRehydratedSessionEvictedAfterTermination(t *testing.T) { _ = resp.Body.Close() } +// TestRehydratedSession_TransientValidateError_Returns503 verifies the fix for +// the wave 4 review finding: a transient Validate error (e.g. a Redis blip) on +// a CACHED rehydrated session must return 503 (retry-able), NOT 404. A 404 would +// signal hard termination and cause the client to re-initialize. The cached +// reconstruction is preserved so the next request (once Validate recovers) is +// served normally. An unknown session (never cached) still gets 404. +func TestRehydratedSession_TransientValidateError_Returns503(t *testing.T) { + t.Parallel() + mgr := &flakySessionManager{sharedSessionManager: *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", + "rehydrated session must be served before the flake") + + // Make Validate return a transient (non-terminated) error once. + mgr.flake(fmt.Errorf("redis blip: connection refused")) + + // The request must be rejected with 503, NOT 404, and the cached session + // must be preserved. + 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.StatusServiceUnavailable, resp.StatusCode, + "a transient Validate error on a cached rehydrated session must return 503, not 404") + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + // The session must still be cached: a subsequent request (Validate now + // succeeding) must be served normally, NOT rebuilt from scratch. + require.Contains(t, listToolNames(t, tsB.URL, sid), "greet", + "cached rehydrated session must survive a transient Validate error") +} + // 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 diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index 5ec3baa..e1997db 100644 --- a/mcpcompat/server/server.go +++ b/mcpcompat/server/server.go @@ -795,12 +795,16 @@ func normalizeObjectSchema(schema any) any { // unset (it serializes as ""). An empty type string is the sentinel for // "no schema declared"; normalize it to "object" so AddTool accepts the // tool (matching mcp-go's leniency for tools with no declared input). - if t, ok := s[schemaTypeKey].(string); ok && t == "" { - // If the schema declares properties/required, this is an object - // schema whose type was simply omitted (spec-loose but common); - // normalize type to "object" preserving the other fields so the tool - // is callable under go-sdk. Otherwise (no properties/required) it's - // a truly empty-declared schema, normalized to the bare object. + // A missing type key with properties/required is a spec-loose but common + // object shape ({"properties":...}) that mcp-go served verbatim and + // callable; add type:"object" so go-sdk's AddTool accepts it. A missing + // type key WITHOUT properties/required (e.g. {$ref}, {oneOf}, {title}) + // passes through verbatim — it is not a "no schema declared" sentinel. + // A non-string type value (e.g. ["object","null"]) also passes through. + t, hasType := s[schemaTypeKey] + typeStr, typeIsStr := t.(string) + if hasType && typeIsStr && typeStr == "" { + // Empty-string type: mcp-go's "no schema declared" sentinel. if _, hasProps := s["properties"]; hasProps { return withTypeObject(s) } @@ -809,11 +813,8 @@ func normalizeObjectSchema(schema any) any { } return map[string]any{schemaTypeKey: schemaTypeObject} } - // No "type" key at all: if the schema declares properties or required, - // treat it as an object schema with type omitted (spec-loose but common) - // and add type:"object" so go-sdk's AddTool accepts it. This matches - // mcp-go, which served such schemas verbatim and callable. - if _, hasType := s[schemaTypeKey]; !hasType { + if !hasType { + // No type key: only normalize if it looks like an object schema. if _, hasProps := s["properties"]; hasProps { return withTypeObject(s) } diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index 45aa5f8..54ee940 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/mcpcompat/server/server_internal_test.go @@ -49,6 +49,44 @@ func TestClientSession_Store(t *testing.T) { assert.Contains(t, cs.GetSessionResources(), "file:///r") } +// TestForgetSession_ClosesNotifChannel verifies that forgetSession closes the +// notification channel, allowing the drain goroutine started in newClientSession +// to exit rather than leaking for the process lifetime. Also verifies it is safe +// to call forgetSession multiple times for the same ID (sync.Once guards the +// close) and that a re-created session for the same ID gets a fresh channel. +func TestForgetSession_ClosesNotifChannel(t *testing.T) { + t.Parallel() + s := NewMCPServer("s", "1") + + cs := s.sessionFor("sess-drain") + ch := cs.notifCh + + // Send a notification to prove the channel is open before forget. + cs.NotificationChannel() <- mcp.JSONRPCNotification{Notification: mcp.Notification{Method: "test"}} + + // forgetSession must close the channel. + s.forgetSession("sess-drain") + + // A send to a closed channel must panic (proving it's closed). + assert.Panics(t, func() { + ch <- mcp.JSONRPCNotification{} + }) + + // Calling forgetSession again for the same ID must not panic (sync.Once + // guards the double-close). The entry was already deleted, so this is a + // no-op. + assert.NotPanics(t, func() { s.forgetSession("sess-drain") }) + + // A new sessionFor for the same ID creates a fresh, open channel — the old + // closed channel is not reused. + cs2 := s.sessionFor("sess-drain") + require.NotNil(t, cs2.notifCh) + assert.NotEqual(t, ch, cs2.notifCh, "a re-created session must get a fresh channel") + // The fresh channel must be usable. + cs2.NotificationChannel() <- mcp.JSONRPCNotification{Notification: mcp.Notification{Method: "test2"}} + s.forgetSession("sess-drain") +} + // TestSetSessionTools_NonObjectSchemaDoesNotPanic verifies that a per-session // overlay tool whose RawInputSchema is a non-object schema ($ref) does NOT crash // the session when synced onto the go-sdk server. go-sdk v1.6.1 AddTool panics diff --git a/mcpcompat/server/session.go b/mcpcompat/server/session.go index 5dbb283..493554c 100644 --- a/mcpcompat/server/session.go +++ b/mcpcompat/server/session.go @@ -73,7 +73,12 @@ type clientSession struct { initialized atomic.Bool registered atomic.Bool notifCh chan mcp.JSONRPCNotification - goSession atomic.Pointer[gosdk.ServerSession] + // notifClose guards against double-close of notifCh: sessionFor can re-create + // an entry for the same ID after forgetSession closed the channel, so a plain + // close would panic on the second forgetSession. The once ensures only the + // first forgetSession closes it. + notifClose sync.Once + 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 @@ -312,10 +317,16 @@ func (s *MCPServer) isLocalSession(id string) bool { // 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. +// mistaken for a live local session. It also closes the notification channel so +// the drain goroutine started in newClientSession exits, preventing a goroutine +// leak per closed session. func (s *MCPServer) forgetSession(id string) { + v, ok := s.sessions.LoadAndDelete(id) s.localSessions.Delete(id) - s.sessions.Delete(id) + if ok { + cs := v.(*clientSession) + cs.notifClose.Do(func() { close(cs.notifCh) }) + } } // bindRehydratedSession binds the clientSession for a session that was created diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index a62a395..9d470bd 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -380,7 +380,23 @@ func (s *StreamableHTTPServer) serveRehydrated(w http.ResponseWriter, r *http.Re // it and drop any locally-cached reconstruction rather than serving it. isTerminated, err := s.sessionIDMgr.Validate(sid) if err != nil { - s.deleteRehydrated(sid) + if isTerminated { + s.deleteRehydrated(sid) + http.Error(w, "Session terminated", http.StatusNotFound) + return + } + // Distinguish a transient store error from a genuinely unknown session. + // If we have a cached reconstruction, the session WAS valid — the error is + // likely transient (e.g. a Redis blip), so preserve the cache and return + // 503 (retry-able), matching the local-session path (issue #156, finding + // 3). A 404 here would signal hard termination and cause the client to + // re-initialize. If we have NO cached reconstruction, the session was + // never seen by this replica — the error is "not found", and 404 is the + // correct response (not a transient issue with a known session). + if s.getRehydrated(sid) != nil { + http.Error(w, "session validation unavailable", http.StatusServiceUnavailable) + return + } http.Error(w, "Invalid session ID", http.StatusNotFound) return } @@ -432,6 +448,16 @@ func (s *StreamableHTTPServer) deleteRehydrated(sid string) { // 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. +// +// Scalability note: rehydrate holds rehydratedMu across buildServer + Connect. +// buildServer copies the tool/resource/prompt maps under s.mu (quick), but +// Connect can block on transport setup, so concurrent rehydrate calls for +// DIFFERENT session IDs serialize on this lock. The double-check at the top +// makes same-sid serialization correct/desired; the cross-sid serialization is +// the incidental cost. Low impact in practice (cross-replica rehydration is +// rare and the first request per sid is the only one that builds), so not worth +// restructuring now — but if rehydration volume ever rises, building the session +// outside the lock and only inserting under it would remove the serialization. func (s *StreamableHTTPServer) rehydrate(r *http.Request, sid string) (*rehydratedSession, error) { s.rehydratedMu.Lock() defer s.rehydratedMu.Unlock() From c469b88696dc343d8e0cc0de82f676d71683c2ed Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Mon, 13 Jul 2026 13:55:41 +0300 Subject: [PATCH 8/8] fix(mcpcompat): address jhrozek's second review (wave 7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses all 5 inline comments from jhrozek's second review (2026-07-13T09:27:55Z), which were missed in wave 6: 1. [BLOCKING] mapConnectError missing JSON-RPC guard (errors.go): mapConnectError had no *jsonrpc.Error guard like mapCallError has. An app-level JSON-RPC error at initialize (200 OK + error body, h.status==0) with "unauthorized" text was misclassified as ErrUnauthorized → false OAuth refresh. Added the symmetric guard: when h.status==0 and err is a *jsonrpc.Error, surface unchanged. Test: TestMapConnectError_DoesNotMisclassifyJSONRPCErrors (3 cases). 2. [Question] SetLevel compat alias (client.go + alias.go): Shim had SetLoggingLevel but not SetLevel/SetLevelRequest/ SetLevelParams, breaking the drop-in contract for callers using c.SetLevel(ctx, mcp.SetLevelRequest{...}). Added SetLevel method (forwards to SetLoggingLevel) and re-exported SetLevelRequest/ SetLevelParams from the mcp package. Test: TestSetLevel_CompatAlias (e2e against a live server). 3. [Compat gap] OnNotification drops resources/updated + elicitation/complete (client.go): ResourceUpdatedHandler and ElicitationCompleteHandler were left nil. Wired both via dispatch, mirroring the progress/logging handlers. Updated OnNotification doc to list all forwarded types. Test: TestResourceUpdatedAndElicitationCompleteHandlers. 4. [Follow-up] No keep-alive (transports.go): Strengthened the WithHeartbeatInterval doc to explain the passive keep-alive is load-bearing (elicitation + all server→client notifications depend on the idle SSE stream staying open) and should be prioritized before internet-facing vMCP use. 5. [Doc note] 503 diverges from mcp-go's 404 on Validate (transports.go): Added a compat-note comment explaining 503 is not mapped by the client classifier (errors.Is(err, ErrSessionTerminated) is false), so callers keying re-init off mcp-go's 404→ErrSessionTerminated path won't re-init on a transient blip. Verification: task lint (0 issues), task test (-race, all green), task license-check (clean). --- mcpcompat/client/client.go | 47 +++++++++++++++---- mcpcompat/client/client_test.go | 37 ++++++++++++++- mcpcompat/client/errors.go | 16 +++++++ mcpcompat/client/errors_internal_test.go | 46 ++++++++++++++++++ .../client/notifications_internal_test.go | 42 +++++++++++++++++ mcpcompat/mcp/alias.go | 11 +++++ mcpcompat/server/transports.go | 23 ++++++++- 7 files changed, 211 insertions(+), 11 deletions(-) diff --git a/mcpcompat/client/client.go b/mcpcompat/client/client.go index 40694fa..e8f9c82 100644 --- a/mcpcompat/client/client.go +++ b/mcpcompat/client/client.go @@ -262,12 +262,11 @@ func (c *Client) Ping(ctx context.Context) error { // SetLoggingLevel sets the server's logging level. This is a renamed and // simplified counterpart to mcp-go's client.Client.SetLevel: rather than // taking an mcp.SetLevelRequest, it accepts the mcp.LoggingLevel directly. -// The SetLevel, SetLevelRequest and SetLevelParams types from mcp-go are -// intentionally NOT re-exported by this shim. The server only delivers -// notifications/message notifications at or above the requested level, so this -// must be called before the OnNotification handler will see logging -// notifications. level is one of the MCP logging levels ("debug", "info", -// "notice", "warning", "error", "critical", "alert", "emergency"). +// The server only delivers notifications/message notifications at or above the +// requested level, so this must be called before the OnNotification handler +// will see logging notifications. level is one of the MCP logging levels +// ("debug", "info", "notice", "warning", "error", "critical", "alert", +// "emergency"). func (c *Client) SetLoggingLevel(ctx context.Context, level mcp.LoggingLevel) error { s, err := c.sessionFor() if err != nil { @@ -276,6 +275,16 @@ func (c *Client) SetLoggingLevel(ctx context.Context, level mcp.LoggingLevel) er return s.SetLoggingLevel(ctx, &gosdk.SetLoggingLevelParams{Level: gosdk.LoggingLevel(level)}) } +// SetLevel is a drop-in compatibility alias for mcp-go's +// client.Client.SetLevel, forwarding to SetLoggingLevel. It is provided so +// downstream code using the upstream idiom +// `c.SetLevel(ctx, mcp.SetLevelRequest{...})` compiles against the shim +// unchanged. SetLevelRequest and SetLevelParams are re-exported from the mcp +// package for the same reason. +func (c *Client) SetLevel(ctx context.Context, request mcp.SetLevelRequest) error { + return c.SetLoggingLevel(ctx, request.Params.Level) +} + // ListTools lists the server's tools. func (c *Client) ListTools(ctx context.Context, request mcp.ListToolsRequest) (*mcp.ListToolsResult, error) { ctx = withErrCapture(ctx) @@ -498,8 +507,15 @@ func (c *Client) ListResourceTemplates( // 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 (method + params) for the list-changed, progress and logging -// notifications, forwarding each to every registered handler. +// values (method + params) for each notification type and forwards each to +// every registered handler. The following notification types are forwarded: +// - notifications/tools/list_changed +// - notifications/prompts/list_changed +// - notifications/resources/list_changed +// - notifications/resources/updated (from resources/subscribe) +// - notifications/elicitation/complete (out-of-band elicitation completion) +// - notifications/progress (server->client progress) +// - notifications/message (server->client logging) // // Server-initiated notifications (including the list_changed notifications, // which the server emits outside any in-flight request) are only delivered if @@ -573,6 +589,21 @@ func (c *Client) installNotificationHandlers(opts *gosdk.ClientOptions) { opts.ResourceListChangedHandler = func(_ context.Context, _ *gosdk.ResourceListChangedRequest) { c.dispatch("notifications/resources/list_changed", nil) } + // Forward server->client resource-updated notifications (the result of a + // resources/subscribe). go-sdk hands the params via + // *ResourceUpdatedNotificationRequest; dispatch synthesizes the + // notifications/resources/updated notification with its params (uri). + opts.ResourceUpdatedHandler = func(_ context.Context, req *gosdk.ResourceUpdatedNotificationRequest) { + c.dispatch("notifications/resources/updated", req.Params) + } + // Forward server->client elicitation-complete notifications (out-of-band + // elicitation completion). go-sdk hands the params via + // *ElicitationCompleteNotificationRequest; dispatch synthesizes the + // notifications/elicitation/complete notification with its params + // (elicitationId). + opts.ElicitationCompleteHandler = func(_ context.Context, req *gosdk.ElicitationCompleteNotificationRequest) { + c.dispatch("notifications/elicitation/complete", req.Params) + } // Forward server->client progress notifications. go-sdk hands the params via // *ProgressNotificationClientRequest; convert them to mcp-go's notification // params (progressToken, progress, total, message) and dispatch. diff --git a/mcpcompat/client/client_test.go b/mcpcompat/client/client_test.go index 120d3d0..7c7af02 100644 --- a/mcpcompat/client/client_test.go +++ b/mcpcompat/client/client_test.go @@ -30,6 +30,9 @@ const echoToolName = "echo" // testClientVersion is the shared client/server version string used in tests. const testClientVersion = "1.0.0" +// testClientName is the shared client name used in InitializeParams across tests. +const testClientName = "test-client" + // 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 { @@ -64,7 +67,7 @@ func TestStreamableClient_EndToEnd(t *testing.T) { initRes, err := c.Initialize(ctx, mcp.InitializeRequest{ Params: mcp.InitializeParams{ ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - ClientInfo: mcp.Implementation{Name: "test-client", Version: testClientVersion}, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, }, }) require.NoError(t, err) @@ -211,7 +214,7 @@ func TestOnNotification_ProgressAndLogging(t *testing.T) { _, err = c.Initialize(ctx, mcp.InitializeRequest{ Params: mcp.InitializeParams{ ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - ClientInfo: mcp.Implementation{Name: "test-client", Version: testClientVersion}, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, }, }) require.NoError(t, err) @@ -254,3 +257,33 @@ func TestOnNotification_ProgressAndLogging(t *testing.T) { defer mu.Unlock() require.GreaterOrEqual(t, len(got), 2, "at least progress and logging notifications expected") } + +// TestSetLevel_CompatAlias verifies that the SetLevel compatibility alias +// (mcp-go's client.Client.SetLevel idiom) delegates to SetLoggingLevel and +// successfully sets the server's logging level. This guards against the +// drop-in contract breaking for downstream code using c.SetLevel(ctx, +// mcp.SetLevelRequest{...}). +func TestSetLevel_CompatAlias(t *testing.T) { + t.Parallel() + ctx := context.Background() + ts := newProgressLoggingServer(t) + + c, err := client.NewStreamableHttpClient(ts.URL) + require.NoError(t, err) + require.NoError(t, c.Start(ctx)) + t.Cleanup(func() { _ = c.Close() }) + + _, err = c.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: testClientName, Version: testClientVersion}, + }, + }) + require.NoError(t, err) + + // SetLevel must work exactly like SetLoggingLevel — it forwards to it. + err = c.SetLevel(ctx, mcp.SetLevelRequest{ + Params: mcp.SetLevelParams{Level: mcp.LoggingLevelDebug}, + }) + assert.NoError(t, err, "SetLevel compat alias must succeed") +} diff --git a/mcpcompat/client/errors.go b/mcpcompat/client/errors.go index 8bc26c9..bddbd75 100644 --- a/mcpcompat/client/errors.go +++ b/mcpcompat/client/errors.go @@ -42,6 +42,14 @@ func enrichWithResponseBody(ctx context.Context, err error) error { // string matching when no body was captured (e.g. a failure before any // response). 401 stays ErrUnauthorized so ToolHive's OAuth refresh flow still // triggers. +// +// JSON-RPC-level errors (a *jsonrpc.Error returned inside a 2xx response) that +// carry no captured HTTP status are application-level initialize rejections, not +// transport failures: they must NOT be reclassified as transport auth/session +// failures based on their text (mirrors mapCallError). Connect runs initialize +// through the same handleSend machinery as the call methods (go-sdk v1.6.1 +// client.go:271), so the same *jsonrpc.Error reaches here with no captured +// status. func mapConnectError(ctx context.Context, err error) error { err = enrichWithResponseBody(ctx, err) h := capturedErr(ctx) @@ -55,6 +63,14 @@ func mapConnectError(ctx context.Context, err error) error { } return transport.NewError(errors.Join(transport.ErrLegacySSEServer, err)) } + var wireErr *jsonrpc.Error + if (h == nil || h.status == 0) && errors.As(err, &wireErr) { + // RPC-level error with no captured HTTP status: an application-level + // initialize rejection (200 OK + JSON-RPC error body), not a transport + // failure. Surface unchanged (mirrors mapCallError) rather than + // string-matching it into a transport auth/legacy-SSE sentinel. + return err + } return mapTransportError(ctx, err, h) } diff --git a/mcpcompat/client/errors_internal_test.go b/mcpcompat/client/errors_internal_test.go index 7f1ec3e..b2e1d79 100644 --- a/mcpcompat/client/errors_internal_test.go +++ b/mcpcompat/client/errors_internal_test.go @@ -123,6 +123,52 @@ func TestMapCallError_DoesNotMisclassifyJSONRPCErrors(t *testing.T) { } } +// TestMapConnectError_DoesNotMisclassifyJSONRPCErrors mirrors the mapCallError +// guard for the connect/initialize path: a JSON-RPC error returned at +// initialize (200 OK + JSON-RPC error body, no captured HTTP status) whose +// message text contains "unauthorized"/"401" must NOT be reclassified as a +// transport auth failure. Connect runs initialize through the same handleSend +// machinery as the call methods (go-sdk v1.6.1 client.go:271), so the same +// *jsonrpc.Error reaches mapConnectError with no captured status. +func TestMapConnectError_DoesNotMisclassifyJSONRPCErrors(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + err error + }{ + { + name: "unauthorized in message", + err: &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "unauthorized to access resource X"}, + }, + { + name: "401 in message", + err: &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "got 401 from downstream"}, + }, + { + name: "method not allowed in message", + err: &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "method not allowed here"}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + // No captured body: this is an RPC-level error at initialize. + ctx := withErrCapture(context.Background()) + mapped := mapConnectError(ctx, tc.err) + require.Error(t, mapped) + assert.False(t, errors.Is(mapped, transport.ErrUnauthorized), + "RPC-level initialize error must not satisfy ErrUnauthorized: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrAuthorizationRequired), + "RPC-level initialize error must not satisfy ErrAuthorizationRequired: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrLegacySSEServer), + "RPC-level initialize error must not satisfy ErrLegacySSEServer: %v", mapped) + assert.False(t, errors.Is(mapped, transport.ErrSessionTerminated), + "RPC-level initialize error must not satisfy ErrSessionTerminated: %v", mapped) + }) + } +} + // TestMapCallError_Transport401 confirms a transport-level 401 (with a captured // HTTP status) through mapCallError still maps to ErrUnauthorized. Scoping, not // removal: real transport 401s must keep triggering ToolHive's OAuth flow. diff --git a/mcpcompat/client/notifications_internal_test.go b/mcpcompat/client/notifications_internal_test.go index 3b389f1..8d1eb2c 100644 --- a/mcpcompat/client/notifications_internal_test.go +++ b/mcpcompat/client/notifications_internal_test.go @@ -303,3 +303,45 @@ func TestListChangedHandlers_MethodStrings(t *testing.T) { "notifications/resources/list_changed", }, got, "list-changed method strings must match the MCP spec names") } + +// TestResourceUpdatedAndElicitationCompleteHandlers verifies that +// installNotificationHandlers wires the ResourceUpdatedHandler and +// ElicitationCompleteHandler so notifications/resources/updated and +// notifications/elicitation/complete are forwarded to the OnNotification +// callback. Without these handlers, a client that subscribed to a resource +// and is waiting for notifications/resources/updated via OnNotification +// receives nothing. +func TestResourceUpdatedAndElicitationCompleteHandlers(t *testing.T) { + t.Parallel() + + c := &Client{} + var ( + mu sync.Mutex + got []string + ) + c.OnNotification(func(n mcp.JSONRPCNotification) { + mu.Lock() + got = append(got, n.Method) + mu.Unlock() + }) + + opts := &gosdk.ClientOptions{} + c.installNotificationHandlers(opts) + + require.NotNil(t, opts.ResourceUpdatedHandler) + require.NotNil(t, opts.ElicitationCompleteHandler) + + opts.ResourceUpdatedHandler(context.Background(), &gosdk.ResourceUpdatedNotificationRequest{ + Params: &gosdk.ResourceUpdatedNotificationParams{URI: "file:///example"}, + }) + opts.ElicitationCompleteHandler(context.Background(), &gosdk.ElicitationCompleteNotificationRequest{ + Params: &gosdk.ElicitationCompleteParams{ElicitationID: "el-1"}, + }) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{ + "notifications/resources/updated", + "notifications/elicitation/complete", + }, got, "resource-updated and elicitation-complete method strings must match the MCP spec names") +} diff --git a/mcpcompat/mcp/alias.go b/mcpcompat/mcp/alias.go index 2003b13..201ccb1 100644 --- a/mcpcompat/mcp/alias.go +++ b/mcpcompat/mcp/alias.go @@ -218,6 +218,17 @@ var NewMetaFromMap = mcpgo.NewMetaFromMap //nolint:revive // name intentionally matches mcp-go for drop-in compatibility. type LoggingLevel = mcpgo.LoggingLevel +// SetLevelRequest is a request from the client to the server to set the logging +// level. Re-exported from mcp-go for drop-in compatibility: downstream code +// using the upstream idiom `c.SetLevel(ctx, mcp.SetLevelRequest{...})` compiles +// against the shim unchanged. The shim's client.Client.SetLevel forwards to +// SetLoggingLevel. +type ( + SetLevelRequest = mcpgo.SetLevelRequest + // SetLevelParams carries the level for a SetLevelRequest. + SetLevelParams = mcpgo.SetLevelParams +) + // MCP logging levels, mirrored from mcp-go. const ( LoggingLevelDebug = mcpgo.LoggingLevelDebug diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index 9d470bd..a18dcac 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -170,9 +170,20 @@ func WithSessionIdManager(manager SessionIdManager) StreamableHTTPOption { // heartbeat was passive (pings written only to an existing GET stream, never // awaited, never terminating). // +// NOTE: the passive keep-alive is more load-bearing than "nice to have" now +// that elicitation and all server→client notifications (list_changed, +// progress, logging, resources/updated, elicitation/complete) structurally +// depend on the client's idle standalone SSE stream staying open. Any +// intermediary (LB, reverse proxy, ToolHive's own proxy) will reap that idle +// stream on timeout — after which elicitation fails (ErrRejected) and +// server-initiated notifications are silently dropped, with no reconnect. It +// also feeds unbounded session growth (abandoned sessions are never reaped). +// // TODO(issue #156): implement a passive keep-alive (SSE comment injection via // a ResponseWriter wrapper around the go-sdk handler) matching mcp-go's design -// and wire this option to it. Until then the value is stored but unused. +// and wire this option to it. This should be prioritized before internet-facing +// vMCP use; a max-session cap / idle sweep would also help. Until then the +// value is stored but unused. func WithHeartbeatInterval(interval time.Duration) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.heartbeat = interval } } @@ -355,6 +366,16 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) // without dropping local state. The local go-sdk session stays // alive; the client retries and, once Validate succeeds again, is // served normally rather than routed through rehydration. + // + // Compat note: mcp-go returned 404 on ANY Validate error; the shim + // returns 503 on a non-terminated (transient) error and reserves + // 404 for genuine termination. 503 is NOT mapped by the client + // classifier (errors.Is(err, ErrSessionTerminated) is false), so a + // caller that keyed re-initialization off mcp-go's 404 → + // ErrSessionTerminated path won't re-init on a transient blip — it + // sees a generic retryable error. This is intentional (not dropping + // local state avoids the split-brain this PR guards against), but + // callers relying on that path should be aware of the divergence. http.Error(w, "session validation unavailable", http.StatusServiceUnavailable) return }