Retain successful response read failures

internal/llm/openai_compatible_client.go

@@ -370,40 +370,69 @@ func tokenCount(value *int) int {
 	return *value
 }
 
+type responseErrorReader struct {
+	reader io.Reader
+	err    error
+}
+
+func (r *responseErrorReader) Read(p []byte) (int, error) {
+	n, err := r.reader.Read(p)
+	if err != nil && !errors.Is(err, io.EOF) {
+		r.err = err
+	}
+	return n, err
+}
+
 func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {
+	observed := &responseErrorReader{reader: body}
 	limited := &io.LimitedReader{
-		R: body,
+		R: observed,
 		N: maxOpenAIChatResponseBytes + 1,
 	}
 	decoder := json.NewDecoder(limited)
 
 	var response openAIChatResponse
 	if err := decoder.Decode(&response); err != nil {
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			return openAIChatResponse{}, &requestFailedError{cause: err}
+		if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
+			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
 		}
 		if limited.N == 0 {
 			return openAIChatResponse{}, openAIChatResponseTooLargeError()
 		}
+		if observed.err != nil {
+			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
+		}
 		return openAIChatResponse{}, fmt.Errorf("%w: failed to decode response", ErrMalformedResponse)
 	}
+	if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
+		return openAIChatResponse{}, &requestFailedError{cause: observed.err}
+	}
 	if limited.N == 0 {
 		return openAIChatResponse{}, openAIChatResponseTooLargeError()
 	}
 
 	var trailing any
 	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			return openAIChatResponse{}, &requestFailedError{cause: err}
+		if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
+			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
 		}
 		if limited.N == 0 {
 			return openAIChatResponse{}, openAIChatResponseTooLargeError()
 		}
+		if observed.err != nil {
+			return openAIChatResponse{}, &requestFailedError{cause: observed.err}
+		}
 		return openAIChatResponse{}, fmt.Errorf("%w: response contains trailing data", ErrMalformedResponse)
 	}
+	if errors.Is(observed.err, context.Canceled) || errors.Is(observed.err, context.DeadlineExceeded) {
+		return openAIChatResponse{}, &requestFailedError{cause: observed.err}
+	}
 	if limited.N == 0 {
 		return openAIChatResponse{}, openAIChatResponseTooLargeError()
 	}
+	if observed.err != nil {
+		return openAIChatResponse{}, &requestFailedError{cause: observed.err}
+	}
 
 	return response, nil
 }

internal/llm/response_read_error_test.go

@@ -0,0 +1,48 @@
+package llm
+
+import (
+	"errors"
+	"io"
+	"strings"
+	"testing"
+)
+
+type responseBytesAndError struct {
+	data []byte
+	err  error
+}
+
+func (r *responseBytesAndError) Read(p []byte) (int, error) {
+	if len(r.data) == 0 {
+		return 0, r.err
+	}
+	n := copy(p, r.data)
+	r.data = r.data[n:]
+	if len(r.data) == 0 {
+		return n, r.err
+	}
+	return n, nil
+}
+
+func TestSuccessfulResponseReadFailuresRetainCause(t *testing.T) {
+	cause := errors.New("private reader diagnostic")
+	for _, test := range []struct {
+		name, prefix string
+		together     bool
+	}{
+		{"partial", `{"choices":[`, false},
+		{"complete then failure", successfulProviderResponse, false},
+		{"bytes and error together", successfulProviderResponse, true},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			var reader io.Reader = io.MultiReader(strings.NewReader(test.prefix), failingReader{err: cause})
+			if test.together {
+				reader = &responseBytesAndError{data: []byte(test.prefix), err: cause}
+			}
+			_, err := decodeOpenAIChatResponse(reader)
+			if !errors.Is(err, ErrTransport) || !errors.Is(err, cause) || errors.Is(err, ErrMalformedResponse) || strings.Contains(err.Error(), cause.Error()) {
+				t.Fatalf("read failure = %v", err)
+			}
+		})
+	}
+}