Skip to content

fix: scope retry health by model and region #1977

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
steebchen wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix: scope retry health by model and region #1977

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b8579f3
fix: retry invalid api key payloads
steebchen Apr 6, 2026
8391b63
fix: isolate key and provider health
steebchen Apr 6, 2026
b49cce3
test: harden flaky gateway specs
steebchen Apr 7, 2026
b45c134
Merge remote-tracking branch 'origin/main' into fix/retry-invalid-api…
steebchen Apr 7, 2026
86a36d2
fix: tighten retry health scoping
steebchen Apr 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions apps/gateway/src/api.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2597,11 +2597,11 @@ describe("api", () => {
expect(logs[0].hasError).toBe(true);
expect(logs[0].errorDetails?.statusCode).toBe(401);

expect(isTrackedKeyHealthy("provider-key-id-stream-auth-error")).toBe(
false,
);
expect(
getTrackedKeyMetrics("provider-key-id-stream-auth-error"),
isTrackedKeyHealthy("provider-key-id-stream-auth-error", "custom"),
).toBe(false);
expect(
getTrackedKeyMetrics("provider-key-id-stream-auth-error", "custom"),
).toMatchObject({
permanentlyBlacklisted: true,
totalRequests: 1,
Expand Down
103 changes: 79 additions & 24 deletions apps/gateway/src/chat/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ import {
MAX_RETRIES,
providerRetryKey,
selectNextProvider,
shouldRetryAlternateKey,
shouldRetryRequest,
} from "./tools/retry-with-fallback.js";
import {
Expand Down Expand Up @@ -1436,7 +1437,6 @@ chat.openapi(completions, async (c) => {
const customProviderKey = await findCustomProviderKey(
project.organizationId,
customProviderName,
requestId,
);
if (!customProviderKey) {
throw new HTTPException(400, {
Expand Down Expand Up @@ -1806,7 +1806,7 @@ chat.openapi(completions, async (c) => {
const providerKey = await findProviderKey(
project.organizationId,
usedProvider,
requestId,
modelInfo.id || stripRegionFromModelName(usedModel, usedRegion),
);
lockedRegion = providerKey
? resolveExplicitRegionFromProviderKey(providerKey)
Expand Down Expand Up @@ -2458,7 +2458,7 @@ chat.openapi(completions, async (c) => {
const providerKey = await findProviderKey(
project.organizationId,
requestedProvider,
requestId,
modelInfo.id || stripRegionFromModelName(usedModel, usedRegion),
);
explicitDirectRegion = providerKey
? resolveExplicitRegionFromProviderKey(providerKey)
Expand Down Expand Up @@ -2696,13 +2696,13 @@ chat.openapi(completions, async (c) => {
providerKey = await findCustomProviderKey(
project.organizationId,
customProviderName,
requestId,
baseModelName,
);
} else {
providerKey = await findProviderKey(
project.organizationId,
usedProvider,
requestId,
baseModelName,
);
Comment on lines 2696 to 2706
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Apr 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Complete the selectionScope propagation for the earlier key lookups.

These lookups are scoped now, but Line 1437 and Lines 1806-1809 / 2457-2460 still call findCustomProviderKey / findProviderKey without the new third argument. That leaves custom-provider validation and direct-provider region locking on global key health, so another model's failures can still hide a healthy key or lock the wrong region.

Also applies to: 2778-2788

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@apps/gateway/src/chat/chat.ts` around lines 2694 - 2704, Calls to
findCustomProviderKey and findProviderKey are missing the new selectionScope
argument in several places; update every invocation (including the earlier call
near the top and the ones that correspond to the referenced blocks) to pass the
current selectionScope value so the key lookups are region/selection-scoped.
Specifically, add selectionScope as the third parameter when calling
findCustomProviderKey(project.organizationId, customProviderName,
selectionScope) and when calling findProviderKey(project.organizationId,
usedProvider, selectionScope) (and similarly for the other direct-provider calls
noted), ensuring you use the same selectionScope variable used elsewhere in this
file so custom-provider validation and region locking are correctly scoped.

}

Expand Down Expand Up @@ -2760,7 +2760,9 @@ chat.openapi(completions, async (c) => {
});
}

const envResult = getProviderEnv(usedProvider);
const envResult = getProviderEnv(usedProvider, {
selectionScope: baseModelName,
});
usedToken = envResult.token;
configIndex = envResult.configIndex;
envVarName = envResult.envVarName;
Expand All @@ -2778,13 +2780,13 @@ chat.openapi(completions, async (c) => {
providerKey = await findCustomProviderKey(
project.organizationId,
customProviderName,
requestId,
baseModelName,
);
} else {
providerKey = await findProviderKey(
project.organizationId,
usedProvider,
requestId,
baseModelName,
);
}

Expand Down Expand Up @@ -2835,7 +2837,9 @@ chat.openapi(completions, async (c) => {
});
}

const envResult = getProviderEnv(usedProvider);
const envResult = getProviderEnv(usedProvider, {
selectionScope: baseModelName,
});
usedToken = envResult.token;
configIndex = envResult.configIndex;
envVarName = envResult.envVarName;
Expand Down Expand Up @@ -4718,10 +4722,21 @@ chat.openapi(completions, async (c) => {

// Report key health for the selected token source
if (envVarName !== undefined) {
reportKeyError(envVarName, configIndex, 0);
reportKeyError(
envVarName,
configIndex,
0,
undefined,
baseModelName,
);
}
if (providerKey?.id) {
reportTrackedKeyError(providerKey.id, 0);
reportTrackedKeyError(
providerKey.id,
0,
undefined,
baseModelName,
);
}

if (willRetrySameProvider && sameProviderRetryContext) {
Expand Down Expand Up @@ -4830,7 +4845,13 @@ chat.openapi(completions, async (c) => {
let sameProviderRetryContext: Awaited<
ReturnType<typeof resolveProviderContext>
> | null = null;
if (isRetryableErrorType(finishReason)) {
if (
shouldRetryAlternateKey(
finishReason,
res.status,
errorResponseText,
)
) {
rememberFailedKey(usedProvider, usedRegion, {
envVarName,
configIndex,
Expand Down Expand Up @@ -4954,13 +4975,15 @@ chat.openapi(completions, async (c) => {
configIndex,
res.status,
errorResponseText,
baseModelName,
);
}
if (providerKey?.id && finishReason !== "content_filter") {
reportTrackedKeyError(
providerKey.id,
res.status,
errorResponseText,
baseModelName,
);
}

Expand Down Expand Up @@ -5099,7 +5122,13 @@ chat.openapi(completions, async (c) => {
let sameProviderRetryContext: Awaited<
ReturnType<typeof resolveProviderContext>
> | null = null;
if (isRetryableErrorType(errorType)) {
if (
shouldRetryAlternateKey(
errorType,
inferredStatusCode,
errorResponseText,
)
) {
rememberFailedKey(usedProvider, usedRegion, {
envVarName,
configIndex,
Expand Down Expand Up @@ -5208,13 +5237,15 @@ chat.openapi(completions, async (c) => {
configIndex,
inferredStatusCode,
errorResponseText,
baseModelName,
);
}
if (providerKey?.id && errorType !== "content_filter") {
reportTrackedKeyError(
providerKey.id,
inferredStatusCode,
errorResponseText,
baseModelName,
);
}

Expand Down Expand Up @@ -7485,16 +7516,27 @@ chat.openapi(completions, async (c) => {
// Report key health for the selected token source
if (envVarName !== undefined) {
if (streamingError !== null) {
reportKeyError(envVarName, configIndex, streamingErrorStatusCode);
reportKeyError(
envVarName,
configIndex,
streamingErrorStatusCode,
undefined,
baseModelName,
);
} else {
reportKeySuccess(envVarName, configIndex);
reportKeySuccess(envVarName, configIndex, baseModelName);
}
}
if (providerKey?.id) {
if (streamingError !== null) {
reportTrackedKeyError(providerKey.id, streamingErrorStatusCode);
reportTrackedKeyError(
providerKey.id,
streamingErrorStatusCode,
undefined,
baseModelName,
);
} else {
reportTrackedKeySuccess(providerKey.id);
reportTrackedKeySuccess(providerKey.id, baseModelName);
}
}

Expand Down Expand Up @@ -7829,10 +7871,10 @@ chat.openapi(completions, async (c) => {

// Report key health for the selected token source
if (envVarName !== undefined) {
reportKeyError(envVarName, configIndex, 0);
reportKeyError(envVarName, configIndex, 0, undefined, baseModelName);
}
if (providerKey?.id) {
reportTrackedKeyError(providerKey.id, 0);
reportTrackedKeyError(providerKey.id, 0, undefined, baseModelName);
}

if (willRetrySameProvider && sameProviderRetryContext) {
Expand Down Expand Up @@ -8188,7 +8230,9 @@ chat.openapi(completions, async (c) => {
let sameProviderRetryContext: Awaited<
ReturnType<typeof resolveProviderContext>
> | null = null;
if (isRetryableErrorType(finishReason)) {
if (
shouldRetryAlternateKey(finishReason, res.status, errorResponseText)
) {
rememberFailedKey(usedProvider, usedRegion, {
envVarName,
configIndex,
Expand Down Expand Up @@ -8313,10 +8357,21 @@ chat.openapi(completions, async (c) => {
// Report key health for the selected token source
// Don't report content_filter as a key error - it's intentional provider behavior
if (envVarName !== undefined && finishReason !== "content_filter") {
reportKeyError(envVarName, configIndex, res.status, errorResponseText);
reportKeyError(
envVarName,
configIndex,
res.status,
errorResponseText,
baseModelName,
);
}
if (providerKey?.id && finishReason !== "content_filter") {
reportTrackedKeyError(providerKey.id, res.status, errorResponseText);
reportTrackedKeyError(
providerKey.id,
res.status,
errorResponseText,
baseModelName,
);
}

if (willRetrySameProvider && sameProviderRetryContext) {
Expand Down Expand Up @@ -9040,10 +9095,10 @@ chat.openapi(completions, async (c) => {
// Report key health for the selected token source
// Note: We don't report empty responses as key errors since they're not upstream errors
if (envVarName !== undefined) {
reportKeySuccess(envVarName, configIndex);
reportKeySuccess(envVarName, configIndex, baseModelName);
}
if (providerKey?.id) {
reportTrackedKeySuccess(providerKey.id);
reportTrackedKeySuccess(providerKey.id, baseModelName);
}

if (cachingEnabled && cacheKey && !stream && !hasEmptyNonStreamingResponse) {
Expand Down
9 changes: 9 additions & 0 deletions apps/gateway/src/chat/tools/get-finish-reason-from-error.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,15 @@ describe("getFinishReasonFromError", () => {
expect(getFinishReasonFromError(403)).toBe("gateway_error");
});

it("returns gateway_error for 400 invalid API key payloads", () => {
expect(
getFinishReasonFromError(
400,
'{"error":{"message":"API key not valid. Please pass a valid API key.","type":"authentication_error","code":"invalid_api_key"}}',
),
).toBe("gateway_error");
});

it("returns client_error when no error text provided for other 4xx", () => {
expect(getFinishReasonFromError(400)).toBe("client_error");
expect(getFinishReasonFromError(422)).toBe("client_error");
Expand Down
10 changes: 8 additions & 2 deletions apps/gateway/src/chat/tools/get-finish-reason-from-error.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import { hasInvalidProviderCredentialError } from "@/lib/provider-auth-errors.js";

/**
* Determines the appropriate finish reason based on HTTP status code and error message
* 5xx status codes indicate upstream provider errors
Expand Down Expand Up @@ -55,8 +57,12 @@ export function getFinishReasonFromError(
return "content_filter";
}

// 401/403 usually indicate invalid or unauthorized provider credentials
if (statusCode === 401 || statusCode === 403) {
// 401/403 and known provider credential payloads indicate bad provider keys.
if (
statusCode === 401 ||
statusCode === 403 ||
hasInvalidProviderCredentialError(errorText)
) {
return "gateway_error";
}

Expand Down
18 changes: 18 additions & 0 deletions apps/gateway/src/chat/tools/get-provider-env.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { afterEach, beforeEach, describe, expect, it } from "vitest";

import { reportKeyError, resetKeyHealth } from "@/lib/api-key-health.js";
import { resetRoundRobinCounters } from "@/lib/round-robin-env.js";

import { getProviderEnv } from "./get-provider-env.js";
Expand All @@ -9,6 +10,7 @@ describe("getProviderEnv", () => {

beforeEach(() => {
resetRoundRobinCounters();
resetKeyHealth();
process.env.LLM_OPENAI_API_KEY = "sk-openai-a,sk-openai-b,sk-openai-c";
});

Expand Down Expand Up @@ -55,4 +57,20 @@ describe("getProviderEnv", () => {
expect(thirdKey.token).toBe("sk-openai-c");
expect(thirdKey.configIndex).toBe(2);
});

it("passes selection scope through to env key health", () => {
reportKeyError("LLM_OPENAI_API_KEY", 0, 500, undefined, "gpt-4");
reportKeyError("LLM_OPENAI_API_KEY", 0, 500, undefined, "gpt-4");
reportKeyError("LLM_OPENAI_API_KEY", 0, 500, undefined, "gpt-4");

const gpt4Selection = getProviderEnv("openai", {
selectionScope: "gpt-4",
});
const claudeSelection = getProviderEnv("openai", {
selectionScope: "claude-3-5-sonnet",
});

expect(gpt4Selection.configIndex).toBe(1);
expect(claudeSelection.configIndex).toBe(0);
});
});
6 changes: 4 additions & 2 deletions apps/gateway/src/chat/tools/get-provider-env.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export interface ProviderEnvResult {
interface GetProviderEnvOptions {
advanceRoundRobin?: boolean;
excludedIndices?: ReadonlySet<number>;
selectionScope?: string;
}

/**
Expand Down Expand Up @@ -62,9 +63,10 @@ export function getProviderEnv(

const advanceRoundRobin = options.advanceRoundRobin ?? true;
const excludedIndices = options.excludedIndices;
const selectionScope = options.selectionScope;
const result = advanceRoundRobin
? getRoundRobinValue(envVar, envValue, excludedIndices)
: peekRoundRobinValue(envVar, envValue, excludedIndices);
? getRoundRobinValue(envVar, envValue, selectionScope, excludedIndices)
: peekRoundRobinValue(envVar, envValue, selectionScope, excludedIndices);

return { token: result.value, configIndex: result.index, envVarName: envVar };
}
Loading
Loading