Skip to content

Commit ee1bb38

Browse files
author
hongwei.dong
committed
fix: clear quota cooldown after capacity changes
1 parent ac4017e commit ee1bb38

5 files changed

Lines changed: 458 additions & 5 deletions

File tree

‎internal/runtime/executor/codex_executor.go‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1402,6 +1402,7 @@ func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*
14021402
auth.Metadata["account_id"] = td.AccountID
14031403
}
14041404
auth.Metadata["email"] = td.Email
1405+
applyCodexCapacityClaims(auth, td.IDToken)
14051406
// Use unified key in files
14061407
auth.Metadata["expired"] = td.Expire
14071408
auth.Metadata["type"] = "codex"
@@ -1410,6 +1411,49 @@ func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*
14101411
return auth, nil
14111412
}
14121413

1414+
func applyCodexCapacityClaims(auth *cliproxyauth.Auth, idToken string) {
1415+
if auth == nil || strings.TrimSpace(idToken) == "" {
1416+
return
1417+
}
1418+
claims, errParse := codexauth.ParseJWTToken(idToken)
1419+
if errParse != nil || claims == nil {
1420+
if errParse != nil {
1421+
log.Warnf("codex executor: failed to parse refreshed capacity claims: %v", errParse)
1422+
}
1423+
return
1424+
}
1425+
1426+
info := claims.CodexAuthInfo
1427+
if auth.Attributes == nil {
1428+
auth.Attributes = make(map[string]string)
1429+
}
1430+
if auth.Metadata == nil {
1431+
auth.Metadata = make(map[string]any)
1432+
}
1433+
if planType := strings.TrimSpace(info.ChatgptPlanType); planType != "" {
1434+
auth.Attributes["plan_type"] = planType
1435+
auth.Metadata["chatgpt_plan_type"] = planType
1436+
} else {
1437+
delete(auth.Attributes, "plan_type")
1438+
delete(auth.Metadata, "chatgpt_plan_type")
1439+
}
1440+
if accountID := strings.TrimSpace(info.ChatgptAccountID); accountID != "" {
1441+
auth.Metadata["chatgpt_account_id"] = accountID
1442+
} else {
1443+
delete(auth.Metadata, "chatgpt_account_id")
1444+
}
1445+
if info.ChatgptSubscriptionActiveStart != nil {
1446+
auth.Metadata["chatgpt_subscription_active_start"] = info.ChatgptSubscriptionActiveStart
1447+
} else {
1448+
delete(auth.Metadata, "chatgpt_subscription_active_start")
1449+
}
1450+
if info.ChatgptSubscriptionActiveUntil != nil {
1451+
auth.Metadata["chatgpt_subscription_active_until"] = info.ChatgptSubscriptionActiveUntil
1452+
} else {
1453+
delete(auth.Metadata, "chatgpt_subscription_active_until")
1454+
}
1455+
}
1456+
14131457
type codexIdentityConfuseState struct {
14141458
enabled bool
14151459
authID string

‎internal/runtime/executor/codex_executor_refresh_test.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package executor
2+
3+
import (
4+
"encoding/base64"
5+
"encoding/json"
6+
"testing"
7+
8+
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
9+
)
10+
11+
func TestApplyCodexCapacityClaimsUpdatesPlanType(t *testing.T) {
12+
auth := &cliproxyauth.Auth{
13+
ID: "codex-auth",
14+
Provider: "codex",
15+
Attributes: map[string]string{
16+
"plan_type": "free",
17+
},
18+
Metadata: map[string]any{
19+
"chatgpt_plan_type": "free",
20+
},
21+
}
22+
23+
applyCodexCapacityClaims(auth, codexCapacityTestJWT(t, map[string]any{
24+
"email": "user@example.com",
25+
"https://api.openai.com/auth": map[string]any{
26+
"chatgpt_account_id": "acct-123",
27+
"chatgpt_plan_type": "plus",
28+
"chatgpt_subscription_active_start": "2026-06-12T00:00:00Z",
29+
"chatgpt_subscription_active_until": "2026-07-12T00:00:00Z",
30+
"chatgpt_subscription_last_checked": "2026-06-12T00:01:00Z",
31+
},
32+
}))
33+
34+
if got := auth.Attributes["plan_type"]; got != "plus" {
35+
t.Fatalf("plan_type attribute = %q, want plus", got)
36+
}
37+
if got := auth.Metadata["chatgpt_plan_type"]; got != "plus" {
38+
t.Fatalf("chatgpt_plan_type metadata = %q, want plus", got)
39+
}
40+
if got := auth.Metadata["chatgpt_account_id"]; got != "acct-123" {
41+
t.Fatalf("chatgpt_account_id metadata = %q, want acct-123", got)
42+
}
43+
if got := auth.Metadata["chatgpt_subscription_active_until"]; got != "2026-07-12T00:00:00Z" {
44+
t.Fatalf("chatgpt_subscription_active_until metadata = %q, want 2026-07-12T00:00:00Z", got)
45+
}
46+
}
47+
48+
func TestApplyCodexCapacityClaimsClearsMissingValues(t *testing.T) {
49+
auth := &cliproxyauth.Auth{
50+
ID: "codex-auth",
51+
Provider: "codex",
52+
Attributes: map[string]string{
53+
"plan_type": "plus",
54+
},
55+
Metadata: map[string]any{
56+
"chatgpt_account_id": "acct-123",
57+
"chatgpt_plan_type": "plus",
58+
"chatgpt_subscription_active_start": "2026-06-12T00:00:00Z",
59+
"chatgpt_subscription_active_until": "2026-07-12T00:00:00Z",
60+
},
61+
}
62+
63+
applyCodexCapacityClaims(auth, codexCapacityTestJWT(t, map[string]any{
64+
"email": "user@example.com",
65+
"https://api.openai.com/auth": map[string]any{},
66+
}))
67+
68+
if _, ok := auth.Attributes["plan_type"]; ok {
69+
t.Fatalf("expected missing plan claim to clear plan_type attribute")
70+
}
71+
for _, key := range []string{
72+
"chatgpt_account_id",
73+
"chatgpt_plan_type",
74+
"chatgpt_subscription_active_start",
75+
"chatgpt_subscription_active_until",
76+
} {
77+
if _, ok := auth.Metadata[key]; ok {
78+
t.Fatalf("expected missing claim to clear metadata %q", key)
79+
}
80+
}
81+
}
82+
83+
func codexCapacityTestJWT(t *testing.T, claims map[string]any) string {
84+
t.Helper()
85+
header, errMarshalHeader := json.Marshal(map[string]string{"alg": "none", "typ": "JWT"})
86+
if errMarshalHeader != nil {
87+
t.Fatalf("marshal jwt header: %v", errMarshalHeader)
88+
}
89+
payload, errMarshalPayload := json.Marshal(claims)
90+
if errMarshalPayload != nil {
91+
t.Fatalf("marshal jwt claims: %v", errMarshalPayload)
92+
}
93+
return base64.RawURLEncoding.EncodeToString(header) + "." +
94+
base64.RawURLEncoding.EncodeToString(payload) + "."
95+
}

‎sdk/cliproxy/auth/conductor.go‎

Lines changed: 135 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"io"
99
"net/http"
1010
"path/filepath"
11+
"reflect"
1112
"sort"
1213
"strconv"
1314
"strings"
@@ -1425,6 +1426,136 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
14251426
return auth.Clone(), nil
14261427
}
14271428

