From b5cc0993dfdc71c624a01b21c85c689ac49f151d Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Wed, 27 May 2026 13:13:22 +0200 Subject: [PATCH] fix(combo): allow rate-limited provider connections after transient 429s --- open-sse/services/combo.ts | 31 +++++++++++++++++++++++++------ src/sse/handlers/chat.ts | 8 ++++++++ src/sse/services/auth.ts | 5 ++++- 3 files changed, 37 insertions(+), 7 deletions(-) diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 743d0cb96d..63bb9f7dd2 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -190,7 +190,10 @@ type ComboLogger = { }; export type SingleModelTarget = - | (ResolvedComboTarget & { modelAbortSignal?: AbortSignal | null }) + | (ResolvedComboTarget & { + allowRateLimitedConnection?: boolean; + modelAbortSignal?: AbortSignal | null; + }) | { modelAbortSignal: AbortSignal }; type HandleSingleModel = ( @@ -201,7 +204,7 @@ type HandleSingleModel = ( type IsModelAvailable = ( modelStr: string, - target?: ResolvedComboTarget + target?: ResolvedComboTarget & { allowRateLimitedConnection?: boolean } ) => Promise | boolean; type ComboRelayOptions = { @@ -3098,6 +3101,7 @@ export async function handleComboChat({ // #1731: Per-set-iteration set of providers whose quota is fully exhausted. // Reset each retry so providers excluded in a previous attempt get another chance. const exhaustedProviders = new Set(); + const transientRateLimitedProviders = new Set(); if (setTry > 0) { log.info("COMBO", `All targets failed — retrying set (${setTry}/${maxSetRetries})`); await new Promise((resolve) => { @@ -3129,6 +3133,11 @@ export async function handleComboChat({ const modelStr = target.modelStr; const provider = target.provider; const profile = await getRuntimeProviderProfile(provider); + const allowRateLimitedConnection = + Boolean(provider && provider !== "unknown") && transientRateLimitedProviders.has(provider); + const targetForAttempt = allowRateLimitedConnection + ? { ...target, allowRateLimitedConnection: true } + : target; // #1731: Skip targets from a provider that already signaled full quota exhaustion this request. if (provider && exhaustedProviders.has(provider)) { @@ -3142,7 +3151,7 @@ export async function handleComboChat({ // Pre-check: skip models where no credentials are available (excluded, rate-limited, or unavailable) if (isModelAvailable) { - const available = await isModelAvailable(modelStr, target); + const available = await isModelAvailable(modelStr, targetForAttempt); if (!available) { log.info("COMBO", `Skipping ${modelStr} — no credentials available or model excluded`); if (i > 0) fallbackCount++; @@ -3232,7 +3241,7 @@ export async function handleComboChat({ } } const result = await handleSingleModelWithTimeout(attemptBody, modelStr, { - ...target, + ...targetForAttempt, failoverBeforeRetry: config.failoverBeforeRetry, }); @@ -3501,6 +3510,8 @@ export async function handleComboChat({ "COMBO", `Provider ${provider} quota exhausted — marking for skip on remaining targets (#1731)` ); + } else if (result.status === 429 && provider && provider !== "unknown") { + transientRateLimitedProviders.add(provider); } // Trigger shared provider circuit breaker for 5xx errors and connection failures. @@ -3690,6 +3701,7 @@ async function handleRoundRobinCombo({ // When a target returns a quota-exhausted 429, remaining targets from the same // provider are skipped to avoid the cascade through N same-provider targets. const exhaustedProviders = new Set(); + const transientRateLimitedProviders = new Set(); // Try each model starting from the round-robin target for (let offset = 0; offset < modelCount; offset++) { @@ -3699,10 +3711,15 @@ async function handleRoundRobinCombo({ const provider = target.provider; const profile = await getRuntimeProviderProfile(provider); const semaphoreKey = `combo:${combo.name}:${target.executionKey}`; + const allowRateLimitedConnection = + Boolean(provider && provider !== "unknown") && transientRateLimitedProviders.has(provider); + const targetForAttempt = allowRateLimitedConnection + ? { ...target, allowRateLimitedConnection: true } + : target; // Pre-check availability if (isModelAvailable) { - const available = await isModelAvailable(modelStr, target); + const available = await isModelAvailable(modelStr, targetForAttempt); if (!available) { log.info("COMBO-RR", `Skipping ${modelStr} — no credentials available or model excluded`); if (offset > 0) fallbackCount++; @@ -3766,7 +3783,7 @@ async function handleRoundRobinCombo({ ); const result = await handleSingleModel(body, modelStr, { - ...target, + ...targetForAttempt, failoverBeforeRetry: config.failoverBeforeRetry, }); @@ -3932,6 +3949,8 @@ async function handleRoundRobinCombo({ if (providerExhausted) { exhaustedProviders.add(provider); log.info("COMBO-RR", `Provider ${provider} quota exhausted — marking for skip (#1731)`); + } else if (result.status === 429 && provider && provider !== "unknown") { + transientRateLimitedProviders.add(provider); } // Transient errors → mark in semaphore so round-robin stops stampeding this target. diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index ca5da75c62..7c2988cfaa 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -438,6 +438,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { const checkModelAvailable = async ( modelString: string, target?: { + allowRateLimitedConnection?: boolean; connectionId?: string | null; allowedConnectionIds?: string[] | null; executionKey?: string | null; @@ -469,6 +470,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { resolvedModel, { sessionKey: sessionAffinityKey, + ...(target?.allowRateLimitedConnection ? { allowRateLimitedConnections: true } : {}), ...(target?.connectionId ? { forcedConnectionId: target.connectionId } : {}), } ); @@ -496,6 +498,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { b: any, m: string, target?: { + allowRateLimitedConnection?: boolean; connectionId?: string | null; executionKey?: string | null; stepId?: string | null; @@ -520,6 +523,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { comboStepId: target?.stepId || null, comboExecutionKey: target?.executionKey || target?.stepId || null, skipUpstreamRetry: target?.failoverBeforeRetry ?? false, + allowRateLimitedConnection: target?.allowRateLimitedConnection === true, preselectedCredentials: comboPreselectedCredentials.get( getComboCredentialCacheKey(m, target) ), @@ -651,6 +655,7 @@ async function handleSingleModelChat( comboStepId?: string | null; comboExecutionKey?: string | null; skipUpstreamRetry?: boolean; + allowRateLimitedConnection?: boolean; preselectedCredentials?: any; cachedSettings?: any; } = {}, @@ -811,6 +816,9 @@ async function handleSingleModelChat( { sessionKey: runtimeOptions.sessionAffinityKey ?? runtimeOptions.sessionId ?? null, excludeConnectionIds: Array.from(excludedConnectionIds), + ...(runtimeOptions.allowRateLimitedConnection + ? { allowRateLimitedConnections: true } + : {}), ...(forceLiveComboTest ? { allowSuppressedConnections: true, diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 06e4e56550..79edfc1169 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -94,6 +94,7 @@ interface RecoverableConnectionState { interface CredentialSelectionOptions { allowSuppressedConnections?: boolean; + allowRateLimitedConnections?: boolean; bypassQuotaPolicy?: boolean; forcedConnectionId?: string | null; excludeConnectionIds?: string[] | null; @@ -846,6 +847,8 @@ export async function getProviderCredentials( } const allowSuppressedConnections = options.allowSuppressedConnections === true; + const allowRateLimitedConnections = + allowSuppressedConnections || options.allowRateLimitedConnections === true; const bypassQuotaPolicy = options.bypassQuotaPolicy === true; const forcedConnectionId = typeof options.forcedConnectionId === "string" && options.forcedConnectionId.trim().length > 0 @@ -973,7 +976,7 @@ export async function getProviderCredentials( return false; } if (!allowSuppressedConnections) { - if (isAccountUnavailable(c.rateLimitedUntil)) return false; + if (!allowRateLimitedConnections && isAccountUnavailable(c.rateLimitedUntil)) return false; if (isTerminalConnectionStatus(c)) return false; if (provider === "codex" && isCodexScopeUnavailable(c, requestedModel)) return false; // Per-model lockout: if this specific model is locked on this connection, skip it