From dd562416274ef7fb7a703351be24926803ea7a3d Mon Sep 17 00:00:00 2001 From: Jacob Clayden Date: Mon, 27 Apr 2026 17:44:56 +0100 Subject: [PATCH 1/2] feat: add Claude proxy account pinning --- cmd/cq/proxy.go | 87 ++++++++- internal/proxy/config.go | 16 ++ internal/proxy/pinned_selector.go | 104 ++++++++++ internal/proxy/pinned_selector_test.go | 257 +++++++++++++++++++++++++ 4 files changed, 461 insertions(+), 3 deletions(-) create mode 100644 internal/proxy/pinned_selector.go create mode 100644 internal/proxy/pinned_selector_test.go diff --git a/cmd/cq/proxy.go b/cmd/cq/proxy.go index 59e815f..1166e58 100644 --- a/cmd/cq/proxy.go +++ b/cmd/cq/proxy.go @@ -22,7 +22,7 @@ import ( func runProxy(args []string) error { if len(args) == 0 { - fmt.Fprintf(os.Stderr, "Usage: cq proxy \n") + fmt.Fprintf(os.Stderr, "Usage: cq proxy \n") return fmt.Errorf("missing subcommand") } switch args[0] { @@ -44,11 +44,55 @@ func runProxy(args []string) error { return err } return runProxyStatus(opts) + case "pin": + return runProxyPin(args[1:]) default: return fmt.Errorf("unknown proxy command: %s", args[0]) } } +func runProxyPin(args []string) error { + cfg, err := proxy.LoadConfig() + if err != nil { + return fmt.Errorf("load config: %w", err) + } + + // cq proxy pin --clear + if len(args) == 1 && args[0] == "--clear" { + cfg.PinnedClaudeAccount = "" + if err := proxy.SaveConfig(cfg); err != nil { + return fmt.Errorf("save config: %w", err) + } + fmt.Println("Pinned Claude account cleared.") + fmt.Println("A running proxy will pick up the change shortly.") + return nil + } + + // cq proxy pin + if len(args) == 1 { + cfg.PinnedClaudeAccount = args[0] + if err := proxy.SaveConfig(cfg); err != nil { + return fmt.Errorf("save config: %w", err) + } + fmt.Printf("Pinned Claude account set to %q.\n", args[0]) + fmt.Println("A running proxy will pick up the change shortly.") + return nil + } + + // cq proxy pin (no args) — show current pin + if len(args) == 0 { + if cfg.PinnedClaudeAccount == "" { + fmt.Println("No pin is active. All Claude requests use automatic account selection.") + } else { + fmt.Printf("Pinned Claude account: %s\n", cfg.PinnedClaudeAccount) + } + return nil + } + + fmt.Fprintf(os.Stderr, "Usage: cq proxy pin [--clear | ]\n") + return fmt.Errorf("unexpected arguments") +} + type proxyCommandOptions struct { Port int } @@ -104,7 +148,14 @@ func runProxyStart(opts proxyCommandOptions) error { claudeProvider := claudeprov.New(refreshClient) quotaCache := proxy.NewQuotaCache(claudeProvider.FetchAccountUsage, cache.DefaultDir()) - selector := proxy.NewAccountSelector(discover, activeEmail, quotaCache) + baseSelector := proxy.NewAccountSelector(discover, activeEmail, quotaCache) + selector := proxy.NewPinnedClaudeSelector(baseSelector, discover, cfg.PinnedClaudeAccount) + if cfg.PinnedClaudeAccount != "" { + fmt.Fprintf(os.Stderr, "cq: pinned claude account: %s\n", cfg.PinnedClaudeAccount) + } + proxyCtx, proxyCancel := context.WithCancel(context.Background()) + defer proxyCancel() + startProxyConfigReload(proxyCtx, selector) accountsMgr := &claudeprov.Accounts{HTTP: refreshClient} switcher := proxy.AccountSwitcher(func(ctx context.Context, email string) error { @@ -272,13 +323,43 @@ func runProxyStart(opts proxyCommandOptions) error { Refresher: proxyRefresher, } - err = srv.ListenAndServe(context.Background()) + err = srv.ListenAndServe(proxyCtx) if headroom != nil { headroom.Stop() } return err } +func startProxyConfigReload(ctx context.Context, selector *proxy.PinnedClaudeSelector) { + go func() { + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + reloadProxyConfig(selector) + } + } + }() +} + +func reloadProxyConfig(selector *proxy.PinnedClaudeSelector) { + defer func() { + if r := recover(); r != nil { + fmt.Fprintf(os.Stderr, "cq: proxy config reload panic: %v\n", r) + } + }() + + cfg, err := proxy.LoadConfig() + if err != nil { + fmt.Fprintf(os.Stderr, "cq: proxy config reload: %v\n", err) + return + } + selector.SetPin(cfg.PinnedClaudeAccount) +} + func runProxyStatus(opts proxyCommandOptions) error { cfg, err := proxy.LoadConfig() if err != nil { diff --git a/internal/proxy/config.go b/internal/proxy/config.go index bb29c7b..84b89ae 100644 --- a/internal/proxy/config.go +++ b/internal/proxy/config.go @@ -33,6 +33,9 @@ type Config struct { // When omitted, cq defaults to cache mode. Explicit "token" preserves the // legacy token-optimised behaviour. HeadroomMode string `json:"headroom_mode,omitempty"` + // PinnedClaudeAccount forces the proxy to route all Claude requests through + // a specific account identified by email or AccountUUID. Omitted when empty. + PinnedClaudeAccount string `json:"pinned_claude_account,omitempty"` } // ResolvedHeadroomMode returns the effective HeadroomMode for this config. @@ -129,6 +132,19 @@ func generateToken() (string, error) { return base64.RawURLEncoding.EncodeToString(buf), nil } +// SaveConfig writes cfg to the standard proxy config path atomically. +func SaveConfig(cfg *Config) error { + if cfg == nil { + return fmt.Errorf("proxy config is nil") + } + saved := *cfg + saved.setDefaults() + if err := saved.validate(); err != nil { + return err + } + return saveConfig(filepath.Join(configDir(), "proxy.json"), &saved) +} + func saveConfig(path string, cfg *Config) error { if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { return fmt.Errorf("create config dir: %w", err) diff --git a/internal/proxy/pinned_selector.go b/internal/proxy/pinned_selector.go new file mode 100644 index 0000000..ca9b03f --- /dev/null +++ b/internal/proxy/pinned_selector.go @@ -0,0 +1,104 @@ +package proxy + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/jacobcxdev/cq/internal/keyring" +) + +// PinnedClaudeSelector wraps an inner ClaudeSelector and, when a pin is set, +// routes all requests through the pinned account (by email or AccountUUID). +// If the pinned account is excluded or unusable, it delegates to the inner +// selector. Thread-safe: pin may be updated while the proxy is running. +type PinnedClaudeSelector struct { + inner ClaudeSelector + discover ClaudeDiscoverer + + mu sync.RWMutex + pin string // email or AccountUUID; empty means no pin +} + +// NewPinnedClaudeSelector creates a PinnedClaudeSelector. +// initialPin may be empty (no pin active). +func NewPinnedClaudeSelector(inner ClaudeSelector, discover ClaudeDiscoverer, initialPin string) *PinnedClaudeSelector { + return &PinnedClaudeSelector{ + inner: inner, + discover: discover, + pin: initialPin, + } +} + +// SetPin atomically updates the active pin. An empty string clears the pin. +func (s *PinnedClaudeSelector) SetPin(pin string) { + s.mu.Lock() + s.pin = pin + s.mu.Unlock() +} + +// Pin returns the current pin value. +func (s *PinnedClaudeSelector) Pin() string { + s.mu.RLock() + defer s.mu.RUnlock() + return s.pin +} + +// Select implements ClaudeSelector. +// +// When a pin is set: +// - Discovers accounts and finds one matching the pin by Email or AccountUUID. +// - If the matched account is in the exclude set, delegates to inner. +// - If matched and usable (non-empty access token, and either unexpired, +// no ExpiresAt, or expired with a refresh token), returns a copy directly. +// - If matched but unusable (no access token, or expired without refresh +// token), returns an error. +// - If not found, returns an error containing the pin value. +// +// When no pin is set, delegates to inner.Select. +func (s *PinnedClaudeSelector) Select(ctx context.Context, exclude ...string) (*keyring.ClaudeOAuth, error) { + s.mu.RLock() + pin := s.pin + s.mu.RUnlock() + + if pin == "" { + return s.inner.Select(ctx, exclude...) + } + + accounts := s.discover() + var matched *keyring.ClaudeOAuth + for i := range accounts { + a := &accounts[i] + if a.Email == pin || a.AccountUUID == pin { + matched = a + break + } + } + + if matched == nil { + return nil, fmt.Errorf("pinned Claude account %q not found", pin) + } + + // If the pinned account is excluded, fall back to the inner selector. + excludeSet := make(map[string]bool, len(exclude)) + for _, e := range exclude { + excludeSet[e] = true + } + if isExcluded(matched, excludeSet) { + return s.inner.Select(ctx, exclude...) + } + + // Usability check: must have an access token, and must be either unexpired, + // have no expiry, or be expired with a refresh token (transport will refresh). + if matched.AccessToken == "" { + return nil, fmt.Errorf("pinned Claude account %q has no access token", pin) + } + now := time.Now().UnixMilli() + if matched.ExpiresAt != 0 && matched.ExpiresAt <= now && matched.RefreshToken == "" { + return nil, fmt.Errorf("pinned Claude account %q token is expired and has no refresh token", pin) + } + + result := *matched + return &result, nil +} diff --git a/internal/proxy/pinned_selector_test.go b/internal/proxy/pinned_selector_test.go new file mode 100644 index 0000000..0aefd06 --- /dev/null +++ b/internal/proxy/pinned_selector_test.go @@ -0,0 +1,257 @@ +package proxy + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/jacobcxdev/cq/internal/keyring" +) + +// innerSelectorFunc adapts a function to the ClaudeSelector interface. +type innerSelectorFunc func(ctx context.Context, exclude ...string) (*keyring.ClaudeOAuth, error) + +func (f innerSelectorFunc) Select(ctx context.Context, exclude ...string) (*keyring.ClaudeOAuth, error) { + return f(ctx, exclude...) +} + +func makePinnedSelector(accounts []keyring.ClaudeOAuth, pin string) (*PinnedClaudeSelector, *innerSelectorFunc) { + inner := innerSelectorFunc(func(ctx context.Context, exclude ...string) (*keyring.ClaudeOAuth, error) { + // Simple inner: return first non-excluded account. + for i := range accounts { + a := &accounts[i] + excludeSet := make(map[string]bool, len(exclude)) + for _, e := range exclude { + excludeSet[e] = true + } + if !isExcluded(a, excludeSet) && a.AccessToken != "" { + result := *a + return &result, nil + } + } + return nil, errors.New("no accounts available") + }) + discover := ClaudeDiscoverer(func() []keyring.ClaudeOAuth { + return accounts + }) + sel := NewPinnedClaudeSelector(inner, discover, pin) + return sel, &inner +} + +func TestPinnedClaudeSelector_NoPinDelegatesToInner(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "a@test.com", AccessToken: "tok-a", ExpiresAt: future}, + {Email: "b@test.com", AccessToken: "tok-b", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // inner returns first non-excluded — a@test.com + if acct.Email != "a@test.com" { + t.Errorf("email = %q, want a@test.com", acct.Email) + } +} + +func TestPinnedClaudeSelector_PinByEmailOverBetterQuota(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "tok-pin", ExpiresAt: future}, + {Email: "other@test.com", AccessToken: "tok-other", ExpiresAt: future + 9999}, + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if acct.Email != "pinned@test.com" { + t.Errorf("email = %q, want pinned@test.com", acct.Email) + } +} + +func TestPinnedClaudeSelector_PinByUUID(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "a@test.com", AccountUUID: "uuid-a", AccessToken: "tok-a", ExpiresAt: future}, + {Email: "b@test.com", AccountUUID: "uuid-b", AccessToken: "tok-b", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "uuid-b") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if acct.AccountUUID != "uuid-b" { + t.Errorf("AccountUUID = %q, want uuid-b", acct.AccountUUID) + } +} + +func TestPinnedClaudeSelector_ExcludedPinDelegatesToInner(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "tok-pin", ExpiresAt: future}, + {Email: "fallback@test.com", AccessToken: "tok-fb", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + // Exclude the pinned account — should fall back to inner selector. + acct, err := sel.Select(context.Background(), "pinned@test.com") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if acct.Email != "fallback@test.com" { + t.Errorf("email = %q, want fallback@test.com", acct.Email) + } +} + +func TestPinnedClaudeSelector_NotFoundReturnsError(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "a@test.com", AccessToken: "tok-a", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "missing@test.com") + + _, err := sel.Select(context.Background()) + if err == nil { + t.Fatal("expected error for missing pin, got nil") + } + if !strings.Contains(err.Error(), `pinned Claude account "missing@test.com" not found`) { + t.Errorf("error = %q, want to contain pin-not-found message", err.Error()) + } +} + +func TestPinnedClaudeSelector_UnusableNoAccessToken(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "", ExpiresAt: future}, // no token + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + _, err := sel.Select(context.Background()) + if err == nil { + t.Fatal("expected error for unusable pin (no token), got nil") + } + if !strings.Contains(err.Error(), "no access token") { + t.Errorf("error = %q, want to contain 'no access token'", err.Error()) + } +} + +func TestPinnedClaudeSelector_UnusableExpiredNoRefresh(t *testing.T) { + past := time.Now().UnixMilli() - 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "tok", ExpiresAt: past, RefreshToken: ""}, + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + _, err := sel.Select(context.Background()) + if err == nil { + t.Fatal("expected error for expired token without refresh, got nil") + } + if !strings.Contains(err.Error(), "expired") { + t.Errorf("error = %q, want to contain 'expired'", err.Error()) + } +} + +func TestPinnedClaudeSelector_ExpiredWithRefreshTokenIsUsable(t *testing.T) { + past := time.Now().UnixMilli() - 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "tok", ExpiresAt: past, RefreshToken: "rt"}, + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if acct.Email != "pinned@test.com" { + t.Errorf("email = %q, want pinned@test.com", acct.Email) + } +} + +func TestPinnedClaudeSelector_NoExpiryIsUsable(t *testing.T) { + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "tok", ExpiresAt: 0}, + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if acct.Email != "pinned@test.com" { + t.Errorf("email = %q, want pinned@test.com", acct.Email) + } +} + +func TestPinnedClaudeSelector_ReturnsCopy(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccessToken: "original", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "pinned@test.com") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatal(err) + } + acct.AccessToken = "mutated" + + acct2, err := sel.Select(context.Background()) + if err != nil { + t.Fatal(err) + } + if acct2.AccessToken != "original" { + t.Error("selector returned reference instead of copy") + } +} + +func TestPinnedClaudeSelector_SetPinUpdatesPin(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "a@test.com", AccessToken: "tok-a", ExpiresAt: future}, + {Email: "b@test.com", AccessToken: "tok-b", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "a@test.com") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatal(err) + } + if acct.Email != "a@test.com" { + t.Errorf("before SetPin: email = %q, want a@test.com", acct.Email) + } + + sel.SetPin("b@test.com") + + acct, err = sel.Select(context.Background()) + if err != nil { + t.Fatal(err) + } + if acct.Email != "b@test.com" { + t.Errorf("after SetPin: email = %q, want b@test.com", acct.Email) + } +} + +func TestPinnedClaudeSelector_ClearPinDelegatesToInner(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "a@test.com", AccessToken: "tok-a", ExpiresAt: future}, + } + sel, _ := makePinnedSelector(accounts, "a@test.com") + sel.SetPin("") + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatal(err) + } + // inner returns a@test.com (only account) + if acct.Email != "a@test.com" { + t.Errorf("email = %q, want a@test.com", acct.Email) + } +} From 7a917c52b6a70680a83fcf45a1557729eb0d6c2a Mon Sep 17 00:00:00 2001 From: Jacob Clayden Date: Mon, 27 Apr 2026 17:50:10 +0100 Subject: [PATCH 2/2] fix: expire Claude proxy pin on quota exhaustion --- cmd/cq/proxy.go | 20 ++++++++++++- internal/proxy/pinned_selector.go | 41 ++++++++++++++++++++++++-- internal/proxy/pinned_selector_test.go | 41 +++++++++++++++++++++++++- 3 files changed, 97 insertions(+), 5 deletions(-) diff --git a/cmd/cq/proxy.go b/cmd/cq/proxy.go index 1166e58..e0c93ac 100644 --- a/cmd/cq/proxy.go +++ b/cmd/cq/proxy.go @@ -149,7 +149,8 @@ func runProxyStart(opts proxyCommandOptions) error { claudeProvider := claudeprov.New(refreshClient) quotaCache := proxy.NewQuotaCache(claudeProvider.FetchAccountUsage, cache.DefaultDir()) baseSelector := proxy.NewAccountSelector(discover, activeEmail, quotaCache) - selector := proxy.NewPinnedClaudeSelector(baseSelector, discover, cfg.PinnedClaudeAccount) + selector := proxy.NewPinnedClaudeSelector(baseSelector, discover, cfg.PinnedClaudeAccount, quotaCache) + selector.SetPinExpireFunc(clearPersistedClaudePin) if cfg.PinnedClaudeAccount != "" { fmt.Fprintf(os.Stderr, "cq: pinned claude account: %s\n", cfg.PinnedClaudeAccount) } @@ -330,6 +331,23 @@ func runProxyStart(opts proxyCommandOptions) error { return err } +func clearPersistedClaudePin(pin string) { + cfg, err := proxy.LoadConfig() + if err != nil { + fmt.Fprintf(os.Stderr, "cq: clear expired claude pin %q: %v\n", pin, err) + return + } + if cfg.PinnedClaudeAccount != pin { + return + } + cfg.PinnedClaudeAccount = "" + if err := proxy.SaveConfig(cfg); err != nil { + fmt.Fprintf(os.Stderr, "cq: clear expired claude pin %q: %v\n", pin, err) + return + } + fmt.Fprintf(os.Stderr, "cq: cleared expired claude pin: %s\n", pin) +} + func startProxyConfigReload(ctx context.Context, selector *proxy.PinnedClaudeSelector) { go func() { ticker := time.NewTicker(5 * time.Second) diff --git a/internal/proxy/pinned_selector.go b/internal/proxy/pinned_selector.go index ca9b03f..96b1304 100644 --- a/internal/proxy/pinned_selector.go +++ b/internal/proxy/pinned_selector.go @@ -14,8 +14,10 @@ import ( // If the pinned account is excluded or unusable, it delegates to the inner // selector. Thread-safe: pin may be updated while the proxy is running. type PinnedClaudeSelector struct { - inner ClaudeSelector - discover ClaudeDiscoverer + inner ClaudeSelector + discover ClaudeDiscoverer + quota QuotaReader + onPinExpire func(string) mu sync.RWMutex pin string // email or AccountUUID; empty means no pin @@ -23,14 +25,22 @@ type PinnedClaudeSelector struct { // NewPinnedClaudeSelector creates a PinnedClaudeSelector. // initialPin may be empty (no pin active). -func NewPinnedClaudeSelector(inner ClaudeSelector, discover ClaudeDiscoverer, initialPin string) *PinnedClaudeSelector { +func NewPinnedClaudeSelector(inner ClaudeSelector, discover ClaudeDiscoverer, initialPin string, quota QuotaReader) *PinnedClaudeSelector { return &PinnedClaudeSelector{ inner: inner, discover: discover, pin: initialPin, + quota: quota, } } +// SetPinExpireFunc configures a callback invoked after the selector clears an exhausted pin. +func (s *PinnedClaudeSelector) SetPinExpireFunc(f func(string)) { + s.mu.Lock() + s.onPinExpire = f + s.mu.Unlock() +} + // SetPin atomically updates the active pin. An empty string clears the pin. func (s *PinnedClaudeSelector) SetPin(pin string) { s.mu.Lock() @@ -88,6 +98,10 @@ func (s *PinnedClaudeSelector) Select(ctx context.Context, exclude ...string) (* if isExcluded(matched, excludeSet) { return s.inner.Select(ctx, exclude...) } + if s.pinExhausted(matched) { + s.expirePin(pin) + return s.inner.Select(ctx, exclude...) + } // Usability check: must have an access token, and must be either unexpired, // have no expiry, or be expired with a refresh token (transport will refresh). @@ -102,3 +116,24 @@ func (s *PinnedClaudeSelector) Select(ctx context.Context, exclude ...string) (* result := *matched return &result, nil } + +func (s *PinnedClaudeSelector) pinExhausted(acct *keyring.ClaudeOAuth) bool { + if s.quota == nil { + return false + } + snap, ok := s.quota.Snapshot(acctIdentifier(acct)) + return ok && snap.Result.MinRemainingPct() == 0 +} + +func (s *PinnedClaudeSelector) expirePin(pin string) { + var onExpire func(string) + s.mu.Lock() + if s.pin == pin { + s.pin = "" + onExpire = s.onPinExpire + } + s.mu.Unlock() + if onExpire != nil { + onExpire(pin) + } +} diff --git a/internal/proxy/pinned_selector_test.go b/internal/proxy/pinned_selector_test.go index 0aefd06..cf52897 100644 --- a/internal/proxy/pinned_selector_test.go +++ b/internal/proxy/pinned_selector_test.go @@ -8,6 +8,7 @@ import ( "time" "github.com/jacobcxdev/cq/internal/keyring" + "github.com/jacobcxdev/cq/internal/quota" ) // innerSelectorFunc adapts a function to the ClaudeSelector interface. @@ -36,7 +37,7 @@ func makePinnedSelector(accounts []keyring.ClaudeOAuth, pin string) (*PinnedClau discover := ClaudeDiscoverer(func() []keyring.ClaudeOAuth { return accounts }) - sel := NewPinnedClaudeSelector(inner, discover, pin) + sel := NewPinnedClaudeSelector(inner, discover, pin, nil) return sel, &inner } @@ -110,6 +111,44 @@ func TestPinnedClaudeSelector_ExcludedPinDelegatesToInner(t *testing.T) { } } +func TestPinnedClaudeSelector_ExhaustedPinClearsAndDelegatesToInner(t *testing.T) { + future := time.Now().UnixMilli() + 3600_000 + accounts := []keyring.ClaudeOAuth{ + {Email: "pinned@test.com", AccountUUID: "uuid-pin", AccessToken: "tok-pin", ExpiresAt: future}, + {Email: "fallback@test.com", AccountUUID: "uuid-fallback", AccessToken: "tok-fb", ExpiresAt: future}, + } + inner := innerSelectorFunc(func(ctx context.Context, exclude ...string) (*keyring.ClaudeOAuth, error) { + return &keyring.ClaudeOAuth{Email: "fallback@test.com", AccessToken: "tok-fb", ExpiresAt: future}, nil + }) + sel := NewPinnedClaudeSelector(inner, func() []keyring.ClaudeOAuth { return accounts }, "pinned@test.com", stubQuotaReader{ + "uuid-pin": { + Result: quota.Result{ + Status: quota.StatusExhausted, + Windows: map[quota.WindowName]quota.Window{ + "5h": {RemainingPct: 0}, + }, + }, + FetchedAt: time.Now(), + }, + }) + expiredPin := "" + sel.SetPinExpireFunc(func(pin string) { expiredPin = pin }) + + acct, err := sel.Select(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if acct.Email != "fallback@test.com" { + t.Errorf("email = %q, want fallback@test.com", acct.Email) + } + if got := sel.Pin(); got != "" { + t.Errorf("pin = %q, want cleared", got) + } + if expiredPin != "pinned@test.com" { + t.Errorf("expired pin = %q, want pinned@test.com", expiredPin) + } +} + func TestPinnedClaudeSelector_NotFoundReturnsError(t *testing.T) { future := time.Now().UnixMilli() + 3600_000 accounts := []keyring.ClaudeOAuth{