@bitkyc08/opencodex 2.84.0 → 2.87.0

package/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).
@@ -25,6 +25,7 @@ import {
   hasHeadroomEvidence,
   isAccountQuotaExhausted,
   rankAccountsByHeadroom,
+  rankAccountsByResetFirst,
   classifyModelFamilyForQuota,
   type QuotaModelFamily,
 } from "./account-quota-rank";
@@ -41,6 +42,7 @@ import type { KiroRefusalKind } from "../adapters/kiro-refusal";
 import { kiroAccountEvidence } from "../providers/kiro-usage";
 import { kiroEvidenceIdentity } from "../providers/kiro-account-state-disk";
 import { kiroAccountSupportsModel } from "../providers/kiro-model-catalog";
+import { antigravityAccountSupportsModel } from "../providers/antigravity-account-roster";
 import { ACCOUNT_QUOTA_TTL_MS } from "../providers/quota-wire";
 import { sweepExpiredOnWrite } from "../lib/state-store-sweeper";
 import { subscribeOAuthAccountPauseChanges } from "../lib/account-selection-events";
@@ -248,11 +250,31 @@ export function eligibleFailoverAccounts(providerName: string, now = Date.now(),
   return eligibleIdsIn(getAccountSet(providerName), providerName, now, family);
 }
 
