Skip to content

Commit 8ed5ec6

Browse files
committed
fix(mcp): queue and drop courtesy replies without blocking client callers
1 parent 85c15e0 commit 8ed5ec6

2 files changed

Lines changed: 114 additions & 20 deletions

File tree

‎internal/mcp/client.go‎

Lines changed: 69 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -49,20 +49,30 @@ type Client struct {
4949
stdin io.WriteCloser
5050
reader *messageReader
5151
writer *messageWriter
52-
mu sync.Mutex
5352
closeMu sync.Mutex
53+
idMu sync.Mutex
5454
nextID int
5555
cleanup func()
5656

57+
writeQueue chan writeOp
58+
writerOnce sync.Once
59+
5760
// dispatchMu guards the response-dispatch state shared with the single
58-
// reader goroutine. It is never held across a blocking read.
61+
// reader goroutine. It is never held across a blocking read or write.
5962
dispatchMu sync.Mutex
6063
readerOnce sync.Once
6164
pending map[int]chan dispatchResult
6265
readErr error
6366
readDone bool
6467
}
6568

69+
type writeOp struct {
70+
message rpcMessage
71+
done chan error
72+
}
73+
74+
const writeQueueCapacity = 32
75+
6676
// dispatchResult carries one matched JSON-RPC response (or a terminal reader
6777
// error) to a waiting caller.
6878
type dispatchResult struct {
@@ -307,20 +317,20 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
307317
return err
308318
}
309319

310-
// Allocate an id, register a response channel, and write the request while
311-
// holding client.mu. The mutex serializes writes and id allocation but is
312-
// released before the (potentially unbounded) wait for the response, so a
313-
// hung server never holds the lock and blocks other callers/Close.
314-
client.mu.Lock()
320+
// Allocate an id and register a response channel. ID allocation and dispatch
321+
// registrations are fast non-blocking operations. Message transmission is
322+
// handled via writeMessage, ensuring a hung server never blocks other callers
323+
// or prevents a caller with a deadline from giving up.
324+
client.idMu.Lock()
315325
id := client.nextID
316326
client.nextID++
327+
client.idMu.Unlock()
317328

