From e41f2951f7439dad189437857dc7f67579c31caa Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 6 May 2026 09:34:40 +0800 Subject: [PATCH] fix: preserve relay body and capture invalid responses body --- relay/channel/openai/chat_via_responses.go | 8 +++-- relay/channel/openai/relay_responses.go | 4 ++- relay/channel/openai/upstream_body_test.go | 40 ++++++++++++++++++++++ relay/common/request_identity.go | 11 ++++++ relay/common/request_identity_test.go | 25 +++++++++++++- 5 files changed, 84 insertions(+), 4 deletions(-) diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 3aaa7c5..a93c390 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -52,7 +52,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } if err := common.Unmarshal(body, &responsesResp); err != nil { - return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + apiErr := types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + apiErr.UpstreamBody = service.TruncateBody(string(body)) + return nil, apiErr } if oaiError := responsesResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { @@ -64,7 +66,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp chatId := helper.GetResponseID(c) chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId) if err != nil { - return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + apiErr := types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + apiErr.UpstreamBody = service.TruncateBody(string(body)) + return nil, apiErr } if usage == nil || usage.TotalTokens == 0 { diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 14107f5..1b88457 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -28,7 +28,9 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } err = common.Unmarshal(responseBody, &responsesResponse) if err != nil { - return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + apiErr := types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + apiErr.UpstreamBody = service.TruncateBody(string(responseBody)) + return nil, apiErr } if oaiError := responsesResponse.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode) diff --git a/relay/channel/openai/upstream_body_test.go b/relay/channel/openai/upstream_body_test.go index 3e6dbae..a43738b 100644 --- a/relay/channel/openai/upstream_body_test.go +++ b/relay/channel/openai/upstream_body_test.go @@ -1,10 +1,16 @@ package openai import ( + "io" + "net/http" + "net/http/httptest" + "strings" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -62,3 +68,37 @@ func TestResponsesStreamEventErrorParsing(t *testing.T) { assert.Nil(t, streamResp.Error) }) } + +func TestOaiResponsesHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"incomplete":`)), + } + + usage, err := OaiResponsesHandler(c, &relaycommon.RelayInfo{}, resp) + require.Nil(t, usage) + require.Error(t, err) + assert.NotEmpty(t, err.UpstreamBody) + assert.Contains(t, err.UpstreamBody, `{"incomplete":`) +} + +func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"incomplete":`)), + } + + usage, err := OaiResponsesToChatHandler(c, &relaycommon.RelayInfo{}, resp) + require.Nil(t, usage) + require.Error(t, err) + assert.NotEmpty(t, err.UpstreamBody) + assert.Contains(t, err.UpstreamBody, `{"incomplete":`) +} diff --git a/relay/common/request_identity.go b/relay/common/request_identity.go index 513199f..824941e 100644 --- a/relay/common/request_identity.go +++ b/relay/common/request_identity.go @@ -1,6 +1,7 @@ package common import ( + "io" "net/http" "strings" @@ -53,14 +54,24 @@ func ExtractTopLevelChatID(c *gin.Context) (string, error) { if c == nil { return "", nil } + if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(c.Request.Header.Get("Content-Type"))), "application/json") { + return "", nil + } bodyStorage, err := commonpkg.GetBodyStorage(c) if err != nil { return "", err } bodyBytes, err := bodyStorage.Bytes() if err != nil || len(bodyBytes) == 0 { + if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr == nil { + c.Request.Body = io.NopCloser(bodyStorage) + } return "", err } + if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr != nil { + return "", seekErr + } + c.Request.Body = io.NopCloser(bodyStorage) var payload map[string]interface{} if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { return "", nil diff --git a/relay/common/request_identity_test.go b/relay/common/request_identity_test.go index 47040c1..e01a44a 100644 --- a/relay/common/request_identity_test.go +++ b/relay/common/request_identity_test.go @@ -2,6 +2,7 @@ package common import ( "bytes" + "io" "net/http" "net/http/httptest" "testing" @@ -14,12 +15,17 @@ func TestExtractChatIDFromReusableBody(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}`)) + body := `{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}` + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(body)) c.Request.Header.Set("Content-Type", "application/json") chatID, err := ExtractTopLevelChatID(c) require.NoError(t, err) require.Equal(t, "chat_abc", chatID) + + bodyAfter, err := io.ReadAll(c.Request.Body) + require.NoError(t, err) + require.Equal(t, body, string(bodyAfter)) } func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) { @@ -87,3 +93,20 @@ func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) { CaptureUpstreamIDFromHTTPResponse(info, resp) require.Equal(t, "ws_upstream_1", info.UpstreamID) } + +func TestExtractTopLevelChatIDSkipsNonJSONBody(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := "--boundary\r\nContent-Disposition: form-data; name=\"chat_id\"\r\n\r\nchat_abc\r\n--boundary--\r\n" + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "multipart/form-data; boundary=boundary") + + chatID, err := ExtractTopLevelChatID(c) + require.NoError(t, err) + require.Empty(t, chatID) + + bodyAfter, err := io.ReadAll(c.Request.Body) + require.NoError(t, err) + require.Equal(t, body, string(bodyAfter)) +}