diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index e91d8e50..318636c1 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -54,8 +54,9 @@ type Connection struct { writer Writer handler Handler - onInternalError func(error) - onDone func() + onInternalError func(error) + onResponseWriteError func(*Request, error) + onDone func() } // inFlightState records the state of the incoming and outgoing calls on a @@ -219,6 +220,11 @@ type ConnectionConfig struct { OnDone func() // optional OnInternalError func(error) // optional + // OnResponseWriteError, if set, is called when the response to an incoming + // call could not be written, for example because the peer went away before + // the handler finished. The response is lost: the peer never receives it. + OnResponseWriteError func(*Request, error) // optional + // PropagateCancellation controls whether cancellation of the context // passed to [NewConnection] is observable by request handlers. // @@ -244,11 +250,12 @@ func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection { } c := &Connection{ - state: inFlightState{closer: cfg.Closer}, - done: make(chan struct{}), - writer: cfg.Writer, - onDone: cfg.OnDone, - onInternalError: cfg.OnInternalError, + state: inFlightState{closer: cfg.Closer}, + done: make(chan struct{}), + writer: cfg.Writer, + onDone: cfg.OnDone, + onInternalError: cfg.OnInternalError, + onResponseWriteError: cfg.OnResponseWriteError, } c.handler = cfg.Bind(c) c.start(ctx, cfg.Reader, cfg.Preempter) @@ -759,6 +766,9 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e delete(s.incomingByID, req.ID) }) writeErr := c.write(notDone{req.ctx}, response) + if writeErr != nil && c.onResponseWriteError != nil { + c.onResponseWriteError(req.Request, writeErr) + } if err == nil { err = writeErr } diff --git a/mcp/transport.go b/mcp/transport.go index 4247bd80..a15d63e4 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -233,6 +233,9 @@ func connect[H handler, State any](ctx context.Context, t Transport, b binder[H, OnInternalError: func(err error) { logger.Error("jsonrpc2 internal error", "error", err) }, + OnResponseWriteError: func(req *jsonrpc2.Request, err error) { + logger.Warn("failed to write response", "method", req.Method, "id", req.ID.Raw(), "error", err) + }, PropagateCancellation: propagateCancellation, }) assert(preempter.conn != nil, "unbound preempter") diff --git a/mcp/transport_test.go b/mcp/transport_test.go index c2fb985d..3d97f532 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -11,10 +11,12 @@ import ( "errors" "fmt" "io" + "log/slog" "strings" "testing" "time" + "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/jsonrpc" ) @@ -132,6 +134,46 @@ func connectionClosedTests() []connectionClosedTest { } } +// TestServerLogsUndeliveredResponse is a regression test for #1147: when the +// client goes away before a tool handler finishes, the server logs that the +// response could not be delivered instead of dropping it silently. +func TestServerLogsUndeliveredResponse(t *testing.T) { + ctx := context.Background() + var buf bytes.Buffer + server := NewServer(testImpl, &ServerOptions{ + Logger: slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn})), + }) + started := make(chan struct{}) + server.AddTool(&Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, _ *CallToolRequest) (*CallToolResult, error) { + close(started) + <-ctx.Done() // cancelled once the client's connection is gone + return &CallToolResult{}, nil + }) + serverIn, clientOut := io.Pipe() + clientIn, serverOut := io.Pipe() + ss, err := server.Connect(ctx, &IOTransport{Reader: serverIn, Writer: serverOut}, nil) + if err != nil { + t.Fatal(err) + } + cs, err := NewClient(testImpl, nil).Connect(ctx, &IOTransport{Reader: clientIn, Writer: clientOut}, nil) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + go cs.CallTool(ctx, &CallToolParams{Name: "slow"}) + <-started + // The client goes away abruptly, as one that gives up on a slow call does. + clientOut.Close() + clientIn.Close() + ss.Wait() + + got := buf.String() + if !strings.Contains(got, `level=WARN msg="failed to write response" method=tools/call id=`) { + t.Errorf("server log does not report the undelivered response:\n%s", got) + } +} + func checkConnectionClosedError(t *testing.T, err error, causes ...error) { t.Helper() for _, want := range append([]error{ErrConnectionClosed}, causes...) {