modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

mcp/streamable: add resumability for the Streamable transport (#133)

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

For #10
This commit is contained in:
Sam Thanawalla 2025-07-22 15:58:29 -04:00 • committed by GitHub
commit a911cd0ffd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 318 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

225
mcp/streamable.go
View file

@ -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"
@ -597,12 +600,39 @@ type StreamableClientTransport struct {
opts StreamableClientTransportOptions
}
// StreamableReconnectOptions defines parameters for client reconnect attempts.
type StreamableReconnectOptions struct {
// 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
// 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
// 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
}
// DefaultReconnectOptions provides sensible defaults for reconnect logic.
var DefaultReconnectOptions = &StreamableReconnectOptions{
MaxRetries: 5,
growFactor: 1.5,
initialDelay: 1 * time.Second,
maxDelay: 30 * time.Second,
}
// 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 *StreamableReconnectOptions
}
// NewStreamableClientTransport returns a new client transport that connects to
@ -628,22 +658,37 @@ 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
}
// 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 *StreamableReconnectOptions
closeOnce sync.Once
closeErr error
ctx context.Context
cancel context.CancelFunc
mu sync.Mutex
protocolVersion string
@ -665,6 +710,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()
@ -745,6 +796,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":
// Section 2.1: The SSE stream is initiated after a POST.
go s.handleSSE(resp)
case "application/json":
// TODO: read the body and send to s.incoming (in a select that also receives from s.done).
@ -757,34 +809,115 @@ func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string
return sessionID, nil
}
func (s *streamableClientConn) handleSSE(resp *http.Response) {
// handleSSE manages the entire lifecycle of an SSE connection. It processes
// an incoming Server-Sent Events stream and automatically handles reconnection
// logic if the stream breaks.
func (s *streamableClientConn) handleSSE(initialResp *http.Response) {
resp := initialResp
var lastEventID string
for {
eventID, clientClosed := s.processStream(resp)
lastEventID = eventID
// If the connection was closed by the client, we're done.
if clientClosed {
return
}
// The stream was interrupted or ended by the server. Attempt to reconnect.
newResp, err := s.reconnect(lastEventID)
if err != nil {
// All reconnection attempts failed. Set the final error, close the
// connection, and exit the goroutine.
s.mu.Lock()
s.err = err
s.mu.Unlock()
s.Close()
return
}
// Reconnection was successful. Continue the loop with the new response.
resp = newResp
}
}
// processStream reads from a single response body, sending events to the
// incoming channel. It returns the ID of the last processed event, any error
// that occurred, and a flag indicating if the connection was closed by the client.
func (s *streamableClientConn) processStream(resp *http.Response) (lastEventID string, clientClosed bool) {
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
return
}
select {
case <-s.done:
return
case s.incoming <- evt.Data:
}
for evt, err := range scanEvents(resp.Body) {
if err != nil {
return lastEventID, false
}
}()
select {
case <-s.done:
case <-done:
if evt.ID != "" {
lastEventID = evt.ID
}
select {
case s.incoming <- evt.Data:
case <-s.done:
// The connection was closed by the client; exit gracefully.
return lastEventID, true
}
}
// The loop finished without an error, indicating the server closed the stream.
// We'll attempt to reconnect, so this is not a client-side close.
return lastEventID, false
}
// reconnect handles the logic of retrying a connection with an exponential
// backoff strategy. It returns a new, valid HTTP response if successful, or
// an error if all retries are exhausted.
func (s *streamableClientConn) reconnect(lastEventID string) (*http.Response, error) {
var finalErr error
for attempt := 0; attempt < s.ReconnectOptions.MaxRetries; attempt++ {
select {
case <-s.done:
return nil, fmt.Errorf("connection closed by client during reconnect")
case <-time.After(calculateReconnectDelay(s.ReconnectOptions, attempt)):
resp, err := s.establishSSE(lastEventID)
if err != nil {
finalErr = err // Store the error and try again.
continue
}
if !isResumable(resp) {
// The server indicated we should not continue.
resp.Body.Close()
return nil, fmt.Errorf("reconnection failed with unresumable status: %s", resp.Status)
}
return resp, nil
}
}
// If the loop completes, all retries have failed.
if finalErr != nil {
return nil, fmt.Errorf("connection failed after %d attempts: %w", s.ReconnectOptions.MaxRetries, finalErr)
}
return nil, fmt.Errorf("connection failed after %d attempts", s.ReconnectOptions.MaxRetries)
}
// isResumable checks if an HTTP response indicates a valid SSE stream that can be processed.
func isResumable(resp *http.Response) bool {
// Per the spec, a 405 response means the server doesn't support SSE streams at this endpoint.
if resp.StatusCode == http.StatusMethodNotAllowed {
return false
}
return strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream")
}
// 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)
@ -803,3 +936,37 @@ func (s *streamableClientConn) Close() error {
})
return s.closeErr
}
// establishSSE establishes the persistent SSE listening stream.
// It is used for reconnect attempts using the Last-Event-ID header to
// resume a broken stream where it left off.
func (s *streamableClientConn) establishSSE(lastEventID string) (*http.Response, error) {
req, err := http.NewRequestWithContext(s.ctx, http.MethodGet, s.url, nil)
if err != nil {
return nil, err
}
s.mu.Lock()
if s._sessionID != "" {
req.Header.Set("Mcp-Session-Id", s._sessionID)
}
s.mu.Unlock()
if lastEventID != "" {
req.Header.Set("Last-Event-ID", lastEventID)
}
req.Header.Set("Accept", "text/event-stream")
return s.client.Do(req)
}
// calculateReconnectDelay calculates a delay using exponential backoff with full jitter.
func calculateReconnectDelay(opts *StreamableReconnectOptions, 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
}

