| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554 |
- //go:build unit
- package service
- import (
- "context"
- "net/http"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/stretchr/testify/require"
- )
- func TestCalculateOpenAI429ResetTime_7dExhausted(t *testing.T) {
- svc := &RateLimitService{}
- // Simulate headers when 7d limit is exhausted (100% used)
- // Primary = 7d (10080 minutes), Secondary = 5h (300 minutes)
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "100")
- headers.Set("x-codex-primary-reset-after-seconds", "384607") // ~4.5 days
- headers.Set("x-codex-primary-window-minutes", "10080") // 7 days
- headers.Set("x-codex-secondary-used-percent", "3")
- headers.Set("x-codex-secondary-reset-after-seconds", "17369") // ~4.8 hours
- headers.Set("x-codex-secondary-window-minutes", "300") // 5 hours
- before := time.Now()
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- after := time.Now()
- if resetAt == nil {
- t.Fatal("expected non-nil resetAt")
- }
- // Should be approximately 384607 seconds from now
- expectedDuration := 384607 * time.Second
- minExpected := before.Add(expectedDuration)
- maxExpected := after.Add(expectedDuration)
- if resetAt.Before(minExpected) || resetAt.After(maxExpected) {
- t.Errorf("resetAt %v not in expected range [%v, %v]", resetAt, minExpected, maxExpected)
- }
- }
- func TestCalculateOpenAI429ResetTime_5hExhausted(t *testing.T) {
- svc := &RateLimitService{}
- // Simulate headers when 5h limit is exhausted (100% used)
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "50")
- headers.Set("x-codex-primary-reset-after-seconds", "500000")
- headers.Set("x-codex-primary-window-minutes", "10080") // 7 days
- headers.Set("x-codex-secondary-used-percent", "100")
- headers.Set("x-codex-secondary-reset-after-seconds", "3600") // 1 hour
- headers.Set("x-codex-secondary-window-minutes", "300") // 5 hours
- before := time.Now()
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- after := time.Now()
- if resetAt == nil {
- t.Fatal("expected non-nil resetAt")
- }
- // Should be approximately 3600 seconds from now
- expectedDuration := 3600 * time.Second
- minExpected := before.Add(expectedDuration)
- maxExpected := after.Add(expectedDuration)
- if resetAt.Before(minExpected) || resetAt.After(maxExpected) {
- t.Errorf("resetAt %v not in expected range [%v, %v]", resetAt, minExpected, maxExpected)
- }
- }
- func TestCalculateOpenAI429ResetTime_NeitherExhausted_ReturnsNil(t *testing.T) {
- svc := &RateLimitService{}
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "80")
- headers.Set("x-codex-primary-reset-after-seconds", "100000")
- headers.Set("x-codex-primary-window-minutes", "10080")
- headers.Set("x-codex-secondary-used-percent", "90")
- headers.Set("x-codex-secondary-reset-after-seconds", "5000")
- headers.Set("x-codex-secondary-window-minutes", "300")
- require.Nil(t, svc.calculateOpenAI429ResetTime(headers))
- }
- func TestCalculateOpenAI429ResetTime_NoCodexHeaders(t *testing.T) {
- svc := &RateLimitService{}
- // No codex headers at all
- headers := http.Header{}
- headers.Set("content-type", "application/json")
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- if resetAt != nil {
- t.Errorf("expected nil resetAt when no codex headers, got %v", resetAt)
- }
- }
- func TestParseOpenAIRateLimitResetTime_OpenCodeGoUsageLimit(t *testing.T) {
- tests := []struct {
- name string
- body string
- want time.Duration
- }{
- {
- name: "days",
- body: `{"type":"error","error":{"type":"GoUsageLimitError","message":"Weekly usage limit reached. Resets in 2 days."}}`,
- want: 48 * time.Hour,
- },
- {
- name: "hours",
- body: `{"type":"error","error":{"type":"GoUsageLimitError","message":"Weekly usage limit reached. Resets in 18 hours."}}`,
- want: 18 * time.Hour,
- },
- {
- name: "hours and minutes",
- body: `{"type":"error","error":{"type":"GoUsageLimitError","message":"5-hour usage limit reached. Resets in 4hr 59min."}}`,
- want: 4*time.Hour + 59*time.Minute,
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- before := time.Now()
- resetAt := parseOpenAIRateLimitResetTime([]byte(tt.body))
- after := time.Now()
- require.NotNil(t, resetAt)
- actual := time.Unix(*resetAt, 0)
- require.False(t, actual.Before(before.Add(tt.want).Truncate(time.Second)))
- require.False(t, actual.After(after.Add(tt.want)))
- })
- }
- }
- func TestParseOpenAIRateLimitResetTime_DoesNotParseUnknownErrorMessage(t *testing.T) {
- body := []byte(`{"error":{"type":"rate_limit_error","message":"Resets in 2 days."}}`)
- require.Nil(t, parseOpenAIRateLimitResetTime(body))
- }
- func TestCalculateOpenAI429ResetTime_ReversedWindowOrder(t *testing.T) {
- svc := &RateLimitService{}
- // Test when OpenAI sends primary as 5h and secondary as 7d (reversed)
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "100") // This is 5h
- headers.Set("x-codex-primary-reset-after-seconds", "3600") // 1 hour
- headers.Set("x-codex-primary-window-minutes", "300") // 5 hours - smaller!
- headers.Set("x-codex-secondary-used-percent", "50")
- headers.Set("x-codex-secondary-reset-after-seconds", "500000")
- headers.Set("x-codex-secondary-window-minutes", "10080") // 7 days - larger!
- before := time.Now()
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- after := time.Now()
- if resetAt == nil {
- t.Fatal("expected non-nil resetAt")
- }
- // Should correctly identify that primary is 5h (smaller window) and use its reset time
- expectedDuration := 3600 * time.Second
- minExpected := before.Add(expectedDuration)
- maxExpected := after.Add(expectedDuration)
- if resetAt.Before(minExpected) || resetAt.After(maxExpected) {
- t.Errorf("resetAt %v not in expected range [%v, %v]", resetAt, minExpected, maxExpected)
- }
- }
- type openAI429SnapshotRepo struct {
- mockAccountRepoForGemini
- rateLimitedID int64
- updatedExtra map[string]any
- bulkUpdatedIDs []int64
- bulkUpdatedPayload AccountBulkUpdate
- }
- func (r *openAI429SnapshotRepo) SetRateLimited(_ context.Context, id int64, _ time.Time) error {
- r.rateLimitedID = id
- return nil
- }
- func (r *openAI429SnapshotRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
- r.updatedExtra = updates
- return nil
- }
- func (r *openAI429SnapshotRepo) BulkUpdate(_ context.Context, ids []int64, updates AccountBulkUpdate) (int64, error) {
- r.bulkUpdatedIDs = append([]int64(nil), ids...)
- r.bulkUpdatedPayload = updates
- return int64(len(ids)), nil
- }
- func TestHandle429_OpenAIPersistsCodexSnapshotImmediately(t *testing.T) {
- repo := &openAI429SnapshotRepo{}
- svc := NewRateLimitService(repo, nil, nil, nil, nil)
- account := &Account{ID: 123, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- 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")
- headers.Set("x-codex-secondary-used-percent", "100")
- headers.Set("x-codex-secondary-reset-after-seconds", "18000")
- headers.Set("x-codex-secondary-window-minutes", "300")
- svc.handle429(context.Background(), account, headers, nil)
- if repo.rateLimitedID != account.ID {
- t.Fatalf("rateLimitedID = %d, want %d", repo.rateLimitedID, account.ID)
- }
- if len(repo.updatedExtra) == 0 {
- t.Fatal("expected codex snapshot to be persisted on 429")
- }
- if got := repo.updatedExtra["codex_5h_used_percent"]; got != 100.0 {
- t.Fatalf("codex_5h_used_percent = %v, want 100", got)
- }
- if got := repo.updatedExtra["codex_7d_used_percent"]; got != 100.0 {
- t.Fatalf("codex_7d_used_percent = %v, want 100", got)
- }
- }
- func TestHandle429_OpenAISyncsObservedPlanType(t *testing.T) {
- repo := &openAI429SnapshotRepo{}
- svc := NewRateLimitService(repo, nil, nil, nil, nil)
- account := &Account{
- ID: 124,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Credentials: map[string]any{"plan_type": "plus"},
- }
- body := []byte(`{"error":{"type":"usage_limit_reached","message":"limit reached","plan_type":"free","resets_at":1777283883}}`)
- svc.handle429(context.Background(), account, http.Header{}, body)
- require.Equal(t, []int64{account.ID}, repo.bulkUpdatedIDs)
- require.Equal(t, "free", repo.bulkUpdatedPayload.Credentials["plan_type"])
- require.Equal(t, "free", account.Credentials["plan_type"])
- require.Equal(t, account.ID, repo.rateLimitedID)
- }
- // TestHandle429_SkipsSparkShadow 外审第8轮 P1:spark 影子的限流状态只由 QueryUsage(/wham/usage
- // codex_bengalfox)维护;/responses 429 携带的 global x-codex-* 不得对影子做任何 DB 限流写入,
- // 否则会把 spark 误耦合到 global codex 窗口、冷却到 global reset。
- func TestHandle429_SkipsSparkShadow(t *testing.T) {
- 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")
- headers.Set("x-codex-secondary-used-percent", "100")
- headers.Set("x-codex-secondary-reset-after-seconds", "18000")
- headers.Set("x-codex-secondary-window-minutes", "300")
- parentID := int64(900)
- shadowRepo := &openAI429SnapshotRepo{}
- shadowSvc := NewRateLimitService(shadowRepo, nil, nil, nil, nil)
- shadow := &Account{
- ID: 901,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- ParentAccountID: &parentID,
- QuotaDimension: QuotaDimensionSpark,
- }
- shadowSvc.handle429(context.Background(), shadow, headers, nil)
- require.Zero(t, shadowRepo.rateLimitedID, "spark shadow must not be SetRateLimited from /responses global 429")
- require.Empty(t, shadowRepo.updatedExtra, "spark shadow must not get a codex snapshot from /responses 429")
- // 反向对照:普通 OpenAI OAuth 账号仍按 global 429 限流。
- normalRepo := &openAI429SnapshotRepo{}
- normalSvc := NewRateLimitService(normalRepo, nil, nil, nil, nil)
- normal := &Account{ID: 902, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
- normalSvc.handle429(context.Background(), normal, headers, nil)
- require.Equal(t, normal.ID, normalRepo.rateLimitedID, "normal OpenAI OAuth account should still be rate limited")
- }
- func TestNormalizedCodexLimits(t *testing.T) {
- // Test the Normalize() method directly
- pUsed := 100.0
- pReset := 384607
- pWindow := 10080
- sUsed := 3.0
- sReset := 17369
- sWindow := 300
- snapshot := &OpenAICodexUsageSnapshot{
- PrimaryUsedPercent: &pUsed,
- PrimaryResetAfterSeconds: &pReset,
- PrimaryWindowMinutes: &pWindow,
- SecondaryUsedPercent: &sUsed,
- SecondaryResetAfterSeconds: &sReset,
- SecondaryWindowMinutes: &sWindow,
- }
- normalized := snapshot.Normalize()
- if normalized == nil {
- t.Fatal("expected non-nil normalized")
- }
- // Primary has larger window (10080 > 300), so primary should be 7d
- if normalized.Used7dPercent == nil || *normalized.Used7dPercent != 100.0 {
- t.Errorf("expected Used7dPercent=100, got %v", normalized.Used7dPercent)
- }
- if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 384607 {
- t.Errorf("expected Reset7dSeconds=384607, got %v", normalized.Reset7dSeconds)
- }
- if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 3.0 {
- t.Errorf("expected Used5hPercent=3, got %v", normalized.Used5hPercent)
- }
- if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 17369 {
- t.Errorf("expected Reset5hSeconds=17369, got %v", normalized.Reset5hSeconds)
- }
- }
- func TestNormalizedCodexLimits_OnlyPrimaryData(t *testing.T) {
- // Test when only primary has data, no window_minutes
- pUsed := 80.0
- pReset := 50000
- snapshot := &OpenAICodexUsageSnapshot{
- PrimaryUsedPercent: &pUsed,
- PrimaryResetAfterSeconds: &pReset,
- // No window_minutes, no secondary data
- }
- normalized := snapshot.Normalize()
- if normalized == nil {
- t.Fatal("expected non-nil normalized")
- }
- // Legacy assumption: primary=7d, secondary=5h
- if normalized.Used7dPercent == nil || *normalized.Used7dPercent != 80.0 {
- t.Errorf("expected Used7dPercent=80, got %v", normalized.Used7dPercent)
- }
- if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 50000 {
- t.Errorf("expected Reset7dSeconds=50000, got %v", normalized.Reset7dSeconds)
- }
- // Secondary (5h) should be nil
- if normalized.Used5hPercent != nil {
- t.Errorf("expected Used5hPercent=nil, got %v", *normalized.Used5hPercent)
- }
- if normalized.Reset5hSeconds != nil {
- t.Errorf("expected Reset5hSeconds=nil, got %v", *normalized.Reset5hSeconds)
- }
- }
- func TestRateLimitService_HandleUpstreamError_403PreservesOriginalUpstreamMessage(t *testing.T) {
- repo := &rateLimitAccountRepoStub{}
- service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
- account := &Account{
- ID: 201,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- }
- shouldDisable := service.HandleUpstreamError(
- context.Background(),
- account,
- 403,
- http.Header{},
- []byte(`{"error":{"message":"workspace forbidden by policy","type":"invalid_request_error"}}`),
- )
- require.True(t, shouldDisable)
- require.Equal(t, 1, repo.setErrorCalls)
- require.Contains(t, repo.lastErrorMsg, "workspace forbidden by policy")
- require.NotContains(t, repo.lastErrorMsg, "account may be suspended or lack permissions")
- }
- func TestRateLimitService_HandleUpstreamError_403FallsBackToRawBody(t *testing.T) {
- repo := &rateLimitAccountRepoStub{}
- service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
- account := &Account{
- ID: 202,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- }
- shouldDisable := service.HandleUpstreamError(
- context.Background(),
- account,
- 403,
- http.Header{},
- []byte(`{"error":{"type":"access_denied","details":{"reason":"ip_blocked"}}}`),
- )
- require.True(t, shouldDisable)
- require.Equal(t, 1, repo.setErrorCalls)
- require.Contains(t, repo.lastErrorMsg, `"access_denied"`)
- require.Contains(t, repo.lastErrorMsg, `"ip_blocked"`)
- require.NotContains(t, repo.lastErrorMsg, "account may be suspended or lack permissions")
- }
- func TestNormalizedCodexLimits_OnlySecondaryData(t *testing.T) {
- // Test when only secondary has data, no window_minutes
- sUsed := 60.0
- sReset := 3000
- snapshot := &OpenAICodexUsageSnapshot{
- SecondaryUsedPercent: &sUsed,
- SecondaryResetAfterSeconds: &sReset,
- // No window_minutes, no primary data
- }
- normalized := snapshot.Normalize()
- if normalized == nil {
- t.Fatal("expected non-nil normalized")
- }
- // Legacy assumption: primary=7d, secondary=5h
- // So secondary goes to 5h
- if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 60.0 {
- t.Errorf("expected Used5hPercent=60, got %v", normalized.Used5hPercent)
- }
- if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 3000 {
- t.Errorf("expected Reset5hSeconds=3000, got %v", normalized.Reset5hSeconds)
- }
- // Primary (7d) should be nil
- if normalized.Used7dPercent != nil {
- t.Errorf("expected Used7dPercent=nil, got %v", *normalized.Used7dPercent)
- }
- }
- func TestNormalizedCodexLimits_BothDataNoWindowMinutes(t *testing.T) {
- // Test when both have data but no window_minutes
- pUsed := 100.0
- pReset := 400000
- sUsed := 50.0
- sReset := 10000
- snapshot := &OpenAICodexUsageSnapshot{
- PrimaryUsedPercent: &pUsed,
- PrimaryResetAfterSeconds: &pReset,
- SecondaryUsedPercent: &sUsed,
- SecondaryResetAfterSeconds: &sReset,
- // No window_minutes
- }
- normalized := snapshot.Normalize()
- if normalized == nil {
- t.Fatal("expected non-nil normalized")
- }
- // Legacy assumption: primary=7d, secondary=5h
- if normalized.Used7dPercent == nil || *normalized.Used7dPercent != 100.0 {
- t.Errorf("expected Used7dPercent=100, got %v", normalized.Used7dPercent)
- }
- if normalized.Reset7dSeconds == nil || *normalized.Reset7dSeconds != 400000 {
- t.Errorf("expected Reset7dSeconds=400000, got %v", normalized.Reset7dSeconds)
- }
- if normalized.Used5hPercent == nil || *normalized.Used5hPercent != 50.0 {
- t.Errorf("expected Used5hPercent=50, got %v", normalized.Used5hPercent)
- }
- if normalized.Reset5hSeconds == nil || *normalized.Reset5hSeconds != 10000 {
- t.Errorf("expected Reset5hSeconds=10000, got %v", normalized.Reset5hSeconds)
- }
- }
- func TestHandle429_AnthropicPlatformUnaffected(t *testing.T) {
- // Verify that Anthropic platform accounts still use the original logic
- // This test ensures we don't break existing Claude account rate limiting
- svc := &RateLimitService{}
- // Simulate Anthropic 429 headers
- headers := http.Header{}
- headers.Set("anthropic-ratelimit-unified-reset", "1737820800") // A future Unix timestamp
- // For Anthropic platform, calculateOpenAI429ResetTime should return nil
- // because it only handles OpenAI platform
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- // Should return nil since there are no x-codex-* headers
- if resetAt != nil {
- t.Errorf("expected nil for Anthropic headers, got %v", resetAt)
- }
- }
- func TestCalculateOpenAI429ResetTime_UserProvidedScenario(t *testing.T) {
- // This is the exact scenario from the user:
- // codex_7d_used_percent: 100
- // codex_7d_reset_after_seconds: 384607 (约4.5天后重置)
- // codex_5h_used_percent: 3
- // codex_5h_reset_after_seconds: 17369 (约4.8小时后重置)
- svc := &RateLimitService{}
- // Simulate headers matching user's data
- // Note: We need to map the canonical 5h/7d back to primary/secondary
- // Based on typical OpenAI behavior: primary=7d (larger window), secondary=5h (smaller window)
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "100")
- headers.Set("x-codex-primary-reset-after-seconds", "384607")
- headers.Set("x-codex-primary-window-minutes", "10080") // 7 days = 10080 minutes
- headers.Set("x-codex-secondary-used-percent", "3")
- headers.Set("x-codex-secondary-reset-after-seconds", "17369")
- headers.Set("x-codex-secondary-window-minutes", "300") // 5 hours = 300 minutes
- before := time.Now()
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- after := time.Now()
- if resetAt == nil {
- t.Fatal("expected non-nil resetAt for user scenario")
- }
- // Should use the 7d reset time (384607 seconds) since 7d limit is exhausted (100%)
- expectedDuration := 384607 * time.Second
- minExpected := before.Add(expectedDuration)
- maxExpected := after.Add(expectedDuration)
- if resetAt.Before(minExpected) || resetAt.After(maxExpected) {
- t.Errorf("resetAt %v not in expected range [%v, %v]", resetAt, minExpected, maxExpected)
- }
- // Verify it's approximately 4.45 days (384607 seconds)
- duration := resetAt.Sub(before)
- actualDays := duration.Hours() / 24.0
- // 384607 / 86400 = ~4.45 days
- if actualDays < 4.4 || actualDays > 4.5 {
- t.Errorf("expected ~4.45 days, got %.2f days", actualDays)
- }
- t.Logf("User scenario: reset_at=%v, duration=%.2f days", resetAt, actualDays)
- }
- func TestCalculateOpenAI429ResetTime_5MinFallbackWhenNoReset(t *testing.T) {
- // Test that we return nil when there's used_percent but no reset_after_seconds
- // This should cause the caller to use the default 5-minute fallback
- svc := &RateLimitService{}
- headers := http.Header{}
- headers.Set("x-codex-primary-used-percent", "100")
- // No reset_after_seconds!
- resetAt := svc.calculateOpenAI429ResetTime(headers)
- // Should return nil since there's no reset time available
- if resetAt != nil {
- t.Errorf("expected nil when no reset_after_seconds, got %v", resetAt)
- }
- }
|