@yansigit/opencodex 2.31.1 → 2.31.3

package/src/oauth/antigravity-routing.ts

@@ -1,3 +1,14 @@
+import { getAccountSet } from "./store";
+import {
+  bindSessionAffinity,
+  buildSessionKeyFromParts,
+  clearAffinityState,
+  clearSessionAffinityForAccount,
+  getSessionAffinity,
+  normalizeAffinityComponent,
+  touchSessionAffinity,
+} from "../routing/account-pool";
+
 /**
  * Process-local Antigravity account health (cooldowns). Stored in an in-memory
  * `Map<string, AntigravityAccountHealth>` for the lifetime of this process only —
@@ -5,6 +16,10 @@
  */
 export type AntigravityCooldownReason = "rate_limited" | "quota_exhausted" | "geo_blocked";
 
+const POOL_KEY_ANTIGRAVITY = "google-antigravity";
+/** Short rate-limit stick-wait: retry the same account instead of hopping. Never wait on quota/geo. */
+export const ANTIGRAVITY_STICK_WAIT_MAX_MS = 5_000;
+
 const DEFAULT_RATE_LIMITED_COOLDOWN_MS = 5_000;
 const MAX_RATE_LIMITED_COOLDOWN_MS = 60_000;
 const DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_MS = 24 * 60 * 60_000;
@@ -118,6 +133,156 @@ export function clearAntigravityAccountCooldown(accountId: string): void {
   accountHealth.delete(accountId);
 }
 
+/** Test helper: reset process-local cooldown state between cases. */
+export function clearAntigravityRoutingHealthForTests(): void {
+  accountHealth.clear();
+}
+
+export function isAntigravityRateLimitStickWait(accountId: string, now = Date.now()): boolean {
+  const health = getAntigravityAccountCooldown(accountId, now);
+  if (!health || health.reason !== "rate_limited") return false;
+  const remaining = health.cooldownUntil - now;
+  return remaining > 0 && remaining <= ANTIGRAVITY_STICK_WAIT_MAX_MS;
+}
+
+/** Milliseconds to wait before retrying the same account on 429, or null when failover should run. */
+export function antigravity429StickWaitMs(accountId: string, now = Date.now()): number | null {
+  if (!isAntigravityRateLimitStickWait(accountId, now)) return null;
+  const health = getAntigravityAccountCooldown(accountId, now);
+  if (!health) return null;
+  const remaining = health.cooldownUntil - now;
+  return remaining > 0 ? remaining : null;
+}
+
+export function isAntigravityAccountEligible(accountId: string, now = Date.now()): boolean {
+  if (!isAntigravityAccountInCooldown(accountId, now)) return true;
+  return isAntigravityRateLimitStickWait(accountId, now);
+}
+
+export function getEligibleAntigravityAccounts(now = Date.now()): string[] {
+  const set = getAccountSet(POOL_KEY_ANTIGRAVITY);
+  if (!set) return [];
+  return set.accounts
+    .filter(account => isAntigravityAccountEligible(account.id, now))
+    .map(account => account.id);
+}
+
+export type AntigravityAccountSelectionReason =
+  | "affinity"
+  | "active"
+  | "failover"
+  | "none"
+  | "all-cooled";
+
+export interface AntigravityAccountSelection {
+  accountId: string | null;
+  reason: AntigravityAccountSelectionReason;
+}
+
+function bindAntigravityAffinityIfPossible(
+  sessionKey: string | null | undefined,
+  accountId: string,
+  now: number,
+): void {
+  bindSessionAffinity(POOL_KEY_ANTIGRAVITY, sessionKey, accountId, now);
+}
+
+/**
+ * Failover-only account pick: stick bound sessions, default new sessions to the store active
+ * account, and hop only when the chosen account is cooled (except short rate-limit stick-wait).
+ * Does not call setActiveAccount.
+ */
+export function resolveAntigravityAccountForSession(
+  sessionKey: string | null | undefined,
+  now = Date.now(),
+): AntigravityAccountSelection {
+  const set = getAccountSet(POOL_KEY_ANTIGRAVITY);
+  if (!set || set.accounts.length === 0) return { accountId: null, reason: "none" };
+
+  const accountIds = set.accounts.map(account => account.id);
+  const activeId = set.activeAccountId;
+  const key = normalizeAffinityComponent(sessionKey);
+
+  if (key) {
+    const affined = getSessionAffinity(POOL_KEY_ANTIGRAVITY, key, now);
+    if (affined) {
+      if (isAntigravityAccountEligible(affined.accountId, now)) {
+        touchSessionAffinity(POOL_KEY_ANTIGRAVITY, key, now);
+        return { accountId: affined.accountId, reason: "affinity" };
+      }
+      const next = nextAntigravityAccount(accountIds, affined.accountId, now);
+      if (next) {
+        bindAntigravityAffinityIfPossible(sessionKey, next, now);
+        return { accountId: next, reason: "failover" };
+      }
+      clearSessionAffinityForAccount(POOL_KEY_ANTIGRAVITY, affined.accountId);
+      const anyCooled = accountIds.some(id => !isAntigravityAccountEligible(id, now));
+      return { accountId: null, reason: anyCooled ? "all-cooled" : "none" };
+    }
+  }
+
+  if (activeId && isAntigravityAccountEligible(activeId, now)) {
+    bindAntigravityAffinityIfPossible(sessionKey, activeId, now);
+    return { accountId: activeId, reason: "active" };
+  }
+
+  const next = nextAntigravityAccount(accountIds, activeId, now);
+  if (next) {
+    bindAntigravityAffinityIfPossible(sessionKey, next, now);
+    return { accountId: next, reason: "failover" };
+  }
+
+  const anyCooled = accountIds.some(id => !isAntigravityAccountEligible(id, now));
+  return { accountId: null, reason: anyCooled ? "all-cooled" : "none" };
+}
+
+export function bindAntigravitySessionAffinity(
+  sessionKey: string | null | undefined,
+  accountId: string,
+  now = Date.now(),
+): void {
+  bindAntigravityAffinityIfPossible(sessionKey, accountId, now);
+}
+
+/**
+ * Pick the next Antigravity account after a rate-limit 429. Cooldown is recorded upstream
+ * (google-http). Does not promote the global active account.
+ */
+export function rotateAntigravityAccountOn429(
+  failedAccountId: string,
+  sessionKey: string | null | undefined,
+  now = Date.now(),
+): string | null {
+  if (antigravity429StickWaitMs(failedAccountId, now) !== null) return null;
+
+  const set = getAccountSet(POOL_KEY_ANTIGRAVITY);
+  if (!set) return null;
+  const accountIds = set.accounts.map(account => account.id);
+
+  clearSessionAffinityForAccount(POOL_KEY_ANTIGRAVITY, failedAccountId);
+
+  const next = nextAntigravityAccount(accountIds, failedAccountId, now);
+  if (!next) return null;
+
+  bindAntigravityAffinityIfPossible(sessionKey, next, now);
+  return next;
+}
+
+/** Test / logout helper. */
+export function clearAntigravityAccountPoolState(): void {
+  clearAffinityState(POOL_KEY_ANTIGRAVITY);
+}
+
+export function antigravitySessionKeyFromParts(input: {
+  sessionIdHeader?: string | null;
+  threadIdHeader?: string | null;
+  promptCacheKey?: string | null;
+  clientThreadId?: string | null;
+  promptCacheKeyIsSharedCohort?: boolean;
+}): string | null {
+  return buildSessionKeyFromParts(input);
+}
+
 export const ANTIGRAVITY_MISSING_PROJECT_MESSAGE =
   "Antigravity requires a discovered Cloud Code Assist project id (re-run `ocx login google-antigravity`).";
 

