From f43cb41be903eefae27a5b43a5eb9e40fdaa4481 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:12:11 +0900 Subject: [PATCH 1/9] feat(oauth): open a turn on the account with known headroom Rotation only reacted to a 429, so a turn could still open on an account a previous probe measured as spent, costing a round trip and one of three rotations to rediscover it. Initial resolution now prefers the ranked account when there is evidence, and returns null otherwise so every provider without per-account quota keeps its current resolution. --- src/oauth/generic-account-failover.ts | 40 +++++++++ src/server/responses/core.ts | 18 +++- tests/kiro-pool-rank.test.ts | 114 ++++++++++++++++++++++++++ 3 files changed, 171 insertions(+), 1 deletion(-) diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index cec3e6321d..32ebd50c94 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -211,6 +211,46 @@ 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; + const set = getAccountSet(providerName); + if (!set || set.accounts.length < 2) return null; + + const active = set.activeAccountId; + // 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 = eligibleFailoverAccounts(providerName, now); + if (eligible.length === 0) return null; + + // Start the ring at the active account so an unranked outcome reproduces today's choice. + const order = set.accounts.map(account => account.id); + 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. + 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); diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 691d75ada4..25f88a68a2 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, @@ -124,6 +125,7 @@ import { GENERIC_OAUTH_MAX_FAILOVERS_PER_REQUEST, isGenericFailoverProvider, isGenericOAuthFailoverEnabled, + preferredInitialAccount, rotateGenericOAuthAccountOn429, } from "../../oauth/generic-account-failover"; import { resolveCopilotApiBaseUrl } from "../../oauth/github-copilot"; @@ -3099,7 +3101,21 @@ 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. + const resolved = preferredAccountId + ? await getValidAccessSnapshotForAccount(route.providerName, preferredAccountId) + : await getValidAccessTokenSnapshot(route.providerName); replayOAuthCredentialSnapshot = { accountId: resolved.accountId, generation: resolved.generation, diff --git a/tests/kiro-pool-rank.test.ts b/tests/kiro-pool-rank.test.ts index 4cef86a93f..ea6dee65ff 100644 --- a/tests/kiro-pool-rank.test.ts +++ b/tests/kiro-pool-rank.test.ts @@ -1,5 +1,15 @@ 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, + preferredInitialAccount, + rotateGenericOAuthAccountOn429, +} from "../src/oauth/generic-account-failover"; +import { getAccountSet, saveCredential, setActiveAccount } from "../src/oauth/store"; +import type { OcxConfig, OcxProviderConfig } from "../src/types"; import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests, @@ -115,3 +125,107 @@ 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 }); + } + }); +}); From db21198c69b5ad9105f57052f0579ad208aeb840 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:19:08 +0900 Subject: [PATCH 2/9] fix(oauth): close three defects in pre-dispatch account selection Review found: an Antigravity account could receive its own bearer beside the configured account's project because the ordinary path only fills an EMPTY project; a quota-less provider could still be redirected once cooling the active account collapsed the candidate list to one; and selection read the uncached credential store twice per request. Evidence is now checked across the whole roster, the roster is read behind the existing presence TTL, and the CCA project is replaced or the request fails closed. --- src/oauth/account-quota-rank.ts | 13 ++++++ src/oauth/generic-account-failover.ts | 60 ++++++++++++++++++++++++--- src/server/responses/core.ts | 21 ++++++++-- tests/generic-oauth-failover.test.ts | 14 +++++++ tests/kiro-pool-rank.test.ts | 49 ++++++++++++++++++++++ 5 files changed, 148 insertions(+), 9 deletions(-) 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 32ebd50c94..12c3e3451a 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. @@ -229,18 +266,27 @@ export function preferredInitialAccount( now = Date.now(), ): string | null { if (!isGenericOAuthFailoverEnabled(config, providerName)) return null; - const set = getAccountSet(providerName); - if (!set || set.accounts.length < 2) 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; - const active = set.activeAccountId; // 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 = eligibleFailoverAccounts(providerName, now); + 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 order = set.accounts.map(account => account.id); 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)); @@ -269,9 +315,11 @@ 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/server/responses/core.ts b/src/server/responses/core.ts index 25f88a68a2..e194625d0e 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -3139,9 +3139,24 @@ 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. + if (preferredAccountId) { + if (!resolved.projectId) { + return formatErrorResponse( + 401, + "authentication_error", + "Selected OAuth account has no Cloud Code Assist project", + ); + } + 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..f2f95f3fc0 100644 --- a/tests/generic-oauth-failover.test.ts +++ b/tests/generic-oauth-failover.test.ts @@ -280,6 +280,20 @@ 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 + 4000); + expect(region).toContain("if (preferredAccountId) {"); + expect(region).toContain("resolved.projectId"); + // ...and it fails closed rather than dispatching with no project at all. + expect(region).toContain("has no Cloud Code Assist project"); + }); }); /** diff --git a/tests/kiro-pool-rank.test.ts b/tests/kiro-pool-rank.test.ts index ea6dee65ff..54ba5320ef 100644 --- a/tests/kiro-pool-rank.test.ts +++ b/tests/kiro-pool-rank.test.ts @@ -228,4 +228,53 @@ describe("pre-dispatch account preference", () => { 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("repeated calls inside the TTL do not re-read 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. Rather than measure atime (which noatime mounts make vacuous), remove the + // store after one warm call: a cached selection still answers, an uncached one cannot. + 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]); + rmSync(join(home, "auth.json"), { force: true }); + for (let i = 0; i < 4; i++) 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 }); + } + }); }); From 68e6bdb65ee469adf92155ece6362f12b249c5d4 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:20:28 +0900 Subject: [PATCH 3/9] fix(oauth): fall back instead of failing when a preferred account has no CCA project The fail-closed 401 turned a preference into a denial: Antigravity tolerates project discovery failing, so an account with no project is an ordinary stored state and choosing it must never break a request that would otherwise have worked. --- src/server/responses/core.ts | 23 ++++++++++++++--------- tests/generic-oauth-failover.test.ts | 11 +++++++---- 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index e194625d0e..b8a94f4ca5 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -3113,9 +3113,19 @@ async function handleResponsesInner( // 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. - const resolved = preferredAccountId + let usedPreferredAccount = preferredAccountId !== null; + let resolved = preferredAccountId ? await getValidAccessSnapshotForAccount(route.providerName, preferredAccountId) : 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, @@ -3145,14 +3155,9 @@ async function handleResponsesInner( // 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. - if (preferredAccountId) { - if (!resolved.projectId) { - return formatErrorResponse( - 401, - "authentication_error", - "Selected OAuth account has no Cloud Code Assist project", - ); - } + // 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 }; diff --git a/tests/generic-oauth-failover.test.ts b/tests/generic-oauth-failover.test.ts index f2f95f3fc0..86ce056815 100644 --- a/tests/generic-oauth-failover.test.ts +++ b/tests/generic-oauth-failover.test.ts @@ -289,10 +289,13 @@ describe("sidecar on429 wiring", () => { const start = coreSource.indexOf("const preferredAccountId ="); expect(start).toBeGreaterThan(-1); const region = coreSource.slice(start, start + 4000); - expect(region).toContain("if (preferredAccountId) {"); - expect(region).toContain("resolved.projectId"); - // ...and it fails closed rather than dispatching with no project at all. - expect(region).toContain("has no Cloud Code Assist project"); + expect(region).toContain("usedPreferredAccount && resolved.projectId"); + // ...and 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"); }); }); From 209bb115b85f75586ddbf2c42344befa72e641a5 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:29:10 +0900 Subject: [PATCH 4/9] fix(oauth): degrade to the active account when a preferred one vanishes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The roster is cached for a short window, so an account can be removed or flagged for reauth after being chosen. Resolving it then throws and the request 401s, even though the active account is still valid — a preference that breaks a working request is worse than no preference. The stale roster is dropped and resolution retried on the active account. --- src/oauth/generic-account-failover.ts | 6 +++++ src/server/responses/core.ts | 21 ++++++++++++--- tests/generic-oauth-failover.test.ts | 15 +++++++---- tests/kiro-pool-rank.test.ts | 39 ++++++++++++++++++++++++++- 4 files changed, 72 insertions(+), 9 deletions(-) diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 12c3e3451a..1b0972e4d8 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -310,6 +310,12 @@ 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) { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index b8a94f4ca5..98384cc499 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -122,6 +122,7 @@ import { import { stampOAuthAccountLabel } from "../../providers/label"; import { failoverAccountSnapshot, + forgetGenericFailoverRoster, GENERIC_OAUTH_MAX_FAILOVERS_PER_REQUEST, isGenericFailoverProvider, isGenericOAuthFailoverEnabled, @@ -3114,9 +3115,23 @@ async function handleResponsesInner( // 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 = preferredAccountId - ? await getValidAccessSnapshotForAccount(route.providerName, preferredAccountId) - : await getValidAccessTokenSnapshot(route.providerName); + let resolved: OAuthAccessSnapshot; + if (preferredAccountId) { + try { + resolved = await getValidAccessSnapshotForAccount(route.providerName, preferredAccountId); + } 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 diff --git a/tests/generic-oauth-failover.test.ts b/tests/generic-oauth-failover.test.ts index 86ce056815..bf17cac044 100644 --- a/tests/generic-oauth-failover.test.ts +++ b/tests/generic-oauth-failover.test.ts @@ -288,14 +288,19 @@ describe("sidecar on429 wiring", () => { // 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 + 4000); + const region = coreSource.slice(start, start + 6000); expect(region).toContain("usedPreferredAccount && resolved.projectId"); - // ...and 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. + // 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 54ba5320ef..58c953ecb3 100644 --- a/tests/kiro-pool-rank.test.ts +++ b/tests/kiro-pool-rank.test.ts @@ -5,10 +5,12 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { clearGenericFailoverHealth, + forgetGenericFailoverRoster, preferredInitialAccount, rotateGenericOAuthAccountOn429, } from "../src/oauth/generic-account-failover"; -import { getAccountSet, saveCredential, setActiveAccount } from "../src/oauth/store"; +import { getAccountSet, removeAccount, saveCredential, setActiveAccount } from "../src/oauth/store"; +import { getValidAccessSnapshotForAccount } from "../src/oauth"; import type { OcxConfig, OcxProviderConfig } from "../src/types"; import { clearAccountQuotaCache, @@ -277,4 +279,39 @@ describe("pre-dispatch account preference", () => { 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]!); + // Still returned: the roster read predates the removal. + expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + // ...and resolving it fails, which is the condition the request path must absorb. + await expect(getValidAccessSnapshotForAccount("xai", ids[1]!)).rejects.toThrow(); + + // After the request path forgets the stale roster, selection stops naming it. + forgetGenericFailoverRoster("xai"); + expect(preferredInitialAccount(config, "xai")).toBeNull(); + 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 }); + } + }); }); From f121bcc916a48ed619318b4bfad7732c960ab7d8 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:34:07 +0900 Subject: [PATCH 5/9] fix(oauth): re-check the chosen account against the live store before redirecting A removed account makes resolution throw, which the caller absorbs, but an account newly flagged needsReauth still has a readable credential: resolution succeeds, no error path fires, and the request dispatches on an account already known to need a fresh login. Only the ranking winner is re-read, so the cost is one store read per actual redirection rather than one per request. --- src/oauth/generic-account-failover.ts | 15 ++++++- tests/kiro-pool-rank.test.ts | 62 +++++++++++++++++++++------ 2 files changed, 62 insertions(+), 15 deletions(-) diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 1b0972e4d8..245605776b 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -294,7 +294,20 @@ export function preferredInitialAccount( const best = rankAccountsByHeadroom(providerName, candidates)[0] ?? null; // Nothing to do when the ranking agrees with the account we would have used anyway. - return best && best !== active ? best : null; + if (!best || best === active) return null; + // The roster above may be up to PRESENCE_CACHE_TTL_MS old, and two kinds of staleness + // matter here. A REMOVED account makes resolution throw, which the caller absorbs. An + // account newly flagged `needsReauth` does NOT throw — its stored credential is still + // readable, so the request would dispatch on an account already known to need a fresh + // login while a healthy active account sat unused. Only the winner is re-checked, so + // this costs one store read per ACTUAL redirection rather than one per request. + const fresh = getAccountSet(providerName)?.accounts.find(account => account.id === best); + if (!fresh || fresh.needsReauth === true) { + roster.delete(providerName); + presence.delete(providerName); + return null; + } + return best; } /** Earliest remaining cooldown, for a client-facing Retry-After when every account is cooled. */ diff --git a/tests/kiro-pool-rank.test.ts b/tests/kiro-pool-rank.test.ts index 58c953ecb3..4b164354d8 100644 --- a/tests/kiro-pool-rank.test.ts +++ b/tests/kiro-pool-rank.test.ts @@ -9,7 +9,13 @@ import { preferredInitialAccount, rotateGenericOAuthAccountOn429, } from "../src/oauth/generic-account-failover"; -import { getAccountSet, removeAccount, saveCredential, setActiveAccount } from "../src/oauth/store"; +import { + getAccountSet, + markAccountNeedsReauth, + removeAccount, + saveCredential, + setActiveAccount, +} from "../src/oauth/store"; import { getValidAccessSnapshotForAccount } from "../src/oauth"; import type { OcxConfig, OcxProviderConfig } from "../src/types"; import { @@ -254,11 +260,11 @@ describe("pre-dispatch account preference", () => { } }); - test("repeated calls inside the TTL do not re-read the credential store", async () => { + test("a non-redirecting request never 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. Rather than measure atime (which noatime mounts make vacuous), remove the - // store after one warm call: a cached selection still answers, an uncached one cannot. + // request. The common case, where the ranking agrees with the active account, must + // answer entirely from cache. Deleting the store proves it: an uncached path could not. home = mkdtempSync(join(tmpdir(), "ocx-predispatch-")); process.env.OPENCODEX_HOME = home; clearGenericFailoverHealth(); @@ -266,11 +272,12 @@ describe("pre-dispatch account preference", () => { 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]); + // Active already holds the most headroom, so there is nothing to redirect. + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { monthlyPercent: 5, updatedAt: Date.now() }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { monthlyPercent: 95, updatedAt: Date.now() }); + expect(preferredInitialAccount(config, "xai")).toBeNull(); rmSync(join(home, "auth.json"), { force: true }); - for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); + for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBeNull(); } finally { clearGenericFailoverHealth(); clearAccountQuotaCache(); @@ -297,15 +304,42 @@ describe("pre-dispatch account preference", () => { expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); await removeAccount("xai", ids[1]!); - // Still returned: the roster read predates the removal. - expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); - // ...and resolving it fails, which is the condition the request path must absorb. + // A redirection re-checks its winner against the live store, so the removed account + // is never named — the request path keeps the healthy active account. + expect(preferredInitialAccount(config, "xai")).toBeNull(); + // ...and had it been named, resolving it would have thrown; the caller absorbs that + // too, but this is the layer that stops it happening at all. await expect(getValidAccessSnapshotForAccount("xai", ids[1]!)).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]); - // After the request path forgets the stale roster, selection stops naming it. - forgetGenericFailoverRoster("xai"); + await markAccountNeedsReauth("xai", ids[1]!, true); expect(preferredInitialAccount(config, "xai")).toBeNull(); - expect(getAccountSet("xai")?.accounts.map(a => a.id)).toEqual([ids[0]!]); + // Proof the flag alone would not have stopped it: the credential still resolves. + await expect(getValidAccessSnapshotForAccount("xai", ids[1]!)).resolves.toBeDefined(); } finally { clearGenericFailoverHealth(); clearAccountQuotaCache(); From 134efaf0645f7139fb67dd07d02cf50014139e43 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:37:15 +0900 Subject: [PATCH 6/9] fix(oauth): reject an unusable preferred account inside the resolver's own read Validating the ranking winner in the selector meant a second uncached credential-file read on every redirected request, which is the steady state of this feature. The check moves into account-scoped resolution, where the store row is already being read: requireUsableAccount rejects a removed or reauth-flagged account and the caller's existing fallback takes the active one. --- src/oauth/generic-account-failover.ts | 21 +++++--------- src/oauth/index.ts | 16 ++++++++--- src/oauth/store.ts | 16 +++++++++++ src/server/responses/core.ts | 10 ++++++- tests/kiro-pool-rank.test.ts | 41 ++++++++++++++++++--------- 5 files changed, 71 insertions(+), 33 deletions(-) diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 245605776b..a36974a629 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -294,20 +294,13 @@ export function preferredInitialAccount( const best = rankAccountsByHeadroom(providerName, candidates)[0] ?? null; // Nothing to do when the ranking agrees with the account we would have used anyway. - if (!best || best === active) return null; - // The roster above may be up to PRESENCE_CACHE_TTL_MS old, and two kinds of staleness - // matter here. A REMOVED account makes resolution throw, which the caller absorbs. An - // account newly flagged `needsReauth` does NOT throw — its stored credential is still - // readable, so the request would dispatch on an account already known to need a fresh - // login while a healthy active account sat unused. Only the winner is re-checked, so - // this costs one store read per ACTUAL redirection rather than one per request. - const fresh = getAccountSet(providerName)?.accounts.find(account => account.id === best); - if (!fresh || fresh.needsReauth === true) { - roster.delete(providerName); - presence.delete(providerName); - return null; - } - return best; + // + // 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. */ 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 98384cc499..9618e0e538 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -3118,7 +3118,15 @@ async function handleResponsesInner( let resolved: OAuthAccessSnapshot; if (preferredAccountId) { try { - resolved = await getValidAccessSnapshotForAccount(route.providerName, preferredAccountId); + // `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 diff --git a/tests/kiro-pool-rank.test.ts b/tests/kiro-pool-rank.test.ts index 4b164354d8..4dfe1491e2 100644 --- a/tests/kiro-pool-rank.test.ts +++ b/tests/kiro-pool-rank.test.ts @@ -260,11 +260,14 @@ describe("pre-dispatch account preference", () => { } }); - test("a non-redirecting request never touches the credential store", async () => { + 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 common case, where the ranking agrees with the active account, must - // answer entirely from cache. Deleting the store proves it: an uncached path could not. + // 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(); @@ -272,11 +275,16 @@ describe("pre-dispatch account preference", () => { try { const ids = await seedAccounts(2); await setActiveAccount("xai", ids[0]!); - // Active already holds the most headroom, so there is nothing to redirect. + // 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() }); - expect(preferredInitialAccount(config, "xai")).toBeNull(); - rmSync(join(home, "auth.json"), { force: true }); for (let i = 0; i < 4; i++) expect(preferredInitialAccount(config, "xai")).toBeNull(); } finally { clearGenericFailoverHealth(); @@ -304,12 +312,13 @@ describe("pre-dispatch account preference", () => { expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); await removeAccount("xai", ids[1]!); - // A redirection re-checks its winner against the live store, so the removed account - // is never named — the request path keeps the healthy active account. - expect(preferredInitialAccount(config, "xai")).toBeNull(); - // ...and had it been named, resolving it would have thrown; the caller absorbs that - // too, but this is the layer that stops it happening at all. - await expect(getValidAccessSnapshotForAccount("xai", ids[1]!)).rejects.toThrow(); + // 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(); @@ -337,9 +346,13 @@ describe("pre-dispatch account preference", () => { expect(preferredInitialAccount(config, "xai")).toBe(ids[1]); await markAccountNeedsReauth("xai", ids[1]!, true); - expect(preferredInitialAccount(config, "xai")).toBeNull(); - // Proof the flag alone would not have stopped it: the credential still resolves. + // 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(); From 420665c8755546d903b3cbf7fea98d3f4961edc3 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 13:48:11 +0900 Subject: [PATCH 7/9] docs(devlog): record the pre-dispatch selection phase and its five review rounds --- .../090_predispatch_selection.md | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 devlog/_plan/260829_kiro_quota_pool/090_predispatch_selection.md 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. From bbc69ef88e3e1b7682d1ffca6b37698a12a42acc Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 14:00:25 +0900 Subject: [PATCH 8/9] docs(devlog): close the pre-request selection gap in the head-to-head --- .../080_head_to_head_result.md | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md b/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md index 06f555b254..ab8936e823 100644 --- a/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md +++ b/devlog/_plan/260829_kiro_quota_pool/080_head_to_head_result.md @@ -24,30 +24,30 @@ AGPL-3.0 reference clone (commit `474df2b`). Behaviour was studied; no code was Stating this plainly, because a comparison that only lists our wins is worthless. -1. **Pre-request selection.** kiro-lb picks an account *before* dispatch with a weighted - race (`kiro/account_manager.py:1183-1208`). Ours ranks only on the 429 recovery path - (`src/oauth/generic-account-failover.ts`), so the first request of a turn can still - land on a spent account. Deferred by design (doc `060`), not solved. -2. **Persistence across restart.** Its quota rows live in SQLite and seed routing at +1. **Persistence across restart.** Its quota rows live in SQLite and seed routing at startup (`kiro/store.py:206-289`). Our caches are process-local, so a restart forgets every measurement until the next probe. -3. **Operations dashboard.** Request-rate charts, per-model token panels, Prometheus +2. **Operations dashboard.** Request-rate charts, per-model token panels, Prometheus export (`kiro/metrics.py`). We render quota bars and a CLI column. -4. **Account onboarding.** Device login for Builder ID, Google and GitHub straight from +3. **Account onboarding.** Device login for Builder ID, Google and GitHub straight from the dashboard (`kiro/device_login.py`). Ours hands off to the Kiro CLI one account at a time. -Item 1 is the one that matters most for the user's "pool 기반 자동 탑재" ask, and it is -the first follow-up work-phase. +**Closed since this was written:** pre-request selection. kiro-lb picks an account before +dispatch with a weighted race (`kiro/account_manager.py:1183-1208`) and this document +originally recorded that as their lead. Doc `090` implements it on our side +(`preferredInitialAccount`), and ours is model-agnostic but evidence-gated and +deterministic where theirs is a random race over stale-by-up-to-15-minutes headroom. ## Honest summary On *correctness of the quota reading* and *safety of the pool machinery* we are ahead: resource selection, overage semantics, trial balances, region validation, credential/route pairing, and stale-state handling are each demonstrably stricter, with tests. On *routing -sophistication* kiro-lb is still ahead on the pre-dispatch axis, and on *operational -surface* it is ahead outright. +sophistication* the two are now comparable — both select before dispatch; ours refuses to +act without evidence, theirs always ranks. On *operational surface* kiro-lb is ahead +outright. -"Better than kiro-lb" is therefore true for the two things this unit set out to do — -display Kiro quota, and make the pool quota-aware — and not yet true as a blanket claim -about the whole gateway. +"Better than kiro-lb" is therefore true for what this unit set out to do — display Kiro +quota, and make the pool quota-aware in both directions — and not a blanket claim about the +whole gateway, which still has a dashboard and restart persistence we do not. From f640bede7de3618487def7a50dd76f4ad621f4e4 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 14:08:14 +0900 Subject: [PATCH 9/9] feat(quota): keep per-account quota across restarts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-memory cache was process-local, so a restart forgot every measurement and the pool opened its next turn blind — the exact state pre-dispatch selection exists to avoid. A small debounced JSON snapshot, modelled on the Codex pool's own, seeds the cache at startup; it stores percentages and reset times only, and a corrupt or stale file loads as empty rather than breaking startup. --- .../100_quota_persistence.md | 59 +++++++++++++ src/providers/account-quota-disk.ts | 79 +++++++++++++++++ src/providers/quota.ts | 35 ++++++++ ...provider-account-quota-persistence.test.ts | 86 +++++++++++++++++++ 4 files changed, 259 insertions(+) create mode 100644 devlog/_plan/260829_kiro_quota_pool/100_quota_persistence.md create mode 100644 src/providers/account-quota-disk.ts create mode 100644 tests/provider-account-quota-persistence.test.ts diff --git a/devlog/_plan/260829_kiro_quota_pool/100_quota_persistence.md b/devlog/_plan/260829_kiro_quota_pool/100_quota_persistence.md new file mode 100644 index 0000000000..6aa38e0761 --- /dev/null +++ b/devlog/_plan/260829_kiro_quota_pool/100_quota_persistence.md @@ -0,0 +1,59 @@ +# 100 — Work-phase 4: quota persistence across restart + +Doc `080` listed three axes where kiro-lb was ahead. Doc `090` closed pre-request +selection. This closes the second: kiro-lb persists quota rows in SQLite and seeds +routing from them at startup (`kiro/store.py:206-289`), while our caches were +process-local — a restart forgot every measurement. + +## Why it matters more now than it did before + +Before pre-dispatch selection, forgetting quota only meant an empty dashboard until the +next probe. Now it means the pool opens its first turn after every restart with no idea +which account has room — precisely the blindness `090` exists to remove. Persistence is +what makes that feature survive a restart rather than warm up from scratch. + +## Design + +`src/providers/account-quota-disk.ts`, modelled directly on the Codex pool's own +snapshot (`src/codex/quota.ts`) rather than inventing a second shape: + +- A single JSON file under `OPENCODEX_HOME`, written atomically, debounced 250ms. +- Keyed exactly like the in-memory cache, so hydration is a direct fill. +- Six-hour maximum age on load. A stale bar is still useful for ORDERING — a wrong + guess costs one 429 that rotation already handles — but a day-old reading of a + monthly window should not outrank a fresh probe. +- Percentages and reset timestamps only. No token, no email, no label; the account id + is the store's own opaque id, which already keys the in-memory cache. +- Corrupt, missing, or future-version files load as empty. A cache must never be able + to break startup. + +Hydration is lazy and once-only, on the first cached read. `clearAccountQuotaCache()` +resets the hydration flag and cancels any pending write, so a cleared cache cannot be +re-seeded from the file it was just cleared of. + +## Accept criteria + +| # | Scenario | Observable proof | +| --- | --- | --- | +| 1 | Write then read in a fresh process | the percentage survives | +| 2 | Snapshot older than six hours | discarded, not loaded | +| 3 | Corrupt JSON | loads empty, does not throw | +| 4 | `version: 2` file | ignored | +| 5 | No file | not an error | +| 6 | Written file inspected | contains percentages; contains no token, email, ARN or secret | +| 7 | Five writes in a burst | one file write, last value wins | +| 8 | Cancelled write | no file created | + +## Verification + +```text +bun x tsc --noEmit -> exit 0 +bun run privacy:scan -> Privacy scan passed +bun test (8 files) -> 208 pass / 0 fail / 634 expect() calls +``` + +## What remains kiro-lb's + +One axis from doc `080`: the operations dashboard — request-rate charts, per-model token +panels, Prometheus export. That is a product surface, not pool machinery, and it is +outside this unit's objective. diff --git a/src/providers/account-quota-disk.ts b/src/providers/account-quota-disk.ts new file mode 100644 index 0000000000..ad6d65982d --- /dev/null +++ b/src/providers/account-quota-disk.ts @@ -0,0 +1,79 @@ +/** + * Last-known per-account provider quota, kept across restarts. + * + * The in-memory cache is process-local, so a restart forgets every measurement and the + * pool opens the next turn with no idea which account has room — exactly the state + * pre-dispatch selection exists to avoid. Codex solved this for its own pool with a small + * disk snapshot (`codex/quota.ts`), and this is the same shape for provider accounts. + * + * What is written is percentages and reset timestamps. No token, no email, no account + * label — the account id is the store's own opaque id, which is already what keys the + * in-memory cache. + */ +import { existsSync, readFileSync } from "node:fs"; +import { join } from "node:path"; +import { atomicWriteFile, getConfigDir } from "../config"; +import type { ProviderQuota } from "./quota-types"; + +const FILENAME = "provider-account-quota-cache.json"; + +/** + * Older than this and the snapshot is discarded on load. + * + * Six hours matches the Codex cache. A stale bar is still better than none for ORDERING — + * it decides which account to try first, and a wrong guess costs one 429 that rotation + * already handles — but a day-old reading of a monthly window has drifted far enough that + * it should not outrank a fresh probe. + */ +const DISK_MAX_AGE_MS = 6 * 60 * 60_000; +const PERSIST_DEBOUNCE_MS = 250; + +type DiskFile = { + version: 1; + /** provider\u0000accountId -> quota, the same key the in-memory cache uses. */ + rows: Record; +}; + +let persistTimer: ReturnType | null = null; + +/** Read the snapshot. Returns an empty map for a missing, corrupt or stale file. */ +export function readPersistedAccountQuotas(now = Date.now()): Map { + const rows = new Map(); + try { + const path = join(getConfigDir(), FILENAME); + if (!existsSync(path)) return rows; + const parsed = JSON.parse(readFileSync(path, "utf8")) as DiskFile; + if (!parsed || parsed.version !== 1 || !parsed.rows || typeof parsed.rows !== "object") return rows; + for (const [key, quota] of Object.entries(parsed.rows)) { + if (!quota || typeof quota !== "object" || typeof quota.updatedAt !== "number") continue; + if (now - quota.updatedAt > DISK_MAX_AGE_MS) continue; + rows.set(key, quota); + } + } catch { + // A corrupt cache must never block routing or the dashboard. + } + return rows; +} + +/** Write the snapshot, debounced. Best-effort: a failed write is not an error. */ +export function schedulePersistAccountQuotas(rows: () => Iterable<[string, ProviderQuota]>): void { + if (persistTimer) clearTimeout(persistTimer); + persistTimer = setTimeout(() => { + persistTimer = null; + try { + const out: Record = {}; + for (const [key, quota] of rows()) out[key] = quota; + const body: DiskFile = { version: 1, rows: out }; + atomicWriteFile(join(getConfigDir(), FILENAME), `${JSON.stringify(body)}\n`); + } catch { + // Best-effort persistence only. + } + }, PERSIST_DEBOUNCE_MS); +} + +/** Test seam: drop any pending write so a suite cannot leak one into the next file. */ +export function cancelPendingAccountQuotaPersist(): void { + if (!persistTimer) return; + clearTimeout(persistTimer); + persistTimer = null; +} diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 55ff088d27..ee472bfa02 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -53,6 +53,11 @@ import { kiroUsageContextForAccount, reconcileKiroAccountUsageState, } from "./kiro-usage"; +import { + cancelPendingAccountQuotaPersist, + readPersistedAccountQuotas, + schedulePersistAccountQuotas, +} from "./account-quota-disk"; export type { ProviderQuota, ProviderQuotaCreditsUsd, ProviderQuotaWindow } from "./quota-types"; @@ -1378,6 +1383,31 @@ type AccountQuotaCacheEntry = { unavailable?: true; }; const accountQuotaCache = new Map(); + +/** + * Seed the cache from the last run, once. + * + * Without this a restart forgets every measurement, so the pool opens its next turn with + * no idea which account has room — the exact blindness pre-dispatch selection exists to + * remove. A hydrated row is still subject to the ordinary TTL, so it orders the first + * request and is replaced by a live probe immediately after. + */ +let diskHydrated = false; +function hydrateAccountQuotaCache(): void { + if (diskHydrated) return; + diskHydrated = true; + for (const [key, quota] of readPersistedAccountQuotas()) { + if (!accountQuotaCache.has(key)) accountQuotaCache.set(key, { ts: quota.updatedAt, quota }); + } +} + +function persistAccountQuotaCache(): void { + schedulePersistAccountQuotas(function* () { + for (const [key, entry] of accountQuotaCache) { + if (entry.quota) yield [key, entry.quota] as [string, ProviderQuota]; + } + }); +} const accountQuotaInflight = new Map>(); let lastReconciledGeneration = 0; let liveAccountQuotaKeys = new Set(); @@ -1475,6 +1505,10 @@ export function clearAccountQuotaCache(provider?: string): void { accountQuotaCache.clear(); accountQuotaInflight.clear(); clearKiroAccountUsageState(); + // A cleared cache must not be re-seeded from the file it was just cleared of, and any + // pending write of the old rows is abandoned. + diskHydrated = false; + cancelPendingAccountQuotaPersist(); return; } const prefix = `${provider}\u0000`; @@ -1486,6 +1520,7 @@ export function clearAccountQuotaCache(provider?: string): void { for (const key of [...accountQuotaInflight.keys()]) { if (key.startsWith(prefix)) accountQuotaInflight.delete(key); } + persistAccountQuotaCache(); } /** diff --git a/tests/provider-account-quota-persistence.test.ts b/tests/provider-account-quota-persistence.test.ts new file mode 100644 index 0000000000..fd98cf141f --- /dev/null +++ b/tests/provider-account-quota-persistence.test.ts @@ -0,0 +1,86 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + cancelPendingAccountQuotaPersist, + readPersistedAccountQuotas, + schedulePersistAccountQuotas, +} from "../src/providers/account-quota-disk"; +import type { ProviderQuota } from "../src/providers/quota-types"; + +const previousHome = process.env.OPENCODEX_HOME; +let home: string; +const FILE = "provider-account-quota-cache.json"; +const KEY = "kiro\u0000acct-a"; + +const settle = () => new Promise(resolve => setTimeout(resolve, 400)); + +beforeEach(() => { + home = mkdtempSync(join(tmpdir(), "ocx-quota-persist-")); + process.env.OPENCODEX_HOME = home; +}); + +afterEach(() => { + cancelPendingAccountQuotaPersist(); + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + rmSync(home, { recursive: true, force: true }); +}); + +const quota = (percent: number, updatedAt = Date.now()): ProviderQuota => ({ + monthlyPercent: percent, + updatedAt, +}); + +describe("provider account quota persistence", () => { + test("rows survive a restart", async () => { + schedulePersistAccountQuotas(() => [[KEY, quota(15)]]); + await settle(); + expect(readPersistedAccountQuotas().get(KEY)?.monthlyPercent).toBe(15); + }); + + test("a stale snapshot is discarded rather than ordering the pool on old data", async () => { + schedulePersistAccountQuotas(() => [[KEY, quota(15, Date.now() - 7 * 60 * 60_000)]]); + await settle(); + expect(readPersistedAccountQuotas().size).toBe(0); + }); + + test("a corrupt file is ignored instead of breaking startup", () => { + writeFileSync(join(home, FILE), "{ not json"); + expect(readPersistedAccountQuotas().size).toBe(0); + }); + + test("a file from a future schema version is ignored", () => { + writeFileSync(join(home, FILE), JSON.stringify({ version: 2, rows: { [KEY]: quota(5) } })); + expect(readPersistedAccountQuotas().size).toBe(0); + }); + + test("a missing file is not an error", () => { + expect(existsSync(join(home, FILE))).toBe(false); + expect(readPersistedAccountQuotas().size).toBe(0); + }); + + test("the snapshot carries percentages only, never credentials or identities", async () => { + schedulePersistAccountQuotas(() => [[KEY, quota(15)]]); + await settle(); + const raw = readFileSync(join(home, FILE), "utf8"); + expect(raw).toContain("monthlyPercent"); + for (const forbidden of ["access", "refresh", "Bearer", "@", "profileArn", "clientSecret"]) { + expect(raw).not.toContain(forbidden); + } + }); + + test("writes are debounced so a burst costs one file write", async () => { + for (let i = 0; i < 5; i++) schedulePersistAccountQuotas(() => [[KEY, quota(i)]]); + await settle(); + expect(readPersistedAccountQuotas().get(KEY)?.monthlyPercent).toBe(4); + }); + + test("a cancelled write leaves no file behind", async () => { + schedulePersistAccountQuotas(() => [[KEY, quota(15)]]); + cancelPendingAccountQuotaPersist(); + await settle(); + expect(existsSync(join(home, FILE))).toBe(false); + }); +});