fix(ai-bedrock, ai-mistral): forward the caller's abort signal to the provider request

.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.

.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.

packages/ai-bedrock/src/adapters/converse-text.ts

@@ -241,12 +241,13 @@ export class BedrockConverseTextAdapter<
 
   protected async sendStream(
     input: ConverseStreamCommandInput,
+    signal?: AbortSignal | null,
   ): Promise<AsyncIterable<ConverseStreamOutput>> {
     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<ConverseCommandOutput> {
     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

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<AbortSignal | undefined>(
+    (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<typeof model> {
+    protected override async getClient(): Promise<BedrockRuntimeClient> {
+      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()
+  })
+})

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<string, unknown>,
     config: MistralClientConfig,
     wrapFetch: FetchWrapper | undefined,
+    signal: AbortSignal | undefined,
   ): AsyncGenerator<MistralRawChunk> {
     const serverURL = (
       config.baseURL ??
@@ -966,6 +980,7 @@ export class MistralTextAdapter<
       method: 'POST',
       headers,
       body: JSON.stringify(body),
+      signal,
     })
 
     if (!response.ok) {

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<AbortSignal | undefined>
+
+beforeEach(() => {
+  let onSend: (signal: AbortSignal | undefined) => void
+  sent = new Promise((resolve) => (onSend = resolve))
+  vi.stubGlobal(
+    'fetch',
+    (input: RequestInfo | URL, init?: RequestInit) =>
+      new Promise<Response>((_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')
+  })
+})