diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index d4b4fbe..361b387 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -778,14 +778,21 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. Usage: &dto.Usage{}, } var err *types.NewAPIError + var normalStreamStarted bool helper.StreamScannerHandler(c, resp, info, func(data string) bool { err = HandleStreamResponseData(c, info, claudeInfo, data) 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 } + normalStreamStarted = true return true }) if err != nil { + helper.ClearEventStreamHeadersIfNotWritten(c) return nil, err } diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go index e34c861..f949113 100644 --- a/relay/channel/claude/relay_claude_test.go +++ b/relay/channel/claude/relay_claude_test.go @@ -1,10 +1,18 @@ package claude import ( + "io" + "net/http" + "net/http/httptest" "strings" "testing" + "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/require" ) 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") } } + +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") +} diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 66f794a..3bc06e0 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1286,6 +1286,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http var imageCount int responseText := strings.Builder{} var streamErr *types.NewAPIError + var normalStreamStarted bool helper.StreamScannerHandler(c, resp, info, func(data string) bool { 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 { 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), types.ErrorCodePromptBlocked, 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 } + normalStreamStarted = true // 统计图片数量 for _, candidate := range geminiResponse.Candidates { for _, part := range candidate.Content.Parts { diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index 7d111bc..021e3af 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -328,6 +328,65 @@ func TestGeminiStreamHandlerReturnsErrorOnPromptBlock(t *testing.T) { 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) { t.Parallel() diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 30abee6..f195899 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -29,7 +29,7 @@ func streamOpenAIErrorFromData(data string, statusCode int) *types.NewAPIError { if openaiErr == nil { return nil } - if statusCode < 100 || statusCode > 599 { + if statusCode < 100 || statusCode > 599 || (statusCode >= 200 && statusCode < 300) { statusCode = http.StatusInternalServerError } apiErr := types.WithOpenAIError(*openaiErr, statusCode) diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index d20934c..e5c0a76 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -122,6 +122,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re var streamItems []string // store stream items var lastStreamData string var streamErr *types.NewAPIError + var normalStreamStarted bool 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 { if apiErr := streamOpenAIErrorFromData(data, resp.StatusCode); apiErr != nil { - if !c.Writer.Written() { + if !normalStreamStarted && !helper.StreamDataWritten(c) { streamErr = apiErr } 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 } + normalStreamStarted = true if lastStreamData != "" { err := HandleStreamFormat(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) if err != nil { diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index e5fed78..a3b1f82 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -85,6 +85,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp var usage = &dto.Usage{} var responseTextBuilder strings.Builder var streamErr *types.NewAPIError + var normalStreamStarted 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 { switch streamResponse.Type { 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 } + normalStreamStarted = true sendResponsesStreamData(c, streamResponse, data) switch streamResponse.Type { case "response.completed": diff --git a/relay/channel/openai/upstream_body_test.go b/relay/channel/openai/upstream_body_test.go index a91c82c..3054ecc 100644 --- a/relay/channel/openai/upstream_body_test.go +++ b/relay/channel/openai/upstream_body_test.go @@ -11,6 +11,7 @@ import ( "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/relay/helper" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" @@ -109,6 +110,7 @@ func TestOaiResponsesStreamHandlerDoesNotForwardErrorEvent(t *testing.T) { usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, 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.NotContains(t, w.Body.String(), "rate_limit_exceeded") assert.Empty(t, w.Body.String()) @@ -148,6 +150,104 @@ func TestOaiStreamHandlerDoesNotForwardInitialErrorEvent(t *testing.T) { 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) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() diff --git a/relay/helper/common.go b/relay/helper/common.go index 8fffb5e..b0445e4 100644 --- a/relay/helper/common.go +++ b/relay/helper/common.go @@ -14,6 +14,40 @@ import ( "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) { defer func() { if r := recover(); r != nil { @@ -55,7 +89,7 @@ func SetEventStreamHeaders(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 } c.Set("event_stream_headers_set", false) @@ -74,6 +108,7 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { } else { c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)}) + MarkStreamDataWritten(c) } _ = FlushWriter(c) 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) { 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)}) + MarkStreamDataWritten(c) _ = FlushWriter(c) } 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("data: %s", data)}) + MarkStreamDataWritten(c) _ = FlushWriter(c) } @@ -101,6 +138,7 @@ func StringData(c *gin.Context, str string) error { } c.Render(-1, common.CustomEvent{Data: "data: " + str}) + MarkStreamDataWritten(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 { return fmt.Errorf("write ping data failed: %w", err) } + MarkStreamBodyWritten(c) return FlushWriter(c) } diff --git a/test/e2e/mock_capture_server.go b/test/e2e/mock_capture_server.go new file mode 100644 index 0000000..eb9a087 --- /dev/null +++ b/test/e2e/mock_capture_server.go @@ -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) +}