From d7faa4911a7484c6081e315c7fa0de019c4d8337 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:28:23 +0100 Subject: [PATCH 01/12] feat(ai, ai-client, adapters): optional abortSignal and timeout for embed, video, files, realtime, decide, rerank Add optional cancel options where the public contracts had none. Every new field is optional, so old calls compile and behave the same. - embed(): abortSignal and timeout, raced like the other activities, and EmbeddingOptions.abortSignal for adapters. - Video: getVideoStatus/getVideo take an optional { abortSignal }. The streaming poll passes it. getVideoJobStatus() takes abortSignal and timeout; an abort rejects and does not fail the job. - Files: uploadFile/getFile/deleteFile and FilesAdapter methods take an optional abortSignal. uploadGeminiFile() takes abortSignal. - Realtime: realtimeToken() and generateToken()/connect() take an optional abortSignal. - decide() and rerank(): timeout plus the abort race. - GenerationMiddlewareContext.signal. - ChatClient client tools get an abortSignal that aborts on stop(). - Adapters forward the signal: openai, lovable, neon, vercel-gateway, cloudflare, cohere, bedrock, gemini, mistral, ollama, grok, byteplus, openrouter, fal, anthropic, elevenlabs. --- .changeset/ai-abort-signal-options.md | 12 + .../ai-anthropic-abort-signal-options.md | 5 + .changeset/ai-bedrock-abort-signal-options.md | 5 + .../ai-byteplus-abort-signal-options.md | 5 + .changeset/ai-client-abort-signal-options.md | 5 + .../ai-cloudflare-abort-signal-options.md | 5 + .changeset/ai-cohere-abort-signal-options.md | 5 + .../ai-elevenlabs-abort-signal-options.md | 5 + .changeset/ai-fal-abort-signal-options.md | 5 + .changeset/ai-gemini-abort-signal-options.md | 5 + .changeset/ai-grok-abort-signal-options.md | 5 + .changeset/ai-lovable-abort-signal-options.md | 5 + .changeset/ai-mistral-abort-signal-options.md | 5 + .changeset/ai-neon-abort-signal-options.md | 5 + .changeset/ai-ollama-abort-signal-options.md | 5 + .changeset/ai-openai-abort-signal-options.md | 5 + .../ai-openrouter-abort-signal-options.md | 5 + .../ai-vercel-gateway-abort-signal-options.md | 5 + docs/advanced/files-api.md | 20 + docs/advanced/otel.md | 2 + docs/api/ai.md | 1 + docs/config.json | 17 +- docs/embeddings.md | 21 ++ docs/evaluate/evaluate.md | 3 +- docs/media/realtime-chat.md | 17 + docs/media/video-generation-advanced.md | 16 + docs/rerank/rerank.md | 3 +- docs/tools/client-tools.md | 20 + packages/ai-anthropic/src/adapters/files.ts | 42 ++- .../ai-anthropic/tests/abort-signal.test.ts | 47 +++ packages/ai-bedrock/src/adapters/embedding.ts | 8 +- .../ai-bedrock/tests/abort-signal.test.ts | 53 +++ packages/ai-byteplus/src/adapters/video.ts | 27 +- .../ai-byteplus/tests/abort-signal.test.ts | 50 +++ packages/ai-client/src/chat-client.ts | 5 + .../chat-client-tool-abort-signal.test.ts | 63 ++++ .../ai-cloudflare/src/adapters/embedding.ts | 10 +- .../ai-cloudflare/tests/abort-signal.test.ts | 57 +++ packages/ai-cohere/src/adapters/embedding.ts | 18 +- packages/ai-cohere/tests/abort-signal.test.ts | 59 +++ packages/ai-elevenlabs/src/realtime/token.ts | 5 +- .../ai-elevenlabs/tests/abort-signal.test.ts | 54 +++ packages/ai-fal/src/adapters/video.ts | 14 +- .../ai-fal/tests/video-abort-signal.test.ts | 38 ++ packages/ai-gemini/src/adapters/embedding.ts | 3 + packages/ai-gemini/src/adapters/files.ts | 38 +- packages/ai-gemini/src/adapters/video.ts | 45 ++- packages/ai-gemini/src/files/index.ts | 37 +- packages/ai-gemini/src/realtime/token.ts | 5 +- packages/ai-gemini/tests/abort-signal.test.ts | 98 +++++ .../tests/upload-gemini-file-abort.test.ts | 79 ++++ packages/ai-grok/src/adapters/files.ts | 33 +- packages/ai-grok/src/adapters/video.ts | 24 +- packages/ai-grok/src/realtime/adapter.ts | 5 +- packages/ai-grok/src/realtime/token.ts | 10 +- packages/ai-grok/tests/abort-signal.test.ts | 74 ++++ packages/ai-lovable/src/adapters/embedding.ts | 6 +- packages/ai-lovable/src/adapters/video.ts | 21 +- .../ai-lovable/tests/abort-signal.test.ts | 71 ++++ packages/ai-mistral/src/adapters/embedding.ts | 6 +- .../ai-mistral/tests/abort-signal.test.ts | 60 +++ packages/ai-neon/src/adapters/embedding.ts | 6 +- packages/ai-neon/tests/abort-signal.test.ts | 57 +++ packages/ai-ollama/src/adapters/embedding.ts | 17 +- packages/ai-ollama/tests/abort-signal.test.ts | 58 +++ packages/ai-openai/src/adapters/embedding.ts | 6 +- packages/ai-openai/src/adapters/files.ts | 37 +- packages/ai-openai/src/adapters/video.ts | 28 +- packages/ai-openai/src/realtime/adapter.ts | 5 +- packages/ai-openai/src/realtime/token.ts | 5 +- packages/ai-openai/tests/abort-signal.test.ts | 98 +++++ packages/ai-openrouter/src/adapters/video.ts | 32 +- .../tests/video-abort-signal.test.ts | 52 +++ .../src/adapters/embedding.ts | 6 +- .../tests/abort-signal.test.ts | 57 +++ packages/ai/src/activities/embed/index.ts | 59 ++- packages/ai/src/activities/evaluate/index.ts | 51 ++- packages/ai/src/activities/files/adapter.ts | 23 +- packages/ai/src/activities/files/index.ts | 39 +- .../ai/src/activities/generateAudio/index.ts | 1 + .../ai/src/activities/generateImage/index.ts | 1 + .../src/activities/generateLiveVideo/index.ts | 1 + .../ai/src/activities/generateSpeech/index.ts | 1 + .../activities/generateTranscription/index.ts | 1 + .../src/activities/generateVideo/adapter.ts | 26 +- .../ai/src/activities/generateVideo/index.ts | 55 ++- .../ai/src/activities/generateVoice/index.ts | 1 + .../ai/src/activities/generateWorld/index.ts | 1 + packages/ai/src/activities/middleware/run.ts | 2 + .../ai/src/activities/middleware/types.ts | 6 + packages/ai/src/activities/rerank/index.ts | 52 ++- packages/ai/src/activities/summarize/index.ts | 1 + packages/ai/src/realtime/index.ts | 10 +- packages/ai/src/realtime/types.ts | 16 +- packages/ai/src/types.ts | 6 + .../ai/tests/abort-signal-options.test.ts | 354 ++++++++++++++++++ packages/ai/tests/evaluate.test.ts | 4 +- packages/ai/tests/rerank.test.ts | 4 +- testing/e2e/global-setup.ts | 45 +++ testing/e2e/src/routeTree.gen.ts | 21 ++ testing/e2e/src/routes/api.embedding-abort.ts | 57 +++ testing/e2e/tests/embedding-abort.spec.ts | 17 + 102 files changed, 2516 insertions(+), 185 deletions(-) create mode 100644 .changeset/ai-abort-signal-options.md create mode 100644 .changeset/ai-anthropic-abort-signal-options.md create mode 100644 .changeset/ai-bedrock-abort-signal-options.md create mode 100644 .changeset/ai-byteplus-abort-signal-options.md create mode 100644 .changeset/ai-client-abort-signal-options.md create mode 100644 .changeset/ai-cloudflare-abort-signal-options.md create mode 100644 .changeset/ai-cohere-abort-signal-options.md create mode 100644 .changeset/ai-elevenlabs-abort-signal-options.md create mode 100644 .changeset/ai-fal-abort-signal-options.md create mode 100644 .changeset/ai-gemini-abort-signal-options.md create mode 100644 .changeset/ai-grok-abort-signal-options.md create mode 100644 .changeset/ai-lovable-abort-signal-options.md create mode 100644 .changeset/ai-mistral-abort-signal-options.md create mode 100644 .changeset/ai-neon-abort-signal-options.md create mode 100644 .changeset/ai-ollama-abort-signal-options.md create mode 100644 .changeset/ai-openai-abort-signal-options.md create mode 100644 .changeset/ai-openrouter-abort-signal-options.md create mode 100644 .changeset/ai-vercel-gateway-abort-signal-options.md create mode 100644 packages/ai-anthropic/tests/abort-signal.test.ts create mode 100644 packages/ai-bedrock/tests/abort-signal.test.ts create mode 100644 packages/ai-byteplus/tests/abort-signal.test.ts create mode 100644 packages/ai-client/tests/chat-client-tool-abort-signal.test.ts create mode 100644 packages/ai-cloudflare/tests/abort-signal.test.ts create mode 100644 packages/ai-cohere/tests/abort-signal.test.ts create mode 100644 packages/ai-elevenlabs/tests/abort-signal.test.ts create mode 100644 packages/ai-fal/tests/video-abort-signal.test.ts create mode 100644 packages/ai-gemini/tests/abort-signal.test.ts create mode 100644 packages/ai-gemini/tests/upload-gemini-file-abort.test.ts create mode 100644 packages/ai-grok/tests/abort-signal.test.ts create mode 100644 packages/ai-lovable/tests/abort-signal.test.ts create mode 100644 packages/ai-mistral/tests/abort-signal.test.ts create mode 100644 packages/ai-neon/tests/abort-signal.test.ts create mode 100644 packages/ai-ollama/tests/abort-signal.test.ts create mode 100644 packages/ai-openai/tests/abort-signal.test.ts create mode 100644 packages/ai-openrouter/tests/video-abort-signal.test.ts create mode 100644 packages/ai-vercel-gateway/tests/abort-signal.test.ts create mode 100644 packages/ai/tests/abort-signal-options.test.ts create mode 100644 testing/e2e/src/routes/api.embedding-abort.ts create mode 100644 testing/e2e/tests/embedding-abort.spec.ts 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-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-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-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..c6383dbdd6 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 and the adapter cancels the provider request: + +```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 d1855bbf8b..ccff66fbfd 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", @@ -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", @@ -941,7 +941,8 @@ { "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" @@ -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..7cde59f265 100644 --- a/docs/media/realtime-chat.md +++ b/docs/media/realtime-chat.md @@ -536,6 +536,23 @@ client.destroy() The React hook disconnects by itself when the component unmounts. With `RealtimeClient`, call `disconnect()` and `destroy()` yourself. +### Cancel the token request + +`realtimeToken()` takes an `abortSignal`. Pass the request signal, so a client that leaves also cancels the provider call: + +```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..626f638121 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: 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/tools/client-tools.md b/docs/tools/client-tools.md index 4c0fdc9d33..24487aad9c 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/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 777991f988..9809f81a99 100644 --- a/packages/ai-byteplus/src/adapters/video.ts +++ b/packages/ai-byteplus/src/adapters/video.ts @@ -213,7 +213,12 @@ export class BytePlusVideoAdapter< init?: Omit, ): Promise<{ response: Response; body: unknown }> { const fetchImpl = this.clientConfig.fetch ?? fetch - const signal = bytePlusTimeoutSignal(this.clientConfig.timeout) + const timeoutSignal = bytePlusTimeoutSignal(this.clientConfig.timeout) + // Keep a caller signal from `init` and the config timeout together. + const signal = + init?.signal && timeoutSignal + ? AbortSignal.any([init.signal, timeoutSignal]) + : (init?.signal ?? timeoutSignal) const response = await fetchImpl(`${this.clientConfig.baseURL}${path}`, { ...init, ...(signal && { signal }), @@ -496,9 +501,13 @@ 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)}`, + signal && { signal }, ) if (!response.ok) { const error = bytePlusArkError(response.status, body, 'video task lookup') @@ -515,10 +524,13 @@ export class BytePlusVideoAdapter< return body as BytePlusVideoTask } - async getVideoStatus(jobId: string): Promise { + 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 @@ -543,10 +555,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..a501970d6b --- /dev/null +++ b/packages/ai-byteplus/tests/abort-signal.test.ts @@ -0,0 +1,50 @@ +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'] 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 }), + ) + }, + ) +}) diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 368dd7373f..6d5dc5c234 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -644,6 +644,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 @@ -3077,6 +3079,8 @@ export class ChatClient< this.continuationGeneration++ const hadLocalStream = this.abortController !== null this.cancelInFlightStream({ setReadyStatus: true }) + this.clientToolAbortController.abort() + this.clientToolAbortController = new AbortController() this.discardPendingSends() this.lastResume = null this.activeInterruptSubmission = undefined @@ -3156,6 +3160,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/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-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..7a5536d086 100644 --- a/packages/ai-fal/src/adapters/video.ts +++ b/packages/ai-fal/src/adapters/video.ts @@ -210,10 +210,14 @@ export class FalVideoAdapter extends BaseVideoAdapter< return getFalVideoDurationOptions(this.model) } - async getVideoStatus(jobId: string): Promise { + 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 +231,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..4dc9fe5c0a --- /dev/null +++ b/packages/ai-fal/tests/video-abort-signal.test.ts @@ -0,0 +1,38 @@ +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 }, + config: () => {}, + }, +})) + +describe('fal video forwards abortSignal to the queue', () => { + it.each(['getVideoStatus', 'getVideo'] 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 4593a8e72f..27cf902fb4 100644 --- a/packages/ai-gemini/src/adapters/video.ts +++ b/packages/ai-gemini/src/adapters/video.ts @@ -507,11 +507,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' } @@ -555,8 +559,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') { @@ -588,11 +593,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( @@ -628,8 +637,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') { @@ -667,17 +677,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 da3a997802..6e34fcb6a3 100644 --- a/packages/ai-grok/src/adapters/video.ts +++ b/packages/ai-grok/src/adapters/video.ts @@ -538,8 +538,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)}`, @@ -550,10 +556,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' } @@ -569,10 +578,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 f68eec3d37..2bdafd729d 100644 --- a/packages/ai-openai/src/adapters/video.ts +++ b/packages/ai-openai/src/adapters/video.ts @@ -171,7 +171,10 @@ export class OpenAIVideoAdapter< */ private getVideosClient(): { create: (req: Record) => Promise<{ id: string }> - retrieve: (id: string) => Promise<{ + retrieve: ( + id: string, + options?: { signal?: AbortSignal }, + ) => Promise<{ id: string status: string progress?: number @@ -179,15 +182,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`. @@ -209,14 +221,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. @@ -231,7 +247,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..76c315a933 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>>() @@ -200,6 +202,7 @@ async function createWebRTCConnection( 'Content-Type': 'application/sdp', }, body: offer.sdp ?? null, + ...(abortSignal && { signal: abortSignal }), }) if (!sdpResponse.ok) { 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-openrouter/src/adapters/video.ts b/packages/ai-openrouter/src/adapters/video.ts index 138dd1cc8d..01f9fec8db 100644 --- a/packages/ai-openrouter/src/adapters/video.ts +++ b/packages/ai-openrouter/src/adapters/video.ts @@ -311,8 +311,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), @@ -322,8 +328,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( @@ -344,7 +355,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( @@ -407,3 +421,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/embed/index.ts b/packages/ai/src/activities/embed/index.ts index f32c07a5af..887b48832a 100644 --- a/packages/ai/src/activities/embed/index.ts +++ b/packages/ai/src/activities/embed/index.ts @@ -7,8 +7,15 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' +import { + abortReasonMessage, + createActivityAbortControls, + isActivityAbortError, + raceWithAbort, +} from '../../utilities/activity-abort' import { createGenerationContext, + runGenerationAbort, runGenerationError, runGenerationFinish, runGenerationStart, @@ -119,6 +126,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 +227,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, @@ -238,13 +262,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 +307,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 +319,17 @@ export async function embed< modelOptions, timestamp: Date.now(), }) - await runGenerationError(middleware, mwCtx, { - error, - duration, - }) + if (isActivityAbortError(error, abortControls.signal)) { + 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..0a3fa3a09e 100644 --- a/packages/ai/src/activities/evaluate/index.ts +++ b/packages/ai/src/activities/evaluate/index.ts @@ -8,7 +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, @@ -165,6 +169,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 @@ -193,7 +203,8 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { // SDK's `RequestAbortedError`, ...). Classifying on `signal.aborted` alone // 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 + // `isActivityAbortError` also matches the `TimeoutError` of `timeout`. + if (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 +475,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 +515,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, @@ -544,14 +554,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 +620,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..57a88ec490 100644 --- a/packages/ai/src/activities/generateVideo/adapter.ts +++ b/packages/ai/src/activities/generateVideo/adapter.ts @@ -112,16 +112,24 @@ 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 /** * Describe the durations this adapter's model accepts. Returns a tagged @@ -234,10 +242,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 44ec2c93b1..f7d3fe7954 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, @@ -683,7 +688,10 @@ async function* runStreamingVideoGeneration< while (Date.now() - startTime < maxDuration) { await sleep(pollingInterval, abortControls.signal) - const statusResult = await adapter.getVideoStatus(jobResult.jobId) + const statusResult = await adapter.getVideoStatus( + jobResult.jobId, + signalOptions, + ) yield { type: 'CUSTOM', @@ -698,7 +706,7 @@ async function* runStreamingVideoGeneration< } if (statusResult.status === 'completed') { - const video = await adapter.getVideo(jobResult.jobId) + const video = await adapter.getVideo(jobResult.jobId, signalOptions) logger.output( `activity=generateVideo jobId=${jobResult.jobId} status=completed`, @@ -838,6 +846,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 } /** @@ -919,7 +938,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() @@ -929,6 +966,7 @@ export async function getVideoJobStatus< const terminalContext = (): GenerationMiddlewareContext => createGenerationContext({ requestId, + signal, activity: 'video', provider: adapter.name, model: adapter.model, @@ -953,7 +991,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') { @@ -962,8 +1003,13 @@ 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. + if (isActivityAbortError(error, signal)) throw error const errorMessage = error instanceof Error ? error.message : 'Failed to get video URL' aiEventClient.emit('video:request:completed', { @@ -1033,6 +1079,7 @@ export async function getVideoJobStatus< try { result = await settleVideoUrl(transformed) } catch (error) { + if (isActivityAbortError(error, signal)) 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..2c8eaaa058 100644 --- a/packages/ai/src/activities/rerank/index.ts +++ b/packages/ai/src/activities/rerank/index.ts @@ -7,7 +7,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 +76,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 @@ -106,7 +116,8 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { // misroute a real failure — e.g. the out-of-range-index throw below — to the // abort hook whenever a shared/long-lived signal happens to already be // aborted, hiding it from `onError` observers. - if (isAbortShapedError(error)) return true + // `isActivityAbortError` also matches the `TimeoutError` of `timeout`. + if (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 +166,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 +175,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', @@ -205,15 +214,19 @@ export async function rerank< const serialized = documents.map(serializeDocument) try { - const result = await adapter.rerank({ - model, - query, - documents: serialized, - topN, - modelOptions, + 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 +276,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 fa83c57f0d..d6882a4cce 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..148b0016d7 --- /dev/null +++ b/packages/ai/tests/abort-signal-options.test.ts @@ -0,0 +1,354 @@ +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('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) + }) +}) 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 af441c6622..5182779aee 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -93,6 +93,7 @@ import { Route as ApiDurableDeliveryRouteImport } from './routes/api.durable-del 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' @@ -653,6 +654,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', @@ -1357,6 +1363,7 @@ export interface FileRoutesByFullPath { '/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 @@ -1561,6 +1568,7 @@ export interface FileRoutesByTo { '/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 @@ -1766,6 +1774,7 @@ export interface FileRoutesById { '/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 @@ -1972,6 +1981,7 @@ export interface FileRouteTypes { | '/api/durable-mcp-input' | '/api/durable-takeover' | '/api/embedding' + | '/api/embedding-abort' | '/api/evaluate-http-status' | '/api/evaluate-refusal' | '/api/evaluate-usage-on-failure' @@ -2176,6 +2186,7 @@ export interface FileRouteTypes { | '/api/durable-mcp-input' | '/api/durable-takeover' | '/api/embedding' + | '/api/embedding-abort' | '/api/evaluate-http-status' | '/api/evaluate-refusal' | '/api/evaluate-usage-on-failure' @@ -2380,6 +2391,7 @@ export interface FileRouteTypes { | '/api/durable-mcp-input' | '/api/durable-takeover' | '/api/embedding' + | '/api/embedding-abort' | '/api/evaluate-http-status' | '/api/evaluate-refusal' | '/api/evaluate-usage-on-failure' @@ -2585,6 +2597,7 @@ export interface RootRouteChildren { ApiDurableMcpInputRoute: typeof ApiDurableMcpInputRoute ApiDurableTakeoverRoute: typeof ApiDurableTakeoverRoute ApiEmbeddingRoute: typeof ApiEmbeddingRoute + ApiEmbeddingAbortRoute: typeof ApiEmbeddingAbortRoute ApiEvaluateHttpStatusRoute: typeof ApiEvaluateHttpStatusRoute ApiEvaluateRefusalRoute: typeof ApiEvaluateRefusalRoute ApiEvaluateUsageOnFailureRoute: typeof ApiEvaluateUsageOnFailureRoute @@ -3291,6 +3304,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' @@ -4267,6 +4287,7 @@ const rootRouteChildren: RootRouteChildren = { ApiDurableMcpInputRoute: ApiDurableMcpInputRoute, ApiDurableTakeoverRoute: ApiDurableTakeoverRoute, ApiEmbeddingRoute: ApiEmbeddingRoute, + ApiEmbeddingAbortRoute: ApiEmbeddingAbortRoute, ApiEvaluateHttpStatusRoute: ApiEvaluateHttpStatusRoute, ApiEvaluateRefusalRoute: ApiEvaluateRefusalRoute, ApiEvaluateUsageOnFailureRoute: ApiEvaluateUsageOnFailureRoute, 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/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, + }) + }) +}) From 506a9ebb699ef9b47d10cfa8c512e2d68f8320ea Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 16:31:51 +0100 Subject: [PATCH 02/12] feat(ai, ai-client, ai-fal, ai-byteplus): cancel hooks for MCP discovery, video jobs, realtime connect, and durable HTTP stop Follow-ups to the optional abortSignal options. Every new field and method is optional, so old code compiles and behaves the same. - MCPToolSource.tools() takes an optional signal. chat() passes the signal of its abortController to MCP discovery. - VideoAdapter has an optional cancelVideoJob(jobId, options?). A streaming generateVideo() calls it (best effort, not awaited) when it aborts after the job exists. fal (fal.queue.cancel) and BytePlus (DELETE task) implement it. - RealtimeClient passes an abortSignal to adapter.connect() and aborts it on disconnect(). A cancelled connect does not set the error state. - fetchServerSentEvents and fetchHttpStream take an optional cancelUrl. ChatClient.stop() then posts { runId, threadId } to it through the new optional connection method cancelRun(). Other aborts do not call it. - Docs for each option, and an E2E spec for cancelUrl (typechecked, CI runs it). --- .changeset/ai-byteplus-cancel-video-job.md | 5 + .changeset/ai-cancel-video-job-mcp-signal.md | 8 ++ .../ai-client-cancel-url-connect-signal.md | 9 ++ .changeset/ai-fal-cancel-video-job.md | 5 + docs/config.json | 8 +- docs/media/realtime-chat.md | 2 + docs/media/video-generation-advanced.md | 26 ++++ docs/media/video-generation.md | 2 +- docs/resumable-streams/advanced.md | 86 +++++++++++- docs/sandbox/takeover.md | 19 +-- docs/tools/mcp-managed.md | 2 +- packages/ai-byteplus/src/adapters/video.ts | 21 +++ .../ai-byteplus/tests/abort-signal.test.ts | 33 ++++- packages/ai-client/src/chat-client.ts | 6 + packages/ai-client/src/connection-adapters.ts | 60 ++++++++ packages/ai-client/src/realtime-client.ts | 7 + .../tests/chat-client-cancel-run.test.ts | 126 +++++++++++++++++ .../realtime-client-connect-signal.test.ts | 67 +++++++++ packages/ai-fal/src/adapters/video.ts | 14 ++ .../ai-fal/tests/video-abort-signal.test.ts | 27 ++-- packages/ai/src/activities/chat/index.ts | 6 +- .../ai/src/activities/chat/mcp/manager.ts | 7 +- packages/ai/src/activities/chat/mcp/types.ts | 9 +- .../src/activities/generateVideo/adapter.ts | 11 ++ .../ai/src/activities/generateVideo/index.ts | 17 +++ .../ai/tests/abort-signal-options.test.ts | 68 +++++++++ packages/ai/tests/chat-mcp-manager.test.ts | 18 +++ packages/ai/tests/chat-mcp.test.ts | 31 +++++ testing/e2e/src/routeTree.gen.ts | 21 +++ .../src/routes/api.durable-delivery-cancel.ts | 130 ++++++++++++++++++ .../tests/delivery-durability-cancel.spec.ts | 62 +++++++++ 31 files changed, 874 insertions(+), 39 deletions(-) create mode 100644 .changeset/ai-byteplus-cancel-video-job.md create mode 100644 .changeset/ai-cancel-video-job-mcp-signal.md create mode 100644 .changeset/ai-client-cancel-url-connect-signal.md create mode 100644 .changeset/ai-fal-cancel-video-job.md create mode 100644 packages/ai-client/tests/chat-client-cancel-run.test.ts create mode 100644 packages/ai-client/tests/realtime-client-connect-signal.test.ts create mode 100644 testing/e2e/src/routes/api.durable-delivery-cancel.ts create mode 100644 testing/e2e/tests/delivery-durability-cancel.spec.ts 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-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-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/docs/config.json b/docs/config.json index ccff66fbfd..bfb3b39fc0 100644 --- a/docs/config.json +++ b/docs/config.json @@ -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", @@ -936,7 +936,7 @@ "label": "Video Generation", "to": "media/video-generation", "addedAt": "2026-04-15", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Video Generation: Advanced", @@ -1494,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", diff --git a/docs/media/realtime-chat.md b/docs/media/realtime-chat.md index 7cde59f265..25ae7e7c04 100644 --- a/docs/media/realtime-chat.md +++ b/docs/media/realtime-chat.md @@ -536,6 +536,8 @@ 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 a client that leaves also cancels the provider call: diff --git a/docs/media/video-generation-advanced.md b/docs/media/video-generation-advanced.md index 626f638121..d0eb35c6ef 100644 --- a/docs/media/video-generation-advanced.md +++ b/docs/media/video-generation-advanced.md @@ -170,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/resumable-streams/advanced.md b/docs/resumable-streams/advanced.md index cb8cfc0717..c9dc7ed0f7 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,90 @@ 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() + +// POST /api/chat +export async function POST(request: Request) { + const { messages, threadId, runId } = await chatParamsFromRequest(request) + const abortController = new AbortController() + 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 }) + } + liveRuns.get(runId)?.abort(RUN_CANCEL_REASON) + 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. +- `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/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-byteplus/src/adapters/video.ts b/packages/ai-byteplus/src/adapters/video.ts index 9809f81a99..22a05de0d9 100644 --- a/packages/ai-byteplus/src/adapters/video.ts +++ b/packages/ai-byteplus/src/adapters/video.ts @@ -524,6 +524,27 @@ export class BytePlusVideoAdapter< return body as BytePlusVideoTask } + /** + * `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 && { signal: options.abortSignal }), + }, + ) + if (!response.ok) { + throw bytePlusArkError(response.status, body, 'video task cancel') + } + } + async getVideoStatus( jobId: string, options?: { abortSignal?: AbortSignal }, diff --git a/packages/ai-byteplus/tests/abort-signal.test.ts b/packages/ai-byteplus/tests/abort-signal.test.ts index a501970d6b..613585ab68 100644 --- a/packages/ai-byteplus/tests/abort-signal.test.ts +++ b/packages/ai-byteplus/tests/abort-signal.test.ts @@ -35,7 +35,7 @@ async function expectAbortReachesFetch( } describe('byteplus adapters forward abortSignal to the request', () => { - it.each(['getVideoStatus', 'getVideo'] as const)( + it.each(['getVideoStatus', 'getVideo', 'cancelVideoJob'] as const)( 'video %s', async (method) => { const { fetch, signals } = hangingFetch() @@ -48,3 +48,34 @@ describe('byteplus adapters forward abortSignal to the request', () => { }, ) }) + +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 6d5dc5c234..8cf95107d1 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -3078,6 +3078,12 @@ 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() diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 00d7b8551d..cb8c028eb6 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -982,6 +982,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 @@ -1130,6 +1138,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 @@ -1189,11 +1199,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 } : {}), } } @@ -1334,6 +1346,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 @@ -1373,6 +1388,49 @@ 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) +} + +/** + * `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', + }, + ) + if (!response.ok) { + throw new Error(`Run cancel failed with status ${response.status}`) + } + }, + } } /** @@ -1469,6 +1527,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 @@ -1629,6 +1688,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 0bac78f317..0801749898 100644 --- a/packages/ai-client/src/realtime-client.ts +++ b/packages/ai-client/src/realtime-client.ts @@ -49,6 +49,8 @@ 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 readonly clientTools: Map @@ -104,6 +106,7 @@ export class RealtimeClient { status: 'connecting', error: null, }) + const connectAbort = (this.connectAbort = new AbortController()) try { // Fetch token from server @@ -121,6 +124,7 @@ export class RealtimeClient { this.connection = await this.options.adapter.connect( this.token, toolsList, + { abortSignal: connectAbort.signal }, ) // A new session starts here. Clear the last transcript before events @@ -150,6 +154,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 @@ -160,6 +166,7 @@ export class RealtimeClient { * Disconnect from the realtime session. */ async disconnect(): Promise { + this.connectAbort?.abort() if (this.tokenRefreshTimeout) { clearTimeout(this.tokenRefreshTimeout) this.tokenRefreshTimeout = null 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..a352106ec8 --- /dev/null +++ b/packages/ai-client/tests/chat-client-cancel-run.test.ts @@ -0,0 +1,126 @@ +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) + }) +}) 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..5064167a82 --- /dev/null +++ b/packages/ai-client/tests/realtime-client-connect-signal.test.ts @@ -0,0 +1,67 @@ +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, + }) + + const pending = client.connect() + await vi.waitFor(() => expect(signals).toHaveLength(1)) + expect(signals[0]!.aborted).toBe(false) + + await client.disconnect() + expect(signals[0]!.aborted).toBe(true) + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(client.getSnapshot().status).toBe('idle') + expect(onError).not.toHaveBeenCalled() + 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-fal/src/adapters/video.ts b/packages/ai-fal/src/adapters/video.ts index 7a5536d086..f123cd794a 100644 --- a/packages/ai-fal/src/adapters/video.ts +++ b/packages/ai-fal/src/adapters/video.ts @@ -210,6 +210,20 @@ export class FalVideoAdapter extends BaseVideoAdapter< return getFalVideoDurationOptions(this.model) } + /** + * 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 }, diff --git a/packages/ai-fal/tests/video-abort-signal.test.ts b/packages/ai-fal/tests/video-abort-signal.test.ts index 4dc9fe5c0a..01d675607c 100644 --- a/packages/ai-fal/tests/video-abort-signal.test.ts +++ b/packages/ai-fal/tests/video-abort-signal.test.ts @@ -19,20 +19,25 @@ const mocks = vi.hoisted(() => { vi.mock('@fal-ai/client', () => ({ fal: { - queue: { status: mocks.hang, result: mocks.hang }, + queue: { status: mocks.hang, result: mocks.hang, cancel: mocks.hang }, config: () => {}, }, })) describe('fal video forwards abortSignal to the queue', () => { - it.each(['getVideoStatus', 'getVideo'] 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' }) - }) + 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/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index c41beeab47..1239f1d0b7 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -6408,7 +6408,7 @@ async function* runChatEngine( const logger = resolveDebugOption(debug) const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover() + const mcpTools = await mcpManager.discover(options.abortController?.signal) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } @@ -7024,7 +7024,7 @@ async function runAgenticStructuredOutput( adapter.combinedStructuredOutputSource?.(options.modelOptions) ?? 'text' const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover() + const mcpTools = await mcpManager.discover(options.abortController?.signal) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } @@ -7386,7 +7386,7 @@ async function* runStreamingStructuredOutputImpl( : undefined const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover() + const mcpTools = await mcpManager.discover(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/generateVideo/adapter.ts b/packages/ai/src/activities/generateVideo/adapter.ts index 57a88ec490..d4a58d82a4 100644 --- a/packages/ai/src/activities/generateVideo/adapter.ts +++ b/packages/ai/src/activities/generateVideo/adapter.ts @@ -131,6 +131,17 @@ export interface VideoAdapter< 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 * union so consumers can render UI / coerce input without provider-specific diff --git a/packages/ai/src/activities/generateVideo/index.ts b/packages/ai/src/activities/generateVideo/index.ts index f7d3fe7954..94fc0f1392 100644 --- a/packages/ai/src/activities/generateVideo/index.ts +++ b/packages/ai/src/activities/generateVideo/index.ts @@ -661,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( @@ -676,6 +678,8 @@ async function* runStreamingVideoGeneration< abortControls.signal, ) + jobId = jobResult.jobId + yield { type: 'CUSTOM', name: 'video:job:created', @@ -777,6 +781,19 @@ async function* runStreamingVideoGeneration< settled = true const elapsed = Date.now() - obsStartTime if (isActivityAbortError(error, abortControls.signal)) { + if (jobId !== undefined) { + // 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, diff --git a/packages/ai/tests/abort-signal-options.test.ts b/packages/ai/tests/abort-signal-options.test.ts index 148b0016d7..a354312bd8 100644 --- a/packages/ai/tests/abort-signal-options.test.ts +++ b/packages/ai/tests/abort-signal-options.test.ts @@ -213,6 +213,74 @@ describe('video job abortSignal', () => { }) }) +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 } 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..75f9f5ff5a 100644 --- a/packages/ai/tests/chat-mcp.test.ts +++ b/packages/ai/tests/chat-mcp.test.ts @@ -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 // -------------------------------------------------------------------------- diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 5182779aee..7009fe0ae2 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -90,6 +90,7 @@ 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' @@ -639,6 +640,11 @@ 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', @@ -1360,6 +1366,7 @@ 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 @@ -1565,6 +1572,7 @@ 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 @@ -1771,6 +1779,7 @@ 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 @@ -1978,6 +1987,7 @@ 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' @@ -2183,6 +2193,7 @@ 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' @@ -2388,6 +2399,7 @@ 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' @@ -2594,6 +2606,7 @@ export interface RootRouteChildren { ApiCrossModelReplayWireRoute: typeof ApiCrossModelReplayWireRoute ApiDevtoolsMemoryRoute: typeof ApiDevtoolsMemoryRoute ApiDurableDeliveryRoute: typeof ApiDurableDeliveryRoute + ApiDurableDeliveryCancelRoute: typeof ApiDurableDeliveryCancelRoute ApiDurableMcpInputRoute: typeof ApiDurableMcpInputRoute ApiDurableTakeoverRoute: typeof ApiDurableTakeoverRoute ApiEmbeddingRoute: typeof ApiEmbeddingRoute @@ -3283,6 +3296,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' @@ -4284,6 +4304,7 @@ const rootRouteChildren: RootRouteChildren = { ApiCrossModelReplayWireRoute: ApiCrossModelReplayWireRoute, ApiDevtoolsMemoryRoute: ApiDevtoolsMemoryRoute, ApiDurableDeliveryRoute: ApiDurableDeliveryRoute, + ApiDurableDeliveryCancelRoute: ApiDurableDeliveryCancelRoute, ApiDurableMcpInputRoute: ApiDurableMcpInputRoute, ApiDurableTakeoverRoute: ApiDurableTakeoverRoute, ApiEmbeddingRoute: ApiEmbeddingRoute, 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/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') + }) +}) From 81b1349676bb0055453a9bb95ad894e6ae1d45c0 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:59:17 +0000 Subject: [PATCH 03/12] ci: apply automated fixes --- testing/e2e/src/routeTree.gen.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 7009fe0ae2..201d996cca 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -640,11 +640,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 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', From 4b9c2c50f66f5dec679a31afffbb2dfd4ab56dd8 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 17:28:48 +0100 Subject: [PATCH 04/12] docs(tools): type-safe signal in the client tool stop sample The client tool context can be undefined, so read abortSignal with ?. Fixes the kiira check. --- docs/tools/client-tools.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tools/client-tools.md b/docs/tools/client-tools.md index 24487aad9c..9f509d735f 100644 --- a/docs/tools/client-tools.md +++ b/docs/tools/client-tools.md @@ -350,7 +350,7 @@ const searchDocs = toolDefinition({ inputSchema: z.object({ query: z.string() }), }).client(async ({ query }, ctx) => { const response = await fetch(`/api/search?q=${encodeURIComponent(query)}`, { - signal: ctx.abortSignal, + signal: ctx?.abortSignal, }); return { results: await response.text() }; }); From adfbea070f9c517ed88db665032d2b53e8a22e38 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:07:44 +0100 Subject: [PATCH 05/12] fix(ai): end a run stopped during MCP discovery through onAbort A source that rejects on abort made discover() throw before TextEngine existed, so middleware onAbort did not run. Discovery now returns no tools when the run signal aborted, and the engine ends the run through its abort path. Other discovery errors still throw. --- packages/ai/src/activities/chat/index.ts | 32 +++++++++++++-- packages/ai/tests/chat-mcp.test.ts | 51 ++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 1239f1d0b7..b3f2cb7331 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -6391,6 +6391,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, @@ -6408,7 +6425,10 @@ async function* runChatEngine( const logger = resolveDebugOption(debug) const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover(options.abortController?.signal) + const mcpTools = await discoverMcpTools( + mcpManager, + options.abortController?.signal, + ) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } @@ -7024,7 +7044,10 @@ async function runAgenticStructuredOutput( adapter.combinedStructuredOutputSource?.(options.modelOptions) ?? 'text' const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover(options.abortController?.signal) + const mcpTools = await discoverMcpTools( + mcpManager, + options.abortController?.signal, + ) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } @@ -7386,7 +7409,10 @@ async function* runStreamingStructuredOutputImpl( : undefined const mcpManager = MCPManager.from(mcp) - const mcpTools = await mcpManager.discover(options.abortController?.signal) + const mcpTools = await discoverMcpTools( + mcpManager, + options.abortController?.signal, + ) if (mcpTools.length > 0) { textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools] } diff --git a/packages/ai/tests/chat-mcp.test.ts b/packages/ai/tests/chat-mcp.test.ts index 75f9f5ff5a..5834959a81 100644 --- a/packages/ai/tests/chat-mcp.test.ts +++ b/packages/ai/tests/chat-mcp.test.ts @@ -610,3 +610,54 @@ 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('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() + }) +}) From 67a0aa817891ff7189012afdffc5d94fdc49042b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:07:45 +0100 Subject: [PATCH 06/12] fix(ai-openai): release the microphone and peer connection when realtime setup fails An aborted or failed SDP request left the mic tracks, the data channel, and the peer connection open. The setup now stops them before it rethrows, like the Grok adapter. --- packages/ai-openai/src/realtime/adapter.ts | 103 ++++++++++-------- .../tests/realtime-connect-teardown.test.ts | 75 +++++++++++++ 2 files changed, 132 insertions(+), 46 deletions(-) create mode 100644 packages/ai-openai/tests/realtime-connect-teardown.test.ts diff --git a/packages/ai-openai/src/realtime/adapter.ts b/packages/ai-openai/src/realtime/adapter.ts index 76c315a933..6d46a9e674 100644 --- a/packages/ai-openai/src/realtime/adapter.ts +++ b/packages/ai-openai/src/realtime/adapter.ts @@ -165,61 +165,72 @@ 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, + // 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, - ...(abortSignal && { signal: abortSignal }), - }) - 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) 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..68b8aedf8e --- /dev/null +++ b/packages/ai-openai/tests/realtime-connect-teardown.test.ts @@ -0,0 +1,75 @@ +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('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() + }) +}) From afc90f1a0f023698c393b7fd8df0264f0f539b82 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:07:46 +0100 Subject: [PATCH 07/12] fix(ai-client): bound the stop() cancel request and stop an earlier connect setup The cancel POST from stop() gets a 10 s deadline, so a cancelUrl that never answers cannot hold a connection. A second RealtimeClient.connect() aborts the setup of the earlier attempt. --- packages/ai-client/src/connection-adapters.ts | 6 +++++ packages/ai-client/src/realtime-client.ts | 2 ++ .../tests/chat-client-cancel-run.test.ts | 14 +++++++++++ .../realtime-client-connect-signal.test.ts | 24 +++++++++++++++++++ 4 files changed, 46 insertions(+) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index cb8c028eb6..6fd47f33a5 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1396,6 +1396,9 @@ export interface FetchConnectionOptions { 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 @@ -1424,6 +1427,9 @@ function fetchCancelRun( }, 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) { diff --git a/packages/ai-client/src/realtime-client.ts b/packages/ai-client/src/realtime-client.ts index 0801749898..dc2c8f2d70 100644 --- a/packages/ai-client/src/realtime-client.ts +++ b/packages/ai-client/src/realtime-client.ts @@ -106,6 +106,8 @@ export class RealtimeClient { status: 'connecting', error: null, }) + // A new connect() owns the setup. Stop the setup of an earlier one. + this.connectAbort?.abort() const connectAbort = (this.connectAbort = new AbortController()) try { diff --git a/packages/ai-client/tests/chat-client-cancel-run.test.ts b/packages/ai-client/tests/chat-client-cancel-run.test.ts index a352106ec8..1ecb304dcc 100644 --- a/packages/ai-client/tests/chat-client-cancel-run.test.ts +++ b/packages/ai-client/tests/chat-client-cancel-run.test.ts @@ -124,3 +124,17 @@ describe('ChatClient.stop() with cancelUrl', () => { 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/realtime-client-connect-signal.test.ts b/packages/ai-client/tests/realtime-client-connect-signal.test.ts index 5064167a82..c934910c80 100644 --- a/packages/ai-client/tests/realtime-client-connect-signal.test.ts +++ b/packages/ai-client/tests/realtime-client-connect-signal.test.ts @@ -50,6 +50,30 @@ describe('RealtimeClient connect abortSignal', () => { 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, + }) + + const first = client.connect() + await vi.waitFor(() => expect(signals).toHaveLength(1)) + const second = client.connect() + await vi.waitFor(() => expect(signals).toHaveLength(2)) + + expect(signals[0]!.aborted).toBe(true) + expect(signals[1]!.aborted).toBe(false) + await expect(first).rejects.toMatchObject({ name: 'AbortError' }) + expect(onError).not.toHaveBeenCalled() + + await client.disconnect() + await expect(second).rejects.toMatchObject({ name: 'AbortError' }) + client.destroy() + }) + it('a connect error that is not an abort still sets the error state', async () => { const client = new RealtimeClient({ getToken: async () => token, From 569ea3f050c9d266867852110e1058aa4e1d7fe2 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:07:49 +0100 Subject: [PATCH 08/12] docs: keep an early Stop in the cancel sample, and narrow two cancel claims --- docs/advanced/files-api.md | 2 +- docs/media/realtime-chat.md | 2 +- docs/resumable-streams/advanced.md | 14 +++++++++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/docs/advanced/files-api.md b/docs/advanced/files-api.md index c6383dbdd6..d87263f3db 100644 --- a/docs/advanced/files-api.md +++ b/docs/advanced/files-api.md @@ -201,7 +201,7 @@ fal storage is upload-only. `falFiles()` has no `get` or `delete`, so `getFile() ### Cancel a file request -`uploadFile()`, `getFile()`, and `deleteFile()` take an `abortSignal`. When it aborts, the call rejects and the adapter cancels the provider 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"; diff --git a/docs/media/realtime-chat.md b/docs/media/realtime-chat.md index 25ae7e7c04..1803c43874 100644 --- a/docs/media/realtime-chat.md +++ b/docs/media/realtime-chat.md @@ -540,7 +540,7 @@ A `disconnect()` while `connect()` still runs cancels the connection setup. `con ### Cancel the token request -`realtimeToken()` takes an `abortSignal`. Pass the request signal, so a client that leaves also cancels the provider call: +`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' diff --git a/docs/resumable-streams/advanced.md b/docs/resumable-streams/advanced.md index c9dc7ed0f7..ec5f8870c2 100644 --- a/docs/resumable-streams/advanced.md +++ b/docs/resumable-streams/advanced.md @@ -181,11 +181,15 @@ 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) @@ -216,7 +220,14 @@ export async function cancelChatRun(request: Request) { if (typeof runId !== 'string') { return new Response('runId must be a string', { status: 400 }) } - liveRuns.get(runId)?.abort(RUN_CANCEL_REASON) + 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 }) } ``` @@ -245,6 +256,7 @@ 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 From 607c239652973cee633eefd13e2b647677bc36cf Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:26:36 +0100 Subject: [PATCH 09/12] test(ai-client): handle the connect() abort rejections at once The first connect() rejects when the second one starts, while the test still waits. Its handler was attached later, so vitest counted an unhandled rejection and the coverage job failed. --- .../realtime-client-connect-signal.test.ts | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/packages/ai-client/tests/realtime-client-connect-signal.test.ts b/packages/ai-client/tests/realtime-client-connect-signal.test.ts index c934910c80..289ef2fc94 100644 --- a/packages/ai-client/tests/realtime-client-connect-signal.test.ts +++ b/packages/ai-client/tests/realtime-client-connect-signal.test.ts @@ -38,13 +38,16 @@ describe('RealtimeClient connect abortSignal', () => { onError, }) - const pending = client.connect() + // 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 expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + await pending expect(client.getSnapshot().status).toBe('idle') expect(onError).not.toHaveBeenCalled() client.destroy() @@ -59,18 +62,23 @@ describe('RealtimeClient connect abortSignal', () => { onError, }) - const first = client.connect() + // 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 = client.connect() + 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 expect(first).rejects.toMatchObject({ name: 'AbortError' }) + await first expect(onError).not.toHaveBeenCalled() await client.disconnect() - await expect(second).rejects.toMatchObject({ name: 'AbortError' }) + await second client.destroy() }) From c550491f43e6b4b3e8b2738b2f3e7b9462758e91 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:34:16 +0100 Subject: [PATCH 10/12] fix(ai): run no pending tools in a run stopped before it starts When MCP discovery fails after the abort, the engine started and then ran pending tool calls before its cancel check, so a pending client tool raised an interrupt after the cancel. The engine now returns before the pending-tool step when the run is cancelled. Its finally block still calls onAbort. --- packages/ai/src/activities/chat/index.ts | 4 +++ packages/ai/tests/chat-mcp.test.ts | 46 +++++++++++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index b3f2cb7331..073a16b74e 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1588,6 +1588,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') { diff --git a/packages/ai/tests/chat-mcp.test.ts b/packages/ai/tests/chat-mcp.test.ts index 5834959a81..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' @@ -643,6 +643,50 @@ describe('chat({ mcp }) abort during discovery', () => { 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() From 2cb834c2b356f99635bcb57f42e2a42335f9609d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:34:17 +0100 Subject: [PATCH 11/12] fix(ai-openai): let an abort stop the wait for the realtime data channel A disconnect() after the SDP step left connect() pending, with the mic and the peer connection open. The wait for the data channel now rejects on abort and closes the connection first. --- packages/ai-openai/src/realtime/adapter.ts | 26 ++++++++++++++- .../tests/realtime-connect-teardown.test.ts | 32 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/packages/ai-openai/src/realtime/adapter.ts b/packages/ai-openai/src/realtime/adapter.ts index 6d46a9e674..644dbca228 100644 --- a/packages/ai-openai/src/realtime/adapter.ts +++ b/packages/ai-openai/src/realtime/adapter.ts @@ -722,7 +722,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/tests/realtime-connect-teardown.test.ts b/packages/ai-openai/tests/realtime-connect-teardown.test.ts index 68b8aedf8e..4db06891f8 100644 --- a/packages/ai-openai/tests/realtime-connect-teardown.test.ts +++ b/packages/ai-openai/tests/realtime-connect-teardown.test.ts @@ -72,4 +72,36 @@ describe('openaiRealtime connect teardown', () => { 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() + }) }) From 3e229e0a2b8ffa8cfed31ec37f2810e34d0450da Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 3 Nov 2026 12:16:11 +0100 Subject: [PATCH 12/12] fix(ai): treat only a real caller abort as an abort A provider can throw an error named TimeoutError when its own fetch timeout fires. The new abort code treated that error as a caller abort. - generateVideo() calls cancelVideoJob() only when the caller signal or the timeout fired. - getVideoJobStatus() rethrows a fetch error only when the signal fired. Without a caller abort, the job is failed and onError runs, as before. - embed() calls onAbort only when the signal fired. - rerank() and decide() treat a TimeoutError as an abort only when the signal fired. Other abort-shaped errors work as before. - embed(), rerank() and decide() clear the timeout timer when onStart throws. rerank() serializes the documents inside the try. - The OpenAI realtime adapter does not ask for the microphone when the signal is already aborted. --- packages/ai-openai/src/realtime/adapter.ts | 3 + .../tests/realtime-connect-teardown.test.ts | 16 +++ packages/ai/src/activities/embed/index.ts | 10 +- packages/ai/src/activities/evaluate/index.ts | 12 +- .../ai/src/activities/generateVideo/index.ts | 9 +- packages/ai/src/activities/rerank/index.ts | 20 ++- .../ai/tests/abort-signal-options.test.ts | 135 ++++++++++++++++++ 7 files changed, 189 insertions(+), 16 deletions(-) diff --git a/packages/ai-openai/src/realtime/adapter.ts b/packages/ai-openai/src/realtime/adapter.ts index 644dbca228..896655f2ad 100644 --- a/packages/ai-openai/src/realtime/adapter.ts +++ b/packages/ai-openai/src/realtime/adapter.ts @@ -168,6 +168,9 @@ async function createWebRTCConnection( // A failed or aborted setup releases the microphone and the peer // connection. Without this, the mic indicator stays on after disconnect(). try { + // 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 { diff --git a/packages/ai-openai/tests/realtime-connect-teardown.test.ts b/packages/ai-openai/tests/realtime-connect-teardown.test.ts index 4db06891f8..ba0689dabc 100644 --- a/packages/ai-openai/tests/realtime-connect-teardown.test.ts +++ b/packages/ai-openai/tests/realtime-connect-teardown.test.ts @@ -62,6 +62,22 @@ describe('openaiRealtime connect teardown', () => { 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 })) diff --git a/packages/ai/src/activities/embed/index.ts b/packages/ai/src/activities/embed/index.ts index 887b48832a..e3041117fc 100644 --- a/packages/ai/src/activities/embed/index.ts +++ b/packages/ai/src/activities/embed/index.ts @@ -10,7 +10,6 @@ import { resolveDebugOption } from '../../logger/resolve' import { abortReasonMessage, createActivityAbortControls, - isActivityAbortError, raceWithAbort, } from '../../utilities/activity-abort' import { @@ -242,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, @@ -319,7 +321,9 @@ export async function embed< modelOptions, timestamp: Date.now(), }) - if (isActivityAbortError(error, abortControls.signal)) { + // 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, diff --git a/packages/ai/src/activities/evaluate/index.ts b/packages/ai/src/activities/evaluate/index.ts index 0a3fa3a09e..3274069237 100644 --- a/packages/ai/src/activities/evaluate/index.ts +++ b/packages/ai/src/activities/evaluate/index.ts @@ -8,6 +8,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' +import { isAbortShapedError } from '../error-payload' import { createActivityAbortControls, isActivityAbortError, @@ -203,8 +204,10 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { // SDK's `RequestAbortedError`, ...). Classifying on `signal.aborted` alone // would misroute a real failure to the abort hook whenever a shared signal // happens to already be aborted, hiding it from `onError` observers. - // `isActivityAbortError` also matches the `TimeoutError` of `timeout`. - if (isActivityAbortError(error)) return true + 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 @@ -537,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, diff --git a/packages/ai/src/activities/generateVideo/index.ts b/packages/ai/src/activities/generateVideo/index.ts index 94fc0f1392..7b198343f1 100644 --- a/packages/ai/src/activities/generateVideo/index.ts +++ b/packages/ai/src/activities/generateVideo/index.ts @@ -781,7 +781,9 @@ async function* runStreamingVideoGeneration< settled = true const elapsed = Date.now() - obsStartTime if (isActivityAbortError(error, abortControls.signal)) { - if (jobId !== undefined) { + // 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 @@ -1026,7 +1028,8 @@ async function runGetVideoJobStatus< ) } catch (error) { // A cancelled call says nothing about the job: rethrow, do not fail it. - if (isActivityAbortError(error, signal)) throw error + // 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', { @@ -1096,7 +1099,7 @@ async function runGetVideoJobStatus< try { result = await settleVideoUrl(transformed) } catch (error) { - if (isActivityAbortError(error, signal)) throw 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/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts index 2c8eaaa058..f4d47e4e90 100644 --- a/packages/ai/src/activities/rerank/index.ts +++ b/packages/ai/src/activities/rerank/index.ts @@ -7,6 +7,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' +import { isAbortShapedError } from '../error-payload' import { createActivityAbortControls, isActivityAbortError, @@ -116,8 +117,10 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { // misroute a real failure — e.g. the out-of-range-index throw below — to the // abort hook whenever a shared/long-lived signal happens to already be // aborted, hiding it from `onError` observers. - // `isActivityAbortError` also matches the `TimeoutError` of `timeout`. - if (isActivityAbortError(error)) return true + 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 @@ -193,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, @@ -209,11 +215,11 @@ 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 { + // 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, diff --git a/packages/ai/tests/abort-signal-options.test.ts b/packages/ai/tests/abort-signal-options.test.ts index a354312bd8..df10ad6999 100644 --- a/packages/ai/tests/abort-signal-options.test.ts +++ b/packages/ai/tests/abort-signal-options.test.ts @@ -420,3 +420,138 @@ describe('decide() and rerank() timeout', () => { 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() + } + }, + ) +})