mcp: cleanup transport after keepalive ping fails

An onClose function is passed to the ServerSession and ClientSession to
help cleanup resources from the caller.

The onClose function will be executed as part of the ServerSession and
ClientSession closure.

Fixes #258

Author: h9jiang

mcp/client.go

@@ -76,9 +76,9 @@ type ClientOptions struct {
 
 // bind implements the binder[*ClientSession] interface, so that Clients can
 // be connected using [connect].
-func (c *Client) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *clientSessionState) *ClientSession {
+func (c *Client) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *clientSessionState, onClose func()) *ClientSession {
 	assert(mcpConn != nil && conn != nil, "nil connection")
-	cs := &ClientSession{conn: conn, mcpConn: mcpConn, client: c}
+	cs := &ClientSession{conn: conn, mcpConn: mcpConn, client: c, onClose: onClose}
 	if state != nil {
 		cs.state = *state
 	}
@@ -130,7 +130,7 @@ func (c *Client) capabilities() *ClientCapabilities {
 // server, calls or notifications will return an error wrapping
 // [ErrConnectionClosed].
 func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptions) (cs *ClientSession, err error) {
-	cs, err = connect(ctx, t, c, (*clientSessionState)(nil))
+	cs, err = connect(ctx, t, c, (*clientSessionState)(nil), nil)
 	if err != nil {
 		return nil, err
 	}
@@ -173,6 +173,8 @@ 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()
+
 	conn            *jsonrpc2.Connection
 	client          *Client
 	keepaliveCancel context.CancelFunc
@@ -208,7 +210,13 @@ func (cs *ClientSession) Close() error {
 	if cs.keepaliveCancel != nil {
 		cs.keepaliveCancel()
 	}
-	return cs.conn.Close()
+	err := cs.conn.Close()
+
+	if cs.onClose != nil {
+		cs.onClose()
+	}
+
+	return err
 }
 
 // Wait waits for the connection to be closed by the server.

mcp/server.go

@@ -698,9 +698,9 @@ func (s *Server) Run(ctx context.Context, t Transport) error {
 
 // bind implements the binder[*ServerSession] interface, so that Servers can
 // be connected using [connect].
-func (s *Server) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *ServerSessionState) *ServerSession {
+func (s *Server) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *ServerSessionState, onClose func()) *ServerSession {
 	assert(mcpConn != nil && conn != nil, "nil connection")
-	ss := &ServerSession{conn: conn, mcpConn: mcpConn, server: s}
+	ss := &ServerSession{conn: conn, mcpConn: mcpConn, server: s, onClose: onClose}
 	if state != nil {
 		ss.state = *state
 	}
@@ -727,6 +727,8 @@ func (s *Server) disconnect(cc *ServerSession) {
 // ServerSessionOptions configures the server session.
 type ServerSessionOptions struct {
 	State *ServerSessionState
+
+	onClose func()
 }
 
 // Connect connects the MCP server over the given transport and starts handling
@@ -739,10 +741,12 @@ type ServerSessionOptions struct {
 // If opts.State is non-nil, it is the initial state for the server.
 func (s *Server) Connect(ctx context.Context, t Transport, opts *ServerSessionOptions) (*ServerSession, error) {
 	var state *ServerSessionState
+	var onClose func()
 	if opts != nil {
 		state = opts.State
+		onClose = opts.onClose
 	}
-	return connect(ctx, t, s, state)
+	return connect(ctx, t, s, state, onClose)
 }
 
 // TODO: (nit) move all ServerSession methods below the ServerSession declaration.
@@ -809,6 +813,8 @@ 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()
+
 	server          *Server
 	conn            *jsonrpc2.Connection
 	mcpConn         Connection
@@ -1043,7 +1049,13 @@ func (ss *ServerSession) Close() error {
 		//    Close is idempotent and conn.Close() handles concurrent calls correctly
 		ss.keepaliveCancel()
 	}
-	return ss.conn.Close()
+	err := ss.conn.Close()
+
+	if ss.onClose != nil {
+		ss.onClose()
+	}
+
+	return err
 }
 
 // Wait waits for the connection to be closed by the client.

mcp/streamable.go

@@ -40,6 +40,8 @@ type StreamableHTTPHandler struct {
 	getServer func(*http.Request) *Server
 	opts      StreamableHTTPOptions
 
+	onTransportDeletion func(sessionID string) // for testing only
+
 	mu sync.Mutex
 	// TODO: we should store the ServerSession along with the transport, because
 	// we need to cancel keepalive requests when closing the transport.
@@ -283,6 +285,19 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			connectOpts = &ServerSessionOptions{
 				State: state,
 			}
+		} else {
+			// Cleanup is only required in stateful mode, as transportation is
+			// not stored in the map otherwise.
+			connectOpts = &ServerSessionOptions{
+				onClose: func() {
+					h.mu.Lock()
+					delete(h.transports, transport.SessionID)
+					h.mu.Unlock()
+					if h.onTransportDeletion != nil {
+						h.onTransportDeletion(transport.SessionID)
+					}
+				},
+			}
 		}
 
 		// Pass req.Context() here, to allow middleware to add context values.

mcp/streamable_test.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"maps"
@@ -332,6 +333,76 @@ func testClientReplay(t *testing.T, test clientReplayTest) {
 	}
 }
 
+func TestServerTransportCleanup(t *testing.T) {
+	server := NewServer(testImpl, &ServerOptions{KeepAlive: 10 * time.Millisecond})
+
+	nClient := 3
+
+	var mu sync.Mutex
+	var id int = -1 // session id starting from "0", "1", "2"...
+	chans := make(map[string]chan struct{}, nClient)
+
+	handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
+		GetSessionID: func() string {
+			mu.Lock()
+			defer mu.Unlock()
+			id++
+			if id == nClient {
+				t.Errorf("creating more than %v session", nClient)
+			}
+			chans[fmt.Sprint(id)] = make(chan struct{}, 1)
+			return fmt.Sprint(id)
+		},
+	})
+
+	handler.onTransportDeletion = func(sessionID string) {
+		chans[sessionID] <- struct{}{}
+	}
+
+	httpServer := httptest.NewServer(handler)
+	defer httpServer.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// Spin up clients connect to the same server but refuse to ping request.
+	for range nClient {
+		client := NewClient(testImpl, nil)
+		pingMiddleware := func(next MethodHandler) MethodHandler {
+			return func(
+				ctx context.Context,
+				method string,
+				req Request,
+			) (Result, error) {
+				if method == "ping" {
+					return &emptyResult{}, errors.New("ping error")
+				}
+				return next(ctx, method, req)
+			}
+		}
+		client.AddReceivingMiddleware(pingMiddleware)
+		clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
+		if err != nil {
+			t.Fatalf("client.Connect() failed: %v", err)
+		}
+		defer clientSession.Close()
+	}
+
+	for _, ch := range chans {
+		select {
+		case <-ctx.Done():
+			t.Errorf("did not capture transport deletion event from all session in 10 seconds")
+		case <-ch: // Received transport deletion signal of this session
+		}
+	}
+
+	handler.mu.Lock()
+	if len(handler.transports) != 0 {
+		t.Errorf("want empty transports map, find %v entries from handler's transports map", len(handler.transports))
+	}
+	handler.mu.Unlock()
+}
+
 // TestServerInitiatedSSE verifies that the persistent SSE connection remains
 // open and can receive server-initiated events.
 func TestServerInitiatedSSE(t *testing.T) {

mcp/transport.go

@@ -122,15 +122,16 @@ func NewInMemoryTransports() (*InMemoryTransport, *InMemoryTransport) {
 }
 
 type binder[T handler, State any] interface {
-	bind(Connection, *jsonrpc2.Connection, State) T
+	// TODO(rfindley): the bind API has gotten too complicated. Simplify.
+	bind(Connection, *jsonrpc2.Connection, State, func()) T
 	disconnect(T)
 }
 
 type handler interface {
 	handle(ctx context.Context, req *jsonrpc.Request) (any, error)
 }
 
-func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, State], s State) (H, error) {
+func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, State], s State, onClose func()) (H, error) {
 	var zero H
 	mcpConn, err := t.Connect(ctx)
 	if err != nil {
@@ -143,7 +144,7 @@ func connect[H handler, State any](ctx context.Context, t Transport, b binder[H,
 		preempter canceller
 	)
 	bind := func(conn *jsonrpc2.Connection) jsonrpc2.Handler {
-		h = b.bind(mcpConn, conn, s)
+		h = b.bind(mcpConn, conn, s, onClose)
 		preempter.conn = conn
 		return jsonrpc2.HandlerFunc(h.handle)
 	}