feat: introduce IOCopyBytesGracefully function for streamlined response body handling

This update adds the IOCopyBytesGracefully function to the common package, which simplifies the process of copying response bodies in the OpenAI handlers. It enhances error handling and ensures proper resource management by encapsulating the logic for setting headers and writing response data. The OpenAI handlers have been refactored to utilize this new function, improving code clarity and maintainability.
This commit is contained in:
CaIon
2025-06-27 22:36:12 +08:00
parent 281054ff4c
commit 9c08d8cf20
2 changed files with 59 additions and 56 deletions

View File

@@ -1,6 +1,12 @@
package common
import "net/http"
import (
"bytes"
"fmt"
"github.com/gin-gonic/gin"
"io"
"net/http"
)
func CloseResponseBodyGracefully(httpResponse *http.Response) {
if httpResponse == nil || httpResponse.Body == nil {
@@ -11,3 +17,40 @@ func CloseResponseBodyGracefully(httpResponse *http.Response) {
SysError("failed to close response body: " + err.Error())
}
}
func IOCopyBytesGracefully(c *gin.Context, src *http.Response, data []byte) {
if src == nil || src.Body == nil {
return
}
defer CloseResponseBodyGracefully(src)
if c.Writer == nil {
return
}
src.Body = io.NopCloser(bytes.NewBuffer(data))
// We shouldn't set the header before we parse the response body, because the parse part may fail.
// And then we will have to send an error response, but in this case, the header has already been set.
// So the httpClient will be confused by the response.
// For example, Postman will report error, and we cannot check the response at all.
for k, v := range src.Header {
// avoid setting Content-Length
if k == "Content-Length" {
continue
}
c.Writer.Header().Set(k, v[0])
}
// set Content-Length header manually
c.Writer.Header().Set("Content-Length", fmt.Sprintf("%d", len(data)))
c.Writer.WriteHeader(src.StatusCode)
c.Writer.WriteHeaderNow()
_, err := io.Copy(c.Writer, src.Body)
if err != nil {
LogError(c, fmt.Sprintf("failed to copy response body: %s", err.Error()))
}
}