package controller import ( "bufio" "context" "errors" "fmt" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "sync" "sync/atomic" "testing" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/config" "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/alicebob/miniredis/v2" "github.com/gin-gonic/gin" "github.com/go-redis/redis/v8" "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) var responsesWSTestUserSequence atomic.Int64 func TestResponsesHTTPPoolRetriesTwiceThenUsesBackup(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {}) oldRetries := common.RetryTimes common.RetryTimes = 0 t.Cleanup(func() { common.RetryTimes = oldRetries }) oldGroups := setting.UserUsableGroups2JSONString() require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) t.Cleanup(func() { require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups)) }) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) if os.Getenv("TEST_PERSONAL_RATIO") == "1" { require.NoError(t, model.SavePersonalGroupRatios(context.Background(), fixture.user.Id, map[string]float64{"default": 0.8, "vip": 0.2})) } service.ResetGroupAvailability() t.Cleanup(service.ResetGroupAvailability) var primaryCalls, backupCalls atomic.Int64 fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) { primaryCalls.Add(1) w.Header().Set("Content-Type", "application/json") w.WriteHeader(500) _, _ = fmt.Fprint(w, `{"error":{"type":"server_error","code":"server_error","message":"test failure"}}`) } backup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { backupCalls.Add(1) w.Header().Set("Content-Type", "application/json") _, _ = fmt.Fprint(w, `{"id":"completed","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}`) })) t.Cleanup(backup.Close) require.NoError(t, model.DB.Model(fixture.channel).Updates(map[string]any{"key": "single-primary-key", "channel_info": `{}`, "auto_ban": 0}).Error) backupChannel := &model.Channel{Name: "backup-pool-test", Key: "test-only", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI, Group: "vip", Models: "ws-billing", BaseURL: &backup.URL, AutoBan: common.GetPointer(0)} require.NoError(t, model.DB.Create(backupChannel).Error) require.NoError(t, model.DB.Create(&model.Ability{ChannelId: backupChannel.Id, Model: "ws-billing", Group: "vip", Enabled: true}).Error) t.Cleanup(func() { require.NoError(t, model.DB.Where("channel_id = ?", backupChannel.Id).Delete(&model.Ability{}).Error) require.NoError(t, model.DB.Delete(backupChannel).Error) }) require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error) body, err := common.Marshal(model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}}) require.NoError(t, err) fixture.token.RoutingPolicy = string(body) fixture.token.Group = "default" fixture.token.CrossGroupRetry = true require.NoError(t, model.DB.Model(fixture.token).Updates(map[string]any{"routing_policy": string(body), "group": "default", "cross_group_retry": true}).Error) require.NoError(t, fixture.token.SelectUpdate()) for range 2 { req, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello"}`)) require.NoError(t, err) req.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key) req.Header.Set("Content-Type", "application/json") response, err := http.DefaultClient.Do(req) require.NoError(t, err) data, err := io.ReadAll(response.Body) require.NoError(t, err) require.NoError(t, response.Body.Close()) require.Equal(t, 200, response.StatusCode, string(data)) select { case <-fixture.httpDone: case <-time.After(3 * time.Second): t.Fatal("HTTP request did not finish") } } require.Equal(t, int64(2), primaryCalls.Load()) require.Equal(t, int64(2), backupCalls.Load()) var logs []model.Log require.NoError(t, model.LOG_DB.Where("token_id = ? AND type = ?", fixture.token.Id, model.LogTypeConsume).Find(&logs).Error) require.Len(t, logs, 2) for _, log := range logs { require.Equal(t, "vip", log.Group) expected := 1000 if os.Getenv("TEST_PERSONAL_RATIO") == "1" { expected = 200 } require.Equal(t, expected, log.Quota) } if os.Getenv("TEST_PERSONAL_RATIO") == "1" { assertResponsesWSAccounting(t, fixture, []int{200, 200}) } } func setupResponsesWSRequestTest(t *testing.T) (*model.User, *model.Token) { t.Helper() previousDB := model.DB previousLogDB := model.LOG_DB previousType := common.MainDatabaseType() previousLogType := common.LogDatabaseType() previousRedis := common.RedisEnabled previousRDB := common.RDB previousMaster, previousMemory, previousSQLite := common.IsMasterNode, common.MemoryCacheEnabled, common.SQLitePath previousEnabled := setting.ModelRequestRateLimitEnabled previousDuration := setting.ModelRequestRateLimitDurationMinutes previousTotal := setting.ModelRequestRateLimitCount previousSuccess := setting.ModelRequestRateLimitSuccessCount setting.ModelRequestRateLimitMutex.Lock() previousGroups := setting.ModelRequestRateLimitGroup setting.ModelRequestRateLimitGroup = nil setting.ModelRequestRateLimitMutex.Unlock() t.Setenv("SQL_DSN", os.Getenv("TEST_RESPONSES_SQL_DSN")) t.Setenv("LOG_SQL_DSN", os.Getenv("TEST_RESPONSES_LOG_SQL_DSN")) common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled = false, false, false common.SQLitePath = filepath.Join(t.TempDir(), "responses.db") require.NoError(t, model.InitDB()) db := model.DB sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) require.NoError(t, model.InitLogDB()) if model.LOG_DB != db { logSQL, err := model.LOG_DB.DB() require.NoError(t, err) logSQL.SetMaxOpenConns(1) t.Cleanup(func() { require.NoError(t, logSQL.Close()) }) } setting.ModelRequestRateLimitEnabled = false t.Cleanup(func() { model.DB = previousDB model.LOG_DB = previousLogDB common.SetDatabaseTypes(previousType, previousLogType) common.RedisEnabled = previousRedis common.RDB = previousRDB common.IsMasterNode, common.MemoryCacheEnabled, common.SQLitePath = previousMaster, previousMemory, previousSQLite setting.ModelRequestRateLimitEnabled = previousEnabled setting.ModelRequestRateLimitDurationMinutes = previousDuration setting.ModelRequestRateLimitCount = previousTotal setting.ModelRequestRateLimitSuccessCount = previousSuccess setting.ModelRequestRateLimitMutex.Lock() setting.ModelRequestRateLimitGroup = previousGroups setting.ModelRequestRateLimitMutex.Unlock() require.NoError(t, sqlDB.Close()) }) require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.PersonalGroupRatio{})) // The shared in-memory limiter outlives each database fixture. Give every // user a separate quota bucket, including when the tests run with -count. user := &model.User{Id: 5062000 + int(responsesWSTestUserSequence.Add(1)), Username: "responses-ws-user", Status: common.UserStatusEnabled, Group: "default", Quota: 1000, AuthVersion: 1} require.NoError(t, db.Create(user).Error) token := &model.Token{UserId: user.Id, Key: "responseswstoken", Status: common.TokenStatusEnabled, ExpiredTime: -1, RemainQuota: 100} require.NoError(t, db.Create(token).Error) t.Cleanup(func() { require.NoError(t, db.Unscoped().Delete(token).Error) require.NoError(t, db.Unscoped().Delete(user).Error) }) return user, token } func newResponsesWSTestRunner(t *testing.T, token *model.Token) (relay.ResponsesWSRequestRunner, *http.Request) { t.Helper() var runner relay.ResponsesWSRequestRunner engine := gin.New() require.NoError(t, engine.SetTrustedProxies([]string{"127.0.0.1"})) engine.GET("/v1/responses", middleware.TokenAuth(), func(c *gin.Context) { runner = newResponsesWSRequestRunner(c) }) request := httptest.NewRequest(http.MethodGet, "/v1/responses", nil) request.RemoteAddr = "127.0.0.1:12345" request.Header.Set("Authorization", "Bearer sk-"+token.Key) request.Header.Set("X-Forwarded-For", "203.0.113.8") request.Header.Set("Connection", "Upgrade") request.Header.Set("Upgrade", "websocket") request.Header.Set("Sec-WebSocket-Key", "handshake-key") request.Header.Set("Sec-WebSocket-Version", "13") request.Header.Set("Sec-WebSocket-Extensions", "permessage-deflate") request.Header.Set("Sec-WebSocket-Protocol", "responses") request.Header.Set("Content-Length", "99") request.Header.Set("Content-Encoding", "gzip") request.Header.Set("X-Application-Header", "preserved") response := httptest.NewRecorder() engine.ServeHTTP(response, request) require.Equal(t, http.StatusOK, response.Code) require.NotNil(t, runner) return runner, request } func TestResponsesWSRequestRunnerRefreshesBillingContextAndCleansBody(t *testing.T) { user, token := setupResponsesWSRequestTest(t) allowedIP := "203.0.113.8" require.NoError(t, model.DB.Model(token).Update("allow_ips", allowedIP).Error) runner, handshake := newResponsesWSTestRunner(t, token) // A later upstream header mutation must not replace the client's credential. handshake.Header.Set("Authorization", "Bearer upstream-secret") handshake.Header.Set("X-Forwarded-For", "198.51.100.99") var storage common.BodyStorage request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"first"}`)) require.Nil(t, runner(request, "ws-first", func(c *gin.Context) *types.NewAPIError { assert.Equal(t, user.Id, c.GetInt("id")) assert.Equal(t, 100, c.GetInt("token_quota")) assert.Equal(t, "203.0.113.8", c.ClientIP()) assert.Equal(t, "ws-first", c.GetString(common.RequestIdKey)) assert.Equal(t, "ws-first", c.Request.Context().Value(common.RequestIdKey)) for _, header := range []string{"Connection", "Upgrade", "Sec-WebSocket-Key", "Sec-WebSocket-Version", "Sec-WebSocket-Extensions", "Sec-WebSocket-Protocol", "Content-Length", "Content-Encoding"} { assert.Empty(t, c.GetHeader(header)) } assert.Equal(t, "preserved", c.GetHeader("X-Application-Header")) c.Set("previous_turn", true) var err error storage, err = common.GetBodyStorage(c) require.NoError(t, err) data, err := storage.Bytes() require.NoError(t, err) assert.JSONEq(t, `{"input":"first"}`, string(data)) return nil })) _, err := storage.Bytes() assert.ErrorIs(t, err, common.ErrStorageClosed) require.NoError(t, model.DB.Model(token).Updates(map[string]any{ "remain_quota": 2, "model_limits_enabled": true, "model_limits": "gpt-5.1", }).Error) require.NoError(t, model.DB.Model(user).Updates(map[string]any{ "group": "updated", "setting": `{"billing_preference":"wallet_only"}`, }).Error) request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"second"}`)) require.Nil(t, runner(request, "ws-second", func(c *gin.Context) *types.NewAPIError { assert.Equal(t, 2, c.GetInt("token_quota")) assert.True(t, c.GetBool("token_model_limit_enabled")) modelLimits, _ := c.Get("token_model_limit") assert.Equal(t, map[string]bool{"gpt-5.1": true}, modelLimits) assert.Equal(t, "updated", common.GetContextKeyString(c, constant.ContextKeyUsingGroup)) userSetting, ok := common.GetContextKeyType[dto.UserSetting](c, constant.ContextKeyUserSetting) require.True(t, ok) assert.Equal(t, "wallet_only", userSetting.BillingPreference) assert.Equal(t, "ws-second", c.GetString(common.RequestIdKey)) _, leaked := c.Get("previous_turn") assert.False(t, leaked) _, leaked = c.Get(common.KeyBodyStorage) assert.False(t, leaked) return nil })) } func TestResponsesWSRequestRunnerRejectsRevokedCredentials(t *testing.T) { for _, tc := range []struct { name string status int update func(*testing.T, *model.User, *model.Token) }{ {name: "disabled token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) { require.NoError(t, model.DB.Model(token).Update("status", common.TokenStatusDisabled).Error) }}, {name: "expired token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) { require.NoError(t, model.DB.Model(token).Update("expired_time", time.Now().Add(-time.Minute).Unix()).Error) }}, {name: "exhausted token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) { require.NoError(t, model.DB.Model(token).Update("remain_quota", 0).Error) }}, {name: "disabled user", status: http.StatusForbidden, update: func(t *testing.T, user *model.User, _ *model.Token) { require.NoError(t, model.DB.Model(user).Update("status", common.UserStatusDisabled).Error) }}, {name: "deleted token", status: http.StatusUnauthorized, update: func(t *testing.T, _ *model.User, token *model.Token) { require.NoError(t, model.DB.Delete(token).Error) }}, {name: "changed IP restriction", status: http.StatusForbidden, update: func(t *testing.T, _ *model.User, token *model.Token) { require.NoError(t, model.DB.Model(token).Update("allow_ips", "198.51.100.0/24").Error) }}, } { t.Run(tc.name, func(t *testing.T) { user, token := setupResponsesWSRequestTest(t) runner, _ := newResponsesWSTestRunner(t, token) require.Nil(t, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "before", func(*gin.Context) *types.NewAPIError { return nil })) tc.update(t, user, token) called := false apiError := runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "after", func(*gin.Context) *types.NewAPIError { called = true return nil }) require.NotNil(t, apiError) assert.Equal(t, tc.status, apiError.StatusCode) assert.False(t, called) assert.NotContains(t, apiError.Error(), token.Key) }) } } func TestResponsesWSRequestRunnerUsesExistingMemoryRateLimit(t *testing.T) { require.NoError(t, i18n.Init()) _, token := setupResponsesWSRequestTest(t) setting.ModelRequestRateLimitEnabled = true setting.ModelRequestRateLimitDurationMinutes = 1 setting.ModelRequestRateLimitCount = 1 setting.ModelRequestRateLimitSuccessCount = 100 runner, _ := newResponsesWSTestRunner(t, token) called := 0 handle := func(*gin.Context) *types.NewAPIError { called++; return nil } require.Nil(t, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "first", handle)) apiError := runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "limited", handle) require.NotNil(t, apiError) assert.Equal(t, http.StatusTooManyRequests, apiError.StatusCode) message := i18n.Translate("", i18n.MsgRateLimitTotalReached, map[string]any{"Minutes": 1, "Max": 1}) assert.Equal(t, common.MessageWithRequestId(message, "limited"), apiError.Error()) assert.Equal(t, 1, called) } func TestResponsesWSRequestRunnerSharesRedisSuccessLimitWithHTTP(t *testing.T) { _, token := setupResponsesWSRequestTest(t) redisServer := miniredis.RunT(t) client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) t.Cleanup(func() { require.NoError(t, client.Close()) }) common.RedisEnabled = true common.RDB = client setting.ModelRequestRateLimitEnabled = true setting.ModelRequestRateLimitDurationMinutes = 1 setting.ModelRequestRateLimitCount = 0 setting.ModelRequestRateLimitSuccessCount = 1 runner, _ := newResponsesWSTestRunner(t, token) failed := types.NewErrorWithStatusCode(errors.New("invalid input"), types.ErrorCodeInvalidRequest, http.StatusBadRequest) assert.Same(t, failed, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "failed", func(*gin.Context) *types.NewAPIError { return failed })) require.Nil(t, runner(httptest.NewRequest(http.MethodPost, "/v1/responses", nil), "successful", func(*gin.Context) *types.NewAPIError { return nil })) engine := gin.New() called := false engine.POST("/v1/responses", middleware.TokenAuth(), middleware.ModelRequestRateLimit(), func(c *gin.Context) { called = true }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) request.Header.Set("Authorization", "Bearer sk-"+token.Key) response := httptest.NewRecorder() engine.ServeHTTP(response, request) assert.Equal(t, http.StatusTooManyRequests, response.Code) assert.False(t, called) } type responsesWSBillingTest struct { httpUpstream http.HandlerFunc httpDone chan struct{} user *model.User token *model.Token channel *model.Channel client *websocket.Conn done chan struct{} upstreamDone chan struct{} connections atomic.Int32 gatewayURL string } func (fixture *responsesWSBillingTest) closeAndWait(t *testing.T) { t.Helper() _ = fixture.client.Close() select { case <-fixture.done: case <-time.After(3 * time.Second): t.Error("gateway handler did not stop after closing its client") } if fixture.connections.Load() > 0 { select { case <-fixture.upstreamDone: case <-time.After(3 * time.Second): t.Error("upstream connection was not closed") } } // Health classification is synchronous at the request boundary; only the // Redis write is asynchronous, see waitPerfCounters. } // waitPerfCounters waits for the asynchronous health samples of the fixture // model to reach wantRequests and returns the request and success counters. func waitPerfCounters(t *testing.T, wantRequests int64) (requests, successes int64) { t.Helper() require.Eventually(t, func() bool { requests, successes = 0, 0 keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result() if err != nil { return false } for _, key := range keys { n, _ := common.RDB.HGet(context.Background(), key, "req").Int64() requests += n n, _ = common.RDB.HGet(context.Background(), key, "ok").Int64() successes += n } return requests >= wantRequests }, 3*time.Second, 10*time.Millisecond) return requests, successes } func newResponsesWSBillingTest(t *testing.T, expression string, handle func(*websocket.Conn, *http.Request), httpEvents ...string) *responsesWSBillingTest { t.Helper() user, token := setupResponsesWSRequestTest(t) previousTimeout := constant.StreamingTimeout constant.StreamingTimeout = 30 t.Cleanup(func() { constant.StreamingTimeout = previousTimeout }) previousBatch, previousLogs, previousCount, previousQuota := common.BatchUpdateEnabled, common.LogConsumeEnabled, constant.CountToken, common.QuotaPerUnit common.BatchUpdateEnabled, common.LogConsumeEnabled, constant.CountToken, common.QuotaPerUnit = false, true, false, 500000 saved := map[string]string{} require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error { switch key { case "billing_setting.billing_mode", "billing_setting.billing_expr", "group_ratio_setting.group_ratio", "perf_metrics_setting.enabled": saved[key] = value } return nil })) t.Cleanup(func() { common.BatchUpdateEnabled, common.LogConsumeEnabled, constant.CountToken, common.QuotaPerUnit = previousBatch, previousLogs, previousCount, previousQuota require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) }) expressions, err := common.Marshal(map[string]string{"ws-billing": expression}) require.NoError(t, err) require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{ "billing_setting.billing_mode": `{"ws-billing":"tiered_expr"}`, "billing_setting.billing_expr": string(expressions), "group_ratio_setting.group_ratio": `{"default":1}`, "perf_metrics_setting.enabled": "true", })) require.NoError(t, model.DB.AutoMigrate(&model.Channel{}, &model.Ability{})) require.NoError(t, model.LOG_DB.AutoMigrate(&model.Log{})) require.NoError(t, model.DB.Model(user).Updates(map[string]any{"quota": 100000, "setting": `{"billing_preference":"wallet_only"}`}).Error) require.NoError(t, model.DB.Model(token).Update("remain_quota", 3000).Error) fixture := &responsesWSBillingTest{user: user, token: token, done: make(chan struct{}), upstreamDone: make(chan struct{}), httpDone: make(chan struct{}, 1)} redisServer := miniredis.RunT(t) redisClient := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) common.RDB, common.RedisEnabled = redisClient, true t.Cleanup(func() { require.NoError(t, redisClient.Close()) }) var upstreamClosed sync.Once upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if !websocket.IsWebSocketUpgrade(r) { if fixture.httpUpstream != nil { fixture.httpUpstream(w, r) return } w.Header().Set("Content-Type", "text/event-stream") for _, event := range httpEvents { _, _ = fmt.Fprintf(w, "data: %s\n\n", event) } _, _ = fmt.Fprint(w, "data: [DONE]\n\n") return } ws, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) if !assert.NoError(t, err) { return } fixture.connections.Add(1) defer ws.Close() defer upstreamClosed.Do(func() { close(fixture.upstreamDone) }) handle(ws, r) })) t.Cleanup(upstream.Close) channel := &model.Channel{Name: "responses-ws-upstream", Key: "upstream-first\nupstream-second", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI, Group: "default", Models: "ws-billing", BaseURL: &upstream.URL, ChannelInfo: model.ChannelInfo{IsMultiKey: true, MultiKeySize: 2, MultiKeyMode: constant.MultiKeyModePolling}} modelMapping := `{"ws-billing":"gpt-4o"}` channel.ModelMapping = &modelMapping channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true}) channel.SetOtherSettings(dto.ChannelOtherSettings{AllowServiceTier: true}) require.NoError(t, model.DB.Create(channel).Error) fixture.channel = channel require.NoError(t, model.DB.Create(&model.Ability{ChannelId: channel.Id, Model: "ws-billing", Group: "default", Enabled: true}).Error) t.Cleanup(func() { require.NoError(t, model.LOG_DB.Where("token_id = ?", token.Id).Delete(&model.Log{}).Error) require.NoError(t, model.DB.Where("channel_id = ?", channel.Id).Delete(&model.Ability{}).Error) require.NoError(t, model.DB.Delete(channel).Error) }) engine := gin.New() engine.GET("/v1/responses", middleware.TokenAuth(), func(c *gin.Context) { defer close(fixture.done) c.Set(common.RequestIdKey, "responses-ws-billing") ResponsesWebSocket(c) }) engine.POST("/v1/responses", middleware.TokenAuth(), middleware.ModelRequestRateLimit(), middleware.Distribute(), func(c *gin.Context) { defer func() { fixture.httpDone <- struct{}{} }() c.Set(common.RequestIdKey, "responses-http-billing") Relay(c, types.RelayFormatOpenAIResponses) }) gateway := httptest.NewServer(engine) fixture.gatewayURL = gateway.URL t.Cleanup(gateway.Close) client, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(gateway.URL, "http")+"/v1/responses", http.Header{"Authorization": []string{"Bearer sk-" + token.Key}}) require.NoError(t, err) fixture.client = client t.Cleanup(func() { fixture.closeAndWait(t) }) return fixture } func TestResponsesInterruptedStreamHealth(t *testing.T) { for _, scenario := range []string{"websocket timeout", "sse timeout", "sse client cancel"} { t.Run(scenario, func(t *testing.T) { events := []string{`{"type":"response.created","response":{"id":"partial","status":"in_progress"}}`, `{"type":"response.output_text.delta","delta":"hello"}`} fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } for _, event := range events { if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) { return } } _, _, _ = ws.ReadMessage() }) constant.StreamingTimeout = 1 finishUpstream := make(chan struct{}) fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") for _, event := range events { _, _ = fmt.Fprintf(w, "data: %s\n\n", event) } w.(http.Flusher).Flush() select { case <-r.Context().Done(): case <-finishUpstream: } } if scenario == "websocket timeout" { require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hello"}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) assert.Equal(t, "response.output_text.delta", readResponsesWSTestEvent(t, fixture.client)["type"]) _, _, err := fixture.client.ReadMessage() require.Error(t, err, "idle timeout closes the incomplete stream") } else { ctx, cancel := context.WithCancel(context.Background()) defer cancel() request, err := http.NewRequestWithContext(ctx, http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello","stream":true}`)) require.NoError(t, err) request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key) request.Header.Set("Content-Type", "application/json") response, err := (&http.Client{Timeout: 3 * time.Second}).Do(request) require.NoError(t, err) if scenario == "sse client cancel" { reader := bufio.NewReader(response.Body) for { line, err := reader.ReadString('\n') require.NoError(t, err) if strings.Contains(line, "response.output_text.delta") { break } } cancel() // Disconnected clients now drain the upstream; let this fixture // finish instead of waiting for the full production drain bound. close(finishUpstream) } else { _, err = io.Copy(io.Discard, response.Body) require.NoError(t, err) } require.NoError(t, response.Body.Close()) select { case <-fixture.httpDone: case <-time.After(3 * time.Second): t.Fatal("interrupted SSE request did not finish") } } fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) if scenario == "sse client cancel" { keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result() require.NoError(t, err) assert.Empty(t, keys, "client cancellation must not affect model health") return } requests, successes := waitPerfCounters(t, 1) assert.Equal(t, int64(1), requests) assert.Zero(t, successes) }) } } func readResponsesWSTestEvent(t *testing.T, client *websocket.Conn) map[string]any { t.Helper() require.NoError(t, client.SetReadDeadline(time.Now().Add(3*time.Second))) _, data, err := client.ReadMessage() require.NoError(t, err) var event map[string]any require.NoError(t, common.Unmarshal(data, &event)) return event } func assertResponsesWSAccounting(t *testing.T, fixture *responsesWSBillingTest, expectedQuotas []int) { t.Helper() var logs []model.Log require.NoError(t, model.LOG_DB.Where("type = ?", model.LogTypeConsume).Order("id").Find(&logs).Error) require.Len(t, logs, len(expectedQuotas)) var charged int for index, quota := range expectedQuotas { assert.Equal(t, quota, logs[index].Quota) charged += quota } require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error) require.NoError(t, model.DB.First(fixture.user, fixture.user.Id).Error) assert.Equal(t, 3000-charged, fixture.token.RemainQuota) assert.Equal(t, charged, fixture.token.UsedQuota) assert.Equal(t, 100000-charged, fixture.user.Quota) assert.Equal(t, charged, fixture.user.UsedQuota) } func TestResponsesWebSocketReusesConnectionAndSettlesEachRequest(t *testing.T) { type upstreamRequest struct { Authorization string Type string `json:"type"` Model string `json:"model"` PreviousResponseID string `json:"previous_response_id"` ServiceTier string `json:"service_tier"` } received := make(chan upstreamRequest, 3) fixture := newResponsesWSBillingTest(t, `param("service_tier") == "priority" ? tier("priority", p * 4) : tier("base", p * 2)`, func(ws *websocket.Conn, r *http.Request) { for index := 1; ; index++ { _, data, err := ws.ReadMessage() if err != nil { return } var event upstreamRequest if !assert.NoError(t, common.Unmarshal(data, &event)) { return } event.Authorization = r.Header.Get("Authorization") received <- event if index == 2 { duplicate := `{"type":"response.completed","response":{"id":"resp_1","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}` if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(duplicate))) { return } } terminal := fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_%d","status":"completed","model":"ws-billing","output":[],"usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`, index) if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(terminal))) { return } } }) client, user, token := fixture.client, fixture.user, fixture.token for index, payload := range []string{ `{"type":"response.create","model":"ws-billing","input":"first","service_tier":"default"}`, `{"type":"response.create","model":"ws-billing","input":"second","previous_response_id":"resp_1","service_tier":"priority"}`, } { require.NoError(t, client.WriteMessage(websocket.TextMessage, []byte(payload))) require.NoError(t, client.SetReadDeadline(time.Now().Add(3*time.Second))) _, data, err := client.ReadMessage() require.NoError(t, err) var terminal struct { Type string `json:"type"` Response struct { ID string `json:"id"` } `json:"response"` } require.NoError(t, common.Unmarshal(data, &terminal)) require.Equal(t, "response.completed", terminal.Type, "unexpected response: %s", data) assert.Equal(t, fmt.Sprintf("resp_%d", index+1), terminal.Response.ID) observed := <-received assert.Equal(t, "Bearer upstream-first", observed.Authorization) assert.Equal(t, "gpt-4o", observed.Model) assert.Equal(t, "response.create", observed.Type) if index == 0 { assert.Empty(t, observed.PreviousResponseID) assert.Equal(t, "default", observed.ServiceTier) } else { assert.Equal(t, "resp_1", observed.PreviousResponseID) assert.Equal(t, "priority", observed.ServiceTier) } } require.NoError(t, client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"exhausted"}`))) require.NoError(t, client.SetReadDeadline(time.Now().Add(3*time.Second))) _, data, err := client.ReadMessage() require.NoError(t, err) var rejection struct { Type string `json:"type"` Status int `json:"status"` } require.NoError(t, common.Unmarshal(data, &rejection)) assert.Equal(t, "error", rejection.Type) assert.Equal(t, http.StatusUnauthorized, rejection.Status) assert.Equal(t, int32(1), fixture.connections.Load()) assert.Empty(t, received, "exhausted token must be rejected before contacting upstream") var logs []model.Log require.NoError(t, model.LOG_DB.Where("type = ?", model.LogTypeConsume).Order("id").Find(&logs).Error) require.Len(t, logs, 2) for index, expectedQuota := range []int{1000, 2000} { assert.Equal(t, expectedQuota, logs[index].Quota) assert.Equal(t, fmt.Sprintf("responses-ws-billing-ws-%d", index), logs[index].RequestId) assert.Equal(t, 1000, logs[index].PromptTokens) assert.Equal(t, 10, logs[index].CompletionTokens) } // Routing decisions are persisted per request: the first create selects a // channel, later creates on the same connection reuse it as a pin. for index, wantDecision := range []string{"attempt:channel_selected", "select:pinned_channel"} { var other struct { AdminInfo struct { RequestPolicy []struct { ChannelID int `json:"channel_id"` Decision struct { Action string `json:"action"` Reason string `json:"reason"` } `json:"decision"` } `json:"request_policy"` } `json:"admin_info"` } require.NoError(t, common.UnmarshalJsonStr(logs[index].Other, &other)) var decisions []string for _, event := range other.AdminInfo.RequestPolicy { assert.Equal(t, fixture.channel.Id, event.ChannelID) decisions = append(decisions, event.Decision.Action+":"+event.Decision.Reason) } assert.Contains(t, decisions, wantDecision, "request %d policy events: %v", index, decisions) } require.NoError(t, model.DB.First(token, token.Id).Error) require.NoError(t, model.DB.First(user, user.Id).Error) assert.Zero(t, token.RemainQuota) assert.Equal(t, 3000, token.UsedQuota) assert.Equal(t, 97000, user.Quota) assert.Equal(t, 3000, user.UsedQuota) } // Both transports must reach the same upstream target with the same // credential placement and settle identically for every supported channel type. func TestResponsesWebSocketDialsNativeResponsesChannelTypes(t *testing.T) { type upstreamTarget struct { Path, Authorization, QueryKey string } queryAuthRoute := dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{ IncomingPath: "/v1/responses", UpstreamPath: "/upstream/responses", Auth: &dto.AdvancedCustomRouteAuth{Type: dto.AdvancedCustomAuthTypeQuery, Name: "api_key", Value: "{api_key}"}, }}}} bearer := func(key string) upstreamTarget { return upstreamTarget{Path: "/v1/responses", Authorization: "Bearer " + key} } for _, tc := range []struct { name string channelType int otherSettings dto.ChannelOtherSettings want func(key string) upstreamTarget }{ {name: "new api", channelType: constant.ChannelTypeNewAPI, want: bearer}, {name: "sub2api", channelType: constant.ChannelTypeSub2API, want: bearer}, {name: "advanced custom query auth", channelType: constant.ChannelTypeAdvancedCustom, otherSettings: queryAuthRoute, want: func(key string) upstreamTarget { return upstreamTarget{Path: "/upstream/responses", QueryKey: key} }}, } { t.Run(tc.name, func(t *testing.T) { const terminal = `{"type":"response.completed","response":{"id":"first","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}` targets := make(chan upstreamTarget, 2) observe := func(r *http.Request) { targets <- upstreamTarget{Path: r.URL.Path, Authorization: r.Header.Get("Authorization"), QueryKey: r.URL.Query().Get("api_key")} } fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, r *http.Request) { observe(r) if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(terminal))) { return } _, _, _ = ws.ReadMessage() }) fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) { observe(r) w.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprintf(w, "data: %s\n\ndata: [DONE]\n\n", terminal) } fixture.channel.SetOtherSettings(tc.otherSettings) require.NoError(t, model.DB.Model(fixture.channel).Updates(map[string]any{"type": tc.channelType, "settings": fixture.channel.OtherSettings}).Error) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"]) assert.Equal(t, tc.want("upstream-first"), <-targets) request, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hi","stream":true}`)) require.NoError(t, err) request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key) request.Header.Set("Content-Type", "application/json") response, err := (&http.Client{Timeout: 3 * time.Second}).Do(request) require.NoError(t, err) body, err := io.ReadAll(response.Body) require.NoError(t, err) require.NoError(t, response.Body.Close()) assert.Equal(t, http.StatusOK, response.StatusCode) assert.Contains(t, string(body), terminal) select { case <-fixture.httpDone: case <-time.After(3 * time.Second): t.Fatal("HTTP request did not finish") } // The polling multi-key channel rotates to its second key for the HTTP request. assert.Equal(t, tc.want("upstream-second"), <-targets) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000, 1000}) }) } } func TestResponsesWebSocketDisconnectSettlesDeliveredOutputOnce(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("output", c * 2)`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } for _, event := range []string{ `{"type":"response.created","response":{"id":"partial","status":"in_progress"}}`, `{"type":"response.output_text.delta","delta":"hello"}`, } { if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) { return } } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi","max_output_tokens":1}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) delta := readResponsesWSTestEvent(t, fixture.client) require.Equal(t, "response.output_text.delta", delta["type"]) assert.Equal(t, "hello", delta["delta"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1}) keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result() require.NoError(t, err) assert.Empty(t, keys, "client cancellation must not affect model health") } func TestResponsesWebSocketCancelErrorDoesNotFinishActiveRequest(t *testing.T) { complete := make(chan struct{}) defer close(complete) allowCreated := make(chan struct{}, 1) defer close(allowCreated) receivedCreate := make(chan struct{}, 1) fixture := newResponsesWSBillingTest(t, `tier("base", p * 2)`, func(ws *websocket.Conn, _ *http.Request) { _, create, err := ws.ReadMessage() if !assert.NoError(t, err) || !assert.Contains(t, string(create), `"type":"response.create"`) { return } receivedCreate <- struct{}{} <-allowCreated if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) { return } _, cancel, err := ws.ReadMessage() if !assert.NoError(t, err) || !assert.JSONEq(t, `{"type":"response.cancel","response_id":"wrong"}`, string(cancel)) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","status":400,"error":{"type":"invalid_request_error","code":"response_not_found","message":"No such response"}}`))) { return } <-complete if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`))) { return } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) select { case <-receivedCreate: case <-time.After(3 * time.Second): t.Fatal("upstream did not receive the initial request") } // Queue cancellation before the upstream accepts the response. The following // conflict is an acknowledgement that the client loop processed both frames. require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"wrong"}`))) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"before-acceptance"}`))) assert.Equal(t, float64(http.StatusConflict), readResponsesWSTestEvent(t, fixture.client)["status"]) allowCreated <- struct{}{} assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) cancelError := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, "error", cancelError["type"]) assert.Equal(t, float64(http.StatusBadRequest), cancelError["status"]) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"overlapping","stream_id":"other","event_id":"busy"}`))) conflict := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, "error", conflict["type"]) assert.Equal(t, float64(http.StatusConflict), conflict["status"]) assert.Equal(t, "other", conflict["stream_id"]) assert.Equal(t, "busy", conflict["event_id"]) complete <- struct{}{} assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) } func TestResponsesWebSocketInitialUpstreamRejectionRefundsReservation(t *testing.T) { for _, tc := range []struct { name, upstream, wantType, wantMessage string status int }{ {name: "structured error", upstream: `{"type":"error","response_id":"rejected","status":400,"error":{"type":"invalid_request_error","code":"invalid_input","message":"Invalid input"}}`, status: http.StatusBadRequest, wantType: "invalid_request_error", wantMessage: "Invalid input"}, // A frame without an error object is still reported as a request error. {name: "bare error", upstream: `{"type":"error","status":500,"message":"upstream rejected"}`, status: http.StatusInternalServerError, wantType: "invalid_request_error", wantMessage: "upstream rejected"}, } { t.Run(tc.name, func(t *testing.T) { preConsumed := make(chan int, 1) tokenID := make(chan int, 1) // Pre-consume no longer estimates completion tokens, so an output-priced expression reserves nothing to refund. fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } var token model.Token if !assert.NoError(t, model.DB.First(&token, <-tokenID).Error) { return } preConsumed <- token.RemainQuota if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(tc.upstream))) { return } _, _, _ = ws.ReadMessage() }) tokenID <- fixture.token.Id // Refund is asynchronous. Observe committed token writes instead of waiting // a fixed delay or returning while its worker still uses the test database. updates := make(chan struct{}, 4) require.NoError(t, model.DB.Callback().Update().After("gorm:commit_or_rollback_transaction").Register("responses-ws-refund", func(tx *gorm.DB) { if tx.Error == nil && tx.Statement.Table == "tokens" { select { case updates <- struct{}{}: default: } } })) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi","max_output_tokens":10}`))) rejection := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, "error", rejection["type"]) assert.Equal(t, float64(tc.status), rejection["status"]) rejectionError, _ := rejection["error"].(map[string]any) assert.Equal(t, tc.wantType, rejectionError["type"]) assert.Equal(t, tc.wantMessage, rejectionError["message"]) assert.Equal(t, 2000, <-preConsumed, "the rejected request reserved quota before contacting upstream") deadline := time.NewTimer(3 * time.Second) defer deadline.Stop() for { require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error) if fixture.token.RemainQuota == 3000 { break } select { case <-updates: case <-deadline.C: t.Fatal("initial rejection did not refund the token reservation") } } fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, nil) }) } } // TEST_RESPONSES_SQL_DSN / TEST_RESPONSES_LOG_SQL_DSN optionally run these // entry-point regressions against isolated real MySQL/PostgreSQL databases. func TestResponsesStreamOutcomesPreserveAccounting(t *testing.T) { for _, transport := range []string{"websocket", "http-sse"} { for _, tc := range []struct { name, expression, terminal string // sseTerminal is the flat error envelope used by the Responses SSE // protocol; the WebSocket protocol nests the error instead. sseTerminal string delta bool failed bool ignored bool }{ {name: "failed-null-fixed", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed","usage":null,"error":{"code":"server_error","message":"sensitive upstream detail"}}}`, failed: true}, {name: "failed-missing-fixed", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed"}}`, failed: true}, {name: "failed-actual-usage", expression: `tier("input", p * 2)`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`, failed: true}, {name: "failed-estimated-output", expression: `tier("output", c * 2000)`, terminal: `{"type":"response.failed","response":{"id":"first","status":"failed","usage":null}}`, delta: true, failed: true}, {name: "error-after-created", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"error","status":500,"error":{"type":"server_error","code":"server_error","message":"Internal server error"}}`, sseTerminal: `{"type":"error","code":"server_error","message":"Internal server error","param":null,"sequence_number":2}`, failed: true}, {name: "business-error-after-created", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"error","status":400,"error":{"type":"invalid_request_error","code":"context_length_exceeded","message":"Input too long"}}`, sseTerminal: `{"type":"error","code":"context_length_exceeded","message":"Input too long","param":null,"sequence_number":2}`, failed: true, ignored: true}, {name: "completed-at-output-limit", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.incomplete","response":{"id":"first","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}}`}, {name: "completed-zero-fixed", expression: `tier("request", fixed(0.002))`, terminal: `{"type":"response.completed","response":{"id":"first","status":"completed","usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`}, } { t.Run(transport+"/"+tc.name, func(t *testing.T) { events := []string{`{"type":"response.created","response":{"id":"first","status":"in_progress"}}`} if tc.delta { events = append(events, `{"type":"response.output_text.delta","delta":"hello"}`) } terminal := tc.terminal if transport == "http-sse" && tc.sseTerminal != "" { terminal = tc.sseTerminal } events = append(events, terminal) fixture := newResponsesWSBillingTest(t, tc.expression, func(ws *websocket.Conn, _ *http.Request) { for turn := 0; ; turn++ { _, request, err := ws.ReadMessage() if err != nil { return } var body map[string]any if !assert.NoError(t, common.Unmarshal(request, &body)) { return } assert.Equal(t, "planner", body["stream_id"]) assert.Equal(t, "gpt-4o", body["model"]) output := events if transport == "http-sse" || turn > 0 { output = []string{`{"type":"response.completed","stream_id":"planner","response":{"id":"second","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}}`} } for _, event := range output { if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) { return } } } }, events...) setting.ModelRequestRateLimitEnabled = true setting.ModelRequestRateLimitDurationMinutes = 1 setting.ModelRequestRateLimitSuccessCount = 1 setting.ModelRequestRateLimitCount = 0 if transport == "http-sse" { request, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hi","stream":true,"max_output_tokens":1}`)) require.NoError(t, err) request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key) request.Header.Set("Content-Type", "application/json") response, err := (&http.Client{Timeout: 3 * time.Second}).Do(request) require.NoError(t, err) body, err := io.ReadAll(response.Body) require.NoError(t, err) require.NoError(t, response.Body.Close()) assert.Equal(t, http.StatusOK, response.StatusCode) if tc.failed { assert.Contains(t, string(body), "model_service_error") assert.Contains(t, string(body), "responses-http-billing") assert.NotContains(t, string(body), "sensitive upstream detail") } else { assert.Contains(t, string(body), terminal) } } else { require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","stream_id":"planner","model":"ws-billing","input":"hi","max_output_tokens":1}`))) for _, expected := range events { actual, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client)) require.NoError(t, err) assert.JSONEq(t, expected, string(actual), "upstream frames must be delivered once and unchanged") } } quotas := []int{1000} if tc.failed { // The first terminal must complete settlement and middleware before // admitting this immediately following request. require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","stream_id":"planner","model":"ws-billing","input":"next","max_output_tokens":1}`))) assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"]) quotas = append(quotas, 1000) } require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","event_id":"limited","stream_id":"planner","model":"ws-billing","input":"limited"}`))) limited := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, float64(http.StatusTooManyRequests), limited["status"]) assert.Equal(t, "planner", limited["stream_id"]) assert.Equal(t, "limited", limited["event_id"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, quotas) var logs []model.Log require.NoError(t, model.LOG_DB.Where("token_id = ? AND type = ?", fixture.token.Id, model.LogTypeConsume).Order("id").Find(&logs).Error) require.Len(t, logs, len(quotas)) var other map[string]any require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other)) stream := other["stream_status"].(map[string]any) if tc.failed { assert.Equal(t, "error", stream["status"]) assert.Equal(t, "failed", stream["response_status"]) } else { assert.Equal(t, "ok", stream["status"]) if tc.name == "completed-at-output-limit" { assert.Equal(t, "incomplete", stream["response_status"]) } else { assert.Equal(t, "completed", stream["response_status"]) } } assert.NotContains(t, logs[0].Other, "sensitive upstream detail") expectedRequests := int64(len(quotas)) if tc.ignored { expectedRequests-- } requests, successes := waitPerfCounters(t, expectedRequests) assert.Equal(t, expectedRequests, requests) assert.Equal(t, int64(1), successes) }) } } } func TestResponsesHTTPHealthCountsFinalResult(t *testing.T) { for _, tc := range []struct { name, code string firstStatus, attempts int success bool ignored bool }{ {"business rejection", "context_length_exceeded", 400, 1, false, true}, {"credentials rejected as 400", "invalid_api_key", 400, 1, false, false}, {"retry succeeds", "server_error", 500, 2, true, false}, } { t.Run(tc.name, func(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {}) oldRetries := common.RetryTimes common.RetryTimes = 1 t.Cleanup(func() { common.RetryTimes = oldRetries }) var attempts atomic.Int64 upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") if attempts.Add(1) == 1 { w.WriteHeader(tc.firstStatus) _, _ = fmt.Fprintf(w, `{"error":{"type":"%s","code":"%s","message":"test rejection"}}`, tc.code, tc.code) return } _, _ = fmt.Fprint(w, `{"id":"completed","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}`) })) t.Cleanup(upstream.Close) require.NoError(t, model.DB.Model(&model.Channel{}).Where("name = ?", "responses-ws-upstream").Update("base_url", upstream.URL).Error) request, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello"}`)) require.NoError(t, err) request.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key) request.Header.Set("Content-Type", "application/json") response, err := http.DefaultClient.Do(request) require.NoError(t, err) _, err = io.Copy(io.Discard, response.Body) require.NoError(t, err) require.NoError(t, response.Body.Close()) select { case <-fixture.httpDone: case <-time.After(3 * time.Second): t.Fatal("HTTP request did not finish") } fixture.closeAndWait(t) assert.Equal(t, int64(tc.attempts), attempts.Load()) if !tc.success { assert.Equal(t, tc.firstStatus, response.StatusCode) } else { assert.Equal(t, http.StatusOK, response.StatusCode) } if tc.ignored { keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result() require.NoError(t, err) assert.Empty(t, keys) return } requests, successes := waitPerfCounters(t, 1) assert.Equal(t, int64(1), requests) if tc.success { assert.Equal(t, int64(1), successes) } else { assert.Zero(t, successes) } }) } } func TestResponsesWebSocketLocalErrorsKeepStreamIdentity(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) { t.Error("local rejection reached upstream") }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","event_id":"invalid","stream_id":"planner","model":"ws-billing","max_output_tokens":18446744073709551615}`))) invalid := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, float64(http.StatusBadRequest), invalid["status"]) assert.Equal(t, "planner", invalid["stream_id"]) assert.Equal(t, "invalid", invalid["event_id"]) fixture.token.Status = common.TokenStatusDisabled require.NoError(t, fixture.token.SelectUpdate()) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","event_id":"revoked","stream_id":"planner","model":"ws-billing","input":"hi"}`))) revoked := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, float64(http.StatusUnauthorized), revoked["status"]) assert.Equal(t, "planner", revoked["stream_id"]) assert.Equal(t, "revoked", revoked["event_id"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, nil) } func TestResponsesWebSocketAmbiguousControlErrorClosesAndSettlesOnce(t *testing.T) { const terminal = `{"type":"error","status":500,"error":{"type":"server_error","code":"server_error","message":"Internal server error"}}` fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) { return } _, cancel, err := ws.ReadMessage() if !assert.NoError(t, err) { return } assert.Contains(t, string(cancel), `"type":"response.cancel"`) if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(terminal))) { return } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"active"}`))) body, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client)) require.NoError(t, err) assert.JSONEq(t, terminal, string(body)) _, _, err = fixture.client.ReadMessage() // A dropped socket would read as a synthetic 1006; 1011 is the gateway's frame. var closeErr *websocket.CloseError require.ErrorAs(t, err, &closeErr, "an unattributed error must close the connection instead of remaining busy") assert.Equal(t, websocket.CloseInternalServerErr, closeErr.Code) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) } // Every way a session ends sends the client a close frame whose code says // why. A dropped socket would read as a synthetic 1006 close error, so a // received 1000, 1001, 1008 or 1011 proves the gateway sent the frame. func TestResponsesWebSocketCloseCodes(t *testing.T) { require.NoError(t, i18n.Init()) const ( created = `{"type":"response.created","response":{"id":"active","status":"in_progress"}}` delta = `{"type":"response.output_text.delta","delta":"hello"}` completed = `{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}` ) for _, tc := range []struct { name string events []string // upstream events after the create, all relayed // After the events, and after the client read them with waitClient, // upstream sends upstreamClose or drops the socket. Otherwise it reads // until the gateway closes the connection. upstreamClose int drop, waitClient bool timeout bool end func(*testing.T, *responsesWSBillingTest) code int reason string upstreamCode int // the close code upstream receives, when checked }{ {name: "upstream closes after completion", events: []string{completed}, upstreamClose: websocket.CloseNormalClosure, code: websocket.CloseNormalClosure, reason: "upstream connection closed"}, {name: "idle upstream goes away", events: []string{completed}, waitClient: true, upstreamClose: websocket.CloseGoingAway, code: websocket.CloseGoingAway, reason: "upstream connection closed"}, {name: "idle upstream drops", events: []string{completed}, waitClient: true, drop: true, code: websocket.CloseInternalServerErr, reason: "upstream connection closed"}, // A normal upstream close still truncates the active response. {name: "upstream closes mid-response", events: []string{created, delta}, upstreamClose: websocket.CloseNormalClosure, code: websocket.CloseInternalServerErr, reason: "upstream connection closed"}, {name: "upstream drops mid-response", events: []string{created, delta}, drop: true, code: websocket.CloseInternalServerErr, reason: "upstream connection closed"}, {name: "upstream response times out", events: []string{created}, timeout: true, code: websocket.CloseInternalServerErr, reason: "upstream response timed out", upstreamCode: websocket.CloseInternalServerErr}, {name: "client closes", events: []string{completed}, end: func(t *testing.T, fixture *responsesWSBillingTest) { require.NoError(t, fixture.client.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))) }, code: websocket.CloseNormalClosure, upstreamCode: websocket.CloseNormalClosure}, {name: "channel disabled", events: []string{completed}, end: func(t *testing.T, fixture *responsesWSBillingTest) { require.Equal(t, 1, service.CloseActiveWebSocketsForChannel(fixture.channel.Id, service.ChannelDisabledCloseReason)) }, code: websocket.ClosePolicyViolation, reason: service.ChannelDisabledCloseReason, upstreamCode: websocket.ClosePolicyViolation}, {name: "token revoked", events: []string{completed}, end: func(t *testing.T, fixture *responsesWSBillingTest) { fixture.token.Status = common.TokenStatusDisabled require.NoError(t, fixture.token.SelectUpdate()) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"again"}`))) assert.Equal(t, float64(http.StatusUnauthorized), readResponsesWSTestEvent(t, fixture.client)["status"]) }, code: websocket.ClosePolicyViolation, upstreamCode: websocket.ClosePolicyViolation}, } { t.Run(tc.name, func(t *testing.T) { clientRead := make(chan struct{}) upstreamCodes := make(chan int, 2) fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { defaultClose := ws.CloseHandler() ws.SetCloseHandler(func(code int, text string) error { upstreamCodes <- code return defaultClose(code, text) }) if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } for _, event := range tc.events { if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) { return } } if tc.waitClient { <-clientRead } if tc.drop { return } if tc.upstreamClose != 0 && !assert.NoError(t, ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(tc.upstreamClose, "upstream shutdown"), time.Now().Add(time.Second))) { return } for { if _, _, err := ws.ReadMessage(); err != nil { return } } }) if tc.timeout { constant.StreamingTimeout = 1 } require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) for _, event := range tc.events { actual, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client)) require.NoError(t, err) assert.JSONEq(t, event, string(actual)) } close(clientRead) if tc.end != nil { tc.end(t, fixture) } require.NoError(t, fixture.client.SetReadDeadline(time.Now().Add(5*time.Second))) _, _, err := fixture.client.ReadMessage() var closeErr *websocket.CloseError require.ErrorAs(t, err, &closeErr) assert.Equal(t, tc.code, closeErr.Code) assert.Equal(t, tc.reason, closeErr.Text) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) if tc.upstreamCode != 0 { require.Len(t, upstreamCodes, 1) assert.Equal(t, tc.upstreamCode, <-upstreamCodes) } }) } } // A control event queued while another is still outstanding is rejected // locally instead of reaching upstream twice. func TestResponsesWebSocketRejectsSecondControlWhileOnePending(t *testing.T) { receivedCancel := make(chan struct{}, 1) release := make(chan struct{}) defer close(release) fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) { return } _, cancel, err := ws.ReadMessage() if !assert.NoError(t, err) || !assert.Contains(t, string(cancel), `"event_id":"first-cancel"`) { return } receivedCancel <- struct{}{} <-release if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`))) { return } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","event_id":"first-cancel","response_id":"active"}`))) select { case <-receivedCancel: case <-time.After(3 * time.Second): t.Fatal("upstream did not receive the first cancel") } require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","event_id":"second-cancel","response_id":"active"}`))) rejected := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, "error", rejected["type"]) assert.Equal(t, float64(http.StatusBadRequest), rejected["status"]) assert.Equal(t, "second-cancel", rejected["event_id"]) rejectedError, _ := rejected["error"].(map[string]any) assert.Contains(t, rejectedError["message"], "already pending") release <- struct{}{} assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) } // A control error resolves the outstanding control, so a corrected cancel is // forwarded instead of being rejected as still pending. func TestResponsesWebSocketAcceptsCancelAfterControlError(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.created","response":{"id":"active","status":"in_progress"}}`))) { return } _, cancel, err := ws.ReadMessage() if !assert.NoError(t, err) || !assert.JSONEq(t, `{"type":"response.cancel","response_id":"wrong"}`, string(cancel)) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"error","status":400,"error":{"type":"invalid_request_error","code":"response_not_found","message":"No such response"}}`))) { return } _, cancel, err = ws.ReadMessage() if !assert.NoError(t, err) || !assert.JSONEq(t, `{"type":"response.cancel","response_id":"active"}`, string(cancel)) { return } if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancelled","response":{"id":"active","status":"cancelled","usage":{"input_tokens":1000,"output_tokens":0,"total_tokens":1000}}}`))) { return } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"wrong"}`))) cancelError := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, "error", cancelError["type"]) assert.Equal(t, float64(http.StatusBadRequest), cancelError["status"]) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.cancel","response_id":"active"}`))) assert.Equal(t, "response.cancelled", readResponsesWSTestEvent(t, fixture.client)["type"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) } // A late error that names the previous response must not end the request // that follows it on the same connection, whether that response completed or // failed. func TestResponsesWebSocketIgnoresLateErrorForPreviousResponse(t *testing.T) { for _, tc := range []struct { name, firstTerminal, firstTerminalType string }{ {name: "after completed", firstTerminalType: "response.completed", firstTerminal: `{"type":"response.completed","response":{"id":"first","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`}, {name: "after error", firstTerminalType: "error", firstTerminal: `{"type":"error","response_id":"first","status":500,"error":{"type":"server_error","code":"server_error","message":"Internal server error"}}`}, } { t.Run(tc.name, func(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { for _, turn := range [][]string{ {`{"type":"response.created","response":{"id":"first","status":"in_progress"}}`, tc.firstTerminal}, { `{"type":"error","response_id":"first","status":500,"error":{"type":"server_error","code":"server_error","message":"late failure"}}`, `{"type":"response.created","response":{"id":"second","status":"in_progress"}}`, `{"type":"response.completed","response":{"id":"second","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`, }, } { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } for _, event := range turn { if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) { return } } } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"first"}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) assert.Equal(t, tc.firstTerminalType, readResponsesWSTestEvent(t, fixture.client)["type"]) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"second"}`))) created := readResponsesWSTestEvent(t, fixture.client) require.Equal(t, "response.created", created["type"], "the late error must not reach the client or end the request: %v", created) assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000, 1000}) }) } } // An error that names a different response than the active one is relayed to // the client and leaves the active generation running. func TestResponsesWebSocketForwardsErrorForOtherResponseWithoutEndingRequest(t *testing.T) { const strayError = `{"type":"error","response_id":"other","status":500,"error":{"type":"server_error","code":"server_error","message":"unrelated failure"}}` fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(ws *websocket.Conn, _ *http.Request) { if _, _, err := ws.ReadMessage(); !assert.NoError(t, err) { return } for _, event := range []string{ `{"type":"response.created","response":{"id":"active","status":"in_progress"}}`, strayError, `{"type":"response.completed","response":{"id":"active","status":"completed","usage":{"input_tokens":1000,"output_tokens":10,"total_tokens":1010}}}`, } { if !assert.NoError(t, ws.WriteMessage(websocket.TextMessage, []byte(event))) { return } } _, _, _ = ws.ReadMessage() }) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) assert.Equal(t, "response.created", readResponsesWSTestEvent(t, fixture.client)["type"]) stray, err := common.Marshal(readResponsesWSTestEvent(t, fixture.client)) require.NoError(t, err) assert.JSONEq(t, strayError, string(stray)) assert.Equal(t, "response.completed", readResponsesWSTestEvent(t, fixture.client)["type"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, []int{1000}) var logs []model.Log require.NoError(t, model.LOG_DB.Where("type = ?", model.LogTypeConsume).Find(&logs).Error) require.Len(t, logs, 1) var other map[string]any require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other)) stream, _ := other["stream_status"].(map[string]any) assert.Equal(t, "ok", stream["status"]) assert.Equal(t, "completed", stream["response_status"]) } // Requests rejected before a RelayInfo exists reach the model health sampler // with the HTTP classification: routing failures count against the model, // business rejections such as a token model limit do not. func TestResponsesWebSocketPreRoutingRejectionsFollowHealthClassification(t *testing.T) { require.NoError(t, i18n.Init()) for _, tc := range []struct { name string reject func(*testing.T, *responsesWSBillingTest) status int sampled bool }{ {name: "no eligible channel", status: http.StatusServiceUnavailable, sampled: true, reject: func(t *testing.T, fixture *responsesWSBillingTest) { require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", fixture.channel.Id).Update("enabled", false).Error) }}, {name: "token model limit", status: http.StatusForbidden, reject: func(t *testing.T, fixture *responsesWSBillingTest) { // Update through the model so the cached token is invalidated. fixture.token.ModelLimitsEnabled = true fixture.token.ModelLimits = "other-model" require.NoError(t, fixture.token.Update()) }}, } { t.Run(tc.name, func(t *testing.T) { fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) { t.Error("rejected request reached upstream") }) tc.reject(t, fixture) require.NoError(t, fixture.client.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"ws-billing","input":"hi"}`))) rejected := readResponsesWSTestEvent(t, fixture.client) assert.Equal(t, "error", rejected["type"]) assert.Equal(t, float64(tc.status), rejected["status"]) fixture.closeAndWait(t) assertResponsesWSAccounting(t, fixture, nil) if !tc.sampled { keys, err := common.RDB.Keys(context.Background(), "perf:ws-billing:*").Result() require.NoError(t, err) assert.Empty(t, keys, "a business rejection must not affect model health") return } requests, successes := waitPerfCounters(t, 1) assert.Equal(t, int64(1), requests) assert.Zero(t, successes) }) } }