diff --git a/packages/ai/README.md b/packages/ai/README.md index 0b2bdc47937..c8dc9ab98f3 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -773,7 +773,9 @@ Provider notes: The promise client exposes the same surface: `ai.video.start(...)` resolves to a handle with `await`, `events`, `result`, `refresh`, `cancel`, and `token`; `ai.video.generate`, `ai.video.resume(model, token)`, and `ai.video.stream` mirror the Effect API. The handle's `status` and `progress` are a snapshot from when it was -created; `refresh()` resolves to a new handle. +created; `refresh()` resolves to a new handle. Every promise method and stream accepts `{ signal }`: like `fetch`, +aborting rejects the Promise or throws from the `for await` loop with `signal.reason` (an `AbortError` `DOMException` +unless `abort(reason)` passed one), while `break` stops a stream without throwing. ```ts import { ai } from "@opencode/ai/promise" diff --git a/packages/ai/docs/media-design.md b/packages/ai/docs/media-design.md index b3c56a1db21..19ccebc9243 100644 --- a/packages/ai/docs/media-design.md +++ b/packages/ai/docs/media-design.md @@ -402,7 +402,7 @@ for await (const event of ai.llm.stream(request)) { … } await ai.dispose() ``` -Streams become `AsyncIterable` via `Stream.toAsyncIterable`. `AIError` is thrown as-is. `AbortSignal` maps to interruption. Nothing in `src/*` except this entrypoint knows about promises. +Streams become `AsyncIterable` via `Stream.toAsyncIterable`. `AIError` is thrown as-is. Aborting an `AbortSignal` interrupts the work and, like `fetch`, rejects the Promise or throws from the stream with `signal.reason` instead of ending the stream as if complete. Nothing in `src/*` except this entrypoint knows about promises. ### Providers diff --git a/packages/ai/src/promise.ts b/packages/ai/src/promise.ts index 3d6ac256244..d4dd8dddc60 100644 --- a/packages/ai/src/promise.ts +++ b/packages/ai/src/promise.ts @@ -42,7 +42,7 @@ export type GenerationHandle = Snapshot & { /** Serializable JSON; pass it back to `resume` from another process. */ readonly token: unknown readonly await: (options?: AwaitOptions & RunOptions) => Promise - /** Status observations until the first terminal one, polling like `await`; abort ends iteration without throwing. */ + /** Status observations until the first terminal one, polling like `await`; abort throws `signal.reason`. */ readonly events: (options?: AwaitOptions & RunOptions) => AsyncIterable /** The result without polling; fails when the generation has not completed. */ readonly result: (options?: RunOptions) => Promise @@ -50,15 +50,16 @@ export type GenerationHandle = Snapshot & { readonly cancel: (options?: RunOptions) => Promise } +// Fails with `signal.reason` so aborted calls reject and aborted streams throw like `fetch`: an `AbortError` by default. const abortEffect = (signal: AbortSignal | undefined) => signal === undefined ? Effect.never - : Effect.callback((resume) => { + : Effect.callback((resume) => { if (signal.aborted) { - resume(Effect.void) + resume(Effect.fail(signal.reason)) return } - const onAbort = () => resume(Effect.void) + const onAbort = () => resume(Effect.fail(signal.reason)) signal.addEventListener("abort", onAbort, { once: true }) return Effect.sync(() => signal.removeEventListener("abort", onAbort)) }) @@ -68,14 +69,14 @@ export const make = (options: Options = {}) => { /** Run any package Effect (for example `LLMClient.compact(...)`) inside this runtime. */ const run = (effect: Effect.Effect, options?: RunOptions) => - runtime.runPromise(effect, { signal: options?.signal }) + runtime.runPromise(Effect.raceFirst(effect, abortEffect(options?.signal))) const iterate = (stream: Stream.Stream, options?: RunOptions): AsyncIterable => Stream.toAsyncIterable( Stream.unwrap( runtime.contextEffect.pipe( Effect.map( - (context): Stream.Stream => + (context): Stream.Stream => stream.pipe(Stream.interruptWhen(abortEffect(options?.signal)), Stream.provideContext(context)), ), ), diff --git a/packages/ai/test/promise.test.ts b/packages/ai/test/promise.test.ts index f93bf784fa1..cc3ede886f0 100644 --- a/packages/ai/test/promise.test.ts +++ b/packages/ai/test/promise.test.ts @@ -22,7 +22,8 @@ const chatBody = sseEvents( /** * Executor layer that answers chat completions with SSE text, image generations with one base64 PNG, Runway video * tasks with a queued submission that succeeds on the second poll, speech with raw audio or SSE audio deltas, OpenAI - * transcription with JSON or SSE text deltas, and AssemblyAI transcripts that complete on the first poll. + * transcription with JSON or SSE text deltas, AssemblyAI transcripts that complete on the first poll, and `slow.test` + * chat completions that send one text delta and never finish. */ const executor = (seen: Array) => RequestExecutor.layer.pipe( @@ -55,6 +56,18 @@ const executor = (seen: Array) => output: "https://replicate.test/a.webp", urls: { get: "https://replicate.test/p_1", cancel: "https://replicate.test/p_1/cancel" }, }) + if (web.url.startsWith("https://slow.test")) + return input.respond( + new ReadableStream({ + start: (controller) => + controller.enqueue( + new TextEncoder().encode( + `data: ${JSON.stringify({ choices: [{ delta: { content: "Hello" } }] })}\n\n`, + ), + ), + }), + { headers: { "content-type": "text/event-stream" } }, + ) if (web.url.endsWith("/chat/completions")) return input.respond(chatBody, { headers: { "content-type": "text/event-stream" } }) if (web.url.endsWith("/audio/speech")) @@ -304,6 +317,77 @@ describe("AI promise client", () => { await ai.dispose() }) + test("aborted calls reject and aborted streams throw with the signal's reason", async () => { + const ai = AI.make({ layer: executor([]) }) + const slow = OpenAI.configure({ apiKey: "test", baseURL: "https://slow.test/v1" }).chat("gpt-4o-mini") + const aborted = new AbortController() + aborted.abort() + const reason = new Error("mine") + + const rejected = await ai.run(Effect.never, { signal: aborted.signal }).catch((error: unknown) => error) + expect(rejected).toBe(aborted.signal.reason) + expect(rejected).toMatchObject({ name: "AbortError" }) + + const inFlight = new AbortController() + setTimeout(() => inFlight.abort(reason), 10) + expect( + await ai.llm + .generate({ model: slow, prompt: "Hello" }, { signal: inFlight.signal }) + .catch((error: unknown) => error), + ).toBe(reason) + + const preAborted = await Array.fromAsync( + ai.speech.stream({ model: openai.speech("gpt-4o-mini-tts"), text: "Hello" }, { signal: aborted.signal }), + ).catch((error: unknown) => error) + expect(preAborted).toBe(aborted.signal.reason) + expect(preAborted).toMatchObject({ name: "AbortError" }) + + const midStream = new AbortController() + const deltas: Array = [] + const midStreamFailure = await Array.fromAsync( + ai.llm.stream({ model: slow, prompt: "Hello" }, { signal: midStream.signal }), + (event) => { + if (!LLMEvent.is.textDelta(event)) return + deltas.push(event.text) + midStream.abort() + }, + ).catch((error: unknown) => error) + expect(deltas).toEqual(["Hello"]) + expect(midStreamFailure).toBe(midStream.signal.reason) + expect(midStreamFailure).toMatchObject({ name: "AbortError" }) + + const model = Runway.configure({ apiKey: "test", baseURL: "https://runway.test/v1" }).video("gen4.5") + const generation = await ai.video.start({ model, prompt: "A kite" }) + const polling = new AbortController() + const events: Array = [] + const eventsFailure = await Array.fromAsync( + generation.events({ poll: { interval: 60_000 }, signal: polling.signal }), + (event) => { + events.push(event.type) + polling.abort(reason) + }, + ).catch((error: unknown) => error) + expect(events).toEqual(["generation-progress"]) + expect(eventsFailure).toBe(reason) + + await ai.dispose() + }) + + test("breaking out of an abortable stream cleans up without throwing", async () => { + const ai = AI.make({ layer: executor([]) }) + const slow = OpenAI.configure({ apiKey: "test", baseURL: "https://slow.test/v1" }).chat("gpt-4o-mini") + const controller = new AbortController() + const deltas: Array = [] + for await (const event of ai.llm.stream({ model: slow, prompt: "Hello" }, { signal: controller.signal })) { + if (!LLMEvent.is.textDelta(event)) continue + deltas.push(event.text) + break + } + controller.abort() + expect(deltas).toEqual(["Hello"]) + await ai.dispose() + }) + test("the default client is created lazily and can be disposed", async () => { expect(typeof AI.ai.llm.generate).toBe("function") expect(typeof AI.ai.image.generate).toBe("function")