smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): completed oneshot transient retry handling

Browse Source
This commit is contained in:
can1357
2026-08-16 02:12:05 +02:00
parent b445134b4d
commit df6d1e1ac5
4 changed files with 188 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/agent/src/compaction/branch-summarization.ts
+12 -6
View File
@@ -334,12 +334,18 @@ export async function generateBranchSummary(
];
// Call LLM for summarization
const response = await instrumentedCompleteSimple(
model,
{ systemPrompt: [SUMMARIZATION_SYSTEM_PROMPT], messages: summarizationMessages },
{ apiKey, signal, maxTokens: 2048, metadata },
{ telemetry: options.telemetry, oneshotKind: "branch_summary", completeImpl: options.completeImpl, retry: {} },
);
let response: AssistantMessage;
try {
response = await instrumentedCompleteSimple(
model,
{ systemPrompt: [SUMMARIZATION_SYSTEM_PROMPT], messages: summarizationMessages },
{ apiKey, signal, maxTokens: 2048, metadata },
{ telemetry: options.telemetry, oneshotKind: "branch_summary", completeImpl: options.completeImpl, retry: {} },
);
} catch (error) {
if (signal.aborted) return { aborted: true };
throw error;
}
// Check if aborted or errored
if (response.stopReason === "aborted") {
packages/agent/test/branch-summarization.test.ts
+55
View File
@@ -230,4 +230,59 @@ describe("branch summarization", () => {
expect(messages.some(m => m.role === "toolResult")).toBe(true);
});
test("returns an aborted result when cancelled during transient retry backoff", async () => {
const reason = new Error("user cancelled branch summary");
let aborted = false;
const signal = {
get aborted() {
return aborted;
},
get reason() {
return aborted ? reason : undefined;
},
addEventListener(type: string, listener: EventListenerOrEventListenerObject) {
if (type !== "abort") return;
aborted = true;
const event = new Event("abort");
if (typeof listener === "function") listener(event);
else listener.handleEvent(event);
},
removeEventListener() {},
} as unknown as AbortSignal;
const entries: SessionEntry[] = [
{
type: "message",
id: "user-1",
parentId: null,
timestamp: new Date(0).toISOString(),
message: { role: "user", content: "Summarize this branch.", timestamp: 0 },
},
];
let calls = 0;
const result = await generateBranchSummary(entries, {
model: MODEL,
apiKey: "test-api-key",
signal,
completeImpl: async () => {
calls += 1;
return {
role: "assistant",
content: [],
api: "mock",
provider: "mock",
model: "mock-model",
usage: ZERO_USAGE,
stopReason: "error",
errorStatus: 529,
errorMessage: "overloaded_error: Overloaded",
timestamp: 1,
};
},
});
expect(calls).toBe(1);
expect(result).toEqual({ aborted: true });
});
});
packages/ai/src/oneshot-retry.ts
+21 -3
View File
@@ -87,6 +87,7 @@ const DEFAULT_BASE_DELAY_MS = 500;
const DEFAULT_MAX_DELAY_MS = 30_000;
/** Cap on pure backoff growth. A provider hint may still exceed this, up to `maxDelayMs`. */
const BACKOFF_CEILING_MS = 8_000;
const RETRY_AFTER_MS_SUFFIX = /(?:^|\s)retry-after-ms=([0-9]+(?:\.[0-9]+)?)(?=\s|$)/i;
function backoffDelayMs(attempt: number, baseDelayMs: number): number {
const growth = Math.min(baseDelayMs * 2 ** (attempt - 1), BACKOFF_CEILING_MS);
@@ -96,8 +97,13 @@ function backoffDelayMs(attempt: number, baseDelayMs: number): number {
}
/** Retryable when the provider says transient, or when it says "wait, then retry". */
function isRetryableOneshotFailure(errorId: number): boolean {
function isRetryableOneshotFailure(errorId: number, errorStatus: number | undefined, errorMessage: string): boolean {
// llama.cpp reports deterministic tool-call JSON parse failures as HTTP 500.
// Replaying the same prompt produces the same malformed output.
if (AIError.LLAMA_CPP_TOOL_CALL_PARSE_PATTERN.test(errorMessage)) return false;
if (AIError.is(errorId, AIError.Flag.ContentBlocked)) return false;
return (
AIError.isTransientStatus(errorStatus) ||
AIError.is(errorId, AIError.Flag.Transient) ||
AIError.is(errorId, AIError.Flag.UsageLimit) ||
AIError.retriable(errorId)
@@ -171,8 +177,10 @@ export async function retryTransientCompletion(
? thrown.message
: String(thrown)
: ((message as AssistantMessage).errorMessage ?? "unknown error");
const errorStatus =
thrown !== undefined ? AIError.status(thrown) : (message as AssistantMessage).errorStatus;
const lastAttempt = attempt >= maxAttempts;
if (lastAttempt || !isRetryableOneshotFailure(errorId)) {
if (lastAttempt || !isRetryableOneshotFailure(errorId, errorStatus, errorMessage)) {
if (thrown !== undefined) throw thrown;
return message as AssistantMessage;
}
@@ -183,7 +191,17 @@ export async function retryTransientCompletion(
// errors (e.g. AnthropicApiError) carry their own headers.
const headers: HeadersLike = thrown !== undefined ? getHeadersFromError(thrown) : options?.getResponseHeaders?.();
const headerHintMs = getRetryAfterMsFromHeaders(headers);
const textHintMs = extractRetryHint(undefined, errorMessage);
const extractedTextHintMs = extractRetryHint(undefined, errorMessage);
const suffixValue = RETRY_AFTER_MS_SUFFIX.exec(errorMessage)?.[1];
const parsedSuffixMs = suffixValue === undefined ? undefined : Number(suffixValue);
const suffixHintMs =
parsedSuffixMs !== undefined && Number.isFinite(parsedSuffixMs) && parsedSuffixMs > 0
? Math.ceil(parsedSuffixMs)
: undefined;
const textHintMs =
extractedTextHintMs === undefined && suffixHintMs === undefined
? undefined
: Math.max(extractedTextHintMs ?? 0, suffixHintMs ?? 0);
const hintMs =
headerHintMs === undefined && textHintMs === undefined
? undefined
packages/ai/test/oneshot-retry.test.ts
+100
View File
@@ -77,6 +77,21 @@ describe("retryTransientCompletion", () => {
expect(final.stopReason).toBe("stop");
});
it("re-issues a status-only 503 error-stop", async () => {
let calls = 0;
const final = await retryTransientCompletion(() => {
calls += 1;
return Promise.resolve(
calls === 1
? message({ stopReason: "error", errorStatus: 503, errorMessage: "request failed" })
: message(),
);
}, fast);
expect(calls).toBe(2);
expect(final.stopReason).toBe("stop");
});
it("returns the failing message unchanged once attempts are exhausted, so caller fallbacks still apply", async () => {
let calls = 0;
const final = await retryTransientCompletion(() => {
@@ -109,6 +124,26 @@ describe("retryTransientCompletion", () => {
expect(final.stopReason).toBe("error");
});
it("does not retry a deterministic llama.cpp tool-call parse failure reported as 500", async () => {
let calls = 0;
const final = await retryTransientCompletion(
() => {
calls += 1;
return Promise.resolve(
message({
stopReason: "error",
errorStatus: 500,
errorMessage: "failed to parse tool call arguments as JSON",
}),
);
},
{ ...fast, maxAttempts: 5 },
);
expect(calls).toBe(1);
expect(final.stopReason).toBe("error");
});
it("retries a thrown transient error and rethrows the last one when exhausted", async () => {
let calls = 0;
const attempt = retryTransientCompletion(() => {
@@ -122,6 +157,22 @@ describe("retryTransientCompletion", () => {
expect(calls).toBe(3);
});
it("retries a thrown status-only 503 error", async () => {
let calls = 0;
const final = await retryTransientCompletion(() => {
calls += 1;
if (calls === 1) {
const error = new Error("request failed") as Error & { status: number };
error.status = 503;
throw error;
}
return Promise.resolve(message());
}, fast);
expect(calls).toBe(2);
expect(final.stopReason).toBe("stop");
});
it("does not retry a thrown non-transient error", async () => {
let calls = 0;
const attempt = retryTransientCompletion(() => {
@@ -228,6 +279,55 @@ describe("retryTransientCompletion", () => {
expect(observedDelay).toBe(120);
});
it("honors the canonical retry-after-ms error-message suffix", async () => {
let calls = 0;
let observedDelay = -1;
await retryTransientCompletion(
() => {
calls += 1;
return Promise.resolve(
calls === 1
? message({
stopReason: "error",
errorStatus: 429,
errorMessage: "rate_limit_error: too many requests retry-after-ms=5",
})
: message(),
);
},
{
maxAttempts: 2,
baseDelayMs: 1,
onRetry: info => {
observedDelay = info.delayMs;
},
},
);
expect(calls).toBe(2);
expect(observedDelay).toBe(5);
});
it("surfaces a canonical retry-after-ms suffix above maxDelayMs", async () => {
let calls = 0;
const final = await retryTransientCompletion(
() => {
calls += 1;
return Promise.resolve(
message({
stopReason: "error",
errorStatus: 429,
errorMessage: "rate_limit_error: too many requests retry-after-ms=12000",
}),
);
},
{ ...fast, maxDelayMs: 1_000 },
);
expect(calls).toBe(1);
expect(final.stopReason).toBe("error");
});
it("surfaces the failure when a retry-after header exceeds maxDelayMs", async () => {
let calls = 0;
const final = await retryTransientCompletion(