mcp: wrap ErrConnectionClosed for notifications and interrupted calls

internal/jsonrpc2/conn.go

@@ -585,7 +585,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter
 		// Retire any outgoing requests that were still in flight: with the Reader no
 		// longer being processed, they necessarily cannot receive a response.
 		for id, ac := range s.outgoingCalls {
-			ac.retire(&Response{ID: id, Error: err})
+			ac.retire(&Response{ID: id, Error: fmt.Errorf("%w: %w", ErrClientClosing, err)})
 		}
 		s.outgoingCalls = nil
 

internal/jsonrpc2/conn_test.go

@@ -108,6 +108,48 @@ func TestShuttingDownWrapsReadError(t *testing.T) {
 	}
 }
 
+func TestReadErrorRetiresCalls(t *testing.T) {
+	for _, readErr := range []error{io.EOF, io.ErrUnexpectedEOF, errors.New("read failed")} {
+		t.Run(readErr.Error(), func(t *testing.T) {
+			ctx := context.Background()
+			failed := make(chan struct{})
+			fail := sync.OnceFunc(func() { close(failed) })
+			w := fakeWriter{out: make(chan Message, 1)}
+			conn := NewConnection(ctx, ConnectionConfig{
+				Reader: readerFunc(func(context.Context) (Message, error) {
+					<-failed
+					return nil, readErr
+				}),
+				Writer: w,
+				Closer: nopCloser{},
+				Bind: func(*Connection) Handler {
+					return HandlerFunc(func(context.Context, *Request) (any, error) {
+						return nil, ErrNotHandled
+					})
+				},
+			})
+			t.Cleanup(func() {
+				fail()
+				conn.Close()
+			})
+
+			call := conn.Call(ctx, "test", nil)
+			<-w.out
+			fail()
+			err := call.Await(ctx, nil)
+			for _, want := range []error{ErrClientClosing, readErr} {
+				if !errors.Is(err, want) {
+					t.Errorf("Await() = %v, want an error wrapping %v", err, want)
+				}
+			}
+		})
+	}
+}
+
+type readerFunc func(context.Context) (Message, error)
+
+func (f readerFunc) Read(ctx context.Context) (Message, error) { return f(ctx) }
+
 // TestCancelFromPeer verifies that CancelFromPeer tells a [ResponseDropper]
 // writer about the call before the handler's context is cancelled, that a
 // local Cancel and a cancellation arriving after the response do not, and that

mcp/cmd_test.go

