From 8cd848313c92edc3ea8ecc913e2fcfd77924a81e Mon Sep 17 00:00:00 2001 From: benjamin Date: Tue, 14 Jul 2026 10:07:48 +0800 Subject: [PATCH] fix(openai): improve reset credit quota detection --- .../service/openai_quota_reset_credits.go | 141 ++++++++++++++ .../openai_quota_reset_credits_test.go | 175 ++++++++++++++++++ .../internal/service/openai_quota_service.go | 87 +++------ .../service/openai_quota_spark_window_test.go | 6 +- 4 files changed, 340 insertions(+), 69 deletions(-) create mode 100644 backend/internal/service/openai_quota_reset_credits.go create mode 100644 backend/internal/service/openai_quota_reset_credits_test.go diff --git a/backend/internal/service/openai_quota_reset_credits.go b/backend/internal/service/openai_quota_reset_credits.go new file mode 100644 index 0000000000..75756976db --- /dev/null +++ b/backend/internal/service/openai_quota_reset_credits.go @@ -0,0 +1,141 @@ +package service + +import ( + "bytes" + "encoding/json" + "strconv" + "strings" +) + +type openAIRateLimitResetCreditDetailPayload struct { + ExpiresAt string `json:"expires_at,omitempty"` + ExpiresAtCamel string `json:"expiresAt,omitempty"` + ResetType string `json:"reset_type,omitempty"` + ResetTypeCamel string `json:"resetType,omitempty"` + Status string `json:"status,omitempty"` +} + +type openAIRateLimitResetCreditDetailsPayload struct { + AvailableCount json.RawMessage `json:"available_count,omitempty"` + AvailableCountCamel json.RawMessage `json:"availableCount,omitempty"` + Credits json.RawMessage `json:"credits,omitempty"` + RateLimitResetCredits json.RawMessage `json:"rate_limit_reset_credits,omitempty"` + Items json.RawMessage `json:"items,omitempty"` + Data json.RawMessage `json:"data,omitempty"` +} + +type openAIRateLimitResetCreditDetails struct { + AvailableCount *int + AvailableCreditCount int + CreditListPresent bool + Credits []OpenAIRateLimitResetCreditDetail +} + +func parseOpenAIRateLimitResetCreditDetails(body []byte) (openAIRateLimitResetCreditDetails, error) { + trimmed := bytes.TrimSpace(body) + if len(trimmed) == 0 { + return openAIRateLimitResetCreditDetails{}, nil + } + + var rawCredits []*openAIRateLimitResetCreditDetailPayload + var availableCount *int + var creditListPresent bool + if trimmed[0] == '[' { + if err := json.Unmarshal(trimmed, &rawCredits); err != nil { + return openAIRateLimitResetCreditDetails{}, err + } + creditListPresent = true + } else { + var payload openAIRateLimitResetCreditDetailsPayload + if err := json.Unmarshal(trimmed, &payload); err != nil { + return openAIRateLimitResetCreditDetails{}, err + } + availableCount = parseOpenAIResetCreditAvailableCount(payload.AvailableCount, payload.AvailableCountCamel) + var err error + rawCredits, creditListPresent, err = firstPresentResetCreditPayload( + payload.Credits, + payload.RateLimitResetCredits, + payload.Items, + payload.Data, + ) + if err != nil { + return openAIRateLimitResetCreditDetails{}, err + } + } + + credits := make([]OpenAIRateLimitResetCreditDetail, 0, len(rawCredits)) + availableCreditCount := 0 + for _, raw := range rawCredits { + if raw == nil { + continue + } + resetType := strings.TrimSpace(raw.ResetType) + if resetType == "" { + resetType = strings.TrimSpace(raw.ResetTypeCamel) + } + if resetType != "" && !strings.EqualFold(resetType, "codex_rate_limits") { + continue + } + if status := strings.TrimSpace(raw.Status); status != "" && !strings.EqualFold(status, "available") { + continue + } + availableCreditCount++ + expiresAt := strings.TrimSpace(raw.ExpiresAt) + if expiresAt == "" { + expiresAt = strings.TrimSpace(raw.ExpiresAtCamel) + } + if expiresAt == "" { + continue + } + credits = append(credits, OpenAIRateLimitResetCreditDetail{ExpiresAt: expiresAt}) + } + return openAIRateLimitResetCreditDetails{ + AvailableCount: availableCount, + AvailableCreditCount: availableCreditCount, + CreditListPresent: creditListPresent, + Credits: credits, + }, nil +} + +func parseOpenAIResetCreditAvailableCount(values ...json.RawMessage) *int { + for _, value := range values { + trimmed := bytes.TrimSpace(value) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + continue + } + + var count int + if trimmed[0] == '"' { + var text string + if err := json.Unmarshal(trimmed, &text); err != nil { + continue + } + parsed, err := strconv.Atoi(strings.TrimSpace(text)) + if err != nil { + continue + } + count = parsed + } else if err := json.Unmarshal(trimmed, &count); err != nil { + continue + } + if count >= 0 { + return &count + } + } + return nil +} + +func firstPresentResetCreditPayload(values ...json.RawMessage) ([]*openAIRateLimitResetCreditDetailPayload, bool, error) { + for _, value := range values { + trimmed := bytes.TrimSpace(value) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + continue + } + var credits []*openAIRateLimitResetCreditDetailPayload + if err := json.Unmarshal(trimmed, &credits); err != nil { + return nil, false, err + } + return credits, true, nil + } + return nil, false, nil +} diff --git a/backend/internal/service/openai_quota_reset_credits_test.go b/backend/internal/service/openai_quota_reset_credits_test.go new file mode 100644 index 0000000000..5d994602e1 --- /dev/null +++ b/backend/internal/service/openai_quota_reset_credits_test.go @@ -0,0 +1,175 @@ +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParseOpenAIRateLimitResetCreditDetails_PreservesAvailableCreditOrder(t *testing.T) { + body := []byte(`{ + "availableCount":"2", + "credits":[ + {"reset_type":"codex_rate_limits","status":"redeemed","expires_at":"2026-07-01T04:05:06Z"}, + {"reset_type":"codex_rate_limits","status":"available","expires_at":"2026-07-04T04:05:06Z"}, + {"resetType":"codex_rate_limits","status":"available","expiresAt":"2026-07-03T04:05:06Z"}, + {"reset_type":"other","status":"available","expires_at":"2026-07-02T04:05:06Z"} + ] + }`) + + details, err := parseOpenAIRateLimitResetCreditDetails(body) + require.NoError(t, err) + require.NotNil(t, details.AvailableCount) + require.Equal(t, 2, *details.AvailableCount) + require.Equal(t, []OpenAIRateLimitResetCreditDetail{ + {ExpiresAt: "2026-07-04T04:05:06Z"}, + {ExpiresAt: "2026-07-03T04:05:06Z"}, + }, details.Credits) +} + +func TestQueryUsageResetCreditCountPrecedence(t *testing.T) { + tests := []struct { + name string + usageBody string + detailBody string + wantCount int + wantCredits int + wantNil bool + }{ + { + name: "detail count creates missing usage credits", + usageBody: `{}`, + detailBody: `{"available_count":3,"credits":[{"expires_at":"2026-07-03T04:05:06Z"}]}`, + wantCount: 3, wantCredits: 1, + }, + { + name: "explicit detail zero overrides usage and records", + usageBody: `{"rate_limit_reset_credits":{"available_count":4}}`, + detailBody: `{"available_count":0,"credits":[{"expires_at":"2026-07-03T04:05:06Z"}]}`, + wantCount: 0, wantCredits: 1, + }, + { + name: "available records override usage when detail count is absent", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[{"expires_at":"2026-07-03T04:05:06Z"},{"expiresAt":"2026-07-04T04:05:06Z"}]}`, + wantCount: 2, wantCredits: 2, + }, + { + name: "empty detail list overrides usage with zero", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[]}`, + wantCount: 0, + }, + { + name: "fully filtered list overrides usage with zero", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[{"reset_type":"codex_rate_limits","status":"redeemed","expires_at":"2026-07-03T04:05:06Z"},{"reset_type":"other","status":"available","expires_at":"2026-07-04T04:05:06Z"}]}`, + wantCount: 0, + }, + { + name: "available records without expiry still count", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[{"status":"available"},{"status":"available","expires_at":"2026-07-04T04:05:06Z"}]}`, + wantCount: 2, wantCredits: 1, + }, + { + name: "shape without count or list preserves usage details", + usageBody: `{"rate_limit_reset_credits":{"available_count":5,"credits":[{"expires_at":"usage-expiry"}]}}`, + detailBody: `{}`, + wantCount: 5, + wantCredits: 1, + }, + { + name: "negative detail count without list preserves usage", + usageBody: `{"rate_limit_reset_credits":{"available_count":4}}`, + detailBody: `{"available_count":-1}`, + wantCount: 4, + }, + { + name: "negative detail count falls back to available records", + usageBody: `{"rate_limit_reset_credits":{"available_count":4}}`, + detailBody: `{"available_count":-1,"credits":[{"status":"available","expires_at":"2026-07-04T04:05:06Z"}]}`, + wantCount: 1, wantCredits: 1, + }, + { + name: "empty object preserves missing usage credits", + usageBody: `{}`, + detailBody: `{}`, + wantNil: true, + }, + { + name: "null body preserves missing usage credits", + usageBody: `{}`, + detailBody: `null`, + wantNil: true, + }, + { + name: "empty body preserves missing usage credits", + usageBody: `{}`, + detailBody: ``, + wantNil: true, + }, + { + name: "null object record is not counted", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `{"credits":[null]}`, + wantCount: 0, + }, + { + name: "null top level record is not counted", + usageBody: `{"rate_limit_reset_credits":{"available_count":7}}`, + detailBody: `[null]`, + wantCount: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := &Account{ + ID: 100, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{ + "chatgpt_account_id": "org-parent123", + }, + } + repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}} + tokenCache := &stubQuotaTokenCache{tokens: map[string]string{ + OpenAITokenCacheKey(account): "fake-token", + }} + tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil) + + var detailCalls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("content-type", "application/json") + switch r.URL.Path { + case "/backend-api/wham/usage": + _, _ = w.Write([]byte(tt.usageBody)) + case "/backend-api/wham/rate-limit-reset-credits": + detailCalls++ + _, _ = w.Write([]byte(tt.detailBody)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv)) + usage, err := svc.QueryUsage(context.Background(), 100) + require.NoError(t, err) + require.NotNil(t, usage) + require.Equal(t, 1, detailCalls) + if tt.wantNil { + require.Nil(t, usage.RateLimitResetCredits) + return + } + require.NotNil(t, usage.RateLimitResetCredits) + require.Equal(t, tt.wantCount, usage.RateLimitResetCredits.AvailableCount) + require.Len(t, usage.RateLimitResetCredits.Credits, tt.wantCredits) + }) + } +} diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index 337f8c1e8f..91b511122e 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -1,11 +1,9 @@ package service import ( - "bytes" "context" "crypto/rand" "encoding/hex" - "encoding/json" "fmt" "log/slog" "net/http" @@ -171,13 +169,26 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* } payload.FetchedAt = time.Now().Unix() - if payload.RateLimitResetCredits != nil && payload.RateLimitResetCredits.AvailableCount > 0 { - payload.RateLimitResetCredits.Credits = s.queryResetCreditDetails(callCtx, client, accessToken, chatGPTAccountID, fedRAMP, accountID) + details := s.queryResetCreditDetails(callCtx, client, accessToken, chatGPTAccountID, fedRAMP, accountID) + if details != nil { + hasDetailCount := details.AvailableCount != nil + if payload.RateLimitResetCredits == nil { + payload.RateLimitResetCredits = &OpenAIRateLimitResetCredits{} + } + if details.CreditListPresent { + payload.RateLimitResetCredits.Credits = details.Credits + } + switch { + case hasDetailCount: + payload.RateLimitResetCredits.AvailableCount = *details.AvailableCount + case details.CreditListPresent: + payload.RateLimitResetCredits.AvailableCount = details.AvailableCreditCount + } } return &payload, nil } -func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) []OpenAIRateLimitResetCreditDetail { +func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) *openAIRateLimitResetCreditDetails { resp, err := client.R(). SetContext(ctx). SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)). @@ -191,12 +202,15 @@ func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client return nil } - credits, err := parseOpenAIRateLimitResetCreditDetails(resp.Bytes()) + details, err := parseOpenAIRateLimitResetCreditDetails(resp.Bytes()) if err != nil { slog.Warn("openai_quota_reset_credit_details_parse_failed", "account_id", accountID, "error", err) return nil } - return credits + if details.AvailableCount == nil && !details.CreditListPresent { + return nil + } + return &details } // ResetCredit consumes one rate_limit_reset_credit for the given OpenAI account. @@ -372,65 +386,6 @@ func generateRedeemRequestID() (string, error) { return fmt.Sprintf("%s-%s-%s-%s-%s", hexStr[0:8], hexStr[8:12], hexStr[12:16], hexStr[16:20], hexStr[20:]), nil } -type openAIRateLimitResetCreditDetailPayload struct { - ExpiresAt string `json:"expires_at,omitempty"` - ExpiresAtCamel string `json:"expiresAt,omitempty"` -} - -type openAIRateLimitResetCreditDetailsPayload struct { - Credits []openAIRateLimitResetCreditDetailPayload `json:"credits,omitempty"` - RateLimitResetCredits []openAIRateLimitResetCreditDetailPayload `json:"rate_limit_reset_credits,omitempty"` - Items []openAIRateLimitResetCreditDetailPayload `json:"items,omitempty"` - Data []openAIRateLimitResetCreditDetailPayload `json:"data,omitempty"` -} - -func parseOpenAIRateLimitResetCreditDetails(body []byte) ([]OpenAIRateLimitResetCreditDetail, error) { - trimmed := bytes.TrimSpace(body) - if len(trimmed) == 0 { - return nil, nil - } - - var rawCredits []openAIRateLimitResetCreditDetailPayload - if trimmed[0] == '[' { - if err := json.Unmarshal(trimmed, &rawCredits); err != nil { - return nil, err - } - } else { - var payload openAIRateLimitResetCreditDetailsPayload - if err := json.Unmarshal(trimmed, &payload); err != nil { - return nil, err - } - rawCredits = firstNonEmptyResetCreditPayload( - payload.Credits, - payload.RateLimitResetCredits, - payload.Items, - payload.Data, - ) - } - - credits := make([]OpenAIRateLimitResetCreditDetail, 0, len(rawCredits)) - for _, raw := range rawCredits { - expiresAt := strings.TrimSpace(raw.ExpiresAt) - if expiresAt == "" { - expiresAt = strings.TrimSpace(raw.ExpiresAtCamel) - } - if expiresAt == "" { - continue - } - credits = append(credits, OpenAIRateLimitResetCreditDetail{ExpiresAt: expiresAt}) - } - return credits, nil -} - -func firstNonEmptyResetCreditPayload(lists ...[]openAIRateLimitResetCreditDetailPayload) []openAIRateLimitResetCreditDetailPayload { - for _, list := range lists { - if len(list) > 0 { - return list - } - } - return nil -} - // buildCodexSparkWindowExtraUpdates extracts Codex Spark usage windows from the // /wham/usage response body's additional_rate_limits, matching the entry with // MeteredFeature == "codex_bengalfox". It produces plain codex_* keys (NOT the diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index c56600d3f5..213669f5e6 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -251,11 +251,11 @@ func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing. t.Run(tt.name, func(t *testing.T) { got, err := parseOpenAIRateLimitResetCreditDetails([]byte(tt.body)) require.NoError(t, err) - require.Len(t, got, len(tt.want)) + require.Len(t, got.Credits, len(tt.want)) for i := range tt.want { - require.Equal(t, tt.want[i], got[i].ExpiresAt) + require.Equal(t, tt.want[i], got.Credits[i].ExpiresAt) } - encoded, err := json.Marshal(got) + encoded, err := json.Marshal(got.Credits) require.NoError(t, err) require.NotContains(t, string(encoded), "secret-id") })