diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index b81a148..66f794a 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1285,6 +1285,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http var usage = &dto.Usage{} var imageCount int responseText := strings.Builder{} + var streamErr *types.NewAPIError helper.StreamScannerHandler(c, resp, info, func(data string) bool { var geminiResponse dto.GeminiChatResponse @@ -1296,6 +1297,13 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) + streamErr = types.NewOpenAIError( + errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason), + types.ErrorCodePromptBlocked, + http.StatusBadRequest, + ) + streamErr.UpstreamBody = service.TruncateBody(data) + return false } // 统计图片数量 @@ -1318,6 +1326,10 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http return callback(data, &geminiResponse) }) + if streamErr != nil { + helper.ClearEventStreamHeadersIfNotWritten(c) + return nil, streamErr + } if imageCount != 0 { if usage.CompletionTokens == 0 { diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index c8f9f83..7d111bc 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -282,6 +282,52 @@ func TestGeminiStreamHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *t require.Equal(t, 110, usage.TotalTokens) } +func TestGeminiStreamHandlerReturnsErrorOnPromptBlock(t *testing.T) { + gin.SetMode(gin.TestMode) + oldStreamingTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 300 + t.Cleanup(func() { + constant.StreamingTimeout = oldStreamingTimeout + }) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + blockReason := "SAFETY" + payload := dto.GeminiChatResponse{ + PromptFeedback: &dto.GeminiChatPromptFeedback{ + BlockReason: &blockReason, + }, + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(body) + "\n\n"))), + } + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gemini-test", + }, + } + + called := false + usage, newAPIError := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { + called = true + return true + }) + + require.Nil(t, usage) + require.NotNil(t, newAPIError) + require.False(t, called) + require.Empty(t, w.Body.String()) + require.Empty(t, w.Header().Get("Content-Type")) + require.Contains(t, newAPIError.UpstreamBody, "SAFETY") +} + func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { t.Parallel() diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 08811a7..30abee6 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -2,6 +2,7 @@ package openai import ( "encoding/json" + "net/http" "strings" "github.com/QuantumNous/new-api/common" @@ -19,6 +20,23 @@ import ( ) // 辅助函数 +func streamOpenAIErrorFromData(data string, statusCode int) *types.NewAPIError { + var errorResp dto.GeneralErrorResponse + if err := common.UnmarshalJsonStr(data, &errorResp); err != nil { + return nil + } + openaiErr := errorResp.TryToOpenAIError() + if openaiErr == nil { + return nil + } + if statusCode < 100 || statusCode > 599 { + statusCode = http.StatusInternalServerError + } + apiErr := types.WithOpenAIError(*openaiErr, statusCode) + apiErr.UpstreamBody = service.TruncateBody(data) + return apiErr +} + func HandleStreamFormat(c *gin.Context, info *relaycommon.RelayInfo, data string, forceFormat bool, thinkToContent bool) error { info.SendResponseCount++ diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 1ef6166..d20934c 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -121,12 +121,22 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re var usage = &dto.Usage{} var streamItems []string // store stream items var lastStreamData string + var streamErr *types.NewAPIError var secondLastStreamData string // 存储倒数第二个stream data，用于音频模型 // 检查是否为音频模型 isAudioModel := strings.Contains(strings.ToLower(model), "audio") helper.StreamScannerHandler(c, resp, info, func(data string) bool { + if apiErr := streamOpenAIErrorFromData(data, resp.StatusCode); apiErr != nil { + if !c.Writer.Written() { + streamErr = apiErr + } else { + logger.LogError(c, "upstream stream error after downstream response was written: "+apiErr.Error()) + } + return false + } + if lastStreamData != "" { err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) if err != nil { @@ -146,6 +156,11 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re }) // 对音频模型，从倒数第二个stream data中提取usage信息 + if streamErr != nil { + helper.ClearEventStreamHeadersIfNotWritten(c) + return nil, streamErr + } + if isAudioModel && secondLastStreamData != "" { var streamResp struct { Usage *dto.Usage `json:"usage"` diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 152c7ef..e5fed78 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -90,6 +90,12 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp var streamResponse dto.ResponsesStreamResponse if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { + switch streamResponse.Type { + case "response.error", "response.failed", "error": + streamErr = handleResponsesStreamError(streamResponse, data) + return false + } + sendResponsesStreamData(c, streamResponse, data) switch streamResponse.Type { case "response.completed": @@ -128,9 +134,6 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } } } - case "response.error", "response.failed", "error": - streamErr = handleResponsesStreamError(streamResponse, data) - return false } } else { logger.LogError(c, "failed to unmarshal stream response: "+err.Error()) @@ -139,6 +142,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp }) if streamErr != nil { + helper.ClearEventStreamHeadersIfNotWritten(c) return nil, streamErr } diff --git a/relay/channel/openai/upstream_body_test.go b/relay/channel/openai/upstream_body_test.go index a43738b..a91c82c 100644 --- a/relay/channel/openai/upstream_body_test.go +++ b/relay/channel/openai/upstream_body_test.go @@ -8,8 +8,10 @@ import ( "testing" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -86,6 +88,66 @@ func TestOaiResponsesHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { assert.Contains(t, err.UpstreamBody, `{"incomplete":`) } +func TestOaiResponsesStreamHandlerDoesNotForwardErrorEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { + constant.StreamingTimeout = oldTimeout + }) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + + errorEvent := `{"type":"error","error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null},"sequence_number":2}` + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), + } + + usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp) + require.Nil(t, usage) + require.Error(t, err) + assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") + assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") + assert.Empty(t, w.Body.String()) + assert.Empty(t, w.Header().Get("Content-Type")) +} + +func TestOaiStreamHandlerDoesNotForwardInitialErrorEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { + constant.StreamingTimeout = oldTimeout + }) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}` + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), + } + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gpt-test", + }, + } + + usage, err := OaiStreamHandler(c, info, resp) + require.Nil(t, usage) + require.Error(t, err) + assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") + assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") + assert.Empty(t, w.Body.String()) + assert.Empty(t, w.Header().Get("Content-Type")) +} + func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() diff --git a/relay/helper/common.go b/relay/helper/common.go index 17ce79d..8fffb5e 100644 --- a/relay/helper/common.go +++ b/relay/helper/common.go @@ -54,6 +54,19 @@ func SetEventStreamHeaders(c *gin.Context) { c.Writer.Header().Set("X-Accel-Buffering", "no") } +func ClearEventStreamHeadersIfNotWritten(c *gin.Context) { + if c == nil || c.Writer == nil || c.Writer.Written() { + return + } + c.Set("event_stream_headers_set", false) + header := c.Writer.Header() + header.Del("Content-Type") + header.Del("Cache-Control") + header.Del("Connection") + header.Del("Transfer-Encoding") + header.Del("X-Accel-Buffering") +} + func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { jsonData, err := common.Marshal(resp) if err != nil {