@@ -7,6 +7,7 @@ package mcp_test
 import (
 	"context"
 	"errors"
+	"io"
 	"log"
 	"os"
 	"os/exec"
@@ -19,6 +20,7 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
@@ -54,9 +56,54 @@ func TestMain(m *testing.M) {
 // serverFuncs defines server functions that may be run as subprocesses via
 // [TestMain].
 var serverFuncs = map[string]func(){
-	"default":       runServer,
-	"cancelContext": runCancelContextServer,
-	"largeFrame":    runLargeFrameServer,
+	"default":        runServer,
+	"cancelContext":  runCancelContextServer,
+	"largeFrame":     runLargeFrameServer,
+	"exitDuringCall": runExitDuringCallServer,
+}
+
+func runExitDuringCallServer() {
+	server := mcp.NewServer(testImpl, nil)
+	mcp.AddTool(server, &mcp.Tool{Name: "exit"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) {
+		os.Exit(3)
+		return nil, nil, nil
+	})
+	if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func TestStdioCallInterruptedByPeerExit(t *testing.T) {
+	requireExec(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	client := mcp.NewClient(testImpl, nil)
+	cs, err := client.Connect(ctx, &mcp.CommandTransport{Command: createServerCommand(t, "exitDuringCall")}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cs.Close()
+
+	for _, name := range []string{"interrupted call", "next call"} {
+		_, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: "exit"})
+		if ctx.Err() != nil {
+			t.Fatalf("%s returned after context cancellation: %v", name, err)
+		}
+		for _, want := range []error{mcp.ErrConnectionClosed, io.EOF} {
+			if !errors.Is(err, want) {
+				t.Errorf("%s = %v, want an error wrapping %v", name, err, want)
+			}
+		}
+		for _, closing := range []error{jsonrpc2.ErrClientClosing, jsonrpc2.ErrServerClosing} {
+			if errors.Is(err, closing) {
+				t.Errorf("%s = %v, must not wrap %v", name, err, closing)
+			}
+		}
+		var wireErr *jsonrpc2.WireError
+		if errors.As(err, &wireErr) && (wireErr.Code == -32003 || wireErr.Code == -32004) {
+			t.Errorf("%s = %v, exposes internal closing code %d", name, err, wireErr.Code)
+		}
+	}
 }
 
 func runServer() {

mcp/shared.go

@@ -160,7 +160,15 @@ func defaultSendingMethodHandler(ctx context.Context, method string, req Request
 	}
 	// Notifications don't have results.
 	if strings.HasPrefix(method, "notifications/") {
-		return nil, req.GetSession().getConn().Notify(ctx, method, params)
+		err := req.GetSession().getConn().Notify(ctx, method, params)
+		if errors.Is(err, jsonrpc2.ErrClientClosing) || errors.Is(err, jsonrpc2.ErrServerClosing) {
+			if cause := withoutClosingErrors(err); cause != nil {
+				err = fmt.Errorf("%w: %w", ErrConnectionClosed, cause)
+			} else {
+				err = ErrConnectionClosed
+			}
+		}
+		return nil, err
 	}
 	// Create the result to unmarshal into.
 	// The concrete type of the result is the return type of the receiving function.

mcp/shared_test.go

@@ -9,13 +9,49 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
 )
 
+func TestNotifyAfterServerClose(t *testing.T) {
+	cs, ss, cleanup := basicConnection(t, nil)
+	defer cleanup()
+	if err := ss.Close(); err != nil {
+		t.Fatal(err)
+	}
+	if err := cs.Wait(); err != nil {
+		t.Fatal(err)
+	}
+
+	err := cs.NotifyProgress(context.Background(), &ProgressNotificationParams{ProgressToken: "test"})
+	checkConnectionClosedError(t, err, io.EOF)
+}
+
+func TestNotifyConnectionClosed(t *testing.T) {
+	for _, test := range connectionClosedTests() {
+		t.Run(test.name, func(t *testing.T) {
+			conn, peer := newTestJSONRPCConnection(t, test.readErr)
+			if test.close {
+				if err := conn.Close(); err != nil {
+					t.Fatal(err)
+				}
+			} else {
+				if err := peer.Close(); err != nil {
+					t.Fatal(err)
+				}
+				_ = conn.Wait()
+			}
+			cs := &ClientSession{client: NewClient(testImpl, nil), conn: conn}
+			err := cs.NotifyProgress(context.Background(), &ProgressNotificationParams{ProgressToken: "test"})
+			checkConnectionClosedError(t, err, test.causes...)
+		})
+	}
+}
+
 func TestValidateRequestMeta(t *testing.T) {
 	tests := []struct {
 		name            string

mcp/transport.go

@@ -311,7 +311,10 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
 		// only for requests believed to still be in progress.
 		return nil
 	case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing):
-		return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err)
+		if cause := withoutClosingErrors(err); cause != nil {
+			return fmt.Errorf("%w: calling %q: %w", ErrConnectionClosed, method, cause)
+		}
+		return fmt.Errorf("%w: calling %q", ErrConnectionClosed, method)
 	case ctx.Err() != nil:
 		// The notifications/cancelled message is best-effort. Retire the call
 		// immediately (so an unresponsive peer cannot delay the eager
@@ -341,6 +344,29 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
 	return nil
 }
 
+// withoutClosingErrors returns err with jsonrpc2.ErrClientClosing and
+// jsonrpc2.ErrServerClosing removed from its tree, or nil if nothing else is
+// left. Those errors are internal: wrapping them would expose their codes.
+func withoutClosingErrors(err error) error {
+	if !errors.Is(err, jsonrpc2.ErrClientClosing) && !errors.Is(err, jsonrpc2.ErrServerClosing) {
+		return err
+	}
+	switch err := err.(type) {
+	case interface{ Unwrap() []error }:
+		var causes []error
+		for _, child := range err.Unwrap() {
+			if cause := withoutClosingErrors(child); cause != nil {
+				causes = append(causes, cause)
+			}
+		}
+		return errors.Join(causes...)
+	case interface{ Unwrap() error }:
+		return withoutClosingErrors(err.Unwrap())
+	default:
+		return nil
+	}
+}
+
 // cancelCall sends a "notifications/cancelled" notification for call and eagerly
 // retires it from conn.
 //