125
mcp/streamable_test.go
View file

@ -10,19 +10,23 @@ import (
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/http/cookiejar"
"net/http/httptest"
"net/http/httputil"
"net/url"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
"github.com/modelcontextprotocol/go-sdk/jsonschema"
)
func TestStreamableTransports(t *testing.T) {
@ -105,6 +109,127 @@ func TestStreamableTransports(t *testing.T) {
}
}
// TestClientReplay verifies that the client can recover from a
// mid-stream network failure and receive replayed messages. It uses a proxy
// that is killed and restarted to simulate a recoverable network outage.
func TestClientReplay(t *testing.T) {
notifications := make(chan string)
// 1. Configure the real MCP server.
server := NewServer(testImpl, nil)
// Use a channel to synchronize the server's message sending with the test's
// proxy-killing action.
serverReadyToKillProxy := make(chan struct{})
serverClosed := make(chan struct{})
server.AddTool(&Tool{Name: "multiMessageTool", InputSchema: &jsonschema.Schema{}},
func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
go func() {
bgCtx := context.Background()
// Send the first two messages immediately.
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg1"})
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg2"})
// Signal the test that it can now kill the proxy.
close(serverReadyToKillProxy)
<-serverClosed
// These messages should be queued for replay by the server after
// the client's connection drops.
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg3"})
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg4"})
}()
return &CallToolResult{}, nil
})
realServer := httptest.NewServer(NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil))
defer realServer.Close()
realServerURL, err := url.Parse(realServer.URL)
if err != nil {
t.Fatalf("Failed to parse real server URL: %v", err)
}
// 2. Configure a proxy that sits between the client and the real server.
proxyHandler := httputil.NewSingleHostReverseProxy(realServerURL)
proxy := httptest.NewServer(proxyHandler)
proxyAddr := proxy.Listener.Addr().String() // Get the address to restart it later.
// 3. Configure the client to connect to the proxy with default options.
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
client := NewClient(testImpl, &ClientOptions{
ProgressNotificationHandler: func(ctx context.Context, cc *ClientSession, params *ProgressNotificationParams) {
notifications <- params.Message
}})
clientSession, err := client.Connect(ctx, NewStreamableClientTransport(proxy.URL, nil))
if err != nil {
t.Fatalf("client.Connect() failed: %v", err)
}
defer clientSession.Close()
clientSession.CallTool(ctx, &CallToolParams{Name: "multiMessageTool"})
// 4. Read and verify messages until the server signals it's ready for the proxy kill.
receivedNotifications := readProgressNotifications(t, ctx, notifications, 2)
wantReceived := []string{"msg1", "msg2"}
if diff := cmp.Diff(wantReceived, receivedNotifications); diff != "" {
t.Errorf("Received notifications mismatch (-want +got):\n%s", diff)
}
select {
case <-serverReadyToKillProxy:
// Server has sent the first two messages and is paused.
case <-ctx.Done():
t.Fatalf("Context timed out before server was ready to kill proxy")
}
// 5. Simulate a total network failure by closing the proxy.
t.Log("--- Killing proxy to simulate network failure ---")
proxy.CloseClientConnections()
proxy.Close()
close(serverClosed)
// 6. Simulate network recovery by restarting the proxy on the same address.
t.Logf("--- Restarting proxy on %s ---", proxyAddr)
listener, err := net.Listen("tcp", proxyAddr)
if err != nil {
t.Fatalf("Failed to listen on proxy address: %v", err)
}
restartedProxy := &http.Server{Handler: proxyHandler}
go restartedProxy.Serve(listener)
defer restartedProxy.Close()
// 7. Continue reading from the same connection object.
// Its internal logic should successfully retry, reconnect to the new proxy,
// and receive the replayed messages.
recoveredNotifications := readProgressNotifications(t, ctx, notifications, 2)
// 8. Verify the correct messages were received on the recovered connection.
wantRecovered := []string{"msg3", "msg4"}
if diff := cmp.Diff(wantRecovered, recoveredNotifications); diff != "" {
t.Errorf("Recovered notifications mismatch (-want +got):\n%s", diff)
}
}
// Helper to read a specific number of progress notifications.
func readProgressNotifications(t *testing.T, ctx context.Context, notifications chan string, count int) []string {
t.Helper()
var collectedNotifications []string
for {
select {
case n := <-notifications:
collectedNotifications = append(collectedNotifications, n)
if len(collectedNotifications) == count {
return collectedNotifications
}
case <-ctx.Done():
if len(collectedNotifications) != count {
t.Fatalf("readProgressNotifications(): did not receive expected notifications, got %d, want %d", len(collectedNotifications), count)
}
return collectedNotifications
}
}
}
func TestStreamableServerTransport(t *testing.T) {
// This test checks detailed behavior of the streamable server transport, by
// faking the behavior of a streamable client using a sequence of HTTP