From 4b57174010fa9c074ed148c661f16828f81a6db0 Mon Sep 17 00:00:00 2001 From: ychampion Date: Thu, 9 Jul 2026 09:00:14 +0000 Subject: [PATCH] internal/jsonrpc2: allow EOF-side responses to drain Treat read-side EOF as a half-close so already-dispatched requests can finish while the connection still rejects new work and propagates real transport failures. Constraint: Preserve explicit close, write-error, and non-EOF read-error shutdown behavior. Rejected: Drain after every read failure | only clean EOF proves the peer stopped sending without invalidating in-flight responses. Confidence: high Scope-risk: narrow Directive: Keep EOF draining covered at the connection boundary rather than transport-specific callers. Tested: Focused race regression 100x, internal/jsonrpc2 race suite 10x, go test -v -race ./..., go test ./..., go vet ./..., gofmt, and diff checks. Not-tested: External conformance locally; the first rebased head passed both CI conformance suites before the test-only race fix. --- internal/jsonrpc2/conn.go | 35 +++++++-- internal/jsonrpc2/conn_test.go | 136 +++++++++++++++++++++++++++++++++ 2 files changed, 164 insertions(+), 7 deletions(-) create mode 100644 internal/jsonrpc2/conn_test.go 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..8101f64a --- /dev/null +++ b/internal/jsonrpc2/conn_test.go @@ -0,0 +1,136 @@ +// 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{} + + conn := NewConnection(context.Background(), ConnectionConfig{ + Reader: reader, + Writer: writer, + Closer: nopCloser{}, + Bind: func(conn *Connection) Handler { + 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") + } +}