From 0d38f0cae6ed76da742cf484147f6d34aebe1ceb Mon Sep 17 00:00:00 2001 From: JUN Date: Thu, 15 Oct 2026 16:44:07 +0900 Subject: [PATCH 1/2] refactor(oauth): share conversation identity extraction --- src/oauth/anthropic-routing.ts | 22 +--------------------- src/oauth/session-affinity-key.ts | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 21 deletions(-) create mode 100644 src/oauth/session-affinity-key.ts diff --git a/src/oauth/anthropic-routing.ts b/src/oauth/anthropic-routing.ts index db6ce859104..c9de9371896 100644 --- a/src/oauth/anthropic-routing.ts +++ b/src/oauth/anthropic-routing.ts @@ -19,7 +19,6 @@ import { clearAllAnthropicFamilyQuota, anthropicModelQuotaFor, anthropicModelPer */ import { clearAllAnthropicRatePauses, anthropicRatePolicyFor, classifyAnthropic429, anthropicRetryAfterMs } from "./anthropic-rate-limit-policy"; import { REFRESH_SKEW_MS } from "./refresh-policy"; -import { createHash } from "node:crypto"; import { captureOAuthAccountSelection, commitOAuthAccountSelection, credentialGeneration, getAccountSet, getAccountCredential, getAccountCredentialWithStatus } from "./store"; import type { OAuthAccessSnapshot } from "./index"; import { getCachedProviderAccountQuota } from "../providers/quota"; @@ -273,26 +272,7 @@ export function formatAnthropicProviderForLog( return `${providerName}-${formatAnthropicAccountOrdinal(accountId)}`; } -export function anthropicSessionKeyFromParts(input: { - sessionIdHeader?: string | null; - threadIdHeader?: string | null; - promptCacheKey?: string | null; - clientThreadId?: string | null; - /** When true, prompt_cache_key is a shared Desktop cohort — ignore it for affinity. */ - promptCacheKeyIsSharedCohort?: boolean; -}): string | null { - const preferred = input.clientThreadId?.trim() - || input.sessionIdHeader?.trim() - || input.threadIdHeader?.trim() - || ""; - if (preferred) { - return preferred.length <= 128 ? preferred : createHash("sha256").update(preferred).digest("hex"); - } - if (input.promptCacheKeyIsSharedCohort) return null; - const cacheKey = input.promptCacheKey?.trim() ?? ""; - if (!cacheKey) return null; - return cacheKey.length <= 128 ? cacheKey : createHash("sha256").update(cacheKey).digest("hex"); -} +export { oauthSessionKeyFromParts as anthropicSessionKeyFromParts } from "./session-affinity-key"; export type AnthropicCooldownRecoveryClaim = Readonly<{ diff --git a/src/oauth/session-affinity-key.ts b/src/oauth/session-affinity-key.ts new file mode 100644 index 00000000000..a97b6e38d3f --- /dev/null +++ b/src/oauth/session-affinity-key.ts @@ -0,0 +1,23 @@ +import { createHash } from "node:crypto"; + +/** Shared conversation identity precedence; Desktop cohort cache keys never bind sessions. */ +export function oauthSessionKeyFromParts(input: { + sessionIdHeader?: string | null; + threadIdHeader?: string | null; + promptCacheKey?: string | null; + clientThreadId?: string | null; + /** When true, prompt_cache_key is a shared Desktop cohort — ignore it for affinity. */ + promptCacheKeyIsSharedCohort?: boolean; +}): string | null { + const preferred = input.clientThreadId?.trim() + || input.sessionIdHeader?.trim() + || input.threadIdHeader?.trim() + || ""; + if (preferred) { + return preferred.length <= 128 ? preferred : createHash("sha256").update(preferred).digest("hex"); + } + if (input.promptCacheKeyIsSharedCohort) return null; + const cacheKey = input.promptCacheKey?.trim() ?? ""; + if (!cacheKey) return null; + return cacheKey.length <= 128 ? cacheKey : createHash("sha256").update(cacheKey).digest("hex"); +} From 6a8556ca6f7b0cea99b0c08a02efd4d0a9985118 Mon Sep 17 00:00:00 2001 From: JUN Date: Thu, 15 Oct 2026 16:44:22 +0900 Subject: [PATCH 2/2] feat(oauth): retain generic pool conversation affinity through recovery Keep automatic affinity-mode choices request-local while preserving the operator selection revision and credential generation fences. Rebind admitted recovery accounts, retain configured new-session strategies, and bound process-local ownership by home, provider and caller. Refs #3375. --- .../docs/reference/configuration/providers.md | 24 +- scripts/test-layout/layout.json | 2 + src/oauth/generic-account-failover.ts | 17 +- src/oauth/generic-session-affinity.ts | 125 ++++++ src/server/responses/request-transport.ts | 57 ++- src/types/provider.ts | 2 + structure/INDEX.md | 7 +- structure/manifest.json | 7 + structure/providers-and-adapters.md | 2 +- structure/providers/anthropic-account-pool.md | 2 + structure/providers/generic-account-pool.md | 25 ++ structure/runtime.md | 2 +- structure/transports/inventory.md | 2 + tests/fixtures/test-layout-expected.json | 2 + tests/oauth/generic-session-affinity.test.ts | 186 +++++++++ .../server-generic-oauth-affinity.test.ts | 386 ++++++++++++++++++ 16 files changed, 826 insertions(+), 22 deletions(-) create mode 100644 src/oauth/generic-session-affinity.ts create mode 100644 structure/providers/generic-account-pool.md create mode 100644 tests/oauth/generic-session-affinity.test.ts create mode 100644 tests/server/server-generic-oauth-affinity.test.ts diff --git a/docs-site/src/content/docs/reference/configuration/providers.md b/docs-site/src/content/docs/reference/configuration/providers.md index 50c3520b0df..32115418b6f 100644 --- a/docs-site/src/content/docs/reference/configuration/providers.md +++ b/docs-site/src/content/docs/reference/configuration/providers.md @@ -1137,6 +1137,7 @@ second account. | --- | --- | --- | --- | | `oauthAccountFailover.enabled?` | `boolean` | presence-driven | Global override for the **pre-dispatch account preference** only. `false` stops a healthy request being steered toward the account with more known headroom. It does **not** disable 429 rotation. | | `providers..oauthAccountFailover.enabled?` | `boolean` | inherits | Per-provider override for the same preference; beats the global setting in either direction. `false` declines the preference for this provider even when the global setting is `true`, and `true` opts this provider in even when the global setting is `false`. Reactive 429 rotation is unaffected either way. | +| `providers..oauthAccountFailover.sessionAffinity?` | `"conversation" \| "none"` | `"none"` | Opt-in conversation pinning for generic OAuth pools. Requires proactive `enabled: true` (provider or inherited global value). Existing sessions keep their account ahead of the new-session strategy until it becomes unavailable. | | `providers..oauthAccountFailover.strategy?` | `"quota" \| "round-robin" \| "fill-first" \| "least-loaded"` (Kiro only) | — | Kiro's `least-loaded` picks the eligible account with the fewest in-flight requests, using account order for ties. It requires `pool.kernel` and proactive preference; other strategies retain their existing rules. | | `providers.kiro.oauthAccountFailover.maxConcurrentPerAccount?` | integer 1–100 | unlimited | A process-local per-account queue holds a slot through response completion or cancellation. A full selected account waits at most 250 ms, then returns 503 `account_capacity` with `Retry-After: 1`. The cap does not move the request to another account. | | `providers..oauthAccountFailover.autoSwitchThreshold?` | `number` | `80` | 0–100 usage percent at which `fill-first` advances off the active account (#695). Set with `ocx account auto-switch threshold `. Read only under `pool.kernel` with `strategy: "fill-first"`; an account with no measured usage counts as under the threshold. | @@ -1166,9 +1167,26 @@ and the `ocx account strategy` / `ocx account auto-switch` / `ocx account sticky preference. `quotaWindow` is not part of the generic contract. Codex (`/api/codex-auth`) and Anthropic (`anthropicAccountPool`) keep their own contracts unchanged. -Deliberately narrower than `anthropicAccountPool`: no session affinity, no quota-ranked -selection, no probe leases. It answers one question — the account that just returned 429 is -cooled, is there another one available. +Set `providers..oauthAccountFailover.sessionAffinity` to `"conversation"` together with +`enabled: true` to keep a conversation on its admitted account. Session identity uses the client +thread marker, session headers, thread header, then `prompt_cache_key`; a shared Desktop cache +cohort is excluded. Bindings are isolated by provider, configuration home and authenticated caller +principal. Unauthenticated local callers share the local scope. No session marker means no binding. + +New sessions use the existing strategy (`pool.kernel` still gates round-robin/fill-first). +Existing sessions stay pinned even when another account has more quota. A pause, reauthentication +requirement, family cooldown or incompatible model roster releases the binding; existing bounded +failure recovery binds the admitted replacement. Kiro's concurrency cap still waits or refuses +rather than moving a healthy bound conversation merely to balance load. + +In this mode automatic choices, including unkeyed requests and failovers, stay request-local: +the dashboard's selected account remains the operator's selection. Request history still attributes +the serving account. Manual selection, same-account reselection and roster/policy revision changes +invalidate provider bindings, including writes from another process. An ordinary refresh keeps the +same account. Fill-first retains its current account in memory without changing the saved selection. +Bindings expire after 24 hours without admission, are bounded to 2,000 entries process-wide +(including fill-first cursors), and are lost on restart. Set `sessionAffinity: "none"` or disable +proactive preference to release them. Codex and Anthropic retain their own affinity contracts. The Codex pool and the Anthropic pool are excluded and keep their own rotation; enabling this changes neither. A provider with a single stored account is a strict no-op, and no cooldown is diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 16227fe5cf1..f555fd7ce0d 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -919,6 +919,8 @@ "antigravity-account-roster.test.ts": "adapters/google", "gemini-inline.test.ts": "images", "gemini-web-search.test.ts": "adapters/google", + "generic-session-affinity.test.ts": "oauth", + "server-generic-oauth-affinity.test.ts": "server", "generic-oauth-failover.test.ts": "oauth", "generic-oauth-refusal-composition.test.ts": "oauth", "github-copilot-auto.test.ts": "providers/github-copilot", "github-copilot-auto-runtime.test.ts": "providers/github-copilot", "github-copilot-catalog-hints.test.ts": "providers/github-copilot", "github-copilot-account-origin.test.ts": "providers/github-copilot", diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index e1a1884037b..b107a1dbbdf 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -7,9 +7,9 @@ import { isBuiltinAnthropicInstanceRow } from "../providers/anthropic-instance"; * exists is Anthropic's — behind its own opt-in. So xAI, Cursor, Kimi, GitHub Copilot, * Antigravity and Nous have no recovery path on a 429 even with several accounts logged in. * - * Deliberately narrower than the Anthropic pool: no session affinity, no quota-ranked selection, - * no probe leases. Those carry provider-specific meaning; this module only answers "the account - * that just 429'd is cooled, is there another one we may use". + * Conversation ownership lives in generic-session-affinity.ts; this module owns account + * eligibility, quota-ranked preference and bounded refusal recovery. Probe leases remain + * provider-specific. * * NOT a home for Codex (`codex/routing.ts` owns quota scopes and probe leases) or Anthropic * (`oauth/anthropic-routing.ts` owns affinity and a fail-closed local-cli credential rule). @@ -416,6 +416,7 @@ function pickFillFirstGenericAccount( now: number, requestedModelId?: string | null, accountRows?: ReadonlyMap, + preferBelowThreshold = false, ): string | null { const stableAll = stableGenericRoster(providerName); if (stableAll.length < 2) return null; @@ -431,6 +432,11 @@ function pickFillFirstGenericAccount( } const start = activeId ? stableAll.indexOf(activeId) : -1; const ring = start >= 0 ? [...stableAll.slice(start + 1), ...stableAll.slice(0, start)] : stableAll; + if (preferBelowThreshold) { + const below = ring.find(id => id !== activeId && eligible.has(id) + && !isOverAutoSwitchThreshold(providerName, id, threshold, requestedModelId, accountRows?.get(id))); + if (below) return below; + } for (const id of ring) { if (id !== activeId && eligible.has(id)) return id; } @@ -634,6 +640,7 @@ export function preferredInitialAccount( providerName: string, now = Date.now(), requestedModelId?: string | null, + activeAccountOverride?: string, ): string | null { // The PROACTIVE predicate, not the reactive one: this steers a request upstream has not // refused, so `oauthAccountFailover.enabled: false` must still be able to refuse it. @@ -642,7 +649,7 @@ export function preferredInitialAccount( // active id separately would delay manual selection and account removal. const selected = getAccountSet(providerName); if (!selected) return null; - const active = selected.activeAccountId; + const active = activeAccountOverride ?? selected.activeAccountId; const accountRows = new Map(selected.accounts.map(account => [account.id, account])); const order = selected.accounts .filter(account => account.paused !== true && account.needsReauth !== true) @@ -694,7 +701,7 @@ export function preferredInitialAccount( return picked && picked !== active ? picked : null; } if (strategy === "fill-first") { - const picked = pickFillFirstGenericAccount(config, providerName, active, now, requestedModelId, accountRows); + const picked = pickFillFirstGenericAccount(config, providerName, active, now, requestedModelId, accountRows, activeAccountOverride !== undefined); return picked && picked !== active ? picked : null; } diff --git a/src/oauth/generic-session-affinity.ts b/src/oauth/generic-session-affinity.ts new file mode 100644 index 00000000000..d4b08d5d040 --- /dev/null +++ b/src/oauth/generic-session-affinity.ts @@ -0,0 +1,125 @@ +/** Request-local conversation ownership for opt-in generic OAuth pools. No credentials are retained. */ +import { createHash } from "node:crypto"; +import { getConfigDir } from "../config/paths"; +import { subscribeAccountSelections } from "../lib/account-selection-events"; +import type { OcxConfig } from "../types"; +import { classifyModelFamilyForQuota } from "./account-quota-rank"; +import { eligibleFailoverAccounts, isGenericFailoverProvider, preferKiroModelSupport, preferredInitialAccount } from "./generic-account-failover"; +import type { OAuthAccessSnapshot } from "./index"; +import { captureOAuthAccountSelection, credentialGeneration, getAccountCredentialWithStatus } from "./store"; +import type { OAuthAccountSelection } from "./types"; + +const IDLE_TTL_MS = 24 * 60 * 60_000; +const MAX_ENTRIES = 2_000; +const bindings = new Map(); +const prefix = (provider: string) => `${JSON.stringify([getConfigDir(), provider])}\u0000`; +const key = (provider: string, session: string) => prefix(provider) + "session:" + createHash("sha256").update(session).digest("hex"); +const cursorKey = (provider: string) => prefix(provider) + "fill-first"; + +/** A persisted operator/roster/policy change revokes old session choices before its writer returns. */ +export function clearGenericSessionAffinity(provider?: string): void { + if (!provider) { bindings.clear(); return; } + const scope = prefix(provider); + for (const id of bindings.keys()) if (id.startsWith(scope)) bindings.delete(id); +} +subscribeAccountSelections(event => { if (event.kind === "oauth") clearGenericSessionAffinity(event.provider); }); + +/** Unlike new-session strategy selection, keeping the surviving account needs no 2+ quorum. */ +export function genericSessionAffinityEnabled(config: OcxConfig, provider: string): boolean { + const row = config.providers[provider]; + const enabled = !!row && row.disabled !== true && provider !== "anthropic2" && isGenericFailoverProvider(provider, row) + && row.oauthAccountFailover?.sessionAffinity === "conversation" + && (row.oauthAccountFailover.enabled ?? config.oauthAccountFailover?.enabled) === true; + if (!enabled) clearGenericSessionAffinity(provider); + return enabled; +} + +export function genericSessionAccountEligible( + config: OcxConfig, provider: string, accountId: string, model?: string | null, now = Date.now(), allowActiveFallback = false, +): boolean { + const eligible = eligibleFailoverAccounts(provider, now, classifyModelFamilyForQuota(provider, model)); + // Existing conversation ownership is independent of transient per-account load. + if (preferKiroModelSupport(provider, eligible, model).includes(accountId)) return true; + // Cached automatic exclusions remain preferences when no alternative exists. Preserve the + // ordinary active/singleton attempt and its upstream refusal instead of inventing a conflict. + if (!allowActiveFallback || eligible.length > 0 || captureOAuthAccountSelection(provider)?.accountId !== accountId) return false; + const row = getAccountCredentialWithStatus(provider, accountId); + return !!row && !row.paused && !row.needsReauth; +} + +/** Peek only. The admitted snapshot, never a speculative token lookup, establishes ownership. */ +export function genericSessionAccount( + config: OcxConfig, provider: string, session: string | null, model?: string | null, now = Date.now(), +): string | null { + if (!genericSessionAffinityEnabled(config, provider) || !session) return null; + return readBinding(config, provider, key(provider, session), model, now); +} + +function readBinding(config: OcxConfig, provider: string, id: string, model: string | null | undefined, now: number): string | null { + const bound = bindings.get(id); + if (!bound) return null; + const current = captureOAuthAccountSelection(provider); + if (!current || current.accountId !== bound.selection.accountId || current.revision !== bound.selection.revision + || now - bound.lastUsedAt >= IDLE_TTL_MS || now < bound.lastUsedAt + || !genericSessionAccountEligible(config, provider, bound.accountId, model, now)) { + bindings.delete(id); + return null; + } + return bound.accountId; +} + +/** New sessions use the configured strategy; an ineligible active account cannot defeat release. */ +export function preferredGenericSessionAccount( + config: OcxConfig, provider: string, session: string | null, model?: string | null, now = Date.now(), +): string | null { + if (!genericSessionAffinityEnabled(config, provider)) return null; + const bound = genericSessionAccount(config, provider, session, model, now); + if (bound) return bound; + const cursor = readBinding(config, provider, cursorKey(provider), model, now); + const active = cursor ?? captureOAuthAccountSelection(provider)?.accountId; + const preferred = preferredInitialAccount(config, provider, now, model, active); + if (preferred) return preferred; + if (active && genericSessionAccountEligible(config, provider, active, model, now)) return active; + return preferKiroModelSupport(provider, + eligibleFailoverAccounts(provider, now, classifyModelFamilyForQuota(provider, model)), model, config)[0] + ?? captureOAuthAccountSelection(provider)?.accountId ?? null; +} + +/** + * Validate against the operator's persisted selection without moving it. ALL automatic choices + * in affinity mode use this receipt, including unkeyed requests and failovers, so concurrent + * conversations cannot invalidate each other's selection revision. Manual A→B→A still fences. + */ +export function admitGenericSessionAccount( + config: OcxConfig, provider: string, expected: OAuthAccountSelection, + snapshot: OAuthAccessSnapshot, model?: string | null, modeEnabled = genericSessionAffinityEnabled(config, provider), +): OAuthAccountSelection | null { + if (!modeEnabled || snapshot.provider !== provider) return null; + const current = captureOAuthAccountSelection(provider); + if (!current || current.accountId !== expected.accountId || current.revision !== expected.revision) return null; + const row = getAccountCredentialWithStatus(provider, snapshot.accountId); + if (!row || row.paused || row.needsReauth || credentialGeneration(row.credential) !== snapshot.generation + || !genericSessionAccountEligible(config, provider, snapshot.accountId, model, Date.now(), true)) return null; + return current; +} + +export function bindGenericSessionAccount( + config: OcxConfig, provider: string, session: string | null, accountId: string, model?: string | null, now = Date.now(), +): void { + if (!genericSessionAffinityEnabled(config, provider) + || !genericSessionAccountEligible(config, provider, accountId, model, now)) return; + for (const [id, entry] of bindings) { + if (now - entry.lastUsedAt >= IDLE_TTL_MS || now < entry.lastUsedAt) bindings.delete(id); + } + const selection = captureOAuthAccountSelection(provider); + if (!selection) return; + const keep = (id: string) => { + bindings.delete(id); + bindings.set(id, { accountId, lastUsedAt: now, selection }); + }; + const hit = session && genericSessionAccount(config, provider, session, model, now) === accountId; + if (!hit && config.pool?.kernel === true && config.providers[provider]?.oauthAccountFailover?.strategy === "fill-first") + keep(cursorKey(provider)); + if (session) keep(key(provider, session)); + while (bindings.size > MAX_ENTRIES) bindings.delete(bindings.keys().next().value!); +} diff --git a/src/server/responses/request-transport.ts b/src/server/responses/request-transport.ts index 1b6cb6f4c44..d0740f944e5 100644 --- a/src/server/responses/request-transport.ts +++ b/src/server/responses/request-transport.ts @@ -27,7 +27,6 @@ import { anthropicRoutingFor, AnthropicAccountCooldownError, formatAnthropicProviderForLog, - anthropicSessionKeyFromParts, } from "../../oauth/anthropic-routing"; import { OAuthAccountPausedError, @@ -49,6 +48,9 @@ import { refusalAwareInitialKiroAccount, noteGenericPoolSelection, } from "../../oauth/generic-account-failover"; +import { admitGenericSessionAccount, bindGenericSessionAccount, genericSessionAccount, genericSessionAccountEligible, genericSessionAffinityEnabled, preferredGenericSessionAccount } from "../../oauth/generic-session-affinity"; +import { oauthSessionKeyFromParts } from "../../oauth/session-affinity-key"; +import { resolveContextPrincipal } from "../auth-cors"; import { tryKiroAlternateAfterTerminalRefresh } from "../../oauth/kiro-terminal-failover"; import { classifyModelFamilyForQuota } from "../../oauth/account-quota-rank"; import { expandInferenceOAuthSendBudget } from "../inference/context"; @@ -120,6 +122,7 @@ export async function prepareResponsesTransport( ) { const { config, logCtx, options, req } = requestContext; const { route, parsed, inboundWire, translatorBudget } = requestState; + const genericAffinityMode = genericSessionAffinityEnabled(config, route.providerName); const anthropicInstance = route.provider.authMode === "oauth" ? configuredAnthropicInstance(config, route.providerName) : undefined; // Refuse custom/orphan B OAuth before any generic resolver can observe its auth row. @@ -217,6 +220,8 @@ export async function prepareResponsesTransport( for (let attempt = 0; attempt < maxSelectionAttempts; attempt++) { requireAnthropicAdmission(); if (candidate.provider !== route.providerName) return null; + if (genericAffinityMode && route.provider.googleMode === "cloud-code-assist" && !candidate.projectId) + throw new OAuthLoginRequiredError(route.providerName); if (anthropicInstance) anthropicRouteDecision = currentAnthropicRouteDecision(); if (!oauthSelection) return null; const proactiveEnabled = anthropicInstance !== undefined @@ -234,7 +239,9 @@ export async function prepareResponsesTransport( if (candidate.provider !== route.providerName) return null; if (anthropicInstance) anthropicRouteDecision = currentAnthropicRouteDecision(); if (anthropicRouteDecision && !routeCandidates(anthropicRouting!.getEligibleAnthropicAccounts(Date.now(), route.modelId), anthropicRouteDecision).includes(candidate.accountId)) return null; - const committed = await commitOAuthAccountSelection(route.providerName, candidate.accountId, { + const committed = genericAffinityMode + ? admitGenericSessionAccount(config, route.providerName, oauthSelection, candidate, route.modelId, genericAffinityMode) + : await commitOAuthAccountSelection(route.providerName, candidate.accountId, { expectedSelection: oauthSelection, expectedCredentialGeneration: candidate.generation, requireUsableAccount: true, @@ -248,14 +255,18 @@ export async function prepareResponsesTransport( )) return null; oauthSelection = committed; servingOAuthSnapshot = candidate; + bindGenericSessionAccount(config, route.providerName, genericSessionKey, candidate.accountId, route.modelId); if (anthropicInstance) replayOAuthCredentialSnapshot = { accountId: candidate.accountId, generation: candidate.generation }; forgetGenericFailoverRoster(route.providerName); return candidate; } // A newer manual choice wins over this request's old proposal, including A→B→A. // Resolve that choice, not the rejected candidate, before trying admission again. + const previousSelection = oauthSelection; oauthSelection = captureOAuthAccountSelection(route.providerName); if (!oauthSelection) return null; + const selectionChanged = oauthSelection.accountId !== previousSelection.accountId + || oauthSelection.revision !== previousSelection.revision; // A revision also changes on per-account policy edits. Re-evaluate the selector // after credential waits even without a model route, rather than reusing stale active. const revisedAnthropic = anthropicInstance !== undefined @@ -263,9 +274,15 @@ export async function prepareResponsesTransport( const revisedAnthropicId = revisedAnthropic?.accountId; if (anthropicInstance !== undefined && !revisedAnthropicId) return null; if (revisedAnthropic) anthropicReason = revisedAnthropic.reason; + const genericAccountId = genericAffinityMode && !selectionChanged + ? (genericSessionAccountEligible(config, route.providerName, candidate.accountId, route.modelId, Date.now(), true) + ? candidate.accountId : (preferredGenericSessionAccount(config, route.providerName, genericSessionKey, route.modelId) + ?? oauthSelection.accountId)) + : oauthSelection.accountId; + if (!genericAccountId && !anthropicInstance) return null; candidate = anthropicInstance !== undefined ? await anthropicRouting!.getAnthropicPoolAccessSnapshot(revisedAnthropicId ?? oauthSelection.accountId) - : await getValidAccessSnapshotForAccount(route.providerName, oauthSelection.accountId, { requireUsableAccount: true }); + : await getValidAccessSnapshotForAccount(route.providerName, genericAccountId!, { requireUsableAccount: true }); if (route.provider.googleMode === "cloud-code-assist" && !candidate.projectId) return null; } return null; @@ -275,7 +292,10 @@ export async function prepareResponsesTransport( // Keyed on the account, not the selection revision: a selection that moved away and back // (A -> B -> A) before the 401 landed still serves the rejected credential, and skipping // the refresh would replay it and spend the one recovery attempt. - const candidate = current?.accountId === sent.accountId ? await forceRefreshOAuthAccessSnapshot(sent) : sent; + const candidate = (current?.accountId === sent.accountId + || (genericAffinityMode && current?.revision === oauthSelection?.revision + && current?.accountId === oauthSelection?.accountId && servingOAuthSnapshot?.accountId === sent.accountId)) + ? await forceRefreshOAuthAccessSnapshot(sent) : sent; const admitted = await commitResolvedOAuthSelection(candidate); if (!admitted) throw new Error("OAuth selection changed during credential recovery"); if (kiroLoadEnabled && options.accountLoad?.lease?.accountId !== admitted.accountId) { @@ -519,6 +539,8 @@ export async function prepareResponsesTransport( return selected?.accountId === binding.selection.accountId && selected?.revision === binding.selection.revision && !!row && !row.paused && !row.needsReauth && row.credential.expires > Date.now() && (anthropicInstance === undefined || !anthropicRouting!.getAnthropicAccountHealthSnapshot(binding.snapshot.accountId) && !anthropicRatePolicy!.anthropicRatePauseUntil(binding.snapshot.accountId) && !anthropicQuota!.anthropicFamilyRejected(binding.snapshot.accountId, route.modelId)) + && (!genericAffinityMode + || genericSessionAccountEligible(config, route.providerName, binding.snapshot.accountId, route.modelId, Date.now(), true)) && credentialGeneration(row.credential) === binding.snapshot.generation; }; /** Bind the negotiated wire adapter and its built requests to the serving credential's authority. */ @@ -592,6 +614,12 @@ export async function prepareResponsesTransport( const accountId = await anthropicRouting!.resolveAnthropicDispatchAccountId(config, anthropicSessionKey, anthropicRouteDecision, route.modelId); candidate = await anthropicRouting!.getAnthropicPoolAccessSnapshot(accountId); } + if (genericAffinityMode + && (!candidate || !genericSessionAccountEligible(config, route.providerName, candidate.accountId, route.modelId, Date.now(), true))) { + const accountId = preferredGenericSessionAccount(config, route.providerName, genericSessionKey, route.modelId) + ?? captureOAuthAccountSelection(route.providerName)?.accountId; + candidate = accountId ? await getValidAccessSnapshotForAccount(route.providerName, accountId, { requireUsableAccount: true }) : undefined; + } if (!candidate || !await applyFailoverSnapshot(candidate, requestParsed, true)) { if (anthropicInstance !== undefined) await anthropicRouting!.resolveAnthropicDispatchAccountId(config, anthropicSessionKey, anthropicRouteDecision, route.modelId); throw new Error("OAuth account selection changed before dispatch"); @@ -780,8 +808,8 @@ export async function prepareResponsesTransport( throw new Error("OAuth account selection changed repeatedly before dispatch"); }; }; - const anthropicSessionKey = anthropicInstance !== undefined && route.provider.authMode === "oauth" - ? anthropicSessionKeyFromParts({ + const oauthSessionKey = route.provider.authMode === "oauth" + ? oauthSessionKeyFromParts({ sessionIdHeader: sessionIdHeaderFromRequest(req.headers), threadIdHeader: req.headers.get("thread-id"), promptCacheKey: typeof parsed.options.promptCacheKey === "string" ? parsed.options.promptCacheKey : null, @@ -789,6 +817,9 @@ export async function prepareResponsesTransport( promptCacheKeyIsSharedCohort: options.promptCacheKeyIsSharedCohort === true, }) : null; + const anthropicSessionKey = anthropicInstance ? oauthSessionKey : null; + const genericSessionKey = genericAffinityMode && oauthSessionKey + ? JSON.stringify([resolveContextPrincipal(req, config, options.admission) ?? "loopback", oauthSessionKey]) : null; if (route.provider.authMode === "oauth") { try { if (anthropicInstance !== undefined && anthropicRouting!.isAnthropicAccountPoolEnabled(config)) { @@ -821,10 +852,17 @@ export async function prepareResponsesTransport( // without quota evidence keeps the resolution it has today. const refusalAwareId = route.providerName === "kiro" && oauthSelection?.accountId ? refusalAwareInitialKiroAccount(config, oauthSelection.accountId, Date.now(), route.modelId) : null; - const preferredAccountId = refusalAwareId ?? (isGenericFailoverProvider(route.providerName, route.provider) + const affinityAccountId = genericSessionAccount(config, route.providerName, genericSessionKey, route.modelId); + const proposedAccountId = affinityAccountId + ?? preferredGenericSessionAccount(config, route.providerName, genericSessionKey, route.modelId) + ?? refusalAwareId ?? (isGenericFailoverProvider(route.providerName, route.provider) ? preferredInitialAccount(config, route.providerName, Date.now(), route.modelId) : null); - let safetyAlternateId: string | null = refusalAwareId; + // An unchanged active choice must retain the ordinary resolver's terminal-refresh + // recovery (not the speculative-preference fallback that resolves the same dead grant). + const preferredAccountId = genericAffinityMode && proposedAccountId === oauthSelection?.accountId + ? null : proposedAccountId; + let safetyAlternateId: string | null = !affinityAccountId && preferredAccountId === refusalAwareId ? refusalAwareId : 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 @@ -926,7 +964,8 @@ export async function prepareResponsesTransport( // Advance the pool cursor only now that this account is actually admitted. The // helper returns immediately unless the kernel is on AND the strategy is // round-robin, so quota and fill-first pools reach it without being touched. - noteGenericPoolSelection(config, route.providerName, resolved.accountId, route.modelId); + if (!affinityAccountId || affinityAccountId !== resolved.accountId) + noteGenericPoolSelection(config, route.providerName, resolved.accountId, route.modelId); } // Anthropic is excluded from isGenericFailoverProvider -- its own pool owns affinity and // a fail-closed local-cli credential rule -- so without this stamp its identity is diff --git a/src/types/provider.ts b/src/types/provider.ts index a06e907a420..fc21280bf3d 100644 --- a/src/types/provider.ts +++ b/src/types/provider.ts @@ -712,6 +712,8 @@ export interface OcxProviderConfig { /** Kiro OAuth only: active serving requests per account, 1..100. */ maxConcurrentPerAccount?: number; enabled?: boolean; + /** Opt-in process-local conversation ownership; requires proactive enablement. */ + sessionAffinity?: "conversation" | "none"; /** * Generic OAuth pool selection strategy (#695). Persisted through the pool-settings * contract. Consumed by the selector only while `pool.kernel` is on; with the flag off diff --git a/structure/INDEX.md b/structure/INDEX.md index 0942d50a14a..6c8fff820fd 100644 --- a/structure/INDEX.md +++ b/structure/INDEX.md @@ -62,6 +62,7 @@ Per-vendor contracts and the adapter authority that constructs them. | Doc | Scope | | --- | --- | +| [`providers/generic-account-pool.md`](providers/generic-account-pool.md) | Opt-in conversation ownership, request-local selection receipts and bounded release for generic OAuth pools. | | [`providers/anthropic-account-thresholds.md`](providers/anthropic-account-thresholds.md) | Account-owned usage thresholds, inheritance, routing boundaries and durable policy changes. | | [`providers-and-adapters.md`](providers-and-adapters.md) | Provider and adapter selection, the adapter inventory, live model discovery, and the hosted-search continuation bridge. | | [`providers/jev-decision.md`](providers/jev-decision.md) | JEV Combo decision methods (TypeSafe, self-hosted System One rows, opencodex models), the decision request path, dashboard surfaces, and the content-free statistics projection. | @@ -177,7 +178,7 @@ A source area can be described by more than one doc, because these docs are orga | `src/lib/windows-owner-acl.ts` | [`remote-link.md`](remote-link.md) | | `src/link/` | [`remote-link.md`](remote-link.md) | | `src/messaging/` | [`local-messaging.md`](local-messaging.md) | -| `src/oauth/` | [`runtime.md`](runtime.md)
[`transports/inventory.md`](transports/inventory.md)
[`providers/anthropic-account-thresholds.md`](providers/anthropic-account-thresholds.md)
[`providers-and-adapters.md`](providers-and-adapters.md)
[`providers/anthropic-account-pool.md`](providers/anthropic-account-pool.md)
[`providers/xai-grok.md`](providers/xai-grok.md) | +| `src/oauth/` | [`providers/generic-account-pool.md`](providers/generic-account-pool.md)
[`runtime.md`](runtime.md)
[`transports/inventory.md`](transports/inventory.md)
[`providers/anthropic-account-thresholds.md`](providers/anthropic-account-thresholds.md)
[`providers-and-adapters.md`](providers-and-adapters.md)
[`providers/anthropic-account-pool.md`](providers/anthropic-account-pool.md)
[`providers/xai-grok.md`](providers/xai-grok.md) | | `src/plugins/` | [`ops/plugins.md`](ops/plugins.md) | | `src/protocols/` | [`data-planes/protocol-paths.md`](data-planes/protocol-paths.md) | | `src/providers/` | [`runtime.md`](runtime.md)
[`subagents.md`](subagents.md)
[`transports/inventory.md`](transports/inventory.md)
[`providers-and-adapters.md`](providers-and-adapters.md)
[`providers/anthropic-account-pool.md`](providers/anthropic-account-pool.md)
[`providers/xai-grok.md`](providers/xai-grok.md) | @@ -194,7 +195,7 @@ A source area can be described by more than one doc, because these docs are orga | `src/router.ts` | [`runtime.md`](runtime.md) | | `src/routing/` | [`catalog.md`](catalog.md) | | `src/routing/history/` | [`dashboard-and-usage.md`](dashboard-and-usage.md) | -| `src/server/` | [`runtime.md`](runtime.md)
[`codex-prompts.md`](codex-prompts.md)
[`catalog.md`](catalog.md)
[`subagents.md`](subagents.md)
[`transports/byte-accounting.md`](transports/byte-accounting.md)
[`transports/responses.md`](transports/responses.md)
[`transports/responses-wire-shapes.md`](transports/responses-wire-shapes.md)
[`transports/responses-failover.md`](transports/responses-failover.md)
[`transports/policy-fallback.md`](transports/policy-fallback.md)
[`transports/streaming-health.md`](transports/streaming-health.md)
[`transports/inventory.md`](transports/inventory.md)
[`data-planes/images.md`](data-planes/images.md)
[`data-planes/inbound-compat.md`](data-planes/inbound-compat.md)
[`providers-and-adapters.md`](providers-and-adapters.md)
[`providers/jev-decision.md`](providers/jev-decision.md)
[`providers/xai-grok.md`](providers/xai-grok.md)
[`adapters/registry.md`](adapters/registry.md)
[`gui-and-management-api.md`](gui-and-management-api.md)
[`dashboard-and-usage.md`](dashboard-and-usage.md)
[`clients/claude-desktop.md`](clients/claude-desktop.md)
[`ops/service-and-sidecars.md`](ops/service-and-sidecars.md) | +| `src/server/` | [`providers/generic-account-pool.md`](providers/generic-account-pool.md)
[`runtime.md`](runtime.md)
[`codex-prompts.md`](codex-prompts.md)
[`catalog.md`](catalog.md)
[`subagents.md`](subagents.md)
[`transports/byte-accounting.md`](transports/byte-accounting.md)
[`transports/responses.md`](transports/responses.md)
[`transports/responses-wire-shapes.md`](transports/responses-wire-shapes.md)
[`transports/responses-failover.md`](transports/responses-failover.md)
[`transports/policy-fallback.md`](transports/policy-fallback.md)
[`transports/streaming-health.md`](transports/streaming-health.md)
[`transports/inventory.md`](transports/inventory.md)
[`data-planes/images.md`](data-planes/images.md)
[`data-planes/inbound-compat.md`](data-planes/inbound-compat.md)
[`providers-and-adapters.md`](providers-and-adapters.md)
[`providers/jev-decision.md`](providers/jev-decision.md)
[`providers/xai-grok.md`](providers/xai-grok.md)
[`adapters/registry.md`](adapters/registry.md)
[`gui-and-management-api.md`](gui-and-management-api.md)
[`dashboard-and-usage.md`](dashboard-and-usage.md)
[`clients/claude-desktop.md`](clients/claude-desktop.md)
[`ops/service-and-sidecars.md`](ops/service-and-sidecars.md) | | `src/server/gui-pair-delivery.ts` | [`remote-link.md`](remote-link.md) | | `src/server/index.ts` | [`adapters/compatibility-lab.md`](adapters/compatibility-lab.md) | | `src/server/management/anthropic-pool-settings.ts` | [`providers/anthropic-account-pool.md`](providers/anthropic-account-pool.md) | @@ -217,7 +218,7 @@ A source area can be described by more than one doc, because these docs are orga | `src/storage/` | [`runtime.md`](runtime.md) | | `src/tray/` | [`runtime.md`](runtime.md) | | `src/types.ts` | [`runtime.md`](runtime.md)
[`config.md`](config.md) | -| `src/types/` | [`providers-and-adapters.md`](providers-and-adapters.md) | +| `src/types/` | [`providers/generic-account-pool.md`](providers/generic-account-pool.md)
[`providers-and-adapters.md`](providers-and-adapters.md) | | `src/update/` | [`runtime.md`](runtime.md)
[`ops/service-and-sidecars.md`](ops/service-and-sidecars.md) | | `src/update/job-supervised-restart.ts` | [`ops/update-supervised-recovery.md`](ops/update-supervised-recovery.md) | | `src/update/job.ts` | [`ops/update-supervised-recovery.md`](ops/update-supervised-recovery.md) | diff --git a/structure/manifest.json b/structure/manifest.json index 5176735ae24..899b9dd9426 100644 --- a/structure/manifest.json +++ b/structure/manifest.json @@ -47,6 +47,13 @@ } ], "docs": [ + { + "path": "providers/generic-account-pool.md", + "tier": 4, + "title": "Generic OAuth Account Pool", + "scope": "Opt-in conversation ownership, request-local selection receipts and bounded release for generic OAuth pools.", + "documents": ["src/oauth/", "src/server/", "src/types/"] + }, { "path": "overview.md", "tier": 1, diff --git a/structure/providers-and-adapters.md b/structure/providers-and-adapters.md index c522806f375..57d1babb422 100644 --- a/structure/providers-and-adapters.md +++ b/structure/providers-and-adapters.md @@ -10,7 +10,7 @@ Per-account usage thresholds follow the [Anthropic account thresholds contract]( A shared-quota Anthropic 429 or classified pre-output account 403 records the served account's cooldown even when the request has used its allowed retry sends. That final account remains excluded on the next request; combo target cooling for 429 is skipped only after the matching account cooldown is present. Exact revoked-token 401s persist generation-fenced reauthentication under [the Anthropic recovery contract](providers/anthropic-account-pool.md#revoked-oauth-access-token-recovery). The Anthropic helper sends share the same routing authority: `getAnthropicSidecarAccessToken` resolves the vision-describe and web-search sidecars' helper model through the same first-match route decision, so a routed send authenticates as the route's own account rather than whatever pool account happens to be active. A strict route with no eligible account fails the helper locally instead of silently falling back to the active outsider, matching the primary-traffic contract; callers without a pool config keep the plain stored-credential path. - +Generic OAuth conversation ownership follows [the generic account-pool contract](providers/generic-account-pool.md#conversation-affinity). ## GitHub Copilot Auto selection `src/providers/github-copilot-transport.ts` retains the credential-specific destination and headers. Automatic selection requires a nonempty catalog whose every valid row explicitly has `model_picker_enabled: false`; missing or malformed permissions retain named routes. Named picker projection removes only explicitly false rows. diff --git a/structure/providers/anthropic-account-pool.md b/structure/providers/anthropic-account-pool.md index cf0c0b9ba8e..01b95f452d0 100644 --- a/structure/providers/anthropic-account-pool.md +++ b/structure/providers/anthropic-account-pool.md @@ -450,3 +450,5 @@ updating live state. A confirmed published write followed by bookkeeping failure pool and returns a fixed warning; an unpublished failure keeps old state. Unknown outcomes require reload and are not represented as a successful save. The GUI checkbox is a saved preference rather than a promise that a route is eligible; explicit rollout opt-outs still apply. + +The shared conversation-key precedence is implemented by `src/oauth/session-affinity-key.ts` and re-exported as `anthropicSessionKeyFromParts`; Anthropic routing semantics remain owned here. [Generic OAuth affinity](generic-account-pool.md#conversation-affinity) uses independent request-local receipts. diff --git a/structure/providers/generic-account-pool.md b/structure/providers/generic-account-pool.md new file mode 100644 index 00000000000..1ba25dc9ce3 --- /dev/null +++ b/structure/providers/generic-account-pool.md @@ -0,0 +1,25 @@ +# Generic OAuth Account Pool + +## Conversation affinity + +`src/oauth/generic-session-affinity.ts` owns opt-in generic-pool conversation bindings. Provider-local +`oauthAccountFailover.sessionAffinity: "conversation"` requires proactive enablement; it does not +change reactive failover activation. Bindings are scoped to configuration home, provider and caller +principal, store only a hashed session key, account ID and selection receipt, and expire after 24 idle +hours within a process-wide 2,000-entry bound. Pause, reauth, model/family ineligibility and persisted +selection-revision changes release them. Cross-process revisions are checked on every lookup. +`src/oauth/session-affinity-key.ts` supplies the existing client-thread/header/cache-key precedence +for both generic and Anthropic callers; shared Desktop cache cohorts do not establish identity. + +`src/server/responses/request-transport.ts` admits every automatic generic affinity-mode choice +without persisting an active-account change, including new/unkeyed sessions and failovers. The +operator's saved selection/revision fences queued sends independently of each conversation's serving +snapshot and credential generation. An unchanged selection with a renewed credential re-resolves +that same account; a changed selection follows the operator. Off-active 401 refresh remains bound +to the sent account. Admitted replacements rebind; affinity hits do not advance round-robin. +Fill-first has a bounded process-local cursor and prefers below-threshold alternatives without +moving the persisted selection. Kiro capacity retains wait/refusal semantics. Codex and Anthropic +are excluded. `tests/oauth/generic-session-affinity.test.ts` and +`tests/server/server-generic-oauth-affinity.test.ts` exercise these boundaries with synthetic accounts. + +`src/types/provider.ts` declares the optional `sessionAffinity` mode; omission keeps existing behavior. diff --git a/structure/runtime.md b/structure/runtime.md index 2a31215cb52..c6ef3294e7f 100644 --- a/structure/runtime.md +++ b/structure/runtime.md @@ -597,4 +597,4 @@ registration succeeds. Busy lease diagnostics name the recorded PID, current liv Bun updater lease and recovery behavior follows the [update transaction contract](ops/service-and-sidecars.md#bun-updater-ownership-transaction). -Companion timeline and filtered totals follow the [companion usage contract](companion.md). [Ongoing priority failback](providers/openai-accounts.md#ongoing-priority-failback) reuses request-triggered quota priming and captured-account dispatch; it adds no periodic worker or mid-request account switch. The serving-install census and bounded foreground delegation in `src/config/serving-runtimes.ts` follow [service command selection](ops/service-and-sidecars.md#background-service-command-selection). The live main-account policy projection is fetched from the attested account-list management API. `src/cli/status-oauth.ts` renders its state and effective thresholds with an optional outside-usage advisory; `src/cli/status.ts` includes the same whitelisted projection in `ocx status --json`. No policy state is inferred from the CLI process's local quota cache when the live read is unavailable. The experimental macOS `ocx chatgpt` launcher, restore and status commands follow the [ChatGPT Desktop contract](clients/chatgpt-desktop.md); the internal stdout filter remains hidden from public capability discovery. +Companion timeline and filtered totals follow the [companion usage contract](companion.md). [Ongoing priority failback](providers/openai-accounts.md#ongoing-priority-failback) reuses request-triggered quota priming and captured-account dispatch; it adds no periodic worker or mid-request account switch. The serving-install census and bounded foreground delegation in `src/config/serving-runtimes.ts` follow [service command selection](ops/service-and-sidecars.md#background-service-command-selection). The live main-account policy projection is fetched from the attested account-list management API. `src/cli/status-oauth.ts` renders its state and effective thresholds with an optional outside-usage advisory; `src/cli/status.ts` includes the same whitelisted projection in `ocx status --json`. No policy state is inferred from the CLI process's local quota cache when the live read is unavailable. The experimental macOS `ocx chatgpt` launcher, restore and status commands follow the [ChatGPT Desktop contract](clients/chatgpt-desktop.md); the internal stdout filter remains hidden from public capability discovery. Generic OAuth conversation ownership follows [the generic account-pool contract](providers/generic-account-pool.md#conversation-affinity). diff --git a/structure/transports/inventory.md b/structure/transports/inventory.md index 896442feb4e..12c7a47728a 100644 --- a/structure/transports/inventory.md +++ b/structure/transports/inventory.md @@ -417,3 +417,5 @@ Only the proxy handshake carries proxy credentials. The raw upgraded socket rema the relay does not decode WebSocket payloads. After CONNECT the paused outer stream resumes before TLS-over-TLS negotiation. Test fixture CA trust is confined to child processes or injected test socket options, never the user's OS trust store. + +Generic OAuth conversation ownership and request-local selection receipts follow [generic OAuth affinity](../providers/generic-account-pool.md#conversation-affinity). diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 210221f4d50..ba3e1339d6c 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -629,6 +629,8 @@ "fulfill-video.test.ts": "videos", "gather-routed-models-single-flight.test.ts": "codex-integration", "gcp-adc.test.ts": "adapters/google", "gemini-37-flash-migration.test.ts": "adapters/google", "gemini-4-argon-antigravity.test.ts": "adapters/google", "antigravity-account-roster.test.ts": "adapters/google", "gemini-inline.test.ts": "images", "gemini-web-search.test.ts": "adapters/google", + "generic-session-affinity.test.ts": "oauth", + "server-generic-oauth-affinity.test.ts": "server", "generic-oauth-failover.test.ts": "oauth", "generic-oauth-refusal-composition.test.ts": "oauth", "github-copilot-auto.test.ts": "providers/github-copilot", "github-copilot-auto-runtime.test.ts": "providers/github-copilot", diff --git a/tests/oauth/generic-session-affinity.test.ts b/tests/oauth/generic-session-affinity.test.ts new file mode 100644 index 00000000000..ede948d3b07 --- /dev/null +++ b/tests/oauth/generic-session-affinity.test.ts @@ -0,0 +1,186 @@ +import { afterEach, beforeEach, expect, test } from "bun:test"; +import { mkdtempSync, readFileSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { admitGenericSessionAccount, bindGenericSessionAccount, clearGenericSessionAffinity, + genericSessionAccount, genericSessionAffinityEnabled, preferredGenericSessionAccount } from "../../src/oauth/generic-session-affinity"; +import { clearGenericFailoverHealth, rotateGenericOAuthAccountOn429 } from "../../src/oauth/generic-account-failover"; +import { getValidAccessSnapshotForAccount } from "../../src/oauth"; +import { captureOAuthAccountSelection, getAccountSet, getAuthStorePath, markAccountNeedsReauth, + saveAccountCredential, saveCredential, setActiveAccount } from "../../src/oauth/store"; +import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests } from "../../src/providers/quota"; +import { acquireAccountLease } from "../../src/oauth/kiro-account-load"; +import { awaitKiroModelRefreshForTests, clearKiroAccountModels, refreshKiroAccountModelsDetached } from "../../src/providers/kiro-model-catalog"; +import { oauthSessionKeyFromParts } from "../../src/oauth/session-affinity-key"; +import type { OcxConfig } from "../../src/types"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; + +let home: string; +let priorHome: string | undefined; +let cfg: OcxConfig; +let ids: string[]; +const now = () => Date.now(); +beforeEach(async () => { + priorHome = process.env.OPENCODEX_HOME; + home = mkdtempSync(join(tmpdir(), "ocx-affinity-unit-")); + process.env.OPENCODEX_HOME = home; + clearGenericSessionAffinity(); + clearGenericFailoverHealth(); + for (const slot of ["a", "b", "c"]) await saveCredential("xai", { + access: `synthetic-${slot}`, refresh: `synthetic-refresh-${slot}`, expires: now() + 3_600_000, + accountId: `synthetic-${slot}`, + }); + ids = getAccountSet("xai")!.accounts.map(row => row.id); + await setActiveAccount("xai", ids[0]!); + cfg = { providers: { xai: { adapter: "openai-chat", baseUrl: "https://api.x.ai/v1", authMode: "oauth", + oauthAccountFailover: { enabled: true, sessionAffinity: "conversation" } } } } as OcxConfig; +}); +afterEach(() => { + clearGenericSessionAffinity(); + clearGenericFailoverHealth(); + clearAccountQuotaCache("xai"); + if (priorHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = priorHome; + removeTreeWithRetry(home); +}); + +test("requires explicit conversation mode and narrow-over-broad proactive permission", () => { + expect(genericSessionAffinityEnabled(cfg, "xai")).toBe(true); + cfg.oauthAccountFailover = { enabled: true }; + cfg.providers.xai!.oauthAccountFailover!.enabled = false; + expect(genericSessionAffinityEnabled(cfg, "xai")).toBe(false); + delete cfg.providers.xai!.oauthAccountFailover!.enabled; + expect(genericSessionAffinityEnabled(cfg, "xai")).toBe(true); + delete cfg.oauthAccountFailover; + expect(genericSessionAffinityEnabled(cfg, "xai")).toBe(false); +}); + +test("disabled mode discards bindings and Codex/Anthropic pools remain excluded", () => { + bindGenericSessionAccount(cfg, "xai", "session", ids[1]!); + cfg.providers.xai!.oauthAccountFailover!.sessionAffinity = "none"; + expect(genericSessionAccount(cfg, "xai", "session")).toBeNull(); + cfg.providers.xai!.oauthAccountFailover!.sessionAffinity = "conversation"; + expect(genericSessionAccount(cfg, "xai", "session")).toBeNull(); + for (const provider of ["openai", "anthropic", "anthropic2"]) { + cfg.providers[provider] = { ...cfg.providers.xai! }; + expect(genericSessionAffinityEnabled(cfg, provider)).toBe(false); + } +}); + +test("ineligible reauth and family cooldown release bindings without an event", async () => { + bindGenericSessionAccount(cfg, "xai", "session", ids[1]!); + await markAccountNeedsReauth("xai", ids[1]!, true); + expect(genericSessionAccount(cfg, "xai", "session")).toBeNull(); + bindGenericSessionAccount(cfg, "xai", "session", ids[0]!); + rotateGenericOAuthAccountOn429(cfg, "xai", ids[0]!, "60"); + expect(genericSessionAccount(cfg, "xai", "session")).toBeNull(); +}); + +test("cross-process manual selection revokes affinity without a local publication", () => { + bindGenericSessionAccount(cfg, "xai", "session", ids[1]!); + const path = getAuthStorePath(); + const store = JSON.parse(readFileSync(path, "utf8")); + store.xai.activeAccountId = ids[2]; + store.xai.selectionRevision = "synthetic-external-revision"; + writeFileSync(path, JSON.stringify(store)); + expect(genericSessionAccount(cfg, "xai", "session")).toBeNull(); + expect(preferredGenericSessionAccount(cfg, "xai", "session")).toBe(ids[2]); +}); + +test("admission preserves the manual receipt and rejects stale credentials and A→B→A", async () => { + const expected = captureOAuthAccountSelection("xai")!; + const snapshot = await getValidAccessSnapshotForAccount("xai", ids[1]!); + expect(admitGenericSessionAccount(cfg, "xai", expected, snapshot)).toEqual(expected); + expect(getAccountSet("xai")!.activeAccountId).toBe(ids[0]); + const row = getAccountSet("xai")!.accounts.find(row => row.id === ids[1])!; + await saveAccountCredential("xai", ids[1]!, { ...row.credential, access: "synthetic-renewed" }); + expect(admitGenericSessionAccount(cfg, "xai", expected, snapshot)).toBeNull(); + const fresh = await getValidAccessSnapshotForAccount("xai", ids[1]!); + expect(admitGenericSessionAccount(cfg, "xai", expected, fresh)).toEqual(expected); + await setActiveAccount("xai", ids[2]!); + await setActiveAccount("xai", ids[0]!); + expect(admitGenericSessionAccount(cfg, "xai", expected, fresh)).toBeNull(); +}); + +test("TTL, total entry limit and config-home/provider scope are bounded", () => { + const clock = now(); + bindGenericSessionAccount(cfg, "xai", "old", ids[1]!, undefined, clock); + expect(genericSessionAccount(cfg, "xai", "old", undefined, clock + 24 * 60 * 60_000)).toBeNull(); + for (let i = 0; i < 2001; i++) bindGenericSessionAccount(cfg, "xai", `session-${i}`, ids[1]!, undefined, clock); + expect(genericSessionAccount(cfg, "xai", "session-0", undefined, clock)).toBeNull(); + expect(genericSessionAccount(cfg, "xai", "session-2000", undefined, clock)).toBe(ids[1]); + cfg.providers.kimi = { ...cfg.providers.xai! }; + expect(genericSessionAccount(cfg, "kimi", "session-2000", undefined, clock)).toBeNull(); + process.env.OPENCODEX_HOME = join(home, "other"); + expect(genericSessionAccount(cfg, "xai", "session-2000", undefined, clock)).toBeNull(); +}); + +test("fill-first advances through three accounts without persisting an automatic active account", () => { + cfg.pool = { kernel: true }; + cfg.providers.xai!.oauthAccountFailover!.strategy = "fill-first"; + cfg.providers.xai!.oauthAccountFailover!.autoSwitchThreshold = 80; + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { weeklyPercent: 90, updatedAt: now() }); + const second = preferredGenericSessionAccount(cfg, "xai", "session-b")!; + bindGenericSessionAccount(cfg, "xai", "session-b", second); + setCachedProviderAccountQuotaForTests("xai", second, { weeklyPercent: 90, updatedAt: now() }); + const third = preferredGenericSessionAccount(cfg, "xai", "session-c"); + expect(third).not.toBe(second); + expect(third).not.toBe(ids[0]); + expect(getAccountSet("xai")!.activeAccountId).toBe(ids[0]); +}); + +test("shared cohort keys do not establish conversation identity", () => { + expect(oauthSessionKeyFromParts({ promptCacheKey: "shared", promptCacheKeyIsSharedCohort: true })).toBeNull(); + expect(oauthSessionKeyFromParts({ sessionIdHeader: "distinct", promptCacheKey: "shared", promptCacheKeyIsSharedCohort: true })).toBe("distinct"); + expect(oauthSessionKeyFromParts({ clientThreadId: "a".repeat(200) })).toHaveLength(64); +}); + + +test("all-cooled accounts retain the ordinary active attempt without a false selection conflict", async () => { + for (const id of ids) rotateGenericOAuthAccountOn429(cfg, "xai", id, "60"); + expect(preferredGenericSessionAccount(cfg, "xai", "session")).toBe(ids[0]); + const selection = captureOAuthAccountSelection("xai")!; + const snapshot = await getValidAccessSnapshotForAccount("xai", ids[0]!); + expect(admitGenericSessionAccount(cfg, "xai", selection, snapshot)).toEqual(selection); + bindGenericSessionAccount(cfg, "xai", "session", ids[0]!); + expect(genericSessionAccount(cfg, "xai", "session")).toBeNull(); +}); + + +test("a held Kiro capacity lease cannot invalidate its own affinity or dispatch admission", async () => { + const priorDiscovery = process.env.OPENCODEX_KIRO_MODEL_DISCOVERY; + process.env.OPENCODEX_KIRO_MODEL_DISCOVERY = "1"; + cfg.providers.kiro = { adapter: "kiro", authMode: "oauth", baseUrl: "https://runtime.us-east-1.kiro.dev", + oauthAccountFailover: { enabled: true, sessionAffinity: "conversation", maxConcurrentPerAccount: 1 } }; + let releaseLease: (() => void) | undefined; + try { + for (const slot of ["a", "b"]) await saveCredential("kiro", { + access: `synthetic-${slot}`, refresh: `synthetic-refresh-${slot}`, expires: now() + 3_600_000, + accountId: `synthetic-${slot}`, kiro: { apiRegion: "us-east-1", + profileArn: `arn:aws:codewhisperer:us-east-1:123456789012:profile/synthetic-${slot}` }, + }); + const rows = getAccountSet("kiro")!.accounts; + const dependencies = { + resolveAddresses: async () => ({ hostname: "runtime.us-east-1.kiro.dev", + addresses: [{ address: "1.1.1.1", family: 4 }], privateNetwork: false }), + pinnedPost: async () => Response.json({ models: [{ modelId: "model-one" }] }), + } as never; + for (const row of rows) refreshKiroAccountModelsDetached(row, cfg.providers.kiro, dependencies); + await Promise.all(rows.map(row => awaitKiroModelRefreshForTests(row.id))); + const bound = rows[0]!.id; + bindGenericSessionAccount(cfg, "kiro", "session", bound, "model-one"); + const lease = await acquireAccountLease("kiro", bound, { maxConcurrentPerAccount: 1 }); + expect(lease).not.toBeNull(); + releaseLease = lease!.release; + expect(genericSessionAccount(cfg, "kiro", "session", "model-one")).toBe(bound); + const receipt = captureOAuthAccountSelection("kiro")!; + const snapshot = await getValidAccessSnapshotForAccount("kiro", bound); + expect(admitGenericSessionAccount(cfg, "kiro", receipt, snapshot, "model-one")).toEqual(receipt); + expect(await acquireAccountLease("kiro", bound, { maxConcurrentPerAccount: 1, waitMs: 1 })).toBeNull(); + } finally { + releaseLease?.(); + clearKiroAccountModels(); + if (priorDiscovery === undefined) delete process.env.OPENCODEX_KIRO_MODEL_DISCOVERY; + else process.env.OPENCODEX_KIRO_MODEL_DISCOVERY = priorDiscovery; + } +}); diff --git a/tests/server/server-generic-oauth-affinity.test.ts b/tests/server/server-generic-oauth-affinity.test.ts new file mode 100644 index 00000000000..4a1f287aa83 --- /dev/null +++ b/tests/server/server-generic-oauth-affinity.test.ts @@ -0,0 +1,386 @@ +import { afterEach, beforeEach, expect, spyOn, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { saveConfig } from "../../src/config"; +import { clearGenericFailoverHealth } from "../../src/oauth/generic-account-failover"; +import * as oauth from "../../src/oauth"; +import { getAccountSet, saveAccountCredential, saveCredential, setAccountPaused, setActiveAccount } from "../../src/oauth/store"; +import { clearPoolRotationState } from "../../src/oauth/pool-kernel"; +import { clearAntigravityAccountRosters, recordAntigravityAccountRoster } from "../../src/providers/antigravity-account-roster"; +import { rotateGenericOAuthAccountOn429 } from "../../src/oauth/generic-account-failover"; +import { subscribeAccountSelections } from "../../src/lib/account-selection-events"; +import { encodeMessage } from "../../src/lib/eventstream-decoder"; +import { handleResponses } from "../../src/server/responses"; +import type { OcxConfig } from "../../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "../helpers/isolated-codex-home"; +import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; +import * as pacing from "../../src/providers/request-pacing"; +import * as retry from "../../src/lib/upstream-retry"; + +const provider = "google-antigravity"; +const origin = "https://daily-cloudcode-pa.googleapis.com"; +let home: string; +let previousHome: string | undefined; +let codexHome: IsolatedCodexHome; +let release: () => void; +let originalFetch: typeof fetch; +let sleep: ReturnType; +let cfg: OcxConfig; +let accounts: Array<{ id: string; auth: string; project: string }>; +let seen: Array<{ auth: string; project: string }>; +let refreshCount: number; +let refuse: ((auth: string) => Response | undefined) | undefined; + +beforeEach(async () => { + previousHome = process.env.OPENCODEX_HOME; + home = mkdtempSync(join(tmpdir(), "ocx-generic-affinity-")); + process.env.OPENCODEX_HOME = home; + codexHome = installIsolatedCodexHome("ocx-affinity-codex-"); + release = acquireOwnedSpendHome(); + clearGenericFailoverHealth(); + clearPoolRotationState(); + originalFetch = globalThis.fetch; + sleep = spyOn(retry, "sleepWithAbort").mockImplementation(async () => {}); + for (const slot of ["one", "two", "three"]) { + await saveCredential(provider, { access: `synthetic-${slot}`, refresh: `synthetic-refresh-${slot}`, + expires: Date.now() + 3_600_000, accountId: `synthetic-${slot}`, projectId: `synthetic-project-${slot}` }); + } + accounts = getAccountSet(provider)!.accounts.map(row => ({ id: row.id, + auth: `Bearer ${row.credential.access}`, project: row.credential.projectId! })); + await setActiveAccount(provider, accounts[0]!.id); + cfg = { defaultProvider: provider, pool: { kernel: true }, providers: { [provider]: { + adapter: "google", authMode: "oauth", baseUrl: origin, googleMode: "cloud-code-assist", + models: ["gemini-3.8-flash"], oauthAccountFailover: { + enabled: true, strategy: "round-robin", sessionAffinity: "conversation", + }, + } } } as unknown as OcxConfig; + saveConfig(cfg); + seen = []; + refuse = undefined; + refreshCount = 0; + globalThis.fetch = (async (input, init) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.href === "https://oauth2.googleapis.com/token") { + refreshCount++; + return Response.json({ access_token: "synthetic-refreshed", expires_in: 3600, token_type: "Bearer" }); + } + if (url.origin !== origin || !["/v1internal:generateContent", "/v1internal:streamGenerateContent"].includes(url.pathname)) + throw new Error("Unexpected external request in affinity regression"); + const auth = new Headers(init?.headers).get("authorization")!; + const body = JSON.parse(String(init?.body)) as { project: string }; + seen.push({ auth, project: body.project }); + return refuse?.(auth) ?? Response.json({ response: { + candidates: [{ content: { role: "model", parts: [{ text: "synthetic result" }] }, finishReason: "STOP" }], + usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 3, totalTokenCount: 8 }, + } }); + }) as typeof fetch; +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + sleep.mockRestore(); + release(); + clearGenericFailoverHealth(); + clearAntigravityAccountRosters(); + clearPoolRotationState(); + codexHome.restore(); + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + removeTreeWithRetry(home); +}); + +async function turn(session?: string, extra: Record = {}, options: Parameters[3] = {}, expectedStatus = 200): Promise { + const response = await handleResponses(new Request("http://127.0.0.1/v1/responses", { + method: "POST", headers: { "content-type": "application/json", ...(session ? { "session-id": session } : {}) }, + body: JSON.stringify({ model: `${provider}/gemini-3.8-flash`, input: "synthetic input", stream: false, ...extra }), + }), cfg, { model: "gemini-3.8-flash", provider }, options); + const status = response.status; + await response.text(); + expect(status).toBe(expectedStatus); + return response; +} + +function assertPairing(): void { + for (const send of seen) expect(send.project).toBe(accounts.find(row => row.auth === send.auth)!.project); +} + +test("interleaved conversations retain their bearer/project without moving manual selection", async () => { + let selections = 0; + const unsubscribe = subscribeAccountSelections(() => { selections++; }); + try { + await turn("conversation-a"); + const first = seen.at(-1)!.auth; + await turn("conversation-b"); + const second = seen.at(-1)!.auth; + expect(second).not.toBe(first); + await turn("conversation-a"); + expect(seen.at(-1)!.auth).toBe(first); + await turn("conversation-b"); + expect(seen.at(-1)!.auth).toBe(second); + expect(getAccountSet(provider)!.activeAccountId).toBe(accounts[0]!.id); + expect(selections).toBe(0); + assertPairing(); + } finally { unsubscribe(); } +}); + +test("429 replacement stays bound after other conversations and cooldown expiry", async () => { + await turn("conversation-a"); + const failed = seen.at(-1)!.auth; + refuse = auth => auth === failed ? Response.json({ error: { code: 429, status: "RESOURCE_EXHAUSTED", + message: "Quota exceeded for quota metric" } }, { status: 429, headers: { "retry-after": "60" } }) : undefined; + await turn("conversation-a"); + const replacement = seen.at(-1)!.auth; + expect(replacement).not.toBe(failed); + refuse = undefined; + clearGenericFailoverHealth(); + await turn("conversation-b"); + await turn("conversation-a"); + expect(seen.at(-1)!.auth).toBe(replacement); + assertPairing(); +}); + +test("manual selection invalidates existing conversations including same-account reselection", async () => { + await turn("conversation-a"); + await turn("conversation-b"); + const selected = accounts.find(row => row.auth !== seen.at(-1)!.auth)!; + await setActiveAccount(provider, selected.id); + cfg.providers[provider]!.oauthAccountFailover!.strategy = "quota"; + await turn("conversation-b"); + expect(seen.at(-1)!.auth).toBe(selected.auth); + await setActiveAccount(provider, selected.id); + await turn("conversation-a"); + expect(seen.at(-1)!.auth).toBe(selected.auth); +}); + +test("a paused bound account is replaced and the survivor remains bound", async () => { + await turn("conversation-a"); + const bound = accounts.find(row => row.auth === seen.at(-1)!.auth)!; + await setAccountPaused(provider, bound.id, true); + await turn("conversation-a"); + const replacement = seen.at(-1)!.auth; + expect(replacement).not.toBe(bound.auth); + await turn("conversation-b"); + await turn("conversation-a"); + expect(seen.at(-1)!.auth).toBe(replacement); +}); + +test("none retains the existing per-request round-robin behavior", async () => { + (cfg.providers[provider]!.oauthAccountFailover as Record).sessionAffinity = "none"; + await turn("conversation-a"); + const first = seen.at(-1)!.auth; + await turn("conversation-a"); + expect(seen.at(-1)!.auth).not.toBe(first); +}); + + +test("401 on an off-active bound account refreshes that account and retains its project", async () => { + await turn("conversation-a"); + await turn("conversation-b"); + const before = seen.at(-1)!; + const bound = accounts.find(row => row.auth === before.auth)!; + expect(bound.id).not.toBe(getAccountSet(provider)!.activeAccountId); + refuse = auth => auth === before.auth ? Response.json({ error: { code: 401, + status: "UNAUTHENTICATED", message: "synthetic rejection" } }, { status: 401 }) : undefined; + await turn("conversation-b"); + expect(refreshCount).toBe(1); + expect(seen.at(-1)!.auth).toBe("Bearer synthetic-refreshed"); + expect(seen.at(-1)!.project).toBe(before.project); + await turn("conversation-b"); + expect(seen.at(-1)!.auth).toBe("Bearer synthetic-refreshed"); +}); + +test.each(["manual", "credential"])("%s change during a bound credential wait preserves the correct authority", async change => { + await turn("conversation-a"); + await turn("conversation-b"); + const bound = accounts.find(row => row.auth === seen.at(-1)!.auth)!; + const ready = Promise.withResolvers(); + const resume = Promise.withResolvers(); + const resolve = oauth.getValidAccessSnapshotForAccount; + let held = false; + const delay = spyOn(oauth, "getValidAccessSnapshotForAccount").mockImplementation(async (...args) => { + const snapshot = await resolve(...args); + if (!held && args[1] === bound.id) { held = true; ready.resolve(); await resume.promise; } + return snapshot; + }); + try { + const pending = turn("conversation-b"); + await ready.promise; + if (change === "manual") { + await setActiveAccount(provider, accounts[0]!.id); + } else { + const row = getAccountSet(provider)!.accounts.find(row => row.id === bound.id)!; + await saveAccountCredential(provider, bound.id, { ...row.credential, access: "synthetic-renewed" }); + } + resume.resolve(); + await pending; + expect(seen.at(-1)!.auth).toBe(change === "manual" ? accounts[0]!.auth : "Bearer synthetic-renewed"); + expect(seen.at(-1)!.project).toBe(change === "manual" ? accounts[0]!.project : bound.project); + } finally { resume.resolve(); delay.mockRestore(); } +}); + + +test("the same session marker is isolated by authenticated principal", async () => { + const admission = (principal: string) => ({ kind: "configured" as const, keyId: "synthetic-scope", + source: "bearer" as const, contextPrincipalId: principal }); + await turn("same-marker", {}, { admission: admission("principal-a") }); + const first = seen.at(-1)!.auth; + await turn("same-marker", {}, { admission: admission("principal-b") }); + const second = seen.at(-1)!.auth; + expect(second).not.toBe(first); + await turn("same-marker", {}, { admission: admission("principal-a") }); + expect(seen.at(-1)!.auth).toBe(first); +}); + +test("unkeyed and shared-cohort turns do not acquire affinity or move the active selection", async () => { + const before = getAccountSet(provider)!.selectionRevision; + await turn(); + const unkeyed = seen.at(-1)!.auth; + await turn(); + expect(seen.at(-1)!.auth).not.toBe(unkeyed); + await turn(undefined, { prompt_cache_key: "shared-cohort" }, { promptCacheKeyIsSharedCohort: true }); + const shared = seen.at(-1)!.auth; + await turn(undefined, { prompt_cache_key: "shared-cohort" }, { promptCacheKeyIsSharedCohort: true }); + expect(seen.at(-1)!.auth).not.toBe(shared); + expect(getAccountSet(provider)!.selectionRevision).toBe(before); +}); + +test.each(["interleave", "manual", "credential", "pause", "disable"])("%s during pacing rebuilds the serving binding correctly", async change => { + await turn("conversation-a"); + const first = seen.at(-1)!.auth; + await turn("conversation-b"); + const bound = accounts.find(row => row.auth === seen.at(-1)!.auth)!; + const ready = Promise.withResolvers(); + const resume = Promise.withResolvers(); + const withSlot = pacing.withProviderRequestSlot; + let held = false; + const delayed: typeof pacing.withProviderRequestSlot = async (...args) => { + if (!held) { held = true; ready.resolve(); await resume.promise; } + return withSlot(...args); + }; + const delay = spyOn(pacing, "withProviderRequestSlot").mockImplementation(delayed); + try { + const pending = turn("conversation-b"); + await ready.promise; + const sendsBefore = seen.length; + if (change === "interleave") { + await turn("conversation-a"); + expect(seen.at(-1)!.auth).toBe(first); + } else if (change === "disable") cfg.providers[provider]!.oauthAccountFailover!.enabled = false; + else if (change === "manual") await setActiveAccount(provider, accounts[0]!.id); + else if (change === "pause") await setAccountPaused(provider, bound.id, true); + else { + const row = getAccountSet(provider)!.accounts.find(row => row.id === bound.id)!; + await saveAccountCredential(provider, bound.id, { ...row.credential, access: "synthetic-renewed" }); + } + resume.resolve(); + await pending; + expect(seen.length).toBeGreaterThan(sendsBefore); + if (change === "manual") expect(seen.at(-1)!.auth).toBe(accounts[0]!.auth); + else if (change === "credential") expect(seen.at(-1)!.auth).toBe("Bearer synthetic-renewed"); + else if (change === "pause") expect(seen.at(-1)!.auth).not.toBe(bound.auth); + else expect(seen.at(-1)!.auth).toBe(bound.auth); + } finally { resume.resolve(); delay.mockRestore(); } +}); + + +test("disabling affinity during a refusal preserves reactive recovery and the manual selection", async () => { + await turn("conversation-a"); + const before = getAccountSet(provider)!.selectionRevision; + const failed = seen.at(-1)!.auth; + refuse = auth => { + if (auth !== failed) return; + cfg.providers[provider]!.oauthAccountFailover!.enabled = false; + return Response.json({ error: { code: 429, status: "RESOURCE_EXHAUSTED", message: "Quota exceeded for quota metric" } }, + { status: 429, headers: { "retry-after": "60" } }); + }; + await turn("conversation-a"); + expect(seen.at(-1)!.auth).not.toBe(failed); + expect(getAccountSet(provider)!.selectionRevision).toBe(before); +}); + + +test("an all-cooled pool retains the upstream 429 and Retry-After rather than a local conflict", async () => { + for (const row of accounts) rotateGenericOAuthAccountOn429(cfg, provider, row.id, "60", Date.now(), "gemini-3.8-flash"); + refuse = () => Response.json({ error: { code: 429, status: "RESOURCE_EXHAUSTED", + message: "Quota exceeded for quota metric" } }, { status: 429, headers: { "retry-after": "60" } }); + const response = await turn("conversation-a", {}, {}, 429); + expect(Number(response.headers.get("retry-after"))).toBeGreaterThan(0); + expect(seen).toHaveLength(1); +}); + +test("existing classified 403 recovery binds the replacement for the next turn", async () => { + await turn("conversation-a"); + const failed = seen.at(-1)!.auth; + refuse = auth => auth === failed ? Response.json({ error: { code: 403, status: "PERMISSION_DENIED", + message: "synthetic validation required", details: [{ reason: "VALIDATION_REQUIRED" }] } }, { status: 403 }) : undefined; + await turn("conversation-a"); + const replacement = seen.at(-1)!.auth; + expect(replacement).not.toBe(failed); + refuse = undefined; + clearGenericFailoverHealth(); + await turn("conversation-b"); + await turn("conversation-a"); + expect(seen.at(-1)!.auth).toBe(replacement); + assertPairing(); +}); + +test("a newly incompatible account model roster releases the conversation to a listing sibling", async () => { + await turn("conversation-a"); + const bound = accounts.find(row => row.auth === seen.at(-1)!.auth)!; + for (const row of accounts) { + const model = row.id === bound.id ? "other-model" : "gemini-3.8-flash"; + recordAntigravityAccountRoster(provider, row.id, [{ id: model, wireModelId: model } as never]); + } + await turn("conversation-a"); + expect(seen.at(-1)!.auth).not.toBe(bound.auth); +}); + + +test("Kiro terminal refresh on the active first choice still recovers to a sibling with affinity enabled", async () => { + for (const slot of ["a", "b"]) await saveCredential("kiro", { + access: `synthetic-${slot}`, refresh: `synthetic-refresh-${slot}`, + expires: slot === "a" ? Date.now() - 1000 : Date.now() + 3_600_000, + accountId: `synthetic-${slot}`, source: "oauth", kiro: { apiRegion: "us-east-1", + profileArn: `arn:aws:codewhisperer:us-east-1:123456789012:profile/synthetic-${slot}` }, + }); + const rows = getAccountSet("kiro")!.accounts; + await setActiveAccount("kiro", rows[0]!.id); + const config = { defaultProvider: "kiro", providers: { kiro: { + adapter: "kiro", authMode: "oauth", baseUrl: "https://runtime.us-east-1.kiro.dev", + models: ["claude-sonnet-4.5"], oauthAccountFailover: { enabled: true, sessionAffinity: "conversation" }, + } } } as OcxConfig; + saveConfig(config); + const before = getAccountSet("kiro")!.selectionRevision; + const sends: string[] = []; + globalThis.fetch = (async (input, init) => { + const url = input instanceof Request ? input.url : String(input); + if (url.endsWith("/refreshToken")) return Response.json({ error: "invalid_grant" }, { status: 400 }); + if (!url.startsWith("https://runtime.") || !url.endsWith(".kiro.dev/")) + throw new Error("Unexpected request in Kiro terminal-refresh regression"); + sends.push(new Headers(init?.headers).get("authorization")!); + return new Response(encodeMessage({ ":message-type": "event", ":event-type": "assistantResponseEvent" }, + new TextEncoder().encode(JSON.stringify({ content: "synthetic result" }))), + { headers: { "content-type": "application/vnd.amazon.eventstream" } }); + }) as typeof fetch; + const response = await handleResponses(new Request("http://127.0.0.1/v1/responses", { + method: "POST", headers: { "content-type": "application/json", "session-id": "kiro-session" }, + body: JSON.stringify({ model: "kiro/claude-sonnet-4.5", input: "synthetic input", stream: false }), + }), config, { provider: "kiro", model: "claude-sonnet-4.5" }); + await response.text(); + expect(response.status).toBe(200); + expect(sends).toEqual(["Bearer synthetic-b"]); + expect(getAccountSet("kiro")!.selectionRevision).toBe(before); +}); + + +test("project-less Antigravity credentials preserve the authentication refusal without dispatch", async () => { + for (const row of getAccountSet(provider)!.accounts) { + const credential = { ...row.credential }; + delete credential.projectId; + await saveAccountCredential(provider, row.id, credential); + } + await turn("conversation-a", {}, {}, 401); + expect(seen).toHaveLength(0); +});