From 14aa6cc7e809525ffb6b40b55c9b40b94bceb296 Mon Sep 17 00:00:00 2001
From: hkfires <10558748+hkfires@users.noreply.github.com>
Date: Mon, 15 Dec 2025 17:45:16 +0800
Subject: [PATCH] fix(api): ensure all response writes are captured for logging

The response writer wrapper has been refactored to more reliably capture response bodies for logging, fixing several edge cases.

- Implements `WriteString` to capture writes from `io.StringWriter`, which were previously missed by the `Write` method override.
- A new `shouldBufferResponseBody` helper centralizes the logic to ensure the body is buffered only when logging is active or for errors when `logOnErrorOnly` is enabled.
- Streaming detection is now more robust. It correctly handles non-streaming error responses (e.g., `application/json`) that are generated for a request that was intended to be streaming.

BREAKING CHANGE: The public methods `Status()`, `Size()`, and `Written()` have been removed from the `ResponseWriterWrapper` as they are no longer required by the new implementation.
---
 internal/api/middleware/response_writer.go | 95 ++++++++++++++--------
 1 file changed, 59 insertions(+), 36 deletions(-)

diff --git a/internal/api/middleware/response_writer.go b/internal/api/middleware/response_writer.go
index b7259bc6..e70556d0 100644
--- a/internal/api/middleware/response_writer.go
+++ b/internal/api/middleware/response_writer.go
@@ -71,22 +71,64 @@ func (w *ResponseWriterWrapper) Write(data []byte) (int, error) {
 	n, err := w.ResponseWriter.Write(data)
 
 	// THEN: Handle logging based on response type
-	if w.isStreaming {
+	if w.isStreaming && w.chunkChannel != nil {
 		// For streaming responses: Send to async logging channel (non-blocking)
-		if w.chunkChannel != nil {
-			select {
-			case w.chunkChannel <- append([]byte(nil), data...): // Non-blocking send with copy
-			default: // Channel full, skip logging to avoid blocking
-			}
+		select {
+		case w.chunkChannel <- append([]byte(nil), data...): // Non-blocking send with copy
+		default: // Channel full, skip logging to avoid blocking
 		}
-	} else {
-		// For non-streaming responses: Buffer complete response
+		return n, err
+	}
+
+	if w.shouldBufferResponseBody() {
 		w.body.Write(data)
 	}
 
 	return n, err
 }
 
+func (w *ResponseWriterWrapper) shouldBufferResponseBody() bool {
+	if w.logger != nil && w.logger.IsEnabled() {
+		return true
+	}
+	if !w.logOnErrorOnly {
+		return false
+	}
+	status := w.statusCode
+	if status == 0 {
+		if statusWriter, ok := w.ResponseWriter.(interface{ Status() int }); ok && statusWriter != nil {
+			status = statusWriter.Status()
+		} else {
+			status = http.StatusOK
+		}
+	}
+	return status >= http.StatusBadRequest
+}
+
+// WriteString wraps the underlying ResponseWriter's WriteString method to capture response data.
+// Some handlers (and fmt/io helpers) write via io.StringWriter; without this override, those writes
+// bypass Write() and would be missing from request logs.
+func (w *ResponseWriterWrapper) WriteString(data string) (int, error) {
+	w.ensureHeadersCaptured()
+
+	// CRITICAL: Write to client first (zero latency)
+	n, err := w.ResponseWriter.WriteString(data)
+
+	// THEN: Capture for logging
+	if w.isStreaming && w.chunkChannel != nil {
+		select {
+		case w.chunkChannel <- []byte(data):
+		default:
+		}
+		return n, err
+	}
+
+	if w.shouldBufferResponseBody() {
+		w.body.WriteString(data)
+	}
+	return n, err
+}
+
 // WriteHeader wraps the underlying ResponseWriter's WriteHeader method.
 // It captures the status code, detects if the response is streaming based on the Content-Type header,
 // and initializes the appropriate logging mechanism (standard or streaming).
@@ -160,12 +202,16 @@ func (w *ResponseWriterWrapper) detectStreaming(contentType string) bool {
 		return true
 	}
 
-	// Check request body for streaming indicators
-	if w.requestInfo.Body != nil {
+	// If a concrete Content-Type is already set (e.g., application/json for error responses),
+	// treat it as non-streaming instead of inferring from the request payload.
+	if strings.TrimSpace(contentType) != "" {
+		return false
+	}
+
+	// Only fall back to request payload hints when Content-Type is not set yet.
+	if w.requestInfo != nil && len(w.requestInfo.Body) > 0 {
 		bodyStr := string(w.requestInfo.Body)
-		if strings.Contains(bodyStr, `"stream": true`) || strings.Contains(bodyStr, `"stream":true`) {
-			return true
-		}
+		return strings.Contains(bodyStr, `"stream": true`) || strings.Contains(bodyStr, `"stream":true`)
 	}
 
 	return false
@@ -335,26 +381,3 @@ func (w *ResponseWriterWrapper) logRequest(statusCode int, headers map[string][]
 		apiResponseErrors,
 	)
 }
-
-// Status returns the HTTP response status code captured by the wrapper.
-// It defaults to 200 if WriteHeader has not been called.
-func (w *ResponseWriterWrapper) Status() int {
-	if w.statusCode == 0 {
-		return 200 // Default status code
-	}
-	return w.statusCode
-}
-
-// Size returns the size of the response body in bytes for non-streaming responses.
-// For streaming responses, it returns -1, as the total size is unknown.
-func (w *ResponseWriterWrapper) Size() int {
-	if w.isStreaming {
-		return -1 // Unknown size for streaming responses
-	}
-	return w.body.Len()
-}
-
-// Written returns true if the response header has been written (i.e., a status code has been set).
-func (w *ResponseWriterWrapper) Written() bool {
-	return w.statusCode != 0
-}