- relay: SSE stream error exposure improvements (Claude/Gemini/OpenAI) - test: add mock capture server for E2E relay capture testing Co-Authored-By: Claude <noreply@anthropic.com>master
| @@ -778,14 +778,21 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. | |||||
| Usage: &dto.Usage{}, | Usage: &dto.Usage{}, | ||||
| } | } | ||||
| var err *types.NewAPIError | var err *types.NewAPIError | ||||
| var normalStreamStarted bool | |||||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | helper.StreamScannerHandler(c, resp, info, func(data string) bool { | ||||
| err = HandleStreamResponseData(c, info, claudeInfo, data) | err = HandleStreamResponseData(c, info, claudeInfo, data) | ||||
| if err != nil { | if err != nil { | ||||
| if normalStreamStarted || helper.StreamDataWritten(c) { | |||||
| logger.LogError(c, "upstream claude stream error after normal stream started: "+err.Error()) | |||||
| err = nil | |||||
| } | |||||
| return false | return false | ||||
| } | } | ||||
| normalStreamStarted = true | |||||
| return true | return true | ||||
| }) | }) | ||||
| if err != nil { | if err != nil { | ||||
| helper.ClearEventStreamHeadersIfNotWritten(c) | |||||
| return nil, err | return nil, err | ||||
| } | } | ||||
| @@ -1,10 +1,18 @@ | |||||
| package claude | package claude | ||||
| import ( | import ( | ||||
| "io" | |||||
| "net/http" | |||||
| "net/http/httptest" | |||||
| "strings" | "strings" | ||||
| "testing" | "testing" | ||||
| "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" | |||||
| "github.com/QuantumNous/new-api/types" | |||||
| "github.com/gin-gonic/gin" | |||||
| "github.com/stretchr/testify/require" | |||||
| ) | ) | ||||
| func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { | func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { | ||||
| @@ -173,3 +181,65 @@ func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) { | |||||
| t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello") | t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello") | ||||
| } | } | ||||
| } | } | ||||
| func TestClaudeStreamHandlerReturnsErrorOnInitialErrorEvent(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/messages", nil) | |||||
| errorEvent := `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` | |||||
| resp := &http.Response{ | |||||
| StatusCode: http.StatusOK, | |||||
| Body: io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")), | |||||
| } | |||||
| info := &relaycommon.RelayInfo{ | |||||
| RelayFormat: types.RelayFormatClaude, | |||||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||||
| UpstreamModelName: "claude-test", | |||||
| }, | |||||
| } | |||||
| usage, err := ClaudeStreamHandler(c, resp, info) | |||||
| require.Nil(t, usage) | |||||
| require.Error(t, err) | |||||
| require.Empty(t, w.Body.String()) | |||||
| } | |||||
| func TestClaudeStreamHandlerDoesNotRetryAfterNormalStreamData(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/messages", nil) | |||||
| normalEvent := `{"type":"message_start","message":{"id":"msg_test","model":"claude-test","usage":{"input_tokens":1,"output_tokens":1}}}` | |||||
| errorEvent := `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` | |||||
| resp := &http.Response{ | |||||
| StatusCode: http.StatusOK, | |||||
| Body: io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")), | |||||
| } | |||||
| info := &relaycommon.RelayInfo{ | |||||
| RelayFormat: types.RelayFormatClaude, | |||||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||||
| UpstreamModelName: "claude-test", | |||||
| }, | |||||
| } | |||||
| usage, err := ClaudeStreamHandler(c, resp, info) | |||||
| require.Nil(t, err) | |||||
| require.NotNil(t, usage) | |||||
| require.Contains(t, w.Body.String(), "message_start") | |||||
| require.NotContains(t, w.Body.String(), "overloaded_error") | |||||
| } | |||||
| @@ -1286,6 +1286,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http | |||||
| var imageCount int | var imageCount int | ||||
| responseText := strings.Builder{} | responseText := strings.Builder{} | ||||
| var streamErr *types.NewAPIError | var streamErr *types.NewAPIError | ||||
| var normalStreamStarted bool | |||||
| 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 | ||||
| @@ -1297,15 +1298,21 @@ 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( | |||||
| apiErr := types.NewOpenAIError( | |||||
| errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason), | errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason), | ||||
| types.ErrorCodePromptBlocked, | types.ErrorCodePromptBlocked, | ||||
| http.StatusBadRequest, | http.StatusBadRequest, | ||||
| ) | ) | ||||
| streamErr.UpstreamBody = service.TruncateBody(data) | |||||
| apiErr.UpstreamBody = service.TruncateBody(data) | |||||
| if !normalStreamStarted && !helper.StreamDataWritten(c) { | |||||
| streamErr = apiErr | |||||
| } else { | |||||
| logger.LogError(c, "upstream gemini stream blocked after normal stream started: "+apiErr.Error()) | |||||
| } | |||||
| return false | return false | ||||
| } | } | ||||
| normalStreamStarted = true | |||||
| // 统计图片数量 | // 统计图片数量 | ||||
| for _, candidate := range geminiResponse.Candidates { | for _, candidate := range geminiResponse.Candidates { | ||||
| for _, part := range candidate.Content.Parts { | for _, part := range candidate.Content.Parts { | ||||
| @@ -328,6 +328,65 @@ func TestGeminiStreamHandlerReturnsErrorOnPromptBlock(t *testing.T) { | |||||
| require.Contains(t, newAPIError.UpstreamBody, "SAFETY") | require.Contains(t, newAPIError.UpstreamBody, "SAFETY") | ||||
| } | } | ||||
| func TestGeminiStreamHandlerDoesNotRetryAfterNormalStreamData(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) | |||||
| normalChunk := dto.GeminiChatResponse{ | |||||
| Candidates: []dto.GeminiChatCandidate{ | |||||
| { | |||||
| Content: dto.GeminiChatContent{ | |||||
| Role: "model", | |||||
| Parts: []dto.GeminiPart{ | |||||
| {Text: "partial"}, | |||||
| }, | |||||
| }, | |||||
| }, | |||||
| }, | |||||
| } | |||||
| blockReason := "SAFETY" | |||||
| blockChunk := dto.GeminiChatResponse{ | |||||
| PromptFeedback: &dto.GeminiChatPromptFeedback{ | |||||
| BlockReason: &blockReason, | |||||
| }, | |||||
| } | |||||
| normalData, err := common.Marshal(normalChunk) | |||||
| require.NoError(t, err) | |||||
| blockData, err := common.Marshal(blockChunk) | |||||
| require.NoError(t, err) | |||||
| resp := &http.Response{ | |||||
| StatusCode: http.StatusOK, | |||||
| Body: io.NopCloser(bytes.NewReader([]byte("data: " + string(normalData) + "\n\ndata: " + string(blockData) + "\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, newAPIError) | |||||
| require.NotNil(t, usage) | |||||
| require.True(t, called) | |||||
| require.Empty(t, w.Body.String()) | |||||
| } | |||||
| func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { | func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { | ||||
| t.Parallel() | t.Parallel() | ||||
| @@ -29,7 +29,7 @@ func streamOpenAIErrorFromData(data string, statusCode int) *types.NewAPIError { | |||||
| if openaiErr == nil { | if openaiErr == nil { | ||||
| return nil | return nil | ||||
| } | } | ||||
| if statusCode < 100 || statusCode > 599 { | |||||
| if statusCode < 100 || statusCode > 599 || (statusCode >= 200 && statusCode < 300) { | |||||
| statusCode = http.StatusInternalServerError | statusCode = http.StatusInternalServerError | ||||
| } | } | ||||
| apiErr := types.WithOpenAIError(*openaiErr, statusCode) | apiErr := types.WithOpenAIError(*openaiErr, statusCode) | ||||
| @@ -122,6 +122,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re | |||||
| var streamItems []string // store stream items | var streamItems []string // store stream items | ||||
| var lastStreamData string | var lastStreamData string | ||||
| var streamErr *types.NewAPIError | var streamErr *types.NewAPIError | ||||
| var normalStreamStarted bool | |||||
| var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 | var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 | ||||
| // 检查是否为音频模型 | // 检查是否为音频模型 | ||||
| @@ -129,14 +130,15 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re | |||||
| 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 apiErr := streamOpenAIErrorFromData(data, resp.StatusCode); apiErr != nil { | ||||
| if !c.Writer.Written() { | |||||
| if !normalStreamStarted && !helper.StreamDataWritten(c) { | |||||
| streamErr = apiErr | streamErr = apiErr | ||||
| } else { | } else { | ||||
| logger.LogError(c, "upstream stream error after downstream response was written: "+apiErr.Error()) | |||||
| logger.LogError(c, "upstream stream error after normal stream started: "+apiErr.Error()) | |||||
| } | } | ||||
| return false | return false | ||||
| } | } | ||||
| normalStreamStarted = true | |||||
| 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 { | ||||
| @@ -85,6 +85,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||||
| var usage = &dto.Usage{} | var usage = &dto.Usage{} | ||||
| var responseTextBuilder strings.Builder | var responseTextBuilder strings.Builder | ||||
| var streamErr *types.NewAPIError | var streamErr *types.NewAPIError | ||||
| var normalStreamStarted bool | |||||
| helper.StreamScannerHandler(c, resp, info, func(data string) bool { | helper.StreamScannerHandler(c, resp, info, func(data string) bool { | ||||
| @@ -92,10 +93,16 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||||
| if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { | if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { | ||||
| switch streamResponse.Type { | switch streamResponse.Type { | ||||
| case "response.error", "response.failed", "error": | case "response.error", "response.failed", "error": | ||||
| streamErr = handleResponsesStreamError(streamResponse, data) | |||||
| apiErr := handleResponsesStreamError(streamResponse, data) | |||||
| if !normalStreamStarted && !helper.StreamDataWritten(c) { | |||||
| streamErr = apiErr | |||||
| } else if apiErr != nil { | |||||
| logger.LogError(c, "upstream responses stream error after normal stream started: "+apiErr.Error()) | |||||
| } | |||||
| return false | return false | ||||
| } | } | ||||
| normalStreamStarted = true | |||||
| sendResponsesStreamData(c, streamResponse, data) | sendResponsesStreamData(c, streamResponse, data) | ||||
| switch streamResponse.Type { | switch streamResponse.Type { | ||||
| case "response.completed": | case "response.completed": | ||||
| @@ -11,6 +11,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/constant" | "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/relay/helper" | |||||
| "github.com/QuantumNous/new-api/types" | "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" | ||||
| @@ -109,6 +110,7 @@ func TestOaiResponsesStreamHandlerDoesNotForwardErrorEvent(t *testing.T) { | |||||
| usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp) | usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp) | ||||
| require.Nil(t, usage) | require.Nil(t, usage) | ||||
| require.Error(t, err) | require.Error(t, err) | ||||
| assert.Equal(t, http.StatusInternalServerError, err.StatusCode) | |||||
| assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") | assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") | ||||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | ||||
| assert.Empty(t, w.Body.String()) | assert.Empty(t, w.Body.String()) | ||||
| @@ -148,6 +150,104 @@ func TestOaiStreamHandlerDoesNotForwardInitialErrorEvent(t *testing.T) { | |||||
| assert.Empty(t, w.Header().Get("Content-Type")) | assert.Empty(t, w.Header().Get("Content-Type")) | ||||
| } | } | ||||
| func TestOaiStreamHandlerReturnsInitialErrorAfterPing(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) | |||||
| require.NoError(t, helper.PingData(c)) | |||||
| 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.Equal(t, http.StatusInternalServerError, err.StatusCode) | |||||
| assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded") | |||||
| assert.Contains(t, w.Body.String(), ": PING") | |||||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||||
| } | |||||
| func TestOaiStreamHandlerDoesNotRetryAfterNormalStreamData(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) | |||||
| normalEvent := `{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"gpt-test","choices":[{"index":0,"delta":{"content":"hi"}}]}` | |||||
| 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: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")), | |||||
| } | |||||
| info := &relaycommon.RelayInfo{ | |||||
| RelayFormat: types.RelayFormatOpenAI, | |||||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||||
| UpstreamModelName: "gpt-test", | |||||
| }, | |||||
| } | |||||
| usage, err := OaiStreamHandler(c, info, resp) | |||||
| require.Nil(t, err) | |||||
| require.NotNil(t, usage) | |||||
| assert.Contains(t, w.Body.String(), "hi") | |||||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||||
| } | |||||
| func TestOaiResponsesStreamHandlerDoesNotRetryAfterNormalStreamData(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) | |||||
| normalEvent := `{"type":"response.completed","response":{"id":"resp_test","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` | |||||
| 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: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")), | |||||
| } | |||||
| info := &relaycommon.RelayInfo{ | |||||
| ChannelMeta: &relaycommon.ChannelMeta{ | |||||
| UpstreamModelName: "gpt-test", | |||||
| }, | |||||
| } | |||||
| usage, err := OaiResponsesStreamHandler(c, info, resp) | |||||
| require.Nil(t, err) | |||||
| require.NotNil(t, usage) | |||||
| assert.Contains(t, w.Body.String(), "response.completed") | |||||
| assert.NotContains(t, w.Body.String(), "rate_limit_exceeded") | |||||
| } | |||||
| func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { | func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) { | ||||
| gin.SetMode(gin.TestMode) | gin.SetMode(gin.TestMode) | ||||
| w := httptest.NewRecorder() | w := httptest.NewRecorder() | ||||
| @@ -14,6 +14,40 @@ import ( | |||||
| "github.com/gorilla/websocket" | "github.com/gorilla/websocket" | ||||
| ) | ) | ||||
| const ( | |||||
| streamBodyWrittenKey = "stream_body_written" | |||||
| streamDataWrittenKey = "stream_data_written" | |||||
| ) | |||||
| func MarkStreamBodyWritten(c *gin.Context) { | |||||
| if c == nil { | |||||
| return | |||||
| } | |||||
| c.Set(streamBodyWrittenKey, true) | |||||
| } | |||||
| func StreamBodyWritten(c *gin.Context) bool { | |||||
| if c == nil { | |||||
| return false | |||||
| } | |||||
| return c.GetBool(streamBodyWrittenKey) | |||||
| } | |||||
| func MarkStreamDataWritten(c *gin.Context) { | |||||
| if c == nil { | |||||
| return | |||||
| } | |||||
| MarkStreamBodyWritten(c) | |||||
| c.Set(streamDataWrittenKey, true) | |||||
| } | |||||
| func StreamDataWritten(c *gin.Context) bool { | |||||
| if c == nil { | |||||
| return false | |||||
| } | |||||
| return c.GetBool(streamDataWrittenKey) | |||||
| } | |||||
| func FlushWriter(c *gin.Context) (err error) { | func FlushWriter(c *gin.Context) (err error) { | ||||
| defer func() { | defer func() { | ||||
| if r := recover(); r != nil { | if r := recover(); r != nil { | ||||
| @@ -55,7 +89,7 @@ func SetEventStreamHeaders(c *gin.Context) { | |||||
| } | } | ||||
| func ClearEventStreamHeadersIfNotWritten(c *gin.Context) { | func ClearEventStreamHeadersIfNotWritten(c *gin.Context) { | ||||
| if c == nil || c.Writer == nil || c.Writer.Written() { | |||||
| if c == nil || c.Writer == nil || StreamBodyWritten(c) { | |||||
| return | return | ||||
| } | } | ||||
| c.Set("event_stream_headers_set", false) | c.Set("event_stream_headers_set", false) | ||||
| @@ -74,6 +108,7 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { | |||||
| } else { | } else { | ||||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | ||||
| c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)}) | c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)}) | ||||
| MarkStreamDataWritten(c) | |||||
| } | } | ||||
| _ = FlushWriter(c) | _ = FlushWriter(c) | ||||
| return nil | return nil | ||||
| @@ -82,12 +117,14 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { | |||||
| func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) { | func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) { | ||||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | ||||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)}) | c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)}) | ||||
| MarkStreamDataWritten(c) | |||||
| _ = FlushWriter(c) | _ = FlushWriter(c) | ||||
| } | } | ||||
| func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) { | func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) { | ||||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) | ||||
| c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)}) | c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)}) | ||||
| MarkStreamDataWritten(c) | |||||
| _ = FlushWriter(c) | _ = FlushWriter(c) | ||||
| } | } | ||||
| @@ -101,6 +138,7 @@ func StringData(c *gin.Context, str string) error { | |||||
| } | } | ||||
| c.Render(-1, common.CustomEvent{Data: "data: " + str}) | c.Render(-1, common.CustomEvent{Data: "data: " + str}) | ||||
| MarkStreamDataWritten(c) | |||||
| return FlushWriter(c) | return FlushWriter(c) | ||||
| } | } | ||||
| @@ -116,6 +154,7 @@ func PingData(c *gin.Context) error { | |||||
| if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil { | if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil { | ||||
| return fmt.Errorf("write ping data failed: %w", err) | return fmt.Errorf("write ping data failed: %w", err) | ||||
| } | } | ||||
| MarkStreamBodyWritten(c) | |||||
| return FlushWriter(c) | return FlushWriter(c) | ||||
| } | } | ||||
| @@ -0,0 +1,208 @@ | |||||
| // +build ignore | |||||
| package main | |||||
| import ( | |||||
| "encoding/json" | |||||
| "fmt" | |||||
| "log" | |||||
| "net/http" | |||||
| "os" | |||||
| "time" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/QuantumNous/new-api/middleware" | |||||
| "github.com/gin-gonic/gin" | |||||
| ) | |||||
| func main() { | |||||
| gin.SetMode(gin.ReleaseMode) | |||||
| middleware.SetCaptureEnabled(1, true) | |||||
| router := gin.New() | |||||
| router.Use(gin.Recovery()) | |||||
| // 模拟 TokenAuth | |||||
| router.Use(func(c *gin.Context) { | |||||
| auth := c.GetHeader("Authorization") | |||||
| if auth == "" || auth == "Bearer disabled" { | |||||
| c.JSON(401, gin.H{"error": gin.H{"message": "Unauthorized"}}) | |||||
| c.Abort() | |||||
| return | |||||
| } | |||||
| c.Set("id", 1) | |||||
| c.Next() | |||||
| }) | |||||
| // 模拟 RequestId | |||||
| router.Use(func(c *gin.Context) { | |||||
| reqID := c.GetHeader(common.RequestIdKey) | |||||
| if reqID == "" { | |||||
| reqID = fmt.Sprintf("req-%d", time.Now().UnixNano()) | |||||
| } | |||||
| c.Set(common.RequestIdKey, reqID) | |||||
| c.Next() | |||||
| }) | |||||
| router.Use(middleware.RelayCaptureMiddleware()) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // OpenAI Chat Completions API | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.POST("/v1/chat/completions", func(c *gin.Context) { | |||||
| stream := c.Query("stream") == "true" | |||||
| if stream { | |||||
| c.Header("Content-Type", "text/event-stream") | |||||
| f := c.Writer.(http.Flusher) | |||||
| for _, s := range []string{ | |||||
| "data: {\"id\":\"chatcmpl-001\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n", | |||||
| "data: {\"id\":\"chatcmpl-001\",\"choices\":[{\"delta\":{\"content\":\"World!\"}}]}\n\n", | |||||
| "data: [DONE]\n\n", | |||||
| } { | |||||
| c.Writer.Write([]byte(s)) | |||||
| f.Flush() | |||||
| } | |||||
| return | |||||
| } | |||||
| c.JSON(200, gin.H{ | |||||
| "id": "chatcmpl-001", | |||||
| "object": "chat.completion", | |||||
| "model": "gpt-4", | |||||
| "choices": []gin.H{{"index": 0, "message": gin.H{"role": "assistant", "content": "Hello from Chat Completions!"}}}, | |||||
| "usage": gin.H{"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10}, | |||||
| }) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // OpenAI Responses API (new) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.POST("/v1/responses", func(c *gin.Context) { | |||||
| c.JSON(200, gin.H{ | |||||
| "id": "resp_abc123", | |||||
| "object": "response", | |||||
| "status": "completed", | |||||
| "output": []gin.H{{ | |||||
| "type": "message", | |||||
| "role": "assistant", | |||||
| "content": []gin.H{{ | |||||
| "type": "output_text", | |||||
| "text": "Hello from Responses API!", | |||||
| }}, | |||||
| }}, | |||||
| "usage": gin.H{"input_tokens": 8, "output_tokens": 4, "total_tokens": 12}, | |||||
| }) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // Claude Messages API (Anthropic) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.POST("/v1/messages", func(c *gin.Context) { | |||||
| c.JSON(200, gin.H{ | |||||
| "id": "msg_01AbCdEfGhIjKlMnOpQrStUv", | |||||
| "type": "message", | |||||
| "role": "assistant", | |||||
| "model": "claude-sonnet-4-6", | |||||
| "content": []gin.H{{"type": "text", "text": "Hello! How can I assist you today?"}}, | |||||
| "stop_reason": "end_turn", | |||||
| "usage": gin.H{"input_tokens": 8, "output_tokens": 12}, | |||||
| }) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // Gemini Native API (/v1beta/models/*) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.POST("/v1beta/models/*path", func(c *gin.Context) { | |||||
| path := c.Param("path") | |||||
| c.JSON(200, gin.H{ | |||||
| "candidates": []gin.H{{ | |||||
| "content": gin.H{ | |||||
| "role": "model", | |||||
| "parts": []gin.H{{"text": fmt.Sprintf("Gemini response for: /%s", path)}}, | |||||
| }, | |||||
| "finishReason": "STOP", | |||||
| }}, | |||||
| "usageMetadata": gin.H{"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15}, | |||||
| }) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // Gemini via /v1/models/* (main route) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.POST("/v1/models/*path", func(c *gin.Context) { | |||||
| path := c.Param("path") | |||||
| c.JSON(200, gin.H{ | |||||
| "candidates": []gin.H{{ | |||||
| "content": gin.H{ | |||||
| "role": "model", | |||||
| "parts": []gin.H{{"text": fmt.Sprintf("Gemini(v1) response for: /%s", path)}}, | |||||
| }, | |||||
| "finishReason": "STOP", | |||||
| }}, | |||||
| }) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // Models list | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.GET("/v1/models", func(c *gin.Context) { | |||||
| c.JSON(200, gin.H{"data": []gin.H{{"id": "gpt-4"}, {"id": "claude-3"}, {"id": "gemini-2.0-flash"}}}) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // Error endpoint for testing error status codes | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.POST("/v1/error", func(c *gin.Context) { | |||||
| code := 500 | |||||
| if c := c.Query("code"); c != "" { | |||||
| fmt.Sscanf(c, "%d", &code) | |||||
| } | |||||
| c.JSON(code, gin.H{"error": gin.H{"message": "test error", "type": "test_error"}}) | |||||
| }) | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| // Capture file management endpoints | |||||
| // ═══════════════════════════════════════════════════════ | |||||
| router.GET("/health", func(c *gin.Context) { | |||||
| c.JSON(200, gin.H{"status": "ok"}) | |||||
| }) | |||||
| router.GET("/capture-files", func(c *gin.Context) { | |||||
| today := time.Now().Format("2006-01-02") | |||||
| dir := fmt.Sprintf("./data/relay-capture/%s", today) | |||||
| entries, err := os.ReadDir(dir) | |||||
| if err != nil { | |||||
| c.JSON(200, gin.H{"files": []string{}, "error": err.Error()}) | |||||
| return | |||||
| } | |||||
| var files []string | |||||
| for _, e := range entries { | |||||
| files = append(files, e.Name()) | |||||
| } | |||||
| c.JSON(200, gin.H{"dir": dir, "files": files}) | |||||
| }) | |||||
| router.GET("/capture-files/:name", func(c *gin.Context) { | |||||
| today := time.Now().Format("2006-01-02") | |||||
| path := fmt.Sprintf("./data/relay-capture/%s/%s", today, c.Param("name")) | |||||
| data, err := os.ReadFile(path) | |||||
| if err != nil { | |||||
| c.JSON(404, gin.H{"error": err.Error()}) | |||||
| return | |||||
| } | |||||
| var record map[string]interface{} | |||||
| json.Unmarshal(data, &record) | |||||
| c.JSON(200, record) | |||||
| }) | |||||
| router.DELETE("/capture-files", func(c *gin.Context) { | |||||
| os.RemoveAll("./data/relay-capture") | |||||
| c.JSON(200, gin.H{"status": "cleaned"}) | |||||
| }) | |||||
| port := "18080" | |||||
| if p := os.Getenv("PORT"); p != "" { | |||||
| port = p | |||||
| } | |||||
| log.Printf("Mock Relay Server on :%s (chat|responses|messages|gemini)", port) | |||||
| http.ListenAndServe(":"+port, router) | |||||
| } | |||||