From 157c17f8e39ee8621ca92716a0a20d9dfc737f71 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:11 +0100 Subject: [PATCH 01/24] fix(ai-claude-code,ai-codex,ai-grok-build): give the chat() request signal to the tool bridge --- ...abort-cli-agent-tool-bridge-claude-code.md | 5 +++ .../abort-cli-agent-tool-bridge-codex.md | 5 +++ .../abort-cli-agent-tool-bridge-grok-build.md | 5 +++ packages/ai-claude-code/src/adapters/text.ts | 4 +- .../ai-claude-code/tests/tool-choice.test.ts | 39 ++++++++++++++++++ packages/ai-codex/src/adapters/text.ts | 4 +- packages/ai-codex/tests/tool-choice.test.ts | 39 ++++++++++++++++++ packages/ai-grok-build/src/adapters/text.ts | 4 +- .../ai-grok-build/tests/tool-choice.test.ts | 40 +++++++++++++++++++ 9 files changed, 142 insertions(+), 3 deletions(-) create mode 100644 .changeset/abort-cli-agent-tool-bridge-claude-code.md create mode 100644 .changeset/abort-cli-agent-tool-bridge-codex.md create mode 100644 .changeset/abort-cli-agent-tool-bridge-grok-build.md diff --git a/.changeset/abort-cli-agent-tool-bridge-claude-code.md b/.changeset/abort-cli-agent-tool-bridge-claude-code.md new file mode 100644 index 0000000000..ce3098bbd6 --- /dev/null +++ b/.changeset/abort-cli-agent-tool-bridge-claude-code.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-claude-code': patch +--- + +The host tool bridge now gets the run signal from `chat()`. Before, it read only `abortController`, which `chat()` does not set, so a stopped run did not abort the bridged tools. diff --git a/.changeset/abort-cli-agent-tool-bridge-codex.md b/.changeset/abort-cli-agent-tool-bridge-codex.md new file mode 100644 index 0000000000..c5664cae66 --- /dev/null +++ b/.changeset/abort-cli-agent-tool-bridge-codex.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-codex': patch +--- + +The host tool bridge now gets the run signal from `chat()`. Before, it read only `abortController`, which `chat()` does not set, so a stopped run did not abort the bridged tools. diff --git a/.changeset/abort-cli-agent-tool-bridge-grok-build.md b/.changeset/abort-cli-agent-tool-bridge-grok-build.md new file mode 100644 index 0000000000..95d420f0b4 --- /dev/null +++ b/.changeset/abort-cli-agent-tool-bridge-grok-build.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-grok-build': patch +--- + +The host tool bridge of the `streaming-json` protocol now gets the run signal from `chat()`. Before, it read only `abortController`, which `chat()` does not set, so a stopped run did not abort the bridged tools. diff --git a/packages/ai-claude-code/src/adapters/text.ts b/packages/ai-claude-code/src/adapters/text.ts index ee5f110dd8..c8706e32fd 100644 --- a/packages/ai-claude-code/src/adapters/text.ts +++ b/packages/ai-claude-code/src/adapters/text.ts @@ -521,7 +521,9 @@ export class ClaudeCodeTextAdapter< ...(permission !== undefined ? { permission } : {}), ...(options.abortController?.signal ? { signal: options.abortController.signal } - : {}), + : options.request?.signal + ? { signal: options.request.signal } + : {}), }) } diff --git a/packages/ai-claude-code/tests/tool-choice.test.ts b/packages/ai-claude-code/tests/tool-choice.test.ts index bbb7cb1bd2..131e6f522a 100644 --- a/packages/ai-claude-code/tests/tool-choice.test.ts +++ b/packages/ai-claude-code/tests/tool-choice.test.ts @@ -160,3 +160,42 @@ describe('claude-code toolChoice', () => { expect(warn).not.toHaveBeenCalled() }) }) + +describe('tool bridge signal', () => { + // chat() passes the run signal as `request.signal` and sets no + // `abortController`, so the bridge must take `request.signal`. + it('gives the chat() request signal to the tool bridge', async () => { + const sbx = await provider.create({}) + try { + const capabilities = { + capabilities: { markProvided: () => {}, has: () => true }, + } as unknown as CapabilityContext + provideSandbox(capabilities, sbx) + let bridgeSignal: AbortSignal | undefined + provideToolBridgeProvisioner(capabilities, { + provision: (_tools, options) => { + bridgeSignal = options.signal + return Promise.reject(new Error('stop after the bridge')) + }, + }) + const adapter = claudeCodeText('haiku', { + claudeExecutable: 'node fake-claude.mjs', + }) + const signal = new AbortController().signal + await collect( + adapter.chatStream({ + threadId: 'thread-1', + model: 'haiku', + messages: [{ role: 'user', content: 'hi' }], + logger: loggerWith(() => {}), + capabilities, + tools, + request: { signal }, + }), + ).catch(() => undefined) + expect(bridgeSignal).toBe(signal) + } finally { + await sbx.destroy() + } + }) +}) diff --git a/packages/ai-codex/src/adapters/text.ts b/packages/ai-codex/src/adapters/text.ts index 5334e855f6..f7175bc7cf 100644 --- a/packages/ai-codex/src/adapters/text.ts +++ b/packages/ai-codex/src/adapters/text.ts @@ -353,7 +353,9 @@ export class CodexTextAdapter< emitCustomEvent: channel.emitCustomEvent, ...(options.abortController?.signal ? { signal: options.abortController.signal } - : {}), + : options.request?.signal + ? { signal: options.request.signal } + : {}), }) } diff --git a/packages/ai-codex/tests/tool-choice.test.ts b/packages/ai-codex/tests/tool-choice.test.ts index ad6e0652a2..b89128b1a8 100644 --- a/packages/ai-codex/tests/tool-choice.test.ts +++ b/packages/ai-codex/tests/tool-choice.test.ts @@ -146,3 +146,42 @@ describe('codex toolChoice', () => { expect(warn).not.toHaveBeenCalled() }) }) + +describe('tool bridge signal', () => { + // chat() passes the run signal as `request.signal` and sets no + // `abortController`, so the bridge must take `request.signal`. + it('gives the chat() request signal to the tool bridge', async () => { + const sbx = await provider.create({}) + try { + const capabilities = { + capabilities: { markProvided: () => {}, has: () => true }, + } as unknown as CapabilityContext + provideSandbox(capabilities, sbx) + let bridgeSignal: AbortSignal | undefined + provideToolBridgeProvisioner(capabilities, { + provision: (_tools, options) => { + bridgeSignal = options.signal + return Promise.reject(new Error('stop after the bridge')) + }, + }) + const adapter = codexText(MODEL, { + codexExecutable: 'node fake-codex.mjs', + }) + const signal = new AbortController().signal + await collect( + adapter.chatStream({ + threadId: 'thread-1', + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + logger: loggerWith(() => {}), + capabilities, + tools, + request: { signal }, + }), + ).catch(() => undefined) + expect(bridgeSignal).toBe(signal) + } finally { + await sbx.destroy() + } + }) +}) diff --git a/packages/ai-grok-build/src/adapters/text.ts b/packages/ai-grok-build/src/adapters/text.ts index c2ae29da5e..3ee4dc8d19 100644 --- a/packages/ai-grok-build/src/adapters/text.ts +++ b/packages/ai-grok-build/src/adapters/text.ts @@ -667,7 +667,9 @@ export class GrokBuildTextAdapter< context: options.context, ...(options.abortController?.signal ? { signal: options.abortController.signal } - : {}), + : options.request?.signal + ? { signal: options.request.signal } + : {}), }) // Grok reads MCP from `/.grok/config.toml`, not `--mcp-config`. await projectGrokMcpBridge(sandbox, cwd, bridge) diff --git a/packages/ai-grok-build/tests/tool-choice.test.ts b/packages/ai-grok-build/tests/tool-choice.test.ts index 4200eceece..07b6b22fe0 100644 --- a/packages/ai-grok-build/tests/tool-choice.test.ts +++ b/packages/ai-grok-build/tests/tool-choice.test.ts @@ -153,3 +153,43 @@ describe('grok-build toolChoice', () => { expect(warn).not.toHaveBeenCalled() }) }) + +describe('tool bridge signal', () => { + // chat() passes the run signal as `request.signal` and sets no + // `abortController`, so the bridge must take `request.signal`. + it('gives the chat() request signal to the tool bridge', async () => { + const sbx = await provider.create({}) + try { + const capabilities = { + capabilities: { markProvided: () => {}, has: () => true }, + } as unknown as CapabilityContext + provideSandbox(capabilities, sbx) + let bridgeSignal: AbortSignal | undefined + provideToolBridgeProvisioner(capabilities, { + provision: (_tools, options) => { + bridgeSignal = options.signal + return Promise.reject(new Error('stop after the bridge')) + }, + }) + const adapter = grokBuildText(MODEL, { + grokExecutable: 'node fake-grok.mjs', + protocol: 'streaming-json', + }) + const signal = new AbortController().signal + await collect( + adapter.chatStream({ + threadId: 'thread-1', + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + logger: loggerWith(() => {}), + capabilities, + tools, + request: { signal }, + }), + ).catch(() => undefined) + expect(bridgeSignal).toBe(signal) + } finally { + await sbx.destroy() + } + }) +}) From 603b2dbec00b486d3f5be178865f247986221ef8 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:13 +0100 Subject: [PATCH 02/24] fix(ai-sandbox): kill the spawnNdjson process on early stop, forward the snapshot tool signal --- .changeset/abort-sandbox-spawn-ndjson.md | 5 ++ packages/ai-sandbox/src/runner.ts | 27 +++++---- packages/ai-sandbox/src/snapshot-tools.ts | 3 +- packages/ai-sandbox/tests/runner.test.ts | 55 ++++++++++++++++++- .../ai-sandbox/tests/snapshot-tools.test.ts | 15 +++++ 5 files changed, 93 insertions(+), 12 deletions(-) create mode 100644 .changeset/abort-sandbox-spawn-ndjson.md diff --git a/.changeset/abort-sandbox-spawn-ndjson.md b/.changeset/abort-sandbox-spawn-ndjson.md new file mode 100644 index 0000000000..1b65bc8d7f --- /dev/null +++ b/.changeset/abort-sandbox-spawn-ndjson.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox': patch +--- + +`spawnNdjson` now kills the agent process when the consumer stops reading early. The `save_sandbox_snapshot` tool now gives the abort signal of the call to `snapshots.save`. diff --git a/packages/ai-sandbox/src/runner.ts b/packages/ai-sandbox/src/runner.ts index a522ee1f08..13c0f4c08d 100644 --- a/packages/ai-sandbox/src/runner.ts +++ b/packages/ai-sandbox/src/runner.ts @@ -415,17 +415,24 @@ export async function* spawnNdjson( } })() - for await (const line of toLines(proc.stdout)) { - const trimmed = line.trim() - if (trimmed === '') continue - let parsed: unknown - try { - parsed = JSON.parse(trimmed) - } catch { - onNonJsonLine?.(trimmed) - continue + // A consumer that stops early (or throws) leaves the agent running. Kill it. + let isRead = false + try { + for await (const line of toLines(proc.stdout)) { + const trimmed = line.trim() + if (trimmed === '') continue + let parsed: unknown + try { + parsed = JSON.parse(trimmed) + } catch { + onNonJsonLine?.(trimmed) + continue + } + yield parsed } - yield parsed + isRead = true + } finally { + if (!isRead) await proc.kill().catch(() => undefined) } const exitCode = await proc.wait() diff --git a/packages/ai-sandbox/src/snapshot-tools.ts b/packages/ai-sandbox/src/snapshot-tools.ts index 8725e42abf..2b3e3d7c61 100644 --- a/packages/ai-sandbox/src/snapshot-tools.ts +++ b/packages/ai-sandbox/src/snapshot-tools.ts @@ -92,13 +92,14 @@ export function createSnapshotTools( required: ['checkpointId', 'label', 'threadId'], additionalProperties: false, }, - }).server(async (input) => { + }).server(async (input, context) => { const label = requiredString(input, 'label') const checkpoint = await snapshots.save({ threadId, runId, label, ...(tenant === undefined ? {} : { tenant }), + ...(context?.abortSignal ? { signal: context.abortSignal } : {}), }) return { checkpointId: checkpoint.id, diff --git a/packages/ai-sandbox/tests/runner.test.ts b/packages/ai-sandbox/tests/runner.test.ts index b9ed3ba2f6..2570e5e77c 100644 --- a/packages/ai-sandbox/tests/runner.test.ts +++ b/packages/ai-sandbox/tests/runner.test.ts @@ -1,5 +1,5 @@ import { readFileSync } from 'node:fs' -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { EventType, memoryStream } from '@tanstack/ai' import { readJournalNdjson, @@ -122,6 +122,59 @@ describe('spawnNdjson', () => { ]) expect(nonJson).toEqual(['Claude Code starting...']) }) + + it('kills the process when the consumer stops early', async () => { + const kill = vi.fn(() => Promise.resolve()) + const base = handleSpawning([]) + const handle: SandboxHandle = { + ...base, + process: { + ...base.process, + spawn: () => + Promise.resolve({ + pid: 1, + stdout: fromChunks(['{"n":1}\n', '{"n":2}\n']), + stderr: fromChunks([]), + stdin: { + write: () => Promise.resolve(), + end: () => Promise.resolve(), + }, + wait: () => Promise.resolve(0), + kill, + }), + }, + } + for await (const event of spawnNdjson(handle, 'agent')) { + expect(event).toEqual({ n: 1 }) + break + } + expect(kill).toHaveBeenCalledTimes(1) + }) + + it('does not kill a process that ended', async () => { + const kill = vi.fn(() => Promise.resolve()) + const base = handleSpawning([]) + const handle: SandboxHandle = { + ...base, + process: { + ...base.process, + spawn: () => + Promise.resolve({ + pid: 1, + stdout: fromChunks(['{"n":1}\n']), + stderr: fromChunks([]), + stdin: { + write: () => Promise.resolve(), + end: () => Promise.resolve(), + }, + wait: () => Promise.resolve(0), + kill, + }), + }, + } + expect(await collect(spawnNdjson(handle, 'agent'))).toEqual([{ n: 1 }]) + expect(kill).not.toHaveBeenCalled() + }) }) /** diff --git a/packages/ai-sandbox/tests/snapshot-tools.test.ts b/packages/ai-sandbox/tests/snapshot-tools.test.ts index 2461b0d1ac..8b65d17c54 100644 --- a/packages/ai-sandbox/tests/snapshot-tools.test.ts +++ b/packages/ai-sandbox/tests/snapshot-tools.test.ts @@ -113,6 +113,21 @@ describe('createSnapshotTools', () => { expect(provider.calls).toMatchObject({ create: 0, resume: 1 }) }) + it('gives the abort signal of the call to the save', async () => { + const { snapshots } = await liveSnapshots() + const save = vi.fn(snapshots.save) + const tool = toolsFor({ ...snapshots, save }).find( + (candidate) => candidate.name === 'save_sandbox_snapshot', + ) + if (tool?.execute === undefined) throw new Error('Missing save tool') + const { signal } = new AbortController() + await tool.execute( + { label: 'release-1' }, + { abortSignal: signal, emitCustomEvent: () => undefined }, + ) + expect(save).toHaveBeenCalledWith(expect.objectContaining({ signal })) + }) + it('rejects save while this thread already has a writer', async () => { const { snapshots } = await liveSnapshots() const writer = await snapshots.checkpoints.acquireWriter(THREAD) From 8850d739de82c62dc65424284672ab7d5fb4969d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:14 +0100 Subject: [PATCH 03/24] fix(ai-acp): cancel the question form when the operation ends --- .changeset/abort-acp-question-elicitation.md | 5 ++ packages/ai-acp/src/agent/index.ts | 23 ++++-- packages/ai-acp/tests/agent.test.ts | 77 ++++++++++++++++++++ 3 files changed, 98 insertions(+), 7 deletions(-) create mode 100644 .changeset/abort-acp-question-elicitation.md diff --git a/.changeset/abort-acp-question-elicitation.md b/.changeset/abort-acp-question-elicitation.md new file mode 100644 index 0000000000..6c04ac4c19 --- /dev/null +++ b/.changeset/abort-acp-question-elicitation.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-acp': patch +--- + +When a turn ends while the editor shows a question form (for example, another client answered it), the agent now sends `$/cancel_request` for that form. Before, the turn waited for the editor and did not report idle. diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts index a23ee5bb71..bee631f464 100644 --- a/packages/ai-acp/src/agent/index.ts +++ b/packages/ai-acp/src/agent/index.ts @@ -513,12 +513,15 @@ export function createAcpAgent( wait, toolCallId, url, + signal, }: { client: AgentContext sessionId: string wait: Wait toolCallId?: string url?: string + /** Sends `$/cancel_request` for the question when it aborts. */ + signal?: AbortSignal }): Promise => { if (!(await canAsk(client, wait, url))) return { status: 'cancelled' } if (wait.kind === 'approval') { @@ -539,12 +542,16 @@ export function createAcpAgent( const scope = { sessionId, ...(toolCallId ? { toolCallId } : {}) } if (wait.kind === 'question' || wait.kind === 'human') { const form = formOf(wait) - const reply = await client.request('elicitation/create', { - ...scope, - mode: 'form', - message: wait.message ?? 'The agent asks for input.', - requestedSchema: form.requestedSchema, - }) + const reply = await client.request( + 'elicitation/create', + { + ...scope, + mode: 'form', + message: wait.message ?? 'The agent asks for input.', + requestedSchema: form.requestedSchema, + }, + { cancellationSignal: signal }, + ) if (reply.action === 'cancel') return { status: 'stop' } if (reply.action !== 'accept') return { status: 'cancelled' } const content = 'content' in reply ? reply.content : undefined @@ -655,7 +662,9 @@ export function createAcpAgent( !turn && kind === 'question' && waitId === wait.waitId, ) if (!isQuestion) return - const reply = await askEditor({ client, sessionId, wait }) + // The operation can end while the editor shows the question (another + // client answered it). Then the editor gets a cancel for it. + const reply = await askEditor({ client, sessionId, wait, signal }) if (signal.aborted) return if (reply.status !== 'resolved') { await session.cancel({ operationId }) diff --git a/packages/ai-acp/tests/agent.test.ts b/packages/ai-acp/tests/agent.test.ts index 89c8d95cc5..9f9e17ce6b 100644 --- a/packages/ai-acp/tests/agent.test.ts +++ b/packages/ai-acp/tests/agent.test.ts @@ -357,6 +357,83 @@ describe('ACP v2 agent', () => { await host.close() }) + it('cancels the form elicitation of a question that another client answered', async () => { + const greet = durableTool( + toolDefinition({ + name: 'greet', + description: 'Ask the user for a name', + inputSchema: { type: 'object', properties: {} }, + }), + async (_args, { step }) => { + const answer = await step.ask({ + name: 'name', + message: 'Name?', + schema: { + type: 'object', + properties: { name: { type: 'string' } }, + required: ['name'], + }, + }) + return { greeted: answer } + }, + ) + const harness = defineHarness({ + name: 'test/acp-answered-elsewhere', + adapter: scripted([toolTurn('greet', '{}'), textTurn('Greeted.')]), + tools: [greet], + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const updates: Array = [] + let asked = false + // An editor that answers the question only when the agent cancels it. + const editor = client() + .onNotification('session/update', ({ params }) => { + updates.push(params.update) + }) + .onRequest('elicitation/create', ({ signal }) => { + asked = true + return new Promise((resolve) => { + signal.addEventListener( + 'abort', + () => resolve({ action: 'cancel' }), + { once: true }, + ) + }) + }) + + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'test-editor', version: '1.0.0' }, + capabilities: { elicitation: { form: {} } }, + }) + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'greet me' }], + }) + await waitFor(() => expect(asked).toBe(true)) + // Another client of the thread answers the question. + const other = await host.open(harness, { threadId: sessionId }) + const question = other + .snapshot() + .pendingWaits.find((wait) => wait.kind === 'question') + if (!question) throw new Error('No question waits') + await other.answer({ + questionId: question.waitId, + value: { name: 'Ada' }, + }) + await waitFor(() => expect(isIdle(updates)).toBe(true)) + await other.close() + }) + + expect(agentText(updates)).toContain('Greeted.') + await host.close() + }) + it('refuses a second connection to one app, and keeps the first', async () => { const greet = durableTool( toolDefinition({ From ffd0aa534d532125be61a768286c316f7821777a Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:16 +0100 Subject: [PATCH 04/24] fix(ai-harness): give search commands the call signal, cancel harnessText turns on early stop --- .changeset/abort-harness-exec-signal.md | 5 ++ .../src/first-party/coding/search.ts | 8 +- .../src/first-party/coding/workspace.ts | 11 ++- packages/ai-harness/src/harness-text.ts | 8 ++ .../ai-harness/tests/coding-search.test.ts | 26 ++++++ .../ai-harness/tests/harness-text.test.ts | 90 +++++++++++++++++++ 6 files changed, 144 insertions(+), 4 deletions(-) create mode 100644 .changeset/abort-harness-exec-signal.md diff --git a/.changeset/abort-harness-exec-signal.md b/.changeset/abort-harness-exec-signal.md new file mode 100644 index 0000000000..eb967247af --- /dev/null +++ b/.changeset/abort-harness-exec-signal.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-harness': patch +--- + +The `list_files` and `grep` commands now get the abort signal of the tool call. `harnessText` now cancels the inner turn, or the remote turn, when the caller stops reading the stream early. diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index 56873d9278..ca1dbb4358 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -95,11 +95,17 @@ async function bundledRg() { } } +/** + * The `exec` wrappers of `hostBackend.exec` that give each command the + * signal of its tool call (see `withCallSignal`). They run on this host too. + */ +export const hostExecs = new WeakSet() + /** The command that starts rg on `backend`, or `undefined` without rg. */ async function findRg(backend: WorkspaceBackend) { // The bundled binary is a host file. A sandbox cannot run it. A bound // workspace is a copy of `hostBackend` with a root, so compare `exec`. - if (backend.exec === hostBackend.exec) { + if (backend.exec === hostBackend.exec || hostExecs.has(backend.exec)) { const bundled = await bundledRg() if (bundled !== undefined) return quoteArg(bundled, process.platform) } diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index b83a4f2a2f..22ef2dc134 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -27,7 +27,7 @@ import { WorkspaceStranded } from './events' import { editTools, fileVersions } from './edit' import { patchPaths, patchTools } from './patch' import { readTools } from './read' -import { searchTools } from './search' +import { hostExecs, searchTools } from './search' import { webTools } from './web' import type { AnyTool, ToolExecutionContext } from '@tanstack/ai' import type { PluginEvent } from '../../extensions' @@ -272,8 +272,8 @@ export function writingRun() { const callingChat = () => calls.getStore()?.chat /** - * `backend`, where each file call gets the abort signal of the tool call - * that runs now, so an aborted call stops its file work too. + * `backend`, where each file call and command gets the abort signal of the + * tool call that runs now, so an aborted call stops its work too. */ function withCallSignal(backend: WorkspaceBackend): WorkspaceBackend { // No options outside a call, or for a call without a signal. @@ -282,8 +282,13 @@ function withCallSignal(backend: WorkspaceBackend): WorkspaceBackend { return chat ? ([{ signal: chat }] as const) : ([] as const) } const { readLines, remove } = backend + // A signal of the caller (like the `bash` stop) wins. + const exec: WorkspaceBackend['exec'] = (command, options) => + backend.exec(command, { ...signal()[0], ...options }) + if (backend.exec === hostBackend.exec) hostExecs.add(exec) return { ...backend, + exec, readFile: (path) => backend.readFile(path, ...signal()), writeFile: (path, data) => backend.writeFile(path, data, ...signal()), stat: (path) => backend.stat(path, ...signal()), diff --git a/packages/ai-harness/src/harness-text.ts b/packages/ai-harness/src/harness-text.ts index 17d7dc1b7d..04271efcd3 100644 --- a/packages/ai-harness/src/harness-text.ts +++ b/packages/ai-harness/src/harness-text.ts @@ -366,6 +366,8 @@ function remoteHarnessText(remote: RemoteHarness): AnyTextAdapter { }).catch(() => {}) } signal?.addEventListener('abort', cancelRemote, { once: true }) + // A caller that stops reading early sends no abort. Cancel then too. + let isEnded = false try { const response = await doFetch(`${base}/run`, { method: 'POST', @@ -408,6 +410,7 @@ function remoteHarnessText(remote: RemoteHarness): AnyTextAdapter { (chunk, isOwn) => FORWARDED.has(chunk.type) && isOwn ? chunk : undefined, ) + isEnded = true if (failed !== undefined) { yield { type: EventType.RUN_ERROR, @@ -437,6 +440,7 @@ function remoteHarnessText(remote: RemoteHarness): AnyTextAdapter { } finally { // A long-lived signal must not keep one listener per call. signal?.removeEventListener('abort', cancelRemote) + if (!isEnded && !signal?.aborted) cancelRemote() } })(), } @@ -523,6 +527,8 @@ export function harnessText( // A signal that aborted already never fires its event. if (signal?.aborted) cancel() else signal?.addEventListener('abort', cancel, { once: true }) + // A caller that stops reading early sends no abort. Cancel then too. + let isEnded = false try { yield { type: EventType.RUN_STARTED, @@ -535,6 +541,7 @@ export function harnessText( (chunk, isOwn) => FORWARDED.has(chunk.type) && isOwn ? chunk : undefined, ) + isEnded = true if (failed !== undefined) { yield { type: EventType.RUN_ERROR, @@ -555,6 +562,7 @@ export function harnessText( } finally { // A long-lived signal must not keep one listener per call. signal?.removeEventListener('abort', cancel) + if (!isEnded && !signal?.aborted) cancel() await session.close() } })(), diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index f4bf473ed9..22abab465b 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -573,3 +573,29 @@ describe('quoteArg', () => { }, ) }) + +describe('the abort signal of a call', () => { + it('goes to each search command', async () => { + const fake = fakeBackend({ 'a.ts': 'one\n' }, (command) => + command === GIT ? ok('a.ts\0') : undefined, + ) + const signals: Array = [] + const backend: WorkspaceBackend = { + ...fake.backend, + exec: (command, options) => { + signals.push(options?.signal) + return fake.backend.exec(command, options) + }, + } + const { tools } = createWorkspaceTools({ root, backend }) + const listFiles = tools.find((each) => each.name === 'list_files')?.execute + if (!listFiles) throw new Error('No tool list_files') + const { signal } = new AbortController() + await listFiles( + {}, + { abortSignal: signal, emitCustomEvent: () => undefined }, + ) + expect(signals.length).toBeGreaterThan(0) + expect(signals.filter((each) => each !== signal)).toEqual([]) + }) +}) diff --git a/packages/ai-harness/tests/harness-text.test.ts b/packages/ai-harness/tests/harness-text.test.ts index eece5b6699..166a8aeb2e 100644 --- a/packages/ai-harness/tests/harness-text.test.ts +++ b/packages/ai-harness/tests/harness-text.test.ts @@ -345,6 +345,96 @@ describe('harnessText', () => { }) }) + it('cancels the remote turn when the caller stops reading early', async () => { + const requests: Array<{ url: string; body: unknown }> = [] + const encoder = new TextEncoder() + const delta = { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'working', + } + const adapter = harnessText({ + url: 'http://remote.test/', + fetch: async (url, init) => { + requests.push({ + url: String(url), + body: typeof init?.body === 'string' ? JSON.parse(init.body) : null, + }) + if (!String(url).endsWith('/run')) return Response.json({}) + // A turn that streams one delta, then runs on. + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode(`data: ${JSON.stringify(delta)}\n\n`), + ) + }, + }), + { status: 200 }, + ) + }, + }) + + for await (const chunk of adapter.chatStream({ + model: adapter.model, + messages: [{ role: 'user', content: 'hello' }], + logger: resolveDebugOption(false), + threadId: 'outer-remote', + runId: 'run-1', + })) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) break + } + + await vi.waitFor(() => expect(requests).toHaveLength(2)) + expect(requests[1]).toEqual({ + url: 'http://remote.test/control', + body: { + threadId: 'outer-remote', + input: { op: 'cancel', operationId: 'run-1' }, + }, + }) + }) + + it('cancels the inner turn when the caller stops reading early', async () => { + const inner = mockAdapter([ + (options) => + (async function* () { + yield { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'working', + timestamp: Date.now(), + } + // A turn that runs until it is cancelled. + const signal: AbortSignal | undefined = options.request?.signal + await new Promise((resolve) => { + if (!signal || signal.aborted) return resolve() + signal.addEventListener('abort', () => resolve(), { once: true }) + }) + })(), + ]) + const studio = defineHarness({ name: 'test/inner', adapter: inner.adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const adapter = harnessText(studio, { host }) + + const stopped = (async () => { + for await (const chunk of adapter.chatStream({ + model: adapter.model, + messages: [{ role: 'user', content: 'hello' }], + logger: resolveDebugOption(false), + threadId: 'outer-early', + })) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) break + } + })() + + await vi.waitFor(() => + expect(inner.calls[0]?.request?.signal?.aborted).toBe(true), + ) + await stopped + await host.close() + }) + it('opens the inner session as the user of each outer harness turn', async () => { const inner = mockAdapter([() => text('one'), () => text('two')]) const seen: Array = [] From 996faee448e160707939e3650cf8bc5225e7ace4 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:17 +0100 Subject: [PATCH 05/24] fix(ai-mcp): handle tasks/cancel for task tools, stop a cancelled elicitation --- .changeset/abort-mcp-task-cancel.md | 5 ++ packages/ai-mcp/src/client.ts | 10 ++- packages/ai-mcp/src/harness-plugin.ts | 7 +- packages/ai-mcp/src/server/create-server.ts | 49 ++++++++++--- packages/ai-mcp/src/server/tasks.ts | 44 ++++++++++-- packages/ai-mcp/tests/harness-elicit.test.ts | 71 +++++++++++++++++++ .../ai-mcp/tests/server/create-server.test.ts | 56 +++++++++++++++ 7 files changed, 221 insertions(+), 21 deletions(-) create mode 100644 .changeset/abort-mcp-task-cancel.md diff --git a/.changeset/abort-mcp-task-cancel.md b/.changeset/abort-mcp-task-cancel.md new file mode 100644 index 0000000000..c28696b72e --- /dev/null +++ b/.changeset/abort-mcp-task-cancel.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-mcp': patch +--- + +`createMCPServer` now answers `tasks/cancel` for an `execution: 'task'` tool on spec 2025. The task status becomes `cancelled`, `ctx.abortSignal` of the tool aborts, and a late result does not replace the cancel. A spec 2025 `elicitation/create` that the server cancels now also stops the question in the harness. diff --git a/packages/ai-mcp/src/client.ts b/packages/ai-mcp/src/client.ts index 53e164b79c..f7e8aeaaf4 100644 --- a/packages/ai-mcp/src/client.ts +++ b/packages/ai-mcp/src/client.ts @@ -310,8 +310,8 @@ class MCPClientImpl implements Omit< ) // A spec 2025 server asks for input with a request to the client. if (onElicit) { - this.#client.setRequestHandler('elicitation/create', (request) => - onElicit(request.params), + this.#client.setRequestHandler('elicitation/create', (request, ctx) => + onElicit(request.params, ctx.mcpReq.signal), ) } } @@ -773,9 +773,13 @@ export async function createMCPClient( return connectTransport(options) } -/** Answers a spec 2025 `elicitation/create` request. */ +/** + * Answers a spec 2025 `elicitation/create` request. `signal` aborts when the + * server cancels the request. + */ export type OnElicit = ( params: ElicitRequest['params'], + signal?: AbortSignal, ) => Promise /** diff --git a/packages/ai-mcp/src/harness-plugin.ts b/packages/ai-mcp/src/harness-plugin.ts index b0a03d5567..3441afca9b 100644 --- a/packages/ai-mcp/src/harness-plugin.ts +++ b/packages/ai-mcp/src/harness-plugin.ts @@ -642,7 +642,7 @@ export function askForInput( * never go to the user of another run. */ export function elicitFor(calling: ReadonlyArray): OnElicit { - return async (params) => { + return async (params, signal) => { const [caller] = calling if (!caller) return { action: 'decline' } const isOneRun = calling.every( @@ -652,7 +652,7 @@ export function elicitFor(calling: ReadonlyArray): OnElicit { other.run.runId === caller.run?.runId, ) if (!isOneRun) return { action: 'decline' } - return elicitResult(await ask(caller.session, params, caller.run)) + return elicitResult(await ask(caller.session, params, caller.run, signal)) } } @@ -665,6 +665,7 @@ async function ask( session: PluginSessionApi, request: unknown, run: RunContext | undefined, + signal?: AbortSignal, ) { if (!run) return { action: 'decline' } const body = isRecord(request) ? request : {} @@ -679,6 +680,8 @@ async function ask( ? body.requestedSchema : undefined, ...(url ? { url } : {}), + // The server cancelled its request: the question stops too. + ...(signal ? { signal } : {}), live: { name: 'mcp-elicitation', run }, }) // A URL request has no content: the answer only says "done". diff --git a/packages/ai-mcp/src/server/create-server.ts b/packages/ai-mcp/src/server/create-server.ts index 66c61fa2fa..474613d00d 100644 --- a/packages/ai-mcp/src/server/create-server.ts +++ b/packages/ai-mcp/src/server/create-server.ts @@ -36,7 +36,7 @@ import type { SampleRequest, ToolInputRequest } from './context' import { isServedOverStdio, rememberServerOptions } from './registry' import { parseToolOutput } from './output' import type { MCPResourceContext, MCPResourceList } from './definitions' -import { getTask, startTask, toCallToolResult } from './tasks' +import { cancelTask, getTask, startTask, toCallToolResult } from './tasks' import { inMemoryTaskStore } from './stores' import type { TaskStore } from './stores' @@ -521,7 +521,7 @@ function serverOptions(era: ProtocolYear, hasTaskTool: boolean) { resources: { listChanged: false }, prompts: { listChanged: false }, ...(hasTaskTool && era === '2025' - ? { tasks: { requests: { tools: { call: {} } } } } + ? { tasks: { requests: { tools: { call: {} } }, cancel: {} } } : {}), }, } @@ -619,12 +619,19 @@ async function runTaskTool( }, authInfo: AuthInfo | undefined, ) { - const ctx = taskContext(input.options.sample, authInfo, input.appContext) - const handle = await startTask(() => runTool(tool, args, ctx), { - store: input.taskStore, - waitUntil: input.options.waitUntil, - owner: input.owner, - }) + const handle = await startTask( + (signal) => + runTool( + tool, + args, + taskContext(input.options.sample, authInfo, input.appContext, signal), + ), + { + store: input.taskStore, + waitUntil: input.options.waitUntil, + owner: input.owner, + }, + ) const polled = await getTask(handle.taskId, input.taskStore, input.owner) if (polled === null) { throw new Error(`Task ${handle.taskId} was not saved.`) @@ -639,10 +646,12 @@ async function runTaskTool( // A task keeps running after tools/call answers with the task handle. // So it must not use that request: its signal, elicitation, or sampling. +// `signal` aborts on `tasks/cancel`. function taskContext( sample: MCPServerOptions['sample'], authInfo: AuthInfo | undefined, appContext: AppContext, + signal: AbortSignal, ) { return { context: { @@ -662,8 +671,7 @@ function taskContext( return sample(request) }, }, - // ponytail: nothing aborts this signal until tasks/cancel is implemented. - abortSignal: new AbortController().signal, + abortSignal: signal, emitCustomEvent() {}, } } @@ -929,6 +937,27 @@ function registerLegacyTaskMethods( return toCallToolResult(polled.record.result) }, ) + server.server.setRequestHandler( + 'tasks/cancel', + { params: taskIdParams }, + async (params) => { + const polled = await getTask(taskIdFrom(params), store, owner) + if (polled === null) { + throw new ProtocolError( + ProtocolErrorCode.InvalidParams, + 'Task not found', + ) + } + // Spec 2025-11-25: a task in a terminal status cannot be cancelled. + if (polled.record.status !== 'working') { + throw new ProtocolError( + ProtocolErrorCode.InvalidParams, + `Task is ${polled.record.status}, so it cannot be cancelled`, + ) + } + return cancelTask(polled.record, store) + }, + ) } function taskIdFrom(params: unknown) { diff --git a/packages/ai-mcp/src/server/tasks.ts b/packages/ai-mcp/src/server/tasks.ts index 0c62e7a278..7479efefe6 100644 --- a/packages/ai-mcp/src/server/tasks.ts +++ b/packages/ai-mcp/src/server/tasks.ts @@ -10,7 +10,7 @@ const toolFailedMessage = 'The tool failed.' * fields that never leave the server: the owner and the tool result. */ type StoredTask = Task & { - status: 'working' | 'completed' | 'failed' + status: 'working' | 'completed' | 'failed' | 'cancelled' /** The auth subject that started the task. Absent without auth. */ owner?: string /** The tool output. Present when `status` is `completed`. */ @@ -30,6 +30,9 @@ type StartTaskOptions = { owner?: string } +/** The abort controllers of the tasks that run in this process, by id. */ +const running = new Map() + /** * Starts a tool run and returns a task id before the run finishes. * @@ -43,6 +46,9 @@ type StartTaskOptions = { * A tool error does not reject this function. * The store records the error on the task as `statusMessage`. * + * `run` gets a signal that `cancelTask` aborts. A cancelled task keeps its + * `cancelled` status: the late result or error of `run` is not saved. + * * @param run - Tool function. It returns the tool result. * @param options - `store` is required. `waitUntil` is optional. * @@ -55,7 +61,7 @@ type StartTaskOptions = { * ``` */ export async function startTask( - run: () => Promise, + run: (signal: AbortSignal) => Promise, options: StartTaskOptions, ) { const taskId = crypto.randomUUID() @@ -69,21 +75,27 @@ export async function startTask( ...(options.owner === undefined ? {} : { owner: options.owner }), } await options.store.set(taskId, working) + const controller = new AbortController() + running.set(taskId, controller) + const settle = async (task: StoredTask) => { + running.delete(taskId) + if (!controller.signal.aborted) await options.store.set(taskId, task) + } - // A tool can throw before it returns a promise. `.then(run)` turns that + // A tool can throw before it returns a promise. `.then` turns that // throw into a rejection, so the store records it. const inflight = Promise.resolve() - .then(run) + .then(() => run(controller.signal)) .then( (result) => - options.store.set(taskId, { + settle({ ...working, status: 'completed', lastUpdatedAt: new Date().toISOString(), result, }), (error: unknown) => - options.store.set(taskId, { + settle({ ...working, status: 'failed', lastUpdatedAt: new Date().toISOString(), @@ -131,6 +143,25 @@ export async function getTask( return { record: value, task: taskView(value) } } +/** + * Cancels the task of `record` (a `working` task from {@link getTask}): + * aborts its run, when it runs in this process, and saves it as + * `cancelled`. Returns the spec 2025-11-25 `Task` for `tasks/cancel`. + */ +export async function cancelTask(record: StoredTask, store: TaskStore) { + // Abort first, so the run does not save a late result over the cancel. + // ponytail: a run on another instance of the server is not aborted. Its + // late result replaces the cancel. Add a store-backed flag if that matters. + running.get(record.taskId)?.abort() + const cancelled: StoredTask = { + ...record, + status: 'cancelled', + lastUpdatedAt: new Date().toISOString(), + } + await store.set(record.taskId, cancelled) + return taskView(cancelled) +} + function taskView(record: StoredTask) { const task = { taskId: record.taskId, @@ -186,6 +217,7 @@ function isStoredTask(value: unknown): value is StoredTask { switch (value.status) { case 'working': case 'failed': + case 'cancelled': return true case 'completed': return 'result' in value diff --git a/packages/ai-mcp/tests/harness-elicit.test.ts b/packages/ai-mcp/tests/harness-elicit.test.ts index 0445e5be2c..e9acd3be6e 100644 --- a/packages/ai-mcp/tests/harness-elicit.test.ts +++ b/packages/ai-mcp/tests/harness-elicit.test.ts @@ -1,4 +1,10 @@ import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { + McpServer, + WebStandardStreamableHTTPServerTransport, +} from '@modelcontextprotocol/server' +import { connectTransport } from '../src/client' import { askForInput, elicitFor, trackCaller } from '../src/harness-plugin' import type { AnyTool } from '@tanstack/ai' import type { PluginSessionApi, RunContext } from '@tanstack/ai-harness' @@ -127,3 +133,68 @@ describe('an MCP elicitation is a live wait of the run that calls', () => { }) }) }) + +describe('a spec 2025 elicitation that the server cancels', () => { + it('gives the request signal to the answer', async () => { + // A server whose tool gives up on its question after 50 ms. + const server = new McpServer({ name: 'trips', version: '1.0.0' }) + server.registerTool( + 'book', + { description: 'Book a trip', inputSchema: z.object({}) }, + async (_args, ctx) => { + await ctx.mcpReq + .elicitInput(form as never, { timeout: 50 }) + .catch(() => undefined) + return { content: [{ type: 'text', text: 'gave up' }] } + }, + ) + const transport = new WebStandardStreamableHTTPServerTransport({ + sessionIdGenerator: () => crypto.randomUUID(), + }) + await server.connect(transport) + let elicitSignal: AbortSignal | undefined + const client = await connectTransport( + { + transport: { + type: 'http', + url: 'http://mcp.test/mcp', + fetch: async (input, init) => + transport.handleRequest(new Request(input, init)), + }, + }, + (_params, signal) => { + elicitSignal = signal + return new Promise((resolve) => { + signal?.addEventListener( + 'abort', + () => resolve({ action: 'cancel' }), + { once: true }, + ) + }) + }, + ) + try { + await client.callTool('book', {}) + await vi.waitFor(() => expect(elicitSignal?.aborted).toBe(true)) + } finally { + await client.close() + await server.close() + } + }) + + it('stops the question when the signal aborts', async () => { + const { session, ask } = askingSession({ city: 'Paris' }) + const calling: Parameters[2] = [] + const onElicit = elicitFor(calling) + const run = runOf('op-agent-1') + const { signal } = new AbortController() + const [tool] = trackCaller( + [toolWith(() => onElicit(form as never, signal))], + session, + calling, + ) + + await tool?.execute?.({}, { context: { run } } as never) + expect(ask).toHaveBeenCalledWith(expect.objectContaining({ signal })) + }) +}) diff --git a/packages/ai-mcp/tests/server/create-server.test.ts b/packages/ai-mcp/tests/server/create-server.test.ts index 3fc6020c1e..b3b341d119 100644 --- a/packages/ai-mcp/tests/server/create-server.test.ts +++ b/packages/ai-mcp/tests/server/create-server.test.ts @@ -369,6 +369,7 @@ describe('createMCPServer', () => { await withClient(server, { era: '2025' }, async (client) => { expect(client.getServerCapabilities()?.tasks).toEqual({ requests: { tools: { call: {} } }, + cancel: {}, }) const listed = await client.listTools() expect(listed.tools[0]?.execution).toEqual({ taskSupport: 'required' }) @@ -407,6 +408,61 @@ describe('createMCPServer', () => { }) }) + it('cancels a running spec 2025 task with tasks/cancel', async () => { + const taskStore = inMemoryTaskStore() + const calls: Array> = [] + let toolSignal: AbortSignal | undefined + const server = createMCPServer({ + name: 'tasks', + version: '1.0.0', + taskStore, + waitUntil(promise) { + calls.push(promise) + }, + tools: [ + toolDefinition({ + name: 'slow', + description: 'Slow work', + execution: 'task', + }).server(async (_args, ctx) => { + toolSignal = ctx?.abortSignal + // Work that runs until the task is cancelled. + await new Promise((resolve) => { + ctx?.abortSignal?.addEventListener('abort', () => resolve(), { + once: true, + }) + }) + return 'late' + }), + ], + }) + + await withClient(server, { era: '2025' }, async (client) => { + const { taskId } = (await taskCall(client, 'slow')).task + await vi.waitFor(() => expect(toolSignal).toBeDefined()) + + const cancelled = await client.request( + { method: 'tasks/cancel', params: { taskId } }, + taskShape, + ) + expect(cancelled).toMatchObject({ taskId, status: 'cancelled' }) + expect(toolSignal?.aborted).toBe(true) + // The late result of the tool does not replace the cancel. + await calls[0] + expect((await taskGet(client, taskId)).status).toBe('cancelled') + // A task in a terminal status cannot be cancelled again. + await expect( + client.request( + { method: 'tasks/cancel', params: { taskId } }, + taskShape, + ), + ).rejects.toThrow() + expect(client.getServerCapabilities()?.tasks).toMatchObject({ + cancel: {}, + }) + }) + }) + it('runs a task tool inline for a spec 2026 request', async () => { // Spec 2026-07-28 has no tasks. The call waits for the tool output. const taskStore = inMemoryTaskStore() From d889d9127ba134e67bb7cd030e3825fce3c52698 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:17 +0100 Subject: [PATCH 06/24] fix(ai-memory): forward the call signal in the hindsight tools --- .changeset/abort-memory-hindsight-tools.md | 5 +++++ .../src/providers/hindsight/index.ts | 10 +++++++--- .../src/providers/hindsight/tools.ts | 12 ++++++++---- .../tests/providers/hindsight.test.ts | 19 +++++++++++++++++++ 4 files changed, 39 insertions(+), 7 deletions(-) create mode 100644 .changeset/abort-memory-hindsight-tools.md diff --git a/.changeset/abort-memory-hindsight-tools.md b/.changeset/abort-memory-hindsight-tools.md new file mode 100644 index 0000000000..4c0e102d82 --- /dev/null +++ b/.changeset/abort-memory-hindsight-tools.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-memory': patch +--- + +The `hindsight_retain`, `hindsight_recall`, and `hindsight_reflect` tools now give the abort signal of the call to the Hindsight client. diff --git a/packages/ai-memory/src/providers/hindsight/index.ts b/packages/ai-memory/src/providers/hindsight/index.ts index a7c2073164..aa3b084840 100644 --- a/packages/ai-memory/src/providers/hindsight/index.ts +++ b/packages/ai-memory/src/providers/hindsight/index.ts @@ -33,14 +33,18 @@ export interface HindsightClientLike { retain: ( bankId: string, text: string, - opts: { context: string; timestamp: Date }, + opts: { context: string; timestamp: Date; signal?: AbortSignal }, ) => Promise recall: ( bankId: string, query: string, - opts: { budget: string }, + opts: { budget: string; signal?: AbortSignal }, ) => Promise - reflect: (bankId: string, query: string) => Promise<{ text?: string }> + reflect: ( + bankId: string, + query: string, + opts?: { signal?: AbortSignal }, + ) => Promise<{ text?: string }> listMemories: ( bankId: string, opts: { limit: number }, diff --git a/packages/ai-memory/src/providers/hindsight/tools.ts b/packages/ai-memory/src/providers/hindsight/tools.ts index bac1179e42..4592787374 100644 --- a/packages/ai-memory/src/providers/hindsight/tools.ts +++ b/packages/ai-memory/src/providers/hindsight/tools.ts @@ -40,7 +40,7 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { required: ['content'], additionalProperties: false, }, - async execute(args) { + async execute(args, ctx) { const content = stringField(args, 'content') const start = Date.now() try { @@ -48,6 +48,7 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { const data = await client.retain(deps.bankId, content, { context: 'chat:tool', timestamp: new Date(), + signal: ctx?.abortSignal, }) deps.onToolRetain?.({ ok: true, @@ -78,13 +79,14 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { required: ['query'], additionalProperties: false, }, - async execute(args) { + async execute(args, ctx) { const query = stringField(args, 'query') const start = Date.now() try { const { client, recallToPrompt } = await deps.getRuntime() const data = await client.recall(deps.bankId, query, { budget: deps.budget, + signal: ctx?.abortSignal, }) const systemPrompt = recallToPrompt(data) const fragments: Array = (data.results ?? []).map( @@ -123,11 +125,13 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { required: ['query'], additionalProperties: false, }, - async execute(args) { + async execute(args, ctx) { const query = stringField(args, 'query') try { const { client } = await deps.getRuntime() - const data = await client.reflect(deps.bankId, query) + const data = await client.reflect(deps.bankId, query, { + signal: ctx?.abortSignal, + }) return data.text ?? '(no reflection)' } catch (err) { return `(reflection failed: ${err instanceof Error ? err.message : String(err)})` diff --git a/packages/ai-memory/tests/providers/hindsight.test.ts b/packages/ai-memory/tests/providers/hindsight.test.ts index 77a9d27355..ed975a147d 100644 --- a/packages/ai-memory/tests/providers/hindsight.test.ts +++ b/packages/ai-memory/tests/providers/hindsight.test.ts @@ -124,4 +124,23 @@ describe('makeHindsightTools', () => { expect(receipts[0]?.ok).toBe(false) expect(receipts[0]?.error).toContain('server down') }) + + it('gives the abort signal of the call to the client', async () => { + const retain = vi.fn(async () => ({ stored: true })) + const recall = vi.fn(async () => ({ results: [] })) + const reflect = vi.fn(async () => ({ text: 'reflection' })) + const tools = makeHindsightTools({ + getRuntime: async () => fakeRuntime({ retain, recall, reflect }), + bankId: 'u1__s1', + budget: 'mid', + }) + const { signal } = new AbortController() + const context = { abortSignal: signal, emitCustomEvent: () => undefined } + await tools[0]?.execute?.({ content: 'x' }, context) + await tools[1]?.execute?.({ query: 'q' }, context) + await tools[2]?.execute?.({ query: 'q' }, context) + for (const method of [retain, recall, reflect]) { + expect(method.mock.lastCall?.at(-1)).toMatchObject({ signal }) + } + }) }) From d9b3c40d309d43cb6affa7e54af95307b13bce28 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:19 +0100 Subject: [PATCH 07/24] fix(ai-sandbox-daytona,ai-sandbox-vercel,ai-sandbox-upstash-box): clean up late sandboxes and kill aborted exec --- .../abort-sandbox-daytona-late-create.md | 5 ++++ .changeset/abort-sandbox-upstash-box-exec.md | 5 ++++ .../abort-sandbox-vercel-late-create.md | 5 ++++ packages/ai-sandbox-daytona/src/provider.ts | 11 ++++++-- .../ai-sandbox-daytona/tests/provider.test.ts | 27 +++++++++++++++++++ packages/ai-sandbox-upstash-box/src/handle.ts | 16 +++++++++++ .../tests/handle.test.ts | 19 +++++++++++++ packages/ai-sandbox-vercel/src/provider.ts | 7 +++++ .../ai-sandbox-vercel/tests/vercel.test.ts | 26 +++++++++++++++++- 9 files changed, 118 insertions(+), 3 deletions(-) create mode 100644 .changeset/abort-sandbox-daytona-late-create.md create mode 100644 .changeset/abort-sandbox-upstash-box-exec.md create mode 100644 .changeset/abort-sandbox-vercel-late-create.md diff --git a/.changeset/abort-sandbox-daytona-late-create.md b/.changeset/abort-sandbox-daytona-late-create.md new file mode 100644 index 0000000000..f74229f1eb --- /dev/null +++ b/.changeset/abort-sandbox-daytona-late-create.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-daytona': patch +--- + +`create` and `restoreSnapshot` now delete a sandbox that the SDK makes after the signal aborted. Before, that sandbox kept running with no owner. diff --git a/.changeset/abort-sandbox-upstash-box-exec.md b/.changeset/abort-sandbox-upstash-box-exec.md new file mode 100644 index 0000000000..45656359e5 --- /dev/null +++ b/.changeset/abort-sandbox-upstash-box-exec.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-upstash-box': patch +--- + +A blocking `exec` with a signal now runs as an `exec.session`, so an abort kills the command in the box. Before, the signal was only checked before the command started. diff --git a/.changeset/abort-sandbox-vercel-late-create.md b/.changeset/abort-sandbox-vercel-late-create.md new file mode 100644 index 0000000000..38062e76f5 --- /dev/null +++ b/.changeset/abort-sandbox-vercel-late-create.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-vercel': patch +--- + +`create` now stops a sandbox that comes after the signal aborted. Before, that sandbox kept running with no owner. diff --git a/packages/ai-sandbox-daytona/src/provider.ts b/packages/ai-sandbox-daytona/src/provider.ts index 1c13ed3893..5f878a68dc 100644 --- a/packages/ai-sandbox-daytona/src/provider.ts +++ b/packages/ai-sandbox-daytona/src/provider.ts @@ -97,7 +97,14 @@ class DaytonaProvider implements SandboxProvider { private async wrapCreated( sandbox: Awaited>, applyEnvSet: boolean, + signal: AbortSignal | undefined, ): Promise { + // The SDK create takes no signal. A sandbox that comes after the abort has + // no owner, so delete it instead of leaving it billed. + if (signal?.aborted) { + await this.daytona.delete(sandbox).catch(() => undefined) + signal.throwIfAborted() + } await sandbox.process.executeCommand(`mkdir -p ${shQuote(this.workdir)}`) return new DaytonaHandle({ sandbox, workdir: this.workdir, applyEnvSet }) } @@ -164,7 +171,7 @@ class DaytonaProvider implements SandboxProvider { ) // Workspace secrets live in Daytona OS env as placeholders. Do not overlay // plaintext via env.set (bootstrap and resume also call env.set). - return this.wrapCreated(sandbox, secrets === undefined) + return this.wrapCreated(sandbox, secrets === undefined, input.signal) } async restoreSnapshot(input: SandboxRestoreInput): Promise { @@ -176,7 +183,7 @@ class DaytonaProvider implements SandboxProvider { ...(secrets !== undefined ? { secrets } : {}), }), ) - return this.wrapCreated(sandbox, secrets === undefined) + return this.wrapCreated(sandbox, secrets === undefined, input.signal) } async resume(input: SandboxResumeInput): Promise { diff --git a/packages/ai-sandbox-daytona/tests/provider.test.ts b/packages/ai-sandbox-daytona/tests/provider.test.ts index bea4b897f4..0273a4fc77 100644 --- a/packages/ai-sandbox-daytona/tests/provider.test.ts +++ b/packages/ai-sandbox-daytona/tests/provider.test.ts @@ -413,3 +413,30 @@ describe('daytonaSandbox create', () => { expect(provider.capabilities().networkPolicy).toBe(true) }) }) + +describe('daytonaSandbox abort', () => { + it.each(['create', 'restoreSnapshot'] as const)( + '%s deletes a sandbox that comes after the abort', + async (method) => { + const abort = new AbortController() + const late = fakeRemoteSandbox('sbx-late') + create.mockImplementation(async () => { + abort.abort() + return late + }) + del.mockResolvedValue(undefined) + const provider = daytonaSandbox({ apiKey: 'test-key' }) + + const made = + method === 'create' + ? provider.create({ signal: abort.signal }) + : provider.restoreSnapshot!({ + snapshotId: 'snap-1', + signal: abort.signal, + }) + + await expect(made).rejects.toThrow() + expect(del).toHaveBeenCalledWith(late) + }, + ) +}) diff --git a/packages/ai-sandbox-upstash-box/src/handle.ts b/packages/ai-sandbox-upstash-box/src/handle.ts index 230b340b8d..a8684949a0 100644 --- a/packages/ai-sandbox-upstash-box/src/handle.ts +++ b/packages/ai-sandbox-upstash-box/src/handle.ts @@ -326,6 +326,22 @@ export class UpstashBoxHandle implements SandboxHandle { opts?: ProcessOptions, ): Promise { opts?.signal?.throwIfAborted() + // `exec.command` cannot be stopped. A session can: its abort kills the + // command in the box. + if (opts?.signal) { + const proc = await this.spawnProcess(command, opts) + const read = async (stream: AsyncIterable) => { + let text = '' + for await (const chunk of stream) text += chunk + return text + } + const [stdout, stderr, exitCode] = await Promise.all([ + read(proc.stdout), + read(proc.stderr), + proc.wait(), + ]) + return { stdout, stderr, exitCode } + } const run = await this.box.exec.command(this.wrap(command, opts)) return { stdout: run.stdout, diff --git a/packages/ai-sandbox-upstash-box/tests/handle.test.ts b/packages/ai-sandbox-upstash-box/tests/handle.test.ts index cf0828cdd0..c6c37e71cc 100644 --- a/packages/ai-sandbox-upstash-box/tests/handle.test.ts +++ b/packages/ai-sandbox-upstash-box/tests/handle.test.ts @@ -161,6 +161,25 @@ describe('UpstashBoxHandle', () => { expect(commands[0]).toBe("cd '/workspace/home' && echo hello") }) + it('kills an exec in the box when its signal aborts', async () => { + const fake = fakeSession() + const { box, sessionOptions } = fakeBox({ session: fake }) + const handle = new UpstashBoxHandle({ box }) + const abort = new AbortController() + + const running = handle.process.exec('sleep 600', { signal: abort.signal }) + await vi.waitFor(() => expect(sessionOptions).toHaveLength(1)) + fake.session.emitStdout('partial') + abort.abort() + + expect(fake.session.kill).toHaveBeenCalledOnce() + await expect(running).resolves.toEqual({ + stdout: 'partial', + stderr: '', + exitCode: 143, + }) + }) + it('maps /workspace cwd and applies env exports in order', async () => { const { box, commands } = fakeBox() const handle = new UpstashBoxHandle({ box }) diff --git a/packages/ai-sandbox-vercel/src/provider.ts b/packages/ai-sandbox-vercel/src/provider.ts index fb4be83951..e1491f6800 100644 --- a/packages/ai-sandbox-vercel/src/provider.ts +++ b/packages/ai-sandbox-vercel/src/provider.ts @@ -129,6 +129,13 @@ class VercelProvider implements SandboxProvider { : {}), ...(input.env ? { env: input.env } : {}), }) + // ponytail: no SDK signal here. It cancels only the request, so a sandbox + // that the server made has no owner. Let the create end, then stop a + // sandbox that comes after the abort instead of leaving it billed. + if (input.signal?.aborted) { + await sandbox.stop().catch(() => undefined) + input.signal.throwIfAborted() + } // Ensure the workspace dir exists via the native (cwd-independent) mkDir — // running a command with a not-yet-existing `cwd` would fail, so we must not // route this through the handle (which runs every command in `workdir`). diff --git a/packages/ai-sandbox-vercel/tests/vercel.test.ts b/packages/ai-sandbox-vercel/tests/vercel.test.ts index b2a533cd34..2f76fd9609 100644 --- a/packages/ai-sandbox-vercel/tests/vercel.test.ts +++ b/packages/ai-sandbox-vercel/tests/vercel.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from 'vitest' -import { APIError } from '@vercel/sandbox' +import { APIError, Sandbox as VercelSandbox } from '@vercel/sandbox' import { journalReadStrategy } from '@tanstack/ai-sandbox' import { vercelSandbox } from '../src/index' import { isDirAlreadyExistsError } from '../src/provider' @@ -196,3 +196,27 @@ describe.skipIf(!hasCreds)('vercel provider (gated on VERCEL_TOKEN)', () => { } }, 180_000) }) + +describe('vercelSandbox create abort', () => { + it('stops a sandbox that comes after the abort', async () => { + const abort = new AbortController() + const stop = vi.fn(() => Promise.resolve()) + const create = vi + .spyOn(VercelSandbox, 'create') + .mockImplementation(async () => { + abort.abort() + return { name: 'sbx-late', mkDir: async () => {}, stop } as never + }) + try { + const provider = vercelSandbox({ + token: 't', + teamId: 't', + projectId: 'p', + }) + await expect(provider.create({ signal: abort.signal })).rejects.toThrow() + expect(stop).toHaveBeenCalledOnce() + } finally { + create.mockRestore() + } + }) +}) From 9766372a8e47e9ad7e964dfc7aee1ed2613583b7 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:20 +0100 Subject: [PATCH 08/24] fix(ai-code-mode): end the program and dispose the isolate when the run aborts --- .changeset/abort-code-mode-run.md | 5 ++ .../ai-code-mode/src/create-code-mode-tool.ts | 12 +++++ .../tests/create-code-mode-tool.test.ts | 51 +++++++++++++++++++ 3 files changed, 68 insertions(+) create mode 100644 .changeset/abort-code-mode-run.md diff --git a/.changeset/abort-code-mode-run.md b/.changeset/abort-code-mode-run.md new file mode 100644 index 0000000000..c00508fe42 --- /dev/null +++ b/.changeset/abort-code-mode-run.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-code-mode': patch +--- + +An abort of the run now ends the `execute_typescript` program: a call that waits for an approval stops, and the isolate is disposed. 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..613407f6d3 100644 --- a/packages/ai-code-mode/src/create-code-mode-tool.ts +++ b/packages/ai-code-mode/src/create-code-mode-tool.ts @@ -216,6 +216,14 @@ export function createCodeModeTool( let isolateContext: IsolateContext | null = null // Aborts when this program ends, so a host call that still waits can stop. const programEnd = new AbortController() + // An abort of the run ends the program: a call that waits (for example + // for an approval) stops, and the isolate goes, so its code stops too. + const runSignal = toolContext?.abortSignal + const onRunAbort = () => { + programEnd.abort() + void isolateContext?.dispose() + } + runSignal?.addEventListener('abort', onRunAbort, { once: true }) // Emit execution started event immediately emitCustomEvent('code_mode:execution_started', { @@ -294,6 +302,9 @@ export function createCodeModeTool( ) } + // An abort that came while the context was made. + if (runSignal?.aborted) onRunAbort() + // Step 5: Execute the code in the sandbox const executionResult = await isolateContext.execute(strippedCode) @@ -364,6 +375,7 @@ export function createCodeModeTool( 'unhandled', ) } finally { + runSignal?.removeEventListener('abort', onRunAbort) // The program has ended (success, error, or timeout). programEnd.abort() // Always clean up the sandbox context diff --git a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts index e482327eaf..2ac9c5f658 100644 --- a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts +++ b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts @@ -681,3 +681,54 @@ describe('createCodeModeTool debug logging', () => { } }) }) + +describe('createCodeModeTool run abort', () => { + it('ends the program and disposes the isolate when the run aborts', async () => { + let programSignal: AbortSignal | undefined + const aroundCalls = vi.fn>( + (_context, signal) => { + programSignal = signal + return (_call, run) => run() + }, + ) + let isExecuting = false + let release = () => {} + const disposed = new Promise((resolve) => { + release = resolve + }) + const dispose = vi.fn(async () => release()) + // An isolate that runs until it is disposed. + const driver: IsolateDriver = { + createContext: async () => ({ + execute: async () => { + isExecuting = true + await disposed + return { + success: false, + error: { name: 'DisposedError', message: 'Context was disposed' }, + logs: [], + } + }, + dispose, + }), + } + const tool = createCodeModeTool({ + driver, + tools: [createMockTool('fetchWeather')], + aroundCalls, + }) + const abort = new AbortController() + + const running = tool.execute!( + { typescriptCode: 'return 1' }, + { abortSignal: abort.signal, emitCustomEvent: vi.fn() }, + ) + await vi.waitFor(() => expect(isExecuting).toBe(true)) + abort.abort() + + // A call that waits for an approval stops with the program. + expect(programSignal?.aborted).toBe(true) + expect(dispose).toHaveBeenCalled() + expect((await running).success).toBe(false) + }) +}) From 803cf31005e01e486cde1a5a1f4c3e2154bddf91 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:20 +0100 Subject: [PATCH 09/24] fix(ai-computer-use): skip queued actions and end waits when the call aborts --- .changeset/abort-computer-use-actions.md | 5 ++ packages/ai-computer-use/src/actions.ts | 40 +++++++++++---- .../src/create-computer-use.ts | 16 ++++-- .../tests/create-computer-use.test.ts | 49 +++++++++++++++++++ testing/e2e/tests/computer-use.spec.ts | 49 +++++++++++++++++++ 5 files changed, 145 insertions(+), 14 deletions(-) create mode 100644 .changeset/abort-computer-use-actions.md diff --git a/.changeset/abort-computer-use-actions.md b/.changeset/abort-computer-use-actions.md new file mode 100644 index 0000000000..058afc2569 --- /dev/null +++ b/.changeset/abort-computer-use-actions.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-computer-use': patch +--- + +An aborted computer action no longer runs when it waits in the queue. A `wait`, `hold_key`, or the next action of a batch now stops when the call aborts. diff --git a/packages/ai-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index a96b690195..71a678f556 100644 --- a/packages/ai-computer-use/src/actions.ts +++ b/packages/ai-computer-use/src/actions.ts @@ -19,8 +19,20 @@ function isRecord(value: unknown): value is Input { return typeof value === 'object' && value !== null && !Array.isArray(value) } -function sleep(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)) +/** Waits `ms`. Rejects with the reason of `signal` when it aborts first. */ +function sleep(ms: number, signal?: AbortSignal) { + 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 }) + }) } function invalid(key: string, expected: string) { @@ -168,6 +180,7 @@ async function runAnthropicAction( driver: ComputerDriver, action: string, input: Input, + signal: AbortSignal | undefined, ) { const click = ANTHROPIC_CLICKS[action] if (click) { @@ -226,11 +239,11 @@ async function runAnthropicAction( case 'hold_key': { const keys = parseKeyCombo(readString(input, 'text')) const ms = readSeconds(input, 'duration') * 1000 - await withKeys(driver, keys, () => sleep(ms)) + await withKeys(driver, keys, () => sleep(ms, signal)) return undefined } case 'wait': - await sleep(readSeconds(input, 'duration') * 1000) + await sleep(readSeconds(input, 'duration') * 1000, signal) return undefined case 'cursor_position': { if (!driver.cursorPosition) { @@ -291,7 +304,11 @@ function readPath(input: Input) { }) } -async function runOpenAIAction(driver: ComputerDriver, input: unknown) { +async function runOpenAIAction( + driver: ComputerDriver, + input: unknown, + signal: AbortSignal | undefined, +) { if (!isRecord(input)) throw invalid('action', 'an object') const type = readString(input, 'type') const modifiers = openaiModifiers(input) @@ -357,7 +374,7 @@ async function runOpenAIAction(driver: ComputerDriver, input: unknown) { await driver.type(readString(input, 'text')) return case 'wait': - await sleep(OPENAI_WAIT_MS) + await sleep(OPENAI_WAIT_MS, signal) return default: throw new Error(`Unknown computer action "${type}".`) @@ -385,18 +402,23 @@ async function runOpenAIAction(driver: ComputerDriver, input: unknown) { export async function runComputerAction( driver: ComputerDriver, input: unknown, + signal?: AbortSignal, ) { if (!isRecord(input)) throw invalid('input', 'an object') const { action, actions } = input if (typeof action === 'string') { - return runAnthropicAction(driver, action, input) + return runAnthropicAction(driver, action, input, signal) } if (Array.isArray(actions)) { - for (const item of actions) await runOpenAIAction(driver, item) + for (const item of actions) { + // An abort stops the batch before its next action. + signal?.throwIfAborted() + await runOpenAIAction(driver, item, signal) + } return undefined } if (isRecord(action)) { - await runOpenAIAction(driver, action) + await runOpenAIAction(driver, action, signal) return undefined } throw invalid('action', 'a string, an object, or an "actions" array') diff --git a/packages/ai-computer-use/src/create-computer-use.ts b/packages/ai-computer-use/src/create-computer-use.ts index ba3130c6c6..7c8e2b32da 100644 --- a/packages/ai-computer-use/src/create-computer-use.ts +++ b/packages/ai-computer-use/src/create-computer-use.ts @@ -257,14 +257,16 @@ export function createComputerUse({ /** Run `work` in the queue of the thread, with the failed-turn rule. */ function runInThread( - threadId: string | undefined, + context: ToolExecutionContext | undefined, work: (thread: ThreadState) => Promise, ) { - const thread = threadOf(threadId ?? SHARED_THREAD) + const thread = threadOf(context?.threadId ?? SHARED_THREAD) clearTimeout(thread.timer) thread.pending += 1 const previous = sharesOneMachine ? sharedQueue : thread.queue const run = previous.then(async () => { + // A call that aborted while it waited in the queue does not run. + context?.abortSignal?.throwIfAborted() if (thread.turnFailed) { throw new Error( 'Not executed: an earlier computer action in this turn failed.', @@ -316,7 +318,7 @@ export function createComputerUse({ description: 'Use a computer with a mouse and a keyboard.', metadata, execute(input: unknown, context?: ToolExecutionContext) { - return runInThread(context?.threadId, async (thread) => { + return runInThread(context, async (thread) => { const checked = readCheckedInput(input) // Fail closed: checks that do not parse also need an approval. const hasChecks = checked === null || checked.checks.length > 0 @@ -329,7 +331,11 @@ export function createComputerUse({ ) } const machine = await machineOf(thread) - const result = await runComputerAction(machine, input) + const result = await runComputerAction( + machine, + input, + context?.abortSignal, + ) if (result === undefined) { await sleep(screenshotDelay) return imageResult(await machine.screenshot()) @@ -340,7 +346,7 @@ export function createComputerUse({ }) const navigate = navigateDefinition.server(({ url }, context) => - runInThread(context?.threadId, async (thread) => { + runInThread(context, async (thread) => { if (!isWebUrl(url)) { throw new Error( 'Refused: navigate needs a full URL that starts with https:// or http://, like https://tanstack.com.', 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..5f4891b726 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,52 @@ describe('createComputerUse navigate', () => { ) }) }) + +describe('the abort signal of a computer action', () => { + /** Run the computer tool as `chat()` does, with `signal`. */ + function executeWith( + computer: Computer, + input: unknown, + toolCallId: string, + signal: AbortSignal, + ) { + const [computerTool] = computer.tools + if (!computerTool.execute) throw new Error('The tool has no execute.') + return computerTool.execute(input, { + toolCallId, + threadId: THREAD, + abortSignal: signal, + emitCustomEvent() {}, + }) + } + + it('does not run a queued action after its call aborts', async () => { + const { gate, open } = createGate() + const computer = setup({ clickGate: gate }) + const first = computer.execute(CLICK, 'call_1') + const abort = new AbortController() + const queued = executeWith(computer, TYPE, 'call_2', abort.signal) + abort.abort() + open() + + await first + await expect(queued).rejects.toThrow() + expect(computer.calls).not.toContain('type') + }) + + it('ends a wait when its call aborts', async () => { + const computer = setup() + const abort = new AbortController() + const waiting = executeWith( + computer, + { action: { type: 'wait' } }, + 'call_1', + abort.signal, + ) + setTimeout(() => abort.abort(), 10) + const started = Date.now() + + await expect(waiting).rejects.toThrow() + expect(Date.now() - started).toBeLessThan(1000) + }) +}) diff --git a/testing/e2e/tests/computer-use.spec.ts b/testing/e2e/tests/computer-use.spec.ts index 5d4cba64ba..17b6f2032f 100644 --- a/testing/e2e/tests/computer-use.spec.ts +++ b/testing/e2e/tests/computer-use.spec.ts @@ -358,6 +358,55 @@ test.describe('computer use with a Playwright page', () => { }) }) + test('a stopped run ends a wait action at once and runs no later action', async ({ + page, + }) => { + const { button } = await openForm(page) + const computerUse = computerUseFor(page) + const model = mockModel([ + openaiStream('resp_cu_stop', [ + { + type: 'computer_call', + id: 'cu_stop', + call_id: 'call_cu_stop', + status: 'completed', + // The OpenAI `wait` takes 2 seconds. The click comes after it. + actions: [ + { type: 'wait' }, + { type: 'click', button: 'left', ...button }, + ], + pending_safety_checks: [], + }, + ]), + openaiStream('resp_cu_stop_2', [OPENAI_ANSWER]), + ]) + const abort = new AbortController() + const started = Date.now() + + // Stop the run when the model asked for the actions. + const stream = chat({ + adapter: openaiAdapter(model), + messages: [USER_MESSAGE], + threadId: 'thread-stop', + tools: computerUse.tools, + middleware: [computerUse.middleware], + interrupts: [computerUse.interrupt], + abortController: abort, + }) + try { + for await (const chunk of stream) { + if (chunk.type === EventType.TOOL_CALL_END) { + setTimeout(() => abort.abort(), 50) + } + } + } catch { + // An aborted run can end with an error. The checks below are the test. + } + + expect(Date.now() - started).toBeLessThan(1500) + await expect(page.locator('#status')).toHaveText('idle') + }) + test('anthropic computer toolset runs each call on the page', async ({ page, }) => { From 1980aa980c9dab83ddf7394741a9b46cb0660229 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:48:22 +0100 Subject: [PATCH 10/24] docs: describe tasks/cancel, code mode abort, and harnessText early stop --- docs/code-mode/code-mode.md | 2 +- docs/config.json | 6 +++--- docs/harness/connect.md | 2 +- docs/mcp/server-tasks.md | 28 ++++++++++++++++++++++++++-- 4 files changed, 31 insertions(+), 7 deletions(-) diff --git a/docs/code-mode/code-mode.md b/docs/code-mode/code-mode.md index 44427cd081..f323bdb0aa 100644 --- a/docs/code-mode/code-mode.md +++ b/docs/code-mode/code-mode.md @@ -141,7 +141,7 @@ Code Mode rejects tools that need approval. A tool marked `needsApproval: true` ### Stop tool calls when the run stops -When a user stops a chat run, a tool call that the sandbox started keeps going, unless the tool listens for the abort. Each `external_*` call gets the `abortSignal` of the chat run in its second argument. Pass it to the slow work: +When a user stops a chat run, Code Mode disposes the isolate, so the program stops. A tool call that the program started keeps going, unless the tool listens for the abort. Each `external_*` call gets the `abortSignal` of the chat run in its second argument. Pass it to the slow work: ```typescript import { toolDefinition } from "@tanstack/ai"; diff --git a/docs/config.json b/docs/config.json index d1855bbf8b..5bc401a7d7 100644 --- a/docs/config.json +++ b/docs/config.json @@ -690,7 +690,7 @@ "label": "MCP Server Tasks", "to": "mcp/server-tasks", "addedAt": "2026-09-22", - "updatedAt": "2026-10-22" + "updatedAt": "2026-11-02" }, { "label": "MCP Server Auth", @@ -1007,7 +1007,7 @@ "label": "Code Mode", "to": "code-mode/code-mode", "addedAt": "2026-04-15", - "updatedAt": "2026-10-29" + "updatedAt": "2026-11-02" }, { "label": "Showing Code Mode in the UI", @@ -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", diff --git a/docs/harness/connect.md b/docs/harness/connect.md index 23362c964b..6e61421531 100644 --- a/docs/harness/connect.md +++ b/docs/harness/connect.md @@ -498,7 +498,7 @@ Each outer thread gets its own inner session, so the harness keeps its own histo - Who runs the inner turn: when the outer `chat()` runs inside a harness turn, the inner session opens as the user of that turn. Else it opens as `user`, for example `harnessText(assistant, { host, user: { id: 'ada' } })`. - Access: the inner thread is `/inner/`. When `host` has `access` and also runs the outer thread, the inner session is checked through the outer thread, so every user who can read a shared outer thread can use its inner session. From app code, `host.open`, `host.submit`, `host.fork`, and `host.transfer` refuse any thread id with `/inner/`. - Usage: the outer run gets the token usage of every inner model call, so the usage totals of the outer session count them. -- Cancel: an abort of the outer call cancels the inner turn. For a remote harness, `harnessText({ url, token })` also tells the remote host to cancel the turn. +- Cancel: an abort of the outer call cancels the inner turn. A caller that stops reading the stream early cancels it too. For a remote harness, `harnessText({ url, token })` also tells the remote host to cancel the turn. If a plugin picks the model of the harness, the harness has no adapter to read its input kinds from. Pass them with `harnessText(assistant, { host, inputModalities: ['text', 'image'] })`. diff --git a/docs/mcp/server-tasks.md b/docs/mcp/server-tasks.md index 3e3dca12e1..9d6fc88f1a 100644 --- a/docs/mcp/server-tasks.md +++ b/docs/mcp/server-tasks.md @@ -66,7 +66,7 @@ The host calls `build_report` and gets a task handle at once. After 5 seconds, t 2. The handler returns a task handle before the work ends. 3. The tool function continues in this process. There is no queue. 4. When the tool function ends, the task store saves the value. -5. The host polls until the status is `completed` or `failed`. +5. The host polls until the status is `completed`, `failed`, or `cancelled`. The default task store keeps each task in memory for this process. A TanStack AI host polls until the task ends. See [MCP Server Tools](../tools/mcp). @@ -76,7 +76,31 @@ The task keeps running after the call returns the task handle. As a result, the - `ctx.context.requestInput` throws an error. Ask for input in a tool without `execution: 'task'`. - `ctx.context.sample` calls the `sample` option of `createMCPServer`. If you do not pass `sample`, it throws an error. -- `ctx.abortSignal` does not abort when the call request ends. +- `ctx.abortSignal` does not abort when the call request ends. It aborts when the host cancels the task. See [Stop a cancelled task](#stop-a-cancelled-task). + +## Stop a cancelled task + +A host sends `tasks/cancel` to stop a task. The task status becomes `cancelled`, and `ctx.abortSignal` of the tool aborts. Give that signal to the slow work, so the work stops too: + +```ts +import { toolDefinition } from '@tanstack/ai' + +const buildReport = toolDefinition({ + name: 'build_report', + description: 'Build a long report', + execution: 'task', +}).server(async (_args, ctx) => { + const response = await fetch('https://reports.example.com/build', { + method: 'POST', + signal: ctx?.abortSignal, + }) + return { text: await response.text() } +}) +``` + +- A task that ended (`completed`, `failed`, or `cancelled`) cannot be cancelled. `tasks/cancel` then returns an error. +- A cancelled task keeps the status `cancelled`. A result that the tool returns after the cancel is not saved. +- Only the process that runs the task can abort its signal. With a shared store on more than one instance, a cancel on another instance saves `cancelled`, but the tool runs on. Its late result then replaces `cancelled`. ## Run tasks on a Worker From ac4f19367fbd06191982aeb698644df8788d12e1 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:56:10 +0100 Subject: [PATCH 11/24] revert(ai-memory): move the hindsight signal change to the feat PR It adds an optional signal to the exported HindsightClientLike type, which is new public API. It belongs in the feat PR, not in this fix PR. --- .changeset/abort-memory-hindsight-tools.md | 5 ----- .../src/providers/hindsight/index.ts | 10 +++------- .../src/providers/hindsight/tools.ts | 12 ++++-------- .../tests/providers/hindsight.test.ts | 19 ------------------- 4 files changed, 7 insertions(+), 39 deletions(-) delete mode 100644 .changeset/abort-memory-hindsight-tools.md diff --git a/.changeset/abort-memory-hindsight-tools.md b/.changeset/abort-memory-hindsight-tools.md deleted file mode 100644 index 4c0e102d82..0000000000 --- a/.changeset/abort-memory-hindsight-tools.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -'@tanstack/ai-memory': patch ---- - -The `hindsight_retain`, `hindsight_recall`, and `hindsight_reflect` tools now give the abort signal of the call to the Hindsight client. diff --git a/packages/ai-memory/src/providers/hindsight/index.ts b/packages/ai-memory/src/providers/hindsight/index.ts index aa3b084840..a7c2073164 100644 --- a/packages/ai-memory/src/providers/hindsight/index.ts +++ b/packages/ai-memory/src/providers/hindsight/index.ts @@ -33,18 +33,14 @@ export interface HindsightClientLike { retain: ( bankId: string, text: string, - opts: { context: string; timestamp: Date; signal?: AbortSignal }, + opts: { context: string; timestamp: Date }, ) => Promise recall: ( bankId: string, query: string, - opts: { budget: string; signal?: AbortSignal }, + opts: { budget: string }, ) => Promise - reflect: ( - bankId: string, - query: string, - opts?: { signal?: AbortSignal }, - ) => Promise<{ text?: string }> + reflect: (bankId: string, query: string) => Promise<{ text?: string }> listMemories: ( bankId: string, opts: { limit: number }, diff --git a/packages/ai-memory/src/providers/hindsight/tools.ts b/packages/ai-memory/src/providers/hindsight/tools.ts index 4592787374..bac1179e42 100644 --- a/packages/ai-memory/src/providers/hindsight/tools.ts +++ b/packages/ai-memory/src/providers/hindsight/tools.ts @@ -40,7 +40,7 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { required: ['content'], additionalProperties: false, }, - async execute(args, ctx) { + async execute(args) { const content = stringField(args, 'content') const start = Date.now() try { @@ -48,7 +48,6 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { const data = await client.retain(deps.bankId, content, { context: 'chat:tool', timestamp: new Date(), - signal: ctx?.abortSignal, }) deps.onToolRetain?.({ ok: true, @@ -79,14 +78,13 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { required: ['query'], additionalProperties: false, }, - async execute(args, ctx) { + async execute(args) { const query = stringField(args, 'query') const start = Date.now() try { const { client, recallToPrompt } = await deps.getRuntime() const data = await client.recall(deps.bankId, query, { budget: deps.budget, - signal: ctx?.abortSignal, }) const systemPrompt = recallToPrompt(data) const fragments: Array = (data.results ?? []).map( @@ -125,13 +123,11 @@ export function makeHindsightTools(deps: HindsightToolDeps): Array { required: ['query'], additionalProperties: false, }, - async execute(args, ctx) { + async execute(args) { const query = stringField(args, 'query') try { const { client } = await deps.getRuntime() - const data = await client.reflect(deps.bankId, query, { - signal: ctx?.abortSignal, - }) + const data = await client.reflect(deps.bankId, query) return data.text ?? '(no reflection)' } catch (err) { return `(reflection failed: ${err instanceof Error ? err.message : String(err)})` diff --git a/packages/ai-memory/tests/providers/hindsight.test.ts b/packages/ai-memory/tests/providers/hindsight.test.ts index ed975a147d..77a9d27355 100644 --- a/packages/ai-memory/tests/providers/hindsight.test.ts +++ b/packages/ai-memory/tests/providers/hindsight.test.ts @@ -124,23 +124,4 @@ describe('makeHindsightTools', () => { expect(receipts[0]?.ok).toBe(false) expect(receipts[0]?.error).toContain('server down') }) - - it('gives the abort signal of the call to the client', async () => { - const retain = vi.fn(async () => ({ stored: true })) - const recall = vi.fn(async () => ({ results: [] })) - const reflect = vi.fn(async () => ({ text: 'reflection' })) - const tools = makeHindsightTools({ - getRuntime: async () => fakeRuntime({ retain, recall, reflect }), - bankId: 'u1__s1', - budget: 'mid', - }) - const { signal } = new AbortController() - const context = { abortSignal: signal, emitCustomEvent: () => undefined } - await tools[0]?.execute?.({ content: 'x' }, context) - await tools[1]?.execute?.({ query: 'q' }, context) - await tools[2]?.execute?.({ query: 'q' }, context) - for (const method of [retain, recall, reflect]) { - expect(method.mock.lastCall?.at(-1)).toMatchObject({ signal }) - } - }) }) From 8d57271f9e9da457c1062ca413149363cee67e81 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 15:56:11 +0100 Subject: [PATCH 12/24] revert(ai-sandbox-upstash-box): drop the session-based exec on abort With a signal, exec ran through a session, which caps each stream at 8 MiB. Inside chat() every tool call has a signal, so every exec would get that cap. exec.command cannot be stopped mid-run, so this row needs a different approach (exec.stream with Run.cancel) checked on a real box. --- .changeset/abort-sandbox-upstash-box-exec.md | 5 ----- packages/ai-sandbox-upstash-box/src/handle.ts | 16 ---------------- .../tests/handle.test.ts | 19 ------------------- 3 files changed, 40 deletions(-) delete mode 100644 .changeset/abort-sandbox-upstash-box-exec.md diff --git a/.changeset/abort-sandbox-upstash-box-exec.md b/.changeset/abort-sandbox-upstash-box-exec.md deleted file mode 100644 index 45656359e5..0000000000 --- a/.changeset/abort-sandbox-upstash-box-exec.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -'@tanstack/ai-sandbox-upstash-box': patch ---- - -A blocking `exec` with a signal now runs as an `exec.session`, so an abort kills the command in the box. Before, the signal was only checked before the command started. diff --git a/packages/ai-sandbox-upstash-box/src/handle.ts b/packages/ai-sandbox-upstash-box/src/handle.ts index a8684949a0..230b340b8d 100644 --- a/packages/ai-sandbox-upstash-box/src/handle.ts +++ b/packages/ai-sandbox-upstash-box/src/handle.ts @@ -326,22 +326,6 @@ export class UpstashBoxHandle implements SandboxHandle { opts?: ProcessOptions, ): Promise { opts?.signal?.throwIfAborted() - // `exec.command` cannot be stopped. A session can: its abort kills the - // command in the box. - if (opts?.signal) { - const proc = await this.spawnProcess(command, opts) - const read = async (stream: AsyncIterable) => { - let text = '' - for await (const chunk of stream) text += chunk - return text - } - const [stdout, stderr, exitCode] = await Promise.all([ - read(proc.stdout), - read(proc.stderr), - proc.wait(), - ]) - return { stdout, stderr, exitCode } - } const run = await this.box.exec.command(this.wrap(command, opts)) return { stdout: run.stdout, diff --git a/packages/ai-sandbox-upstash-box/tests/handle.test.ts b/packages/ai-sandbox-upstash-box/tests/handle.test.ts index c6c37e71cc..cf0828cdd0 100644 --- a/packages/ai-sandbox-upstash-box/tests/handle.test.ts +++ b/packages/ai-sandbox-upstash-box/tests/handle.test.ts @@ -161,25 +161,6 @@ describe('UpstashBoxHandle', () => { expect(commands[0]).toBe("cd '/workspace/home' && echo hello") }) - it('kills an exec in the box when its signal aborts', async () => { - const fake = fakeSession() - const { box, sessionOptions } = fakeBox({ session: fake }) - const handle = new UpstashBoxHandle({ box }) - const abort = new AbortController() - - const running = handle.process.exec('sleep 600', { signal: abort.signal }) - await vi.waitFor(() => expect(sessionOptions).toHaveLength(1)) - fake.session.emitStdout('partial') - abort.abort() - - expect(fake.session.kill).toHaveBeenCalledOnce() - await expect(running).resolves.toEqual({ - stdout: 'partial', - stderr: '', - exitCode: 143, - }) - }) - it('maps /workspace cwd and applies env exports in order', async () => { const { box, commands } = fakeBox() const handle = new UpstashBoxHandle({ box }) From 8486aed55a5009ea46680c90f762193ce9eb31ba Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 17:18:42 +0100 Subject: [PATCH 13/24] fix(ai-code-mode, ai-computer-use, ai-harness, ai-sandbox-daytona, ai-sandbox-vercel): close abort gaps found in review - code mode: do not run the code when the run aborts while the context is made; dispose the isolate once and await that same promise. - computer use: an action does not start when its call aborts while the machine starts; the screenshot wait after an action ends on abort. - harness: the shared rg probe no longer takes the call signal, so one aborted call cannot cache 'no rg' for every later call. - daytona, vercel: check the abort again after the workdir is made, and delete or stop the sandbox. - e2e: the stop test now fails on any error that is not the abort. - docs: code mode says which drivers stop at once; add cancelled to the task status tables. --- docs/code-mode/code-mode.md | 2 +- docs/mcp/server-tasks.md | 3 +- .../ai-code-mode/src/create-code-mode-tool.ts | 34 +++++++--- .../tests/create-code-mode-tool.test.ts | 62 +++++++++++++++++++ packages/ai-computer-use/src/actions.ts | 4 +- .../src/create-computer-use.ts | 10 +-- .../tests/create-computer-use.test.ts | 31 ++++++++++ .../src/first-party/coding/search.ts | 9 ++- .../ai-harness/tests/coding-search.test.ts | 42 ++++++++++++- packages/ai-sandbox-daytona/src/provider.ts | 6 +- .../ai-sandbox-daytona/tests/provider.test.ts | 16 +++++ packages/ai-sandbox-vercel/src/provider.ts | 6 +- .../ai-sandbox-vercel/tests/vercel.test.ts | 22 +++++++ testing/e2e/tests/computer-use.spec.ts | 7 ++- 14 files changed, 231 insertions(+), 23 deletions(-) diff --git a/docs/code-mode/code-mode.md b/docs/code-mode/code-mode.md index f323bdb0aa..33518fd236 100644 --- a/docs/code-mode/code-mode.md +++ b/docs/code-mode/code-mode.md @@ -141,7 +141,7 @@ Code Mode rejects tools that need approval. A tool marked `needsApproval: true` ### Stop tool calls when the run stops -When a user stops a chat run, Code Mode disposes the isolate, so the program stops. A tool call that the program started keeps going, unless the tool listens for the abort. Each `external_*` call gets the `abortSignal` of the chat run in its second argument. Pass it to the slow work: +When a user stops a chat run, Code Mode disposes the isolate. With the Node and E2B drivers the program stops at once. With other drivers, the running code can continue until it ends or until its timeout. A tool call that the program started keeps going, unless the tool listens for the abort. Each `external_*` call gets the `abortSignal` of the chat run in its second argument. Pass it to the slow work: ```typescript import { toolDefinition } from "@tanstack/ai"; diff --git a/docs/mcp/server-tasks.md b/docs/mcp/server-tasks.md index 9d6fc88f1a..2a8680ecab 100644 --- a/docs/mcp/server-tasks.md +++ b/docs/mcp/server-tasks.md @@ -196,7 +196,7 @@ On spec 2025, the call response is the `CreateTaskResult` from the spec. The `ta | Field | Value | | --- | --- | | `taskId` | The id from the call. | -| `status` | `working`, `completed`, or `failed`. | +| `status` | `working`, `completed`, `failed`, or `cancelled`. | | `ttl` | `null`. | | `createdAt` | The time of the first `set`. | | `lastUpdatedAt` | The time of the last `set`. | @@ -209,6 +209,7 @@ On spec 2025, the call response is the `CreateTaskResult` from the spec. The `ta | `working` | The tool function is still at work. | | `completed` | The store holds the tool result. | | `failed` | The tool function threw an error. | +| `cancelled` | A host sent `tasks/cancel`. The tool's `ctx.abortSignal` aborted. | If the tool function throws, the task keeps the text of that error in `statusMessage`. If that error has no text, the task keeps `The tool failed.` 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 613407f6d3..380f960906 100644 --- a/packages/ai-code-mode/src/create-code-mode-tool.ts +++ b/packages/ai-code-mode/src/create-code-mode-tool.ts @@ -217,11 +217,21 @@ export function createCodeModeTool( // Aborts when this program ends, so a host call that still waits can stop. const programEnd = new AbortController() // An abort of the run ends the program: a call that waits (for example - // for an approval) stops, and the isolate goes, so its code stops too. + // for an approval) stops, and the isolate is disposed. How fast running + // code stops depends on the driver. const runSignal = toolContext?.abortSignal + let disposing: Promise | undefined + const disposeContext = () => { + if (isolateContext && !disposing) { + disposing = isolateContext.dispose() + // Marked as seen; the finally block still awaits it. + disposing.catch(() => {}) + } + return disposing + } const onRunAbort = () => { programEnd.abort() - void isolateContext?.dispose() + void disposeContext() } runSignal?.addEventListener('abort', onRunAbort, { once: true }) @@ -302,8 +312,20 @@ export function createCodeModeTool( ) } - // An abort that came while the context was made. - if (runSignal?.aborted) onRunAbort() + // An abort that came while the context was made: do not run the code. + if (runSignal?.aborted) { + onRunAbort() + return finish( + { + success: false, + error: { + message: 'The run was aborted before the code ran', + name: 'AbortError', + }, + }, + 'create-context', + ) + } // Step 5: Execute the code in the sandbox const executionResult = await isolateContext.execute(strippedCode) @@ -379,9 +401,7 @@ export function createCodeModeTool( // The program has ended (success, error, or timeout). programEnd.abort() // Always clean up the sandbox context - if (isolateContext) { - await isolateContext.dispose() - } + await disposeContext() } }, ) diff --git a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts index 2ac9c5f658..1fdd89a224 100644 --- a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts +++ b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts @@ -731,4 +731,66 @@ describe('createCodeModeTool run abort', () => { expect(dispose).toHaveBeenCalled() expect((await running).success).toBe(false) }) + + it('does not run the code when the run aborts while the context is made', async () => { + const execute = vi.fn() + const dispose = vi.fn(async () => {}) + const abort = new AbortController() + const driver: IsolateDriver = { + createContext: async () => { + abort.abort() + return { execute, dispose } + }, + } + const tool = createCodeModeTool({ + driver, + tools: [createMockTool('fetchWeather')], + }) + + const result = await tool.execute!( + { typescriptCode: 'return 1' }, + { abortSignal: abort.signal, emitCustomEvent: vi.fn() }, + ) + + expect(execute).not.toHaveBeenCalled() + expect(result.success).toBe(false) + expect(dispose).toHaveBeenCalledTimes(1) + }) + + it('disposes the isolate once when the run aborts', async () => { + let release = () => {} + const disposed = new Promise((resolve) => { + release = resolve + }) + const dispose = vi.fn(async () => release()) + const driver: IsolateDriver = { + createContext: async () => ({ + execute: async () => { + await disposed + return { + success: false, + error: { name: 'DisposedError', message: 'Context was disposed' }, + logs: [], + } + }, + dispose, + }), + } + const tool = createCodeModeTool({ + driver, + tools: [createMockTool('fetchWeather')], + }) + const abort = new AbortController() + + const running = tool.execute!( + { typescriptCode: 'return 1' }, + { abortSignal: abort.signal, emitCustomEvent: vi.fn() }, + ) + await vi.waitFor(() => expect(dispose).not.toHaveBeenCalled()) + await new Promise((resolve) => setTimeout(resolve, 10)) + abort.abort() + await running + + expect(dispose).toHaveBeenCalledTimes(1) + }) }) diff --git a/packages/ai-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index 71a678f556..5f114673dc 100644 --- a/packages/ai-computer-use/src/actions.ts +++ b/packages/ai-computer-use/src/actions.ts @@ -20,7 +20,7 @@ function isRecord(value: unknown): value is Input { } /** Waits `ms`. Rejects with the reason of `signal` when it aborts first. */ -function sleep(ms: number, signal?: AbortSignal) { +export function sleep(ms: number, signal?: AbortSignal) { return new Promise((resolve, reject) => { signal?.throwIfAborted() const onAbort = () => { @@ -405,6 +405,8 @@ export async function runComputerAction( signal?: AbortSignal, ) { if (!isRecord(input)) throw invalid('input', 'an object') + // The call can abort while its machine starts. + signal?.throwIfAborted() const { action, actions } = input if (typeof action === 'string') { return runAnthropicAction(driver, action, input, signal) diff --git a/packages/ai-computer-use/src/create-computer-use.ts b/packages/ai-computer-use/src/create-computer-use.ts index 7c8e2b32da..9260a9b49a 100644 --- a/packages/ai-computer-use/src/create-computer-use.ts +++ b/packages/ai-computer-use/src/create-computer-use.ts @@ -1,7 +1,7 @@ import { brandProviderTool, toolDefinition } from '@tanstack/ai' import { arrayBufferToBase64 } from '@tanstack/ai-utils' import { z } from 'zod' -import { runComputerAction } from './actions' +import { runComputerAction, sleep } from './actions' import { computerSafetyCheck, safetyCheckSchema } from './safety-check' import type { ChatMiddleware, @@ -108,10 +108,6 @@ function isWebUrl(url: string) { } } -function sleep(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)) -} - function imageResult(bytes: Uint8Array) { const part: ImagePart = { type: 'image', @@ -337,7 +333,7 @@ export function createComputerUse({ context?.abortSignal, ) if (result === undefined) { - await sleep(screenshotDelay) + await sleep(screenshotDelay, context?.abortSignal) return imageResult(await machine.screenshot()) } return result.type === 'image' ? imageResult(result.bytes) : result.text @@ -357,7 +353,7 @@ export function createComputerUse({ throw new Error('The computer driver has no navigate method.') } await machine.navigate(url) - await sleep(screenshotDelay) + await sleep(screenshotDelay, context?.abortSignal) return imageResult(await machine.screenshot()) }), ) 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 5f4891b726..13912c9a99 100644 --- a/packages/ai-computer-use/tests/create-computer-use.test.ts +++ b/packages/ai-computer-use/tests/create-computer-use.test.ts @@ -979,4 +979,35 @@ describe('the abort signal of a computer action', () => { await expect(waiting).rejects.toThrow() expect(Date.now() - started).toBeLessThan(1000) }) + + it('does not run an action when its call aborts while the machine starts', async () => { + const abort = new AbortController() + const { factory, machines } = createFactory({ + // `create` runs while the machine starts: abort the call there. + failCreate: () => { + abort.abort() + return false + }, + }) + const computer = setupWith(factory) + + await expect( + executeWith(computer, CLICK, 'call_1', abort.signal), + ).rejects.toThrow() + expect(machines.get(THREAD)).toEqual([]) + }) + + it('ends the screenshot wait after an action when its call aborts', async () => { + const { driver, calls } = createDriver() + const computer = setupWith(driver, { screenshotDelay: 5000 }) + const abort = new AbortController() + const running = executeWith(computer, CLICK, 'call_1', abort.signal) + await vi.waitFor(() => expect(calls).toContain('click:end')) + const started = Date.now() + abort.abort() + + await expect(running).rejects.toThrow() + expect(Date.now() - started).toBeLessThan(1000) + expect(calls).not.toContain('screenshot') + }) }) diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index ca1dbb4358..e07666c016 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -109,10 +109,17 @@ async function findRg(backend: WorkspaceBackend) { const bundled = await bundledRg() if (bundled !== undefined) return quoteArg(bundled, process.platform) } - const probe = await backend.exec('rg --version', { timeoutMs: TIMEOUT_MS }) + // Every later call reuses this result, so the probe does not take the + // signal of the call that runs it: an aborted call would cache "no rg". + const probe = await backend.exec('rg --version', { + timeoutMs: TIMEOUT_MS, + signal: NEVER_ABORTS, + }) return probe.exitCode === 0 ? 'rg' : undefined } +const NEVER_ABORTS = new AbortController().signal + /** Run an rg command in `cwd`, and give its output. */ async function runRg(backend: WorkspaceBackend, command: string, cwd: string) { const { exitCode, stdout, stderr } = await backend.exec(command, { diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index 22abab465b..f74853339d 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -583,7 +583,8 @@ describe('the abort signal of a call', () => { const backend: WorkspaceBackend = { ...fake.backend, exec: (command, options) => { - signals.push(options?.signal) + // The rg probe is shared by every call, so it has its own signal. + if (command !== 'rg --version') signals.push(options?.signal) return fake.backend.exec(command, options) }, } @@ -598,4 +599,43 @@ describe('the abort signal of a call', () => { expect(signals.length).toBeGreaterThan(0) expect(signals.filter((each) => each !== signal)).toEqual([]) }) + + it('still finds rg after a call that aborted during the rg probe', async () => { + const fake = fakeBackend({ 'a.ts': 'one\n' }, (command) => + command.startsWith('rg --files ') ? ok('./a.ts\0') : undefined, + ) + const backend: WorkspaceBackend = { + ...fake.backend, + exec: (command, options) => { + fake.commands.push(command) + // An aborted command ends with a failure, like `hostBackend.exec`. + if (command === 'rg --version') { + return Promise.resolve( + options?.signal?.aborted + ? { exitCode: 1, stdout: '', stderr: '' } + : ok('ripgrep 14.1.0'), + ) + } + return fake.backend.exec(command, options) + }, + } + const { tools } = createWorkspaceTools({ root, backend }) + const listFiles = tools.find((each) => each.name === 'list_files')?.execute + if (!listFiles) throw new Error('No tool list_files') + const aborted = new AbortController() + aborted.abort() + await Promise.resolve( + listFiles( + {}, + { abortSignal: aborted.signal, emitCustomEvent: () => undefined }, + ), + ).catch(() => undefined) + + expect(await listFiles({}, { emitCustomEvent: () => undefined })).toBe( + 'a.ts', + ) + expect( + fake.commands.filter((c) => c.startsWith('rg --files ')), + ).not.toEqual([]) + }) }) diff --git a/packages/ai-sandbox-daytona/src/provider.ts b/packages/ai-sandbox-daytona/src/provider.ts index 5f878a68dc..dfe9ab414d 100644 --- a/packages/ai-sandbox-daytona/src/provider.ts +++ b/packages/ai-sandbox-daytona/src/provider.ts @@ -101,11 +101,15 @@ class DaytonaProvider implements SandboxProvider { ): Promise { // The SDK create takes no signal. A sandbox that comes after the abort has // no owner, so delete it instead of leaving it billed. - if (signal?.aborted) { + const deleteIfAborted = async () => { + if (!signal?.aborted) return await this.daytona.delete(sandbox).catch(() => undefined) signal.throwIfAborted() } + await deleteIfAborted() await sandbox.process.executeCommand(`mkdir -p ${shQuote(this.workdir)}`) + // The abort can also come while the workdir is made. + await deleteIfAborted() return new DaytonaHandle({ sandbox, workdir: this.workdir, applyEnvSet }) } diff --git a/packages/ai-sandbox-daytona/tests/provider.test.ts b/packages/ai-sandbox-daytona/tests/provider.test.ts index 0273a4fc77..f45736173d 100644 --- a/packages/ai-sandbox-daytona/tests/provider.test.ts +++ b/packages/ai-sandbox-daytona/tests/provider.test.ts @@ -439,4 +439,20 @@ describe('daytonaSandbox abort', () => { expect(del).toHaveBeenCalledWith(late) }, ) + + it('deletes the sandbox when the abort comes while its workdir is made', async () => { + const abort = new AbortController() + const sandbox = fakeRemoteSandbox('sbx-setup', { + executeCommand: vi.fn(async () => { + abort.abort() + return { result: '', exitCode: 0 } + }), + }) + create.mockResolvedValue(sandbox) + del.mockResolvedValue(undefined) + const provider = daytonaSandbox({ apiKey: 'test-key' }) + + await expect(provider.create({ signal: abort.signal })).rejects.toThrow() + expect(del).toHaveBeenCalledWith(sandbox) + }) }) diff --git a/packages/ai-sandbox-vercel/src/provider.ts b/packages/ai-sandbox-vercel/src/provider.ts index e1491f6800..e25fc1d53f 100644 --- a/packages/ai-sandbox-vercel/src/provider.ts +++ b/packages/ai-sandbox-vercel/src/provider.ts @@ -132,10 +132,12 @@ class VercelProvider implements SandboxProvider { // ponytail: no SDK signal here. It cancels only the request, so a sandbox // that the server made has no owner. Let the create end, then stop a // sandbox that comes after the abort instead of leaving it billed. - if (input.signal?.aborted) { + const stopIfAborted = async () => { + if (!input.signal?.aborted) return await sandbox.stop().catch(() => undefined) input.signal.throwIfAborted() } + await stopIfAborted() // Ensure the workspace dir exists via the native (cwd-independent) mkDir — // running a command with a not-yet-existing `cwd` would fail, so we must not // route this through the handle (which runs every command in `workdir`). @@ -149,6 +151,8 @@ class VercelProvider implements SandboxProvider { } catch (error) { if (!isDirAlreadyExistsError(error)) throw error } + // The abort can also come while the workdir is made. + await stopIfAborted() return new VercelHandle({ sandbox, workdir: this.workdir, diff --git a/packages/ai-sandbox-vercel/tests/vercel.test.ts b/packages/ai-sandbox-vercel/tests/vercel.test.ts index 2f76fd9609..f160777e17 100644 --- a/packages/ai-sandbox-vercel/tests/vercel.test.ts +++ b/packages/ai-sandbox-vercel/tests/vercel.test.ts @@ -219,4 +219,26 @@ describe('vercelSandbox create abort', () => { create.mockRestore() } }) + + it('stops the sandbox when the abort comes while its workdir is made', async () => { + const abort = new AbortController() + const stop = vi.fn(() => Promise.resolve()) + const mkDir = vi.fn(async () => { + abort.abort() + }) + const create = vi + .spyOn(VercelSandbox, 'create') + .mockResolvedValue({ name: 'sbx-setup', mkDir, stop } as never) + try { + const provider = vercelSandbox({ + token: 't', + teamId: 't', + projectId: 'p', + }) + await expect(provider.create({ signal: abort.signal })).rejects.toThrow() + expect(stop).toHaveBeenCalledOnce() + } finally { + create.mockRestore() + } + }) }) diff --git a/testing/e2e/tests/computer-use.spec.ts b/testing/e2e/tests/computer-use.spec.ts index 17b6f2032f..d440536fc2 100644 --- a/testing/e2e/tests/computer-use.spec.ts +++ b/testing/e2e/tests/computer-use.spec.ts @@ -399,10 +399,13 @@ test.describe('computer use with a Playwright page', () => { setTimeout(() => abort.abort(), 50) } } - } catch { - // An aborted run can end with an error. The checks below are the test. + } catch (error) { + // An aborted run can end with an error. Any other error fails the test. + if (!abort.signal.aborted) throw error } + // The stop really happened, so the timing check below tests it. + expect(abort.signal.aborted).toBe(true) expect(Date.now() - started).toBeLessThan(1500) await expect(page.locator('#status')).toHaveText('idle') }) From 349955dc32f1e1f23644eaa082c8fa8c35c5f21f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:14:51 +0100 Subject: [PATCH 14/24] fix(ai-harness): start no search work for a call that aborted already --- .../src/first-party/coding/search.ts | 7 ++++-- .../ai-harness/tests/coding-search.test.ts | 25 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index e07666c016..34bd98a52d 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -439,7 +439,9 @@ export function searchTools(env: ToolEnv) { }, }, replay: 'safe', - }).server(async (args: unknown) => { + }).server(async (args: unknown, context) => { + // An aborted call starts no search work, not even the shared rg probe. + context?.abortSignal?.throwIfAborted() const base = await env.reach( optionalString(args, 'path') ?? '.', 'list_files', @@ -472,7 +474,8 @@ export function searchTools(env: ToolEnv) { required: ['pattern'], }, replay: 'safe', - }).server(async (args: unknown) => { + }).server(async (args: unknown, context) => { + context?.abortSignal?.throwIfAborted() const pattern = stringArg(args, 'pattern') if (/[\r\n]/.test(pattern)) { throw new Error('The pattern must be one line. grep tests each line.') diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index f74853339d..5b6cee40d6 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -638,4 +638,29 @@ describe('the abort signal of a call', () => { fake.commands.filter((c) => c.startsWith('rg --files ')), ).not.toEqual([]) }) + + it.each([ + ['list_files', {}], + ['grep', { pattern: 'one' }], + ])( + '%s starts no command for a call that aborted already', + async (name, args) => { + const fake = fakeBackend({ 'a.ts': 'one\n' }) + const { tools } = createWorkspaceTools({ root, backend: fake.backend }) + const execute = tools.find((each) => each.name === name)?.execute + if (!execute) throw new Error(`No tool ${name}`) + const aborted = new AbortController() + aborted.abort() + + await expect( + Promise.resolve( + execute(args, { + abortSignal: aborted.signal, + emitCustomEvent: () => undefined, + }), + ), + ).rejects.toThrow() + expect(fake.commands).toEqual([]) + }, + ) }) From 8a72ccaf3081cdcfd5b2ae096625aad74b3abdf1 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:14:53 +0100 Subject: [PATCH 15/24] fix(ai-code-mode): make no isolate when the run aborts while the snippets load --- .../ai-code-mode/src/create-code-mode-tool.ts | 25 +++++++++++-------- .../tests/create-code-mode-tool.test.ts | 21 ++++++++++++++++ 2 files changed, 36 insertions(+), 10 deletions(-) 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 380f960906..0a612fbcad 100644 --- a/packages/ai-code-mode/src/create-code-mode-tool.ts +++ b/packages/ai-code-mode/src/create-code-mode-tool.ts @@ -234,6 +234,17 @@ export function createCodeModeTool( void disposeContext() } runSignal?.addEventListener('abort', onRunAbort, { once: true }) + const abortedBeforeRun = () => + finish( + { + success: false, + error: { + message: 'The run was aborted before the code ran', + name: 'AbortError', + }, + }, + 'create-context', + ) // Emit execution started event immediately emitCustomEvent('code_mode:execution_started', { @@ -289,6 +300,9 @@ export function createCodeModeTool( aroundCalls?.(toolContext, programEnd.signal), ) + // An abort that came while the snippets loaded: make no context. + if (runSignal?.aborted) return abortedBeforeRun() + // Step 4: Create sandbox context with event-aware bindings try { isolateContext = await driver.createContext({ @@ -315,16 +329,7 @@ export function createCodeModeTool( // An abort that came while the context was made: do not run the code. if (runSignal?.aborted) { onRunAbort() - return finish( - { - success: false, - error: { - message: 'The run was aborted before the code ran', - name: 'AbortError', - }, - }, - 'create-context', - ) + return abortedBeforeRun() } // Step 5: Execute the code in the sandbox diff --git a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts index 1fdd89a224..a6088d3aac 100644 --- a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts +++ b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts @@ -757,6 +757,27 @@ describe('createCodeModeTool run abort', () => { expect(dispose).toHaveBeenCalledTimes(1) }) + it('does not make the context when the run aborts while the snippets load', async () => { + const createContext = vi.fn() + const abort = new AbortController() + const tool = createCodeModeTool({ + driver: { createContext }, + tools: [createMockTool('fetchWeather')], + getSnippetBindings: async () => { + abort.abort() + return {} + }, + }) + + const result = await tool.execute!( + { typescriptCode: 'return 1' }, + { abortSignal: abort.signal, emitCustomEvent: vi.fn() }, + ) + + expect(createContext).not.toHaveBeenCalled() + expect(result.success).toBe(false) + }) + it('disposes the isolate once when the run aborts', async () => { let release = () => {} const disposed = new Promise((resolve) => { From f65da49c53d7185cecb43b85992a899e309200f3 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:14:55 +0100 Subject: [PATCH 16/24] fix(ai-computer-use): stop navigate and compound actions between steps on abort --- .changeset/abort-computer-use-actions.md | 2 +- packages/ai-computer-use/src/actions.ts | 47 ++++++--- .../src/create-computer-use.ts | 2 + .../ai-computer-use/tests/actions.test.ts | 99 +++++++++++++++++++ .../tests/create-computer-use.test.ts | 29 ++++++ 5 files changed, 164 insertions(+), 15 deletions(-) diff --git a/.changeset/abort-computer-use-actions.md b/.changeset/abort-computer-use-actions.md index 058afc2569..17e3f2fb11 100644 --- a/.changeset/abort-computer-use-actions.md +++ b/.changeset/abort-computer-use-actions.md @@ -2,4 +2,4 @@ '@tanstack/ai-computer-use': patch --- -An aborted computer action no longer runs when it waits in the queue. A `wait`, `hold_key`, or the next action of a batch now stops when the call aborts. +An aborted computer action no longer runs when it waits in the queue. A `wait`, `hold_key`, a repeated key press, a drag, a `navigate`, or the next action of a batch now stops when the call aborts. Held keys and the mouse button are still released. diff --git a/packages/ai-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index 5f114673dc..cba0470d55 100644 --- a/packages/ai-computer-use/src/actions.ts +++ b/packages/ai-computer-use/src/actions.ts @@ -92,9 +92,13 @@ function readSeconds(input: Input, key: string) { return value } -/** Presses `keys` in order, runs `action`, then releases the keys that went down. */ +/** + * Presses `keys` in order, runs `action`, then releases the keys that went + * down. An abort of `signal` stops it between steps. The keys still go up. + */ async function withKeys( driver: ComputerDriver, + signal: AbortSignal | undefined, keys: Array, action: () => Promise, ) { @@ -104,6 +108,7 @@ async function withKeys( for (const key of keys) { await driver.keyDown(key) pressed.push(key) + signal?.throwIfAborted() } await action() } finally { @@ -118,13 +123,22 @@ async function withKeys( if (releaseError !== undefined) throw releaseError } -async function drag(driver: ComputerDriver, path: Array) { +async function drag( + driver: ComputerDriver, + path: Array, + signal: AbortSignal | undefined, +) { const [start, ...rest] = path if (start === undefined) throw invalid('path', 'a non-empty array') await driver.move(start.x, start.y) + signal?.throwIfAborted() await driver.mouseDown('left') try { - for (const point of rest) await driver.move(point.x, point.y) + for (const point of rest) { + // An abort stops the drag between moves. The mouse still goes up. + signal?.throwIfAborted() + await driver.move(point.x, point.y) + } } finally { await driver.mouseUp('left') } @@ -185,7 +199,7 @@ async function runAnthropicAction( const click = ANTHROPIC_CLICKS[action] if (click) { const { x, y } = await anthropicPoint(driver, input) - await withKeys(driver, anthropicModifiers(input), () => + await withKeys(driver, signal, anthropicModifiers(input), () => driver.click(x, y, click), ) return undefined @@ -198,7 +212,10 @@ async function runAnthropicAction( const combos = readString(input, 'text').split(/\s+/).filter(Boolean) const repeat = readRepeat(input) for (let i = 0; i < repeat; i++) { - for (const combo of combos) await driver.press(parseKeyCombo(combo)) + for (const combo of combos) { + signal?.throwIfAborted() + await driver.press(parseKeyCombo(combo)) + } } return undefined } @@ -215,8 +232,8 @@ async function runAnthropicAction( readCoordinate(input, 'start_coordinate'), readCoordinate(input, 'coordinate'), ] - await withKeys(driver, anthropicModifiers(input), () => - drag(driver, path), + await withKeys(driver, signal, anthropicModifiers(input), () => + drag(driver, path, signal), ) return undefined } @@ -231,7 +248,7 @@ async function runAnthropicAction( const pixels = readNumber(input, 'scroll_amount') * SCROLL_STEP const direction = readString(input, 'scroll_direction') const { deltaX, deltaY } = scrollDelta(direction, pixels) - await withKeys(driver, anthropicModifiers(input), () => + await withKeys(driver, signal, anthropicModifiers(input), () => driver.scroll(x, y, deltaX, deltaY), ) return undefined @@ -239,7 +256,7 @@ async function runAnthropicAction( case 'hold_key': { const keys = parseKeyCombo(readString(input, 'text')) const ms = readSeconds(input, 'duration') * 1000 - await withKeys(driver, keys, () => sleep(ms, signal)) + await withKeys(driver, signal, keys, () => sleep(ms, signal)) return undefined } case 'wait': @@ -326,7 +343,7 @@ async function runOpenAIAction( } const x = readNumber(input, 'x') const y = readNumber(input, 'y') - await withKeys(driver, modifiers, () => + await withKeys(driver, signal, modifiers, () => driver.click(x, y, { button, count: 1 }), ) return @@ -334,14 +351,16 @@ async function runOpenAIAction( case 'double_click': { const x = readNumber(input, 'x') const y = readNumber(input, 'y') - await withKeys(driver, modifiers, () => + await withKeys(driver, signal, modifiers, () => driver.click(x, y, { button: 'left', count: 2 }), ) return } case 'drag': { const path = readPath(input) - await withKeys(driver, modifiers, () => drag(driver, path)) + await withKeys(driver, signal, modifiers, () => + drag(driver, path, signal), + ) return } case 'keypress': { @@ -357,7 +376,7 @@ async function runOpenAIAction( case 'move': { const x = readNumber(input, 'x') const y = readNumber(input, 'y') - await withKeys(driver, modifiers, () => driver.move(x, y)) + await withKeys(driver, signal, modifiers, () => driver.move(x, y)) return } case 'scroll': { @@ -365,7 +384,7 @@ async function runOpenAIAction( const y = readNumber(input, 'y') const deltaX = readNumber(input, 'scroll_x') const deltaY = readNumber(input, 'scroll_y') - await withKeys(driver, modifiers, () => + await withKeys(driver, signal, modifiers, () => driver.scroll(x, y, deltaX, deltaY), ) return diff --git a/packages/ai-computer-use/src/create-computer-use.ts b/packages/ai-computer-use/src/create-computer-use.ts index 9260a9b49a..2c5ece98b4 100644 --- a/packages/ai-computer-use/src/create-computer-use.ts +++ b/packages/ai-computer-use/src/create-computer-use.ts @@ -352,6 +352,8 @@ export function createComputerUse({ if (machine.navigate === undefined) { throw new Error('The computer driver has no navigate method.') } + // The call can abort while its machine starts. + context?.abortSignal?.throwIfAborted() await machine.navigate(url) await sleep(screenshotDelay, context?.abortSignal) return imageResult(await machine.screenshot()) diff --git a/packages/ai-computer-use/tests/actions.test.ts b/packages/ai-computer-use/tests/actions.test.ts index b2717b09a8..3f5a789856 100644 --- a/packages/ai-computer-use/tests/actions.test.ts +++ b/packages/ai-computer-use/tests/actions.test.ts @@ -502,3 +502,102 @@ describe('OpenAI actions', () => { it('throws for an input with no action', async () => { await expect(run({ foo: 1 })).rejects.toThrow('"action"') }) + +describe('an abort during a compound action', () => { + /** + * A driver that records its calls. Call number `nth` of `method` aborts + * the signal while it runs, as a slow driver call that the user stops. + */ + function abortingDuring(method: 'press' | 'keyDown' | 'move', nth = 1) { + const { driver, calls } = createDriver() + const controller = new AbortController() + const original = driver[method] as ( + ...args: Array + ) => Promise + let count = 0 + driver[method] = (async (...args: Array) => { + await original(...args) + count += 1 + if (count === nth) controller.abort() + }) as never + return { driver, calls, signal: controller.signal } + } + + it('stops a repeated key press', async () => { + const { driver, calls, signal } = abortingDuring('press') + await expect( + runComputerAction( + driver, + { action: 'key', text: 'a', repeat: 5 }, + signal, + ), + ).rejects.toThrow() + expect(calls.filter(([name]) => name === 'press')).toHaveLength(1) + }) + + it('does not run the action after a modifier key, and releases the key', async () => { + const { driver, calls, signal } = abortingDuring('keyDown') + await expect( + runComputerAction( + driver, + { + action: { + type: 'click', + button: 'left', + x: 1, + y: 2, + keys: ['shift', 'ctrl'], + }, + }, + signal, + ), + ).rejects.toThrow() + expect(calls.map(([name]) => name)).toEqual(['keyDown', 'keyUp']) + }) + + it('does not press the mouse when the drag aborts at its first move', async () => { + const { driver, calls, signal } = abortingDuring('move') + await expect( + runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + ], + }, + }, + signal, + ), + ).rejects.toThrow() + expect(calls.map(([name]) => name)).toEqual(['move']) + }) + + it('stops a drag between two moves, and releases the mouse', async () => { + const { driver, calls, signal } = abortingDuring('move', 2) + await expect( + runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + { x: 2, y: 2 }, + ], + }, + }, + signal, + ), + ).rejects.toThrow() + expect(calls.map(([name]) => name)).toEqual([ + 'move', + 'mouseDown', + 'move', + 'mouseUp', + ]) + }) +}) 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 13912c9a99..f79c3e2ffc 100644 --- a/packages/ai-computer-use/tests/create-computer-use.test.ts +++ b/packages/ai-computer-use/tests/create-computer-use.test.ts @@ -997,6 +997,35 @@ describe('the abort signal of a computer action', () => { expect(machines.get(THREAD)).toEqual([]) }) + it('does not navigate when its call aborts while the machine starts', async () => { + const abort = new AbortController() + const { factory, machines } = createFactory({ + navigate: true, + // `create` runs while the machine starts: abort the call there. + failCreate: () => { + abort.abort() + return false + }, + }) + const [, navigateTool] = setupWith(factory).tools + if (!navigateTool?.execute) throw new Error('There is no navigate tool.') + + await expect( + Promise.resolve( + navigateTool.execute( + { url: 'https://tanstack.com' }, + { + toolCallId: 'call_nav', + threadId: THREAD, + abortSignal: abort.signal, + emitCustomEvent() {}, + }, + ), + ), + ).rejects.toThrow() + expect(machines.get(THREAD)).toEqual([]) + }) + it('ends the screenshot wait after an action when its call aborts', async () => { const { driver, calls } = createDriver() const computer = setupWith(driver, { screenshotDelay: 5000 }) From db2a554ce928bcca333aadfdcd442e2f86935142 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:14:57 +0100 Subject: [PATCH 17/24] docs(mcp): say that a task cancel aborts only the task in this process --- docs/mcp/server-tasks.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/mcp/server-tasks.md b/docs/mcp/server-tasks.md index 2a8680ecab..fd377be252 100644 --- a/docs/mcp/server-tasks.md +++ b/docs/mcp/server-tasks.md @@ -209,7 +209,7 @@ On spec 2025, the call response is the `CreateTaskResult` from the spec. The `ta | `working` | The tool function is still at work. | | `completed` | The store holds the tool result. | | `failed` | The tool function threw an error. | -| `cancelled` | A host sent `tasks/cancel`. The tool's `ctx.abortSignal` aborted. | +| `cancelled` | A host sent `tasks/cancel`. `ctx.abortSignal` of the tool aborts only in the process that runs the task. With a shared store, see [Stop a cancelled task](#stop-a-cancelled-task). | If the tool function throws, the task keeps the text of that error in `statusMessage`. If that error has no text, the task keeps `The tool failed.` From a52b1a4e3d4b24b71e0e907fa26d5611eed57c26 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:41:14 +0100 Subject: [PATCH 18/24] fix(ai-computer-use): reject a drag that aborts during its last move --- packages/ai-computer-use/src/actions.ts | 5 ++-- .../ai-computer-use/tests/actions.test.ts | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/packages/ai-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index cba0470d55..5a963dcb18 100644 --- a/packages/ai-computer-use/src/actions.ts +++ b/packages/ai-computer-use/src/actions.ts @@ -135,9 +135,10 @@ async function drag( await driver.mouseDown('left') try { for (const point of rest) { - // An abort stops the drag between moves. The mouse still goes up. - signal?.throwIfAborted() await driver.move(point.x, point.y) + // An abort stops the drag after a move, also the last one. The mouse + // still goes up. + signal?.throwIfAborted() } } finally { await driver.mouseUp('left') diff --git a/packages/ai-computer-use/tests/actions.test.ts b/packages/ai-computer-use/tests/actions.test.ts index 3f5a789856..3700a1f2b1 100644 --- a/packages/ai-computer-use/tests/actions.test.ts +++ b/packages/ai-computer-use/tests/actions.test.ts @@ -575,6 +575,31 @@ describe('an abort during a compound action', () => { expect(calls.map(([name]) => name)).toEqual(['move']) }) + it('rejects a drag that aborts during its final move, and releases the mouse', async () => { + const { driver, calls, signal } = abortingDuring('move', 2) + await expect( + runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + ], + }, + }, + signal, + ), + ).rejects.toThrow() + expect(calls.map(([name]) => name)).toEqual([ + 'move', + 'mouseDown', + 'move', + 'mouseUp', + ]) + }) + it('stops a drag between two moves, and releases the mouse', async () => { const { driver, calls, signal } = abortingDuring('move', 2) await expect( From f9e5f596e26110a7bd14de619241ed1c8526e70b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:41:16 +0100 Subject: [PATCH 19/24] fix(ai-code-mode): load no snippets when the run aborts during transpile --- .../ai-code-mode/src/create-code-mode-tool.ts | 3 +++ .../tests/create-code-mode-tool.test.ts | 22 +++++++++++++++++++ 2 files changed, 25 insertions(+) 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 0a612fbcad..898569f0ce 100644 --- a/packages/ai-code-mode/src/create-code-mode-tool.ts +++ b/packages/ai-code-mode/src/create-code-mode-tool.ts @@ -274,6 +274,9 @@ export function createCodeModeTool( ) } + // An abort that came while the code was transpiled: load no snippets. + if (runSignal?.aborted) return abortedBeforeRun() + // Step 2: Get dynamic snippet bindings if available const snippetBindings = getSnippetBindings ? await getSnippetBindings() diff --git a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts index a6088d3aac..126f48bc2e 100644 --- a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts +++ b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts @@ -778,6 +778,28 @@ describe('createCodeModeTool run abort', () => { expect(result.success).toBe(false) }) + it('loads no snippets when the run aborts while the code is transpiled', async () => { + const getSnippetBindings = vi.fn(async () => ({})) + const abort = new AbortController() + const tool = createCodeModeTool({ + driver: { createContext: vi.fn() }, + tools: [createMockTool('fetchWeather')], + transpile: async (code) => { + abort.abort() + return code + }, + getSnippetBindings, + }) + + const result = await tool.execute!( + { typescriptCode: 'return 1' }, + { abortSignal: abort.signal, emitCustomEvent: vi.fn() }, + ) + + expect(getSnippetBindings).not.toHaveBeenCalled() + expect(result.success).toBe(false) + }) + it('disposes the isolate once when the run aborts', async () => { let release = () => {} const disposed = new Promise((resolve) => { From d7b49de8a888a5a7e93fa9d249a7fa9a90bbd978 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:41:18 +0100 Subject: [PATCH 20/24] fix(ai-harness): reject a search that the abort stops --- .../src/first-party/coding/search.ts | 4 +++ .../ai-harness/tests/coding-search.test.ts | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index 34bd98a52d..9662f90c29 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -449,6 +449,9 @@ export function searchTools(env: ToolEnv) { ) const isWanted = globMatcher(optionalString(args, 'pattern')) const { files, partial } = await filesUnder(base) + // A stopped command ends with a failure code, which can read as "no + // files". An aborted search is not a result. + context?.abortSignal?.throwIfAborted() const listed = files.filter(isWanted) const shown = listed .slice(0, MAX_FILES) @@ -496,6 +499,7 @@ export function searchTools(env: ToolEnv) { partial, files: listed, } = await grepUnder(base, pattern, isWanted, isProtected !== undefined) + context?.abortSignal?.throwIfAborted() // The note counts the protected files in the list, not in the hits, // so it does not tell whether a protected file has the pattern. const skipped = isProtected diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index 5b6cee40d6..c88ed8aa02 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -663,4 +663,36 @@ describe('the abort signal of a call', () => { expect(fake.commands).toEqual([]) }, ) + + it.each([ + ['list_files', {}], + ['grep', { pattern: 'one' }], + ])('%s rejects a search that the abort stops', async (name, args) => { + const fake = fakeBackend({ 'a.ts': 'one\n' }) + const controller = new AbortController() + const backend: WorkspaceBackend = { + ...fake.backend, + // The abort comes while the search command runs. A stopped command + // ends with a failure code, like on `hostBackend`. + exec: (command, options) => { + if (command === 'rg --version') { + return fake.backend.exec(command, options) + } + controller.abort() + return Promise.resolve({ exitCode: 1, stdout: '', stderr: '' }) + }, + } + const { tools } = createWorkspaceTools({ root, backend }) + const execute = tools.find((each) => each.name === name)?.execute + if (!execute) throw new Error(`No tool ${name}`) + + await expect( + Promise.resolve( + execute(args, { + abortSignal: controller.signal, + emitCustomEvent: () => undefined, + }), + ), + ).rejects.toThrow() + }) }) From 905c1e5c0d81509d89d9b52423148f7c8b5573f8 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:52:59 +0100 Subject: [PATCH 21/24] fix(ai-computer-use): make no drag move after an abort during mouseDown The signal can abort while the button goes down. Check it right after mouseDown, inside the try, so the button still goes up. --- packages/ai-computer-use/src/actions.ts | 2 ++ .../ai-computer-use/tests/actions.test.ts | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/packages/ai-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index 5a963dcb18..352442e3c0 100644 --- a/packages/ai-computer-use/src/actions.ts +++ b/packages/ai-computer-use/src/actions.ts @@ -134,6 +134,8 @@ async function drag( signal?.throwIfAborted() await driver.mouseDown('left') try { + // The abort can come while the button goes down. + signal?.throwIfAborted() for (const point of rest) { await driver.move(point.x, point.y) // An abort stops the drag after a move, also the last one. The mouse diff --git a/packages/ai-computer-use/tests/actions.test.ts b/packages/ai-computer-use/tests/actions.test.ts index 3700a1f2b1..ab0ef8ca57 100644 --- a/packages/ai-computer-use/tests/actions.test.ts +++ b/packages/ai-computer-use/tests/actions.test.ts @@ -600,6 +600,30 @@ describe('an abort during a compound action', () => { ]) }) + it('makes no move after a drag aborts during mouseDown, and releases the mouse', async () => { + const { driver, calls, signal } = abortingDuring('mouseDown', 1) + await expect( + runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + ], + }, + }, + signal, + ), + ).rejects.toThrow() + expect(calls.map(([name]) => name)).toEqual([ + 'move', + 'mouseDown', + 'mouseUp', + ]) + }) + it('stops a drag between two moves, and releases the mouse', async () => { const { driver, calls, signal } = abortingDuring('move', 2) await expect( From f2412897f2da853902f70c0a8e6074a64ef84713 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 2 Nov 2026 18:53:22 +0100 Subject: [PATCH 22/24] test(ai-computer-use): let abortingDuring take mouseDown --- packages/ai-computer-use/tests/actions.test.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/ai-computer-use/tests/actions.test.ts b/packages/ai-computer-use/tests/actions.test.ts index ab0ef8ca57..e0c45529e2 100644 --- a/packages/ai-computer-use/tests/actions.test.ts +++ b/packages/ai-computer-use/tests/actions.test.ts @@ -508,7 +508,10 @@ describe('an abort during a compound action', () => { * A driver that records its calls. Call number `nth` of `method` aborts * the signal while it runs, as a slow driver call that the user stops. */ - function abortingDuring(method: 'press' | 'keyDown' | 'move', nth = 1) { + function abortingDuring( + method: 'press' | 'keyDown' | 'move' | 'mouseDown', + nth = 1, + ) { const { driver, calls } = createDriver() const controller = new AbortController() const original = driver[method] as ( From 41df967db410ec05073430215b22dcde776280dc Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 3 Nov 2026 12:17:02 +0100 Subject: [PATCH 23/24] fix(ai-computer-use): finish a drag that started before the abort Each mouseUp is a drop. An abort after mouseDown made a click on the start element, and an abort between moves dropped the item at a point in the middle. The drag now checks the signal only before mouseDown. --- .changeset/abort-computer-use-actions.md | 2 +- packages/ai-computer-use/src/actions.ts | 12 +- .../ai-computer-use/tests/actions.test.ts | 105 +++++++----------- 3 files changed, 44 insertions(+), 75 deletions(-) diff --git a/.changeset/abort-computer-use-actions.md b/.changeset/abort-computer-use-actions.md index 17e3f2fb11..b5a5263f9f 100644 --- a/.changeset/abort-computer-use-actions.md +++ b/.changeset/abort-computer-use-actions.md @@ -2,4 +2,4 @@ '@tanstack/ai-computer-use': patch --- -An aborted computer action no longer runs when it waits in the queue. A `wait`, `hold_key`, a repeated key press, a drag, a `navigate`, or the next action of a batch now stops when the call aborts. Held keys and the mouse button are still released. +An aborted computer action no longer runs when it waits in the queue. A `wait`, `hold_key`, a repeated key press, a `navigate`, or the next action of a batch now stops when the call aborts. Held keys are still released. A drag does not start after an abort. A drag that started goes to its end point, because a release in the middle drops the item at the wrong point. diff --git a/packages/ai-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index 352442e3c0..84ad1bd900 100644 --- a/packages/ai-computer-use/src/actions.ts +++ b/packages/ai-computer-use/src/actions.ts @@ -132,16 +132,12 @@ async function drag( if (start === undefined) throw invalid('path', 'a non-empty array') await driver.move(start.x, start.y) signal?.throwIfAborted() + // After the button goes down, an abort does not stop the drag. Each + // mouseUp is a drop, so a stop in the middle drops the item at the wrong + // point. await driver.mouseDown('left') try { - // The abort can come while the button goes down. - signal?.throwIfAborted() - for (const point of rest) { - await driver.move(point.x, point.y) - // An abort stops the drag after a move, also the last one. The mouse - // still goes up. - signal?.throwIfAborted() - } + for (const point of rest) await driver.move(point.x, point.y) } finally { await driver.mouseUp('left') } diff --git a/packages/ai-computer-use/tests/actions.test.ts b/packages/ai-computer-use/tests/actions.test.ts index e0c45529e2..ca54e26c86 100644 --- a/packages/ai-computer-use/tests/actions.test.ts +++ b/packages/ai-computer-use/tests/actions.test.ts @@ -578,78 +578,51 @@ describe('an abort during a compound action', () => { expect(calls.map(([name]) => name)).toEqual(['move']) }) - it('rejects a drag that aborts during its final move, and releases the mouse', async () => { - const { driver, calls, signal } = abortingDuring('move', 2) - await expect( - runComputerAction( - driver, - { - action: { - type: 'drag', - path: [ - { x: 0, y: 0 }, - { x: 1, y: 1 }, - ], - }, - }, - signal, - ), - ).rejects.toThrow() - expect(calls.map(([name]) => name)).toEqual([ - 'move', - 'mouseDown', - 'move', - 'mouseUp', - ]) - }) - - it('makes no move after a drag aborts during mouseDown, and releases the mouse', async () => { - const { driver, calls, signal } = abortingDuring('mouseDown', 1) - await expect( - runComputerAction( - driver, - { - action: { - type: 'drag', - path: [ - { x: 0, y: 0 }, - { x: 1, y: 1 }, - ], - }, + it('finishes a drag that aborts while the button goes down', async () => { + const { driver, calls, signal } = abortingDuring('mouseDown') + await runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + ], }, - signal, - ), - ).rejects.toThrow() - expect(calls.map(([name]) => name)).toEqual([ - 'move', - 'mouseDown', - 'mouseUp', + }, + signal, + ) + expect(calls).toEqual([ + ['move', 0, 0], + ['mouseDown', 'left'], + ['move', 1, 1], + ['mouseUp', 'left'], ]) }) - it('stops a drag between two moves, and releases the mouse', async () => { + it('finishes a drag that aborts between two moves', async () => { const { driver, calls, signal } = abortingDuring('move', 2) - await expect( - runComputerAction( - driver, - { - action: { - type: 'drag', - path: [ - { x: 0, y: 0 }, - { x: 1, y: 1 }, - { x: 2, y: 2 }, - ], - }, + await runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + { x: 2, y: 2 }, + ], }, - signal, - ), - ).rejects.toThrow() - expect(calls.map(([name]) => name)).toEqual([ - 'move', - 'mouseDown', - 'move', - 'mouseUp', + }, + signal, + ) + expect(calls).toEqual([ + ['move', 0, 0], + ['mouseDown', 'left'], + ['move', 1, 1], + ['move', 2, 2], + ['mouseUp', 'left'], ]) }) }) From ac6478202d7a80ffdcf7332ff152e9b5ea2f2c58 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 3 Nov 2026 12:17:03 +0100 Subject: [PATCH 24/24] fix(ai-mcp): keep the result of a task that ends during tasks/cancel The cancel handler read the task as working and then waited. When the run ended in that gap, cancelTask saved cancelled over completed. Now cancelTask reads the task again when no run is in this process, and keeps a terminal status. The handler then returns an error. --- packages/ai-mcp/src/server/create-server.ts | 10 +++- packages/ai-mcp/src/server/tasks.ts | 13 ++++- .../ai-mcp/tests/server/create-server.test.ts | 58 +++++++++++++++++++ 3 files changed, 79 insertions(+), 2 deletions(-) diff --git a/packages/ai-mcp/src/server/create-server.ts b/packages/ai-mcp/src/server/create-server.ts index 474613d00d..0c8ec5aad4 100644 --- a/packages/ai-mcp/src/server/create-server.ts +++ b/packages/ai-mcp/src/server/create-server.ts @@ -955,7 +955,15 @@ function registerLegacyTaskMethods( `Task is ${polled.record.status}, so it cannot be cancelled`, ) } - return cancelTask(polled.record, store) + const task = await cancelTask(polled.record, store) + // The task can end while cancelTask runs. Then it keeps its status. + if (task.status !== 'cancelled') { + throw new ProtocolError( + ProtocolErrorCode.InvalidParams, + `Task is ${task.status}, so it cannot be cancelled`, + ) + } + return task }, ) } diff --git a/packages/ai-mcp/src/server/tasks.ts b/packages/ai-mcp/src/server/tasks.ts index 7479efefe6..1d2a345c88 100644 --- a/packages/ai-mcp/src/server/tasks.ts +++ b/packages/ai-mcp/src/server/tasks.ts @@ -147,12 +147,23 @@ export async function getTask( * Cancels the task of `record` (a `working` task from {@link getTask}): * aborts its run, when it runs in this process, and saves it as * `cancelled`. Returns the spec 2025-11-25 `Task` for `tasks/cancel`. + * When the task got a terminal status after the caller read `record`, the + * task does not change, and the result is that terminal `Task`. */ export async function cancelTask(record: StoredTask, store: TaskStore) { // Abort first, so the run does not save a late result over the cancel. // ponytail: a run on another instance of the server is not aborted. Its // late result replaces the cancel. Add a store-backed flag if that matters. - running.get(record.taskId)?.abort() + const controller = running.get(record.taskId) + controller?.abort() + if (controller === undefined) { + // The run can end after the caller read `record`. Keep its result. + // ponytail: a re-read, not a compare-and-set. KeyedStore has no CAS. + const current = await store.get(record.taskId) + if (isStoredTask(current) && current.status !== 'working') { + return taskView(current) + } + } const cancelled: StoredTask = { ...record, status: 'cancelled', diff --git a/packages/ai-mcp/tests/server/create-server.test.ts b/packages/ai-mcp/tests/server/create-server.test.ts index b3b341d119..5f58d332a3 100644 --- a/packages/ai-mcp/tests/server/create-server.test.ts +++ b/packages/ai-mcp/tests/server/create-server.test.ts @@ -463,6 +463,64 @@ describe('createMCPServer', () => { }) }) + it('keeps the result of a task that completes while tasks/cancel reads it', async () => { + const inner = inMemoryTaskStore() + const calls: Array> = [] + let finish: (() => void) | undefined + let raceNextRead = false + const taskStore = { + ...inner, + // After it reads the `working` task, the run completes and saves + // its result. Then the stale read goes to tasks/cancel. + async get(id: string) { + const value = await inner.get(id) + if (raceNextRead) { + raceNextRead = false + finish?.() + await calls[0] + } + return value + }, + } + const server = createMCPServer({ + name: 'tasks', + version: '1.0.0', + taskStore, + waitUntil(promise) { + calls.push(promise) + }, + tools: [ + toolDefinition({ + name: 'slow', + description: 'Slow work', + execution: 'task', + }).server(async () => { + await new Promise((resolve) => { + finish = resolve + }) + return 'done' + }), + ], + }) + + await withClient(server, { era: '2025' }, async (client) => { + const { taskId } = (await taskCall(client, 'slow')).task + await vi.waitFor(() => expect(finish).toBeDefined()) + + raceNextRead = true + await expect( + client.request( + { method: 'tasks/cancel', params: { taskId } }, + taskShape, + ), + ).rejects.toThrow('completed') + expect((await taskGet(client, taskId)).status).toBe('completed') + expect(await taskResultOf(client, taskId)).toMatchObject({ + content: [{ type: 'text', text: 'done' }], + }) + }) + }) + it('runs a task tool inline for a spec 2026 request', async () => { // Spec 2026-07-28 has no tasks. The call waits for the tool output. const taskStore = inMemoryTaskStore()