[AI Assistance] Cleanly handle abort in case of paused side effect

Bug: 380394373
Change-Id: I122f644d31746bf4a2f4c7951d39d8075b761be6
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6099097
Auto-Submit: Samiya Caur <[email protected]>
Commit-Queue: Alex Rudenko <[email protected]>
Reviewed-by: Alex Rudenko <[email protected]>

Author: Samiya Caur

front_end/panels/ai_assistance/agents/AiAgent.ts

@@ -339,7 +339,8 @@ export abstract class AiAgent<T> {
     return request;
   }
 
-  handleAction(action: string): AsyncGenerator<SideEffectResponse, ActionResponse, void>;
+  handleAction(action: string, options?: {signal?: AbortSignal}):
+      AsyncGenerator<SideEffectResponse, ActionResponse, void>;
   handleAction(): never {
     throw new Error('Unexpected action found');
   }
@@ -510,10 +511,7 @@ STOP`;
         debugLog('Error calling the AIDA API', err);
 
         if (err instanceof Host.AidaClient.AidaAbortError) {
-          const response = {
-            type: ResponseType.ERROR,
-            error: ErrorType.ABORT,
-          } as const;
+          const response = this.#createAbortResponse();
           this.#addHistory(response);
           yield response;
           break;
@@ -597,7 +595,13 @@ STOP`;
       }
 
       if (action) {
-        const result = yield* this.handleAction(action);
+        const result = yield* this.handleAction(action, {signal: options.signal});
+        if (options?.signal?.aborted) {
+          const response = this.#createAbortResponse();
+          this.#addHistory(response);
+          yield response;
+          break;
+        }
         this.#addHistory(result);
         query = `OBSERVATION: ${result.output}`;
         // Capture history state for the next iteration query.
@@ -631,6 +635,13 @@ STOP`;
       yield entry;
     }
   }
+
+  #createAbortResponse(): ResponseData {
+    return {
+      type: ResponseType.ERROR,
+      error: ErrorType.ABORT,
+    };
+  }
 }
 
 export function isDebugMode(): boolean {

front_end/panels/ai_assistance/agents/StylingAgent.test.ts

@@ -11,7 +11,7 @@ import {
 } from '../../../testing/EnvironmentHelpers.js';
 import * as AiAssistance from '../ai_assistance.js';
 
-const {StylingAgent} = AiAssistance;
+const {StylingAgent, ErrorType} = AiAssistance;
 
 describeWithEnvironment('StylingAgent', () => {
   function mockHostConfig(
@@ -693,7 +693,7 @@ c`;
 
           count++;
         }
-        const execJs = sinon.mock().twice();
+        const execJs = sinon.mock().once();
         execJs.onCall(0).throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
         const agent = new StylingAgent({
           aidaClient: mockAidaClient(generateActionAndAnswer),
@@ -710,6 +710,42 @@ c`;
         assert.strictEqual(actionStep.output, 'Error: User denied code execution with side effects.');
         assert.strictEqual(execJs.getCalls().length, 1);
       });
+
+      it('returns error when side effect is aborted', async () => {
+        const promise = Promise.withResolvers();
+        const stub = sinon.stub().returns(promise);
+        async function* generateAction() {
+          yield {
+            explanation: `ACTION
+            $0.style.backgroundColor = 'red'
+            STOP`,
+            metadata: {},
+            completed: true,
+          };
+        }
+        const execJs = sinon.mock().once().throws(
+            new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        const agent = new StylingAgent({
+          aidaClient: mockAidaClient(generateAction),
+          createExtensionScope,
+          confirmSideEffectForTest: stub,
+          execJs,
+        });
+        const controller = new AbortController();
+
+        const agentPromise = Array.fromAsync(
+            agent.run('test', {selected: new AiAssistance.NodeContext(element), signal: controller.signal}));
+        await stub.calledOnce;
+        const awaitTimeout = (delay: number) => new Promise(resolve => setTimeout(resolve, delay));
+        await awaitTimeout(10).then(() => controller.abort());
+        const responses = await agentPromise;
+
+        const errorStep = responses.at(-1) as AiAssistance.ErrorResponse;
+        assert.strictEqual(execJs.getCalls().length, 1);
+        assert.exists(errorStep);
+        assert.strictEqual(errorStep.error, ErrorType.ABORT);
+        await promise.promise.then(value => assert.strictEqual(value, false));
+      });
     });
 
     describe('long `Observation` text handling', () => {

front_end/panels/ai_assistance/agents/StylingAgent.ts

@@ -582,7 +582,9 @@ export class StylingAgent extends AiAgent<SDK.DOMModel.DOMNode> {
     return output.trim();
   }
 
-  override async * handleAction(action: string): AsyncGenerator<SideEffectResponse, ActionResponse, void> {
+  override async *
+      handleAction(action: string, options?: {signal?: AbortSignal}):
+          AsyncGenerator<SideEffectResponse, ActionResponse, void> {
     debugLog(`Action to execute: ${action}`);
     if (this.executionMode === Root.Runtime.HostConfigFreestylerExecutionMode.NO_SCRIPTS) {
       return {
@@ -621,6 +623,10 @@ export class StylingAgent extends AiAgent<SDK.DOMModel.DOMNode> {
 
         const sideEffectConfirmationPromiseWithResolvers = this.#confirmSideEffect<boolean>();
 
+        options?.signal?.addEventListener('abort', () => {
+          sideEffectConfirmationPromiseWithResolvers.resolve(false);
+        }, {once: true});
+
         yield {
           type: ResponseType.SIDE_EFFECT,
           code: action,