mcp: drain SSE messages when closing

mcp/sse.go

@@ -127,10 +127,9 @@ type SSEServerTransport struct {
 	// the transport is connected.
 	incoming chan jsonrpc.Message
 
-	// We must guard both pushes to the incoming queue and writes to the response
-	// writer, because incoming POST requests are arbitrarily concurrent and we
-	// need to ensure we don't write push to the queue, or write to the
-	// ResponseWriter, after the session GET request exits.
+	// Incoming POST requests are arbitrarily concurrent. The mutex guards writes
+	// to the ResponseWriter and the closed state, which is checked around queue
+	// pushes so messages are not accepted after the session GET request exits.
 	mu     sync.Mutex    // also guards writes to Response
 	closed bool          // set when the stream is closed
 	done   chan struct{} // closed when the connection is closed
@@ -163,8 +162,25 @@ func (t *SSEServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request)
 			return
 		}
 	}
+	t.mu.Lock()
+	closed := t.closed
+	t.mu.Unlock()
+	if closed {
+		http.Error(w, "session closed", http.StatusBadRequest)
+		return
+	}
 	select {
 	case t.incoming <- msg:
+		t.mu.Lock()
+		closed = t.closed
+		if closed {
+			t.drainIncoming()
+		}
+		t.mu.Unlock()
+		if closed {
+			http.Error(w, "session closed", http.StatusBadRequest)
+			return
+		}
 		w.WriteHeader(http.StatusAccepted)
 	case <-t.done:
 		http.Error(w, "session closed", http.StatusBadRequest)
@@ -305,10 +321,22 @@ func (s *sseServerConn) SessionID() string { return "" }
 
 // Read implements jsonrpc2.Reader.
 func (s *sseServerConn) Read(ctx context.Context) (jsonrpc.Message, error) {
+	s.t.mu.Lock()
+	closed := s.t.closed
+	s.t.mu.Unlock()
+	if closed {
+		return nil, io.EOF
+	}
 	select {
 	case <-ctx.Done():
 		return nil, ctx.Err()
 	case msg := <-s.t.incoming:
+		s.t.mu.Lock()
+		closed := s.t.closed
+		s.t.mu.Unlock()
+		if closed {
+			return nil, io.EOF
+		}
 		return msg, nil
 	case <-s.t.done:
 		return nil, io.EOF
@@ -340,18 +368,30 @@ func (s *sseServerConn) Write(ctx context.Context, msg jsonrpc.Message) error {
 	return err
 }
 
+// drainIncoming must be called with t.mu held.
+func (t *SSEServerTransport) drainIncoming() {
+	for {
+		select {
+		case <-t.incoming:
+		default:
+			return
+		}
+	}
+}
+
 // Close implements io.Closer, and closes the session.
 //
 // It must be safe to call Close more than once, as the close may
 // asynchronously be initiated by either the server closing its connection, or
 // by the hanging GET exiting.
 func (s *sseServerConn) Close() error {
 	s.t.mu.Lock()
-	defer s.t.mu.Unlock()
 	if !s.t.closed {
 		s.t.closed = true
 		close(s.t.done)
+		s.t.drainIncoming()
 	}
+	s.t.mu.Unlock()
 	return nil
 }
 

mcp/sse_test.go

@@ -7,16 +7,20 @@ package mcp
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"slices"
+	"strings"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
 )
 
 // TestSSEServerTransport_SupportedVersions verifies that the deprecated
@@ -190,6 +194,45 @@ func TestSSEServer(t *testing.T) {
 	}
 }
 
+func TestSSEServerConnCloseDrainsIncoming(t *testing.T) {
+	transport := &SSEServerTransport{Response: httptest.NewRecorder()}
+	connection, err := transport.Connect(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	transport.incoming <- &jsonrpc.Request{Method: "ping"}
+	if err := connection.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if _, err := connection.Read(ctx); !errors.Is(err, io.EOF) {
+		t.Fatalf("Read after Close error = %v, want io.EOF", err)
+	}
+}
+
+func TestSSEServerTransportServeHTTPAfterClose(t *testing.T) {
+	transport := &SSEServerTransport{Response: httptest.NewRecorder()}
+	connection, err := transport.Connect(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := connection.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	req := httptest.NewRequest(http.MethodPost, "/messages", strings.NewReader(
+		`{"jsonrpc":"2.0","id":1,"method":"ping","params":{}}`,
+	))
+	recorder := httptest.NewRecorder()
+	transport.ServeHTTP(recorder, req)
+	if got, want := recorder.Code, http.StatusBadRequest; got != want {
+		t.Fatalf("POST after Close status = %d, want %d", got, want)
+	}
+}
+
 // roundTripperFunc is a helper to create a custom RoundTripper
 type roundTripperFunc func(*http.Request) (*http.Response, error)