fix(sse): don't cool down the connection on a self-inflicted upstream timeout (504) (#5064)

Integrated into release/v3.8.38 (leva 5)

CHANGELOG.md

@@ -25,6 +25,7 @@ _In development — bullets added per PR; finalized at release._
 - **oauth(kiro):** support Kiro IDC (organization) token import — when the `~/.aws/sso/cache` token carries a `clientIdHash`, auto-import now reads the linked client registration file to obtain `clientId`/`clientSecret`, probes the Kiro IDE `profile.json` for `profileArn` (ARN region normalized to `us-east-1` for the runtime gateway), and refreshes via the regional AWS OIDC endpoint instead of the social path; the import schema and modal forward these credentials so manual imports also work for IDC tokens. (thanks @enjoyer-hub)
 - **fix(translator):** preserve client `cache_control` breakpoints when routing Claude-format requests (e.g. Claude Code) to Alibaba DashScope's OpenAI-compatible providers (`alibaba` / `alibaba-cn`). The Claude→OpenAI translation previously stripped the markers from the system and message text blocks, so DashScope's explicit caching never engaged and every request was a cache miss. Cache hints now survive when preservation is requested for caching-capable OpenAI-format providers. (thanks @sacrtap)
 - **fix(tts):** resolve Gemini TTS models from catalog and add `gemini-3.1-flash-tts-preview` as the new default Vertex TTS model. (thanks @nguyenha935)
+- **fix(sse): don't cool down a healthy connection on a self-inflicted upstream timeout (504)** — when OmniRoute's own deadline elapses (surfaced as `TimeoutError`/`BodyTimeoutError` → 504), the connection is no longer disabled/failed-over, so a slow-but-healthy provider isn't penalised for our timeout. Genuine upstream 5xx/429 still trigger cooldown; antigravity keeps its own policy. (thanks @costaeder)
 
 ---
 

open-sse/handlers/chatCore.ts

@@ -2668,8 +2668,15 @@ export async function handleChatCore({
         ? "Request aborted"
         : formatProviderError(error, provider, model, failureStatus);
     const upstreamErrorCode = getUpstreamErrorIdentifier(error);
+    // Tag our own deadline timeouts (fetch-start TimeoutError / body BodyTimeoutError,
+    // both surfaced as a 504) as "upstream_timeout" so the cooldown layer can tell a
+    // slow-but-not-failed request apart from a real provider 5xx. (Antigravity already
+    // tags its pre-response timeout via the code below.)
+    const isOwnDeadlineTimeout =
+      failureStatus === HTTP_STATUS.GATEWAY_TIMEOUT &&
+      (error.name === "TimeoutError" || error.name === "BodyTimeoutError");
     const upstreamErrorType =
-      upstreamErrorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE
+      upstreamErrorCode === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE || isOwnDeadlineTimeout
         ? "upstream_timeout"
         : failureStatus === 401
           ? "authentication_error"

open-sse/handlers/chatCore/cooldownClassification.ts

@@ -0,0 +1,26 @@
+import { HTTP_STATUS } from "../../config/constants.ts";
+
+/**
+ * Whether a failed single-model attempt is a *self-inflicted* upstream timeout — i.e.
+ * OmniRoute's own deadline (fetch-start `TimeoutError`, body `BodyTimeoutError`, or the
+ * combo-per-model timeout) fired while the upstream was still processing the request,
+ * surfaced as a 504 tagged `errorType: "upstream_timeout"`.
+ *
+ * Such a timeout is NOT a provider rejection — the connection is healthy, we just gave
+ * up waiting — so the caller must skip the connection cooldown for it. Cooling the
+ * connection down on our own timeout penalises a healthy account and, when a provider
+ * has a single connection, blocks every subsequent request behind a self-inflicted
+ * cooldown. Antigravity keeps its own pre-response-timeout cooldown policy and is
+ * therefore excluded here.
+ */
+export function isSelfInflictedUpstreamTimeout(
+  status: number,
+  errorType: string | undefined | null,
+  provider: string
+): boolean {
+  return (
+    status === HTTP_STATUS.GATEWAY_TIMEOUT &&
+    errorType === "upstream_timeout" &&
+    provider !== "antigravity"
+  );
+}

src/sse/handlers/chat.ts

@@ -19,6 +19,7 @@ import {
 import { getModelInfo, getComboForModel } from "../services/model";
 import { resolveBareModelToConnectionDefault } from "@omniroute/open-sse/services/model.ts";
 import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
+import { isSelfInflictedUpstreamTimeout } from "@omniroute/open-sse/handlers/chatCore/cooldownClassification.ts";
 import { applyNoThinkingAlias } from "@omniroute/open-sse/utils/noThinkingAlias.ts";
 import { handleComboChat } from "@omniroute/open-sse/services/combo.ts";
 import { resolveComboConfig } from "@omniroute/open-sse/services/comboConfig.ts";
@@ -1502,7 +1503,14 @@ async function handleSingleModelChat(
         (credentials.providerSpecificData?.extraApiKeys as string[] | undefined) ?? [];
       const hasExtraKeys = extraKeys.length > 0 || connectionHasExtraKeys(credentials.connectionId);
       const is401 = result.status === 401;
-      const skipConnectionDisable = is401 && hasExtraKeys;
+      // Our own deadline timeout (fetch-start / body / combo-per-model, surfaced as a 504
+      // tagged "upstream_timeout") fired on a slow-but-not-failed upstream — the request
+      // was still being processed. The connection is healthy, so don't cool it down: a
+      // self-inflicted-timeout cooldown penalises a healthy account and, when a provider
+      // has a single connection, blocks every subsequent request.
+      const skipConnectionDisable =
+        (is401 && hasExtraKeys) ||
+        isSelfInflictedUpstreamTimeout(result.status, result.errorType, provider);
 
       const { shouldFallback, cooldownMs } = skipConnectionDisable
         ? { shouldFallback: false, cooldownMs: 0 }

tests/unit/self-inflicted-upstream-timeout.test.ts

@@ -0,0 +1,27 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+
+const { isSelfInflictedUpstreamTimeout } = await import(
+  "../../open-sse/handlers/chatCore/cooldownClassification.ts"
+);
+
+test("504 + upstream_timeout on a non-antigravity provider is self-inflicted (skip cooldown)", () => {
+  assert.equal(isSelfInflictedUpstreamTimeout(504, "upstream_timeout", "claude"), true);
+  assert.equal(isSelfInflictedUpstreamTimeout(504, "upstream_timeout", "openai"), true);
+});
+
+test("antigravity keeps its own pre-response-timeout cooldown policy", () => {
+  assert.equal(isSelfInflictedUpstreamTimeout(504, "upstream_timeout", "antigravity"), false);
+});
+
+test("a real provider 5xx / 429 is NOT a self-inflicted timeout", () => {
+  assert.equal(isSelfInflictedUpstreamTimeout(502, "server_error", "claude"), false);
+  assert.equal(isSelfInflictedUpstreamTimeout(500, undefined, "claude"), false);
+  assert.equal(isSelfInflictedUpstreamTimeout(429, "rate_limit", "claude"), false);
+});
+
+test("a 504 without the upstream_timeout tag is NOT self-inflicted", () => {
+  assert.equal(isSelfInflictedUpstreamTimeout(504, undefined, "claude"), false);
+  assert.equal(isSelfInflictedUpstreamTimeout(504, null, "claude"), false);
+  assert.equal(isSelfInflictedUpstreamTimeout(504, "authentication_error", "claude"), false);
+});