package/src/oauth/cursor-routing.ts

@@ -0,0 +1,252 @@
+/**
+ * Opt-in Cursor OAuth account pool (shared account-pool kernel).
+ *
+ * Default OFF. When enabled with ≥2 eligible OAuth accounts:
+ * - Sticky session affinity keyed by `_clientThreadId` (not prompt_cache_key)
+ * - Bounded 429/auth failover (cap 3; pre-commit only in core.ts)
+ * - No new-session spreading (failover-only, like Antigravity)
+ *
+ * Does not instantiate or wire `CursorCredentialRouter` weighted RR (#2334).
+ * Kernel must not call setActiveAccount — callers bind `_cursorIdentityScope` only.
+ */
+import { fallbackCodexAccountLogLabel } from "../codex/account-label";
+import {
+  ACCOUNT_POOL_MAX_FAILOVERS,
+  bindSessionAffinity,
+  clearAccountPoolState,
+  clearSessionAffinityForAccount,
+  getSessionAffinity,
+  isAccountPoolEligible,
+  isRateLimitStickWait,
+  normalizeAffinityComponent,
+  recordPoolAccountCooldown,
+  rotatePoolAccountOn429,
+  rotatePoolAccountOnAuth,
+  touchSessionAffinity,
+  type AccountPoolPlugin,
+} from "../routing/account-pool";
+import type { OcxConfig } from "../types";
+import { getAccountSet } from "./store";
+
+export const POOL_KEY_CURSOR = "cursor";
+export const CURSOR_POOL_MAX_FAILOVERS_PER_REQUEST = ACCOUNT_POOL_MAX_FAILOVERS;
+
+export interface CursorAccountPoolConfig {
+  enabled?: boolean;
+}
+
+const cursorPoolPlugin: AccountPoolPlugin = {
+  poolKey: POOL_KEY_CURSOR,
+  sessionKeyFromRequest: cursorSessionKeyFromParts,
+  listEligibleAccountIds(now) {
+    return getEligibleCursorAccountIds(now);
+  },
+};
+
+export function cursorAccountPoolConfig(config: OcxConfig): CursorAccountPoolConfig {
+  const raw = config.cursorAccountPool;
+  if (!raw || typeof raw !== "object") return {};
+  return raw;
+}
+
+export function isCursorAccountPoolEnabled(config: OcxConfig): boolean {
+  return cursorAccountPoolConfig(config).enabled === true;
+}
+
+/** Pool is active only when explicitly enabled and at least two OAuth accounts exist. */
+export function isCursorAccountPoolActive(config: OcxConfig, _now = Date.now()): boolean {
+  if (!isCursorAccountPoolEnabled(config)) return false;
+  const set = getAccountSet(POOL_KEY_CURSOR);
+  if (!set) return false;
+  return set.accounts.filter(account => account.needsReauth !== true).length >= 2;
+}
+
+/**
+ * Sticky key for Cursor OAuth pooling — `_clientThreadId` only (not prompt_cache_key).
+ */
+export function cursorSessionKeyFromParts(input: {
+  clientThreadId?: string | null;
+}): string | null {
+  const key = normalizeAffinityComponent(input.clientThreadId);
+  return key || null;
+}
+
+function isCursorAccountEligible(accountId: string, now: number): boolean {
+  return isAccountPoolEligible(POOL_KEY_CURSOR, accountId, now, {
+    allowStickWait: isRateLimitStickWait(POOL_KEY_CURSOR, accountId, now),
+  });
+}
+
+function getEligibleCursorAccountIds(now: number): string[] {
+  const set = getAccountSet(POOL_KEY_CURSOR);
+  if (!set) return [];
+  return set.accounts
+    .filter(account => account.needsReauth !== true && isCursorAccountEligible(account.id, now))
+    .map(account => account.id);
+}
+
+function nextCursorAccount(
+  accountIds: readonly string[],
+  afterId: string | undefined,
+  now: number,
+): string | undefined {
+  if (accountIds.length === 0) return undefined;
+  const activeIndex = afterId === undefined ? -1 : accountIds.indexOf(afterId);
+  const startIndex = activeIndex < 0 ? 0 : activeIndex + 1;
+  for (let offset = 0; offset < accountIds.length; offset += 1) {
+    const accountId = accountIds[(startIndex + offset) % accountIds.length]!;
+    if (afterId !== undefined && accountId === afterId) continue;
+    if (isCursorAccountEligible(accountId, now)) return accountId;
+  }
+  return undefined;
+}
+
+function bindCursorAffinityIfPossible(
+  sessionKey: string | null | undefined,
+  accountId: string,
+  now: number,
+): void {
+  bindSessionAffinity(POOL_KEY_CURSOR, sessionKey, accountId, now);
+}
+
+export type CursorAccountSelectionReason =
+  | "affinity"
+  | "active"
+  | "failover"
+  | "none"
+  | "all-cooled"
+  | "pool-disabled";
+
+export interface CursorAccountSelection {
+  accountId: string | null;
+  reason: CursorAccountSelectionReason;
+}
+
+/**
+ * Resolve which Cursor OAuth account should serve this session.
+ * When the pool is inactive, returns the store active account.
+ */
+export function resolveCursorAccountForSession(
+  sessionKey: string | null | undefined,
+  config: OcxConfig,
+  now = Date.now(),
+): CursorAccountSelection {
+  const set = getAccountSet(POOL_KEY_CURSOR);
+  if (!set || set.accounts.length === 0) return { accountId: null, reason: "none" };
+
+  if (!isCursorAccountPoolActive(config, now)) {
+    return { accountId: set.activeAccountId, reason: "pool-disabled" };
+  }
+
+  const accountIds = getEligibleCursorAccountIds(now);
+  const activeId = set.activeAccountId;
+  const key = normalizeAffinityComponent(sessionKey);
+
+  if (key) {
+    const affined = getSessionAffinity(POOL_KEY_CURSOR, key, now);
+    if (affined && isCursorAccountEligible(affined.accountId, now)) {
+      touchSessionAffinity(POOL_KEY_CURSOR, key, now);
+      return { accountId: affined.accountId, reason: "affinity" };
+    }
+    if (affined) {
+      const next = nextCursorAccount(accountIds, affined.accountId, now);
+      if (next) {
+        bindCursorAffinityIfPossible(sessionKey, next, now);
+        return { accountId: next, reason: "failover" };
+      }
+      clearSessionAffinityForAccount(POOL_KEY_CURSOR, affined.accountId);
+    }
+  }
+
+  if (activeId && accountIds.includes(activeId)) {
+    bindCursorAffinityIfPossible(sessionKey, activeId, now);
+    return { accountId: activeId, reason: "active" };
+  }
+
+  const next = nextCursorAccount(accountIds, activeId, now);
+  if (next) {
+    bindCursorAffinityIfPossible(sessionKey, next, now);
+    return { accountId: next, reason: "failover" };
+  }
+
+  const anyCooled = set.accounts.some(account => !isCursorAccountEligible(account.id, now));
+  return { accountId: null, reason: anyCooled ? "all-cooled" : "none" };
+}
+
+export function bindCursorSessionAffinity(
+  sessionKey: string | null | undefined,
+  accountId: string,
+  now = Date.now(),
+): void {
+  bindCursorAffinityIfPossible(sessionKey, accountId, now);
+}
+
+export function rotateCursorAccountOn429(
+  config: OcxConfig,
+  failedAccountId: string,
+  retryAfterHeader: string | null | undefined,
+  sessionKey?: string | null,
+  now = Date.now(),
+): string | null {
+  if (!isCursorAccountPoolActive(config, now)) return null;
+  const next = rotatePoolAccountOn429(
+    cursorPoolPlugin,
+    failedAccountId,
+    sessionKey ?? null,
+    retryAfterHeader ?? null,
+    now,
+  );
+  if (next) {
+    console.warn(
+      `[cursor-pool] 429 on ${formatCursorAccountOrdinal(failedAccountId)}; failing over to ${formatCursorAccountOrdinal(next)}`,
+    );
+  }
+  return next;
+}
+
+export function rotateCursorAccountOnAuth(
+  config: OcxConfig,
+  failedAccountId: string,
+  sessionKey: string | null | undefined,
+  now = Date.now(),
+): string | null {
+  if (!isCursorAccountPoolActive(config, now)) return null;
+  const next = rotatePoolAccountOnAuth(
+    cursorPoolPlugin,
+    failedAccountId,
+    sessionKey ?? null,
+    now,
+  );
+  if (next) {
+    console.warn(
+      `[cursor-pool] auth failure on ${formatCursorAccountOrdinal(failedAccountId)}; failing over to ${formatCursorAccountOrdinal(next)}`,
+    );
+  }
+  return next;
+}
+
+/** Billing exhaustion — long cooldown; must not enter the 429 hop carousel. */
+export function recordCursorAccountBillingCooldown(
+  accountId: string,
+  retryAfterHeader: string | null | undefined,
+  now = Date.now(),
+): void {
+  recordPoolAccountCooldown(POOL_KEY_CURSOR, accountId, "billing", retryAfterHeader, now);
+}
+
+/** Test / logout helper. */
+export function clearCursorAccountPoolState(): void {
+  clearAccountPoolState(POOL_KEY_CURSOR);
+}
+
+export function formatCursorAccountOrdinal(accountId: string): string {
+  return fallbackCodexAccountLogLabel(accountId);
+}
+
+export function formatCursorProviderForLog(
+  providerName: string,
+  accountId: string | null | undefined,
+): string {
+  if (!accountId) return providerName;
+  return `${providerName}-${formatCursorAccountOrdinal(accountId)}`;
+}

