fix(a2a-server): ensure task cancellation aborts execution loop (#28316)

2026-07-14 03:50:49 -07:00 · 2026-07-13 17:22:39 +00:00
parent 42ee2b74c7
commit f5e58ff18b
3 changed files with 515 additions and 169 deletions
@@ -14,6 +14,12 @@ import type {
 import { EventEmitter } from 'node:events';
 import { requestStorage } from '../http/requestStorage.js';

+vi.mock('../utils/path_utils.js', () => ({
+  validateWorkspacePath: vi
+    .fn()
+    .mockImplementation(async (path?: string) => path || process.cwd()),
+}));
+
 // Mocks for constructor dependencies
 vi.mock('../config/config.js', () => ({
  loadConfig: vi.fn().mockReturnValue({
@@ -300,4 +306,125 @@ describe('CoderAgentExecutor', () => {
      true,
    );
  });
+
+  it('cancelTask should abort the active execution loop', async () => {
+    const abortSpy = vi.spyOn(AbortController.prototype, 'abort');
+    const taskId = 'test-task-to-cancel';
+    const contextId = 'test-context';
+
+    const mockSocket = new EventEmitter();
+    (requestStorage.getStore as Mock).mockReturnValue({
+      req: { socket: mockSocket },
+    });
+
+    const requestContext = {
+      userMessage: {
+        messageId: 'msg-1',
+        taskId,
+        contextId,
+        parts: [{ kind: 'text', text: 'a long running prompt' }],
+        metadata: {
+          coderAgent: { kind: 'agent-settings', workspacePath: '/tmp' },
+        },
+      },
+    } as unknown as RequestContext;
+
+    // Don't await this, let it run in the background.
+    let primaryError: Error | null = null;
+    const primaryPromise = executor.execute(requestContext, mockEventBus);
+    primaryPromise.catch((err) => {
+      primaryError = err as Error;
+    });
+
+    // Poll until the task is registered in the executor to avoid flaky timeouts in slow CI environments.
+    let attempts = 0;
+    while (!executor.getTask(taskId)) {
+      if (primaryError) {
+        throw new Error(`Primary execution failed early: ${primaryError}`);
+      }
+      if (attempts++ > 100) {
+        // 100 * 5ms = 500ms timeout
+        throw new Error('Timed out waiting for task to be registered');
+      }
+      await new Promise((resolve) => setTimeout(resolve, 5));
+    }
+
+    const wrapper = executor.getTask(taskId);
+    expect(wrapper).toBeDefined();
+    const setTaskStateSpy = vi
+      .spyOn(wrapper!.task, 'setTaskStateAndPublishUpdate')
+      .mockImplementation((newState) => {
+        // Make the mock realistic: actually update the state when called.
+        wrapper!.task.taskState = newState;
+      });
+
+    // Now, cancel the task.
+    await executor.cancelTask(taskId, mockEventBus);
+
+    // Verify that the abort method on the controller was called and state was updated.
+    expect(abortSpy).toHaveBeenCalledOnce();
+    expect(setTaskStateSpy).toHaveBeenCalledWith(
+      'canceled',
+      expect.any(Object),
+      'Task canceled by user request.',
+      undefined,
+      true,
+    );
+
+    // Clean up the test by allowing the promise to resolve.
+    // The abort call should have unblocked the acceptUserMessage generator.
+    await primaryPromise;
+
+    // Verify task is evicted from cache
+    expect(executor.getTask(taskId)).toBeUndefined();
+
+    abortSpy.mockRestore();
+  });
+
+  it('cancelTask should explicitly save task state to TaskStore and evict task during active aborts', async () => {
+    const taskId = 'test-task-active-abort-save';
+    const contextId = 'test-context';
+
+    const mockSocket = new EventEmitter();
+    (requestStorage.getStore as Mock).mockReturnValue({
+      req: { socket: mockSocket },
+    });
+
+    const requestContext = {
+      userMessage: {
+        messageId: 'msg-1',
+        taskId,
+        contextId,
+        parts: [{ kind: 'text', text: 'a long running prompt' }],
+        metadata: {
+          coderAgent: { kind: 'agent-settings', workspacePath: '/tmp' },
+        },
+      },
+    } as unknown as RequestContext;
+
+    const primaryPromise = executor.execute(requestContext, mockEventBus);
+
+    // Wait for task to be registered
+    let attempts = 0;
+    while (!executor.getTask(taskId)) {
+      if (attempts++ > 100) {
+        throw new Error('Timed out waiting for task to be registered');
+      }
+      await new Promise((resolve) => setTimeout(resolve, 5));
+    }
+
+    const wrapper = executor.getTask(taskId)!;
+    const saveSpy = vi.spyOn(mockTaskStore, 'save');
+
+    // Now, cancel the task.
+    await executor.cancelTask(taskId, mockEventBus);
+
+    // Verify that the task state was saved to TaskStore during cancelTask
+    expect(saveSpy).toHaveBeenCalled();
+    expect(wrapper.task.dispose).toHaveBeenCalled();
+    expect(executor.getTask(taskId)).toBeUndefined();
+
+    // Clean up the test by allowing the promise to resolve.
+    await primaryPromise;
+  });
 });
@@ -4,7 +4,7 @@
 * SPDX-License-Identifier: Apache-2.0
 */

-import type { Message, Task as SDKTask } from '@a2a-js/sdk';
+import type { Task as SDKTask } from '@a2a-js/sdk';
 import type {
  TaskStore,
  AgentExecutor,
@@ -42,6 +42,7 @@ import { loadExtensions } from '../config/extension.js';
 import { Task } from './task.js';
 import { requestStorage } from '../http/requestStorage.js';
 import { pushTaskStateFailed } from '../utils/executor_utils.js';
+import { validateWorkspacePath } from '../utils/path_utils.js';

 /**
 * Provides a wrapper for Task. Passes data from Task to SDKTask.
@@ -90,6 +91,12 @@ export class CoderAgentExecutor implements AgentExecutor {
  private tasks: Map<string, TaskWrapper> = new Map();
  // Track tasks with an active execution loop.
  private executingTasks = new Set<string>();
+  private activeAbortControllers = new Map<string, Set<AbortController>>();
+  // Track tasks currently initializing to prevent race conditions.
+  private initializingTasks = new Set<string>();
+  private initializationPromises = new Map<string, Promise<TaskWrapper>>();
+  // Track explicitly canceled task IDs to handle cancellation during initialization.
+  private explicitlyCanceledTasks = new Set<string>();

  constructor(private taskStore?: TaskStore) {}

@@ -126,7 +133,30 @@ export class CoderAgentExecutor implements AgentExecutor {
      );
    }

-    const agentSettings = persistedState._agentSettings;
+    let agentSettings;
+    try {
+      agentSettings = {
+        ...(persistedState._agentSettings ?? {}),
+        workspacePath: await validateWorkspacePath(
+          persistedState._agentSettings?.workspacePath,
+        ),
+        isTrusted: false,
+      };
+    } catch (error) {
+      logger.error(
+        `[CoderAgentExecutor] Invalid workspace path in persisted state for task ${sdkTask.id}:`,
+        error,
+      );
+      if (eventBus) {
+        void pushTaskStateFailed(
+          error,
+          eventBus,
+          sdkTask.id,
+          sdkTask.contextId,
+        );
+      }
+      throw error; // Re-throw to be caught by caller
+    }
    const config = await this.getConfig(agentSettings, sdkTask.id);
    const contextId: string =
      getContextIdFromMetadata(metadata) || sdkTask.contextId;
@@ -180,6 +210,17 @@ export class CoderAgentExecutor implements AgentExecutor {
    return Array.from(this.tasks.values());
  }

+  private cleanupAndEvictTask(taskId: string) {
+    const wrapper = this.tasks.get(taskId);
+    if (wrapper) {
+      logger.info(
+        `[CoderAgentExecutor] Task ${taskId} reached terminal state ${wrapper.task.taskState}. Evicting and disposing.`,
+      );
+      wrapper.task.dispose();
+      this.tasks.delete(taskId);
+    }
+  }
+
  cancelTask = async (
    taskId: string,
    eventBus: ExecutionEventBus,
@@ -187,6 +228,51 @@ export class CoderAgentExecutor implements AgentExecutor {
    logger.info(
      `[CoderAgentExecutor] Received cancel request for task ${taskId}`,
    );
+
+    const abortControllers = this.activeAbortControllers.get(taskId);
+    if (abortControllers && abortControllers.size > 0) {
+      this.explicitlyCanceledTasks.add(taskId);
+      logger.info(
+        `[CoderAgentExecutor] Aborting ${abortControllers.size} active execution loop(s) for task ${taskId}.`,
+      );
+      // Abort first to ensure loops are stopped.
+      for (const controller of Array.from(abortControllers)) {
+        controller.abort();
+      }
+
+      // Then, attempt to update state and persist.
+      const wrapper = this.tasks.get(taskId);
+      if (wrapper) {
+        const { task } = wrapper;
+        task.cancelPendingTools('Task canceled by user request.');
+        task.setTaskStateAndPublishUpdate(
+          'canceled',
+          { kind: CoderAgentEvent.StateChangeEvent },
+          'Task canceled by user request.',
+          undefined,
+          true,
+        );
+        try {
+          await this.taskStore?.save(wrapper.toSDKTask());
+          logger.info(
+            `[CoderAgentExecutor] Task ${taskId} state CANCELED saved during active abort.`,
+          );
+        } catch (saveError) {
+          logger.error(
+            `[CoderAgentExecutor] Failed to save task ${taskId} state during active abort:`,
+            saveError,
+          );
+        }
+        this.cleanupAndEvictTask(taskId);
+      }
+      return;
+    }
+
+    // If there is no active execution loop, the task is idle.
+    // We can clean it up directly.
+    logger.info(
+      `[CoderAgentExecutor] No active execution for task ${taskId}. Cleaning up directly.`,
+    );
    const wrapper = this.tasks.get(taskId);

    if (!wrapper) {
@@ -244,7 +330,7 @@ export class CoderAgentExecutor implements AgentExecutor {

    try {
      logger.info(
-        `[CoderAgentExecutor] Initiating cancellation for task ${taskId}.`,
+        `[CoderAgentExecutor] Initiating cancellation for idle task ${taskId}.`,
      );
      task.cancelPendingTools('Task canceled by user request.');

@@ -265,8 +351,7 @@ export class CoderAgentExecutor implements AgentExecutor {
      logger.info(`[CoderAgentExecutor] Task ${taskId} state CANCELED saved.`);

      // Cleanup listener subscriptions to avoid memory leaks.
-      wrapper.task.dispose();
-      this.tasks.delete(taskId);
+      this.cleanupAndEvictTask(taskId);
    } catch (error) {
      const errorMessage =
        error instanceof Error ? error.message : 'Unknown error';
@@ -332,175 +417,240 @@ export class CoderAgentExecutor implements AgentExecutor {
    const abortController = new AbortController();
    const abortSignal = abortController.signal;

-    if (store) {
-      // Grab the raw socket from the request object
-      const socket = store.req.socket;
-      const onSocketEnd = () => {
-        logger.info(
-          `[CoderAgentExecutor] Socket ended for message ${userMessage.messageId} (task ${taskId}). Aborting execution loop.`,
-        );
-        if (!abortController.signal.aborted) {
-          abortController.abort();
-        }
-        // Clean up the listener to prevent memory leaks
-        socket.removeListener('end', onSocketEnd);
-      };
-
-      // Listen on the socket's 'end' event (remote closed the connection)
-      socket.on('end', onSocketEnd);
-      socket.once('close', () => {
-        socket.removeListener('end', onSocketEnd);
-      });
-
-      // It's also good practice to remove the listener if the task completes successfully
-      abortSignal.addEventListener('abort', () => {
-        socket.removeListener('end', onSocketEnd);
-      });
-      logger.info(
-        `[CoderAgentExecutor] Socket close handler set up for task ${taskId}.`,
-      );
+    if (!this.activeAbortControllers.has(taskId)) {
+      this.activeAbortControllers.set(taskId, new Set());
    }
+    this.activeAbortControllers.get(taskId)!.add(abortController);

-    let wrapper: TaskWrapper | undefined = this.tasks.get(taskId);
+    let proceedToMainLoop = false;
+    let wrapper: TaskWrapper | undefined;
+    let isPrimaryExecution = false;

-    if (wrapper) {
-      wrapper.task.eventBus = eventBus;
-      logger.info(`[CoderAgentExecutor] Task ${taskId} found in memory cache.`);
-    } else if (sdkTask) {
-      logger.info(
-        `[CoderAgentExecutor] Task ${taskId} found in TaskStore. Reconstructing...`,
-      );
-      try {
-        wrapper = await this.reconstruct(sdkTask, eventBus);
-      } catch (e) {
-        logger.error(
-          `[CoderAgentExecutor] Failed to hydrate task ${taskId}:`,
-          e,
-        );
-        const stateChange: StateChange = {
-          kind: CoderAgentEvent.StateChangeEvent,
+    try {
+      if (store) {
+        const socket = store.req.socket;
+        const onSocketEnd = () => {
+          logger.info(
+            `[CoderAgentExecutor] Socket ended for message ${userMessage.messageId} (task ${taskId}). Aborting execution loop.`,
+          );
+          if (!abortController.signal.aborted) {
+            abortController.abort();
+          }
+          socket.removeListener('end', onSocketEnd);
        };
-        eventBus.publish({
-          kind: 'status-update',
-          taskId,
-          contextId: sdkTask.contextId,
-          status: {
-            state: 'failed',
-            message: {
-              kind: 'message',
-              role: 'agent',
-              parts: [
-                {
-                  kind: 'text',
-                  text: 'Internal error: Task state lost or corrupted.',
-                },
-              ],
-              messageId: uuidv4(),
-              taskId,
-              contextId: sdkTask.contextId,
-            } as Message,
-          },
-          final: true,
-          metadata: { coderAgent: stateChange },
-        });
-        return;
-      }
-    } else {
-      logger.info(`[CoderAgentExecutor] Creating new task ${taskId}.`);
-      const agentSettings = getAgentSettingsFromMetadata(userMessage.metadata);
-      try {
-        wrapper = await this.createTask(
-          taskId,
-          contextId,
-          agentSettings,
-          eventBus,
+        socket.on('end', onSocketEnd);
+        socket.once('close', () => socket.removeListener('end', onSocketEnd));
+        abortSignal.addEventListener('abort', () =>
+          socket.removeListener('end', onSocketEnd),
        );
-      } catch (error) {
-        logger.error(
-          `[CoderAgentExecutor] Error creating task ${taskId}:`,
-          error,
-        );
-        // eslint-disable-next-line @typescript-eslint/no-floating-promises
-        pushTaskStateFailed(error, eventBus, taskId, contextId);
-        return;
-      }
-      const newTaskSDK = wrapper.toSDKTask();
-      eventBus.publish({
-        ...newTaskSDK,
-        kind: 'task',
-        status: { state: 'submitted', timestamp: new Date().toISOString() },
-        history: [userMessage],
-      });
-      try {
-        await this.taskStore?.save(newTaskSDK);
-        logger.info(`[CoderAgentExecutor] New task ${taskId} saved to store.`);
-      } catch (saveError) {
-        logger.error(
-          `[CoderAgentExecutor] Failed to save new task ${taskId} to store:`,
-          saveError,
+        logger.info(
+          `[CoderAgentExecutor] Socket close handler set up for task ${taskId}.`,
        );
      }
-    }

-    if (!wrapper) {
-      logger.error(
-        `[CoderAgentExecutor] Task ${taskId} is unexpectedly undefined after load/create.`,
-      );
-      return;
+      // Check if the task is currently initializing
+      if (this.initializingTasks.has(taskId)) {
+        logger.info(
+          `[CoderAgentExecutor] Task ${taskId} is currently initializing. Waiting for initialization to complete.`,
+        );
+        const initPromise = this.initializationPromises.get(taskId);
+        if (initPromise) {
+          try {
+            wrapper = await initPromise;
+          } catch {
+            logger.error(
+              `[CoderAgentExecutor] Failed to wait for task ${taskId} initialization.`,
+            );
+            return;
+          }
+        }
+      }
+
+      if (!wrapper) {
+        this.initializingTasks.add(taskId);
+        const initPromise = (async () => {
+          let initializedWrapper: TaskWrapper | undefined;
+          initializedWrapper = this.tasks.get(taskId);
+
+          if (initializedWrapper) {
+            initializedWrapper.task.eventBus = eventBus;
+            logger.info(
+              `[CoderAgentExecutor] Task ${taskId} found in memory cache.`,
+            );
+          } else if (sdkTask) {
+            logger.info(
+              `[CoderAgentExecutor] Task ${taskId} found in TaskStore. Reconstructing...`,
+            );
+            try {
+              initializedWrapper = await this.reconstruct(sdkTask, eventBus);
+            } catch (e) {
+              logger.error(
+                `[CoderAgentExecutor] Aborting execution due to failed task reconstruction for task ${taskId}:`,
+                e,
+              );
+              throw e;
+            }
+          } else {
+            let agentSettings: AgentSettings;
+            try {
+              const rawAgentSettings = getAgentSettingsFromMetadata(
+                userMessage.metadata,
+              );
+              const validatedWorkspacePath = await validateWorkspacePath(
+                rawAgentSettings?.workspacePath,
+              );
+              agentSettings = {
+                kind: CoderAgentEvent.StateAgentSettingsEvent,
+                ...(rawAgentSettings || {}),
+                workspacePath: validatedWorkspacePath,
+                isTrusted: false,
+              };
+              initializedWrapper = await this.createTask(
+                taskId,
+                contextId,
+                agentSettings,
+                eventBus,
+              );
+            } catch (error) {
+              logger.error(
+                `[CoderAgentExecutor] Error creating task ${taskId}:`,
+                error,
+              );
+              void pushTaskStateFailed(error, eventBus, taskId, contextId);
+              throw error;
+            }
+            const newTaskSDK = initializedWrapper.toSDKTask();
+            eventBus.publish({
+              ...newTaskSDK,
+              kind: 'task',
+              status: {
+                state: 'submitted',
+                timestamp: new Date().toISOString(),
+              },
+              history: [userMessage],
+            });
+            try {
+              await this.taskStore?.save(newTaskSDK);
+              logger.info(
+                `[CoderAgentExecutor] New task ${taskId} saved to store.`,
+              );
+            } catch (saveError) {
+              logger.error(
+                `[CoderAgentExecutor] Failed to save new task ${taskId} to store:`,
+                saveError,
+              );
+            }
+          }
+          return initializedWrapper;
+        })();
+
+        this.initializationPromises.set(taskId, initPromise);
+
+        try {
+          wrapper = await initPromise;
+        } catch {
+          // Error is already handled/logged inside the promise
+          return;
+        } finally {
+          this.initializingTasks.delete(taskId);
+          this.initializationPromises.delete(taskId);
+        }
+      }
+
+      if (!wrapper) {
+        logger.error(
+          `[CoderAgentExecutor] Task ${taskId} is unexpectedly undefined after load/create.`,
+        );
+        return;
+      }
+
+      const currentTask = wrapper.task;
+
+      if (['canceled', 'failed', 'completed'].includes(currentTask.taskState)) {
+        logger.warn(
+          `[CoderAgentExecutor] Attempted to execute task ${taskId} which is already in state ${currentTask.taskState}. Ignoring.`,
+        );
+        return;
+      }
+
+      if (abortSignal.aborted) {
+        logger.warn(
+          `[CoderAgentExecutor] Task ${taskId} was aborted during initialization.`,
+        );
+        const isExplicitCancel = this.explicitlyCanceledTasks.has(taskId);
+        const finalState = isExplicitCancel ? 'canceled' : 'input-required';
+        const message = isExplicitCancel
+          ? 'Task canceled by user request.'
+          : 'Execution aborted by client.';
+        currentTask.setTaskStateAndPublishUpdate(
+          finalState,
+          { kind: CoderAgentEvent.StateChangeEvent },
+          message,
+          undefined,
+          true,
+        );
+        try {
+          await this.taskStore?.save(wrapper.toSDKTask());
+        } catch (saveError) {
+          logger.error(
+            `[CoderAgentExecutor] Failed to save task ${taskId} state:`,
+            saveError,
+          );
+        }
+        if (isExplicitCancel) {
+          this.cleanupAndEvictTask(taskId);
+        }
+        return;
+      }
+
+      if (this.executingTasks.has(taskId)) {
+        logger.info(
+          `[CoderAgentExecutor] Task ${taskId} has a pending execution. Processing message and yielding.`,
+        );
+        currentTask.eventBus = eventBus;
+        try {
+          for await (const _ of currentTask.acceptUserMessage(
+            requestContext,
+            abortController.signal,
+          )) {
+            logger.info(
+              `[CoderAgentExecutor] Processing user message ${userMessage.messageId} in secondary execution loop for task ${taskId}.`,
+            );
+          }
+        } catch (error) {
+          if (!abortController.signal.aborted) {
+            throw error;
+          }
+          logger.info(
+            `[CoderAgentExecutor] Secondary execution loop for task ${taskId} was aborted.`,
+          );
+        }
+        return;
+      }
+
+      isPrimaryExecution = true;
+
+      proceedToMainLoop = true;
+    } finally {
+      this.explicitlyCanceledTasks.delete(taskId);
+      if (!proceedToMainLoop) {
+        const controllers = this.activeAbortControllers.get(taskId);
+        if (controllers) {
+          controllers.delete(abortController);
+          if (controllers.size === 0) {
+            this.activeAbortControllers.delete(taskId);
+          }
+        }
+      }
    }

    const currentTask = wrapper.task;
-
-    if (['canceled', 'failed', 'completed'].includes(currentTask.taskState)) {
-      logger.warn(
-        `[CoderAgentExecutor] Attempted to execute task ${taskId} which is already in state ${currentTask.taskState}. Ignoring.`,
-      );
-      return;
-    }
-
-    if (this.executingTasks.has(taskId)) {
-      logger.info(
-        `[CoderAgentExecutor] Task ${taskId} has a pending execution. Processing message and yielding.`,
-      );
-      currentTask.eventBus = eventBus;
-      for await (const _ of currentTask.acceptUserMessage(
-        requestContext,
-        abortController.signal,
-      )) {
-        logger.info(
-          `[CoderAgentExecutor] Processing user message ${userMessage.messageId} in secondary execution loop for task ${taskId}.`,
-        );
-      }
-      // End this execution-- the original/source will be resumed.
-      return;
-    }
-
-    // Check if this is the primary/initial execution for this task
-    const isPrimaryExecution = !this.executingTasks.has(taskId);
-
-    if (!isPrimaryExecution) {
-      logger.info(
-        `[CoderAgentExecutor] Primary execution already active for task ${taskId}. Starting secondary loop for message ${userMessage.messageId}.`,
-      );
-      currentTask.eventBus = eventBus;
-      for await (const _ of currentTask.acceptUserMessage(
-        requestContext,
-        abortController.signal,
-      )) {
-        logger.info(
-          `[CoderAgentExecutor] Processing user message ${userMessage.messageId} in secondary execution loop for task ${taskId}.`,
-        );
-      }
-      // End this execution-- the original/source will be resumed.
-      return;
-    }
-
-    logger.info(
-      `[CoderAgentExecutor] Starting main execution for message ${userMessage.messageId} for task ${taskId}.`,
-    );
-    this.executingTasks.add(taskId);
-
    try {
+      logger.info(
+        `[CoderAgentExecutor] Starting main execution for message ${userMessage.messageId} for task ${taskId}.`,
+      );
+      this.executingTasks.add(taskId);
+
      let agentTurnActive = true;
      logger.info(`[CoderAgentExecutor] Task ${taskId}: Processing user turn.`);
      let agentEvents = currentTask.acceptUserMessage(
@@ -509,6 +659,12 @@ export class CoderAgentExecutor implements AgentExecutor {
      );

      while (agentTurnActive) {
+        if (abortSignal.aborted) {
+          logger.info(
+            `[CoderAgentExecutor] Task ${taskId} aborted before turn. Exiting loop.`,
+          );
+          throw new Error('Execution aborted');
+        }
        logger.info(
          `[CoderAgentExecutor] Task ${taskId}: Processing agent turn (LLM stream).`,
        );
@@ -555,7 +711,6 @@ export class CoderAgentExecutor implements AgentExecutor {
          const completedTools = currentTask.getAndClearCompletedTools();

          if (completedTools.length > 0) {
-            // If all completed tool calls were canceled, manually add them to history and set state to input-required, final:true
            if (completedTools.every((tool) => tool.status === 'cancelled')) {
              logger.info(
                `[CoderAgentExecutor] Task ${taskId}: All tool calls were cancelled. Updating history and ending agent turn.`,
@@ -581,7 +736,6 @@ export class CoderAgentExecutor implements AgentExecutor {
                completedTools,
                abortSignal,
              );
-              // Continue the loop to process the LLM response to the tool results.
            }
          } else {
            logger.info(
@@ -644,6 +798,13 @@ export class CoderAgentExecutor implements AgentExecutor {
      }
    } finally {
      if (isPrimaryExecution) {
+        const controllers = this.activeAbortControllers.get(taskId);
+        if (controllers) {
+          controllers.delete(abortController);
+          if (controllers.size === 0) {
+            this.activeAbortControllers.delete(taskId);
+          }
+        }
        this.executingTasks.delete(taskId);
        logger.info(
          `[CoderAgentExecutor] Saving final state for task ${taskId}.`,
@@ -661,11 +822,7 @@ export class CoderAgentExecutor implements AgentExecutor {
        if (
          ['canceled', 'failed', 'completed'].includes(currentTask.taskState)
        ) {
-          logger.info(
-            `[CoderAgentExecutor] Task ${taskId} reached terminal state ${currentTask.taskState}. Evicting and disposing.`,
-          );
-          wrapper.task.dispose();
-          this.tasks.delete(taskId);
+          this.cleanupAndEvictTask(taskId);
        }
      }
    }
@@ -0,0 +1,62 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import { resolveToRealPath, isSubpath } from '@google/gemini-cli-core';
+
+/**
+ * Validates a workspace path to prevent path traversal attacks.
+ *
+ * @param workspacePath The path to validate.
+ * @param allowedRoot The root directory the path must be within. Defaults to CWD.
+ * @returns The resolved, safe path.
+ * @throws An error if the path is invalid or outside the allowed root.
+ */
+export async function validateWorkspacePath(
+  workspacePath?: string,
+  allowedRoot: string = process.cwd(),
+): Promise<string> {
+  const trimmedPath = workspacePath?.trim();
+  if (!trimmedPath) {
+    return resolveToRealPath(allowedRoot);
+  }
+
+  if (trimmedPath.includes('\0')) {
+    throw new Error('Security violation: Null byte detected in path.');
+  }
+
+  try {
+    const canonicalAllowedRoot = resolveToRealPath(allowedRoot);
+    const resolvedWorkspacePath = path.resolve(
+      canonicalAllowedRoot,
+      trimmedPath,
+    );
+    const canonicalWorkspacePath = resolveToRealPath(resolvedWorkspacePath);
+
+    // Check if the resolved path is within the allowed root directory
+    if (
+      canonicalWorkspacePath !== canonicalAllowedRoot &&
+      !isSubpath(canonicalAllowedRoot, canonicalWorkspacePath)
+    ) {
+      throw new Error(
+        `Security violation: The path "${trimmedPath}" is outside the allowed root directory.`,
+      );
+    }
+
+    const stats = await fs.promises.stat(canonicalWorkspacePath);
+    if (!stats.isDirectory()) {
+      throw new Error(`The path "${trimmedPath}" is not a directory.`);
+    }
+
+    return canonicalWorkspacePath;
+  } catch (e) {
+    if (e instanceof Error && 'code' in e && e.code === 'ENOENT') {
+      throw new Error(`The path "${trimmedPath}" does not exist.`);
+    }
+    throw e; // Re-throw other errors
+  }
+}