diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index 03d273a51..e3ca37e57 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -432,8 +432,16 @@ export function streamSimple( 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( // 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; diff --git a/packages/ai/src/utils/abort.ts b/packages/ai/src/utils/abort.ts index 212741f6a..54d9f0da5 100644 --- a/packages/ai/src/utils/abort.ts +++ b/packages/ai/src/utils/abort.ts @@ -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(promise: Promise, signal: AbortSignal | undefined): Promise { + if (!signal) return promise; + if (signal.aborted) return Promise.reject(signal.reason ?? new Error("Request was aborted")); + const { promise: aborted, reject } = Promise.withResolvers(); + 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)); +} diff --git a/packages/ai/src/utils/abortable-iterator.ts b/packages/ai/src/utils/abortable-iterator.ts deleted file mode 100644 index ce983a424..000000000 --- a/packages/ai/src/utils/abortable-iterator.ts +++ /dev/null @@ -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(iterable: AsyncIterable, signal?: AbortSignal): AsyncGenerator { - 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>) => - 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 } | { 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" }); - } - } -} diff --git a/packages/ai/src/utils/event-stream.ts b/packages/ai/src/utils/event-stream.ts index 1eff70487..f4819d98f 100644 --- a/packages/ai/src/utils/event-stream.ts +++ b/packages/ai/src/utils/event-stream.ts @@ -5,6 +5,8 @@ export class EventStream implements AsyncIterable { queue: T[] = []; waiting: Array<{ resolve: (value: IteratorResult) => void; reject: (err: unknown) => void }> = []; done = false; + /** True once finalResultPromise has been resolved or rejected. */ + resultSettled = false; #failed = false; #error: unknown = undefined; finalResultPromise: Promise; @@ -30,6 +32,7 @@ export class EventStream implements AsyncIterable { if (this.isComplete(event)) { this.done = true; + this.resultSettled = true; this.resolveFinalResult(this.extractResult(event)); } @@ -54,7 +57,13 @@ export class EventStream implements AsyncIterable { 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 implements AsyncIterable { 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( (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 } - | { 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 } + | { 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" }); } } diff --git a/packages/ai/src/utils/retry-after.ts b/packages/ai/src/utils/retry-after.ts index 86bdac6c8..d226211b6 100644 --- a/packages/ai/src/utils/retry-after.ts +++ b/packages/ai/src/utils/retry-after.ts @@ -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); diff --git a/packages/ai/src/utils/retry.ts b/packages/ai/src/utils/retry.ts index 8a5539186..ed56b519b 100644 --- a/packages/ai/src/utils/retry.ts +++ b/packages/ai/src/utils/retry.ts @@ -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( // 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; diff --git a/packages/ai/src/utils/stream-markup-healing.ts b/packages/ai/src/utils/stream-markup-healing.ts index 7114b9019..3598c9868 100644 --- a/packages/ai/src/utils/stream-markup-healing.ts +++ b/packages/ai/src/utils/stream-markup-healing.ts @@ -36,6 +36,8 @@ const DSML_PARAMETER_OPEN_RE = new RegExp( "y", ); const DSML_PARAMETER_CLOSE_RE = new RegExp(``, "y"); +/** Canonical DSML section-open shape; `|` positions accept either pipe variant. */ +const DSML_SECTION_OPEN_TEMPLATE = "<|DSML|tool_calls>"; const THINK_OPEN = ""; const THINK_CLOSE = ""; @@ -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; diff --git a/packages/ai/src/utils/validation.ts b/packages/ai/src/utils/validation.ts index 506c72978..9ec7e8d93 100644 --- a/packages/ai/src/utils/validation.ts +++ b/packages/ai/src/utils/validation.ts @@ -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 = {}; + 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 diff --git a/packages/ai/test/abortable-iterator.test.ts b/packages/ai/test/abortable-iterator.test.ts deleted file mode 100644 index 4afbeb3b5..000000000 --- a/packages/ai/test/abortable-iterator.test.ts +++ /dev/null @@ -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(handlers: { next: () => Promise>; onReturn?: () => void }): AsyncIterable { - return { - [Symbol.asyncIterator](): AsyncIterator { - return { - next: handlers.next, - async return(): Promise> { - 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({ - next: async () => { - nextCalls += 1; - if (nextCalls === 1) return { done: false, value: 1 }; - const { promise } = Promise.withResolvers>(); - 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({ - next: async () => { - const { promise } = Promise.withResolvers>(); - 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({ - 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({ - 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({ - 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({ - 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(); -}); diff --git a/packages/ai/test/copilot-retry.test.ts b/packages/ai/test/copilot-retry.test.ts index c172483e7..79b79d85f 100644 --- a/packages/ai/test/copilot-retry.test.ts +++ b/packages/ai/test/copilot-retry.test.ts @@ -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 }).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(); diff --git a/packages/ai/test/event-stream.test.ts b/packages/ai/test/event-stream.test.ts index 0d9c97a8b..c26db24a0 100644 --- a/packages/ai/test/event-stream.test.ts +++ b/packages/ai/test/event-stream.test.ts @@ -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); + }); }); diff --git a/packages/ai/test/stream-markup-healing.test.ts b/packages/ai/test/stream-markup-healing.test.ts index 9d48abe82..f11930a92 100644 --- a/packages/ai/test/stream-markup-healing.test.ts +++ b/packages/ai/test/stream-markup-healing.test.ts @@ -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', + ), + ).toBe(""); + expect(healing.drainCompleted()).toHaveLength(1); + }); }); describe("StreamMarkupHealing thinking pattern", () => {