package/src/oauth/index.ts

@@ -1073,6 +1073,9 @@ export function upsertOAuthProvider(config: OcxConfig, provider: string): void {
   if (existing?.commandCodeVersion !== undefined) {
     next.commandCodeVersion = existing.commandCodeVersion;
   }
+  if (existing?.projectContext !== undefined) {
+    next.projectContext = existing.projectContext;
+  }
   // User-configured price overlays are operator data, not preset state; a
   // re-login, add-account, or reauth must not silently drop them from the
   // Logs/Usage estimates.

package/src/providers/cursor-pool.ts

@@ -5,9 +5,9 @@
  * weighted round-robin selection with per-credential auth-failure cooldown
  * and one-retry failover on a different account before surfacing the error.
  *
- * OpenCodex already has JWT-based multi-account identification (src/oauth/cursor.ts)
- * and Anthropic-specific 429 rotation; this module adds Cursor-aware weighted
- * routing on top of those primitives.
+ * **Not wired in opencodex.** Session-pinned OAuth pooling uses
+ * `src/oauth/cursor-routing.ts` on the shared account-pool kernel instead (#2334).
+ * This module remains for unit tests only until removed.
  */
 
 export interface CursorCredential {

package/src/routing/account-pool/affinity.ts

@@ -0,0 +1,125 @@
+import { createHash } from "node:crypto";
+import { retainedUtf8Bytes } from "../../lib/admission";
+
+export const AFFINITY_IDLE_TTL_MS = 24 * 60 * 60_000;
+export const MAX_AFFINITY_ENTRIES = 2_000;
+export const MAX_AFFINITY_COMPONENT_BYTES = 512;
+
+interface AffinityEntry {
+  accountId: string;
+  lastUsedAt: number;
+}
+
+const affinityByPool = new Map<string, Map<string, AffinityEntry>>();
+
+export function normalizeAffinityComponent(value: string | null | undefined): string {
+  const normalized = value?.trim() ?? "";
+  return normalized && retainedUtf8Bytes(normalized) <= MAX_AFFINITY_COMPONENT_BYTES ? normalized : "";
+}
+
+/**
+ * Build a sticky session key from request headers.
+ * Prefer true session/thread ids; ignore Desktop shared cache-cohort prompt_cache_key.
+ */
+export function buildSessionKeyFromParts(input: {
+  sessionIdHeader?: string | null;
+  threadIdHeader?: string | null;
+  clientThreadId?: string | null;
+  promptCacheKey?: string | null;
+  promptCacheKeyIsSharedCohort?: boolean;
+}): string | null {
+  const preferred = (
+    input.clientThreadId
+    ?? input.sessionIdHeader
+    ?? 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");
+}
+
+function getPoolMap(poolKey: string): Map<string, AffinityEntry> {
+  let map = affinityByPool.get(poolKey);
+  if (!map) {
+    map = new Map();
+    affinityByPool.set(poolKey, map);
+  }
+  return map;
+}
+
+function pruneExpiredAffinity(poolKey: string, now: number): void {
+  const map = affinityByPool.get(poolKey);
+  if (!map) return;
+  for (const [key, entry] of map) {
+    if (now - entry.lastUsedAt > AFFINITY_IDLE_TTL_MS) map.delete(key);
+  }
+  if (map.size <= MAX_AFFINITY_ENTRIES) return;
+  const sorted = [...map.entries()].sort((a, b) => a[1].lastUsedAt - b[1].lastUsedAt);
+  const drop = map.size - MAX_AFFINITY_ENTRIES;
+  for (let i = 0; i < drop; i++) map.delete(sorted[i]![0]);
+}
+
+export function getSessionAffinity(
+  poolKey: string,
+  sessionKey: string | null | undefined,
+  now: number,
+): AffinityEntry | null {
+  pruneExpiredAffinity(poolKey, now);
+  const key = normalizeAffinityComponent(sessionKey);
+  if (!key) return null;
+  const entry = getPoolMap(poolKey).get(key);
+  if (!entry) return null;
+  if (now - entry.lastUsedAt > AFFINITY_IDLE_TTL_MS) {
+    getPoolMap(poolKey).delete(key);
+    return null;
+  }
+  return entry;
+}
+
+export function bindSessionAffinity(
+  poolKey: string,
+  sessionKey: string | null | undefined,
+  accountId: string,
+  now: number,
+): void {
+  const key = normalizeAffinityComponent(sessionKey);
+  if (!key || !normalizeAffinityComponent(accountId)) return;
+  getPoolMap(poolKey).set(key, { accountId, lastUsedAt: now });
+  pruneExpiredAffinity(poolKey, now);
+}
+
+export function touchSessionAffinity(
+  poolKey: string,
+  sessionKey: string | null | undefined,
+  now: number,
+): void {
+  const key = normalizeAffinityComponent(sessionKey);
+  if (!key) return;
+  const entry = getPoolMap(poolKey).get(key);
+  if (entry) entry.lastUsedAt = now;
+}
+
+export function clearSessionAffinityForAccount(poolKey: string, accountId: string): void {
+  const map = affinityByPool.get(poolKey);
+  if (!map) return;
+  for (const [key, entry] of map) {
+    if (entry.accountId === accountId) map.delete(key);
+  }
+}
+
+export function clearAffinityState(poolKey?: string): void {
+  if (poolKey === undefined) {
+    affinityByPool.clear();
+    return;
+  }
+  affinityByPool.delete(poolKey);
+}
+
+export function affinitySizeForTests(poolKey: string): number {
+  return affinityByPool.get(poolKey)?.size ?? 0;
+}