From a33e12619ee3dd7b034a952a51feae9cec69759c Mon Sep 17 00:00:00 2001 From: Tam Nhu Tran Date: Thu, 11 Jun 2026 12:09:42 -0400 Subject: [PATCH] feat(cliproxy): managed drain order for account pools - ccs cliproxy accounts order : show effective drain order exactly as the selector resolves it (priority bucket desc, then file-name tie-break, Windows-only case folding to match upstream) - --by-tier derives priorities from tier metadata where present; claude pools state tier-unknown and fall back to file order - --set for manual order; duplicate IDs rejected; priority always >= 1 - dual write path: management-API PATCH when proxy running (avoids MarkResult persist clobber), atomic direct write when stopped - drift warning when stored order no longer matches auth files; usage-attribution stability covered by regression test Part of #1464 account pools (phase 4). --- .../accounts/__tests__/drain-order.test.ts | 593 +++++++++++++++++ src/cliproxy/accounts/drain-order.ts | 366 +++++++++++ src/cliproxy/accounts/registry.ts | 51 +- src/cliproxy/accounts/types.ts | 22 + src/cliproxy/routing/pool-opt-in-prompt.ts | 19 +- src/commands/cliproxy/help-subcommand.ts | 13 + src/commands/cliproxy/index.ts | 16 + src/commands/cliproxy/order-subcommand.ts | 596 ++++++++++++++++++ 8 files changed, 1666 insertions(+), 10 deletions(-) create mode 100644 src/cliproxy/accounts/__tests__/drain-order.test.ts create mode 100644 src/cliproxy/accounts/drain-order.ts create mode 100644 src/commands/cliproxy/order-subcommand.ts diff --git a/src/cliproxy/accounts/__tests__/drain-order.test.ts b/src/cliproxy/accounts/__tests__/drain-order.test.ts new file mode 100644 index 00000000..02a0e013 --- /dev/null +++ b/src/cliproxy/accounts/__tests__/drain-order.test.ts @@ -0,0 +1,593 @@ +/** + * Tests for drain order priority computation, write path, and attribution stability. + */ +import { describe, expect, it } from 'bun:test'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import { runWithScopedCcsHome } from '../../../utils/config-manager'; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +async function withIsolatedHome(fn: (homeDir: string) => Promise | T): Promise { + const homeDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-drain-order-')); + try { + return await runWithScopedCcsHome(homeDir, () => fn(homeDir)); + } finally { + fs.rmSync(homeDir, { recursive: true, force: true }); + } +} + +/** Create a minimal auth JSON file in the auth dir. */ +function writeAuthFile( + authDir: string, + fileName: string, + fields: Record = {} +): void { + fs.mkdirSync(authDir, { recursive: true, mode: 0o700 }); + const content = JSON.stringify({ type: 'antigravity', email: fileName, ...fields }, null, 2); + fs.writeFileSync(path.join(authDir, fileName), content, { mode: 0o600 }); +} + +function readAuthFile(authDir: string, fileName: string): Record { + return JSON.parse(fs.readFileSync(path.join(authDir, fileName), 'utf-8')) as Record< + string, + unknown + >; +} + +// --------------------------------------------------------------------------- +// Import helpers (avoid module-level imports that capture paths at load time) +// --------------------------------------------------------------------------- + +async function loadDrainOrder() { + return import(`../drain-order?drain-order=${Date.now()}`); +} + +async function loadRegistry() { + return import(`../registry?drain-order-registry=${Date.now()}`); +} + +async function loadStatsTransformer() { + return import(`../../../cliproxy/services/stats-fetcher?drain-order-stats-fetcher=${Date.now()}`); +} + +async function loadOrderSubcommand() { + return import( + `../../../commands/cliproxy/order-subcommand?drain-order-order-subcommand=${Date.now()}` + ); +} + +// --------------------------------------------------------------------------- +// Priority computation tests +// --------------------------------------------------------------------------- + +describe('rankToPriority', () => { + it('maps rank 0 (highest) to total priority', async () => { + const { rankToPriority } = await loadDrainOrder(); + expect(rankToPriority(0, 3)).toBe(3); + }); + + it('maps last rank to MIN_PRIORITY (1)', async () => { + const { rankToPriority, MIN_PRIORITY } = await loadDrainOrder(); + expect(rankToPriority(2, 3)).toBe(MIN_PRIORITY); + expect(rankToPriority(2, 3)).toBeGreaterThanOrEqual(1); + }); + + it('never returns 0', async () => { + const { rankToPriority } = await loadDrainOrder(); + for (let total = 1; total <= 5; total++) { + for (let rank = 0; rank < total; rank++) { + expect(rankToPriority(rank, total)).toBeGreaterThanOrEqual(1); + } + } + }); + + it('single account gets priority 1', async () => { + const { rankToPriority } = await loadDrainOrder(); + expect(rankToPriority(0, 1)).toBe(1); + }); +}); + +describe('computeManualDrainOrder', () => { + it('assigns descending priorities to specified accounts', async () => { + const { computeManualDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'a@x.com', tokenFile: 'a.json' }, + { accountId: 'b@x.com', tokenFile: 'b.json' }, + { accountId: 'c@x.com', tokenFile: 'c.json' }, + ]; + const entries = computeManualDrainOrder(['a@x.com', 'b@x.com', 'c@x.com'], accounts); + const byId = new Map( + entries.map((e: { accountId: string; priority: number }) => [e.accountId, e.priority]) + ); + expect(byId.get('a@x.com')).toBeGreaterThan(byId.get('b@x.com')); + expect(byId.get('b@x.com')).toBeGreaterThan(byId.get('c@x.com')); + expect(byId.get('c@x.com')).toBeGreaterThanOrEqual(1); + }); + + it('throws for unknown account ID', async () => { + const { computeManualDrainOrder } = await loadDrainOrder(); + const accounts = [{ accountId: 'a@x.com', tokenFile: 'a.json' }]; + expect(() => computeManualDrainOrder(['z@x.com'], accounts)).toThrow(); + }); + + it('throws for duplicate account ID in the order list', async () => { + const { computeManualDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'a@x.com', tokenFile: 'a.json' }, + { accountId: 'b@x.com', tokenFile: 'b.json' }, + ]; + expect(() => computeManualDrainOrder(['a@x.com', 'b@x.com', 'a@x.com'], accounts)).toThrow( + /[Dd]uplicate/ + ); + }); + + it('unspecified accounts get MIN_PRIORITY', async () => { + const { computeManualDrainOrder, MIN_PRIORITY } = await loadDrainOrder(); + const accounts = [ + { accountId: 'a@x.com', tokenFile: 'a.json' }, + { accountId: 'b@x.com', tokenFile: 'b.json' }, + ]; + const entries = computeManualDrainOrder(['a@x.com'], accounts); + const bEntry = entries.find((e: { accountId: string }) => e.accountId === 'b@x.com'); + expect(bEntry?.priority).toBe(MIN_PRIORITY); + }); + + it('is idempotent - same inputs produce same priorities', async () => { + const { computeManualDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'a@x.com', tokenFile: 'a.json' }, + { accountId: 'b@x.com', tokenFile: 'b.json' }, + ]; + const entries1 = computeManualDrainOrder(['a@x.com', 'b@x.com'], accounts); + const entries2 = computeManualDrainOrder(['a@x.com', 'b@x.com'], accounts); + expect(JSON.stringify(entries1)).toBe(JSON.stringify(entries2)); + }); + + it('partial --set: every specified account has strictly higher priority than every unspecified account', async () => { + // Regression for the tie at priority 1 when last specified == MIN_PRIORITY == unspecified. + // With a pool of 5, specifying 2 should give the last specified priority > 1. + const { computeManualDrainOrder, MIN_PRIORITY } = await loadDrainOrder(); + const accounts = [ + { accountId: 'a@x.com', tokenFile: 'a.json' }, + { accountId: 'b@x.com', tokenFile: 'b.json' }, + { accountId: 'c@x.com', tokenFile: 'c.json' }, + { accountId: 'd@x.com', tokenFile: 'd.json' }, + { accountId: 'e@x.com', tokenFile: 'e.json' }, + ]; + const entries = computeManualDrainOrder(['a@x.com', 'b@x.com'], accounts); + const byId = new Map( + entries.map((e: { accountId: string; priority: number }) => [e.accountId, e.priority]) + ); + const specifiedPriorities = ['a@x.com', 'b@x.com'].map((id) => byId.get(id) as number); + const unspecifiedPriorities = ['c@x.com', 'd@x.com', 'e@x.com'].map( + (id) => byId.get(id) as number + ); + const minSpecified = Math.min(...specifiedPriorities); + const maxUnspecified = Math.max(...unspecifiedPriorities); + expect(minSpecified).toBeGreaterThan(maxUnspecified); + // Unspecified accounts still get MIN_PRIORITY + for (const p of unspecifiedPriorities) { + expect(p).toBe(MIN_PRIORITY); + } + }); +}); + +describe('computeTierDrainOrder', () => { + it('ultra > pro > free priority ordering', async () => { + const { computeTierDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'free@x.com', tokenFile: 'free.json', tier: 'free' as const }, + { accountId: 'ultra@x.com', tokenFile: 'ultra.json', tier: 'ultra' as const }, + { accountId: 'pro@x.com', tokenFile: 'pro.json', tier: 'pro' as const }, + ]; + const entries = computeTierDrainOrder(accounts); + const byId = new Map( + entries.map((e: { accountId: string; priority: number }) => [e.accountId, e.priority]) + ); + expect(byId.get('ultra@x.com')).toBeGreaterThan(byId.get('pro@x.com')); + expect(byId.get('pro@x.com')).toBeGreaterThan(byId.get('free@x.com')); + expect(byId.get('free@x.com')).toBeGreaterThanOrEqual(1); + }); + + it('unknown tier gets lowest priority (MIN_PRIORITY)', async () => { + const { computeTierDrainOrder, MIN_PRIORITY } = await loadDrainOrder(); + const accounts = [ + { accountId: 'ultra@x.com', tokenFile: 'ultra.json', tier: 'ultra' as const }, + { accountId: 'unknown@x.com', tokenFile: 'unknown.json', tier: 'unknown' as const }, + ]; + const entries = computeTierDrainOrder(accounts); + const unknownEntry = entries.find( + (e: { accountId: string }) => e.accountId === 'unknown@x.com' + ); + expect(unknownEntry?.priority).toBe(MIN_PRIORITY); + }); + + it('accounts with same tier get equal priority', async () => { + const { computeTierDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'a@x.com', tokenFile: 'a.json', tier: 'pro' as const }, + { accountId: 'b@x.com', tokenFile: 'b.json', tier: 'pro' as const }, + ]; + const entries = computeTierDrainOrder(accounts); + const [a, b] = entries as Array<{ accountId: string; priority: number }>; + expect(a.priority).toBe(b.priority); + }); + + it('tie-break by tokenFile (plain byte order, not locale) matches selector contract', async () => { + // The upstream Go selector breaks priority ties by Auth.ID asc (lowercased file path). + // For the agy prefix-named fleet, file names differ from emails, so verify we sort + // by tokenFile not accountId. + const { computeTierDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'z@x.com', tokenFile: 'antigravity-aaa.json', tier: 'pro' as const }, + { accountId: 'a@x.com', tokenFile: 'antigravity-zzz.json', tier: 'pro' as const }, + ]; + // Both same tier -> same priority; sort by tokenFile should put 'aaa' before 'zzz' + const entries = computeTierDrainOrder(accounts) as Array<{ + accountId: string; + tokenFile: string; + priority: number; + }>; + // Confirm equal priorities + expect(entries[0].priority).toBe(entries[1].priority); + // Sort as the display layer does: priority desc, tokenFile lowercase asc + const sorted = entries + .slice() + .sort( + (a, b) => + b.priority - a.priority || + (a.tokenFile.toLowerCase() < b.tokenFile.toLowerCase() ? -1 : 1) + ); + // antigravity-aaa.json sorts before antigravity-zzz.json (byte/file order) + expect(sorted[0].tokenFile).toBe('antigravity-aaa.json'); + expect(sorted[1].tokenFile).toBe('antigravity-zzz.json'); + // accountId order is the reverse, confirming file order != email order for this fleet + expect(sorted[0].accountId).toBe('z@x.com'); + expect(sorted[1].accountId).toBe('a@x.com'); + }); + + it('tierDerived is true only for accounts with a real tier rank', async () => { + const { computeTierDrainOrder } = await loadDrainOrder(); + const accounts = [ + { accountId: 'ultra@x.com', tokenFile: 'ultra.json', tier: 'ultra' as const }, + { accountId: 'unknown@x.com', tokenFile: 'unknown.json', tier: 'unknown' as const }, + ]; + const entries = computeTierDrainOrder(accounts) as Array<{ + accountId: string; + tierDerived: boolean; + }>; + const ultraEntry = entries.find((e) => e.accountId === 'ultra@x.com'); + const unknownEntry = entries.find((e) => e.accountId === 'unknown@x.com'); + expect(ultraEntry?.tierDerived).toBe(true); + expect(unknownEntry?.tierDerived).toBe(false); + }); +}); + +// --------------------------------------------------------------------------- +// Tie-break key (platform-conditional, mirrors upstream synthesizer) +// --------------------------------------------------------------------------- + +describe('tieBreakKey', () => { + it('preserves case on non-Windows (byte order, upper before lower)', async () => { + const { tieBreakKey } = await loadOrderSubcommand(); + // On non-Windows the selector compares raw Auth.ID bytes: 'B' (0x42) < 'a' (0x61). + const files = ['antigravity-aaa.json', 'antigravity-Bbb.json']; + const sorted = files + .slice() + .sort((a, b) => (tieBreakKey(a, 'linux') < tieBreakKey(b, 'linux') ? -1 : 1)); + expect(sorted).toEqual(['antigravity-Bbb.json', 'antigravity-aaa.json']); + }); + + it('lowercases on Windows so mixed case sorts case-insensitively', async () => { + const { tieBreakKey } = await loadOrderSubcommand(); + const files = ['antigravity-aaa.json', 'antigravity-Bbb.json']; + const sorted = files + .slice() + .sort((a, b) => (tieBreakKey(a, 'win32') < tieBreakKey(b, 'win32') ? -1 : 1)); + expect(sorted).toEqual(['antigravity-aaa.json', 'antigravity-Bbb.json']); + }); +}); + +// --------------------------------------------------------------------------- +// Direct file write tests +// --------------------------------------------------------------------------- + +describe('writeAuthFilePriorityDirect', () => { + it('writes priority field to auth JSON', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json', { email: 'test@x.com' }); + + const { writeAuthFilePriorityDirect } = await loadDrainOrder(); + writeAuthFilePriorityDirect('test.json', 3); + + const data = readAuthFile(authDir, 'test.json'); + expect(data.priority).toBe(3); + }); + }); + + it('preserves all other fields', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json', { email: 'test@x.com', token: 'abc123' }); + + const { writeAuthFilePriorityDirect } = await loadDrainOrder(); + writeAuthFilePriorityDirect('test.json', 2); + + const data = readAuthFile(authDir, 'test.json'); + expect(data.email).toBe('test@x.com'); + expect(data.token).toBe('abc123'); + expect(data.priority).toBe(2); + }); + }); + + it('rejects priority 0 (management layer treats 0 as delete)', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json'); + + const { writeAuthFilePriorityDirect } = await loadDrainOrder(); + expect(() => writeAuthFilePriorityDirect('test.json', 0)).toThrow(); + }); + }); + + it('rejects negative priority', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json'); + + const { writeAuthFilePriorityDirect } = await loadDrainOrder(); + expect(() => writeAuthFilePriorityDirect('test.json', -1)).toThrow(); + }); + }); + + it('throws when file does not exist', async () => { + await withIsolatedHome(async (_homeDir) => { + const { writeAuthFilePriorityDirect } = await loadDrainOrder(); + expect(() => writeAuthFilePriorityDirect('nonexistent.json', 1)).toThrow(); + }); + }); +}); + +describe('readAuthFilePriority', () => { + it('returns the priority from a file that has one', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json', { priority: 5 }); + + const { readAuthFilePriority } = await loadDrainOrder(); + expect(readAuthFilePriority('test.json')).toBe(5); + }); + }); + + it('returns undefined for file with no priority', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json', { email: 'x@y.com' }); + + const { readAuthFilePriority } = await loadDrainOrder(); + expect(readAuthFilePriority('test.json')).toBeUndefined(); + }); + }); + + it('returns undefined for missing file', async () => { + await withIsolatedHome(async (_homeDir) => { + const { readAuthFilePriority } = await loadDrainOrder(); + expect(readAuthFilePriority('ghost.json')).toBeUndefined(); + }); + }); + + it('returns undefined for priority 0 (invalid sentinel)', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'test.json', { priority: 0 }); + + const { readAuthFilePriority } = await loadDrainOrder(); + expect(readAuthFilePriority('test.json')).toBeUndefined(); + }); + }); +}); + +// --------------------------------------------------------------------------- +// applyDrainOrder (direct write path, proxy stopped) +// --------------------------------------------------------------------------- + +describe('applyDrainOrder - direct write path (proxy stopped)', () => { + it('writes priorities to all files when proxy stopped', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'a.json', { email: 'a@x.com' }); + writeAuthFile(authDir, 'b.json', { email: 'b@x.com' }); + + const { applyDrainOrder } = await loadDrainOrder(); + const entries = [ + { + accountId: 'a@x.com', + tokenFile: 'a.json', + priority: 2, + tierDerived: false, + currentPriority: undefined, + }, + { + accountId: 'b@x.com', + tokenFile: 'b.json', + priority: 1, + tierDerived: false, + currentPriority: undefined, + }, + ]; + const result = await applyDrainOrder(entries, false); + + expect(result.usedManagementApi).toBe(false); + expect(result.written).toHaveLength(2); + expect(result.failed).toHaveLength(0); + + const a = readAuthFile(authDir, 'a.json'); + const b = readAuthFile(authDir, 'b.json'); + expect(a.priority).toBe(2); + expect(b.priority).toBe(1); + }); + }); + + it('skips entries where priority already matches (idempotency)', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'a.json', { email: 'a@x.com', priority: 3 }); + + const { applyDrainOrder } = await loadDrainOrder(); + const entries = [ + { + accountId: 'a@x.com', + tokenFile: 'a.json', + priority: 3, + tierDerived: false, + currentPriority: undefined, + }, + ]; + const result = await applyDrainOrder(entries, false); + + expect(result.skipped).toHaveLength(1); + expect(result.written).toHaveLength(0); + }); + }); + + it('records failures for missing files', async () => { + await withIsolatedHome(async (_homeDir) => { + const { applyDrainOrder } = await loadDrainOrder(); + const entries = [ + { + accountId: 'ghost@x.com', + tokenFile: 'ghost.json', + priority: 2, + tierDerived: false, + currentPriority: undefined, + }, + ]; + const result = await applyDrainOrder(entries, false); + expect(result.failed).toHaveLength(1); + expect(result.written).toHaveLength(0); + }); + }); +}); + +// --------------------------------------------------------------------------- +// Registry drain order persistence +// --------------------------------------------------------------------------- + +describe('registry drain order persistence', () => { + it('saveDrainOrderConfig persists and loadDrainOrderConfig retrieves', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'antigravity-a.json', { email: 'a@x.com', type: 'antigravity' }); + + const { registerAccount } = await loadRegistry(); + registerAccount('agy', 'antigravity-a.json', 'a@x.com'); + + const { saveDrainOrderConfig, loadDrainOrderConfig } = await loadRegistry(); + const config = { mode: 'manual' as const, orderedIds: ['a@x.com'] }; + const saved = saveDrainOrderConfig('agy', config); + expect(saved).toBe(true); + + const loaded = loadDrainOrderConfig('agy'); + expect(loaded?.mode).toBe('manual'); + expect(loaded?.orderedIds).toEqual(['a@x.com']); + }); + }); + + it('clearDrainOrderConfig removes persisted config', async () => { + await withIsolatedHome(async (homeDir) => { + const authDir = path.join(homeDir, '.ccs', 'cliproxy', 'auth'); + writeAuthFile(authDir, 'antigravity-a.json', { email: 'a@x.com', type: 'antigravity' }); + + const { registerAccount } = await loadRegistry(); + registerAccount('agy', 'antigravity-a.json', 'a@x.com'); + + const { saveDrainOrderConfig, loadDrainOrderConfig, clearDrainOrderConfig } = + await loadRegistry(); + saveDrainOrderConfig('agy', { mode: 'tier' }); + const cleared = clearDrainOrderConfig('agy'); + expect(cleared).toBe(true); + + expect(loadDrainOrderConfig('agy')).toBeUndefined(); + }); + }); + + it('returns false for saveDrainOrderConfig when provider not in registry', async () => { + await withIsolatedHome(async (_homeDir) => { + const { saveDrainOrderConfig } = await loadRegistry(); + // No accounts registered for 'gemini' -> returns false + const result = saveDrainOrderConfig('gemini', { mode: 'tier' }); + expect(result).toBe(false); + }); + }); +}); + +// --------------------------------------------------------------------------- +// Attribution stability: auth_index ordering survives priority rewrite +// --------------------------------------------------------------------------- + +describe('attribution stability - auth_index unaffected by priority rewrite', () => { + it('buildAuthIndexToAccountMap produces stable output regardless of priority values', async () => { + // Simulate the auth files list returned by /v0/management/auth-files with various priorities. + // auth_index is assigned by CLIProxy at load time based on file discovery order, + // NOT by priority. So priority rewrites must not change auth_index values. + const { buildAuthIndexToAccountMap } = await loadStatsTransformer(); + + const authFilesBeforeRewrite = [ + { auth_index: 0, email: 'a@x.com', provider: 'antigravity' }, + { auth_index: 1, email: 'b@x.com', provider: 'antigravity' }, + { auth_index: 2, email: 'c@x.com', provider: 'antigravity' }, + ]; + + const authFilesAfterRewrite = [ + // Priority changed but auth_index values are assigned by CLIProxy independently + { auth_index: 0, email: 'a@x.com', provider: 'antigravity', priority: 3 }, + { auth_index: 1, email: 'b@x.com', provider: 'antigravity', priority: 2 }, + { auth_index: 2, email: 'c@x.com', provider: 'antigravity', priority: 1 }, + ]; + + const mapBefore = buildAuthIndexToAccountMap(authFilesBeforeRewrite); + const mapAfter = buildAuthIndexToAccountMap(authFilesAfterRewrite); + + // auth_index -> email mapping must be identical after priority rewrite + expect(mapBefore.get('0')).toBe('a@x.com'); + expect(mapBefore.get('1')).toBe('b@x.com'); + expect(mapBefore.get('2')).toBe('c@x.com'); + + expect(mapAfter.get('0')).toBe(mapBefore.get('0')); + expect(mapAfter.get('1')).toBe(mapBefore.get('1')); + expect(mapAfter.get('2')).toBe(mapBefore.get('2')); + }); + + it('buildAuthIndexToAccountMap skips entries without auth_index', async () => { + const { buildAuthIndexToAccountMap } = await loadStatsTransformer(); + + const authFiles = [ + { email: 'no-index@x.com', provider: 'antigravity' }, // no auth_index + { auth_index: 5, email: 'has-index@x.com', provider: 'antigravity' }, + ]; + + const map = buildAuthIndexToAccountMap(authFiles); + expect(map.size).toBe(1); + expect(map.get('5')).toBe('has-index@x.com'); + expect(map.has('undefined')).toBe(false); + }); + + it('buildAuthIndexToAccountMap handles both numeric and string auth_index keys', async () => { + const { buildAuthIndexToAccountMap } = await loadStatsTransformer(); + + const authFiles = [ + { auth_index: 0, email: 'numeric@x.com', provider: 'antigravity' }, + { auth_index: '1', email: 'string@x.com', provider: 'antigravity' }, + ]; + + const map = buildAuthIndexToAccountMap(authFiles); + // Both stored as String(auth_index) + expect(map.get('0')).toBe('numeric@x.com'); + expect(map.get('1')).toBe('string@x.com'); + }); +}); diff --git a/src/cliproxy/accounts/drain-order.ts b/src/cliproxy/accounts/drain-order.ts new file mode 100644 index 00000000..43a63183 --- /dev/null +++ b/src/cliproxy/accounts/drain-order.ts @@ -0,0 +1,366 @@ +/** + * Drain Order Management + * + * Manages CLIProxy account drain order via top-level "priority" field in auth JSONs. + * + * Architecture: + * - Priority >= 1 always (management layer treats 0 as delete-the-attribute) + * - Write path is dual: management API when proxy running, direct file when stopped + * - Tier-aware defaults only where AccountTier metadata exists (agy/gemini) + * - Claude pools: tier is unknown -> stable file order + manual --set required + * - Selector drain order: priority bucket desc, then Auth.ID asc + */ + +import * as fs from 'fs'; +import * as path from 'path'; +import { getAuthDir } from '../config/config-generator'; +import { + getProxyTarget, + buildProxyUrl, + buildManagementHeaders, +} from '../proxy/proxy-target-resolver'; +import type { AccountTier } from './types'; + +/** Minimum valid priority value. Management layer treats 0 as delete. */ +export const MIN_PRIORITY = 1; + +/** Management API path for patching auth file fields */ +const AUTH_FILES_FIELDS_PATH = '/v0/management/auth-files/fields'; + +/** Timeout for management API calls in ms */ +const MGMT_TIMEOUT_MS = 5000; + +/** + * Tier rank for priority derivation. + * Higher rank = higher priority = drained first. + * ultra > pro > free, unknown = no rank (returns undefined). + */ +const TIER_RANK: Record, number> = { + ultra: 3, + pro: 2, + free: 1, +}; + +/** + * Compute 1-based priority from a 0-based rank among N accounts. + * Rank 0 = highest priority = assigned priority N. + * Rank N-1 = lowest priority = assigned priority 1. + * + * Never returns 0 (management layer treats 0 as delete). + */ +export function rankToPriority(rank: number, total: number): number { + const raw = total - rank; + return Math.max(MIN_PRIORITY, raw); +} + +/** Tier rank for a known tier, undefined for unknown. */ +export function tierRank(tier: AccountTier | undefined): number | undefined { + if (!tier || tier === 'unknown') return undefined; + return TIER_RANK[tier]; +} + +/** + * Result of drain order computation for a single account. + */ +export interface DrainOrderEntry { + /** Account ID */ + accountId: string; + /** Token file name (without directory) */ + tokenFile: string; + /** Computed priority (>= 1, higher = drained earlier) */ + priority: number; + /** Tier if known */ + tier?: AccountTier; + /** Whether tier was used to derive this priority */ + tierDerived: boolean; + /** Current priority read from auth file (undefined = not yet set) */ + currentPriority: number | undefined; +} + +/** + * Input record for priority computation. + */ +export interface DrainOrderInput { + accountId: string; + tokenFile: string; + tier?: AccountTier; +} + +/** + * Compute drain order priorities for a list of accounts using manual ordering. + * + * @param orderedIds Account IDs in desired drain order (first = highest priority) + * @param allAccounts All accounts for the provider (some may not be in orderedIds) + * @returns DrainOrderEntry array ordered by resulting priority desc + */ +export function computeManualDrainOrder( + orderedIds: string[], + allAccounts: DrainOrderInput[] +): DrainOrderEntry[] { + const accountMap = new Map(allAccounts.map((a) => [a.accountId, a])); + const result: DrainOrderEntry[] = []; + + // Reject duplicate IDs: a repeated account in --set is ambiguous and would + // otherwise be assigned two different priorities. + const seen = new Set(); + for (const id of orderedIds) { + if (seen.has(id)) { + throw new Error(`Duplicate account ID in order: ${id}`); + } + seen.add(id); + } + + // Validate that all specified IDs exist + for (const id of orderedIds) { + if (!accountMap.has(id)) { + throw new Error(`Account ID not found: ${id}`); + } + } + + const specifiedSet = new Set(orderedIds); + // Use total + 1 so rank (N-1) maps to priority 2, keeping every specified + // account strictly above the MIN_PRIORITY floor assigned to unspecified ones. + const total = orderedIds.length + 1; + + // Assign priorities to specified accounts (first = highest priority) + for (let rank = 0; rank < orderedIds.length; rank++) { + const id = orderedIds[rank]; + const account = accountMap.get(id); + if (!account) continue; + + result.push({ + accountId: id, + tokenFile: account.tokenFile, + priority: rankToPriority(rank, total), + tier: account.tier, + tierDerived: false, + currentPriority: undefined, + }); + } + + // Remaining accounts not in the specified list get priority 1 (lowest) + for (const account of allAccounts) { + if (!specifiedSet.has(account.accountId)) { + result.push({ + accountId: account.accountId, + tokenFile: account.tokenFile, + priority: MIN_PRIORITY, + tier: account.tier, + tierDerived: false, + currentPriority: undefined, + }); + } + } + + return result; +} + +/** + * Compute drain order priorities for accounts using tier metadata. + * Only valid where tier metadata exists (agy/gemini providers). + * + * Accounts with unknown tier are sorted last (priority MIN_PRIORITY). + * Accounts within the same tier bucket are given equal priorities. + * + * @param accounts Accounts with optional tier metadata + * @returns DrainOrderEntry array + */ +export function computeTierDrainOrder(accounts: DrainOrderInput[]): DrainOrderEntry[] { + // Group by tier rank + const withRank: Array<{ input: DrainOrderInput; rank: number }> = accounts.map((a) => ({ + input: a, + rank: tierRank(a.tier) ?? 0, + })); + + // Get unique ranks descending + const uniqueRanks = [...new Set(withRank.map((x) => x.rank))].sort((a, b) => b - a); + const numRanks = uniqueRanks.length; + + // Assign priority buckets: highest rank -> highest priority (numRanks), lowest -> MIN_PRIORITY + const rankToPriorityMap = new Map(); + uniqueRanks.forEach((rank, idx) => { + // idx 0 = highest rank, gets highest priority + const priority = Math.max(MIN_PRIORITY, numRanks - idx); + rankToPriorityMap.set(rank, priority); + }); + + return withRank.map(({ input, rank }) => ({ + accountId: input.accountId, + tokenFile: input.tokenFile, + priority: rankToPriorityMap.get(rank) ?? MIN_PRIORITY, + tier: input.tier, + tierDerived: rank > 0, // only true when tier contributed a real rank + currentPriority: undefined, + })); +} + +/** + * Read the current priority from an auth JSON file. + * Returns undefined if file missing, unreadable, or has no priority field. + */ +export function readAuthFilePriority(tokenFile: string): number | undefined { + try { + const authDir = getAuthDir(); + const filePath = path.join(authDir, tokenFile); + const content = fs.readFileSync(filePath, 'utf-8'); + const data = JSON.parse(content) as { priority?: unknown }; + if (typeof data.priority === 'number' && Number.isFinite(data.priority) && data.priority >= 1) { + return data.priority; + } + return undefined; + } catch { + return undefined; + } +} + +/** + * Write priority directly to an auth JSON file (proxy must be stopped). + * Validates priority >= MIN_PRIORITY before writing. + * Preserves all other fields; uses atomic temp-rename. + */ +export function writeAuthFilePriorityDirect(tokenFile: string, priority: number): void { + if (priority < MIN_PRIORITY) { + throw new Error(`Priority must be >= ${MIN_PRIORITY}, got ${priority}`); + } + + const authDir = getAuthDir(); + const filePath = path.join(authDir, tokenFile); + + if (!fs.existsSync(filePath)) { + throw new Error(`Auth file not found: ${tokenFile}`); + } + + const content = fs.readFileSync(filePath, 'utf-8'); + const data = JSON.parse(content) as Record; + + data.priority = priority; + + const tempPath = `${filePath}.tmp.${process.pid}`; + fs.writeFileSync(tempPath, JSON.stringify(data, null, 2) + '\n', { mode: 0o600 }); + fs.renameSync(tempPath, filePath); +} + +/** + * Write priority via CLIProxy management API (proxy must be running). + * Uses PATCH /v0/management/auth-files/fields. + * + * @param tokenFile File name (basename, e.g. "antigravity-foo.json") + * @param priority Priority value >= MIN_PRIORITY + * @returns true on success, false on failure + */ +export async function writeAuthFilePriorityViaApi( + tokenFile: string, + priority: number +): Promise { + if (priority < MIN_PRIORITY) { + throw new Error(`Priority must be >= ${MIN_PRIORITY}, got ${priority}`); + } + + const target = getProxyTarget(); + const url = buildProxyUrl(target, AUTH_FILES_FIELDS_PATH); + const headers = buildManagementHeaders(target, { 'Content-Type': 'application/json' }); + + const body = JSON.stringify({ name: tokenFile, priority }); + + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), MGMT_TIMEOUT_MS); + + try { + const response = await fetch(url, { + method: 'PATCH', + headers, + body, + signal: controller.signal, + }); + + return response.ok; + } catch { + return false; + } finally { + clearTimeout(timeoutId); + } +} + +/** + * Result of a drain order apply operation. + */ +export interface ApplyDrainOrderResult { + /** Entries that were written successfully */ + written: DrainOrderEntry[]; + /** Entries that were skipped (priority unchanged) */ + skipped: DrainOrderEntry[]; + /** Entries that failed to write */ + failed: Array<{ entry: DrainOrderEntry; reason: string }>; + /** Whether the proxy was running during the write (management API path) */ + usedManagementApi: boolean; +} + +/** + * Apply drain order priorities to auth files. + * + * Write path selection: + * - Proxy running (HTTP health check passes): use management API + * - Proxy stopped: direct file write + * + * Idempotent: skips entries where priority already equals the target value. + * + * @param entries Drain order entries with target priorities + * @param proxyRunning Whether the proxy is currently running + * @returns Result summary + */ +export async function applyDrainOrder( + entries: DrainOrderEntry[], + proxyRunning: boolean +): Promise { + const result: ApplyDrainOrderResult = { + written: [], + skipped: [], + failed: [], + usedManagementApi: proxyRunning, + }; + + for (const entry of entries) { + // Idempotency: read current value and skip if already set + const currentPriority = readAuthFilePriority(entry.tokenFile); + if (currentPriority === entry.priority) { + result.skipped.push({ ...entry, currentPriority }); + continue; + } + + if (proxyRunning) { + const ok = await writeAuthFilePriorityViaApi(entry.tokenFile, entry.priority); + if (ok) { + result.written.push({ ...entry, currentPriority }); + } else { + result.failed.push({ + entry: { ...entry, currentPriority }, + reason: 'management API PATCH failed', + }); + } + } else { + try { + writeAuthFilePriorityDirect(entry.tokenFile, entry.priority); + result.written.push({ ...entry, currentPriority }); + } catch (err) { + result.failed.push({ + entry: { ...entry, currentPriority }, + reason: (err as Error).message, + }); + } + } + } + + return result; +} + +/** + * Read current drain order priorities for a set of accounts. + * Annotates each DrainOrderEntry with its current priority from disk. + * + * @param entries Entries to annotate (modifies currentPriority in-place) + */ +export function annotateCurrentPriorities(entries: DrainOrderEntry[]): void { + for (const entry of entries) { + entry.currentPriority = readAuthFilePriority(entry.tokenFile); + } +} diff --git a/src/cliproxy/accounts/registry.ts b/src/cliproxy/accounts/registry.ts index b666c1e9..06ca3dff 100644 --- a/src/cliproxy/accounts/registry.ts +++ b/src/cliproxy/accounts/registry.ts @@ -10,7 +10,7 @@ import { CLIProxyProvider } from '../types'; import { PROVIDER_CAPABILITIES } from '../provider-capabilities'; import { PROVIDER_TYPE_VALUES } from '../auth/auth-types'; import { getAuthDir, getCliproxyDir } from '../config/config-generator'; -import { AccountsRegistry, AccountInfo, PROVIDERS_WITHOUT_EMAIL } from './types'; +import { AccountsRegistry, AccountInfo, PROVIDERS_WITHOUT_EMAIL, DrainOrderConfig } from './types'; import { getAccountsRegistryPath, getPausedDir, @@ -269,6 +269,10 @@ function populateRegistryFromTokenFiles( lastUsedAt: hydratedAccount?.lastUsedAt || (token.stats.mtime || token.stats.birthtime || new Date()).toISOString(), + // Preserve tier from the hydrated registry entry. Without this, populate + // would silently strip tier metadata written by the quota command, breaking + // --by-tier ordering for file-copied fleets that are re-hydrated on read. + tier: hydratedAccount?.tier, }; if (token.paused) { @@ -820,3 +824,48 @@ export function discoverExistingAccounts(): void { populateRegistryFromTokenFiles(registry); }); } + +/** + * Persist drain order configuration for a provider. + * The config survives re-auth and registry sync; priorities are NOT automatically + * re-applied after sync. Re-run `ccs cliproxy accounts order --by-tier` + * or `--set` after adding or re-authing accounts to re-apply. + */ +export function saveDrainOrderConfig( + provider: CLIProxyProvider, + config: DrainOrderConfig +): boolean { + return mutateAccountsRegistry((registry) => { + const providerAccounts = registry.providers[provider]; + if (!providerAccounts) { + return false; + } + providerAccounts.drainOrder = config; + return true; + }); +} + +/** + * Load persisted drain order configuration for a provider. + * Returns undefined if none stored (implies file order). + */ +export function loadDrainOrderConfig(provider: CLIProxyProvider): DrainOrderConfig | undefined { + return withAccountsRegistryLock(() => { + const registry = readAccountsRegistryFromDisk(); + return registry.providers[provider]?.drainOrder; + }); +} + +/** + * Clear drain order configuration for a provider (revert to file order). + */ +export function clearDrainOrderConfig(provider: CLIProxyProvider): boolean { + return mutateAccountsRegistry((registry) => { + const providerAccounts = registry.providers[provider]; + if (!providerAccounts) { + return false; + } + delete providerAccounts.drainOrder; + return true; + }); +} diff --git a/src/cliproxy/accounts/types.ts b/src/cliproxy/accounts/types.ts index 3d098f54..f64b2cd6 100644 --- a/src/cliproxy/accounts/types.ts +++ b/src/cliproxy/accounts/types.ts @@ -41,12 +41,34 @@ export interface AccountInfo { projectId?: string; } +/** + * Drain order mode for a provider's account pool. + * - "manual": user-specified ordered list of account IDs + * - "tier": auto-derived from AccountTier (ultra > pro > free); unknown = last + * - "file": stable file-system order (default, no priority writes) + */ +export type DrainOrderMode = 'manual' | 'tier' | 'file'; + +/** Persisted drain order configuration for a provider */ +export interface DrainOrderConfig { + /** How priorities were last set */ + mode: DrainOrderMode; + /** + * Ordered account IDs for manual mode. + * First = drained first (highest priority). + * Stale IDs (accounts removed since last --set) are silently ignored on re-apply. + */ + orderedIds?: string[]; +} + /** Provider accounts configuration */ export interface ProviderAccounts { /** Default account ID for this provider */ default: string; /** Map of account ID to account metadata */ accounts: Record>; + /** Persisted drain order configuration (optional; absent = file order) */ + drainOrder?: DrainOrderConfig; } /** Accounts registry structure */ diff --git a/src/cliproxy/routing/pool-opt-in-prompt.ts b/src/cliproxy/routing/pool-opt-in-prompt.ts index 70a2a361..e4181a3a 100644 --- a/src/cliproxy/routing/pool-opt-in-prompt.ts +++ b/src/cliproxy/routing/pool-opt-in-prompt.ts @@ -155,28 +155,29 @@ export async function maybeOfferPoolRouting( return { prompted: false, enabled: false, skipped: true, skipReason: 'dismissed' }; } - // Remote / Docker target: print hint, do not prompt + // Remote / Docker target: print hint, do not prompt. + // Check before account count so remote targets with locally-unknown registries + // still get the hint (the remote CLIProxy holds the authoritative account list). const target = getProxyTarget(); if (target.isRemote) { printRemoteHint(provider); return { prompted: false, enabled: false, skipped: true, skipReason: 'remote-target' }; } - // Non-TTY (piped input / CI): skip silently - if (!process.stdin.isTTY || !process.stderr.isTTY) { - return { prompted: false, enabled: false, skipped: true, skipReason: 'non-tty' }; - } - // Verify the actual post-add count is >= 2. registerAccount deduplicates by // email/token-file so re-authenticating the single existing account keeps the - // count at 1 — not a real 1->2 transition. Checking here (after TTY and remote - // guards) ensures the re-auth path correctly falls back to not-at-transition - // without triggering the prompt for a single-account user. + // count at 1 — not a real 1->2 transition. Checking here (before TTY) ensures + // re-auth dedup always returns not-at-transition even in non-TTY/CI sessions. const accountCountAfter = getProviderAccounts(provider).length; if (accountCountAfter < 2) { return { prompted: false, enabled: false, skipped: true, skipReason: 'not-at-transition' }; } + // Non-TTY (piped input / CI): skip silently + if (!process.stdin.isTTY || !process.stderr.isTTY) { + return { prompted: false, enabled: false, skipped: true, skipReason: 'non-tty' }; + } + // Gather all providers with 2+ accounts for disclosure const multiAccountProviders = getMultiAccountProviders(); const providerList = diff --git a/src/commands/cliproxy/help-subcommand.ts b/src/commands/cliproxy/help-subcommand.ts index 1bc7ab10..4fcb7e6f 100644 --- a/src/commands/cliproxy/help-subcommand.ts +++ b/src/commands/cliproxy/help-subcommand.ts @@ -66,6 +66,19 @@ export async function showHelp(): Promise { ['pool', 'Show pool routing status (fill-first + affinity + 429 cooldown)'], ['pool --enable', 'Enable pool routing (writes cooling/affinity/retry-cap to config)'], ['pool --disable', 'Disable pool routing and restore non-pool config'], + [ + 'accounts order ', + 'Show effective drain order (priority bucket desc, then ID asc)', + ], + [ + 'accounts order --by-tier', + 'Set drain order from tier metadata (ultra > pro > free)', + ], + [ + 'accounts order --set a,b,c', + 'Set manual drain order (comma-separated account IDs)', + ], + ['accounts order --reset', 'Revert drain order to stable file order'], ], ], [ diff --git a/src/commands/cliproxy/index.ts b/src/commands/cliproxy/index.ts index 34a3ca78..b580fdda 100644 --- a/src/commands/cliproxy/index.ts +++ b/src/commands/cliproxy/index.ts @@ -49,6 +49,7 @@ import { handleCatalogJson, } from './catalog-subcommand'; import { handlePoolSubcommand } from './pool-subcommand'; +import { handleOrderSubcommand } from './order-subcommand'; /** * Parse --backend flag from args @@ -192,6 +193,21 @@ export async function handleCliproxyCommand(args: string[]): Promise { return; } + if (command === 'accounts') { + const subcommand = remainingArgs[1]; + if (subcommand === 'order') { + await handleOrderSubcommand(remainingArgs.slice(2)); + return; + } + // Unknown (or missing) accounts subcommand: report and show help. + // 'order' is currently the only accounts subcommand. + console.error(`[X] Unknown accounts subcommand: ${subcommand ?? '(none)'}`); + console.error(' Usage: ccs cliproxy accounts order '); + process.exitCode = 1; + await showHelp(); + return; + } + if (command === 'routing') { const subcommand = remainingArgs[1]; if (subcommand === 'set') { diff --git a/src/commands/cliproxy/order-subcommand.ts b/src/commands/cliproxy/order-subcommand.ts new file mode 100644 index 00000000..446db691 --- /dev/null +++ b/src/commands/cliproxy/order-subcommand.ts @@ -0,0 +1,596 @@ +/** + * CLIProxy Accounts Drain Order Subcommand + * + * Handles: + * ccs cliproxy accounts order + * ccs cliproxy accounts order --by-tier + * ccs cliproxy accounts order --set a@x.com,b@y.com,... + * ccs cliproxy accounts order --reset + */ + +import { initUI, header, subheader, color, dim, ok, fail, warn, info } from '../../utils/ui'; +import { extractOption, hasAnyFlag } from '../arg-extractor'; +import { + saveDrainOrderConfig, + loadDrainOrderConfig, + clearDrainOrderConfig, +} from '../../cliproxy/accounts/registry'; +import { getProviderAccounts } from '../../cliproxy/accounts/query'; +import { + computeManualDrainOrder, + computeTierDrainOrder, + applyDrainOrder, + annotateCurrentPriorities, + readAuthFilePriority, + type DrainOrderEntry, + type DrainOrderInput, +} from '../../cliproxy/accounts/drain-order'; +import { detectRunningProxy } from '../../cliproxy/proxy/proxy-detector'; +import { getProxyTarget } from '../../cliproxy/proxy/proxy-target-resolver'; +import { mapExternalProviderName } from '../../cliproxy/provider-capabilities'; +import type { CLIProxyProvider } from '../../cliproxy/types'; +import type { AccountTier } from '../../cliproxy/accounts/types'; +import { resolveLifecyclePort } from '../../cliproxy/config/port-manager'; + +/** Providers where tier metadata is expected and --by-tier is meaningful. */ +const TIER_AWARE_PROVIDERS = new Set(['agy', 'gemini']); + +/** + * Tie-break key for a token file, matching the upstream Go selector. + * The selector breaks priority ties by Auth.ID asc. Upstream (synthesizer + * file.go:120-122) lowercases Auth.ID only on Windows; on other platforms it + * compares by raw byte order. We mirror that platform-conditional behaviour so + * the displayed order matches what CLIProxy actually drains. + * + * @param platform process.platform; defaults to the current platform. Exposed + * for tests so non-Windows byte-order behaviour can be asserted deterministically. + */ +export function tieBreakKey( + tokenFile: string, + platform: NodeJS.Platform = process.platform +): string { + return platform === 'win32' ? tokenFile.toLowerCase() : tokenFile; +} + +function formatTierLabel(tier: AccountTier | undefined, tierDerived: boolean): string { + if (!tier || tier === 'unknown') { + return dim('unknown'); + } + const label = + tier === 'ultra' ? color(tier, 'success') : tier === 'pro' ? color(tier, 'info') : dim(tier); + return tierDerived ? label : dim(tier); +} + +function printOrderTable( + entries: DrainOrderEntry[], + showCurrentPriority: boolean, + sortByFilePriority: boolean = false +): void { + const rows = entries.slice().sort((a, b) => { + if (sortByFilePriority) { + // Sort by currentPriority (the selector's actual input), treating undefined as 0. + // This is the priority the selector actually sees on disk, not the computed config priority. + const aCurrent = a.currentPriority ?? 0; + const bCurrent = b.currentPriority ?? 0; + return bCurrent - aCurrent || (tieBreakKey(a.tokenFile) < tieBreakKey(b.tokenFile) ? -1 : 1); + } + // Tie-break by tokenFile ascending, matching the Go selector. Upstream sorts + // by Auth.ID byte order and only lowercases on Windows (see tieBreakKey). + // localeCompare is intentionally avoided - Go sorts by byte value, not locale. + return ( + b.priority - a.priority || (tieBreakKey(a.tokenFile) < tieBreakKey(b.tokenFile) ? -1 : 1) + ); + }); + + const maxId = Math.max(...rows.map((r) => r.accountId.length), 'Account'.length); + const maxPri = Math.max(...rows.map((r) => String(r.priority).length), 'Priority'.length); + + console.log( + ` ${'#'.padStart(3)} ${color('Account'.padEnd(maxId), 'command')} ${'Priority'.padStart(maxPri)} Tier` + ); + console.log(` ${'---'} ${'-'.repeat(maxId)} ${'-'.repeat(maxPri)} ----`); + + rows.forEach((entry, idx) => { + const pos = String(idx + 1).padStart(3); + const id = entry.accountId.padEnd(maxId); + const pri = String(entry.priority).padStart(maxPri); + const tierLabel = formatTierLabel(entry.tier, entry.tierDerived); + const currentNote = + showCurrentPriority && entry.currentPriority !== undefined + ? dim(` (file: ${entry.currentPriority})`) + : ''; + console.log( + ` ${pos} ${color(id, 'command')} ${color(pri, 'info')} ${tierLabel}${currentNote}` + ); + }); +} + +/** + * Show effective drain order for a provider without making any changes. + */ +async function handleOrderShow(provider: CLIProxyProvider): Promise { + // Use getProviderAccounts() so auth files not yet in accounts.json + // (e.g. file-copied fleets) are visible and tier metadata is preserved. + const allAccounts = getProviderAccounts(provider); + + if (allAccounts.length === 0) { + console.log(warn(`No accounts found for provider: ${provider}`)); + console.log(''); + return; + } + + const accounts: DrainOrderInput[] = allAccounts + .filter((a) => !a.paused) + .map((a) => ({ + accountId: a.id, + tokenFile: a.tokenFile, + tier: a.tier, + })); + + if (accounts.length === 0) { + console.log(warn(`All accounts for ${provider} are paused.`)); + console.log(''); + return; + } + + const drainCfg = loadDrainOrderConfig(provider); + let entries: DrainOrderEntry[]; + let modeLabel: string; + + // Manual mode is only usable when at least one stored ID still maps to a live + // account. If every stored ID is stale, fall back to the file-order view. + let manualValidIds: string[] | undefined; + if (drainCfg?.mode === 'manual' && drainCfg.orderedIds && drainCfg.orderedIds.length > 0) { + const existingIds = new Set(accounts.map((a) => a.accountId)); + manualValidIds = drainCfg.orderedIds.filter((id) => existingIds.has(id)); + } + + if (manualValidIds && manualValidIds.length > 0) { + entries = computeManualDrainOrder(manualValidIds, accounts); + modeLabel = `manual (${color('--set', 'command')})`; + } else if (drainCfg?.mode === 'tier') { + entries = computeTierDrainOrder(accounts); + modeLabel = `tier-derived (${color('--by-tier', 'command')})`; + } else { + // File order: no priority writes; show stable file order. Reached when no + // config is stored, or when a stored manual order has only stale IDs. + printFileOrderView(provider, accounts); + return; + } + + annotateCurrentPriorities(entries); + + // Detect drift: any account where the computed config priority does not match + // what is actually on disk (the selector's real input). Missing file priority + // (undefined) treated as 0 by the selector, so it also counts as drift when + // the config says something higher. + const hasDrift = entries.some((e) => (e.currentPriority ?? 0) !== e.priority); + + console.log(` Mode: ${modeLabel}`); + if (hasDrift) { + console.log(''); + console.log( + warn( + `Drift detected: stored order does not match auth files; re-run --set/--by-tier to re-apply.` + ) + ); + } + console.log(''); + // Show rows in selector pick order (currentPriority desc, tokenFile asc on tie). + // This reflects what CLIProxy will actually drain, not just the intended config. + console.log(subheader('Effective drain order (selector pick order, highest priority first):')); + printOrderTable(entries, true, true); + console.log(''); + console.log( + dim(` To update: ccs cliproxy accounts order ${provider} --set `) + ); + if (TIER_AWARE_PROVIDERS.has(provider)) { + console.log(dim(` Tier-based: ccs cliproxy accounts order ${provider} --by-tier`)); + } + console.log(dim(` To reset: ccs cliproxy accounts order ${provider} --reset`)); + console.log(''); +} + +function readFilePriorityNote(tokenFile: string): string { + try { + const p = readAuthFilePriority(tokenFile); + return p !== undefined ? dim(` [priority: ${p}]`) : ''; + } catch { + return ''; + } +} + +/** + * Render the stable file-order view (no priority writes). Used when no drain + * order config is stored, or when a stored manual order has only stale IDs. + */ +function printFileOrderView(provider: CLIProxyProvider, accounts: DrainOrderInput[]): void { + if (!TIER_AWARE_PROVIDERS.has(provider)) { + console.log( + info( + `Tier unknown for ${provider} accounts - using file order.\n` + + ` Use --set to specify manual order.` + ) + ); + } + console.log(` Mode: ${dim('file order (no priority set)')}`); + console.log(''); + console.log(subheader('Active accounts (file order):')); + // Sort by tokenFile ascending to match the Go selector's stable order. + // Upstream sorts by Auth.ID byte order and only lowercases on Windows (see tieBreakKey). + const sortedByFile = accounts + .slice() + .sort((a, b) => (tieBreakKey(a.tokenFile) < tieBreakKey(b.tokenFile) ? -1 : 1)); + sortedByFile.forEach((a, idx) => { + const pri = readFilePriorityNote(a.tokenFile); + console.log(` ${String(idx + 1).padStart(3)} ${color(a.accountId, 'command')}${pri}`); + }); + console.log(''); + console.log( + dim(` To set manual order: ccs cliproxy accounts order ${provider} --set a@x.com,b@y.com`) + ); + if (TIER_AWARE_PROVIDERS.has(provider)) { + console.log( + dim(` To use tier-based order: ccs cliproxy accounts order ${provider} --by-tier`) + ); + } + console.log(''); +} + +/** + * Apply tier-derived drain order for a provider. + */ +async function handleOrderByTier(provider: CLIProxyProvider): Promise { + if (!TIER_AWARE_PROVIDERS.has(provider)) { + console.log( + warn( + `Tier metadata is not available for ${provider} accounts.\n` + + ` Tier is tracked for: ${[...TIER_AWARE_PROVIDERS].join(', ')}\n` + + ` Use --set for manual order instead.` + ) + ); + console.log(''); + process.exitCode = 1; + return; + } + + // Use getProviderAccounts() so auth files not yet in accounts.json + // (e.g. file-copied fleets) are visible and tier metadata is preserved. + const allProviderAccounts = getProviderAccounts(provider); + + if (allProviderAccounts.length === 0) { + console.log(warn(`No accounts found for provider: ${provider}`)); + console.log(''); + process.exitCode = 1; + return; + } + + const accounts: DrainOrderInput[] = allProviderAccounts + .filter((a) => !a.paused) + .map((a) => ({ + accountId: a.id, + tokenFile: a.tokenFile, + tier: a.tier, + })); + + if (accounts.length === 0) { + console.log(warn(`All ${provider} accounts are paused.`)); + console.log(''); + process.exitCode = 1; + return; + } + + const allUnknown = accounts.every((a) => !a.tier || a.tier === 'unknown'); + if (allUnknown) { + console.log( + warn( + `All ${provider} accounts have unknown tier.\n` + + ` Run quota to populate tier metadata, or use --set for manual order.` + ) + ); + console.log(''); + process.exitCode = 1; + return; + } + + const entries = computeTierDrainOrder(accounts); + + console.log(subheader('Computed tier-based drain order:')); + printOrderTable(entries, false); + console.log(''); + + // Check proxy state for write path selection. + // Use the configured local port (not the hardcoded default) for detection. + // Remote targets: refuse with guidance - detection and current-priority reads + // require management API access that is not wired in v1. + const proxyTarget = getProxyTarget(); + if (proxyTarget.isRemote) { + console.log( + warn( + `Remote proxy target detected. Drain order management for remote proxies is not\n` + + ` supported in v1. Run this command on the host running CLIProxy instead.` + ) + ); + console.log(''); + process.exitCode = 1; + return; + } + + const localPort = resolveLifecyclePort(); + const proxyStatus = await detectRunningProxy(localPort); + + // Ambiguous state: proxy process is alive but not yet responding to HTTP. + // A direct file write while CLIProxy is serving could be silently clobbered + // by the MarkResult-persist path. Refuse rather than risk priority loss. + if (proxyStatus.running && !proxyStatus.verified) { + console.log(fail('Proxy state ambiguous - retry in a moment or stop the proxy first.')); + console.log(''); + process.exitCode = 1; + return; + } + + const proxyRunning = proxyStatus.running && proxyStatus.verified; + + if (proxyRunning) { + console.log(info('Proxy is running - writing via management API (avoids clobber race).')); + } else { + console.log(info('Proxy is stopped - writing directly to auth files.')); + } + console.log(''); + + const result = await applyDrainOrder(entries, proxyRunning); + + if (result.written.length > 0) { + console.log(ok(`Set priorities for ${result.written.length} account(s).`)); + } + if (result.skipped.length > 0) { + console.log(info(`Skipped ${result.skipped.length} account(s) (priority already correct).`)); + } + if (result.failed.length > 0) { + for (const f of result.failed) { + console.log(fail(`Failed to set priority for ${f.entry.accountId}: ${f.reason}`)); + } + process.exitCode = 1; + } + + if (result.failed.length === 0) { + // Persist mode + const persisted = saveDrainOrderConfig(provider, { mode: 'tier' }); + if (persisted) { + console.log( + ok( + `Drain order mode saved as "tier". Re-run "ccs cliproxy accounts order ${provider} --by-tier" after adding or re-authing accounts to re-apply.` + ) + ); + } else { + console.log( + warn( + `Order applied to auth files but mode not persisted - no registered accounts for ${provider}; run ccs cliproxy quota to register, then re-run` + ) + ); + } + } + console.log(''); +} + +/** + * Apply manual drain order specified as comma-separated account IDs. + */ +async function handleOrderSet(provider: CLIProxyProvider, setArg: string): Promise { + const orderedIds = setArg + .split(',') + .map((s) => s.trim()) + .filter((s) => s.length > 0); + + if (orderedIds.length === 0) { + console.log(fail('--set requires a comma-separated list of account IDs.')); + console.log(''); + process.exitCode = 1; + return; + } + + // Use getProviderAccounts() so auth files not yet in accounts.json + // (e.g. file-copied fleets) are visible and tier metadata is preserved. + const allProviderAccounts = getProviderAccounts(provider); + + if (allProviderAccounts.length === 0) { + console.log(warn(`No accounts found for provider: ${provider}`)); + console.log(''); + process.exitCode = 1; + return; + } + + const allAccounts: DrainOrderInput[] = allProviderAccounts + .filter((a) => !a.paused) + .map((a) => ({ + accountId: a.id, + tokenFile: a.tokenFile, + tier: a.tier, + })); + + let entries: DrainOrderEntry[]; + try { + entries = computeManualDrainOrder(orderedIds, allAccounts); + } catch (err) { + console.log(fail(`${(err as Error).message}`)); + console.log(''); + process.exitCode = 1; + return; + } + + console.log(subheader('Computed manual drain order:')); + printOrderTable(entries, false); + console.log(''); + + // Use the configured local port (not the hardcoded default) for detection. + // Remote targets: refuse with guidance. + const proxyTarget = getProxyTarget(); + if (proxyTarget.isRemote) { + console.log( + warn( + `Remote proxy target detected. Drain order management for remote proxies is not\n` + + ` supported in v1. Run this command on the host running CLIProxy instead.` + ) + ); + console.log(''); + process.exitCode = 1; + return; + } + + const localPort = resolveLifecyclePort(); + const proxyStatus = await detectRunningProxy(localPort); + + // Ambiguous state: proxy process is alive but not yet responding to HTTP. + // A direct file write while CLIProxy is serving could be silently clobbered + // by the MarkResult-persist path. Refuse rather than risk priority loss. + if (proxyStatus.running && !proxyStatus.verified) { + console.log(fail('Proxy state ambiguous - retry in a moment or stop the proxy first.')); + console.log(''); + process.exitCode = 1; + return; + } + + const proxyRunning = proxyStatus.running && proxyStatus.verified; + + if (proxyRunning) { + console.log(info('Proxy is running - writing via management API (avoids clobber race).')); + } else { + console.log(info('Proxy is stopped - writing directly to auth files.')); + } + console.log(''); + + const result = await applyDrainOrder(entries, proxyRunning); + + if (result.written.length > 0) { + console.log(ok(`Set priorities for ${result.written.length} account(s).`)); + } + if (result.skipped.length > 0) { + console.log(info(`Skipped ${result.skipped.length} account(s) (priority already correct).`)); + } + if (result.failed.length > 0) { + for (const f of result.failed) { + console.log(fail(`Failed to set priority for ${f.entry.accountId}: ${f.reason}`)); + } + process.exitCode = 1; + } + + if (result.failed.length === 0) { + const persisted = saveDrainOrderConfig(provider, { mode: 'manual', orderedIds }); + if (persisted) { + console.log( + ok( + `Drain order mode saved as "manual". Re-run "ccs cliproxy accounts order ${provider} --set " after adding or re-authing accounts to re-apply.` + ) + ); + } else { + console.log( + warn( + `Order applied to auth files but mode not persisted - no registered accounts for ${provider}; run ccs cliproxy quota to register, then re-run` + ) + ); + } + } + console.log(''); +} + +/** + * Reset drain order to file order (removes persisted config, no priority writes). + */ +async function handleOrderReset(provider: CLIProxyProvider): Promise { + const cleared = clearDrainOrderConfig(provider); + if (cleared) { + console.log(ok(`Drain order reset to file order for ${provider}.`)); + console.log( + dim( + ' Note: existing priority values in auth files are not removed.\n' + + ' CLIProxy will continue using them until they are overwritten or files are re-created.' + ) + ); + } else { + console.log(info(`No drain order config found for ${provider}. Already in file order.`)); + } + console.log(''); +} + +function printOrderHelp(provider?: string): void { + const prov = provider ?? ''; + console.log(subheader('Usage:')); + console.log(` ${color(`ccs cliproxy accounts order ${prov}`, 'command')}`); + console.log(` ${color(`ccs cliproxy accounts order ${prov} --by-tier`, 'command')}`); + console.log(` ${color(`ccs cliproxy accounts order ${prov} --set a@x.com,b@y.com`, 'command')}`); + console.log(` ${color(`ccs cliproxy accounts order ${prov} --reset`, 'command')}`); + console.log(''); + console.log(subheader('Options:')); + console.log(` ${dim('(no flags)')} Show effective drain order`); + console.log( + ` ${color('--by-tier', 'command')} Derive order from tier metadata (ultra > pro > free)` + ); + console.log( + ` ${color('--set', 'command')} Set manual order (comma-separated account IDs)` + ); + console.log(` ${color('--reset', 'command')} Revert to stable file order`); + console.log(''); + console.log(dim(' Tier-based ordering is available for: agy, gemini')); + console.log(dim(' Claude accounts have unknown tier; use --set for manual order.')); + console.log(''); +} + +/** + * Main handler for `ccs cliproxy accounts order [flags]` + */ +export async function handleOrderSubcommand(args: string[]): Promise { + await initUI(); + console.log(''); + console.log(header('CLIProxy Drain Order')); + console.log(''); + + if (hasAnyFlag(args, ['--help', '-h']) || args.length === 0) { + printOrderHelp(); + return; + } + + // First positional arg is provider name + const providerRaw = args[0]; + if (!providerRaw || providerRaw.startsWith('-')) { + console.log(fail('Provider name required. Usage: ccs cliproxy accounts order ')); + console.log(''); + printOrderHelp(); + process.exitCode = 1; + return; + } + + const provider = mapExternalProviderName(providerRaw) as CLIProxyProvider | null; + if (!provider) { + console.log(fail(`Unknown provider: ${providerRaw}`)); + console.log(''); + process.exitCode = 1; + return; + } + + const subArgs = args.slice(1); + + if (hasAnyFlag(subArgs, ['--reset'])) { + await handleOrderReset(provider); + return; + } + + if (hasAnyFlag(subArgs, ['--by-tier'])) { + await handleOrderByTier(provider); + return; + } + + const extracted = extractOption(subArgs, ['--set']); + if (extracted.found) { + if (extracted.missingValue || !extracted.value) { + console.log(fail('--set requires a value: --set a@x.com,b@y.com,...')); + console.log(''); + process.exitCode = 1; + return; + } + await handleOrderSet(provider, extracted.value); + return; + } + + // Default: show current order + await handleOrderShow(provider); +}