cneicy
a5f7b8e586
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
448 lines
22 KiB
Go
448 lines
22 KiB
Go
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))
|
|
}
|
|
}
|
|
}
|