diff --git a/.changeset/bedrock-converse-forward-abort-signal.md b/.changeset/bedrock-converse-forward-abort-signal.md new file mode 100644 index 0000000000..408f6cbc5d --- /dev/null +++ b/.changeset/bedrock-converse-forward-abort-signal.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-bedrock': patch +--- + +Aborting a chat now cancels the in-flight Converse request to Bedrock, for both streaming and structured output. diff --git a/.changeset/mistral-forward-abort-signal.md b/.changeset/mistral-forward-abort-signal.md new file mode 100644 index 0000000000..25455c365e --- /dev/null +++ b/.changeset/mistral-forward-abort-signal.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-mistral': patch +--- + +Aborting a chat now cancels the in-flight request to Mistral, for both streaming and structured output. diff --git a/packages/ai-bedrock/src/adapters/converse-text.ts b/packages/ai-bedrock/src/adapters/converse-text.ts index c4400e9741..331f2ef2a4 100644 --- a/packages/ai-bedrock/src/adapters/converse-text.ts +++ b/packages/ai-bedrock/src/adapters/converse-text.ts @@ -241,12 +241,13 @@ export class BedrockConverseTextAdapter< protected async sendStream( input: ConverseStreamCommandInput, + signal?: AbortSignal | null, ): Promise> { const { ConverseStreamCommand } = await this.importBedrockRuntime() const client = await this.getClient() const command = new ConverseStreamCommand(input) addRestoreDocuments(command, input) - const res = await client.send(command) + const res = await client.send(command, { abortSignal: signal ?? undefined }) if (!res.stream) { throw new Error('Bedrock Converse: empty stream response') } @@ -255,12 +256,13 @@ export class BedrockConverseTextAdapter< protected async send( input: ConverseCommandInput, + signal?: AbortSignal | null, ): Promise { const { ConverseCommand } = await this.importBedrockRuntime() const client = await this.getClient() const command = new ConverseCommand(input) addRestoreDocuments(command, input) - return client.send(command) + return client.send(command, { abortSignal: signal ?? undefined }) } // --------------------------------------------------------------------------- @@ -281,7 +283,7 @@ export class BedrockConverseTextAdapter< if (!input.toolConfig && input.messages) { input.messages = toolBlocksToText(input.messages) } - const stream = await this.sendStream(input) + const stream = await this.sendStream(input, options.request?.signal) for await (const chunk of processConverseStream( stream, () => this.generateId(), @@ -344,7 +346,7 @@ export class BedrockConverseTextAdapter< ...this.buildInput(requestOptions), toolConfig: buildStructuredToolConfig(outputSchema), } - const res = await this.send(input) + const res = await this.send(input, chatOptions.request?.signal) // A forced tool call cut off at the output cap carries partial input, // which would read like a schema failure (issue #1426). if (res.stopReason === 'max_tokens') { @@ -408,7 +410,7 @@ export class BedrockConverseTextAdapter< ...this.buildInput(requestOptions), toolConfig: buildStructuredToolConfig(outputSchema), } - const stream = await this.sendStream(input) + const stream = await this.sendStream(input, chatOptions.request?.signal) // The forced tool streams its `input` as partial-JSON fragments inside // `contentBlockDelta.delta.toolUse.input`. We surface them as diff --git a/packages/ai-bedrock/tests/converse/abort-signal-wire.test.ts b/packages/ai-bedrock/tests/converse/abort-signal-wire.test.ts new file mode 100644 index 0000000000..3e55acc2a5 --- /dev/null +++ b/packages/ai-bedrock/tests/converse/abort-signal-wire.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from 'vitest' +import { BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime' +import { chat } from '@tanstack/ai' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { BedrockConverseTextAdapter } from '../../src/adapters/converse-text' + +const logger = resolveDebugOption(false) +const model = 'us.amazon.nova-pro-v1:0' +const messages = [{ role: 'user' as const, content: 'Hi' }] +const outputSchema = { type: 'object', properties: {} } + +/** + * An adapter whose SDK client sends through a handler that never answers. + * `sent` is the abort signal that the handler got for its one request, and + * only that signal can end the request. + */ +function wireAdapter() { + let onSend: (signal: AbortSignal | undefined) => void + const sent = new Promise( + (resolve) => (onSend = resolve), + ) + const client = new BedrockRuntimeClient({ + region: 'us-east-1', + maxAttempts: 1, + credentials: { accessKeyId: 'test', secretAccessKey: 'test' }, + requestHandler: { + handle(_request: unknown, options?: { abortSignal?: AbortSignal }) { + return new Promise((_resolve, reject) => { + const signal = options?.abortSignal + onSend(signal) + signal?.addEventListener('abort', () => reject(signal.reason)) + }) + }, + }, + }) + class WireAdapter extends BedrockConverseTextAdapter { + protected override async getClient(): Promise { + return client + } + } + return { adapter: new WireAdapter({ apiKey: 'test-key' }, model), sent } +} + +describe('Bedrock Converse abort signal', () => { + it('aborts the stream request when the caller aborts', async () => { + const { adapter, sent } = wireAdapter() + const abortController = new AbortController() + const run = (async () => { + for await (const _ of chat({ + adapter, + threadId: 'thread-1', + messages, + abortController, + })) { + // drain + } + })() + + const signal = await sent + abortController.abort() + + expect(signal?.aborted).toBe(true) + await run + }) + + it('aborts the structured output request when the caller aborts', async () => { + const { adapter, sent } = wireAdapter() + const abortController = new AbortController() + const run = chat({ + adapter, + threadId: 'thread-1', + messages, + outputSchema, + abortController, + }) + + const signal = await sent + abortController.abort() + + expect(signal?.aborted).toBe(true) + await expect(run).rejects.toThrow() + }) + + // `chat()` runs `structuredOutputStream` on Converse, so call it directly. + it('aborts a structuredOutput call when the caller aborts', async () => { + const { adapter, sent } = wireAdapter() + const abortController = new AbortController() + const run = adapter.structuredOutput({ + chatOptions: { + threadId: 'thread-1', + logger, + model, + messages, + request: { signal: abortController.signal }, + }, + outputSchema, + }) + + const signal = await sent + abortController.abort() + + expect(signal?.aborted).toBe(true) + await expect(run).rejects.toThrow() + }) +}) diff --git a/packages/ai-mistral/src/adapters/text.ts b/packages/ai-mistral/src/adapters/text.ts index 652bc4ca3f..3cf5480b0e 100644 --- a/packages/ai-mistral/src/adapters/text.ts +++ b/packages/ai-mistral/src/adapters/text.ts @@ -356,6 +356,7 @@ export class MistralTextAdapter< body, this.rawConfig, options.wrapFetch, + options.request?.signal ?? undefined, ) for await (const chunk of this.processMistralStreamChunks( stream, @@ -456,17 +457,29 @@ export class MistralTextAdapter< : chatOptions.wrapFetch ? createMistralClient(this.rawConfig, next) : this.client - const response = await client.chat.complete({ + const params = { ...nonStreamParams, responseFormat: { - type: 'json_schema', + type: 'json_schema' as const, jsonSchema: { name: 'structured_output', schemaDefinition: jsonSchema, strict, }, }, - }) + } + // @mistralai/mistralai 2.2.0 defaults timeoutMs to 30_000 (funcs/chatComplete.js) + // and skips it when the request has a signal (lib/sdks.js). See #1517. + const signal = chatOptions.request?.signal + const timeoutMs = this.rawConfig.timeoutMs || 30_000 + const response = signal + ? await client.chat.complete(params, { + signal: + timeoutMs > 0 + ? AbortSignal.any([signal, AbortSignal.timeout(timeoutMs)]) + : signal, + }) + : await client.chat.complete(params) // A response cut off at the output cap is a truncated JSON document; // report it before the parse error (issue #1426). @@ -939,6 +952,7 @@ export class MistralTextAdapter< body: Record, config: MistralClientConfig, wrapFetch: FetchWrapper | undefined, + signal: AbortSignal | undefined, ): AsyncGenerator { const serverURL = ( config.baseURL ?? @@ -966,6 +980,7 @@ export class MistralTextAdapter< method: 'POST', headers, body: JSON.stringify(body), + signal, }) if (!response.ok) { diff --git a/packages/ai-mistral/tests/abort-signal-wire.test.ts b/packages/ai-mistral/tests/abort-signal-wire.test.ts new file mode 100644 index 0000000000..e619e5afd8 --- /dev/null +++ b/packages/ai-mistral/tests/abort-signal-wire.test.ts @@ -0,0 +1,89 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { chat } from '@tanstack/ai' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { createMistralText } from '../src/adapters/text' + +const logger = resolveDebugOption(false) +const model = 'mistral-large-latest' +const messages = [{ role: 'user' as const, content: 'Hi' }] +const outputSchema = { type: 'object', properties: {} } + +// The signal of the request that reached `fetch`. Mistral never answers it, +// so only this signal can end it. +let sent: Promise + +beforeEach(() => { + let onSend: (signal: AbortSignal | undefined) => void + sent = new Promise((resolve) => (onSend = resolve)) + vi.stubGlobal( + 'fetch', + (input: RequestInfo | URL, init?: RequestInit) => + new Promise((_resolve, reject) => { + // The stream path sends a URL and an init. The SDK sends a Request. + const signal = + input instanceof Request ? input.signal : (init?.signal ?? undefined) + onSend(signal) + signal?.addEventListener('abort', () => reject(signal.reason)) + }), + ) +}) + +afterEach(() => { + vi.unstubAllGlobals() +}) + +describe('Mistral abort signal', () => { + it('aborts the stream request when the caller aborts', async () => { + const abortController = new AbortController() + const run = (async () => { + for await (const _ of chat({ + adapter: createMistralText(model, 'test-key'), + threadId: 'thread-1', + messages, + abortController, + })) { + // drain + } + })() + + const signal = await sent + abortController.abort() + + expect(signal?.aborted).toBe(true) + await run + }) + + it('aborts the structured output request when the caller aborts', async () => { + const abortController = new AbortController() + const run = chat({ + adapter: createMistralText(model, 'test-key'), + threadId: 'thread-1', + messages, + outputSchema, + abortController, + }) + + const signal = await sent + abortController.abort() + + expect(signal?.aborted).toBe(true) + await expect(run).rejects.toThrow() + }) + + it('keeps timeoutMs on a structured output call with a signal', async () => { + const run = createMistralText(model, 'test-key', { + timeoutMs: 20, + }).structuredOutput({ + chatOptions: { + threadId: 'thread-1', + logger, + model, + messages, + request: { signal: new AbortController().signal }, + }, + outputSchema, + }) + + await expect(run).rejects.toThrow('Request timed out') + }) +})