1429+
// ShouldInheritModelStates reports whether an active auth update should keep
1430+
// previous per-model runtime availability. Capacity identity changes, such as a
1431+
// Codex plan upgrade, must clear stale quota cooldowns so the account can be
1432+
// retried immediately with its new limits.
1433+
func ShouldInheritModelStates(existing, incoming *Auth) bool {
1434+
if existing == nil || incoming == nil {
1435+
return false
1436+
}
1437+
if existing.Disabled || existing.Status == StatusDisabled || incoming.Disabled || incoming.Status == StatusDisabled {
1438+
return false
1439+
}
1440+
if len(incoming.ModelStates) > 0 || len(existing.ModelStates) == 0 {
1441+
return false
1442+
}
1443+
return capacityIdentitySignatureEqual(authCapacityIdentitySignature(existing), authCapacityIdentitySignature(incoming))
1444+
}
1445+
1446+
func shouldClearModelStates(existing, incoming *Auth) bool {
1447+
if existing == nil || incoming == nil || len(existing.ModelStates) == 0 || len(incoming.ModelStates) == 0 {
1448+
return false
1449+
}
1450+
if existing.Disabled || existing.Status == StatusDisabled || incoming.Disabled || incoming.Status == StatusDisabled {
1451+
return true
1452+
}
1453+
return !capacityIdentitySignatureEqual(authCapacityIdentitySignature(existing), authCapacityIdentitySignature(incoming))
1454+
}
1455+
1456+
var capacityIdentityKeys = map[string]struct{}{
1457+
"account_id": {},
1458+
"account_uuid": {},
1459+
"accountid": {},
1460+
"balance": {},
1461+
"chatgpt_account_id": {},
1462+
"chatgpt_plan_type": {},
1463+
"chatgpt_subscription_active_start": {},
1464+
"chatgpt_subscription_active_until": {},
1465+
"credit": {},
1466+
"credits": {},
1467+
"organization_id": {},
1468+
"org_id": {},
1469+
"plan": {},
1470+
"plan_type": {},
1471+
"project_id": {},
1472+
"quota": {},
1473+
"quota_limit": {},
1474+
"service_tier": {},
1475+
"subscription": {},
1476+
"subscription_plan": {},
1477+
"subscription_status": {},
1478+
"team_id": {},
1479+
"tier": {},
1480+
}
1481+
1482+
func authCapacityIdentitySignature(auth *Auth) map[string]string {
1483+
if auth == nil {
1484+
return nil
1485+
}
1486+
signature := make(map[string]string)
1487+
for key, value := range auth.Attributes {
1488+
normalized := strings.ToLower(strings.TrimSpace(key))
1489+
if _, ok := capacityIdentityKeys[normalized]; !ok {
1490+
continue
1491+
}
1492+
val := strings.TrimSpace(value)
1493+
if val == "" {
1494+
continue
1495+
}
1496+
signature["attr:"+normalized] = val
1497+
}
1498+
for key, value := range auth.Metadata {
1499+
normalized := strings.ToLower(strings.TrimSpace(key))
1500+
if _, ok := capacityIdentityKeys[normalized]; !ok {
1501+
continue
1502+
}
1503+
val := capacityIdentityValue(value)
1504+
if val == "" {
1505+
continue
1506+
}
1507+
signature["meta:"+normalized] = val
1508+
}
1509+
return signature
1510+
}
1511+
1512+
func capacityIdentityValue(value any) string {
1513+
if value == nil {
1514+
return ""
1515+
}
1516+
val := reflect.ValueOf(value)
1517+
for val.Kind() == reflect.Ptr {
1518+
if val.IsNil() {
1519+
return ""
1520+
}
1521+
val = val.Elem()
1522+
}
1523+
underlying := val.Interface()
1524+
switch typed := underlying.(type) {
1525+
case string:
1526+
return strings.TrimSpace(typed)
1527+
case []byte:
1528+
return strings.TrimSpace(string(typed))
1529+
case json.Number:
1530+
return strings.TrimSpace(typed.String())
1531+
case bool:
1532+
if typed {
1533+
return "true"
1534+
}
1535+
return "false"
1536+
case time.Time:
1537+
return typed.Format(time.RFC3339)
1538+
default:
1539+
raw, err := json.Marshal(typed)
1540+
if err != nil {
1541+
return ""
1542+
}
1543+
return strings.TrimSpace(string(raw))
1544+
}
1545+
}
1546+
1547+
func capacityIdentitySignatureEqual(left, right map[string]string) bool {
1548+
if len(left) != len(right) {
1549+
return false
1550+
}
1551+
for key, leftValue := range left {
1552+
if right[key] != leftValue {
1553+
return false
1554+
}
1555+
}
1556+
return true
1557+
}
1558+
14281559
// Update replaces an existing auth entry and notifies hooks.
14291560
func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
14301561
if auth == nil || auth.ID == "" {
@@ -1443,10 +1574,10 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
14431574
auth.Success = existing.Success
14441575
auth.Failed = existing.Failed
14451576
auth.recentRequests = existing.recentRequests
1446-
if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled {
1447-
if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 {
1448-
auth.ModelStates = existing.ModelStates
1449-
}
1577+
if ShouldInheritModelStates(existing, auth) {
1578+
auth.ModelStates = existing.ModelStates
1579+
} else if shouldClearModelStates(existing, auth) {
1580+
auth.ModelStates = nil
14501581
}
14511582
auth.EnsureIndex()
14521583
authClone := auth.Clone()

0 commit comments

Comments
 (0)