575 lines
28 KiB
Go
575 lines
28 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/model"
|
|
"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/relay/helper"
|
|
"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: "suffix", models: []string{"other-requested"}, returned: "other-requested"},
|
|
{name: "provider path", models: []string{"vendor/requested"}, returned: "vendor/requested"},
|
|
{name: "provider path case differs", models: []string{"vendor/REQUESTED"}, returned: "vendor/REQUESTED"},
|
|
{name: "provider path with dated mapped model", models: []string{"vendor/mapped-2026-09-01"}, returned: "vendor/mapped-2026-09-01", mismatch: true},
|
|
{name: "nested provider path", models: []string{"accounts/vendor/models/requested"}, returned: "accounts/vendor/models/requested"},
|
|
{name: "provider path reverse prefix still warns", models: []string{"vendor/request"}, returned: "vendor/request", mismatch: true},
|
|
{name: "provider path other model still warns", models: []string{"vendor/other"}, returned: "vendor/other", mismatch: true},
|
|
{name: "provider path only still warns", models: []string{"vendor/"}, returned: "vendor/", 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,
|
|
}, stored.ResponseModel)
|
|
assert.NotContains(t, other.JSONString(), "mismatch")
|
|
assert.Equal(t, tc.mismatch, stored.ResponseModel.Mismatch())
|
|
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 TestResponseModelExpectedProviderPath(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
requested string
|
|
upstream string
|
|
returned string
|
|
mismatch bool
|
|
}{
|
|
{name: "requested carries provider path", requested: "vendor/requested", returned: "requested", mismatch: true},
|
|
{name: "same provider path on both sides", requested: "vendor/requested", returned: "vendor/requested-2026-09-01"},
|
|
{name: "provider path alone does not match", requested: "vendor/requested", returned: "vendor", mismatch: true},
|
|
{name: "different model behind same provider still warns", requested: "vendor/requested", returned: "vendor/other", mismatch: true},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
info := &relaycommon.RelayInfo{
|
|
OriginModelName: tc.requested,
|
|
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream},
|
|
}
|
|
info.ObserveResponseModel(tc.returned)
|
|
require.NotNil(t, info.ResponseModel)
|
|
assert.Equal(t, tc.returned, info.ResponseModel.ReturnedModel)
|
|
assert.Equal(t, tc.mismatch, 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"}, info.ResponseModel)
|
|
assert.True(t, info.ResponseModel.Mismatch())
|
|
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"}}, nil)
|
|
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Usage: &dto.Usage{InputTokens: 2, OutputTokens: 3}}}, nil)
|
|
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))
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestUpstreamTimingResponsesGuardAndScope(t *testing.T) {
|
|
service.InitHttpClient()
|
|
previous := constant.StreamingTimeout
|
|
constant.StreamingTimeout = 30
|
|
t.Cleanup(func() { constant.StreamingTimeout = previous })
|
|
for _, tc := range []struct {
|
|
name, group, model string
|
|
guard bool
|
|
}{
|
|
{"guard accepted", "gpt-s-unstable", "gpt-6-astra", true},
|
|
{"guard rejected", "gpt-s-unstable", "gpt-5.6-luna", true},
|
|
{"without guard", "gpt-s-unstable", "gpt-6-astra", false},
|
|
{"other group", "gpt-v-pro", "gpt-6-astra", true},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
payload := ": keepalive\n\ndata: {\"type\":\"response.in_progress\",\"response\":{}}\n\n" +
|
|
"data: {\"type\":\"response.created\",\"response\":{\"model\":\"" + tc.model + "\"}}\n\n"
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, _ = io.Copy(io.Discard, r.Body)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
_, _ = io.WriteString(w, payload)
|
|
}))
|
|
defer server.Close()
|
|
writer := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(writer)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}"))
|
|
info := &relaycommon.RelayInfo{
|
|
StartTime: time.Now(), UsingGroup: tc.group, AstraAutoSwitch: tc.guard,
|
|
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}, IsStream: true,
|
|
}
|
|
if tc.guard {
|
|
info.ResponseGuard = &relaycommon.ResponseModelGuard{}
|
|
}
|
|
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(), "model confirmation must not deliver buffered data")
|
|
if tc.guard && tc.model != "gpt-6-astra" {
|
|
require.Error(t, err)
|
|
require.Nil(t, resp)
|
|
snapshot := info.UpstreamTiming.Snapshot()
|
|
events := snapshot["events_ms"].(map[string]float64)
|
|
assert.Contains(t, events, "first_sse_data")
|
|
assert.NotContains(t, events, "model_confirmed")
|
|
assert.NotContains(t, events, "first_downstream_flush")
|
|
return
|
|
}
|
|
require.NoError(t, err)
|
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(resp.Body)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, payload, string(body), "instrumentation preserves every upstream byte")
|
|
resp.Body = io.NopCloser(strings.NewReader(string(body)))
|
|
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
|
|
assert.NoError(t, helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.created"}, data))
|
|
})
|
|
assert.Contains(t, writer.Body.String(), tc.model)
|
|
other := model.NewLogOther()
|
|
service.AppendRelayLogAdminInfo(c, info, other)
|
|
logData := other.Snapshot()
|
|
assert.NotContains(t, logData, "upstream_timing")
|
|
admin := logData["admin_info"].(map[string]any)
|
|
if tc.group != "gpt-s-unstable" {
|
|
assert.Nil(t, info.UpstreamTiming)
|
|
assert.NotContains(t, admin, "upstream_timing")
|
|
return
|
|
}
|
|
require.NotNil(t, info.UpstreamTiming)
|
|
snapshot := admin["upstream_timing"].(map[string]any)
|
|
events := snapshot["events_ms"].(map[string]float64)
|
|
ordered := []string{"upstream_start", "get_connection", "connection_ready", "request_sent", "response_first_byte", "response_headers", "first_sse_data"}
|
|
if tc.guard {
|
|
ordered = append(ordered, "model_confirmed")
|
|
}
|
|
ordered = append(ordered, "first_downstream_flush")
|
|
for i, key := range ordered {
|
|
require.Contains(t, events, key)
|
|
if i > 0 {
|
|
assert.GreaterOrEqual(t, events[key], events[ordered[i-1]], key)
|
|
}
|
|
}
|
|
assert.EqualValues(t, 2, snapshot["request_body_bytes"])
|
|
assert.Equal(t, "HTTP/1.1", snapshot["protocol"])
|
|
assert.Equal(t, 1, snapshot["write_count"])
|
|
assert.Equal(t, false, snapshot["write_failed"])
|
|
})
|
|
}
|
|
}
|