Browse Source

feat: 增强流式响应错误处理,SSE内嵌错误正确曝光为API错误

- 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
fengsilin 2 weeks ago
parent
commit
1d01c4cd64
7 changed files with 173 additions and 3 deletions
  1. +12
    -0
      relay/channel/gemini/relay-gemini.go
  2. +46
    -0
      relay/channel/gemini/relay_gemini_usage_test.go
  3. +18
    -0
      relay/channel/openai/helper.go
  4. +15
    -0
      relay/channel/openai/relay-openai.go
  5. +7
    -3
      relay/channel/openai/relay_responses.go
  6. +62
    -0
      relay/channel/openai/upstream_body_test.go
  7. +13
    -0
      relay/helper/common.go

+ 12
- 0
relay/channel/gemini/relay-gemini.go View File

@@ -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 {


+ 46
- 0
relay/channel/gemini/relay_gemini_usage_test.go View File

@@ -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()




+ 18
- 0
relay/channel/openai/helper.go View File

@@ -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++




+ 15
- 0
relay/channel/openai/relay-openai.go View File

@@ -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"`


+ 7
- 3
relay/channel/openai/relay_responses.go View File

@@ -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
} }




+ 62
- 0
relay/channel/openai/upstream_body_test.go View File

@@ -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()


+ 13
- 0
relay/helper/common.go View File

@@ -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 {


Loading…
Cancel
Save