mcp/streamable: add resumability for the Streamable transport

This CL implements a retry mechanism to resume SSE streams to recover
from network failures.

Author: samthanawalla

mcp/streamable.go

@@ -9,11 +9,14 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"math"
+	"math/rand/v2"
 	"net/http"
 	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
@@ -594,12 +597,39 @@ type StreamableClientTransport struct {
 	opts StreamableClientTransportOptions
 }
 
+// StreamableClientReconnectOptions defines parameters for client reconnect attempts.
+type StreamableClientReconnectOptions struct {
+	// InitialDelay is the base delay for the first reconnect attempt.
+	InitialDelay time.Duration
+
+	// MaxDelay caps the backoff delay, preventing it from growing indefinitely.
+	MaxDelay time.Duration
+
+	// GrowFactor is the multiplicative factor by which the delay increases after each attempt.
+	// A value of 1.0 results in a constant delay, while a value of 2.0 would double it each time.
+	// It must be 1.0 or greater if MaxRetries is greater than 0.
+	GrowFactor float64
+
+	// MaxRetries is the maximum number of times to attempt a reconnect before giving up.
+	// A value of 0 or less means never retry.
+	MaxRetries int
+}
+
+// DefaultReconnectOptions provides sensible defaults for reconnect logic.
+var DefaultReconnectOptions = &StreamableClientReconnectOptions{
+	InitialDelay: 1 * time.Second,
+	MaxDelay:     30 * time.Second,
+	GrowFactor:   1.5,
+	MaxRetries:   5,
+}
+
 // StreamableClientTransportOptions provides options for the
 // [NewStreamableClientTransport] constructor.
 type StreamableClientTransportOptions struct {
 	// HTTPClient is the client to use for making HTTP requests. If nil,
 	// http.DefaultClient is used.
-	HTTPClient *http.Client
+	HTTPClient       *http.Client
+	ReconnectOptions *StreamableClientReconnectOptions
 }
 
 // NewStreamableClientTransport returns a new client transport that connects to
@@ -625,22 +655,42 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er
 	if client == nil {
 		client = http.DefaultClient
 	}
