Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
3796 lines
137 KiB
Go
3796 lines
137 KiB
Go
package service
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/model"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||
"github.com/cespare/xxhash/v2"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/stretchr/testify/require"
|
||
"github.com/tidwall/gjson"
|
||
)
|
||
|
||
// 编译期接口断言
|
||
var _ AccountRepository = (*stubOpenAIAccountRepo)(nil)
|
||
var _ GatewayCache = (*stubGatewayCache)(nil)
|
||
|
||
type stubOpenAIAccountRepo struct {
|
||
AccountRepository
|
||
accounts []Account
|
||
}
|
||
|
||
type tempUnschedulableOpenAIAccountRepo struct {
|
||
stubOpenAIAccountRepo
|
||
modelRateLimitAccountID int64
|
||
modelRateLimitKey string
|
||
}
|
||
|
||
func (r *tempUnschedulableOpenAIAccountRepo) SetModelRateLimit(_ context.Context, accountID int64, modelKey string, _ time.Time, _ ...string) error {
|
||
r.modelRateLimitAccountID = accountID
|
||
r.modelRateLimitKey = modelKey
|
||
return nil
|
||
}
|
||
|
||
type snapshotUpdateAccountRepo struct {
|
||
stubOpenAIAccountRepo
|
||
updateExtraCalls chan map[string]any
|
||
}
|
||
|
||
func (r *snapshotUpdateAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
|
||
if r.updateExtraCalls != nil {
|
||
copied := make(map[string]any, len(updates))
|
||
for k, v := range updates {
|
||
copied[k] = v
|
||
}
|
||
r.updateExtraCalls <- copied
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (r stubOpenAIAccountRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
|
||
for i := range r.accounts {
|
||
if r.accounts[i].ID == id {
|
||
return &r.accounts[i], nil
|
||
}
|
||
}
|
||
return nil, errors.New("account not found")
|
||
}
|
||
|
||
func (r stubOpenAIAccountRepo) GetByIDs(ctx context.Context, ids []int64) ([]*Account, error) {
|
||
if len(ids) == 0 {
|
||
return []*Account{}, nil
|
||
}
|
||
index := make(map[int64]*Account, len(r.accounts))
|
||
for i := range r.accounts {
|
||
account := &r.accounts[i]
|
||
index[account.ID] = account
|
||
}
|
||
out := make([]*Account, 0, len(ids))
|
||
seen := make(map[int64]struct{}, len(ids))
|
||
for _, id := range ids {
|
||
if _, ok := seen[id]; ok {
|
||
continue
|
||
}
|
||
seen[id] = struct{}{}
|
||
if account, ok := index[id]; ok {
|
||
out = append(out, account)
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
func (r stubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
|
||
var result []Account
|
||
for _, acc := range r.accounts {
|
||
if acc.Platform == platform {
|
||
result = append(result, acc)
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
func (r stubOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||
var result []Account
|
||
for _, acc := range r.accounts {
|
||
if acc.Platform == platform {
|
||
result = append(result, acc)
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
func (r stubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||
return r.ListSchedulableByPlatform(ctx, platform)
|
||
}
|
||
|
||
func TestOpenAIGatewayService_ForwardAsAnthropic_CapacityShedReturnsRequestScopedFailoverWithoutCommit(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
body := []byte(`{"model":"gpt-5.4","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
|
||
c.Request.Header.Set("Content-Type", "application/json")
|
||
|
||
upstreamBody := []byte(`{"error":{"message":"Our servers are currently overloaded. Please try again later."}}`)
|
||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||
{
|
||
StatusCode: http.StatusBadRequest,
|
||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||
Body: io.NopCloser(bytes.NewReader(upstreamBody)),
|
||
},
|
||
{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
`data: {"type":"response.output_text.delta","delta":"ok"}`,
|
||
"",
|
||
`data: {"type":"response.completed","response":{"id":"resp_second","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
},
|
||
}}
|
||
repo := &tempUnschedulableOpenAIAccountRepo{}
|
||
rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||
svc := &OpenAIGatewayService{
|
||
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{
|
||
Enabled: false, AllowInsecureHTTP: true,
|
||
}}},
|
||
httpUpstream: upstream,
|
||
rateLimitService: rateLimitService,
|
||
}
|
||
account := &Account{
|
||
ID: 5099, Name: "temporary-unschedulable", Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey, Concurrency: 1,
|
||
Credentials: map[string]any{
|
||
"api_key": "sk-test",
|
||
"base_url": "http://upstream.example",
|
||
"model_provider": "env-openai",
|
||
"temp_unschedulable_enabled": true,
|
||
"temp_unschedulable_rules": []any{map[string]any{
|
||
"error_code": float64(http.StatusBadRequest),
|
||
"keywords": []any{"our servers are currently overloaded", "please try again later"},
|
||
"duration_minutes": float64(1),
|
||
}},
|
||
},
|
||
}
|
||
|
||
_, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
|
||
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusBadRequest, failoverErr.StatusCode)
|
||
require.True(t, failoverErr.ShouldRetryNextAccount())
|
||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||
require.True(t, failoverErr.RequestScopedTransient)
|
||
require.Zero(t, repo.modelRateLimitAccountID, "request-scoped capacity shedding must not change account health")
|
||
require.Empty(t, repo.modelRateLimitKey)
|
||
require.False(t, IsResponseCommitted(c))
|
||
require.Equal(t, http.StatusOK, rec.Code)
|
||
require.Empty(t, rec.Body.String())
|
||
|
||
secondRec := httptest.NewRecorder()
|
||
secondContext, _ := gin.CreateTestContext(secondRec)
|
||
secondContext.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
|
||
secondContext.Request.Header.Set("Content-Type", "application/json")
|
||
secondAccount := *account
|
||
secondAccount.ID = 5100
|
||
secondAccount.Name = "healthy-failover-account"
|
||
result, secondErr := svc.ForwardAsAnthropic(context.Background(), secondContext, &secondAccount, body, "", "")
|
||
require.NoError(t, secondErr)
|
||
require.NotNil(t, result)
|
||
require.Equal(t, "resp_second", result.ResponseID)
|
||
require.NotEmpty(t, secondRec.Body.String())
|
||
}
|
||
|
||
func TestFailoverOpenAIUpstreamHTTPError_NilContextSkipsTempUnschedulablePolicy(t *testing.T) {
|
||
repo := &tempUnschedulableOpenAIAccountRepo{}
|
||
svc := &OpenAIGatewayService{
|
||
rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil),
|
||
}
|
||
account := &Account{
|
||
ID: 5099, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{
|
||
"temp_unschedulable_enabled": true,
|
||
"temp_unschedulable_rules": []any{map[string]any{
|
||
"error_code": float64(http.StatusBadRequest),
|
||
"keywords": []any{"custom temporary outage"},
|
||
"duration_minutes": float64(1),
|
||
}},
|
||
},
|
||
}
|
||
body := []byte(`{"error":{"message":"Custom temporary outage."}}`)
|
||
resp := &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{}}
|
||
|
||
got := svc.failoverOpenAIUpstreamHTTPError(
|
||
context.Background(), nil, account, resp, body,
|
||
"Custom temporary outage.", "gpt-5.4",
|
||
)
|
||
|
||
require.Nil(t, got)
|
||
require.Zero(t, repo.modelRateLimitAccountID)
|
||
require.Empty(t, repo.modelRateLimitKey)
|
||
}
|
||
|
||
type groupAwareStubOpenAIAccountRepo struct {
|
||
stubOpenAIAccountRepo
|
||
}
|
||
|
||
func (r groupAwareStubOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
|
||
var result []Account
|
||
for _, acc := range r.accounts {
|
||
if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
|
||
result = append(result, acc)
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
func (r groupAwareStubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
|
||
var result []Account
|
||
for _, acc := range r.accounts {
|
||
if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
|
||
result = append(result, acc)
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
type stubConcurrencyCache struct {
|
||
ConcurrencyCache
|
||
loadBatchErr error
|
||
loadMap map[int64]*AccountLoadInfo
|
||
acquireResults map[int64]bool
|
||
waitCounts map[int64]int
|
||
skipDefaultLoad bool
|
||
}
|
||
|
||
type cancelReadCloser struct{}
|
||
|
||
func (c cancelReadCloser) Read(p []byte) (int, error) { return 0, context.Canceled }
|
||
func (c cancelReadCloser) Close() error { return nil }
|
||
|
||
type errReadCloser struct {
|
||
err error
|
||
}
|
||
|
||
func (r errReadCloser) Read([]byte) (int, error) { return 0, r.err }
|
||
func (r errReadCloser) Close() error { return nil }
|
||
|
||
type openAIStreamReadThenErrorCloser struct {
|
||
reader *strings.Reader
|
||
err error
|
||
}
|
||
|
||
func (r *openAIStreamReadThenErrorCloser) Read(p []byte) (int, error) {
|
||
if r.reader != nil && r.reader.Len() > 0 {
|
||
return r.reader.Read(p)
|
||
}
|
||
return 0, r.err
|
||
}
|
||
|
||
func (r *openAIStreamReadThenErrorCloser) Close() error { return nil }
|
||
|
||
type failingGinWriter struct {
|
||
gin.ResponseWriter
|
||
failAfter int
|
||
writes int
|
||
}
|
||
|
||
func (w *failingGinWriter) Write(p []byte) (int, error) {
|
||
if w.writes >= w.failAfter {
|
||
return 0, errors.New("write failed")
|
||
}
|
||
w.writes++
|
||
return w.ResponseWriter.Write(p)
|
||
}
|
||
|
||
func (c stubConcurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
|
||
if c.acquireResults != nil {
|
||
if result, ok := c.acquireResults[accountID]; ok {
|
||
return result, nil
|
||
}
|
||
}
|
||
return true, nil
|
||
}
|
||
|
||
func (c stubConcurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
|
||
return nil
|
||
}
|
||
|
||
func (c stubConcurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
|
||
if c.loadBatchErr != nil {
|
||
return nil, c.loadBatchErr
|
||
}
|
||
out := make(map[int64]*AccountLoadInfo, len(accounts))
|
||
if c.skipDefaultLoad && c.loadMap != nil {
|
||
for _, acc := range accounts {
|
||
if load, ok := c.loadMap[acc.ID]; ok {
|
||
out[acc.ID] = load
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
for _, acc := range accounts {
|
||
if c.loadMap != nil {
|
||
if load, ok := c.loadMap[acc.ID]; ok {
|
||
out[acc.ID] = load
|
||
continue
|
||
}
|
||
}
|
||
out[acc.ID] = &AccountLoadInfo{AccountID: acc.ID, LoadRate: 0}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHash_Priority(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
bodyWithKey := []byte(`{"prompt_cache_key":"ses_aaa"}`)
|
||
|
||
// 1) session_id header wins
|
||
c.Request.Header.Set("session_id", "sess-123")
|
||
c.Request.Header.Set("conversation_id", "conv-456")
|
||
h1 := svc.GenerateSessionHash(c, bodyWithKey)
|
||
if h1 == "" {
|
||
t.Fatalf("expected non-empty hash")
|
||
}
|
||
|
||
// 2) conversation_id used when session_id absent
|
||
c.Request.Header.Del("session_id")
|
||
h2 := svc.GenerateSessionHash(c, bodyWithKey)
|
||
if h2 == "" {
|
||
t.Fatalf("expected non-empty hash")
|
||
}
|
||
if h1 == h2 {
|
||
t.Fatalf("expected different hashes for different keys")
|
||
}
|
||
|
||
// 3) prompt_cache_key used when both headers absent
|
||
c.Request.Header.Del("conversation_id")
|
||
h3 := svc.GenerateSessionHash(c, bodyWithKey)
|
||
if h3 == "" {
|
||
t.Fatalf("expected non-empty hash")
|
||
}
|
||
if h2 == h3 {
|
||
t.Fatalf("expected different hashes for different keys")
|
||
}
|
||
|
||
// 4) empty when no signals
|
||
h4 := svc.GenerateSessionHash(c, []byte(`{}`))
|
||
if h4 != "" {
|
||
t.Fatalf("expected empty hash when no signals")
|
||
}
|
||
}
|
||
|
||
func TestOpenAIGatewayService_ClientSessionHeaderPriority(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
recorder := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(recorder)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
c.Set("api_key", &APIKey{ID: 901, Group: &Group{Platform: PlatformGrok}})
|
||
|
||
headers := []struct {
|
||
name string
|
||
value string
|
||
}{
|
||
{name: "session_id", value: "generic-session"},
|
||
{name: "conversation_id", value: "generic-conversation"},
|
||
{name: openCodeSessionAffinityHeader, value: "opencode-affinity"},
|
||
{name: openCodeSessionIDHeader, value: "opencode-session-id"},
|
||
{name: openCodeNativeSessionHeader, value: "opencode-native-session"},
|
||
{name: codeBuddyConversationHeader, value: "codebuddy-conversation"},
|
||
{name: grokConversationIDHeader, value: "grok-conversation"},
|
||
}
|
||
for _, header := range headers {
|
||
c.Request.Header.Set(header.name, header.value)
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
body := []byte(`{"prompt_cache_key":"body-session"}`)
|
||
for _, header := range headers {
|
||
require.Equal(t, header.value, svc.ExtractSessionID(c, body), header.name)
|
||
require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String(header.value)), svc.GenerateExplicitSessionHash(c, body), header.name)
|
||
if header.name != grokConversationIDHeader {
|
||
require.Equal(t, header.value, explicitOpenAISessionID(c, body), header.name)
|
||
}
|
||
c.Request.Header.Del(header.name)
|
||
}
|
||
require.Equal(t, "body-session", svc.ExtractSessionID(c, body))
|
||
}
|
||
|
||
func TestOpenAIGatewayService_ClientSessionHeadersIgnorePerRequestIDs(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
recorder := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(recorder)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
for name, value := range map[string]string{
|
||
"X-Conversation-Request-ID": "request-rotates-every-turn",
|
||
"X-Conversation-Message-ID": "message-rotates-every-turn",
|
||
"X-Request-ID": "generic-request-id",
|
||
} {
|
||
c.Request.Header.Set(name, value)
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
require.Empty(t, explicitOpenAIHeaderSessionID(c))
|
||
require.Empty(t, svc.ExtractSessionID(c, nil))
|
||
require.Empty(t, svc.GenerateExplicitSessionHash(c, nil))
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHash_UsesXXHash64(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
|
||
c.Request.Header.Set("session_id", "sess-fixed-value")
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
got := svc.GenerateSessionHash(c, nil)
|
||
want := fmt.Sprintf("%016x", xxhash.Sum64String("sess-fixed-value"))
|
||
require.Equal(t, want, got)
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHash_AttachesLegacyHashToContext(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
|
||
c.Request.Header.Set("session_id", "sess-legacy-check")
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
sessionHash := svc.GenerateSessionHash(c, nil)
|
||
require.NotEmpty(t, sessionHash)
|
||
require.NotNil(t, c.Request)
|
||
require.NotNil(t, c.Request.Context())
|
||
require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
|
||
}
|
||
|
||
func TestExtractOpenAIResponseIDFromJSONBytes(t *testing.T) {
|
||
require.Equal(t, "resp_json", extractOpenAIResponseIDFromJSONBytes([]byte(`{"id":"resp_json"}`)))
|
||
require.Equal(t, "resp_sse", extractOpenAIResponseIDFromJSONBytes([]byte(`{"type":"response.completed","response":{"id":"resp_sse"}}`)))
|
||
require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`{"response":{}}`)))
|
||
require.Empty(t, extractOpenAIResponseIDFromJSONBytes([]byte(`not-json`)))
|
||
}
|
||
|
||
// 复现 #4386:gpt-image-2 /v1/images/edits 的 usage 携带 input_tokens_details.image_tokens,
|
||
// 提取器须将图片输入 token 单独填入 ImageInputTokens(此前被丢弃并入 InputTokens 按文本价计费)。
|
||
func TestExtractOpenAIUsage_CapturesImageInputTokens(t *testing.T) {
|
||
body := []byte(`{"usage":{"input_tokens":371,"input_tokens_details":{"image_tokens":352,"text_tokens":19},"output_tokens":439,"output_tokens_details":{"image_tokens":439,"text_tokens":0},"total_tokens":810}}`)
|
||
usage, ok := extractOpenAIUsageFromJSONBytes(body)
|
||
require.True(t, ok)
|
||
require.Equal(t, 371, usage.InputTokens)
|
||
require.Equal(t, 352, usage.ImageInputTokens)
|
||
require.Equal(t, 439, usage.OutputTokens)
|
||
require.Equal(t, 439, usage.ImageOutputTokens)
|
||
|
||
// prompt_tokens_details 回退路径(部分上游用 prompt_tokens 口径)。
|
||
promptStyle := []byte(`{"usage":{"prompt_tokens":100,"prompt_tokens_details":{"image_tokens":80}}}`)
|
||
pu, ok := extractOpenAIUsageFromJSONBytes(promptStyle)
|
||
require.True(t, ok)
|
||
require.Equal(t, 100, pu.InputTokens)
|
||
require.Equal(t, 80, pu.ImageInputTokens)
|
||
|
||
// 纯文本请求:无 image_tokens 时 ImageInputTokens 为 0,行为不变。
|
||
textOnly := []byte(`{"usage":{"input_tokens":50,"output_tokens":10}}`)
|
||
tu, ok := extractOpenAIUsageFromJSONBytes(textOnly)
|
||
require.True(t, ok)
|
||
require.Zero(t, tu.ImageInputTokens)
|
||
}
|
||
|
||
func TestExtractOpenAIUsage_ReadsClineDataEnvelope(t *testing.T) {
|
||
body := []byte(`{"data":{"choices":[{"message":{"content":"OK"}}],"usage":{"prompt_tokens":8,"completion_tokens":27,"total_tokens":35,"prompt_tokens_details":{"cached_tokens":4}}},"success":true}`)
|
||
|
||
usage, ok := extractOpenAIUsageFromJSONBytes(body)
|
||
|
||
require.True(t, ok)
|
||
require.Equal(t, 8, usage.InputTokens)
|
||
require.Equal(t, 27, usage.OutputTokens)
|
||
require.Equal(t, 4, usage.CacheReadInputTokens)
|
||
}
|
||
|
||
func TestExtractOpenAIUsage_ReadsWrappedResponsesDataEnvelope(t *testing.T) {
|
||
body := []byte(`{"data":{"response":{"usage":{"input_tokens":11,"output_tokens":5,"total_tokens":16,"input_tokens_details":{"cached_tokens":2}}}}}`)
|
||
|
||
usage, ok := extractOpenAIUsageFromJSONBytes(body)
|
||
|
||
require.True(t, ok)
|
||
require.Equal(t, 11, usage.InputTokens)
|
||
require.Equal(t, 5, usage.OutputTokens)
|
||
require.Equal(t, 2, usage.CacheReadInputTokens)
|
||
}
|
||
|
||
func TestExtractOpenAIUsage_PreservesResponseUsagePriority(t *testing.T) {
|
||
body := []byte(`{"data":{"usage":{"prompt_tokens":100,"completion_tokens":50}},"response":{"usage":{"input_tokens":11,"output_tokens":5}}}`)
|
||
|
||
usage, ok := extractOpenAIUsageFromJSONBytes(body)
|
||
|
||
require.True(t, ok)
|
||
require.Equal(t, 11, usage.InputTokens)
|
||
require.Equal(t, 5, usage.OutputTokens)
|
||
}
|
||
|
||
func TestOpenAIGatewayService_BindHTTPResponseAccount(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
groupID := int64(4201)
|
||
c.Set("api_key", &APIKey{ID: 501, GroupID: &groupID})
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
account := &Account{ID: 37001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
svc.bindHTTPResponseAccount(context.Background(), c, account, "resp_http_001")
|
||
|
||
got, err := svc.getOpenAIWSStateStore().GetResponseAccount(context.Background(), groupID, "resp_http_001")
|
||
require.NoError(t, err)
|
||
require.Equal(t, account.ID, got)
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateExplicitSessionHash_SkipsContentFallback(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
svc := &OpenAIGatewayService{}
|
||
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
|
||
|
||
t.Run("stateless image body stays unstuck", func(t *testing.T) {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
||
|
||
require.Empty(t, svc.GenerateExplicitSessionHash(c, body))
|
||
require.Empty(t, openAILegacySessionHashFromContext(c.Request.Context()))
|
||
})
|
||
|
||
t.Run("prompt_cache_key is explicit", func(t *testing.T) {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
||
|
||
got := svc.GenerateExplicitSessionHash(c, []byte(`{"model":"gpt-image-2","prompt_cache_key":"image-session"}`))
|
||
require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String("image-session")), got)
|
||
require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
|
||
})
|
||
|
||
t.Run("header overrides body", func(t *testing.T) {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
||
c.Request.Header.Set("session_id", "header-session")
|
||
|
||
got := svc.GenerateExplicitSessionHash(c, []byte(`{"prompt_cache_key":"body-session"}`))
|
||
require.Equal(t, fmt.Sprintf("%016x", xxhash.Sum64String("header-session")), got)
|
||
})
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHashWithFallback(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
seed := "openai_ws_ingress:9:100:200"
|
||
|
||
got := svc.GenerateSessionHashWithFallback(c, []byte(`{}`), seed)
|
||
want := fmt.Sprintf("%016x", xxhash.Sum64String(seed))
|
||
require.Equal(t, want, got)
|
||
require.NotEmpty(t, openAILegacySessionHashFromContext(c.Request.Context()))
|
||
|
||
empty := svc.GenerateSessionHashWithFallback(c, []byte(`{}`), " ")
|
||
require.Equal(t, "", empty)
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHash_ContentFallback(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":"Hello"}]}`)
|
||
|
||
hash := svc.GenerateSessionHash(c, body)
|
||
require.NotEmpty(t, hash, "content-based fallback should produce a hash")
|
||
|
||
hash2 := svc.GenerateSessionHash(c, body)
|
||
require.Equal(t, hash, hash2, "same content should produce same hash")
|
||
|
||
bodyExtended := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi!"},{"role":"user","content":"How are you?"}]}`)
|
||
hashExtended := svc.GenerateSessionHash(c, bodyExtended)
|
||
require.Equal(t, hash, hashExtended, "hash should be stable across later turns")
|
||
|
||
bodyDifferent := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Different question"}]}`)
|
||
hashDifferent := svc.GenerateSessionHash(c, bodyDifferent)
|
||
require.NotEqual(t, hash, hashDifferent, "different content should produce different hash")
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHash_ExplicitSignalWinsOverContent(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Hello"}]}`)
|
||
|
||
contentHash := svc.GenerateSessionHash(c, body)
|
||
require.NotEmpty(t, contentHash)
|
||
|
||
c.Request.Header.Set("session_id", "explicit-session")
|
||
explicitHash := svc.GenerateSessionHash(c, body)
|
||
require.NotEmpty(t, explicitHash)
|
||
require.NotEqual(t, contentHash, explicitHash, "explicit session_id should override content fallback")
|
||
}
|
||
|
||
func TestOpenAIGatewayService_GenerateSessionHash_EmptyBodyStillEmpty(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/chat/completions", nil)
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
require.Empty(t, svc.GenerateSessionHash(c, []byte(`{}`)))
|
||
require.Empty(t, svc.GenerateSessionHash(c, nil))
|
||
}
|
||
|
||
func (c stubConcurrencyCache) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
|
||
if c.waitCounts != nil {
|
||
if count, ok := c.waitCounts[accountID]; ok {
|
||
return count, nil
|
||
}
|
||
}
|
||
return 0, nil
|
||
}
|
||
|
||
type stubGatewayCache struct {
|
||
sessionBindings map[string]int64
|
||
deletedSessions map[string]int
|
||
}
|
||
|
||
func (c *stubGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
|
||
if id, ok := c.sessionBindings[sessionHash]; ok {
|
||
return id, nil
|
||
}
|
||
return 0, errors.New("not found")
|
||
}
|
||
|
||
func (c *stubGatewayCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
|
||
if c.sessionBindings == nil {
|
||
c.sessionBindings = make(map[string]int64)
|
||
}
|
||
c.sessionBindings[sessionHash] = accountID
|
||
return nil
|
||
}
|
||
|
||
func (c *stubGatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
|
||
return nil
|
||
}
|
||
|
||
func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
|
||
if c.sessionBindings == nil {
|
||
return nil
|
||
}
|
||
if c.deletedSessions == nil {
|
||
c.deletedSessions = make(map[string]int)
|
||
}
|
||
c.deletedSessions[sessionHash]++
|
||
delete(c.sessionBindings, sessionHash)
|
||
return nil
|
||
}
|
||
|
||
func (c *stubGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
|
||
return nil
|
||
}
|
||
func (c *stubGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
|
||
return nil, nil
|
||
}
|
||
func (c *stubGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
|
||
return true, nil
|
||
}
|
||
|
||
func (c *stubGatewayCache) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
|
||
return nil
|
||
}
|
||
|
||
func (c *stubGatewayCache) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
|
||
return nil
|
||
}
|
||
func (c *stubGatewayCache) GetReasoningContent(_ context.Context, _ string) (string, error) {
|
||
return "", ErrReasoningContentNotFound
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
|
||
now := time.Now()
|
||
resetAt := now.Add(10 * time.Minute)
|
||
groupID := int64(1)
|
||
|
||
rateLimited := Account{
|
||
ID: 1,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Concurrency: 1,
|
||
Priority: 0,
|
||
RateLimitResetAt: &resetAt,
|
||
}
|
||
available := Account{
|
||
ID: 2,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Concurrency: 1,
|
||
Priority: 1,
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: stubOpenAIAccountRepo{accounts: []Account{rateLimited, available}},
|
||
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.2", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil {
|
||
t.Fatalf("expected selection with account")
|
||
}
|
||
if selection.Account.ID != available.ID {
|
||
t.Fatalf("expected account %d, got %d", available.ID, selection.Account.ID)
|
||
}
|
||
if selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_ImageRateLimitSkipsOnlyImageRequests(t *testing.T) {
|
||
future := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
|
||
groupID := int64(1)
|
||
|
||
imageLimited := Account{
|
||
ID: 1,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Concurrency: 1,
|
||
Priority: 0,
|
||
Extra: map[string]any{
|
||
modelRateLimitsKey: map[string]any{
|
||
openAIImageGenerationRateLimitKey: map[string]any{
|
||
"rate_limit_reset_at": future,
|
||
},
|
||
},
|
||
},
|
||
}
|
||
available := Account{
|
||
ID: 2,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Concurrency: 1,
|
||
Priority: 1,
|
||
}
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: stubOpenAIAccountRepo{accounts: []Account{imageLimited, available}},
|
||
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
|
||
}
|
||
|
||
imageSelection, err := svc.SelectAccountWithLoadAwareness(WithOpenAIImageGenerationIntent(context.Background()), &groupID, "", "gpt-5.4", nil)
|
||
require.NoError(t, err)
|
||
require.NotNil(t, imageSelection)
|
||
require.Equal(t, available.ID, imageSelection.Account.ID)
|
||
if imageSelection.ReleaseFunc != nil {
|
||
imageSelection.ReleaseFunc()
|
||
}
|
||
|
||
textSelection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.4", nil)
|
||
require.NoError(t, err)
|
||
require.NotNil(t, textSelection)
|
||
require.Equal(t, imageLimited.ID, textSelection.Account.ID)
|
||
if textSelection.ReleaseFunc != nil {
|
||
textSelection.ReleaseFunc()
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulableWhenNoConcurrencyService(t *testing.T) {
|
||
now := time.Now()
|
||
resetAt := now.Add(10 * time.Minute)
|
||
groupID := int64(1)
|
||
|
||
rateLimited := Account{
|
||
ID: 1,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Concurrency: 1,
|
||
Priority: 0,
|
||
RateLimitResetAt: &resetAt,
|
||
}
|
||
available := Account{
|
||
ID: 2,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Concurrency: 1,
|
||
Priority: 1,
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: stubOpenAIAccountRepo{accounts: []Account{rateLimited, available}},
|
||
// concurrencyService is nil, forcing the non-load-batch selection path.
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.2", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil {
|
||
t.Fatalf("expected selection with account")
|
||
}
|
||
if selection.Account.ID != available.ID {
|
||
t.Fatalf("expected account %d, got %d", available.ID, selection.Account.ID)
|
||
}
|
||
if selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_StickyUnschedulableClearsSession(t *testing.T) {
|
||
sessionHash := "session-1"
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusDisabled, Schedulable: true, Concurrency: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
|
||
}
|
||
if acc == nil || acc.ID != 2 {
|
||
t.Fatalf("expected account 2, got %+v", acc)
|
||
}
|
||
if cache.deletedSessions["openai:"+sessionHash] != 1 {
|
||
t.Fatalf("expected sticky session to be deleted")
|
||
}
|
||
if cache.sessionBindings["openai:"+sessionHash] != 2 {
|
||
t.Fatalf("expected sticky session to bind to account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_StickyOutsideGroupClearsSession(t *testing.T) {
|
||
sessionHash := "session-outside-group"
|
||
groupID := int64(1001)
|
||
repo := groupAwareStubOpenAIAccountRepo{
|
||
stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
|
||
},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), &groupID, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
|
||
}
|
||
if acc == nil || acc.ID != 2 {
|
||
t.Fatalf("expected account 2, got %+v", acc)
|
||
}
|
||
if cache.deletedSessions["openai:"+sessionHash] != 1 {
|
||
t.Fatalf("expected sticky session to be deleted")
|
||
}
|
||
if cache.sessionBindings["openai:"+sessionHash] != 2 {
|
||
t.Fatalf("expected sticky session to bind to account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_StickyUnschedulableClearsSession(t *testing.T) {
|
||
sessionHash := "session-2"
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusDisabled, Schedulable: true, Concurrency: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
|
||
t.Fatalf("expected account 2, got %+v", selection)
|
||
}
|
||
if cache.deletedSessions["openai:"+sessionHash] != 1 {
|
||
t.Fatalf("expected sticky session to be deleted")
|
||
}
|
||
if cache.sessionBindings["openai:"+sessionHash] != 2 {
|
||
t.Fatalf("expected sticky session to bind to account 2")
|
||
}
|
||
if selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_StickyOutsideGroupClearsSession(t *testing.T) {
|
||
sessionHash := "session-load-outside-group"
|
||
groupID := int64(1002)
|
||
repo := groupAwareStubOpenAIAccountRepo{
|
||
stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, AccountGroups: []AccountGroup{{GroupID: groupID}}},
|
||
},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
|
||
t.Fatalf("expected account 2, got %+v", selection)
|
||
}
|
||
if cache.deletedSessions["openai:"+sessionHash] != 1 {
|
||
t.Fatalf("expected sticky session to be deleted")
|
||
}
|
||
if cache.sessionBindings["openai:"+sessionHash] != 2 {
|
||
t.Fatalf("expected sticky session to bind to account 2")
|
||
}
|
||
if selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_NoModelSupport(t *testing.T) {
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{
|
||
ID: 1,
|
||
Platform: PlatformOpenAI,
|
||
Status: StatusActive,
|
||
Schedulable: true,
|
||
Credentials: map[string]any{"model_mapping": map[string]any{"gpt-3.5-turbo": "gpt-3.5-turbo"}},
|
||
},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "gpt-4", nil)
|
||
if err == nil {
|
||
t.Fatalf("expected error for unsupported model")
|
||
}
|
||
if acc != nil {
|
||
t.Fatalf("expected nil account for unsupported model")
|
||
}
|
||
if !strings.Contains(err.Error(), "supporting model") {
|
||
t.Fatalf("unexpected error: %v", err)
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorFallback(t *testing.T) {
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
loadBatchErr: errors.New("load batch failed"),
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "fallback", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil {
|
||
t.Fatalf("expected selection")
|
||
}
|
||
if selection.Account.ID != 2 {
|
||
t.Fatalf("expected account 2, got %d", selection.Account.ID)
|
||
}
|
||
if cache.sessionBindings["openai:fallback"] != 2 {
|
||
t.Fatalf("expected sticky session updated")
|
||
}
|
||
if selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_NoSlotFallbackWait(t *testing.T) {
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
acquireResults: map[int64]bool{1: false},
|
||
loadMap: map[int64]*AccountLoadInfo{
|
||
1: {AccountID: 1, LoadRate: 10},
|
||
},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.WaitPlan == nil {
|
||
t.Fatalf("expected wait plan fallback")
|
||
}
|
||
if selection.Account == nil || selection.Account.ID != 1 {
|
||
t.Fatalf("expected account 1")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_SetsStickyBinding(t *testing.T) {
|
||
sessionHash := "bind"
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
|
||
}
|
||
if acc == nil || acc.ID != 1 {
|
||
t.Fatalf("expected account 1")
|
||
}
|
||
if cache.sessionBindings["openai:"+sessionHash] != 1 {
|
||
t.Fatalf("expected sticky session binding")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_StickyWaitPlan(t *testing.T) {
|
||
sessionHash := "sticky-wait"
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
acquireResults: map[int64]bool{1: false},
|
||
waitCounts: map[int64]int{1: 0},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.WaitPlan == nil {
|
||
t.Fatalf("expected sticky wait plan")
|
||
}
|
||
if selection.Account == nil || selection.Account.ID != 1 {
|
||
t.Fatalf("expected account 1")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_PrefersLowerLoad(t *testing.T) {
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
loadMap: map[int64]*AccountLoadInfo{
|
||
1: {AccountID: 1, LoadRate: 80},
|
||
2: {AccountID: 2, LoadRate: 10},
|
||
},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "load", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
|
||
t.Fatalf("expected account 2")
|
||
}
|
||
if cache.sessionBindings["openai:load"] != 2 {
|
||
t.Fatalf("expected sticky session updated")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_StickyExcludedFallback(t *testing.T) {
|
||
sessionHash := "excluded"
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
excluded := map[int64]struct{}{1: {}}
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", excluded)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
|
||
}
|
||
if acc == nil || acc.ID != 2 {
|
||
t.Fatalf("expected account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_StickyNonOpenAI(t *testing.T) {
|
||
sessionHash := "non-openai"
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformAnthropic, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 2},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{
|
||
sessionBindings: map[string]int64{"openai:" + sessionHash: 1},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, sessionHash, "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
|
||
}
|
||
if acc == nil || acc.ID != 2 {
|
||
t.Fatalf("expected account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_NoAccounts(t *testing.T) {
|
||
repo := stubOpenAIAccountRepo{accounts: []Account{}}
|
||
cache := &stubGatewayCache{}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "", nil)
|
||
if err == nil {
|
||
t.Fatalf("expected error for no accounts")
|
||
}
|
||
if acc != nil {
|
||
t.Fatalf("expected nil account")
|
||
}
|
||
if !strings.Contains(err.Error(), "no available OpenAI accounts") {
|
||
t.Fatalf("unexpected error: %v", err)
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_NoCandidates(t *testing.T) {
|
||
groupID := int64(1)
|
||
resetAt := time.Now().Add(1 * time.Hour)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, RateLimitResetAt: &resetAt},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
|
||
if err == nil {
|
||
t.Fatalf("expected error for no candidates")
|
||
}
|
||
if selection != nil {
|
||
t.Fatalf("expected nil selection")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_AllFullWaitPlan(t *testing.T) {
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
loadMap: map[int64]*AccountLoadInfo{
|
||
1: {AccountID: 1, LoadRate: 100},
|
||
},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.WaitPlan == nil {
|
||
t.Fatalf("expected wait plan")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorNoAcquire(t *testing.T) {
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
loadBatchErr: errors.New("load batch failed"),
|
||
acquireResults: map[int64]bool{1: false},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.WaitPlan == nil {
|
||
t.Fatalf("expected wait plan")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_MissingLoadInfo(t *testing.T) {
|
||
groupID := int64(1)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
loadMap: map[int64]*AccountLoadInfo{
|
||
1: {AccountID: 1, LoadRate: 50},
|
||
},
|
||
skipDefaultLoad: true,
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
|
||
t.Fatalf("expected account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountForModelWithExclusions_LeastRecentlyUsed(t *testing.T) {
|
||
oldTime := time.Now().Add(-2 * time.Hour)
|
||
newTime := time.Now().Add(-1 * time.Hour)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Priority: 1, LastUsedAt: &newTime},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Priority: 1, LastUsedAt: &oldTime},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
}
|
||
|
||
acc, err := svc.SelectAccountForModelWithExclusions(context.Background(), nil, "", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountForModelWithExclusions error: %v", err)
|
||
}
|
||
if acc == nil || acc.ID != 2 {
|
||
t.Fatalf("expected account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAISelectAccountWithLoadAwareness_PreferNeverUsed(t *testing.T) {
|
||
groupID := int64(1)
|
||
lastUsed := time.Now().Add(-1 * time.Hour)
|
||
repo := stubOpenAIAccountRepo{
|
||
accounts: []Account{
|
||
{ID: 1, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, LastUsedAt: &lastUsed},
|
||
{ID: 2, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1},
|
||
},
|
||
}
|
||
cache := &stubGatewayCache{}
|
||
concurrencyCache := stubConcurrencyCache{
|
||
loadMap: map[int64]*AccountLoadInfo{
|
||
1: {AccountID: 1, LoadRate: 10},
|
||
2: {AccountID: 2, LoadRate: 10},
|
||
},
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{
|
||
accountRepo: repo,
|
||
cache: cache,
|
||
concurrencyService: NewConcurrencyService(concurrencyCache),
|
||
}
|
||
|
||
selection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-4", nil)
|
||
if err != nil {
|
||
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
|
||
}
|
||
if selection == nil || selection.Account == nil || selection.Account.ID != 2 {
|
||
t.Fatalf("expected account 2")
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingTimeout(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 1,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
start := time.Now()
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, start, "model", "model")
|
||
_ = pw.Close()
|
||
_ = pr.Close()
|
||
|
||
if err == nil || !strings.Contains(err.Error(), "stream data interval timeout") {
|
||
t.Fatalf("expected stream timeout error, got %v", err)
|
||
}
|
||
if !strings.Contains(rec.Body.String(), "\"type\":\"error\"") || !strings.Contains(rec.Body.String(), "stream_timeout") {
|
||
t.Fatalf("expected OpenAI-compatible error SSE event, got %q", rec.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingContextCanceledReturnsIncompleteErrorWithoutInjectingErrorEvent(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
cancel()
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: cancelReadCloser{},
|
||
Header: http.Header{},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
|
||
if err == nil || !strings.Contains(err.Error(), "stream usage incomplete") {
|
||
t.Fatalf("expected incomplete stream error, got %v", err)
|
||
}
|
||
if strings.Contains(rec.Body.String(), "event: error") || strings.Contains(rec.Body.String(), "stream_read_error") {
|
||
t.Fatalf("expected no injected SSE error event, got %q", rec.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingReadErrorBeforeOutputReturnsFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: errReadCloser{err: io.ErrUnexpectedEOF},
|
||
Header: http.Header{"X-Request-Id": []string{"rid-disconnect"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
}
|
||
|
||
func TestOpenAIStreamingPostOutputDisconnectQuarantinesSharedProxyWithoutSameStreamFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
proxyID := int64(4698)
|
||
account := &Account{
|
||
ID: 469801,
|
||
Name: "oauth-on-shared-proxy",
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeOAuth,
|
||
ProxyID: &proxyID,
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
}}}
|
||
// collapseInterval 0: the two loop iterations below record within the
|
||
// production collapse window and must count as distinct failure events here.
|
||
svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
|
||
failureThreshold: 2,
|
||
failureWindow: time.Minute,
|
||
quarantineTTL: 10 * time.Minute,
|
||
maxEntries: 16,
|
||
})
|
||
|
||
for _, readErr := range []error{
|
||
io.ErrUnexpectedEOF,
|
||
errors.New("http2: client connection lost"),
|
||
} {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: &openAIStreamReadThenErrorCloser{
|
||
reader: strings.NewReader(strings.Join([]string{
|
||
"event: response.output_text.delta",
|
||
`data: {"type":"response.output_text.delta","delta":"partial"}`,
|
||
"",
|
||
}, "\n")),
|
||
err: readErr,
|
||
},
|
||
Header: http.Header{"X-Request-Id": []string{"rid-proxy-disconnect"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "gpt-5.6-sol", "gpt-5.6-sol")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr), "post-output disconnect must not fail over inside the same stream")
|
||
require.Contains(t, rec.Body.String(), "partial")
|
||
}
|
||
|
||
scheduler := &defaultOpenAIAccountScheduler{service: svc}
|
||
compatible, reason := scheduler.isAccountRequestCompatibleReason(context.Background(), account, OpenAIAccountScheduleRequest{})
|
||
require.False(t, compatible, "the next request must exclude accounts sharing the quarantined proxy")
|
||
require.Equal(t, "proxy_stream_quarantined", reason)
|
||
}
|
||
|
||
func TestOpenAIStreamingTerminalAndClientCancellationDoNotQuarantineProxy(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
proxyID := int64(4699)
|
||
account := &Account{ID: 469901, Name: "oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, ProxyID: &proxyID}
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}}
|
||
|
||
terminalRecorder := httptest.NewRecorder()
|
||
terminalCtx, _ := gin.CreateTestContext(terminalRecorder)
|
||
terminalCtx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||
terminalResp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: &openAIStreamReadThenErrorCloser{
|
||
reader: strings.NewReader(strings.Join([]string{
|
||
"event: response.completed",
|
||
`data: {"type":"response.completed","response":{"status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":3,"total_tokens":8}}}`,
|
||
"",
|
||
}, "\n")),
|
||
err: io.ErrUnexpectedEOF,
|
||
},
|
||
Header: http.Header{},
|
||
}
|
||
_, err := svc.handleStreamingResponse(terminalCtx.Request.Context(), terminalResp, terminalCtx, account, time.Now(), "model", "model")
|
||
require.NoError(t, err)
|
||
|
||
for range 2 {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
cancel()
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: &openAIStreamReadThenErrorCloser{
|
||
reader: strings.NewReader("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
|
||
err: context.Canceled,
|
||
},
|
||
Header: http.Header{},
|
||
}
|
||
_, err = svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
}
|
||
|
||
scheduler := &defaultOpenAIAccountScheduler{service: svc}
|
||
compatible, reason := scheduler.isAccountRequestCompatibleReason(context.Background(), account, OpenAIAccountScheduleRequest{})
|
||
require.True(t, compatible)
|
||
require.Empty(t, reason)
|
||
}
|
||
|
||
func TestOpenAIStreamingResponseFailedBeforeOutputReturnsFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.in_progress",
|
||
`data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","error":{"message":"An error occurred while processing your request."}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-failed"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||
require.False(t, failoverErr.RetryableOnSameAccount)
|
||
require.Contains(t, string(failoverErr.ResponseBody), "An error occurred while processing your request")
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
}
|
||
|
||
func TestOpenAIStreamingResponseFailedBeforeOutputCapacityErrorReturnsFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.in_progress",
|
||
`data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","error":{"message":"Selected model is at capacity. Please try a different model.","type":"invalid_request_error"}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-capacity-failed"}},
|
||
}
|
||
|
||
account := &Account{
|
||
ID: 1,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Name: "pool-account",
|
||
Credentials: map[string]any{
|
||
"pool_mode": true,
|
||
},
|
||
}
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||
require.Contains(t, string(failoverErr.ResponseBody), "Selected model is at capacity")
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
}
|
||
|
||
func TestOpenAIStreamingResponseFailedBeforeOutputServerOverloadedCodeReturnsFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"id":"resp_1","error":{"code":"server_is_overloaded","message":"Please retry later."}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-overloaded-failed"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||
require.Contains(t, string(failoverErr.ResponseBody), "Please retry later")
|
||
// 容量降载是请求级信号:非池模式账号也要先在同账号重试,且不得据此临时封禁账号。
|
||
// 否则单个被降载的请求会把整池账号逐个消耗掉,而降载因素在每个账号上都相同。
|
||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||
require.True(t, failoverErr.RequestScopedTransient)
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
}
|
||
|
||
func TestOpenAIStreamingResponseFailedBeforeOutputRateLimitUsesPoolRetryPolicy(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{
|
||
"X-Request-Id": []string{"rid-rate-limit-failed"},
|
||
"Retry-After": []string{"1"},
|
||
},
|
||
}
|
||
account := &Account{
|
||
ID: 1,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Name: "pool-account",
|
||
Credentials: map[string]any{
|
||
"pool_mode": true,
|
||
"pool_mode_retry_count": float64(1),
|
||
"pool_mode_retry_status_codes": []any{float64(http.StatusTooManyRequests)},
|
||
},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||
require.Equal(t, "1", failoverErr.ResponseHeaders.Get("Retry-After"))
|
||
require.Equal(t, "rate_limit_error", gjson.GetBytes(failoverErr.ResponseBody, "error.type").String())
|
||
require.Contains(t, string(failoverErr.ResponseBody), "Concurrency limit exceeded")
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
|
||
opsVal, ok := c.Get(OpsUpstreamErrorsKey)
|
||
require.True(t, ok)
|
||
opsEvents, ok := opsVal.([]*OpsUpstreamErrorEvent)
|
||
require.True(t, ok)
|
||
require.NotEmpty(t, opsEvents)
|
||
require.Equal(t, http.StatusTooManyRequests, opsEvents[len(opsEvents)-1].UpstreamStatusCode)
|
||
}
|
||
|
||
// 流内 rate limit 只产生 failover 错误,不写账号级限流/封禁状态:
|
||
// HTTP 200 流的 x-codex-* 头是正常配额快照,不能按 429 头驱动账号冷却。
|
||
func TestOpenAIStreamingResponseFailedRateLimitDoesNotBlockAccountScheduling(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"rate_limit_exceeded","message":"Concurrency limit exceeded for account, please retry later"}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{
|
||
"X-Codex-Primary-Used-Percent": []string{"12"},
|
||
"X-Codex-Primary-Reset-After-Seconds": []string{"604800"},
|
||
"Retry-After": []string{"1"},
|
||
},
|
||
}
|
||
account := &Account{
|
||
ID: 11,
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeOAuth,
|
||
Name: "oauth-account",
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||
require.False(t, failoverErr.RetryableOnSameAccount)
|
||
require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
|
||
}
|
||
|
||
func TestOpenAIStreamingResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
longInstructions := strings.Repeat("You are GPT-5.1 running in the Codex CLI. ", 20)
|
||
failedPayload := fmt.Sprintf(
|
||
`{"type":"response.failed","response":{"id":"resp_failed","object":"response","created_at":1782446336,"status":"failed","instructions":%q,"output":[{"type":"message","content":[{"type":"output_text","text":"large"}]}],"usage":{"input_tokens":123,"output_tokens":0},"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."}}}`,
|
||
longInstructions,
|
||
)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_failed"}}`,
|
||
"",
|
||
"event: response.output_text.delta",
|
||
`data: {"type":"response.output_text.delta","delta":"partial"}`,
|
||
"",
|
||
"event: response.failed",
|
||
"data: " + failedPayload,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-failed-after-output"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
|
||
body := rec.Body.String()
|
||
require.Contains(t, body, "event: response.failed")
|
||
require.Contains(t, body, "context_length_exceeded")
|
||
require.Contains(t, body, `"type":"invalid_request_error"`)
|
||
require.Contains(t, body, "Your input exceeds the context window")
|
||
require.NotContains(t, body, "You are GPT-5.1 running in the Codex CLI")
|
||
require.NotContains(t, body, `"instructions"`)
|
||
require.NotContains(t, body, `"output"`)
|
||
require.NotContains(t, body, `"usage"`)
|
||
}
|
||
|
||
func TestOpenAIStreamingContextWindowResponseFailedBeforeOutputPassesThrough(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","error":{"type":"upstream_error","message":"Your input exceeds the context window of this model. Please adjust your input and try again.","code":null}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-context-window-failed"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr))
|
||
require.True(t, c.Writer.Written())
|
||
require.Contains(t, rec.Body.String(), "response.failed")
|
||
require.Contains(t, rec.Body.String(), `"type":"upstream_error"`)
|
||
require.Contains(t, rec.Body.String(), "Your input exceeds the context window")
|
||
}
|
||
|
||
func TestOpenAIStreamingContextWindowResponseFailedBeforeOutputAppliesPassthroughRule(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
rule := newNonFailoverPassthroughRule(http.StatusBadRequest, "context_length_exceeded", http.StatusBadRequest, "")
|
||
rule.Platforms = []string{PlatformOpenAI}
|
||
rule.PassthroughBody = true
|
||
rule.CustomMessage = nil
|
||
ruleSvc := &ErrorPassthroughService{}
|
||
ruleSvc.setLocalCache([]*model.ErrorPassthroughRule{rule})
|
||
BindErrorPassthroughService(c, ruleSvc)
|
||
|
||
upstreamMessage := "Your input exceeds the context window of this model. Please adjust your input and try again."
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"` + upstreamMessage + `"}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-context-window-passthrough-rule"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr))
|
||
require.True(t, IsResponseCommitted(c))
|
||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||
body := rec.Body.String()
|
||
require.Equal(t, "upstream_error", gjson.Get(body, "error.type").String())
|
||
require.Equal(t, upstreamMessage, gjson.Get(body, "error.message").String())
|
||
require.NotContains(t, body, "response.failed")
|
||
require.NotContains(t, body, "Upstream request failed")
|
||
// 命中透传规则也应记录 ops 上游错误事件(对齐 CC/Messages 与 antigravity 先例)。
|
||
opsVal, opsRecorded := c.Get(OpsUpstreamErrorsKey)
|
||
require.True(t, opsRecorded, "passthrough hit should record an ops upstream error event")
|
||
opsEvents, _ := opsVal.([]*OpsUpstreamErrorEvent)
|
||
require.NotEmpty(t, opsEvents)
|
||
}
|
||
|
||
func TestOpenAIStreamingPreambleOnlyMissingTerminalReturnsFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.in_progress",
|
||
`data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-missing-terminal"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
}
|
||
|
||
func TestOpenAIStreamingPreambleKeepaliveUsesDownstreamIdle(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
// Keepalive is based on *downstream* idle time (last flush to client),
|
||
// not upstream event cadence. Interval is seconds (config unit).
|
||
StreamKeepaliveInterval: 1,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
// Emit preamble/progress quickly so clientOutputStarted is true, then
|
||
// leave a real downstream idle gap longer than keepaliveInterval so the
|
||
// ticker can write ":\n\n". Frequent upstream ticks used to refresh
|
||
// lastDownstreamWriteAt and flake on loaded CI runners.
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n"))
|
||
time.Sleep(50 * time.Millisecond)
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{\"id\":\"resp_1\"}}\n\n"))
|
||
time.Sleep(1300 * time.Millisecond)
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}}\n\n"))
|
||
}()
|
||
|
||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
_ = pr.Close()
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.Contains(t, rec.Body.String(), ":\n\n")
|
||
require.Contains(t, rec.Body.String(), "response.completed")
|
||
}
|
||
|
||
func TestOpenAIStreamingNormalizesTerminalOutputFromDeltas(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
`data: {"type":"response.created","response":{"id":"resp_sdk_parse"}}`,
|
||
"",
|
||
`data: {"type":"response.output_text.delta","delta":"pon"}`,
|
||
"",
|
||
`data: {"type":"response.output_text.delta","delta":"g"}`,
|
||
"",
|
||
`data: {"type":"response.completed","response":{"id":"resp_sdk_parse","status":"completed","output":null,"usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-sdk-parse"}},
|
||
}
|
||
|
||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
|
||
terminalType, terminalPayload, ok := extractOpenAISSETerminalEvent(rec.Body.String())
|
||
require.True(t, ok)
|
||
require.Equal(t, "response.completed", terminalType)
|
||
output := gjson.GetBytes(terminalPayload, "response.output")
|
||
require.True(t, output.IsArray())
|
||
require.Len(t, output.Array(), 1)
|
||
require.Equal(t, "pong", gjson.GetBytes(terminalPayload, "response.output.0.content.0.text").String())
|
||
}
|
||
|
||
func TestOpenAIStreamingNormalizesTerminalOutputToEmptyArray(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
`data: {"type":"response.completed","response":{"id":"resp_empty","status":"completed","output":null,"usage":{"input_tokens":1,"output_tokens":0}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-empty-output"}},
|
||
}
|
||
|
||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
|
||
terminalType, terminalPayload, ok := extractOpenAISSETerminalEvent(rec.Body.String())
|
||
require.True(t, ok)
|
||
require.Equal(t, "response.completed", terminalType)
|
||
output := gjson.GetBytes(terminalPayload, "response.output")
|
||
require.True(t, output.IsArray())
|
||
require.Len(t, output.Array(), 0)
|
||
}
|
||
|
||
func TestOpenAIStreamingPolicyResponseFailedBeforeOutputPassesThrough(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","error":{"type":"safety_error","message":"This request has been flagged for potentially high-risk cyber activity."}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-policy-failed"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr))
|
||
require.True(t, c.Writer.Written())
|
||
require.Contains(t, rec.Body.String(), "response.failed")
|
||
require.Contains(t, rec.Body.String(), "high-risk cyber activity")
|
||
}
|
||
|
||
func TestOpenAIStreamingClientDisconnectDrainsUpstreamUsage(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
c.Writer = &failingGinWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.in_progress\",\"response\":{}}\n\n"))
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
|
||
}()
|
||
|
||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
|
||
_ = pr.Close()
|
||
if err != nil {
|
||
t.Fatalf("expected nil error, got %v", err)
|
||
}
|
||
if result == nil || result.usage == nil {
|
||
t.Fatalf("expected usage result")
|
||
}
|
||
if result.usage.InputTokens != 3 || result.usage.OutputTokens != 5 || result.usage.CacheReadInputTokens != 1 {
|
||
t.Fatalf("unexpected usage: %+v", *result.usage)
|
||
}
|
||
if strings.Contains(rec.Body.String(), "event: error") || strings.Contains(rec.Body.String(), "write_failed") {
|
||
t.Fatalf("expected no injected SSE error event, got %q", rec.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingMissingTerminalEventReturnsIncompleteError(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\",\"output_index\":0}\n\n"))
|
||
}()
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
|
||
_ = pr.Close()
|
||
if err == nil || !strings.Contains(err.Error(), "missing terminal event") {
|
||
t.Fatalf("expected missing terminal event error, got %v", err)
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughMissingTerminalEventReturnsIncompleteError(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\",\"output_index\":0}\n\n"))
|
||
}()
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
|
||
_ = pr.Close()
|
||
if err == nil || !strings.Contains(err.Error(), "missing terminal event") {
|
||
t.Fatalf("expected missing terminal event error, got %v", err)
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughPostOutputDisconnectQuarantinesSharedProxy(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
proxyID := int64(4698)
|
||
account := &Account{ID: 469804, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, ProxyID: &proxyID}
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}}
|
||
// collapseInterval 0: the loop below records within the production collapse
|
||
// window and must count as distinct failure events here.
|
||
svc.openaiProxyStreamCircuit = newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
|
||
failureThreshold: 2,
|
||
failureWindow: time.Minute,
|
||
quarantineTTL: 10 * time.Minute,
|
||
maxEntries: 16,
|
||
})
|
||
|
||
for _, readErr := range []error{io.ErrUnexpectedEOF, errors.New("http2: client connection lost")} {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: &openAIStreamReadThenErrorCloser{
|
||
reader: strings.NewReader("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
|
||
err: readErr,
|
||
},
|
||
Header: http.Header{"X-Request-Id": []string{"rid-passthrough-proxy-disconnect"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, account, time.Now(), "model", "model")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr), "post-output disconnect must not fail over inside the same stream")
|
||
require.Contains(t, rec.Body.String(), "partial")
|
||
}
|
||
|
||
require.True(t, svc.isOpenAIProxyStreamQuarantined(context.Background(), account))
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughResponseFailedBeforeOutputReturnsFailover(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","error":{"message":"upstream processing failed"}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-passthrough-failed"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.ErrorAs(t, err, &failoverErr)
|
||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||
require.Contains(t, string(failoverErr.ResponseBody), "upstream processing failed")
|
||
require.False(t, c.Writer.Written())
|
||
require.Empty(t, rec.Body.String())
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughContextWindowResponseFailedBeforeOutputAppliesPassthroughRule(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
rule := newNonFailoverPassthroughRule(http.StatusBadRequest, "input exceeds the context window", http.StatusBadRequest, "")
|
||
rule.Platforms = []string{PlatformOpenAI}
|
||
rule.PassthroughBody = true
|
||
rule.CustomMessage = nil
|
||
ruleSvc := &ErrorPassthroughService{}
|
||
ruleSvc.setLocalCache([]*model.ErrorPassthroughRule{rule})
|
||
BindErrorPassthroughService(c, ruleSvc)
|
||
|
||
upstreamMessage := "Your input exceeds the context window of this model. Please adjust your input and try again."
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"` + upstreamMessage + `"}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-pass-context-window-passthrough-rule"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr))
|
||
require.True(t, IsResponseCommitted(c))
|
||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||
body := rec.Body.String()
|
||
require.Equal(t, "upstream_error", gjson.Get(body, "error.type").String())
|
||
require.Equal(t, upstreamMessage, gjson.Get(body, "error.message").String())
|
||
require.NotContains(t, body, "response.failed")
|
||
require.NotContains(t, body, "Upstream request failed")
|
||
// 命中透传规则也应记录 ops 上游错误事件(对齐 CC/Messages 与 antigravity 先例)。
|
||
opsVal, opsRecorded := c.Get(OpsUpstreamErrorsKey)
|
||
require.True(t, opsRecorded, "passthrough hit should record an ops upstream error event")
|
||
opsEvents, _ := opsVal.([]*OpsUpstreamErrorEvent)
|
||
require.NotEmpty(t, opsEvents)
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughContextWindowResponseFailedBeforeOutputWithoutRulePassesThrough(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"id":"resp_1","error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-pass-context-window-no-rule"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr))
|
||
body := rec.Body.String()
|
||
require.Contains(t, body, "event: response.failed")
|
||
require.Contains(t, body, "context_length_exceeded")
|
||
require.Contains(t, body, "Your input exceeds the context window")
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughResponseFailedAfterOutputSanitizesVerboseResponseForClient(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
longInstructions := strings.Repeat("You are GPT-5.1 running in the Codex CLI. ", 20)
|
||
failedPayload := fmt.Sprintf(
|
||
`{"type":"response.failed","response":{"id":"resp_pass_failed","object":"response","created_at":1782446336,"status":"failed","instructions":%q,"output":[{"type":"message","content":[{"type":"output_text","text":"large"}]}],"usage":{"input_tokens":123,"output_tokens":0},"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."}}}`,
|
||
longInstructions,
|
||
)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"resp_pass_failed"}}`,
|
||
"",
|
||
"event: response.output_text.delta",
|
||
`data: {"type":"response.output_text.delta","delta":"partial"}`,
|
||
"",
|
||
"event: response.failed",
|
||
"data: " + failedPayload,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-pass-failed-after-output"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "", "")
|
||
require.Error(t, err)
|
||
|
||
body := rec.Body.String()
|
||
require.Contains(t, body, "event: response.failed")
|
||
require.Contains(t, body, "context_length_exceeded")
|
||
require.Contains(t, body, `"type":"invalid_request_error"`)
|
||
require.Contains(t, body, "Your input exceeds the context window")
|
||
require.NotContains(t, body, "You are GPT-5.1 running in the Codex CLI")
|
||
require.NotContains(t, body, `"instructions"`)
|
||
require.NotContains(t, body, `"output"`)
|
||
require.NotContains(t, body, `"usage"`)
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughResponseDoneWithoutDoneMarkerStillSucceeds(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.done\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
|
||
}()
|
||
|
||
result, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
|
||
_ = pr.Close()
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.NotNil(t, result.usage)
|
||
require.Equal(t, 2, result.usage.InputTokens)
|
||
require.Equal(t, 3, result.usage.OutputTokens)
|
||
require.Equal(t, 1, result.usage.CacheReadInputTokens)
|
||
}
|
||
|
||
func TestOpenAIStreamingPassthroughResponseIncompleteWithoutDoneMarkerStillSucceeds(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.incomplete\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n"))
|
||
}()
|
||
|
||
result, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "", "")
|
||
_ = pr.Close()
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.NotNil(t, result.usage)
|
||
require.Equal(t, 2, result.usage.InputTokens)
|
||
require.Equal(t, 3, result.usage.OutputTokens)
|
||
require.Equal(t, 1, result.usage.CacheReadInputTokens)
|
||
}
|
||
|
||
func TestOpenAIStreamingTooLong(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: 64 * 1024,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
// 写入超过 MaxLineSize 的单行数据,触发 ErrTooLong
|
||
payload := "data: " + strings.Repeat("a", 128*1024) + "\n"
|
||
_, _ = pw.Write([]byte(payload))
|
||
}()
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 2}, time.Now(), "model", "model")
|
||
_ = pr.Close()
|
||
|
||
if !errors.Is(err, bufio.ErrTooLong) {
|
||
t.Fatalf("expected ErrTooLong, got %v", err)
|
||
}
|
||
if !strings.Contains(rec.Body.String(), "\"type\":\"error\"") || !strings.Contains(rec.Body.String(), "response_too_large") {
|
||
t.Fatalf("expected OpenAI-compatible error SSE event, got %q", rec.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestOpenAINonStreamingContentTypePassThrough(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
body := []byte(`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(bytes.NewReader(body)),
|
||
Header: http.Header{"Content-Type": []string{"application/vnd.test+json"}},
|
||
}
|
||
|
||
_, err := svc.handleNonStreamingResponse(c.Request.Context(), resp, c, &Account{}, "model", "model")
|
||
if err != nil {
|
||
t.Fatalf("handleNonStreamingResponse error: %v", err)
|
||
}
|
||
|
||
if !strings.Contains(rec.Header().Get("Content-Type"), "application/vnd.test+json") {
|
||
t.Fatalf("expected Content-Type passthrough, got %q", rec.Header().Get("Content-Type"))
|
||
}
|
||
}
|
||
|
||
func TestOpenAINonStreamingContentTypeDefault(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
body := []byte(`{"usage":{"input_tokens":1,"output_tokens":2,"input_tokens_details":{"cached_tokens":0}}}`)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(bytes.NewReader(body)),
|
||
Header: http.Header{},
|
||
}
|
||
|
||
_, err := svc.handleNonStreamingResponse(c.Request.Context(), resp, c, &Account{}, "model", "model")
|
||
if err != nil {
|
||
t.Fatalf("handleNonStreamingResponse error: %v", err)
|
||
}
|
||
|
||
if !strings.Contains(rec.Header().Get("Content-Type"), "application/json") {
|
||
t.Fatalf("expected default Content-Type, got %q", rec.Header().Get("Content-Type"))
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingHeadersOverride(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
ResponseHeaders: config.ResponseHeaderConfig{Enabled: false},
|
||
},
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{
|
||
"Cache-Control": []string{"upstream"},
|
||
"X-Request-Id": []string{"req-123"},
|
||
"Content-Type": []string{"application/custom"},
|
||
},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{}}\n\n"))
|
||
}()
|
||
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
|
||
_ = pr.Close()
|
||
if err != nil {
|
||
t.Fatalf("handleStreamingResponse error: %v", err)
|
||
}
|
||
|
||
if rec.Header().Get("Cache-Control") != "no-cache" {
|
||
t.Fatalf("expected Cache-Control override, got %q", rec.Header().Get("Cache-Control"))
|
||
}
|
||
if rec.Header().Get("Content-Type") != "text/event-stream" {
|
||
t.Fatalf("expected Content-Type override, got %q", rec.Header().Get("Content-Type"))
|
||
}
|
||
if rec.Header().Get("X-Request-Id") != "req-123" {
|
||
t.Fatalf("expected X-Request-Id passthrough, got %q", rec.Header().Get("X-Request-Id"))
|
||
}
|
||
}
|
||
|
||
func TestOpenAIStreamingReuseScannerBufferAndStillWorks(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Gateway: config.GatewayConfig{
|
||
StreamDataIntervalTimeout: 0,
|
||
StreamKeepaliveInterval: 0,
|
||
MaxLineSize: defaultMaxLineSize,
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
pr, pw := io.Pipe()
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: pr,
|
||
Header: http.Header{},
|
||
}
|
||
|
||
go func() {
|
||
defer func() { _ = pw.Close() }()
|
||
_, _ = pw.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"input_tokens_details\":{\"cached_tokens\":3}}}}\n\n"))
|
||
}()
|
||
|
||
result, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1}, time.Now(), "model", "model")
|
||
_ = pr.Close()
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.NotNil(t, result.usage)
|
||
require.Equal(t, 1, result.usage.InputTokens)
|
||
require.Equal(t, 2, result.usage.OutputTokens)
|
||
require.Equal(t, 3, result.usage.CacheReadInputTokens)
|
||
}
|
||
|
||
func TestOpenAIInvalidBaseURLWhenAllowlistDisabled(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
account := &Account{
|
||
Platform: PlatformOpenAI,
|
||
Type: AccountTypeAPIKey,
|
||
Credentials: map[string]any{"base_url": "://invalid-url"},
|
||
}
|
||
|
||
_, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte("{}"), "token", false, "", false)
|
||
if err == nil {
|
||
t.Fatalf("expected error for invalid base_url when allowlist disabled")
|
||
}
|
||
}
|
||
|
||
func TestOpenAIValidateUpstreamBaseURLDisabledRequiresHTTPS(t *testing.T) {
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
if _, err := svc.validateUpstreamBaseURL("http://not-https.example.com"); err == nil {
|
||
t.Fatalf("expected http to be rejected when allow_insecure_http is false")
|
||
}
|
||
normalized, err := svc.validateUpstreamBaseURL("https://example.com")
|
||
if err != nil {
|
||
t.Fatalf("expected https to be allowed when allowlist disabled, got %v", err)
|
||
}
|
||
if normalized != "https://example.com" {
|
||
t.Fatalf("expected raw url passthrough, got %q", normalized)
|
||
}
|
||
}
|
||
|
||
func TestOpenAIValidateUpstreamBaseURLDisabledAllowsHTTP(t *testing.T) {
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{
|
||
Enabled: false,
|
||
AllowInsecureHTTP: true,
|
||
},
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
normalized, err := svc.validateUpstreamBaseURL("http://not-https.example.com")
|
||
if err != nil {
|
||
t.Fatalf("expected http allowed when allow_insecure_http is true, got %v", err)
|
||
}
|
||
if normalized != "http://not-https.example.com" {
|
||
t.Fatalf("expected raw url passthrough, got %q", normalized)
|
||
}
|
||
}
|
||
|
||
func TestOpenAIValidateUpstreamBaseURLEnabledEnforcesAllowlist(t *testing.T) {
|
||
cfg := &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{
|
||
Enabled: true,
|
||
UpstreamHosts: []string{"example.com"},
|
||
},
|
||
},
|
||
}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
|
||
if _, err := svc.validateUpstreamBaseURL("https://example.com"); err != nil {
|
||
t.Fatalf("expected allowlisted host to pass, got %v", err)
|
||
}
|
||
if _, err := svc.validateUpstreamBaseURL("https://evil.com"); err == nil {
|
||
t.Fatalf("expected non-allowlisted host to fail")
|
||
}
|
||
}
|
||
|
||
func TestOpenAIUpdateCodexUsageSnapshotFromHeaders(t *testing.T) {
|
||
repo := &snapshotUpdateAccountRepo{updateExtraCalls: make(chan map[string]any, 1)}
|
||
svc := &OpenAIGatewayService{accountRepo: repo}
|
||
headers := http.Header{}
|
||
headers.Set("x-codex-primary-used-percent", "12")
|
||
headers.Set("x-codex-secondary-used-percent", "34")
|
||
headers.Set("x-codex-primary-window-minutes", "300")
|
||
headers.Set("x-codex-secondary-window-minutes", "10080")
|
||
headers.Set("x-codex-primary-reset-after-seconds", "600")
|
||
headers.Set("x-codex-secondary-reset-after-seconds", "86400")
|
||
|
||
svc.UpdateCodexUsageSnapshotFromHeaders(context.Background(), 123, headers)
|
||
|
||
select {
|
||
case updates := <-repo.updateExtraCalls:
|
||
require.Equal(t, 12.0, updates["codex_5h_used_percent"])
|
||
require.Equal(t, 34.0, updates["codex_7d_used_percent"])
|
||
require.Equal(t, 600, updates["codex_5h_reset_after_seconds"])
|
||
require.Equal(t, 86400, updates["codex_7d_reset_after_seconds"])
|
||
case <-time.After(2 * time.Second):
|
||
t.Fatal("expected UpdateExtra to be called")
|
||
}
|
||
}
|
||
|
||
func TestOpenAIResponsesRequestPathSuffix(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
|
||
tests := []struct {
|
||
name string
|
||
path string
|
||
want string
|
||
}{
|
||
{name: "exact v1 responses", path: "/v1/responses", want: ""},
|
||
{name: "compact v1 responses", path: "/v1/responses/compact", want: "/compact"},
|
||
{name: "compact alias responses", path: "/responses/compact/", want: "/compact"},
|
||
{name: "nested suffix", path: "/openai/v1/responses/compact/detail", want: "/compact/detail"},
|
||
{name: "unrelated path", path: "/v1/chat/completions", want: ""},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
c.Request = httptest.NewRequest(http.MethodPost, tt.path, nil)
|
||
require.Equal(t, tt.want, openAIResponsesRequestPathSuffix(c))
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestNormalizeOpenAICompactRequestBodyPreservesCurrentCodexPayloadFields(t *testing.T) {
|
||
body := []byte(`{"model":"gpt-5.5","input":[{"type":"message","role":"user","content":"compact me"}],"instructions":"compact-test","tools":[{"type":"function","name":"shell"}],"parallel_tool_calls":true,"reasoning":{"effort":"high"},"text":{"verbosity":"low"},"previous_response_id":"resp_123","store":true,"stream":true,"prompt_cache_key":"cache_123"}`)
|
||
|
||
normalized, changed, err := normalizeOpenAICompactRequestBody(body)
|
||
|
||
require.NoError(t, err)
|
||
require.True(t, changed)
|
||
require.Equal(t, "gpt-5.5", gjson.GetBytes(normalized, "model").String())
|
||
require.True(t, gjson.GetBytes(normalized, "tools").Exists())
|
||
require.True(t, gjson.GetBytes(normalized, "parallel_tool_calls").Bool())
|
||
require.Equal(t, "high", gjson.GetBytes(normalized, "reasoning.effort").String())
|
||
require.Equal(t, "low", gjson.GetBytes(normalized, "text.verbosity").String())
|
||
require.Equal(t, "resp_123", gjson.GetBytes(normalized, "previous_response_id").String())
|
||
require.False(t, gjson.GetBytes(normalized, "store").Exists())
|
||
require.False(t, gjson.GetBytes(normalized, "stream").Exists())
|
||
require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists())
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesCompactPath(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
account := &Account{Type: AccountTypeOAuth}
|
||
|
||
req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
|
||
require.NoError(t, err)
|
||
require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
|
||
require.Equal(t, "application/json", req.Header.Get("Accept"))
|
||
require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
|
||
require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
|
||
require.NotEmpty(t, req.Header.Get("Session_Id"))
|
||
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesExplicitAPIKeyBetaHeader(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
|
||
c.Request.Header.Set("OpenAI-Beta", "api-key-specific-beta")
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||
},
|
||
}}
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
|
||
require.NoError(t, err)
|
||
require.Equal(t, "api-key-specific-beta", req.Header.Get("OpenAI-Beta"), "OAuth-only backport must not alter API-key passthrough headers")
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestCompactForcesJSONAcceptForOAuth(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
account := &Account{
|
||
Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
|
||
}
|
||
|
||
req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", true)
|
||
require.NoError(t, err)
|
||
require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
|
||
require.Equal(t, "application/json", req.Header.Get("Accept"))
|
||
require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
|
||
require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
|
||
require.NotEmpty(t, req.Header.Get("Session_Id"))
|
||
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestOAuthMessagesBridgeUsesSessionOnly(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
body := []byte(`{"model":"gpt-5.5","prompt_cache_key":"anthropic-metadata-session-1","input":[{"type":"message","role":"developer","content":[{"type":"input_text","text":"<sub2api-claude-code-todo-guard>"}]},{"type":"message","role":"user","content":"hello"}]}`)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||
c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
|
||
c.Request.Header.Set("originator", "codex_cli_rs")
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
account := &Account{
|
||
Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
|
||
}
|
||
|
||
req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, body, "token", true, "anthropic-metadata-session-1", false)
|
||
require.NoError(t, err)
|
||
require.NotEmpty(t, req.Header.Get("Session_Id"))
|
||
require.Empty(t, req.Header.Get("Conversation_Id"))
|
||
require.Empty(t, req.Header.Get("OpenAI-Beta"))
|
||
require.Empty(t, req.Header.Get("originator"))
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestPreservesCompactPathForAPIKeyBaseURL(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/responses/compact", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||
},
|
||
}}
|
||
account := &Account{
|
||
Type: AccountTypeAPIKey,
|
||
Platform: PlatformOpenAI,
|
||
Credentials: map[string]any{"base_url": "https://example.com/v1"},
|
||
}
|
||
|
||
req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", false)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "https://example.com/v1/responses/compact", req.URL.String())
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestPreservesCodexIdentityHeaders(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
|
||
c.Request.Header.Set("X-Codex-Window-ID", "window-http")
|
||
c.Request.Header.Set("X-Codex-Installation-ID", "installation-http")
|
||
c.Request.Header.Set("X-Test", "blocked")
|
||
|
||
body := []byte(`{"model":"gpt-5","input":"hello"}`)
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{
|
||
Security: config.SecurityConfig{
|
||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||
},
|
||
}}
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, body, "token", false, "", true)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "window-http", req.Header.Get("X-Codex-Window-ID"))
|
||
require.Equal(t, "installation-http", req.Header.Get("X-Codex-Installation-ID"))
|
||
require.Empty(t, req.Header.Get("X-Test"))
|
||
require.True(t, openai.EvaluateEngineFingerprint(req.Header, body, openai.DefaultEngineFingerprintSignals))
|
||
}
|
||
|
||
func TestOpenAIBuildUpstreamRequestOAuthOfficialClientOriginatorCompatibility(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
|
||
// 强制统一出口:客户端自报的 originator / User-Agent 都不参与上游身份构造,
|
||
// 一律改写为网关规范身份,天然满足 originator 与 UA 首段配套的上游校验(issue #3901)。
|
||
tests := []struct {
|
||
name string
|
||
userAgent string
|
||
originator string
|
||
}{
|
||
{name: "official desktop ua", userAgent: "Codex Desktop/1.2.3"},
|
||
{
|
||
name: "mismatched originator",
|
||
userAgent: "codex_vscode/0.140.2 (Mac OS X 14.0; arm64) vscode (codex_vscode; 0.140.2)",
|
||
originator: "codex_cli_rs",
|
||
},
|
||
{
|
||
name: "tui identity",
|
||
userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
|
||
originator: "codex-tui",
|
||
},
|
||
{name: "official originator without ua", originator: "codex_vscode"},
|
||
{name: "third-party ua", userAgent: "luna/1.2.0"},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
|
||
if tt.userAgent != "" {
|
||
c.Request.Header.Set("User-Agent", tt.userAgent)
|
||
}
|
||
if tt.originator != "" {
|
||
c.Request.Header.Set("originator", tt.originator)
|
||
}
|
||
|
||
svc := &OpenAIGatewayService{}
|
||
account := &Account{
|
||
Type: AccountTypeOAuth,
|
||
Credentials: map[string]any{"chatgpt_account_id": "chatgpt-acc"},
|
||
}
|
||
|
||
isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator"))
|
||
req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", isCodexCLI)
|
||
require.NoError(t, err)
|
||
require.Equal(t, openai.CodexDefaultOriginator, req.Header.Get("originator"))
|
||
require.Equal(t, codexCLIUserAgent, req.Header.Get("User-Agent"))
|
||
require.Equal(t, codexCLIVersion, req.Header.Get("version"))
|
||
})
|
||
}
|
||
}
|
||
|
||
// ==================== P1-08 修复:model 替换性能优化测试 ====================
|
||
|
||
// ==================== P1-08 修复:model 替换性能优化测试 =============
|
||
func TestReplaceModelInSSELine(t *testing.T) {
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
tests := []struct {
|
||
name string
|
||
line string
|
||
from string
|
||
to string
|
||
expected string
|
||
}{
|
||
{
|
||
name: "顶层 model 字段替换",
|
||
line: `data: {"id":"chatcmpl-123","model":"gpt-4o","choices":[]}`,
|
||
from: "gpt-4o",
|
||
to: "my-custom-model",
|
||
expected: `data: {"id":"chatcmpl-123","model":"my-custom-model","choices":[]}`,
|
||
},
|
||
{
|
||
name: "嵌套 response.model 替换",
|
||
line: `data: {"type":"response","response":{"id":"resp-1","model":"gpt-4o","output":[]}}`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: {"type":"response","response":{"id":"resp-1","model":"my-model","output":[]}}`,
|
||
},
|
||
{
|
||
name: "model 不匹配时不替换",
|
||
line: `data: {"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: {"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
|
||
},
|
||
{
|
||
name: "无 model 字段时不替换",
|
||
line: `data: {"id":"chatcmpl-123","choices":[]}`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: {"id":"chatcmpl-123","choices":[]}`,
|
||
},
|
||
{
|
||
name: "空 data 行",
|
||
line: `data: `,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: `,
|
||
},
|
||
{
|
||
name: "[DONE] 行",
|
||
line: `data: [DONE]`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: [DONE]`,
|
||
},
|
||
{
|
||
name: "非 data: 前缀行",
|
||
line: `event: message`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `event: message`,
|
||
},
|
||
{
|
||
name: "非法 JSON 不替换",
|
||
line: `data: {invalid json}`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: {invalid json}`,
|
||
},
|
||
{
|
||
name: "无空格 data: 格式",
|
||
line: `data:{"id":"x","model":"gpt-4o"}`,
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: `data: {"id":"x","model":"my-model"}`,
|
||
},
|
||
{
|
||
name: "model 名含特殊字符",
|
||
line: `data: {"model":"org/model-v2.1-beta"}`,
|
||
from: "org/model-v2.1-beta",
|
||
to: "custom/alias",
|
||
expected: `data: {"model":"custom/alias"}`,
|
||
},
|
||
{
|
||
name: "空行",
|
||
line: "",
|
||
from: "gpt-4o",
|
||
to: "my-model",
|
||
expected: "",
|
||
},
|
||
{
|
||
name: "保持其他字段不变",
|
||
line: `data: {"id":"abc","object":"chat.completion.chunk","model":"gpt-4o","created":1234567890,"choices":[{"index":0,"delta":{"content":"hi"}}]}`,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: `data: {"id":"abc","object":"chat.completion.chunk","model":"alias","created":1234567890,"choices":[{"index":0,"delta":{"content":"hi"}}]}`,
|
||
},
|
||
{
|
||
name: "顶层优先于嵌套:同时存在两个 model",
|
||
line: `data: {"model":"gpt-4o","response":{"model":"gpt-4o"}}`,
|
||
from: "gpt-4o",
|
||
to: "replaced",
|
||
expected: `data: {"model":"replaced","response":{"model":"gpt-4o"}}`,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
got := svc.replaceModelInSSELine(tt.line, tt.from, tt.to)
|
||
require.Equal(t, tt.expected, got)
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestReplaceModelInSSEBody(t *testing.T) {
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
tests := []struct {
|
||
name string
|
||
body string
|
||
from string
|
||
to string
|
||
expected string
|
||
}{
|
||
{
|
||
name: "多行 SSE body 替换",
|
||
body: "data: {\"model\":\"gpt-4o\",\"choices\":[]}\n\ndata: {\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n",
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: "data: {\"model\":\"alias\",\"choices\":[]}\n\ndata: {\"model\":\"alias\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n",
|
||
},
|
||
{
|
||
name: "无需替换的 body",
|
||
body: "data: {\"model\":\"gpt-3.5-turbo\"}\n\ndata: [DONE]\n",
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: "data: {\"model\":\"gpt-3.5-turbo\"}\n\ndata: [DONE]\n",
|
||
},
|
||
{
|
||
name: "混合 event 和 data 行",
|
||
body: "event: message\ndata: {\"model\":\"gpt-4o\"}\n\n",
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: "event: message\ndata: {\"model\":\"alias\"}\n\n",
|
||
},
|
||
{
|
||
name: "空 body",
|
||
body: "",
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: "",
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
got := svc.replaceModelInSSEBody(tt.body, tt.from, tt.to)
|
||
require.Equal(t, tt.expected, got)
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestReplaceModelInResponseBody(t *testing.T) {
|
||
svc := &OpenAIGatewayService{}
|
||
|
||
tests := []struct {
|
||
name string
|
||
body string
|
||
from string
|
||
to string
|
||
expected string
|
||
}{
|
||
{
|
||
name: "替换顶层 model",
|
||
body: `{"id":"chatcmpl-123","model":"gpt-4o","choices":[]}`,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: `{"id":"chatcmpl-123","model":"alias","choices":[]}`,
|
||
},
|
||
{
|
||
name: "model 不匹配不替换",
|
||
body: `{"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: `{"id":"chatcmpl-123","model":"gpt-3.5-turbo","choices":[]}`,
|
||
},
|
||
{
|
||
name: "无 model 字段不替换",
|
||
body: `{"id":"chatcmpl-123","choices":[]}`,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: `{"id":"chatcmpl-123","choices":[]}`,
|
||
},
|
||
{
|
||
name: "非法 JSON 返回原值",
|
||
body: `not json`,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: `not json`,
|
||
},
|
||
{
|
||
name: "空 body 返回原值",
|
||
body: ``,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: ``,
|
||
},
|
||
{
|
||
name: "保持嵌套结构不变",
|
||
body: `{"model":"gpt-4o","usage":{"prompt_tokens":10,"completion_tokens":20},"choices":[{"message":{"role":"assistant","content":"hello"}}]}`,
|
||
from: "gpt-4o",
|
||
to: "alias",
|
||
expected: `{"model":"alias","usage":{"prompt_tokens":10,"completion_tokens":20},"choices":[{"message":{"role":"assistant","content":"hello"}}]}`,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
got := svc.replaceModelInResponseBody([]byte(tt.body), tt.from, tt.to)
|
||
require.Equal(t, tt.expected, string(got))
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestExtractOpenAISSEDataLine(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
line string
|
||
wantData string
|
||
wantOK bool
|
||
}{
|
||
{name: "标准格式", line: `data: {"type":"x"}`, wantData: `{"type":"x"}`, wantOK: true},
|
||
{name: "无空格格式", line: `data:{"type":"x"}`, wantData: `{"type":"x"}`, wantOK: true},
|
||
{name: "纯空数据", line: `data: `, wantData: ``, wantOK: true},
|
||
{name: "非 data 行", line: `event: message`, wantData: ``, wantOK: false},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
got, ok := extractOpenAISSEDataLine(tt.line)
|
||
require.Equal(t, tt.wantOK, ok)
|
||
require.Equal(t, tt.wantData, got)
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestParseSSEUsage_SelectiveParsing(t *testing.T) {
|
||
svc := &OpenAIGatewayService{}
|
||
usage := &OpenAIUsage{InputTokens: 9, OutputTokens: 8, CacheReadInputTokens: 7}
|
||
|
||
// 非 completed 事件,不应覆盖 usage
|
||
svc.parseSSEUsage(`{"type":"response.in_progress","response":{"usage":{"input_tokens":1,"output_tokens":2}}}`, usage)
|
||
require.Equal(t, 9, usage.InputTokens)
|
||
require.Equal(t, 8, usage.OutputTokens)
|
||
require.Equal(t, 7, usage.CacheReadInputTokens)
|
||
|
||
// completed 事件,应提取 usage
|
||
svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":2}}}}`, usage)
|
||
require.Equal(t, 3, usage.InputTokens)
|
||
require.Equal(t, 5, usage.OutputTokens)
|
||
require.Equal(t, 2, usage.CacheReadInputTokens)
|
||
|
||
// done 事件同样可能携带最终 usage
|
||
svc.parseSSEUsage(`{"type":"response.done","response":{"usage":{"input_tokens":13,"output_tokens":15,"input_tokens_details":{"cached_tokens":4}}}}`, usage)
|
||
require.Equal(t, 13, usage.InputTokens)
|
||
require.Equal(t, 15, usage.OutputTokens)
|
||
require.Equal(t, 4, usage.CacheReadInputTokens)
|
||
|
||
// failed 事件在部分上游路径也会携带已消耗 usage,应与 WS/passthrough 保持一致
|
||
svc.parseSSEUsage(`{"type":"response.failed","response":{"usage":{"input_tokens":17,"output_tokens":19,"input_tokens_details":{"cached_tokens":6}}}}`, usage)
|
||
require.Equal(t, 17, usage.InputTokens)
|
||
require.Equal(t, 19, usage.OutputTokens)
|
||
require.Equal(t, 6, usage.CacheReadInputTokens)
|
||
|
||
svc.parseSSEUsage(`{"type":"response.completed","response":{"usage":{"prompt_tokens":21,"completion_tokens":8,"prompt_tokens_details":{"cached_tokens":6}}}}`, usage)
|
||
require.Equal(t, 21, usage.InputTokens)
|
||
require.Equal(t, 8, usage.OutputTokens)
|
||
require.Equal(t, 6, usage.CacheReadInputTokens)
|
||
}
|
||
|
||
func TestExtractOpenAIUsageFromJSONBytes_AcceptsResponseAndChatUsageShapes(t *testing.T) {
|
||
usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":9,"output_tokens":5,"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":4}}}`))
|
||
require.True(t, ok)
|
||
require.Equal(t, 9, usage.InputTokens)
|
||
require.Equal(t, 5, usage.OutputTokens)
|
||
require.Equal(t, 2, usage.CacheReadInputTokens)
|
||
require.Equal(t, 4, usage.CacheCreationInputTokens)
|
||
|
||
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":13,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":4,"cache_creation_tokens":3}}}}`))
|
||
require.True(t, ok)
|
||
require.Equal(t, 13, usage.InputTokens)
|
||
require.Equal(t, 7, usage.OutputTokens)
|
||
require.Equal(t, 4, usage.CacheReadInputTokens)
|
||
require.Equal(t, 3, usage.CacheCreationInputTokens)
|
||
|
||
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":11,"output_tokens":2,"cache_write_input_tokens":6}}`))
|
||
require.True(t, ok)
|
||
require.Equal(t, 6, usage.CacheCreationInputTokens)
|
||
|
||
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":7}}}`))
|
||
require.True(t, ok)
|
||
require.Equal(t, 7, usage.CacheCreationInputTokens, "官方嵌套字段应优先于兼容顶层别名")
|
||
|
||
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":0}}}`))
|
||
require.True(t, ok)
|
||
require.Zero(t, usage.CacheCreationInputTokens, "官方嵌套字段显式为零时仍应优先于兼容顶层别名")
|
||
|
||
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_read_input_tokens":19,"input_tokens_details":{"cached_tokens":0}}}`))
|
||
require.True(t, ok)
|
||
require.Zero(t, usage.CacheReadInputTokens, "官方嵌套缓存读取字段显式为零时仍应优先于兼容顶层别名")
|
||
}
|
||
|
||
func TestExtractCodexFinalResponse_SampleReplay(t *testing.T) {
|
||
body := strings.Join([]string{
|
||
`event: message`,
|
||
`data: {"type":"response.in_progress","response":{"id":"resp_1"}}`,
|
||
`data: {"type":"response.completed","response":{"id":"resp_1","model":"gpt-4o","usage":{"input_tokens":11,"output_tokens":22,"input_tokens_details":{"cached_tokens":3}}}}`,
|
||
`data: [DONE]`,
|
||
}, "\n")
|
||
|
||
finalResp, ok := extractCodexFinalResponse(body)
|
||
require.True(t, ok)
|
||
require.Contains(t, string(finalResp), `"id":"resp_1"`)
|
||
require.Contains(t, string(finalResp), `"input_tokens":11`)
|
||
}
|
||
|
||
func TestHandleSSEToJSON_CompletedEventReturnsJSON(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
}
|
||
body := []byte(strings.Join([]string{
|
||
`data: {"type":"response.in_progress","response":{"id":"resp_2"}}`,
|
||
`data: {"type":"response.completed","response":{"id":"resp_2","model":"gpt-4o","usage":{"input_tokens":7,"output_tokens":9,"input_tokens_details":{"cached_tokens":1}}}}`,
|
||
`data: [DONE]`,
|
||
}, "\n"))
|
||
|
||
usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, usage)
|
||
require.Equal(t, 7, usage.InputTokens)
|
||
require.Equal(t, 9, usage.OutputTokens)
|
||
require.Equal(t, 1, usage.CacheReadInputTokens)
|
||
// Header 可能由上游 Content-Type 透传;关键是 body 已转换为最终 JSON 响应。
|
||
require.NotContains(t, rec.Body.String(), "event:")
|
||
require.Contains(t, rec.Body.String(), `"id":"resp_2"`)
|
||
require.NotContains(t, rec.Body.String(), "data:")
|
||
}
|
||
|
||
func TestHandleNonStreamingResponse_APIKeyFallsBackToSSEBodyWhenContentTypeIsWrong(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
`data: {"type":"response.output_text.delta","delta":"hel"}`,
|
||
`data: {"type":"response.output_text.delta","delta":"lo"}`,
|
||
`data: {"type":"response.completed","response":{"id":"resp_api_key_sse","object":"response","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}`,
|
||
`data: [DONE]`,
|
||
}, "\n"))),
|
||
}
|
||
account := &Account{ID: 1, Type: AccountTypeAPIKey}
|
||
|
||
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.Equal(t, 3, result.InputTokens)
|
||
require.Equal(t, 2, result.OutputTokens)
|
||
require.NotContains(t, rec.Body.String(), "data:")
|
||
require.Equal(t, "resp_api_key_sse", gjson.Get(rec.Body.String(), "id").String())
|
||
require.Equal(t, "hello", gjson.Get(rec.Body.String(), "output.0.content.0.text").String())
|
||
}
|
||
|
||
func TestHandleNonStreamingResponse_OAuthJSONBodyWithDataEventTextKeepsJSONUsage(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
// Plain JSON compact response whose output text happens to contain the
|
||
// literal substrings "data:" and "event:" (e.g. echoing shell/log output).
|
||
// This must NOT be misdetected as SSE framing: it has a top-level usage
|
||
// object and no upstream text/event-stream Content-Type.
|
||
jsonBody := `{"id":"resp_oauth_compact","object":"response","model":"gpt-5.4","status":"completed",` +
|
||
`"output":[{"type":"message","content":[{"type":"output_text",` +
|
||
`"text":"processing data: 1,2,3 then event: click finished"}]}],` +
|
||
`"usage":{"input_tokens":11,"output_tokens":22,"total_tokens":33}}`
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||
Body: io.NopCloser(strings.NewReader(jsonBody)),
|
||
}
|
||
account := &Account{ID: 146, Type: AccountTypeOAuth}
|
||
|
||
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.4", "gpt-5.4")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
require.Equal(t, 11, result.InputTokens)
|
||
require.Equal(t, 22, result.OutputTokens)
|
||
// Response must remain the original JSON body (not routed through the SSE
|
||
// path, which would rewrite/lose the body or usage).
|
||
require.Equal(t, "application/json", rec.Header().Get("Content-Type"))
|
||
require.Equal(t, "resp_oauth_compact", gjson.Get(rec.Body.String(), "id").String())
|
||
require.Equal(t, int64(33), gjson.Get(rec.Body.String(), "usage.total_tokens").Int())
|
||
require.Contains(t, rec.Body.String(), "processing data: 1,2,3 then event: click finished")
|
||
}
|
||
|
||
func TestHandleNonStreamingResponse_ObservesUpstreamModelBeforeClientRewrite(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||
Body: io.NopCloser(strings.NewReader(
|
||
`{"id":"resp_model_audit","object":"response","model":"gpt-5.5","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}`,
|
||
)),
|
||
}
|
||
account := &Account{ID: 1, Type: AccountTypeAPIKey}
|
||
|
||
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.6-sol", "gpt-5.5")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, result)
|
||
|
||
// 客户端仍看到自己请求的模型名,审计观察器则保留改写前的上游声明。
|
||
require.Equal(t, "gpt-5.6-sol", gjson.Get(rec.Body.String(), "model").String())
|
||
require.Equal(t, "gpt-5.5", observedUpstreamResponseModel(c))
|
||
require.False(t, observedUpstreamResponseModelConflict(c))
|
||
}
|
||
|
||
func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
}
|
||
body := []byte(strings.Join([]string{
|
||
`data: {"type":"response.output_item.done","item":{"id":"ig_123","type":"image_generation_call","status":"generating","result":"aGVsbG8=","revised_prompt":"draw a cat","output_format":"png"}}`,
|
||
`data: {"type":"response.completed","response":{"id":"resp_img","model":"gpt-5.4","output":[],"usage":{"input_tokens":7,"output_tokens":9,"output_tokens_details":{"image_tokens":4}}}}`,
|
||
`data: [DONE]`,
|
||
}, "\n"))
|
||
|
||
usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-5.4", "gpt-5.4")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, usage)
|
||
require.Equal(t, 4, usage.ImageOutputTokens)
|
||
require.NotContains(t, rec.Body.String(), "data:")
|
||
require.Equal(t, "image_generation_call", gjson.Get(rec.Body.String(), "output.0.type").String())
|
||
require.Equal(t, "completed", gjson.Get(rec.Body.String(), "output.0.status").String())
|
||
require.Equal(t, "aGVsbG8=", gjson.Get(rec.Body.String(), "output.0.result").String())
|
||
require.Equal(t, "draw a cat", gjson.Get(rec.Body.String(), "output.0.revised_prompt").String())
|
||
}
|
||
|
||
func TestHandleSSEToJSON_NoFinalResponseKeepsSSEBody(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
}
|
||
body := []byte(strings.Join([]string{
|
||
`data: {"type":"response.in_progress","response":{"id":"resp_3"}}`,
|
||
`data: [DONE]`,
|
||
}, "\n"))
|
||
|
||
usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
|
||
require.NoError(t, err)
|
||
require.NotNil(t, usage)
|
||
require.Equal(t, 0, usage.InputTokens)
|
||
require.Contains(t, rec.Header().Get("Content-Type"), "text/event-stream")
|
||
require.Contains(t, rec.Body.String(), `data: {"type":"response.in_progress"`)
|
||
}
|
||
|
||
func TestHandleSSEToJSON_ResponseFailedReturnsProtocolError(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||
}
|
||
body := []byte(strings.Join([]string{
|
||
`data: {"type":"response.failed","error":{"message":"upstream rejected request"}}`,
|
||
`data: [DONE]`,
|
||
}, "\n"))
|
||
|
||
usage, err := svc.handleSSEToJSON(resp, c, nil, body, "gpt-4o", "gpt-4o")
|
||
require.Nil(t, usage)
|
||
require.Error(t, err)
|
||
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||
require.Contains(t, rec.Body.String(), "upstream rejected request")
|
||
require.Contains(t, rec.Header().Get("Content-Type"), "application/json")
|
||
}
|
||
|
||
func TestOpenAICompatSSEFrameParserResetsEventTypeAtFrameBoundary(t *testing.T) {
|
||
var parser openAICompatSSEFrameParser
|
||
|
||
frame, ok := parser.AddLine("event: response.created")
|
||
require.False(t, ok)
|
||
require.Empty(t, frame)
|
||
|
||
frame, ok = parser.AddLine(`data: {"response":{"id":"resp_1"}}`)
|
||
require.False(t, ok)
|
||
require.Empty(t, frame)
|
||
|
||
frame, ok = parser.AddLine("")
|
||
require.True(t, ok)
|
||
require.Equal(t, "response.created", frame.EventType)
|
||
require.JSONEq(t, `{"response":{"id":"resp_1"}}`, frame.Data)
|
||
|
||
frame, ok = parser.AddLine(`data: {"delta":"ok"}`)
|
||
require.False(t, ok)
|
||
require.Empty(t, frame.EventType)
|
||
|
||
frame, ok = parser.AddLine("")
|
||
require.True(t, ok)
|
||
require.Empty(t, frame.EventType)
|
||
require.JSONEq(t, `{"delta":"ok"}`, frame.Data)
|
||
}
|
||
|
||
func TestStreamingPassthroughCyberPolicyMarksAndPassesThrough(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"r1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","response":{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid-cyber"}},
|
||
}
|
||
|
||
_, err := svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, time.Now(), "m", "m")
|
||
require.Error(t, err)
|
||
var failoverErr *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &failoverErr), "cyber must NOT failover")
|
||
require.Contains(t, rec.Body.String(), "cyber_policy", "response.failed passed through to client")
|
||
mark := GetOpsCyberPolicy(c)
|
||
require.NotNil(t, mark)
|
||
require.Equal(t, "flagged for cyber policy", mark.Message)
|
||
}
|
||
|
||
func TestHandleStreamingResponseCyberPolicyMarks(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
cfg := &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}
|
||
svc := &OpenAIGatewayService{cfg: cfg}
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusOK,
|
||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||
"event: response.created",
|
||
`data: {"type":"response.created","response":{"id":"r1"}}`,
|
||
"",
|
||
"event: response.failed",
|
||
`data: {"type":"response.failed","error":{"code":"cyber_policy","message":"flagged"}}`,
|
||
"",
|
||
}, "\n"))),
|
||
Header: http.Header{"X-Request-Id": []string{"rid"}},
|
||
}
|
||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, time.Now(), "m", "m")
|
||
require.Error(t, err)
|
||
var fo *UpstreamFailoverError
|
||
require.False(t, errors.As(err, &fo))
|
||
require.NotNil(t, GetOpsCyberPolicy(c))
|
||
}
|
||
|
||
func TestHandleErrorResponseCyberPolicyPassthrough(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
cyberBody := `{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}`
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusBadRequest,
|
||
Header: http.Header{"Content-Type": []string{"application/json"}, "X-Request-Id": []string{"rid"}},
|
||
Body: io.NopCloser(strings.NewReader(cyberBody)),
|
||
}
|
||
_, err := svc.handleErrorResponse(context.Background(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, nil)
|
||
require.Error(t, err)
|
||
require.Equal(t, http.StatusBadRequest, rec.Code, "passthrough upstream 400, not rewrapped 502")
|
||
require.Contains(t, rec.Body.String(), "cyber_policy", "client sees original cyber body")
|
||
require.NotContains(t, rec.Body.String(), "Upstream request failed", "must not 502-rewrap")
|
||
mark := GetOpsCyberPolicy(c)
|
||
require.NotNil(t, mark)
|
||
require.Equal(t, http.StatusBadRequest, mark.UpstreamStatus)
|
||
}
|
||
|
||
func TestHandleCompatErrorResponseCyberPolicyEarlyReturn(t *testing.T) {
|
||
gin.SetMode(gin.TestMode)
|
||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||
rec := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(rec)
|
||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||
cyberBody := `{"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}`
|
||
resp := &http.Response{
|
||
StatusCode: http.StatusBadRequest,
|
||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||
Body: io.NopCloser(strings.NewReader(cyberBody)),
|
||
}
|
||
var gotStatus int
|
||
var gotType, gotMsg string
|
||
writeError := func(_ *gin.Context, statusCode int, errType, message string) {
|
||
gotStatus, gotType, gotMsg = statusCode, errType, message
|
||
}
|
||
// cyber 命中应早返回(写兼容错误 + 不冷却账号),而非落到通用 "Upstream request failed"。
|
||
_, err := svc.handleCompatErrorResponse(resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "a"}, writeError)
|
||
require.Error(t, err)
|
||
require.Equal(t, http.StatusBadRequest, gotStatus)
|
||
require.Equal(t, "invalid_request_error", gotType)
|
||
require.Contains(t, gotMsg, "flagged for cyber policy")
|
||
require.NotContains(t, gotMsg, "Upstream request failed")
|
||
require.NotNil(t, GetOpsCyberPolicy(c))
|
||
}
|