diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 73187963c72..091706ff605 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -1402,6 +1402,7 @@ func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (* auth.Metadata["account_id"] = td.AccountID } auth.Metadata["email"] = td.Email + applyCodexCapacityClaims(auth, td.IDToken) // Use unified key in files auth.Metadata["expired"] = td.Expire auth.Metadata["type"] = "codex" @@ -1410,6 +1411,46 @@ func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (* return auth, nil } +func applyCodexCapacityClaims(auth *cliproxyauth.Auth, idToken string) { + if auth == nil || strings.TrimSpace(idToken) == "" { + return + } + claims, errParse := codexauth.ParseJWTToken(idToken) + if errParse != nil || claims == nil { + if errParse != nil { + log.Warnf("codex executor: failed to parse refreshed capacity claims: %v", errParse) + } + return + } + + info := claims.CodexAuthInfo + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + if planType := strings.TrimSpace(info.ChatgptPlanType); planType != "" { + auth.Attributes["plan_type"] = planType + auth.Metadata["chatgpt_plan_type"] = planType + } + if accountID := strings.TrimSpace(info.ChatgptAccountID); accountID != "" { + auth.Metadata["chatgpt_account_id"] = accountID + } else { + delete(auth.Metadata, "chatgpt_account_id") + } + if info.ChatgptSubscriptionActiveStart != nil { + auth.Metadata["chatgpt_subscription_active_start"] = info.ChatgptSubscriptionActiveStart + } else { + delete(auth.Metadata, "chatgpt_subscription_active_start") + } + if info.ChatgptSubscriptionActiveUntil != nil { + auth.Metadata["chatgpt_subscription_active_until"] = info.ChatgptSubscriptionActiveUntil + } else { + delete(auth.Metadata, "chatgpt_subscription_active_until") + } +} + type codexIdentityConfuseState struct { enabled bool authID string diff --git a/internal/runtime/executor/codex_executor_refresh_test.go b/internal/runtime/executor/codex_executor_refresh_test.go new file mode 100644 index 00000000000..880298ad9c0 --- /dev/null +++ b/internal/runtime/executor/codex_executor_refresh_test.go @@ -0,0 +1,97 @@ +package executor + +import ( + "encoding/base64" + "encoding/json" + "testing" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func TestApplyCodexCapacityClaimsUpdatesPlanType(t *testing.T) { + auth := &cliproxyauth.Auth{ + ID: "codex-auth", + Provider: "codex", + Attributes: map[string]string{ + "plan_type": "free", + }, + Metadata: map[string]any{ + "chatgpt_plan_type": "free", + }, + } + + applyCodexCapacityClaims(auth, codexCapacityTestJWT(t, map[string]any{ + "email": "user@example.com", + "https://api.openai.com/auth": map[string]any{ + "chatgpt_account_id": "acct-123", + "chatgpt_plan_type": "plus", + "chatgpt_subscription_active_start": "2026-06-12T00:00:00Z", + "chatgpt_subscription_active_until": "2026-07-12T00:00:00Z", + "chatgpt_subscription_last_checked": "2026-06-12T00:01:00Z", + }, + })) + + if got := auth.Attributes["plan_type"]; got != "plus" { + t.Fatalf("plan_type attribute = %q, want plus", got) + } + if got := auth.Metadata["chatgpt_plan_type"]; got != "plus" { + t.Fatalf("chatgpt_plan_type metadata = %q, want plus", got) + } + if got := auth.Metadata["chatgpt_account_id"]; got != "acct-123" { + t.Fatalf("chatgpt_account_id metadata = %q, want acct-123", got) + } + if got := auth.Metadata["chatgpt_subscription_active_until"]; got != "2026-07-12T00:00:00Z" { + t.Fatalf("chatgpt_subscription_active_until metadata = %q, want 2026-07-12T00:00:00Z", got) + } +} + +func TestApplyCodexCapacityClaimsPreservesMissingPlanAndClearsMissingValues(t *testing.T) { + auth := &cliproxyauth.Auth{ + ID: "codex-auth", + Provider: "codex", + Attributes: map[string]string{ + "plan_type": "plus", + }, + Metadata: map[string]any{ + "chatgpt_account_id": "acct-123", + "chatgpt_plan_type": "plus", + "chatgpt_subscription_active_start": "2026-06-12T00:00:00Z", + "chatgpt_subscription_active_until": "2026-07-12T00:00:00Z", + }, + } + + applyCodexCapacityClaims(auth, codexCapacityTestJWT(t, map[string]any{ + "email": "user@example.com", + "https://api.openai.com/auth": map[string]any{}, + })) + + if got := auth.Attributes["plan_type"]; got != "plus" { + t.Fatalf("expected missing plan claim to preserve plan_type attribute, got %q", got) + } + if got := auth.Metadata["chatgpt_plan_type"]; got != "plus" { + t.Fatalf("expected missing plan claim to preserve chatgpt_plan_type metadata, got %q", got) + } + for _, key := range []string{ + "chatgpt_account_id", + "chatgpt_subscription_active_start", + "chatgpt_subscription_active_until", + } { + if _, ok := auth.Metadata[key]; ok { + t.Fatalf("expected missing claim to clear metadata %q", key) + } + } +} + +func codexCapacityTestJWT(t *testing.T, claims map[string]any) string { + t.Helper() + header, errMarshalHeader := json.Marshal(map[string]string{"alg": "none", "typ": "JWT"}) + if errMarshalHeader != nil { + t.Fatalf("marshal jwt header: %v", errMarshalHeader) + } + payload, errMarshalPayload := json.Marshal(claims) + if errMarshalPayload != nil { + t.Fatalf("marshal jwt claims: %v", errMarshalPayload) + } + return base64.RawURLEncoding.EncodeToString(header) + "." + + base64.RawURLEncoding.EncodeToString(payload) + "." +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 08b81dadc06..dc5a86ddb58 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "path/filepath" + "reflect" "sort" "strconv" "strings" @@ -1425,6 +1426,178 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { return auth.Clone(), nil } +// ShouldInheritModelStates reports whether an active auth update should keep +// previous per-model runtime availability. Capacity identity changes, such as a +// Codex plan upgrade, must clear stale quota cooldowns so the account can be +// retried immediately with its new limits. +func ShouldInheritModelStates(existing, incoming *Auth) bool { + if existing == nil || incoming == nil { + return false + } + if existing.Disabled || existing.Status == StatusDisabled || incoming.Disabled || incoming.Status == StatusDisabled { + return false + } + if len(incoming.ModelStates) > 0 || len(existing.ModelStates) == 0 { + return false + } + return !capacityIdentityChanged(existing, incoming) +} + +func shouldClearModelStates(existing, incoming *Auth) bool { + if existing == nil || incoming == nil || len(existing.ModelStates) == 0 || len(incoming.ModelStates) == 0 { + return false + } + if existing.Disabled || existing.Status == StatusDisabled || incoming.Disabled || incoming.Status == StatusDisabled { + return true + } + return capacityIdentityChanged(existing, incoming) +} + +var capacityIdentityKeys = map[string]struct{}{ + "account_id": {}, + "account_uuid": {}, + "accountid": {}, + "balance": {}, + "chatgpt_account_id": {}, + "chatgpt_plan_type": {}, + "chatgpt_subscription_active_start": {}, + "chatgpt_subscription_active_until": {}, + "credit": {}, + "credits": {}, + "organization_id": {}, + "org_id": {}, + "plan": {}, + "plan_type": {}, + "project_id": {}, + "quota": {}, + "quota_limit": {}, + "service_tier": {}, + "subscription": {}, + "subscription_plan": {}, + "subscription_status": {}, + "team_id": {}, + "tier": {}, +} + +func authCapacityIdentitySignature(auth *Auth) map[string]string { + if auth == nil { + return nil + } + signature := make(map[string]string) + for key, value := range auth.Attributes { + normalized := strings.ToLower(strings.TrimSpace(key)) + canonical := capacityIdentityCanonicalKey(normalized) + if canonical == "" { + continue + } + val := strings.TrimSpace(value) + if val == "" { + continue + } + signature[canonical] = val + } + for key, value := range auth.Metadata { + normalized := strings.ToLower(strings.TrimSpace(key)) + canonical := capacityIdentityCanonicalKey(normalized) + if canonical == "" { + continue + } + val := capacityIdentityValue(value) + if val == "" { + continue + } + if _, exists := signature[canonical]; exists { + continue + } + signature[canonical] = val + } + return signature +} + +func capacityIdentityCanonicalKey(normalized string) string { + if _, ok := capacityIdentityKeys[normalized]; !ok { + return "" + } + switch normalized { + case "account_id", "account_uuid", "accountid", "chatgpt_account_id": + return "account_id" + case "chatgpt_plan_type", "plan", "plan_type", "subscription_plan": + return "plan_type" + case "chatgpt_subscription_active_start": + return "subscription_active_start" + case "chatgpt_subscription_active_until": + return "subscription_active_until" + case "credit", "credits": + return "credits" + case "organization_id", "org_id": + return "org_id" + default: + return normalized + } +} + +func capacityIdentityValue(value any) string { + if value == nil { + return "" + } + val := reflect.ValueOf(value) + for val.Kind() == reflect.Ptr { + if val.IsNil() { + return "" + } + val = val.Elem() + } + underlying := val.Interface() + switch typed := underlying.(type) { + case string: + return strings.TrimSpace(typed) + case []byte: + return strings.TrimSpace(string(typed)) + case json.Number: + return strings.TrimSpace(typed.String()) + case bool: + if typed { + return "true" + } + return "false" + case time.Time: + return typed.Format(time.RFC3339) + default: + raw, err := json.Marshal(typed) + if err != nil { + return "" + } + return strings.TrimSpace(string(raw)) + } +} + +func capacityIdentityChanged(existing, incoming *Auth) bool { + existingSignature := authCapacityIdentitySignature(existing) + if len(existingSignature) == 0 { + return false + } + incomingSignature := authCapacityIdentitySignature(incoming) + for key, existingValue := range existingSignature { + incomingValue, ok := incomingSignature[key] + if !ok || incomingValue != existingValue { + return true + } + } + return false +} + +func capacityIdentitySignatureEqual(left, right map[string]string) bool { + if len(left) != len(right) { + return false + } + for key, leftValue := range left { + if right[key] != leftValue { + return false + } + } + return true +} + // Update replaces an existing auth entry and notifies hooks. func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { if auth == nil || auth.ID == "" { @@ -1443,10 +1616,10 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { auth.Success = existing.Success auth.Failed = existing.Failed auth.recentRequests = existing.recentRequests - if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { - if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { - auth.ModelStates = existing.ModelStates - } + if ShouldInheritModelStates(existing, auth) { + auth.ModelStates = existing.ModelStates + } else if shouldClearModelStates(existing, auth) { + clearRuntimeModelStates(auth) } auth.EnsureIndex() authClone := auth.Clone() @@ -1462,6 +1635,20 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { return auth.Clone(), nil } +func clearRuntimeModelStates(auth *Auth) { + if auth == nil { + return + } + auth.ModelStates = nil + clearAggregatedAvailability(auth) + if !auth.Disabled && auth.Status != StatusDisabled { + if auth.Status == StatusError { + auth.Status = StatusActive + } + auth.StatusMessage = "" + } +} + // Remove deletes an auth from runtime state without persisting. // Disk and token-store deletion must be handled by the caller. func (m *Manager) Remove(ctx context.Context, id string) { diff --git a/sdk/cliproxy/auth/conductor_update_test.go b/sdk/cliproxy/auth/conductor_update_test.go index 7dd44ff801e..4857cee6961 100644 --- a/sdk/cliproxy/auth/conductor_update_test.go +++ b/sdk/cliproxy/auth/conductor_update_test.go @@ -3,6 +3,7 @@ package auth import ( "context" "testing" + "time" ) func TestManager_Update_PreservesModelStates(t *testing.T) { @@ -202,3 +203,255 @@ func TestManager_Update_ActiveInheritsModelStates(t *testing.T) { t.Fatalf("expected BackoffLevel to be %d, got %d", backoffLevel, state.Quota.BackoffLevel) } } + +func TestManager_Update_PlanChangeDoesNotInheritQuotaModelStates(t *testing.T) { + m := NewManager(nil, nil, nil) + + model := "codex-model" + nextRetry := time.Now().Add(30 * time.Minute) + + if _, err := m.Register(context.Background(), &Auth{ + ID: "auth-plan-upgrade", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "plan_type": "free", + }, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: nextRetry, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: nextRetry, + BackoffLevel: 5, + }, + }, + }, + }); err != nil { + t.Fatalf("register auth: %v", err) + } + + if _, err := m.Update(context.Background(), &Auth{ + ID: "auth-plan-upgrade", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "plan_type": "plus", + }, + }); err != nil { + t.Fatalf("update auth: %v", err) + } + + updated, ok := m.GetByID("auth-plan-upgrade") + if !ok || updated == nil { + t.Fatalf("expected auth to be present") + } + if len(updated.ModelStates) != 0 { + t.Fatalf("expected plan change to clear stale ModelStates, got %d entries", len(updated.ModelStates)) + } +} + +func TestManager_Update_PlanChangeClearsClonedQuotaModelStates(t *testing.T) { + m := NewManager(nil, nil, nil) + + model := "codex-model" + nextRetry := time.Now().Add(30 * time.Minute) + + registered, err := m.Register(context.Background(), &Auth{ + ID: "auth-cloned-plan-upgrade", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "plan_type": "free", + }, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: nextRetry, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: nextRetry, + BackoffLevel: 5, + }, + }, + }, + }) + if err != nil { + t.Fatalf("register auth: %v", err) + } + + refreshed := registered.Clone() + refreshed.Attributes["plan_type"] = "plus" + refreshed.Status = StatusError + refreshed.StatusMessage = "quota exhausted" + refreshed.Unavailable = true + refreshed.NextRetryAfter = nextRetry + refreshed.Quota = QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: nextRetry, + BackoffLevel: 5, + } + if len(refreshed.ModelStates) == 0 { + t.Fatalf("expected cloned refresh auth to carry old ModelStates") + } + if _, err := m.Update(context.Background(), refreshed); err != nil { + t.Fatalf("update auth: %v", err) + } + + updated, ok := m.GetByID("auth-cloned-plan-upgrade") + if !ok || updated == nil { + t.Fatalf("expected auth to be present") + } + if len(updated.ModelStates) != 0 { + t.Fatalf("expected cloned plan change to clear stale ModelStates, got %d entries", len(updated.ModelStates)) + } + if updated.Unavailable || !updated.NextRetryAfter.IsZero() || updated.Quota.Exceeded || updated.Quota.BackoffLevel != 0 || updated.StatusMessage != "" { + t.Fatalf("expected cloned plan change to clear aggregate cooldown, got unavailable=%v next=%v quota=%+v message=%q", updated.Unavailable, updated.NextRetryAfter, updated.Quota, updated.StatusMessage) + } + if updated.Status != StatusActive { + t.Fatalf("expected cloned plan change to restore active status, got %s", updated.Status) + } +} + +func TestManager_Update_SamePlanInheritsQuotaModelStates(t *testing.T) { + m := NewManager(nil, nil, nil) + + model := "codex-model" + nextRetry := time.Now().Add(30 * time.Minute) + + if _, err := m.Register(context.Background(), &Auth{ + ID: "auth-same-plan", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "plan_type": "plus", + }, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: nextRetry, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: nextRetry, + BackoffLevel: 5, + }, + }, + }, + }); err != nil { + t.Fatalf("register auth: %v", err) + } + + if _, err := m.Update(context.Background(), &Auth{ + ID: "auth-same-plan", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "plan_type": "plus", + }, + }); err != nil { + t.Fatalf("update auth: %v", err) + } + + updated, ok := m.GetByID("auth-same-plan") + if !ok || updated == nil { + t.Fatalf("expected auth to be present") + } + if len(updated.ModelStates) == 0 { + t.Fatalf("expected same plan to inherit ModelStates") + } + state := updated.ModelStates[model] + if state == nil || !state.Quota.Exceeded || state.Quota.BackoffLevel != 5 { + t.Fatalf("expected inherited quota state, got %+v", state) + } +} + +func TestManager_Update_SamePlanWithBackfilledCapacityMetadataKeepsClonedModelStates(t *testing.T) { + m := NewManager(nil, nil, nil) + + model := "codex-model" + nextRetry := time.Now().Add(30 * time.Minute) + + registered, err := m.Register(context.Background(), &Auth{ + ID: "auth-backfilled-same-plan", + Provider: "codex", + Status: StatusActive, + Attributes: map[string]string{ + "plan_type": "plus", + }, + Metadata: map[string]any{ + "email": "user@example.com", + }, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusError, + StatusMessage: "quota exhausted", + Unavailable: true, + NextRetryAfter: nextRetry, + Quota: QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: nextRetry, + BackoffLevel: 5, + }, + }, + }, + }) + if err != nil { + t.Fatalf("register auth: %v", err) + } + + refreshed := registered.Clone() + refreshed.Metadata["account_id"] = "acct-123" + refreshed.Metadata["chatgpt_account_id"] = "acct-123" + refreshed.Metadata["chatgpt_plan_type"] = "plus" + refreshed.Metadata["chatgpt_subscription_active_until"] = "2026-07-12T00:00:00Z" + if _, err := m.Update(context.Background(), refreshed); err != nil { + t.Fatalf("update auth: %v", err) + } + + updated, ok := m.GetByID("auth-backfilled-same-plan") + if !ok || updated == nil { + t.Fatalf("expected auth to be present") + } + state := updated.ModelStates[model] + if state == nil || !state.Quota.Exceeded || state.Quota.BackoffLevel != 5 { + t.Fatalf("expected same plan metadata backfill to keep cloned quota state, got %+v", state) + } +} + +func TestCapacityIdentitySignatureNormalizesPointersAndIgnoresEmptyValues(t *testing.T) { + planType := "plus" + subscriptionUntil := time.Date(2026, 6, 12, 0, 0, 0, 0, time.UTC) + + left := authCapacityIdentitySignature(&Auth{ + Attributes: map[string]string{ + "plan_type": "", + }, + Metadata: map[string]any{ + "chatgpt_plan_type": &planType, + "chatgpt_subscription_active_until": &subscriptionUntil, + "subscription_status": "", + }, + }) + right := authCapacityIdentitySignature(&Auth{ + Metadata: map[string]any{ + "chatgpt_plan_type": "plus", + "chatgpt_subscription_active_until": "2026-06-12T00:00:00Z", + }, + }) + + if !capacityIdentitySignatureEqual(left, right) { + t.Fatalf("expected normalized capacity signatures to match, left=%v right=%v", left, right) + } +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index d3cd9a4b63d..7d116fdc9aa 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -688,7 +688,7 @@ func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth if !existing.Disabled && existing.Status != coreauth.StatusDisabled && !auth.Disabled && auth.Status != coreauth.StatusDisabled { auth.LastRefreshedAt = existing.LastRefreshedAt auth.NextRefreshAfter = existing.NextRefreshAfter - if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { + if coreauth.ShouldInheritModelStates(existing, auth) { auth.ModelStates = existing.ModelStates } }