package relay import ( "context" "fmt" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel/claude" "github.com/QuantumNous/new-api/relay/channel/gemini" "github.com/QuantumNous/new-api/relay/channel/openai" relaycommon "github.com/QuantumNous/new-api/relay/common" relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestResponseGuardBeforeDelivery(t *testing.T) { for _, tc := range []struct{ name, content, body, reason string }{ {"astra", "application/json", `{"model":"gpt-6-astra","output":[]}`, "accepted"}, {"luna", "application/json", `{"model":"gpt-5.6-luna","usage":{"input_tokens":8,"output_tokens":2}}`, "model_mismatch"}, {"missing", "application/json", `{"output":[]}`, "missing_model"}, {"not a prefix match", "application/json", `{"model":"gpt-6-astra-other"}`, "model_mismatch"}, {"sse astra", "text/event-stream", "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-6-astra\"}}\n\ndata: [DONE]\n\n", "accepted"}, {"sse luna", "text/event-stream", "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n", "model_mismatch"}, {"sse no model", "text/event-stream", "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden\"}\n\ndata: [DONE]\n\n", "missing_model"}, {"sse buffer bound", "text/event-stream", ":" + strings.Repeat("x", 1<<20), "confirmation_buffer_limit"}, } { t.Run(tc.name, func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", tc.content) w.Header().Set("X-Codex-Turn-State", "must-not-leak") _, _ = io.WriteString(w, tc.body) })) defer server.Close() writer := httptest.NewRecorder() c, _ := gin.CreateTestContext(writer) c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}")) info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}, IsStream: strings.Contains(tc.content, "event-stream")} req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}")) require.NoError(t, err) resp, err := channel.DoRequest(c, req, info) assert.False(t, c.Writer.Written()) assert.Empty(t, writer.Header().Get("X-Codex-Turn-State")) assert.Equal(t, tc.reason, info.ResponseGuard.Last().Reason) if tc.reason == "accepted" { require.NoError(t, err) defer resp.Body.Close() body, err := io.ReadAll(resp.Body) require.NoError(t, err) assert.Equal(t, tc.body, string(body)) } else { require.Error(t, err) assert.Nil(t, resp) assert.Equal(t, http.StatusConflict, info.ResponseGuard.Last().APIError().StatusCode) } }) } } func TestResponseGuardAcceptsModelAfterFormerTenSecondCutoff(t *testing.T) { previous := constant.StreamingTimeout constant.StreamingTimeout = 30 t.Cleanup(func() { constant.StreamingTimeout = previous }) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { select { case <-time.After(11 * time.Second): w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"model":"gpt-6-astra","output":[]}`) case <-r.Context().Done(): } })) defer server.Close() c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest("POST", "/v1/responses", nil) info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}} req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}")) require.NoError(t, err) resp, err := channel.DoRequest(c, req, info) require.NoError(t, err) defer resp.Body.Close() assert.Equal(t, "accepted", info.ResponseGuard.Last().Reason) assert.False(t, c.Writer.Written()) } func TestResponseGuardOnlyEnablesForAstraRequests(t *testing.T) { tests := []struct { name string path string model string enabled bool want bool }{ {name: "astra enabled", path: "/v1/responses", model: "gpt-6-astra", enabled: true, want: true}, {name: "direct luna unaffected", path: "/v1/responses", model: "gpt-5.6-luna", enabled: true}, {name: "token opt out", path: "/v1/responses", model: "gpt-6-astra"}, {name: "other endpoint", path: "/v1/chat/completions", model: "gpt-6-astra", enabled: true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { info := &relaycommon.RelayInfo{ OriginModelName: tt.model, AstraAutoSwitch: tt.enabled, ChannelMeta: &relaycommon.ChannelMeta{ ChannelSetting: dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(!tt.enabled)}, }, } assert.Equal(t, tt.want, shouldEnableAstraResponseGuard(tt.path, info)) }) } } func TestResponseGuardTokenOptOutAllowsMismatchDespiteChannelSetting(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{"model":"gpt-5.6-luna","output":[]}`) })) defer server.Close() c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}")) guard := &relaycommon.ResponseModelGuard{} _, previous := guard.Begin(context.Background(), 1, 1) previous.Close() previous.Reason = "http_error" info := &relaycommon.RelayInfo{ ResponseGuard: guard, UpstreamRequests: 1, ChannelMeta: &relaycommon.ChannelMeta{ ChannelId: 2, ChannelSetting: dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(true)}, }, } req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}")) require.NoError(t, err) resp, err := channel.DoRequest(c, req, info) require.NoError(t, err) defer resp.Body.Close() body, err := io.ReadAll(resp.Body) require.NoError(t, err) assert.Contains(t, string(body), `"model":"gpt-5.6-luna"`) assert.Len(t, guard.Attempts, 1) assert.Equal(t, "http_error", previous.Reason) } func TestResponseGuardPublicErrorPrefersObservedMismatch(t *testing.T) { guard := &relaycommon.ResponseModelGuard{} _, mismatch := guard.Begin(context.Background(), 1, 1) mismatch.Returned = "gpt-5.6-luna" require.Error(t, mismatch.Reject("model_mismatch")) t.Cleanup(mismatch.Close) err := guard.APIError() require.NotNil(t, err) assert.Equal(t, relaycommon.ResponseGuardErrorCode, err.GetErrorCode()) assert.Equal(t, http.StatusConflict, err.StatusCode) assert.Contains(t, err.Error(), "gpt-6-astra") assert.Contains(t, err.Error(), "gpt-5.6-luna") } func TestResponseGuardLateMismatchAndCancellation(t *testing.T) { for _, cancelAfterAccept := range []bool{false, true} { t.Run(fmt.Sprint(cancelAfterAccept), func(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() guard := &relaycommon.ResponseModelGuard{} _, attempt := guard.Begin(ctx, 1, 1) defer attempt.Close() first := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_ok\",\"model\":\"gpt-6-astra\"}}\n\n" model := "gpt-5.6-luna" if cancelAfterAccept { model = "gpt-6-astra" } last := fmt.Sprintf("data: {\"type\":\"response.completed\",\"response\":{\"model\":%q,\"usage\":{\"input_tokens\":8,\"output_tokens\":2}}}\n\n", model) resp := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(first + last))} require.NoError(t, attempt.Confirm(resp, ctx)) defer resp.Body.Close() if cancelAfterAccept { cancel() } body, err := io.ReadAll(resp.Body) if cancelAfterAccept { require.NoError(t, err) assert.Equal(t, first+last, string(body)) } else { require.Error(t, err) assert.Equal(t, first, string(body)) assert.Equal(t, "model_mismatch", attempt.Reason) } require.NotNil(t, attempt.Usage) assert.Equal(t, 8, attempt.Usage.Input) }) } } func TestResponseGuardRejectsCancelledClientAndExhaustedBudget(t *testing.T) { for _, exhausted := range []bool{false, true} { t.Run(fmt.Sprint(exhausted), func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) ctx, cancel := context.WithCancel(context.Background()) defer cancel() c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}")).WithContext(ctx) info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{}} if exhausted { info.UpstreamRequests = 4 } else { cancel() } req, err := http.NewRequest("POST", "http://127.0.0.1:1", strings.NewReader("{}")) require.NoError(t, err) resp, err := channel.DoRequest(c, req, info) require.Error(t, err) assert.Nil(t, resp) assert.Empty(t, info.ResponseGuard.Attempts) }) } } func TestResponseModelComparisonAndLog(t *testing.T) { for _, tc := range []struct { name string models []string returned string mismatch bool }{ {name: "absent", models: []string{"", " "}}, {name: "requested model", models: []string{"requested"}, returned: "requested"}, {name: "mapped model", models: []string{"mapped"}, returned: "mapped"}, {name: "different model", models: []string{"other"}, returned: "other", mismatch: true}, {name: "case differs", models: []string{"Requested"}, returned: "Requested"}, {name: "requested prefix", models: []string{"requested-2026-09-01"}, returned: "requested-2026-09-01"}, {name: "mapped case differs", models: []string{"MAPPED"}, returned: "MAPPED"}, {name: "mapped prefix", models: []string{"mapped-2026-09-01"}, returned: "mapped-2026-09-01"}, {name: "reverse prefix still warns", models: []string{"request"}, returned: "request", mismatch: true}, {name: "substring still warns", models: []string{"other-requested"}, returned: "other-requested", mismatch: true}, {name: "compatible difference survives matching frames", models: []string{"mapped", "Requested", "", "requested"}, returned: "Requested"}, {name: "warning supersedes compatible difference", models: []string{"Requested", "other"}, returned: "other", mismatch: true}, {name: "compatible difference cannot erase warning", models: []string{"other", "Requested"}, returned: "other", mismatch: true}, {name: "mismatch survives later frames", models: []string{"mapped", "other", "", "requested"}, returned: "other", mismatch: true}, } { t.Run(tc.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) info := &relaycommon.RelayInfo{ OriginModelName: "requested", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped", IsModelMapped: true}, } for _, name := range tc.models { info.ObserveResponseModel(name) } other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1) var stored struct { ResponseModel *relaycommon.ResponseModel `json:"response_model"` } require.NoError(t, common.UnmarshalJsonStr(other.JSONString(), &stored)) if tc.returned == "" { assert.Nil(t, stored.ResponseModel) return } require.NotNil(t, stored.ResponseModel) assert.Equal(t, &relaycommon.ResponseModel{ RequestedModel: "requested", UpstreamModel: "mapped", ReturnedModel: tc.returned, Mismatch: tc.mismatch, }, stored.ResponseModel) assert.Equal(t, "mapped", info.UpstreamModelName) }) } } func TestResponseModelLogOmitsUnchangedModel(t *testing.T) { for _, tc := range []struct { name string upstream string returned string mapped bool record bool }{ {name: "same model", upstream: "requested", returned: "requested"}, {name: "no upstream name", returned: "requested"}, {name: "mapped model", upstream: "mapped", returned: "mapped", mapped: true, record: true}, {name: "mapped response echoes request", upstream: "mapped", returned: "requested", mapped: true, record: true}, {name: "prefix difference", upstream: "requested", returned: "requested-2026-09-01", record: true}, {name: "case difference", upstream: "requested", returned: "REQUESTED", record: true}, {name: "mismatch", upstream: "requested", returned: "other", record: true}, } { t.Run(tc.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) info := &relaycommon.RelayInfo{ OriginModelName: "requested", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream, IsModelMapped: tc.mapped}, } info.ObserveResponseModel(tc.returned) other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1) require.NotNil(t, info.ResponseModel) assert.Equal(t, float64(1), other.Snapshot()["model_ratio"]) if tc.record { assert.Equal(t, *info.ResponseModel, other.Snapshot()["response_model"]) } else { assert.NotContains(t, other.Snapshot(), "response_model") } }) } } func TestResponseModelEmptyExpectedNamesDoNotMatchEveryPrefix(t *testing.T) { for _, tc := range []struct{ requested, upstream string }{ {}, {upstream: "mapped"}, {requested: "requested"}, } { t.Run(tc.requested+"/"+tc.upstream, func(t *testing.T) { info := &relaycommon.RelayInfo{ OriginModelName: tc.requested, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream}, } info.ObserveResponseModel("other") require.NotNil(t, info.ResponseModel) assert.True(t, info.ResponseModel.Mismatch) }) } } func TestResponseModelRetryResetsObservation(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) common.SetContextKey(c, constant.ContextKeyOriginalModel, "requested") info := &relaycommon.RelayInfo{OriginModelName: "requested"} info.InitChannelMeta(c) info.ObserveResponseModel("other") require.True(t, info.ResponseModel.Mismatch) info.InitChannelMeta(c) assert.Nil(t, info.ResponseModel) info.ObserveResponseModel("requested") assert.False(t, info.ResponseModel.Mismatch) } func TestResponseModelHandlersCaptureBeforeConversion(t *testing.T) { gin.SetMode(gin.TestMode) previousTimeout := constant.StreamingTimeout constant.StreamingTimeout = 30 t.Cleanup(func() { constant.StreamingTimeout = previousTimeout }) chat := `{"id":"chat_1","model":"returned","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}` chatStream := "data: " + `{"id":"chat_1","model":"returned","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}` + "\n\ndata: " + `{"id":"chat_1","choices":[],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}` + "\n\ndata: [DONE]\n\n" responses := `{"id":"resp_1","model":"returned","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}` responsesStream := "data: " + `{"type":"response.created","response":{"id":"resp_1","model":"returned"}}` + "\n\ndata: " + `{"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}}` + "\n\n" claudeBody := `{"id":"msg_1","type":"message","model":"returned","role":"assistant","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","usage":{"input_tokens":2,"output_tokens":3}}` claudeStream := "data: " + `{"type":"message_start","message":{"id":"msg_1","type":"message","model":"returned","role":"assistant","content":[],"usage":{"input_tokens":2,"output_tokens":0}}}` + "\n\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}` + "\n\ndata: " + `{"type":"message_stop"}` + "\n\n" geminiBody := `{"modelVersion":"returned","candidates":[{"content":{"role":"model","parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3,"totalTokenCount":5}}` geminiStream := "data: " + geminiBody + "\n\n" for _, tc := range []struct { name string body string format types.RelayFormat stream bool handler func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.NewAPIError) }{ {"chat", chat, types.RelayFormatOpenAI, false, openai.OpenaiHandler}, {"chat stream", chatStream, types.RelayFormatOpenAI, true, openai.OaiStreamHandler}, {"responses", responses, types.RelayFormatOpenAIResponses, false, openai.OaiResponsesHandler}, {"responses stream", responsesStream, types.RelayFormatOpenAIResponses, true, openai.OaiResponsesStreamHandler}, {"responses to chat", responses, types.RelayFormatOpenAI, false, openai.OaiResponsesToChatHandler}, {"responses to chat stream", responsesStream, types.RelayFormatOpenAI, true, openai.OaiResponsesToChatStreamHandler}, {"responses buffered to chat", responsesStream, types.RelayFormatOpenAI, false, openai.OaiResponsesToChatBufferedStreamHandler}, {"chat to responses", chat, types.RelayFormatOpenAIResponses, false, openai.OaiChatToResponsesHandler}, {"chat to responses stream", chatStream, types.RelayFormatOpenAIResponses, true, openai.OaiChatToResponsesStreamHandler}, {"claude", claudeBody, types.RelayFormatClaude, false, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { return claude.ClaudeHandler(c, resp, info) }}, {"claude stream", claudeStream, types.RelayFormatClaude, true, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { return claude.ClaudeStreamHandler(c, resp, info) }}, {"claude to responses stream", claudeStream, types.RelayFormatOpenAIResponses, true, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { return claude.ClaudeResponsesStreamHandler(c, resp, info) }}, {"gemini", geminiBody, types.RelayFormatGemini, false, gemini.GeminiTextGenerationHandler}, {"gemini stream", geminiStream, types.RelayFormatGemini, true, gemini.GeminiTextGenerationStreamHandler}, {"gemini to chat", geminiBody, types.RelayFormatOpenAI, false, gemini.GeminiChatHandler}, {"gemini to responses", geminiBody, types.RelayFormatOpenAIResponses, false, gemini.GeminiResponsesHandler}, {"gemini to responses stream", geminiStream, types.RelayFormatOpenAIResponses, true, gemini.GeminiResponsesStreamHandler}, } { t.Run(tc.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) info := &relaycommon.RelayInfo{ OriginModelName: "requested", RelayFormat: tc.format, RelayMode: relayconstant.RelayModeChatCompletions, IsStream: tc.stream, DisablePing: true, StartTime: time.Now(), ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"}, ThinkingContentInfo: relaycommon.ThinkingContentInfo{IsFirstThinkingContent: true}, } usage, apiErr := tc.handler(c, info, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(tc.body))}) require.Nil(t, apiErr) require.NotNil(t, usage) assert.Equal(t, 5, usage.TotalTokens) require.NotNil(t, info.ResponseModel) assert.Equal(t, &relaycommon.ResponseModel{RequestedModel: "requested", UpstreamModel: "mapped", ReturnedModel: "returned", Mismatch: true}, info.ResponseModel) other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1) assert.Equal(t, *info.ResponseModel, other.Snapshot()["response_model"]) }) } } func TestResponseModelSharedResponsesAccumulator(t *testing.T) { // WebSocket turns and HTTP SSE use this same accumulator. A model-free // terminal must not erase the model declared before an interrupted stream. info := &relaycommon.RelayInfo{OriginModelName: "requested", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"}} accumulator := service.NewResponsesUsageAccumulator(info) accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.created", Response: &dto.OpenAIResponsesResponse{Model: "returned"}}) accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Usage: &dto.Usage{InputTokens: 2, OutputTokens: 3}}}) assert.Equal(t, 5, accumulator.Finish().TotalTokens) require.NotNil(t, info.ResponseModel) assert.True(t, info.ResponseModel.Mismatch) assert.Equal(t, "returned", info.ResponseModel.ReturnedModel) } func TestResponseModelDoesNotRecordSynthesizedModel(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) info := &relaycommon.RelayInfo{ OriginModelName: "requested", RelayFormat: types.RelayFormatOpenAI, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"}, } _, apiErr := openai.OaiResponsesToChatBufferedStreamHandler(c, info, &http.Response{ StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")), }) require.Nil(t, apiErr) assert.Nil(t, info.ResponseModel) } func TestAstraDetectionCanBeDisabledIndependently(t *testing.T) { for _, disabled := range []bool{false, true} { info := &relaycommon.RelayInfo{OriginModelName: "gpt-6-astra", AstraDetectionDisabled: disabled, AstraAutoSwitch: true, ChannelMeta: &relaycommon.ChannelMeta{}} info.ObserveResponseModel("gpt-5.6-luna") if disabled { require.Nil(t, info.ResponseModel) require.False(t, shouldEnableAstraResponseGuard("/v1/responses", info)) } else { require.NotNil(t, info.ResponseModel) require.True(t, info.ResponseModel.Mismatch) require.True(t, shouldEnableAstraResponseGuard("/v1/responses", info)) } } }