Skip to content

fix(credentials): do not cool down keys for at-capacity errors #6

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

Merged
yelixir-dev merged 1 commit into from
Oct 2, 2026
Merged

fix(credentials): do not cool down keys for at-capacity errors #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
59 changes: 38 additions & 21 deletions src/commandcode.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -388,6 +388,13 @@ export class CommandCodeClient implements CommandCodeUpstream {
model: body.params.model,
ignoreCooldown: true,
});
} else if (attempt === 0 && error instanceof NoAvailableCommandCodeCredentialError) {
// A cooldown is a preference, not a hard gate: when every credential is
// cooling down, serve with one anyway instead of failing the request.
credential = await this.router.select({
model: body.params.model,
ignoreCooldown: true,
});
} else {
if (lastError instanceof Error) throw lastError;
throw error;
Expand All @@ -405,7 +412,7 @@ export class CommandCodeClient implements CommandCodeUpstream {
else this.router.recordFailure(credential.id, { statusCode });
finalized = true;
};
const finalizeCallerAbort = () => {
const finalizeRelease = () => {
if (finalized) return;
this.router.release(credential.id);
finalized = true;
Expand All @@ -419,14 +426,17 @@ export class CommandCodeClient implements CommandCodeUpstream {
response.statusText,
await responseBody(response),
);
finalizeFailure(response.status);
lastError = error;
if (
attempt < maxAttempts - 1 &&
shouldRetry(response.status) &&
!effectiveSignal.aborted
) {
if (isFatalCredFailure(response.status)) fatalIds.add(credential.id);
const fatal = isFatalCredFailure(response.status);
const willRetry =
attempt < maxAttempts - 1 && shouldRetry(response.status) && !effectiveSignal.aborted;
// Only credential-scoped failures may start a cooldown. Provider-scoped
// statuses like 429/5xx hit every credential at once, so cooling them
// down would bench the whole pool over a single upstream incident.
if (fatal) finalizeFailure(response.status);
else finalizeRelease();
if (willRetry) {
if (fatal) fatalIds.add(credential.id);
else retryableFailed.add(credential.id);
await retryBackoff(attempt, this.config.commandCodeRetryBackoffMs ?? 250);
continue;
Expand All @@ -439,9 +449,10 @@ export class CommandCodeClient implements CommandCodeUpstream {
response.statusText,
"Upstream response body is empty",
);
finalizeFailure(response.status);
lastError = error;
if (attempt < maxAttempts - 1 && !effectiveSignal.aborted) {
const willRetry = attempt < maxAttempts - 1 && !effectiveSignal.aborted;
finalizeRelease();
if (willRetry) {
retryableFailed.add(credential.id);
await retryBackoff(attempt, this.config.commandCodeRetryBackoffMs ?? 250);
continue;
Expand All @@ -453,20 +464,21 @@ export class CommandCodeClient implements CommandCodeUpstream {
for await (const event of parseCommandCodeStream(response.body)) {
const statusCode = errorStatusCode(event);
if (event.type === "error") {
finalizeFailure(statusCode);
lastError = new CommandCodeHttpError(
statusCode ?? 502,
"CommandCode stream error",
event,
);
if (
const fatal = statusCode !== undefined && isFatalCredFailure(statusCode);
const willRetry =
!emittedVisibleEvent &&
attempt < maxAttempts - 1 &&
shouldRetry(statusCode) &&
!effectiveSignal.aborted
) {
if (statusCode !== undefined && isFatalCredFailure(statusCode))
fatalIds.add(credential.id);
!effectiveSignal.aborted;
if (fatal) finalizeFailure(statusCode);
else finalizeRelease();
if (willRetry) {
if (fatal) fatalIds.add(credential.id);
else retryableFailed.add(credential.id);
await retryBackoff(attempt, this.config.commandCodeRetryBackoffMs ?? 250);
continue attemptLoop;
Expand All @@ -481,12 +493,17 @@ export class CommandCodeClient implements CommandCodeUpstream {
return;
} catch (error) {
const statusCode = errorStatusCodeFromUnknown(error);
if (signal?.aborted === true) finalizeCallerAbort();
else finalizeFailure(statusCode);
lastError = error;
if (attempt < maxAttempts - 1 && shouldRetry(statusCode) && !effectiveSignal.aborted) {
if (statusCode !== undefined && isFatalCredFailure(statusCode))
fatalIds.add(credential.id);
const fatal = statusCode !== undefined && isFatalCredFailure(statusCode);
const willRetry =
signal?.aborted !== true &&
attempt < maxAttempts - 1 &&
shouldRetry(statusCode) &&
!effectiveSignal.aborted;
if (fatal) finalizeFailure(statusCode);
else finalizeRelease();
if (willRetry) {
if (fatal) fatalIds.add(credential.id);
else retryableFailed.add(credential.id);
await retryBackoff(attempt, this.config.commandCodeRetryBackoffMs ?? 250);
continue;
Expand Down
16 changes: 13 additions & 3 deletions src/provider.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,11 @@ export class CommandCodeProviderClient {
else this.router.recordFailure(credential.id, { statusCode });
finalized = true;
};
const finalizeRelease = () => {
if (finalized) return;
this.router.release(credential.id);
finalized = true;
};

try {
const response = await this.fetchChat(body, credential, effectiveSignal);
Expand All @@ -152,14 +157,18 @@ export class CommandCodeProviderClient {
response.statusText,
await responseBody(response),
);
finalizeFailure(response.status);
const fatal = isFatalCredFailure(response.status);
// Only credential-scoped failures may start a cooldown; provider-scoped
// statuses like 429/5xx hit every credential at once.
if (fatal) finalizeFailure(response.status);
else finalizeRelease();
lastError = error;
if (
attempt < maxAttempts - 1 &&
shouldRetryStatus(response.status) &&
!effectiveSignal.aborted
) {
if (isFatalCredFailure(response.status)) fatalIds.add(credential.id);
if (fatal) fatalIds.add(credential.id);
else retryableFailed.add(credential.id);
await retryBackoff(attempt, this.config.commandCodeRetryBackoffMs ?? 250);
continue;
Expand All @@ -171,7 +180,8 @@ export class CommandCodeProviderClient {
} catch (error) {
if (error instanceof CommandCodeHttpError && finalized) throw error;
const statusCode = error instanceof CommandCodeHttpError ? error.status : undefined;
finalizeFailure(statusCode);
if (statusCode !== undefined && isFatalCredFailure(statusCode)) finalizeFailure(statusCode);
else finalizeRelease();
lastError = error;
if (
attempt < maxAttempts - 1 &&
Expand Down
139 changes: 139 additions & 0 deletions tests/commandcode.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {
parseCommandCodeStream,
} from "../src/commandcode.js";
import type { BridgeConfig, CommandCodeEvent, CommandCodeGenerateBody } from "../src/types.js";
import type { CommandCodeCredentialRouter } from "../src/credential-router.js";

function streamFromChunks(chunks: string[]): ReadableStream<Uint8Array> {
const encoder = new TextEncoder();
Expand Down Expand Up @@ -434,6 +435,144 @@ describe("CommandCode client credential routing", () => {
expect(events).toContainEqual(expect.objectContaining({ type: "text-delta", text: "ok" }));
});

it("does not cool down a credential when a retried upstream error eventually succeeds", async () => {
const postResponses = [
new Response('data: {"type":"error","message":"temporary upstream error"}\n', {
status: 200,
}),
new Response(
'data: {"type":"text-delta","text":"ok"}\ndata: {"type":"finish","finishReason":"stop"}\n',
{ status: 200 },
),
];
const fetchMock = vi.fn<typeof fetch>(async (input, init) => {
const billing = billingResponse(String(input));
if (billing) return billing;
if (init?.method === "POST") return postResponses.shift()!;
throw new Error(`Unexpected fetch ${String(input)}`);
});
vi.stubGlobal("fetch", fetchMock);

const client = new CommandCodeClient({
...baseConfig,
commandCodeCredentials: [{ id: "alpha", apiKey: "alpha-secret", weight: 1 }],
});
const events = await collectEvents(client.generate(generateBody));

const posts = postCalls(fetchMock);
expect(posts).toHaveLength(2);
expect((posts[1]?.[1] as RequestInit).headers).toMatchObject({
Authorization: "Bearer alpha-secret",
});
expect(events).toContainEqual(expect.objectContaining({ type: "text-delta", text: "ok" }));

const diagnostics = await client.getCredentialDiagnostics();
expect(diagnostics[0]?.disabledUntil).toBeNull();
});

it("serves through a credential cooldown instead of failing the request", async () => {
const fetchMock = vi.fn<typeof fetch>(async (input, init) => {
const billing = billingResponse(String(input));
if (billing) return billing;
if (init?.method === "POST") {
return new Response(
'data: {"type":"text-delta","text":"ok"}\ndata: {"type":"finish","finishReason":"stop"}\n',
{ status: 200 },
);
}
throw new Error(`Unexpected fetch ${String(input)}`);
});
vi.stubGlobal("fetch", fetchMock);

const client = new CommandCodeClient({
...baseConfig,
commandCodeCredentials: [{ id: "alpha", apiKey: "alpha-secret", weight: 1 }],
});

await collectEvents(client.generate(generateBody));
const router = (client as unknown as { router: CommandCodeCredentialRouter }).router;
router.recordFailure("alpha", { statusCode: 429 });
const diagnostics = await client.getCredentialDiagnostics();
expect(diagnostics[0]?.disabledUntil).not.toBeNull();

const events = await collectEvents(client.generate(generateBody));
expect(events).toContainEqual(expect.objectContaining({ type: "text-delta", text: "ok" }));
const posts = postCalls(fetchMock);
expect(posts).toHaveLength(2);
expect((posts[1]?.[1] as RequestInit).headers).toMatchObject({
Authorization: "Bearer alpha-secret",
});
});

it("does not cool down the credential when the provider reports at capacity without a status", async () => {
let postCount = 0;
const atCapacity =
'data: {"type":"error","message":"The request limited providers for this model and they are currently at capacity. Providers considered: deepseek."}\n';
const fetchMock = vi.fn<typeof fetch>(async (input, init) => {
const billing = billingResponse(String(input));
if (billing) return billing;
if (init?.method === "POST") {
postCount += 1;
if (postCount === 1) return new Response(atCapacity, { status: 200 });
return new Response(
'data: {"type":"text-delta","text":"ok"}\ndata: {"type":"finish","finishReason":"stop"}\n',
{ status: 200 },
);
}
throw new Error(`Unexpected fetch ${String(input)}`);
});
vi.stubGlobal("fetch", fetchMock);

const client = new CommandCodeClient({
...baseConfig,
commandCodeCredentials: [{ id: "alpha", apiKey: "alpha-secret", weight: 1 }],
commandCodeRetryMaxAttempts: 1,
});

const events = await collectEvents(client.generate(generateBody));
expect(events).toContainEqual(expect.objectContaining({ type: "error" }));
const diagnostics = await client.getCredentialDiagnostics();
expect(diagnostics[0]?.disabledUntil).toBeNull();

const retryEvents = await collectEvents(client.generate(generateBody));
expect(retryEvents).toContainEqual(expect.objectContaining({ type: "text-delta", text: "ok" }));
expect(postCalls(fetchMock)).toHaveLength(2);
});

it("does not cool down the credential when a provider-scoped 429 exhausts retries", async () => {
let postCount = 0;
const fetchMock = vi.fn<typeof fetch>(async (input, init) => {
const billing = billingResponse(String(input));
if (billing) return billing;
if (init?.method === "POST") {
postCount += 1;
if (postCount === 1) return new Response("at capacity", { status: 429 });
return new Response(
'data: {"type":"text-delta","text":"ok"}\ndata: {"type":"finish","finishReason":"stop"}\n',
{ status: 200 },
);
}
throw new Error(`Unexpected fetch ${String(input)}`);
});
vi.stubGlobal("fetch", fetchMock);

const client = new CommandCodeClient({
...baseConfig,
commandCodeCredentials: [{ id: "alpha", apiKey: "alpha-secret", weight: 1 }],
commandCodeRetryMaxAttempts: 1,
});

await expect(collectEvents(client.generate(generateBody))).rejects.toBeInstanceOf(
CommandCodeHttpError,
);
const diagnostics = await client.getCredentialDiagnostics();
expect(diagnostics[0]?.disabledUntil).toBeNull();

const events = await collectEvents(client.generate(generateBody));
expect(events).toContainEqual(expect.objectContaining({ type: "text-delta", text: "ok" }));
expect(postCalls(fetchMock)).toHaveLength(2);
});

it("releases the credential without cooldown when the caller aborts mid-stream", async () => {
let hanging = true;
let postStarted: () => void = () => undefined;
Expand Down
Loading