andreivince · Sep 20, 2026 · Sep 21, 2026 · Sep 23, 2026 · Sep 28, 2026 · Sep 30, 2026
@@ -15,6 +15,7 @@ import {
   type Mocked,
 } from 'vitest';
 import { randomUUID } from 'node:crypto';
+import { getEventListeners } from 'node:events';
 
 vi.mock('node:crypto', () => ({
   randomUUID: vi.fn(),
@@ -1465,6 +1466,275 @@ describe('Scheduler (Orchestrator)', () => {
   });
 
   describe('Cleanup', () => {
+    async function startBlockedBatch(requests: ToolCallRequestInfo[] = [req1]) {
+      let finishExecution!: () => void;
+      let notifyStarted!: () => void;
+      const executionFinished = new Promise<void>((resolve) => {
+        finishExecution = resolve;
+      });
+      const executionStarted = new Promise<void>((resolve) => {
+        notifyStarted = resolve;
+      });
+
+      mockExecutor.execute.mockImplementationOnce(async () => {
+        notifyStarted();
+        await executionFinished;
+        return {
+          status: CoreToolCallStatus.Success,
+          request: req1,
+          tool: mockTool,
+          invocation: mockInvocation as unknown as AnyToolInvocation,
+          response: {
+            callId: req1.callId,
+            responseParts: [],
+            resultDisplay: 'Completed',
+            error: undefined,
+            errorType: undefined,
+          },
+        };
+      });
+
+      const activeBatch = scheduler.schedule(requests, signal);
+      await executionStarted;
+      return { activeBatch, finishExecution };
+    }
+
+    it('should reject queued batches and remove abort listeners on dispose()', async () => {
+      const { activeBatch, finishExecution } = await startBlockedBatch();
+      const queuedSignals = [
+        new AbortController().signal,
+        new AbortController().signal,
+      ];
+      const queuedBatches = queuedSignals.map((queuedSignal, index) =>
+        scheduler.schedule(
+          { ...req2, callId: `queued-${index}` },
+          queuedSignal,
+        ),
+      );
+      const onSettled = vi.fn();
+      const queuedResults = Promise.allSettled(queuedBatches).then(onSettled);
+
+      try {
+        for (const queuedSignal of queuedSignals) {
+          expect(getEventListeners(queuedSignal, 'abort')).toHaveLength(1);
+        }
+
+        scheduler.dispose();
+        scheduler.dispose();
+
+        // Let promise reactions run while the active execution stays blocked.
+        await new Promise<void>((resolve) => setImmediate(resolve));
+
+        expect(onSettled).toHaveBeenCalledExactlyOnceWith([
+          { status: 'rejected', reason: new Error('Scheduler disposed') },
+          { status: 'rejected', reason: new Error('Scheduler disposed') },
+        ]);
+        for (const queuedSignal of queuedSignals) {
+          expect(getEventListeners(queuedSignal, 'abort')).toHaveLength(0);
+          expect(queuedSignal.aborted).toBe(false);
+        }
+        expect(signal.aborted).toBe(false);
+        expect(mockStateManager.cancelAllQueued).not.toHaveBeenCalled();
+        expect(mockExecutor.execute).toHaveBeenCalledTimes(1);
+      } finally {
+        finishExecution();
+        await activeBatch;
+        await queuedResults;
+      }
+
+      expect(mockExecutor.execute).toHaveBeenCalledTimes(1);
+      expect(mockStateManager.finalizeCall).toHaveBeenCalledExactlyOnceWith(
+        req1.callId,
+      );
+    });
+
+    describe('Disposal within a batch', () => {
+      beforeEach(async () => {
+        scheduler.dispose();
+        const actual =
+          await vi.importActual<typeof import('./state-manager.js')>(
+            './state-manager.js',
+          );
+        vi.mocked(SchedulerStateManager).mockImplementation(
+          (...args) => new actual.SchedulerStateManager(...args),
+        );
+        scheduler = new Scheduler({
+          context: mockConfig,
+          messageBus: mockMessageBus,
+          getPreferredEditor,
+          schedulerId: 'root',
+        });
+      });
+
+      afterEach(() => scheduler.dispose());
+
+      it('should finish executing tools and cancel remaining serial tools', async () => {
+        const { activeBatch, finishExecution } = await startBlockedBatch([
+          req1,
+          req2,
+          { ...req2, callId: 'call-3' },
+        ]);
+        scheduler.dispose();
+        finishExecution();
+
+        const results = await activeBatch;
+
+        expect(mockExecutor.execute).toHaveBeenCalledTimes(1);
+        expect(
+          results.map((call) => [call.request.callId, call.status]),
+        ).toEqual([
+          ['call-1', CoreToolCallStatus.Success],
+          ['call-2', CoreToolCallStatus.Cancelled],
+          ['call-3', CoreToolCallStatus.Cancelled],
+        ]);
+        expect(results[1].response.responseParts).toEqual([
+          expect.objectContaining({
+            functionResponse: expect.objectContaining({
+              response: {
+                error: expect.stringContaining('Scheduler disposed'),
+              },
+            }),
+          }),
+        ]);
+        expect(signal.aborted).toBe(false);
+        expect(scheduler.completedCalls).toEqual([]);
+      });
+
+      it('should preserve a sandbox error without requesting expansion after disposal', async () => {
+        const response: ToolCallResponseInfo = {
+          callId: req1.callId,
+          responseParts: [],
+          resultDisplay: 'Sandbox denied the command',
+          error: new Error(
+            JSON.stringify({
+              rootCommand: 'echo',
+              additionalPermissions: {},
+            }),
+          ),
+          errorType: ToolErrorType.SANDBOX_EXPANSION_REQUIRED,
+        };
+        vi.mocked(mockTool.build).mockReturnValue({
+          ...mockInvocation,
+          getDescription: () => 'Test tool',
+        } as unknown as AnyToolInvocation);
+        mockExecutor.execute.mockImplementationOnce(async ({ call }) => {
+          scheduler.dispose();
+          return { ...call, status: CoreToolCallStatus.Error, response };
+        });
+
+        const results = await scheduler.schedule([req1, req2], signal);
+
+        expect(resolveConfirmation).not.toHaveBeenCalled();
+        expect(mockExecutor.execute).toHaveBeenCalledTimes(1);
+        expect(
+          results.map((call) => [call.request.callId, call.status]),
+        ).toEqual([
+          ['call-1', CoreToolCallStatus.Error],
+          ['call-2', CoreToolCallStatus.Cancelled],
+        ]);
+        expect(results[0].response).toEqual(response);
+        expect(signal.aborted).toBe(false);
+        expect(scheduler.completedCalls).toEqual([]);
+      });
+
+      it.each([
+        ['policy', PolicyDecision.ALLOW],
+        ['policy', PolicyDecision.ASK_USER],
+        ['confirmation', PolicyDecision.ASK_USER],
+      ] as const)(
+        'should not execute tools when disposed during %s validation (%s)',
+        async (phase, decision) => {
+          let release!: () => void;
+          let entered!: () => void;
+          const gate = new Promise<void>((resolve) => {
+            release = resolve;
+          });
+          const started = new Promise<void>((resolve) => {
+            entered = resolve;
+          });
+          if (phase === 'policy') {
+            vi.mocked(checkPolicy).mockImplementationOnce(async () => {
+              entered();
+              await gate;
+              return { decision };
+            });
+          } else {
+            vi.mocked(checkPolicy).mockResolvedValueOnce({
+              decision: PolicyDecision.ASK_USER,
+            });
+            vi.mocked(resolveConfirmation).mockImplementationOnce(async () => {
+              entered();
+              await gate;
+              return { outcome: ToolConfirmationOutcome.ProceedOnce };
+            });
+          }
+          const batch = scheduler.schedule([req1, req2], signal);
+          await started;
+          scheduler.dispose();
+          release();
+
+          const results = await batch;
+
+          if (phase === 'policy') {
+            expect(resolveConfirmation).not.toHaveBeenCalled();
+          }
+          expect(mockExecutor.execute).not.toHaveBeenCalled();
+          expect(
+            results.map((call) => [call.request.callId, call.status]),
+          ).toEqual([
+            ['call-1', CoreToolCallStatus.Cancelled],
+            ['call-2', CoreToolCallStatus.Cancelled],
+          ]);
+          expect(signal.aborted).toBe(false);
+          expect(scheduler.completedCalls).toEqual([]);
+        },
+      );
+    });
+
+    it('should reject new batches after dispose()', async () => {
+      scheduler.dispose();
+
+      await expect(scheduler.schedule(req1, signal)).rejects.toThrow(
+        'Scheduler disposed',
+      );
+      expect(mockStateManager.enqueue).not.toHaveBeenCalled();
+      expect(mockExecutor.execute).not.toHaveBeenCalled();
+      expect(getEventListeners(signal, 'abort')).toHaveLength(0);
+    });
+
+    it('should preserve caller cancellation for queued batches before dispose()', async () => {
+      const { activeBatch, finishExecution } = await startBlockedBatch();
+      const queuedController = new AbortController();
+      const queuedBatch = scheduler.schedule(req2, queuedController.signal);
+      const queuedResult = Promise.allSettled([queuedBatch]);
+
+      try {
+        queuedController.abort();
+
+        await expect(queuedResult).resolves.toEqual([
+          {
+            status: 'rejected',
+            reason: new Error('Tool call cancelled while in queue.'),
+          },
+        ]);
+        expect(
+          getEventListeners(queuedController.signal, 'abort'),
+        ).toHaveLength(0);
+
+        await expect(
+          scheduler.schedule(req2, queuedController.signal),
+        ).rejects.toThrow('Operation cancelled');
+
+        scheduler.dispose();
+        expect(signal.aborted).toBe(false);
+      } finally {
+        finishExecution();
+        await activeBatch;
+      }
+
+      expect(mockExecutor.execute).toHaveBeenCalledTimes(1);
+    });
+
     it('should unregister McpProgress listener on dispose()', () => {
       const onSpy = vi.spyOn(coreEvents, 'on');
       const offSpy = vi.spyOn(coreEvents, 'off');

@@ -156,6 +156,10 @@ export class Scheduler {
   dispose(): void {
     coreEvents.off(CoreEvent.McpProgress, this.handleMcpProgress);
     this.disposeController.abort();
+
+    for (const request of this.requestQueue.splice(0)) {
+      request.reject(new Error('Scheduler disposed'));
+    }
   }
 
   private readonly handleMcpProgress = (payload: McpProgressPayload) => {
@@ -220,6 +224,10 @@ export class Scheduler {
         sessionId: this.context.config.getSessionId(),
       },
       async ({ metadata: spanMetadata }) => {
+        if (this.disposeController.signal.aborted) {
+          throw new Error('Scheduler disposed');
+        }
+
         const requests = Array.isArray(request) ? request : [request];
 
         spanMetadata.input = requests;
@@ -454,7 +462,11 @@ export class Scheduler {
    * @returns true if the loop should continue, false if it should terminate.
    */
   private async _processNextItem(signal: AbortSignal): Promise<boolean> {
-    if (signal.aborted || this.isCancelling) {
+    if (
+      signal.aborted ||
+      this.isCancelling ||
+      this.disposeController.signal.aborted
+    ) {
       // Finalize active calls that are terminal
       const activeCalls = this.state.allActiveCalls;
       for (const call of activeCalls) {
@@ -463,7 +475,11 @@ export class Scheduler {
         }
       }
 
-      this.state.cancelAllQueued('Operation cancelled');
+      this.state.cancelAllQueued(
+        this.disposeController.signal.aborted
+          ? 'Scheduler disposed'
+          : 'Operation cancelled',
+      );
       return false;
     }
 
@@ -706,6 +722,14 @@ export class Scheduler {
     let lastDetails: SerializableConfirmationDetails | undefined;
 
     if (decision === PolicyDecision.ASK_USER) {
+      if (this.disposeController.signal.aborted) {
+        this.state.updateStatus(
+          callId,
+          CoreToolCallStatus.Cancelled,
+          'Scheduler disposed',
+        );
+        return;
+      }
       const forcedDecision =
         hookDecision === 'ask' ||
         (policyDecision === PolicyDecision.ALLOW && hasTaintRisk)
@@ -765,11 +789,13 @@ export class Scheduler {
     signal: AbortSignal,
   ): Promise<boolean> {
     const callId = toolCall.request.callId;
-    if (signal.aborted) {
+    if (signal.aborted || this.disposeController.signal.aborted) {
       this.state.updateStatus(
         callId,
         CoreToolCallStatus.Cancelled,
-        'Operation cancelled',
+        this.disposeController.signal.aborted
+          ? 'Scheduler disposed'
+          : 'Operation cancelled',
       );
       return false;
     }
@@ -877,7 +903,7 @@ export class Scheduler {
       sandboxDetailsStr = result.response.error?.message || '';
     }
 
-    if (isSandboxError) {
+    if (isSandboxError && !this.disposeController.signal.aborted) {
       try {
         // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
         const parsedError = JSON.parse(sandboxDetailsStr) as {