diff --git a/pkg/transport/proxy/streamable/initialize_cache.go b/pkg/transport/proxy/streamable/initialize_cache.go new file mode 100644 index 0000000000..654c22e07d --- /dev/null +++ b/pkg/transport/proxy/streamable/initialize_cache.go @@ -0,0 +1,180 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "context" + "encoding/json" + "errors" + "sync/atomic" + + "golang.org/x/exp/jsonrpc2" + "golang.org/x/sync/singleflight" + + sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// errUnexpectedInitializeReply is returned when the backend answers initialize +// with something that is not a *jsonrpc2.Response, so it cannot be re-addressed +// to the calling client. +var errUnexpectedInitializeReply = errors.New("backend returned an unexpected reply to initialize") + +// initializeCache holds the InitializeResult the backend produced for the +// first client handshake, so later handshakes are answered locally instead of +// reaching the backend again. +// +// A stdio MCP server is one process and therefore exactly one MCP session, +// while this proxy fans an arbitrary number of HTTP client sessions onto it. +// Forwarding every client's initialize to that single session means only the +// first can succeed: go-sdk v1.7+ rejects a second handshake on an +// already-initialized session with `duplicate "initialize" received`, which +// takes down every client after the first -- and, behind vMCP, every health +// check and tool call as well (see #5890, #1982). +// +// Answering from a cache restores the behaviour go-sdk v1.6.1 and earlier had, +// where a repeated initialize returned the original result, but places it in +// the component that actually owns the multiplexing rather than asking every +// MCP server to tolerate duplicate handshakes. +// +// The cached value is process-scoped rather than session-scoped: an +// InitializeResult carries the backend's protocol version, capabilities, +// instructions and server info, none of which vary by client. +type initializeCache struct { + // flight collapses concurrent handshakes into a single upstream call whose + // outcome every waiter shares. Serializing them behind a mutex instead + // would queue each waiter's own round trip: with an unresponsive backend + // the Nth client waits N*requestTimeout (60s each) before it may even + // start. + flight singleflight.Group + + // result is the raw InitializeResult of the first successful handshake, + // nil until one completes. + // + // UNSYNCHRONISED BY DESIGN: only the flight callback in interceptInitialize + // may read or write this field. singleflight runs at most one callback per + // key at a time and a later callback starts only after the previous one has + // completed through the group's own mutex, so callback-only access can + // never overlap. Reading it anywhere else -- a status endpoint, a helper + // accessor -- reintroduces a data race, which is why no accessor exists. + // Once assigned it is never mutated, so the same bytes may be handed to + // concurrent responses. + result json.RawMessage + + // initializedForwarded records whether a notifications/initialized has + // already been passed to the backend. This deliberately does NOT use the + // flight: singleflight de-duplicates rather than serializes, so two + // concurrent claims would share one callback's answer and both forward. + // CompareAndSwap gives exactly one winner. + initializedForwarded atomic.Bool +} + +// handshake is the id-free half of an initialize outcome, shared by a flight +// with all of its waiters. Exactly one of the fields is set. Sharing this +// rather than a whole *jsonrpc2.Response means each caller builds its own +// reply, so a waiter cannot receive the forwarding caller's JSON-RPC id. +type handshake struct { + result json.RawMessage + rpcErr error +} + +// initializeFlightKey is the sole singleflight key: the cached handshake is +// process-scoped, so all callers share one flight. +const initializeFlightKey = "initialize" + +// newInitializeCache creates an empty initializeCache. +func newInitializeCache() *initializeCache { + return &initializeCache{} +} + +// interceptInitialize answers a client's initialize request: the first caller's +// request is forwarded to the backend and its result cached, and every later +// caller is served that cached result without an upstream round trip. +// +// Only a successful handshake is cached. A transport failure or a JSON-RPC +// error from the backend leaves the cache empty so the next client retries the +// handshake for real, rather than pinning every future client to one bad +// response. +// +// Concurrent first handshakes are collapsed by singleflight, so exactly one +// reaches the backend and the rest share its outcome. The upstream call runs on +// a context detached from the caller's request (context.WithoutCancel): the +// result belongs to every waiter, so the first caller disconnecting must not +// cancel the handshake for the others. forwardUpstream still applies +// p.requestTimeout to that detached context, so it cannot run unbounded. +// +// Every response is rebuilt with the calling request's own JSON-RPC id. A +// shared flight result carries the id of whichever caller performed the +// forward, and returning it verbatim would hand later waiters an id they never +// sent. +// +// One client-visible consequence: the protocol version negotiated by the first +// client is replayed to every later client. A client that asked for a +// different version receives the first client's instead, and per the MCP +// lifecycle spec should disconnect if it cannot speak it. This is inherent to +// collapsing N client sessions onto one backend session, and matches what +// go-sdk v1.6.1 and earlier did when replaying a cached result. +func (p *HTTPProxy) interceptInitialize( + ctx context.Context, sessID string, req *jsonrpc2.Request, +) jsonrpc2.Message { + c := p.initialize + v, err, _ := c.flight.Do(initializeFlightKey, func() (any, error) { + // A handshake already succeeded: replay it, no backend round trip. + if c.result != nil { + return handshake{result: c.result}, nil + } + + // Detached from the caller's context: this result belongs to every + // waiter, so whichever caller performs the forward must not cancel the + // handshake for the others by disconnecting. forwardUpstream still + // bounds it with p.requestTimeout. + msg, ferr := p.forwardUpstream(context.WithoutCancel(ctx), sessID, req) + if ferr != nil { + return nil, ferr // nothing cached; the next client retries for real + } + + resp, ok := msg.(*jsonrpc2.Response) + if !ok { + return nil, errUnexpectedInitializeReply + } + if resp.Error != nil { + // Shared with this flight's waiters but deliberately not cached, so + // a transient failure cannot pin every future client to one + // response. + return handshake{rpcErr: resp.Error}, nil + } + if len(resp.Result) == 0 { + return nil, errUnexpectedInitializeReply // no result and no error + } + + c.result = resp.Result + return handshake{result: resp.Result}, nil + }) + if err != nil { + return upstreamErrorResponse(req.ID, err) + } + + h := v.(handshake) // the callback returns a handshake or a non-nil error + return &jsonrpc2.Response{ID: req.ID, Result: h.result, Error: h.rpcErr} +} + +// claimInitializedForward reports whether a client's notifications/initialized +// should be passed to the backend, and records that it has been. +// +// The backend completes a single handshake and so expects a single initialized +// notification; the first client's is forwarded and every later one is +// swallowed. Swallowed notifications are still acknowledged with 202, which is +// all the sending client expects -- notifications carry no response. +func (p *HTTPProxy) claimInitializedForward() bool { + return p.initialize.initializedForwarded.CompareAndSwap(false, true) +} + +// isInitializedNotification reports whether msg is a client's +// notifications/initialized (a JSON-RPC notification, so no id). +func isInitializedNotification(msg jsonrpc2.Message) bool { + req, ok := msg.(*jsonrpc2.Request) + if !ok { + return false + } + return req.ID.Raw() == nil && req.Method == string(sdkmcp.MethodNotificationInitialized) +} diff --git a/pkg/transport/proxy/streamable/initialize_cache_test.go b/pkg/transport/proxy/streamable/initialize_cache_test.go new file mode 100644 index 0000000000..e6127ecf0f --- /dev/null +++ b/pkg/transport/proxy/streamable/initialize_cache_test.go @@ -0,0 +1,385 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/exp/jsonrpc2" +) + +// fakeStdioBackend stands in for the single stdio MCP session behind the +// proxy: it drains the proxy's outbound message channel, records every method +// that reaches it, and answers id-bearing requests via respond. +type fakeStdioBackend struct { + mu sync.Mutex + methods []string +} + +func (b *fakeStdioBackend) record(method string) { + b.mu.Lock() + defer b.mu.Unlock() + b.methods = append(b.methods, method) +} + +// count returns how many times method reached the backend. +func (b *fakeStdioBackend) count(method string) int { + b.mu.Lock() + defer b.mu.Unlock() + n := 0 + for _, m := range b.methods { + if m == method { + n++ + } + } + return n +} + +// startFakeStdioBackend wires a fakeStdioBackend to proxy and returns it. +// respond is consulted only for id-bearing requests; notifications are +// recorded and dropped, as a real backend would handle them. +func startFakeStdioBackend( + ctx context.Context, proxy *HTTPProxy, respond func(*jsonrpc2.Request) jsonrpc2.Message, +) *fakeStdioBackend { + backend := &fakeStdioBackend{} + go func() { + for { + select { + case msg := <-proxy.GetMessageChannel(): + req, ok := msg.(*jsonrpc2.Request) + if !ok { + continue + } + backend.record(req.Method) + if !req.ID.IsValid() { + continue + } + if reply := respond(req); reply != nil { + _ = proxy.ForwardResponseToClients(ctx, reply) + } + case <-ctx.Done(): + return + } + } + }() + return backend +} + +// testProtocolVersion is the protocol version the fake backend negotiates. +const testProtocolVersion = "2025-11-25" + +// initializeResultFor builds a realistic InitializeResult response for req. +func initializeResultFor(req *jsonrpc2.Request) jsonrpc2.Message { + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{ + "protocolVersion": testProtocolVersion, + "capabilities": map[string]any{"tools": map[string]any{}}, + "serverInfo": map[string]any{"name": "fake-stdio-server", "version": "1.0.0"}, + }, nil) + return resp +} + +// strictInitializeResponder returns a backend responder that reproduces +// go-sdk v1.7+ server-session semantics: the first initialize is answered with +// an InitializeResult and every later one is rejected with the same +// `duplicate "initialize" received` error ServerSession.initialize returns +// (go-sdk mcp/server.go:1978). go-sdk rejects a repeated +// notifications/initialized the same way (mcp/server.go:1424). +// +// Tests use this rather than a permissive stub so they assert the outcome that +// actually matters -- a later client's handshake succeeds -- instead of only +// counting how many messages reached the backend. A duplicate that escapes the +// cache fails the handshake here exactly as it does against a real server. +func strictInitializeResponder() func(*jsonrpc2.Request) jsonrpc2.Message { + var mu sync.Mutex + var handshakeDone bool + return func(req *jsonrpc2.Request) jsonrpc2.Message { + if req.Method != "initialize" { + resp, _ := jsonrpc2.NewResponse(req.ID, map[string]any{}, nil) + return resp + } + mu.Lock() + defer mu.Unlock() + if handshakeDone { + resp, _ := jsonrpc2.NewResponse(req.ID, nil, + jsonrpc2.NewError(-32603, `duplicate "initialize" received`)) + return resp + } + handshakeDone = true + return initializeResultFor(req) + } +} + +// startInitTestProxy starts a proxy on a free port and returns it with its +// /mcp endpoint URL. +func startInitTestProxy(t *testing.T) (*HTTPProxy, string) { + t.Helper() + port := getFreePort(t) + proxy := NewHTTPProxy("localhost", port, nil, nil) + require.NoError(t, proxy.Start(t.Context())) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = proxy.Stop(stopCtx) + }) + // Give the listener a moment to accept connections. + time.Sleep(100 * time.Millisecond) + return proxy, fmt.Sprintf("http://localhost:%d%s", port, StreamableHTTPEndpoint) +} + +// postInitialize sends an initialize request and returns the decoded JSON-RPC +// response body. +func postInitialize(t *testing.T, url string, id string) map[string]any { + t.Helper() + body := fmt.Sprintf( + `{"jsonrpc":"2.0","id":%q,"method":"initialize","params":{"protocolVersion":%q}}`, + id, testProtocolVersion) + resp, err := http.Post(url, "application/json", bytes.NewReader([]byte(body))) //nolint:gosec // test-local URL + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + var decoded map[string]any + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + return decoded +} + +// TestEveryClientHandshakeSucceedsAgainstStrictBackend verifies the behaviour +// #5890 reports: with a backend that rejects a repeated initialize, every +// client after the first used to fail. Only the first handshake may reach the +// single stdio session; the rest are answered from the cached InitializeResult +// with their own JSON-RPC id echoed back. +func TestEveryClientHandshakeSucceedsAgainstStrictBackend(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + backend := startFakeStdioBackend(t.Context(), proxy, strictInitializeResponder()) + + var results []any + for i := range 3 { + decoded := postInitialize(t, url, fmt.Sprintf("client-%d", i)) + assert.Equal(t, fmt.Sprintf("client-%d", i), decoded["id"], + "each client must get its own JSON-RPC id back, not the first client's") + assert.Nil(t, decoded["error"], + "no client may see the backend's duplicate-initialize rejection") + require.NotNil(t, decoded["result"], "handshake should succeed") + results = append(results, decoded["result"]) + } + + assert.Equal(t, 1, backend.count("initialize"), + "only the first handshake may reach the single stdio session") + assert.Equal(t, results[0], results[1], "cached result must match the original") + assert.Equal(t, results[0], results[2], "cached result must match the original") +} + +// TestConcurrentInitializeReachesBackendOnce verifies that simultaneous +// handshakes are single-flighted: without the lock held across the first +// upstream round trip, every concurrent client would see an empty cache and +// race a second initialize to the backend. +func TestConcurrentInitializeReachesBackendOnce(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + strict := strictInitializeResponder() + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + // Slow the first handshake so the others pile up behind it. + time.Sleep(150 * time.Millisecond) + return strict(req) + }) + + const clients = 5 + var wg sync.WaitGroup + for i := range clients { + wg.Go(func() { + id := fmt.Sprintf("client-%d", i) + decoded := postInitialize(t, url, id) + assert.Nil(t, decoded["error"], + "a handshake that raced past the cache would be rejected by the backend") + assert.NotNil(t, decoded["result"], "every concurrent handshake should succeed") + // Waiters share one singleflight result, which carries the id of + // whichever caller performed the forward. Each response must be + // re-addressed to its own caller. + assert.Equal(t, id, decoded["id"], + "a waiter must never receive another caller's JSON-RPC id") + }) + } + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent handshakes to complete") + } + + assert.Equal(t, 1, backend.count("initialize"), + "concurrent handshakes must collapse into a single upstream initialize") +} + +// TestFailedInitializeIsNotCached verifies that a backend error is not stored, +// so a transient failure cannot pin every future client to one bad response. +func TestFailedInitializeIsNotCached(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + + var attempts int + var mu sync.Mutex + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + mu.Lock() + attempts++ + first := attempts == 1 + mu.Unlock() + if first { + resp, _ := jsonrpc2.NewResponse(req.ID, nil, jsonrpc2.NewError(-32603, "backend not ready")) + return resp + } + return initializeResultFor(req) + }) + + failed := postInitialize(t, url, "client-0") + assert.NotNil(t, failed["error"], "first handshake should surface the backend error") + + succeeded := postInitialize(t, url, "client-1") + assert.NotNil(t, succeeded["result"], "second handshake should be retried against the backend") + + assert.Equal(t, 2, backend.count("initialize"), + "a failed handshake must not be cached; the next client must retry for real") +} + +// TestDuplicateInitializedNotificationIsNotForwarded verifies that only the +// first notifications/initialized reaches the backend. The backend completed a +// single handshake, so a second initialized notification is out of lifecycle +// for it; later ones are acknowledged locally instead. +func TestDuplicateInitializedNotificationIsNotForwarded(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + backend := startFakeStdioBackend(t.Context(), proxy, strictInitializeResponder()) + + for range 3 { + resp, err := http.Post(url, "application/json", //nolint:gosec // test-local URL + bytes.NewReader([]byte(`{"jsonrpc":"2.0","method":"notifications/initialized"}`))) + require.NoError(t, err) + require.Equal(t, http.StatusAccepted, resp.StatusCode, + "every client's initialized notification must still be acknowledged") + _ = resp.Body.Close() + } + + // The backend goroutine records asynchronously; give it a moment to drain. + assert.Eventually(t, func() bool { + return backend.count("notifications/initialized") >= 1 + }, 2*time.Second, 20*time.Millisecond, "the first notification should reach the backend") + + assert.Equal(t, 1, backend.count("notifications/initialized"), + "only one initialized notification may reach the single stdio session") +} + +// TestInitializeSurvivesFirstCallerDisconnect verifies that a client giving up +// mid-handshake does not cancel the shared upstream call for everyone else. +// +// The handshake is collapsed by singleflight, and singleflight runs the shared +// function with whichever context the first caller supplied. Passing that +// context straight through would let one client's disconnect fail every waiter, +// so interceptInitialize detaches it with context.WithoutCancel. +func TestInitializeSurvivesFirstCallerDisconnect(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + strict := strictInitializeResponder() + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + // Long enough that the first caller aborts while the handshake is in + // flight, short enough to keep the test quick. + time.Sleep(400 * time.Millisecond) + return strict(req) + }) + + // First caller: gives up well before the backend replies. + var wg sync.WaitGroup + wg.Go(func() { + ctx, cancel := context.WithTimeout(t.Context(), 80*time.Millisecond) + defer cancel() + body := fmt.Sprintf( + `{"jsonrpc":"2.0","id":%q,"method":"initialize","params":{"protocolVersion":%q}}`, + "quitter", testProtocolVersion) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte(body))) + if !assert.NoError(t, err) { + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + resp, err := http.DefaultClient.Do(req) //nolint:bodyclose // err path returns no body + assert.Error(t, err, "the first caller is expected to abort mid-handshake") + if err == nil { + _ = resp.Body.Close() + } + }) + + // Give the first caller time to enter the flight and then abort. + time.Sleep(150 * time.Millisecond) + + // Second caller arrives while the shared handshake is still running. + decoded := postInitialize(t, url, "survivor") + assert.Nil(t, decoded["error"], + "the shared handshake must not be cancelled by the first caller disconnecting") + require.NotNil(t, decoded["result"], "second caller should still complete its handshake") + assert.Equal(t, "survivor", decoded["id"]) + + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for the aborting caller to finish") + } + + assert.Equal(t, 1, backend.count("initialize"), + "the disconnect must not cause a second upstream handshake") +} + +// TestMalformedInitializeReplyBecomesInternalError pins an intentional +// behaviour change made when the flight started sharing an id-free handshake +// rather than a whole response. +// +// A JSON-RPC success carrying neither an error nor a result is malformed for +// initialize -- InitializeResult requires protocolVersion, capabilities and +// serverInfo. It used to be forwarded to the client verbatim; it is now +// reported as an internal error, and nothing is cached, so the next client +// still gets a real handshake attempt. +func TestMalformedInitializeReplyBecomesInternalError(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + var mu sync.Mutex + var calls int + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + mu.Lock() + calls++ + first := calls == 1 + mu.Unlock() + if first { + // Success envelope with no result and no error. + return &jsonrpc2.Response{ID: req.ID} + } + return initializeResultFor(req) + }) + + malformed := postInitialize(t, url, "client-0") + require.NotNil(t, malformed["error"], "a result-less success must not be replayed to the client") + assert.Equal(t, "client-0", malformed["id"]) + + recovered := postInitialize(t, url, "client-1") + assert.Nil(t, recovered["error"], "the malformed reply must not be cached") + require.NotNil(t, recovered["result"], "the next client should get a real handshake") + + assert.Equal(t, 2, backend.count("initialize"), + "a malformed reply must leave the cache empty so the next client retries") +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy.go b/pkg/transport/proxy/streamable/streamable_proxy.go index f656b96015..5fb812eadf 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy.go +++ b/pkg/transport/proxy/streamable/streamable_proxy.go @@ -108,6 +108,11 @@ type HTTPProxy struct { // cross-session leakage. See dispatcher_routing.go. routing *sessionRouter + // initialize caches the backend's InitializeResult so only the first + // client handshake reaches the single stdio session behind this proxy and + // every later one is answered locally. See initialize_cache.go. + initialize *initializeCache + // uriLocks serializes each uri's resources/subscribe|unsubscribe ref-count // decision with the upstream forward it may trigger, so concurrent // subscribe/unsubscribe calls for the SAME uri from different sessions can @@ -233,6 +238,7 @@ func NewHTTPProxy( sessionTTL: session.DefaultSessionTTL, serverStreams: newServerStreamRegistry(), routing: newSessionRouter(), + initialize: newInitializeCache(), uriLocks: newKeyedMutex(), standaloneSSE: true, } @@ -1170,6 +1176,12 @@ func (p *HTTPProxy) interceptSessionScopedRequest( } switch req.Method { + case string(sdkmcp.MethodInitialize): + // Served from the shared InitializeResult cache after the first + // handshake, so the single stdio session behind this proxy never + // receives a second initialize. See initialize_cache.go. + return p.interceptInitialize(ctx, sessID, req), true + case methodResourcesSubscribe: uri, ok := extractStringParam(req.Params, "uri") if !ok || uri == "" { @@ -1544,6 +1556,14 @@ func (p *HTTPProxy) handleNotificationOrClientResponse(w http.ResponseWriter, se if sessID != "" { p.sessionManager.Get(sessID) } + // The backend completed a single handshake and expects a single + // notifications/initialized. Later clients' copies are acknowledged + // here without being forwarded, mirroring the initialize interception + // that answered their handshakes from cache. + if isInitializedNotification(msg) && !p.claimInitializedForward() { + w.WriteHeader(http.StatusAccepted) + return true + } if err := p.SendMessageToDestination(msg); err != nil { slog.Error("failed to send message to destination", "error", err) }