feat: enhance chat handling with cached settings and deduplicate quota fetches in reset-aware strategy

2026-07-26 09:52:11 +03:00 · 2026-05-08 16:57:25 +00:00
parent 1e80366b42
commit 3662f90095
5 changed files with 131 additions and 48 deletions
--- 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);
--- a/open-sse/services/combo.ts
+++ b/open-sse/services/combo.ts
@@ -849,6 +849,7 @@ function scoreResetAwareQuota(quota: unknown, config: ReturnType<typeof resolveR
 async function getQuotaAwareConnectionsForTarget(
  target: ResolvedComboTarget,
  connectionCache: Map<string, Array<Record<string, unknown>>>,
+  connectionLoadPromises: Map<string, Promise<Array<Record<string, unknown>>>>,
  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<Record<string, unknown>>)
-        : [];
-      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<Record<string, unknown>>)
+              : [];
+            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<string, Array<Record<string, unknown>>>();
+  const connectionLoadPromises = new Map<string, Promise<Array<Record<string, unknown>>>>();
+  const quotaPromises = new Map<string, Promise<unknown>>();
  const connectionById = new Map<string, Record<string, unknown>>();
  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 };
--- 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<string, any>();
+    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();

--- 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,
--- 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({