fix(ai): hardened core stream, retry, and abort infrastructure
abort-aware auth retry loop preserving resolver errors; EventStream.end() can no longer strand .result(); Copilot retry honors Retry-After; DSML hold-back only triggers on real section prefixes and marks capped params explicitly; idle-iterator hoists racers with bounded reaction retention; validation errors truncate embedded args; removed dead leaky iterateUntilAbort.
This commit is contained in:
can1357
@@ -432,8 +432,16 @@ export function streamSimple<TApi extends Api>(
|
||||
let lastKey: string | undefined;
|
||||
try {
|
||||
lastKey = (await apiKeyResolver({ lastChance: false, error: undefined, signal })) || undefined;
|
||||
} catch {
|
||||
lastKey = undefined;
|
||||
} catch (error) {
|
||||
// A thrown resolver is a broker/OAuth/network failure, not a missing
|
||||
// key — surface the cause instead of masking it as "No API key".
|
||||
outer.fail(
|
||||
new Error(
|
||||
`Failed to resolve API key for provider ${model.provider}: ${error instanceof Error ? error.message : String(error)}`,
|
||||
{ cause: error },
|
||||
),
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (lastKey === undefined) {
|
||||
outer.fail(new Error(`No API key for provider: ${model.provider}`));
|
||||
@@ -446,6 +454,9 @@ export function streamSimple<TApi extends Api>(
|
||||
// resolver yields the same key it just tried or `undefined`; the
|
||||
// final step's attempt clears the capture flag so it emits directly.
|
||||
for (let step = 0; step < AUTH_RETRY_STEPS.length; step++) {
|
||||
// Caller aborted between attempts: don't mint a fresh token or fire
|
||||
// another doomed request — emit the captured failure instead.
|
||||
if (signal?.aborted) break;
|
||||
const nextKey = await resolveRetryKey(apiKeyResolver, AUTH_RETRY_STEPS[step]!, failure.error, signal);
|
||||
if (nextKey === undefined || nextKey === lastKey) continue;
|
||||
lastKey = nextKey;
|
||||
|
||||
@@ -49,3 +49,17 @@ export function createAbortSourceTracker(callerSignal?: AbortSignal): AbortSourc
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Race a shared promise against a caller's AbortSignal without coupling the
|
||||
* underlying work to that signal. The shared promise keeps running (and caches
|
||||
* its result) even when an individual caller bails out.
|
||||
*/
|
||||
export function raceWithSignal<T>(promise: Promise<T>, signal: AbortSignal | undefined): Promise<T> {
|
||||
if (!signal) return promise;
|
||||
if (signal.aborted) return Promise.reject(signal.reason ?? new Error("Request was aborted"));
|
||||
const { promise: aborted, reject } = Promise.withResolvers<never>();
|
||||
const onAbort = () => reject(signal.reason ?? new Error("Request was aborted"));
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
return Promise.race([promise, aborted]).finally(() => signal.removeEventListener("abort", onAbort));
|
||||
}
|
||||
|
||||
@@ -1,69 +0,0 @@
|
||||
function abortReason(signal: AbortSignal): Error {
|
||||
const reason = signal.reason;
|
||||
if (reason instanceof Error) return reason;
|
||||
if (typeof reason === "string") return new Error(reason);
|
||||
return new Error("Request was aborted");
|
||||
}
|
||||
|
||||
/**
|
||||
* Iterates a provider stream until it yields, ends, errors, or the caller aborts.
|
||||
*/
|
||||
export async function* iterateUntilAbort<T>(iterable: AsyncIterable<T>, signal?: AbortSignal): AsyncGenerator<T> {
|
||||
const iterator = iterable[Symbol.asyncIterator]();
|
||||
const closeIterator = (): void => {
|
||||
const returnPromise = iterator.return?.();
|
||||
if (returnPromise) {
|
||||
void returnPromise.catch(() => {});
|
||||
}
|
||||
};
|
||||
|
||||
if (signal?.aborted) {
|
||||
closeIterator();
|
||||
throw abortReason(signal);
|
||||
}
|
||||
|
||||
const withResult = (promise: Promise<IteratorResult<T>>) =>
|
||||
promise.then(
|
||||
result => ({ kind: "next" as const, result }),
|
||||
error => ({ kind: "error" as const, error }),
|
||||
);
|
||||
|
||||
while (true) {
|
||||
if (signal?.aborted) {
|
||||
closeIterator();
|
||||
throw abortReason(signal);
|
||||
}
|
||||
const racers: Array<
|
||||
Promise<{ kind: "next"; result: IteratorResult<T> } | { kind: "error"; error: unknown } | { kind: "abort" }>
|
||||
> = [withResult(iterator.next())];
|
||||
let abortListener: (() => void) | undefined;
|
||||
let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
|
||||
if (signal) {
|
||||
const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
|
||||
resolveAbort = resolve;
|
||||
abortListener = () => resolve({ kind: "abort" });
|
||||
signal.addEventListener("abort", abortListener, { once: true });
|
||||
racers.push(promise);
|
||||
}
|
||||
|
||||
try {
|
||||
const outcome = await Promise.race(racers);
|
||||
if (outcome.kind === "abort") {
|
||||
closeIterator();
|
||||
throw abortReason(signal!);
|
||||
}
|
||||
if (outcome.kind === "error") {
|
||||
throw outcome.error;
|
||||
}
|
||||
if (outcome.result.done) {
|
||||
return;
|
||||
}
|
||||
yield outcome.result.value;
|
||||
} finally {
|
||||
if (abortListener && signal) {
|
||||
signal.removeEventListener("abort", abortListener);
|
||||
}
|
||||
resolveAbort?.({ kind: "abort" });
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,8 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
|
||||
queue: T[] = [];
|
||||
waiting: Array<{ resolve: (value: IteratorResult<T>) => void; reject: (err: unknown) => void }> = [];
|
||||
done = false;
|
||||
/** True once finalResultPromise has been resolved or rejected. */
|
||||
resultSettled = false;
|
||||
#failed = false;
|
||||
#error: unknown = undefined;
|
||||
finalResultPromise: Promise<R>;
|
||||
@@ -30,6 +32,7 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
|
||||
|
||||
if (this.isComplete(event)) {
|
||||
this.done = true;
|
||||
this.resultSettled = true;
|
||||
this.resolveFinalResult(this.extractResult(event));
|
||||
}
|
||||
|
||||
@@ -54,7 +57,13 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
|
||||
end(result?: R): void {
|
||||
this.done = true;
|
||||
if (result !== undefined) {
|
||||
this.resultSettled = true;
|
||||
this.resolveFinalResult(result);
|
||||
} else if (!this.resultSettled) {
|
||||
// end() without a terminal value must still settle result() —
|
||||
// otherwise complete()/result() awaits hang forever.
|
||||
this.resultSettled = true;
|
||||
this.rejectFinalResult(new Error("Stream ended without a final result"));
|
||||
}
|
||||
// Notify all waiting consumers that we're done
|
||||
while (this.waiting.length > 0) {
|
||||
@@ -75,6 +84,7 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
|
||||
this.done = true;
|
||||
this.#failed = true;
|
||||
this.#error = err;
|
||||
this.resultSettled = true;
|
||||
this.rejectFinalResult(err);
|
||||
while (this.waiting.length > 0) {
|
||||
const waiter = this.waiting.shift()!;
|
||||
@@ -126,6 +136,7 @@ export class AssistantMessageEventStream extends EventStream<AssistantMessageEve
|
||||
// Completion resolves the final result and still emits the terminal event.
|
||||
if (this.isComplete(event)) {
|
||||
this.done = true;
|
||||
this.resultSettled = true;
|
||||
this.resolveFinalResult(this.extractResult(event));
|
||||
}
|
||||
|
||||
@@ -135,7 +146,13 @@ export class AssistantMessageEventStream extends EventStream<AssistantMessageEve
|
||||
override end(result?: AssistantMessage): void {
|
||||
this.done = true;
|
||||
if (result !== undefined) {
|
||||
this.resultSettled = true;
|
||||
this.resolveFinalResult(result);
|
||||
} else if (!this.resultSettled) {
|
||||
// Mirror the base class: a result-less end() must not leave
|
||||
// result() pending forever.
|
||||
this.resultSettled = true;
|
||||
this.rejectFinalResult(new Error("Stream ended without a final result"));
|
||||
}
|
||||
this.endWaiting();
|
||||
}
|
||||
|
||||
@@ -2,6 +2,8 @@ import { $env } from "@oh-my-pi/pi-utils";
|
||||
|
||||
const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 120_000;
|
||||
const DEFAULT_STREAM_FIRST_EVENT_TIMEOUT_MS = 100_000;
|
||||
/** Re-mint persistent race promises every N iterations (see hoisted-racer comment). */
|
||||
const RACER_REMINT_INTERVAL = 1024;
|
||||
|
||||
function normalizeIdleTimeoutMs(value: string | undefined, fallback: number): number | undefined {
|
||||
if (value === undefined) return fallback;
|
||||
@@ -164,109 +166,184 @@ export async function* iterateWithIdleTimeout<T>(
|
||||
(firstItemTimeoutMs === undefined || firstItemTimeoutMs <= 0) &&
|
||||
(options.idleTimeoutMs === undefined || options.idleTimeoutMs <= 0);
|
||||
|
||||
while (true) {
|
||||
let activeTimeoutMs: number | undefined;
|
||||
if (awaitingFirstItem) {
|
||||
if (firstItemDeadlineMs !== undefined) {
|
||||
activeTimeoutMs = firstItemDeadlineMs - Date.now();
|
||||
if (activeTimeoutMs <= 0) {
|
||||
options.onFirstItemTimeout?.();
|
||||
closeIterator();
|
||||
throw new Error(options.firstItemErrorMessage ?? options.errorMessage);
|
||||
}
|
||||
}
|
||||
} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
|
||||
activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
|
||||
if (activeTimeoutMs <= 0) {
|
||||
options.onIdle?.();
|
||||
closeIterator();
|
||||
throw new Error(options.errorMessage);
|
||||
}
|
||||
// Persistent racers, hoisted out of the per-item loop. The abort promise can
|
||||
// only ever resolve once (abort latches), and a timeout resolution always
|
||||
// precedes a throw — so neither needs per-item re-creation. This keeps the
|
||||
// token hot path free of timer create/destroy and listener churn.
|
||||
//
|
||||
// Each Promise.race() call still attaches a reaction record to every pending
|
||||
// racer, and those records live until the racer settles — so a never-firing
|
||||
// abort/timeout promise would accumulate one record per streamed item for
|
||||
// the stream's whole life. The loop re-mints both promises every
|
||||
// RACER_REMINT_INTERVAL iterations to keep that retention bounded; the
|
||||
// listener and timer callbacks resolve through late-bound variables so a
|
||||
// re-mint never strands them.
|
||||
let abortPromise: Promise<{ kind: "abort" }> | undefined;
|
||||
let abortListener: (() => void) | undefined;
|
||||
let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
|
||||
if (abortSignal) {
|
||||
const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
|
||||
resolveAbort = resolve;
|
||||
abortListener = () => resolveAbort?.({ kind: "abort" });
|
||||
abortSignal.addEventListener("abort", abortListener, { once: true });
|
||||
abortPromise = promise;
|
||||
}
|
||||
|
||||
let timeoutPromise: Promise<{ kind: "timeout" }> | undefined;
|
||||
let resolveTimeout: ((value: { kind: "timeout" }) => void) | undefined;
|
||||
let timeoutFired = false;
|
||||
let timer: NodeJS.Timeout | undefined;
|
||||
let timerFireAtMs = Infinity;
|
||||
|
||||
const currentDeadlineMs = (): number | undefined => {
|
||||
if (awaitingFirstItem) return firstItemDeadlineMs;
|
||||
if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
|
||||
return lastProgressAt + options.idleTimeoutMs;
|
||||
}
|
||||
|
||||
const nextResultPromise = withRacy(iterator.next());
|
||||
|
||||
const racers: Array<
|
||||
Promise<
|
||||
| { kind: "next"; result: IteratorResult<T> }
|
||||
| { kind: "error"; error: unknown }
|
||||
| { kind: "timeout" }
|
||||
| { kind: "abort" }
|
||||
>
|
||||
> = [nextResultPromise];
|
||||
|
||||
let timer: NodeJS.Timeout | undefined;
|
||||
let resolveTimeout: ((value: { kind: "timeout" }) => void) | undefined;
|
||||
const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
|
||||
if (enforceTimeout) {
|
||||
return undefined;
|
||||
};
|
||||
const onTimerFire = (): void => {
|
||||
timer = undefined;
|
||||
timerFireAtMs = Infinity;
|
||||
const deadlineMs = currentDeadlineMs();
|
||||
if (deadlineMs === undefined) return;
|
||||
const remainingMs = deadlineMs - Date.now();
|
||||
if (remainingMs > 0) {
|
||||
// Progress moved the deadline since this timer was armed — re-arm for
|
||||
// the remainder. One stale wake per idle period, not one per item.
|
||||
timerFireAtMs = deadlineMs;
|
||||
timer = setTimeout(onTimerFire, remainingMs);
|
||||
return;
|
||||
}
|
||||
timeoutFired = true;
|
||||
resolveTimeout?.({ kind: "timeout" });
|
||||
};
|
||||
const armTimer = (deadlineMs: number): void => {
|
||||
if (timeoutPromise === undefined || timeoutFired) {
|
||||
// A fired-but-unconsumed resolution (the item won the same race) is
|
||||
// stale — racing it again would fake a timeout, so mint a fresh one.
|
||||
const { promise, resolve } = Promise.withResolvers<{ kind: "timeout" }>();
|
||||
timeoutPromise = promise;
|
||||
resolveTimeout = resolve;
|
||||
timer = setTimeout(() => resolve({ kind: "timeout" }), activeTimeoutMs);
|
||||
racers.push(promise);
|
||||
timeoutFired = false;
|
||||
}
|
||||
|
||||
let abortListener: (() => void) | undefined;
|
||||
let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
|
||||
if (abortSignal) {
|
||||
const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
|
||||
resolveAbort = resolve;
|
||||
abortListener = () => resolve({ kind: "abort" });
|
||||
abortSignal.addEventListener("abort", abortListener, { once: true });
|
||||
racers.push(promise);
|
||||
if (timer !== undefined) {
|
||||
// An armed timer firing at or before the new deadline re-arms itself.
|
||||
if (timerFireAtMs <= deadlineMs) return;
|
||||
clearTimeout(timer);
|
||||
}
|
||||
timerFireAtMs = deadlineMs;
|
||||
timer = setTimeout(onTimerFire, Math.max(0, deadlineMs - Date.now()));
|
||||
};
|
||||
|
||||
// Tracks whether this iteration handed an item to the consumer and resumed
|
||||
// normally. Any other exit — internal throw, `done` return, or the consumer
|
||||
// abandoning us via `.return()`/`.throw()` at the `yield` below — must close
|
||||
// the upstream iterator so the underlying SSE body / SDK stream (and its
|
||||
// socket) is released instead of being left suspended.
|
||||
let continuing = false;
|
||||
try {
|
||||
const outcome = await Promise.race(racers);
|
||||
if (outcome.kind === "abort") {
|
||||
closeIterator();
|
||||
throw abortReason(abortSignal!);
|
||||
}
|
||||
if (outcome.kind === "timeout") {
|
||||
if (!awaitingFirstItem) {
|
||||
options.onIdle?.();
|
||||
} else {
|
||||
options.onFirstItemTimeout?.();
|
||||
try {
|
||||
let raceCount = 0;
|
||||
while (true) {
|
||||
if (++raceCount % RACER_REMINT_INTERVAL === 0) {
|
||||
if (abortPromise !== undefined && !abortSignal!.aborted) {
|
||||
const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
|
||||
resolveAbort = resolve;
|
||||
abortPromise = promise;
|
||||
}
|
||||
if (timeoutPromise !== undefined && !timeoutFired) {
|
||||
const { promise, resolve } = Promise.withResolvers<{ kind: "timeout" }>();
|
||||
resolveTimeout = resolve;
|
||||
timeoutPromise = promise;
|
||||
}
|
||||
closeIterator();
|
||||
throw new Error(
|
||||
!awaitingFirstItem ? options.errorMessage : (options.firstItemErrorMessage ?? options.errorMessage),
|
||||
);
|
||||
}
|
||||
if (outcome.kind === "error") {
|
||||
throw outcome.error;
|
||||
let activeTimeoutMs: number | undefined;
|
||||
if (awaitingFirstItem) {
|
||||
if (firstItemDeadlineMs !== undefined) {
|
||||
activeTimeoutMs = firstItemDeadlineMs - Date.now();
|
||||
if (activeTimeoutMs <= 0) {
|
||||
options.onFirstItemTimeout?.();
|
||||
closeIterator();
|
||||
throw new Error(options.firstItemErrorMessage ?? options.errorMessage);
|
||||
}
|
||||
}
|
||||
} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
|
||||
activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
|
||||
if (activeTimeoutMs <= 0) {
|
||||
options.onIdle?.();
|
||||
closeIterator();
|
||||
throw new Error(options.errorMessage);
|
||||
}
|
||||
}
|
||||
if (outcome.result.done) {
|
||||
markFirstItemReceived();
|
||||
return;
|
||||
|
||||
const nextResultPromise = withRacy(iterator.next());
|
||||
|
||||
const racers: Array<
|
||||
Promise<
|
||||
| { kind: "next"; result: IteratorResult<T> }
|
||||
| { kind: "error"; error: unknown }
|
||||
| { kind: "timeout" }
|
||||
| { kind: "abort" }
|
||||
>
|
||||
> = [nextResultPromise];
|
||||
|
||||
const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
|
||||
if (enforceTimeout) {
|
||||
armTimer(Date.now() + activeTimeoutMs!);
|
||||
racers.push(timeoutPromise!);
|
||||
}
|
||||
const item = outcome.result.value;
|
||||
// Non-progress items (e.g. provider keepalives, synthetic `start` events that
|
||||
// arrive before the model has produced any tokens) MUST NOT flip us out of
|
||||
// `awaitingFirstItem`. Otherwise the next iteration switches from the (longer)
|
||||
// first-item watchdog to the (shorter) idle watchdog while we're still waiting
|
||||
// on the model's first real output.
|
||||
if (isProgressItem(item)) {
|
||||
markFirstItemReceived();
|
||||
lastProgressAt = Date.now();
|
||||
if (abortPromise) {
|
||||
racers.push(abortPromise);
|
||||
}
|
||||
yield item;
|
||||
continuing = true;
|
||||
} finally {
|
||||
if (!continuing) closeIterator();
|
||||
if (timer !== undefined) clearTimeout(timer);
|
||||
// Resolve dangling promises so the racers don't leak (Promise.race is one-shot).
|
||||
resolveTimeout?.({ kind: "timeout" });
|
||||
if (abortListener && abortSignal) {
|
||||
abortSignal.removeEventListener("abort", abortListener);
|
||||
|
||||
// Tracks whether this iteration handed an item to the consumer and resumed
|
||||
// normally. Any other exit — internal throw, `done` return, or the consumer
|
||||
// abandoning us via `.return()`/`.throw()` at the `yield` below — must close
|
||||
// the upstream iterator so the underlying SSE body / SDK stream (and its
|
||||
// socket) is released instead of being left suspended.
|
||||
let continuing = false;
|
||||
try {
|
||||
const outcome = await Promise.race(racers);
|
||||
if (outcome.kind === "abort") {
|
||||
closeIterator();
|
||||
throw abortReason(abortSignal!);
|
||||
}
|
||||
if (outcome.kind === "timeout") {
|
||||
if (!awaitingFirstItem) {
|
||||
options.onIdle?.();
|
||||
} else {
|
||||
options.onFirstItemTimeout?.();
|
||||
}
|
||||
closeIterator();
|
||||
throw new Error(
|
||||
!awaitingFirstItem ? options.errorMessage : (options.firstItemErrorMessage ?? options.errorMessage),
|
||||
);
|
||||
}
|
||||
if (outcome.kind === "error") {
|
||||
throw outcome.error;
|
||||
}
|
||||
if (outcome.result.done) {
|
||||
markFirstItemReceived();
|
||||
return;
|
||||
}
|
||||
const item = outcome.result.value;
|
||||
// Non-progress items (e.g. provider keepalives, synthetic `start` events that
|
||||
// arrive before the model has produced any tokens) MUST NOT flip us out of
|
||||
// `awaitingFirstItem`. Otherwise the next iteration switches from the (longer)
|
||||
// first-item watchdog to the (shorter) idle watchdog while we're still waiting
|
||||
// on the model's first real output.
|
||||
if (isProgressItem(item)) {
|
||||
markFirstItemReceived();
|
||||
lastProgressAt = Date.now();
|
||||
}
|
||||
yield item;
|
||||
continuing = true;
|
||||
} finally {
|
||||
if (!continuing) closeIterator();
|
||||
}
|
||||
resolveAbort?.({ kind: "abort" });
|
||||
}
|
||||
} finally {
|
||||
if (timer !== undefined) clearTimeout(timer);
|
||||
// Settle the persistent racers so the final Promise.race releases them.
|
||||
resolveTimeout?.({ kind: "timeout" });
|
||||
if (abortListener && abortSignal) {
|
||||
abortSignal.removeEventListener("abort", abortListener);
|
||||
}
|
||||
resolveAbort?.({ kind: "abort" });
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -28,7 +28,7 @@ export function getRetryAfterMsFromHeaders(headers: HeadersLike): number | undef
|
||||
return Math.max(...candidates);
|
||||
}
|
||||
|
||||
function getHeadersFromError(error: unknown): HeadersLike {
|
||||
export function getHeadersFromError(error: unknown): HeadersLike {
|
||||
if (!error || typeof error !== "object") return undefined;
|
||||
const record = error as { headers?: unknown; response?: { headers?: unknown }; cause?: unknown };
|
||||
const direct = extractHeaders(record.headers) ?? extractHeaders(record.response?.headers);
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { scheduler } from "node:timers/promises";
|
||||
import { extractHttpStatusFromError, isRetryableError } from "@oh-my-pi/pi-utils";
|
||||
import { getHeadersFromError, getRetryAfterMsFromHeaders } from "./retry-after";
|
||||
|
||||
/**
|
||||
* GitHub Copilot intermittently rejects preview models (gpt-5.3-codex,
|
||||
@@ -24,6 +25,8 @@ export function isCopilotTransientModelError(error: unknown): boolean {
|
||||
|
||||
const COPILOT_MODEL_RETRY_MAX_ATTEMPTS = 3;
|
||||
const COPILOT_MODEL_RETRY_BASE_DELAY_MS = 400;
|
||||
/** Longest server-requested backoff we are willing to sit out before giving up. */
|
||||
const COPILOT_RETRY_AFTER_MAX_WAIT_MS = 30_000;
|
||||
|
||||
/**
|
||||
* Wrap an initial Copilot request so transient `model_not_supported` 400s are
|
||||
@@ -49,9 +52,23 @@ export async function callWithCopilotModelRetry<T>(
|
||||
// guaranteed-dead attempt — surface the original error, not the
|
||||
// scheduler's AbortError.
|
||||
if (options.signal?.aborted) throw error;
|
||||
if (!isCopilotTransientModelError(error) && !isRetryableError(error)) throw error;
|
||||
const transientModelError = isCopilotTransientModelError(error);
|
||||
if (!transientModelError && !isRetryableError(error)) throw error;
|
||||
if (attempt === COPILOT_MODEL_RETRY_MAX_ATTEMPTS - 1) break;
|
||||
await scheduler.wait(retryBaseDelayMs * (attempt + 1), { signal: options.signal });
|
||||
let delayMs = retryBaseDelayMs * (attempt + 1);
|
||||
if (!transientModelError) {
|
||||
const status = extractHttpStatusFromError(error);
|
||||
if (status !== undefined) {
|
||||
// Status-bearing retryable errors (429/5xx) are only re-sent when
|
||||
// the server told us when to come back — a blind fixed-delay retry
|
||||
// of a rate limit just burns the remaining attempts. Status-less
|
||||
// transport blips (socket close, h2 reset) keep the linear backoff.
|
||||
const retryAfterMs = getRetryAfterMsFromHeaders(getHeadersFromError(error));
|
||||
if (retryAfterMs === undefined || retryAfterMs > COPILOT_RETRY_AFTER_MAX_WAIT_MS) throw error;
|
||||
delayMs = Math.max(delayMs, retryAfterMs);
|
||||
}
|
||||
}
|
||||
await scheduler.wait(delayMs, { signal: options.signal });
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
|
||||
@@ -36,6 +36,8 @@ const DSML_PARAMETER_OPEN_RE = new RegExp(
|
||||
"y",
|
||||
);
|
||||
const DSML_PARAMETER_CLOSE_RE = new RegExp(`</${DSML_PIPE}DSML${DSML_PIPE}parameter>`, "y");
|
||||
/** Canonical DSML section-open shape; `|` positions accept either pipe variant. */
|
||||
const DSML_SECTION_OPEN_TEMPLATE = "<|DSML|tool_calls>";
|
||||
|
||||
const THINK_OPEN = "<think>";
|
||||
const THINK_CLOSE = "</think>";
|
||||
@@ -81,6 +83,7 @@ type XmlToolState =
|
||||
readonly paramName: string;
|
||||
readonly isString: boolean;
|
||||
value: string;
|
||||
truncated?: boolean;
|
||||
};
|
||||
|
||||
type ThinkingTag = { readonly open: string; readonly close: string };
|
||||
@@ -429,12 +432,25 @@ export class StreamMarkupHealing {
|
||||
continue;
|
||||
}
|
||||
} else if (this.#tryMatch(config.parameterClose)) {
|
||||
state.args[state.paramName] = coerceXmlParamValue(state.value, state.isString);
|
||||
// A capped value executes with silently corrupted input unless the
|
||||
// truncation is made explicit — the marker fails JSON params loudly
|
||||
// and tells the model/tool what happened to string params.
|
||||
const paramValue = state.truncated
|
||||
? `${state.value}\n…[parameter truncated: exceeded ${MAX_XML_PARAM_VALUE_LENGTH} bytes]`
|
||||
: state.value;
|
||||
state.args[state.paramName] = coerceXmlParamValue(paramValue, state.isString);
|
||||
config.setState({ kind: "invoke", name: state.invokeName, args: state.args });
|
||||
continue;
|
||||
}
|
||||
|
||||
if (this.#startsWithPartialXmlTag()) break;
|
||||
if (state.kind === "idle") {
|
||||
// In idle, a bare `<` is legitimate output (`a < b`, generics, JSX).
|
||||
// Only hold back tails that could still grow into the DSML
|
||||
// section-open tag; everything else flows through immediately.
|
||||
if (this.#startsWithPartialDsmlSectionOpen()) break;
|
||||
} else if (this.#startsWithPartialXmlTag()) {
|
||||
break;
|
||||
}
|
||||
|
||||
const ch = this.#buffer[this.#offset]!;
|
||||
this.#offset += 1;
|
||||
@@ -443,11 +459,15 @@ export class StreamMarkupHealing {
|
||||
continue;
|
||||
}
|
||||
if (state.kind === "parameter") {
|
||||
if (state.value.length >= MAX_XML_PARAM_VALUE_LENGTH) {
|
||||
config.setState({ kind: "idle" });
|
||||
continue;
|
||||
if (state.value.length < MAX_XML_PARAM_VALUE_LENGTH) {
|
||||
state.value += ch;
|
||||
} else {
|
||||
// Beyond the cap the value stops growing, but we stay in
|
||||
// `parameter` state so the rest of the envelope — including its
|
||||
// close tags — is still swallowed instead of leaking into
|
||||
// visible text. The close handler appends an explicit marker.
|
||||
state.truncated = true;
|
||||
}
|
||||
state.value += ch;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -511,6 +531,21 @@ export class StreamMarkupHealing {
|
||||
return true;
|
||||
}
|
||||
|
||||
#startsWithPartialDsmlSectionOpen(): boolean {
|
||||
const tailLength = this.#buffer.length - this.#offset;
|
||||
if (tailLength === 0 || tailLength >= DSML_SECTION_OPEN_TEMPLATE.length) return false;
|
||||
for (let i = 0; i < tailLength; i++) {
|
||||
const ch = this.#buffer[this.#offset + i]!;
|
||||
const expected = DSML_SECTION_OPEN_TEMPLATE[i]!;
|
||||
if (expected === "|") {
|
||||
if (ch !== "|" && ch !== "｜") return false;
|
||||
} else if (ch !== expected) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
#bufferIsPrefixOf(token: string, remainingLength: number): boolean {
|
||||
for (let i = 0; i < remainingLength; i++) {
|
||||
if (this.#buffer[this.#offset + i] !== token[i]) return false;
|
||||
|
||||
@@ -979,6 +979,23 @@ export function validateToolCall(tools: Tool[], toolCall: ToolCall): ToolCall["a
|
||||
return validateToolArguments(tool, toolCall);
|
||||
}
|
||||
|
||||
/** Cap per-field string lengths when embedding received args in an error message. */
|
||||
const MAX_ERROR_ARG_STRING_LENGTH = 256;
|
||||
|
||||
function truncateArgsForError(value: unknown): unknown {
|
||||
if (typeof value === "string") {
|
||||
if (value.length <= MAX_ERROR_ARG_STRING_LENGTH) return value;
|
||||
return `${value.slice(0, MAX_ERROR_ARG_STRING_LENGTH)}… [truncated ${value.length - MAX_ERROR_ARG_STRING_LENGTH} chars]`;
|
||||
}
|
||||
if (Array.isArray(value)) return value.map(truncateArgsForError);
|
||||
if (value !== null && typeof value === "object") {
|
||||
const out: Record<string, unknown> = {};
|
||||
for (const [key, entry] of Object.entries(value)) out[key] = truncateArgsForError(entry);
|
||||
return out;
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates tool call arguments against the tool's schema (Zod or plain JSON
|
||||
* Schema). Applies LLM-quirk coercions (numeric strings, JSON-string
|
||||
@@ -1025,12 +1042,15 @@ export function validateToolArguments(tool: Tool, toolCall: ToolCall): ToolCall[
|
||||
// existing tests; the detailed body is informational.
|
||||
const errors = result.messages.join("\n") || "Unknown validation error";
|
||||
|
||||
// Truncate long per-field strings: the full payload (potentially hundreds
|
||||
// of KB for write/edit-class calls) would otherwise round-trip back to the
|
||||
// model inside the tool error.
|
||||
const receivedArgs = changed
|
||||
? {
|
||||
original: originalArgs,
|
||||
normalized: normalizedArgs,
|
||||
original: truncateArgsForError(originalArgs),
|
||||
normalized: truncateArgsForError(normalizedArgs),
|
||||
}
|
||||
: originalArgs;
|
||||
: truncateArgsForError(originalArgs);
|
||||
|
||||
const errorMessage = `Validation failed for tool "${
|
||||
toolCall.name
|
||||
|
||||
@@ -1,139 +0,0 @@
|
||||
import { afterEach, describe, expect, it, vi } from "bun:test";
|
||||
import { iterateUntilAbort } from "@oh-my-pi/pi-ai/utils/abortable-iterator";
|
||||
|
||||
function makeSource<T>(handlers: { next: () => Promise<IteratorResult<T>>; onReturn?: () => void }): AsyncIterable<T> {
|
||||
return {
|
||||
[Symbol.asyncIterator](): AsyncIterator<T> {
|
||||
return {
|
||||
next: handlers.next,
|
||||
async return(): Promise<IteratorResult<T>> {
|
||||
handlers.onReturn?.();
|
||||
return { done: true, value: undefined as unknown as T };
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
describe("iterateUntilAbort", () => {
|
||||
it("observes aborts that happen between yielded items and calls iterator.return()", async () => {
|
||||
const controller = new AbortController();
|
||||
let nextCalls = 0;
|
||||
let returnCalled = false;
|
||||
const source = makeSource<number>({
|
||||
next: async () => {
|
||||
nextCalls += 1;
|
||||
if (nextCalls === 1) return { done: false, value: 1 };
|
||||
const { promise } = Promise.withResolvers<IteratorResult<number>>();
|
||||
return promise;
|
||||
},
|
||||
onReturn: () => {
|
||||
returnCalled = true;
|
||||
},
|
||||
});
|
||||
const iterator = iterateUntilAbort(source, controller.signal);
|
||||
|
||||
await expect(iterator.next()).resolves.toEqual({ done: false, value: 1 });
|
||||
controller.abort();
|
||||
await expect(iterator.next()).rejects.toThrow(/abort/i);
|
||||
expect(nextCalls).toBe(1);
|
||||
expect(returnCalled).toBe(true);
|
||||
});
|
||||
|
||||
it("observes aborts that fire DURING an in-flight iterator.next()", async () => {
|
||||
const controller = new AbortController();
|
||||
let returnCalled = false;
|
||||
const source = makeSource<number>({
|
||||
next: async () => {
|
||||
const { promise } = Promise.withResolvers<IteratorResult<number>>();
|
||||
return promise; // never resolves
|
||||
},
|
||||
onReturn: () => {
|
||||
returnCalled = true;
|
||||
},
|
||||
});
|
||||
const iterator = iterateUntilAbort(source, controller.signal);
|
||||
|
||||
const pending = iterator.next();
|
||||
setTimeout(() => controller.abort(new Error("torn down")), 5);
|
||||
|
||||
await expect(pending).rejects.toThrow(/torn down/);
|
||||
expect(returnCalled).toBe(true);
|
||||
});
|
||||
|
||||
it("rejects immediately when the signal is already aborted before the first next()", async () => {
|
||||
const controller = new AbortController();
|
||||
controller.abort(new Error("preflight"));
|
||||
let returnCalled = false;
|
||||
const source = makeSource<number>({
|
||||
next: async () => ({ done: false, value: 1 }),
|
||||
onReturn: () => {
|
||||
returnCalled = true;
|
||||
},
|
||||
});
|
||||
|
||||
const iterator = iterateUntilAbort(source, controller.signal);
|
||||
await expect(iterator.next()).rejects.toThrow(/preflight/);
|
||||
expect(returnCalled).toBe(true);
|
||||
});
|
||||
|
||||
it("yields every item and terminates cleanly when the source completes naturally", async () => {
|
||||
const items = [1, 2, 3];
|
||||
let i = 0;
|
||||
const source = makeSource<number>({
|
||||
next: async () =>
|
||||
i < items.length
|
||||
? { done: false, value: items[i++]! }
|
||||
: { done: true, value: undefined as unknown as number },
|
||||
});
|
||||
|
||||
const collected: number[] = [];
|
||||
for await (const item of iterateUntilAbort(source)) {
|
||||
collected.push(item);
|
||||
}
|
||||
expect(collected).toEqual(items);
|
||||
});
|
||||
|
||||
it("propagates errors from the underlying iterator.next()", async () => {
|
||||
const source = makeSource<number>({
|
||||
next: async () => {
|
||||
throw new Error("upstream blew up");
|
||||
},
|
||||
});
|
||||
|
||||
await expect(async () => {
|
||||
for await (const _ of iterateUntilAbort(source)) {
|
||||
// no body
|
||||
}
|
||||
}).toThrow("upstream blew up");
|
||||
});
|
||||
|
||||
it("does not leak abort listeners across iterations", async () => {
|
||||
const controller = new AbortController();
|
||||
const addSpy = vi.spyOn(controller.signal, "addEventListener");
|
||||
const removeSpy = vi.spyOn(controller.signal, "removeEventListener");
|
||||
|
||||
const items = [1, 2, 3, 4, 5];
|
||||
let i = 0;
|
||||
const source = makeSource<number>({
|
||||
next: async () =>
|
||||
i < items.length
|
||||
? { done: false, value: items[i++]! }
|
||||
: { done: true, value: undefined as unknown as number },
|
||||
});
|
||||
|
||||
for await (const _ of iterateUntilAbort(source, controller.signal)) {
|
||||
// no body
|
||||
}
|
||||
// Every addEventListener("abort", ...) must be paired with a removeEventListener
|
||||
// call (no leaks across iterations).
|
||||
const adds = addSpy.mock.calls.filter(([type]) => type === "abort").length;
|
||||
const removes = removeSpy.mock.calls.filter(([type]) => type === "abort").length;
|
||||
expect(adds).toBe(removes);
|
||||
expect(adds).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
@@ -120,6 +120,57 @@ describe("callWithCopilotModelRetry", () => {
|
||||
expect(calls).toBe(1);
|
||||
});
|
||||
|
||||
it("does not blind-retry a 429 that carries no Retry-After guidance", async () => {
|
||||
let calls = 0;
|
||||
const err = copilotError({ status: 429, message: "rate limited" });
|
||||
await expect(
|
||||
callWithCopilotModelRetry(
|
||||
async () => {
|
||||
calls += 1;
|
||||
throw err;
|
||||
},
|
||||
{ provider: "github-copilot", retryBaseDelayMs: 0 },
|
||||
),
|
||||
).rejects.toBe(err);
|
||||
expect(calls).toBe(1);
|
||||
});
|
||||
|
||||
it("honors Retry-After on a 429 and retries", async () => {
|
||||
let calls = 0;
|
||||
const result = await callWithCopilotModelRetry(
|
||||
async () => {
|
||||
calls += 1;
|
||||
if (calls === 1) {
|
||||
const err = copilotError({ status: 429, message: "rate limited" });
|
||||
(err as unknown as { headers: Record<string, string> }).headers = { "retry-after": "0.01" };
|
||||
throw err;
|
||||
}
|
||||
return "ok" as const;
|
||||
},
|
||||
{ provider: "github-copilot", retryBaseDelayMs: 0 },
|
||||
);
|
||||
expect(result).toBe("ok");
|
||||
expect(calls).toBe(2);
|
||||
});
|
||||
|
||||
it("still retries status-less transport blips with the linear backoff", async () => {
|
||||
let calls = 0;
|
||||
const result = await callWithCopilotModelRetry(
|
||||
async () => {
|
||||
calls += 1;
|
||||
if (calls === 1) {
|
||||
throw new Error(
|
||||
'HTTP2StreamReset fetching "https://api.example.com/x". For more information, pass `verbose: true` in the second argument to fetch()',
|
||||
);
|
||||
}
|
||||
return "ok" as const;
|
||||
},
|
||||
{ provider: "github-copilot", retryBaseDelayMs: 0 },
|
||||
);
|
||||
expect(result).toBe("ok");
|
||||
expect(calls).toBe(2);
|
||||
});
|
||||
|
||||
it("stops retrying when the caller aborts during backoff", async () => {
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
|
||||
@@ -33,4 +33,18 @@ describe("AssistantMessageEventStream", () => {
|
||||
expect(stream.queue[0]).toMatchObject({ type: "text_delta", delta: "a" });
|
||||
expect(stream.queue[1]).toMatchObject({ type: "text_delta", delta: "b" });
|
||||
});
|
||||
|
||||
it("rejects result() when ended without a terminal value", async () => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
stream.end();
|
||||
await expect(stream.result()).rejects.toThrow(/ended without a final result/);
|
||||
});
|
||||
|
||||
it("keeps the pushed terminal result when end() follows a done event", async () => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
const message = createPartial("final");
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
stream.end();
|
||||
await expect(stream.result()).resolves.toBe(message);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -218,6 +218,25 @@ describe("StreamMarkupHealing DSML envelope pattern", () => {
|
||||
expect(calls[0].name).toBe("bash");
|
||||
expect(JSON.parse(calls[0].arguments)).toEqual({ cmd: "ls -la" });
|
||||
});
|
||||
|
||||
it("passes a bare '<' in idle prose through without holding it back", () => {
|
||||
const healing = new StreamMarkupHealing({ pattern: "dsml" });
|
||||
// No '>' anywhere in the tail — the old any-'<' hold-back froze display here.
|
||||
expect(healing.feed("if a < b:\n return a")).toBe("if a < b:\n return a");
|
||||
});
|
||||
|
||||
it("still holds back a tail that is a partial DSML section-open tag", () => {
|
||||
const healing = new StreamMarkupHealing({ pattern: "dsml" });
|
||||
expect(healing.feed("run ")).toBe("run ");
|
||||
expect(healing.feed("<｜DSML｜tool")).toBe("");
|
||||
expect(healing.feed("_calls>")).toBe("");
|
||||
expect(
|
||||
healing.feed(
|
||||
'<｜DSML｜invoke name="bash"><｜DSML｜parameter name="cmd">ls</｜DSML｜parameter></｜DSML｜invoke></｜DSML｜tool_calls>',
|
||||
),
|
||||
).toBe("");
|
||||
expect(healing.drainCompleted()).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("StreamMarkupHealing thinking pattern", () => {
|
||||
|
||||
Reference in New Issue
Block a user