318329
responses := make(chan dispatchResult, 1)
319330
client.dispatchMu.Lock()
320331
if client.readDone {
321332
readErr := client.readErr
322333
client.dispatchMu.Unlock()
323-
client.mu.Unlock()
324334
if readErr != nil {
325335
return readErr
326336
}
@@ -329,16 +339,14 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
329339
client.pending[id] = responses
330340
client.dispatchMu.Unlock()
331341

332-
if err := client.writer.write(rpcMessage{
342+
if err := client.writeMessage(ctx, rpcMessage{
333343
ID: id,
334344
Method: method,
335345
Params: rawParams,
336346
}); err != nil {
337347
client.removePending(id)
338-
client.mu.Unlock()
339348
return err
340349
}
341-
client.mu.Unlock()
342350

343351
select {
344352
case <-ctx.Done():
@@ -361,6 +369,43 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
361369
}
362370
}
363371

372+
// ensureWriter lazily starts the single writer goroutine.
373+
func (client *Client) ensureWriter() {
374+
client.writerOnce.Do(func() {
375+
if client.writeQueue == nil {
376+
client.writeQueue = make(chan writeOp, writeQueueCapacity)
377+
}
378+
go client.writeLoop()
379+
})
380+
}
381+
382+
func (client *Client) writeLoop() {
383+
for op := range client.writeQueue {
384+
err := client.writer.write(op.message)
385+
if op.done != nil {
386+
op.done <- err
387+
}
388+
}
389+
}
390+
391+
func (client *Client) writeMessage(ctx context.Context, message rpcMessage) error {
392+
client.ensureWriter()
393+
done := make(chan error, 1)
394+
op := writeOp{message: message, done: done}
395+
select {
396+
case <-ctx.Done():
397+
return ctx.Err()
398+
case client.writeQueue <- op:
399+
}
400+
401+
select {
402+
case <-ctx.Done():
403+
return ctx.Err()
404+
case err := <-done:
405+
return err
406+
}
407+
}
408+
364409
// ensureReader lazily starts the single reader goroutine. It runs once per
365410
// client; subsequent calls are no-ops.
366411
func (client *Client) ensureReader() {
@@ -389,21 +434,25 @@ func (client *Client) readLoop() {
389434
// It must never be routed as a response to a pending client request.
390435
if message.Method != "" {
391436
if message.ID != nil && jsonRPCIDEchoable(message.ID) {
392-
// Send the courtesy -32601 reply asynchronously off the read loop so
393-
// an undrained server stdin pipe never stalls readLoop or holds client.mu.
437+
// Send the courtesy -32601 reply via the bounded writer queue. If the
438+
// queue is saturated (e.g. an undrained server pipe), drop the reply
439+
// immediately so it never stalls readLoop, holds a mutex, or blocks callers.
440+
client.ensureWriter()
394441
id := message.ID
395442
method := message.Method
396-
go func() {
397-
client.mu.Lock()
398-
defer client.mu.Unlock()
399-
_ = client.writer.write(rpcMessage{
443+
courtesy := writeOp{
444+
message: rpcMessage{
400445
ID: id,
401446
Error: &rpcError{
402447
Code: -32601,
403448
Message: fmt.Sprintf("Method %q not supported", method),
404449
},
405-
})
406-
}()
450+
},
451+
}
452+
select {
453+
case client.writeQueue <- courtesy:
454+
default:
455+
}
407456
}
408457
continue
409458
}
@@ -519,7 +568,7 @@ func (client *Client) notify(method string, params any) error {
519568
if err != nil {
520569
return err
521570
}
522-
return client.writer.write(rpcMessage{
571+
return client.writeMessage(context.Background(), rpcMessage{
523572
Method: method,
524573
Params: rawParams,
525574
})

‎internal/mcp/client_test.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -987,3 +987,48 @@ func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) {
987987
// Expected: dropped cleanly
988988
}
989989
}
990+
991+
// TestStdioClientUndrainedServerDoesNotBlockCallerWithDeadline verifies that
992+
// when the server's input pipe is completely blocked and courtesy replies are
993+
// queued/dropped, a caller invoking request() with a deadline aborts cleanly
994+
// when the context expires rather than hanging indefinitely on a write or mutex.
995+
func TestStdioClientUndrainedServerDoesNotBlockCallerWithDeadline(t *testing.T) {
996+
inReader, inWriter := io.Pipe()
997+
outReader, outWriter := io.Pipe()
998+
999+
client := &Client{
1000+
reader: newMessageReader(inReader),
1001+
writer: newMessageWriter(outWriter),
1002+
pending: make(map[int]chan dispatchResult),
1003+
}
1004+
defer func() {
1005+
_ = inWriter.Close()
1006+
_ = outReader.Close()
1007+
}()
1008+
1009+
client.ensureReader()
1010+
1011+
// Flood server requests to saturate write queue while outReader is NOT drained.
1012+
for i := 1; i <= 50; i++ {
1013+
serverReq := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"roots/list","params":{}}`+"\n", i+100)
1014+
_, err := inWriter.Write([]byte(serverReq))
1015+
if err != nil {
1016+
t.Fatalf("failed to write server request: %v", err)
1017+
}
1018+
}
1019+
1020+
// Caller with short deadline invokes request()
1021+
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
1022+
defer cancel()
1023+
1024+
start := time.Now()
1025+
err := client.request(ctx, "tools/list", map[string]any{}, nil)
1026+
elapsed := time.Since(start)
1027+
1028+
if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {
1029+
t.Fatalf("expected deadline exceeded, got: %v", err)
1030+
}
1031+
if elapsed > 1*time.Second {
1032+
t.Fatalf("request took too long to abort on deadline: %v", elapsed)
1033+
}
1034+
}

0 commit comments

Comments
 (0)