From d774ccecac46f7315a794d5ce33653eaac3a8192 Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Tue, 11 Aug 2026 14:51:35 +0200 Subject: [PATCH] fix(routing): account for active OAuth sessions (#8940) --- open-sse/services/autoCombo/modePacks.ts | 6 + open-sse/services/autoCombo/scoring.ts | 6 + open-sse/services/combo.ts | 16 +- .../services/combo/applyStrategyOrdering.ts | 5 +- .../services/combo/promptCacheAffinity.ts | 64 ++++++-- .../services/combo/resolveAutoStrategy.ts | 15 +- open-sse/services/combo/targetResolution.ts | 4 +- open-sse/services/combo/types.ts | 6 +- open-sse/services/oauthSessionOccupancy.ts | 114 ++++++++++++++ src/i18n/messages/de.json | 1 + src/i18n/messages/en.json | 1 + src/lib/combos/intelligentRouting.ts | 6 + src/lib/db/providers/lazyConnectionView.ts | 21 +-- src/shared/validation/schemas/combo.ts | 1 + src/sse/handlers/chat.ts | 144 +++++++++++------- src/sse/services/auth.ts | 65 +++++++- tests/unit/oauth-session-occupancy.test.ts | 61 ++++++++ tests/unit/prompt-cache-affinity.test.ts | 41 +++++ tests/unit/sse-auth.test.ts | 41 +++++ 19 files changed, 517 insertions(+), 101 deletions(-) create mode 100644 open-sse/services/oauthSessionOccupancy.ts create mode 100644 tests/unit/oauth-session-occupancy.test.ts diff --git a/open-sse/services/autoCombo/modePacks.ts b/open-sse/services/autoCombo/modePacks.ts index 7344496d10..5c0bfd54cc 100644 --- a/open-sse/services/autoCombo/modePacks.ts +++ b/open-sse/services/autoCombo/modePacks.ts @@ -24,6 +24,7 @@ export const MODE_PACKS: Record = { tierAffinity: 0, specificityMatch: 0, contextAffinity: 0.01, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }, @@ -39,6 +40,7 @@ export const MODE_PACKS: Record = { tierAffinity: 0, specificityMatch: 0, contextAffinity: 0.0, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }, @@ -54,6 +56,7 @@ export const MODE_PACKS: Record = { tierAffinity: 0, specificityMatch: 0, contextAffinity: 0.0, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }, @@ -69,6 +72,7 @@ export const MODE_PACKS: Record = { tierAffinity: 0, specificityMatch: 0, contextAffinity: 0.0, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }, @@ -85,6 +89,7 @@ export const MODE_PACKS: Record = { tierAffinity: 0, specificityMatch: 0, contextAffinity: 0.0, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }, @@ -105,6 +110,7 @@ export const MODE_PACKS: Record = { tierAffinity: 0, specificityMatch: 0, contextAffinity: 0.03, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }, diff --git a/open-sse/services/autoCombo/scoring.ts b/open-sse/services/autoCombo/scoring.ts index a8758df00c..f64ab607f2 100644 --- a/open-sse/services/autoCombo/scoring.ts +++ b/open-sse/services/autoCombo/scoring.ts @@ -20,6 +20,7 @@ export interface ScoringFactors { specificityMatch: number; contextAffinity: number; cacheAffinity?: number; + sessionAvailability?: number; resetWindowAffinity: number; connectionDensity: number; } @@ -36,6 +37,7 @@ export interface ScoringWeights { specificityMatch: number; contextAffinity: number; cacheAffinity?: number; + sessionAvailability?: number; resetWindowAffinity: number; connectionDensity: number; } @@ -52,6 +54,7 @@ export const DEFAULT_WEIGHTS: ScoringWeights = { specificityMatch: 0.05, contextAffinity: 0.05, cacheAffinity: 0, + sessionAvailability: 0.05, resetWindowAffinity: 0, connectionDensity: 0.05, }; @@ -101,6 +104,7 @@ export interface ProviderCandidate { contextAffinity?: number; /** Score [0..1] for the account selected by the stable prompt-cache key. */ cacheAffinity?: number; + sessionAvailability?: number; /** Score [0..1] for quota reset-window preference; sooner selected reset windows score higher. */ resetWindowAffinity?: number; connectionPoolSize?: number; @@ -135,6 +139,7 @@ export function calculateScore(factors: ScoringFactors, weights: ScoringWeights) (weights.specificityMatch ?? 0) * factors.specificityMatch + (weights.contextAffinity ?? 0) * factors.contextAffinity + (weights.cacheAffinity ?? 0) * (factors.cacheAffinity ?? 0) + + (weights.sessionAvailability ?? 0) * (factors.sessionAvailability ?? 1) + (weights.resetWindowAffinity ?? 0) * factors.resetWindowAffinity + (weights.connectionDensity ?? 0) * factors.connectionDensity ); @@ -260,6 +265,7 @@ export function calculateFactors( specificityMatch: calculateSpecificityMatch(candidate, manifestHint), contextAffinity: clamp01(candidate.contextAffinity ?? 0.5), cacheAffinity: clamp01(candidate.cacheAffinity ?? 0), + sessionAvailability: clamp01(candidate.sessionAvailability ?? 1), resetWindowAffinity: clamp01(candidate.resetWindowAffinity ?? 0.5), connectionDensity: clamp01(((candidate.connectionPoolSize ?? 1) - 1) / 10), }; diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 8b23ad3d33..cfeaf247ca 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -74,6 +74,7 @@ import { notifyWebhookEvent } from "../../src/lib/webhookDispatcher"; import { type ProviderCandidate } from "./autoCombo/scoring.ts"; import { estimateTokens } from "./contextManager.ts"; import { getSessionConnection } from "./sessionManager.ts"; +import { getOAuthSessionAvailability } from "./oauthSessionOccupancy.ts"; import { applySessionStickiness, normalizeStickinessMessages, @@ -451,6 +452,9 @@ export async function buildAutoCandidates( let quotaCutoffReason: string | undefined; const fetcher = getQuotaFetcher(provider); const connection = target.connectionId ? connectionById.get(target.connectionId) : undefined; + const authType = typeof connection?.authType === "string" ? connection.authType : null; + const sessionAvailability = + authType === "oauth" ? getOAuthSessionAvailability(target.connectionId, sessionId) : 1; // Gate the terminal-status cutoff behind the same opt-in as the quota-percent // cutoff (#4483): when quota cutoff is disabled, a connection in a terminal // testStatus must still fall through to normal connection-cooldown / model-lockout @@ -525,6 +529,7 @@ export async function buildAutoCandidates( accountTier: "standard" as const, quotaResetIntervalSecs: 86400, contextAffinity, + sessionAvailability, resetWindowAffinity, quotaCutoffBlocked, quotaCutoffReason, @@ -532,6 +537,7 @@ export async function buildAutoCandidates( statusPenaltyReason, connectionPoolSize: connectionPoolCounts.get(provider) ?? 1, connectionId: target.connectionId ?? undefined, + authType, }; }) ); @@ -705,6 +711,7 @@ export async function handleComboChat({ signal, hiddenModelsByProvider, clientManagedResponsesContext, + relayOptions, }); } @@ -2336,6 +2343,7 @@ async function handleRoundRobinCombo({ nesting = null, hiddenModelsByProvider = getHiddenModelsByProvider(), clientManagedResponsesContext, + relayOptions, }: HandleRoundRobinOptions): Promise { const config = settings ? resolveComboConfig(combo, settings) @@ -2563,7 +2571,13 @@ async function handleRoundRobinCombo({ // stickiness engages on the /v1/responses surface, not just Chat Completions. normalizeStickinessMessages(body as { messages?: unknown; input?: unknown }) ); - const rrAffinity = applyPromptCacheAffinity(filteredTargets, body, rrAffinityEnabled); + const rrAffinity = applyPromptCacheAffinity( + filteredTargets, + body, + rrAffinityEnabled, + "global", + relayOptions?.sessionId + ); if (rrAffinity.applied) { const stickyFirst = _rrSessionSticky.stuck ? _rrSessionSticky.targets[0] : null; filteredTargets = stickyFirst diff --git a/open-sse/services/combo/applyStrategyOrdering.ts b/open-sse/services/combo/applyStrategyOrdering.ts index 34514c881d..38f07fbcdc 100644 --- a/open-sse/services/combo/applyStrategyOrdering.ts +++ b/open-sse/services/combo/applyStrategyOrdering.ts @@ -26,6 +26,7 @@ export interface ApplyStrategyOrderingDeps { body: Record; log: ComboLogger; apiKeyAllowedConnections: string[] | null; + sessionKey?: string | null; } /** @@ -45,7 +46,7 @@ export async function applyStrategyOrdering( initialOrderedTargets: ResolvedComboTarget[], deps: ApplyStrategyOrderingDeps ): Promise { - const { combo, config, body, log, apiKeyAllowedConnections } = deps; + const { combo, config, body, log, apiKeyAllowedConnections, sessionKey } = deps; let orderedTargets = initialOrderedTargets; if (strategy === "lkgp") { @@ -205,7 +206,7 @@ export async function applyStrategyOrdering( if (resolvePromptCacheAffinityKey(body)) { orderedTargets = await expandPromptCacheAffinityTargets(orderedTargets); } - const affinity = applyPromptCacheAffinity(orderedTargets, body, true, "global"); + const affinity = applyPromptCacheAffinity(orderedTargets, body, true, "global", sessionKey); orderedTargets = affinity.targets; log.info( "COMBO", diff --git a/open-sse/services/combo/promptCacheAffinity.ts b/open-sse/services/combo/promptCacheAffinity.ts index 4433e7b69d..f7675f0675 100644 --- a/open-sse/services/combo/promptCacheAffinity.ts +++ b/open-sse/services/combo/promptCacheAffinity.ts @@ -6,10 +6,12 @@ import { import { getCachedProviderConnections } from "../../../src/lib/db/readCache"; import { parseModel } from "../model.ts"; import type { ResolvedComboTarget } from "./types.ts"; +import { getOAuthSessionAvailability } from "../oauthSessionOccupancy.ts"; interface PromptCacheAffinityTarget { executionKey: string; connectionId?: string | null; + authType?: string | null; } export type PromptCacheAffinitySource = "explicit" | "prefix"; @@ -126,6 +128,23 @@ function rendezvousScore(key: string, identity: string): bigint { return BigInt(`0x${digest.slice(0, 32)}`); } +const MAX_RENDEZVOUS_HIGH_BITS = (1n << 64n) - 1n; + +function normalizedRendezvousScore(key: string, identity: string): number { + return Number(rendezvousScore(key, identity) >> 64n) / Number(MAX_RENDEZVOUS_HIGH_BITS); +} + +function combinedAffinityScore( + key: string, + target: PromptCacheAffinityTarget, + sessionKey?: string | null +): number { + const cacheScore = normalizedRendezvousScore(key, promptCacheTargetIdentity(target)); + const availability = + target.authType === "oauth" ? getOAuthSessionAvailability(target.connectionId, sessionKey) : 1; + return cacheScore * 0.75 + availability * 0.25; +} + /** * Return a normalized cache-locality score for auto-combo scoring. The target * selected by rendezvous hashing receives 1; all other accounts receive 0. @@ -133,15 +152,16 @@ function rendezvousScore(key: string, identity: string): bigint { */ export function calculatePromptCacheAffinityScores( targets: PromptCacheAffinityTarget[], - body: Record | null | undefined + body: Record | null | undefined, + sessionKey?: string | null ): Map { const resolution = resolvePromptCacheAffinityKey(body); if (!resolution || targets.length === 0) return new Map(); let winnerIdentity = ""; - let winnerScore = -1n; + let winnerScore = -1; for (const target of targets) { const identity = promptCacheTargetIdentity(target); - const score = rendezvousScore(resolution.key, identity); + const score = combinedAffinityScore(resolution.key, target, sessionKey); if (score > winnerScore || (score === winnerScore && identity < winnerIdentity)) { winnerIdentity = identity; winnerScore = score; @@ -166,15 +186,13 @@ export async function expandPromptCacheAffinityTargets( ): Promise { const providers = Array.from( new Set( - targets - .filter((target) => !target.connectionId) - .map( - (target) => - target.provider || - parseModel(target.modelStr).provider || - parseModel(target.modelStr).providerAlias || - "unknown" - ) + targets.map( + (target) => + target.provider || + parseModel(target.modelStr).provider || + parseModel(target.modelStr).providerAlias || + "unknown" + ) ) ); const connectionsByProvider = new Map>>(); @@ -201,7 +219,18 @@ export function expandPromptCacheAffinityTargetsFromConnections( const expandedTargets: ResolvedComboTarget[] = []; for (const target of targets) { if (target.connectionId) { - expandedTargets.push(target); + const provider = + target.provider || + parseModel(target.modelStr).provider || + parseModel(target.modelStr).providerAlias || + "unknown"; + const connection = (connectionsByProvider.get(provider) || []).find( + (candidate) => candidate?.id === target.connectionId + ); + expandedTargets.push({ + ...target, + authType: typeof connection?.authType === "string" ? connection.authType : target.authType, + }); continue; } const parsed = parseModel(target.modelStr); @@ -227,9 +256,13 @@ export function expandPromptCacheAffinityTargetsFromConnections( continue; } for (const connectionId of scopedConnectionIds) { + const connection = (connectionsByProvider.get(provider) || []).find( + (candidate) => candidate?.id === connectionId + ); expandedTargets.push({ ...target, connectionId, + authType: typeof connection?.authType === "string" ? connection.authType : null, executionKey: `${target.executionKey}@${connectionId}`, }); } @@ -291,7 +324,8 @@ export function applyPromptCacheAffinity( targets: ResolvedComboTarget[], body: Record | null | undefined, enabled: boolean = true, - scope: "model" | "global" = "global" + scope: "model" | "global" = "global", + sessionKey?: string | null ): PromptCacheAffinityResult { const resolution = enabled ? resolvePromptCacheAffinityKey(body) : null; if (!resolution || targets.length <= 1) { @@ -307,7 +341,7 @@ export function applyPromptCacheAffinity( target, index, identity: promptCacheTargetIdentity(target), - score: rendezvousScore(resolution.key, promptCacheTargetIdentity(target)), + score: combinedAffinityScore(resolution.key, target, sessionKey), baseModel: scope === "model" ? getBaseModelIdentity(target) : null, })); diff --git a/open-sse/services/combo/resolveAutoStrategy.ts b/open-sse/services/combo/resolveAutoStrategy.ts index 6e0ecad032..1d6d2bf156 100644 --- a/open-sse/services/combo/resolveAutoStrategy.ts +++ b/open-sse/services/combo/resolveAutoStrategy.ts @@ -1,4 +1,8 @@ -import { errorResponse, unavailableResponse, errorResponseWithComboDiagnostics } from "../../utils/error.ts"; +import { + errorResponse, + unavailableResponse, + errorResponseWithComboDiagnostics, +} from "../../utils/error.ts"; import { BudgetExceededError, selectProvider as selectAutoProvider } from "../autoCombo/engine.ts"; import { resolveRequestModePack, @@ -118,8 +122,7 @@ export async function resolveAutoStrategyOrder( // registry/capability rows honestly report toolCalling:false. const filtered = eligibleTargets.filter( (target) => - supportsToolCalling(target.modelStr) || - providerSupportsEmulatedToolCalling(target.provider) + supportsToolCalling(target.modelStr) || providerSupportsEmulatedToolCalling(target.provider) ); if (filtered.length > 0) { eligibleTargets = filtered; @@ -287,7 +290,11 @@ export async function resolveAutoStrategyOrder( resetWindowConfig, autoCandidateResilienceSettings ); - const cacheAffinityScores = calculatePromptCacheAffinityScores(candidates, body); + const cacheAffinityScores = calculatePromptCacheAffinityScores( + candidates, + body, + relayOptions?.sessionId + ); for (const candidate of candidates) { candidate.cacheAffinity = cacheAffinityScores.get(promptCacheTargetIdentity(candidate)) ?? 0; } diff --git a/open-sse/services/combo/targetResolution.ts b/open-sse/services/combo/targetResolution.ts index 8ef9e2bc97..40ac191ad6 100644 --- a/open-sse/services/combo/targetResolution.ts +++ b/open-sse/services/combo/targetResolution.ts @@ -453,6 +453,7 @@ async function orderByStrategy( body, log, apiKeyAllowedConnections: deps.apiKeyAllowedConnections, + sessionKey: deps.relayOptions?.sessionId, }); return { orderedTargets, autoUsedExplicitRouter: false }; } @@ -682,7 +683,8 @@ async function applyPromptCacheStage( promptCacheAffinityTargets, body, promptCacheAffinityEnabled, - isDeterministicStrategy ? "model" : "global" + isDeterministicStrategy ? "model" : "global", + deps.relayOptions?.sessionId ); if (!promptCacheAffinity.applied) return orderedTargets; const protectedOriginal = diff --git a/open-sse/services/combo/types.ts b/open-sse/services/combo/types.ts index 71b4e849ec..9f9f31c4b4 100644 --- a/open-sse/services/combo/types.ts +++ b/open-sse/services/combo/types.ts @@ -114,10 +114,7 @@ export type HandleComboChatOptions = { clientManagedResponsesContext?: boolean; }; -export type HandleRoundRobinOptions = Omit< - HandleComboChatOptions, - "relayOptions" | "apiKeyAllowedConnections" ->; +export type HandleRoundRobinOptions = Omit; export type HistoricalLatencyStatsEntry = { totalRequests?: number; @@ -167,6 +164,7 @@ export type ResolvedComboTarget = { executionKey: string; modelStr: string; provider: string; + authType?: string | null; providerId: string | null; connectionId: string | null; allowedConnectionIds?: string[] | null; diff --git a/open-sse/services/oauthSessionOccupancy.ts b/open-sse/services/oauthSessionOccupancy.ts new file mode 100644 index 0000000000..e9d48386c8 --- /dev/null +++ b/open-sse/services/oauthSessionOccupancy.ts @@ -0,0 +1,114 @@ +const DEFAULT_LEASE_MS = 10 * 60_000; + +interface SessionLease { + requests: number; + expiresAt: number; +} + +const occupancy = new Map>(); + +function prune(now = Date.now()): void { + for (const [connectionId, sessions] of occupancy) { + for (const [sessionKey, lease] of sessions) { + if (lease.expiresAt <= now) sessions.delete(sessionKey); + } + if (sessions.size === 0) occupancy.delete(connectionId); + } +} + +export function getForeignOAuthSessionCount( + connectionId: string | null | undefined, + sessionKey: string | null | undefined, + now = Date.now() +): number { + if (!connectionId) return 0; + prune(now); + const sessions = occupancy.get(connectionId); + if (!sessions) return 0; + let count = 0; + for (const key of sessions.keys()) { + if (!sessionKey || key !== sessionKey) count++; + } + return count; +} + +export function getOAuthSessionAvailability( + connectionId: string | null | undefined, + sessionKey: string | null | undefined, + now = Date.now() +): number { + return 1 / (1 + getForeignOAuthSessionCount(connectionId, sessionKey, now)); +} + +export function reserveOAuthSession( + connectionId: string, + sessionKey: string, + leaseMs = DEFAULT_LEASE_MS, + now = Date.now() +): () => void { + if (!connectionId || !sessionKey) return () => {}; + prune(now); + const sessions = occupancy.get(connectionId) ?? new Map(); + const current = sessions.get(sessionKey); + sessions.set(sessionKey, { + requests: (current?.requests ?? 0) + 1, + expiresAt: now + Math.max(1, leaseMs), + }); + occupancy.set(connectionId, sessions); + + let released = false; + return () => { + if (released) return; + released = true; + const activeSessions = occupancy.get(connectionId); + const active = activeSessions?.get(sessionKey); + if (!activeSessions || !active) return; + if (active.requests <= 1) activeSessions.delete(sessionKey); + else activeSessions.set(sessionKey, { ...active, requests: active.requests - 1 }); + if (activeSessions.size === 0) occupancy.delete(connectionId); + }; +} + +export function wrapResponseWithOAuthSessionRelease( + response: Response, + release: () => void +): Response { + if (!response.body) { + release(); + return response; + } + const reader = response.body.getReader(); + const body = new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + release(); + controller.close(); + return; + } + controller.enqueue(value); + } catch (error) { + release(); + controller.error(error); + } + }, + async cancel(reason) { + release(); + try { + await reader.cancel(reason); + } catch { + // The upstream stream is already closing; the lease has still been released. + } + }, + }); + return new Response(body, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }); +} + +export function _clearOAuthSessionOccupancyForTest(): void { + occupancy.clear(); +} diff --git a/src/i18n/messages/de.json b/src/i18n/messages/de.json index 00f9c2053a..d993330be8 100644 --- a/src/i18n/messages/de.json +++ b/src/i18n/messages/de.json @@ -3324,6 +3324,7 @@ "weightStability": "Stability", "weightTierPriority": "Tier", "weightCacheAffinity": "Cache-Treffer-Affinität", + "weightSessionAvailability": "Session-Verfügbarkeit", "reviewIntelligentTitle": "Intelligent Routing Config", "strategyRecommendations": { "priority": { diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 8af252a53f..bfae814002 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -3338,6 +3338,7 @@ "weightStability": "Stability", "weightTierPriority": "Tier", "weightCacheAffinity": "Cache Hit Affinity", + "weightSessionAvailability": "Session Availability", "reviewIntelligentTitle": "Intelligent Routing Config", "strategyRecommendations": { "priority": { diff --git a/src/lib/combos/intelligentRouting.ts b/src/lib/combos/intelligentRouting.ts index 13cf531348..976e13e047 100644 --- a/src/lib/combos/intelligentRouting.ts +++ b/src/lib/combos/intelligentRouting.ts @@ -17,6 +17,7 @@ export type IntelligentRoutingWeights = { specificityMatch: number; contextAffinity: number; cacheAffinity: number; + sessionAvailability: number; resetWindowAffinity: number; }; @@ -52,6 +53,7 @@ export const DEFAULT_INTELLIGENT_WEIGHTS: IntelligentRoutingWeights = { specificityMatch: 0.05, contextAffinity: 0.08, cacheAffinity: 0, + sessionAvailability: 0.05, resetWindowAffinity: 0, }; @@ -82,6 +84,7 @@ export const FACTOR_LABELS: Record = { specificityMatch: "Specificity", contextAffinity: "Context Affinity", cacheAffinity: "Cache Hit Affinity", + sessionAvailability: "Session Availability", resetWindowAffinity: "Reset Window", }; @@ -158,6 +161,9 @@ export function normalizeIntelligentRoutingConfig(config: unknown): IntelligentR toFiniteNumber(rawWeights.contextAffinity) ?? DEFAULT_INTELLIGENT_WEIGHTS.contextAffinity, cacheAffinity: toFiniteNumber(rawWeights.cacheAffinity) ?? DEFAULT_INTELLIGENT_WEIGHTS.cacheAffinity, + sessionAvailability: + toFiniteNumber(rawWeights.sessionAvailability) ?? + DEFAULT_INTELLIGENT_WEIGHTS.sessionAvailability, resetWindowAffinity: toFiniteNumber(rawWeights.resetWindowAffinity) ?? DEFAULT_INTELLIGENT_WEIGHTS.resetWindowAffinity, diff --git a/src/lib/db/providers/lazyConnectionView.ts b/src/lib/db/providers/lazyConnectionView.ts index 7a52b1936d..db7f215ee7 100644 --- a/src/lib/db/providers/lazyConnectionView.ts +++ b/src/lib/db/providers/lazyConnectionView.ts @@ -14,9 +14,7 @@ import { decrypt } from "../encryption"; type JsonRecord = Record; function asRecord(value: unknown): JsonRecord { - return value && typeof value === "object" && !Array.isArray(value) - ? (value as JsonRecord) - : {}; + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; } function toStringOrNull(value: unknown): string | null { @@ -41,6 +39,7 @@ function toNullableNumber(value: unknown): number | null { export interface ProviderConnectionView { id: string; provider: string; + authType: string | null; email: string | null; isActive: boolean; rateLimitedUntil: string | null; @@ -79,6 +78,7 @@ export function toProviderConnection(value: unknown): ProviderConnectionView { return { id: toStringOrNull(row.id) || "", provider: toStringOrNull(row.provider) || "", + authType: toStringOrNull(row.authType), email: toStringOrNull(row.email), isActive: row.isActive === true, rateLimitedUntil: toStringOrNull(row.rateLimitedUntil), @@ -98,9 +98,7 @@ export function toProviderConnection(value: unknown): ProviderConnectionView { lastErrorType: toStringOrNull(row.lastErrorType), lastErrorSource: toStringOrNull(row.lastErrorSource), errorCode: - typeof row.errorCode === "string" || typeof row.errorCode === "number" - ? row.errorCode - : null, + typeof row.errorCode === "string" || typeof row.errorCode === "number" ? row.errorCode : null, backoffLevel: toNumber(row.backoffLevel, 0), maxConcurrent: toNullableNumber(row.maxConcurrent), quotaWindowThresholds, @@ -115,9 +113,7 @@ export function toProviderConnection(value: unknown): ProviderConnectionView { * * Non-credential reads hit the already-coerced view directly at zero cost. */ -export function createLazyConnectionView( - row: Record -): ProviderConnectionView { +export function createLazyConnectionView(row: Record): ProviderConnectionView { const base = toProviderConnection(row); let decrypted: Record | undefined; @@ -153,9 +149,7 @@ const CREDENTIAL_FIELDS = new Set(["apiKey", "accessToken", "refreshToken", "idT * without any caller changes. The typed createLazyConnectionView remains * available for new code that wants a structured view. */ -export function createLazyRowProxy( - row: Record -): Record { +export function createLazyRowProxy(row: Record): Record { let decrypted: Record | undefined; const ensureDecrypted = () => { @@ -179,8 +173,7 @@ export function createLazyRowProxy( return () => { const result: Record = {}; for (const key of Object.keys(target)) { - result[key] = - CREDENTIAL_FIELDS.has(key) ? ensureDecrypted()[key] : target[key]; + result[key] = CREDENTIAL_FIELDS.has(key) ? ensureDecrypted()[key] : target[key]; } return result; }; diff --git a/src/shared/validation/schemas/combo.ts b/src/shared/validation/schemas/combo.ts index 71d2a992de..f9b7e3c832 100644 --- a/src/shared/validation/schemas/combo.ts +++ b/src/shared/validation/schemas/combo.ts @@ -88,6 +88,7 @@ export const scoringWeightsSchema = z specificityMatch: z.number().min(0).max(1).optional().default(0.05), contextAffinity: z.number().min(0).max(1).optional().default(0.08), cacheAffinity: z.number().min(0).max(1).optional().default(0), + sessionAvailability: z.number().min(0).max(1).optional().default(0.05), resetWindowAffinity: z.number().min(0).max(1).optional().default(0), }) .optional(); diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 6785f074c3..34c02c3bf1 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -93,6 +93,7 @@ import { } from "./chatPredicates"; import { markAntigravityMissingCloudCodeProject } from "@omniroute/open-sse/services/antigravityProjectPersistence.ts"; import { connectionHasExtraKeys } from "@omniroute/open-sse/services/apiKeyRotator.ts"; +import { wrapResponseWithOAuthSessionRelease } from "@omniroute/open-sse/services/oauthSessionOccupancy.ts"; import { extractReasoningIntent, type ExtractedReasoningIntent, @@ -787,7 +788,12 @@ async function handleChatImplementation( ); if (!creds || creds.allRateLimited) return false; - comboPreselectedCredentials.set(getComboCredentialCacheKey(modelString, target), creds); + // OAuth selection must happen atomically with occupancy reservation in the + // actual dispatch. Availability preflight may finish well before a combo + // target runs, so caching OAuth credentials here would reintroduce a race. + if (creds.authType !== "oauth") { + comboPreselectedCredentials.set(getComboCredentialCacheKey(modelString, target), creds); + } return true; }; @@ -801,21 +807,12 @@ async function handleChatImplementation( // Per-request Auto-Combo controls (#6023 / #6024 / #6025 / #3470): steer an // `auto` combo on this single request without mutating its stored config. const perRequestAutoControls = resolveRequestAutoControls(request.headers); - const relayOptions = - combo.strategy === "context-relay" || - bypassProviderQuotaPolicy || - Object.keys(perRequestAutoControls).length > 0 - ? { - ...(combo.strategy === "context-relay" - ? { - sessionId, - config: relayConfig, - } - : {}), - ...(bypassProviderQuotaPolicy ? { bypassProviderQuotaPolicy: true } : {}), - ...perRequestAutoControls, - } - : undefined; + const relayOptions = { + sessionId, + ...(combo.strategy === "context-relay" ? { config: relayConfig } : {}), + ...(bypassProviderQuotaPolicy ? { bypassProviderQuotaPolicy: true } : {}), + ...perRequestAutoControls, + }; telemetry.endPhase(); // Context-relay keeps generation in combo.ts, but handoff injection lives here @@ -859,9 +856,12 @@ async function handleChatImplementation( comboExecutionKey: target?.executionKey || target?.stepId || null, skipUpstreamRetry: target?.failoverBeforeRetry ?? false, allowRateLimitedConnection: target?.allowRateLimitedConnection === true, - preselectedCredentials: comboPreselectedCredentials.get( - getComboCredentialCacheKey(m, target) - ), + preselectedCredentials: (() => { + const key = getComboCredentialCacheKey(m, target); + const credentials = comboPreselectedCredentials.get(key); + comboPreselectedCredentials.delete(key); + return credentials; + })(), cachedSettings: settings, providerId: target?.providerId ?? null, correlationId: reqId, @@ -903,6 +903,11 @@ async function handleChatImplementation( correlationId: reqId, }); + for (const credentials of comboPreselectedCredentials.values()) { + credentials.releaseOAuthSession?.(); + } + comboPreselectedCredentials.clear(); + // ── Global Fallback Provider (#689) ──────────────────────────────────── // If combo exhausted all models, try the global fallback before giving up. if ( @@ -1280,13 +1285,17 @@ async function handleSingleModelChat( // re-attempt to exactly one for the whole request. Declared outside both retry // loops so it can never reset and loop. let streamEarlyEofRetries = 0; + const occupancySessionKey = + runtimeOptions.sessionAffinityKey ?? runtimeOptions.sessionId ?? `request:${randomUUID()}`; + let initialPreselectedCredentials = runtimeOptions.preselectedCredentials; requestAttemptLoop: while (true) { const excludedConnectionIds = new Set(); let lastError = requestRetryLastError; let lastStatus = requestRetryLastStatus; let lastCooldownMs = requestRetryLastCooldownMs; - let preselectedCredentials = runtimeOptions.preselectedCredentials; + let preselectedCredentials = initialPreselectedCredentials; + initialPreselectedCredentials = null; while (true) { const credentials = @@ -1298,7 +1307,8 @@ async function handleSingleModelChat( effectiveAllowedConnections, model, { - sessionKey: runtimeOptions.sessionAffinityKey ?? runtimeOptions.sessionId ?? null, + sessionKey: occupancySessionKey, + reserveOAuthSession: true, excludeConnectionIds: Array.from(excludedConnectionIds), ...(runtimeOptions.allowRateLimitedConnection ? { allowRateLimitedConnections: true } @@ -1404,6 +1414,7 @@ async function handleSingleModelChat( } const accountId = credentials.connectionId.slice(0, 8); + const releaseOAuthSession = credentials.releaseOAuthSession ?? (() => {}); log.info("AUTH", `Using ${provider} account: ${accountId}...`); // #474: when the request used a bare model name (no "/" — e.g. an alias // that resolved to "auto") and the selected connection declares a @@ -1425,7 +1436,10 @@ async function handleSingleModelChat( reasoningDecision: runtimeOptions.reasoningDecision, requestRoutingTags: runtimeOptions.reasoningRequestTags, }); - if (connectionRouting.response) return connectionRouting.response; + if (connectionRouting.response) { + releaseOAuthSession(); + return connectionRouting.response; + } requestBody = connectionRouting.body; } let injectedHandoff = null; @@ -1450,7 +1464,13 @@ async function handleSingleModelChat( ); } } - const refreshedCredentials = await checkAndRefreshToken(provider, credentials); + let refreshedCredentials; + try { + refreshedCredentials = await checkAndRefreshToken(provider, credentials); + } catch (error) { + releaseOAuthSession(); + throw error; + } const storeEnabled = isOpenAIResponsesStoreEnabled( refreshedCredentials?.providerSpecificData ?? credentials?.providerSpecificData ); @@ -1480,52 +1500,64 @@ async function handleSingleModelChat( refreshedCredentials ); } - const proxyInfo = await safeResolveProxy(credentials.connectionId, apiKeyInfo?.id, provider); + let proxyInfo; + try { + proxyInfo = await safeResolveProxy(credentials.connectionId, apiKeyInfo?.id, provider); + } catch (error) { + releaseOAuthSession(); + throw error; + } // #5217: sink for the proxy the executor pins internally (e.g. OpencodeExecutor // rotation) so the egress log below reflects the real egress, not "direct". const appliedProxySink: { proxy: unknown } = { proxy: null }; const proxyStartTime = Date.now(); - // 4. Execute chat via core after breaker gate checks (with optional TLS tracking) if (telemetry) telemetry.startPhase("connect"); const dispatchClientRawRequest = resolveDispatchClientRawRequest( clientRawRequest, runtimeOptions.modelAbortSignal ); - const execution = await executeChatWithBreaker({ - bypassCircuitBreaker: forceLiveComboTest || hasForcedConnection, - breaker, - body: requestBody, - provider, - model: effectiveModel, - refreshedCredentials, - proxyInfo, - appliedProxySink, - log, - clientRawRequest: dispatchClientRawRequest, - credentials, - apiKeyInfo, - userAgent, - comboName, - comboStrategy, - isCombo, - comboStepId: runtimeOptions.comboStepId ?? null, - comboExecutionKey: runtimeOptions.comboExecutionKey ?? runtimeOptions.comboStepId ?? null, - extendedContext, - modelApiFormat: apiFormat, - modelTargetFormat: targetFormat, - providerProfile, - cachedSettings: runtimeOptions.cachedSettings, - skipUpstreamRetry: runtimeOptions.skipUpstreamRetry ?? false, - correlationId: runtimeOptions?.correlationId ?? null, - modelPinned: runtimeOptions?.modelPinned ?? false, - routingComboId: runtimeOptions?.routingComboId ?? null, - }); + let execution: Awaited>; + try { + execution = await executeChatWithBreaker({ + bypassCircuitBreaker: forceLiveComboTest || hasForcedConnection, + breaker, + body: requestBody, + provider, + model: effectiveModel, + refreshedCredentials, + proxyInfo, + appliedProxySink, + log, + clientRawRequest: dispatchClientRawRequest, + credentials, + apiKeyInfo, + userAgent, + comboName, + comboStrategy, + isCombo, + comboStepId: runtimeOptions.comboStepId ?? null, + comboExecutionKey: runtimeOptions.comboExecutionKey ?? runtimeOptions.comboStepId ?? null, + extendedContext, + modelApiFormat: apiFormat, + modelTargetFormat: targetFormat, + providerProfile, + cachedSettings: runtimeOptions.cachedSettings, + skipUpstreamRetry: runtimeOptions.skipUpstreamRetry ?? false, + correlationId: runtimeOptions?.correlationId ?? null, + modelPinned: runtimeOptions?.modelPinned ?? false, + routingComboId: runtimeOptions?.routingComboId ?? null, + }); + } catch (error) { + releaseOAuthSession(); + throw error; + } if (telemetry) telemetry.endPhase(); if ("localResourcePressureResult" in execution) { return execution.localResourcePressureResult.response; } const { result, tlsFingerprintUsed } = execution; + if (!result.success) releaseOAuthSession(); const proxyLatency = Date.now() - proxyStartTime; const providerAlias = PROVIDER_ID_TO_ALIAS[provider] || provider; @@ -1560,6 +1592,10 @@ async function handleSingleModelChat( } if (telemetry) telemetry.startPhase("finalize"); if (telemetry) telemetry.endPhase(); + if (requestBody.stream === true) { + return wrapResponseWithOAuthSessionRelease(result.response, releaseOAuthSession); + } + releaseOAuthSession(); return result.response; } diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index ef6117ba32..b289a358fa 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -97,6 +97,11 @@ import { getResource404Bypass } from "./requestResourceHealth"; import * as log from "../utils/logger"; import { fisherYatesShuffle, getNextFromDeckSync } from "@/shared/utils/shuffleDeck"; import { readHeaderValue, type AuthRequestHeaders } from "./headerReader.ts"; +import { + getOAuthSessionAvailability, + reserveOAuthSession, +} from "@omniroute/open-sse/services/oauthSessionOccupancy.ts"; + type JsonRecord = Record; interface RecoverableConnectionState { connectionId: string; @@ -115,6 +120,7 @@ interface CredentialSelectionOptions { excludeConnectionIds?: string[] | null; sessionKey?: string | null; sessionAffinityTtlMs?: number | null; + reserveOAuthSession?: boolean; } interface CooldownInspectionState { connection: ProviderConnectionView; @@ -1505,7 +1511,15 @@ export async function getProviderCredentials( }; } - const orderedConnections = withQuota; + const orderedConnections = [...withQuota].sort((a, b) => { + if (a.authType !== "oauth" || b.authType !== "oauth") return 0; + const priorityDelta = (a.priority || 999) - (b.priority || 999); + if (priorityDelta !== 0) return priorityDelta; + return ( + getOAuthSessionAvailability(b.id, options.sessionKey) - + getOAuthSessionAvailability(a.id, options.sessionKey) + ); + }); const providerStrategyOverrides = (settings.providerStrategies || {}) as Record< string, @@ -1686,6 +1700,26 @@ export async function getProviderCredentials( connection = orderedConnections[0]; } + if (options.reserveOAuthSession === true && connection?.authType === "oauth") { + const selectedPriority = connection.priority || 999; + const selectedAvailability = getOAuthSessionAvailability(connection.id, options.sessionKey); + const moreAvailablePeer = [...orderedConnections] + .filter( + (candidate) => + candidate.authType === "oauth" && (candidate.priority || 999) <= selectedPriority + 1 + ) + .sort( + (a, b) => + getOAuthSessionAvailability(b.id, options.sessionKey) - + getOAuthSessionAvailability(a.id, options.sessionKey) + ) + .find( + (candidate) => + getOAuthSessionAvailability(candidate.id, options.sessionKey) > selectedAvailability + ); + if (moreAvailablePeer) connection = moreAvailablePeer; + } + if (provider === "antigravity" && connection) { log.info( "AUTH", @@ -1699,6 +1733,11 @@ export async function getProviderCredentials( syncHealthFromDB(connection.id, apiKeyHealth); } + const releaseOAuthSession = + options.reserveOAuthSession === true && connection.authType === "oauth" && options.sessionKey + ? reserveOAuthSession(connection.id, options.sessionKey) + : undefined; + return { apiKey: connection.apiKey, accessToken: connection.accessToken, @@ -1720,6 +1759,7 @@ export async function getProviderCredentials( // connectionId name. id: connection.id, provider: connection.provider, + authType: connection.authType, email: connection.email, connectionId: connection.id, // Include current status for optimization check @@ -1734,6 +1774,7 @@ export async function getProviderCredentials( // getProviderCredentialsWithQuotaPreflight can see them. Without this, // user-set cutoffs would silently never enforce. quotaWindowThresholds: connection.quotaWindowThresholds ?? null, + ...(releaseOAuthSession ? { releaseOAuthSession } : {}), }; } finally { selectionLock.release(); @@ -1821,7 +1862,11 @@ export async function getProviderCredentialsWithQuotaPreflight( return credentials; } - const connectionId = credentials.connectionId; + const selectedCredentials = credentials as typeof credentials & { + connectionId?: string; + releaseOAuthSession?: () => void; + }; + const connectionId = selectedCredentials.connectionId; if (!connectionId) { return credentials; } @@ -1891,14 +1936,22 @@ export async function getProviderCredentialsWithQuotaPreflight( const modelAwarePreflight = provider === "codex" || provider === "openrouter"; const preflightCredentials = requestedModel && modelAwarePreflight ? { ...credentials, requestedModel } : credentials; - const preflight = await preflightQuota(provider, connectionId, preflightCredentials, { - resolveMinRemainingPercent, - resolveWarnRemainingPercent: () => warnThresholdPercent, - }); + let preflight; + try { + preflight = await preflightQuota(provider, connectionId, preflightCredentials, { + resolveMinRemainingPercent, + resolveWarnRemainingPercent: () => warnThresholdPercent, + }); + } catch (error) { + selectedCredentials.releaseOAuthSession?.(); + throw error; + } if (preflight.proceed) { return credentials; } + selectedCredentials.releaseOAuthSession?.(); + const unavailableUntil = await markQuotaPreflightAccountUnavailable( provider, connectionId, diff --git a/tests/unit/oauth-session-occupancy.test.ts b/tests/unit/oauth-session-occupancy.test.ts new file mode 100644 index 0000000000..f891f97ec0 --- /dev/null +++ b/tests/unit/oauth-session-occupancy.test.ts @@ -0,0 +1,61 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import { + _clearOAuthSessionOccupancyForTest, + getForeignOAuthSessionCount, + getOAuthSessionAvailability, + reserveOAuthSession, + wrapResponseWithOAuthSessionRelease, +} from "../../open-sse/services/oauthSessionOccupancy.ts"; + +test.beforeEach(() => _clearOAuthSessionOccupancyForTest()); + +test("counts foreign sessions while ignoring parallel requests from the same session", () => { + const releaseA1 = reserveOAuthSession("account-a", "session-a", 60_000, 1_000); + const releaseA2 = reserveOAuthSession("account-a", "session-a", 60_000, 1_000); + + assert.equal(getForeignOAuthSessionCount("account-a", "session-a", 1_001), 0); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b", 1_001), 1); + assert.equal(getOAuthSessionAvailability("account-a", "session-b", 1_001), 0.5); + + releaseA1(); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b", 1_001), 1); + releaseA2(); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b", 1_001), 0); +}); + +test("expired leases fail open", () => { + reserveOAuthSession("account-a", "session-a", 100, 1_000); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b", 1_099), 1); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b", 1_100), 0); +}); + +test("response wrapper releases only after streaming completes", async () => { + const release = reserveOAuthSession("account-a", "session-a", 60_000); + const wrapped = wrapResponseWithOAuthSessionRelease( + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data")); + controller.close(); + }, + }) + ), + release + ); + + assert.equal(getForeignOAuthSessionCount("account-a", "session-b"), 1); + assert.equal(await wrapped.text(), "data"); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b"), 0); +}); + +test("response wrapper releases on cancellation", async () => { + const release = reserveOAuthSession("account-a", "session-a", 60_000); + const wrapped = wrapResponseWithOAuthSessionRelease( + new Response(new ReadableStream({ pull() {} })), + release + ); + await wrapped.body?.cancel("client disconnected"); + assert.equal(getForeignOAuthSessionCount("account-a", "session-b"), 0); +}); diff --git a/tests/unit/prompt-cache-affinity.test.ts b/tests/unit/prompt-cache-affinity.test.ts index 4d9f8bc593..7e25456f02 100644 --- a/tests/unit/prompt-cache-affinity.test.ts +++ b/tests/unit/prompt-cache-affinity.test.ts @@ -9,6 +9,10 @@ import { } from "../../open-sse/services/combo/promptCacheAffinity.ts"; import type { ResolvedComboTarget } from "../../open-sse/services/combo/types.ts"; import { applyStrategyOrdering } from "../../open-sse/services/combo/applyStrategyOrdering.ts"; +import { + _clearOAuthSessionOccupancyForTest, + reserveOAuthSession, +} from "../../open-sse/services/oauthSessionOccupancy.ts"; function target( executionKey: string, @@ -132,3 +136,40 @@ test("cache-optimized strategy routes a stable prompt key to the same account", const second = await applyStrategyOrdering("cache-optimized", [...targets].reverse(), deps); assert.equal(first[0].connectionId, second[0].connectionId); }); + +test("foreign OAuth session softly redirects cache affinity while the same session stays local", () => { + _clearOAuthSessionOccupancyForTest(); + const oauthTargets = [ + { ...target("step-a", "account-a"), authType: "oauth" }, + { ...target("step-b", "account-b"), authType: "oauth" }, + ]; + const fixture = Array.from({ length: 10_000 }, (_, index) => { + const body = { prompt_cache_key: `occupied-cache-key-${index}` }; + const baseline = applyPromptCacheAffinity(oauthTargets, body, true, "global", "session-a").targets; + const occupied = baseline[0]; + const alternative = baseline[1]; + const release = reserveOAuthSession(occupied.connectionId!, "session-a"); + const foreignFirst = applyPromptCacheAffinity(oauthTargets, body, true, "global", "session-b").targets[0]; + release(); + return foreignFirst.connectionId === alternative.connectionId + ? { body, occupied, alternative } + : null; + }).find(Boolean); + assert.ok(fixture, "test fixture must find a close rendezvous score"); + const { body, occupied, alternative } = fixture; + const release = reserveOAuthSession(occupied.connectionId!, "session-a"); + + assert.equal( + applyPromptCacheAffinity(oauthTargets, body, true, "global", "session-a").targets[0].connectionId, + occupied.connectionId, + "the owning session keeps its cache-local account" + ); + assert.equal( + applyPromptCacheAffinity(oauthTargets, body, true, "global", "session-b").targets[0].connectionId, + alternative.connectionId, + "a foreign session prefers the free OAuth account" + ); + + release(); + _clearOAuthSessionOccupancyForTest(); +}); diff --git a/tests/unit/sse-auth.test.ts b/tests/unit/sse-auth.test.ts index 94b884c9bc..5a81d97b64 100644 --- a/tests/unit/sse-auth.test.ts +++ b/tests/unit/sse-auth.test.ts @@ -15,6 +15,7 @@ const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); const auth = await import("../../src/sse/services/auth.ts"); const quotaCache = await import("../../src/domain/quotaCache.ts"); const fallback = await import("../../open-sse/services/accountFallback.ts"); +const oauthOccupancy = await import("../../open-sse/services/oauthSessionOccupancy.ts"); async function resetStorage() { core.resetDbInstance(); @@ -64,6 +65,7 @@ async function flushWrites() { } test.beforeEach(async () => { + oauthOccupancy._clearOAuthSessionOccupancyForTest(); await resetStorage(); }); @@ -525,6 +527,45 @@ test("getProviderCredentials keeps separate codex affinity per session", async ( assert.equal(sessionB2.connectionId, second.id); }); +test("concurrent OAuth selections reserve different available accounts atomically", async () => { + await settingsDb.updateSettings({ + fallbackStrategy: "fill-first", + codexSessionAffinityTtlMs: 0, + }); + const first = await seedConnection("codex", { + authType: "oauth", + name: "codex-occupancy-a", + priority: 1, + }); + const second = await seedConnection("codex", { + authType: "oauth", + name: "codex-occupancy-b", + priority: 1, + }); + assert.ok(second.priority <= first.priority + 1); + + const [sessionA, sessionB] = await Promise.all([ + auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "occupancy-session-a", + reserveOAuthSession: true, + }), + auth.getProviderCredentials("codex", null, null, "gpt-5.5", { + sessionKey: "occupancy-session-b", + reserveOAuthSession: true, + }), + ]); + + assert.equal(sessionA.authType, "oauth"); + assert.equal(typeof sessionA.releaseOAuthSession, "function"); + assert.equal( + oauthOccupancy.getForeignOAuthSessionCount(sessionA.connectionId, "occupancy-session-b"), + 1 + ); + assert.notEqual(sessionA.connectionId, sessionB.connectionId); + sessionA.releaseOAuthSession?.(); + sessionB.releaseOAuthSession?.(); +}); + test("getProviderCredentials rebinds codex session when affinity connection is excluded", async () => { await settingsDb.updateSettings({ fallbackStrategy: "round-robin",