fix(combo): propagate selected connection ID to fallback error responses for correct model lockout (#4809)

Integrated into release/v3.8.36

config/quality/file-size-baseline.json

@@ -217,7 +217,7 @@
     "src/shared/constants/sidebarVisibility.ts": 1100,
     "src/shared/services/cliRuntime.ts": 1090,
     "src/shared/validation/schemas.ts": 2523,
-    "src/sse/handlers/chat.ts": 1516,
+    "src/sse/handlers/chat.ts": 1521,
     "src/sse/services/auth.ts": 2289
   },
   "testCap": 800,

src/lib/resilience/modelLockoutSettings.ts

@@ -30,11 +30,7 @@ function toInteger(
   const min = options.min ?? 0;
   const max = options.max ?? Number.MAX_SAFE_INTEGER;
   const parsed =
-    typeof value === "number"
-      ? value
-      : typeof value === "string"
-        ? Number(value)
-        : fallback;
+    typeof value === "number" ? value : typeof value === "string" ? Number(value) : fallback;
   return Number.isFinite(parsed) ? Math.max(min, Math.min(max, Math.trunc(parsed))) : fallback;
 }
 
@@ -65,10 +61,14 @@ export function resolveModelLockoutSettings(
     process.execArgv.includes("--test") ||
     process.argv.some((arg) => typeof arg === "string" && arg.includes("test"));
 
-  const baseCooldownMs = toInteger(raw.baseCooldownMs, DEFAULT_MODEL_LOCKOUT_SETTINGS.baseCooldownMs, {
-    min: isTest ? 0 : 5_000,
-    max: 600_000,
-  });
+  const baseCooldownMs = toInteger(
+    raw.baseCooldownMs,
+    DEFAULT_MODEL_LOCKOUT_SETTINGS.baseCooldownMs,
+    {
+      min: isTest ? 0 : 5_000,
+      max: 600_000,
+    }
+  );
   const maxCooldownMs = Math.max(
     toInteger(raw.maxCooldownMs, DEFAULT_MODEL_LOCKOUT_SETTINGS.maxCooldownMs, {
       min: isTest ? 0 : 5_000,

src/sse/handlers/chat.ts

@@ -1055,7 +1055,7 @@ async function handleSingleModelChat(
           breaker._onFailure();
         }
 
-        return handleNoCredentials(
+        const noCredsRes = handleNoCredentials(
           credentials,
           excludedConnectionIds.size > 0 ? Array.from(excludedConnectionIds)[0] : null,
           provider,
@@ -1063,6 +1063,11 @@ async function handleSingleModelChat(
           lastError,
           lastStatus
         );
+        const lastFailedConnectionId =
+          excludedConnectionIds.size > 0
+            ? Array.from(excludedConnectionIds)[excludedConnectionIds.size - 1]
+            : null;
+        return withSelectedConnectionHeader(noCredsRes, lastFailedConnectionId);
       }
 
       const accountId = credentials.connectionId.slice(0, 8);

tests/unit/rate-limit-queue-timeout-lockout.test.ts

@@ -0,0 +1,140 @@
+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(), "omr-rl-queue-timeout-lockout-"));
+process.env.DATA_DIR = TEST_DATA_DIR;
+process.env.API_KEY_SECRET = "test-rl-lockout-secret";
+
+const core = await import("../../src/lib/db/core.ts");
+const providersDb = await import("../../src/lib/db/providers.ts");
+const { handleComboChat } = await import("../../open-sse/services/combo.ts");
+const { getModelLockoutInfo, isModelLocked } =
+  await import("../../open-sse/services/accountFallback.ts");
+
+function createLog() {
+  return {
+    info: () => {},
+    warn: () => {},
+    error: () => {},
+    debug: () => {},
+  };
+}
+
+async function seedConnection(provider: string, overrides: any = {}): Promise<any> {
+  return providersDb.createProviderConnection({
+    provider,
+    authType: "apikey",
+    name: overrides.name || `${provider}-${Math.random().toString(16).slice(2, 8)}`,
+    apiKey: overrides.apiKey || `sk-test-${Math.random().toString(16).slice(2, 8)}`,
+    isActive: true,
+    testStatus: "active",
+    rateLimitedUntil: null,
+    backoffLevel: overrides.backoffLevel || 0,
+    providerSpecificData: overrides.providerSpecificData || {},
+  });
+}
+
+function errorResponseWithoutConnectionId(status: number) {
+  return new Response(JSON.stringify({ error: { message: "Local queue timeout" } }), {
+    status,
+    headers: { "content-type": "application/json" },
+  });
+}
+
+function errorResponseWithConnectionId(status: number, connectionId: string) {
+  return new Response(JSON.stringify({ error: { message: "Local queue timeout" } }), {
+    status,
+    headers: {
+      "content-type": "application/json",
+      "X-OmniRoute-Selected-Connection-Id": connectionId,
+    },
+  });
+}
+
+test.afterEach(() => {
+  core.resetDbInstance();
+  fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
+});
+
+test("RATE_LIMIT_QUEUE_TIMEOUT lockout behaves correctly depending on connection ID header", async () => {
+  const provider = "openai";
+  const model = "gpt-4";
+
+  // Seed the connection first!
+  const connection = await seedConnection(provider);
+  const connectionId = connection.id;
+
+  const customSettings = {
+    modelLockout: {
+      enabled: true,
+      errorCodes: [502, 520],
+      baseCooldownMs: 3000,
+      maxCooldownMs: 5000,
+      maxBackoffSteps: 10,
+      useExponentialBackoff: true,
+    },
+  };
+
+  const logs = createLog();
+
+  // Scenario 1: Response lacks connection ID header (Current buggy behavior)
+  await handleComboChat({
+    body: {},
+    combo: {
+      name: "test-combo",
+      strategy: "priority",
+      models: [`${provider}/${model}`],
+      config: { maxRetries: 0, retryDelayMs: 0, fallbackDelayMs: 0 },
+    },
+    handleSingleModel: async () => {
+      return errorResponseWithoutConnectionId(502);
+    },
+    isModelAvailable: async () => true,
+    log: logs as any,
+    settings: customSettings,
+    allCombos: null,
+  });
+
+  // Verify that the model is NOT locked for the actual connection connectionId
+  const lockedForConnBuggy = isModelLocked(provider, connectionId, model);
+  assert.equal(
+    lockedForConnBuggy,
+    false,
+    "Model should not be locked for actual connection when header is missing"
+  );
+
+  // Verify that it got locked under the empty string connectionId ""
+  const lockedForEmptyBuggy = isModelLocked(provider, "", model);
+  assert.equal(
+    lockedForEmptyBuggy,
+    true,
+    "Model is incorrectly locked under empty string connectionId when header is missing"
+  );
+
+  const { clearAllModelLockouts } = await import("../../open-sse/services/accountFallback.ts");
+  clearAllModelLockouts();
+
+  await handleComboChat({
+    body: {},
+    combo: {
+      name: "test-combo",
+      strategy: "priority",
+      models: [`${provider}/${model}`],
+      config: { maxRetries: 0, retryDelayMs: 0, fallbackDelayMs: 0 },
+    },
+    handleSingleModel: async () => {
+      return errorResponseWithConnectionId(502, connectionId);
+    },
+    isModelAvailable: async () => true,
+    log: logs as any,
+    settings: customSettings,
+    allCombos: null,
+  });
+
+  // Verify that the model IS locked for the actual connection
+  const lockedForConnFixed = isModelLocked(provider, connectionId, model);
+  assert.equal(lockedForConnFixed, true, "Model should be locked for the correct connection ID");
+});