From 9c08d8cf2068e4eb468593fd219e79587ff4ec4e Mon Sep 17 00:00:00 2001 From: CaIon <1808837298@qq.com> Date: Fri, 27 Jun 2025 22:36:12 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20introduce=20IOCopyBytesGrac?= =?UTF-8?q?efully=20function=20for=20streamlined=20response=20body=20handl?= =?UTF-8?q?ing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- common/http.go | 45 +++++++++++++++++- relay/channel/openai/relay-openai.go | 70 ++++++---------------------- 2 files changed, 59 insertions(+), 56 deletions(-) diff --git a/common/http.go b/common/http.go index db8579d0..315b19af 100644 --- a/common/http.go +++ b/common/http.go @@ -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())) + } +} diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 191e5ec6..05112f84 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -251,22 +251,8 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI responseBody = claudeRespStr } - // Reset response body - resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) - // 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 resp.Header { - c.Writer.Header().Set(k, v[0]) - } - c.Writer.WriteHeader(resp.StatusCode) - _, err = io.Copy(c.Writer, resp.Body) - if err != nil { - //return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil - common.SysError("error copying response body: " + err.Error()) - } - common.CloseResponseBodyGracefully(resp) + common.IOCopyBytesGracefully(c, resp, responseBody) + return nil, &simpleResponse.Usage } @@ -304,24 +290,9 @@ func OpenaiSTTHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil } common.CloseResponseBodyGracefully(resp) - if err != nil { - return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil - } - // Reset response body - resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) - // 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 resp.Header { - c.Writer.Header().Set(k, v[0]) - } - c.Writer.WriteHeader(resp.StatusCode) - _, err = io.Copy(c.Writer, resp.Body) - if err != nil { - return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil - } - common.CloseResponseBodyGracefully(resp) + + // 写入新的 response body + common.IOCopyBytesGracefully(c, resp, responseBody) usage := &dto.Usage{} usage.PromptTokens = audioTokens @@ -602,33 +573,22 @@ func OpenaiHandlerWithUsage(c *gin.Context, resp *http.Response, info *relaycomm if err != nil { return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil } - common.CloseResponseBodyGracefully(resp) - // Reset response body - resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) - // 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 resp.Header { - c.Writer.Header().Set(k, v[0]) - } - // reset content length - c.Writer.Header().Set("Content-Length", fmt.Sprintf("%d", len(responseBody))) - c.Writer.WriteHeader(resp.StatusCode) - _, err = io.Copy(c.Writer, resp.Body) - if err != nil { - common.SysError("error copying response body: " + err.Error()) - } - common.CloseResponseBodyGracefully(resp) - // Once we've written to the client, we should not return errors anymore - // because the upstream has already consumed resources and returned content - // We should still perform billing even if parsing fails var usageResp dto.SimpleResponse err = json.Unmarshal(responseBody, &usageResp) if err != nil { return service.OpenAIErrorWrapper(err, "parse_response_body_failed", http.StatusInternalServerError), nil } + + // 关闭旧的 response body（已被读取，再次读取会导致错误） + common.CloseResponseBodyGracefully(resp) + + // 写入新的 response body + common.IOCopyBytesGracefully(c, resp, responseBody) + + // Once we've written to the client, we should not return errors anymore + // because the upstream has already consumed resources and returned content + // We should still perform billing even if parsing fails // format if usageResp.InputTokens > 0 { usageResp.PromptTokens += usageResp.InputTokens