Skip to content

Commit 39d4587

Browse files
authored
fix(upstream): serialize session re-initialization with in-flight requests (#1650)
## Root cause Spec 113-e (#1501) re-initializes a terminated Streamable HTTP session in place: `managed.reinitSession -> runReinit -> core.(*Client).ReinitializeSession` calls mcp-go `Initialize` on the **shared** mcp-go client. `Initialize` writes the client's negotiated protocol state (`applyNegotiatedVersion`, server/client capabilities, client info) without synchronization. Every other request on the same client reads that state: before it is sent (`sendRequest -> applyRequestMeta -> isModern`), and for tools/call again on the response (`callToolOnce`). So one caller re-initializing while another call is in flight, or just starting, is a data race in production, not only in the test. The CI "Unit Tests (shuffle)" lane (`go test -race -shuffle=on`) caught it in `TestSessionReinit_ConcurrentCallsSingleFlight`: ``` Read: mcp-go client.(*Client).isModern <- applyRequestMeta <- sendRequest <- callToolOnce <- CallTool <- core.(*Client).CallTool (internal/upstream/core/client.go:613) <- managed.callToolWithSession (internal/upstream/managed/session_reinit.go:323) Write: mcp-go client.(*Client).applyNegotiatedVersion <- initializeLegacy <- Initialize <- core.(*Client).ReinitializeSession (internal/upstream/core/session_reinit.go:57) <- managed.runReinit (:197) <- reinitSession (:156) <- callToolWithSession (:327) ``` The same race exists for tools/list, ping, prompts/list and prompts/get, because they use the same client. ## Fix (mcp-go is not modified) **core** - A per-client **session gate** (`golang.org/x/sync/semaphore`). The package was already in the module graph and is promoted from indirect to direct. `CallTool`, `ListTools`, `Ping`, `ListPrompts` and `GetPrompt` hold the gate shared for the whole mcp-go call. `ReinitializeSession` holds it exclusively around `Initialize` only. - Both sides honour their context: a re-init waiting behind a long call gives up at its 30s flight deadline instead of hanging. - The gate is FIFO (writer-preferring). - No self-deadlock: the managed layer triggers a re-init only after `core.CallTool` has returned and released the gate. - On tools/call, the gate is taken before the Spec 117 dispatch re-check and before `MarkDispatched`, so waiting at the gate sends nothing. - A **session generation**, bumped under the exclusive gate on every re-init. A request pinned to a generation with `WithSessionGeneration` is refused with `ErrSessionChanged` if a re-init ran after the pin was taken. Nothing is sent in that case. **managed** - tools/call sends go through `sendVerified`. It pins each send to the generation its identity and epoch checks ran against. A call that queued behind `Initialize` is therefore re-checked against the re-listed toolset (FR-082/FR-083a), instead of going out on a session whose tools/list has not been verified yet. The loop is bounded (`maxSessionAdmits`). - `publishBaseline` publishes the identity baseline and its verified generation together. A listing that a re-init overtook is dropped. If a re-init's verifying tools/list failed, its generation is listed again before any call is checked against it. ## Evidence - **Failing first:** the new `TestSessionReinit_RequestsDoNotRaceInitialize` (concurrent re-init + tools/call + tools/list + ping) reported `WARNING: DATA RACE` on every run before the fix, with the same stacks as CI. After the fix it passes. Each new overtaken-call or stale-listing test was also checked to fail with its guard disabled. - `go test -race -shuffle=on ./internal/upstream/managed/ -run TestSessionReinit -count=50`: ok - `go test -race ./internal/upstream/...` and `-tags server`: ok - The CI internal/server race suite (`-tags server`, unit-tests.yml `-skip` regex): all 13 packages ok, no races. - `golangci-lint` (`.github/.golangci.yml`), bare and with `--build-tags server`: no findings in the changed files. - **Live check:** a `-race` build of `mcpproxy serve` on a scratch port, data dir and config, plus a legacy Streamable HTTP upstream that issues session ids and answers 404 for unknown ones. Over 3 rounds, 40 concurrent `call_tool_read` calls each went through mcpproxy `/mcp`, and the upstream terminated all sessions before and in the middle of each burst. Result: **120/120 calls ok, 0 failed, 3 in-place re-inits, 0 `DATA RACE`** in the proxy log, and the server stayed connected and healthy. - Cross-review: 4 rounds of opencode `gpt-6.1-sol`. Rounds 1–3 each found one genuine High issue in the fix as it stood, and each was fixed: 1. A call queued behind `Initialize` skipped the identity re-check. 2. A failed verifying listing left the new session unverified. 3. A stale listing could overwrite a verified baseline. Round 4: **VERDICT: CLEAN**. No new dependencies.
1 parent bb7ddcc commit 39d4587

7 files changed

Lines changed: 558 additions & 57 deletions

File tree

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ require (
3939
go.opentelemetry.io/otel/trace v1.46.0
4040
go.uber.org/zap v1.28.0
4141
golang.org/x/mod v0.41.0
42+
golang.org/x/sync v0.23.0
4243
golang.org/x/sys v0.48.0
4344
golang.org/x/term v0.46.0
4445
golang.org/x/text v0.42.0
@@ -145,7 +146,6 @@ require (
145146
go.uber.org/multierr v1.10.0 // indirect
146147
go.yaml.in/yaml/v3 v3.0.5 // indirect
147148
golang.org/x/net v0.58.0 // indirect
148-
golang.org/x/sync v0.23.0 // indirect
149149
golang.org/x/tools v0.49.0 // indirect
150150
google.golang.org/genproto/googleapis/api v0.0.0-20260819154853-08b0e4226688 // indirect
151151
google.golang.org/genproto/googleapis/rpc v0.0.0-20260819154853-08b0e4226688 // indirect

‎internal/upstream/core/client.go‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,12 @@ type Client struct {
9494
// when multiple requests are in-flight simultaneously
9595
sseRequestMu sync.Mutex
9696

97+
// sessionGate keeps requests off the mcp-go client while a session
98+
// re-initialize rewrites its protocol state (see ReinitializeSession).
99+
sessionGate sessionGate
100+
// sessionGen counts session re-initializes (see SessionGeneration).
101+
sessionGen atomic.Uint64
102+
97103
// retryAfter collects the `Retry-After` hints this upstream's HTTP/SSE
98104
// responses carry. mcp-go flattens a 429 into an error string long before
99105
// the connection state machine sees it, so the hint is captured by a
@@ -345,6 +351,11 @@ func (c *Client) Ping(ctx context.Context) error {
345351
// pin MUST come with an era-aware replacement, or spec 074's health loop
346352
// goes blind; that is tracked as a required task of the pin lift, not as
347353
// cleanup.
354+
release, err := c.beginRequest(ctx, string(mcp.MethodPing))
355+
if err != nil {
356+
return err
357+
}
358+
defer release()
348359
//nolint:staticcheck // SA1019: valid on the legacy era this client is pinned to; see above.
349360
return client.Ping(ctx)
350361
}
@@ -388,8 +399,13 @@ func (c *Client) ListTools(ctx context.Context) ([]*config.ToolMetadata, error)
388399
c.logger.Info("Making direct tools list call to upstream server",
389400
zap.String("server", c.config.Name))
390401

402+
release, err := c.beginRequest(ctx, string(mcp.MethodToolsList))
403+
if err != nil {
404+
return nil, err
405+
}
391406
listReq := mcp.ListToolsRequest{}
392407
toolsResult, err := client.ListTools(ctx, listReq)
408+
release()
393409
if err != nil {
394410
// Debug, not Error: both callers above (managed.Client.ListTools and
395411
// upstream.Manager.discoverTools) log this same failure, so logging it
@@ -526,11 +542,21 @@ func (c *Client) CallTool(ctx context.Context, toolName string, args map[string]
526542
zap.String("server", c.config.Name),
527543
zap.String("tool", toolName))
528544
}
545+
// A session re-initialize in progress on this client finishes before the
546+
// call is handed to mcp-go, and none starts until it returns (the gate is
547+
// held across the whole mcp-go call, which reads the client's protocol
548+
// state on the response too). Waiting here sends nothing.
549+
release, err := c.beginRequest(ctx, string(mcp.MethodToolsCall))
550+
if err != nil {
551+
return nil, err
552+
}
553+
defer release()
554+
529555
// Spec 117 FR-018: the caller's policy verdicts are re-read here, after
530-
// the SSE request lock (a call may have waited behind a long in-flight
531-
// call while a capture failure marked the server review-dirty) and before
532-
// every send, including a session re-init's retry. Nothing is sent when
533-
// a check refuses.
556+
// the SSE request lock and the session gate (a call may have waited behind
557+
// a long in-flight call while a capture failure marked the server
558+
// review-dirty) and before every send, including a session re-init's
559+
// retry. Nothing is sent when a check refuses.
534560
if err := RunDispatchRecheck(ctx); err != nil {
535561
return nil, err
536562
}

‎internal/upstream/core/prompts.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,12 @@ func (c *Client) ListPrompts(ctx context.Context) ([]mcp.Prompt, error) {
7474
// ListPromptsByPage instead and stop at maxListPromptsPages pages, at
7575
// maxListPromptsItems accumulated prompts (aligned with the per-server trim in
7676
// Manager.ListPrompts), or on ctx expiry — whichever comes first (Finding F14).
77+
release, err := c.beginRequest(ctx, string(mcp.MethodPromptsList))
78+
if err != nil {
79+
return nil, err
80+
}
81+
defer release()
82+
7783
listReq := mcp.ListPromptsRequest{}
7884
var prompts []mcp.Prompt
7985
for page := 1; ; page++ {
@@ -152,7 +158,12 @@ func (c *Client) GetPrompt(ctx context.Context, name string, args map[string]str
152158
request.Params.Name = name
153159
request.Params.Arguments = args
154160

161+
release, err := c.beginRequest(ctx, string(mcp.MethodPromptsGet))
162+
if err != nil {
163+
return nil, err
164+
}
155165
result, err := upstreamClient.GetPrompt(ctx, request)
166+
release()
156167
if err != nil {
157168
return nil, fmt.Errorf("GetPrompt failed for '%s': %w", name, err)
158169
}
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package core
2+
3+
import (
4+
"context"
5+
"errors"
6+
"testing"
7+
"time"
8+
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
// A re-initialize waits for requests already on the client and gives up at its
13+
// own deadline instead of hanging behind a long in-flight call.
14+
func TestSessionGate_ExclusiveWaitsForSharedAndHonoursContext(t *testing.T) {
15+
var g sessionGate
16+
release, err := g.acquireShared(context.Background())
17+
require.NoError(t, err)
18+
19+
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
20+
defer cancel()
21+
_, err = g.acquireExclusive(ctx)
22+
require.True(t, errors.Is(err, context.DeadlineExceeded), "got %v", err)
23+
24+
release()
25+
excl, err := g.acquireExclusive(context.Background())
26+
require.NoError(t, err)
27+
excl()
28+
}
29+
30+
// A request queued behind a re-initialize returns when its caller goes away,
31+
// and is admitted once the re-initialize releases the gate.
32+
func TestSessionGate_SharedQueuesBehindExclusive(t *testing.T) {
33+
c := &Client{}
34+
g := &c.sessionGate
35+
excl, err := g.acquireExclusive(context.Background())
36+
require.NoError(t, err)
37+
38+
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
39+
defer cancel()
40+
_, err = c.beginRequest(ctx, "tools/call")
41+
require.True(t, errors.Is(err, context.DeadlineExceeded), "got %v", err)
42+
43+
admitted := make(chan struct{})
44+
go func() {
45+
rel, err := g.acquireShared(context.Background())
46+
if err == nil {
47+
rel()
48+
}
49+
close(admitted)
50+
}()
51+
select {
52+
case <-admitted:
53+
t.Fatal("request admitted while a re-initialize holds the gate")
54+
case <-time.After(50 * time.Millisecond):
55+
}
56+
excl()
57+
select {
58+
case <-admitted:
59+
case <-time.After(5 * time.Second):
60+
t.Fatal("request not admitted after the re-initialize released the gate")
61+
}
62+
}

‎internal/upstream/core/session_reinit.go‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,104 @@ package core
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
7+
"math"
8+
"sync"
69

710
"github.com/mark3labs/mcp-go/client/transport"
811
"github.com/mark3labs/mcp-go/mcp"
912
"go.uber.org/zap"
13+
"golang.org/x/sync/semaphore"
1014
)
1115

16+
// sessionGate serializes a session re-initialize with the requests sent on the
17+
// same mcp-go client. mcp-go's Initialize writes the client's negotiated
18+
// protocol state (protocol version, server and client capabilities, client
19+
// info) without synchronization, and every request reads that state: before it
20+
// is sent (sendRequest -> applyRequestMeta/applyStandardHeaders) and, for
21+
// tools/call, again on the response. So a request holds the gate shared for
22+
// the whole mcp-go call and a re-initialize holds it exclusively for the
23+
// Initialize handshake only.
24+
//
25+
// It is a weighted semaphore rather than a sync.RWMutex so both sides honour
26+
// their context: a re-initialize waiting behind a long in-flight call gives up
27+
// at the flight's deadline instead of hanging, and a request queued behind a
28+
// re-initialize returns when its caller goes away, with nothing sent. Like
29+
// sync.RWMutex it is writer-preferring (FIFO): once a re-initialize is
30+
// waiting, new requests queue behind it rather than starving it.
31+
//
32+
// The zero value is ready; the semaphore is created on first use because core
33+
// clients are also built as struct literals.
34+
type sessionGate struct {
35+
once sync.Once
36+
sem *semaphore.Weighted
37+
}
38+
39+
const sessionGateWeight = math.MaxInt64
40+
41+
func (g *sessionGate) weighted() *semaphore.Weighted {
42+
g.once.Do(func() { g.sem = semaphore.NewWeighted(sessionGateWeight) })
43+
return g.sem
44+
}
45+
46+
// acquireShared admits one request; release it when the mcp-go call returns.
47+
func (g *sessionGate) acquireShared(ctx context.Context) (release func(), err error) {
48+
sem := g.weighted()
49+
if err := sem.Acquire(ctx, 1); err != nil {
50+
return nil, err
51+
}
52+
return func() { sem.Release(1) }, nil
53+
}
54+
55+
// acquireExclusive waits for every admitted request to finish and holds new
56+
// ones back until release is called.
57+
func (g *sessionGate) acquireExclusive(ctx context.Context) (release func(), err error) {
58+
sem := g.weighted()
59+
if err := sem.Acquire(ctx, sessionGateWeight); err != nil {
60+
return nil, err
61+
}
62+
return func() { sem.Release(sessionGateWeight) }, nil
63+
}
64+
65+
// ErrSessionChanged is returned, with nothing sent, by a request whose context
66+
// was pinned to a session generation (WithSessionGeneration) when a session
67+
// re-initialize ran after the pin was taken. The caller re-validates what it
68+
// checked against the old session and sends again.
69+
var ErrSessionChanged = errors.New("upstream session was re-initialized before the request was sent")
70+
71+
type sessionGenerationKey struct{}
72+
73+
// WithSessionGeneration pins the requests made with ctx to session generation
74+
// gen (from SessionGeneration): such a request is refused with
75+
// ErrSessionChanged, before anything is sent, if a session re-initialize has
76+
// run since. The managed layer pins tools/call this way so a call checked
77+
// against one session's toolset is never sent on a session whose toolset has
78+
// not been verified yet.
79+
func WithSessionGeneration(ctx context.Context, gen uint64) context.Context {
80+
return context.WithValue(ctx, sessionGenerationKey{}, gen)
81+
}
82+
83+
// SessionGeneration counts the session re-initializes run on this client. It
84+
// changes while ReinitializeSession holds the session gate exclusively, so a
85+
// request admitted to the gate sees a stable value for its whole send.
86+
func (c *Client) SessionGeneration() uint64 { return c.sessionGen.Load() }
87+
88+
// beginRequest admits a request to the mcp-go client, waiting out a session
89+
// re-initialize in progress. A context error is wrapped so it still matches
90+
// context.Canceled / context.DeadlineExceeded.
91+
func (c *Client) beginRequest(ctx context.Context, method string) (release func(), err error) {
92+
release, err = c.sessionGate.acquireShared(ctx)
93+
if err != nil {
94+
return nil, fmt.Errorf("%s not sent: waiting for session re-initialize: %w", method, err)
95+
}
96+
if gen, ok := ctx.Value(sessionGenerationKey{}).(uint64); ok && gen != c.sessionGen.Load() {
97+
release()
98+
return nil, fmt.Errorf("%s not sent: %w", method, ErrSessionChanged)
99+
}
100+
return release, nil
101+
}
102+
12103
// Spec 113-e (G8): helpers that let the managed layer re-initialize a
13104
// Streamable HTTP MCP session in place after the upstream answered HTTP 404
14105
// (mcp-go: transport.ErrSessionTerminated) for a session id it no longer
@@ -54,7 +145,18 @@ func (c *Client) ReinitializeSession(ctx context.Context) error {
54145
req.Params.ClientInfo = mcp.Implementation{Name: "mcpproxy-go", Version: "1.0.0"}
55146
req.Params.Capabilities = mcp.ClientCapabilities{}
56147

148+
// Initialize rewrites the shared client's protocol state, so it runs only
149+
// once every request already on this client has returned, and holds new
150+
// ones until it is done.
151+
release, err := c.sessionGate.acquireExclusive(ctx)
152+
if err != nil {
153+
return fmt.Errorf("session re-initialize: waiting for in-flight requests: %w", err)
154+
}
155+
// Initialize mutates the client's session even when it fails, so the
156+
// generation moves either way.
157+
c.sessionGen.Add(1)
57158
res, err := cl.Initialize(ctx, req)
159+
release()
58160
if err != nil {
59161
return fmt.Errorf("session re-initialize failed: %w", err)
60162
}

0 commit comments

Comments
 (0)