Merge pull request #3638 from infinityf4p/show-reset-credit-expiration

feat(openai): show reset credit expiration
2026-09-24 16:05:44 +08:00 · 2026-07-02 17:39:17 +08:00
parent 6bc6998b0e dfb36e45f1
commit 537bd50469
7 changed files with 442 additions and 4 deletions
@@ -1,9 +1,11 @@
 package service

 import (
+	"bytes"
 	"context"
 	"crypto/rand"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"log/slog"
 	"net/http"
@@ -11,6 +13,7 @@ import (
 	"time"

 	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
+	"github.com/imroc/req/v3"
 )

 // ErrSparkShadowResetNotSupported is returned when ResetCredit is called on a
@@ -23,8 +26,10 @@ var ErrSparkShadowResetNotSupported = infraerrors.New(http.StatusConflict, "SPAR
 // Endpoints used by the OpenAI/ChatGPT/Codex quota query and reset feature.
 const (
 	chatGPTUsageURL             = "https://chatgpt.com/backend-api/wham/usage"
+	chatGPTRateLimitCreditsURL  = "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits"
 	chatGPTRateLimitResetURL    = "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume"
 	openaiQuotaUpstreamTimeout  = 20 * time.Second
+	openaiQuotaCodexBeta        = "codex-1"
 	openaiQuotaCodexOriginator  = "Codex Desktop"
 	openaiQuotaCodexLanguageTag = "zh-CN"
 	openaiQuotaSecFetchSite     = "none"
@@ -57,10 +62,17 @@ type OpenAIAdditionalRateLimit struct {
 	RateLimit      *OpenAIRateLimit `json:"rate_limit,omitempty"`
 }

+// OpenAIRateLimitResetCreditDetail is the sanitized metadata surfaced for one
+// available reset credit. Do not add upstream ids or tokens here.
+type OpenAIRateLimitResetCreditDetail struct {
+	ExpiresAt string `json:"expires_at,omitempty"`
+}
+
 // OpenAIRateLimitResetCredits captures the "available_count" surfaced for the
 // rate_limit_reset_credit grant type, which the reset action consumes.
 type OpenAIRateLimitResetCredits struct {
-	AvailableCount int `json:"available_count"`
+	AvailableCount int                                `json:"available_count"`
+	Credits        []OpenAIRateLimitResetCreditDetail `json:"credits,omitempty"`
 }

 // OpenAIQuotaUsage is the typed projection of /wham/usage we expose to the UI.
@@ -159,9 +171,34 @@ 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)
+	}
 	return &payload, nil
 }

+func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) []OpenAIRateLimitResetCreditDetail {
+	resp, err := client.R().
+		SetContext(ctx).
+		SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)).
+		Get(chatGPTRateLimitCreditsURL)
+	if err != nil {
+		slog.Warn("openai_quota_reset_credit_details_failed", "account_id", accountID, "error", err)
+		return nil
+	}
+	if !resp.IsSuccessState() {
+		slog.Warn("openai_quota_reset_credit_details_failed", "account_id", accountID, "status", resp.StatusCode)
+		return nil
+	}
+
+	credits, 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
+}
+
 // ResetCredit consumes one rate_limit_reset_credit for the given OpenAI account.
 // The redeem_request_id is auto-generated (uuid-like) — upstream uses it for
 // idempotency. Returns the consumed credit metadata so the UI can refresh.
@@ -306,6 +343,7 @@ func buildCodexCommonHeaders(accessToken, chatGPTAccountID string, fedRAMP bool)
 	headers := map[string]string{
 		"authorization":      "Bearer " + accessToken,
 		"chatgpt-account-id": chatGPTAccountID,
+		"openai-beta":        openaiQuotaCodexBeta,
 		"oai-language":       openaiQuotaCodexLanguageTag,
 		"originator":         openaiQuotaCodexOriginator,
 		"accept":             "application/json",
@@ -334,6 +372,65 @@ 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
@@ -214,6 +214,154 @@ func TestPrepareUpstreamCallShadowResolve(t *testing.T) {
 		"prepareUpstreamCall should use parent's chatgpt_account_id after shadow resolve")
 }

