- 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 usage = &dto.Usage{} | ||||
| var imageCount int | var imageCount int | ||||
| responseText := strings.Builder{} | responseText := strings.Builder{} | ||||
| var streamErr *types.NewAPIError | |||||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | helper.StreamScannerHandler(c, resp, info, func(data string) bool { | ||||
| var geminiResponse dto.GeminiChatResponse | 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 { | 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)) | 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) | return callback(data, &geminiResponse) | ||||
| }) | }) | ||||
| if streamErr != nil { | |||||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||||
| return nil, streamErr | |||||
| } | |||||
| if imageCount != 0 { | if imageCount != 0 { | ||||
| if usage.CompletionTokens == 0 { | if usage.CompletionTokens == 0 { | ||||
| @@ -282,6 +282,52 @@ func TestGeminiStreamHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *t | |||||
| require.Equal(t, 110, usage.TotalTokens) | 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) { | func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { | ||||
| t.Parallel() | t.Parallel() | ||||
| @@ -2,6 +2,7 @@ package openai | |||||
| import ( | import ( | ||||
| "encoding/json" | "encoding/json" | ||||
| "net/http" | |||||
| "strings" | "strings" | ||||
| "github.com/QuantumNous/new-api/common" | "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 { | func HandleStreamFormat(c *gin.Context, info *relaycommon.RelayInfo, data string, forceFormat bool, thinkToContent bool) error { | ||||
| info.SendResponseCount++ | info.SendResponseCount++ | ||||
| @@ -121,12 +121,22 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re | |||||
| var usage = &dto.Usage{} | var usage = &dto.Usage{} | ||||
| var streamItems []string // store stream items | var streamItems []string // store stream items | ||||
| var lastStreamData string | var lastStreamData string | ||||
| var streamErr *types.NewAPIError | |||||
| var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 | var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 | ||||
| // 检查是否为音频模型 | // 检查是否为音频模型 | ||||
| isAudioModel := strings.Contains(strings.ToLower(model), "audio") | isAudioModel := strings.Contains(strings.ToLower(model), "audio") | ||||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | 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 != "" { | if lastStreamData != "" { | ||||
| err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) | err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) | ||||
| if err != nil { | if err != nil { | ||||
| @@ -146,6 +156,11 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re | |||||
| }) | }) | ||||
| // 对音频模型,从倒数第二个stream data中提取usage信息 | // 对音频模型,从倒数第二个stream data中提取usage信息 | ||||
| if streamErr != nil { | |||||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||||
| return nil, streamErr | |||||
| } | |||||
| if isAudioModel && secondLastStreamData != "" { | if isAudioModel && secondLastStreamData != "" { | ||||
| var streamResp struct { | var streamResp struct { | ||||
| Usage *dto.Usage `json:"usage"` | Usage *dto.Usage `json:"usage"` | ||||
| @@ -90,6 +90,12 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||||
| var streamResponse dto.ResponsesStreamResponse | var streamResponse dto.ResponsesStreamResponse | ||||
| if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { | 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) | sendResponsesStreamData(c, streamResponse, data) | ||||
| switch streamResponse.Type { | switch streamResponse.Type { | ||||
| case "response.completed": | 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 { | } else { | ||||
| logger.LogError(c, "failed to unmarshal stream response: "+err.Error()) | 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 { | if streamErr != nil { | ||||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||||
| return nil, streamErr | return nil, streamErr | ||||
| } | } | ||||
| @@ -8,8 +8,10 @@ import ( | |||||
| "testing" | "testing" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/constant" | |||||
| "github.com/QuantumNous/new-api/dto" | "github.com/QuantumNous/new-api/dto" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | relaycommon "github.com/QuantumNous/new-api/relay/common" | ||||
| "github.com/QuantumNous/new-api/types" | |||||
| "github.com/gin-gonic/gin" | "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" | ||||
| @@ -86,6 +88,66 @@ func TestOaiResponsesHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { | |||||
| assert.Contains(t, err.UpstreamBody, `{"incomplete":`) | 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) { | func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { | ||||
| gin.SetMode(gin.TestMode) | gin.SetMode(gin.TestMode) | ||||
| w := httptest.NewRecorder() | w := httptest.NewRecorder() | ||||
| @@ -54,6 +54,19 @@ func SetEventStreamHeaders(c *gin.Context) { | |||||
| c.Writer.Header().Set("X-Accel-Buffering", "no") | 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 { | func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { | ||||
| jsonData, err := common.Marshal(resp) | jsonData, err := common.Marshal(resp) | ||||
| if err != nil { | if err != nil { | ||||