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/.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/.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/.changeset/abort-computer-use-actions.md b/.changeset/abort-computer-use-actions.md new file mode 100644 index 0000000000..b5a5263f9f --- /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`, 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/.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/.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/.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-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/.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/docs/code-mode/code-mode.md b/docs/code-mode/code-mode.md index 44427cd081..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, 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. 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/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..fd377be252 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 @@ -172,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`. | @@ -185,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`. `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.` 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({ 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-code-mode/src/create-code-mode-tool.ts b/packages/ai-code-mode/src/create-code-mode-tool.ts index f709a69433..898569f0ce 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,35 @@ 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 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 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', { @@ -245,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() @@ -271,6 +303,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({ @@ -294,6 +329,12 @@ export function createCodeModeTool( ) } + // An abort that came while the context was made: do not run the code. + if (runSignal?.aborted) { + onRunAbort() + return abortedBeforeRun() + } + // Step 5: Execute the code in the sandbox const executionResult = await isolateContext.execute(strippedCode) @@ -364,12 +405,11 @@ 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 - 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 e482327eaf..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 @@ -681,3 +681,159 @@ 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) + }) + + 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('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('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) => { + 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-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-computer-use/src/actions.ts b/packages/ai-computer-use/src/actions.ts index a96b690195..84ad1bd900 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. */ +export 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) { @@ -80,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, ) { @@ -92,6 +108,7 @@ async function withKeys( for (const key of keys) { await driver.keyDown(key) pressed.push(key) + signal?.throwIfAborted() } await action() } finally { @@ -106,10 +123,18 @@ 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() + // 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 { for (const point of rest) await driver.move(point.x, point.y) @@ -168,11 +193,12 @@ async function runAnthropicAction( driver: ComputerDriver, action: string, input: Input, + signal: AbortSignal | undefined, ) { 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 @@ -185,7 +211,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 } @@ -202,8 +231,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 } @@ -218,7 +247,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 @@ -226,11 +255,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, signal, 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 +320,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) @@ -309,7 +342,7 @@ async function runOpenAIAction(driver: ComputerDriver, input: unknown) { } 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 @@ -317,14 +350,16 @@ async function runOpenAIAction(driver: ComputerDriver, input: unknown) { 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': { @@ -340,7 +375,7 @@ async function runOpenAIAction(driver: ComputerDriver, input: unknown) { 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': { @@ -348,7 +383,7 @@ async function runOpenAIAction(driver: ComputerDriver, input: unknown) { 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 @@ -357,7 +392,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 +420,25 @@ 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') + // The call can abort while its machine starts. + signal?.throwIfAborted() 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..2c5ece98b4 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', @@ -257,14 +253,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 +314,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,9 +327,13 @@ 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) + await sleep(screenshotDelay, context?.abortSignal) return imageResult(await machine.screenshot()) } return result.type === 'image' ? imageResult(result.bytes) : result.text @@ -340,7 +342,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.', @@ -350,8 +352,10 @@ 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) + 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..ca54e26c86 100644 --- a/packages/ai-computer-use/tests/actions.test.ts +++ b/packages/ai-computer-use/tests/actions.test.ts @@ -502,3 +502,127 @@ 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' | 'mouseDown', + 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('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, + ) + expect(calls).toEqual([ + ['move', 0, 0], + ['mouseDown', 'left'], + ['move', 1, 1], + ['mouseUp', 'left'], + ]) + }) + + it('finishes a drag that aborts between two moves', async () => { + const { driver, calls, signal } = abortingDuring('move', 2) + await runComputerAction( + driver, + { + action: { + type: 'drag', + path: [ + { x: 0, y: 0 }, + { x: 1, y: 1 }, + { x: 2, y: 2 }, + ], + }, + }, + signal, + ) + expect(calls).toEqual([ + ['move', 0, 0], + ['mouseDown', 'left'], + ['move', 1, 1], + ['move', 2, 2], + ['mouseUp', 'left'], + ]) + }) +}) 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..f79c3e2ffc 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,112 @@ 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) + }) + + 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('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 }) + 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-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() + } + }) +}) diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index 56873d9278..9662f90c29 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -95,18 +95,31 @@ 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) } - 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, { @@ -426,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', @@ -434,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) @@ -459,7 +477,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.') @@ -480,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/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..c88ed8aa02 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -573,3 +573,126 @@ 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) => { + // 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) + }, + } + 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([]) + }) + + 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([]) + }) + + 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([]) + }, + ) + + 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() + }) +}) 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 = [] 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..0c8ec5aad4 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,35 @@ 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`, + ) + } + 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 + }, + ) } function taskIdFrom(params: unknown) { diff --git a/packages/ai-mcp/src/server/tasks.ts b/packages/ai-mcp/src/server/tasks.ts index 0c62e7a278..1d2a345c88 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,36 @@ 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`. + * 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. + 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', + lastUpdatedAt: new Date().toISOString(), + } + await store.set(record.taskId, cancelled) + return taskView(cancelled) +} + function taskView(record: StoredTask) { const task = { taskId: record.taskId, @@ -186,6 +228,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..5f58d332a3 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,119 @@ 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('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() diff --git a/packages/ai-sandbox-daytona/src/provider.ts b/packages/ai-sandbox-daytona/src/provider.ts index 1c13ed3893..dfe9ab414d 100644 --- a/packages/ai-sandbox-daytona/src/provider.ts +++ b/packages/ai-sandbox-daytona/src/provider.ts @@ -97,8 +97,19 @@ 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. + 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 }) } @@ -164,7 +175,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 +187,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..f45736173d 100644 --- a/packages/ai-sandbox-daytona/tests/provider.test.ts +++ b/packages/ai-sandbox-daytona/tests/provider.test.ts @@ -413,3 +413,46 @@ 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) + }, + ) + + 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 fb4be83951..e25fc1d53f 100644 --- a/packages/ai-sandbox-vercel/src/provider.ts +++ b/packages/ai-sandbox-vercel/src/provider.ts @@ -129,6 +129,15 @@ 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. + 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`). @@ -142,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 b2a533cd34..f160777e17 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,49 @@ 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() + } + }) + + 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/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) diff --git a/testing/e2e/tests/computer-use.spec.ts b/testing/e2e/tests/computer-use.spec.ts index 5d4cba64ba..d440536fc2 100644 --- a/testing/e2e/tests/computer-use.spec.ts +++ b/testing/e2e/tests/computer-use.spec.ts @@ -358,6 +358,58 @@ 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 (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') + }) + test('anthropic computer toolset runs each call on the page', async ({ page, }) => {