diff --git a/README.md b/README.md index a7ff5d44a47e..45d3c067fb1b 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,12 @@ gateway: - 修改 Compose `.env` 后需要重建容器以更新环境变量。调大等待时间时,需要确认客户端及反向代理的首字节超时。 - 升级自动应用 `237_add_api_key_concurrency_limit.sql`,旧 Key 默认为 `0`。回退二进制不会撤销数据库新增字段。 +### 中转 API Key 的上游 429 + +- OpenAI 平台普通 API Key 收到上游 429 时,使用管理端已有的“429 兜底冷却”设置(默认启用、5 秒),不再把透传的 ChatGPT/Codex 周窗口、`resets_at` 或使用率快照写成整个 Key 的配额状态。关闭该设置时不新增本地兜底冷却,上游错误仍按原有响应与故障转移逻辑处理。 +- OAuth / Setup Token 保留真实配额窗口;Ollama Cloud、国产供应商和 OpenCode Go 的已识别专用配额逻辑,以及池模式、自定义错误策略和显式临时停调规则,仍按原有优先级处理。 +- 管理端测试 API Key 时,不再用透传的 ChatGPT 429 重置时间更新 Key 的调度状态。升级不会自动清理此前已写入的长冷却;核实状态后可使用管理端恢复功能处理。 + ## 技术栈 | 组件 | 技术 | diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index a5186da3e15f..e677312fe720 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -2641,6 +2641,12 @@ func (s *AccountTestService) reconcileOpenAI429State(ctx context.Context, accoun return } + // A relayed ChatGPT quota window is not the API key's own quota. Manual + // tests must not turn it into a persistent key-wide scheduling block. + if !account.IsOpenAIOAuthLike() { + return + } + persistOpenAI429PlanType(ctx, s.accountRepo, account, body) var resetAt *time.Time diff --git a/backend/internal/service/account_test_service_openai_test.go b/backend/internal/service/account_test_service_openai_test.go index 83ac5a40d77e..08061723ad99 100644 --- a/backend/internal/service/account_test_service_openai_test.go +++ b/backend/internal/service/account_test_service_openai_test.go @@ -444,9 +444,9 @@ func TestAccountTestService_OpenAIAPIKey429DoesNotSyncUpstreamPlanType(t *testin } require.Equal(t, AccountTypeAPIKey, account.Type) require.Equal(t, "test-key", account.GetCredential("api_key")) - require.Equal(t, account.ID, repo.rateLimitedID) - require.NotNil(t, account.RateLimitResetAt) - require.Equal(t, int64(4102444800), account.RateLimitResetAt.Unix()) + require.Zero(t, repo.rateLimitedID) + require.Nil(t, account.RateLimitedAt) + require.Nil(t, account.RateLimitResetAt) }) } } diff --git a/backend/internal/service/openai_account_runtime_block_fastpath_test.go b/backend/internal/service/openai_account_runtime_block_fastpath_test.go index d8b0447a492e..dc966253d14f 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -371,8 +371,10 @@ func TestOpenAI429FastPath_OpenCodeGoUsageLimitUsesMessageResetDuration(t *testi repo := &rateLimit429AccountRepoStub{} rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) svc := &OpenAIGatewayService{rateLimitService: rateLimitService} - rateLimitService.SetAccountRuntimeBlocker(svc) - account := &Account{ID: 44, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + blocker := &runtimeBlockRecorder{} + rateLimitService.SetAccountRuntimeBlocker(blocker) + // Only an identified provider owns this quota; generic relayed OpenAI keys use fallback. + account := &Account{ID: 44, Platform: PlatformOpenCodeGo, Type: AccountTypeAPIKey} body := []byte(`{"type":"error","error":{"type":"GoUsageLimitError","message":"5-hour usage limit reached. Resets in 4hr 59min. To continue using this model now, enable usage from your available balance: https://opencode.ai/workspace/wrk_test/go"},"metadata":{"workspace":"wrk_test","limitName":"5 hour"}}`) before := time.Now() @@ -391,7 +393,8 @@ func TestOpenAI429FastPath_OpenCodeGoUsageLimitUsesMessageResetDuration(t *testi expectedResetAfter := 4*time.Hour + 59*time.Minute require.False(t, repo.lastRateLimitReset.Before(before.Add(expectedResetAfter-time.Second))) require.False(t, repo.lastRateLimitReset.After(after.Add(expectedResetAfter))) - require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, []*Account{account}, blocker.accounts) + require.Equal(t, []time.Time{repo.lastRateLimitReset}, blocker.until) } // TestOpenAI429FastPath_SkipsSparkShadow 外审第8轮 P1:spark 影子被选中后若 /responses 返回 429, diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index 5a3e939b2be1..7022286f0d3f 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -1635,10 +1635,10 @@ func TestOpenAIGatewayService_OpenAIPassthrough_RetryableStatusesTriggerFailover return fmt.Sprintf(`{"error":{"message":"The usage limit has been reached","type":"usage_limit_reached","resets_at":%d}}`, resetAt) }(), expectFailover: true, - assertRepo: func(t *testing.T, repo *openAIPassthroughFailoverRepo, _ time.Time) { + assertRepo: func(t *testing.T, repo *openAIPassthroughFailoverRepo, start time.Time) { require.Len(t, repo.rateLimitCalls, 1) require.Empty(t, repo.overloadCalls) - require.True(t, time.Until(repo.rateLimitCalls[0]) > 24*time.Hour) + requireOpenAIWSQuotaCooldown(t, AccountTypeAPIKey, repo.rateLimitCalls[0], start, time.Now(), time.Time{}) }, }, { diff --git a/backend/internal/service/openai_ws_ratelimit_signal_test.go b/backend/internal/service/openai_ws_ratelimit_signal_test.go index 13346cc46029..aff71075646b 100644 --- a/backend/internal/service/openai_ws_ratelimit_signal_test.go +++ b/backend/internal/service/openai_ws_ratelimit_signal_test.go @@ -84,159 +84,189 @@ func (r *openAICodexExtraListRepo) ListWithFilters(_ context.Context, params pag } func TestOpenAIGatewayService_Forward_WSv2ErrorEventUsageLimitPersistsRateLimit(t *testing.T) { - gin.SetMode(gin.TestMode) - - resetAt := time.Now().Add(2 * time.Hour).Unix() - upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} - wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) - if err != nil { - t.Errorf("upgrade websocket failed: %v", err) - return - } - defer func() { _ = conn.Close() }() - - var req map[string]any - if err := conn.ReadJSON(&req); err != nil { - t.Errorf("read ws request failed: %v", err) - return - } - _ = conn.WriteJSON(map[string]any{ - "type": "error", - "error": map[string]any{ - "code": "rate_limit_exceeded", - "type": "usage_limit_reached", - "message": "The usage limit has been reached", - "resets_at": resetAt, - }, + for _, accountType := range []string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeSetupToken} { + t.Run(accountType, func(t *testing.T) { + gin.SetMode(gin.TestMode) + + resetAt := time.Now().Add(2 * time.Hour).Unix() + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket failed: %v", err) + return + } + defer func() { _ = conn.Close() }() + + var req map[string]any + if err := conn.ReadJSON(&req); err != nil { + t.Errorf("read ws request failed: %v", err) + return + } + _ = conn.WriteJSON(map[string]any{ + "type": "error", + "error": map[string]any{ + "code": "rate_limit_exceeded", + "type": "usage_limit_reached", + "message": "The usage limit has been reached", + "resets_at": resetAt, + }, + }) + })) + defer wsServer.Close() + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") + + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_should_not_run"}`)), + }, + } + + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + account := Account{ + ID: 501, + Name: "openai-ws-rate-limit-event", + Platform: PlatformOpenAI, + Type: accountType, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: openAIWSQuotaTestCredentials(accountType, wsServer.URL), + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + }, + } + repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} + rateSvc := &RateLimitService{accountRepo: repo} + svc := &OpenAIGatewayService{ + accountRepo: repo, + rateLimitService: rateSvc, + httpUpstream: upstream, + cache: &stubGatewayCache{}, + cfg: cfg, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + + pool := newOpenAIWSConnPool(cfg) + defer pool.Close() + pool.setClientDialerForTest(&openAIWSQuotaLocalDialer{url: "ws" + strings.TrimPrefix(wsServer.URL, "http")}) + svc.openaiWSPool = pool + rateSvc.SetAccountRuntimeBlocker(svc) + + body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`) + before := time.Now() + result, err := svc.Forward(context.Background(), c, &account, body) + after := time.Now() + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.Nil(t, upstream.lastReq, "WS 限流 error event 不应回退到同账号 HTTP") + require.Len(t, repo.rateLimitCalls, 1) + requireOpenAIWSQuotaCooldown(t, accountType, repo.rateLimitCalls[0], before, after, time.Unix(resetAt, 0)) + require.Empty(t, repo.updateExtra, "body-only errors have no Codex snapshot") + requireOpenAIWSQuotaRuntimeBlock(t, svc, &account, repo.rateLimitCalls[0]) }) - })) - defer wsServer.Close() - - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) - c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") - - upstream := &httpUpstreamRecorder{ - resp: &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_should_not_run"}`)), - }, - } - - cfg := newOpenAIWSV2TestConfig() - cfg.Security.URLAllowlist.Enabled = false - cfg.Security.URLAllowlist.AllowInsecureHTTP = true - - account := Account{ - ID: 501, - Name: "openai-ws-rate-limit-event", - Platform: PlatformOpenAI, - Type: AccountTypeAPIKey, - Status: StatusActive, - Schedulable: true, - Concurrency: 1, - Credentials: map[string]any{ - "api_key": "sk-test", - "base_url": wsServer.URL, - }, - Extra: map[string]any{ - "responses_websockets_v2_enabled": true, - }, - } - repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} - rateSvc := &RateLimitService{accountRepo: repo} - svc := &OpenAIGatewayService{ - accountRepo: repo, - rateLimitService: rateSvc, - httpUpstream: upstream, - cache: &stubGatewayCache{}, - cfg: cfg, - openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), - toolCorrector: NewCodexToolCorrector(), } - - body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`) - result, err := svc.Forward(context.Background(), c, &account, body) - require.Error(t, err) - require.Nil(t, result) - require.Equal(t, http.StatusTooManyRequests, rec.Code) - require.Nil(t, upstream.lastReq, "WS 限流 error event 不应回退到同账号 HTTP") - require.Len(t, repo.rateLimitCalls, 1) - require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second) } func TestOpenAIGatewayService_Forward_WSv2Handshake429PersistsRateLimit(t *testing.T) { - gin.SetMode(gin.TestMode) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("x-codex-primary-used-percent", "100") - w.Header().Set("x-codex-primary-reset-after-seconds", "7200") - w.Header().Set("x-codex-primary-window-minutes", "10080") - w.Header().Set("x-codex-secondary-used-percent", "3") - w.Header().Set("x-codex-secondary-reset-after-seconds", "1800") - w.Header().Set("x-codex-secondary-window-minutes", "300") - w.WriteHeader(http.StatusTooManyRequests) - _, _ = w.Write([]byte(`{"error":{"type":"rate_limit_exceeded","message":"rate limited"}}`)) - })) - defer server.Close() - - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) - c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") - - upstream := &httpUpstreamRecorder{ - resp: &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_should_not_run"}`)), - }, - } - - cfg := newOpenAIWSV2TestConfig() - cfg.Security.URLAllowlist.Enabled = false - cfg.Security.URLAllowlist.AllowInsecureHTTP = true - - account := Account{ - ID: 502, - Name: "openai-ws-rate-limit-handshake", - Platform: PlatformOpenAI, - Type: AccountTypeAPIKey, - Status: StatusActive, - Schedulable: true, - Concurrency: 1, - Credentials: map[string]any{ - "api_key": "sk-test", - "base_url": server.URL, - }, - Extra: map[string]any{ - "responses_websockets_v2_enabled": true, - }, - } - repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} - rateSvc := &RateLimitService{accountRepo: repo} - svc := &OpenAIGatewayService{ - accountRepo: repo, - rateLimitService: rateSvc, - httpUpstream: upstream, - cache: &stubGatewayCache{}, - cfg: cfg, - openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), - toolCorrector: NewCodexToolCorrector(), + for _, accountType := range []string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeSetupToken} { + t.Run(accountType, func(t *testing.T) { + gin.SetMode(gin.TestMode) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("x-codex-primary-used-percent", "100") + w.Header().Set("x-codex-primary-reset-after-seconds", "7200") + w.Header().Set("x-codex-primary-window-minutes", "10080") + w.Header().Set("x-codex-secondary-used-percent", "3") + w.Header().Set("x-codex-secondary-reset-after-seconds", "1800") + w.Header().Set("x-codex-secondary-window-minutes", "300") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"type":"rate_limit_exceeded","message":"rate limited"}}`)) + })) + defer server.Close() + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("User-Agent", "unit-test-agent/1.0") + + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_http_should_not_run"}`)), + }, + } + + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + account := Account{ + ID: 502, + Name: "openai-ws-rate-limit-handshake", + Platform: PlatformOpenAI, + Type: accountType, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: openAIWSQuotaTestCredentials(accountType, server.URL), + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + }, + } + repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} + rateSvc := &RateLimitService{accountRepo: repo} + svc := &OpenAIGatewayService{ + accountRepo: repo, + rateLimitService: rateSvc, + httpUpstream: upstream, + cache: &stubGatewayCache{}, + cfg: cfg, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + } + + pool := newOpenAIWSConnPool(cfg) + defer pool.Close() + pool.setClientDialerForTest(&openAIWSQuotaLocalDialer{url: "ws" + strings.TrimPrefix(server.URL, "http")}) + svc.openaiWSPool = pool + rateSvc.SetAccountRuntimeBlocker(svc) + + body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`) + before := time.Now() + result, err := svc.Forward(context.Background(), c, &account, body) + after := time.Now() + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusTooManyRequests, rec.Code) + require.Nil(t, upstream.lastReq, "WS 握手 429 不应回退到同账号 HTTP") + require.Len(t, repo.rateLimitCalls, 1) + if accountType == AccountTypeAPIKey { + requireOpenAIWSQuotaCooldown(t, accountType, repo.rateLimitCalls[0], before, after, time.Time{}) + require.Empty(t, repo.updateExtra, "relayed Codex quota must not become the key's snapshot") + } else { + require.False(t, repo.rateLimitCalls[0].Before(before.Add(2*time.Hour-time.Second))) + require.False(t, repo.rateLimitCalls[0].After(after.Add(2*time.Hour))) + require.NotEmpty(t, repo.updateExtra, "OAuth-like handshake 429 must persist its real quota snapshot") + require.Contains(t, repo.updateExtra[0], "codex_usage_updated_at") + require.Equal(t, 100.0, repo.updateExtra[0]["codex_7d_used_percent"]) + require.Equal(t, 3.0, repo.updateExtra[0]["codex_5h_used_percent"]) + } + requireOpenAIWSQuotaRuntimeBlock(t, svc, &account, repo.rateLimitCalls[0]) + }) } - - body := []byte(`{"model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`) - result, err := svc.Forward(context.Background(), c, &account, body) - require.Error(t, err) - require.Nil(t, result) - require.Equal(t, http.StatusTooManyRequests, rec.Code) - require.Nil(t, upstream.lastReq, "WS 握手 429 不应回退到同账号 HTTP") - require.Len(t, repo.rateLimitCalls, 1) - require.NotEmpty(t, repo.updateExtra, "握手 429 的 x-codex 头应立即落库") - require.Contains(t, repo.updateExtra[0], "codex_usage_updated_at") } func TestOpenAIGatewayService_Forward_WSv2Handshake502RecordsModelTransient(t *testing.T) { @@ -286,111 +316,120 @@ func TestOpenAIGatewayService_Forward_WSv2Handshake502RecordsModelTransient(t *t } func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_ErrorEventUsageLimitPersistsRateLimit(t *testing.T) { - gin.SetMode(gin.TestMode) - - cfg := newOpenAIWSV2TestConfig() - cfg.Security.URLAllowlist.Enabled = false - cfg.Security.URLAllowlist.AllowInsecureHTTP = true - cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 - cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 - cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 - cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 - cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 - cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 - cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 - - resetAt := time.Now().Add(90 * time.Minute).Unix() - captureConn := &openAIWSCaptureConn{ - events: [][]byte{ - []byte(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":PLACEHOLDER}}`), - }, - } - captureConn.events[0] = []byte(strings.ReplaceAll(string(captureConn.events[0]), "PLACEHOLDER", strconv.FormatInt(resetAt, 10))) - captureDialer := &openAIWSCaptureDialer{conn: captureConn} - pool := newOpenAIWSConnPool(cfg) - pool.setClientDialerForTest(captureDialer) - - account := Account{ - ID: 503, - Name: "openai-ingress-rate-limit", - Platform: PlatformOpenAI, - Type: AccountTypeAPIKey, - Status: StatusActive, - Schedulable: true, - Concurrency: 1, - Credentials: map[string]any{ - "api_key": "sk-test", - }, - Extra: map[string]any{ - "responses_websockets_v2_enabled": true, - }, - } - repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} - rateSvc := &RateLimitService{accountRepo: repo} - svc := &OpenAIGatewayService{ - accountRepo: repo, - rateLimitService: rateSvc, - httpUpstream: &httpUpstreamRecorder{}, - cache: &stubGatewayCache{}, - cfg: cfg, - openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), - toolCorrector: NewCodexToolCorrector(), - openaiWSPool: pool, - } - - serverErrCh := make(chan error, 1) - wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) - if err != nil { - serverErrCh <- err - return - } - defer func() { _ = conn.CloseNow() }() - - rec := httptest.NewRecorder() - ginCtx, _ := gin.CreateTestContext(rec) - req := r.Clone(r.Context()) - req.Header = req.Header.Clone() - req.Header.Set("User-Agent", "unit-test-agent/1.0") - ginCtx.Request = req - - readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) - msgType, firstMessage, readErr := conn.Read(readCtx) - cancel() - if readErr != nil { - serverErrCh <- readErr - return - } - if msgType != coderws.MessageText && msgType != coderws.MessageBinary { - serverErrCh <- io.ErrUnexpectedEOF - return - } - - serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, &account, "sk-test", firstMessage, nil) - })) - defer wsServer.Close() - - dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) - clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) - cancelDial() - require.NoError(t, err) - defer func() { _ = clientConn.CloseNow() }() - - writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) - err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) - cancelWrite() - require.NoError(t, err) - - select { - case serverErr := <-serverErrCh: - require.Error(t, serverErr) - var failoverErr *UpstreamFailoverError - require.ErrorAs(t, serverErr, &failoverErr) - require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) - require.Len(t, repo.rateLimitCalls, 1) - require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second) - case <-time.After(5 * time.Second): - t.Fatal("等待 ingress websocket 结束超时") + for _, accountType := range []string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeSetupToken} { + t.Run(accountType, func(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := newOpenAIWSV2TestConfig() + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + resetAt := time.Now().Add(90 * time.Minute).Unix() + captureConn := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":PLACEHOLDER}}`), + }, + } + captureConn.events[0] = []byte(strings.ReplaceAll(string(captureConn.events[0]), "PLACEHOLDER", strconv.FormatInt(resetAt, 10))) + captureDialer := &openAIWSCaptureDialer{conn: captureConn} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(captureDialer) + defer pool.Close() + + account := Account{ + ID: 503, + Name: "openai-ingress-rate-limit", + Platform: PlatformOpenAI, + Type: accountType, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: openAIWSQuotaTestCredentials(accountType, ""), + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + }, + } + repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}} + rateSvc := &RateLimitService{accountRepo: repo} + svc := &OpenAIGatewayService{ + accountRepo: repo, + rateLimitService: rateSvc, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + cfg: cfg, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + rateSvc.SetAccountRuntimeBlocker(svc) + + serverErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- io.ErrUnexpectedEOF + return + } + + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, &account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + before := time.Now() + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + + select { + case serverErr := <-serverErrCh: + after := time.Now() + require.Error(t, serverErr) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, serverErr, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Len(t, repo.rateLimitCalls, 1) + requireOpenAIWSQuotaCooldown(t, accountType, repo.rateLimitCalls[0], before, after, time.Unix(resetAt, 0)) + require.Empty(t, repo.updateExtra, "body-only errors have no Codex snapshot") + requireOpenAIWSQuotaRuntimeBlock(t, svc, &account, repo.rateLimitCalls[0]) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + }) } } @@ -579,3 +618,40 @@ func TestOpenAIWSRateLimitFailoverError_OAuthKeepsSameAccountDeadline(t *testing require.True(t, apiKeyErr.SameAccountRetryDeadline.IsZero()) require.Zero(t, apiKeyErr.SameAccountRetryDelay) } + +// openAIWSQuotaLocalDialer routes every dial to the test server, including +// OAuth-like accounts whose normal endpoint cannot be overridden by base_url. +type openAIWSQuotaLocalDialer struct{ url string } + +func (d *openAIWSQuotaLocalDialer) Dial(ctx context.Context, _ string, headers http.Header, _ string) (openAIWSClientConn, int, http.Header, error) { + return newDefaultOpenAIWSClientDialer().Dial(ctx, d.url, headers, "") +} + +func openAIWSQuotaTestCredentials(accountType, baseURL string) map[string]any { + credentials := map[string]any{"base_url": baseURL} + if accountType == AccountTypeAPIKey { + credentials["api_key"] = "sk-test" + } else { + credentials["access_token"] = "test-access-token" + credentials["chatgpt_account_id"] = "test-chatgpt-account" + } + return credentials +} + +func requireOpenAIWSQuotaCooldown(t *testing.T, accountType string, got, before, after, quotaReset time.Time) { + t.Helper() + if accountType == AccountTypeAPIKey { + cooldown := time.Duration(defaultRateLimit429CooldownSeconds) * time.Second + require.False(t, got.Before(before.Add(cooldown)), "API key must receive the configured fallback cooldown") + require.False(t, got.After(after.Add(cooldown)), "API key must not inherit a relayed quota window") + } else { + require.Equal(t, quotaReset, got, "OAuth-like credentials must retain their real quota reset") + } +} + +func requireOpenAIWSQuotaRuntimeBlock(t *testing.T, svc *OpenAIGatewayService, account *Account, resetAt time.Time) { + t.Helper() + until, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID) + require.True(t, ok, "429 must block scheduling at runtime as well as in the repository") + require.Equal(t, resetAt, until) +} diff --git a/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go b/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go index 0b8ada3b4756..5b54a6049dd5 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go +++ b/backend/internal/service/openai_ws_v2_passthrough_lifecycle_test.go @@ -219,57 +219,115 @@ func startPassthroughLifecycleServerWithHooks( } func TestPassthroughLifecycle_LaterTurnPreOutputRateLimitRequestsReconnect(t *testing.T) { - gin.SetMode(gin.TestMode) - controlCtx, cancelControl := context.WithCancelCause(context.Background()) - defer cancelControl(context.Canceled) - upstream := newStagedPassthroughConn() - account := passthroughLifecycleAccount() - repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{*account}}} - svc := newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream) - svc.accountRepo = repo - svc.rateLimitService = &RateLimitService{accountRepo: repo} - - server, serverErr := startPassthroughLifecycleServer(t, controlCtx, svc, account) - defer server.Close() - clientConn := dialPassthroughLifecycleClient(t, server) - defer func() { _ = clientConn.CloseNow() }() - firstRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second) - require.Equal(t, "response.create", gjson.GetBytes(firstRequest, "type").String()) - - upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) - completed, err := readPassthroughLifecycleFrame(t, clientConn, time.Second) - require.NoError(t, err) - require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + for _, accountType := range []string{AccountTypeAPIKey, AccountTypeOAuth, AccountTypeSetupToken} { + t.Run(accountType, func(t *testing.T) { + gin.SetMode(gin.TestMode) + controlCtx, cancelControl := context.WithCancelCause(context.Background()) + defer cancelControl(context.Canceled) + upstream := newStagedPassthroughConn() + account := passthroughLifecycleAccount() + account.Type = accountType + account.Credentials = openAIWSQuotaTestCredentials(accountType, "") + if account.IsOpenAIOAuthLike() { + account.Extra = map[string]any{"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough} + } + repo := &openAIWSRateLimitSignalRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{*account}}} + cfg := passthroughLifecycleConfig() + cfg.Gateway.OpenAIWS.OAuthEnabled = true + svc := newPassthroughLifecycleService(cfg, upstream) + svc.accountRepo = repo + svc.rateLimitService = &RateLimitService{accountRepo: repo} + + svc.rateLimitService.SetAccountRuntimeBlocker(svc) + type turnSettlement struct { + turn int + result *OpenAIForwardResult + err error + upstreamClosed bool + } + settlements := make(chan turnSettlement, 3) + admissions := make(chan int, 2) + var heldSlots atomic.Int32 + // The ingress handler has already admitted the first turn. + heldSlots.Store(1) + hooks := &OpenAIWSIngressHooks{ + BeforeTurn: func(turn int) error { + heldSlots.Add(1) + admissions <- turn + return nil + }, + AfterTurn: func(turn int, result *OpenAIForwardResult, err error) { + heldSlots.Add(-1) + closed := false + select { + case <-upstream.closed: + closed = true + default: + } + settlements <- turnSettlement{turn, result, err, closed} + }, + } + server, serverErr := startPassthroughLifecycleServer(t, controlCtx, svc, account, hooks) + defer server.Close() + clientConn := dialPassthroughLifecycleClient(t, server) + defer func() { _ = clientConn.CloseNow() }() + firstRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second) + require.Equal(t, "response.create", gjson.GetBytes(firstRequest, "type").String()) - writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) - err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) - cancelWrite() - require.NoError(t, err) - secondRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second) - require.Equal(t, "response.create", gjson.GetBytes(secondRequest, "type").String()) + upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`) + completed, err := readPassthroughLifecycleFrame(t, clientConn, time.Second) + require.NoError(t, err) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) - resetAt := time.Now().Add(90 * time.Minute).Unix() - upstream.Send(fmt.Sprintf(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":%d}}`, resetAt)) - _, err = readPassthroughLifecycleFrame(t, clientConn, time.Second) - var websocketCloseErr coderws.CloseError - require.ErrorAs(t, err, &websocketCloseErr) - require.Equal(t, coderws.StatusTryAgainLater, websocketCloseErr.Code) - require.Equal(t, "upstream rate limit exceeded; please reconnect", websocketCloseErr.Reason) - require.Len(t, repo.rateLimitCalls, 1) - require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second) + writeCtx, cancelWrite := context.WithTimeout(context.Background(), time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + secondRequest := requirePassthroughUpstreamWrite(t, upstream, time.Second) + require.Equal(t, "response.create", gjson.GetBytes(secondRequest, "type").String()) + + resetAt := time.Now().Add(90 * time.Minute).Unix() + before := time.Now() + upstream.Send(fmt.Sprintf(`{"type":"error","error":{"code":"rate_limit_exceeded","type":"usage_limit_reached","message":"The usage limit has been reached","resets_at":%d}}`, resetAt)) + _, err = readPassthroughLifecycleFrame(t, clientConn, time.Second) + var websocketCloseErr coderws.CloseError + require.ErrorAs(t, err, &websocketCloseErr) + require.Equal(t, coderws.StatusTryAgainLater, websocketCloseErr.Code) + require.Equal(t, "upstream rate limit exceeded; please reconnect", websocketCloseErr.Reason) - select { - case err := <-serverErr: - var closeErr *OpenAIWSClientCloseError - require.ErrorAs(t, err, &closeErr) - require.Equal(t, coderws.StatusTryAgainLater, closeErr.StatusCode()) - case <-time.After(time.Second): - t.Fatal("later-turn rate limit did not terminate passthrough") - } - select { - case replay := <-upstream.writes: - t.Fatalf("later-turn reconnect must not replay the retained first request: %s", replay) - default: + select { + case err := <-serverErr: + var closeErr *OpenAIWSClientCloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusTryAgainLater, closeErr.StatusCode()) + case <-time.After(time.Second): + t.Fatal("later-turn rate limit did not terminate passthrough") + } + after := time.Now() + require.Len(t, repo.rateLimitCalls, 1) + requireOpenAIWSQuotaCooldown(t, accountType, repo.rateLimitCalls[0], before, after, time.Unix(resetAt, 0)) + require.Empty(t, repo.updateExtra, "body-only 429 must not invent a Codex snapshot") + requireOpenAIWSQuotaRuntimeBlock(t, svc, account, repo.rateLimitCalls[0]) + require.Len(t, admissions, 1, "only the second turn may consume another admission") + require.Equal(t, 2, <-admissions) + require.Zero(t, heldSlots.Load(), "both admitted turns must release their concurrency slots") + require.Len(t, settlements, 2, "each admitted turn must settle exactly once") + first := <-settlements + require.Equal(t, 1, first.turn) + require.NoError(t, first.err) + require.NotNil(t, first.result) + require.Equal(t, OpenAIUsage{InputTokens: 1, OutputTokens: 1}, first.result.Usage) + second := <-settlements + require.Equal(t, 2, second.turn) + require.Error(t, second.err) + require.Nil(t, second.result, "rejected pre-output turn must not reuse the first turn's usage") + require.True(t, second.upstreamClosed, "close upstream before releasing the rejected turn's slots") + select { + case replay := <-upstream.writes: + t.Fatalf("later-turn reconnect must not replay the retained first request: %s", replay) + default: + } + }) } } diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index a7a361ceb11f..3d5f2716f515 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -1287,6 +1287,15 @@ func (s *RateLimitService) handle429Cooldown(ctx context.Context, account *Accou } // 1. OpenAI 平台:优先尝试解析 x-codex-* 响应头(用于 rate_limit_exceeded) if account.Platform == PlatformOpenAI { + // An API key may relay a different ChatGPT account's quota error. Its + // reset time and Codex snapshot do not describe this key's capacity. + // Keep real 429 backoff without blocking the entire key for that window. + // Provider-specific quota handling above and explicit account policies + // remain authoritative. + if account.Type == AccountTypeAPIKey { + s.apply429FallbackRateLimit(ctx, account, "openai_apikey_upstream_429") + return + } persistOpenAI429PlanType(ctx, s.accountRepo, account, responseBody) s.persistOpenAICodexSnapshot(ctx, account, headers) notifyOpenAIAutoReset(account.ID) diff --git a/backend/internal/service/ratelimit_service_openai_apikey_quota_test.go b/backend/internal/service/ratelimit_service_openai_apikey_quota_test.go new file mode 100644 index 000000000000..1d53307f869e --- /dev/null +++ b/backend/internal/service/ratelimit_service_openai_apikey_quota_test.go @@ -0,0 +1,111 @@ +//go:build unit + +package service + +import ( + "context" + "encoding/json" + "net/http" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestHandle429_OpenAIAPIKeyIgnoresRelayedQuotaWindows(t *testing.T) { + for _, source := range []string{"body_reset_at", "body_reset_seconds", "codex_headers", "anthropic_headers", "opencode_body", "ordinary_429"} { + t.Run(source, func(t *testing.T) { + repo := &openAI429SnapshotRepo{} + svc := NewRateLimitService(repo, nil, nil, nil, nil) + blocker := &runtimeBlockRecorder{} + svc.SetAccountRuntimeBlocker(blocker) + account := &Account{ID: 125, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + headers := http.Header{} + body := []byte(`{"error":{"type":"rate_limit_error","message":"slow down"}}`) + switch source { + case "body_reset_at": + body = []byte(`{"error":{"type":"usage_limit_reached","plan_type":"pro","resets_at":4102444800}}`) + case "body_reset_seconds": + body = []byte(`{"error":{"type":"rate_limit_exceeded","resets_in_seconds":604800}}`) + case "opencode_body": + body = []byte(`{"error":{"type":"GoUsageLimitError","message":"5-hour usage limit reached. Resets in 4hr 59min."}}`) + case "codex_headers": + headers.Set("x-codex-primary-used-percent", "100") + headers.Set("x-codex-primary-reset-after-seconds", "604800") + headers.Set("x-codex-primary-window-minutes", "10080") + case "anthropic_headers": + headers.Set("anthropic-ratelimit-unified-reset", "4102444800") + } + + before := time.Now() + svc.HandleUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, body) + after := time.Now() + + cooldown := time.Duration(defaultRateLimit429CooldownSeconds) * time.Second + require.Equal(t, account.ID, repo.rateLimitedID) + require.False(t, repo.rateLimitResetAt.Before(before.Add(cooldown))) + require.False(t, repo.rateLimitResetAt.After(after.Add(cooldown))) + require.Empty(t, repo.updatedExtra, "relayed Codex quota must not become the key's snapshot") + require.Empty(t, repo.bulkUpdatedIDs) + require.Equal(t, []*Account{account}, blocker.accounts) + require.Equal(t, []string{"429_fallback"}, blocker.reasons) + require.Equal(t, repo.rateLimitResetAt, blocker.until[0]) + }) + } +} + +func TestHandle429_OpenAIAPIKeyQuotaRespectsFallbackSettings(t *testing.T) { + for _, enabled := range []bool{true, false} { + name := "disabled" + if enabled { + name = "enabled" + } + t.Run(name, func(t *testing.T) { + repo := &openAI429SnapshotRepo{} + settingsRepo := newMockSettingRepo() + data, err := json.Marshal(RateLimit429CooldownSettings{Enabled: enabled, CooldownSeconds: 12}) + require.NoError(t, err) + settingsRepo.data[SettingKeyRateLimit429CooldownSettings] = string(data) + svc := NewRateLimitService(repo, nil, nil, nil, nil) + svc.SetSettingService(NewSettingService(settingsRepo, &config.Config{})) + blocker := &runtimeBlockRecorder{} + svc.SetAccountRuntimeBlocker(blocker) + account := &Account{ID: 126, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + body := []byte(`{"error":{"type":"usage_limit_reached","plan_type":"pro","resets_at":4102444800}}`) + + before := time.Now() + svc.handle429(context.Background(), account, http.Header{}, body) + after := time.Now() + + if enabled { + require.Equal(t, account.ID, repo.rateLimitedID) + require.False(t, repo.rateLimitResetAt.Before(before.Add(12*time.Second))) + require.False(t, repo.rateLimitResetAt.After(after.Add(12*time.Second))) + require.Len(t, blocker.accounts, 1) + } else { + require.Zero(t, repo.rateLimitedID) + require.Empty(t, blocker.accounts) + } + require.Empty(t, repo.updatedExtra) + }) + } +} + +func TestAccountTestService_OpenAIAPIKey429HeadersDoNotBlockScheduling(t *testing.T) { + repo := &openAIAccountTestRepo{} + svc := &AccountTestService{accountRepo: repo} + account := &Account{ID: 127, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusError, ErrorMessage: "existing error"} + headers := http.Header{} + headers.Set("x-codex-primary-used-percent", "100") + headers.Set("x-codex-primary-reset-after-seconds", "604800") + headers.Set("x-codex-primary-window-minutes", "10080") + + svc.reconcileOpenAI429State(context.Background(), account, headers, nil) + + require.Zero(t, repo.rateLimitedID) + require.Zero(t, repo.clearedErrorID) + require.Nil(t, account.RateLimitResetAt) + require.Equal(t, StatusError, account.Status) + require.Equal(t, "existing error", account.ErrorMessage) +} diff --git a/backend/internal/service/ratelimit_service_openai_test.go b/backend/internal/service/ratelimit_service_openai_test.go index 7d583e5ea778..02ecc8ebf156 100644 --- a/backend/internal/service/ratelimit_service_openai_test.go +++ b/backend/internal/service/ratelimit_service_openai_test.go @@ -242,7 +242,9 @@ func TestHandle429_OpenAIAPIKeyDoesNotSyncUpstreamPlanType(t *testing.T) { account := &Account{ID: 125, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: credentials} body := []byte(`{"error":{"type":"` + errorType + `","plan_type":"pro","resets_at":4102444800}}`) + before := time.Now() svc.handle429(context.Background(), account, http.Header{}, body) + after := time.Now() require.Empty(t, repo.bulkUpdatedIDs, "upstream plans must not be persisted for API keys") require.Nil(t, repo.bulkUpdatedPayload.Credentials) @@ -253,7 +255,9 @@ func TestHandle429_OpenAIAPIKeyDoesNotSyncUpstreamPlanType(t *testing.T) { require.Equal(t, AccountTypeAPIKey, account.Type) require.Equal(t, "test-key", account.GetCredential("api_key")) require.Equal(t, account.ID, repo.rateLimitedID, "429 cooldown must still be persisted") - require.Equal(t, int64(4102444800), repo.rateLimitResetAt.Unix()) + cooldown := time.Duration(defaultRateLimit429CooldownSeconds) * time.Second + require.False(t, repo.rateLimitResetAt.Before(before.Add(cooldown))) + require.False(t, repo.rateLimitResetAt.After(after.Add(cooldown))) }) } }