+func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) {
+	tests := []struct {
+		name string
+		body string
+		want []string
+	}{
+		{
+			name: "credits",
+			body: `{"credits":[{"id":"secret-id","expires_at":"2026-07-03T04:05:06Z"}]}`,
+			want: []string{"2026-07-03T04:05:06Z"},
+		},
+		{
+			name: "rate limit reset credits",
+			body: `{"rate_limit_reset_credits":[{"expiresAt":"2026-07-04T04:05:06Z"}]}`,
+			want: []string{"2026-07-04T04:05:06Z"},
+		},
+		{
+			name: "items",
+			body: `{"items":[{"expires_at":"2026-07-05T04:05:06Z"}]}`,
+			want: []string{"2026-07-05T04:05:06Z"},
+		},
+		{
+			name: "data",
+			body: `{"data":[{"expires_at":"2026-07-06T04:05:06Z"}]}`,
+			want: []string{"2026-07-06T04:05:06Z"},
+		},
+		{
+			name: "array",
+			body: `[{"expires_at":"2026-07-07T04:05:06Z"}]`,
+			want: []string{"2026-07-07T04:05:06Z"},
+		},
+	}
+
+	for _, tt := range tests {
+		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))
+			for i := range tt.want {
+				require.Equal(t, tt.want[i], got[i].ExpiresAt)
+			}
+			encoded, err := json.Marshal(got)
+			require.NoError(t, err)
+			require.NotContains(t, string(encoded), "secret-id")
+		})
+	}
+}
+
+func TestQueryUsageIncludesResetCreditExpirations_EndToEnd(t *testing.T) {
+	ctx := context.Background()
+	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 capturedBeta string
+	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":
+			_ = json.NewEncoder(w).Encode(OpenAIQuotaUsage{
+				RateLimitResetCredits: &OpenAIRateLimitResetCredits{AvailableCount: 2},
+			})
+		case "/backend-api/wham/rate-limit-reset-credits":
+			detailCalls++
+			capturedBeta = r.Header.Get("OpenAI-Beta")
+			require.Equal(t, "org-parent123", r.Header.Get("ChatGPT-Account-ID"))
+			_, _ = w.Write([]byte(`{"credits":[{"id":"secret-credit-id","expires_at":"2026-07-03T04:05:06Z"},{"expiresAt":"2026-07-04T04:05:06Z"}]}`))
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer srv.Close()
+
+	svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
+	usage, err := svc.QueryUsage(ctx, 100)
+	require.NoError(t, err)
+	require.NotNil(t, usage)
+	require.NotNil(t, usage.RateLimitResetCredits)
+	require.Equal(t, 2, usage.RateLimitResetCredits.AvailableCount)
+	require.Equal(t, 1, detailCalls)
+	require.Equal(t, openaiQuotaCodexBeta, capturedBeta)
+	require.Equal(t, []OpenAIRateLimitResetCreditDetail{
+		{ExpiresAt: "2026-07-03T04:05:06Z"},
+		{ExpiresAt: "2026-07-04T04:05:06Z"},
+	}, usage.RateLimitResetCredits.Credits)
+
+	encoded, err := json.Marshal(usage)
+	require.NoError(t, err)
+	require.NotContains(t, string(encoded), "secret-credit-id")
+}
+
+func TestQueryUsageResetCreditDetails401NonFatal(t *testing.T) {
+	ctx := context.Background()
+	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":
+			_ = json.NewEncoder(w).Encode(OpenAIQuotaUsage{
+				RateLimitResetCredits: &OpenAIRateLimitResetCredits{AvailableCount: 1},
+			})
+		case "/backend-api/wham/rate-limit-reset-credits":
+			detailCalls++
+			w.WriteHeader(http.StatusUnauthorized)
+			_, _ = w.Write([]byte(`{"error":"unauthorized","id":"secret-error-id"}`))
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer srv.Close()
+
+	svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
+	usage, err := svc.QueryUsage(ctx, 100)
+	require.NoError(t, err)
+	require.NotNil(t, usage)
+	require.NotNil(t, usage.RateLimitResetCredits)
+	require.Equal(t, 1, usage.RateLimitResetCredits.AvailableCount)
+	require.Equal(t, 1, detailCalls)
+	require.Empty(t, usage.RateLimitResetCredits.Credits)
+}
+
 // TestResetCreditGetByIDError_FailsClosed 验证守卫「失败关闭」语义：
 // 当守卫的 GetByID 发生瞬时错误时，ResetCredit 必须立即返回该错误，
 // 不得旁路进入 prepareUpstreamCall（否则影子账号会借 resolve 路径操作母账号）。