Explorar el Código

fix: preserve relay body and capture invalid responses body

master
fengsilin hace 2 días
padre
commit
e41f2951f7
Se han modificado 5 ficheros con 84 adiciones y 4 borrados
  1. +6
    -2
      relay/channel/openai/chat_via_responses.go
  2. +3
    -1
      relay/channel/openai/relay_responses.go
  3. +40
    -0
      relay/channel/openai/upstream_body_test.go
  4. +11
    -0
      relay/common/request_identity.go
  5. +24
    -1
      relay/common/request_identity_test.go

+ 6
- 2
relay/channel/openai/chat_via_responses.go Ver fichero

@@ -52,7 +52,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
} }


if err := common.Unmarshal(body, &responsesResp); err != nil { 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 != "" { 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) chatId := helper.GetResponseID(c)
chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId) chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId)
if err != nil { 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 { if usage == nil || usage.TotalTokens == 0 {


+ 3
- 1
relay/channel/openai/relay_responses.go Ver fichero

@@ -28,7 +28,9 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
} }
err = common.Unmarshal(responseBody, &responsesResponse) err = common.Unmarshal(responseBody, &responsesResponse)
if err != nil { 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 != "" { if oaiError := responsesResponse.GetOpenAIError(); oaiError != nil && oaiError.Type != "" {
apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode) apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode)


+ 40
- 0
relay/channel/openai/upstream_body_test.go Ver fichero

@@ -1,10 +1,16 @@
package openai package openai


import ( import (
"io"
"net/http"
"net/http/httptest"
"strings"
"testing" "testing"


"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto" "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/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@@ -62,3 +68,37 @@ func TestResponsesStreamEventErrorParsing(t *testing.T) {
assert.Nil(t, streamResp.Error) 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":`)
}

+ 11
- 0
relay/common/request_identity.go Ver fichero

@@ -1,6 +1,7 @@
package common package common


import ( import (
"io"
"net/http" "net/http"
"strings" "strings"


@@ -53,14 +54,24 @@ func ExtractTopLevelChatID(c *gin.Context) (string, error) {
if c == nil { if c == nil {
return "", 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) bodyStorage, err := commonpkg.GetBodyStorage(c)
if err != nil { if err != nil {
return "", err return "", err
} }
bodyBytes, err := bodyStorage.Bytes() bodyBytes, err := bodyStorage.Bytes()
if err != nil || len(bodyBytes) == 0 { if err != nil || len(bodyBytes) == 0 {
if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr == nil {
c.Request.Body = io.NopCloser(bodyStorage)
}
return "", err return "", err
} }
if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr != nil {
return "", seekErr
}
c.Request.Body = io.NopCloser(bodyStorage)
var payload map[string]interface{} var payload map[string]interface{}
if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil {
return "", nil return "", nil


+ 24
- 1
relay/common/request_identity_test.go Ver fichero

@@ -2,6 +2,7 @@ package common


import ( import (
"bytes" "bytes"
"io"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"testing" "testing"
@@ -14,12 +15,17 @@ func TestExtractChatIDFromReusableBody(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
w := httptest.NewRecorder() w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w) 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") c.Request.Header.Set("Content-Type", "application/json")


chatID, err := ExtractTopLevelChatID(c) chatID, err := ExtractTopLevelChatID(c)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "chat_abc", chatID) 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) { func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) {
@@ -87,3 +93,20 @@ func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) {
CaptureUpstreamIDFromHTTPResponse(info, resp) CaptureUpstreamIDFromHTTPResponse(info, resp)
require.Equal(t, "ws_upstream_1", info.UpstreamID) 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))
}

Cargando…
Cancelar
Guardar