diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 3fdc212c58..bc5310a4ea 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -971,6 +971,7 @@ export async function handleChatCore({ comboStepId = null, comboExecutionKey = null, disableEmergencyFallback = false, + cachedSettings = null, }) { let { provider, model, extendedContext } = modelInfo; const requestedModel = @@ -1424,7 +1425,7 @@ export async function handleChatCore({ nativeCodexPassthrough && isCompactResponsesEndpoint(endpointPath) ? false : resolveStreamFlag(body?.stream, acceptHeader); - const settings = await getCachedSettings(); + const settings = cachedSettings ?? (await getCachedSettings()); credentials = applyCodexGlobalFastServiceTier(provider, credentials, settings); effectiveServiceTier = resolveEffectiveServiceTier(body); setGeminiThoughtSignatureMode(settings.antigravitySignatureCacheMode); diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index e413029366..3ff14b78e9 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -849,6 +849,7 @@ function scoreResetAwareQuota(quota: unknown, config: ReturnType>>, + connectionLoadPromises: Map>>>, comboName: string, log: { warn?: (...args: unknown[]) => void } ) { @@ -861,25 +862,35 @@ async function getQuotaAwareConnectionsForTarget( return cached.connections; } - try { - const connections = await getProviderConnections({ provider, isActive: true }); - const activeConnections = Array.isArray(connections) - ? (connections as Array>) - : []; - connectionCache.set(provider, activeConnections); - resetAwareConnectionCache.set(provider, { - connections: activeConnections, - fetchedAt: Date.now(), - }); - } catch (error) { - log.warn?.("COMBO", "Reset-aware failed to load quota-aware connections.", { - comboName, - err: error, - operation: "getProviderConnections", + if (!connectionLoadPromises.has(provider)) { + connectionLoadPromises.set( provider, - }); - connectionCache.set(provider, []); + (async () => { + try { + const connections = await getProviderConnections({ provider, isActive: true }); + const activeConnections = Array.isArray(connections) + ? (connections as Array>) + : []; + resetAwareConnectionCache.set(provider, { + connections: activeConnections, + fetchedAt: Date.now(), + }); + return activeConnections; + } catch (error) { + log.warn?.("COMBO", "Reset-aware failed to load quota-aware connections.", { + comboName, + err: error, + operation: "getProviderConnections", + provider, + }); + return []; + } + })() + ); } + + const connections = await connectionLoadPromises.get(provider)!; + connectionCache.set(provider, connections); } return connectionCache.get(provider) || []; } @@ -957,12 +968,20 @@ async function orderTargetsByResetAwareQuota( const config = resolveResetAwareConfig(configSource); const connectionCache = new Map>>(); + const connectionLoadPromises = new Map>>>(); + const quotaPromises = new Map>(); const connectionById = new Map>(); const expandedTargets: ResolvedComboTarget[] = []; const targetsWithConnections = await Promise.all( targets.map(async (target) => ({ - connections: await getQuotaAwareConnectionsForTarget(target, connectionCache, comboName, log), + connections: await getQuotaAwareConnectionsForTarget( + target, + connectionCache, + connectionLoadPromises, + comboName, + log + ), target, })) ); @@ -1008,17 +1027,23 @@ async function orderTargetsByResetAwareQuota( const provider = getResetAwareProvider(target); const fetcher = provider ? getQuotaFetcher(provider) : null; if (fetcher && provider && target.connectionId) { - try { - quota = await fetcher(target.connectionId, connectionById.get(target.connectionId)); - } catch (error) { - log.warn?.("COMBO", "Reset-aware quota fetch failed.", { - comboName, - connectionId: target.connectionId, - err: error, - operation: "quotaFetch", - provider, - }); + const quotaKey = `${provider}:${target.connectionId}`; + if (!quotaPromises.has(quotaKey)) { + quotaPromises.set( + quotaKey, + fetcher(target.connectionId, connectionById.get(target.connectionId)).catch((error) => { + log.warn?.("COMBO", "Reset-aware quota fetch failed.", { + comboName, + connectionId: target.connectionId, + err: error, + operation: "quotaFetch", + provider, + }); + return null; + }) + ); } + quota = await quotaPromises.get(quotaKey)!; } const { score } = scoreResetAwareQuota(quota, config); return { target, score, index }; diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 0d13eda7cb..de643574ef 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -287,9 +287,18 @@ export async function handleChat(request: any, clientRawRequest: any = null) { // Pre-check function used by combo routing. For explicit combo live tests, // avoid pre-skipping so each model gets a real execution attempt. + const comboPreselectedCredentials = new Map(); + const getComboCredentialCacheKey = ( + modelString: string, + target?: { connectionId?: string | null; executionKey?: string | null } + ) => `${target?.executionKey || target?.connectionId || ""}:${modelString}`; const checkModelAvailable = async ( modelString: string, - target?: { connectionId?: string | null; allowedConnectionIds?: string[] | null } + target?: { + connectionId?: string | null; + allowedConnectionIds?: string[] | null; + executionKey?: string | null; + } ) => { if (isComboLiveTest) return true; @@ -321,6 +330,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { ); if (!creds || creds.allRateLimited) return false; + comboPreselectedCredentials.set(getComboCredentialCacheKey(modelString, target), creds); return true; }; @@ -362,6 +372,10 @@ export async function handleChat(request: any, clientRawRequest: any = null) { allowedConnectionIds: target?.allowedConnectionIds ?? null, comboStepId: target?.stepId || null, comboExecutionKey: target?.executionKey || target?.stepId || null, + preselectedCredentials: comboPreselectedCredentials.get( + getComboCredentialCacheKey(m, target) + ), + cachedSettings: settings, }, combo.strategy, true @@ -481,6 +495,8 @@ async function handleSingleModelChat( allowedConnectionIds?: string[] | null; comboStepId?: string | null; comboExecutionKey?: string | null; + preselectedCredentials?: any; + cachedSettings?: any; } = {}, comboStrategy: string | null = null, isCombo: boolean = false @@ -523,7 +539,7 @@ async function handleSingleModelChat( const userAgent = request?.headers?.get("user-agent") || ""; const baseRetrySettings = resolveCooldownAwareRetrySettings( - await getCachedSettings().catch(() => ({})) + runtimeOptions.cachedSettings ?? (await getCachedSettings().catch(() => ({}))) ); const disableCooldownAwareRetry = isCombo || forceLiveComboTest || runtimeOptions.emergencyFallbackTried === true; @@ -557,26 +573,31 @@ async function handleSingleModelChat( let lastError = requestRetryLastError; let lastStatus = requestRetryLastStatus; let lastCooldownMs = requestRetryLastCooldownMs; + let preselectedCredentials = runtimeOptions.preselectedCredentials; while (true) { - const credentials = await getProviderCredentialsWithQuotaPreflight( - provider, - null, - effectiveAllowedConnections, - model, - { - excludeConnectionIds: Array.from(excludedConnectionIds), - ...(forceLiveComboTest - ? { - allowSuppressedConnections: true, - bypassQuotaPolicy: true, + const credentials = + preselectedCredentials && excludedConnectionIds.size === 0 + ? preselectedCredentials + : await getProviderCredentialsWithQuotaPreflight( + provider, + null, + effectiveAllowedConnections, + model, + { + excludeConnectionIds: Array.from(excludedConnectionIds), + ...(forceLiveComboTest + ? { + allowSuppressedConnections: true, + bypassQuotaPolicy: true, + } + : {}), + ...(runtimeOptions.forcedConnectionId + ? { forcedConnectionId: runtimeOptions.forcedConnectionId } + : {}), } - : {}), - ...(runtimeOptions.forcedConnectionId - ? { forcedConnectionId: runtimeOptions.forcedConnectionId } - : {}), - } - ); + ); + preselectedCredentials = null; if (!credentials || "allRateLimited" in credentials) { if (credentials?.allRateLimited) { @@ -710,6 +731,7 @@ async function handleSingleModelChat( comboExecutionKey: runtimeOptions.comboExecutionKey ?? runtimeOptions.comboStepId ?? null, extendedContext, providerProfile, + cachedSettings: runtimeOptions.cachedSettings, }); if (telemetry) telemetry.endPhase(); diff --git a/src/sse/handlers/chatHelpers.ts b/src/sse/handlers/chatHelpers.ts index 3079acce78..f01a6e2073 100644 --- a/src/sse/handlers/chatHelpers.ts +++ b/src/sse/handlers/chatHelpers.ts @@ -186,6 +186,7 @@ export async function executeChatWithBreaker({ comboExecutionKey, extendedContext, providerProfile, + cachedSettings, }: any): Promise<{ result: any; tlsFingerprintUsed: boolean }> { let tlsFingerprintUsed = false; @@ -206,6 +207,7 @@ export async function executeChatWithBreaker({ isCombo, comboStepId, comboExecutionKey, + cachedSettings, onCredentialsRefreshed: async (newCreds: any) => { await updateProviderCredentials(credentials.connectionId, { accessToken: newCreds.accessToken, diff --git a/tests/unit/combo-strategies.test.ts b/tests/unit/combo-strategies.test.ts index 2937ee3595..1669ad4368 100644 --- a/tests/unit/combo-strategies.test.ts +++ b/tests/unit/combo-strategies.test.ts @@ -375,6 +375,39 @@ test("reset-aware strategy uses registered quota fetchers for non-Codex provider assert.equal(await selectedConnectionFor(combo), soon); }); +test("reset-aware strategy deduplicates quota fetches for repeated connection targets", async () => { + const provider = `dedupe-provider-${randomUUID()}`; + const connectionId = `shared-${randomUUID()}`; + let fetchCount = 0; + + registerQuotaFetcher(provider, async (id) => { + fetchCount++; + assert.equal(id, connectionId); + return { + used: 20, + total: 100, + percentUsed: 0.2, + resetAt: new Date(Date.now() + 24 * 3600 * 1000).toISOString(), + }; + }); + + const combo = { + name: `reset-aware-dedupe-${randomUUID()}`, + strategy: "reset-aware", + models: ["model-a", "model-b"].map((model, index) => ({ + kind: "model", + provider, + providerId: provider, + model, + connectionId, + id: `dedupe-${index}`, + })), + }; + + assert.equal(await selectedConnectionFor(combo), connectionId); + assert.equal(fetchCount, 1); +}); + test("reset-aware strategy respects API-key allowed connections during expansion", async () => { const provider = `limited-provider-${randomUUID()}`; const disallowed = await providersDb.createProviderConnection({