-/** Model evidence only ranks accounts already eligible and below any configured cap. */
+/** Whether an account's own model listing includes `model`; undefined when unknown. */
+function accountListsModel(providerName: string, accountId: string, model: string): boolean | undefined {
+  if (providerName === "kiro") return kiroAccountSupportsModel(accountId, model);
+  if (providerName === "google-antigravity") return antigravityAccountSupportsModel(providerName, accountId, model);
+  return undefined;
+}
+
+/**
+ * Model evidence only ranks accounts already eligible and below any configured cap.
+ *
+ * Kiro and Antigravity both gate models per account. For Antigravity (Claude 4.6 vs 5.5, Argon)
+ * an account whose discovered roster lists the model wins; an account whose roster omits it is
+ * skipped while any other remains; an account with no roster yet is unknown and stays usable.
+ */
 export function preferKiroModelSupport(
   providerName: string, ids: string[], model?: string | null, config?: OcxConfig,
 ): string[] {
-  if (providerName !== "kiro" || !model) return ids;
+  if (!model) return ids;
+  if (providerName === "google-antigravity") {
+    const positive = ids.filter(id => accountListsModel(providerName, id, model) === true);
+    if (positive.length) return positive;
+    const notListedAgainst = ids.filter(id => accountListsModel(providerName, id, model) !== false);
+    return notListedAgainst.length ? notListedAgainst : ids;
+  }
+  if (providerName !== "kiro") return ids;
   const cap = config?.providers?.kiro?.oauthAccountFailover?.maxConcurrentPerAccount;
   const withRoom = typeof cap === "number" && Number.isInteger(cap) && cap >= 1 && cap <= 100
     ? ids.filter(id => accountInFlight("kiro", id) < cap) : ids;
@@ -330,7 +352,7 @@ export function rotateAntigravityAccountOnAuthRefusal(
 }
 
 /** Generic pool strategies the kernel can actually run. `quota` IS the pre-kernel path. */
-type ActiveGenericStrategy = "round-robin" | "fill-first" | "least-loaded";
+type ActiveGenericStrategy = "round-robin" | "fill-first" | "reset-first" | "least-loaded";
 
 /** Matches the Codex and Anthropic pools; the DTO still reports `null` for "not stored". */
 const DEFAULT_GENERIC_AUTO_SWITCH_THRESHOLD = 80;
@@ -344,10 +366,13 @@ const DEFAULT_GENERIC_AUTO_SWITCH_THRESHOLD = 80;
  * branch instead of a four-way one.
  */
 function activeGenericStrategy(config: OcxConfig, providerName: string): ActiveGenericStrategy | null {
+  // All pool strategies run behind the kernel flag; `quota` is the pre-kernel path.
+  // Letting `reset-first` bypass it would change 429 selection for pools that never
+  // opted into kernel-gated strategies.
   if (config.pool?.kernel !== true) return null;
   const raw = config.providers?.[providerName]?.oauthAccountFailover?.strategy;
   if (providerName === "kiro" && raw === "least-loaded") return raw;
-  return raw === "round-robin" || raw === "fill-first" ? raw : null;
+  return raw === "round-robin" || raw === "fill-first" || raw === "reset-first" ? raw : null;
 }
 
 function genericStickyLimit(config: OcxConfig, providerName: string): number {
@@ -395,6 +420,7 @@ function pickFillFirstGenericAccount(
   now: number,
   requestedModelId?: string | null,
   accountRows?: ReadonlyMap<string, ProviderAccount>,
+  preferBelowThreshold = false,
 ): string | null {
   const stableAll = stableGenericRoster(providerName);
   if (stableAll.length < 2) return null;
@@ -410,6 +436,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;
   }
@@ -553,6 +584,10 @@ export function rotateGenericOAuthAccountOnRefusal(
     }
     return null;
   }
+  if (strategy === "reset-first") {
+    return rankAccountsByResetFirst(providerName, candidates, requestedModelId, now,
+      new Map(set.accounts.map(account => [account.id, account])))[0] ?? null;
+  }
   // With no quota evidence this returns the ring untouched, so providers without
   // per-account quota keep exactly the traversal they have today.
   return rankAccountsByHeadroom(providerName, candidates, requestedModelId,
@@ -613,6 +648,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.
@@ -621,7 +657,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)
@@ -631,10 +667,13 @@ export function preferredInitialAccount(
   const modelEligible = preferKiroModelSupport(providerName,
     eligibleFailoverAccounts(providerName, now, classifyModelFamilyForQuota(providerName, requestedModelId)),
     requestedModelId, config);
-  const activeLacks = providerName === "kiro" && !!active && !!requestedModelId
-    && kiroAccountSupportsModel(active, requestedModelId) === false;
+  const activeLacks = (providerName === "kiro" || providerName === "google-antigravity")
+    && !!active && !!requestedModelId
+    && accountListsModel(providerName, active, requestedModelId) === false;
   const listingSibling = activeLacks
-    ? modelEligible.find(id => id !== active && kiroModelListingHasRoom(config, id, requestedModelId!))
+    ? modelEligible.find(id => id !== active && (providerName === "kiro"
+      ? kiroModelListingHasRoom(config, id, requestedModelId!)
+      : accountListsModel(providerName, id, requestedModelId!) === true))
     : undefined;
   if (listingSibling) return listingSibling;
 
@@ -670,9 +709,20 @@ 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;
   }
+  if (strategy === "reset-first") {
+    // Same candidate set as round-robin: paused, reauth-needed and cooled accounts are out,
+    // and an account whose model listing omits the requested model loses to one that has it.
+    // Reset ordering needs quota evidence; without it the ordinary active-account path runs.
+    const family = classifyModelFamilyForQuota(providerName, requestedModelId);
+    const eligibleNow = preferKiroModelSupport(providerName,
+      eligibleFailoverAccounts(providerName, now, family), requestedModelId, config);
+    if (eligibleNow.length === 0 || !hasHeadroomEvidence(providerName, eligibleNow, requestedModelId, accountRows)) return null;
+    const best = rankAccountsByResetFirst(providerName, eligibleNow, requestedModelId, now, accountRows)[0] ?? null;
+    return best && best !== active ? best : null;
+  }
 
   const activeRow = selected.accounts.find(account => account.id === active);
   if (activeRow && activeRow.paused !== true && activeRow.needsReauth !== true

package/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<string, { accountId: string; lastUsedAt: number; selection: OAuthAccountSelection }>();
+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!);
+}