diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 38148e3edd9..e14b791a8c6 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -6,10 +6,6 @@ import { updateProviderConnection, getSettings, getCachedSettings, - getSessionAccountAffinity, - upsertSessionAccountAffinity, - touchSessionAccountAffinity, - deleteSessionAccountAffinity, } from "@/lib/localDb"; import { DEFAULT_QUOTA_THRESHOLD_PERCENT, @@ -60,6 +56,12 @@ import { WEB_COOKIE_PROVIDERS, } from "@/shared/constants/providers"; import { isModelExcludedByConnection } from "@/domain/connectionModelRules"; +import { + applySessionAffinityPin, + formatSessionKeyForLog, + resolveSessionAffinityTtlMs, + selectSessionAffinityConnection, +} from "./sessionAffinityPin"; import { isNoAuthProviderBlockedBySettings } from "./noAuthProviderSettings"; import { resolveAccountProxiesFromRegistry } from "./noAuthProxyResolution"; import * as log from "../utils/logger"; @@ -304,10 +306,6 @@ export function extractSessionAffinityKey( return `input:sha256:${createHash("sha256").update(inputText.slice(0, 4096)).digest("hex")}`; } -function formatSessionKeyForLog(sessionKey: string): string { - return `${sessionKey.slice(0, 18)}...`; -} - function getCodexLimitPolicy(providerSpecificData: JsonRecord): { use5h: boolean; useWeekly: boolean; @@ -713,69 +711,6 @@ function compareP2CConnections( return a.id.localeCompare(b.id); } -function compareLruConnections(a: ProviderConnectionView, b: ProviderConnectionView): number { - if (!a.lastUsedAt && !b.lastUsedAt) return (a.priority || 999) - (b.priority || 999); - if (!a.lastUsedAt) return -1; - if (!b.lastUsedAt) return 1; - const recencyDelta = new Date(a.lastUsedAt).getTime() - new Date(b.lastUsedAt).getTime(); - if (recencyDelta !== 0) return recencyDelta; - if ((a.consecutiveUseCount || 0) !== (b.consecutiveUseCount || 0)) { - return (a.consecutiveUseCount || 0) - (b.consecutiveUseCount || 0); - } - return (a.priority || 999) - (b.priority || 999); -} - -async function selectSessionAffinityConnection( - provider: string, - sessionKey: string | null | undefined, - connections: ProviderConnectionView[], - ttlMs = 0 -): Promise { - if (!sessionKey || connections.length === 0 || ttlMs <= 0) return null; - - const existing = getSessionAccountAffinity(sessionKey, provider, ttlMs); - if (existing) { - const connection = connections.find((candidate) => candidate.id === existing.connectionId); - if (connection) { - touchSessionAccountAffinity(sessionKey, provider, Date.now(), ttlMs); - await updateProviderConnection(connection.id, { - lastUsedAt: new Date().toISOString(), - consecutiveUseCount: (connection.consecutiveUseCount || 0) + 1, - }); - log.info( - "AUTH", - `session_key=${formatSessionKeyForLog(sessionKey)} -> connection ${connection.id.slice( - 0, - 8 - )} (affinity)` - ); - return connection; - } - - deleteSessionAccountAffinity(sessionKey, provider); - log.info( - "AUTH", - `affinity cleared for session_key=${formatSessionKeyForLog(sessionKey)} provider=${provider}` - ); - } - - const connection = [...connections].sort(compareLruConnections)[0] ?? null; - if (!connection) return null; - - upsertSessionAccountAffinity(sessionKey, provider, connection.id, Date.now(), ttlMs); - await updateProviderConnection(connection.id, { - lastUsedAt: new Date().toISOString(), - consecutiveUseCount: 1, - }); - log.info( - "AUTH", - `new affinity created for session_key=${formatSessionKeyForLog( - sessionKey - )} -> connection ${connection.id.slice(0, 8)}` - ); - return connection; -} - /** * Sentinel connection id used for the synthetic credentials of no-auth / * keyless providers. It is NOT a real DB row, so it @@ -1083,7 +1018,7 @@ export async function getProviderCredentials( const allowRateLimitedConnections = allowSuppressedConnections || options.allowRateLimitedConnections === true; const bypassQuotaPolicy = options.bypassQuotaPolicy === true; - const forcedConnectionId = + let forcedConnectionId = typeof options.forcedConnectionId === "string" && options.forcedConnectionId.trim().length > 0 ? options.forcedConnectionId.trim() : null; @@ -1092,6 +1027,11 @@ export async function getProviderCredentials( options.excludeConnectionIds ); + // Fetched early so the session-affinity-pin override (#5903) can consult + // the TTL before forcedConnectionId narrows the connection pool. + const settings = await getSettings(); + const sessionAffinityTtlMs = resolveSessionAffinityTtlMs(provider, options, settings); + // Fix #922: Check for aliases (nvidia/nvidia_nim) to ensure credentials are found const providersToSearch = await getProviderSearchPool(provider); const connectionResults = await Promise.all( @@ -1106,6 +1046,24 @@ export async function getProviderCredentials( if (allowedConnections && allowedConnections.length > 0) { connections = connections.filter((conn) => allowedConnections.includes(conn.id)); } + + // #5903: an active session-affinity pin outranks a per-request reset-aware + // forcedConnectionId (see sessionAffinityPin leaf for the full rationale). + forcedConnectionId = + applySessionAffinityPin({ + forcedConnectionId, + options, + sessionAffinityTtlMs, + connections, + provider, + requestedModel, + excludedConnectionIds, + isTerminalConnectionStatus, + isCodexScopeUnavailable, + isQuotaPolicyBlocked: (c) => + evaluateQuotaLimitPolicy(provider, c as ProviderConnectionView, requestedModel).blocked, + }) ?? forcedConnectionId; + if (forcedConnectionId) { connections = connections.filter((conn) => conn.id === forcedConnectionId); } @@ -1483,18 +1441,7 @@ export async function getProviderCredentials( const orderedConnections = withQuota; - const settings = await getSettings(); const strategy = settings.fallbackStrategy || "fill-first"; - const sessionAffinityTtlMs = - provider === "codex" - ? Number.isFinite(Number(options.sessionAffinityTtlMs)) && - Number(options.sessionAffinityTtlMs) > 0 - ? Number(options.sessionAffinityTtlMs) - : Number.isFinite(Number(settings.codexSessionAffinityTtlMs)) && - Number(settings.codexSessionAffinityTtlMs) > 0 - ? Number(settings.codexSessionAffinityTtlMs) - : 0 - : 0; let connection; const affinityConnection = await selectSessionAffinityConnection( diff --git a/src/sse/services/sessionAffinityPin.ts b/src/sse/services/sessionAffinityPin.ts new file mode 100644 index 00000000000..f8f4749d480 --- /dev/null +++ b/src/sse/services/sessionAffinityPin.ts @@ -0,0 +1,247 @@ +/** + * #5903 — session-affinity-pin resolution + TTL, extracted from auth.ts as a + * pure leaf so the frozen god-file `auth.ts` does not grow. + * + * Problem: reset-aware (and other quota-scoring) combo strategies recompute a + * "winner" connection on every request and hand it to getProviderCredentials + * as `forcedConnectionId`. That id narrows the connection pool to exactly one + * connection BEFORE session affinity is consulted, so an existing pin pointing + * at a previously-selected account is never found and gets silently + * deleted/re-pinned to the fresh winner — breaking "same session -> reuse + * pinned account". + * + * Fix: when an active, non-expired affinity pin already exists for this + * (session, provider) AND the pinned connection is still eligible, the pin wins + * over the freshly recomputed `forcedConnectionId`. If the pin is ineligible + * (rate-limited / exhausted / model-locked / etc.) the caller keeps its forced + * connection, so the existing 429-driven `deleteSessionAccountAffinity` + * failover still owns rotating away from a pin that stops working. + * + * This module stays decoupled from auth.ts internals: the three predicates that + * live in (or would cause a cycle back into) auth.ts — + * `isTerminalConnectionStatus`, `isCodexScopeUnavailable`, and the quota-policy + * check wrapping `evaluateQuotaLimitPolicy` — are injected as callbacks. + */ + +import { + getSessionAccountAffinity, + upsertSessionAccountAffinity, + touchSessionAccountAffinity, + deleteSessionAccountAffinity, +} from "@/lib/db/sessionAccountAffinity"; +import { updateProviderConnection } from "@/lib/db/providers"; +import { isModelExcludedByConnection } from "@/domain/connectionModelRules"; +import { isAccountQuotaExhausted } from "@/domain/quotaCache"; +import { + isAccountUnavailable, + isModelLocked, +} from "@omniroute/open-sse/services/accountFallback.ts"; +import * as log from "../utils/logger"; + +/** Minimal structural view of a provider connection this module reads. */ +export interface AffinityPinConnection { + id: string; + testStatus?: string | null; + rateLimitedUntil?: string | null; + providerSpecificData?: unknown; +} + +/** Fields the LRU tie-break / session-affinity selection reads. */ +export interface SessionAffinityConnection { + id: string; + lastUsedAt?: string | null; + consecutiveUseCount?: number | null; + priority?: number | null; +} + +export function formatSessionKeyForLog(sessionKey: string): string { + return `${sessionKey.slice(0, 18)}...`; +} + +function compareLruConnections(a: SessionAffinityConnection, b: SessionAffinityConnection): number { + if (!a.lastUsedAt && !b.lastUsedAt) return (a.priority || 999) - (b.priority || 999); + if (!a.lastUsedAt) return -1; + if (!b.lastUsedAt) return 1; + const recencyDelta = new Date(a.lastUsedAt).getTime() - new Date(b.lastUsedAt).getTime(); + if (recencyDelta !== 0) return recencyDelta; + if ((a.consecutiveUseCount || 0) !== (b.consecutiveUseCount || 0)) { + return (a.consecutiveUseCount || 0) - (b.consecutiveUseCount || 0); + } + return (a.priority || 999) - (b.priority || 999); +} + +/** + * Session-affinity account selection (moved from auth.ts alongside the #5903 + * pin-override so all session-affinity logic lives in one leaf). Reuses an + * active pin when its connection is in the pool; otherwise picks the LRU + * connection and creates a fresh pin. Behavior byte-identical to the original. + */ +export async function selectSessionAffinityConnection( + provider: string, + sessionKey: string | null | undefined, + connections: T[], + ttlMs = 0 +): Promise { + if (!sessionKey || connections.length === 0 || ttlMs <= 0) return null; + + const existing = getSessionAccountAffinity(sessionKey, provider, ttlMs); + if (existing) { + const connection = connections.find((candidate) => candidate.id === existing.connectionId); + if (connection) { + touchSessionAccountAffinity(sessionKey, provider, Date.now(), ttlMs); + await updateProviderConnection(connection.id, { + lastUsedAt: new Date().toISOString(), + consecutiveUseCount: (connection.consecutiveUseCount || 0) + 1, + }); + log.info( + "AUTH", + `session_key=${formatSessionKeyForLog(sessionKey)} -> connection ${connection.id.slice( + 0, + 8 + )} (affinity)` + ); + return connection; + } + + deleteSessionAccountAffinity(sessionKey, provider); + log.info( + "AUTH", + `affinity cleared for session_key=${formatSessionKeyForLog(sessionKey)} provider=${provider}` + ); + } + + const connection = [...connections].sort(compareLruConnections)[0] ?? null; + if (!connection) return null; + + upsertSessionAccountAffinity(sessionKey, provider, connection.id, Date.now(), ttlMs); + await updateProviderConnection(connection.id, { + lastUsedAt: new Date().toISOString(), + consecutiveUseCount: 1, + }); + log.info( + "AUTH", + `new affinity created for session_key=${formatSessionKeyForLog( + sessionKey + )} -> connection ${connection.id.slice(0, 8)}` + ); + return connection; +} + +/** Subset of credential-selection options the pin resolution consults. */ +export interface AffinityPinOptions { + sessionKey?: string | null; + allowSuppressedConnections?: boolean; + allowRateLimitedConnections?: boolean; + bypassQuotaPolicy?: boolean; + sessionAffinityTtlMs?: number | null; +} + +/** Settings subset needed to resolve the codex session-affinity TTL. */ +export interface AffinityPinSettings { + codexSessionAffinityTtlMs?: number | null; +} + +/** + * Resolve the effective session-affinity TTL. Only codex opts in today: an + * explicit per-request override wins, else the persisted codex setting, else 0 + * (disabled). Kept here so auth.ts can reuse it at both the pin-override site + * and the downstream `selectSessionAffinityConnection` site with one call. + */ +export function resolveSessionAffinityTtlMs( + provider: string, + options: AffinityPinOptions, + settings: AffinityPinSettings +): number { + if (provider !== "codex") return 0; + const override = Number(options.sessionAffinityTtlMs); + if (Number.isFinite(override) && override > 0) return override; + const configured = Number(settings.codexSessionAffinityTtlMs); + if (Number.isFinite(configured) && configured > 0) return configured; + return 0; +} + +/** + * Predicates supplied by the caller because they either live in auth.ts or + * would introduce a circular import if pulled in directly. + */ +export interface AffinityPinPredicates { + /** auth.ts::isTerminalConnectionStatus (banned/expired/credits_exhausted). */ + isTerminalConnectionStatus: (connection: AffinityPinConnection) => boolean; + /** auth.ts::isCodexScopeUnavailable (codex per-scope cooldown). */ + isCodexScopeUnavailable: ( + connection: AffinityPinConnection, + requestedModel: string | null + ) => boolean; + /** Wraps auth.ts::evaluateQuotaLimitPolicy(...).blocked for one connection. */ + isQuotaPolicyBlocked: (connection: AffinityPinConnection) => boolean; +} + +export interface ApplySessionAffinityPinParams extends AffinityPinPredicates { + forcedConnectionId: string | null; + options: AffinityPinOptions; + sessionAffinityTtlMs: number; + connections: AffinityPinConnection[]; + provider: string; + requestedModel: string | null; + excludedConnectionIds: Set; +} + +/** + * Mirrors the eligibility predicates applied later in getProviderCredentials + * (availableConnections filter + quota policy + quota exhaustion) but scoped to + * a single candidate connection. Pure/read-only. + */ +function isConnectionEligibleForAffinityPin( + connection: AffinityPinConnection, + params: ApplySessionAffinityPinParams +): boolean { + const { provider, requestedModel, options } = params; + const allowSuppressed = options.allowSuppressedConnections === true; + const allowRateLimited = allowSuppressed || options.allowRateLimitedConnections === true; + if (params.excludedConnectionIds.has(connection.id)) return false; + if ( + requestedModel && + isModelExcludedByConnection(requestedModel, connection.providerSpecificData) + ) { + return false; + } + if (!allowSuppressed) { + if (!allowRateLimited && isAccountUnavailable(connection.rateLimitedUntil)) return false; + if (params.isTerminalConnectionStatus(connection)) return false; + if (provider === "codex" && params.isCodexScopeUnavailable(connection, requestedModel)) { + return false; + } + if (requestedModel && isModelLocked(provider, connection.id, requestedModel)) return false; + } + if (isAccountQuotaExhausted(connection.id)) return false; + if (options.bypassQuotaPolicy !== true && params.isQuotaPolicyBlocked(connection)) return false; + return true; +} + +/** + * If an active, non-expired affinity pin exists for (sessionKey, provider) and + * the pinned connection is present-and-eligible in the current pool, returns + * that pinned connectionId (which should override `forcedConnectionId`) and + * logs the override. Returns null when the caller should keep its + * `forcedConnectionId` — no session, TTL disabled, no pin, pin already equals + * the forced id, pin absent from pool, or pin ineligible. + */ +export function applySessionAffinityPin(params: ApplySessionAffinityPinParams): string | null { + const { forcedConnectionId, options, sessionAffinityTtlMs, connections, provider } = params; + const sessionKey = options.sessionKey; + if (!forcedConnectionId || !sessionKey || sessionAffinityTtlMs <= 0) return null; + + const pinned = getSessionAccountAffinity(sessionKey, provider, sessionAffinityTtlMs); + if (!pinned || pinned.connectionId === forcedConnectionId) return null; + + const pinnedConnection = connections.find((conn) => conn.id === pinned.connectionId); + if (!pinnedConnection || !isConnectionEligibleForAffinityPin(pinnedConnection, params)) { + return null; + } + + log.info( + "AUTH", + `session affinity pin ${pinned.connectionId.slice(0, 8)}... overrides forcedConnectionId ${forcedConnectionId.slice(0, 8)}... (#5903)` + ); + return pinned.connectionId; +} diff --git a/tests/unit/codex-session-affinity-reset-aware-5903.test.ts b/tests/unit/codex-session-affinity-reset-aware-5903.test.ts new file mode 100644 index 00000000000..02f4ef6f1fb --- /dev/null +++ b/tests/unit/codex-session-affinity-reset-aware-5903.test.ts @@ -0,0 +1,181 @@ +// #5903: Codex session affinity must win over a per-request reset-aware +// re-scoring. The reset-aware combo strategy (open-sse/services/combo/quotaStrategies.ts) +// recomputes its "winner" connection on every request and hands it to +// getProviderCredentials as forcedConnectionId (src/sse/handlers/chat.ts). +// Before the fix, forcedConnectionId narrowed the connection pool BEFORE +// session affinity was consulted, so a fresh quota-scoring winner silently +// evicted the existing pin (deleteSessionAccountAffinity) on every request — +// breaking "same session -> reuse pinned account". +// +// This test drives auth.getProviderCredentials directly (the same call shape +// chat.ts uses: sessionKey + forcedConnectionId together) to reproduce the +// bug without needing the full combo/quota-scoring machinery. + +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-codex-affinity-5903-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = process.env.API_KEY_SECRET || "codex-affinity-5903-test-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const settingsDb = await import("../../src/lib/db/settings.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); +const affinityDb = await import("../../src/lib/db/sessionAccountAffinity.ts"); +const auth = await import("../../src/sse/services/auth.ts"); + +async function resetStorage() { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +async function seedConnection(provider: string, overrides: any = {}) { + return providersDb.createProviderConnection({ + provider, + authType: overrides.authType || "oauth", + name: overrides.name || `${provider}-${Math.random().toString(16).slice(2, 8)}`, + accessToken: overrides.accessToken || `at-${Math.random().toString(16).slice(2, 10)}`, + refreshToken: overrides.refreshToken, + isActive: overrides.isActive ?? true, + testStatus: overrides.testStatus || "active", + priority: overrides.priority, + providerSpecificData: overrides.providerSpecificData || {}, + }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("codex session affinity wins over a per-request reset-aware forcedConnectionId (#5903)", async () => { + await settingsDb.updateSettings({ + fallbackStrategy: "reset-aware", + codexSessionAffinityTtlMs: 60_000, + }); + + const connectionA = await seedConnection("codex", { name: "codex-reset-aware-a" }); + const connectionB = await seedConnection("codex", { name: "codex-reset-aware-b" }); + + // Request 1: reset-aware quota scoring picks A as the winner for session S. + const request1 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-S", + forcedConnectionId: connectionA.id, + }); + assert.equal(request1?.connectionId, connectionA.id, "request 1 should pin to the scored winner A"); + assert.equal( + affinityDb.getSessionAccountAffinity("session-S", "codex", 60_000)?.connectionId, + connectionA.id, + "affinity row must be created for session-S pointing at A" + ); + + // Request 2: quota state shifted and reset-aware now scores B higher for + // the SAME session. Without the fix, forcedConnectionId=B narrows the pool + // to just B before affinity is checked, evicting the A pin and re-pinning + // to B. With the fix, the existing active pin (A) must win. + const request2 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-S", + forcedConnectionId: connectionB.id, + }); + assert.equal( + request2?.connectionId, + connectionA.id, + "request 2 must still use the pinned connection A, not the freshly re-scored B" + ); + assert.equal( + affinityDb.getSessionAccountAffinity("session-S", "codex", 60_000)?.connectionId, + connectionA.id, + "affinity row for session-S must remain pinned to A after re-scoring" + ); + + // A brand-new session (S2) has no existing pin, so the freshly re-scored + // winner (B) must be honored and a NEW pin created for S2. + const request3 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-S2", + forcedConnectionId: connectionB.id, + }); + assert.equal(request3?.connectionId, connectionB.id, "a new session must honor the fresh re-scored pick"); + assert.equal( + affinityDb.getSessionAccountAffinity("session-S2", "codex", 60_000)?.connectionId, + connectionB.id, + "a new affinity row for session-S2 must be created pointing at B" + ); + + // Session S must remain unaffected by S2's independent pin. + assert.equal( + affinityDb.getSessionAccountAffinity("session-S", "codex", 60_000)?.connectionId, + connectionA.id, + "session-S pin must stay isolated from session-S2" + ); +}); + +test("reset-aware forcedConnectionId is honored when the pinned connection becomes ineligible (#5903)", async () => { + await settingsDb.updateSettings({ + fallbackStrategy: "reset-aware", + codexSessionAffinityTtlMs: 60_000, + }); + + const connectionA = await seedConnection("codex", { name: "codex-reset-aware-ineligible-a" }); + const connectionB = await seedConnection("codex", { name: "codex-reset-aware-ineligible-b" }); + + const request1 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-failover", + forcedConnectionId: connectionA.id, + }); + assert.equal(request1?.connectionId, connectionA.id); + + // A becomes rate-limited (e.g. 429 handled by markAccountUnavailable in + // production). Reset-aware re-scores and now forces B. The pin (A) is no + // longer eligible, so the freshly forced B must be used instead of + // failing the whole request. + await providersDb.updateProviderConnection(connectionA.id, { + rateLimitedUntil: new Date(Date.now() + 60_000).toISOString(), + }); + + const request2 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-failover", + forcedConnectionId: connectionB.id, + }); + assert.equal( + request2?.connectionId, + connectionB.id, + "an ineligible pin must fall through to the freshly forced connection" + ); +}); + +test("no session affinity configured: reset-aware forcedConnectionId applies exactly as before (#5903)", async () => { + await settingsDb.updateSettings({ + fallbackStrategy: "reset-aware", + codexSessionAffinityTtlMs: 0, + }); + + const connectionA = await seedConnection("codex", { name: "codex-no-affinity-a" }); + const connectionB = await seedConnection("codex", { name: "codex-no-affinity-b" }); + + const request1 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-no-ttl", + forcedConnectionId: connectionA.id, + }); + assert.equal(request1?.connectionId, connectionA.id); + + const request2 = await auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "session-no-ttl", + forcedConnectionId: connectionB.id, + }); + assert.equal( + request2?.connectionId, + connectionB.id, + "with affinity disabled (ttl=0) each request must honor the fresh forcedConnectionId" + ); +});