diff --git a/devlog/_plan/260829_kiro_quota_pool/090_predispatch_selection.md b/devlog/_plan/260829_kiro_quota_pool/090_predispatch_selection.md new file mode 100644 index 0000000000..b73548fe19 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/090_predispatch_selection.md @@ -0,0 +1,87 @@ +# 090 — Work-phase 3: pre-dispatch account selection + +Doc `080` recorded kiro-lb as ahead on one axis that matters directly to the user's ask: +it picks an account *before* dispatch, while we only reordered the 429 recovery path. This +phase closes that gap. Branch `codex/kiro-pool-predispatch`, off merged `dev` `d82b3049d`. + +## What changed + +`preferredInitialAccount(config, provider)` answers "which account should open this turn". +The initial OAuth resolution in `src/server/responses/core.ts` consults it and, when it +names an account, resolves that account's snapshot instead of the active one. + +It is a **preference, not a gate**. A null answer means "use the active account", and null +is returned for: rotation disabled, fewer than two accounts, no quota evidence anywhere on +the roster, every candidate cooled, or the ranking simply agreeing with the active account. +A provider with no per-account quota therefore behaves exactly as before. + +## Five review rounds + +An independent reviewer failed this four times before passing. Each finding was real, and +three of them were defects I would not have found by testing the happy path. + +### Round 1 — three blockers + +1. **Antigravity could pair B's bearer with A's project.** The ordinary path fills the CCA + project only when it is *empty* (`!route.provider.project`), so a preferred account + installed its own bearer beside the configured account's project — #2841 in its + original shape, at a site nobody had reason to look at. +2. **A quota-less provider could still be redirected.** Cooling the active account collapses + the eligible list to one candidate, and ranking a single candidate returns it unchanged. + That *looks* like a ranked answer while nothing was ever measured. Evidence is now + checked across the whole roster, before eligibility narrows anything. +3. **Two uncached credential-file reads per request.** `loadAuthStore` chmods the config + dir, chmods the secret, and re-parses the whole file on every call — the exact cost the + neighbouring `PRESENCE_CACHE_TTL_MS` comment exists to warn about. + +### Round 2 — the fail-closed 401 was worse than the bug + +My first Antigravity fix returned 401 when a preferred account had no project. But +Antigravity tolerates project discovery failing, so a project-less account is an ordinary +stored state: a *preference* had been given the power to break a request that would +otherwise have worked. It now falls back to the active account. + +### Round 3 — a removed account became a 401 + +The roster is cached for two seconds, so an account can be deleted after being chosen. +Resolving it throws, and that throw reached the client as 401 while a healthy active +account sat unused. The reviewer reproduced it exactly. Resolution failures now drop the +stale roster and retry on the active account. + +### Round 4 — the one a catch could not catch + +The sharpest finding. An account newly flagged `needsReauth` **does not throw**: its +credential is still readable, so resolution succeeds and no error path fires. The request +would dispatch on an account already known to need a fresh login. + +My first fix re-read the store to validate the winner — and reopened blocker 3, because the +steady state of this feature is a pool where one account consistently ranks higher, so +"validate only on redirect" is "validate on every request". + +### Round 5 — atomic validation, then PASS + +The check belongs where the store row is *already* being read. +`getAccountCredentialWithStatus` returns credential and `needsReauth` from one read, and +`requireUsableAccount` makes account-scoped resolution reject an unusable account from +inside it. Selection now performs no store read at all; the caller's existing fallback +handles the rejection. Zero added I/O on the redirect path, both stale classes closed. + +## Verification + +```text +bun x tsc --noEmit -> exit 0 +bun run privacy:scan -> Privacy scan passed +bun test (11 files) -> 181 pass / 0 fail / 656 expect() calls +core-lab-boundary -> pass, no new src/lab/ reach +``` + +Tests worth naming, because each encodes a defect above: a redirecting selection with +`auth.json` deleted still answers (proves the cache); a reauth-flagged account resolves +plainly but rejects under `requireUsableAccount` (proves why a catch was insufficient); and +cooling the *active* account of a quota-less provider still returns null. + +## Result + +The "pre-request selection" row moves out of doc `080`'s "they are ahead" column. Two rows +remain there honestly: kiro-lb persists quota across restart, and it has a real operations +dashboard. Neither is in scope here. diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts index 3cc0e29e2d..b3b4e0bd00 100644 --- a/src/oauth/account-quota-rank.ts +++ b/src/oauth/account-quota-rank.ts @@ -75,6 +75,19 @@ export function rankAccountsByHeadroom(provider: string, ring: readonly string[] .map(entry => entry.id); } +/** + * Do we hold any measurement at all for these accounts? + * + * Ranking a single candidate is trivially the identity, which makes it useless as an + * evidence test: a caller that has already filtered its list down to one account would be + * told "ranked" when nothing was measured. Pre-dispatch selection asks this first so it + * can decline to act on a roster it knows nothing about. + */ +export function hasHeadroomEvidence(provider: string, ids: readonly string[]): boolean { + return ids.some(id => + headroomOf(provider, id) !== null + || (provider === "kiro" && getKiroAccountExhaustion(`${provider}\u0000${id}`) !== null)); +} /** * How long to cool an account that just 429'd, when we know its allowance is spent. * diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index cec3e6321d..a36974a629 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -16,7 +16,7 @@ */ import { getAccountSet } from "./store"; import { getValidAccessSnapshotForAccount, type OAuthAccessSnapshot } from "./index"; -import { exhaustedCooldownMs, rankAccountsByHeadroom } from "./account-quota-rank"; +import { exhaustedCooldownMs, hasHeadroomEvidence, rankAccountsByHeadroom } from "./account-quota-rank"; import { parseRetryAfterMs } from "../combos/failover"; import { sweepExpiredOnWrite } from "../lib/state-store-sweeper"; import type { OcxConfig, OcxProviderConfig } from "../types"; @@ -61,12 +61,29 @@ interface PresenceEntry { readAt: number; } +/** + * Ordered roster plus the active id, for the pre-dispatch preference. + * + * Same reasoning as the presence cache: `getAccountSet` reads through `loadAuthStore`, + * which chmods and re-parses the whole credential file on every call. Selection needs the + * ORDER and the active id, which the presence count cannot supply, so it gets its own + * TTL-bounded row. Ids and an active pointer only — never a credential. + */ +interface RosterEntry { + ids: string[]; + activeId: string | null; + readAt: number; +} + /** Process-local, like the Anthropic pool's: a restart is allowed to forget a cooldown. */ const health = new Map(); /** Provider -> recent eligible-account count. TTL-bounded; never holds credential material. */ const presence = new Map(); +/** Provider -> recently read roster. TTL-bounded; never holds credential material. */ +const roster = new Map(); + const healthKey = (provider: string, accountId: string) => `${provider}\u0000${accountId}`; function isCooled(provider: string, accountId: string, now: number): boolean { @@ -101,6 +118,24 @@ function eligibleAccountCount(providerName: string, now: number): number { return eligible; } +/** + * Roster ids and the active pointer, read at most once per TTL window. + * + * `needsReauth` accounts are excluded for the same reason the presence count excludes + * them: a revoked credential cannot serve the request we are about to send. + */ +function cachedRoster(providerName: string, now: number): { ids: string[]; activeId: string | null } { + const cached = roster.get(providerName); + if (cached && now >= cached.readAt && now - cached.readAt < PRESENCE_CACHE_TTL_MS) { + return { ids: cached.ids, activeId: cached.activeId }; + } + const set = getAccountSet(providerName); + const ids = set ? set.accounts.filter(a => a.needsReauth !== true).map(a => a.id) : []; + const activeId = set?.activeAccountId ?? null; + roster.set(providerName, { ids, activeId, readAt: now }); + return { ids, activeId }; +} + /** * Presence IS consent (#2568d). * @@ -184,6 +219,8 @@ export function rotateGenericOAuthAccountOn429( // A rotation means the roster in use just changed; do not answer the next activation question // from a count read before the failure. presence.delete(providerName); + // Same for the selection roster: the next request must not pick from a pre-failure read. + roster.delete(providerName); // Deterministic: start after the failed account so repeated 429s walk the roster instead of // hammering whichever id happens to sort first. The ring is built BEFORE ranking — ranking // the store's own order would change which account a quota-less provider rotates to. @@ -211,6 +248,61 @@ export async function failoverAccountSnapshot( return getValidAccessSnapshotForAccount(providerName, accountId); } +/** + * Which account should serve the FIRST attempt of a request. + * + * Rotation only ever ran after a 429, so a turn still opened on whichever account happened + * to be active — including one a previous probe already measured as spent. That costs a + * full upstream round trip and one of three rotations to rediscover what the cache knew. + * + * Returns null whenever the ordinary active-account path should be used unchanged: no + * quorum, rotation disabled, a single account, or no quota evidence to act on. This is a + * preference, never a gate — a cooled or unmeasured account is still perfectly usable, so + * an empty answer means "carry on", not "refuse". + */ +export function preferredInitialAccount( + config: OcxConfig, + providerName: string, + now = Date.now(), +): string | null { + if (!isGenericOAuthFailoverEnabled(config, providerName)) return null; + // This runs on the initial resolution of EVERY request, and `loadAuthStore` has no + // cache: each call chmods the config dir, chmods the secret, reads the whole file and + // normalizes it (store.ts:136-151). So the store is consulted at most ONCE here, behind + // the same TTL the presence check uses, and never at all for a single-account provider. + const { ids: order, activeId: active } = cachedRoster(providerName, now); + if (order.length < 2) return null; + + // Evidence is required BEFORE eligibility narrows the field. Without this, a provider + // with no quota data at all could still be redirected: cool the active account with a + // 429 and the eligible list collapses to one candidate, which any ranking returns + // unchanged — an answer that looks ranked but was never measured. The no-op guarantee + // for quota-less providers has to be checked on the full roster. + if (!hasHeadroomEvidence(providerName, order)) return null; + + // Cooldowns are respected here, unlike in the presence count: this picks the account to + // send to right now, and one inside its 429 window is the single candidate we hold + // positive evidence against. + const eligible = order.filter(id => !isCooled(providerName, id, now)); + if (eligible.length === 0) return null; + + // Start the ring at the active account so an unranked outcome reproduces today's choice. + const start = active ? order.indexOf(active) : -1; + const ring = start >= 0 ? [...order.slice(start), ...order.slice(0, start)] : order; + const candidates = ring.filter(id => eligible.includes(id)); + if (candidates.length === 0) return null; + + const best = rankAccountsByHeadroom(providerName, candidates)[0] ?? null; + // Nothing to do when the ranking agrees with the account we would have used anyway. + // + // The roster may be up to PRESENCE_CACHE_TTL_MS old, so this answer is a PREFERENCE the + // caller must be able to abandon: it resolves the account with `requireUsableAccount`, + // which rejects a removed or reauth-flagged account inside the store read it was already + // performing, and falls back to the active account. Validating here instead would mean a + // second uncached read of the credential file on every redirected request. + return best && best !== active ? best : null; +} + /** Earliest remaining cooldown, for a client-facing Retry-After when every account is cooled. */ export function genericFailoverRetryAfterSeconds(providerName: string, now = Date.now()): number | null { const set = getAccountSet(providerName); @@ -224,14 +316,22 @@ export function genericFailoverRetryAfterSeconds(providerName: string, now = Dat return earliest === null ? null : Math.max(1, Math.ceil((earliest - now) / 1000)); } +/** Test seam and manual-recovery hook. */ +export function forgetGenericFailoverRoster(providerName: string): void { + roster.delete(providerName); + presence.delete(providerName); +} + /** Test seam and manual-recovery hook. */ export function clearGenericFailoverHealth(providerName?: string): void { if (!providerName) { health.clear(); presence.clear(); + roster.clear(); return; } presence.delete(providerName); + roster.delete(providerName); for (const key of [...health.keys()]) { if (key.startsWith(`${providerName}\u0000`)) health.delete(key); } diff --git a/src/oauth/index.ts b/src/oauth/index.ts index 1b6f4690db..7a356df8c8 100644 --- a/src/oauth/index.ts +++ b/src/oauth/index.ts @@ -4,7 +4,7 @@ import type { OcxConfig, OcxProviderConfig, RefreshPolicy } from "../types"; import { loadConfig, resolveEnvValue, saveConfig } from "../config"; import { maskEmail } from "../lib/privacy"; import { KiroTokenRefreshError, environmentKiroRoutingMetadata, loginKiro, refreshKiroToken, settleKiroLoginTransaction } from "./kiro"; -import { getAccountCredential, getAccountSet, removeAccount, saveAccountCredential, saveCredential, setActiveAccount, getCredential, credentialGeneration, createOAuthRefreshIntentLock, mergeAccountCredential, markAccountNeedsReauthIfGeneration, readOAuthRefreshIntent, writeOAuthRefreshIntent, markOAuthRefreshIntentStaleOwner, clearOAuthRefreshIntent, normalizeAuthStoreBuffer, OAuthMutationBusyError } from "./store"; +import { getAccountCredential, getAccountCredentialWithStatus, getAccountSet, removeAccount, saveAccountCredential, saveCredential, setActiveAccount, getCredential, credentialGeneration, createOAuthRefreshIntentLock, mergeAccountCredential, markAccountNeedsReauthIfGeneration, readOAuthRefreshIntent, writeOAuthRefreshIntent, markOAuthRefreshIntentStaleOwner, clearOAuthRefreshIntent, normalizeAuthStoreBuffer, OAuthMutationBusyError } from "./store"; import { loginXai, refreshXaiToken, XAI_LOCAL_CLI_DETACH_WARNING, XaiTokenRequestError } from "./xai"; import { ANTHROPIC_OAUTH_BETA, AnthropicTokenError, loginAnthropic, refreshAnthropicToken } from "./anthropic"; import { loginKimi, refreshKimiToken } from "./kimi"; @@ -436,11 +436,18 @@ async function resolveAccessSnapshotForAccount( provider: string, accountId: string, rejectedGeneration?: string, + requireUsableAccount = false, ): Promise { const def = OAUTH_PROVIDERS[provider]; if (!def) throw new UnsupportedOAuthProviderError(provider); - const cred = getAccountCredential(provider, accountId); - if (!cred) throw new OAuthLoginRequiredError(provider); + // One store read answers both questions. A caller that opts in gets the account REJECTED + // when it needs reauthentication, which a bare credential read cannot detect: a revoked + // account keeps a readable credential, so resolution would otherwise succeed and the + // request would dispatch on an account already known to need a fresh login. + const row = getAccountCredentialWithStatus(provider, accountId); + if (!row) throw new OAuthLoginRequiredError(provider); + if (requireUsableAccount && row.needsReauth) throw new OAuthLoginRequiredError(provider); + const cred = row.credential; const current = accessSnapshot(provider, accountId, cred); if (rejectedGeneration !== undefined && current.generation !== rejectedGeneration) return current; if (rejectedGeneration === undefined && cred.expires > Date.now() + REFRESH_SKEW_MS) return current; @@ -529,8 +536,9 @@ export async function getValidAccessTokenForAccount(provider: string, accountId: export async function getValidAccessSnapshotForAccount( provider: string, accountId: string, + opts: { requireUsableAccount?: boolean } = {}, ): Promise { - return resolveAccessSnapshotForAccount(provider, accountId); + return resolveAccessSnapshotForAccount(provider, accountId, undefined, opts.requireUsableAccount === true); } /** Terminal refresh failures (revoked/rotated-away grants) — retrying cannot succeed. */ diff --git a/src/oauth/store.ts b/src/oauth/store.ts index 494389a18c..755d86f3e2 100644 --- a/src/oauth/store.ts +++ b/src/oauth/store.ts @@ -641,6 +641,22 @@ export function getAccountCredential(provider: string, accountId: string): OAuth return loadAuthStore()[provider]?.accounts.find(a => a.id === accountId)?.credential ?? null; } +/** + * Credential plus the account's reauth flag, from ONE store read. + * + * A caller that checks `needsReauth` separately pays a second `loadAuthStore`, which + * chmods and re-parses the whole credential file. Returning both together lets an + * account-scoped resolver reject a revoked account without that extra read. + */ +export function getAccountCredentialWithStatus( + provider: string, + accountId: string, +): { credential: OAuthCredentials; needsReauth: boolean } | null { + const account = loadAuthStore()[provider]?.accounts.find(a => a.id === accountId); + if (!account?.credential) return null; + return { credential: account.credential, needsReauth: account.needsReauth === true }; +} + /** Persist a refreshed credential for a SPECIFIC account without touching activeAccountId. */ export async function saveAccountCredential( provider: string, diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 691d75ada4..9618e0e538 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -101,6 +101,7 @@ import type { import { forceRefreshOAuthAccessSnapshot, getValidAccessTokenForAccount, + getValidAccessSnapshotForAccount, getValidAccessTokenSnapshot, publicOAuthAuthenticationErrorMessage, type OAuthAccessSnapshot, @@ -121,9 +122,11 @@ import { import { stampOAuthAccountLabel } from "../../providers/label"; import { failoverAccountSnapshot, + forgetGenericFailoverRoster, GENERIC_OAUTH_MAX_FAILOVERS_PER_REQUEST, isGenericFailoverProvider, isGenericOAuthFailoverEnabled, + preferredInitialAccount, rotateGenericOAuthAccountOn429, } from "../../oauth/generic-account-failover"; import { resolveCopilotApiBaseUrl } from "../../oauth/github-copilot"; @@ -3099,7 +3102,53 @@ async function handleResponsesInner( route.provider = { ...route.provider, apiKey: accessToken }; logCtx.provider = formatAnthropicProviderForLog("anthropic", selection.accountId, config); } else { - const resolved = await getValidAccessTokenSnapshot(route.providerName); + // Prefer the account with known headroom BEFORE the first attempt. Rotation alone + // only reacts to a 429, so a turn could open on an account a previous probe already + // measured as spent. A null answer means "use the active account", so every provider + // without quota evidence keeps the resolution it has today. + const preferredAccountId = isGenericFailoverProvider(route.providerName, route.provider) + ? preferredInitialAccount(config, route.providerName) + : null; + // Resolved account-scoped, NOT through failoverAccountSnapshot: that helper marks a + // rotation site, and rotation sites must apply their credential through + // applyFailoverSnapshot's pairing rules. This is initial resolution — the code below + // already pairs the snapshot's Kiro metadata, Copilot origin and Antigravity project + // with this same bearer, exactly as it does for the active account. + let usedPreferredAccount = preferredAccountId !== null; + let resolved: OAuthAccessSnapshot; + if (preferredAccountId) { + try { + // `requireUsableAccount` makes a removed OR reauth-flagged account throw from + // inside the resolver's own store read. Without it a revoked account resolves + // successfully — its credential is still readable — and the request would + // dispatch on an account already known to need a fresh login. + resolved = await getValidAccessSnapshotForAccount( + route.providerName, + preferredAccountId, + { requireUsableAccount: true }, + ); + } catch { + // The roster is read behind a short TTL, so a preferred account can be removed + // or flagged for reauth in the window after it was cached. Resolving it then + // throws, and a PREFERENCE that turns a healthy request into a 401 is worse + // than no preference at all — the active account is still perfectly usable. + // Drop the stale roster so the next request re-reads it, and carry on. + forgetGenericFailoverRoster(route.providerName); + usedPreferredAccount = false; + resolved = await getValidAccessTokenSnapshot(route.providerName); + } + } else { + resolved = await getValidAccessTokenSnapshot(route.providerName); + } + // A Cloud Code Assist account needs its own project. Antigravity's refresh path + // tolerates project discovery failing, so a stored account can legitimately have + // none — and a PREFERENCE must never turn a working request into an error. Fall + // back to the ordinary active-account resolution instead, which is exactly what + // would have happened had the preference never existed. + if (usedPreferredAccount && route.provider.googleMode === "cloud-code-assist" && !resolved.projectId) { + resolved = await getValidAccessTokenSnapshot(route.providerName); + usedPreferredAccount = false; + } replayOAuthCredentialSnapshot = { accountId: resolved.accountId, generation: resolved.generation, @@ -3123,9 +3172,19 @@ async function handleResponsesInner( // Antigravity (cloud-code-assist) needs the discovered Cloud Code Assist project id in the // CCA envelope. Keep it paired with the token snapshot so an account rotation cannot mix // a fresh token with project metadata re-read from a different credential generation. - if (route.provider.googleMode === "cloud-code-assist" && !route.provider.project) { - const projectId = resolved.projectId; - if (projectId) route.provider = { ...route.provider, project: projectId }; + if (route.provider.googleMode === "cloud-code-assist") { + // When pre-dispatch chose a DIFFERENT account, the configured project belongs to + // the account we did not use, and `!route.provider.project` would skip right past + // it — installing B's bearer alongside A's project. That is the #2841 pairing bug + // in its original shape, so the preferred-account path replaces the project + // unconditionally and refuses to dispatch at all if the chosen account has none. + // A project-less preferred account already fell back above, so by here the + // preferred path always has one. + if (usedPreferredAccount && resolved.projectId) { + route.provider = { ...route.provider, project: resolved.projectId }; + } else if (!route.provider.project && resolved.projectId) { + route.provider = { ...route.provider, project: resolved.projectId }; + } } } } catch (err) { diff --git a/tests/generic-oauth-failover.test.ts b/tests/generic-oauth-failover.test.ts index a9401efd5e..bf17cac044 100644 --- a/tests/generic-oauth-failover.test.ts +++ b/tests/generic-oauth-failover.test.ts @@ -280,6 +280,28 @@ describe("sidecar on429 wiring", () => { // Kiro routing metadata still travels with its own token. expect(body).toContain("_kiroAuthContext"); }); + + test("pre-dispatch selection replaces the CCA project instead of inheriting one", () => { + // The same pairing rule as the rotation helper, at the OTHER site that can change which + // account serves a request. The ordinary path is guarded by `!route.provider.project`, + // so without an explicit branch a preferred account would install its own bearer next + // to the configured account's project — #2841 in its original shape. + const start = coreSource.indexOf("const preferredAccountId ="); + expect(start).toBeGreaterThan(-1); + const region = coreSource.slice(start, start + 6000); + expect(region).toContain("usedPreferredAccount && resolved.projectId"); + // A project-less preferred account falls BACK to the ordinary active-account resolution + // rather than erroring: a preference must never turn a working request into a failure, + // and Antigravity tolerates project discovery failing, so an account with no project is + // an ordinary stored state. + expect(region).toContain("usedPreferredAccount = false"); + expect(region).not.toContain("has no Cloud Code Assist project"); + // Both fallbacks — a project-less account and an unresolvable one — must reach the SAME + // active-account resolution, so neither can dispatch on a half-applied identity. + const fallbacks = region.match(/usedPreferredAccount = false;/g) ?? []; + expect(fallbacks.length).toBe(2); + expect(region).toContain("forgetGenericFailoverRoster(route.providerName)"); + }); }); /** diff --git a/tests/kiro-pool-rank.test.ts b/tests/kiro-pool-rank.test.ts index 4cef86a93f..4dfe1491e2 100644 --- a/tests/kiro-pool-rank.test.ts +++ b/tests/kiro-pool-rank.test.ts @@ -1,5 +1,23 @@ import { afterEach, describe, expect, test } from "bun:test"; import { exhaustedCooldownMs, rankAccountsByHeadroom } from "../src/oauth/account-quota-rank"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + clearGenericFailoverHealth, + forgetGenericFailoverRoster, + preferredInitialAccount, + rotateGenericOAuthAccountOn429, +} from "../src/oauth/generic-account-failover"; +import { + getAccountSet, + markAccountNeedsReauth, + removeAccount, + saveCredential, + setActiveAccount, +} from "../src/oauth/store"; +import { getValidAccessSnapshotForAccount } from "../src/oauth"; +import type { OcxConfig, OcxProviderConfig } from "../src/types"; import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests, @@ -115,3 +133,232 @@ describe("exhaustion cooldown", () => { expect(exhaustedCooldownMs("xai", "a")).toBeNull(); }); }); + +describe("pre-dispatch account preference", () => { + const OAUTH_PROVIDER = { + adapter: "openai-chat", + baseUrl: "https://api.x.ai/v1", + authMode: "oauth", + } as unknown as OcxProviderConfig; + + const config = { providers: { xai: OAUTH_PROVIDER } } as unknown as OcxConfig; + const originalHome = process.env.OPENCODEX_HOME; + let home: string; + + async function seedAccounts(count: number): Promise { + for (let i = 0; i < count; i++) { + await saveCredential("xai", { + access: `access-${i}`, + refresh: `refresh-${i}`, + expires: Date.now() + 3_600_000, + accountId: `uuid-${i}`, + } as never, { addAccount: true }); + } + return getAccountSet("xai")?.accounts.map(a => a.id) ?? []; + } + + test("the account with more headroom is chosen before the first request", async () => { + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); + expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("no quota evidence leaves the active account alone", async () => { + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + // Null means "use the ordinary active-account path", so nothing changes for a + // provider that reports no per-account quota. + expect(preferredInitialAccount(config, "xai")).toBeNull(); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("a single account is never redirected", async () => { + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(1); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 99, updatedAt: Date.now() }); + expect(preferredInitialAccount(config, "xai")).toBeNull(); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("an account cooled by a recent 429 is not chosen to open the next request", async () => { + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + // Cool the roomier account: positive evidence against it outweighs its headroom. + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 50, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 1, updatedAt: Date.now() }); + rotateGenericOAuthAccountOn429(config, "xai", ids[1]!, null); + expect(preferredInitialAccount(config, "xai")).toBeNull(); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("a quota-less provider is never redirected, even when the ACTIVE account is cooled", async () => { + // The inverse of the case above, and the one that actually broke the no-op guarantee: + // cooling the active account collapses the eligible list to a single candidate, which + // any ranking returns unchanged. That looks like a ranked answer but nothing was ever + // measured, so evidence has to be checked against the whole roster first. + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + rotateGenericOAuthAccountOn429(config, "xai", ids[0]!, null); + expect(preferredInitialAccount(config, "xai")).toBeNull(); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("neither a redirecting nor a non-redirecting selection touches the credential store", async () => { + // loadAuthStore chmods the config dir, chmods the secret, and re-parses the whole + // credential file on every call — and this runs on the initial resolution of EVERY + // request. The steady state of this feature is a pool where one account consistently + // ranks higher, so the REDIRECTING path must be cached too — validating the winner here + // would put a second uncached read in front of every such request. Deleting the store + // proves it: an uncached path could not answer at all. Staleness is caught at + // resolution instead, inside a store read the resolver already performs. + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + // Redirecting: the other account holds more headroom on every call. + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); + expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + rmSync(join(home, "auth.json"), { force: true }); + for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + + // Non-redirecting: the active account already ranks best. + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 5, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 95, updatedAt: Date.now() }); + for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBeNull(); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("a preferred account removed inside the TTL degrades to the active account", async () => { + // The roster is cached for a short window, so an account can be removed after it was + // chosen. Resolving it then throws, and the request path must fall back to the active + // account rather than 401 — a preference must never break a request that would have + // worked without it. + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); + expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + + await removeAccount("xai", ids[1]!); + // Selection is a cached PREFERENCE, so it may still name the removed account... + expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + // ...and resolution is where that is caught. The request path absorbs this throw and + // falls back to the active account. + await expect( + getValidAccessSnapshotForAccount("xai", ids[1]!, { requireUsableAccount: true }), + ).rejects.toThrow(); + expect(getAccountSet("xai")?.accounts.map(a => a.id)).toEqual([ids[0]!]); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); + + test("a preferred account flagged for reauth inside the TTL is not selected", async () => { + // The dangerous variant of stale roster data: unlike a removal, a needsReauth account + // still has a readable credential, so resolution SUCCEEDS and no error path fires. The + // request would dispatch on an account already known to need a fresh login while a + // healthy active account sat unused. + home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + try { + const ids = await seedAccounts(2); + await setActiveAccount("xai", ids[0]!); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 95, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 5, updatedAt: Date.now() }); + expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + + await markAccountNeedsReauth("xai", ids[1]!, true); + // An ordinary resolve SUCCEEDS — the credential is still readable — which is exactly + // why the flag must be checked inside the resolver rather than trusted to throw. + await expect(getValidAccessSnapshotForAccount("xai", ids[1]!)).resolves.toBeDefined(); + // With the opt-in the request path uses, it is rejected and the caller falls back. + await expect( + getValidAccessSnapshotForAccount("xai", ids[1]!, { requireUsableAccount: true }), + ).rejects.toThrow(); + } finally { + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + rmSync(home, { recursive: true, force: true }); + } + }); +});