address review comments

Author: findleyr

mcp/client.go

@@ -11,6 +11,7 @@ import (
 	"iter"
 	"slices"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/jsonschema-go/jsonschema"
@@ -177,7 +178,11 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio
 // Call [ClientSession.Close] to close the connection, or await server
 // termination with [ClientSession.Wait].
 type ClientSession struct {
-	onClose func()
+	// Ensure that onClose is called at most once.
+	// We defensively use an atomic CompareAndSwap rather than a sync.Once, in case the
+	// onClose callback triggers a re-entrant call to Close.
+	calledOnClose atomic.Bool
+	onClose       func()
 
 	conn            *jsonrpc2.Connection
 	client          *Client
@@ -205,6 +210,8 @@ func (cs *ClientSession) ID() string {
 // Close performs a graceful close of the connection, preventing new requests
 // from being handled, and waiting for ongoing requests to return. Close then
 // terminates the connection.
+//
+// Close is idempotent and concurrency safe.
 func (cs *ClientSession) Close() error {
 	// Note: keepaliveCancel access is safe without a mutex because:
 	// 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls)
@@ -216,7 +223,7 @@ func (cs *ClientSession) Close() error {
 	}
 	err := cs.conn.Close()
 
-	if cs.onClose != nil {
+	if cs.onClose != nil && cs.calledOnClose.CompareAndSwap(false, true) {
 		cs.onClose()
 	}
 

mcp/server.go

@@ -19,6 +19,7 @@ import (
 	"reflect"
 	"slices"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/jsonschema-go/jsonschema"
@@ -920,7 +921,11 @@ func newServerRequest[P Params](ss *ServerSession, params P) *ServerRequest[P] {
 // Call [ServerSession.Close] to close the connection, or await client
 // termination with [ServerSession.Wait].
 type ServerSession struct {
-	onClose func()
+	// Ensure that onClose is called at most once.
+	// We defensively use an atomic CompareAndSwap rather than a sync.Once, in case the
+	// onClose callback triggers a re-entrant call to Close.
+	calledOnClose atomic.Bool
+	onClose       func()
 
 	server          *Server
 	conn            *jsonrpc2.Connection
@@ -1185,6 +1190,8 @@ func (ss *ServerSession) setLevel(_ context.Context, params *SetLoggingLevelPara
 // Close performs a graceful shutdown of the connection, preventing new
 // requests from being handled, and waiting for ongoing requests to return.
 // Close then terminates the connection.
+//
+// Close is idempotent and concurrency safe.
 func (ss *ServerSession) Close() error {
 	if ss.keepaliveCancel != nil {
 		// Note: keepaliveCancel access is safe without a mutex because:
@@ -1196,7 +1203,7 @@ func (ss *ServerSession) Close() error {
 	}
 	err := ss.conn.Close()
 
-	if ss.onClose != nil {
+	if ss.onClose != nil && ss.calledOnClose.CompareAndSwap(false, true) {
 		ss.onClose()
 	}
 

mcp/streamable.go

@@ -54,11 +54,14 @@ type sessionInfo struct {
 	// If timeout is set, automatically close the session after an idle period.
 	timeout time.Duration
 	timerMu sync.Mutex
+	refs    int // reference count
 	timer   *time.Timer
 }
 
-// resetTimeout resets the inactivity timer.
-func (i *sessionInfo) resetTimeout() {
+// startPOST signals that a POST request for this session is starting (which
+// carries a client->server message), stopping the session timeout if it was
+// running.
+func (i *sessionInfo) startPOST() {
 	if i.timeout <= 0 {
 		return
 	}
@@ -67,13 +70,36 @@ func (i *sessionInfo) resetTimeout() {
 	defer i.timerMu.Unlock()
 
 	if i.timer == nil {
+		return // timer stopped permanently
+	}
+	if i.refs == 0 {
+		i.timer.Stop()
+	}
+	i.refs++
+}
+
+// endPOST sigals that a request for this session is ending, starting the
+// timeout if there are no other requests running.
+func (i *sessionInfo) endPOST() {
+	if i.timeout <= 0 {
 		return
 	}
-	// Reset the timer if we successfully stopped it.
-	i.timer.Reset(i.timeout)
+
+	i.timerMu.Lock()
+	defer i.timerMu.Unlock()
+
+	if i.timer == nil {
+		return // timer stopped permanently
+	}
+
+	i.refs--
+	assert(i.refs >= 0, "negative ref count")
+	if i.refs == 0 {
+		i.timer.Reset(i.timeout)
+	}
 }
 
-// stopTimer stops the inactivity timer.
+// stopTimer stops the inactivity timer permanently.
 func (i *sessionInfo) stopTimer() {
 	i.timerMu.Lock()
 	defer i.timerMu.Unlock()
@@ -374,13 +400,17 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			session:   session,
 			transport: transport,
 		}
+
 		if h.opts.Stateless {
 			// Stateless mode: close the session when the request exits.
 			defer session.Close() // close the fake session after handling the request
 		} else {
 			// Otherwise, save the transport so that it can be reused
 
 			// Clean up the session when it times out.
+			//
+			// Note that the timer here may fire multiple times, but
+			// sessInfo.session.Close is idempotent.
 			if h.opts.SessionTimeout > 0 {
 				sessInfo.timeout = h.opts.SessionTimeout
 				sessInfo.timer = time.AfterFunc(sessInfo.timeout, func() {
@@ -393,7 +423,11 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		}
 	}
 
-	sessInfo.resetTimeout()
+	if req.Method == http.MethodPost {
+		sessInfo.startPOST()
+		defer sessInfo.endPOST()
+	}
+
 	sessInfo.transport.ServeHTTP(w, req)
 }
 

mcp/streamable_client_test.go

@@ -34,7 +34,6 @@ type streamableResponse struct {
 	body                string // or ""
 	optional            bool   // if set, request need not be sent
 	wantProtocolVersion string // if "", unchecked
-	callback            func() // if set, called after the request is handled
 }
 
 type fakeResponses map[streamableRequestKey]*streamableResponse
@@ -96,9 +95,6 @@ func (s *fakeStreamableServer) ServeHTTP(w http.ResponseWriter, req *http.Reques
 		http.Error(w, "no response", http.StatusInternalServerError)
 		return
 	}
-	if resp.callback != nil {
-		defer resp.callback()
-	}
 	for k, v := range resp.header {
 		w.Header().Set(k, v)
 	}
@@ -411,3 +407,38 @@ func TestStreamableClientStrictness(t *testing.T) {
 		})
 	}
 }
+
+func TestStreamableClientUnresumableRequest(t *testing.T) {
+	// This test verifies that the client fails fast when making a request that
+	// is unresumable, because it does not contain any events.
+	ctx := context.Background()
+	fake := &fakeStreamableServer{
+		t: t,
+		responses: fakeResponses{
+			{"POST", "", methodInitialize}: {
+				header: header{
+					"Content-Type":  "text/event-stream",
+					sessionIDHeader: "123",
+				},
+				body: "",
+			},
+			{"DELETE", "123", ""}: {optional: true},
+		},
+	}
+	httpServer := httptest.NewServer(fake)
+	defer httpServer.Close()
+
+	transport := &StreamableClientTransport{Endpoint: httpServer.URL}
+	client := NewClient(testImpl, nil)
+	cs, err := client.Connect(ctx, transport, nil)
+	if err == nil {
+		cs.Close()
+		t.Fatalf("Connect succeeded unexpectedly")
+	}
+	// This may be a bit of a change detector, but for now check that we're
+	// actually exercising the early failure codepath.
+	msg := "terminated without response"
+	if !strings.Contains(err.Error(), msg) {
+		t.Errorf("Connect: got error %v, want containing %q", err, msg)
+	}
+}