diff --git a/nodejs/README.md b/nodejs/README.md index bc91cd793..65d3f6909 100644 --- a/nodejs/README.md +++ b/nodejs/README.md @@ -280,7 +280,11 @@ unsubscribe(); ##### `abort(): Promise` -Abort the currently processing message in this session. +Abort the currently processing message in this session. This also aborts the `AbortSignal` passed to any in-flight tool handlers (see [Cancelling Tool Handlers](#cancelling-tool-handlers)). + +##### `cancelToolCall(toolCallId: string): boolean` + +Cooperatively cancel a single in-flight tool handler by aborting the `AbortSignal` on its `ToolInvocation`, without aborting the broader agentic loop. Returns `true` if a matching in-flight tool call was found and signaled, `false` otherwise. ##### `getEvents(): Promise` @@ -503,6 +507,35 @@ defineTool("lookup_issue", { }); ``` +#### Cancelling Tool Handlers + +Long-running tool handlers can opt in to cooperative cancellation. Each handler's `ToolInvocation` carries a standard [`AbortSignal`](https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal) that aborts when `session.abort()` (which cancels the whole agentic loop) or `session.cancelToolCall(toolCallId)` (which cancels a single in-flight handler) is invoked. Forward it to abortable APIs or check `signal.aborted`: + +```ts +defineTool("fetch_data", { + description: "Fetch a large payload", + parameters: z.object({ url: z.string() }), + handler: async ({ url }, { signal }) => { + // The fetch is aborted automatically when the session/tool is cancelled + const res = await fetch(url, { signal }); + return await res.text(); + }, +}); +``` + +Cancel a specific in-flight handler without aborting the rest of the turn: + +```ts +session.on("tool.execution_start", (event) => { + setTimeout(() => { + // Returns true if a matching in-flight handler was signaled + session.cancelToolCall(event.data.toolCallId); + }, 5000); +}); +``` + +Handlers that ignore the signal continue to run to completion, so existing handlers keep working unchanged. + ### Commands Register slash commands so that users of the CLI's TUI can invoke custom actions via `/commandName`. Each command has a `name`, optional `description`, and a `handler` called when the user executes it. diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index 8ae19755a..96ee22fbe 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -121,6 +121,7 @@ export class CopilotSession { private typedEventHandlers: Map void>> = new Map(); private toolHandlers: Map = new Map(); + private inFlightToolCalls: Map = new Map(); private canvases: Map = new Map(); private commandHandlers: Map = new Map(); private permissionHandler?: PermissionHandler; @@ -563,12 +564,15 @@ export class CopilotSession { traceparent?: string, tracestate?: string ): Promise { + const abortController = new AbortController(); + this.inFlightToolCalls.set(toolCallId, abortController); try { const rawResult = await handler(args, { sessionId: this.sessionId, toolCallId, toolName, arguments: args, + signal: abortController.signal, traceparent, tracestate, }); @@ -593,6 +597,12 @@ export class CopilotSession { } // Connection lost or RPC error — nothing we can do } + } finally { + // Only clear if this is still the controller for this toolCallId; + // guards against a recycled toolCallId from a later invocation. + if (this.inFlightToolCalls.get(toolCallId) === abortController) { + this.inFlightToolCalls.delete(toolCallId); + } } } @@ -1170,6 +1180,9 @@ export class CopilotSession { * ``` */ async disconnect(): Promise { + // Abort any in-flight tool handlers so they can release resources. + this._abortInFlightToolCalls(); + this.inFlightToolCalls.clear(); await this.connection.sendRequest("session.destroy", { sessionId: this.sessionId, }); @@ -1209,11 +1222,70 @@ export class CopilotSession { * ``` */ async abort(): Promise { + // Cooperatively cancel any in-flight tool handlers that opted in to the + // AbortSignal exposed on their ToolInvocation. Handlers that ignore the + // signal continue to run to completion. + this._abortInFlightToolCalls(); await this.connection.sendRequest("session.abort", { sessionId: this.sessionId, }); } + /** + * Cooperatively cancels a single in-flight tool handler by aborting the + * `AbortSignal` on its `ToolInvocation`, without aborting the broader + * agentic loop. + * + * This only affects handlers that opted in to the signal (e.g. by passing + * it to `fetch`, `child_process.spawn`, or checking `signal.aborted`). + * Handlers that ignore the signal continue to run to completion. + * + * @param toolCallId - The `toolCallId` of the in-flight tool invocation to cancel + * @returns `true` if a matching in-flight tool call was found and signaled, `false` otherwise + * + * @example + * ```typescript + * const session = await client.createSession({ + * tools: [ + * defineTool("fetch_data", { + * handler: async (args, { signal }) => { + * const res = await fetch(args.url, { signal }); + * return await res.text(); + * }, + * }), + * ], + * }); + * + * session.on((event) => { + * if (event.type === "tool.execution_start") { + * // Cancel a specific tool call after a deadline + * setTimeout(() => session.cancelToolCall(event.data.toolCallId), 5000); + * } + * }); + * ``` + */ + cancelToolCall(toolCallId: string): boolean { + const controller = this.inFlightToolCalls.get(toolCallId); + if (!controller) { + return false; + } + // Remove the entry up front so a subsequent cancelToolCall (or the + // handler's own cleanup) for the same id is a no-op and returns false. + this.inFlightToolCalls.delete(toolCallId); + controller.abort(); + return true; + } + + /** + * Aborts the AbortSignal for every in-flight tool handler. + * @internal + */ + private _abortInFlightToolCalls(): void { + for (const controller of this.inFlightToolCalls.values()) { + controller.abort(); + } + } + /** * Change the model for this session. * The new model takes effect for the next message. Conversation history is preserved. diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts index bad1c33ad..ae92f268e 100644 --- a/nodejs/src/types.ts +++ b/nodejs/src/types.ts @@ -475,6 +475,15 @@ export interface ToolInvocation { toolCallId: string; toolName: string; arguments: unknown; + /** + * An `AbortSignal` that aborts when `session.abort()` or + * `session.cancelToolCall(toolCallId)` is invoked while this handler is + * in flight. Handlers may opt in to cooperative cancellation by forwarding + * it to abortable APIs (`fetch(url, { signal })`, `child_process.spawn`, + * etc.) or by checking `signal.aborted`. Handlers that ignore it continue + * to run to completion, preserving existing behavior. + */ + signal: AbortSignal; /** W3C Trace Context traceparent from the CLI's execute_tool span. */ traceparent?: string; /** W3C Trace Context tracestate from the CLI's execute_tool span. */ diff --git a/nodejs/test/e2e/abort.e2e.test.ts b/nodejs/test/e2e/abort.e2e.test.ts index 89877387c..803c1ec67 100644 --- a/nodejs/test/e2e/abort.e2e.test.ts +++ b/nodejs/test/e2e/abort.e2e.test.ts @@ -99,6 +99,11 @@ describe("Abort", async () => { releaseToolResolve = resolve; }); + let signalAbortedResolve!: () => void; + const signalAborted = new Promise((resolve) => { + signalAbortedResolve = resolve; + }); + const session = await client.createSession({ onPermissionRequest: approveAll, tools: [ @@ -107,8 +112,15 @@ describe("Abort", async () => { parameters: z.object({ value: z.string().describe("Value to analyze"), }), - handler: async ({ value }) => { + handler: async ({ value }, { signal }) => { toolStartedResolve(value); + if (signal.aborted) { + signalAbortedResolve(); + } else { + signal.addEventListener("abort", () => signalAbortedResolve(), { + once: true, + }); + } return await releaseTool; }, }), @@ -127,6 +139,9 @@ describe("Abort", async () => { // Abort while the tool is running await session.abort(); + // The handler's AbortSignal should fire as a result of session.abort() + await withTimeout(signalAborted, 10_000, "tool handler AbortSignal"); + // Release the tool so its task doesn't leak releaseToolResolve("RELEASED_AFTER_ABORT"); @@ -153,4 +168,75 @@ describe("Abort", async () => { await session.disconnect(); }); + + it( + "should cancel a single tool call via cancelToolCall", + { timeout: TEST_TIMEOUT_MS }, + async () => { + let toolCallIdResolve!: (value: string) => void; + const toolCallIdReady = new Promise((resolve) => { + toolCallIdResolve = resolve; + }); + + let releaseToolResolve!: (value: string) => void; + const releaseTool = new Promise((resolve) => { + releaseToolResolve = resolve; + }); + + let signalAbortedResolve!: () => void; + const signalAborted = new Promise((resolve) => { + signalAbortedResolve = resolve; + }); + + const session = await client.createSession({ + onPermissionRequest: approveAll, + tools: [ + defineTool("slow_analysis", { + description: "A slow analysis tool that blocks until released", + parameters: z.object({ + value: z.string().describe("Value to analyze"), + }), + handler: async ({ value: _value }, { signal, toolCallId }) => { + toolCallIdResolve(toolCallId); + if (signal.aborted) { + signalAbortedResolve(); + } else { + signal.addEventListener("abort", () => signalAbortedResolve(), { + once: true, + }); + } + return await releaseTool; + }, + }), + ], + }); + + // Fire-and-forget + void session.send({ + prompt: "Use slow_analysis with value 'test_cancel'. Wait for the result.", + }); + + // Wait for the tool to start executing and capture its toolCallId + const toolCallId = await withTimeout( + toolCallIdReady, + 60_000, + "slow_analysis toolCallId" + ); + + // Unknown toolCallIds return false + expect(session.cancelToolCall("nonexistent-tool-call-id")).toBe(false); + + // Cancelling the in-flight tool call returns true and fires its signal + expect(session.cancelToolCall(toolCallId)).toBe(true); + await withTimeout(signalAborted, 10_000, "tool handler AbortSignal via cancelToolCall"); + + // A second cancel of the same (now-removed) call returns false + expect(session.cancelToolCall(toolCallId)).toBe(false); + + // Release the tool so its task doesn't leak + releaseToolResolve("RELEASED_AFTER_CANCEL"); + + await session.disconnect(); + } + ); }); diff --git a/test/snapshots/abort/should_cancel_a_single_tool_call_via_canceltoolcall.yaml b/test/snapshots/abort/should_cancel_a_single_tool_call_via_canceltoolcall.yaml new file mode 100644 index 000000000..6af99ad56 --- /dev/null +++ b/test/snapshots/abort/should_cancel_a_single_tool_call_via_canceltoolcall.yaml @@ -0,0 +1,24 @@ +models: + - claude-sonnet-4.5 +conversations: + - messages: + - role: system + content: ${system} + - role: user + content: Use slow_analysis with value 'test_cancel'. Wait for the result. + - role: assistant + content: I'll call the slow_analysis tool with the value 'test_cancel' and wait for it to complete. + - role: assistant + tool_calls: + - id: toolcall_0 + type: function + function: + name: report_intent + arguments: '{"intent":"Running slow analysis test"}' + - role: assistant + tool_calls: + - id: toolcall_1 + type: function + function: + name: slow_analysis + arguments: '{"value":"test_cancel"}'