diff --git a/.changeset/ai-abort-signal-options.md b/.changeset/ai-abort-signal-options.md new file mode 100644 index 0000000000..aea95be939 --- /dev/null +++ b/.changeset/ai-abort-signal-options.md @@ -0,0 +1,12 @@ +--- +'@tanstack/ai': minor +--- + +New optional cancel options. Old calls work the same. + +- `embed()` takes `abortSignal` and `timeout`. The effective signal reaches the adapter as `EmbeddingOptions.abortSignal`. +- `getVideoJobStatus()` takes `abortSignal` and `timeout`. The video adapter contract adds an optional `{ abortSignal }` to `getVideoStatus(jobId, options?)` and `getVideo(jobId, options?)`. The streaming `generateVideo()` poll passes its signal to both. +- `uploadFile()`, `getFile()`, and `deleteFile()` take `abortSignal`. `FilesAdapter.upload/get/delete` take an optional `{ abortSignal }`. +- `realtimeToken()` takes `abortSignal`. `RealtimeTokenAdapter.generateToken(options?)` and `RealtimeAdapter.connect(token, clientTools?, options?)` take an optional `{ abortSignal }`. +- `decide()` and `rerank()` take `timeout`, and race the adapter call against the signal like the other activities. +- `GenerationMiddlewareContext` has an optional `signal`: the effective activity signal. diff --git a/.changeset/ai-anthropic-abort-signal-options.md b/.changeset/ai-anthropic-abort-signal-options.md new file mode 100644 index 0000000000..5c3eb0ef43 --- /dev/null +++ b/.changeset/ai-anthropic-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-anthropic': minor +--- + +Forward the new optional `abortSignal` of the files adapter to `upload`, `get`, and `delete`. diff --git a/.changeset/ai-bedrock-abort-signal-options.md b/.changeset/ai-bedrock-abort-signal-options.md new file mode 100644 index 0000000000..7b2a7b39ca --- /dev/null +++ b/.changeset/ai-bedrock-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-bedrock': minor +--- + +Forward the new optional embedding `abortSignal` to each `InvokeModel` call. diff --git a/.changeset/ai-byteplus-abort-signal-options.md b/.changeset/ai-byteplus-abort-signal-options.md new file mode 100644 index 0000000000..5847552a95 --- /dev/null +++ b/.changeset/ai-byteplus-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-byteplus': minor +--- + +Forward the new optional `abortSignal` of `getVideoStatus` and `getVideo` to the task lookup. A caller signal now works together with the config `timeout`. diff --git a/.changeset/ai-byteplus-cancel-video-job.md b/.changeset/ai-byteplus-cancel-video-job.md new file mode 100644 index 0000000000..9d7e1fedbd --- /dev/null +++ b/.changeset/ai-byteplus-cancel-video-job.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-byteplus': minor +--- + +`byteplusVideo` implements the optional `cancelVideoJob(jobId, options?)` with `DELETE /contents/generations/tasks/{id}`. Ark cancels only a queued task. A streaming `generateVideo()` calls it when the caller aborts. diff --git a/.changeset/ai-cancel-video-job-mcp-signal.md b/.changeset/ai-cancel-video-job-mcp-signal.md new file mode 100644 index 0000000000..3215ac74fc --- /dev/null +++ b/.changeset/ai-cancel-video-job-mcp-signal.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai': minor +--- + +New optional cancel hooks. Old code works the same. + +- `MCPToolSource.tools()` takes an optional `signal`. `chat()` passes the signal of its `abortController` to each MCP source during tool discovery. +- `VideoAdapter` has an optional `cancelVideoJob(jobId, options?)` method. A streaming `generateVideo()` calls it (best effort) when its `abortSignal` or `timeout` aborts after the provider accepted the job. A failed cancel is logged and does not change the abort result. diff --git a/.changeset/ai-client-abort-signal-options.md b/.changeset/ai-client-abort-signal-options.md new file mode 100644 index 0000000000..1693d89808 --- /dev/null +++ b/.changeset/ai-client-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': minor +--- + +A client tool now gets `abortSignal` in its execute context. It aborts when you call `stop()` on the `ChatClient`. diff --git a/.changeset/ai-client-cancel-url-connect-signal.md b/.changeset/ai-client-cancel-url-connect-signal.md new file mode 100644 index 0000000000..5917a7f0f0 --- /dev/null +++ b/.changeset/ai-client-cancel-url-connect-signal.md @@ -0,0 +1,9 @@ +--- +'@tanstack/ai-client': minor +--- + +New optional cancel features. Old code works the same. + +- `fetchServerSentEvents` and `fetchHttpStream` take an optional `cancelUrl`. With it, `ChatClient.stop()` sends `POST cancelUrl` with `{ runId, threadId }`, so the server can cancel a durable run. Only `stop()` sends it, not an unmount or a reload. +- Connection adapters can have an optional `cancelRun({ runId, threadId })` method. `stop()` calls it. +- `RealtimeClient` passes an `abortSignal` to `adapter.connect()`. A `disconnect()` during `connect()` aborts the setup, and the status goes to `idle`, not `error`. diff --git a/.changeset/ai-cloudflare-abort-signal-options.md b/.changeset/ai-cloudflare-abort-signal-options.md new file mode 100644 index 0000000000..88f7da8305 --- /dev/null +++ b/.changeset/ai-cloudflare-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-cloudflare': minor +--- + +Forward the new optional embedding `abortSignal` to the Workers AI REST request. diff --git a/.changeset/ai-cohere-abort-signal-options.md b/.changeset/ai-cohere-abort-signal-options.md new file mode 100644 index 0000000000..21eefb8107 --- /dev/null +++ b/.changeset/ai-cohere-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-cohere': minor +--- + +Forward the new optional embedding `abortSignal` to the embed request and to an image URL download. diff --git a/.changeset/ai-elevenlabs-abort-signal-options.md b/.changeset/ai-elevenlabs-abort-signal-options.md new file mode 100644 index 0000000000..ea76a48ea7 --- /dev/null +++ b/.changeset/ai-elevenlabs-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-elevenlabs': minor +--- + +Forward the new optional `abortSignal` of `generateToken()` to the signed URL request. diff --git a/.changeset/ai-fal-abort-signal-options.md b/.changeset/ai-fal-abort-signal-options.md new file mode 100644 index 0000000000..da32b1ec45 --- /dev/null +++ b/.changeset/ai-fal-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-fal': minor +--- + +Forward the new optional `abortSignal` of `getVideoStatus` and `getVideo` to `fal.queue.status` and `fal.queue.result`. diff --git a/.changeset/ai-fal-cancel-video-job.md b/.changeset/ai-fal-cancel-video-job.md new file mode 100644 index 0000000000..a9bfdbee61 --- /dev/null +++ b/.changeset/ai-fal-cancel-video-job.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-fal': minor +--- + +`falVideo` implements the optional `cancelVideoJob(jobId, options?)` with `fal.queue.cancel`. A streaming `generateVideo()` calls it when the caller aborts. diff --git a/.changeset/ai-gemini-abort-signal-options.md b/.changeset/ai-gemini-abort-signal-options.md new file mode 100644 index 0000000000..12a3860069 --- /dev/null +++ b/.changeset/ai-gemini-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-gemini': minor +--- + +Forward the new optional `abortSignal` in embeddings, Veo and Omni video status and download, files `get` and `delete`, and the realtime token request. `uploadGeminiFile()` takes `abortSignal`: it cancels the status requests and the wait between polls. diff --git a/.changeset/ai-grok-abort-signal-options.md b/.changeset/ai-grok-abort-signal-options.md new file mode 100644 index 0000000000..b4c8802f9f --- /dev/null +++ b/.changeset/ai-grok-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-grok': minor +--- + +Forward the new optional `abortSignal` in video status and download, files upload/get/delete, the realtime token request, and the realtime WebRTC SDP request. diff --git a/.changeset/ai-lovable-abort-signal-options.md b/.changeset/ai-lovable-abort-signal-options.md new file mode 100644 index 0000000000..3a186651e7 --- /dev/null +++ b/.changeset/ai-lovable-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-lovable': minor +--- + +Forward the new optional `abortSignal` in embeddings and in video status and download. diff --git a/.changeset/ai-mistral-abort-signal-options.md b/.changeset/ai-mistral-abort-signal-options.md new file mode 100644 index 0000000000..ee09c68d31 --- /dev/null +++ b/.changeset/ai-mistral-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-mistral': minor +--- + +Forward the new optional embedding `abortSignal` to the embeddings request. diff --git a/.changeset/ai-neon-abort-signal-options.md b/.changeset/ai-neon-abort-signal-options.md new file mode 100644 index 0000000000..fcbf72ef8a --- /dev/null +++ b/.changeset/ai-neon-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-neon': minor +--- + +Forward the new optional embedding `abortSignal` to the embeddings request. diff --git a/.changeset/ai-ollama-abort-signal-options.md b/.changeset/ai-ollama-abort-signal-options.md new file mode 100644 index 0000000000..8bf513963a --- /dev/null +++ b/.changeset/ai-ollama-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-ollama': minor +--- + +Forward the new optional embedding `abortSignal` to the embed request. An injected `Ollama` client cannot take a per-call signal, so `embed()` only rejects on abort for it. diff --git a/.changeset/ai-openai-abort-signal-options.md b/.changeset/ai-openai-abort-signal-options.md new file mode 100644 index 0000000000..aa7efb1671 --- /dev/null +++ b/.changeset/ai-openai-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-openai': minor +--- + +Forward the new optional `abortSignal` to the provider request in embeddings, video status and download, files upload/get/delete, the realtime token request, and the realtime WebRTC SDP request. diff --git a/.changeset/ai-openrouter-abort-signal-options.md b/.changeset/ai-openrouter-abort-signal-options.md new file mode 100644 index 0000000000..5484647132 --- /dev/null +++ b/.changeset/ai-openrouter-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-openrouter': minor +--- + +Forward the new optional `abortSignal` of `getVideoStatus` and `getVideo` to the generation lookup and the content download. diff --git a/.changeset/ai-vercel-gateway-abort-signal-options.md b/.changeset/ai-vercel-gateway-abort-signal-options.md new file mode 100644 index 0000000000..a5abaa0cf4 --- /dev/null +++ b/.changeset/ai-vercel-gateway-abort-signal-options.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-vercel-gateway': minor +--- + +Forward the new optional embedding `abortSignal` to the embeddings request. diff --git a/docs/advanced/files-api.md b/docs/advanced/files-api.md index 332f0aacda..d87263f3db 100644 --- a/docs/advanced/files-api.md +++ b/docs/advanced/files-api.md @@ -199,6 +199,26 @@ await deleteFile({ adapter: files, id: handle }); fal storage is upload-only. `falFiles()` has no `get` or `delete`, so `getFile()` and `deleteFile()` throw with it. +### Cancel a file request + +`uploadFile()`, `getFile()`, and `deleteFile()` take an `abortSignal`. When it aborts, the call rejects. The adapter passes the signal to the provider request when the provider SDK can cancel it: + +```typescript +import { uploadFile } from "@tanstack/ai"; +import { openaiFiles } from "@tanstack/ai-openai"; + +export async function POST(request: Request) { + const handle = await uploadFile({ + adapter: openaiFiles(), + input: await request.blob(), + abortSignal: request.signal, + }); + return Response.json(handle); +} +``` + +The Gemini and fal SDKs do not cancel a running upload. With those providers, the call still rejects at once, but the upload can finish on the provider. + ### Grok (xAI): handles are public URLs `grokFiles()` uploads to the xAI Files API. Then it mints a [public URL](https://docs.x.ai/developers/files/public-urls) for the file and uses that URL as the reference of the handle. diff --git a/docs/advanced/otel.md b/docs/advanced/otel.md index 300260b2c5..cb3ec3f48d 100644 --- a/docs/advanced/otel.md +++ b/docs/advanced/otel.md @@ -314,6 +314,8 @@ The four extension points also apply to media spans. Media spans report `kind: ' The shared lifecycle hooks (`onStart`, `onUsage`, `onFinish`, `onAbort`, `onError`) use the `GenerationMiddlewareContext`, which is the same for each activity. That is why one instance works everywhere. For a custom backend, implement the base `GenerationMiddleware` contract. Its hooks get the `GenerationMiddlewareContext` and fire for each activity, chat included. +In a custom middleware, `ctx.signal` is the effective abort signal of the activity call: the caller `abortSignal` and the `timeout` together. It is `undefined` when the call has neither. Pass it to async work in a result transform, such as an upload, so that a cancelled call also stops that work. + The `GenerationMiddleware` types come from the package root. The `otelMiddleware` value lives on the `@tanstack/ai/middlewares/otel` subpath, so an import of `@tanstack/ai` never needs the optional `@opentelemetry/api` peer. ### Capturing media content diff --git a/docs/api/ai.md b/docs/api/ai.md index 6fbeb450b8..7fd608dfce 100644 --- a/docs/api/ai.md +++ b/docs/api/ai.md @@ -454,6 +454,7 @@ console.log(result.queue.value, result.urgency.value, result.refund.value); | `state` | Required. The content that every question judges: a string, an object, or an array. An array is one state, not a batch. | | `questions` | Required. A map of `choice`, `score`, and `boolean` questions. The key `meta` is reserved. | | `abortSignal` | Cancels the request. | +| `timeout` | Stops the request after this many milliseconds. | | `modelOptions` | Provider options. | | `middleware` | Observe-only generation middleware. | | `debug` | Debug logging. | diff --git a/docs/config.json b/docs/config.json index 0fa0528cb7..1957e47cb1 100644 --- a/docs/config.json +++ b/docs/config.json @@ -434,7 +434,7 @@ "label": "OpenTelemetry", "to": "advanced/otel", "addedAt": "2026-05-08", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Compaction", @@ -626,7 +626,7 @@ "label": "Managed MCP with chat()", "to": "tools/mcp-managed", "addedAt": "2026-06-05", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Manual MCP", @@ -751,7 +751,7 @@ "label": "Advanced", "to": "resumable-streams/advanced", "addedAt": "2026-08-04", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "WebSockets", @@ -884,7 +884,7 @@ "label": "Realtime Voice Chat", "to": "media/realtime-chat", "addedAt": "2026-04-15", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Text-to-Speech", @@ -936,12 +936,13 @@ "label": "Video Generation", "to": "media/video-generation", "addedAt": "2026-04-15", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Video Generation: Advanced", "to": "media/video-generation-advanced", - "addedAt": "2026-10-21" + "addedAt": "2026-10-21", + "updatedAt": "2026-11-02" }, { "label": "Live Generation", @@ -971,7 +972,7 @@ "label": "Embeddings", "to": "embeddings", "addedAt": "2026-07-10", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" } ], "tab": "guides" @@ -983,7 +984,7 @@ "label": "Reranking", "to": "rerank/rerank", "addedAt": "2026-06-25", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" } ], "tab": "guides" @@ -995,7 +996,7 @@ "label": "Evaluate", "to": "evaluate/evaluate", "addedAt": "2026-09-18", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" } ], "tab": "guides" @@ -1493,7 +1494,7 @@ "label": "Takeover & Detached Runs", "to": "sandbox/takeover", "addedAt": "2026-08-04", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Takeover Internals", @@ -1609,7 +1610,7 @@ "label": "Files API", "to": "advanced/files-api", "addedAt": "2026-09-11", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Per-Model Type Safety", @@ -2005,7 +2006,7 @@ "label": "@tanstack/ai", "to": "api/ai", "addedAt": "2026-04-15", - "updatedAt": "2026-10-21" + "updatedAt": "2026-11-02" }, { "label": "@tanstack/ai-client", diff --git a/docs/embeddings.md b/docs/embeddings.md index 9ac3e8a916..ba48aae936 100644 --- a/docs/embeddings.md +++ b/docs/embeddings.md @@ -277,6 +277,27 @@ const result = await embed({ console.log(result.usage?.promptTokens); ``` +### Cancel or time out + +Pass `abortSignal` to cancel the call, or `timeout` in milliseconds to stop it after a fixed time. The adapter cancels the provider request, and `embed()` rejects. A timeout rejects with a `TimeoutError`. + +```typescript +import { embed } from "@tanstack/ai"; +import { openaiEmbedding } from "@tanstack/ai-openai"; + +export async function POST(request: Request) { + const result = await embed({ + adapter: openaiEmbedding("text-embedding-3-small"), + input: "a red guitar", + abortSignal: request.signal, + timeout: 10_000, + }); + return Response.json(result.embeddings[0]?.vector); +} +``` + +When you pass both, the first one to fire wins. Middleware gets `onAbort`, not `onError`. + ### Errors `embed()` rejects with the provider error. Middleware `onError` hooks run before the rejection reaches your code. diff --git a/docs/evaluate/evaluate.md b/docs/evaluate/evaluate.md index dcdb8bd778..60003eeda7 100644 --- a/docs/evaluate/evaluate.md +++ b/docs/evaluate/evaluate.md @@ -313,6 +313,7 @@ On the TypeSafe wire, the `boolean` question type is named `noul`. | `state` | `string \| object \| array` | The data that every question judges. Required. | | `questions` | `Record` | Map of `choice`, `score`, and `boolean` questions. Required. | | `abortSignal` | `AbortSignal` | Cancel the request | +| `timeout` | `number` | Stop the request after this many milliseconds | | `modelOptions` | provider options | Provider-specific options | | `middleware` | `Array` | Observe-only lifecycle hooks (usage, finish, error, abort) | | `debug` | `DebugOption` | Debug logging | @@ -325,7 +326,7 @@ On the TypeSafe wire, the `boolean` question type is named `noul`. ### Cancel a request -Pass an `abortSignal`: +Pass an `abortSignal`. To stop after a fixed time, pass `timeout` in milliseconds instead. A timeout rejects with a `TimeoutError`. ```typescript import { decide, boolean } from '@tanstack/ai' diff --git a/docs/media/realtime-chat.md b/docs/media/realtime-chat.md index cb0210ace3..1803c43874 100644 --- a/docs/media/realtime-chat.md +++ b/docs/media/realtime-chat.md @@ -536,6 +536,25 @@ client.destroy() The React hook disconnects by itself when the component unmounts. With `RealtimeClient`, call `disconnect()` and `destroy()` yourself. +A `disconnect()` while `connect()` still runs cancels the connection setup. `connect()` then rejects with an `AbortError`, and the status is `idle`, not `error`. The OpenAI and Grok adapters stop their setup requests. + +### Cancel the token request + +`realtimeToken()` takes an `abortSignal`. Pass the request signal, so the provider call stops when the HTTP request to your route is aborted: + +```typescript +import { realtimeToken } from '@tanstack/ai' +import { openaiRealtimeToken } from '@tanstack/ai-openai' + +export async function POST(request: Request) { + const token = await realtimeToken({ + adapter: openaiRealtimeToken({ model: 'gpt-realtime-2.1' }), + abortSignal: request.signal, + }) + return Response.json(token) +} +``` + ### Tips - Mint tokens only on the server. Never send an API key to the browser. diff --git a/docs/media/video-generation-advanced.md b/docs/media/video-generation-advanced.md index 85f82bc876..d0eb35c6ef 100644 --- a/docs/media/video-generation-advanced.md +++ b/docs/media/video-generation-advanced.md @@ -65,6 +65,22 @@ console.log('Video ready:', videoUrl) If you use middleware, such as `withGenerationPersistence`, pass the same `middleware` and `threadId` to `generateVideo()` and to `getVideoJobStatus()`. The submit call only opens the run. The poll that sees the finished job closes the run and records the result. +To cancel one poll, pass `abortSignal`. To limit it, pass `timeout` in milliseconds. The limit covers the status request and the video download. + +```typescript +import { getVideoJobStatus } from '@tanstack/ai' +import { openaiVideo } from '@tanstack/ai-openai' + +const status = await getVideoJobStatus({ + adapter: openaiVideo('sora-2'), + jobId: 'video_123', + timeout: 15_000, +}) +console.log(status.status) +``` + +A cancelled poll rejects. It does not change the job. The provider job keeps running, and a later poll can still read it. + ### Server function that polls (fetcher) A TanStack Start server function can run the full loop and return the finished result as JSON: @@ -154,6 +170,32 @@ export function VideoGenerator() { } ``` +## Cancel the provider job + +Pass `abortSignal` (or `timeout`) to a streaming `generateVideo()`. If it aborts after the provider accepts the job, `generateVideo()` asks the provider to cancel the job: + +```typescript +import { generateVideo, toServerSentEventsResponse } from '@tanstack/ai' +import { falVideo } from '@tanstack/ai-fal' + +export function POST(request: Request) { + const stream = generateVideo({ + adapter: falVideo('fal-ai/veo3.1'), + prompt: 'A timelapse of a city skyline at dusk', + threadId: 'thread-1', + stream: true, + abortSignal: request.signal, + }) + return toServerSentEventsResponse(stream) +} +``` + +- fal.ai and BytePlus have a cancel API. fal can refuse a request that already runs. BytePlus cancels only a queued task. +- The other adapters have no cancel API, so their job keeps running. +- The cancel is best effort. A failed cancel goes to the error log, and the run still ends as aborted. +- To cancel a job that you poll yourself, call `adapter.cancelVideoJob?.(jobId)`. +- In your own adapter, add the optional `cancelVideoJob(jobId, options?)` method when the provider can cancel a job. + ## Image-to-video support What each provider does with image parts in the `prompt`. The roles are in [Role hints](./video-generation#role-hints). diff --git a/docs/media/video-generation.md b/docs/media/video-generation.md index 026b451289..bee6f21e73 100644 --- a/docs/media/video-generation.md +++ b/docs/media/video-generation.md @@ -144,7 +144,7 @@ const stream = generateVideo({ }) ``` -`stop()` on the hook only closes the stream in the browser. The provider job can continue to run. +`stop()` on the hook only closes the stream in the browser. The provider job can continue to run. To cancel the job on fal.ai or BytePlus too, see [Cancel the provider job](./video-generation-advanced#cancel-the-provider-job). The `jobId`, `videoStatus`, `onJobCreated`, and `onStatusUpdate` hook values are listed in [Generation Hooks](./generation-hooks#usegeneratevideo). The other transports (a server function that returns JSON, or one that returns an SSE `Response`) are in [Generations](./generations#transports-in-full). diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md index 9c79946611..eafd9b2912 100644 --- a/docs/rerank/rerank.md +++ b/docs/rerank/rerank.md @@ -207,6 +207,7 @@ pnpm dev | `documents` | `Array` | The candidate documents. Required. | | `topN` | `number` | Return only the top N results | | `abortSignal` | `AbortSignal` | Cancel the request | +| `timeout` | `number` | Stop the request after this many milliseconds | | `modelOptions` | provider options | Provider-specific options (see below) | | `middleware` | `Array` | Observe-only lifecycle hooks (usage, finish, error, abort) | @@ -262,7 +263,7 @@ interface RerankResult { ### Cancel a request -Pass an `abortSignal`: +Pass an `abortSignal`. To stop after a fixed time, pass `timeout` in milliseconds instead. A timeout rejects with a `TimeoutError`. ```typescript import { rerank } from '@tanstack/ai' diff --git a/docs/resumable-streams/advanced.md b/docs/resumable-streams/advanced.md index cb8cfc0717..ec5f8870c2 100644 --- a/docs/resumable-streams/advanced.md +++ b/docs/resumable-streams/advanced.md @@ -142,7 +142,7 @@ If the replay ends on a client tool, the client sends the tool result after the A run ends early only on a real cancel or a failure: -- **Cancel**: you pass an `AbortController` to the response as `abortController`, then abort it. For example, from a Stop button, or by forwarding `request.signal` to a controller of your own. This stops the producer and appends a terminal `RUN_ERROR`. A plain client disconnect does not do this. Pass a controller if you want a disconnect to stop the run too. +- **Cancel**: you pass an `AbortController` to the response as `abortController`, then abort it. For example, from a Stop button, or by forwarding `request.signal` to a controller of your own. This stops the producer and appends a terminal `RUN_ERROR`. A plain client disconnect does not do this. Pass a controller if you want a disconnect to stop the run too. To abort it from `stop()`, see [Cancel the run on stop()](#cancel-the-run-on-stop). - **Provider failure**: the model stream throws, and the error is appended as a terminal `RUN_ERROR`. In both cases, the producer writes a terminal and then awaits `close()` on exit. A reconnecting or joining client then sees a terminal and does not hang. @@ -163,6 +163,102 @@ function respond(request: Request, stream: AsyncIterable) { A durable source must end with its own terminal event (`RUN_FINISHED` or `RUN_ERROR`). On normal completion, core terminalizes the log only if the source emitted a terminal. Without one, a durable consumer reconnects once, gets nothing new, and fails with `DurableStreamIncompleteError`. `chat()` always emits `RUN_FINISHED`, so this affects only streams that you build yourself. +### Cancel the run on stop() + +Set `cancelUrl` on the connection. Then `stop()` also sends `POST cancelUrl` with the JSON body `{ runId, threadId }`, and your route aborts the run. + +1. Keep the controller of each run, and abort it from a cancel route: + +```ts +import { + RUN_CANCEL_REASON, + chat, + chatParamsFromRequest, + memoryStream, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' + +// The runs that this process drives, by run id. +const liveRuns = new Map() +// Cancels that arrived before their chat request, by run id. +const earlyCancels = new Set() + +// POST /api/chat +export async function POST(request: Request) { + const { messages, threadId, runId } = await chatParamsFromRequest(request) + const abortController = new AbortController() + // Stop was pressed before this request arrived: start the run aborted. + if (earlyCancels.delete(runId)) abortController.abort(RUN_CANCEL_REASON) + liveRuns.set(runId, abortController) + const forget = () => { + liveRuns.delete(runId) + } + const stream = chat({ + adapter: openaiText('gpt-6.1-sol'), + messages, + threadId, + runId, + abortController, + middleware: [ + { name: 'forget-run', onFinish: forget, onAbort: forget, onError: forget }, + ], + }) + return toServerSentEventsResponse(stream, { + abortController, + durability: { adapter: memoryStream(request) }, + }) +} + +// POST /api/chat/cancel +export async function cancelChatRun(request: Request) { + const body: unknown = await request.json() + if (typeof body !== 'object' || body === null || !('runId' in body)) { + return new Response('runId is required', { status: 400 }) + } + const runId = body.runId + if (typeof runId !== 'string') { + return new Response('runId must be a string', { status: 400 }) + } + const run = liveRuns.get(runId) + if (run) { + run.abort(RUN_CANCEL_REASON) + } else { + // The chat request can still be on its way. Keep the cancel for a minute. + earlyCancels.add(runId) + setTimeout(() => earlyCancels.delete(runId), 60_000) + } + return new Response(null, { status: 204 }) +} +``` + +2. Set `cancelUrl` on the client: + +```tsx +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' + +export function Chat({ threadId }: { threadId: string }) { + const chat = useChat({ + threadId, + connection: fetchServerSentEvents('/api/chat', { + cancelUrl: '/api/chat/cancel', + }), + }) + + return +} +``` + +Press Stop while a reply streams. The run ends with a terminal `RUN_ERROR` in the log, and a reload no longer resumes it. + +Behavior to know: + +- Only `stop()` sends the cancel. An unmount, a reload, or a lost connection does not, so the run keeps going for them. +- The cancel request has the same `headers` and `credentials` as the chat request. Authorize it like the chat route. +- `fetchHttpStream` takes `cancelUrl` too. +- A cancel can arrive before its chat request. The sample keeps it in `earlyCancels`, so the run starts aborted. +- `liveRuns` is in the memory of one process. The cancel request must reach the process that drives the run. For runs on other hosts, see [Detach vs cancel](../sandbox/takeover#detach-vs-cancel). + ### memoryStream in production In one live process, `memoryStream` survives a client disconnect. The producer outlives the socket and keeps writing into the log, so a later reconnect or `joinRun` resumes a run that still goes on. Two limits keep it to development and single-process deployments: diff --git a/docs/sandbox/takeover.md b/docs/sandbox/takeover.md index a817e29814..6875f7ad06 100644 --- a/docs/sandbox/takeover.md +++ b/docs/sandbox/takeover.md @@ -317,7 +317,7 @@ export async function POST(request: Request) { `requestRunCancel` writes no status. Only the driver knows when the agent is dead and the sandbox is gone, so only the driver writes `'aborted'`. -2. Make the Stop button call the endpoint. `chat.stop()` alone sends the server nothing, and a durable run keeps working, because that looks like a refresh. +2. Set `cancelUrl` on the connection. Then `chat.stop()` also posts `{ runId, threadId }` to the endpoint. Without it, `chat.stop()` sends the server nothing, and a durable run keeps working, because that looks like a refresh. ```tsx import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' @@ -325,21 +325,12 @@ import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' export function StoppableChat({ threadId }: { threadId: string }) { const chat = useChat({ threadId, - connection: fetchServerSentEvents('/api/chat'), + connection: fetchServerSentEvents('/api/chat', { + cancelUrl: '/api/chat/cancel', + }), }) - async function stopRun(): Promise { - // Local: stop rendering the stream now. - chat.stop() - // Remote: tell the server that this is a cancel, not a lost connection. - await fetch('/api/chat/cancel', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ threadId }), - }) - } - - return + return } ``` diff --git a/docs/tools/client-tools.md b/docs/tools/client-tools.md index 4c0fdc9d33..9f509d735f 100644 --- a/docs/tools/client-tools.md +++ b/docs/tools/client-tools.md @@ -336,6 +336,26 @@ const chat = useChat(chatOptions); Does the server also need a value from the client? Send it with `forwardedProps`, check it in your route, and put it in the server `chat({ context })`. See [Runtime Context](../advanced/runtime-context). +### Stop a client tool when the chat stops + +The second argument also has `abortSignal`. It aborts when you call `stop()` on the chat. Pass it to `fetch` and other long work: + +```typescript +import { toolDefinition } from "@tanstack/ai"; +import { z } from "zod"; + +const searchDocs = toolDefinition({ + name: "search_docs", + description: "Search the docs index of the app", + inputSchema: z.object({ query: z.string() }), +}).client(async ({ query }, ctx) => { + const response = await fetch(`/api/search?q=${encodeURIComponent(query)}`, { + signal: ctx?.abortSignal, + }); + return { results: await response.text() }; +}); +``` + ### One definition, two implementations A definition can have a `.server()` and a `.client()` implementation. The tool that you pass to `chat()` decides where it runs: diff --git a/docs/tools/mcp-managed.md b/docs/tools/mcp-managed.md index 2259b84846..dcacdabf60 100644 --- a/docs/tools/mcp-managed.md +++ b/docs/tools/mcp-managed.md @@ -67,7 +67,7 @@ That is all. There is no `try/finally`. Your `useChat` code does not change. ## How it works -- Before the first model call, `chat()` calls `tools()` on each client at the same time and merges the results into one tool set. +- Before the first model call, `chat()` calls `tools()` on each client at the same time and merges the results into one tool set. When you pass `abortController` to `chat()`, each `tools()` call gets its signal as `signal`. - Tools from `mcp.clients` merge with any tools that you pass in `tools`. You can use both in one call. - `connection: 'close'` (the default) closes each client when the run ends: on success, on error, and on abort. - `connection: 'keep-alive'` never closes the clients. You own them. diff --git a/packages/ai-anthropic/src/adapters/files.ts b/packages/ai-anthropic/src/adapters/files.ts index 133859ddad..a3323f77bc 100644 --- a/packages/ai-anthropic/src/adapters/files.ts +++ b/packages/ai-anthropic/src/adapters/files.ts @@ -31,30 +31,50 @@ export class AnthropicFilesAdapter extends BaseFilesAdapter<'anthropic'> { this.client = createAnthropicClient(config) } - async upload(input: FileUploadInput): Promise> { + async upload( + input: FileUploadInput, + options?: { abortSignal?: AbortSignal }, + ): Promise> { const { blob, mimeType, filename } = normalizeFileUploadInput(input) const file = await toFile(blob, filename, { ...(mimeType ? { type: mimeType } : {}), }) - const result = await this.client.beta.files.upload({ - file, - betas: [FILES_API_BETA], - }) + const result = await this.client.beta.files.upload( + { file, betas: [FILES_API_BETA] }, + requestOptions(options), + ) return toFileHandle(result) } - async get(id: string): Promise> { - const result = await this.client.beta.files.retrieveMetadata(id, { - betas: [FILES_API_BETA], - }) + async get( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise> { + const result = await this.client.beta.files.retrieveMetadata( + id, + { betas: [FILES_API_BETA] }, + requestOptions(options), + ) return toFileHandle(result) } - async delete(id: string): Promise { - await this.client.beta.files.delete(id, { betas: [FILES_API_BETA] }) + async delete( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + await this.client.beta.files.delete( + id, + { betas: [FILES_API_BETA] }, + requestOptions(options), + ) } } +/** SDK request options that carry the caller's abort signal, if any. */ +function requestOptions(options?: { abortSignal?: AbortSignal }) { + return options?.abortSignal ? { signal: options.abortSignal } : undefined +} + function toFileHandle(file: BetaFileMetadata): FileHandle<'anthropic'> { return { id: file.id, diff --git a/packages/ai-anthropic/tests/abort-signal.test.ts b/packages/ai-anthropic/tests/abort-signal.test.ts new file mode 100644 index 0000000000..297334a3b3 --- /dev/null +++ b/packages/ai-anthropic/tests/abort-signal.test.ts @@ -0,0 +1,47 @@ +import { describe, expect, it, vi } from 'vitest' +import { createAnthropicFiles } from '../src/adapters/files' + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('anthropic adapters forward abortSignal to the request', () => { + it.each(['upload', 'get', 'delete'] as const)('files %s', async (method) => { + const { fetch, signals } = hangingFetch() + const files = createAnthropicFiles('k', { fetch, maxRetries: 0 }) + await expectAbortReachesFetch(signals, (abortSignal) => + method === 'upload' + ? files.upload({ data: 'AAAA', mimeType: 'image/png' }, { abortSignal }) + : files[method]('file-1', { abortSignal }), + ) + }) +}) diff --git a/packages/ai-bedrock/src/adapters/embedding.ts b/packages/ai-bedrock/src/adapters/embedding.ts index 564a678a04..2ba7c509d1 100644 --- a/packages/ai-bedrock/src/adapters/embedding.ts +++ b/packages/ai-bedrock/src/adapters/embedding.ts @@ -170,6 +170,7 @@ export class BedrockEmbeddingAdapter< protected async invokeModel( modelId: string, body: Record, + abortSignal?: AbortSignal, ): Promise { const { InvokeModelCommand } = await this.importBedrockRuntime() const client = await this.getClient() @@ -180,6 +181,7 @@ export class BedrockEmbeddingAdapter< accept: 'application/json', body: JSON.stringify(body), }), + abortSignal ? { abortSignal } : {}, ) return JSON.parse(new TextDecoder().decode(response.body)) } @@ -254,7 +256,7 @@ export class BedrockEmbeddingAdapter< if (dimensions !== undefined) body.dimensions = dimensions if (normalize !== undefined) body.normalize = normalize return readTitanEmbeddingBody( - await this.invokeModel(model, body), + await this.invokeModel(model, body, options.abortSignal), `${this.name} ${model}`, ) }, @@ -299,7 +301,7 @@ export class BedrockEmbeddingAdapter< const image = item.images[0] if (image) body.inputImage = toTitanInputImage(image, model) return readTitanEmbeddingBody( - await this.invokeModel(model, body), + await this.invokeModel(model, body, options.abortSignal), `${this.name} ${model}`, ) }, @@ -343,7 +345,7 @@ export class BedrockEmbeddingAdapter< } if (truncate !== undefined) body.truncate = truncate return readCohereEmbeddingBody( - await this.invokeModel(model, body), + await this.invokeModel(model, body, options.abortSignal), `${this.name} ${model}`, ) }, diff --git a/packages/ai-bedrock/tests/abort-signal.test.ts b/packages/ai-bedrock/tests/abort-signal.test.ts new file mode 100644 index 0000000000..70e4a4cdf2 --- /dev/null +++ b/packages/ai-bedrock/tests/abort-signal.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { BedrockEmbeddingAdapter } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** Records the `send` options; hangs until their `abortSignal` aborts. */ +class HangingClientAdapter extends BedrockEmbeddingAdapter<'amazon.titan-embed-text-v2:0'> { + signals: Array = [] + + protected override importBedrockRuntime(): any { + return Promise.resolve({ + InvokeModelCommand: class { + constructor(readonly input: unknown) {} + }, + }) + } + + protected override getClient(): any { + return Promise.resolve({ + send: (_command: unknown, options?: { abortSignal?: AbortSignal }) => { + const signal = options?.abortSignal + this.signals.push(signal) + if (!signal) return Promise.reject(new Error('request had no signal')) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }, + }) + } +} + +describe('bedrock embedding forwards abortSignal to InvokeModel', () => { + it('aborts the in-flight send', async () => { + const adapter = new HangingClientAdapter( + { region: 'us-east-1' }, + 'amazon.titan-embed-text-v2:0', + ) + const controller = new AbortController() + const pending = adapter.createEmbeddings({ + model: 'amazon.titan-embed-text-v2:0', + input: ['a'], + logger, + abortSignal: controller.signal, + }) + await vi.waitFor(() => expect(adapter.signals).toHaveLength(1)) + expect(adapter.signals[0]).toBe(controller.signal) + controller.abort() + await expect(pending).rejects.toThrow() + }) +}) diff --git a/packages/ai-byteplus/src/adapters/video.ts b/packages/ai-byteplus/src/adapters/video.ts index 8fc66966be..5c40731f79 100644 --- a/packages/ai-byteplus/src/adapters/video.ts +++ b/packages/ai-byteplus/src/adapters/video.ts @@ -210,7 +210,7 @@ export class BytePlusVideoAdapter< private async request( path: string, - init?: Omit, + init?: Omit, abortSignal?: AbortSignal, ): Promise<{ response: Response; body: unknown }> { const fetchImpl = this.clientConfig.fetch ?? fetch @@ -498,9 +498,14 @@ export class BytePlusVideoAdapter< * interpret; the honest answer is that the response was not a task at all, * so say so while the body is still in hand. */ - private async retrieveTask(jobId: string): Promise { + private async retrieveTask( + jobId: string, + signal?: AbortSignal, + ): Promise { const { response, body } = await this.request( `${TASKS_PATH}/${encodeURIComponent(jobId)}`, + undefined, + signal, ) if (!response.ok) { const error = bytePlusArkError(response.status, body, 'video task lookup') @@ -517,10 +522,32 @@ export class BytePlusVideoAdapter< return body as BytePlusVideoTask } - async getVideoStatus(jobId: string): Promise { + /** + * `DELETE /contents/generations/tasks/{id}`. Ark can cancel only a + * `queued` task, and the call fails for a `running` one. On a finished + * task, Ark deletes the task record. + */ + async cancelVideoJob( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const { response, body } = await this.request( + `${TASKS_PATH}/${encodeURIComponent(jobId)}`, + { method: 'DELETE' }, + options?.abortSignal, + ) + if (!response.ok) { + throw bytePlusArkError(response.status, body, 'video task cancel') + } + } + + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { let task: BytePlusVideoTask try { - task = await this.retrieveTask(jobId) + task = await this.retrieveTask(jobId, options?.abortSignal) } catch (error) { // A task record lives 7 days from creation; past that the id 404s. Keep // Ark's own code/message: a 404 from a wrong baseURL, a proxy, or a @@ -545,10 +572,13 @@ export class BytePlusVideoAdapter< } } - override async getVideo(jobId: string): Promise { + override async getVideo( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { let task: BytePlusVideoTask try { - task = await this.retrieveTask(jobId) + task = await this.retrieveTask(jobId, options?.abortSignal) } catch (error) { // See getVideoStatus: Ark's detail distinguishes an expired id from a // misrouted request. diff --git a/packages/ai-byteplus/tests/abort-signal.test.ts b/packages/ai-byteplus/tests/abort-signal.test.ts new file mode 100644 index 0000000000..613585ab68 --- /dev/null +++ b/packages/ai-byteplus/tests/abort-signal.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it, vi } from 'vitest' +import { createBytePlusVideo } from '../src/adapters/video' + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('byteplus adapters forward abortSignal to the request', () => { + it.each(['getVideoStatus', 'getVideo', 'cancelVideoJob'] as const)( + 'video %s', + async (method) => { + const { fetch, signals } = hangingFetch() + const adapter = createBytePlusVideo('seedance-1-5-pro-251215', 'k', { + fetch, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter[method]('job-1', { abortSignal }), + ) + }, + ) +}) + +describe('byteplus cancelVideoJob', () => { + it('sends DELETE for the task and resolves on success', async () => { + const calls: Array<{ url: string; method?: string }> = [] + const fetch = async (input: unknown, init?: RequestInit) => { + calls.push({ url: String(input), method: init?.method }) + return new Response('{}', { status: 200 }) + } + const adapter = createBytePlusVideo('seedance-1-5-pro-251215', 'k', { + fetch, + }) + await adapter.cancelVideoJob('job-1') + expect(calls).toHaveLength(1) + expect(calls[0]!.method).toBe('DELETE') + expect(calls[0]!.url).toMatch(/\/contents\/generations\/tasks\/job-1$/) + }) + + it('rejects when Ark refuses the cancel', async () => { + const fetch = async () => + new Response( + JSON.stringify({ error: { code: 'x', message: 'task is running' } }), + { status: 400 }, + ) + const adapter = createBytePlusVideo('seedance-1-5-pro-251215', 'k', { + fetch, + }) + await expect(adapter.cancelVideoJob('job-1')).rejects.toThrow( + /task is running/, + ) + }) +}) diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 381b9bdfe0..dce5dab219 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -660,6 +660,8 @@ export class ChatClient< // Flag to deduplicate continuation checks during action draining private continuationPending = false private subscriptionAbortController: AbortController | null = null + /** Aborts the `abortSignal` of running client tools on `stop()`. */ + private clientToolAbortController = new AbortController() private processingResolve: (() => void) | null = null /** * `connect()` adapters push the full HTTP body into the subscribe queue, then @@ -3103,7 +3105,15 @@ export class ChatClient< // This aborts the local request. A durable server run keeps going. this.continuationGeneration++ const hadLocalStream = this.abortController !== null + const stoppedRunId = hadLocalStream ? this.currentRunId : null + if (stoppedRunId) { + // Opt-in server cancel (`cancelUrl`). Best effort: the local stop is done. + const run = { runId: stoppedRunId, threadId: this.threadId } + void (async () => this.connection.cancelRun?.(run))().catch(() => {}) + } this.cancelInFlightStream({ setReadyStatus: true }) + this.clientToolAbortController.abort() + this.clientToolAbortController = new AbortController() this.discardPendingSends() this.lastResume = null this.activeInterruptSubmission = undefined @@ -3183,6 +3193,7 @@ export class ChatClient< toolCallId: args.toolCallId, threadId: this.threadId, context: context as TContext, + abortSignal: this.clientToolAbortController.signal, emitCustomEvent: () => {}, }) await this.addToolResultForClientTool( diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 3dcc9bb0ea..257a7023ff 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1014,6 +1014,14 @@ export interface ConnectConnectionAdapter { abortSignal: AbortSignal | undefined, runContext: RunAgentInputContext, ) => AsyncIterable + /** + * Optional. Ask the server to cancel a run. `ChatClient.stop()` calls it + * (best effort, not awaited) for the run that it stops. Other aborts, such + * as an unmount, do not call it, so a durable run keeps going for them. + * `fetchServerSentEvents` and `fetchHttpStream` provide it when you set + * `cancelUrl`. + */ + cancelRun?: (run: { runId: string; threadId: string }) => Promise /** * Fetch server-driven hydration for a generation `threadId`: the last * generation's resume snapshot, plus a cursor to a run still generating if @@ -1162,6 +1170,8 @@ export interface SubscribeConnectionAdapter { abortSignal: AbortSignal | undefined, runContext: RunAgentInputContext, ) => Promise + /** See {@link ConnectConnectionAdapter.cancelRun}. */ + cancelRun?: (run: { runId: string; threadId: string }) => Promise /** * Re-attach to an existing run by id, replaying its stream from the start off * the server's delivery-durability sink. Present only when the underlying @@ -1221,11 +1231,13 @@ export function normalizeConnectionAdapter( const hydrate = (connection as SubscribeConnectionAdapter).hydrate?.bind( connection, ) + const cancelRun = connection.cancelRun?.bind(connection) return { subscribe: connection.subscribe.bind(connection), send: connection.send.bind(connection), ...(joinRun ? { joinRun } : {}), ...(hydrate ? { hydrate } : {}), + ...(cancelRun ? { cancelRun } : {}), } } @@ -1375,6 +1387,9 @@ export function normalizeConnectionAdapter( } await waitUntilSubscriberIdle(abortSignal) }, + ...(connection.cancelRun + ? { cancelRun: connection.cancelRun.bind(connection) } + : {}), // Expose joinRun only when the underlying connection is resumable. Require // a real function — `'joinRun' in connection` is true for // `{ joinRun: undefined }`, which would wrap a non-callable and throw on @@ -1414,6 +1429,55 @@ export interface FetchConnectionOptions { fetchClient?: typeof globalThis.fetch /** Bounding for resumable-SSE reconnection (throttle delay + attempt ceiling). */ reconnect?: ReconnectOptions + /** + * Opt in to server cancel. `ChatClient.stop()` then sends + * `POST cancelUrl` with the JSON body `{ runId, threadId }` and the same + * `headers` and `credentials`. Your route aborts the run. + */ + cancelUrl?: string | (() => string) +} + +/** How long the cancel POST from `stop()` can take before it is aborted. */ +const CANCEL_RUN_TIMEOUT_MS = 10_000 + +/** + * `cancelRun` for the fetch adapters. Present only when `cancelUrl` can be + * set: a static options object with `cancelUrl`, or an options function + * (resolved on `stop()`; it does nothing when it returns no `cancelUrl`). + */ +function fetchCancelRun( + options: + | FetchConnectionOptions + | (() => FetchConnectionOptions | Promise), +): Pick { + if (typeof options !== 'function' && options.cancelUrl === undefined) { + return {} + } + return { + async cancelRun({ runId, threadId }) { + const resolved = typeof options === 'function' ? await options() : options + const { cancelUrl } = resolved + if (cancelUrl === undefined) return + const response = await (resolved.fetchClient ?? fetch)( + typeof cancelUrl === 'function' ? cancelUrl() : cancelUrl, + { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...mergeHeaders(resolved.headers), + }, + body: JSON.stringify({ runId, threadId }), + credentials: resolved.credentials || 'same-origin', + // `stop()` does not wait for this request. The deadline stops a + // `cancelUrl` that never answers from holding a connection. + signal: AbortSignal.timeout(CANCEL_RUN_TIMEOUT_MS), + }, + ) + if (!response.ok) { + throw new Error(`Run cancel failed with status ${response.status}`) + } + }, + } } /** @@ -1510,6 +1574,7 @@ export function fetchServerSentEvents( | (() => FetchConnectionOptions | Promise) = {}, ): ResumableConnectConnectionAdapter { return { + ...fetchCancelRun(options), async *connect(messages, data, abortSignal, runContext) { // Resolve URL and options if they are functions const resolvedUrl = typeof url === 'function' ? url() : url @@ -1670,6 +1735,7 @@ export function fetchHttpStream( | (() => FetchConnectionOptions | Promise) = {}, ): ResumableConnectConnectionAdapter { return { + ...fetchCancelRun(options), async *connect(messages, data, abortSignal, runContext) { // Resolve URL and options if they are functions const resolvedUrl = typeof url === 'function' ? url() : url diff --git a/packages/ai-client/src/realtime-client.ts b/packages/ai-client/src/realtime-client.ts index 9c2d0f7614..a7e750ff8b 100644 --- a/packages/ai-client/src/realtime-client.ts +++ b/packages/ai-client/src/realtime-client.ts @@ -49,9 +49,10 @@ const TOKEN_REFRESH_BUFFER_MS = 60_000 export class RealtimeClient { private readonly options: RealtimeClientOptions private connection: RealtimeConnection | null = null + /** Aborts the `adapter.connect()` setup when `disconnect()` runs. */ + private connectAbort: AbortController | 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() @@ -105,13 +106,14 @@ export class RealtimeClient { status: 'connecting', error: null, }) - // disconnect() bumps this. A connect() it overtook must stop. - const generation = ++this.connectGeneration + // A new connect() owns the setup. Stop the setup of an earlier one. + this.connectAbort?.abort() + const connectAbort = (this.connectAbort = new AbortController()) try { // Fetch token from server const token = await this.options.getToken() - if (generation !== this.connectGeneration) return + connectAbort.signal.throwIfAborted() this.token = token // Schedule token refresh @@ -126,12 +128,13 @@ export class RealtimeClient { const connection = await this.options.adapter.connect( this.token, toolsList, + { abortSignal: connectAbort.signal }, ) - if (generation !== this.connectGeneration) { + if (connectAbort.signal.aborted) { // A failed close of a connection that was never used is not an // error of the client. await connection.disconnect().catch(() => {}) - return + connectAbort.signal.throwIfAborted() } this.connection = connection @@ -153,7 +156,7 @@ export class RealtimeClient { // Start audio capture if configured if (this.options.autoCapture) { await connection.startAudioCapture() - if (generation !== this.connectGeneration) { + if (connectAbort.signal.aborted) { // 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) { @@ -161,7 +164,7 @@ export class RealtimeClient { // error of the live session. await connection.disconnect().catch(() => {}) } - return + connectAbort.signal.throwIfAborted() } } @@ -172,6 +175,8 @@ export class RealtimeClient { this.options.onConnect?.() } catch (error) { const err = error instanceof Error ? error : new Error(String(error)) + // `disconnect()` cancelled this connect: the state is already idle. + if (connectAbort.signal.aborted) throw err this.updateState({ status: 'error', error: err }) this.options.onError?.(err) throw err @@ -182,7 +187,7 @@ export class RealtimeClient { * Disconnect from the realtime session. */ async disconnect(): Promise { - this.connectGeneration++ + this.connectAbort?.abort() if (this.tokenRefreshTimeout) { clearTimeout(this.tokenRefreshTimeout) this.tokenRefreshTimeout = null @@ -470,11 +475,11 @@ export class RealtimeClient { } private async refreshToken(): Promise { - const generation = this.connectGeneration + const connectSignal = this.connectAbort?.signal try { const token = await this.options.getToken() // disconnect() ran during the refresh: do not re-arm the timer. - if (generation !== this.connectGeneration) return + if (connectSignal?.aborted) return this.token = token this.scheduleTokenRefresh() this.connection?.updateToken?.(this.token) diff --git a/packages/ai-client/tests/abort-signal-reach.test.ts b/packages/ai-client/tests/abort-signal-reach.test.ts index 6700f012e2..98c9cc57e8 100644 --- a/packages/ai-client/tests/abort-signal-reach.test.ts +++ b/packages/ai-client/tests/abort-signal-reach.test.ts @@ -296,6 +296,9 @@ describe('ChatClient.stop() (L1)', () => { ) }) +// An overtaken or cancelled connect() rejects with this. +const ABORT = { name: 'AbortError' } + describe('RealtimeClient.disconnect() during connect() (L4)', () => { // `gate` is the connect() step that waits while disconnect() runs. for (const gate of ['getToken', 'adapter.connect', 'startAudioCapture']) { @@ -330,7 +333,7 @@ describe('RealtimeClient.disconnect() during connect() (L4)', () => { await new Promise((resolve) => setTimeout(resolve, 0)) await client.disconnect() release() - await connecting.catch(() => {}) + await expect(connecting).rejects.toMatchObject(ABORT) const opened = adapterConnect.mock.calls.length > 0 // Either it never connects, or it closes what it opened. @@ -368,11 +371,13 @@ describe('RealtimeClient.disconnect() during connect() (L4)', () => { }) const connecting = client.connect() + // disconnect() cancels this connect(), so it rejects with AbortError. + const rejected = expect(connecting).rejects.toMatchObject(ABORT) await vi.waitFor(() => expect(client.status).toBe('connecting')) await new Promise((resolve) => setTimeout(resolve, 0)) await client.disconnect() release() - await connecting + await rejected expect(connection.disconnect).toHaveBeenCalled() expect(client.status).toBe('idle') @@ -405,11 +410,13 @@ describe('RealtimeClient.disconnect() during connect() (L4)', () => { }) const firstConnect = client.connect() + // The second connect() cancels the first, so it rejects with AbortError. + const rejected = expect(firstConnect).rejects.toMatchObject(ABORT) await vi.waitFor(() => expect(adapterConnect).toHaveBeenCalledTimes(1)) await new Promise((resolve) => setTimeout(resolve, 0)) await client.connect() release() - await firstConnect + await rejected expect(first.disconnect).toHaveBeenCalled() expect(second.disconnect).not.toHaveBeenCalled() @@ -465,12 +472,13 @@ describe('RealtimeClient keeps the live session when a second connect() wins', ( }) const firstConnect = client.connect() + const rejected = expect(firstConnect).rejects.toMatchObject(ABORT) await vi.waitFor(() => expect(Reflect.get(client, 'connection')).toBe(first), ) await client.connect() release() - await firstConnect + await rejected expect(first.disconnect).toHaveBeenCalled() expect(client.status).toBe('connected') @@ -546,10 +554,11 @@ describe('RealtimeClient.disconnect() during a token refresh', () => { }) const first = client.connect() + const rejected = expect(first).rejects.toMatchObject(ABORT) await vi.waitFor(() => expect(getToken).toHaveBeenCalledTimes(1)) await client.connect() release() - await first + await rejected await client.disconnect() expect(vi.getTimerCount()).toBe(0) diff --git a/packages/ai-client/tests/chat-client-cancel-run.test.ts b/packages/ai-client/tests/chat-client-cancel-run.test.ts new file mode 100644 index 0000000000..1ecb304dcc --- /dev/null +++ b/packages/ai-client/tests/chat-client-cancel-run.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it, vi } from 'vitest' +import { ChatClient } from '../src/chat-client' +import { + fetchHttpStream, + fetchServerSentEvents, +} from '../src/connection-adapters' + +/** + * A fetch that keeps the chat POST open until its signal aborts, and records + * every other request (the cancel POST) with a 204 reply. + */ +function fakeFetch(cancelStatus = 204) { + const chatSignals: Array = [] + const cancels: Array<{ url: string; init: RequestInit }> = [] + const fetchClient = vi.fn( + (input: RequestInfo | URL, init?: RequestInit): Promise => { + const url = String(input) + if (url.endsWith('/cancel')) { + cancels.push({ url, init: init ?? {} }) + return Promise.resolve(new Response(null, { status: cancelStatus })) + } + const signal = init?.signal + if (!signal) return Promise.reject(new Error('chat POST had no signal')) + chatSignals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }, + ) + return { fetchClient, chatSignals, cancels } +} + +describe('ChatClient.stop() with cancelUrl', () => { + it.each([ + ['fetchServerSentEvents', fetchServerSentEvents], + ['fetchHttpStream', fetchHttpStream], + ] as const)('%s posts { runId, threadId } to cancelUrl', async (_, make) => { + const { fetchClient, chatSignals, cancels } = fakeFetch() + const client = new ChatClient({ + threadId: 'thread-1', + connection: make('/api/chat', { + fetchClient, + cancelUrl: '/api/chat/cancel', + headers: { Authorization: 'Bearer t' }, + }), + }) + void client.sendMessage('hi') + await vi.waitFor(() => expect(chatSignals).toHaveLength(1)) + const runId = client.getCurrentRunId() + expect(runId).toEqual(expect.any(String)) + + client.stop() + + expect(chatSignals[0]!.aborted).toBe(true) + await vi.waitFor(() => expect(cancels).toHaveLength(1)) + const { init } = cancels[0]! + expect(init.method).toBe('POST') + expect(init.headers).toMatchObject({ Authorization: 'Bearer t' }) + expect(JSON.parse(String(init.body))).toEqual({ + runId, + threadId: 'thread-1', + }) + }) + + it('reads cancelUrl from an options function', async () => { + const { fetchClient, chatSignals, cancels } = fakeFetch() + const client = new ChatClient({ + threadId: 'thread-1', + connection: fetchServerSentEvents('/api/chat', () => ({ + fetchClient, + cancelUrl: () => '/api/chat/cancel', + })), + }) + void client.sendMessage('hi') + await vi.waitFor(() => expect(chatSignals).toHaveLength(1)) + client.stop() + await vi.waitFor(() => expect(cancels).toHaveLength(1)) + }) + + it('a failed cancel does not throw from stop()', async () => { + const { fetchClient, chatSignals, cancels } = fakeFetch(500) + const client = new ChatClient({ + threadId: 'thread-1', + connection: fetchServerSentEvents('/api/chat', { + fetchClient, + cancelUrl: '/api/chat/cancel', + }), + }) + void client.sendMessage('hi') + await vi.waitFor(() => expect(chatSignals).toHaveLength(1)) + expect(() => client.stop()).not.toThrow() + await vi.waitFor(() => expect(cancels).toHaveLength(1)) + expect(client.getSnapshot().status).toBe('ready') + }) + + it('an unmount (unsubscribe) does not cancel the run', async () => { + const { fetchClient, chatSignals, cancels } = fakeFetch() + const client = new ChatClient({ + threadId: 'thread-1', + connection: fetchServerSentEvents('/api/chat', { + fetchClient, + cancelUrl: '/api/chat/cancel', + }), + }) + void client.sendMessage('hi') + await vi.waitFor(() => expect(chatSignals).toHaveLength(1)) + client.unsubscribe() + expect(chatSignals[0]!.aborted).toBe(true) + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(cancels).toHaveLength(0) + }) + + it('without cancelUrl, stop() sends no extra request', async () => { + const { fetchClient, chatSignals } = fakeFetch() + const connection = fetchServerSentEvents('/api/chat', { fetchClient }) + expect(connection.cancelRun).toBeUndefined() + const client = new ChatClient({ threadId: 'thread-1', connection }) + void client.sendMessage('hi') + await vi.waitFor(() => expect(chatSignals).toHaveLength(1)) + client.stop() + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(fetchClient).toHaveBeenCalledTimes(1) + }) +}) + +describe('cancelRun request deadline', () => { + it('sends the cancel POST with a signal, so a dead cancelUrl cannot hold it', async () => { + const { fetchClient, cancels } = fakeFetch() + const connection = fetchServerSentEvents('/api/chat', { + fetchClient, + cancelUrl: '/api/chat/cancel', + }) + await connection.cancelRun?.({ runId: 'run-1', threadId: 'thread-1' }) + expect(cancels).toHaveLength(1) + expect(cancels[0]!.init.signal).toBeInstanceOf(AbortSignal) + expect(cancels[0]!.init.signal?.aborted).toBe(false) + }) +}) diff --git a/packages/ai-client/tests/chat-client-tool-abort-signal.test.ts b/packages/ai-client/tests/chat-client-tool-abort-signal.test.ts new file mode 100644 index 0000000000..b1bd9153e7 --- /dev/null +++ b/packages/ai-client/tests/chat-client-tool-abort-signal.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it, vi } from 'vitest' +import { ChatClient } from '../src/chat-client' +import { createToolCallChunks } from './test-utils' +import type { ConnectConnectionAdapter } from '../src/connection-adapters' + +describe('client tool abortSignal', () => { + function setup() { + const signals: Array = [] + let call = 0 + const adapter: ConnectConnectionAdapter = { + async *connect() { + call++ + yield* createToolCallChunks([ + { id: `tc-${call}`, name: 'slow_tool', arguments: '{}' }, + ]) + }, + } + const client = new ChatClient({ + threadId: 'thread-1', + connection: adapter, + tools: [ + { + __toolSide: 'client' as const, + name: 'slow_tool', + description: 'Waits until it is cancelled', + execute: ( + _input: unknown, + context?: { abortSignal?: AbortSignal }, + ) => { + const signal = context?.abortSignal + signals.push(signal) + return new Promise((_, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }, + }, + ], + }) + return { client, signals } + } + + it('aborts the signal of a running client tool on stop()', async () => { + const { client, signals } = setup() + void client.sendMessage('go') + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]?.aborted).toBe(false) + client.stop() + expect(signals[0]?.aborted).toBe(true) + }) + + it('gives a later tool call a fresh signal after stop()', async () => { + const { client, signals } = setup() + void client.sendMessage('go') + await vi.waitFor(() => expect(signals).toHaveLength(1)) + client.stop() + void client.sendMessage('again') + await vi.waitFor(() => expect(signals).toHaveLength(2)) + expect(signals[1]?.aborted).toBe(false) + client.stop() + }) +}) diff --git a/packages/ai-client/tests/realtime-client-connect-signal.test.ts b/packages/ai-client/tests/realtime-client-connect-signal.test.ts new file mode 100644 index 0000000000..289ef2fc94 --- /dev/null +++ b/packages/ai-client/tests/realtime-client-connect-signal.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, it, vi } from 'vitest' +import { RealtimeClient } from '../src' +import type { RealtimeAdapter } from '@tanstack/ai/client' + +const token = { + provider: 'test', + token: 't', + expiresAt: Date.now() + 3_600_000, + config: {}, +} + +/** An adapter whose connect waits until its `abortSignal` aborts. */ +function hangingAdapter() { + const signals: Array = [] + const adapter: RealtimeAdapter = { + provider: 'test', + connect: (_token, _tools, options) => { + const signal = options?.abortSignal + signals.push(signal) + if (!signal) return Promise.reject(new Error('connect had no signal')) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }, + } + return { adapter, signals } +} + +describe('RealtimeClient connect abortSignal', () => { + it('disconnect() aborts a connect that is in progress', async () => { + const { adapter, signals } = hangingAdapter() + const onError = vi.fn() + const client = new RealtimeClient({ + getToken: async () => token, + adapter, + onError, + }) + + // Handle the rejection now: it comes while the test still awaits. + const pending = expect(client.connect()).rejects.toMatchObject({ + name: 'AbortError', + }) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + + await client.disconnect() + expect(signals[0]!.aborted).toBe(true) + await pending + expect(client.getSnapshot().status).toBe('idle') + expect(onError).not.toHaveBeenCalled() + client.destroy() + }) + + it('a second connect() aborts the setup of the first', async () => { + const { adapter, signals } = hangingAdapter() + const onError = vi.fn() + const client = new RealtimeClient({ + getToken: async () => token, + adapter, + onError, + }) + + // Handle each rejection now: it comes while the test still awaits. + const first = expect(client.connect()).rejects.toMatchObject({ + name: 'AbortError', + }) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + const second = expect(client.connect()).rejects.toMatchObject({ + name: 'AbortError', + }) + await vi.waitFor(() => expect(signals).toHaveLength(2)) + + expect(signals[0]!.aborted).toBe(true) + expect(signals[1]!.aborted).toBe(false) + await first + expect(onError).not.toHaveBeenCalled() + + await client.disconnect() + await second + client.destroy() + }) + + it('a connect error that is not an abort still sets the error state', async () => { + const client = new RealtimeClient({ + getToken: async () => token, + adapter: { + provider: 'test', + connect: async () => { + throw new Error('refused') + }, + }, + }) + await expect(client.connect()).rejects.toThrow('refused') + expect(client.getSnapshot().status).toBe('error') + client.destroy() + }) +}) diff --git a/packages/ai-cloudflare/src/adapters/embedding.ts b/packages/ai-cloudflare/src/adapters/embedding.ts index 715a2c5b06..cd66449d60 100644 --- a/packages/ai-cloudflare/src/adapters/embedding.ts +++ b/packages/ai-cloudflare/src/adapters/embedding.ts @@ -43,10 +43,12 @@ export class CloudflareEmbeddingAdapter< `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, { provider: this.name, model }, ) - const output = (await runModel(this.cfConfig, model, { - ...options.modelOptions, - text: texts, - })) as { data?: Array> } + const output = (await runModel( + this.cfConfig, + model, + { ...options.modelOptions, text: texts }, + { signal: options.abortSignal }, + )) as { data?: Array> } if (!Array.isArray(output.data) || output.data.length !== texts.length) { throw new Error( `Workers AI ${model} returned ${output.data?.length ?? 0} embeddings for ${texts.length} inputs`, diff --git a/packages/ai-cloudflare/tests/abort-signal.test.ts b/packages/ai-cloudflare/tests/abort-signal.test.ts new file mode 100644 index 0000000000..b3691161ac --- /dev/null +++ b/packages/ai-cloudflare/tests/abort-signal.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createCloudflareEmbedding } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('cloudflare adapters forward abortSignal to the request', () => { + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + const adapter = createCloudflareEmbedding('@cf/baai/bge-m3', { + accountId: 'acc', + apiKey: 'k', + fetch, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: '@cf/baai/bge-m3', + input: ['a'], + logger, + abortSignal, + }), + ) + }) +}) diff --git a/packages/ai-cohere/src/adapters/embedding.ts b/packages/ai-cohere/src/adapters/embedding.ts index 142eb06ba8..9d457de1db 100644 --- a/packages/ai-cohere/src/adapters/embedding.ts +++ b/packages/ai-cohere/src/adapters/embedding.ts @@ -66,10 +66,16 @@ async function fetchWithTimeout( ): Promise { const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), timeoutMs) + // Also abort when the caller's signal aborts. + const callerSignal = init?.signal + const onAbort = () => controller.abort(callerSignal?.reason) + if (callerSignal?.aborted) onAbort() + else callerSignal?.addEventListener('abort', onAbort, { once: true }) try { return await fetch(url, { ...init, signal: controller.signal }) } finally { clearTimeout(timeoutId) + callerSignal?.removeEventListener('abort', onAbort) } } @@ -158,7 +164,9 @@ export class CohereEmbeddingAdapter< for (const image of item.images) { content.push({ type: 'image_url', - image_url: { url: await this.resolveImageUrl(image) }, + image_url: { + url: await this.resolveImageUrl(image, options.abortSignal), + }, }) } return { content } @@ -198,6 +206,7 @@ export class CohereEmbeddingAdapter< ...transport.headers, }, body: JSON.stringify(body), + ...(options.abortSignal && { signal: options.abortSignal }), }, timeoutMs, ) @@ -266,7 +275,10 @@ export class CohereEmbeddingAdapter< * remote image URLs, so everything is normalized to a `data:` URI unless * the caller already provided one. */ - protected async resolveImageUrl(image: ImagePart): Promise { + protected async resolveImageUrl( + image: ImagePart, + signal?: AbortSignal, + ): Promise { const source = image.source if (isFileSource(source)) throw unsupportedFileSourceError(this.name) @@ -292,7 +304,7 @@ export class CohereEmbeddingAdapter< const response = await fetchWithTimeout( source.value, - undefined, + signal ? { signal } : undefined, this.clientConfig.timeout ?? DEFAULT_TIMEOUT_MS, ) if (!response.ok) { diff --git a/packages/ai-cohere/tests/abort-signal.test.ts b/packages/ai-cohere/tests/abort-signal.test.ts new file mode 100644 index 0000000000..72f20cbfb1 --- /dev/null +++ b/packages/ai-cohere/tests/abort-signal.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createCohereEmbedding } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('cohere adapters forward abortSignal to the request', () => { + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + try { + const adapter = createCohereEmbedding('embed-v4.0', 'k') + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'embed-v4.0', + input: ['a'], + modelOptions: { inputType: 'search_document' }, + logger, + abortSignal, + }), + ) + } finally { + vi.unstubAllGlobals() + } + }) +}) diff --git a/packages/ai-elevenlabs/src/realtime/token.ts b/packages/ai-elevenlabs/src/realtime/token.ts index 58d2345a63..c97e73ac78 100644 --- a/packages/ai-elevenlabs/src/realtime/token.ts +++ b/packages/ai-elevenlabs/src/realtime/token.ts @@ -38,12 +38,15 @@ export function elevenlabsRealtimeToken( return { provider: 'elevenlabs', - async generateToken(): Promise { + async generateToken(request?: { + abortSignal?: AbortSignal + }): Promise { const { overrides } = options const agentId = options.agentId ?? getElevenLabsAgentIdFromEnv() const response = await client.conversationalAi.conversations.getSignedUrl( { agentId }, + request?.abortSignal && { abortSignal: request.abortSignal }, ) // Signed URLs are valid for 30 minutes diff --git a/packages/ai-elevenlabs/tests/abort-signal.test.ts b/packages/ai-elevenlabs/tests/abort-signal.test.ts new file mode 100644 index 0000000000..074ad2c9b0 --- /dev/null +++ b/packages/ai-elevenlabs/tests/abort-signal.test.ts @@ -0,0 +1,54 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { elevenlabsRealtimeToken } from '../src/realtime/token' + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (input: unknown, init?: RequestInit): Promise => { + // Some SDKs pass a `Request`, others pass `init`. + const signal = + init?.signal ?? (input instanceof Request ? input.signal : undefined) + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('elevenlabs realtime token forwards abortSignal', () => { + afterEach(() => { + vi.unstubAllGlobals() + vi.unstubAllEnvs() + }) + + it('generateToken', async () => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + vi.stubEnv('ELEVENLABS_API_KEY', 'k') + const adapter = elevenlabsRealtimeToken({ agentId: 'agent-1' }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.generateToken({ abortSignal }), + ) + }) +}) diff --git a/packages/ai-fal/src/adapters/video.ts b/packages/ai-fal/src/adapters/video.ts index b996168a6e..f123cd794a 100644 --- a/packages/ai-fal/src/adapters/video.ts +++ b/packages/ai-fal/src/adapters/video.ts @@ -210,10 +210,28 @@ export class FalVideoAdapter extends BaseVideoAdapter< return getFalVideoDurationOptions(this.model) } - async getVideoStatus(jobId: string): Promise { + /** + * Cancels a queued request with `fal.queue.cancel`. fal can refuse when the + * request already runs, and then this rejects. + */ + async cancelVideoJob( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + await fal.queue.cancel(this.model, { + requestId: jobId, + ...(options?.abortSignal ? { abortSignal: options.abortSignal } : {}), + }) + } + + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { const statusResponse = (await fal.queue.status(this.model, { requestId: jobId, logs: true, + ...(options?.abortSignal ? { abortSignal: options.abortSignal } : {}), })) as FalStatusResponse return { @@ -227,13 +245,19 @@ export class FalVideoAdapter extends BaseVideoAdapter< } } - override async getVideo(jobId: string): Promise { + override async getVideo( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { let result try { result = await fal.queue.result(this.model, { requestId: jobId, + ...(options?.abortSignal ? { abortSignal: options.abortSignal } : {}), }) } catch (error: unknown) { + // A cancelled call is not a provider failure: keep the abort error. + if (options?.abortSignal?.aborted) throw error // fal.ai may report COMPLETED status but throw on result fetch // (e.g. 422 validation errors). Extract the detailed error info. const err = error as { body?: { detail?: unknown }; message?: string } diff --git a/packages/ai-fal/tests/video-abort-signal.test.ts b/packages/ai-fal/tests/video-abort-signal.test.ts new file mode 100644 index 0000000000..01d675607c --- /dev/null +++ b/packages/ai-fal/tests/video-abort-signal.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it, vi } from 'vitest' +import { falVideo } from '../src/adapters/video' + +/** Records the queue options; hangs until their `abortSignal` aborts. */ +const mocks = vi.hoisted(() => { + const signals: Array = [] + const hang = (_endpoint: string, options: { abortSignal?: AbortSignal }) => { + const signal = options.abortSignal + signals.push(signal) + if (!signal) return Promise.reject(new Error('request had no signal')) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { signals, hang } +}) + +vi.mock('@fal-ai/client', () => ({ + fal: { + queue: { status: mocks.hang, result: mocks.hang, cancel: mocks.hang }, + config: () => {}, + }, +})) + +describe('fal video forwards abortSignal to the queue', () => { + it.each(['getVideoStatus', 'getVideo', 'cancelVideoJob'] as const)( + '%s', + async (method) => { + mocks.signals.length = 0 + const adapter = falVideo('fal-ai/veo3/image-to-video', { apiKey: 'k' }) + const controller = new AbortController() + const pending = adapter[method]('job-1', { + abortSignal: controller.signal, + }) + await vi.waitFor(() => expect(mocks.signals).toHaveLength(1)) + expect(mocks.signals[0]).toBe(controller.signal) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + }, + ) +}) diff --git a/packages/ai-gemini/src/adapters/embedding.ts b/packages/ai-gemini/src/adapters/embedding.ts index f73a80c350..3ac5f9af62 100644 --- a/packages/ai-gemini/src/adapters/embedding.ts +++ b/packages/ai-gemini/src/adapters/embedding.ts @@ -69,6 +69,9 @@ export class GeminiEmbeddingAdapter< if (modelOptions?.title !== undefined) { config.title = modelOptions.title } + if (options.abortSignal) { + config.abortSignal = options.abortSignal + } logger.request( `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, diff --git a/packages/ai-gemini/src/adapters/files.ts b/packages/ai-gemini/src/adapters/files.ts index 98d784f9de..8a202dc240 100644 --- a/packages/ai-gemini/src/adapters/files.ts +++ b/packages/ai-gemini/src/adapters/files.ts @@ -24,21 +24,40 @@ export class GeminiFilesAdapter extends BaseFilesAdapter<'gemini'> { this.client = createGeminiClient(config) } - async upload(input: FileUploadInput): Promise> { + async upload( + input: FileUploadInput, + options?: { abortSignal?: AbortSignal }, + ): Promise> { const { blob, mimeType } = normalizeFileUploadInput(input) + const abortSignal = options?.abortSignal const file = await this.client.files.upload({ file: blob, - ...(mimeType ? { config: { mimeType } } : {}), + ...(mimeType || abortSignal + ? { + config: { + ...(mimeType ? { mimeType } : {}), + ...(abortSignal ? { abortSignal } : {}), + }, + } + : {}), }) return toFileHandle(file) } - async get(id: string): Promise> { - return toFileHandle(await this.client.files.get({ name: id })) + async get( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise> { + return toFileHandle( + await this.client.files.get({ name: id, ...signalConfig(options) }), + ) } - async delete(id: string): Promise { - await this.client.files.delete({ name: id }) + async delete( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + await this.client.files.delete({ name: id, ...signalConfig(options) }) } } @@ -82,3 +101,10 @@ export function geminiFiles( ): GeminiFilesAdapter { return createGeminiFiles(getGeminiApiKeyFromEnv(), config) } + +/** SDK config that carries the caller's abort signal, if any. */ +function signalConfig(options?: { abortSignal?: AbortSignal }) { + return options?.abortSignal + ? { config: { abortSignal: options.abortSignal } } + : {} +} diff --git a/packages/ai-gemini/src/adapters/video.ts b/packages/ai-gemini/src/adapters/video.ts index 538f2fd5bb..3d53f81560 100644 --- a/packages/ai-gemini/src/adapters/video.ts +++ b/packages/ai-gemini/src/adapters/video.ts @@ -535,11 +535,15 @@ export class GeminiVideoAdapter< return { image, lastFrame, referenceImages } } - async getVideoStatus(jobId: string): Promise { + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const signal = options?.abortSignal if (isInteractionsVideoModel(this.model)) { - return await this.getInteractionsVideoStatus(jobId) + return await this.getInteractionsVideoStatus(jobId, signal) } - const operation = await this.getOperation(jobId) + const operation = await this.getOperation(jobId, signal) if (!operation.done) { return { jobId, status: 'processing' } @@ -583,8 +587,9 @@ export class GeminiVideoAdapter< */ private async getInteractionsVideoStatus( jobId: string, + signal?: AbortSignal, ): Promise { - const interaction = await this.getInteraction(jobId) + const interaction = await this.getInteraction(jobId, signal) const status = interaction.status if (status === 'in_progress') { @@ -616,11 +621,15 @@ export class GeminiVideoAdapter< } } - override async getVideo(jobId: string): Promise { + override async getVideo( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const signal = options?.abortSignal if (isInteractionsVideoModel(this.model)) { - return await this.getInteractionsVideoUrl(jobId) + return await this.getInteractionsVideoUrl(jobId, signal) } - const operation = await this.getOperation(jobId) + const operation = await this.getOperation(jobId, signal) if (!operation.done) { throw new Error( @@ -656,8 +665,9 @@ export class GeminiVideoAdapter< */ private async getInteractionsVideoUrl( jobId: string, + signal?: AbortSignal, ): Promise { - const interaction = await this.getInteraction(jobId) + const interaction = await this.getInteraction(jobId, signal) const status = interaction.status if (status === 'in_progress') { @@ -695,17 +705,28 @@ export class GeminiVideoAdapter< * instance (it calls `_fromAPIResponse` on it), so reconstruct one from * the job ID rather than passing an object literal. */ - private async getOperation(jobId: string): Promise { + private async getOperation( + jobId: string, + signal?: AbortSignal, + ): Promise { const operation = new GenerateVideosOperation() operation.name = jobId - return await this.client.operations.getVideosOperation({ operation }) + return await this.client.operations.getVideosOperation({ + operation, + ...(signal && { config: { abortSignal: signal } }), + }) } /** * Fetch an Omni background interaction by id. */ - private async getInteraction(jobId: string): Promise { - return await this.client.interactions.get(jobId) + private async getInteraction( + jobId: string, + signal?: AbortSignal, + ): Promise { + return signal + ? await this.client.interactions.get(jobId, undefined, { signal }) + : await this.client.interactions.get(jobId) } } diff --git a/packages/ai-gemini/src/files/index.ts b/packages/ai-gemini/src/files/index.ts index 2190c337a4..6bc420ee7f 100644 --- a/packages/ai-gemini/src/files/index.ts +++ b/packages/ai-gemini/src/files/index.ts @@ -33,6 +33,12 @@ export interface GeminiUploadFileOptions { pollIntervalMs?: number /** Max time to wait for processing, in ms. Default `300000` (5 min). */ timeoutMs?: number + /** + * Cancels the status requests and the wait between polls. The call + * rejects when the signal aborts. The Gemini SDK does not cancel the + * upload request itself. + */ + abortSignal?: AbortSignal } /** @@ -54,13 +60,19 @@ export async function uploadGeminiFile( mimeType, pollIntervalMs = 5000, timeoutMs = 300_000, + abortSignal, } = options const client = createGeminiClient({ apiKey }) let uploaded = await client.files.upload({ file, - ...(mimeType && { config: { mimeType } }), + ...((mimeType || abortSignal) && { + config: { + ...(mimeType && { mimeType }), + ...(abortSignal && { abortSignal }), + }, + }), }) const fileName = uploaded.name @@ -75,8 +87,11 @@ export async function uploadGeminiFile( `Gemini file processing timed out after ${timeoutMs}ms (${fileName}).`, ) } - await new Promise((resolve) => setTimeout(resolve, pollIntervalMs)) - uploaded = await client.files.get({ name: fileName }) + await sleep(pollIntervalMs, abortSignal) + uploaded = await client.files.get({ + name: fileName, + ...(abortSignal && { config: { abortSignal } }), + }) } if (uploaded.state === FileState.FAILED) { @@ -112,3 +127,19 @@ export function geminiVideoPart( ...(metadata && { metadata }), } } + +/** Wait `ms`, or reject with the abort reason when `signal` aborts first. */ +function sleep(ms: number, signal: AbortSignal | undefined): Promise { + return new Promise((resolve, reject) => { + signal?.throwIfAborted() + const onAbort = () => { + clearTimeout(timer) + reject(signal?.reason) + } + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort) + resolve() + }, ms) + signal?.addEventListener('abort', onAbort, { once: true }) + }) +} diff --git a/packages/ai-gemini/src/realtime/token.ts b/packages/ai-gemini/src/realtime/token.ts index fbb978047f..01e707f6db 100644 --- a/packages/ai-gemini/src/realtime/token.ts +++ b/packages/ai-gemini/src/realtime/token.ts @@ -37,7 +37,9 @@ export function geminiRealtimeToken( return { provider: 'gemini', - async generateToken(): Promise { + async generateToken(request?: { + abortSignal?: AbortSignal + }): Promise { // Computed per call so a reused adapter doesn't mint tokens with a // fixed (and eventually past) expiry. Defaults to 30 minutes. const expireTime = options.expiresAt ?? Date.now() + 30 * 60 * 1000 @@ -50,6 +52,7 @@ export function geminiRealtimeToken( httpOptions: { apiVersion: 'v1alpha', }, + ...(request?.abortSignal && { abortSignal: request.abortSignal }), }, }) diff --git a/packages/ai-gemini/tests/abort-signal.test.ts b/packages/ai-gemini/tests/abort-signal.test.ts new file mode 100644 index 0000000000..d08ef15eec --- /dev/null +++ b/packages/ai-gemini/tests/abort-signal.test.ts @@ -0,0 +1,98 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createGeminiEmbedding } from '../src/adapters/embedding' +import { createGeminiVideo } from '../src/adapters/video' +import { createGeminiFiles } from '../src/adapters/files' +import { geminiRealtimeToken } from '../src/realtime/token' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (input: unknown, init?: RequestInit): Promise => { + // The Interactions client passes a `Request`, the rest pass `init`. + const signal = + init?.signal ?? (input instanceof Request ? input.signal : undefined) + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('gemini adapters forward abortSignal to the request', () => { + let signals: Array + + beforeEach(() => { + const hanging = hangingFetch() + signals = hanging.signals + vi.stubGlobal('fetch', hanging.fetch) + vi.stubEnv('GOOGLE_API_KEY', 'k') + }) + + afterEach(() => { + vi.unstubAllGlobals() + vi.unstubAllEnvs() + }) + + it('embedding', async () => { + const adapter = createGeminiEmbedding('gemini-embedding-001', 'k') + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input: ['a'], + logger, + abortSignal, + }), + ) + }) + + it.each([ + ['veo', 'veo-3.1-generate-preview'], + ['omni', 'gemini-omni-1.1-flash'], + ] as const)('video %s status and download', async (_label, model) => { + for (const method of ['getVideoStatus', 'getVideo'] as const) { + signals.length = 0 + const adapter = createGeminiVideo(model, 'k') + await expectAbortReachesFetch(signals, (abortSignal) => + adapter[method]('operations/job-1', { abortSignal }), + ) + } + }) + + // `upload` is not here: the Gemini SDK does not send a signal for uploads. + it.each(['get', 'delete'] as const)('files %s', async (method) => { + const files = createGeminiFiles('k') + await expectAbortReachesFetch(signals, (abortSignal) => + files[method]('files/abc', { abortSignal }), + ) + }) + + it('realtime token', async () => { + await expectAbortReachesFetch(signals, (abortSignal) => + geminiRealtimeToken().generateToken({ abortSignal }), + ) + }) +}) diff --git a/packages/ai-gemini/tests/upload-gemini-file-abort.test.ts b/packages/ai-gemini/tests/upload-gemini-file-abort.test.ts new file mode 100644 index 0000000000..e8260dc26c --- /dev/null +++ b/packages/ai-gemini/tests/upload-gemini-file-abort.test.ts @@ -0,0 +1,79 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { uploadGeminiFile } from '../src/files' +import type * as Utils from '../src/utils' + +const mocks = vi.hoisted(() => ({ upload: vi.fn(), get: vi.fn() })) + +vi.mock('../src/utils', async (importOriginal) => ({ + ...(await importOriginal()), + createGeminiClient: () => ({ + files: { upload: mocks.upload, get: mocks.get }, + }), +})) + +/** Resolves never; rejects with the reason when `signal` aborts. */ +function untilAbort(signal: AbortSignal | undefined): Promise { + return new Promise((_, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) +} + +describe('uploadGeminiFile abortSignal', () => { + beforeEach(() => { + mocks.upload.mockReset() + mocks.get.mockReset() + mocks.upload.mockResolvedValue({ name: 'files/a', state: 'PROCESSING' }) + mocks.get.mockImplementation( + (params: { config?: { abortSignal?: AbortSignal } }) => + untilAbort(params.config?.abortSignal), + ) + }) + + it('stops the wait between polls', async () => { + const controller = new AbortController() + const pending = uploadGeminiFile(new Blob(['x']), { + apiKey: 'k', + mimeType: 'video/mp4', + pollIntervalMs: 60_000, + abortSignal: controller.signal, + }) + await vi.waitFor(() => expect(mocks.upload).toHaveBeenCalled()) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(mocks.get).not.toHaveBeenCalled() + }) + + it('passes the signal to the status request', async () => { + const controller = new AbortController() + const pending = uploadGeminiFile(new Blob(['x']), { + apiKey: 'k', + mimeType: 'video/mp4', + pollIntervalMs: 1, + abortSignal: controller.signal, + }) + await vi.waitFor(() => expect(mocks.get).toHaveBeenCalled()) + expect(mocks.get.mock.calls[0]![0].config.abortSignal).toBe( + controller.signal, + ) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + }) + + it('sends no signal when none is set', async () => { + mocks.upload.mockResolvedValue({ + name: 'files/a', + state: 'ACTIVE', + uri: 'https://x/files/a', + }) + await uploadGeminiFile(new Blob(['x']), { + apiKey: 'k', + mimeType: 'video/mp4', + }) + expect(mocks.upload).toHaveBeenCalledWith({ + file: expect.any(Blob), + config: { mimeType: 'video/mp4' }, + }) + }) +}) diff --git a/packages/ai-grok/src/adapters/files.ts b/packages/ai-grok/src/adapters/files.ts index 39e8ddfa99..fd763a1430 100644 --- a/packages/ai-grok/src/adapters/files.ts +++ b/packages/ai-grok/src/adapters/files.ts @@ -68,15 +68,19 @@ export class GrokFilesAdapter extends BaseFilesAdapter<'grok'> { this.expiresAfter = expiresAfter } - async upload(input: FileUploadInput): Promise> { + async upload( + input: FileUploadInput, + options?: { abortSignal?: AbortSignal }, + ): Promise> { + const signal = options?.abortSignal const { blob, mimeType, filename } = normalizeFileUploadInput(input) const file = await toFile(blob, filename, { ...(mimeType ? { type: mimeType } : {}), }) - const uploaded = await this.client.files.create({ - file, - purpose: 'assistants', - }) + const uploaded = await this.client.files.create( + { file, purpose: 'assistants' }, + signal && { signal }, + ) // Minting the public URL is idempotent: asking again for a file that // already has one returns the same URL. @@ -87,6 +91,7 @@ export class GrokFilesAdapter extends BaseFilesAdapter<'grok'> { this.expiresAfter !== undefined ? { expires_after: this.expiresAfter } : {}, + ...(signal && { signal }), }, ) if (!publicUrl.public_url) { @@ -116,13 +121,17 @@ export class GrokFilesAdapter extends BaseFilesAdapter<'grok'> { * it. After {@link revokePublicUrl}, a `get()` gives the file a public URL * again. Do not call it for a file whose URL you revoked. */ - async get(id: string): Promise> { - const file = await this.client.files.retrieve(id) + async get( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise> { + const signal = options?.abortSignal + const file = await this.client.files.retrieve(id, signal && { signal }) // `retrieve` reports the stored object, not its public URL. Re-minting is // idempotent, so this returns the URL the upload already created. const publicUrl = await this.client.post( `/files/${id}/public-url`, - { body: {} }, + { body: {}, ...(signal && { signal }) }, ) return { id: file.id, @@ -136,8 +145,12 @@ export class GrokFilesAdapter extends BaseFilesAdapter<'grok'> { } } - async delete(id: string): Promise { - await this.client.files.delete(id) + async delete( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const signal = options?.abortSignal + await this.client.files.delete(id, signal && { signal }) } /** diff --git a/packages/ai-grok/src/adapters/video.ts b/packages/ai-grok/src/adapters/video.ts index 508956ae2e..9dad696464 100644 --- a/packages/ai-grok/src/adapters/video.ts +++ b/packages/ai-grok/src/adapters/video.ts @@ -544,8 +544,14 @@ export class GrokVideoAdapter< } } - private async retrieveJob(jobId: string): Promise { - const response = await this.request(`/videos/${jobId}`) + private async retrieveJob( + jobId: string, + signal?: AbortSignal, + ): Promise { + const response = await this.request( + `/videos/${jobId}`, + signal && { signal }, + ) if (!response.ok) { const error = new Error( `grok: video status request failed (${response.status} ${response.statusText}): ${await this.errorMessage(response)}`, @@ -556,10 +562,13 @@ export class GrokVideoAdapter< return (await response.json()) as GrokVideoStatusResponse } - async getVideoStatus(jobId: string): Promise { + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { let response: GrokVideoStatusResponse try { - response = await this.retrieveJob(jobId) + response = await this.retrieveJob(jobId, options?.abortSignal) } catch (error) { if ((error as { status?: number }).status === 404) { return { jobId, status: 'failed', error: 'Job not found' } @@ -575,10 +584,13 @@ export class GrokVideoAdapter< } } - override async getVideo(jobId: string): Promise { + override async getVideo( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { let response: GrokVideoStatusResponse try { - response = await this.retrieveJob(jobId) + response = await this.retrieveJob(jobId, options?.abortSignal) } catch (error) { if ((error as { status?: number }).status === 404) { throw new Error(`Video job not found: ${jobId}`) diff --git a/packages/ai-grok/src/realtime/adapter.ts b/packages/ai-grok/src/realtime/adapter.ts index a26456d6cc..f7762ed27a 100644 --- a/packages/ai-grok/src/realtime/adapter.ts +++ b/packages/ai-grok/src/realtime/adapter.ts @@ -89,6 +89,7 @@ export function grokRealtime( async connect( token: RealtimeToken, _clientTools?: ReadonlyArray, + options?: { abortSignal?: AbortSignal }, ): Promise { const model = token.config.model ?? GROK_DEFAULT_REALTIME_MODEL logger.request(`activity=realtime provider=grok model=${model}`, { @@ -97,7 +98,7 @@ export function grokRealtime( }) if (connectionMode === 'webrtc') { - return createWebRTCConnection(token, logger) + return createWebRTCConnection(token, logger, options?.abortSignal) } const error = new Error('WebSocket connection mode not yet implemented') logger.errors('grok.realtime fatal', { @@ -115,6 +116,7 @@ export function grokRealtime( async function createWebRTCConnection( token: RealtimeToken, logger: InternalLogger, + abortSignal?: AbortSignal, ): Promise { const model = token.config.model ?? GROK_DEFAULT_REALTIME_MODEL const eventHandlers = new Map>>() @@ -498,6 +500,7 @@ async function createWebRTCConnection( 'Content-Type': 'application/sdp', }, body: offer.sdp, + ...(abortSignal && { signal: abortSignal }), }) if (!sdpResponse.ok) { diff --git a/packages/ai-grok/src/realtime/token.ts b/packages/ai-grok/src/realtime/token.ts index b574c50888..a4eaabab3c 100644 --- a/packages/ai-grok/src/realtime/token.ts +++ b/packages/ai-grok/src/realtime/token.ts @@ -41,7 +41,9 @@ export function grokRealtimeToken( return { provider: 'grok', - async generateToken(): Promise { + async generateToken(request?: { + abortSignal?: AbortSignal + }): Promise { const model: GrokRealtimeModel = options.model ?? GROK_DEFAULT_REALTIME_MODEL @@ -67,6 +69,11 @@ export function grokRealtimeToken( controller.abort(new Error('Grok realtime token request timed out')), DEFAULT_TOKEN_FETCH_TIMEOUT_MS, ) + // Also abort when the caller's signal aborts. + const callerSignal = request?.abortSignal + const onAbort = () => controller.abort(callerSignal?.reason) + if (callerSignal?.aborted) onAbort() + else callerSignal?.addEventListener('abort', onAbort, { once: true }) try { const response = await fetch(GROK_REALTIME_CLIENT_SECRETS_URL, { @@ -127,6 +134,7 @@ export function grokRealtimeToken( throw error } finally { clearTimeout(timeout) + callerSignal?.removeEventListener('abort', onAbort) } }, } diff --git a/packages/ai-grok/tests/abort-signal.test.ts b/packages/ai-grok/tests/abort-signal.test.ts new file mode 100644 index 0000000000..4cb950b4f1 --- /dev/null +++ b/packages/ai-grok/tests/abort-signal.test.ts @@ -0,0 +1,74 @@ +import { describe, expect, it, vi } from 'vitest' +import { createGrokVideo } from '../src/adapters/video' +import { createGrokFiles } from '../src/adapters/files' +import { grokRealtimeToken } from '../src/realtime/token' + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('grok adapters forward abortSignal to the request', () => { + it.each(['getVideoStatus', 'getVideo'] as const)( + 'video %s', + async (method) => { + const { fetch, signals } = hangingFetch() + const adapter = createGrokVideo('grok-imagine-video', 'k', { fetch }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter[method]('job-1', { abortSignal }), + ) + }, + ) + + it.each(['upload', 'get', 'delete'] as const)('files %s', async (method) => { + const { fetch, signals } = hangingFetch() + const files = createGrokFiles('k', { fetch, maxRetries: 0 }) + await expectAbortReachesFetch(signals, (abortSignal) => + method === 'upload' + ? files.upload({ data: 'AAAA', mimeType: 'image/png' }, { abortSignal }) + : files[method]('file-1', { abortSignal }), + ) + }) + + it('realtime token', async () => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + vi.stubEnv('XAI_API_KEY', 'k') + try { + await expectAbortReachesFetch(signals, (abortSignal) => + grokRealtimeToken().generateToken({ abortSignal }), + ) + } finally { + vi.unstubAllGlobals() + vi.unstubAllEnvs() + } + }) +}) diff --git a/packages/ai-lovable/src/adapters/embedding.ts b/packages/ai-lovable/src/adapters/embedding.ts index e87dcf4ec4..6147bf6faa 100644 --- a/packages/ai-lovable/src/adapters/embedding.ts +++ b/packages/ai-lovable/src/adapters/embedding.ts @@ -59,7 +59,11 @@ export class LovableEmbeddingAdapter< { provider: this.name, model }, ) - const response = await this.client.embeddings.create(request) + const response = options.abortSignal + ? await this.client.embeddings.create(request, { + signal: options.abortSignal, + }) + : await this.client.embeddings.create(request) const usage: TokenUsage = { promptTokens: response.usage.prompt_tokens, diff --git a/packages/ai-lovable/src/adapters/video.ts b/packages/ai-lovable/src/adapters/video.ts index c3149b4f2f..6cbedd2c39 100644 --- a/packages/ai-lovable/src/adapters/video.ts +++ b/packages/ai-lovable/src/adapters/video.ts @@ -149,9 +149,15 @@ export class LovableVideoAdapter< } } - async getVideoStatus(jobId: string): Promise { + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const requestOptions = options?.abortSignal && { + signal: options.abortSignal, + } try { - const response = await this.client.videos.retrieve(jobId) + const response = await this.client.videos.retrieve(jobId, requestOptions) return { jobId, status: this.mapStatus(response.status), @@ -170,9 +176,13 @@ export class LovableVideoAdapter< override async getVideo( jobId: string, + options?: { abortSignal?: AbortSignal }, ): Promise { + const requestOptions = options?.abortSignal && { + signal: options.abortSignal, + } try { - const videoInfo = await this.client.videos.retrieve(jobId) + const videoInfo = await this.client.videos.retrieve(jobId, requestOptions) const directUrl = videoResourceUrl(videoInfo) if (directUrl) { return { @@ -186,7 +196,10 @@ export class LovableVideoAdapter< // No URL: hand the download stream to generation middleware instead of // buffering the whole video. - const response = await this.client.videos.downloadContent(jobId) + const response = await this.client.videos.downloadContent( + jobId, + requestOptions, + ) if (!response.body) throw new Error('Video download returned no body') return { jobId, diff --git a/packages/ai-lovable/tests/abort-signal.test.ts b/packages/ai-lovable/tests/abort-signal.test.ts new file mode 100644 index 0000000000..c8a4ac3f29 --- /dev/null +++ b/packages/ai-lovable/tests/abort-signal.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createLovableEmbedding } from '../src/adapters/embedding' +import { createLovableVideo } from '../src/adapters/video' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('lovable adapters forward abortSignal to the request', () => { + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + const adapter = createLovableEmbedding('google/gemini-embedding-001', 'k', { + fetch, + maxRetries: 0, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'google/gemini-embedding-001', + input: ['a'], + logger, + abortSignal, + }), + ) + }) + + it.each(['getVideoStatus', 'getVideo'] as const)( + 'video %s', + async (method) => { + const { fetch, signals } = hangingFetch() + const adapter = createLovableVideo('google/veo-3.1-lite', 'k', { + fetch, + maxRetries: 0, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter[method]('job-1', { abortSignal }), + ) + }, + ) +}) diff --git a/packages/ai-mistral/src/adapters/embedding.ts b/packages/ai-mistral/src/adapters/embedding.ts index cbb733a976..10fd14020b 100644 --- a/packages/ai-mistral/src/adapters/embedding.ts +++ b/packages/ai-mistral/src/adapters/embedding.ts @@ -85,7 +85,11 @@ export class MistralEmbeddingAdapter< { provider: this.name, model }, ) - const response = await this.client.embeddings.create(request) + const response = options.abortSignal + ? await this.client.embeddings.create(request, { + fetchOptions: { signal: options.abortSignal }, + }) + : await this.client.embeddings.create(request) const usage: TokenUsage = { promptTokens: response.usage.promptTokens ?? 0, diff --git a/packages/ai-mistral/tests/abort-signal.test.ts b/packages/ai-mistral/tests/abort-signal.test.ts new file mode 100644 index 0000000000..dc2e940379 --- /dev/null +++ b/packages/ai-mistral/tests/abort-signal.test.ts @@ -0,0 +1,60 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createMistralEmbedding } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (input: unknown, init?: RequestInit): Promise => { + // Some SDKs pass a `Request`, others pass `init`. + const signal = + init?.signal ?? (input instanceof Request ? input.signal : undefined) + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('mistral adapters forward abortSignal to the request', () => { + afterEach(() => { + vi.unstubAllGlobals() + }) + + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + const adapter = createMistralEmbedding('mistral-embed', 'k') + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'mistral-embed', + input: ['a'], + logger, + abortSignal, + }), + ) + }) +}) diff --git a/packages/ai-neon/src/adapters/embedding.ts b/packages/ai-neon/src/adapters/embedding.ts index a3d1ec57c6..9c957a0c3c 100644 --- a/packages/ai-neon/src/adapters/embedding.ts +++ b/packages/ai-neon/src/adapters/embedding.ts @@ -66,7 +66,11 @@ export class NeonEmbeddingAdapter< { provider: this.name, model }, ) - const response = await this.client.embeddings.create(request) + const response = options.abortSignal + ? await this.client.embeddings.create(request, { + signal: options.abortSignal, + }) + : await this.client.embeddings.create(request) const usage: TokenUsage = { promptTokens: response.usage.prompt_tokens, diff --git a/packages/ai-neon/tests/abort-signal.test.ts b/packages/ai-neon/tests/abort-signal.test.ts new file mode 100644 index 0000000000..4dde8214e0 --- /dev/null +++ b/packages/ai-neon/tests/abort-signal.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createNeonEmbedding } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('neon adapters forward abortSignal to the request', () => { + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + const adapter = createNeonEmbedding('bge-large-en', 'k', { + baseURL: 'https://gateway.test/v1', + fetch, + maxRetries: 0, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'bge-large-en', + input: ['a'], + logger, + abortSignal, + }), + ) + }) +}) diff --git a/packages/ai-ollama/src/adapters/embedding.ts b/packages/ai-ollama/src/adapters/embedding.ts index 4b56e2e04d..aaeec64d6b 100644 --- a/packages/ai-ollama/src/adapters/embedding.ts +++ b/packages/ai-ollama/src/adapters/embedding.ts @@ -48,6 +48,8 @@ export class OllamaEmbeddingAdapter< readonly name = 'ollama' as const protected client: Ollama + /** The config of the adapter's own client. An injected client has none. */ + private readonly clientConfig: OllamaEmbeddingConfig | undefined constructor( hostOrClientOrConfig: string | Ollama | OllamaEmbeddingConfig | undefined, @@ -58,12 +60,14 @@ export class OllamaEmbeddingAdapter< typeof hostOrClientOrConfig === 'string' || hostOrClientOrConfig === undefined ) { - this.client = createOllamaClient({ host: hostOrClientOrConfig }) + this.clientConfig = { host: hostOrClientOrConfig } + this.client = createOllamaClient(this.clientConfig) } else if ('embed' in hostOrClientOrConfig) { // Ollama client instance (has an embed method) this.client = hostOrClientOrConfig } else { // OllamaEmbeddingConfig object + this.clientConfig = hostOrClientOrConfig this.client = createOllamaClient(hostOrClientOrConfig) } } @@ -102,7 +106,16 @@ export class OllamaEmbeddingAdapter< { provider: this.name, model }, ) - const response: EmbedResponse = await this.client.embed(request) + // `ollama` takes no per-call signal, so a call with a signal uses a + // client whose fetch carries it. An injected client cannot take one. + const signal = options.abortSignal + const client = + signal && this.clientConfig + ? createOllamaClient(this.clientConfig, (input, init) => + fetch(input, { ...init, signal }), + ) + : this.client + const response: EmbedResponse = await client.embed(request) // Include usage only when Ollama actually reported a prompt token count. const promptEvalCount = extractPromptEvalCount(response) diff --git a/packages/ai-ollama/tests/abort-signal.test.ts b/packages/ai-ollama/tests/abort-signal.test.ts new file mode 100644 index 0000000000..9c7c83b3a9 --- /dev/null +++ b/packages/ai-ollama/tests/abort-signal.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createOllamaEmbedding } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('ollama adapters forward abortSignal to the request', () => { + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + try { + const adapter = createOllamaEmbedding('nomic-embed-text') + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'nomic-embed-text', + input: ['a'], + logger, + abortSignal, + }), + ) + } finally { + vi.unstubAllGlobals() + } + }) +}) diff --git a/packages/ai-openai/src/adapters/embedding.ts b/packages/ai-openai/src/adapters/embedding.ts index 7072d1ce7b..88ec354fc2 100644 --- a/packages/ai-openai/src/adapters/embedding.ts +++ b/packages/ai-openai/src/adapters/embedding.ts @@ -76,7 +76,11 @@ export class OpenAIEmbeddingAdapter< { provider: this.name, model }, ) - const response = await this.client.embeddings.create(request) + const response = options.abortSignal + ? await this.client.embeddings.create(request, { + signal: options.abortSignal, + }) + : await this.client.embeddings.create(request) const usage: TokenUsage = { promptTokens: response.usage.prompt_tokens, diff --git a/packages/ai-openai/src/adapters/files.ts b/packages/ai-openai/src/adapters/files.ts index bcc2fb270f..e2b16a5d81 100644 --- a/packages/ai-openai/src/adapters/files.ts +++ b/packages/ai-openai/src/adapters/files.ts @@ -35,27 +35,46 @@ export class OpenAIFilesAdapter extends BaseFilesAdapter<'openai'> { this.purpose = purpose ?? 'user_data' } - async upload(input: FileUploadInput): Promise> { + async upload( + input: FileUploadInput, + options?: { abortSignal?: AbortSignal }, + ): Promise> { const { blob, mimeType, filename } = normalizeFileUploadInput(input) const file = await toFile(blob, filename, { ...(mimeType ? { type: mimeType } : {}), }) - const result = await this.client.files.create({ - file, - purpose: this.purpose, - }) + const result = await this.client.files.create( + { file, purpose: this.purpose }, + requestOptions(options), + ) return toFileHandle(result) } - async get(id: string): Promise> { - return toFileHandle(await this.client.files.retrieve(id)) + async get( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise> { + const request = requestOptions(options) + return toFileHandle( + request + ? await this.client.files.retrieve(id, request) + : await this.client.files.retrieve(id), + ) } - async delete(id: string): Promise { - await this.client.files.delete(id) + async delete( + id: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + await this.client.files.delete(id, requestOptions(options)) } } +/** SDK request options that carry the caller's abort signal, if any. */ +function requestOptions(options?: { abortSignal?: AbortSignal }) { + return options?.abortSignal ? { signal: options.abortSignal } : undefined +} + function toFileHandle(file: FileObject): FileHandle<'openai'> { return { id: file.id, diff --git a/packages/ai-openai/src/adapters/video.ts b/packages/ai-openai/src/adapters/video.ts index 7c83365a69..eb97eb83a9 100644 --- a/packages/ai-openai/src/adapters/video.ts +++ b/packages/ai-openai/src/adapters/video.ts @@ -177,7 +177,10 @@ export class OpenAIVideoAdapter< req: Record, options?: { signal?: AbortSignal | undefined }, ) => Promise<{ id: string }> - retrieve: (id: string) => Promise<{ + retrieve: ( + id: string, + options?: { signal?: AbortSignal }, + ) => Promise<{ id: string status: string progress?: number @@ -185,15 +188,24 @@ export class OpenAIVideoAdapter< expires_at?: number error?: { message?: string } }> - downloadContent: (id: string) => Promise + downloadContent: ( + id: string, + options?: { signal?: AbortSignal }, + ) => Promise } { return (this.client as { videos: any }).videos } - async getVideoStatus(jobId: string): Promise { + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const requestOptions = options?.abortSignal && { + signal: options.abortSignal, + } try { const videosClient = this.getVideosClient() - const response = await videosClient.retrieve(jobId) + const response = await videosClient.retrieve(jobId, requestOptions) // `VideoStatusResult` declares optional fields without `| undefined`; // spread conditionally so we omit absent fields rather than assigning // `undefined`. @@ -215,14 +227,18 @@ export class OpenAIVideoAdapter< override async getVideo( jobId: string, + options?: { abortSignal?: AbortSignal }, ): Promise { + const requestOptions = options?.abortSignal && { + signal: options.abortSignal, + } try { const videosClient = this.getVideosClient() // Prefer retrieve() because many openai-compatible backends (and the // aimock test harness) return the URL directly on the video resource // and do not implement a separate /content endpoint. - const videoInfo = await videosClient.retrieve(jobId) + const videoInfo = await videosClient.retrieve(jobId, requestOptions) if (videoInfo.url) { // `VideoUrlResult.expiresAt` is `expiresAt?: Date` without `| undefined`; // omit the field when the API didn't return an expiry. @@ -237,7 +253,7 @@ export class OpenAIVideoAdapter< // No URL: hand the download stream to generation middleware (which // stores it and sets the URL) instead of buffering the whole video. - const response = await videosClient.downloadContent(jobId) + const response = await videosClient.downloadContent(jobId, requestOptions) if (!response.body) throw new Error('Video download returned no body') return { jobId, diff --git a/packages/ai-openai/src/realtime/adapter.ts b/packages/ai-openai/src/realtime/adapter.ts index 798b885d13..896655f2ad 100644 --- a/packages/ai-openai/src/realtime/adapter.ts +++ b/packages/ai-openai/src/realtime/adapter.ts @@ -48,6 +48,7 @@ export function openaiRealtime( async connect( token: RealtimeToken, _clientTools?: ReadonlyArray, + options?: { abortSignal?: AbortSignal }, ): Promise { const model = token.config.model ?? 'gpt-realtime-2.1' logger.request(`activity=realtime provider=openai model=${model}`, { @@ -56,7 +57,7 @@ export function openaiRealtime( }) if (connectionMode === 'webrtc') { - return createWebRTCConnection(token, logger) + return createWebRTCConnection(token, logger, options?.abortSignal) } const error = new Error('WebSocket connection mode not yet implemented') logger.errors('openai.realtime fatal', { @@ -74,6 +75,7 @@ export function openaiRealtime( async function createWebRTCConnection( token: RealtimeToken, logger: InternalLogger, + abortSignal?: AbortSignal, ): Promise { const eventHandlers = new Map>>() @@ -163,60 +165,75 @@ async function createWebRTCConnection( } } - // IMPORTANT: Request microphone access and add audio track BEFORE creating offer - // OpenAI's Realtime API requires an audio track in the SDP offer + // A failed or aborted setup releases the microphone and the peer + // connection. Without this, the mic indicator stays on after disconnect(). try { - localStream = await navigator.mediaDevices.getUserMedia({ - audio: { - echoCancellation: true, - noiseSuppression: true, - sampleRate: 24000, + // A disconnect during the token fetch already aborted the signal. Do not + // ask for the microphone in that case. + abortSignal?.throwIfAborted() + // IMPORTANT: Request microphone access and add audio track BEFORE creating offer + // OpenAI's Realtime API requires an audio track in the SDP offer + try { + localStream = await navigator.mediaDevices.getUserMedia({ + audio: { + echoCancellation: true, + noiseSuppression: true, + sampleRate: 24000, + }, + }) + + // Add audio track to peer connection + for (const track of localStream.getAudioTracks()) { + pc.addTrack(track, localStream) + } + } catch (error) { + throw new Error( + `Microphone access required for realtime voice: ${error instanceof Error ? error.message : error}`, + ) + } + + // Create and set local description (now includes audio track) + const offer = await pc.createOffer() + await pc.setLocalDescription(offer) + + // Send SDP to OpenAI's GA `/calls` endpoint and get the answer. The model + // is bound to the ephemeral token (minted via `/v1/realtime/client_secrets`), + // so it must NOT be passed as a query param — GA rejects `?model=` with a + // 400. `offer.sdp` is `string | undefined` per the WebRTC type definitions; + // coerce to `null`, which `RequestInit.body` accepts. + const sdpResponse = await fetch(`${OPENAI_REALTIME_URL}/calls`, { + method: 'POST', + headers: { + Authorization: `Bearer ${token.token}`, + 'Content-Type': 'application/sdp', }, + body: offer.sdp ?? null, + ...(abortSignal && { signal: abortSignal }), }) - // Add audio track to peer connection - for (const track of localStream.getAudioTracks()) { - pc.addTrack(track, localStream) + if (!sdpResponse.ok) { + const errorText = await sdpResponse.text() + const error = new Error( + `Failed to establish WebRTC connection: ${sdpResponse.status} - ${errorText}`, + ) + logger.errors('openai.realtime fatal', { + error, + source: 'openai.realtime', + }) + throw error } - } catch (error) { - throw new Error( - `Microphone access required for realtime voice: ${error instanceof Error ? error.message : error}`, - ) - } - - // Create and set local description (now includes audio track) - const offer = await pc.createOffer() - await pc.setLocalDescription(offer) - - // Send SDP to OpenAI's GA `/calls` endpoint and get the answer. The model - // is bound to the ephemeral token (minted via `/v1/realtime/client_secrets`), - // so it must NOT be passed as a query param — GA rejects `?model=` with a - // 400. `offer.sdp` is `string | undefined` per the WebRTC type definitions; - // coerce to `null`, which `RequestInit.body` accepts. - const sdpResponse = await fetch(`${OPENAI_REALTIME_URL}/calls`, { - method: 'POST', - headers: { - Authorization: `Bearer ${token.token}`, - 'Content-Type': 'application/sdp', - }, - body: offer.sdp ?? null, - }) - if (!sdpResponse.ok) { - const errorText = await sdpResponse.text() - const error = new Error( - `Failed to establish WebRTC connection: ${sdpResponse.status} - ${errorText}`, - ) - logger.errors('openai.realtime fatal', { - error, - source: 'openai.realtime', - }) + const answerSdp = await sdpResponse.text() + await pc.setRemoteDescription({ type: 'answer', sdp: answerSdp }) + } catch (error) { + for (const track of localStream?.getTracks() ?? []) track.stop() + localStream = null + channel.close() + dataChannel = null + pc.close() throw error } - const answerSdp = await sdpResponse.text() - await pc.setRemoteDescription({ type: 'answer', sdp: answerSdp }) - // Set up input audio analysis now that we have the stream setupInputAudioAnalysis(localStream) @@ -708,7 +725,31 @@ async function createWebRTCConnection( // Wait for the data channel to be open before returning the connection. // This ensures session.update (tools, instructions, etc.) can be sent immediately. - await dataChannelReady + // An abort during this wait releases the mic, the peer connection, and the + // audio context, and rejects, so connect() does not stay pending. + try { + await waitUnlessAborted(dataChannelReady, abortSignal) + } catch (error) { + await connection.disconnect() + throw error + } return connection } + +/** Resolve with `promise`, or reject with the abort reason when `signal` aborts first. */ +function waitUnlessAborted( + promise: Promise, + signal: AbortSignal | undefined, +): Promise { + if (!signal) return promise + return new Promise((resolve, reject) => { + const onAbort = () => reject(signal.reason) + if (signal.aborted) return onAbort() + signal.addEventListener('abort', onAbort, { once: true }) + void promise.then(() => { + signal.removeEventListener('abort', onAbort) + resolve() + }) + }) +} diff --git a/packages/ai-openai/src/realtime/token.ts b/packages/ai-openai/src/realtime/token.ts index 22d475e46a..298dcd1f95 100644 --- a/packages/ai-openai/src/realtime/token.ts +++ b/packages/ai-openai/src/realtime/token.ts @@ -84,7 +84,9 @@ export function openaiRealtimeToken( return { provider: 'openai', - async generateToken(): Promise { + async generateToken(request?: { + abortSignal?: AbortSignal + }): Promise { const model: OpenAIRealtimeModel = options.model ?? 'gpt-realtime-2.1' // Only the model is sent server-side; all other session config @@ -97,6 +99,7 @@ export function openaiRealtimeToken( 'Content-Type': 'application/json', }, body: JSON.stringify(buildClientSecretRequest(model)), + ...(request?.abortSignal && { signal: request.abortSignal }), }) if (!response.ok) { diff --git a/packages/ai-openai/tests/abort-signal.test.ts b/packages/ai-openai/tests/abort-signal.test.ts new file mode 100644 index 0000000000..c96eecefe9 --- /dev/null +++ b/packages/ai-openai/tests/abort-signal.test.ts @@ -0,0 +1,98 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createOpenaiEmbedding } from '../src/adapters/embedding' +import { createOpenaiVideo } from '../src/adapters/video' +import { createOpenaiFiles } from '../src/adapters/files' +import { openaiRealtimeToken } from '../src/realtime/token' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('openai adapters forward abortSignal to the request', () => { + afterEach(() => { + vi.unstubAllGlobals() + vi.unstubAllEnvs() + }) + + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + const adapter = createOpenaiEmbedding('text-embedding-3-small', 'k', { + fetch, + maxRetries: 0, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'text-embedding-3-small', + input: ['a'], + logger, + abortSignal, + }), + ) + }) + + it.each(['getVideoStatus', 'getVideo'] as const)( + 'video %s', + async (method) => { + const { fetch, signals } = hangingFetch() + const adapter = createOpenaiVideo('sora-2', 'k', { + fetch, + maxRetries: 0, + }) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter[method]('video_1', { abortSignal }), + ) + }, + ) + + it.each(['upload', 'get', 'delete'] as const)('files %s', async (method) => { + const { fetch, signals } = hangingFetch() + const files = createOpenaiFiles('k', { fetch, maxRetries: 0 }) + await expectAbortReachesFetch(signals, (abortSignal) => + method === 'upload' + ? files.upload({ data: 'AAAA', mimeType: 'image/png' }, { abortSignal }) + : files[method]('file-1', { abortSignal }), + ) + }) + + it('realtime token', async () => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + vi.stubEnv('OPENAI_API_KEY', 'k') + const adapter = openaiRealtimeToken() + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.generateToken({ abortSignal }), + ) + }) +}) diff --git a/packages/ai-openai/tests/realtime-connect-teardown.test.ts b/packages/ai-openai/tests/realtime-connect-teardown.test.ts new file mode 100644 index 0000000000..ba0689dabc --- /dev/null +++ b/packages/ai-openai/tests/realtime-connect-teardown.test.ts @@ -0,0 +1,123 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { openaiRealtime } from '../src/realtime/adapter' +import type { RealtimeToken } from '@tanstack/ai' + +const token: RealtimeToken = { + provider: 'openai', + token: 'ek_test', + expiresAt: Date.now() + 60_000, + config: { model: 'gpt-realtime-2.1' }, +} + +/** Fake browser WebRTC and microphone globals that record what was released. */ +function stubBrowser() { + const track = { kind: 'audio', stop: vi.fn() } + const stream = { getAudioTracks: () => [track], getTracks: () => [track] } + const channel = { close: vi.fn() } + const pc = { + createDataChannel: () => channel, + addTrack: vi.fn(), + createOffer: async () => ({ type: 'offer', sdp: 'v=0' }), + setLocalDescription: async () => {}, + close: vi.fn(), + } + vi.stubGlobal( + 'RTCPeerConnection', + vi.fn(function RTCPeerConnection() { + return pc + }), + ) + vi.stubGlobal('navigator', { + mediaDevices: { getUserMedia: async () => stream }, + }) + return { track, channel, pc } +} + +describe('openaiRealtime connect teardown', () => { + afterEach(() => { + vi.unstubAllGlobals() + }) + + it('releases the microphone and the peer connection when the SDP request aborts', async () => { + const { track, channel, pc } = stubBrowser() + vi.stubGlobal( + 'fetch', + (_url: string, init?: RequestInit) => + new Promise((_, reject) => { + init?.signal?.addEventListener('abort', () => + reject(init.signal?.reason), + ) + }), + ) + const controller = new AbortController() + const pending = openaiRealtime().connect(token, undefined, { + abortSignal: controller.signal, + }) + await vi.waitFor(() => expect(pc.addTrack).toHaveBeenCalled()) + controller.abort() + + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(track.stop).toHaveBeenCalled() + expect(channel.close).toHaveBeenCalled() + expect(pc.close).toHaveBeenCalled() + }) + + it('does not ask for the microphone when the signal is already aborted', async () => { + const { pc } = stubBrowser() + const getUserMedia = vi.fn() + vi.stubGlobal('navigator', { mediaDevices: { getUserMedia } }) + const controller = new AbortController() + controller.abort() + + await expect( + openaiRealtime().connect(token, undefined, { + abortSignal: controller.signal, + }), + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(getUserMedia).not.toHaveBeenCalled() + expect(pc.close).toHaveBeenCalled() + }) + + it('releases them when the SDP request fails', async () => { + const { track, pc } = stubBrowser() + vi.stubGlobal('fetch', async () => new Response('bad', { status: 500 })) + + await expect(openaiRealtime().connect(token)).rejects.toThrow( + 'Failed to establish WebRTC connection', + ) + expect(track.stop).toHaveBeenCalled() + expect(pc.close).toHaveBeenCalled() + }) + + it('releases everything when disconnect() aborts the wait for the data channel', async () => { + const { track, channel, pc } = stubBrowser() + Object.assign(pc, { setRemoteDescription: async () => {} }) + const audioContext = { + state: 'running', + createAnalyser: () => ({ connect: vi.fn() }), + createMediaStreamSource: () => ({ connect: vi.fn() }), + close: vi.fn(async () => {}), + } + const AudioContextMock = vi.fn(function AudioContext() { + return audioContext + }) + vi.stubGlobal('AudioContext', AudioContextMock) + vi.stubGlobal('fetch', async () => new Response('v=0', { status: 200 })) + const controller = new AbortController() + // The data channel never opens, so connect() waits until the abort. + const settled = expect( + openaiRealtime().connect(token, undefined, { + abortSignal: controller.signal, + }), + ).rejects.toMatchObject({ name: 'AbortError' }) + // The SDP step is done once the input audio analysis starts. + await vi.waitFor(() => expect(AudioContextMock).toHaveBeenCalled()) + controller.abort() + + await settled + expect(track.stop).toHaveBeenCalled() + expect(channel.close).toHaveBeenCalled() + expect(pc.close).toHaveBeenCalled() + expect(audioContext.close).toHaveBeenCalled() + }) +}) diff --git a/packages/ai-openrouter/src/adapters/video.ts b/packages/ai-openrouter/src/adapters/video.ts index 68968bdc2e..66b1e8bc3c 100644 --- a/packages/ai-openrouter/src/adapters/video.ts +++ b/packages/ai-openrouter/src/adapters/video.ts @@ -312,8 +312,14 @@ export class OpenRouterVideoAdapter< return getVideoDurationOptions(this.model) } - async getVideoStatus(jobId: string): Promise { - const response = await this.client.videoGeneration.getGeneration({ jobId }) + async getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise { + const response = await this.client.videoGeneration.getGeneration( + { jobId }, + ...requestArgs(options?.abortSignal), + ) return { jobId, status: mapStatus(response.status), @@ -323,8 +329,13 @@ export class OpenRouterVideoAdapter< override async getVideo( jobId: string, + options?: { abortSignal?: AbortSignal }, ): Promise { - const response = await this.client.videoGeneration.getGeneration({ jobId }) + const reqArgs = requestArgs(options?.abortSignal) + const response = await this.client.videoGeneration.getGeneration( + { jobId }, + ...reqArgs, + ) const status = mapStatus(response.status) if (status === 'failed') { throw new Error( @@ -345,7 +356,10 @@ export class OpenRouterVideoAdapter< // accepts `video/mp4` and streams the bytes. let stream: ReadableStream try { - stream = await this.client.videoGeneration.getVideoContent({ jobId }) + stream = await this.client.videoGeneration.getVideoContent( + { jobId }, + ...reqArgs, + ) } catch (error) { const detail = error instanceof Error ? error.message : String(error) throw new Error( @@ -408,3 +422,13 @@ export function openRouterVideo( const apiKey = getOpenRouterApiKeyFromEnv() return createOpenRouterVideo(model, apiKey, config) } + +/** + * SDK request options that carry the caller's abort signal. Empty without a + * signal, so the SDK call keeps its one-argument form. + */ +function requestArgs( + signal: AbortSignal | undefined, +): [] | [{ fetchOptions: { signal: AbortSignal } }] { + return signal ? [{ fetchOptions: { signal } }] : [] +} diff --git a/packages/ai-openrouter/tests/video-abort-signal.test.ts b/packages/ai-openrouter/tests/video-abort-signal.test.ts new file mode 100644 index 0000000000..89f8ff8a87 --- /dev/null +++ b/packages/ai-openrouter/tests/video-abort-signal.test.ts @@ -0,0 +1,52 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { createOpenRouterVideo } from '../src/adapters/video' + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (input: unknown, init?: RequestInit): Promise => { + // Some SDKs pass a `Request`, others pass `init`. + const signal = + init?.signal ?? (input instanceof Request ? input.signal : undefined) + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('openrouter video forwards abortSignal to the request', () => { + afterEach(() => { + vi.unstubAllGlobals() + }) + + it.each(['getVideoStatus', 'getVideo'] as const)('%s', async (method) => { + const { fetch, signals } = hangingFetch() + vi.stubGlobal('fetch', fetch) + const adapter = createOpenRouterVideo('bytedance/seedance-2.0', 'k') + await expectAbortReachesFetch(signals, (abortSignal) => + adapter[method]('job-1', { abortSignal }), + ) + }) +}) diff --git a/packages/ai-vercel-gateway/src/adapters/embedding.ts b/packages/ai-vercel-gateway/src/adapters/embedding.ts index 0bafb78c95..c1a3fcf8bc 100644 --- a/packages/ai-vercel-gateway/src/adapters/embedding.ts +++ b/packages/ai-vercel-gateway/src/adapters/embedding.ts @@ -71,7 +71,11 @@ export class VercelGatewayEmbeddingAdapter< { provider: this.name, model }, ) - const response = await this.client.embeddings.create(request) + const response = options.abortSignal + ? await this.client.embeddings.create(request, { + signal: options.abortSignal, + }) + : await this.client.embeddings.create(request) const usage: TokenUsage = { promptTokens: response.usage.prompt_tokens, diff --git a/packages/ai-vercel-gateway/tests/abort-signal.test.ts b/packages/ai-vercel-gateway/tests/abort-signal.test.ts new file mode 100644 index 0000000000..983f3c5b9f --- /dev/null +++ b/packages/ai-vercel-gateway/tests/abort-signal.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it, vi } from 'vitest' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createVercelGatewayEmbedding } from '../src/adapters/embedding' + +const logger = resolveDebugOption(false) + +/** + * A fetch that fails at once when it gets no signal, and otherwise waits + * until that signal aborts. `signals` records what each request got. + */ +function hangingFetch() { + const signals: Array = [] + const fetch = (_input: unknown, init?: RequestInit): Promise => { + const signal = init?.signal + if (!signal) return Promise.reject(new Error('request had no signal')) + signals.push(signal) + return new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + } + return { fetch, signals } +} + +/** Start `call`, wait for its request, abort the caller, expect a rejection. */ +async function expectAbortReachesFetch( + signals: Array, + call: (abortSignal: AbortSignal) => Promise, +) { + const controller = new AbortController() + const pending = call(controller.signal) + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toThrow() + expect(signals[0]!.aborted).toBe(true) +} + +describe('vercel-gateway adapters forward abortSignal to the request', () => { + it('embedding', async () => { + const { fetch, signals } = hangingFetch() + const adapter = createVercelGatewayEmbedding( + 'alibaba/qwen3-embedding-0.6b', + 'k', + { fetch, maxRetries: 0 }, + ) + await expectAbortReachesFetch(signals, (abortSignal) => + adapter.createEmbeddings({ + model: 'alibaba/qwen3-embedding-0.6b', + input: ['a'], + logger, + abortSignal, + }), + ) + }) +}) diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 379235ed4f..93c41c9d10 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1594,6 +1594,10 @@ class TextEngine< return } + // A run stopped before it started runs no pending tools and raises no + // interrupt. The `finally` below reports it through `onAbort`. + if (this.isCancelled()) return + const hasResumedTools = this.getPendingToolCallsFromMessages().length > 0 const pendingPhase = yield* this.checkForPendingToolCalls() if (pendingPhase === 'stop' || pendingPhase === 'wait') { @@ -6397,6 +6401,23 @@ function runStreamingText( return stream } +/** + * MCP tools for one run. When the run is stopped during discovery, it returns + * no tools: the engine then ends the run through its abort path, so middleware + * gets `onAbort`. Other discovery errors still throw. + */ +async function discoverMcpTools( + mcpManager: MCPManager, + signal: AbortSignal | undefined, +) { + try { + return await mcpManager.discover(signal) + } catch (error) { + if (signal?.aborted) return [] + throw error + } +} + async function* runChatEngine( options: RuntimeTextActivityOptions, engineRef: DeliveryEngineRef, @@ -6414,7 +6435,10 @@ async function* runChatEngine( const logger = resolveDebugOption(debug) const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover() + const mcpTools = await discoverMcpTools( + mcpManager, + options.abortController?.signal, + ) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } @@ -7030,7 +7054,10 @@ async function runAgenticStructuredOutput( adapter.combinedStructuredOutputSource?.(options.modelOptions) ?? 'text' const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover() + const mcpTools = await discoverMcpTools( + mcpManager, + options.abortController?.signal, + ) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } @@ -7392,7 +7419,10 @@ async function* runStreamingStructuredOutputImpl( : undefined const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover() + const mcpTools = await discoverMcpTools( + mcpManager, + options.abortController?.signal, + ) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } diff --git a/packages/ai/src/activities/chat/mcp/manager.ts b/packages/ai/src/activities/chat/mcp/manager.ts index bac5e22ec4..ec0583e017 100644 --- a/packages/ai/src/activities/chat/mcp/manager.ts +++ b/packages/ai/src/activities/chat/mcp/manager.ts @@ -70,12 +70,15 @@ export class MCPManager { * Discover + merge tools from all sources. Throws on a fatal discovery error * (no `onDiscoveryError`, or it re-threw) or a duplicate tool name; in that * case it first closes any connected sources when the policy is 'close'. + * `signal` (the chat run signal) goes to each source's `tools()`. */ - async discover(): Promise> { + async discover(signal?: AbortSignal): Promise> { if (this.#sources.length === 0) return [] try { const settled = await Promise.allSettled( - this.#sources.map((s) => s.tools({ lazy: this.#lazyTools })), + this.#sources.map((s) => + s.tools({ lazy: this.#lazyTools, ...(signal && { signal }) }), + ), ) const tools: Array = [] const zipped = this.#sources.map( diff --git a/packages/ai/src/activities/chat/mcp/types.ts b/packages/ai/src/activities/chat/mcp/types.ts index 96c7911caa..713be05dcb 100644 --- a/packages/ai/src/activities/chat/mcp/types.ts +++ b/packages/ai/src/activities/chat/mcp/types.ts @@ -26,7 +26,14 @@ export interface MCPToolSource { // Keep the options shape in sync with ai-mcp's `ToolsOptions` — extra // optional fields added there still match structurally, but chat() only // forwards what is declared here. - tools: (options?: { lazy?: boolean }) => Promise> + /** + * Lists the tools. `chat()` passes the run's `signal` so that a cancelled + * run also cancels discovery. A source that ignores `signal` still works. + */ + tools: (options?: { + lazy?: boolean + signal?: AbortSignal + }) => Promise> close: () => Promise /** * Reads an MCP resource by URI. Used by the chat manager to eagerly fetch diff --git a/packages/ai/src/activities/embed/index.ts b/packages/ai/src/activities/embed/index.ts index f32c07a5af..e3041117fc 100644 --- a/packages/ai/src/activities/embed/index.ts +++ b/packages/ai/src/activities/embed/index.ts @@ -7,8 +7,14 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' +import { + abortReasonMessage, + createActivityAbortControls, + raceWithAbort, +} from '../../utilities/activity-abort' import { createGenerationContext, + runGenerationAbort, runGenerationError, runGenerationFinish, runGenerationStart, @@ -119,6 +125,18 @@ export type EmbedOptions< * `GenerationMiddleware` contract for a custom backend. */ middleware?: Array + /** + * Maximum duration of this activity invocation in milliseconds. + * No SDK-wide default — choose a value suitable for the provider and job. + * Composed with {@link abortSignal}; the first abort wins. + */ + timeout?: number + /** + * Caller cancellation signal (request disconnects, job/runtime cancellation). + * Composed with {@link timeout} into an effective signal forwarded to the + * adapter. Request-specific — not stored on global provider client config. + */ + abortSignal?: AbortSignal } & ({} extends EmbedProviderOptionsForModel ? { /** Provider-specific options for embedding generation */ modelOptions?: EmbedProviderOptionsForModel< @@ -208,9 +226,14 @@ export async function embed< : [options.input] const { textInputCount, imageInputCount } = countEmbeddingInputModalities(inputItems) + const abortControls = createActivityAbortControls({ + timeout: options.timeout, + abortSignal: options.abortSignal, + }) const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'embedding', provider: adapter.name, model, @@ -218,7 +241,10 @@ export async function embed< createId, }) - await runGenerationStart(middleware, mwCtx) + await runGenerationStart(middleware, mwCtx).catch((error: unknown) => { + abortControls.clear() + throw error + }) aiEventClient.emit('embedding:request:started', { requestId, @@ -238,13 +264,18 @@ export async function embed< }) try { - const result = await adapter.createEmbeddings({ - model, - input: inputItems, - dimensions: options.dimensions, - modelOptions, - logger, - }) + const result = await raceWithAbort( + adapter.createEmbeddings({ + model, + input: inputItems, + dimensions: options.dimensions, + modelOptions, + logger, + ...(abortControls.signal ? { abortSignal: abortControls.signal } : {}), + }), + abortControls.signal, + ) + abortControls.clear() const duration = Date.now() - startTime aiEventClient.emit('embedding:request:completed', { @@ -278,6 +309,7 @@ export async function embed< return result } catch (error) { + abortControls.clear() const duration = Date.now() - startTime const err = error as Error aiEventClient.emit('embedding:request:error', { @@ -289,10 +321,19 @@ export async function embed< modelOptions, timestamp: Date.now(), }) - await runGenerationError(middleware, mwCtx, { - error, - duration, - }) + // Only a real caller abort or timeout is an abort. A provider error named + // `TimeoutError` (for example a client-side fetch timeout) is an error. + if (abortControls.signal?.aborted) { + await runGenerationAbort(middleware, mwCtx, { + reason: abortReasonMessage(error, abortControls.signal), + duration, + }) + } else { + await runGenerationError(middleware, mwCtx, { + error, + duration, + }) + } logger.errors('embed activity failed', { error, source: 'embed', diff --git a/packages/ai/src/activities/evaluate/index.ts b/packages/ai/src/activities/evaluate/index.ts index 3c2b94c5b9..3274069237 100644 --- a/packages/ai/src/activities/evaluate/index.ts +++ b/packages/ai/src/activities/evaluate/index.ts @@ -9,6 +9,11 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' import { isAbortShapedError } from '../error-payload' +import { + createActivityAbortControls, + isActivityAbortError, + raceWithAbort, +} from '../../utilities/activity-abort' import { createGenerationContext, runGenerationAbort, @@ -165,6 +170,12 @@ export interface EvaluateActivityOptions< modelOptions?: EvaluateProviderOptions /** Forwarded to the provider request for cancellation. */ abortSignal?: AbortSignal + /** + * Maximum duration of this activity invocation in milliseconds. + * No SDK-wide default. Composed with {@link abortSignal}; the first abort + * wins. + */ + timeout?: number /** * Observe-only middleware notified on start, usage, success, abort, and * error. Pass `otelMiddleware()` to emit OpenTelemetry spans, or implement @@ -194,6 +205,9 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { // would misroute a real failure to the abort hook whenever a shared signal // happens to already be aborted, hiding it from `onError` observers. if (isAbortShapedError(error)) return true + // A `TimeoutError` is an abort only when our signal fired (the `timeout` + // option). A provider's own fetch timeout stays an error. + if (signal?.aborted && isActivityAbortError(error)) return true // Fall back to signal state only for non-Error throws we can't otherwise // identify; a real Error with a non-abort name is never an abort. return error instanceof Error ? false : signal?.aborted === true @@ -464,6 +478,7 @@ export function boolean(options: { * @param options.questions Questions built with `choice`, `score`, `boolean`. * @param options.modelOptions Provider-specific options. * @param options.abortSignal Cancels the in-flight request. + * @param options.timeout Maximum duration in milliseconds. * @param options.middleware Observe-only generation middleware. * @param options.debug Debug logging option. * @@ -503,23 +518,21 @@ export async function decide< TAdapter extends EvaluateAdapter>, TQuestions extends Record, >(options: EvaluateActivityOptions) { - const { - adapter, - state, - questions, - modelOptions, - abortSignal, - middleware, - debug, - } = options + const { adapter, state, questions, modelOptions, middleware, debug } = options const model = adapter.model const keys = assertQuestions(questions) const requestId = createId('evaluate') const startTime = Date.now() const logger: InternalLogger = resolveDebugOption(debug) + const abortControls = createActivityAbortControls({ + timeout: options.timeout, + abortSignal: options.abortSignal, + }) + const abortSignal = abortControls.signal const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'evaluate', provider: adapter.name, model, @@ -527,7 +540,10 @@ export async function decide< createId, }) - await runGenerationStart(middleware, mwCtx) + await runGenerationStart(middleware, mwCtx).catch((error: unknown) => { + abortControls.clear() + throw error + }) aiEventClient.emit('evaluate:request:started', { requestId, @@ -544,14 +560,18 @@ export async function decide< }) try { - const result = await adapter.evaluate({ - model, - state, - questions, - modelOptions, + const result = await raceWithAbort( + adapter.evaluate({ + model, + state, + questions, + modelOptions, + abortSignal, + logger, + }), abortSignal, - logger, - }) + ) + abortControls.clear() let answers try { @@ -606,6 +626,7 @@ export async function decide< ...(result.provider !== undefined && { provider: result.provider }), }) } catch (error) { + abortControls.clear() const duration = Date.now() - startTime try { if (isEvaluateError(error) && error.usage) { diff --git a/packages/ai/src/activities/files/adapter.ts b/packages/ai/src/activities/files/adapter.ts index caef9c7189..e8ac32a51f 100644 --- a/packages/ai/src/activities/files/adapter.ts +++ b/packages/ai/src/activities/files/adapter.ts @@ -71,13 +71,25 @@ export interface FileHandle { * `TName` is the provider name literal (`'openai'`, `'gemini'`, ...); concrete * adapters bind it so the handles they issue carry their provenance in the * type system. + * + * Each method takes an optional `{ abortSignal }`. Adapters should forward it + * to the provider request when supported. */ export interface FilesAdapter { readonly kind: 'files' readonly name: TName - upload: (input: FileUploadInput) => Promise> - get?: (id: string) => Promise> - delete?: (id: string) => Promise + upload: ( + input: FileUploadInput, + options?: { abortSignal?: AbortSignal }, + ) => Promise> + get?: ( + id: string, + options?: { abortSignal?: AbortSignal }, + ) => Promise> + delete?: ( + id: string, + options?: { abortSignal?: AbortSignal }, + ) => Promise } export type AnyFilesAdapter = FilesAdapter @@ -116,5 +128,8 @@ export abstract class BaseFilesAdapter< readonly kind = 'files' as const abstract readonly name: TName - abstract upload(input: FileUploadInput): Promise> + abstract upload( + input: FileUploadInput, + options?: { abortSignal?: AbortSignal }, + ): Promise> } diff --git a/packages/ai/src/activities/files/index.ts b/packages/ai/src/activities/files/index.ts index 80d59dada0..0c18b17452 100644 --- a/packages/ai/src/activities/files/index.ts +++ b/packages/ai/src/activities/files/index.ts @@ -7,6 +7,7 @@ * error when the selected provider has no lifecycle API. */ +import { raceWithAbort } from '../../utilities/activity-abort' import type { ContentPartFileSource } from '../../types' import type { FileHandle, FileUploadInput, FilesAdapter } from './adapter' @@ -27,8 +28,30 @@ export const kind = 'files' as const export async function uploadFile(options: { adapter: FilesAdapter & { kind: typeof kind } input: FileUploadInput + /** Cancels the upload. The call rejects when the signal aborts. */ + abortSignal?: AbortSignal }): Promise> { - return options.adapter.upload(options.input) + const { adapter } = options + return withAbort( + adapter.upload.bind(adapter), + options.input, + options.abortSignal, + ) +} + +/** + * Run an adapter call with `{ abortSignal }` only when there is a signal, and + * reject on abort even if the adapter ignores the signal. + */ +function withAbort( + call: (arg: TArg, options?: { abortSignal?: AbortSignal }) => Promise, + arg: TArg, + abortSignal: AbortSignal | undefined, +): Promise { + return raceWithAbort( + abortSignal ? call(arg, { abortSignal }) : call(arg), + abortSignal, + ) } /** @@ -51,15 +74,18 @@ export async function getFile(options: { // `NoInfer` so `TName` comes from the adapter only. Otherwise a foreign // handle widens it to a union and the call compiles. id: string | FileHandle> + /** Cancels the request. The call rejects when the signal aborts. */ + abortSignal?: AbortSignal }): Promise> { const { adapter } = options - if (!adapter.get) { + const get = adapter.get?.bind(adapter) + if (!get) { throw new Error( `${adapter.name}: files adapter does not support get() — this provider ` + `has no file-retrieval API.`, ) } - return adapter.get(toLifecycleId(options.id)) + return withAbort(get, toLifecycleId(options.id), options.abortSignal) } /** @@ -72,15 +98,18 @@ export async function getFile(options: { export async function deleteFile(options: { adapter: FilesAdapter & { kind: typeof kind } id: string | FileHandle> + /** Cancels the request. The call rejects when the signal aborts. */ + abortSignal?: AbortSignal }): Promise { const { adapter } = options - if (!adapter.delete) { + const remove = adapter.delete?.bind(adapter) + if (!remove) { throw new Error( `${adapter.name}: files adapter does not support delete() — this ` + `provider has no file-deletion API.`, ) } - return adapter.delete(toLifecycleId(options.id)) + return withAbort(remove, toLifecycleId(options.id), options.abortSignal) } /** diff --git a/packages/ai/src/activities/generateAudio/index.ts b/packages/ai/src/activities/generateAudio/index.ts index ab470ba09b..770dcd2105 100644 --- a/packages/ai/src/activities/generateAudio/index.ts +++ b/packages/ai/src/activities/generateAudio/index.ts @@ -203,6 +203,7 @@ async function runGenerateAudio< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'audio', provider: adapter.name, model, diff --git a/packages/ai/src/activities/generateImage/index.ts b/packages/ai/src/activities/generateImage/index.ts index 33c5258fc2..9ad4539fc8 100644 --- a/packages/ai/src/activities/generateImage/index.ts +++ b/packages/ai/src/activities/generateImage/index.ts @@ -298,6 +298,7 @@ async function runGenerateImage< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'image', provider: adapter.name, model, diff --git a/packages/ai/src/activities/generateLiveVideo/index.ts b/packages/ai/src/activities/generateLiveVideo/index.ts index a09a0aff0f..f5de242d19 100644 --- a/packages/ai/src/activities/generateLiveVideo/index.ts +++ b/packages/ai/src/activities/generateLiveVideo/index.ts @@ -201,6 +201,7 @@ async function runGenerateLiveVideo< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'liveVideo', provider: adapter.name, model, diff --git a/packages/ai/src/activities/generateSpeech/index.ts b/packages/ai/src/activities/generateSpeech/index.ts index d6f084a328..9f279fde1b 100644 --- a/packages/ai/src/activities/generateSpeech/index.ts +++ b/packages/ai/src/activities/generateSpeech/index.ts @@ -302,6 +302,7 @@ async function runGenerateSpeech< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'tts', provider: adapter.name, model, diff --git a/packages/ai/src/activities/generateTranscription/index.ts b/packages/ai/src/activities/generateTranscription/index.ts index 92aabb1428..4ea736ce65 100644 --- a/packages/ai/src/activities/generateTranscription/index.ts +++ b/packages/ai/src/activities/generateTranscription/index.ts @@ -250,6 +250,7 @@ async function runGenerateTranscription< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'transcription', provider: adapter.name, model, diff --git a/packages/ai/src/activities/generateVideo/adapter.ts b/packages/ai/src/activities/generateVideo/adapter.ts index 937febeaa8..d4a58d82a4 100644 --- a/packages/ai/src/activities/generateVideo/adapter.ts +++ b/packages/ai/src/activities/generateVideo/adapter.ts @@ -112,16 +112,35 @@ export interface VideoAdapter< ) => Promise /** - * Get the current status of a video generation job. + * Get the current status of a video generation job. Adapters should forward + * `options.abortSignal` to the provider request when supported. */ - getVideoStatus: (jobId: string) => Promise + getVideoStatus: ( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ) => Promise /** * Get the finished video: a public URL when the provider has one, or the * download stream for generation middleware to host. Call only after - * status is 'completed'. + * status is 'completed'. Adapters should forward `options.abortSignal` to + * the provider request (and the download) when supported. */ - getVideo: (jobId: string) => Promise + getVideo: ( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ) => Promise + + /** + * Optional. Ask the provider to stop a job that is not finished. Implement + * it only when the provider has a real cancel API. Streaming + * `generateVideo()` calls it (best effort) when its `abortSignal` or + * `timeout` aborts after the job exists. A rejection is logged, not thrown. + */ + cancelVideoJob?: ( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ) => Promise /** * Describe the durations this adapter's model accepts. Returns a tagged @@ -234,10 +253,16 @@ export abstract class BaseVideoAdapter< >, ): Promise - abstract getVideoStatus(jobId: string): Promise + abstract getVideoStatus( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise /** See {@link VideoAdapter.getVideo}. */ - abstract getVideo(jobId: string): Promise + abstract getVideo( + jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise /** * Default implementation returns `{ kind: 'none' }`. Adapters that have diff --git a/packages/ai/src/activities/generateVideo/index.ts b/packages/ai/src/activities/generateVideo/index.ts index 3cfd2c3612..760fface19 100644 --- a/packages/ai/src/activities/generateVideo/index.ts +++ b/packages/ai/src/activities/generateVideo/index.ts @@ -495,6 +495,7 @@ async function runCreateVideoJob< const contextFor = (runId?: string): GenerationMiddlewareContext => createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'video', provider: adapter.name, model, @@ -613,6 +614,9 @@ async function* runStreamingVideoGeneration< timeout, abortSignal: callerAbortSignal, }) + const signalOptions = abortControls.signal + ? { abortSignal: abortControls.signal } + : undefined const providerName = (adapter as { name?: string; provider?: string }).provider ?? (adapter as { name?: string }).name ?? @@ -629,6 +633,7 @@ async function* runStreamingVideoGeneration< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'video', provider: adapter.name, model, @@ -656,6 +661,8 @@ async function* runStreamingVideoGeneration< // fired, so the `finally` below can fire one on abandonment without // double-firing. let settled = false + // Set once the provider has a job, so an abort can ask it to stop the job. + let jobId: string | undefined try { // Create the video generation job const jobResult = await raceWithAbort( @@ -671,6 +678,8 @@ async function* runStreamingVideoGeneration< abortControls.signal, ) + jobId = jobResult.jobId + yield { type: 'CUSTOM', name: 'video:job:created', @@ -684,7 +693,7 @@ async function* runStreamingVideoGeneration< await sleep(pollingInterval, abortControls.signal) const statusResult = await raceWithAbort( - adapter.getVideoStatus(jobResult.jobId), + adapter.getVideoStatus(jobResult.jobId, signalOptions), abortControls.signal, ) @@ -703,7 +712,7 @@ async function* runStreamingVideoGeneration< if (statusResult.status === 'completed') { // The caller can abort when it reads the status above. abortControls.signal?.throwIfAborted() - const download = adapter.getVideo(jobResult.jobId) + const download = adapter.getVideo(jobResult.jobId, signalOptions) const video = await raceWithAbort(download, abortControls.signal).catch( (error: unknown) => { // Nobody reads a video stream that comes after the abort, so @@ -787,6 +796,21 @@ async function* runStreamingVideoGeneration< settled = true const elapsed = Date.now() - obsStartTime if (isActivityAbortError(error, abortControls.signal)) { + // Cancel only on a real caller abort or timeout. A provider error + // named `TimeoutError` does not mean the caller wants the job stopped. + if (jobId !== undefined && abortControls.signal?.aborted) { + // Best effort: the run already ends as aborted, so a failed cancel + // is only logged. Not awaited, so a slow provider does not delay it. + const cancelledJobId = jobId + void (async () => adapter.cancelVideoJob?.(cancelledJobId))().catch( + (cancelError: unknown) => + logger.errors('generateVideo cancelVideoJob failed', { + jobId: cancelledJobId, + error: cancelError, + source: 'generateVideo', + }), + ) + } await runGenerationAbort(middleware, mwCtx, { reason: abortReasonMessage(error, abortControls.signal), duration: elapsed, @@ -856,6 +880,17 @@ export interface VideoJobStatusOptions< * charged for the wait. */ middleware?: Array + /** + * Maximum duration of this status call (status poll plus the video fetch) + * in milliseconds. Composed with {@link abortSignal}; the first abort wins. + */ + timeout?: number + /** + * Caller cancellation signal. Composed with {@link timeout} and forwarded to + * the adapter's status and video requests. An abort rejects the call; it + * does not mark the job or the run as failed. + */ + abortSignal?: AbortSignal } /** @@ -937,7 +972,25 @@ export interface VideoJobStatusResult { export async function getVideoJobStatus< TAdapter extends VideoAdapter, >(options: VideoJobStatusOptions): Promise { + const abortControls = createActivityAbortControls({ + timeout: options.timeout, + abortSignal: options.abortSignal, + }) + try { + return await runGetVideoJobStatus(options, abortControls.signal) + } finally { + abortControls.clear() + } +} + +async function runGetVideoJobStatus< + TAdapter extends VideoAdapter, +>( + options: VideoJobStatusOptions, + signal: AbortSignal | undefined, +): Promise { const { adapter, jobId, middleware } = options + const signalOptions = signal ? { abortSignal: signal } : undefined const requestId = createId('video-status') const startTime = Date.now() @@ -947,6 +1000,7 @@ export async function getVideoJobStatus< const terminalContext = (): GenerationMiddlewareContext => createGenerationContext({ requestId, + signal, activity: 'video', provider: adapter.name, model: adapter.model, @@ -971,7 +1025,10 @@ export async function getVideoJobStatus< }) // Get status first - const statusResult = await adapter.getVideoStatus(jobId) + const statusResult = await raceWithAbort( + adapter.getVideoStatus(jobId, signalOptions), + signal, + ) // If completed, also get the URL if (statusResult.status === 'completed') { @@ -980,8 +1037,14 @@ export async function getVideoJobStatus< // surface as itself, not be relabelled "failed to get video URL" and then // re-reported to the very middleware that threw. try { - video = await adapter.getVideo(jobId) + video = await raceWithAbort( + adapter.getVideo(jobId, signalOptions), + signal, + ) } catch (error) { + // A cancelled call says nothing about the job: rethrow, do not fail it. + // Without a caller abort, a provider `TimeoutError` still fails the job. + if (signal?.aborted) throw error const errorMessage = error instanceof Error ? error.message : 'Failed to get video URL' aiEventClient.emit('video:request:completed', { @@ -1051,6 +1114,7 @@ export async function getVideoJobStatus< try { result = await settleVideoUrl(transformed) } catch (error) { + if (signal?.aborted) throw error // The download broke mid-stream: the job is terminal, so fail the run. await runGenerationError(middleware, mwCtx, { error, diff --git a/packages/ai/src/activities/generateVoice/index.ts b/packages/ai/src/activities/generateVoice/index.ts index 6e1e67e77b..699f0bb454 100644 --- a/packages/ai/src/activities/generateVoice/index.ts +++ b/packages/ai/src/activities/generateVoice/index.ts @@ -237,6 +237,7 @@ async function runGenerateVoice< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'voice', provider: adapter.name, model, diff --git a/packages/ai/src/activities/generateWorld/index.ts b/packages/ai/src/activities/generateWorld/index.ts index 57c4204e11..9b9daa1279 100644 --- a/packages/ai/src/activities/generateWorld/index.ts +++ b/packages/ai/src/activities/generateWorld/index.ts @@ -203,6 +203,7 @@ async function runGenerateWorld< const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'world', provider: adapter.name, model, diff --git a/packages/ai/src/activities/middleware/run.ts b/packages/ai/src/activities/middleware/run.ts index a1793cda93..d5671f0a27 100644 --- a/packages/ai/src/activities/middleware/run.ts +++ b/packages/ai/src/activities/middleware/run.ts @@ -23,6 +23,7 @@ export function createGenerationContext(args: { threadId?: string runId?: string artifactInputs?: unknown + signal?: AbortSignal createId: (prefix: string) => string }): GenerationMiddlewareContext { return { @@ -38,6 +39,7 @@ export function createGenerationContext(args: { context: undefined, resultTransforms: [], artifactInputs: args.artifactInputs, + ...(args.signal && { signal: args.signal }), } } diff --git a/packages/ai/src/activities/middleware/types.ts b/packages/ai/src/activities/middleware/types.ts index 11ad54ace7..e61af694d3 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -86,6 +86,12 @@ export interface GenerationMiddlewareContext { createId: (prefix: string) => string /** Runtime context provided by the activity options, if any. */ context: TContext + /** + * The activity's effective abort signal (caller `abortSignal` plus + * `timeout`), when the call has one. Pass it to async work in a result + * transform, such as an upload, so a cancelled call stops that work too. + */ + signal?: AbortSignal /** * Result transforms registered by middleware during this activity call. * Transforms run after the raw adapter result exists and before the final diff --git a/packages/ai/src/activities/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts index 15ad1d16c2..f4d47e4e90 100644 --- a/packages/ai/src/activities/rerank/index.ts +++ b/packages/ai/src/activities/rerank/index.ts @@ -8,6 +8,11 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' import { isAbortShapedError } from '../error-payload' +import { + createActivityAbortControls, + isActivityAbortError, + raceWithAbort, +} from '../../utilities/activity-abort' import { createGenerationContext, runGenerationAbort, @@ -72,6 +77,12 @@ export interface RerankActivityOptions< modelOptions?: RerankProviderOptions /** Forwarded to the provider request for cancellation. */ abortSignal?: AbortSignal + /** + * Maximum duration of this activity invocation in milliseconds. + * No SDK-wide default. Composed with {@link abortSignal}; the first abort + * wins. + */ + timeout?: number /** * Observe-only middleware notified on start, usage, success, abort, and * error. Pass `otelMiddleware()` to emit OpenTelemetry spans, or implement @@ -107,6 +118,9 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { // abort hook whenever a shared/long-lived signal happens to already be // aborted, hiding it from `onError` observers. if (isAbortShapedError(error)) return true + // A `TimeoutError` is an abort only when our signal fired (the `timeout` + // option). A provider's own fetch timeout stays an error. + if (signal?.aborted && isActivityAbortError(error)) return true // Fall back to signal state only for non-Error throws we can't otherwise // identify; a real Error with a non-abort name is never an abort. return error instanceof Error ? false : signal?.aborted === true @@ -155,15 +169,7 @@ export async function rerank< >( options: RerankActivityOptions, ): Promise> { - const { - adapter, - query, - documents, - topN, - modelOptions, - abortSignal, - middleware, - } = options + const { adapter, query, documents, topN, modelOptions, middleware } = options const model = adapter.model const requestId = createId('rerank') const startTime = Date.now() @@ -172,9 +178,15 @@ export async function rerank< if (documents.length === 0) { throw new Error('rerank() requires at least one document') } + const abortControls = createActivityAbortControls({ + timeout: options.timeout, + abortSignal: options.abortSignal, + }) + const abortSignal = abortControls.signal const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, // `rerank` joins the GenerationActivity union; otel maps it to its own // gen_ai.operation.name. activity: 'rerank', @@ -184,7 +196,10 @@ export async function rerank< createId, }) - await runGenerationStart(middleware, mwCtx) + await runGenerationStart(middleware, mwCtx).catch((error: unknown) => { + abortControls.clear() + throw error + }) aiEventClient.emit('rerank:request:started', { requestId, @@ -200,20 +215,24 @@ export async function rerank< documentCount: documents.length, }) - // Serialize once; reuse for the request only. Original documents are mapped - // back by index below so the caller's element type is preserved. - const serialized = documents.map(serializeDocument) - try { - const result = await adapter.rerank({ - model, - query, - documents: serialized, - topN, - modelOptions, + // Serialize once; reuse for the request only. Original documents are + // mapped back by index below so the caller's element type is preserved. + // Inside the `try`, so a throw here also clears the timeout timer. + const serialized = documents.map(serializeDocument) + const result = await raceWithAbort( + adapter.rerank({ + model, + query, + documents: serialized, + topN, + modelOptions, + abortSignal, + logger, + }), abortSignal, - logger, - }) + ) + abortControls.clear() const ranking = result.ranking.map((r) => { const document = documents[r.index] @@ -263,6 +282,7 @@ export async function rerank< usage: result.usage, } } catch (error) { + abortControls.clear() const duration = Date.now() - startTime if (isAbortError(error, abortSignal)) { await runGenerationAbort(middleware, mwCtx, { diff --git a/packages/ai/src/activities/summarize/index.ts b/packages/ai/src/activities/summarize/index.ts index 09c28ab810..58f44fbb18 100644 --- a/packages/ai/src/activities/summarize/index.ts +++ b/packages/ai/src/activities/summarize/index.ts @@ -255,6 +255,7 @@ async function runSummarize( const mwCtx = createGenerationContext({ requestId, + signal: abortControls.signal, activity: 'summarize', provider: adapter.name, model, diff --git a/packages/ai/src/realtime/index.ts b/packages/ai/src/realtime/index.ts index bba5ad67cb..34e603be37 100644 --- a/packages/ai/src/realtime/index.ts +++ b/packages/ai/src/realtime/index.ts @@ -1,3 +1,4 @@ +import { raceWithAbort } from '../utilities/activity-abort' import type { RealtimeToken, RealtimeTokenOptions } from './types' export { createRealtimeEventEmitter } from './event-emitter' @@ -29,6 +30,11 @@ export type * from './types' export async function realtimeToken( options: RealtimeTokenOptions, ): Promise { - const { adapter } = options - return adapter.generateToken() + const { adapter, abortSignal } = options + return raceWithAbort( + abortSignal + ? adapter.generateToken({ abortSignal }) + : adapter.generateToken(), + abortSignal, + ) } diff --git a/packages/ai/src/realtime/types.ts b/packages/ai/src/realtime/types.ts index f519d2bd97..aae8a4547f 100644 --- a/packages/ai/src/realtime/types.ts +++ b/packages/ai/src/realtime/types.ts @@ -77,8 +77,13 @@ export interface RealtimeToken { export interface RealtimeTokenAdapter { /** Provider identifier */ provider: string - /** Generate an ephemeral token for client use */ - generateToken: () => Promise + /** + * Generate an ephemeral token for client use. Adapters should forward + * `options.abortSignal` to the provider request when supported. + */ + generateToken: (options?: { + abortSignal?: AbortSignal + }) => Promise } /** @@ -87,6 +92,10 @@ export interface RealtimeTokenAdapter { export interface RealtimeTokenOptions { /** The token adapter to use */ adapter: RealtimeTokenAdapter + /** + * Cancels the token request. The call rejects when the signal aborts. + */ + abortSignal?: AbortSignal } // ============================================================================ @@ -324,11 +333,14 @@ export interface RealtimeAdapter { * Create a connection using the provided token * @param token - The ephemeral token from the server * @param clientTools - Optional client-side tools to register with the provider + * @param options - Optional `abortSignal` that cancels the connection setup + * requests, when the provider supports it * @returns A connection instance */ connect: ( token: RealtimeToken, clientTools?: ReadonlyArray, + options?: { abortSignal?: AbortSignal }, ) => Promise } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index a873a8a2d7..758fdf4553 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -3355,6 +3355,12 @@ export interface EmbeddingOptions { * blocks. */ logger: InternalLogger + /** + * Effective abort signal composed by the activity from caller `abortSignal` + * and/or `timeout`. Adapters should forward this to the provider SDK when + * supported. Request-specific — never store on a global client config. + */ + abortSignal?: AbortSignal } /** diff --git a/packages/ai/tests/abort-signal-options.test.ts b/packages/ai/tests/abort-signal-options.test.ts new file mode 100644 index 0000000000..df10ad6999 --- /dev/null +++ b/packages/ai/tests/abort-signal-options.test.ts @@ -0,0 +1,557 @@ +import { describe, expect, it, vi } from 'vitest' +import { + choice, + decide, + deleteFile, + embed, + generateVideo, + getFile, + getVideoJobStatus, + realtimeToken, + rerank, + uploadFile, +} from '../src/index' +import { BaseEmbeddingAdapter } from '../src/activities/embed/adapter' +import type { EvaluateAdapter } from '../src/activities/evaluate/adapter' +import type { RerankAdapter } from '../src/activities/rerank/adapter' +import type { FilesAdapter } from '../src/activities/files/adapter' +import type { + GenerationAbortInfo, + GenerationMiddleware, + GenerationMiddlewareContext, +} from '../src/activities/middleware' +import type { RealtimeTokenAdapter } from '../src/realtime/types' +import type { + EmbeddingOptions, + EmbeddingResult, + StreamChunk, + VideoStatusResult, +} from '../src/types' + +/** A promise that settles only when `signal` aborts (never, without one). */ +function untilAbort(signal: AbortSignal | undefined): Promise { + return new Promise((_, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) +} + +function recorder() { + const starts: Array = [] + const aborts: Array = [] + const errors: Array = [] + const middleware: GenerationMiddleware = { + name: 'rec', + onStart: (ctx) => { + starts.push(ctx) + }, + onAbort: (_ctx, info) => { + aborts.push(info) + }, + onError: (_ctx, info) => { + errors.push(info.error) + }, + } + return { starts, aborts, errors, middleware } +} + +class HangingEmbeddingAdapter extends BaseEmbeddingAdapter<'mock-embed'> { + readonly name = 'mock' + calls: Array = [] + constructor() { + super('mock-embed', {}) + } + createEmbeddings(options: EmbeddingOptions): Promise { + this.calls.push(options) + if (!options.abortSignal) { + return Promise.resolve({ + id: 'e1', + model: options.model, + embeddings: [{ vector: [1], index: 0 }], + }) + } + return untilAbort(options.abortSignal) + } +} + +describe('embed() abortSignal and timeout', () => { + it('forwards the signal to the adapter and rejects on abort', async () => { + const adapter = new HangingEmbeddingAdapter() + const controller = new AbortController() + const rec = recorder() + const pending = embed({ + adapter, + input: 'a', + abortSignal: controller.signal, + middleware: [rec.middleware], + }) + await vi.waitFor(() => expect(adapter.calls).toHaveLength(1)) + expect(adapter.calls[0]!.abortSignal?.aborted).toBe(false) + expect(rec.starts[0]!.signal).toBe(adapter.calls[0]!.abortSignal) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(adapter.calls[0]!.abortSignal?.aborted).toBe(true) + expect(rec.aborts).toHaveLength(1) + expect(rec.errors).toHaveLength(0) + }) + + it('rejects with a TimeoutError when timeout elapses', async () => { + const adapter = new HangingEmbeddingAdapter() + await expect( + embed({ adapter, input: 'a', timeout: 5 }), + ).rejects.toMatchObject({ name: 'TimeoutError' }) + expect(adapter.calls[0]!.abortSignal?.aborted).toBe(true) + }) + + it('passes no signal when neither option is set', async () => { + const adapter = new HangingEmbeddingAdapter() + const rec = recorder() + const result = await embed({ + adapter, + input: 'a', + middleware: [rec.middleware], + }) + expect(result.embeddings).toHaveLength(1) + expect('abortSignal' in adapter.calls[0]!).toBe(false) + expect(rec.starts[0]!.signal).toBeUndefined() + }) +}) + +function videoAdapter(behavior: 'hang' | 'complete') { + return { + kind: 'video' as const, + name: 'test-video', + model: 'test-model', + '~types': {} as any, + availableDurations: () => ({ kind: 'none' as const }), + snapDuration: () => undefined, + createVideoJob: vi.fn(async () => ({ + jobId: 'job-1', + model: 'test-model', + })), + getVideoStatus: vi.fn( + ( + _jobId: string, + options?: { abortSignal?: AbortSignal }, + ): Promise => + behavior === 'hang' + ? untilAbort(options?.abortSignal) + : Promise.resolve({ jobId: 'job-1', status: 'completed' }), + ), + getVideo: vi.fn( + async (_jobId: string, _options?: { abortSignal?: AbortSignal }) => ({ + jobId: 'job-1', + url: 'https://example.com/v.mp4', + }), + ), + } +} + +describe('video job abortSignal', () => { + it('getVideoJobStatus forwards the signal and rejects on abort', async () => { + const adapter = videoAdapter('hang') + const controller = new AbortController() + const pending = getVideoJobStatus({ + adapter, + jobId: 'job-1', + abortSignal: controller.signal, + }) + await vi.waitFor(() => expect(adapter.getVideoStatus).toHaveBeenCalled()) + const passed = adapter.getVideoStatus.mock.calls[0]![1]?.abortSignal + expect(passed?.aborted).toBe(false) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(passed?.aborted).toBe(true) + }) + + it('getVideoJobStatus timeout rejects instead of reporting a failed job', async () => { + const adapter = videoAdapter('hang') + await expect( + getVideoJobStatus({ adapter, jobId: 'job-1', timeout: 5 }), + ).rejects.toMatchObject({ name: 'TimeoutError' }) + }) + + it('getVideoJobStatus passes the signal to getVideo, and none without it', async () => { + const adapter = videoAdapter('complete') + const controller = new AbortController() + await getVideoJobStatus({ + adapter, + jobId: 'job-1', + abortSignal: controller.signal, + }) + expect(adapter.getVideo.mock.calls[0]![1]?.abortSignal).toBeDefined() + + const plain = videoAdapter('complete') + const result = await getVideoJobStatus({ adapter: plain, jobId: 'job-1' }) + expect(result.status).toBe('completed') + expect(plain.getVideoStatus.mock.calls[0]![1]).toBeUndefined() + expect(plain.getVideo.mock.calls[0]![1]).toBeUndefined() + }) + + it('the streaming poll passes the signal to status and download', async () => { + const adapter = videoAdapter('complete') + const controller = new AbortController() + const chunks: Array = [] + for await (const chunk of generateVideo({ + adapter, + prompt: 'a cat', + threadId: 't1', + stream: true, + pollingInterval: 1, + abortSignal: controller.signal, + })) { + chunks.push(chunk) + } + expect(chunks.at(-1)?.type).toBe('RUN_FINISHED') + expect(adapter.getVideoStatus.mock.calls[0]![1]?.abortSignal).toBe( + controller.signal, + ) + expect(adapter.getVideo.mock.calls[0]![1]?.abortSignal).toBe( + controller.signal, + ) + }) +}) + +describe('generateVideo cancelVideoJob', () => { + async function abortAfterJob( + cancelVideoJob?: (jobId: string) => Promise, + ) { + const adapter = { ...videoAdapter('hang'), cancelVideoJob } + const controller = new AbortController() + const chunks: Array = [] + for await (const chunk of generateVideo({ + adapter, + prompt: 'a cat', + threadId: 't1', + stream: true, + pollingInterval: 1, + abortSignal: controller.signal, + debug: false, + })) { + chunks.push(chunk) + if (chunk.type === 'CUSTOM' && chunk.name === 'video:job:created') { + void vi + .waitFor(() => expect(adapter.getVideoStatus).toHaveBeenCalled()) + .then(() => controller.abort()) + } + } + return chunks + } + + it('asks the provider to cancel the job when the caller aborts', async () => { + const cancel = vi.fn(async (_jobId: string) => {}) + const chunks = await abortAfterJob(cancel) + expect(chunks.at(-1)?.type).toBe('RUN_ERROR') + expect(cancel).toHaveBeenCalledWith('job-1') + }) + + it('a failed cancel does not change the abort result', async () => { + const cancel = vi.fn(async (_jobId: string) => { + throw new Error('cancel failed') + }) + const chunks = await abortAfterJob(cancel) + const last = chunks.at(-1) + expect(last?.type).toBe('RUN_ERROR') + expect(last && 'message' in last && last.message).not.toContain( + 'cancel failed', + ) + expect(cancel).toHaveBeenCalledTimes(1) + }) + + it('an adapter without cancelVideoJob still aborts', async () => { + const chunks = await abortAfterJob(undefined) + expect(chunks.at(-1)?.type).toBe('RUN_ERROR') + }) + + it('a finished run does not cancel', async () => { + const cancel = vi.fn(async (_jobId: string) => {}) + const adapter = { ...videoAdapter('complete'), cancelVideoJob: cancel } + for await (const _chunk of generateVideo({ + adapter, + prompt: 'a cat', + threadId: 't1', + stream: true, + pollingInterval: 1, + abortSignal: new AbortController().signal, + })) { + // drain + } + expect(cancel).not.toHaveBeenCalled() + }) +}) + +describe('files abortSignal', () => { + function filesAdapter() { + const handle = { id: 'f1', provider: 'mock' as const } + const signals: Array = [] + const argCounts: Array = [] + const call = (...args: Array) => { + argCounts.push(args.length) + const options = args[1] as { abortSignal?: AbortSignal } | undefined + signals.push(options?.abortSignal) + return options?.abortSignal + ? untilAbort(options.abortSignal) + : Promise.resolve(handle) + } + const adapter: FilesAdapter<'mock'> & { kind: 'files' } = { + kind: 'files', + name: 'mock', + upload: (...args) => call(...args), + get: (...args) => call(...args), + delete: (...args) => call(...args).then(() => undefined), + } + return { adapter, signals, argCounts } + } + + it.each(['upload', 'get', 'delete'] as const)( + '%s forwards the signal and rejects on abort', + async (method) => { + const { adapter, signals } = filesAdapter() + const controller = new AbortController() + const pending = + method === 'upload' + ? uploadFile({ + adapter, + input: new Blob(['x']), + abortSignal: controller.signal, + }) + : method === 'get' + ? getFile({ adapter, id: 'f1', abortSignal: controller.signal }) + : deleteFile({ adapter, id: 'f1', abortSignal: controller.signal }) + expect(signals[0]).toBe(controller.signal) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + }, + ) + + it('calls the adapter without options when no signal is set', async () => { + const { adapter, argCounts } = filesAdapter() + await uploadFile({ adapter, input: new Blob(['x']) }) + await getFile({ adapter, id: 'f1' }) + await deleteFile({ adapter, id: 'f1' }) + expect(argCounts).toEqual([1, 1, 1]) + }) +}) + +describe('realtimeToken() abortSignal', () => { + it('forwards the signal and rejects on abort', async () => { + let received: AbortSignal | undefined + const adapter: RealtimeTokenAdapter = { + provider: 'mock', + generateToken: (options) => { + received = options?.abortSignal + return untilAbort(options?.abortSignal) + }, + } + const controller = new AbortController() + const pending = realtimeToken({ adapter, abortSignal: controller.signal }) + expect(received).toBe(controller.signal) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + }) + + it('calls generateToken with no arguments when no signal is set', async () => { + const generateToken = vi.fn(async () => ({ + provider: 'mock', + token: 't', + expiresAt: 0, + config: {}, + })) + await realtimeToken({ adapter: { provider: 'mock', generateToken } }) + expect(generateToken).toHaveBeenCalledWith() + }) +}) + +describe('decide() and rerank() timeout', () => { + it('decide() rejects on timeout and reports onAbort', async () => { + const rec = recorder() + let received: AbortSignal | undefined + const adapter: EvaluateAdapter = { + kind: 'evaluate', + name: 'mock', + model: 'mock-model', + '~types': { providerOptions: {} }, + evaluate: (opts) => { + received = opts.abortSignal + // Ignores the signal: the race still rejects. + return new Promise(() => {}) + }, + } + await expect( + decide({ + adapter, + state: { text: 'hi' }, + questions: { + kind: choice({ instructions: 'Pick one', options: { a: 'A' } }), + }, + timeout: 5, + middleware: [rec.middleware], + }), + ).rejects.toMatchObject({ name: 'TimeoutError' }) + expect(received?.aborted).toBe(true) + expect(rec.aborts).toHaveLength(1) + expect(rec.errors).toHaveLength(0) + }) + + it('rerank() rejects on timeout and reports onAbort', async () => { + const rec = recorder() + let received: AbortSignal | undefined + const adapter: RerankAdapter = { + kind: 'rerank', + name: 'mock', + model: 'mock-model', + '~types': { providerOptions: {} }, + rerank: (opts) => { + received = opts.abortSignal + return new Promise(() => {}) + }, + } + await expect( + rerank({ + adapter, + query: 'q', + documents: ['a'], + timeout: 5, + middleware: [rec.middleware], + }), + ).rejects.toMatchObject({ name: 'TimeoutError' }) + expect(received?.aborted).toBe(true) + expect(rec.aborts).toHaveLength(1) + }) +}) + +/** An error a provider SDK throws on its own fetch timeout. */ +function providerTimeout(): Error { + const error = new Error('provider timed out') + error.name = 'TimeoutError' + return error +} + +describe('a provider TimeoutError without a caller abort', () => { + it('generateVideo does not cancel the job', async () => { + const cancel = vi.fn(async (_jobId: string) => {}) + const adapter = { ...videoAdapter('complete'), cancelVideoJob: cancel } + adapter.getVideoStatus.mockRejectedValueOnce(providerTimeout()) + for await (const _chunk of generateVideo({ + adapter, + prompt: 'a cat', + threadId: 't1', + stream: true, + pollingInterval: 1, + debug: false, + })) { + // drain + } + expect(adapter.getVideoStatus).toHaveBeenCalled() + expect(cancel).not.toHaveBeenCalled() + }) + + it('getVideoJobStatus reports a failed job and calls onError', async () => { + const adapter = videoAdapter('complete') + adapter.getVideo.mockRejectedValueOnce(providerTimeout()) + const rec = recorder() + const result = await getVideoJobStatus({ + adapter, + jobId: 'job-1', + middleware: [rec.middleware], + }) + expect(result.status).toBe('failed') + expect(rec.errors).toHaveLength(1) + expect(rec.aborts).toHaveLength(0) + }) + + it('embed calls onError, not onAbort', async () => { + const adapter = new HangingEmbeddingAdapter() + adapter.createEmbeddings = () => Promise.reject(providerTimeout()) + const rec = recorder() + await expect( + embed({ adapter, input: 'a', middleware: [rec.middleware] }), + ).rejects.toThrow('provider timed out') + expect(rec.errors).toHaveLength(1) + expect(rec.aborts).toHaveLength(0) + }) + + it('rerank calls onError, not onAbort', async () => { + const adapter: RerankAdapter = { + kind: 'rerank', + name: 'mock', + model: 'mock-model', + '~types': { providerOptions: {} }, + rerank: () => Promise.reject(providerTimeout()), + } + const rec = recorder() + await expect( + rerank({ + adapter, + query: 'q', + documents: ['a'], + middleware: [rec.middleware], + }), + ).rejects.toThrow('provider timed out') + expect(rec.errors).toHaveLength(1) + expect(rec.aborts).toHaveLength(0) + }) +}) + +describe('the timeout timer when onStart throws', () => { + const failingStart: GenerationMiddleware = { + name: 'fail-start', + onStart: () => { + throw new Error('onStart failed') + }, + } + + it.each(['embed', 'rerank', 'decide'] as const)( + '%s clears the timer', + async (activity) => { + vi.useFakeTimers() + try { + const pending = + activity === 'embed' + ? embed({ + adapter: new HangingEmbeddingAdapter(), + input: 'a', + timeout: 60_000, + middleware: [failingStart], + }) + : activity === 'rerank' + ? rerank({ + adapter: { + kind: 'rerank', + name: 'mock', + model: 'mock-model', + '~types': { providerOptions: {} }, + rerank: () => new Promise(() => {}), + }, + query: 'q', + documents: ['a'], + timeout: 60_000, + middleware: [failingStart], + }) + : decide({ + adapter: { + kind: 'evaluate', + name: 'mock', + model: 'mock-model', + '~types': { providerOptions: {} }, + evaluate: () => new Promise(() => {}), + }, + state: { text: 'hi' }, + questions: { + kind: choice({ + instructions: 'Pick one', + options: { a: 'A' }, + }), + }, + timeout: 60_000, + middleware: [failingStart], + }) + await expect(pending).rejects.toThrow('onStart failed') + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }, + ) +}) diff --git a/packages/ai/tests/chat-mcp-manager.test.ts b/packages/ai/tests/chat-mcp-manager.test.ts index d236bf828c..d47d37a386 100644 --- a/packages/ai/tests/chat-mcp-manager.test.ts +++ b/packages/ai/tests/chat-mcp-manager.test.ts @@ -85,6 +85,24 @@ describe('MCPManager', () => { expect(spyA).toHaveBeenCalledWith({ lazy: true }) }) + it('discover(signal) forwards the signal to every source', async () => { + const a = source([tool('a')]) + const spyA = vi.spyOn(a, 'tools') + const controller = new AbortController() + await MCPManager.from({ clients: [a] }).discover(controller.signal) + expect(spyA).toHaveBeenCalledWith({ + lazy: false, + signal: controller.signal, + }) + }) + + it('discover() without a signal sends no signal key', async () => { + const a = source([tool('a')]) + const spyA = vi.spyOn(a, 'tools') + await MCPManager.from({ clients: [a] }).discover() + expect(spyA).toHaveBeenCalledWith({ lazy: false }) + }) + it('discover() throws MCPDuplicateToolNameError on collision', async () => { const m = MCPManager.from({ clients: [source([tool('x')]), source([tool('x')])], diff --git a/packages/ai/tests/chat-mcp.test.ts b/packages/ai/tests/chat-mcp.test.ts index a7189ca5d7..696e6c201b 100644 --- a/packages/ai/tests/chat-mcp.test.ts +++ b/packages/ai/tests/chat-mcp.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { chat } from '../src/activities/chat/index' import { MCPDuplicateToolNameError } from '../src/activities/chat/mcp/manager' -import { collectChunks, createMockAdapter, ev } from './test-utils' +import { clientTool, collectChunks, createMockAdapter, ev } from './test-utils' import type { StreamChunk } from '../src/types' import type { MCPToolSource } from '../src/activities/chat/mcp/types' import type { ServerTool } from '../src/activities/chat/tools/tool-definition' @@ -298,6 +298,37 @@ describe('chat({ mcp })', () => { expect(toolsSpy).toHaveBeenCalledWith({ lazy: true }) }) + it('passes the run abort signal to source.tools()', async () => { + const toolsSpy = vi.fn().mockResolvedValue([]) + const source: MCPToolSource = { tools: toolsSpy, close: async () => {} } + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.textStart(), + ev.textContent('ok'), + ev.textEnd(), + ev.runFinished('stop'), + ], + ], + }) + const abortController = new AbortController() + + const stream = chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'hi' }], + mcp: { clients: [source] }, + abortController, + }) + await collectChunks(stream as AsyncIterable) + + expect(toolsSpy).toHaveBeenCalledWith({ + lazy: false, + signal: abortController.signal, + }) + }) + // -------------------------------------------------------------------------- // Case 7a: onDiscoveryError returning skips failed source, run proceeds // -------------------------------------------------------------------------- @@ -579,3 +610,98 @@ describe('chat({ mcp }) — structured-output runners', () => { expect(mcpExecuteSpy).toHaveBeenCalledTimes(1) }) }) + +describe('chat({ mcp }) abort during discovery', () => { + /** A source whose discovery the caller aborts; it then rejects like an SDK. */ + function abortingSource(abortController: AbortController): MCPToolSource { + return { + tools: (options) => { + abortController.abort() + return Promise.reject(options?.signal?.reason) + }, + close: async () => {}, + } + } + + it('ends the stream through the abort path and calls onAbort', async () => { + const { adapter } = createMockAdapter({ iterations: [] }) + const abortController = new AbortController() + const onAbort = vi.fn() + const onError = vi.fn() + + const stream = chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'hi' }], + mcp: { clients: [abortingSource(abortController)] }, + middleware: [{ name: 'observer', onAbort, onError }], + abortController, + }) + await collectChunks(stream as AsyncIterable) + + expect(onAbort).toHaveBeenCalledOnce() + expect(onError).not.toHaveBeenCalled() + }) + + it('raises no tool interrupt for a pending client tool after the abort', async () => { + const { adapter, calls } = createMockAdapter({ iterations: [] }) + const abortController = new AbortController() + const onAbort = vi.fn() + + const stream = chat({ + threadId: 'thread-1', + adapter, + messages: [ + { role: 'user', content: 'Notify me' }, + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'call_client', + type: 'function' as const, + function: { + name: 'showNotification', + arguments: '{"message":"done"}', + }, + }, + ], + }, + ], + tools: [clientTool('showNotification')], + mcp: { clients: [abortingSource(abortController)] }, + middleware: [{ name: 'observer', onAbort }], + abortController, + }) + const chunks = await collectChunks(stream as AsyncIterable) + + expect(onAbort).toHaveBeenCalledOnce() + expect( + chunks.some( + (c) => + c.type === 'RUN_FINISHED' && + 'outcome' in c && + c.outcome?.type === 'interrupt', + ), + ).toBe(false) + expect(calls).toHaveLength(0) + }) + + it('calls onAbort on the structured output path too', async () => { + const { adapter } = createMockAdapter({ iterations: [] }) + const abortController = new AbortController() + const onAbort = vi.fn() + + await chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'hi' }], + mcp: { clients: [abortingSource(abortController)] }, + middleware: [{ name: 'observer', onAbort }], + outputSchema: z.object({ ok: z.boolean() }), + abortController, + }).catch(() => undefined) + + expect(onAbort).toHaveBeenCalledOnce() + }) +}) diff --git a/packages/ai/tests/evaluate.test.ts b/packages/ai/tests/evaluate.test.ts index 587dabdff5..7a36387b3e 100644 --- a/packages/ai/tests/evaluate.test.ts +++ b/packages/ai/tests/evaluate.test.ts @@ -690,7 +690,9 @@ describe('decide()', () => { const { middleware, events } = recordingMiddleware() const controller = new AbortController() const adapter = mockEvaluateAdapter(async () => { - controller.abort() + // Abort just after the genuine error settles the call. An abort that + // comes first wins the activity's abort race instead. + setTimeout(() => controller.abort(), 0) throw new Error('genuine provider failure') }) diff --git a/packages/ai/tests/rerank.test.ts b/packages/ai/tests/rerank.test.ts index 7c89c8ae3f..8ebf64de3c 100644 --- a/packages/ai/tests/rerank.test.ts +++ b/packages/ai/tests/rerank.test.ts @@ -218,7 +218,9 @@ describe('rerank() activity', () => { const { middleware, events } = recordingMiddleware() const controller = new AbortController() const adapter = mockRerankAdapter(async () => { - controller.abort() + // Abort just after the genuine error settles the call. An abort that + // comes first wins the activity's abort race instead. + setTimeout(() => controller.abort(), 0) throw new Error('genuine provider failure') }) diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts index 3a1a613bac..03e7b08737 100644 --- a/testing/e2e/global-setup.ts +++ b/testing/e2e/global-setup.ts @@ -189,6 +189,11 @@ export default async function globalSetup() { // Gateway adapter keeps only the origin of its baseURL, so that mount sits on // the real evaluate path. Both answer only requests carrying the // `x-evaluate-http-status` header and fall through otherwise. + // embed() abort: `/v1/embeddings` never answers, and `/status` reports + // whether the client closed the request. The embedding-abort spec checks + // that `embed({ timeout })` cancels the real HTTP request. + mock.mount('/embedding-abort', embeddingAbortMount()) + mock.mount('/evaluate-http-status', evaluateHttpStatusMount()) mock.mount('/v4/ai/evaluation-model', evaluateHttpStatusMount()) @@ -2353,6 +2358,46 @@ function evaluateHttpStatusMount(): Mountable { } } +/** + * Never answers `POST /v1/embeddings`, and records per `X-Test-Id` whether the + * client closed that request before a reply. `GET /status?testId=` reads it. + */ +function embeddingAbortMount(): Mountable { + const closedByClient = new Map() + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + pathname: string, + ): Promise { + if (req.method === 'GET' && pathname === '/status') { + const testId = new URL(req.url ?? '', 'http://x').searchParams.get( + 'testId', + ) + res.setHeader('Content-Type', 'application/json') + res.end( + JSON.stringify({ + closedByClient: closedByClient.get(testId ?? '') ?? false, + }), + ) + return true + } + if (req.method !== 'POST' || pathname !== '/v1/embeddings') return false + const testId = String(req.headers['x-test-id'] ?? '') + closedByClient.set(testId, false) + await drainBody(req) + res.on('close', () => { + if (!res.writableEnded) closedByClient.set(testId, true) + }) + // Safety net so a broken client cannot hold the socket forever. + setTimeout(() => { + if (!res.writableEnded) res.end() + }, 30_000).unref() + return true + }, + } +} + /** * Mounts the OpenAI Decisions `/v1/decisions` response with a `refusal` answer * and billed usage. The OpenAI evaluate adapter turns it into an diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index d34374909e..96f4648d3e 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -91,9 +91,11 @@ import { Route as ApiCopilotWireRouteImport } from './routes/api.copilot-wire' import { Route as ApiCrossModelReplayWireRouteImport } from './routes/api.cross-model-replay-wire' import { Route as ApiDevtoolsMemoryRouteImport } from './routes/api.devtools-memory' import { Route as ApiDurableDeliveryRouteImport } from './routes/api.durable-delivery' +import { Route as ApiDurableDeliveryCancelRouteImport } from './routes/api.durable-delivery-cancel' import { Route as ApiDurableMcpInputRouteImport } from './routes/api.durable-mcp-input' import { Route as ApiDurableTakeoverRouteImport } from './routes/api.durable-takeover' import { Route as ApiEmbeddingRouteImport } from './routes/api.embedding' +import { Route as ApiEmbeddingAbortRouteImport } from './routes/api.embedding-abort' import { Route as ApiEvaluateHttpStatusRouteImport } from './routes/api.evaluate-http-status' import { Route as ApiEvaluateRefusalRouteImport } from './routes/api.evaluate-refusal' import { Route as ApiEvaluateUsageOnFailureRouteImport } from './routes/api.evaluate-usage-on-failure' @@ -645,6 +647,12 @@ const ApiDurableDeliveryRoute = ApiDurableDeliveryRouteImport.update({ path: '/api/durable-delivery', getParentRoute: () => rootRouteImport, } as any) +const ApiDurableDeliveryCancelRoute = + ApiDurableDeliveryCancelRouteImport.update({ + id: '/api/durable-delivery-cancel', + path: '/api/durable-delivery-cancel', + getParentRoute: () => rootRouteImport, + } as any) const ApiDurableMcpInputRoute = ApiDurableMcpInputRouteImport.update({ id: '/api/durable-mcp-input', path: '/api/durable-mcp-input', @@ -660,6 +668,11 @@ const ApiEmbeddingRoute = ApiEmbeddingRouteImport.update({ path: '/api/embedding', getParentRoute: () => rootRouteImport, } as any) +const ApiEmbeddingAbortRoute = ApiEmbeddingAbortRouteImport.update({ + id: '/api/embedding-abort', + path: '/api/embedding-abort', + getParentRoute: () => rootRouteImport, +} as any) const ApiEvaluateHttpStatusRoute = ApiEvaluateHttpStatusRouteImport.update({ id: '/api/evaluate-http-status', path: '/api/evaluate-http-status', @@ -1367,9 +1380,11 @@ export interface FileRoutesByFullPath { '/api/cross-model-replay-wire': typeof ApiCrossModelReplayWireRoute '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute + '/api/durable-delivery-cancel': typeof ApiDurableDeliveryCancelRoute '/api/durable-mcp-input': typeof ApiDurableMcpInputRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute '/api/embedding': typeof ApiEmbeddingRoute + '/api/embedding-abort': typeof ApiEmbeddingAbortRoute '/api/evaluate-http-status': typeof ApiEvaluateHttpStatusRoute '/api/evaluate-refusal': typeof ApiEvaluateRefusalRoute '/api/evaluate-usage-on-failure': typeof ApiEvaluateUsageOnFailureRoute @@ -1573,9 +1588,11 @@ export interface FileRoutesByTo { '/api/cross-model-replay-wire': typeof ApiCrossModelReplayWireRoute '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute + '/api/durable-delivery-cancel': typeof ApiDurableDeliveryCancelRoute '/api/durable-mcp-input': typeof ApiDurableMcpInputRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute '/api/embedding': typeof ApiEmbeddingRoute + '/api/embedding-abort': typeof ApiEmbeddingAbortRoute '/api/evaluate-http-status': typeof ApiEvaluateHttpStatusRoute '/api/evaluate-refusal': typeof ApiEvaluateRefusalRoute '/api/evaluate-usage-on-failure': typeof ApiEvaluateUsageOnFailureRoute @@ -1780,9 +1797,11 @@ export interface FileRoutesById { '/api/cross-model-replay-wire': typeof ApiCrossModelReplayWireRoute '/api/devtools-memory': typeof ApiDevtoolsMemoryRoute '/api/durable-delivery': typeof ApiDurableDeliveryRoute + '/api/durable-delivery-cancel': typeof ApiDurableDeliveryCancelRoute '/api/durable-mcp-input': typeof ApiDurableMcpInputRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute '/api/embedding': typeof ApiEmbeddingRoute + '/api/embedding-abort': typeof ApiEmbeddingAbortRoute '/api/evaluate-http-status': typeof ApiEvaluateHttpStatusRoute '/api/evaluate-refusal': typeof ApiEvaluateRefusalRoute '/api/evaluate-usage-on-failure': typeof ApiEvaluateUsageOnFailureRoute @@ -1988,9 +2007,11 @@ export interface FileRouteTypes { | '/api/cross-model-replay-wire' | '/api/devtools-memory' | '/api/durable-delivery' + | '/api/durable-delivery-cancel' | '/api/durable-mcp-input' | '/api/durable-takeover' | '/api/embedding' + | '/api/embedding-abort' | '/api/evaluate-http-status' | '/api/evaluate-refusal' | '/api/evaluate-usage-on-failure' @@ -2194,9 +2215,11 @@ export interface FileRouteTypes { | '/api/cross-model-replay-wire' | '/api/devtools-memory' | '/api/durable-delivery' + | '/api/durable-delivery-cancel' | '/api/durable-mcp-input' | '/api/durable-takeover' | '/api/embedding' + | '/api/embedding-abort' | '/api/evaluate-http-status' | '/api/evaluate-refusal' | '/api/evaluate-usage-on-failure' @@ -2400,9 +2423,11 @@ export interface FileRouteTypes { | '/api/cross-model-replay-wire' | '/api/devtools-memory' | '/api/durable-delivery' + | '/api/durable-delivery-cancel' | '/api/durable-mcp-input' | '/api/durable-takeover' | '/api/embedding' + | '/api/embedding-abort' | '/api/evaluate-http-status' | '/api/evaluate-refusal' | '/api/evaluate-usage-on-failure' @@ -2607,9 +2632,11 @@ export interface RootRouteChildren { ApiCrossModelReplayWireRoute: typeof ApiCrossModelReplayWireRoute ApiDevtoolsMemoryRoute: typeof ApiDevtoolsMemoryRoute ApiDurableDeliveryRoute: typeof ApiDurableDeliveryRoute + ApiDurableDeliveryCancelRoute: typeof ApiDurableDeliveryCancelRoute ApiDurableMcpInputRoute: typeof ApiDurableMcpInputRoute ApiDurableTakeoverRoute: typeof ApiDurableTakeoverRoute ApiEmbeddingRoute: typeof ApiEmbeddingRoute + ApiEmbeddingAbortRoute: typeof ApiEmbeddingAbortRoute ApiEvaluateHttpStatusRoute: typeof ApiEvaluateHttpStatusRoute ApiEvaluateRefusalRoute: typeof ApiEvaluateRefusalRoute ApiEvaluateUsageOnFailureRoute: typeof ApiEvaluateUsageOnFailureRoute @@ -3303,6 +3330,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiDurableDeliveryRouteImport parentRoute: typeof rootRouteImport } + '/api/durable-delivery-cancel': { + id: '/api/durable-delivery-cancel' + path: '/api/durable-delivery-cancel' + fullPath: '/api/durable-delivery-cancel' + preLoaderRoute: typeof ApiDurableDeliveryCancelRouteImport + parentRoute: typeof rootRouteImport + } '/api/durable-mcp-input': { id: '/api/durable-mcp-input' path: '/api/durable-mcp-input' @@ -3324,6 +3358,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiEmbeddingRouteImport parentRoute: typeof rootRouteImport } + '/api/embedding-abort': { + id: '/api/embedding-abort' + path: '/api/embedding-abort' + fullPath: '/api/embedding-abort' + preLoaderRoute: typeof ApiEmbeddingAbortRouteImport + parentRoute: typeof rootRouteImport + } '/api/evaluate-http-status': { id: '/api/evaluate-http-status' path: '/api/evaluate-http-status' @@ -4305,9 +4346,11 @@ const rootRouteChildren: RootRouteChildren = { ApiCrossModelReplayWireRoute: ApiCrossModelReplayWireRoute, ApiDevtoolsMemoryRoute: ApiDevtoolsMemoryRoute, ApiDurableDeliveryRoute: ApiDurableDeliveryRoute, + ApiDurableDeliveryCancelRoute: ApiDurableDeliveryCancelRoute, ApiDurableMcpInputRoute: ApiDurableMcpInputRoute, ApiDurableTakeoverRoute: ApiDurableTakeoverRoute, ApiEmbeddingRoute: ApiEmbeddingRoute, + ApiEmbeddingAbortRoute: ApiEmbeddingAbortRoute, ApiEvaluateHttpStatusRoute: ApiEvaluateHttpStatusRoute, ApiEvaluateRefusalRoute: ApiEvaluateRefusalRoute, ApiEvaluateUsageOnFailureRoute: ApiEvaluateUsageOnFailureRoute, diff --git a/testing/e2e/src/routes/api.durable-delivery-cancel.ts b/testing/e2e/src/routes/api.durable-delivery-cancel.ts new file mode 100644 index 0000000000..c29b1dcf5e --- /dev/null +++ b/testing/e2e/src/routes/api.durable-delivery-cancel.ts @@ -0,0 +1,130 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + RUN_CANCEL_REASON, + chatParamsFromRequest, + memoryStream, + toServerSentEventsResponse, +} from '@tanstack/ai' +import type { StreamChunk } from '@tanstack/ai' + +/** + * Provider-free harness for the opt-in durable HTTP cancel (`cancelUrl`). + * + * - `POST` → a slow durable run. Its controller is kept in `liveRuns`. + * - `POST ?op=cancel` with `{ runId }` → aborts that run with + * `RUN_CANCEL_REASON`, like an app's cancel route. + * - `GET ?runId=…` → `{ ending }`: `'finished'`, the abort reason, or `null` + * while the run is still going. + */ +const liveRuns = new Map() +const endings = new Map() + +function sleep(ms: number, signal: AbortSignal): Promise { + return new Promise((resolve) => { + const timer = setTimeout(resolve, ms) + signal.addEventListener( + 'abort', + () => { + clearTimeout(timer) + resolve() + }, + { once: true }, + ) + }) +} + +/** Ten content chunks, 100ms apart. Stops early when `signal` aborts. */ +function slowRun( + threadId: string, + runId: string, + signal: AbortSignal, +): AsyncIterable { + return (async function* () { + try { + yield { + type: 'RUN_STARTED', + threadId, + runId, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'TEXT_MESSAGE_START', + messageId: 'm', + role: 'assistant', + timestamp: Date.now(), + } as StreamChunk + for (let i = 1; i <= 10; i++) { + await sleep(100, signal) + if (signal.aborted) return + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId: 'm', + model: 'fixed', + delta: String(i), + content: String(i), + timestamp: Date.now(), + } as StreamChunk + } + yield { + type: 'TEXT_MESSAGE_END', + messageId: 'm', + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'RUN_FINISHED', + threadId, + runId, + model: 'fixed', + finishReason: 'stop', + timestamp: Date.now(), + } as StreamChunk + endings.set(runId, 'finished') + } finally { + if (signal.aborted) endings.set(runId, String(signal.reason)) + liveRuns.delete(runId) + } + })() +} + +export const Route = createFileRoute('/api/durable-delivery-cancel')({ + server: { + handlers: { + POST: async ({ request }) => { + const url = new URL(request.url) + if (url.searchParams.get('op') === 'cancel') { + const body: unknown = await request.json() + if (typeof body !== 'object' || body === null || !('runId' in body)) { + return new Response('runId is required', { status: 400 }) + } + const runId = body.runId + if (typeof runId !== 'string') { + return new Response('runId must be a string', { status: 400 }) + } + liveRuns.get(runId)?.abort(RUN_CANCEL_REASON) + return new Response(null, { status: 204 }) + } + + const { threadId, runId } = await chatParamsFromRequest(request) + const abortController = new AbortController() + liveRuns.set(runId, abortController) + // memoryStream keys the log by `?runId`. + url.searchParams.set('runId', runId) + return toServerSentEventsResponse( + slowRun(threadId, runId, abortController.signal), + { + abortController, + durability: { + adapter: memoryStream( + new Request(url, { headers: request.headers }), + ), + }, + }, + ) + }, + GET: ({ request }) => { + const runId = new URL(request.url).searchParams.get('runId') ?? '' + return Response.json({ ending: endings.get(runId) ?? null }) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.embedding-abort.ts b/testing/e2e/src/routes/api.embedding-abort.ts new file mode 100644 index 0000000000..b33451128b --- /dev/null +++ b/testing/e2e/src/routes/api.embedding-abort.ts @@ -0,0 +1,57 @@ +import { createFileRoute } from '@tanstack/react-router' +import { embed } from '@tanstack/ai' +import { createOpenaiEmbedding } from '@tanstack/ai-openai' + +const LLMOCK_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' +const DUMMY_KEY = 'sk-e2e-test-dummy-key' + +/** + * Calls `embed({ timeout })` against a mount that never answers, then asks the + * mount whether the client closed the HTTP request. A closed request proves + * the timeout signal reached the provider fetch. + */ +export const Route = createFileRoute('/api/embedding-abort')({ + server: { + handlers: { + POST: async ({ request }) => { + const testId = new URL(request.url).searchParams.get('testId') ?? '' + const base = `${LLMOCK_BASE}/embedding-abort` + let errorName: string | undefined + try { + await embed({ + adapter: createOpenaiEmbedding( + 'text-embedding-3-small', + DUMMY_KEY, + { + baseURL: `${base}/v1`, + defaultHeaders: { 'X-Test-Id': testId }, + maxRetries: 0, + }, + ), + input: 'a red guitar', + timeout: 300, + }) + } catch (error) { + errorName = error instanceof Error ? error.name : String(error) + } + + // The socket close can land just after the rejection. Poll briefly. + let closedByClient = false + for (let i = 0; i < 20 && !closedByClient; i++) { + const status = await fetch( + `${base}/status?testId=${encodeURIComponent(testId)}`, + ) + const body: unknown = await status.json() + closedByClient = + typeof body === 'object' && + body !== null && + 'closedByClient' in body && + body.closedByClient === true + if (!closedByClient) await new Promise((r) => setTimeout(r, 50)) + } + + return Response.json({ errorName, closedByClient }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/delivery-durability-cancel.spec.ts b/testing/e2e/tests/delivery-durability-cancel.spec.ts new file mode 100644 index 0000000000..d179598272 --- /dev/null +++ b/testing/e2e/tests/delivery-durability-cancel.spec.ts @@ -0,0 +1,62 @@ +import { RUN_CANCEL_REASON } from '@tanstack/ai' +import { ChatClient, fetchServerSentEvents } from '@tanstack/ai-client' +import { test, expect } from './fixtures' +import type { APIRequestContext } from '@playwright/test' + +const ROUTE = 'http://localhost:3010/api/durable-delivery-cancel' + +function startRun(testId: string) { + const client = new ChatClient({ + threadId: `thread-durable-cancel-${testId}`, + connection: fetchServerSentEvents(ROUTE, { + cancelUrl: `${ROUTE}?op=cancel`, + }), + }) + void client.sendMessage('go') + return client +} + +async function ending( + request: APIRequestContext, + runId: string, +): Promise { + const response = await request.get( + `${ROUTE}?runId=${encodeURIComponent(runId)}`, + ) + const body = await response.json() + return typeof body === 'object' && body !== null && 'ending' in body + ? body.ending + : undefined +} + +test.describe('durable HTTP delivery: cancelUrl', () => { + test('stop() cancels the server run', async ({ request, testId }) => { + const client = startRun(testId) + try { + await expect.poll(() => client.getSnapshot().status).toBe('streaming') + const runId = client.getCurrentRunId() + expect(runId).toEqual(expect.any(String)) + + client.stop() + + await expect + .poll(() => ending(request, runId ?? '')) + .toBe(RUN_CANCEL_REASON) + } finally { + client.unsubscribe() + } + }) + + test('an unsubscribe without stop() lets the run finish', async ({ + request, + testId, + }) => { + const client = startRun(testId) + await expect.poll(() => client.getSnapshot().status).toBe('streaming') + const runId = client.getCurrentRunId() + + client.unsubscribe() + + await expect.poll(() => ending(request, runId ?? '')).toBe('finished') + }) +}) diff --git a/testing/e2e/tests/embedding-abort.spec.ts b/testing/e2e/tests/embedding-abort.spec.ts new file mode 100644 index 0000000000..f8fed24803 --- /dev/null +++ b/testing/e2e/tests/embedding-abort.spec.ts @@ -0,0 +1,17 @@ +import { test, expect } from './fixtures' + +test.describe('embed() timeout', () => { + test('rejects and cancels the provider request', async ({ + request, + testId, + }) => { + const response = await request.post( + `/api/embedding-abort?testId=${encodeURIComponent(testId)}`, + ) + expect(response.ok()).toBe(true) + expect(await response.json()).toEqual({ + errorName: 'TimeoutError', + closedByClient: true, + }) + }) +})