Improve in-flight cancel latency via run abort signal propagation

src/backends/native/agent.test.ts

@@ -24,10 +24,11 @@ describe('NativeAgent', () => {
     const response = await agent.process('Hi');
 
     expect(response).toBe('Hello!');
-    expect(mockClient.chat).toHaveBeenCalledWith({
+    expect(mockClient.chat).toHaveBeenCalledWith(expect.objectContaining({
       messages: [{ role: 'user', content: 'Hi' }],
       system: 'You are helpful.',
-    });
+      signal: expect.any(AbortSignal),
+    }));
 
     const history = agent.getHistory();
     expect(history).toHaveLength(2);

src/backends/native/agent.ts

@@ -83,6 +83,7 @@ export class NativeAgent {
   private _lastToolFingerprint?: string;
   private _cancelRequested = false;
   private _runInProgress = false;
+  private _runAbortController?: AbortController;
   private modelTimeoutMs: number;
 
   constructor(config: NativeAgentConfig) {
@@ -106,6 +107,7 @@ export class NativeAgent {
 
   async process(userMessage: string, attachments?: Attachment[]): Promise<string> {
     this._cancelRequested = false;
+    this._runAbortController = new AbortController();
     if ('clearAbort' in this.modelClient && typeof this.modelClient.clearAbort === 'function') {
       this.modelClient.clearAbort();
     }
@@ -144,6 +146,7 @@ export class NativeAgent {
     } finally {
       this._runInProgress = false;
       this._cancelRequested = false;
+      this._runAbortController = undefined;
     }
   }
 
@@ -353,7 +356,9 @@ export class NativeAgent {
             }
             : undefined;
 
-          const result = await toolExecutor.execute(internalName, tc.args, perCallContext);
+          const result = await toolExecutor.execute(internalName, tc.args, perCallContext, {
+            signal: this._runAbortController?.signal,
+          });
 
           this.onToolUse?.({ type: 'end', tool: internalName, result });
 
@@ -426,11 +431,22 @@ export class NativeAgent {
   }
 
   private async chatWithRouter(request: ChatRequest): Promise<ChatResponse> {
+    const runSignal = this._runAbortController?.signal;
+    const requestSignal = request.signal;
+    const signal = runSignal && requestSignal
+      ? AbortSignal.any([runSignal, requestSignal])
+      : (runSignal ?? requestSignal);
+
+    const requestWithSignal = signal
+      ? { ...request, signal }
+      : request;
+
     const requestPromise = 'getClient' in this.modelClient
-      ? (this.modelClient as ModelRouter).chat(request, this.currentTier)
-      : this.modelClient.chat(request);
+      ? (this.modelClient as ModelRouter).chat(requestWithSignal, this.currentTier)
+      : this.modelClient.chat(requestWithSignal);
 
     let timer: NodeJS.Timeout | undefined;
+    let abortCleanup: (() => void) | undefined;
     const timeoutPromise = new Promise<never>((_, reject) => {
       timer = setTimeout(() => {
         const error = new Error(`Model request timed out after ${this.modelTimeoutMs}ms`);
@@ -439,13 +455,31 @@ export class NativeAgent {
       }, this.modelTimeoutMs);
       timer.unref?.();
     });
+    const abortPromise = signal
+      ? new Promise<never>((_, reject) => {
+        if (signal.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);
+        };
+        signal.addEventListener('abort', onAbort, { once: true });
+        abortCleanup = () => signal.removeEventListener('abort', onAbort);
+      })
+      : null;
 
     try {
-      return await Promise.race([requestPromise, timeoutPromise]);
+      return await Promise.race([requestPromise, timeoutPromise, ...(abortPromise ? [abortPromise] : [])]);
     } finally {
       if (timer) {
         clearTimeout(timer);
       }
+      abortCleanup?.();
     }
   }
 
@@ -544,6 +578,7 @@ export class NativeAgent {
   cancel(): void {
     if (this._runInProgress) {
       this._cancelRequested = true;
+      this._runAbortController?.abort();
       if ('requestAbort' in this.modelClient && typeof this.modelClient.requestAbort === 'function') {
         this.modelClient.requestAbort();
       }
@@ -555,7 +590,7 @@ export class NativeAgent {
   }
 
   private throwIfCancelled(): void {
-    if (!this._cancelRequested) {
+    if (!this._cancelRequested && !this._runAbortController?.signal.aborted) {
       return;
     }