From aae898fa3027eb3a094359ac5a4a1e64f5c354eb Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Tue, 5 Aug 2025 13:42:22 +0000
Subject: [PATCH 1/2] .github: add race and 1.25 tests

This CL makes the following changes to our Github workflows, motivated
by #227:

- Add -race tests.
- Add go1.25rc2 to the build matrix.
- Check out code before setting up Go, so that the setup step can use
  dependency caching.
---
 .github/workflows/test.yml | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index f2f8cb8e..d5009944 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -1,11 +1,11 @@
 name: Test
 on:
   # Manual trigger
-  workflow_dispatch:  
+  workflow_dispatch:
   push:
     branches: main
   pull_request:
-  
+
 permissions:
   contents: read
 
@@ -13,10 +13,10 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-    - name: Set up Go
-      uses: actions/setup-go@v5
     - name: Check out code
       uses: actions/checkout@v4
+    - name: Set up Go
+      uses: actions/setup-go@v5
     - name: Check formatting
       run: |
         unformatted=$(gofmt -l .)
@@ -26,18 +26,30 @@ jobs:
           exit 1
         fi
         echo "All Go files are properly formatted"
-  
+
   test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        go: [ '1.23', '1.24' ]
+        go: [ '1.23', '1.24', '1.25.0-rc.2' ]
     steps:
+    - name: Check out code
+      uses: actions/checkout@v4
     - name: Set up Go
       uses: actions/setup-go@v5
       with:
         go-version: ${{ matrix.go }}
-    - name: Check out code
-      uses: actions/checkout@v4
     - name: Test
       run: go test -v ./...
+
+  race-test:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v4
+    - name: Set up Go
+      uses: actions/setup-go@v5
+      with:
+        go-version: '1.24'
+    - name: Test with -race
+      run: go test -v -race ./...

From 0499eb55e07e2594e884f0828b3a0a102803afc9 Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Tue, 5 Aug 2025 14:50:21 +0000
Subject: [PATCH 2/2] 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 3f53a689..86af05ce 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