diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 50f013f5b85..b964831c233 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -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((resolve) => { + finishExecution = resolve; + }); + const executionStarted = new Promise((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((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( + './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((resolve) => { + release = resolve; + }); + const started = new Promise((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'); diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 0dddaee8b4c..5306841a5ac 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -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 { - 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 { 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 {