diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 89307c3f4..e91d8e50e 100644 --- a/internal/jsonrpc2/conn.go +++ b/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 diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go index b4a2ec1d3..9e77bf0a3 100644 --- a/internal/jsonrpc2/conn_test.go +++ b/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 diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index be021c8ed..01c171bee 100644 --- a/mcp/cmd_test.go +++ b/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() { diff --git a/mcp/shared.go b/mcp/shared.go index dcc69b482..e845128b9 100644 --- a/mcp/shared.go +++ b/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. diff --git a/mcp/shared_test.go b/mcp/shared_test.go index e9dc19066..acd1d4d09 100644 --- a/mcp/shared_test.go +++ b/mcp/shared_test.go @@ -9,6 +9,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "strings" "testing" @@ -16,6 +17,41 @@ import ( "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 diff --git a/mcp/transport.go b/mcp/transport.go index 2b04ae9c6..0329626a7 100644 --- a/mcp/transport.go +++ b/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. // diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 90f7cb628..cd844de3c 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -19,6 +19,190 @@ import ( "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) +func TestCallErrorsDoNotWrapConnectionClosed(t *testing.T) { + tests := []struct { + name string + rpcError error + result Result + cancel bool + want error + }{ + {name: "context cancellation", cancel: true, want: context.Canceled}, + {name: "JSON-RPC error", rpcError: jsonrpc2.ErrInvalidParams, want: jsonrpc2.ErrInvalidParams}, + {name: "decoding EOF", result: &unmarshalErrorResult{err: io.EOF}, want: io.EOF}, + {name: "decoding unexpected EOF", result: &unmarshalErrorResult{err: io.ErrUnexpectedEOF}, want: io.ErrUnexpectedEOF}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + conn, peer := newTestJSONRPCConnection(t, nil) + done := make(chan error, 1) + go func() { + done <- call(ctx, conn, "test", nil, test.result) + }() + msg, err := peer.Read(ctx) + if err != nil { + t.Fatal(err) + } + if test.cancel { + cancel() + } else { + resp := &jsonrpc.Response{ + ID: msg.(*jsonrpc.Request).ID, + Result: []byte("{}"), + Error: test.rpcError, + } + if err := peer.Write(ctx, resp); err != nil { + t.Fatal(err) + } + } + err = <-done + if !errors.Is(err, test.want) { + t.Errorf("call() = %v, want an error wrapping %v", err, test.want) + } + if errors.Is(err, ErrConnectionClosed) { + t.Errorf("call() = %v, must not wrap ErrConnectionClosed", err) + } + }) + } +} + +func TestCallConnectionClosed(t *testing.T) { + for _, test := range connectionClosedTests() { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + conn, peer := newTestJSONRPCConnection(t, test.readErr) + if test.close { + if err := conn.Close(); err != nil { + t.Fatal(err) + } + } else { + done := make(chan error, 1) + go func() { done <- call(ctx, conn, "test", nil, nil) }() + if _, err := peer.Read(ctx); err != nil { + t.Fatal(err) + } + if err := peer.Close(); err != nil { + t.Fatal(err) + } + checkConnectionClosedError(t, <-done, test.causes...) + _ = conn.Wait() + } + err := call(ctx, conn, "test", nil, nil) + checkConnectionClosedError(t, err, test.causes...) + if err != nil && !strings.Contains(err.Error(), `calling "test"`) { + t.Errorf("call() = %v, want the method name in the error", err) + } + }) + } +} + +type connectionClosedTest struct { + name string + readErr error + causes []error + close bool +} + +func connectionClosedTests() []connectionClosedTest { + readErr := errors.New("read failed") + wrappedReadErr := fmt.Errorf("reading: %w", readErr) + sameMessageErr := errors.New("client is closing") + return []connectionClosedTest{ + {name: "local close", close: true}, + {name: "EOF", readErr: io.EOF, causes: []error{io.EOF}}, + {name: "unexpected EOF", readErr: io.ErrUnexpectedEOF, causes: []error{io.ErrUnexpectedEOF}}, + {name: "read error", readErr: readErr, causes: []error{readErr}}, + {name: "wrapped read error", readErr: wrappedReadErr, causes: []error{wrappedReadErr, readErr}}, + {name: "client closing", readErr: jsonrpc2.NewError(-32003, "client is closing")}, + {name: "server closing", readErr: jsonrpc2.NewError(-32004, "server is closing")}, + { + name: "wrapped closing error", + readErr: fmt.Errorf("outer: %w", fmt.Errorf("%w: %w", jsonrpc2.ErrServerClosing, wrappedReadErr)), + causes: []error{wrappedReadErr, readErr}, + }, + { + name: "joined closing errors", + readErr: errors.Join(jsonrpc2.ErrClientClosing, errors.Join(jsonrpc2.ErrServerClosing, readErr), io.ErrUnexpectedEOF), + causes: []error{readErr, io.ErrUnexpectedEOF}, + }, + {name: "wire error cause", readErr: jsonrpc2.ErrInvalidParams, causes: []error{jsonrpc2.ErrInvalidParams}}, + {name: "same message cause", readErr: sameMessageErr, causes: []error{sameMessageErr}}, + } +} + +func checkConnectionClosedError(t *testing.T, err error, causes ...error) { + t.Helper() + for _, want := range append([]error{ErrConnectionClosed}, causes...) { + if !errors.Is(err, want) { + t.Errorf("error = %v, want an error wrapping %v", err, want) + } + } + for _, closing := range []error{jsonrpc2.ErrClientClosing, jsonrpc2.ErrServerClosing} { + if errors.Is(err, closing) { + t.Errorf("error = %v, must not wrap %v", err, closing) + } + } + var wireErr *jsonrpc2.WireError + if errors.As(err, &wireErr) && (wireErr.Code == -32003 || wireErr.Code == -32004) { + t.Errorf("error = %v, exposes internal closing code %d", err, wireErr.Code) + } +} + +type unmarshalErrorResult struct { + ResultBase + err error +} + +func (r *unmarshalErrorResult) UnmarshalJSON([]byte) error { return r.err } + +func newTestJSONRPCConnection(t *testing.T, readErr error) (*jsonrpc2.Connection, Connection) { + t.Helper() + ctx := context.Background() + clientT, serverT := NewInMemoryTransports() + peer, err := serverT.Connect(ctx) + if err != nil { + t.Fatal(err) + } + client, err := clientT.Connect(ctx) + if err != nil { + peer.Close() + t.Fatal(err) + } + if readErr != nil { + client = &readErrorConnection{Connection: client, err: readErr} + } + conn := jsonrpc2.NewConnection(ctx, jsonrpc2.ConnectionConfig{ + Reader: client, + Writer: client, + Closer: client, + Bind: func(*jsonrpc2.Connection) jsonrpc2.Handler { + return jsonrpc2.HandlerFunc(func(context.Context, *jsonrpc2.Request) (any, error) { + return nil, jsonrpc2.ErrNotHandled + }) + }, + }) + t.Cleanup(func() { + peer.Close() + conn.Close() + }) + return conn, peer +} + +type readErrorConnection struct { + Connection + err error +} + +func (c *readErrorConnection) Read(ctx context.Context) (jsonrpc.Message, error) { + msg, err := c.Connection.Read(ctx) + if err != nil { + return nil, c.err + } + return msg, nil +} + func TestBatchFraming(t *testing.T) { // This test checks that the ndjsonFramer can read and write JSON batches. //