From 69b71517e72fcdd4cebc63e9a68d313304c4a57b Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sat, 29 Aug 2026 02:11:52 +0200 Subject: [PATCH 01/11] security(mcp): prevent server requests from matching client pending responses --- internal/mcp/client.go | 16 ++++++++++ internal/mcp/client_test.go | 63 +++++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 064e7f213..f7bae0bc6 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -385,6 +385,22 @@ func (client *Client) readLoop() { client.failAll(err) return } + // A message with a Method is a server-initiated request or notification. + // It must never be routed as a response to a pending client request. + if message.Method != "" { + if message.ID != nil { + client.mu.Lock() + _ = client.writer.write(rpcMessage{ + ID: message.ID, + Error: &rpcError{ + Code: -32601, + Message: fmt.Sprintf("Method %q not supported", message.Method), + }, + }) + client.mu.Unlock() + } + continue + } if message.ID == nil { continue } diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 8a66b7f69..a15a7b545 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -842,3 +842,66 @@ func TestBoundedBufferCapsRetainedBytes(t *testing.T) { t.Fatalf("retained %q, want %q (capped at 8 bytes, head kept)", got, "hellowor") } } + +func TestStdioClientIgnoresServerInitiatedRequestsInPendingResponses(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + + // Drain output responses from client + go func() { + buf := make([]byte, 1024) + for { + if _, err := outReader.Read(buf); err != nil { + return + } + } + }() + + client.ensureReader() + + // Register a pending response for ID 1 + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + // Simulate server sending a request with ID 1 ("roots/list") + serverReq := `{"jsonrpc":"2.0","id":1,"method":"roots/list","params":{}}` + "\n" + go func() { + _, _ = inWriter.Write([]byte(serverReq)) + }() + + // The pending channel for client ID 1 should NOT receive the server request. + select { + case res := <-responses: + t.Fatalf("pending request 1 received server request: %#v", res.message) + case <-time.After(100 * time.Millisecond): + // Expected: server request was not misdelivered to pending client request + } + + // Now send the actual response for ID 1 + serverResp := `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}` + "\n" + go func() { + _, _ = inWriter.Write([]byte(serverResp)) + }() + + select { + case res := <-responses: + if res.message.Method != "" || len(res.message.Result) == 0 { + t.Fatalf("expected valid response, got %#v", res.message) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("timed out waiting for actual response") + } +} + From 1fed7e37d0fff28bad7d26a68be9dc3fd97f93b7 Mon Sep 17 00:00:00 2001 From: cl-ment Date: Mon, 24 Aug 2026 12:48:34 +0200 Subject: [PATCH 02/11] fix(mcp): send courtesy -32601 reply asynchronously off read loop --- internal/mcp/client.go | 24 ++++++++++++-------- internal/mcp/client_test.go | 44 +++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 9 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index f7bae0bc6..e4bbdc774 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -389,15 +389,21 @@ func (client *Client) readLoop() { // It must never be routed as a response to a pending client request. if message.Method != "" { if message.ID != nil { - client.mu.Lock() - _ = client.writer.write(rpcMessage{ - ID: message.ID, - Error: &rpcError{ - Code: -32601, - Message: fmt.Sprintf("Method %q not supported", message.Method), - }, - }) - client.mu.Unlock() + // Send the courtesy -32601 reply asynchronously off the read loop so + // an undrained server stdin pipe never stalls readLoop or holds client.mu. + id := message.ID + method := message.Method + go func() { + client.mu.Lock() + defer client.mu.Unlock() + _ = client.writer.write(rpcMessage{ + ID: id, + Error: &rpcError{ + Code: -32601, + Message: fmt.Sprintf("Method %q not supported", method), + }, + }) + }() } continue } diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index a15a7b545..58ce25ad4 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -905,3 +905,47 @@ func TestStdioClientIgnoresServerInitiatedRequestsInPendingResponses(t *testing. } } +// TestStdioClientUndrainedServerDoesNotStallReadLoop proves that when a server +// stops draining its stdin, an asynchronous courtesy -32601 reply does not stall +// the read loop or block client.mu from servicing legitimate responses. +func TestStdioClientUndrainedServerDoesNotStallReadLoop(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + + client.ensureReader() + + // 1. Register a pending response for call ID 1 + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + // 2. Server sends request ID 2 (we DO NOT drain outReader so server stdin pipe is blocked) + serverReq := `{"jsonrpc":"2.0","id":2,"method":"roots/list","params":{}}` + "\n" + go func() { + _, _ = inWriter.Write([]byte(serverReq)) + // 3. Immediately send response for ID 1 + serverResp := `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}` + "\n" + _, _ = inWriter.Write([]byte(serverResp)) + }() + + // 4. Verify the response for ID 1 is delivered without being stalled by the undrained pipe + select { + case res := <-responses: + if res.message.Method != "" || len(res.message.Result) == 0 { + t.Fatalf("expected valid response for ID 1, got %#v", res.message) + } + case <-time.After(1 * time.Second): + t.Fatal("readLoop stalled on undrained error write; pending response was not delivered") + } +} From 512723a253da506792a248ebdba7d4ac3ec173e1 Mon Sep 17 00:00:00 2001 From: cl-ment Date: Mon, 24 Aug 2026 18:31:20 +0200 Subject: [PATCH 03/11] fix(mcp): restore jsonRPCIDEchoable and validate request IDs before responding --- internal/mcp/client.go | 23 +++++++++++++++++++++- internal/mcp/client_test.go | 38 +++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index e4bbdc774..20bd9a88a 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -388,7 +388,7 @@ func (client *Client) readLoop() { // A message with a Method is a server-initiated request or notification. // It must never be routed as a response to a pending client request. if message.Method != "" { - if message.ID != nil { + if message.ID != nil && jsonRPCIDEchoable(message.ID) { // Send the courtesy -32601 reply asynchronously off the read loop so // an undrained server stdin pipe never stalls readLoop or holds client.mu. id := message.ID @@ -493,6 +493,27 @@ func rpcIDMatches(value any, id int) bool { } } +// jsonRPCIDEchoable reports whether id is a valid JSON-RPC 2.0 identifier type +// (string, integer, or float with no fractional part) that is safe to echo back. +func jsonRPCIDEchoable(id any) bool { + if id == nil { + return false + } + switch v := id.(type) { + case string: + return true + case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64: + return true + case float64: + return v == float64(int64(v)) + case json.Number: + _, err := v.Int64() + return err == nil + default: + return false + } +} + func (client *Client) notify(method string, params any) error { rawParams, err := json.Marshal(params) if err != nil { diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 58ce25ad4..c53a6dc19 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -949,3 +949,41 @@ func TestStdioClientUndrainedServerDoesNotStallReadLoop(t *testing.T) { t.Fatal("readLoop stalled on undrained error write; pending response was not delivered") } } + +func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + + client.ensureReader() + + // Server sends request with boolean ID (invalid JSON-RPC id) + serverReq := `{"jsonrpc":"2.0","id":true,"method":"roots/list","params":{}}` + "\n" + go func() { + _, _ = inWriter.Write([]byte(serverReq)) + }() + + // Read should not produce any output response for invalid ID + buf := make([]byte, 1024) + readChan := make(chan int, 1) + go func() { + n, _ := outReader.Read(buf) + readChan <- n + }() + + select { + case n := <-readChan: + t.Fatalf("unexpected reply for invalid ID: %s", string(buf[:n])) + case <-time.After(100 * time.Millisecond): + // Expected: dropped cleanly + } +} From 80e71d1aacffd900c044925cc1fddada2f84789f Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Thu, 27 Aug 2026 15:12:20 +0200 Subject: [PATCH 04/11] fix(mcp): queue and drop courtesy replies without blocking client callers --- internal/mcp/client.go | 89 ++++++++++++++++++++++++++++--------- internal/mcp/client_test.go | 45 +++++++++++++++++++ 2 files changed, 114 insertions(+), 20 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 20bd9a88a..86d8ad0bb 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -49,13 +49,16 @@ type Client struct { stdin io.WriteCloser reader *messageReader writer *messageWriter - mu sync.Mutex closeMu sync.Mutex + idMu sync.Mutex nextID int cleanup func() + writeQueue chan writeOp + writerOnce sync.Once + // dispatchMu guards the response-dispatch state shared with the single - // reader goroutine. It is never held across a blocking read. + // reader goroutine. It is never held across a blocking read or write. dispatchMu sync.Mutex readerOnce sync.Once pending map[int]chan dispatchResult @@ -63,6 +66,13 @@ type Client struct { readDone bool } +type writeOp struct { + message rpcMessage + done chan error +} + +const writeQueueCapacity = 32 + // dispatchResult carries one matched JSON-RPC response (or a terminal reader // error) to a waiting caller. type dispatchResult struct { @@ -307,20 +317,20 @@ func (client *Client) request(ctx context.Context, method string, params any, ta return err } - // Allocate an id, register a response channel, and write the request while - // holding client.mu. The mutex serializes writes and id allocation but is - // released before the (potentially unbounded) wait for the response, so a - // hung server never holds the lock and blocks other callers/Close. - client.mu.Lock() + // Allocate an id and register a response channel. ID allocation and dispatch + // registrations are fast non-blocking operations. Message transmission is + // handled via writeMessage, ensuring a hung server never blocks other callers + // or prevents a caller with a deadline from giving up. + client.idMu.Lock() id := client.nextID client.nextID++ + client.idMu.Unlock() responses := make(chan dispatchResult, 1) client.dispatchMu.Lock() if client.readDone { readErr := client.readErr client.dispatchMu.Unlock() - client.mu.Unlock() if readErr != nil { return readErr } @@ -329,16 +339,14 @@ func (client *Client) request(ctx context.Context, method string, params any, ta client.pending[id] = responses client.dispatchMu.Unlock() - if err := client.writer.write(rpcMessage{ + if err := client.writeMessage(ctx, rpcMessage{ ID: id, Method: method, Params: rawParams, }); err != nil { client.removePending(id) - client.mu.Unlock() return err } - client.mu.Unlock() select { case <-ctx.Done(): @@ -361,6 +369,43 @@ func (client *Client) request(ctx context.Context, method string, params any, ta } } +// ensureWriter lazily starts the single writer goroutine. +func (client *Client) ensureWriter() { + client.writerOnce.Do(func() { + if client.writeQueue == nil { + client.writeQueue = make(chan writeOp, writeQueueCapacity) + } + go client.writeLoop() + }) +} + +func (client *Client) writeLoop() { + for op := range client.writeQueue { + err := client.writer.write(op.message) + if op.done != nil { + op.done <- err + } + } +} + +func (client *Client) writeMessage(ctx context.Context, message rpcMessage) error { + client.ensureWriter() + done := make(chan error, 1) + op := writeOp{message: message, done: done} + select { + case <-ctx.Done(): + return ctx.Err() + case client.writeQueue <- op: + } + + select { + case <-ctx.Done(): + return ctx.Err() + case err := <-done: + return err + } +} + // ensureReader lazily starts the single reader goroutine. It runs once per // client; subsequent calls are no-ops. func (client *Client) ensureReader() { @@ -389,21 +434,25 @@ func (client *Client) readLoop() { // It must never be routed as a response to a pending client request. if message.Method != "" { if message.ID != nil && jsonRPCIDEchoable(message.ID) { - // Send the courtesy -32601 reply asynchronously off the read loop so - // an undrained server stdin pipe never stalls readLoop or holds client.mu. + // Send the courtesy -32601 reply via the bounded writer queue. If the + // queue is saturated (e.g. an undrained server pipe), drop the reply + // immediately so it never stalls readLoop, holds a mutex, or blocks callers. + client.ensureWriter() id := message.ID method := message.Method - go func() { - client.mu.Lock() - defer client.mu.Unlock() - _ = client.writer.write(rpcMessage{ + courtesy := writeOp{ + message: rpcMessage{ ID: id, Error: &rpcError{ Code: -32601, Message: fmt.Sprintf("Method %q not supported", method), }, - }) - }() + }, + } + select { + case client.writeQueue <- courtesy: + default: + } } continue } @@ -519,7 +568,7 @@ func (client *Client) notify(method string, params any) error { if err != nil { return err } - return client.writer.write(rpcMessage{ + return client.writeMessage(context.Background(), rpcMessage{ Method: method, Params: rawParams, }) diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index c53a6dc19..b7d529e05 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -987,3 +987,48 @@ func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) { // Expected: dropped cleanly } } + +// TestStdioClientUndrainedServerDoesNotBlockCallerWithDeadline verifies that +// when the server's input pipe is completely blocked and courtesy replies are +// queued/dropped, a caller invoking request() with a deadline aborts cleanly +// when the context expires rather than hanging indefinitely on a write or mutex. +func TestStdioClientUndrainedServerDoesNotBlockCallerWithDeadline(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + + client.ensureReader() + + // Flood server requests to saturate write queue while outReader is NOT drained. + for i := 1; i <= 50; i++ { + serverReq := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"roots/list","params":{}}`+"\n", i+100) + _, err := inWriter.Write([]byte(serverReq)) + if err != nil { + t.Fatalf("failed to write server request: %v", err) + } + } + + // Caller with short deadline invokes request() + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + err := client.request(ctx, "tools/list", map[string]any{}, nil) + elapsed := time.Since(start) + + if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) { + t.Fatalf("expected deadline exceeded, got: %v", err) + } + if elapsed > 1*time.Second { + t.Fatalf("request took too long to abort on deadline: %v", elapsed) + } +} From 3a43d567d5363311e34e03c19ccf8f568e93da27 Mon Sep 17 00:00:00 2001 From: KRATOS <84986124+gnanam1990@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:43:23 +0530 Subject: [PATCH 05/11] fix(mcp): drop canceled queued requests --- internal/mcp/client.go | 17 +++- internal/mcp/client_test.go | 167 +++++++++++++++++++++++++++++------- internal/mcp/hang_test.go | 4 +- 3 files changed, 153 insertions(+), 35 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 86d8ad0bb..9c732a193 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -67,6 +67,7 @@ type Client struct { } type writeOp struct { + ctx context.Context message rpcMessage done chan error } @@ -319,8 +320,8 @@ func (client *Client) request(ctx context.Context, method string, params any, ta // Allocate an id and register a response channel. ID allocation and dispatch // registrations are fast non-blocking operations. Message transmission is - // handled via writeMessage, ensuring a hung server never blocks other callers - // or prevents a caller with a deadline from giving up. + // handled via writeMessage so a caller with a canceled context can stop + // waiting even when the peer is not draining its input. client.idMu.Lock() id := client.nextID client.nextID++ @@ -381,6 +382,16 @@ func (client *Client) ensureWriter() { func (client *Client) writeLoop() { for op := range client.writeQueue { + if op.ctx != nil { + select { + case <-op.ctx.Done(): + if op.done != nil { + op.done <- op.ctx.Err() + } + continue + default: + } + } err := client.writer.write(op.message) if op.done != nil { op.done <- err @@ -391,7 +402,7 @@ func (client *Client) writeLoop() { func (client *Client) writeMessage(ctx context.Context, message rpcMessage) error { client.ensureWriter() done := make(chan error, 1) - op := writeOp{message: message, done: done} + op := writeOp{ctx: ctx, message: message, done: done} select { case <-ctx.Done(): return ctx.Err() diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index b7d529e05..918cba457 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -1,6 +1,7 @@ package mcp import ( + "bytes" "context" "encoding/json" "errors" @@ -857,16 +858,6 @@ func TestStdioClientIgnoresServerInitiatedRequestsInPendingResponses(t *testing. _ = outReader.Close() }() - // Drain output responses from client - go func() { - buf := make([]byte, 1024) - for { - if _, err := outReader.Read(buf); err != nil { - return - } - } - }() - client.ensureReader() // Register a pending response for ID 1 @@ -877,23 +868,39 @@ func TestStdioClientIgnoresServerInitiatedRequestsInPendingResponses(t *testing. // Simulate server sending a request with ID 1 ("roots/list") serverReq := `{"jsonrpc":"2.0","id":1,"method":"roots/list","params":{}}` + "\n" + if _, err := inWriter.Write([]byte(serverReq)); err != nil { + t.Fatalf("write server request: %v", err) + } + + // The server request must take the separate request path: it receives a + // method-not-found response and cannot resolve the pending client call. + courtesyResult := make(chan dispatchResult, 1) go func() { - _, _ = inWriter.Write([]byte(serverReq)) + message, err := newMessageReader(outReader).read() + courtesyResult <- dispatchResult{message: message, err: err} }() - - // The pending channel for client ID 1 should NOT receive the server request. + select { + case result := <-courtesyResult: + if result.err != nil { + t.Fatalf("read courtesy response: %v", result.err) + } + if !rpcIDMatches(result.message.ID, 1) || result.message.Error == nil || result.message.Error.Code != -32601 { + t.Fatalf("courtesy response = %#v, want id 1 and error -32601", result.message) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for method-not-found response") + } select { case res := <-responses: t.Fatalf("pending request 1 received server request: %#v", res.message) - case <-time.After(100 * time.Millisecond): - // Expected: server request was not misdelivered to pending client request + default: } // Now send the actual response for ID 1 serverResp := `{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}` + "\n" - go func() { - _, _ = inWriter.Write([]byte(serverResp)) - }() + if _, err := inWriter.Write([]byte(serverResp)); err != nil { + t.Fatalf("write server response: %v", err) + } select { case res := <-responses: @@ -907,7 +914,7 @@ func TestStdioClientIgnoresServerInitiatedRequestsInPendingResponses(t *testing. // TestStdioClientUndrainedServerDoesNotStallReadLoop proves that when a server // stops draining its stdin, an asynchronous courtesy -32601 reply does not stall -// the read loop or block client.mu from servicing legitimate responses. +// the read loop from servicing legitimate responses. func TestStdioClientUndrainedServerDoesNotStallReadLoop(t *testing.T) { inReader, inWriter := io.Pipe() outReader, outWriter := io.Pipe() @@ -968,23 +975,123 @@ func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) { // Server sends request with boolean ID (invalid JSON-RPC id) serverReq := `{"jsonrpc":"2.0","id":true,"method":"roots/list","params":{}}` + "\n" + if _, err := inWriter.Write([]byte(serverReq)); err != nil { + t.Fatalf("write invalid server request: %v", err) + } + + // A following valid request acts as an ordering barrier: receiving its + // response proves readLoop already processed the invalid request. Because the + // writer is serial, any incorrect reply to the boolean ID would arrive first. + validReq := `{"jsonrpc":"2.0","id":7,"method":"roots/list","params":{}}` + "\n" + if _, err := inWriter.Write([]byte(validReq)); err != nil { + t.Fatalf("write valid server request: %v", err) + } + readResult := make(chan dispatchResult, 1) go func() { - _, _ = inWriter.Write([]byte(serverReq)) + message, err := newMessageReader(outReader).read() + readResult <- dispatchResult{message: message, err: err} }() - // Read should not produce any output response for invalid ID - buf := make([]byte, 1024) - readChan := make(chan int, 1) + select { + case result := <-readResult: + if result.err != nil { + t.Fatalf("read valid request response: %v", result.err) + } + if !rpcIDMatches(result.message.ID, 7) || result.message.Error == nil || result.message.Error.Code != -32601 { + t.Fatalf("response = %#v, want only valid id 7 and error -32601", result.message) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for valid request response") + } +} + +type gatedCaptureWriter struct { + started chan struct{} + release chan struct{} + startedOnce sync.Once + releaseOnce sync.Once + mu sync.Mutex + buffer bytes.Buffer +} + +func newGatedCaptureWriter() *gatedCaptureWriter { + return &gatedCaptureWriter{ + started: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (writer *gatedCaptureWriter) Write(p []byte) (int, error) { + writer.startedOnce.Do(func() { close(writer.started) }) + <-writer.release + writer.mu.Lock() + defer writer.mu.Unlock() + return writer.buffer.Write(p) +} + +func (writer *gatedCaptureWriter) Release() { + writer.releaseOnce.Do(func() { close(writer.release) }) +} + +func (writer *gatedCaptureWriter) Bytes() []byte { + writer.mu.Lock() + defer writer.mu.Unlock() + return append([]byte(nil), writer.buffer.Bytes()...) +} + +func TestStdioClientDoesNotWriteCanceledQueuedRequest(t *testing.T) { + reader := newBlockingReader() + defer reader.Close() + + output := newGatedCaptureWriter() + defer output.Release() + client := &Client{ + reader: newMessageReader(reader), + writer: newMessageWriter(output), + pending: make(map[int]chan dispatchResult), + } + + blockerDone := make(chan error, 1) go func() { - n, _ := outReader.Read(buf) - readChan <- n + blockerDone <- client.writeMessage(context.Background(), rpcMessage{Method: "notifications/blocker"}) }() - select { - case n := <-readChan: - t.Fatalf("unexpected reply for invalid ID: %s", string(buf[:n])) - case <-time.After(100 * time.Millisecond): - // Expected: dropped cleanly + case <-output.started: + case <-time.After(time.Second): + t.Fatal("initial write did not reach the blocked transport") + } + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + if err := client.request(ctx, "tools/call", map[string]any{"name": "side_effect"}, nil); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("request() error = %v, want context deadline exceeded", err) + } + + output.Release() + if err := <-blockerDone; err != nil { + t.Fatalf("release initial write: %v", err) + } + if err := client.writeMessage(context.Background(), rpcMessage{Method: "notifications/sentinel"}); err != nil { + t.Fatalf("write sentinel: %v", err) + } + + messages := newMessageReader(bytes.NewReader(output.Bytes())) + first, err := messages.read() + if err != nil { + t.Fatalf("read initial message: %v", err) + } + if first.Method != "notifications/blocker" { + t.Fatalf("first method = %q, want notifications/blocker", first.Method) + } + second, err := messages.read() + if err != nil { + t.Fatalf("read sentinel message: %v", err) + } + if second.Method != "notifications/sentinel" { + t.Fatalf("second method = %q, want notifications/sentinel", second.Method) + } + if extra, err := messages.read(); !errors.Is(err, io.EOF) { + t.Fatalf("unexpected queued message %#v, read error = %v", extra, err) } } diff --git a/internal/mcp/hang_test.go b/internal/mcp/hang_test.go index d62600745..7510cb4c8 100644 --- a/internal/mcp/hang_test.go +++ b/internal/mcp/hang_test.go @@ -82,7 +82,7 @@ func (reader *blockingReader) Close() error { // A hung stdio server must not block Client.request forever: a cancelled // per-call context must unblock the wait and surface ctx.Err(), and it must -// not hold client.mu (a second caller must still be able to proceed). +// not hold shared client state that prevents a second caller from proceeding. func TestClientRequestUnblocksOnContextCancel(t *testing.T) { reader := newBlockingReader() defer reader.Close() @@ -125,7 +125,7 @@ func TestClientRequestUnblocksOnContextCancel(t *testing.T) { t.Fatalf("second request() error = %v, want context.Canceled", err) } case <-time.After(3 * time.Second): - t.Fatal("second request() blocked — client.mu was held across the hung read") + t.Fatal("second request() blocked behind the hung read") } } From 6ed0f3f7fb1804efdf49b188ceba4e6663754432 Mon Sep 17 00:00:00 2001 From: KRATOS <84986124+gnanam1990@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:58:13 +0530 Subject: [PATCH 06/11] fix(mcp): preserve valid numeric request IDs --- internal/mcp/client.go | 11 ++++-- internal/mcp/client_test.go | 67 +++++++++++++++++++++++++++++++++++++ internal/mcp/hang_test.go | 33 ++++++++++++++---- 3 files changed, 102 insertions(+), 9 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 9c732a193..fd02fec6a 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "math" "os" "os/exec" "strconv" @@ -554,7 +555,7 @@ func rpcIDMatches(value any, id int) bool { } // jsonRPCIDEchoable reports whether id is a valid JSON-RPC 2.0 identifier type -// (string, integer, or float with no fractional part) that is safe to echo back. +// (string or finite number) that is safe to echo back. func jsonRPCIDEchoable(id any) bool { if id == nil { return false @@ -565,9 +566,13 @@ func jsonRPCIDEchoable(id any) bool { case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64: return true case float64: - return v == float64(int64(v)) + return !math.IsNaN(v) && !math.IsInf(v, 0) case json.Number: - _, err := v.Int64() + parsed, err := v.Float64() + if err != nil || math.IsNaN(parsed) || math.IsInf(parsed, 0) { + return false + } + _, err = json.Marshal(v) return err == nil default: return false diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 918cba457..5eedfd9b8 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "math" "net/http" "net/http/httptest" "os" @@ -1005,6 +1006,72 @@ func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) { } } +func TestStdioClientRepliesToValidNonIntegerServerRequestIDs(t *testing.T) { + requests := []struct { + wire string + want float64 + }{ + {wire: `{"jsonrpc":"2.0","id":1.5,"method":"roots/list","params":{}}` + "\n", want: 1.5}, + {wire: `{"jsonrpc":"2.0","id":2e2,"method":"roots/list","params":{}}` + "\n", want: 200}, + } + for _, request := range requests { + t.Run(fmt.Sprint(request.want), func(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + t.Cleanup(func() { + _ = inWriter.Close() + _ = outReader.Close() + }) + client.ensureReader() + if _, err := inWriter.Write([]byte(request.wire)); err != nil { + t.Fatalf("write server request: %v", err) + } + result := make(chan dispatchResult, 1) + go func() { + response, err := newMessageReader(outReader).read() + result <- dispatchResult{message: response, err: err} + }() + select { + case response := <-result: + if response.err != nil { + t.Fatalf("read method-not-found response: %v", response.err) + } + id, ok := response.message.ID.(float64) + if !ok || id != request.want || response.message.Error == nil || response.message.Error.Code != -32601 { + t.Fatalf("response = %#v, want id %v and error -32601", response.message, request.want) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for method-not-found response") + } + }) + } +} + +func TestJSONRPCIDEchoableAcceptsFiniteJSONNumbers(t *testing.T) { + for _, test := range []struct { + name string + id any + want bool + }{ + {name: "fractional float", id: 1.5, want: true}, + {name: "exponent json number", id: json.Number("2e2"), want: true}, + {name: "not a number", id: math.NaN(), want: false}, + {name: "positive infinity", id: math.Inf(1), want: false}, + {name: "invalid json number", id: json.Number("not-a-number"), want: false}, + } { + t.Run(test.name, func(t *testing.T) { + if got := jsonRPCIDEchoable(test.id); got != test.want { + t.Fatalf("jsonRPCIDEchoable(%v) = %v, want %v", test.id, got, test.want) + } + }) + } +} + type gatedCaptureWriter struct { started chan struct{} release chan struct{} diff --git a/internal/mcp/hang_test.go b/internal/mcp/hang_test.go index 7510cb4c8..bf8d39091 100644 --- a/internal/mcp/hang_test.go +++ b/internal/mcp/hang_test.go @@ -62,6 +62,15 @@ type blockingReader struct { release chan struct{} } +type signalingWriter struct { + writes chan struct{} +} + +func (writer *signalingWriter) Write(p []byte) (int, error) { + writer.writes <- struct{}{} + return len(p), nil +} + func newBlockingReader() *blockingReader { return &blockingReader{release: make(chan struct{})} } @@ -86,10 +95,11 @@ func (reader *blockingReader) Close() error { func TestClientRequestUnblocksOnContextCancel(t *testing.T) { reader := newBlockingReader() defer reader.Close() + writes := make(chan struct{}, 2) client := &Client{ reader: newMessageReader(reader), - writer: newMessageWriter(io.Discard), + writer: newMessageWriter(&signalingWriter{writes: writes}), nextID: 1, } @@ -99,6 +109,11 @@ func TestClientRequestUnblocksOnContextCancel(t *testing.T) { done <- client.request(ctx, "tools/list", map[string]any{}, nil) }() + select { + case <-writes: + case <-time.After(time.Second): + t.Fatal("first request did not reach the transport") + } // The request is now parked waiting for a response that never comes. cancel() @@ -111,15 +126,21 @@ func TestClientRequestUnblocksOnContextCancel(t *testing.T) { t.Fatal("request() hung on a non-responsive server") } - // The lock must be free: a second request under an already-cancelled - // context must return immediately rather than block. - cancelled, cancel2 := context.WithCancel(context.Background()) - cancel2() + // The shared state must be free: prove a second live request reaches the + // transport, then cancel it and verify cancellation releases the caller. + secondCtx, cancel2 := context.WithCancel(context.Background()) + defer cancel2() second := make(chan error, 1) go func() { - second <- client.request(cancelled, "tools/list", map[string]any{}, nil) + second <- client.request(secondCtx, "tools/list", map[string]any{}, nil) }() select { + case <-writes: + case <-time.After(time.Second): + t.Fatal("second request did not reach the transport") + } + cancel2() + select { case err := <-second: if !errors.Is(err, context.Canceled) { t.Fatalf("second request() error = %v, want context.Canceled", err) From bbb5142d718282cf80cca821a71a4964c97636a9 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sat, 29 Aug 2026 10:15:00 +0200 Subject: [PATCH 07/11] fix(mcp): treat JSON method member presence as the request discriminator An explicit empty method string is a request, not a response. Courtesy -32601 is still dropped when the write queue is full: failAll on that path stopped the reader and blocked inbound pipe writes (TestStdioClientUndrainedServerDoesNotBlockCallerWithDeadline). --- internal/mcp/client.go | 5 +---- internal/mcp/protocol.go | 30 ++++++++++++++++++++------ internal/mcp/protocol_presence_test.go | 26 ++++++++++++++++++++++ 3 files changed, 51 insertions(+), 10 deletions(-) create mode 100644 internal/mcp/protocol_presence_test.go diff --git a/internal/mcp/client.go b/internal/mcp/client.go index fd02fec6a..a3078dc90 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -444,11 +444,8 @@ func (client *Client) readLoop() { } // A message with a Method is a server-initiated request or notification. // It must never be routed as a response to a pending client request. - if message.Method != "" { + if message.methodPresent || message.Method != "" { if message.ID != nil && jsonRPCIDEchoable(message.ID) { - // Send the courtesy -32601 reply via the bounded writer queue. If the - // queue is saturated (e.g. an undrained server pipe), drop the reply - // immediately so it never stalls readLoop, holds a mutex, or blocks callers. client.ensureWriter() id := message.ID method := message.Method diff --git a/internal/mcp/protocol.go b/internal/mcp/protocol.go index fc900a60c..29541c870 100644 --- a/internal/mcp/protocol.go +++ b/internal/mcp/protocol.go @@ -17,12 +17,30 @@ import ( const maxMessageBytes = 64 * 1024 * 1024 type rpcMessage struct { - JSONRPC string `json:"jsonrpc,omitempty"` - ID any `json:"id,omitempty"` - Method string `json:"method,omitempty"` - Params json.RawMessage `json:"params,omitempty"` - Result json.RawMessage `json:"result,omitempty"` - Error *rpcError `json:"error,omitempty"` + JSONRPC string `json:"jsonrpc,omitempty"` + ID any `json:"id,omitempty"` + Method string `json:"method,omitempty"` + Params json.RawMessage `json:"params,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *rpcError `json:"error,omitempty"` + methodPresent bool `json:"-"` +} + +func (m *rpcMessage) UnmarshalJSON(data []byte) error { + var probe struct { + Method *string `json:"method"` + } + if err := json.Unmarshal(data, &probe); err != nil { + return err + } + type wire rpcMessage + var w wire + if err := json.Unmarshal(data, &w); err != nil { + return err + } + *m = rpcMessage(w) + m.methodPresent = probe.Method != nil + return nil } type rpcError struct { diff --git a/internal/mcp/protocol_presence_test.go b/internal/mcp/protocol_presence_test.go new file mode 100644 index 000000000..1a8c02687 --- /dev/null +++ b/internal/mcp/protocol_presence_test.go @@ -0,0 +1,26 @@ +package mcp + +import ( + "encoding/json" + "testing" +) + +func TestRPCMessageMethodPresence(t *testing.T) { + var req rpcMessage + if err := json.Unmarshal([]byte(`{"jsonrpc":"2.0","id":1,"method":""}`), &req); err != nil { + t.Fatal(err) + } + if !req.methodPresent { + t.Fatal(`{"method":""} must set methodPresent so it is not routed as a response`) + } + var resp rpcMessage + if err := json.Unmarshal([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`), &resp); err != nil { + t.Fatal(err) + } + if resp.methodPresent { + t.Fatal("response without method member must not set methodPresent") + } + if resp.Method != "" { + t.Fatalf("response Method = %q", resp.Method) + } +} From cc5ca6dbf17fb3a69aeee678ccdb3f0c15f50e28 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sat, 29 Aug 2026 10:21:14 +0200 Subject: [PATCH 08/11] fix(mcp): count dropped courtesy replies when the write queue is full The inbound reader must not stop. An empty select default is now an atomic counter rather than a silent discard. --- internal/mcp/client.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index a3078dc90..e99139307 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -13,6 +13,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/Gitlawb/zero/internal/execution" @@ -55,8 +56,9 @@ type Client struct { nextID int cleanup func() - writeQueue chan writeOp - writerOnce sync.Once + writeQueue chan writeOp + writerOnce sync.Once + droppedCourtesy atomic.Uint64 // dispatchMu guards the response-dispatch state shared with the single // reader goroutine. It is never held across a blocking read or write. @@ -461,6 +463,7 @@ func (client *Client) readLoop() { select { case client.writeQueue <- courtesy: default: + client.droppedCourtesy.Add(1) } } continue From f0b6b521aec6c6ea1ae95e5d885a251ab905bea0 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 17:01:52 +0200 Subject: [PATCH 09/11] fix(mcp): retain method presence, writer lifecycle, and original JSON-RPC IDs --- internal/mcp/client.go | 184 ++++++++++-- internal/mcp/client_test.go | 370 ++++++++++++++++++++++++- internal/mcp/protocol.go | 32 ++- internal/mcp/protocol_presence_test.go | 7 + 4 files changed, 549 insertions(+), 44 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index e99139307..92b536398 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -13,7 +13,6 @@ import ( "strconv" "strings" "sync" - "sync/atomic" "time" "github.com/Gitlawb/zero/internal/execution" @@ -56,9 +55,13 @@ type Client struct { nextID int cleanup func() - writeQueue chan writeOp - writerOnce sync.Once - droppedCourtesy atomic.Uint64 + writeMu sync.Mutex + writeQueue chan writeOp + writeClosed bool + writeSenders sync.WaitGroup + writerStop chan struct{} + writerDone chan struct{} + courtesyOverflow []writeOp // dispatchMu guards the response-dispatch state shared with the single // reader goroutine. It is never held across a blocking read or write. @@ -77,6 +80,8 @@ type writeOp struct { const writeQueueCapacity = 32 +var errMCPClientClosed = errors.New("MCP client closed") + // dispatchResult carries one matched JSON-RPC response (or a terminal reader // error) to a waiting caller. type dispatchResult struct { @@ -266,7 +271,8 @@ func (client *Client) Close() error { // Fail any callers still waiting on a response. The blocking read in the // reader goroutine is released below when stdin closes and the process // exits (or is killed), EOFing stdout. - client.failAll(errors.New("MCP client closed")) + client.failAll(errMCPClientClosed) + client.beginWriterShutdown() var err error stdin := client.stdin @@ -306,6 +312,7 @@ func (client *Client) Close() error { client.cleanup() client.cleanup = nil } + client.finishWriterShutdown() return err } @@ -373,18 +380,81 @@ func (client *Client) request(ctx context.Context, method string, params any, ta } } -// ensureWriter lazily starts the single writer goroutine. func (client *Client) ensureWriter() { - client.writerOnce.Do(func() { - if client.writeQueue == nil { - client.writeQueue = make(chan writeOp, writeQueueCapacity) + _ = client.startWriter() +} + +func (client *Client) startWriter() error { + client.writeMu.Lock() + defer client.writeMu.Unlock() + if client.writeClosed { + return errMCPClientClosed + } + if client.writeQueue != nil { + return nil + } + client.writeQueue = make(chan writeOp, writeQueueCapacity) + client.writerStop = make(chan struct{}) + client.writerDone = make(chan struct{}) + go client.writeLoop() + return nil +} + +func (client *Client) writerStopped() bool { + if client.writerStop == nil { + return false + } + select { + case <-client.writerStop: + return true + default: + return false + } +} + +func (client *Client) beginWriterShutdown() { + client.writeMu.Lock() + defer client.writeMu.Unlock() + if client.writeClosed { + return + } + client.writeClosed = true + if client.writerStop != nil { + close(client.writerStop) + } + for _, op := range client.courtesyOverflow { + if op.done != nil { + op.done <- errMCPClientClosed } - go client.writeLoop() - }) + } + client.courtesyOverflow = nil +} + +func (client *Client) finishWriterShutdown() { + client.writeMu.Lock() + queue := client.writeQueue + done := client.writerDone + client.writeQueue = nil + client.writeMu.Unlock() + if queue == nil { + return + } + client.writeSenders.Wait() + close(queue) + if done != nil { + <-done + } } func (client *Client) writeLoop() { + defer close(client.writerDone) for op := range client.writeQueue { + if client.writerStopped() { + if op.done != nil { + op.done <- errMCPClientClosed + } + continue + } if op.ctx != nil { select { case <-op.ctx.Done(): @@ -399,22 +469,88 @@ func (client *Client) writeLoop() { if op.done != nil { op.done <- err } + client.drainCourtesyOverflow() + } +} + +func (client *Client) drainCourtesyOverflow() { + client.writeMu.Lock() + defer client.writeMu.Unlock() + if client.writeClosed || client.writeQueue == nil { + client.courtesyOverflow = nil + return + } + for len(client.courtesyOverflow) > 0 { + select { + case client.writeQueue <- client.courtesyOverflow[0]: + client.courtesyOverflow = client.courtesyOverflow[1:] + default: + return + } + } +} + +func (client *Client) enqueueCourtesy(message rpcMessage) { + if err := client.startWriter(); err != nil { + return + } + op := writeOp{message: message} + client.writeMu.Lock() + defer client.writeMu.Unlock() + if client.writeClosed || client.writeQueue == nil { + return + } + select { + case client.writeQueue <- op: + default: + client.courtesyOverflow = append(client.courtesyOverflow, op) } } func (client *Client) writeMessage(ctx context.Context, message rpcMessage) error { - client.ensureWriter() + if err := ctx.Err(); err != nil { + return err + } + if err := client.startWriter(); err != nil { + return err + } done := make(chan error, 1) op := writeOp{ctx: ctx, message: message, done: done} + + client.writeMu.Lock() + if client.writeClosed { + client.writeMu.Unlock() + return errMCPClientClosed + } + stop := client.writerStop + queue := client.writeQueue + client.writeSenders.Add(1) + client.writeMu.Unlock() + select { case <-ctx.Done(): + client.writeSenders.Done() return ctx.Err() - case client.writeQueue <- op: + case <-stop: + client.writeSenders.Done() + return errMCPClientClosed + case queue <- op: + client.writeSenders.Done() } select { case <-ctx.Done(): return ctx.Err() + case <-stop: + select { + case err := <-done: + if err != nil { + return err + } + return errMCPClientClosed + case <-ctx.Done(): + return ctx.Err() + } case err := <-done: return err } @@ -448,23 +584,13 @@ func (client *Client) readLoop() { // It must never be routed as a response to a pending client request. if message.methodPresent || message.Method != "" { if message.ID != nil && jsonRPCIDEchoable(message.ID) { - client.ensureWriter() - id := message.ID - method := message.Method - courtesy := writeOp{ - message: rpcMessage{ - ID: id, - Error: &rpcError{ - Code: -32601, - Message: fmt.Sprintf("Method %q not supported", method), - }, + client.enqueueCourtesy(rpcMessage{ + ID: message.ID, + Error: &rpcError{ + Code: -32601, + Message: fmt.Sprintf("Method %q not supported", message.Method), }, - } - select { - case client.writeQueue <- courtesy: - default: - client.droppedCourtesy.Add(1) - } + }) } continue } diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 5eedfd9b8..969041be9 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -1,6 +1,7 @@ package mcp import ( + "bufio" "bytes" "context" "encoding/json" @@ -12,6 +13,7 @@ import ( "net/http/httptest" "os" "os/exec" + "runtime" "strings" "sync" "sync/atomic" @@ -1008,14 +1010,15 @@ func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) { func TestStdioClientRepliesToValidNonIntegerServerRequestIDs(t *testing.T) { requests := []struct { + name string wire string - want float64 + want string }{ - {wire: `{"jsonrpc":"2.0","id":1.5,"method":"roots/list","params":{}}` + "\n", want: 1.5}, - {wire: `{"jsonrpc":"2.0","id":2e2,"method":"roots/list","params":{}}` + "\n", want: 200}, + {name: "fractional", wire: `{"jsonrpc":"2.0","id":1.5,"method":"roots/list","params":{}}` + "\n", want: "1.5"}, + {name: "exponent", wire: `{"jsonrpc":"2.0","id":2e2,"method":"roots/list","params":{}}` + "\n", want: "2e2"}, } for _, request := range requests { - t.Run(fmt.Sprint(request.want), func(t *testing.T) { + t.Run(request.name, func(t *testing.T) { inReader, inWriter := io.Pipe() outReader, outWriter := io.Pipe() client := &Client{ @@ -1041,9 +1044,12 @@ func TestStdioClientRepliesToValidNonIntegerServerRequestIDs(t *testing.T) { if response.err != nil { t.Fatalf("read method-not-found response: %v", response.err) } - id, ok := response.message.ID.(float64) - if !ok || id != request.want || response.message.Error == nil || response.message.Error.Code != -32601 { - t.Fatalf("response = %#v, want id %v and error -32601", response.message, request.want) + got, err := json.Marshal(response.message.ID) + if err != nil { + t.Fatalf("marshal id: %v", err) + } + if string(got) != request.want || response.message.Error == nil || response.message.Error.Code != -32601 { + t.Fatalf("response = %#v, want id %s and error -32601", response.message, request.want) } case <-time.After(time.Second): t.Fatal("timed out waiting for method-not-found response") @@ -1206,3 +1212,353 @@ func TestStdioClientUndrainedServerDoesNotBlockCallerWithDeadline(t *testing.T) t.Fatalf("request took too long to abort on deadline: %v", elapsed) } } + +func TestStdioClientEmptyMethodDoesNotCompletePending(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + client.ensureReader() + + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":""}` + "\n")); err != nil { + t.Fatalf("write empty-method frame: %v", err) + } + + courtesyResult := make(chan dispatchResult, 1) + go func() { + message, err := newMessageReader(outReader).read() + courtesyResult <- dispatchResult{message: message, err: err} + }() + select { + case result := <-courtesyResult: + if result.err != nil { + t.Fatalf("read courtesy response: %v", result.err) + } + if !rpcIDMatches(result.message.ID, 1) || result.message.Error == nil || result.message.Error.Code != -32601 { + t.Fatalf("courtesy response = %#v, want id 1 and error -32601", result.message) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for method-not-found response") + } + select { + case res := <-responses: + t.Fatalf("pending request 1 completed by empty-method frame: %#v", res.message) + default: + } + + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}` + "\n")); err != nil { + t.Fatalf("write server response: %v", err) + } + select { + case res := <-responses: + if res.message.Method != "" || len(res.message.Result) == 0 { + t.Fatalf("expected valid response, got %#v", res.message) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("timed out waiting for actual response") + } +} + +func TestStdioClientNullMethodDoesNotCompletePending(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + client.ensureReader() + + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":null}` + "\n")); err != nil { + t.Fatalf("write null-method frame: %v", err) + } + + courtesyResult := make(chan dispatchResult, 1) + go func() { + message, err := newMessageReader(outReader).read() + courtesyResult <- dispatchResult{message: message, err: err} + }() + select { + case result := <-courtesyResult: + if result.err != nil { + t.Fatalf("read courtesy response: %v", result.err) + } + if !rpcIDMatches(result.message.ID, 1) || result.message.Error == nil || result.message.Error.Code != -32601 { + t.Fatalf("courtesy response = %#v, want id 1 and error -32601", result.message) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for method-not-found response") + } + select { + case res := <-responses: + t.Fatalf("pending request 1 completed by null-method frame: %#v", res.message) + default: + } + + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}` + "\n")); err != nil { + t.Fatalf("write server response: %v", err) + } + select { + case res := <-responses: + if res.message.Method != "" || len(res.message.Result) == 0 { + t.Fatalf("expected valid response, got %#v", res.message) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("timed out waiting for actual response") + } +} + +func TestStdioClientCourtesyReplySurvivesFullWriteQueue(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + client.ensureReader() + + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + const extra = 8 + n := writeQueueCapacity + extra + for i := 0; i < n; i++ { + frame := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"roots/list","params":{}}`+"\n", i+100) + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatalf("write server request %d: %v", i, err) + } + } + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"ok":true}}` + "\n")); err != nil { + t.Fatalf("write pending response: %v", err) + } + select { + case res := <-responses: + if res.err != nil || len(res.message.Result) == 0 { + t.Fatalf("pending response = %#v err=%v", res.message, res.err) + } + case <-time.After(2 * time.Second): + t.Fatal("readLoop stalled while the write queue was full") + } + + drainDone := make(chan []rpcMessage, 1) + go func() { + reader := newMessageReader(outReader) + var got []rpcMessage + for len(got) < n { + message, err := reader.read() + if err != nil { + drainDone <- got + return + } + got = append(got, message) + } + drainDone <- got + }() + select { + case got := <-drainDone: + if len(got) != n { + t.Fatalf("courtesy replies = %d, want %d (queue-full requests dropped)", len(got), n) + } + seen := make(map[int]bool, n) + for _, message := range got { + id, ok := rpcMessageID(message.ID) + if !ok || message.Error == nil || message.Error.Code != -32601 { + t.Fatalf("courtesy reply = %#v, want -32601", message) + } + seen[id] = true + } + for i := 0; i < n; i++ { + if !seen[i+100] { + t.Fatalf("missing courtesy reply for id %d after output resumed", i+100) + } + } + case <-time.After(3 * time.Second): + t.Fatal("timed out draining courtesy replies after output resumed") + } +} + +func TestWriterLoopExitsOnRepeatedClose(t *testing.T) { + before := runtime.NumGoroutine() + for i := 0; i < 25; i++ { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + copied := make(chan struct{}) + go func() { + _, _ = io.Copy(io.Discard, outReader) + close(copied) + }() + client := &Client{ + stdin: outWriter, + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + } + client.ensureWriter() + if err := client.writeMessage(context.Background(), rpcMessage{Method: "notifications/ping"}); err != nil { + t.Fatalf("writeMessage: %v", err) + } + if err := client.Close(); err != nil && !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("Close: %v", err) + } + _ = inWriter.Close() + <-copied + } + runtime.GC() + time.Sleep(20 * time.Millisecond) + after := runtime.NumGoroutine() + if after > before+8 { + t.Fatalf("goroutines leaked: before=%d after=%d", before, after) + } +} + +func TestWriteMessageReleasedOnClose(t *testing.T) { + reader := newBlockingReader() + defer reader.Close() + output := newGatedCaptureWriter() + client := &Client{ + reader: newMessageReader(reader), + writer: newMessageWriter(output), + } + + blockerDone := make(chan error, 1) + go func() { + blockerDone <- client.writeMessage(context.Background(), rpcMessage{Method: "notifications/blocker"}) + }() + select { + case <-output.started: + case <-time.After(time.Second): + t.Fatal("initial write did not reach the transport") + } + + queuedDone := make(chan error, 1) + go func() { + queuedDone <- client.writeMessage(context.Background(), rpcMessage{Method: "notifications/queued"}) + }() + deadline := time.Now().Add(time.Second) + for len(client.writeQueue) == 0 { + if time.Now().After(deadline) { + t.Fatal("queued write did not enqueue") + } + time.Sleep(time.Millisecond) + } + + closeDone := make(chan error, 1) + go func() { + closeDone <- client.Close() + }() + select { + case <-client.writerStop: + case <-time.After(time.Second): + t.Fatal("Close did not begin writer shutdown") + } + output.Release() + select { + case err := <-closeDone: + if err != nil && !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("Close: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Close hung waiting for the writer worker") + } + + select { + case err := <-queuedDone: + if err == nil { + t.Fatal("queued write succeeded after Close") + } + case <-time.After(2 * time.Second): + t.Fatal("queued write was not released on Close") + } + select { + case <-blockerDone: + case <-time.After(2 * time.Second): + t.Fatal("blocked write was not released on Close") + } + + messages := newMessageReader(bytes.NewReader(output.Bytes())) + for { + message, err := messages.read() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("read captured output: %v", err) + } + if message.Method == "notifications/queued" { + t.Fatal("queued request was written after Close") + } + } +} + +func TestStdioClientPreservesLargeNumericRequestID(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + client.ensureReader() + + const rawID = "9007199254740993" + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":` + rawID + `,"method":"roots/list","params":{}}` + "\n")); err != nil { + t.Fatalf("write server request: %v", err) + } + + lineDone := make(chan struct { + line string + err error + }, 1) + go func() { + line, err := bufio.NewReader(outReader).ReadString('\n') + lineDone <- struct { + line string + err error + }{line: line, err: err} + }() + select { + case got := <-lineDone: + if got.err != nil { + t.Fatalf("read courtesy response: %v", got.err) + } + if !strings.Contains(got.line, rawID) { + t.Fatalf("serialized courtesy response %q does not preserve id %s", got.line, rawID) + } + if !strings.Contains(got.line, `"-32601"`) && !strings.Contains(got.line, `-32601`) { + t.Fatalf("serialized courtesy response %q missing -32601", got.line) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for courtesy response") + } +} diff --git a/internal/mcp/protocol.go b/internal/mcp/protocol.go index 29541c870..061561876 100644 --- a/internal/mcp/protocol.go +++ b/internal/mcp/protocol.go @@ -2,6 +2,7 @@ package mcp import ( "bufio" + "bytes" "encoding/json" "errors" "fmt" @@ -27,19 +28,34 @@ type rpcMessage struct { } func (m *rpcMessage) UnmarshalJSON(data []byte) error { - var probe struct { - Method *string `json:"method"` - } + var probe map[string]json.RawMessage if err := json.Unmarshal(data, &probe); err != nil { return err } - type wire rpcMessage - var w wire - if err := json.Unmarshal(data, &w); err != nil { + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + var wire struct { + JSONRPC string `json:"jsonrpc"` + ID any `json:"id"` + Params json.RawMessage `json:"params"` + Result json.RawMessage `json:"result"` + Error *rpcError `json:"error"` + } + if err := decoder.Decode(&wire); err != nil { return err } - *m = rpcMessage(w) - m.methodPresent = probe.Method != nil + m.JSONRPC = wire.JSONRPC + m.ID = wire.ID + m.Params = wire.Params + m.Result = wire.Result + m.Error = wire.Error + if raw, ok := probe["method"]; ok { + m.methodPresent = true + var method string + if err := json.Unmarshal(raw, &method); err == nil { + m.Method = method + } + } return nil } diff --git a/internal/mcp/protocol_presence_test.go b/internal/mcp/protocol_presence_test.go index 1a8c02687..d2566b503 100644 --- a/internal/mcp/protocol_presence_test.go +++ b/internal/mcp/protocol_presence_test.go @@ -23,4 +23,11 @@ func TestRPCMessageMethodPresence(t *testing.T) { if resp.Method != "" { t.Fatalf("response Method = %q", resp.Method) } + var nullMethod rpcMessage + if err := json.Unmarshal([]byte(`{"jsonrpc":"2.0","id":1,"method":null}`), &nullMethod); err != nil { + t.Fatal(err) + } + if !nullMethod.methodPresent { + t.Fatal(`{"method":null} must set methodPresent so it is not routed as a response`) + } } From caa4cfc67f495cd05363557bb9bcc7614e3e961a Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 17:58:08 +0200 Subject: [PATCH 10/11] fix(mcp): bound courtesy overflow, pin writer queue, match exponent IDs --- internal/mcp/client.go | 54 ++++++++++------- internal/mcp/client_test.go | 118 ++++++++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 21 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 92b536398..8c0f11982 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -79,6 +79,7 @@ type writeOp struct { } const writeQueueCapacity = 32 +const courtesyOverflowCap = writeQueueCapacity var errMCPClientClosed = errors.New("MCP client closed") @@ -393,10 +394,11 @@ func (client *Client) startWriter() error { if client.writeQueue != nil { return nil } - client.writeQueue = make(chan writeOp, writeQueueCapacity) + queue := make(chan writeOp, writeQueueCapacity) + client.writeQueue = queue client.writerStop = make(chan struct{}) client.writerDone = make(chan struct{}) - go client.writeLoop() + go client.writeLoop(queue) return nil } @@ -446,9 +448,12 @@ func (client *Client) finishWriterShutdown() { } } -func (client *Client) writeLoop() { +func (client *Client) writeLoop(queue <-chan writeOp) { defer close(client.writerDone) - for op := range client.writeQueue { + if queue == nil { + return + } + for op := range queue { if client.writerStopped() { if op.done != nil { op.done <- errMCPClientClosed @@ -503,6 +508,9 @@ func (client *Client) enqueueCourtesy(message rpcMessage) { select { case client.writeQueue <- op: default: + if len(client.courtesyOverflow) >= courtesyOverflowCap { + return + } client.courtesyOverflow = append(client.courtesyOverflow, op) } } @@ -637,6 +645,20 @@ func (client *Client) failAll(err error) { // rpcMessageID extracts the integer id from a JSON-RPC id value across the // numeric/string encodings a server may use. +func jsonNumberAsInt(n json.Number) (int64, bool) { + if parsed, err := n.Int64(); err == nil { + return parsed, true + } + f, err := n.Float64() + if err != nil || math.IsNaN(f) || math.IsInf(f, 0) || f != math.Trunc(f) { + return 0, false + } + if f > float64(math.MaxInt64) || f < float64(math.MinInt64) { + return 0, false + } + return int64(f), true +} + func rpcMessageID(value any) (int, bool) { switch typed := value.(type) { case int: @@ -644,10 +666,13 @@ func rpcMessageID(value any) (int, bool) { case int64: return int(typed), true case float64: + if math.IsNaN(typed) || math.IsInf(typed, 0) || typed != math.Trunc(typed) { + return 0, false + } return int(typed), true case json.Number: - parsed, err := typed.Int64() - if err != nil { + parsed, ok := jsonNumberAsInt(typed) + if !ok { return 0, false } return int(parsed), true @@ -663,21 +688,8 @@ func rpcMessageID(value any) (int, bool) { } func rpcIDMatches(value any, id int) bool { - switch typed := value.(type) { - case int: - return typed == id - case int64: - return typed == int64(id) - case float64: - return typed == float64(id) - case json.Number: - parsed, err := typed.Int64() - return err == nil && parsed == int64(id) - case string: - return typed == strconv.Itoa(id) - default: - return false - } + got, ok := rpcMessageID(value) + return ok && got == id } // jsonRPCIDEchoable reports whether id is a valid JSON-RPC 2.0 identifier type diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 969041be9..938344f7f 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -1562,3 +1562,121 @@ func TestStdioClientPreservesLargeNumericRequestID(t *testing.T) { t.Fatal("timed out waiting for courtesy response") } } + +func TestRPCMessageIDAcceptsExponentForm(t *testing.T) { + got, ok := rpcMessageID(json.Number("1e0")) + if !ok || got != 1 { + t.Fatalf("rpcMessageID(1e0) = %d, %v, want 1, true", got, ok) + } + got, ok = rpcMessageID(json.Number("2e2")) + if !ok || got != 200 { + t.Fatalf("rpcMessageID(2e2) = %d, %v, want 200, true", got, ok) + } + if _, ok := rpcMessageID(json.Number("1.5")); ok { + t.Fatal("fractional json.Number should not match") + } + if rpcIDMatches(json.Number("1e0"), 1) != true { + t.Fatal("rpcIDMatches(1e0, 1) = false") + } +} + +func TestStdioClientMatchesExponentFormResponseID(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + client.ensureReader() + + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1e0,"result":{"ok":true}}` + "\n")); err != nil { + t.Fatalf("write exponent-id response: %v", err) + } + select { + case res := <-responses: + if res.err != nil || len(res.message.Result) == 0 { + t.Fatalf("pending 1 not completed by id 1e0: %#v err=%v", res.message, res.err) + } + case <-time.After(time.Second): + t.Fatal("response id 1e0 did not resolve pending request 1") + } +} + +func TestStdioClientCourtesyOverflowIsBounded(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + client := &Client{ + reader: newMessageReader(inReader), + writer: newMessageWriter(outWriter), + pending: make(map[int]chan dispatchResult), + } + defer func() { + _ = inWriter.Close() + _ = outReader.Close() + }() + client.ensureReader() + + responses := make(chan dispatchResult, 1) + client.dispatchMu.Lock() + client.pending[1] = responses + client.dispatchMu.Unlock() + + flood := writeQueueCapacity + courtesyOverflowCap + 64 + for i := 0; i < flood; i++ { + frame := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"roots/list","params":{}}`+"\n", i+100) + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatalf("write server request %d: %v", i, err) + } + } + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"ok":true}}` + "\n")); err != nil { + t.Fatalf("write pending response: %v", err) + } + select { + case res := <-responses: + if res.err != nil { + t.Fatalf("pending stalled: %v", res.err) + } + case <-time.After(2 * time.Second): + t.Fatal("readLoop stalled during courtesy flood") + } + + client.writeMu.Lock() + n := len(client.courtesyOverflow) + client.writeMu.Unlock() + if n > courtesyOverflowCap { + t.Fatalf("courtesyOverflow = %d, want <= %d", n, courtesyOverflowCap) + } +} + +func TestWriterLoopCloseBeforeScheduleDoesNotHang(t *testing.T) { + for i := 0; i < 50; i++ { + client := &Client{ + writer: newMessageWriter(io.Discard), + } + if err := client.startWriter(); err != nil { + t.Fatalf("startWriter: %v", err) + } + done := make(chan error, 1) + go func() { + done <- client.Close() + }() + select { + case err := <-done: + if err != nil && !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("Close: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Close hung: writeLoop ranged over a nil queue") + } + } +} From 984f4b2a143d42be0db9a2fc7622adcea9ffe4ec Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 20:11:37 +0200 Subject: [PATCH 11/11] fix(mcp): skip SSE frames that carry a method member Treat method presence the same on stdio and SSE so empty or null method never completes a pending client request. --- internal/mcp/client.go | 2 +- internal/mcp/network_client.go | 5 ++- internal/mcp/network_client_test.go | 57 +++++++++++++++++++++++++++++ internal/mcp/protocol.go | 4 ++ 4 files changed, 66 insertions(+), 2 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 8c0f11982..932b4c371 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -590,7 +590,7 @@ func (client *Client) readLoop() { } // A message with a Method is a server-initiated request or notification. // It must never be routed as a response to a pending client request. - if message.methodPresent || message.Method != "" { + if message.isRequestOrNotification() { if message.ID != nil && jsonRPCIDEchoable(message.ID) { client.enqueueCourtesy(rpcMessage{ ID: message.ID, diff --git a/internal/mcp/network_client.go b/internal/mcp/network_client.go index b422b3c28..b8b829af2 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -539,6 +539,9 @@ func (client *remoteSSEClient) deliverEventMessage(value string) error { if err := decoder.Decode(&message); err != nil { return fmt.Errorf("decode MCP SSE stream message: %w", err) } + if message.isRequestOrNotification() { + return nil + } key := rpcResponseKey(message.ID) if key == "" { return nil @@ -633,7 +636,7 @@ func decodeSSERPCMessage(reader io.Reader) (rpcMessage, error) { // those — the response has no method — and keep scanning. Previously the // first message event was returned unconditionally, so a leading // notification surfaced to the caller as an id mismatch and failed the call. - if candidate.Method != "" { + if candidate.isRequestOrNotification() { return true } decoded = candidate diff --git a/internal/mcp/network_client_test.go b/internal/mcp/network_client_test.go index fba92f1ec..e4b406447 100644 --- a/internal/mcp/network_client_test.go +++ b/internal/mcp/network_client_test.go @@ -339,3 +339,60 @@ func TestDecodeSSERPCMessageSkipsNotifications(t *testing.T) { t.Fatalf("expected a result payload, got %#v", msg) } } + +func TestDecodeSSERPCMessageSkipsEmptyAndNullMethod(t *testing.T) { + stream := "event: message\n" + + `data: {"jsonrpc":"2.0","id":7,"method":""}` + "\n\n" + + "event: message\n" + + `data: {"jsonrpc":"2.0","id":7,"method":null}` + "\n\n" + + "event: message\n" + + `data: {"jsonrpc":"2.0","id":7,"result":{"ok":true}}` + "\n\n" + + msg, err := decodeSSERPCMessage(strings.NewReader(stream)) + if err != nil { + t.Fatalf("decodeSSERPCMessage: %v", err) + } + if msg.isRequestOrNotification() { + t.Fatalf("empty/null method must not be treated as the response, got method %q present=%v", msg.Method, msg.methodPresent) + } + if !rpcIDMatches(msg.ID, 7) { + t.Fatalf("expected response id 7, got %#v", msg.ID) + } +} + +func TestDeliverEventMessageSkipsMethodPresence(t *testing.T) { + client := &remoteSSEClient{pending: map[string]chan ssePendingResponse{}} + key := rpcResponseKey(1) + pending := make(chan ssePendingResponse, 1) + client.pending[key] = pending + + if err := client.deliverEventMessage(`{"jsonrpc":"2.0","id":1,"method":""}`); err != nil { + t.Fatalf("deliverEventMessage empty method: %v", err) + } + if err := client.deliverEventMessage(`{"jsonrpc":"2.0","id":1,"method":null}`); err != nil { + t.Fatalf("deliverEventMessage null method: %v", err) + } + select { + case got := <-pending: + t.Fatalf("method presence must not complete pending, got %#v", got) + default: + } + if _, ok := client.pending[key]; !ok { + t.Fatal("deliverEventMessage must not delete pending for a request/notification") + } + + if err := client.deliverEventMessage(`{"jsonrpc":"2.0","id":1,"result":{"ok":true}}`); err != nil { + t.Fatalf("deliverEventMessage response: %v", err) + } + select { + case got := <-pending: + if got.err != nil { + t.Fatalf("true response: %v", got.err) + } + if got.message.isRequestOrNotification() { + t.Fatal("true response must not carry a method") + } + default: + t.Fatal("true response must complete pending") + } +} diff --git a/internal/mcp/protocol.go b/internal/mcp/protocol.go index 061561876..0f388e20c 100644 --- a/internal/mcp/protocol.go +++ b/internal/mcp/protocol.go @@ -27,6 +27,10 @@ type rpcMessage struct { methodPresent bool `json:"-"` } +func (m rpcMessage) isRequestOrNotification() bool { + return m.methodPresent || m.Method != "" +} + func (m *rpcMessage) UnmarshalJSON(data []byte) error { var probe map[string]json.RawMessage if err := json.Unmarshal(data, &probe); err != nil {