Improve in-flight cancel latency via run abort signal propagation

src/tools/executor.ts

@@ -24,6 +24,10 @@ export interface ToolExecutionObserverEvent {
   timestampSeconds: number;
 }
 
+export interface ToolExecuteOptions {
+  signal?: AbortSignal;
+}
+
 export class ToolExecutor {
   private registry: ToolRegistry;
   private hooks: HookEngine;
@@ -64,7 +68,12 @@ export class ToolExecutor {
     return base;
   }
 
-  async execute(toolName: string, args: unknown, context?: ToolPolicyContext): Promise<ToolResult> {
+  async execute(
+    toolName: string,
+    args: unknown,
+    context?: ToolPolicyContext,
+    options?: ToolExecuteOptions,
+  ): Promise<ToolResult> {
     const executionId = randomUUID();
     const executionEnvironment = this.resolveEffectiveExecutionEnvironment(toolName, context);
     const skillName = context?.skillName;
@@ -279,31 +288,56 @@ export class ToolExecutor {
     });
 
     let timeoutHandle: NodeJS.Timeout | undefined;
-    const abortController = new AbortController();
+    const timeoutAbortController = new AbortController();
+    const externalSignal = options?.signal;
+    const combinedSignal = externalSignal
+      ? AbortSignal.any([externalSignal, timeoutAbortController.signal])
+      : timeoutAbortController.signal;
+    let externalAbortCleanup: (() => void) | undefined;
+
     try {
+      const externalAbortPromise = externalSignal
+        ? new Promise<ToolResult>((_, reject) => {
+          if (externalSignal.aborted) {
+            const error = new Error('Operation cancelled by user.');
+            error.name = 'AbortError';
+            reject(error);
+            return;
+          }
+          const onAbort = () => {
+            const error = new Error('Operation cancelled by user.');
+            error.name = 'AbortError';
+            reject(error);
+          };
+          externalSignal.addEventListener('abort', onAbort, { once: true });
+          externalAbortCleanup = () => externalSignal.removeEventListener('abort', onAbort);
+        })
+        : null;
+
       const result = await Promise.race([
         (async () => {
           if (executionEnvironment === 'sandbox' && this.sandboxManager) {
             const sandboxSessionId = context?.sessionId ?? `${context?.channel ?? 'unknown'}:${context?.sender ?? 'unknown'}`;
             const sandbox = await this.sandboxManager.getOrCreate(sandboxSessionId);
             if (toolName === 'shell.exec') {
-              return createSandboxedShellTool(sandbox).execute(args, { signal: abortController.signal });
+              return createSandboxedShellTool(sandbox).execute(args, { signal: combinedSignal });
             }
             if (toolName === 'process.start') {
-              return createSandboxedProcessStartTool(sandbox).execute(args, { signal: abortController.signal });
+              return createSandboxedProcessStartTool(sandbox).execute(args, { signal: combinedSignal });
             }
           }
-          return tool.execute(args, { signal: abortController.signal });
+          return tool.execute(args, { signal: combinedSignal });
         })(),
         new Promise<ToolResult>((_, reject) => {
           timeoutHandle = setTimeout(
             () => {
-              abortController.abort();
+              timeoutAbortController.abort();
               reject(new Error(`Tool '${toolName}' timed out after ${this.defaultTimeoutMs}ms`));
             },
             this.defaultTimeoutMs,
           );
         }),
+        ...(externalAbortPromise ? [externalAbortPromise] : []),
       ]);
 
       const duration = Date.now() - startTime;
@@ -357,6 +391,10 @@ export class ToolExecutor {
         timestampSeconds: Math.floor(Date.now() / 1000),
       });
 
+      if (externalSignal?.aborted && this.isAbortError(error)) {
+        throw error;
+      }
+
       return {
         success: false,
         output: '',
@@ -366,9 +404,14 @@ export class ToolExecutor {
       if (timeoutHandle) {
         clearTimeout(timeoutHandle);
       }
+      externalAbortCleanup?.();
     }
   }
 
+  private isAbortError(error: unknown): boolean {
+    return error instanceof Error && error.name === 'AbortError';
+  }
+
   private notifyExecutionObserver(event: ToolExecutionObserverEvent): void {
     if (!this.executionObserver) {
       return;