From f4a9396942843dbbe20cb3d4d1de962e720ce887 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Tue, 5 Aug 2025 14:50:21 +0000 Subject: [PATCH] mcp: fix a bug causing premature termination of streams As reported in #227, there was a race where streams could be terminated before the final reply was written. Fix this by checking the number of outstanding requests in the same critical section that acquires outgoing messages to write. Fixes #227 --- mcp/streamable.go | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/mcp/streamable.go b/mcp/streamable.go index 3f53a68..86af05c 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -267,9 +267,7 @@ type stream struct { // streamRequests is the set of unanswered incoming RPCs for the stream. // - // Lifecycle: requests values persist until the requests have been - // replied to by the server. Notably, NOT until they are sent to an HTTP - // response, as delivery is not guaranteed. + // Requests persist until their response data has been added to outgoing. requests map[jsonrpc.ID]struct{} } @@ -482,6 +480,7 @@ stream: t.mu.Lock() outgoing := stream.outgoing stream.outgoing = nil + nOutstanding := len(stream.requests) t.mu.Unlock() for _, data := range outgoing { @@ -493,9 +492,6 @@ stream: } } - t.mu.Lock() - nOutstanding := len(stream.requests) - t.mu.Unlock() // If all requests have been handled and replied to, we should terminate this connection. // "After the JSON-RPC response has been sent, the server SHOULD close the SSE stream." // ยง6.4, https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#sending-messages-to-the-server