From 32dd9f39c966eac0d3dbc6d8ad71354a3ad19075 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:53:06 +0100 Subject: [PATCH 01/16] feat(ai-mcp, ai-harness, ai-sandbox, ai-memory, ai-code-mode, ai-computer-use): optional abort signals for agents, sandboxes, MCP Add optional cancel signals where these contracts had none. Every new field is optional, so old calls compile and behave the same. - Code Mode: IsolateContext.execute(code, { signal }). execute_typescript passes the tool signal. The node, cloudflare, and e2b drivers stop the run. - Computer use: ComputerDriverFactory.create({ threadId, signal }), passed to the sandbox provider by ai-computer-sandbox. - MCP client: signal on connect, tools, resources, readResource, resourceTemplates, prompts, getPrompt, and the lazy list in callTool. createMcpAppCallHandler(req, { signal }) and createMcpAppBridge().callTool(input, { signal }). - Harness client: an optional last { signal } on each request method. - Sandbox: signal on git clone/push/pull, snapshot, fork, env.set, and ports.connect. The exec-backed git forwards it. - Memory: recall(scope, query, { signal }) and Embedder.embed(text, { signal }); memoryMiddleware passes the chat run signal. --- .changeset/ai-client-abort-signal-agents.md | 5 + .../ai-code-mode-abort-signal-agents.md | 5 + ...ai-computer-sandbox-abort-signal-agents.md | 5 + .../ai-computer-use-abort-signal-agents.md | 5 + .changeset/ai-harness-abort-signal-agents.md | 5 + ...-isolate-cloudflare-abort-signal-agents.md | 5 + .../ai-isolate-e2b-abort-signal-agents.md | 5 + .../ai-isolate-node-abort-signal-agents.md | 5 + .changeset/ai-mcp-abort-signal-agents.md | 10 + .changeset/ai-memory-abort-signal-agents.md | 5 + .changeset/ai-sandbox-abort-signal-agents.md | 5 + .../ai-sandbox-blaxel-abort-signal-agents.md | 5 + docs/code-mode/code-mode-isolates.md | 10 +- docs/computer-use/custom-driver.md | 2 +- docs/config.json | 14 +- docs/harness/coding-tools-backends.md | 2 + docs/harness/connect.md | 13 + docs/mcp/apps.md | 3 +- docs/memory/adapters.md | 4 +- docs/tools/mcp-manual.md | 21 + packages/ai-client/src/mcp-app-bridge.ts | 28 +- .../tests/mcp-app-bridge-signal.test.ts | 37 ++ .../ai-code-mode/src/create-code-mode-tool.ts | 5 +- packages/ai-code-mode/src/types.ts | 8 +- .../ai-code-mode/tests/execute-signal.test.ts | 46 ++ packages/ai-computer-sandbox/src/driver.ts | 4 +- .../ai-computer-sandbox/tests/driver.test.ts | 14 + .../src/create-computer-use.ts | 8 +- packages/ai-computer-use/src/types.ts | 10 +- .../tests/create-computer-use.test.ts | 27 + packages/ai-harness/src/client.ts | 483 ++++++++++++------ .../ai-harness/tests/client-signal.test.ts | 70 +++ .../src/isolate-driver.ts | 6 +- .../tests/isolate-driver.test.ts | 27 + packages/ai-isolate-e2b/src/isolate-driver.ts | 12 +- .../tests/isolate-driver.test.ts | 31 ++ .../ai-isolate-node/src/isolate-context.ts | 31 +- .../tests/isolate-driver.test.ts | 28 + packages/ai-mcp/src/apps/call-handler.ts | 16 +- packages/ai-mcp/src/client.ts | 88 +++- packages/ai-mcp/src/types.ts | 7 + .../ai-mcp/tests/apps/call-handler.test.ts | 22 + packages/ai-mcp/tests/client-signal.test.ts | 38 ++ packages/ai-memory/src/internal/store.ts | 12 +- packages/ai-memory/src/middleware.ts | 6 +- .../src/providers/in-memory/index.ts | 3 +- .../ai-memory/src/providers/mem0/index.ts | 4 +- .../ai-memory/src/providers/redis/index.ts | 3 +- packages/ai-memory/src/types.ts | 9 +- .../ai-memory/tests/recall-signal.test.ts | 74 +++ packages/ai-sandbox-blaxel/src/handle.ts | 6 +- packages/ai-sandbox/src/contracts.ts | 41 +- packages/ai-sandbox/src/git-exec.ts | 22 +- packages/ai-sandbox/tests/git-exec.test.ts | 27 + 54 files changed, 1129 insertions(+), 258 deletions(-) create mode 100644 .changeset/ai-client-abort-signal-agents.md create mode 100644 .changeset/ai-code-mode-abort-signal-agents.md create mode 100644 .changeset/ai-computer-sandbox-abort-signal-agents.md create mode 100644 .changeset/ai-computer-use-abort-signal-agents.md create mode 100644 .changeset/ai-harness-abort-signal-agents.md create mode 100644 .changeset/ai-isolate-cloudflare-abort-signal-agents.md create mode 100644 .changeset/ai-isolate-e2b-abort-signal-agents.md create mode 100644 .changeset/ai-isolate-node-abort-signal-agents.md create mode 100644 .changeset/ai-mcp-abort-signal-agents.md create mode 100644 .changeset/ai-memory-abort-signal-agents.md create mode 100644 .changeset/ai-sandbox-abort-signal-agents.md create mode 100644 .changeset/ai-sandbox-blaxel-abort-signal-agents.md create mode 100644 packages/ai-client/tests/mcp-app-bridge-signal.test.ts create mode 100644 packages/ai-code-mode/tests/execute-signal.test.ts create mode 100644 packages/ai-harness/tests/client-signal.test.ts create mode 100644 packages/ai-mcp/tests/client-signal.test.ts create mode 100644 packages/ai-memory/tests/recall-signal.test.ts diff --git a/.changeset/ai-client-abort-signal-agents.md b/.changeset/ai-client-abort-signal-agents.md new file mode 100644 index 0000000000..f3cac932b8 --- /dev/null +++ b/.changeset/ai-client-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': minor +--- + +`createMcpAppBridge().callTool(input, { signal })` passes the signal to the call request. diff --git a/.changeset/ai-code-mode-abort-signal-agents.md b/.changeset/ai-code-mode-abort-signal-agents.md new file mode 100644 index 0000000000..fad03dddda --- /dev/null +++ b/.changeset/ai-code-mode-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-code-mode': minor +--- + +`IsolateContext.execute(code, options?)` takes an optional `{ signal }`. `execute_typescript` passes the abort signal of the tool call, so a driver can stop the run. diff --git a/.changeset/ai-computer-sandbox-abort-signal-agents.md b/.changeset/ai-computer-sandbox-abort-signal-agents.md new file mode 100644 index 0000000000..388770ac59 --- /dev/null +++ b/.changeset/ai-computer-sandbox-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-computer-sandbox': minor +--- + +Pass the `create()` signal to `provider.create()`. diff --git a/.changeset/ai-computer-use-abort-signal-agents.md b/.changeset/ai-computer-use-abort-signal-agents.md new file mode 100644 index 0000000000..8097a358a7 --- /dev/null +++ b/.changeset/ai-computer-use-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-computer-use': minor +--- + +`ComputerDriverFactory.create({ threadId, signal })` gets the abort signal of the tool call that needs the machine. diff --git a/.changeset/ai-harness-abort-signal-agents.md b/.changeset/ai-harness-abort-signal-agents.md new file mode 100644 index 0000000000..9e5e244c35 --- /dev/null +++ b/.changeset/ai-harness-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-harness': minor +--- + +Each request method of `createHarnessClient` takes an optional last argument `{ signal }`: the control inputs, `snapshot`, `transcript`, `agentTranscript`, `describe`, `putMedia`, `mediaUrl`, `loadMedia`, and `sessions.*`. diff --git a/.changeset/ai-isolate-cloudflare-abort-signal-agents.md b/.changeset/ai-isolate-cloudflare-abort-signal-agents.md new file mode 100644 index 0000000000..4f3c13857e --- /dev/null +++ b/.changeset/ai-isolate-cloudflare-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-isolate-cloudflare': minor +--- + +Forward the `execute()` signal to the worker request. diff --git a/.changeset/ai-isolate-e2b-abort-signal-agents.md b/.changeset/ai-isolate-e2b-abort-signal-agents.md new file mode 100644 index 0000000000..81de234349 --- /dev/null +++ b/.changeset/ai-isolate-e2b-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-isolate-e2b': minor +--- + +Kill the sandbox process group when the `execute()` signal aborts. The run returns an `AbortError` failure. diff --git a/.changeset/ai-isolate-node-abort-signal-agents.md b/.changeset/ai-isolate-node-abort-signal-agents.md new file mode 100644 index 0000000000..9c8c84b763 --- /dev/null +++ b/.changeset/ai-isolate-node-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-isolate-node': minor +--- + +Stop a running script when the `execute()` signal aborts. The run returns an `AbortError` failure. diff --git a/.changeset/ai-mcp-abort-signal-agents.md b/.changeset/ai-mcp-abort-signal-agents.md new file mode 100644 index 0000000000..f60b8cb9d3 --- /dev/null +++ b/.changeset/ai-mcp-abort-signal-agents.md @@ -0,0 +1,10 @@ +--- +'@tanstack/ai-mcp': minor +--- + +New optional signals. Old calls work the same. + +- `createMCPClient({ signal })` cancels the connect handshake. +- `tools({ signal })`, `resources({ signal })`, `readResource(uri, { signal })`, `resourceTemplates({ signal })`, `prompts({ signal })`, and `getPrompt(name, args, { signal })` send the signal to the MCP SDK. +- The lazy tool list in `callTool` uses the signal of the call. +- The handler from `createMcpAppCallHandler` takes `(req, { signal })` and passes it to the connect, the tool list, and the tool call. diff --git a/.changeset/ai-memory-abort-signal-agents.md b/.changeset/ai-memory-abort-signal-agents.md new file mode 100644 index 0000000000..2975acbc5c --- /dev/null +++ b/.changeset/ai-memory-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-memory': minor +--- + +`MemoryAdapter.recall(scope, query, { signal })` and `Embedder.embed(text, { signal })` take an optional signal. `memoryMiddleware` passes the chat run signal. `inMemory()`, `redis()`, and `mem0()` forward it. diff --git a/.changeset/ai-sandbox-abort-signal-agents.md b/.changeset/ai-sandbox-abort-signal-agents.md new file mode 100644 index 0000000000..68f0af1068 --- /dev/null +++ b/.changeset/ai-sandbox-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox': minor +--- + +New optional signals on the sandbox contract: `git.clone({ signal })`, `git.push(dir, { signal })`, `git.pull(dir, { signal })`, `snapshot(label, { signal })`, `fork({ signal })`, `env.set(vars, { signal })`, and `ports.connect(port, { signal })`. The exec-backed git passes the signal to each git command. diff --git a/.changeset/ai-sandbox-blaxel-abort-signal-agents.md b/.changeset/ai-sandbox-blaxel-abort-signal-agents.md new file mode 100644 index 0000000000..4dbf765c4e --- /dev/null +++ b/.changeset/ai-sandbox-blaxel-abort-signal-agents.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-blaxel': minor +--- + +Pass the `git.push` and `git.pull` signal to the git command. diff --git a/docs/code-mode/code-mode-isolates.md b/docs/code-mode/code-mode-isolates.md index 22bdeeac00..ecf9504dcf 100644 --- a/docs/code-mode/code-mode-isolates.md +++ b/docs/code-mode/code-mode-isolates.md @@ -388,7 +388,10 @@ interface IsolateConfig { } interface IsolateContext { - execute(code: string): Promise> + execute( + code: string, + options?: { signal?: AbortSignal }, + ): Promise> dispose(): Promise } @@ -399,3 +402,8 @@ interface ExecutionResult { error?: NormalizedError } ``` + +Code Mode passes the abort signal of the tool call as `options.signal`. When it aborts, stop the run and return a failure. The built-in drivers do this: + +- The Node, Cloudflare, and E2B drivers stop the run at once. +- The QuickJS and Daytona drivers run until their `timeout`. diff --git a/docs/computer-use/custom-driver.md b/docs/computer-use/custom-driver.md index 4f6f3bd65d..caaa38a5de 100644 --- a/docs/computer-use/custom-driver.md +++ b/docs/computer-use/custom-driver.md @@ -191,7 +191,7 @@ The factory has these members: | `display` | The screen size of every machine. Each driver from `create` has this `display`. | | `environment` | The kind of every machine. | | `capabilities` | `{ zoom, cursorPosition, navigate }`: the optional methods that every driver from `create` has. The tool is built before any machine exists, so the factory tells it. | -| `create({ threadId })` | Makes the driver of one thread. If it throws, the model gets the error. The next model turn calls `create` again. | +| `create({ threadId, signal })` | Makes the driver of one thread. `signal` is the abort signal of the tool call that needs the machine. Pass it to your setup requests. If `create` throws, the model gets the error. The next model turn calls `create` again. | | `close()` | Optional. Frees what the threads share. `computerUse.close()` calls it after it closes every thread. | When a thread closes or goes idle, `createComputerUse` calls `close()` on the driver of that thread. diff --git a/docs/config.json b/docs/config.json index d1855bbf8b..f7985474fc 100644 --- a/docs/config.json +++ b/docs/config.json @@ -632,7 +632,7 @@ "label": "Manual MCP", "to": "tools/mcp-manual", "addedAt": "2026-06-05", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "MCP Type Generation", @@ -660,7 +660,7 @@ "label": "MCP Apps", "to": "mcp/apps", "addedAt": "2026-06-24", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Serve Tools over HTTP", @@ -1025,7 +1025,7 @@ "label": "Code Mode Isolate Drivers", "to": "code-mode/code-mode-isolates", "addedAt": "2026-04-15", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Lazy Tools", @@ -1061,7 +1061,7 @@ "label": "Custom Driver", "to": "computer-use/custom-driver", "addedAt": "2026-10-20", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" } ], "tab": "guides" @@ -1092,7 +1092,7 @@ "label": "Connect clients", "to": "harness/connect", "addedAt": "2026-10-13", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Write a client in any language", @@ -1229,7 +1229,7 @@ "label": "Sandbox and extend the coding tools", "to": "harness/coding-tools-backends", "addedAt": "2026-10-13", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Run the coding tools for many threads", @@ -1556,7 +1556,7 @@ "label": "Adapters", "to": "memory/adapters", "addedAt": "2026-08-04", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "Custom Adapter", diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md index 6d579fc546..613c3c2ac2 100644 --- a/docs/harness/coding-tools-backends.md +++ b/docs/harness/coding-tools-backends.md @@ -62,6 +62,8 @@ export const sandboxed = defineHarness({ To give the agent the project folder on your machine, [copy it into the container](#copy-your-project-into-the-container). +To stop a slow clone, pass `signal`: `sandbox.git.clone({ url, signal })`. `git.push(dir, { signal })` and `git.pull(dir, { signal })` take it too. The signal stops the git command when it aborts. + 3. When you are done, stop the sandbox with `await sandbox.destroy()`. Every provider works. Get the handle from `create()`, or from `resume()` for a sandbox that runs already. See [Providers](../sandbox/providers). To create a sandbox for each thread, and resume it on another server, see [Run the coding tools for many threads](./coding-workspaces). diff --git a/docs/harness/connect.md b/docs/harness/connect.md index 23362c964b..3ab9da2a4b 100644 --- a/docs/harness/connect.md +++ b/docs/harness/connect.md @@ -188,6 +188,19 @@ Without `from`, `events()` starts at the newest event, so it sends no old events - To read every event from the start of the session, pass `from: '0'`. Do this when you rebuild a transcript from the events. - If the snapshot answer has no `cursor` (the URL is not a harness handler), `events()` throws and does not try again. +Each request method of the client takes an optional last argument `{ signal }`. When the signal aborts, the request stops and the call rejects: + +```ts group=harness-connect-client +const controller = new AbortController() +setTimeout(() => controller.abort(), 5000) + +const snapshot = await client.snapshot({ signal: controller.signal }) +await client.send('Summarize my open tickets.', undefined, { + signal: controller.signal, +}) +console.log(snapshot.cursor) +``` + To upload files and show the media that your agents make, read [Send and show media](./media). ## Send a message to a running agent diff --git a/docs/mcp/apps.md b/docs/mcp/apps.md index 40f6326e20..94135d750d 100644 --- a/docs/mcp/apps.md +++ b/docs/mcp/apps.md @@ -211,7 +211,8 @@ export const Route = createFileRoute('/api/mcp-apps/call')({ POST: async ({ request }) => { // body: { threadId, serverId, toolName, args?, messageId? } const body = await request.json() - const result = await handler(body) + // The signal stops the connect and the tool call if the widget leaves. + const result = await handler(body, { signal: request.signal }) return new Response(JSON.stringify(result), { headers: { 'Content-Type': 'application/json' }, }) diff --git a/docs/memory/adapters.md b/docs/memory/adapters.md index 5cdcd3fb39..09661b6c9b 100644 --- a/docs/memory/adapters.md +++ b/docs/memory/adapters.md @@ -15,7 +15,7 @@ keywords: - options --- -Each memory adapter implements `recall` and `save`. Swap one for another in `memoryMiddleware`. +Each memory adapter implements `recall` and `save`. Swap one for another in `memoryMiddleware`. `memoryMiddleware` calls `recall(scope, query, { signal })` with the abort signal of the chat run. The built-in adapters pass it to their requests when the backend client can take it. | Adapter | Import | Use it for | Extra install | | --- | --- | --- | --- | @@ -64,7 +64,7 @@ const memory = inMemory({ ``` - **`extract`** returns `ExtractedFact[]` (`{ text, kind?, importance?, metadata? }`), or `undefined` to save nothing extra. Put an LLM-based fact extractor here. The adapter does not depend on a model. -- **`embedder`** embeds the query on `recall` and the stored text on `save`. With no embedder, the score uses word overlap, recency, and importance only. +- **`embedder`** embeds the query on `recall` and the stored text on `save`. With no embedder, the score uses word overlap, recency, and importance only. On `recall`, `embed(text, { signal })` gets the abort signal of the chat run. Pass it to your embedding request, so a stopped run does not wait for it. ## `inMemory()` diff --git a/docs/tools/mcp-manual.md b/docs/tools/mcp-manual.md index 2eb72cabcc..b0dbc52097 100644 --- a/docs/tools/mcp-manual.md +++ b/docs/tools/mcp-manual.md @@ -328,6 +328,27 @@ const stream = chat({ controller.abort() ``` +Outside a chat run, pass `signal` to a client call yourself. Each of these calls takes it: + +```ts +import { createMCPClient } from '@tanstack/ai-mcp' + +const controller = new AbortController() +const { signal } = controller + +const mcp = await createMCPClient({ + transport: { type: 'http', url: process.env.MCP_URL! }, + signal, // cancels the connect only +}) + +const tools = await mcp.tools({ signal }) +const resources = await mcp.resources({ signal }) +const prompts = await mcp.prompts({ signal }) +console.log(tools.length, resources.length, prompts.length) +``` + +`readResource(uri, { signal })`, `resourceTemplates({ signal })`, and `getPrompt(name, args, { signal })` take it too. The client sends the signal to the MCP SDK, which sends `notifications/cancelled` to the server. + ## Full server and client example This TanStack Start route connects to two MCP servers and streams the response to the browser. diff --git a/packages/ai-client/src/mcp-app-bridge.ts b/packages/ai-client/src/mcp-app-bridge.ts index 7de435c653..531748b9bb 100644 --- a/packages/ai-client/src/mcp-app-bridge.ts +++ b/packages/ai-client/src/mcp-app-bridge.ts @@ -12,17 +12,20 @@ export interface CreateMcpAppBridgeOptions { } export interface McpAppBridge { - callTool: (input: { - serverId?: string - toolName: string - args?: Record - /** - * Reserved — forwarded to the call handler for correlation purposes but - * not consumed by the handler. Accepted on the wire; the handler does not - * read it (mirrors the `meta` convention on `UIResourcePart`). - */ - messageId?: string - }) => Promise + callTool: ( + input: { + serverId?: string + toolName: string + args?: Record + /** + * Reserved — forwarded to the call handler for correlation purposes but + * not consumed by the handler. Accepted on the wire; the handler does not + * read it (mirrors the `meta` convention on `UIResourcePart`). + */ + messageId?: string + }, + options?: { signal?: AbortSignal }, + ) => Promise sendPrompt: (text: string) => Promise openLink: (url: string) => { isError: boolean } } @@ -61,8 +64,9 @@ export function createMcpAppBridge( const doFetch = fetchImpl ?? fetch return { - async callTool(input) { + async callTool(input, callOptions) { const response = await doFetch(callEndpoint, { + ...(callOptions?.signal && { signal: callOptions.signal }), method: 'POST', headers: { 'content-type': 'application/json' }, body: JSON.stringify({ diff --git a/packages/ai-client/tests/mcp-app-bridge-signal.test.ts b/packages/ai-client/tests/mcp-app-bridge-signal.test.ts new file mode 100644 index 0000000000..1b4356f091 --- /dev/null +++ b/packages/ai-client/tests/mcp-app-bridge-signal.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it, vi } from 'vitest' +import { createMcpAppBridge } from '../src/mcp-app-bridge' + +/** Waits until the request signal aborts; fails at once without a signal. */ +function hangingFetch() { + return vi.fn( + (_input: RequestInfo | URL, init?: RequestInit): Promise => { + const signal = init?.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('createMcpAppBridge callTool signal', () => { + it('passes the signal to the call request and rejects on abort', async () => { + const fetchImpl = hangingFetch() + const bridge = createMcpAppBridge({ + threadId: 't1', + callEndpoint: '/api/mcp-call', + chat: { sendMessage: async () => {} }, + fetchImpl, + }) + const controller = new AbortController() + const pending = bridge.callTool( + { toolName: 'place_order' }, + { signal: controller.signal }, + ) + expect(fetchImpl.mock.calls[0]![1]?.signal).toBe(controller.signal) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + }) +}) diff --git a/packages/ai-code-mode/src/create-code-mode-tool.ts b/packages/ai-code-mode/src/create-code-mode-tool.ts index f709a69433..25a919dd43 100644 --- a/packages/ai-code-mode/src/create-code-mode-tool.ts +++ b/packages/ai-code-mode/src/create-code-mode-tool.ts @@ -295,7 +295,10 @@ export function createCodeModeTool( } // Step 5: Execute the code in the sandbox - const executionResult = await isolateContext.execute(strippedCode) + const signal = toolContext?.abortSignal + const executionResult = signal + ? await isolateContext.execute(strippedCode, { signal }) + : await isolateContext.execute(strippedCode) // Emit console logs as custom events if (executionResult.logs && executionResult.logs.length > 0) { diff --git a/packages/ai-code-mode/src/types.ts b/packages/ai-code-mode/src/types.ts index f69f072303..63b50d79fd 100644 --- a/packages/ai-code-mode/src/types.ts +++ b/packages/ai-code-mode/src/types.ts @@ -47,9 +47,13 @@ export interface IsolateConfig { */ export interface IsolateContext { /** - * Execute generated code and return results + * Execute generated code and return results. Drivers stop the run when + * `options.signal` aborts, if their runtime can. */ - execute: (code: string) => Promise> + execute: ( + code: string, + options?: { signal?: AbortSignal }, + ) => Promise> /** * Clean up sandbox resources diff --git a/packages/ai-code-mode/tests/execute-signal.test.ts b/packages/ai-code-mode/tests/execute-signal.test.ts new file mode 100644 index 0000000000..d39dc84335 --- /dev/null +++ b/packages/ai-code-mode/tests/execute-signal.test.ts @@ -0,0 +1,46 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { toolDefinition } from '@tanstack/ai' +import { createCodeModeTool } from '../src/create-code-mode-tool' +import type { IsolateContext, IsolateDriver } from '../src/types' + +function setup() { + const context: IsolateContext = { + execute: vi.fn().mockResolvedValue({ success: true, value: 1, logs: [] }), + dispose: vi.fn().mockResolvedValue(undefined), + } + const driver: IsolateDriver = { + createContext: vi.fn().mockResolvedValue(context), + } + const tool = createCodeModeTool({ + driver, + tools: [ + toolDefinition({ + name: 'lookup', + description: 'Look up a value', + inputSchema: z.object({ q: z.string() }), + }).server(async () => 'x'), + ], + }) + return { tool, context } +} + +describe('execute_typescript passes the tool signal to the isolate', () => { + it('forwards context.abortSignal to execute()', async () => { + const { tool, context } = setup() + const controller = new AbortController() + await tool.execute!( + { typescriptCode: 'return 1' }, + { abortSignal: controller.signal, emitCustomEvent() {} }, + ) + expect(context.execute).toHaveBeenCalledWith(expect.any(String), { + signal: controller.signal, + }) + }) + + it('calls execute() with the code only when there is no signal', async () => { + const { tool, context } = setup() + await tool.execute!({ typescriptCode: 'return 1' }) + expect(vi.mocked(context.execute).mock.calls[0]).toHaveLength(1) + }) +}) diff --git a/packages/ai-computer-sandbox/src/driver.ts b/packages/ai-computer-sandbox/src/driver.ts index 1c4dcafbba..7c134d66db 100644 --- a/packages/ai-computer-sandbox/src/driver.ts +++ b/packages/ai-computer-sandbox/src/driver.ts @@ -253,8 +253,8 @@ export function sandboxComputer({ display, environment: 'linux', capabilities: { zoom: false, cursorPosition: true, navigate: false }, - async create() { - const sandbox = await provider.create({}) + async create({ signal }) { + const sandbox = await provider.create(signal ? { signal } : {}) try { const driver = await createSandboxComputerDriver({ sandbox, display }) const { width, height } = driver.display diff --git a/packages/ai-computer-sandbox/tests/driver.test.ts b/packages/ai-computer-sandbox/tests/driver.test.ts index f656121d49..1e667f7715 100644 --- a/packages/ai-computer-sandbox/tests/driver.test.ts +++ b/packages/ai-computer-sandbox/tests/driver.test.ts @@ -284,3 +284,17 @@ describe('sandboxComputer', () => { error.mockRestore() }) }) + +describe('sandboxComputer create signal', () => { + it('passes the signal to provider.create, and {} without one', async () => { + const provider = fakeProvider([fakeSandbox(), fakeSandbox()]) + const factory = sandboxComputer({ provider }) + const controller = new AbortController() + await factory.create({ threadId: 'a', signal: controller.signal }) + await factory.create({ threadId: 'b' }) + expect(vi.mocked(provider.create).mock.calls).toEqual([ + [{ signal: controller.signal }], + [{}], + ]) + }) +}) diff --git a/packages/ai-computer-use/src/create-computer-use.ts b/packages/ai-computer-use/src/create-computer-use.ts index ba3130c6c6..744c859800 100644 --- a/packages/ai-computer-use/src/create-computer-use.ts +++ b/packages/ai-computer-use/src/create-computer-use.ts @@ -243,10 +243,10 @@ export function createComputerUse({ await machine?.close?.() } - function machineOf(thread: ThreadState) { + function machineOf(thread: ThreadState, signal?: AbortSignal) { if (!('create' in driver)) return driver thread.driver ??= driver - .create({ threadId: thread.id }) + .create({ threadId: thread.id, ...(signal && { signal }) }) .catch((error: unknown) => { // The next call of the thread calls create again. thread.driver = undefined @@ -328,7 +328,7 @@ export function createComputerUse({ 'Refused: this computer action has safety checks that nobody approved. Add the `interrupt` from createComputerUse to chat({ interrupts }) and the `middleware` to chat({ middleware }).', ) } - const machine = await machineOf(thread) + const machine = await machineOf(thread, context?.abortSignal) const result = await runComputerAction(machine, input) if (result === undefined) { await sleep(screenshotDelay) @@ -346,7 +346,7 @@ export function createComputerUse({ 'Refused: navigate needs a full URL that starts with https:// or http://, like https://tanstack.com.', ) } - const machine = await machineOf(thread) + const machine = await machineOf(thread, context?.abortSignal) if (machine.navigate === undefined) { throw new Error('The computer driver has no navigate method.') } diff --git a/packages/ai-computer-use/src/types.ts b/packages/ai-computer-use/src/types.ts index 0aecbc2e3c..34af2a0dfa 100644 --- a/packages/ai-computer-use/src/types.ts +++ b/packages/ai-computer-use/src/types.ts @@ -79,8 +79,14 @@ export interface ComputerDriverFactory { environment: ComputerEnvironment /** The optional methods that every created driver has. */ capabilities: { zoom: boolean; cursorPosition: boolean; navigate: boolean } - /** Make the machine for one thread. */ - create: (context: { threadId: string }) => Promise + /** + * Make the machine for one thread. `signal` is the abort signal of the + * tool call that needs the machine. Stop the setup when it aborts. + */ + create: (context: { + threadId: string + signal?: AbortSignal + }) => Promise /** * Free what the threads share, like the browser. * `createComputerUse(...).close()` with no thread id calls this. diff --git a/packages/ai-computer-use/tests/create-computer-use.test.ts b/packages/ai-computer-use/tests/create-computer-use.test.ts index 909cfcb21b..35d0ce5f7c 100644 --- a/packages/ai-computer-use/tests/create-computer-use.test.ts +++ b/packages/ai-computer-use/tests/create-computer-use.test.ts @@ -931,3 +931,30 @@ describe('createComputerUse navigate', () => { ) }) }) + +describe('driver factory signal', () => { + it('passes the abort signal of the tool call to create()', async () => { + const signals: Array = [] + const factory: ComputerDriverFactory = { + display: { width: 1280, height: 800 }, + environment: 'linux', + capabilities: { zoom: false, cursorPosition: false, navigate: false }, + async create({ signal }) { + signals.push(signal) + return createDriver().driver + }, + } + const [computerTool] = createComputerUse({ + driver: factory, + screenshotDelay: 0, + }).tools + const controller = new AbortController() + await computerTool.execute!(CLICK, { + toolCallId: 'call_1', + threadId: THREAD, + abortSignal: controller.signal, + emitCustomEvent() {}, + }) + expect(signals).toEqual([controller.signal]) + }) +}) diff --git a/packages/ai-harness/src/client.ts b/packages/ai-harness/src/client.ts index ca74ca18ca..8535979a55 100644 --- a/packages/ai-harness/src/client.ts +++ b/packages/ai-harness/src/client.ts @@ -167,15 +167,22 @@ export interface HarnessClient { * `as: 'operator'` is the reply of an operator while a handoff to a * person is open: no model runs. The sender needs the `operator` action. */ - send: (message: UserInput, options?: ClientSendOptions) => Promise + send: ( + message: UserInput, + options?: ClientSendOptions, + request?: { signal?: AbortSignal }, + ) => Promise /** * Start a turn from the stored transcript, with no new message. See * `session.continueWithoutMessage`. */ - continueWithoutMessage: (options?: { - inputId?: string - ephemeral?: ReadonlyArray - }) => Promise + continueWithoutMessage: ( + options?: { + inputId?: string + ephemeral?: ReadonlyArray + }, + request?: { signal?: AbortSignal }, + ) => Promise /** * Answer the waits of `snapshot().pendingWaits` that have an * `interrupt`, as `session.resolve` does. The waits of the last turn (the @@ -188,7 +195,10 @@ export interface HarnessClient { * await client.resolve({ resume: [{ interruptId, status: 'resolved' }] }) * ``` */ - resolve: (options: ResolveOptions) => Promise + resolve: ( + options: ResolveOptions, + request?: { signal?: AbortSignal }, + ) => Promise /** * Stop the operation `operationId`, or the running turn without it, as * `session.cancel` does. @@ -198,7 +208,10 @@ export interface HarnessClient { * await client.cancel({ operationId }) * ``` */ - cancel: (options?: CancelOptions) => Promise + cancel: ( + options?: CancelOptions, + request?: { signal?: AbortSignal }, + ) => Promise /** * Cancel an input that waits (see `snapshot().waitingInputs`). An input * that started is refused with `not_waiting`: use `cancel`. @@ -208,15 +221,21 @@ export interface HarnessClient { * await client.cancelInput({ inputId }) * ``` */ - cancelInput: (options: { inputId: string }) => Promise + cancelInput: ( + options: { inputId: string }, + request?: { signal?: AbortSignal }, + ) => Promise /** * Move an input that waits: `'steer'` joins the running turn, `'queue'` * runs it as its own turn later. */ - changeInputDelivery: (change: { - inputId: string - delivery: WaitingInput['delivery'] - }) => Promise + changeInputDelivery: ( + change: { + inputId: string + delivery: WaitingInput['delivery'] + }, + request?: { signal?: AbortSignal }, + ) => Promise agents: ClientAgentHandles /** * Answer a wait of `snapshot().pendingWaits` that no turn holds, as @@ -224,7 +243,10 @@ export interface HarnessClient { * or an agent run, an MCP question, or a `human` wait. The waits of the * stopped turn take `resolve`. */ - answer: (answer: { questionId: string; value: unknown }) => Promise + answer: ( + answer: { questionId: string; value: unknown }, + request?: { signal?: AbortSignal }, + ) => Promise /** * Run the plugin command `name`. Its result arrives as a * `harness.command.result` event. @@ -234,7 +256,10 @@ export interface HarnessClient { * await client.runCommand({ name: 'mode', input: 'plan' }) * ``` */ - runCommand: (options: CommandOptions) => Promise + runCommand: ( + options: CommandOptions, + request?: { signal?: AbortSignal }, + ) => Promise /** * Run a harness tool with no model turn, as `session.runTool` does. The * harness must list it in `expose.tools`. `args` is typed from the @@ -245,12 +270,15 @@ export interface HarnessClient { * await client.runTool({ name: 'sync_inbox', args: { folder: 'INBOX' } }) * ``` */ - runTool: ['name']>(options: { - name: TName - args?: HarnessToolInput - meta?: Record - inputId?: string - }) => Promise + runTool: ['name']>( + options: { + name: TName + args?: HarnessToolInput + meta?: Record + inputId?: string + }, + request?: { signal?: AbortSignal }, + ) => Promise /** * Add host records to the session log, as `session.appendRecords` does. The * harness must list each record type in `expose.records`. With `inputId`, @@ -261,15 +289,21 @@ export interface HarnessClient { * await client.appendRecords({ records: [{ type: 'app.signal', text: 'Saved.' }] }) * ``` */ - appendRecords: (options: { - records: ReadonlyArray - inputId?: string - }) => Promise + appendRecords: ( + options: { + records: ReadonlyArray + inputId?: string + }, + request?: { signal?: AbortSignal }, + ) => Promise /** Change a session setting. */ - setPluginSetting: (change: { - key: string - value: unknown - }) => Promise + setPluginSetting: ( + change: { + key: string + value: unknown + }, + request?: { signal?: AbortSignal }, + ) => Promise /** * Change the stored settings of the thread. See `session.updateThreadSettings`. * @@ -278,7 +312,10 @@ export interface HarnessClient { * await client.updateThreadSettings({ settings: { model: 'strong' } }) * ``` */ - updateThreadSettings: (options: ConfigureOptions) => Promise + updateThreadSettings: ( + options: ConfigureOptions, + request?: { signal?: AbortSignal }, + ) => Promise /** * Start a fresh model context from the next turn. The model sees `note` * first, when you give one. The transcript keeps every message. @@ -288,7 +325,10 @@ export interface HarnessClient { * await client.reset({ note: 'The user wants a new plan.' }) * ``` */ - reset: (options?: ResetOptions) => Promise + reset: ( + options?: ResetOptions, + request?: { signal?: AbortSignal }, + ) => Promise /** * Go back to the message `messageId`, as `session.revert` does. The * harness must list `'undo'` in `expose.commands`. @@ -298,9 +338,12 @@ export interface HarnessClient { * await client.revert({ messageId }) * ``` */ - revert: (options: { messageId: string }) => Promise + revert: ( + options: { messageId: string }, + request?: { signal?: AbortSignal }, + ) => Promise /** End the revert that stands, as `session.unrevert` does. */ - unrevert: () => Promise + unrevert: (request?: { signal?: AbortSignal }) => Promise /** * The session events from `from` (exclusive). Without `from`, the events * start at the cursor of the snapshot: only new events come. Pass @@ -323,9 +366,11 @@ export interface HarnessClient { * again to reconnect: the current statuses come first again. */ hostEvents: (options?: { signal?: AbortSignal }) => AsyncIterable - snapshot: () => Promise + snapshot: (request?: { signal?: AbortSignal }) => Promise /** The saved messages of the thread. */ - transcript: () => Promise> + transcript: (request?: { + signal?: AbortSignal + }) => Promise> /** * The saved messages of a child agent run of the thread, by its * `subagentRunId` or by its `runId`. Empty for an unknown id. A child with @@ -339,12 +384,13 @@ export interface HarnessClient { */ agentTranscript: ( run: { subagentRunId: string } | { runId: string }, + request?: { signal?: AbortSignal }, ) => Promise> /** * The commands, settings, and tools of the session. Only the commands in * `expose.commands` and the config keys in `expose.config`. */ - describe: () => Promise + describe: (request?: { signal?: AbortSignal }) => Promise /** * Store a file in the media store of the thread. Send the record in a * prompt with `mediaPart(record)`. Throws when the handler refuses the @@ -357,6 +403,7 @@ export interface HarnessClient { */ putMedia: ( file: MediaUpload, + request?: { signal?: AbortSignal }, ) => Promise /** * A signed URL for a media file, for ``, `