- Gemini stream: 捕获 promptFeedback.blockReason 后停止流并返回 OpenAPI 错误 - OpenAI stream: 检测 SSE 帧中的 error 响应,在首帧写入前转换为 API 错误 - Responses stream: 处理 response.error / response.failed / error 类型帧 - 新增 ClearEventStreamHeadersIfNotWritten 工具方法,错误转换时清除流式头 - 新增 streamOpenAIErrorFromData,从 SSE data 帧提取 OpenAPI 错误结构 - 上游错误体会截断后存入 UpstreamBody 供排障 Co-Authored-By: Claude <noreply@anthropic.com>master
| @@ -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 { | |||
| @@ -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() | |||
| @@ -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++ | |||
| @@ -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"` | |||
| @@ -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 | |||
| } | |||
| @@ -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() | |||
| @@ -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 { | |||