-	return &streamableClientConn{
-		url:      t.url,
-		client:   client,
-		incoming: make(chan []byte, 100),
-		done:     make(chan struct{}),
-	}, nil
+	reconnOpts := t.opts.ReconnectOptions
+	if reconnOpts == nil {
+		reconnOpts = DefaultReconnectOptions
+	}
+	// A factor less than 1.0 would cause the delay to shrink, which defeats
+	// the purpose of backoff.
+	if reconnOpts.MaxRetries > 0 && reconnOpts.GrowFactor < 1.0 {
+		return nil, fmt.Errorf("invalid grow factor, cannot be less than 1.0: %v", reconnOpts.GrowFactor)
+	}
+	// Create a new cancellable context that will manage the connection's lifecycle.
+	// This is crucial for cleanly shutting down the background SSE listener by
+	// cancelling its blocking network operations, which prevents hangs on exit.
+	connCtx, cancel := context.WithCancel(context.Background())
+	conn := &streamableClientConn{
+		url:              t.url,
+		client:           client,
+		incoming:         make(chan []byte, 100),
+		done:             make(chan struct{}),
+		ReconnectOptions: reconnOpts,
+		ctx:              connCtx,
+		cancel:           cancel,
+	}
+	return conn, nil
 }
 
 type streamableClientConn struct {
-	url      string
-	client   *http.Client
-	incoming chan []byte
-	done     chan struct{}
+	url              string
+	client           *http.Client
+	incoming         chan []byte
+	done             chan struct{}
+	ReconnectOptions *StreamableClientReconnectOptions
 
 	closeOnce sync.Once
 	closeErr  error
+	ctx       context.Context
+	cancel    context.CancelFunc
 
 	mu              sync.Mutex
 	protocolVersion string
@@ -662,6 +712,12 @@ func (c *streamableClientConn) SessionID() string {
 
 // Read implements the [Connection] interface.
 func (s *streamableClientConn) Read(ctx context.Context) (jsonrpc.Message, error) {
+	s.mu.Lock()
+	err := s.err
+	s.mu.Unlock()
+	if err != nil {
+		return nil, err
+	}
 	select {
 	case <-ctx.Done():
 		return nil, ctx.Err()
@@ -704,11 +760,19 @@ func (s *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
 	if sessionID == "" {
 		// locked
 		s._sessionID = gotSessionID
+		// With the session now established, launch the persistent background listener for server-pushed events.
+		go s.establishSSE(&startSSEState{})
 	}
 
 	return nil
 }
 
+// startSSEState holds the state for initiating an SSE stream.
+type startSSEState struct {
+	lastEventID string
+	attempt     int
+}
+
 func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string, msg jsonrpc.Message) (string, error) {
 	data, err := jsonrpc2.EncodeMessage(msg)
 	if err != nil {
@@ -742,7 +806,7 @@ func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string
 	sessionID = resp.Header.Get(sessionIDHeader)
 	switch ct := resp.Header.Get("Content-Type"); ct {
 	case "text/event-stream":
-		go s.handleSSE(resp)
+		go s.handleSSE(resp, &startSSEState{})
 	case "application/json":
 		// TODO: read the body and send to s.incoming (in a select that also receives from s.done).
 		resp.Body.Close()
@@ -754,17 +818,20 @@ func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string
 	return sessionID, nil
 }
 
-func (s *streamableClientConn) handleSSE(resp *http.Response) {
+// handleSSE processes an incoming Server-Sent Events stream, pushing received messages to the client's channel.
+// If the stream breaks, it uses the last received event ID to automatically trigger the reconnect logic.
+func (s *streamableClientConn) handleSSE(resp *http.Response, opts *startSSEState) {
 	defer resp.Body.Close()
 
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
 		for evt, err := range scanEvents(resp.Body) {
 			if err != nil {
-				// TODO: surface this error; possibly break the stream
+				s.scheduleReconnect(opts)
 				return
 			}
+			opts.lastEventID = evt.id
 			select {
 			case <-s.done:
 				return
@@ -782,6 +849,8 @@ func (s *streamableClientConn) handleSSE(resp *http.Response) {
 // Close implements the [Connection] interface.
 func (s *streamableClientConn) Close() error {
 	s.closeOnce.Do(func() {
+		// Cancel any hanging network requests.
+		s.cancel()
 		close(s.done)
 
 		req, err := http.NewRequest(http.MethodDelete, s.url, nil)
@@ -800,3 +869,87 @@ func (s *streamableClientConn) Close() error {
 	})
 	return s.closeErr
 }
+
+// establishSSE creates and manages the persistent SSE listening stream.
+// It is used for both the initial connection and all subsequent reconnect attempts,
+// using the Last-Event-ID header to resume a broken stream where it left off.
+// On a successful response, it delegates to handleSSE to process events;
+// on failure, it triggers the client's reconnect logic.
+func (s *streamableClientConn) establishSSE(opts *startSSEState) {
+	select {
+	case <-s.done:
+		return
+	default:
+	}
+
+	req, err := http.NewRequestWithContext(s.ctx, http.MethodGet, s.url, nil)
+	if err != nil {
+		return
+	}
+	s.mu.Lock()
+	if s._sessionID != "" {
+		req.Header.Set("Mcp-Session-Id", s._sessionID)
+	}
+	s.mu.Unlock()
+	if opts.lastEventID != "" {
+		req.Header.Set("Last-Event-ID", opts.lastEventID)
+	}
+	req.Header.Set("Accept", "text/event-stream")
+
+	resp, err := s.client.Do(req)
+	if err != nil {
+		// On connection error, schedule a retry.
+		s.scheduleReconnect(opts)
+		return
+	}
+
+	// Per the spec, a 405 response means the server doesn't support SSE streams at this endpoint.
+	if resp.StatusCode == http.StatusMethodNotAllowed {
+		resp.Body.Close()
+		return
+	}
+
+	if !strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") {
+		resp.Body.Close()
+		s.scheduleReconnect(opts)
+		return
+	}
+
+	s.handleSSE(resp, opts)
+}
+
+// scheduleReconnect schedules the next SSE reconnect attempt after a delay.
+func (s *streamableClientConn) scheduleReconnect(opts *startSSEState) {
+	reconnOpts := s.ReconnectOptions
+	if opts.attempt >= reconnOpts.MaxRetries {
+		s.mu.Lock()
+		s.err = fmt.Errorf("connection failed after %d attempts", reconnOpts.MaxRetries)
+		s.mu.Unlock()
+		s.Close() // Close the connection to unblock any readers.
+		return
+	}
+
+	delay := calculateReconnectDelay(reconnOpts, opts.attempt)
+
+	select {
+	case <-s.done:
+		return
+	case <-time.After(delay):
+		opts.attempt++
+		s.establishSSE(opts)
+	}
+}
+
+// calculateReconnectDelay calculates a delay using exponential backoff with full jitter.
+func calculateReconnectDelay(opts *StreamableClientReconnectOptions, attempt int) time.Duration {
+	// Calculate the exponential backoff using the grow factor.
+	backoffDuration := time.Duration(float64(opts.InitialDelay) * math.Pow(opts.GrowFactor, float64(attempt)))
+
+	// Cap the backoffDuration at maxDelay.
+	backoffDuration = min(backoffDuration, opts.MaxDelay)
+
+	// Use a full jitter using backoffDuration
+	jitter := rand.N(backoffDuration)
+
+	return backoffDuration + jitter
+}