From b9f1d832c80b644eec2e997e85a6105b9d0c0b5d Mon Sep 17 00:00:00 2001 From: Anjaligarhwal Date: Wed, 8 Apr 2026 08:35:53 +0530 Subject: [PATCH] fix(core): dispose Scheduler to prevent McpProgress listener leak (#24870) --- packages/cli/src/nonInteractiveCli.test.ts | 1 + packages/cli/src/nonInteractiveCli.ts | 4 +- .../src/nonInteractiveCliAgentSession.test.ts | 1 + .../cli/src/nonInteractiveCliAgentSession.ts | 4 +- .../core/src/agents/agent-scheduler.test.ts | 52 +++++++++++++++++++ packages/core/src/agents/agent-scheduler.ts | 6 ++- 6 files changed, 65 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 855707de9e..5d0c3d1016 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -71,6 +71,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { Scheduler: class { schedule = mockSchedulerSchedule; cancelAll = vi.fn(); + dispose = vi.fn(); }, isTelemetrySdkInitialized: vi.fn().mockReturnValue(true), ChatRecordingService: MockChatRecordingService, diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 26daaf66a1..dc5255edee 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -187,6 +187,7 @@ export async function runNonInteractive( }; let errorToHandle: unknown | undefined; + let scheduler: Scheduler | undefined; try { consolePatcher.patch(); @@ -215,7 +216,7 @@ export async function runNonInteractive( }); const geminiClient = config.getGeminiClient(); - const scheduler = new Scheduler({ + scheduler = new Scheduler({ context: config, messageBus: config.getMessageBus(), getPreferredEditor: () => undefined, @@ -528,6 +529,7 @@ export async function runNonInteractive( // Cleanup stdin cancellation before other cleanup cleanupStdinCancellation(); + scheduler?.dispose(); consolePatcher.cleanup(); coreEvents.off(CoreEvent.UserFeedback, handleUserFeedback); } diff --git a/packages/cli/src/nonInteractiveCliAgentSession.test.ts b/packages/cli/src/nonInteractiveCliAgentSession.test.ts index 617f80aca6..923109643c 100644 --- a/packages/cli/src/nonInteractiveCliAgentSession.test.ts +++ b/packages/cli/src/nonInteractiveCliAgentSession.test.ts @@ -71,6 +71,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { Scheduler: class { schedule = mockSchedulerSchedule; cancelAll = vi.fn(); + dispose = vi.fn(); }, isTelemetrySdkInitialized: vi.fn().mockReturnValue(true), ChatRecordingService: MockChatRecordingService, diff --git a/packages/cli/src/nonInteractiveCliAgentSession.ts b/packages/cli/src/nonInteractiveCliAgentSession.ts index fe5fbceba2..7f36ce6cf5 100644 --- a/packages/cli/src/nonInteractiveCliAgentSession.ts +++ b/packages/cli/src/nonInteractiveCliAgentSession.ts @@ -184,6 +184,7 @@ export async function runNonInteractive({ }; let errorToHandle: unknown | undefined; + let scheduler: Scheduler | undefined; let abortSession = () => {}; try { consolePatcher.patch(); @@ -215,7 +216,7 @@ export async function runNonInteractive({ }); const geminiClient = config.getGeminiClient(); - const scheduler = new Scheduler({ + scheduler = new Scheduler({ context: config, messageBus: config.getMessageBus(), getPreferredEditor: () => undefined, @@ -612,6 +613,7 @@ export async function runNonInteractive({ cleanupStdinCancellation(); abortController.signal.removeEventListener('abort', abortSession); + scheduler?.dispose(); consolePatcher.cleanup(); coreEvents.off(CoreEvent.UserFeedback, handleUserFeedback); } diff --git a/packages/core/src/agents/agent-scheduler.test.ts b/packages/core/src/agents/agent-scheduler.test.ts index 5d5b6569af..8ac15f181e 100644 --- a/packages/core/src/agents/agent-scheduler.test.ts +++ b/packages/core/src/agents/agent-scheduler.test.ts @@ -15,6 +15,7 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js'; vi.mock('../scheduler/scheduler.js', () => ({ Scheduler: vi.fn().mockImplementation(() => ({ schedule: vi.fn().mockResolvedValue([{ status: 'success' }]), + dispose: vi.fn(), })), })); @@ -125,6 +126,57 @@ describe('agent-scheduler', () => { expect(schedulerConfig.toolRegistry).not.toBe(mainRegistry); }); + it('should dispose the scheduler after schedule completes', async () => { + const mockConfig = { + getPromptRegistry: vi.fn(), + getResourceRegistry: vi.fn(), + messageBus: mockMessageBus, + toolRegistry: mockToolRegistry, + } as unknown as Mocked; + + const options = { + schedulerId: 'subagent-1', + toolRegistry: mockToolRegistry as unknown as ToolRegistry, + signal: new AbortController().signal, + }; + + await scheduleAgentTools(mockConfig as unknown as Config, [], options); + + const schedulerInstance = vi.mocked(Scheduler).mock.results[0].value; + expect(schedulerInstance.dispose).toHaveBeenCalledOnce(); + }); + + it('should dispose the scheduler even when schedule throws', async () => { + const scheduleError = new Error('schedule failed'); + vi.mocked(Scheduler).mockImplementationOnce( + () => + ({ + schedule: vi.fn().mockRejectedValue(scheduleError), + dispose: vi.fn(), + }) as unknown as Scheduler, + ); + + const mockConfig = { + getPromptRegistry: vi.fn(), + getResourceRegistry: vi.fn(), + messageBus: mockMessageBus, + toolRegistry: mockToolRegistry, + } as unknown as Mocked; + + const options = { + schedulerId: 'subagent-1', + toolRegistry: mockToolRegistry as unknown as ToolRegistry, + signal: new AbortController().signal, + }; + + await expect( + scheduleAgentTools(mockConfig as unknown as Config, [], options), + ).rejects.toThrow('schedule failed'); + + const schedulerInstance = vi.mocked(Scheduler).mock.results[0].value; + expect(schedulerInstance.dispose).toHaveBeenCalledOnce(); + }); + it('should create an AgentLoopContext that has a defined .config property', async () => { const mockConfig = { getPromptRegistry: vi.fn(), diff --git a/packages/core/src/agents/agent-scheduler.ts b/packages/core/src/agents/agent-scheduler.ts index 8bed1de00b..09b32980a9 100644 --- a/packages/core/src/agents/agent-scheduler.ts +++ b/packages/core/src/agents/agent-scheduler.ts @@ -85,5 +85,9 @@ export async function scheduleAgentTools( onWaitingForConfirmation, }); - return scheduler.schedule(requests, signal); + try { + return await scheduler.schedule(requests, signal); + } finally { + scheduler.dispose(); + } }