Add Claude proxy account pinning

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 <start|install|uninstall|restart|status>\n")
+		fmt.Fprintf(os.Stderr, "Usage: cq proxy <start|install|uninstall|restart|status|pin>\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 <email-or-uuid>
+	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 | <email-or-account-uuid>]\n")
+	return fmt.Errorf("unexpected arguments")
+}
+
 type proxyCommandOptions struct {
 	Port int
 }
@@ -104,7 +148,15 @@ 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, quotaCache)
+	selector.SetPinExpireFunc(clearPersistedClaudePin)
+	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 +324,60 @@ func runProxyStart(opts proxyCommandOptions) error {
 		Refresher:             proxyRefresher,
 	}
 
-	err = srv.ListenAndServe(context.Background())
+	err = srv.ListenAndServe(proxyCtx)
 	if headroom != nil {
 		headroom.Stop()
 	}
 	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)
+		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 {

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)

internal/proxy/pinned_selector.go

@@ -0,0 +1,139 @@
+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
+	quota       QuotaReader
+	onPinExpire func(string)
+
+	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, 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()
+	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...)
+	}
+	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).
+	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
+}
+
+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)
+	}
+}