diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 94da6ae3..36319acd 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -199,6 +199,21 @@ type Writer interface { Write(context.Context, Message) error } +// NonFatal marks an error returned by [Writer.Write] as a failure of the +// written message only: the connection remains usable. Unlike [ErrRejected], +// the caller receives err itself. +func NonFatal(err error) error { + if err == nil { + return nil + } + return &nonFatalError{err} +} + +type nonFatalError struct{ err error } + +func (e *nonFatalError) Error() string { return e.err.Error() } +func (e *nonFatalError) Unwrap() error { return e.err } + // A ConnectionConfig configures a bidirectional jsonrpc2 connection. type ConnectionConfig struct { Reader Reader // required @@ -769,6 +784,9 @@ func (c *Connection) write(ctx context.Context, msg Message) error { if err == nil { err = c.writer.Write(ctx, msg) } + if nf, ok := err.(*nonFatalError); ok { + return nf.err + } // For cancelled or rejected requests, we don't set the writeErr (which would // break the connection). They can just be returned to the caller. diff --git a/mcp/streamable.go b/mcp/streamable.go index a7ab2275..1080315c 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -2325,7 +2325,8 @@ type streamableClientConn struct { // Logical reads are distributed across multiple http requests. Whenever any // of them fails to process their response, we must break the connection, by - // failing the pending Read. + // failing the pending Read. Without a session, only that call is ended + // instead (see [streamableClientConn.endCall]). // // Achieve this by storing the failure message, and signalling when reads are // broken. See also [streamableClientConn.fail] and @@ -2437,6 +2438,14 @@ func (c *streamableClientConn) fail(err error) { } } +// endCall ends call with err, leaving the connection usable. +func (c *streamableClientConn) endCall(call *jsonrpc.Request, err error) { + select { + case c.incoming <- &jsonrpc2.Response{ID: call.ID, Error: err}: + case <-c.done: + } +} + func (c *streamableClientConn) failure() error { select { case <-c.failed: @@ -2568,6 +2577,8 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e // plain text 404 or 400), which checkResponse cannot classify as a // per-call rejection on its own. err = fmt.Errorf("%w: %w", err, jsonrpc2.ErrRejected) + } else if !errors.Is(err, jsonrpc2.ErrRejected) && c.SessionID() == "" { + return jsonrpc2.NonFatal(err) } else if !errors.Is(err, jsonrpc2.ErrRejected) { // Only fail the connection for non-transient errors. // Transient errors (wrapped with ErrRejected) should not break the connection. @@ -2612,7 +2623,7 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e contentType := baseMediaType(resp.Header.Get("Content-Type")) switch contentType { case "application/json": - go c.handleJSON(requestSummary, resp) + go c.handleJSON(requestSummary, resp, forCall) case "text/event-stream": var forCall *jsonrpc.Request @@ -2627,7 +2638,11 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e default: resp.Body.Close() - return fmt.Errorf("%s: unsupported content type %q", requestSummary, contentType) + err := fmt.Errorf("%s: unsupported content type %q", requestSummary, contentType) + if c.SessionID() == "" { + return jsonrpc2.NonFatal(err) + } + return err } return nil } @@ -2693,16 +2708,26 @@ func protocolVersionFromMessage(msg jsonrpc.Message) string { return v } -func (c *streamableClientConn) handleJSON(requestSummary string, resp *http.Response) { +func (c *streamableClientConn) handleJSON(requestSummary string, resp *http.Response, forCall *jsonrpc.Request) { body, err := io.ReadAll(resp.Body) resp.Body.Close() if err != nil { - c.fail(fmt.Errorf("%s: failed to read body: %v", requestSummary, err)) + err = fmt.Errorf("%s: failed to read body: %v", requestSummary, err) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } return } msg, err := jsonrpc.DecodeMessage(body) if err != nil { - c.fail(fmt.Errorf("%s: failed to decode response: %v", requestSummary, err)) + err = fmt.Errorf("%s: failed to decode response: %v", requestSummary, err) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } return } select { @@ -2751,7 +2776,12 @@ func (c *streamableClientConn) handleSSE(ctx context.Context, requestSummary str retriesWithoutProgress++ if retriesWithoutProgress > c.maxRetries { if ctx.Err() == nil { - c.fail(fmt.Errorf("%s: exceeded %d retries without progress (session ID: %v)", requestSummary, c.maxRetries, c.sessionID)) + err := fmt.Errorf("%s: exceeded %d retries without progress (session ID: %v)", requestSummary, c.maxRetries, c.sessionID) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } } return } @@ -2761,17 +2791,26 @@ func (c *streamableClientConn) handleSSE(ctx context.Context, requestSummary str newResp, err := c.connectSSE(ctx, lastEventID, reconnectDelay, false) if err != nil { // If the client didn't cancel this request, any failure to execute it - // breaks the logical MCP session. + // breaks the logical MCP session, if there is one. if ctx.Err() == nil { - // All reconnection attempts failed: fail the connection. - c.fail(fmt.Errorf("%s: failed to reconnect (session ID: %v): %v", requestSummary, c.sessionID, err)) + // All reconnection attempts failed. + err = fmt.Errorf("%s: failed to reconnect (session ID: %v): %v", requestSummary, c.sessionID, err) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } } return } resp = newResp if err := c.checkResponse(ctx, requestSummary, resp); err != nil { - c.fail(err) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } return } } @@ -2834,9 +2873,14 @@ func (c *streamableClientConn) processStream(ctx context.Context, requestSummary } // Malformed events are hard errors that indicate corrupted data or protocol - // violations. These should fail the connection permanently. + // violations, so the stream is not resumed. if errors.Is(err, errMalformedEvent) { - c.fail(fmt.Errorf("%s: %v", requestSummary, err)) + err = fmt.Errorf("%s: %v", requestSummary, err) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } return "", 0, true } @@ -2868,7 +2912,12 @@ func (c *streamableClientConn) processStream(ctx context.Context, requestSummary msg, err := jsonrpc.DecodeMessage(evt.Data) if err != nil { - c.fail(fmt.Errorf("%s: failed to decode event: %v", requestSummary, err)) + err = fmt.Errorf("%s: failed to decode event: %v", requestSummary, err) + if forCall != nil && c.SessionID() == "" { + c.endCall(forCall, err) + } else { + c.fail(err) + } return "", 0, true } @@ -2897,14 +2946,7 @@ func (c *streamableClientConn) processStream(ctx context.Context, requestSummary // Note that this is different from the cancellation case above, since the // caller is still waiting for a response that will never come. if lastEventID == "" && forCall != nil { - errmsg := &jsonrpc2.Response{ - ID: forCall.ID, - Error: fmt.Errorf("request terminated without response"), - } - select { - case c.incoming <- errmsg: - case <-c.done: - } + c.endCall(forCall, fmt.Errorf("request terminated without response")) } return lastEventID, reconnectDelay, false } diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index ae2319b2..89653f2a 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -1854,24 +1854,23 @@ func TestStreamableClientHandlerErrorPropagation(t *testing.T) { }, { // Same server response, but the user opted out via - // MCPGODEBUG=noprotocolerrorbody=1. Pre-fix behavior: the - // session is torn down. + // MCPGODEBUG=noprotocolerrorbody=1. The body is ignored, but + // without a session the failure still fails only the call. name: "400 with JSON-RPC error body (noprotocolerrorbody=1)", callStatus: http.StatusBadRequest, callBody: jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"), disableBodyDecode: true, wantErrRejected: false, - wantSessionAlive: false, + wantSessionAlive: true, }, { - // Legacy server returns plain-text 400 with no JSON-RPC - // body: cannot be classified as a per-call rejection, so the - // session is still torn down. + // Plain-text 400 with no JSON-RPC body: without a session, + // only the call fails. name: "400 with plain-text body", callStatus: http.StatusBadRequest, callBody: "Bad Request", wantErrRejected: false, - wantSessionAlive: false, + wantSessionAlive: true, }, { // SEP-2575: server returns -32601 for an unimplemented @@ -1884,13 +1883,13 @@ func TestStreamableClientHandlerErrorPropagation(t *testing.T) { wantSessionAlive: true, }, { - // A bare 404 with no JSON-RPC body means the session has - // been terminated. Session must not survive. - name: "404 with empty body (terminated session)", + // A bare 404 with no JSON-RPC body: there is no session to + // terminate, so only the call fails. + name: "404 with empty body", callStatus: http.StatusNotFound, callBody: "", wantErrRejected: false, - wantSessionAlive: false, + wantSessionAlive: true, }, }