From 2e575f12aee5b84cb2c73838074e56bdb9c5ffbb Mon Sep 17 00:00:00 2001 From: amirejaz Date: Thu, 30 Jul 2026 22:11:08 +0100 Subject: [PATCH 1/4] Serve repeat initialize from a cached result A stdio MCP server is one process and therefore exactly one MCP session, while the streamable proxy fans an arbitrary number of HTTP client sessions onto it. Every client's initialize was forwarded to that single session, so only the first could succeed: go-sdk v1.7+ rejects a second handshake with `duplicate "initialize" received`. Behind vMCP that takes down far more than the reported health check -- every call verb builds a fresh client and handshakes, so tool calls fail too once the aggregator has claimed the one handshake at startup. Cache the InitializeResult from the first handshake and answer later ones locally, restoring what go-sdk v1.6.1 and earlier did when it replayed a repeated initialize, but in the component that actually owns the multiplexing. The cache lock is held across the first upstream round trip so concurrent handshakes single-flight instead of racing a second initialize to the backend; forwardUpstream's timeout bounds it. Errors are never cached, so a transient failure cannot pin every later client to one bad response. Swallow notifications/initialized after the first for the same reason: the backend completed one handshake and expects one such notification. Later ones are still acknowledged with 202. Server capabilities in the replayed result do not vary by client, and client capabilities are inert here because server-initiated requests are already rejected by the shared proxy. The negotiated protocol version is the one genuinely shared value: a client that asked for a different one receives the first client's. Closes #1982 Co-Authored-By: Claude Opus 5 (1M context) --- .../proxy/streamable/initialize_cache.go | 125 +++++++++ .../proxy/streamable/initialize_cache_test.go | 249 ++++++++++++++++++ .../proxy/streamable/streamable_proxy.go | 20 ++ 3 files changed, 394 insertions(+) create mode 100644 pkg/transport/proxy/streamable/initialize_cache.go create mode 100644 pkg/transport/proxy/streamable/initialize_cache_test.go diff --git a/pkg/transport/proxy/streamable/initialize_cache.go b/pkg/transport/proxy/streamable/initialize_cache.go new file mode 100644 index 0000000000..9dc9b5face --- /dev/null +++ b/pkg/transport/proxy/streamable/initialize_cache.go @@ -0,0 +1,125 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package streamable + +import ( + "context" + "encoding/json" + "sync" + + "golang.org/x/exp/jsonrpc2" + + sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// 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 { + // mu guards both fields below and additionally serializes the first + // upstream handshake (see HTTPProxy.interceptInitialize), so a concurrent + // second initialize waits for the first to finish and then reads the + // cache instead of racing another handshake to the backend. + mu sync.Mutex + + // result is the raw InitializeResult of the first successful handshake, + // nil until one completes. Once assigned it is never mutated, so it is + // safe to hand the same bytes to concurrent responses. + result json.RawMessage + + // initializedForwarded records whether a notifications/initialized has + // already been passed to the backend. + initializedForwarded bool +} + +// 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. +// +// mu is held across the upstream round trip on the first call. That is what +// makes the forward single-flight, and it is bounded: forwardUpstream applies +// p.requestTimeout, so an unresponsive backend cannot block other sessions' +// handshakes indefinitely -- the same reasoning uriLocks relies on for +// resources/subscribe. +// +// 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 { + p.initialize.mu.Lock() + defer p.initialize.mu.Unlock() + + if p.initialize.result != nil { + return &jsonrpc2.Response{ID: req.ID, Result: p.initialize.result} + } + + msg, err := p.forwardUpstream(ctx, sessID, req) + if err != nil { + return upstreamErrorResponse(req.ID, err) + } + if resp, ok := msg.(*jsonrpc2.Response); ok && resp.Error == nil && len(resp.Result) > 0 { + p.initialize.result = resp.Result + } + return msg +} + +// 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 { + p.initialize.mu.Lock() + defer p.initialize.mu.Unlock() + + if p.initialize.initializedForwarded { + return false + } + p.initialize.initializedForwarded = true + return 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..9504fafa37 --- /dev/null +++ b/pkg/transport/proxy/streamable/initialize_cache_test.go @@ -0,0 +1,249 @@ +// 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 +} + +// 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 +} + +// TestInitializeReachesBackendOnceAndIsCachedThereafter verifies that only the +// first client handshake is forwarded to the single stdio session, and that +// every later client is answered from the cached InitializeResult with its own +// JSON-RPC id echoed back. +// +// Without the cache each handshake reaches the backend, and go-sdk v1.7+ +// rejects every one after the first with `duplicate "initialize" received` +// (#5890). +func TestInitializeReachesBackendOnceAndIsCachedThereafter(t *testing.T) { + t.Parallel() + + proxy, url := startInitTestProxy(t) + backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { + return initializeResultFor(req) + }) + + 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") + 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) + 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 initializeResultFor(req) + }) + + const clients = 5 + var wg sync.WaitGroup + for i := range clients { + wg.Go(func() { + decoded := postInitialize(t, url, fmt.Sprintf("client-%d", i)) + assert.NotNil(t, decoded["result"], "every concurrent handshake should succeed") + }) + } + 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, func(req *jsonrpc2.Request) jsonrpc2.Message { + return initializeResultFor(req) + }) + + 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") +} 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) } From a5303d5d3970d3b6447e1eeb386f6ff79abbdd47 Mon Sep 17 00:00:00 2001 From: amirejaz Date: Thu, 30 Jul 2026 22:28:55 +0100 Subject: [PATCH 2/4] Assert handshakes against a duplicate-rejecting backend The initialize cache tests used a permissive stub that accepted every handshake, so they could only count how many messages reached the backend -- they could never produce the error the cache exists to prevent. Give the fake backend go-sdk v1.7+ session semantics: answer the first initialize and reject every later one with the same `duplicate "initialize" received` error ServerSession.initialize returns. The tests now assert the outcome that matters, that a later client's handshake succeeds, and fail with the literal error from #5890 when the cache is removed rather than only with a mismatched count. Co-Authored-By: Claude Opus 5 (1M context) --- .../proxy/streamable/initialize_cache_test.go | 61 ++++++++++++++----- 1 file changed, 45 insertions(+), 16 deletions(-) diff --git a/pkg/transport/proxy/streamable/initialize_cache_test.go b/pkg/transport/proxy/streamable/initialize_cache_test.go index 9504fafa37..70be404bfc 100644 --- a/pkg/transport/proxy/streamable/initialize_cache_test.go +++ b/pkg/transport/proxy/streamable/initialize_cache_test.go @@ -88,6 +88,37 @@ func initializeResultFor(req *jsonrpc2.Request) jsonrpc2.Message { 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) { @@ -122,27 +153,24 @@ func postInitialize(t *testing.T, url string, id string) map[string]any { return decoded } -// TestInitializeReachesBackendOnceAndIsCachedThereafter verifies that only the -// first client handshake is forwarded to the single stdio session, and that -// every later client is answered from the cached InitializeResult with its own -// JSON-RPC id echoed back. -// -// Without the cache each handshake reaches the backend, and go-sdk v1.7+ -// rejects every one after the first with `duplicate "initialize" received` -// (#5890). -func TestInitializeReachesBackendOnceAndIsCachedThereafter(t *testing.T) { +// 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, func(req *jsonrpc2.Request) jsonrpc2.Message { - return initializeResultFor(req) - }) + 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"]) } @@ -161,10 +189,11 @@ 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 initializeResultFor(req) + return strict(req) }) const clients = 5 @@ -172,6 +201,8 @@ func TestConcurrentInitializeReachesBackendOnce(t *testing.T) { for i := range clients { wg.Go(func() { decoded := postInitialize(t, url, fmt.Sprintf("client-%d", i)) + 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") }) } @@ -226,9 +257,7 @@ func TestDuplicateInitializedNotificationIsNotForwarded(t *testing.T) { t.Parallel() proxy, url := startInitTestProxy(t) - backend := startFakeStdioBackend(t.Context(), proxy, func(req *jsonrpc2.Request) jsonrpc2.Message { - return initializeResultFor(req) - }) + backend := startFakeStdioBackend(t.Context(), proxy, strictInitializeResponder()) for range 3 { resp, err := http.Post(url, "application/json", //nolint:gosec // test-local URL From 5d17dea02953eb8c87fa7bef59eb043dd1966ca5 Mon Sep 17 00:00:00 2001 From: amirejaz Date: Fri, 31 Jul 2026 20:16:43 +0100 Subject: [PATCH 3/4] Collapse concurrent handshakes with singleflight Holding the cache mutex across the upstream round trip serialized concurrent first handshakes: each waiter ran its own forward once the previous released the lock, so against an unresponsive backend the Nth client waited N*requestTimeout -- 60s each -- before it could even start. Per-call boundedness is not queue boundedness. singleflight collapses them into one upstream call whose outcome every waiter shares, so a slow backend costs one timeout rather than one per client. The mutex now only guards the cached result and the initialized-notification flag, and is never held across a network call. Two consequences the shared flight introduces, both handled: Run the upstream call on a context detached from the caller's request. singleflight invokes the shared function with whichever context the first caller supplied, so its disconnect would otherwise cancel the handshake for every waiter -- and worse, the cancelled call still reaches the backend, so the next attempt hits an already-initialized session and fails with `duplicate "initialize" received`, reintroducing the bug this cache exists to fix. forwardUpstream still applies requestTimeout to the detached context. Rebuild every response with the calling request's own JSON-RPC id. A shared result carries the id of whichever caller performed the forward, and returning it verbatim would hand later waiters an id they never sent. Co-Authored-By: Claude Opus 5 (1M context) --- .../proxy/streamable/initialize_cache.go | 97 +++++++++++++++---- .../proxy/streamable/initialize_cache_test.go | 70 ++++++++++++- 2 files changed, 148 insertions(+), 19 deletions(-) diff --git a/pkg/transport/proxy/streamable/initialize_cache.go b/pkg/transport/proxy/streamable/initialize_cache.go index 9dc9b5face..2c4aed75a8 100644 --- a/pkg/transport/proxy/streamable/initialize_cache.go +++ b/pkg/transport/proxy/streamable/initialize_cache.go @@ -6,13 +6,20 @@ package streamable import ( "context" "encoding/json" + "errors" "sync" "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. @@ -34,10 +41,14 @@ import ( // InitializeResult carries the backend's protocol version, capabilities, // instructions and server info, none of which vary by client. type initializeCache struct { - // mu guards both fields below and additionally serializes the first - // upstream handshake (see HTTPProxy.interceptInitialize), so a concurrent - // second initialize waits for the first to finish and then reads the - // cache instead of racing another handshake to the backend. + // flight collapses concurrent first 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 + + // mu guards the two fields below. It is never held across an upstream call. mu sync.Mutex // result is the raw InitializeResult of the first successful handshake, @@ -50,11 +61,32 @@ type initializeCache struct { initializedForwarded bool } +// 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{} } +// cachedResult returns the cached InitializeResult, or nil if no handshake has +// succeeded yet. +func (c *initializeCache) cachedResult() json.RawMessage { + c.mu.Lock() + defer c.mu.Unlock() + return c.result +} + +// storeResult records the InitializeResult of the first successful handshake. +func (c *initializeCache) storeResult(result json.RawMessage) { + c.mu.Lock() + defer c.mu.Unlock() + if c.result == nil { + c.result = result + } +} + // 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. @@ -64,11 +96,17 @@ func newInitializeCache() *initializeCache { // handshake for real, rather than pinning every future client to one bad // response. // -// mu is held across the upstream round trip on the first call. That is what -// makes the forward single-flight, and it is bounded: forwardUpstream applies -// p.requestTimeout, so an unresponsive backend cannot block other sessions' -// handshakes indefinitely -- the same reasoning uriLocks relies on for -// resources/subscribe. +// 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 @@ -79,21 +117,44 @@ func newInitializeCache() *initializeCache { func (p *HTTPProxy) interceptInitialize( ctx context.Context, sessID string, req *jsonrpc2.Request, ) jsonrpc2.Message { - p.initialize.mu.Lock() - defer p.initialize.mu.Unlock() - - if p.initialize.result != nil { - return &jsonrpc2.Response{ID: req.ID, Result: p.initialize.result} + if cached := p.initialize.cachedResult(); cached != nil { + return &jsonrpc2.Response{ID: req.ID, Result: cached} } - msg, err := p.forwardUpstream(ctx, sessID, req) + v, err, _ := p.initialize.flight.Do(initializeFlightKey, func() (any, error) { + // A handshake may have completed between the read above and entering + // the flight. + if cached := p.initialize.cachedResult(); cached != nil { + return cached, nil + } + msg, ferr := p.forwardUpstream(context.WithoutCancel(ctx), sessID, req) + if ferr != nil { + return nil, ferr + } + if resp, ok := msg.(*jsonrpc2.Response); ok && resp.Error == nil && len(resp.Result) > 0 { + p.initialize.storeResult(resp.Result) + return resp.Result, nil + } + // A JSON-RPC error (or an unexpected shape) is shared with this + // flight's waiters but deliberately not cached, so the next client + // retries the handshake for real rather than being pinned to one bad + // response. + return msg, nil + }) if err != nil { return upstreamErrorResponse(req.ID, err) } - if resp, ok := msg.(*jsonrpc2.Response); ok && resp.Error == nil && len(resp.Result) > 0 { - p.initialize.result = resp.Result + + switch out := v.(type) { + case json.RawMessage: + return &jsonrpc2.Response{ID: req.ID, Result: out} + case *jsonrpc2.Response: + return &jsonrpc2.Response{ID: req.ID, Result: out.Result, Error: out.Error} + default: + // Not a response we can re-address to this caller; surface an error + // rather than replying with another caller's id. + return upstreamErrorResponse(req.ID, errUnexpectedInitializeReply) } - return msg } // claimInitializedForward reports whether a client's notifications/initialized diff --git a/pkg/transport/proxy/streamable/initialize_cache_test.go b/pkg/transport/proxy/streamable/initialize_cache_test.go index 70be404bfc..d12ca36103 100644 --- a/pkg/transport/proxy/streamable/initialize_cache_test.go +++ b/pkg/transport/proxy/streamable/initialize_cache_test.go @@ -200,10 +200,16 @@ func TestConcurrentInitializeReachesBackendOnce(t *testing.T) { var wg sync.WaitGroup for i := range clients { wg.Go(func() { - decoded := postInitialize(t, url, fmt.Sprintf("client-%d", i)) + 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{}) @@ -276,3 +282,65 @@ func TestDuplicateInitializedNotificationIsNotForwarded(t *testing.T) { 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") +} From eab13acfa6b341b925d626ab5dd1ed6be584d290 Mon Sep 17 00:00:00 2001 From: amirejaz Date: Fri, 31 Jul 2026 21:57:09 +0100 Subject: [PATCH 4/4] Confine handshake cache to the flight callback The mutex existed only because the cached result was also read on a fast path outside the flight. 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 a field touched solely inside the callback can never have overlapping accesses. Dropping the fast path makes the re-check already in the callback the only cache read, and the lock guards nothing. This matches the shape of the repo's other singleflight user, pkg/cache.ValidatingCache: all cache access lives inside the flight and the callback shares a small result struct rather than a whole value. Share an id-free handshake{result, rpcErr} instead of a *jsonrpc2.Response, so each caller builds its own reply. A waiter can no longer receive the forwarding caller's JSON-RPC id by construction, which retires the type switch and its default branch. initializedForwarded becomes an atomic.Bool. It 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. One intentional behaviour change: a success carrying neither error nor result was forwarded verbatim and is now an internal error. That envelope is malformed for initialize, which requires protocolVersion, capabilities and serverInfo. Nothing is cached, so the next client still gets a real handshake attempt. Co-Authored-By: Claude Opus 5 (1M context) --- .../proxy/streamable/initialize_cache.go | 126 +++++++++--------- .../proxy/streamable/initialize_cache_test.go | 39 ++++++ 2 files changed, 99 insertions(+), 66 deletions(-) diff --git a/pkg/transport/proxy/streamable/initialize_cache.go b/pkg/transport/proxy/streamable/initialize_cache.go index 2c4aed75a8..654c22e07d 100644 --- a/pkg/transport/proxy/streamable/initialize_cache.go +++ b/pkg/transport/proxy/streamable/initialize_cache.go @@ -7,7 +7,7 @@ import ( "context" "encoding/json" "errors" - "sync" + "sync/atomic" "golang.org/x/exp/jsonrpc2" "golang.org/x/sync/singleflight" @@ -41,24 +41,41 @@ var errUnexpectedInitializeReply = errors.New("backend returned an unexpected re // InitializeResult carries the backend's protocol version, capabilities, // instructions and server info, none of which vary by client. type initializeCache struct { - // flight collapses concurrent first 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 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 - // mu guards the two fields below. It is never held across an upstream call. - mu sync.Mutex - // result is the raw InitializeResult of the first successful handshake, - // nil until one completes. Once assigned it is never mutated, so it is - // safe to hand the same bytes to concurrent responses. + // 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. - initializedForwarded bool + // 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 @@ -70,23 +87,6 @@ func newInitializeCache() *initializeCache { return &initializeCache{} } -// cachedResult returns the cached InitializeResult, or nil if no handshake has -// succeeded yet. -func (c *initializeCache) cachedResult() json.RawMessage { - c.mu.Lock() - defer c.mu.Unlock() - return c.result -} - -// storeResult records the InitializeResult of the first successful handshake. -func (c *initializeCache) storeResult(result json.RawMessage) { - c.mu.Lock() - defer c.mu.Unlock() - if c.result == nil { - c.result = result - } -} - // 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. @@ -117,44 +117,45 @@ func (c *initializeCache) storeResult(result json.RawMessage) { func (p *HTTPProxy) interceptInitialize( ctx context.Context, sessID string, req *jsonrpc2.Request, ) jsonrpc2.Message { - if cached := p.initialize.cachedResult(); cached != nil { - return &jsonrpc2.Response{ID: req.ID, Result: cached} - } - - v, err, _ := p.initialize.flight.Do(initializeFlightKey, func() (any, error) { - // A handshake may have completed between the read above and entering - // the flight. - if cached := p.initialize.cachedResult(); cached != nil { - return cached, nil + 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 + 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 resp, ok := msg.(*jsonrpc2.Response); ok && resp.Error == nil && len(resp.Result) > 0 { - p.initialize.storeResult(resp.Result) - return resp.Result, nil + if len(resp.Result) == 0 { + return nil, errUnexpectedInitializeReply // no result and no error } - // A JSON-RPC error (or an unexpected shape) is shared with this - // flight's waiters but deliberately not cached, so the next client - // retries the handshake for real rather than being pinned to one bad - // response. - return msg, nil + + c.result = resp.Result + return handshake{result: resp.Result}, nil }) if err != nil { return upstreamErrorResponse(req.ID, err) } - switch out := v.(type) { - case json.RawMessage: - return &jsonrpc2.Response{ID: req.ID, Result: out} - case *jsonrpc2.Response: - return &jsonrpc2.Response{ID: req.ID, Result: out.Result, Error: out.Error} - default: - // Not a response we can re-address to this caller; surface an error - // rather than replying with another caller's id. - return upstreamErrorResponse(req.ID, errUnexpectedInitializeReply) - } + 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 @@ -165,14 +166,7 @@ func (p *HTTPProxy) interceptInitialize( // swallowed. Swallowed notifications are still acknowledged with 202, which is // all the sending client expects -- notifications carry no response. func (p *HTTPProxy) claimInitializedForward() bool { - p.initialize.mu.Lock() - defer p.initialize.mu.Unlock() - - if p.initialize.initializedForwarded { - return false - } - p.initialize.initializedForwarded = true - return true + return p.initialize.initializedForwarded.CompareAndSwap(false, true) } // isInitializedNotification reports whether msg is a client's diff --git a/pkg/transport/proxy/streamable/initialize_cache_test.go b/pkg/transport/proxy/streamable/initialize_cache_test.go index d12ca36103..e6127ecf0f 100644 --- a/pkg/transport/proxy/streamable/initialize_cache_test.go +++ b/pkg/transport/proxy/streamable/initialize_cache_test.go @@ -344,3 +344,42 @@ func TestInitializeSurvivesFirstCallerDisconnect(t *testing.T) { 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") +}