diff --git a/.changeset/client-abort-reaches-request.md b/.changeset/client-abort-reaches-request.md new file mode 100644 index 0000000000..43e466f904 --- /dev/null +++ b/.changeset/client-abort-reaches-request.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': patch +--- + +`stop()` no longer waits for a client tool that never settles. Fetch and XHR connection adapters now abort on the client's signal and on your own `options.signal`. `stream()` and `rpcStream()` end on abort even if the source ignores the signal. `RealtimeClient.disconnect()` during `connect()` or a token refresh no longer leaves a connection or a token refresh timer behind. diff --git a/.changeset/client-cancel-like-stop.md b/.changeset/client-cancel-like-stop.md new file mode 100644 index 0000000000..a02c24fa38 --- /dev/null +++ b/.changeset/client-cancel-like-stop.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': patch +--- + +A cancel now stops the request. `stop()` and an abort of the connection `options.signal` reach the fetch, XHR, `stream()`, and `rpcStream()` requests. A realtime `disconnect()` during `connect()` does not stop a pending `getToken()` or `adapter.connect()` call: it drops the late result and closes the connection that comes back. An abort of `options.signal` now ends the chat like `stop()`, not as a finished run. A realtime connection that a later `connect()` replaced no longer changes the status of the live session. diff --git a/.changeset/core-abort-reaches-provider.md b/.changeset/core-abort-reaches-provider.md new file mode 100644 index 0000000000..a0b066ba8e --- /dev/null +++ b/.changeset/core-abort-reaches-provider.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +A cancel now stops the provider request, not only the caller. `summarize()` (also with `stream: true`, which now passes `abortSignal` and `timeout` to the adapter and reports a cancel to middleware `onAbort`, not `onError`), middleware `ctx.abort()` when the caller passed an `abortController` (adapter request, structured output, and `ctx.signal`), subagent children, `generateVideo` status polls, and aborted `toWebSocket` turns now stop when the run is cancelled. diff --git a/docs/advanced/middleware-advanced.md b/docs/advanced/middleware-advanced.md index c69c4038fe..42ab4db7e4 100644 --- a/docs/advanced/middleware-advanced.md +++ b/docs/advanced/middleware-advanced.md @@ -838,7 +838,7 @@ Each hook gets a `ChatMiddlewareContext` as its first argument. The fields you u | `chunkIndex` | `number` | How many chunks the run sent so far | | `callMetadata` | `Readonly` | The `metadata.tanstack` of the current model call (`source`, `stopReason`, `responseId`, and more). It changes as chunks arrive. Each assistant message of the call gets a copy. | | `messages` | `ReadonlyArray` | The current messages (read-only) | -| `signal` | `AbortSignal \| undefined` | The external abort signal | +| `signal` | `AbortSignal \| undefined` | Aborts when the caller aborts the run or a middleware calls `ctx.abort()` | | `abort(reason?)` | `function` | Abort the run from middleware | | `retry()` | `function` | Call the model again. Works only in [`onModelError`](#onmodelerror). | | `emitCustomEvent(name, value, options?)` | `function` | Push a `CUSTOM` chunk onto the stream now. See [Live custom events](#live-custom-events). | diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 368dd7373f..381b9bdfe0 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -110,6 +110,22 @@ type SchedulerWithYield = { yield?: () => Promise } +/** + * Wait for `promise`, but reject with an AbortError once `signal` aborts. + * A client tool that never settles must not block stop(). + */ +function untilAborted(promise: Promise, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const onAbort = () => + reject(Object.assign(new Error('Aborted'), { name: 'AbortError' })) + if (signal.aborted) return onAbort() + signal.addEventListener('abort', onAbort, { once: true }) + promise.then(resolve, reject).finally(() => { + signal.removeEventListener('abort', onAbort) + }) + }) +} + function yieldToHost(): Promise { const { scheduler } = globalThis as typeof globalThis & { scheduler?: SchedulerWithYield @@ -2884,7 +2900,10 @@ export class ChatClient< // Wait for pending client tool executions if (this.pendingToolExecutions.size > 0) { - await Promise.all(this.pendingToolExecutions.values()) + await untilAborted( + Promise.all(this.pendingToolExecutions.values()), + signal, + ) } // Finalize (idempotent — may already be done by RUN_FINISHED handler). @@ -2899,7 +2918,10 @@ export class ChatClient< this.raiseLocalApprovals() // Wait for the tools that finalizeStream() started, as above. if (this.pendingToolExecutions.size > 0) { - await Promise.all(this.pendingToolExecutions.values()) + await untilAborted( + Promise.all(this.pendingToolExecutions.values()), + signal, + ) } streamCompletedSuccessfully = true } catch (err: unknown) { @@ -2913,6 +2935,11 @@ export class ChatClient< ) runTerminalEventEmitted = true } + // The abort came from the user's own `options.signal`, not from + // stop(). End the same way stop() does. + if (!signal.aborted && generation === this.streamGeneration) { + this.stop() + } return false } if (error instanceof ByokMissingError) { diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 00d7b8551d..3dcc9bb0ea 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -328,6 +328,38 @@ export function createReconnectTracker( } } +/** + * A signal that aborts when either one does. The client's signal (stop()) and + * the user's own `options.signal` must both stop the request. + */ +function eitherSignal( + a: AbortSignal | undefined, + b: AbortSignal | undefined, +): AbortSignal | undefined { + if (!a || !b) return a ?? b + const signal = a.aborted ? a : b.aborted ? b : anySignal(a, b) + requestSignals.set(a, signal) + return signal +} + +/** + * The client's signal -> the signal its request follows. `send()` reads it, + * so an abort of `options.signal` is a stop, not an end or an error. + */ +const requestSignals = new WeakMap() + +function anySignal(a: AbortSignal, b: AbortSignal): AbortSignal { + // `AbortSignal.any` links weakly, so a long-lived `options.signal` keeps + // nothing from a finished request. + if (typeof AbortSignal.any === 'function') return AbortSignal.any([a, b]) + // ponytail: older runtimes keep these listeners on a long-lived signal + // until it aborts. + const controller = new AbortController() + a.addEventListener('abort', () => controller.abort(a.reason), { once: true }) + b.addEventListener('abort', () => controller.abort(b.reason), { once: true }) + return controller.signal +} + /** Resolve after `ms`, or immediately once `signal` aborts. Never rejects. */ function abortableDelay(ms: number, signal?: AbortSignal): Promise { if (ms <= 0 || signal?.aborted) return Promise.resolve() @@ -1278,6 +1310,10 @@ export function normalizeConnectionAdapter( let hasTerminalEvent = false let upstreamThreadId: string | undefined let upstreamRunId: string | undefined + const aborted = () => + abortSignal?.aborted === true || + (abortSignal !== undefined && + requestSignals.get(abortSignal)?.aborted === true) try { const stream = connection.connect( messages, @@ -1303,7 +1339,12 @@ export function normalizeConnectionAdapter( // The event payload may carry an upstream/provider runId when one was // observed, but stamp the caller's request runId so getChunkRunId() // correlates to activeRunIds / currentRunId (same as real stream chunks). - if (!abortSignal?.aborted && !hasTerminalEvent) { + // Only the user's `options.signal` aborted: end like the client's own + // stop(), which rejects with an AbortError. + if (aborted() && !abortSignal?.aborted) { + throw Object.assign(new Error('Aborted'), { name: 'AbortError' }) + } + if (!aborted() && !hasTerminalEvent) { push( withTanstackMetadata( { @@ -1318,7 +1359,7 @@ export function normalizeConnectionAdapter( ) } } catch (err) { - if (!abortSignal?.aborted && !hasTerminalEvent) { + if (!aborted() && !hasTerminalEvent) { const message = err instanceof Error ? err.message : 'Unknown error in connect()' const synthetic: RunErrorEvent = { @@ -1500,7 +1541,7 @@ export function fetchServerSentEvents( // `RequestInit.signal` is typed `AbortSignal | null` (no `undefined` // under `exactOptionalPropertyTypes`), so spread it conditionally // rather than passing `undefined` explicitly. - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) // POST URL is byte-identical to a plain request; the run id (when set) // rides in the X-Run-Id header so durability can key the log by it // without changing the request URL existing clients rely on. @@ -1543,7 +1584,7 @@ export function fetchServerSentEvents( ...mergeHeaders(resolvedOptions.headers), } const fetchClient = resolvedOptions.fetchClient ?? fetch - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) yield* resumableStream( fetchEventSource( @@ -1660,7 +1701,7 @@ export function fetchHttpStream( // `RequestInit.signal` is typed `AbortSignal | null` (no `undefined` // under `exactOptionalPropertyTypes`), so spread it conditionally // rather than passing `undefined` explicitly. - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) // POST URL is byte-identical to a plain request; the run id (when set) // rides in the X-Run-Id header so durability can key the log by it // without changing the request URL existing clients rely on. @@ -1700,7 +1741,7 @@ export function fetchHttpStream( ...mergeHeaders(resolvedOptions.headers), } const fetchClient = resolvedOptions.fetchClient ?? fetch - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) yield* resumableStream( fetchEventSource( @@ -1907,7 +1948,7 @@ export function fetchJson( runId: string, abortSignal: AbortSignal | undefined, ) { - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) return { fetchClient: resolvedOptions.fetchClient ?? fetch, url: resolvedUrl, @@ -2257,7 +2298,7 @@ export function xhrServerSentEvents( async *connect(messages, data, abortSignal, runContext) { const resolvedUrl = typeof url === 'function' ? url() : url const resolvedOptions = await resolveXhrConnectionOptions(options) - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) // POST URL is byte-identical to a plain request; the run id (when set) // rides in the X-Run-Id header so durability can key the log by it // without changing the request URL existing clients rely on. @@ -2279,7 +2320,7 @@ export function xhrServerSentEvents( async *joinRun(runId, abortSignal) { const resolvedUrl = typeof url === 'function' ? url() : url const resolvedOptions = await resolveXhrConnectionOptions(options) - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) const joinUrl = withSearchParams(resolvedUrl, { offset: '-1', runId }) yield* resumableStream( xhrEventSource( @@ -2341,7 +2382,7 @@ export function xhrHttpStream( async *connect(messages, data, abortSignal, runContext) { const resolvedUrl = typeof url === 'function' ? url() : url const resolvedOptions = await resolveXhrConnectionOptions(options) - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) // POST URL is byte-identical to a plain request; the run id (when set) // rides in the X-Run-Id header so durability can key the log by it // without changing the request URL existing clients rely on. @@ -2363,7 +2404,7 @@ export function xhrHttpStream( async *joinRun(runId, abortSignal) { const resolvedUrl = typeof url === 'function' ? url() : url const resolvedOptions = await resolveXhrConnectionOptions(options) - const signal = abortSignal || resolvedOptions.signal + const signal = eitherSignal(abortSignal, resolvedOptions.signal) const joinUrl = withSearchParams(resolvedUrl, { offset: '-1', runId }) yield* resumableStream( xhrEventSource( @@ -2896,12 +2937,13 @@ export function stream( async *connect(messages, data, abortSignal, runContext) { // Pass messages as-is (UIMessages with parts preserved) // Server-side chat() handles conversion to ModelMessages - yield* streamFactory( + const source = streamFactory( messages, data, abortSignal, buildRunAgentInputBody(messages, data, runContext, {}), ) + yield* abortSignal ? abortableIterable(source, abortSignal) : source }, ...(handlers?.hydrate ? { hydrate: handlers.hydrate } : {}), ...(handlers?.hydrateGeneration @@ -2987,7 +3029,11 @@ async function* abortableIterable( yield result.value } } finally { - await iterator.return?.() + const closing = iterator.return?.() + // A source stuck in an await runs return() only after that await ends, + // so do not wait for it once the signal has aborted. + if (signal.aborted) void Promise.resolve(closing).catch(() => {}) + else await closing } } @@ -3034,12 +3080,13 @@ export function rpcStream( async *connect(messages, data, abortSignal, runContext) { // Pass messages as-is (UIMessages with parts preserved) // Server-side chat() handles conversion to ModelMessages - yield* rpcCall( + const source = rpcCall( messages, data, abortSignal, buildRunAgentInputBody(messages, data, runContext, {}), ) + yield* abortSignal ? abortableIterable(source, abortSignal) : source }, ...(handlers?.hydrate ? { hydrate: handlers.hydrate } : {}), ...(handlers?.hydrateGeneration diff --git a/packages/ai-client/src/realtime-client.ts b/packages/ai-client/src/realtime-client.ts index 0bac78f317..9c2d0f7614 100644 --- a/packages/ai-client/src/realtime-client.ts +++ b/packages/ai-client/src/realtime-client.ts @@ -51,6 +51,7 @@ export class RealtimeClient { private connection: RealtimeConnection | null = null private token: RealtimeToken | null = null private tokenRefreshTimeout: ReturnType | null = null + private connectGeneration = 0 private readonly clientTools: Map private readonly stateChangeCallbacks: Set = new Set() @@ -104,10 +105,14 @@ export class RealtimeClient { status: 'connecting', error: null, }) + // disconnect() bumps this. A connect() it overtook must stop. + const generation = ++this.connectGeneration try { // Fetch token from server - this.token = await this.options.getToken() + const token = await this.options.getToken() + if (generation !== this.connectGeneration) return + this.token = token // Schedule token refresh this.scheduleTokenRefresh() @@ -118,10 +123,17 @@ export class RealtimeClient { ? Array.from(this.clientTools.values()) : undefined - this.connection = await this.options.adapter.connect( + const connection = await this.options.adapter.connect( this.token, toolsList, ) + if (generation !== this.connectGeneration) { + // A failed close of a connection that was never used is not an + // error of the client. + await connection.disconnect().catch(() => {}) + return + } + this.connection = connection // A new session starts here. Clear the last transcript before events // flow, so a greeting that arrives while audio capture starts is kept. @@ -140,7 +152,17 @@ export class RealtimeClient { // Start audio capture if configured if (this.options.autoCapture) { - await this.connection.startAudioCapture() + await connection.startAudioCapture() + if (generation !== this.connectGeneration) { + // A later connect() replaced this connection: close it. After + // disconnect(), `this.connection` is null and it is closed already. + if (this.connection && this.connection !== connection) { + // A failed close of the replaced connection is not an + // error of the live session. + await connection.disconnect().catch(() => {}) + } + return + } } this.updateState({ @@ -160,6 +182,7 @@ export class RealtimeClient { * Disconnect from the realtime session. */ async disconnect(): Promise { + this.connectGeneration++ if (this.tokenRefreshTimeout) { clearTimeout(this.tokenRefreshTimeout) this.tokenRefreshTimeout = null @@ -435,6 +458,8 @@ export class RealtimeClient { private scheduleTokenRefresh(): void { if (!this.token) return + // A connect() that replaced an earlier one must not leave its timer. + if (this.tokenRefreshTimeout) clearTimeout(this.tokenRefreshTimeout) const timeUntilExpiry = this.token.expiresAt - Date.now() const refreshIn = Math.max(0, timeUntilExpiry - TOKEN_REFRESH_BUFFER_MS) @@ -445,8 +470,12 @@ export class RealtimeClient { } private async refreshToken(): Promise { + const generation = this.connectGeneration try { - this.token = await this.options.getToken() + const token = await this.options.getToken() + // disconnect() ran during the refresh: do not re-arm the timer. + if (generation !== this.connectGeneration) return + this.token = token this.scheduleTokenRefresh() this.connection?.updateToken?.(this.token) // Note: Some providers may require reconnection with new token @@ -459,18 +488,25 @@ export class RealtimeClient { } private subscribeToConnectionEvents(): void { - if (!this.connection) return + const connection = this.connection + if (!connection) return + // A later connect() can replace this connection. Its late events (such as + // `idle` from its disconnect) must not change the live session. + const on: RealtimeConnection['on'] = (event, handler) => + connection.on(event, (payload) => { + if (this.connection === connection) handler(payload) + }) // Status changes this.unsubscribers.push( - this.connection.on('status_change', ({ status }) => { + on('status_change', ({ status }) => { this.updateState({ status }) }), ) // Mode changes this.unsubscribers.push( - this.connection.on('mode_change', ({ mode }) => { + on('mode_change', ({ mode }) => { this.updateState({ mode }) }), ) @@ -479,7 +515,7 @@ export class RealtimeClient { // User transcripts are added as messages when final (no separate message_complete for user input) // Assistant transcripts are streamed, final message comes via message_complete this.unsubscribers.push( - this.connection.on('transcript', ({ role, transcript, isFinal }) => { + on('transcript', ({ role, transcript, isFinal }) => { if (role === 'user') { this.updateState({ pendingUserTranscript: isFinal ? null : transcript, @@ -504,7 +540,7 @@ export class RealtimeClient { // Tool calls this.unsubscribers.push( - this.connection.on('tool_call', ({ toolCallId, toolName, input }) => { + on('tool_call', ({ toolCallId, toolName, input }) => { if (!toolCallId) { console.error( '[RealtimeClient] tool_call missing toolCallId', @@ -544,7 +580,7 @@ export class RealtimeClient { // Message complete this.unsubscribers.push( - this.connection.on('message_complete', ({ message }) => { + on('message_complete', ({ message }) => { // Replace pending message with final version if needed const existingIndex = this.state.messages.findIndex( (m) => m.id === message.id, @@ -561,7 +597,7 @@ export class RealtimeClient { // Interruption this.unsubscribers.push( - this.connection.on('interrupted', ({ messageId }) => { + on('interrupted', ({ messageId }) => { if (messageId) { const newMessages = this.state.messages.map((m) => m.id === messageId ? { ...m, interrupted: true } : m, @@ -578,20 +614,20 @@ export class RealtimeClient { // Errors this.unsubscribers.push( - this.connection.on('error', ({ error }) => { + on('error', ({ error }) => { this.updateState({ error }) this.options.onError?.(error) }), ) this.unsubscribers.push( - this.connection.on('go_away', ({ timeLeft }) => { + on('go_away', ({ timeLeft }) => { this.options.onGoAway?.(timeLeft) }), ) this.unsubscribers.push( - this.connection.on('usage', (usage) => { + on('usage', (usage) => { this.options.onUsage?.(usage) }), ) diff --git a/packages/ai-client/tests/abort-signal-reach.test.ts b/packages/ai-client/tests/abort-signal-reach.test.ts new file mode 100644 index 0000000000..6700f012e2 --- /dev/null +++ b/packages/ai-client/tests/abort-signal-reach.test.ts @@ -0,0 +1,560 @@ +/** + * A client cancel (`stop()`, the user's own `options.signal`, `disconnect()`) + * must reach the work it should stop. The fake requests below hang until + * their signal aborts, like a real fetch that keeps the run going. + */ +import { getEventListeners } from 'node:events' +import { describe, expect, it, vi } from 'vitest' +import { ChatClient } from '../src/chat-client' +import { RealtimeClient } from '../src/realtime-client' +import { + fetchHttpStream, + fetchJson, + fetchServerSentEvents, + rpcStream, + stream, + xhrHttpStream, + xhrServerSentEvents, +} from '../src/connection-adapters' +import { createToolCallChunks } from './test-utils' +import type { ConnectConnectionAdapter } from '../src/connection-adapters' +import type { StreamChunk } from '@tanstack/ai/client' + +const RUN_CONTEXT = { threadId: 'thread-1', runId: 'run-1' } + +/** 'done' when `work` settles within `ms`, else 'hung'. */ +function settlesWithin(work: Promise, ms = 500) { + return Promise.race([ + work.then( + () => 'done', + () => 'done', + ), + new Promise((resolve) => setTimeout(() => resolve('hung'), ms)), + ]) +} + +async function drain(iterable: AsyncIterable) { + for await (const _ of iterable) { + // drain + } +} + +/** A fetch that hangs until its signal aborts (rejects at once without one). */ +function hangingFetch() { + const seen: { signal?: AbortSignal } = {} + const fetchClient = ((_url: RequestInfo | URL, init?: RequestInit) => { + const signal = init?.signal ?? undefined + seen.signal = signal + if (!signal) { + return Promise.reject(new Error('request sent without an abort signal')) + } + return new Promise((_, reject) => { + const fail = () => + reject(Object.assign(new Error('aborted'), { name: 'AbortError' })) + if (signal.aborted) fail() + else signal.addEventListener('abort', fail, { once: true }) + }) + }) as typeof fetch + return { fetchClient, aborted: () => seen.signal?.aborted === true } +} + +/** An XHR whose response never comes; `abort()` ends it. */ +function hangingXhr() { + const xhr = { + status: 0, + statusText: '', + responseText: '', + withCredentials: false, + onprogress: null as null | ((ev: unknown) => void), + onload: null as null | ((ev: unknown) => void), + onerror: null as null | ((ev: unknown) => void), + onabort: null as null | ((ev: unknown) => void), + onloadend: null as null | ((ev: unknown) => void), + open: () => {}, + send: () => {}, + setRequestHeader: () => {}, + getResponseHeader: () => null, + abort: vi.fn(() => { + xhr.onabort?.({}) + xhr.onloadend?.({}) + }), + } + return { + xhrFactory: () => xhr as unknown as XMLHttpRequest, + aborted: () => xhr.abort.mock.calls.length > 0, + } +} + +describe("connection adapters follow the user's options.signal (L2)", () => { + const cases = [ + ['fetchServerSentEvents', fetchServerSentEvents], + ['fetchHttpStream', fetchHttpStream], + ['fetchJson', fetchJson], + ] as const + for (const [name, factory] of cases) { + for (const method of ['connect', 'joinRun'] as const) { + it(`${name}.${method}`, async () => { + const user = new AbortController() + const client = new AbortController() + const fake = hangingFetch() + const adapter = factory('/api/chat', { + fetchClient: fake.fetchClient, + signal: user.signal, + }) + const run = drain( + method === 'connect' + ? adapter.connect([], undefined, client.signal, RUN_CONTEXT) + : adapter.joinRun!('run-1', client.signal), + ) + setTimeout(() => user.abort(), 10) + expect(await settlesWithin(run)).toBe('done') + expect(fake.aborted()).toBe(true) + }) + } + } + + const xhrCases = [ + ['xhrServerSentEvents', xhrServerSentEvents], + ['xhrHttpStream', xhrHttpStream], + ] as const + for (const [name, factory] of xhrCases) { + for (const method of ['connect', 'joinRun'] as const) { + it(`${name}.${method}`, async () => { + const user = new AbortController() + const client = new AbortController() + const fake = hangingXhr() + const adapter = factory('/api/chat', { + xhrFactory: fake.xhrFactory, + signal: user.signal, + }) + const run = drain( + method === 'connect' + ? adapter.connect([], undefined, client.signal, RUN_CONTEXT) + : adapter.joinRun!('run-1', client.signal), + ) + setTimeout(() => user.abort(), 10) + expect(await settlesWithin(run)).toBe('done') + expect(fake.aborted()).toBe(true) + }) + } + } + + it('leaves no listener on a reused options.signal after each request', async () => { + const user = new AbortController() + const fetchClient = (() => + Promise.resolve( + new Response('', { headers: { 'content-type': 'text/event-stream' } }), + )) as typeof fetch + const adapter = fetchServerSentEvents('/api/chat', { + fetchClient, + signal: user.signal, + }) + for (let run = 0; run < 3; run++) { + const client = new AbortController() + await drain(adapter.connect([], undefined, client.signal, RUN_CONTEXT)) + } + expect(getEventListeners(user.signal, 'abort')).toHaveLength(0) + }) +}) + +describe('ChatClient ends a user abort of options.signal like stop()', () => { + /** An SSE reply that sends some text, then waits until its signal aborts. */ + const partialFetch = (() => (_url: RequestInfo | URL, init?: RequestInit) => { + const encoder = new TextEncoder() + const body = new ReadableStream({ + start(controller) { + for (const event of [ + { type: 'RUN_STARTED', ...RUN_CONTEXT, timestamp: 1 }, + { + type: 'TEXT_MESSAGE_START', + messageId: 'm1', + role: 'assistant', + timestamp: 1, + }, + { + type: 'TEXT_MESSAGE_CONTENT', + messageId: 'm1', + delta: 'Hel', + timestamp: 1, + }, + ]) { + controller.enqueue( + encoder.encode(`data: ${JSON.stringify(event)} + +`), + ) + } + init?.signal?.addEventListener('abort', () => + controller.error( + Object.assign(new Error('aborted'), { name: 'AbortError' }), + ), + ) + }, + }) + return Promise.resolve( + new Response(body, { + headers: { 'content-type': 'text/event-stream' }, + }), + ) + }) as unknown as () => typeof fetch + + it.each(['options.signal', 'stop()'] as const)( + 'cancel with %s', + async (cancel) => { + const user = new AbortController() + const onError = vi.fn() + const onFinish = vi.fn() + const client = new ChatClient({ + threadId: 'thread-1', + connection: fetchServerSentEvents('/api/chat', { + fetchClient: partialFetch(), + signal: user.signal, + }), + onError, + onFinish, + }) + const send = client.sendMessage('Go') + await vi.waitFor(() => + expect(JSON.stringify(client.getMessages())).toContain('Hel'), + ) + if (cancel === 'stop()') client.stop() + else user.abort() + + expect(await settlesWithin(send)).toBe('done') + expect(client.getStatus()).toBe('ready') + expect(client.getIsLoading()).toBe(false) + expect(onError).not.toHaveBeenCalled() + expect(onFinish).not.toHaveBeenCalled() + }, + ) +}) + +describe('stream() and rpcStream() stop on abort (L3)', () => { + // A source that ignores its signal and never ends. + const ignoresSignal = () => + (async function* (): AsyncGenerator { + await new Promise(() => {}) + })() + + for (const [name, adapter] of [ + ['stream', stream(ignoresSignal)], + ['rpcStream', rpcStream(ignoresSignal)], + ] as const) { + it(name, async () => { + const client = new AbortController() + const run = drain( + adapter.connect([], undefined, client.signal, RUN_CONTEXT), + ) + setTimeout(() => client.abort(), 10) + expect(await settlesWithin(run)).toBe('done') + }) + } +}) + +describe('ChatClient.stop() (L1)', () => { + // 'stop' starts the tool while chunks stream. 'tool_calls' starts it in + // finalizeStream(), after the stream ends. + it.each(['stop', 'tool_calls'] as const)( + 'does not wait for a client tool that never ends (finishReason %s)', + async (finishReason) => { + let toolStarted = false + const chunks = createToolCallChunks([ + { id: 'tc-1', name: 'stuck', arguments: '{}' }, + ]) + chunks[chunks.length - 1] = { + ...chunks[chunks.length - 1]!, + finishReason, + } as StreamChunk + const connection: ConnectConnectionAdapter = { + async *connect(_messages, _data, abortSignal) { + for (const chunk of chunks) { + if (abortSignal?.aborted) return + yield chunk + } + }, + } + const client = new ChatClient({ + threadId: 'thread-1', + connection, + tools: [ + { + __toolSide: 'client' as const, + name: 'stuck', + description: 'Never finishes', + execute: () => { + toolStarted = true + return new Promise(() => {}) + }, + }, + ], + }) + const send = client.sendMessage('Go') + await vi.waitFor(() => expect(toolStarted).toBe(true)) + client.stop() + expect(await settlesWithin(send)).toBe('done') + }, + ) +}) + +describe('RealtimeClient.disconnect() during connect() (L4)', () => { + // `gate` is the connect() step that waits while disconnect() runs. + for (const gate of ['getToken', 'adapter.connect', 'startAudioCapture']) { + it(`stays disconnected when disconnect() runs during ${gate}`, async () => { + let release!: () => void + const held = new Promise((resolve) => { + release = resolve + }) + const wait = (step: string) => (step === gate ? held : undefined) + const connection = { + on: () => () => {}, + updateSession: () => {}, + disconnect: vi.fn(async () => {}), + startAudioCapture: async () => { + await wait('startAudioCapture') + }, + } + const adapterConnect = vi.fn(async () => { + await wait('adapter.connect') + return connection + }) + const client = new RealtimeClient({ + getToken: async () => { + await wait('getToken') + return { token: 't', expiresAt: Date.now() + 60_000 } as any + }, + adapter: { connect: adapterConnect } as any, + }) + + const connecting = client.connect() + await vi.waitFor(() => expect(client.status).toBe('connecting')) + await new Promise((resolve) => setTimeout(resolve, 0)) + await client.disconnect() + release() + await connecting.catch(() => {}) + + const opened = adapterConnect.mock.calls.length > 0 + // Either it never connects, or it closes what it opened. + expect(!opened || connection.disconnect.mock.calls.length > 0).toBe(true) + expect(client.status).toBe('idle') + // No token refresh is armed for a session that does not exist. + expect(Reflect.get(client, 'tokenRefreshTimeout')).toBeNull() + }) + } + + it('stays idle when the late connection fails to close', async () => { + let release!: () => void + const held = new Promise((resolve) => { + release = resolve + }) + const connection = { + on: () => () => {}, + updateSession: () => {}, + disconnect: vi.fn(async () => { + throw new Error('close failed') + }), + startAudioCapture: async () => {}, + } + const onError = vi.fn() + const client = new RealtimeClient({ + getToken: async () => + ({ token: 't', expiresAt: Date.now() + 60_000 }) as any, + adapter: { + connect: async () => { + await held + return connection + }, + } as any, + onError, + }) + + const connecting = client.connect() + await vi.waitFor(() => expect(client.status).toBe('connecting')) + await new Promise((resolve) => setTimeout(resolve, 0)) + await client.disconnect() + release() + await connecting + + expect(connection.disconnect).toHaveBeenCalled() + expect(client.status).toBe('idle') + expect(client.error).toBeNull() + expect(onError).not.toHaveBeenCalled() + }) + + it('closes a connection that a second connect() replaced during audio startup', async () => { + let release!: () => void + const held = new Promise((resolve) => { + release = resolve + }) + const makeConnection = (startAudioCapture: () => Promise) => ({ + on: () => () => {}, + updateSession: () => {}, + disconnect: vi.fn(async () => {}), + startAudioCapture, + }) + const first = makeConnection(() => held) + const second = makeConnection(async () => {}) + const adapterConnect = vi + .fn() + .mockResolvedValueOnce(first) + .mockResolvedValueOnce(second) + const client = new RealtimeClient({ + getToken: async () => + ({ token: 't', expiresAt: Date.now() + 60_000 }) as any, + adapter: { connect: adapterConnect } as any, + autoCapture: true, + }) + + const firstConnect = client.connect() + await vi.waitFor(() => expect(adapterConnect).toHaveBeenCalledTimes(1)) + await new Promise((resolve) => setTimeout(resolve, 0)) + await client.connect() + release() + await firstConnect + + expect(first.disconnect).toHaveBeenCalled() + expect(second.disconnect).not.toHaveBeenCalled() + await client.disconnect() + }) +}) + +describe('RealtimeClient keeps the live session when a second connect() wins', () => { + // The first connection emits `idle` when it closes (like the OpenAI + // adapter), and its close can fail. + it.each(['emits idle', 'throws'] as const)( + 'the replaced connection %s on close', + async (onClose) => { + let release!: () => void + const held = new Promise((resolve) => { + release = resolve + }) + const makeConnection = (startAudioCapture: () => Promise) => { + const handlers = new Map void>>() + const connection = { + on: (event: string, handler: (payload: any) => void) => { + const set = handlers.get(event) ?? new Set() + set.add(handler) + handlers.set(event, set) + return () => set.delete(handler) + }, + emit: (event: string, payload: unknown) => + handlers.get(event)?.forEach((handler) => handler(payload)), + updateSession: () => {}, + disconnect: vi.fn(async () => {}), + startAudioCapture, + } + return connection + } + const first = makeConnection(() => held) + first.disconnect.mockImplementation(async () => { + if (onClose === 'throws') throw new Error('close failed') + first.emit('status_change', { status: 'idle' }) + }) + const second = makeConnection(async () => {}) + const onError = vi.fn() + const client = new RealtimeClient({ + getToken: async () => + ({ token: 't', expiresAt: Date.now() + 60_000 }) as any, + adapter: { + connect: vi + .fn() + .mockResolvedValueOnce(first) + .mockResolvedValueOnce(second), + } as any, + autoCapture: true, + onError, + }) + + const firstConnect = client.connect() + await vi.waitFor(() => + expect(Reflect.get(client, 'connection')).toBe(first), + ) + await client.connect() + release() + await firstConnect + + expect(first.disconnect).toHaveBeenCalled() + expect(client.status).toBe('connected') + expect(client.error).toBeNull() + expect(onError).not.toHaveBeenCalled() + await client.disconnect() + }, + ) +}) + +describe('RealtimeClient.disconnect() during a token refresh', () => { + it('does not re-arm the refresh timer', async () => { + let release!: () => void + const held = new Promise((resolve) => { + release = resolve + }) + // First token expires now, so the refresh starts at once and then waits. + const getToken = vi + .fn() + .mockResolvedValueOnce({ token: 't1', expiresAt: Date.now() }) + .mockImplementationOnce(async () => { + await held + return { token: 't2', expiresAt: Date.now() } + }) + const client = new RealtimeClient({ + getToken, + adapter: { + connect: async () => ({ + on: () => () => {}, + updateSession: () => {}, + disconnect: async () => {}, + startAudioCapture: async () => {}, + }), + } as any, + }) + + await client.connect() + await vi.waitFor(() => expect(getToken).toHaveBeenCalledTimes(2)) + await client.disconnect() + release() + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(Reflect.get(client, 'tokenRefreshTimeout')).toBeNull() + }) + + it('leaves no refresh timer of a connect() that a later connect() replaced', async () => { + vi.useFakeTimers() + try { + let release!: () => void + const held = new Promise((resolve) => { + release = resolve + }) + const connection = { + on: () => () => {}, + updateSession: () => {}, + disconnect: async () => {}, + startAudioCapture: async () => {}, + } + const getToken = vi.fn( + async () => ({ token: 't', expiresAt: Date.now() + 120_000 }) as any, + ) + const client = new RealtimeClient({ + getToken, + adapter: { + connect: vi + .fn() + .mockImplementationOnce(async () => { + await held + return connection + }) + .mockResolvedValueOnce(connection), + } as any, + }) + + const first = client.connect() + await vi.waitFor(() => expect(getToken).toHaveBeenCalledTimes(1)) + await client.connect() + release() + await first + await client.disconnect() + + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }) +}) diff --git a/packages/ai-client/tests/connection-adapters-abort.test.ts b/packages/ai-client/tests/connection-adapters-abort.test.ts index 766a6a38f9..85cb621165 100644 --- a/packages/ai-client/tests/connection-adapters-abort.test.ts +++ b/packages/ai-client/tests/connection-adapters-abort.test.ts @@ -58,9 +58,11 @@ describe('Connection Adapters - Abort Signal Handling', () => { expect(fetchCall?.[1]?.signal).toBe(abortSignal) }) - it('should use provided abortSignal over options.signal', async () => { - const providedSignal = new AbortController().signal - const optionsSignal = new AbortController().signal + it('should combine provided abortSignal with options.signal', async () => { + const provided = new AbortController() + const providedSignal = provided.signal + const optionsController = new AbortController() + const optionsSignal = optionsController.signal const mockResponse = { ok: true, @@ -88,8 +90,10 @@ describe('Connection Adapters - Abort Signal Handling', () => { // Consume all chunks } - const fetchCall = fetchMock.mock.calls[0] - expect(fetchCall?.[1]?.signal).toBe(providedSignal) + const fetchSignal: AbortSignal = fetchMock.mock.calls[0]?.[1]?.signal + expect(fetchSignal.aborted).toBe(false) + optionsController.abort() + expect(fetchSignal.aborted).toBe(true) }) it('should stop reading stream when aborted', async () => { diff --git a/packages/ai-gemini/tests/gemini-adapter.test.ts b/packages/ai-gemini/tests/gemini-adapter.test.ts index d857fb4621..7126e9f987 100644 --- a/packages/ai-gemini/tests/gemini-adapter.test.ts +++ b/packages/ai-gemini/tests/gemini-adapter.test.ts @@ -241,7 +241,12 @@ describe('GeminiAdapter through AI', () => { expect(mocks.generateContentStreamSpy).toHaveBeenCalledTimes(1) const [payload] = mocks.generateContentStreamSpy.mock.calls[0]! - expect(payload.config.abortSignal).toBe(abortController.signal) + // chat() links the caller signal with middleware ctx.abort(), so the SDK + // gets a linked signal that aborts with the caller. + const sdkSignal: AbortSignal = payload.config.abortSignal + expect(sdkSignal.aborted).toBe(false) + abortController.abort() + expect(sdkSignal.aborted).toBe(true) }) it('omits config.abortSignal when no abort controller is supplied', async () => { diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index 6ed92991d2..a103544ff9 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -1355,6 +1355,7 @@ export function createSyntheticSubagentTools( const toolContext = context as | { toolCallId?: string + abortSignal?: AbortSignal [EMIT_STREAM_CHUNK]?: (chunk: StreamChunk) => void } | undefined @@ -1419,7 +1420,8 @@ export function createSyntheticSubagentTools( } active += 1 const sink = createSubagentSink() - const link = linkAbort(parent.abortSignal) + // The tool context signal also fires on a middleware ctx.abort(). + const link = linkAbort(toolContext?.abortSignal ?? parent.abortSignal) let timedOut = false const timer = timeout === undefined diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index c41beeab47..379235ed4f 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1061,6 +1061,8 @@ class TextEngine< private readonly streamId: string private readonly effectiveRequest?: Request | RequestInit private readonly effectiveSignal?: AbortSignal + /** The caller signal plus ctx.abort(), only when there is a caller signal. */ + private readonly linkedSignal?: AbortSignal private messages: Array private activities: Array = [] @@ -1275,9 +1277,6 @@ class TextEngine< >(this.tools) this.requestId = this.createId('chat') this.streamId = this.createId('stream') - this.effectiveRequest = config.params.abortController - ? { signal: config.params.abortController.signal } - : undefined this.effectiveSignal = config.params.abortController?.signal this.threadId = config.params.threadId this.runIdOverride = config.params.runId @@ -1296,6 +1295,13 @@ class TextEngine< this.effectiveSignal, this.middlewareAbortController.signal, ) + // With a caller signal, the provider request and ctx.signal follow + // ctx.abort() too. Without one they get no signal, as before: some SDKs + // drop their own timeout when a request has a signal. + this.linkedSignal = this.effectiveSignal ? this.toolAbortSignal : undefined + this.effectiveRequest = this.linkedSignal + ? { signal: this.linkedSignal } + : undefined this.middlewareCtx = { requestId: this.requestId, streamId: this.streamId, @@ -1312,7 +1318,7 @@ class TextEngine< iteration: 0, chunkIndex: 0, callMetadata: this.currentCallMetadata, - signal: this.effectiveSignal, + signal: this.linkedSignal, abort: (reason?: string) => { this.abortReason = reason this.middlewareAbortController?.abort(reason) diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 76e78c24ac..f31a3dc91e 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -234,7 +234,7 @@ export interface ChatMiddlewareContext { * message of the call gets a copy of it. */ callMetadata: Readonly - /** Abort signal from the chat request */ + /** Aborts when the caller aborts the chat request or a middleware calls `abort()`. */ signal?: AbortSignal /** Abort the chat run with a reason */ abort: (reason?: string) => void diff --git a/packages/ai/src/activities/generateVideo/index.ts b/packages/ai/src/activities/generateVideo/index.ts index 44ec2c93b1..3cfd2c3612 100644 --- a/packages/ai/src/activities/generateVideo/index.ts +++ b/packages/ai/src/activities/generateVideo/index.ts @@ -683,7 +683,10 @@ async function* runStreamingVideoGeneration< while (Date.now() - startTime < maxDuration) { await sleep(pollingInterval, abortControls.signal) - const statusResult = await adapter.getVideoStatus(jobResult.jobId) + const statusResult = await raceWithAbort( + adapter.getVideoStatus(jobResult.jobId), + abortControls.signal, + ) yield { type: 'CUSTOM', @@ -698,7 +701,22 @@ async function* runStreamingVideoGeneration< } if (statusResult.status === 'completed') { - const video = await adapter.getVideo(jobResult.jobId) + // The caller can abort when it reads the status above. + abortControls.signal?.throwIfAborted() + const download = adapter.getVideo(jobResult.jobId) + const video = await raceWithAbort(download, abortControls.signal).catch( + (error: unknown) => { + // Nobody reads a video stream that comes after the abort, so + // cancel its body to close the download. + download.then( + (late) => { + if ('body' in late) late.body?.cancel().catch(() => undefined) + }, + () => undefined, + ) + throw error + }, + ) logger.output( `activity=generateVideo jobId=${jobResult.jobId} status=completed`, diff --git a/packages/ai/src/activities/summarize/chat-stream-summarize.ts b/packages/ai/src/activities/summarize/chat-stream-summarize.ts index 3e7f2fa967..e958f2469c 100644 --- a/packages/ai/src/activities/summarize/chat-stream-summarize.ts +++ b/packages/ai/src/activities/summarize/chat-stream-summarize.ts @@ -407,6 +407,11 @@ export class ChatStreamSummarizeAdapter< // `runId` absent when unset, under `exactOptionalPropertyTypes`. ...(options.runId !== undefined ? { runId: options.runId } : {}), threadId: options.threadId, + // Without this, a cancel or timeout rejects summarize() but the wrapped + // provider request keeps running. + ...(options.abortSignal + ? { request: { signal: options.abortSignal } } + : {}), } } diff --git a/packages/ai/src/activities/summarize/index.ts b/packages/ai/src/activities/summarize/index.ts index fa83c57f0d..09c28ab810 100644 --- a/packages/ai/src/activities/summarize/index.ts +++ b/packages/ai/src/activities/summarize/index.ts @@ -404,11 +404,26 @@ async function* runStreamingSummarize( // Use real streaming if the adapter supports it if (adapter.summarizeStream) { - yield* runNativeSummarizeStream( - options, - adapter.summarizeStream(summarizeOptions), - logger, - ) + // Same abort/timeout composition as the non-streaming path. + const abortControls = createActivityAbortControls({ + timeout: options.timeout, + abortSignal: options.abortSignal, + }) + try { + yield* runNativeSummarizeStream( + options, + adapter.summarizeStream({ + ...summarizeOptions, + ...(abortControls.signal + ? { abortSignal: abortControls.signal } + : {}), + }), + logger, + abortControls.signal, + ) + } finally { + abortControls.clear() + } return } @@ -447,6 +462,7 @@ async function* runNativeSummarizeStream( options: SummarizeActivityOptions, true>, stream: AsyncIterable, logger: InternalLogger, + signal: AbortSignal | undefined, ): AsyncIterable { const { adapter, middleware, modelOptions } = options const mwCtx = createGenerationContext({ @@ -488,17 +504,31 @@ async function* runNativeSummarizeStream( yield chunk } if (!settled) { - await runGenerationFinish(middleware, mwCtx, { - duration: Date.now() - startTime, - }) settled = true + // An adapter can end an aborted request with a RUN_ERROR chunk instead + // of a throw. The consumer got that chunk; middleware gets onAbort. + if (signal?.aborted) { + await runGenerationAbort(middleware, mwCtx, { + reason: abortReasonMessage(undefined, signal), + duration: Date.now() - startTime, + }) + } else { + await runGenerationFinish(middleware, mwCtx, { + duration: Date.now() - startTime, + }) + } } } catch (error) { settled = true - await runGenerationError(middleware, mwCtx, { - error, - duration: Date.now() - startTime, - }) + const duration = Date.now() - startTime + if (isActivityAbortError(error, signal)) { + await runGenerationAbort(middleware, mwCtx, { + reason: abortReasonMessage(error, signal), + duration, + }) + } else { + await runGenerationError(middleware, mwCtx, { error, duration }) + } logger.errors('summarize activity failed', { error, source: 'summarize', diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 290bac4afa..399fa44156 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -272,6 +272,8 @@ export function toWebSocketStream( } } else { for await (const chunk of init.onRun(ctx)) { + // Ends the onRun stream (iterator.return()) even if it ignores ctx.signal. + if (turnAbort.signal.aborted) break socket.send(encodeWsFrame(chunk, undefined)) } } diff --git a/packages/ai/tests/abort-signal-reach.test.ts b/packages/ai/tests/abort-signal-reach.test.ts new file mode 100644 index 0000000000..558cfc4b17 --- /dev/null +++ b/packages/ai/tests/abort-signal-reach.test.ts @@ -0,0 +1,549 @@ +/** + * A cancel (caller signal, timeout, middleware `ctx.abort()`) must reach the + * provider request, not only the caller. Each fake below rejects at once when + * it gets no signal, and otherwise hangs until its signal aborts — like a real + * fetch that keeps running (and billing) until something aborts it. + */ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { chat } from '../src/activities/chat/index' +import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { summarize } from '../src/activities/summarize/index' +import { ChatStreamSummarizeAdapter } from '../src/activities/summarize/chat-stream-summarize' +import { generateVideo } from '../src/activities/generateVideo/index' +import { toWebSocketStream } from '../src/stream-to-websocket' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import type { ChatStreamCapable } from '../src/activities/summarize/chat-stream-summarize' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { WebSocketLike } from '../src/stream-to-websocket' +import type { StreamChunk } from '../src/types' + +/** A provider request: no signal fails at once, else hangs until abort. */ +function hangUntilAborted(signal: AbortSignal | undefined): Promise { + if (!signal) { + return Promise.reject(new Error('request sent without an abort signal')) + } + return new Promise((_, reject) => { + const fail = () => + reject(Object.assign(new Error('aborted'), { name: 'AbortError' })) + if (signal.aborted) fail() + else signal.addEventListener('abort', fail, { once: true }) + }) +} + +/** 'done' when `work` settles within `ms`, else 'hung'. */ +function settlesWithin(work: Promise, ms = 500) { + return Promise.race([ + work.then( + () => 'done', + () => 'done', + ), + new Promise((resolve) => setTimeout(() => resolve('hung'), ms)), + ]) +} + +function summarizeAdapter(seen: { signal?: AbortSignal }) { + const text: ChatStreamCapable = { + chatStream: (options) => + (async function* () { + seen.signal = options.request?.signal ?? undefined + yield* [] + await hangUntilAborted(seen.signal) + })(), + } + return new ChatStreamSummarizeAdapter( + text, + 'test-model', + 'openai', + ) +} + +describe('summarize() abort reaches the wrapped chatStream', () => { + it('non-stream: caller abort aborts the request (C1)', async () => { + const seen: { signal?: AbortSignal } = {} + const controller = new AbortController() + const run = summarize({ + adapter: summarizeAdapter(seen), + threadId: 'thread-1', + text: 'long text', + abortSignal: controller.signal, + }) + setTimeout(() => controller.abort(), 10) + await expect(run).rejects.toThrow() + expect(seen.signal?.aborted).toBe(true) + }) + + it('stream: caller abort aborts the request (C2)', async () => { + const seen: { signal?: AbortSignal } = {} + const controller = new AbortController() + const stream = summarize({ + adapter: summarizeAdapter(seen), + threadId: 'thread-1', + text: 'long text', + stream: true, + abortSignal: controller.signal, + }) + setTimeout(() => controller.abort(), 10) + await expect(collectChunks(stream)).rejects.toThrow() + expect(seen.signal?.aborted).toBe(true) + }) + + it('stream: timeout aborts the request (C2)', async () => { + const seen: { signal?: AbortSignal } = {} + const stream = summarize({ + adapter: summarizeAdapter(seen), + threadId: 'thread-1', + text: 'long text', + stream: true, + timeout: 10, + }) + await expect(collectChunks(stream)).rejects.toThrow() + expect(seen.signal?.aborted).toBe(true) + }) + + it('stream: caller abort reports onAbort, not onError', async () => { + const onAbort = vi.fn() + const onError = vi.fn() + const controller = new AbortController() + const stream = summarize({ + adapter: summarizeAdapter({}), + threadId: 'thread-1', + text: 'long text', + stream: true, + abortSignal: controller.signal, + middleware: [{ name: 'observe', onAbort, onError }], + }) + setTimeout(() => controller.abort(), 10) + await expect(collectChunks(stream)).rejects.toThrow() + expect(onError).not.toHaveBeenCalled() + expect(onAbort).toHaveBeenCalledTimes(1) + }) + + it('stream: a plain Error after the abort still reports onAbort', async () => { + // Many SDKs throw a plain Error (no AbortError name) once the request is + // cancelled, so only the aborted signal shows it was a cancel. + const text: ChatStreamCapable = { + chatStream: (options) => + (async function* () { + yield* [] + await hangUntilAborted(options.request?.signal ?? undefined).catch( + () => { + throw new Error('socket closed') + }, + ) + })(), + } + const onAbort = vi.fn() + const onError = vi.fn() + const controller = new AbortController() + const stream = summarize({ + adapter: new ChatStreamSummarizeAdapter( + text, + 'test-model', + 'openai', + ), + threadId: 'thread-1', + text: 'long text', + stream: true, + abortSignal: controller.signal, + middleware: [{ name: 'observe', onAbort, onError }], + }) + setTimeout(() => controller.abort(), 10) + await expect(collectChunks(stream)).rejects.toThrow() + expect(onError).not.toHaveBeenCalled() + expect(onAbort).toHaveBeenCalledTimes(1) + }) + + it('stream: an abort that comes as a RUN_ERROR chunk reports onAbort', async () => { + // Like openai-base: a failed request becomes a RUN_ERROR chunk. + const text: ChatStreamCapable = { + chatStream: (options) => + (async function* () { + yield ev.runStarted() + try { + await hangUntilAborted(options.request?.signal ?? undefined) + } catch { + yield ev.runError('Request was aborted.') + } + })(), + } + const onAbort = vi.fn() + const onFinish = vi.fn() + const controller = new AbortController() + const stream = summarize({ + adapter: new ChatStreamSummarizeAdapter( + text, + 'test-model', + 'openai', + ), + threadId: 'thread-1', + text: 'long text', + stream: true, + abortSignal: controller.signal, + middleware: [{ name: 'observe', onAbort, onFinish }], + }) + setTimeout(() => controller.abort(), 10) + const chunks = await collectChunks(stream).catch(() => []) + expect(onFinish).not.toHaveBeenCalled() + expect(onAbort).toHaveBeenCalledTimes(1) + // The error chunk still reaches the consumer. + expect( + chunks.some( + (chunk) => + chunk.type === 'RUN_ERROR' && + chunk.message === 'Request was aborted.', + ), + ).toBe(true) + }) +}) + +describe('middleware ctx.abort() reaches the provider', () => { + it('aborts the chatStream request (C3)', async () => { + let seen: AbortSignal | undefined + const { adapter } = createMockAdapter({ + chatStreamFn: (options) => + (async function* () { + seen = options.request?.signal + yield ev.runStarted() + // Like a real adapter: a failed request becomes a RUN_ERROR chunk. + try { + await hangUntilAborted(seen) + } catch (error) { + yield ev.runError((error as Error).message) + } + })(), + }) + const onAbort = vi.fn() + const onError = vi.fn() + const middleware: ChatMiddleware = { + name: 'aborter', + onStart: (ctx) => { + setTimeout(() => ctx.abort('stop'), 10) + }, + onAbort, + onError, + } + const run = collectChunks( + chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Hi' }], + middleware: [middleware], + abortController: new AbortController(), + }) as AsyncIterable, + ) + expect(await settlesWithin(run)).toBe('done') + expect(seen?.aborted).toBe(true) + expect(onAbort).toHaveBeenCalledOnce() + expect(onError).not.toHaveBeenCalled() + }) + + it('aborts the structured output request (C3)', async () => { + let seen: AbortSignal | undefined + let abortRun: ((reason?: string) => void) | undefined + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('{}'), + ev.textEnd(), + ev.runFinished('stop'), + ], + ], + structuredOutput: (options) => { + seen = options.chatOptions.request?.signal + abortRun?.('stop') + return hangUntilAborted(seen) + }, + }) + const run = chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Hi' }], + outputSchema: z.object({ a: z.string() }), + middleware: [ + { + name: 'aborter', + onStart: (ctx) => { + abortRun = ctx.abort + }, + }, + ], + abortController: new AbortController(), + }) + expect(await settlesWithin(run)).toBe('done') + expect(seen?.aborted).toBe(true) + }) + + it('aborts ctx.signal (C4)', async () => { + let signal: AbortSignal | undefined + const { adapter } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished('stop')]], + }) + await collectChunks( + chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Hi' }], + middleware: [ + { + name: 'aborter', + onStart: (ctx) => { + signal = ctx.signal + ctx.abort('stop') + }, + }, + ], + abortController: new AbortController(), + }) as AsyncIterable, + ) + expect(signal?.aborted).toBe(true) + }) + + it('gives the provider no signal when the caller gives none', async () => { + // Some SDKs drop their own timeout when a request has a signal. + let request: unknown = 'not called' + let ctxSignal: AbortSignal | undefined + const { adapter } = createMockAdapter({ + chatStreamFn: (options) => + (async function* () { + request = options.request + yield ev.runStarted() + yield ev.runFinished('stop') + })(), + }) + await collectChunks( + chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Hi' }], + middleware: [ + { + name: 'observe', + onStart: (ctx) => { + ctxSignal = ctx.signal + }, + }, + ], + }) as AsyncIterable, + ) + expect(request).toBeUndefined() + expect(ctxSignal).toBeUndefined() + }) + + it('stops a subagent child (C5)', async () => { + let outcome = 'running' + const worker = defineAgent({ + name: 'worker', + description: 'Waits for a stop', + run: (ctx) => + new Promise((resolve) => { + ctx.abortSignal?.addEventListener('abort', () => { + outcome = 'stopped' + resolve('stopped') + }) + setTimeout(() => { + if (outcome === 'running') outcome = 'never stopped' + resolve('done') + }, 300) + }), + }) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('c1', 'worker'), + ev.toolArgs('c1', '{}'), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) + await collectChunks( + chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [worker] }, + middleware: [ + { + name: 'aborter', + onBeforeToolCall: (ctx) => { + setTimeout(() => ctx.abort('stop'), 10) + return undefined + }, + }, + ], + }) as AsyncIterable, + ) + expect(outcome).toBe('stopped') + }) +}) + +describe('generateVideo() polling', () => { + // `stuck` is the adapter call that hangs while the caller aborts. + for (const stuck of ['getVideoStatus', 'getVideo'] as const) { + it(`stops on abort while ${stuck} hangs (C8)`, async () => { + const controller = new AbortController() + const hang = () => { + setTimeout(() => controller.abort(), 10) + return new Promise(() => {}) + } + const adapter = { + kind: 'video' as const, + name: 'mock', + model: 'mock-video-model', + createVideoJob: async () => ({ + jobId: 'job-1', + model: 'mock-video-model', + }), + getVideoStatus: () => + stuck === 'getVideoStatus' + ? hang() + : Promise.resolve({ status: 'completed' as const }), + getVideo: () => + stuck === 'getVideo' + ? hang() + : Promise.resolve({ url: 'https://example.com/v.mp4' }), + } + const stream = generateVideo({ + threadId: 'thread-1', + adapter: adapter as any, + prompt: 'a cat', + stream: true, + pollingInterval: 1, + abortSignal: controller.signal, + }) + expect(await settlesWithin(collectChunks(stream))).toBe('done') + }) + } + + it('starts no download when the caller aborts on the completed status', async () => { + const controller = new AbortController() + const getVideo = vi.fn(async () => ({ url: 'https://example.com/v.mp4' })) + const adapter = { + kind: 'video' as const, + name: 'mock', + model: 'mock-video-model', + createVideoJob: async () => ({ + jobId: 'job-1', + model: 'mock-video-model', + }), + getVideoStatus: async () => ({ status: 'completed' as const }), + getVideo, + } + const stream = generateVideo({ + threadId: 'thread-1', + adapter: adapter as any, + prompt: 'a cat', + stream: true, + pollingInterval: 1, + abortSignal: controller.signal, + }) + let sawCompleted = false + // The stream ends the run itself; it may not throw. + await (async () => { + for await (const chunk of stream as AsyncIterable) { + if ( + chunk.type === 'CUSTOM' && + chunk.name === 'video:status' && + (chunk.value as { status?: string }).status === 'completed' + ) { + sawCompleted = true + controller.abort() + } + } + })().catch(() => undefined) + expect(sawCompleted).toBe(true) + expect(getVideo).not.toHaveBeenCalled() + }) + + it('cancels the body of a video stream that comes after the abort', async () => { + const controller = new AbortController() + const cancel = vi.fn() + let deliver = () => {} + const adapter = { + kind: 'video' as const, + name: 'mock', + model: 'mock-video-model', + createVideoJob: async () => ({ + jobId: 'job-1', + model: 'mock-video-model', + }), + getVideoStatus: async () => ({ status: 'completed' as const }), + // The download answers only after the caller aborted. + getVideo: () => + new Promise((resolve) => { + deliver = () => + resolve({ + jobId: 'job-1', + contentType: 'video/mp4', + body: new ReadableStream({ cancel }), + }) + setTimeout(() => controller.abort(), 10) + }), + } + const stream = generateVideo({ + threadId: 'thread-1', + adapter: adapter as any, + prompt: 'a cat', + stream: true, + pollingInterval: 1, + abortSignal: controller.signal, + }) + await collectChunks(stream).catch(() => undefined) + deliver() + await vi.waitFor(() => expect(cancel).toHaveBeenCalled()) + }) +}) + +describe('toWebSocketStream() (non-durable)', () => { + it('stops pumping a turn after an abort frame (C11)', async () => { + const sent: Array = [] + const handlers: Record void>> = {} + const socket: WebSocketLike = { + send: (data) => sent.push(data), + close: () => {}, + addEventListener: (type, handler) => { + ;(handlers[type] ??= []).push(handler) + }, + } + const emit = (data: string) => { + for (const h of handlers.message ?? []) h({ data }) + } + let closed = false + toWebSocketStream(socket, new Request('https://x/api/chat'), { + // Ignores ctx.signal, like a hand-rolled stream. + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + try { + yield ev.runStarted(runId, threadId) + await new Promise((resolve) => setTimeout(resolve, 20)) + yield ev.textContent('late') + await new Promise((resolve) => setTimeout(resolve, 20)) + yield ev.textContent('later') + } finally { + closed = true + } + })(), + }) + emit( + JSON.stringify({ + threadId: 't', + runId: 'run-1', + messages: [], + tools: [], + context: [], + forwardedProps: {}, + state: {}, + }), + ) + await new Promise((resolve) => setTimeout(resolve, 5)) + emit(JSON.stringify({ type: 'abort', runId: 'run-1' })) + await new Promise((resolve) => setTimeout(resolve, 80)) + const deltas = sent.map((s) => JSON.parse(s).delta).filter(Boolean) + expect(deltas).toEqual([]) + expect(closed).toBe(true) + }) +}) diff --git a/testing/e2e/fixtures/summarize-abort/basic.json b/testing/e2e/fixtures/summarize-abort/basic.json new file mode 100644 index 0000000000..9b9d801724 --- /dev/null +++ b/testing/e2e/fixtures/summarize-abort/basic.json @@ -0,0 +1,14 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[summarize-abort] A slow summary that the timeout must stop" + }, + "response": { + "content": "This summary streams very slowly so that a short timeout ends the run before the model finishes and the request to the provider is aborted." + }, + "chunkSize": 2, + "streamingProfile": { "tps": 5 } + } + ] +} diff --git a/testing/e2e/src/routes/api.summarize.ts b/testing/e2e/src/routes/api.summarize.ts index 7e8040d7bf..b138fafd97 100644 --- a/testing/e2e/src/routes/api.summarize.ts +++ b/testing/e2e/src/routes/api.summarize.ts @@ -186,6 +186,7 @@ export const Route = createFileRoute('/api/summarize')({ testId, aimockPort, maxLength, + timeout, } = data as { text: string provider: Provider @@ -193,6 +194,7 @@ export const Route = createFileRoute('/api/summarize')({ testId?: string aimockPort?: number maxLength?: number + timeout?: number } try { @@ -216,6 +218,7 @@ export const Route = createFileRoute('/api/summarize')({ adapter, text, maxLength, + timeout, stream: false, }) return new Response(JSON.stringify({ summary }), { @@ -227,6 +230,7 @@ export const Route = createFileRoute('/api/summarize')({ adapter, text, maxLength, + timeout, stream: true, }) return toServerSentEventsResponse(stream) diff --git a/testing/e2e/tests/summarize-abort.spec.ts b/testing/e2e/tests/summarize-abort.spec.ts new file mode 100644 index 0000000000..16999f0e56 --- /dev/null +++ b/testing/e2e/tests/summarize-abort.spec.ts @@ -0,0 +1,30 @@ +import { test, expect } from './fixtures' + +// fixtures/summarize-abort/basic.json streams at 5 tokens per second, so the +// full summary takes many seconds. A 300 ms timeout must end the run first. +test('streaming summarize stops at its timeout', async ({ + request, + aimockPort, + testId, +}) => { + const started = Date.now() + const response = await request.post('/api/summarize', { + data: { + provider: 'openai', + text: '[summarize-abort] A slow summary that the timeout must stop', + stream: true, + timeout: 300, + aimockPort, + testId, + }, + }) + const body = await response.text() + + const elapsed = Date.now() - started + expect(body).toContain('RUN_ERROR') + expect(body).not.toContain('RUN_FINISHED') + // An error that is not the timeout (a failed connection, a bad route) + // comes at once. The timeout ends the run only after 300 ms. + expect(elapsed).toBeGreaterThanOrEqual(250) + expect(elapsed).toBeLessThan(5000) +})