diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 4994c63b..3f193271 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -168,6 +168,21 @@ func (s *inFlightState) shuttingDown(errClosing error) error { return nil } +func (s *inFlightState) responseWriteErr(errClosing error) error { + if s.connClosing { + return errClosing + } + if s.writeErr != nil { + // A broken writer prevents responses from reaching the peer. A read-side + // EOF alone does not: already-received requests can still drain responses. + return fmt.Errorf("%w: %v", errClosing, s.writeErr) + } + if s.readErr != nil && !errors.Is(s.readErr, io.EOF) { + return fmt.Errorf("%w: %v", errClosing, s.readErr) + } + return nil +} + // incomingRequest is used to track an incoming request as it is being handled type incomingRequest struct { *Request // the request being processed @@ -549,12 +564,14 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter } s.outgoingCalls = nil - // Cancel any incoming requests still in flight: with the reader gone we - // cannot receive cancellation notifications, and likely cannot write a - // response either, so parked handlers have nothing useful left to do. - // Mirrors the equivalent cleanup on write failure. - for _, r := range s.incomingByID { - r.cancel() + if !errors.Is(err, io.EOF) { + // Cancel any incoming requests still in flight: with the reader gone we + // cannot receive cancellation notifications, and likely cannot write a + // response either, so parked handlers have nothing useful left to do. + // Mirrors the equivalent cleanup on write failure. + for _, r := range s.incomingByID { + r.cancel() + } } }) } @@ -756,7 +773,11 @@ func (c *Connection) write(ctx context.Context, msg Message) error { if req, ok := msg.(*Request); ok && !req.IsCall() && s.outgoingNotifications > 0 { return } - err = s.shuttingDown(ErrServerClosing) + if _, ok := msg.(*Response); ok { + err = s.responseWriteErr(ErrServerClosing) + } else { + err = s.shuttingDown(ErrServerClosing) + } }) if err == nil { err = c.writer.Write(ctx, msg) diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go new file mode 100644 index 00000000..2dde505f --- /dev/null +++ b/internal/jsonrpc2/conn_test.go @@ -0,0 +1,138 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package jsonrpc2 + +import ( + "context" + "errors" + "io" + "sync" + "testing" + "time" +) + +func TestReadEOFAllowsInFlightResponseToDrain(t *testing.T) { + call, err := NewCall(Int64ID(1), "ping", nil) + if err != nil { + t.Fatal(err) + } + + reader := &scriptedReader{messages: []Message{call}} + writer := &recordingWriter{} + + var conn *Connection + conn = NewConnection(context.Background(), ConnectionConfig{ + Reader: reader, + Writer: writer, + Closer: nopCloser{}, + Bind: func(c *Connection) Handler { + conn = c + return HandlerFunc(func(context.Context, *Request) (any, error) { + waitForReadErr(t, conn) + return "pong", nil + }) + }, + OnInternalError: func(err error) { + t.Errorf("internal error: %v", err) + }, + }) + + waitForConnection(t, conn) + + responses := writer.messages() + if len(responses) != 1 { + t.Fatalf("writer saw %d messages, want 1", len(responses)) + } + response, ok := responses[0].(*Response) + if !ok { + t.Fatalf("writer saw %T, want *Response", responses[0]) + } + if response.ID != Int64ID(1) { + t.Fatalf("response ID = %v, want 1", response.ID.Raw()) + } + if string(response.Result) != `"pong"` { + t.Fatalf("response result = %s, want %q", response.Result, `"pong"`) + } +} + +type scriptedReader struct { + messages []Message + index int +} + +func (r *scriptedReader) Read(context.Context) (Message, error) { + if r.index >= len(r.messages) { + return nil, io.EOF + } + msg := r.messages[r.index] + r.index++ + return msg, nil +} + +type recordingWriter struct { + mu sync.Mutex + written []Message +} + +func (w *recordingWriter) Write(_ context.Context, msg Message) error { + w.mu.Lock() + defer w.mu.Unlock() + w.written = append(w.written, msg) + return nil +} + +func (w *recordingWriter) messages() []Message { + w.mu.Lock() + defer w.mu.Unlock() + return append([]Message(nil), w.written...) +} + +type nopCloser struct{} + +func (nopCloser) Close() error { return nil } + +func waitForReadErr(t *testing.T, conn *Connection) { + t.Helper() + + deadline := time.After(2 * time.Second) + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + + for { + conn.stateMu.Lock() + err := conn.state.readErr + conn.stateMu.Unlock() + if err != nil { + if !errors.Is(err, io.EOF) { + t.Fatalf("read error = %v, want EOF", err) + } + return + } + + select { + case <-deadline: + t.Fatal("timed out waiting for read EOF") + case <-ticker.C: + } + } +} + +func waitForConnection(t *testing.T, conn *Connection) { + t.Helper() + + done := make(chan error, 1) + go func() { + done <- conn.Wait() + }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Wait returned error: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for connection shutdown") + } +}