diff --git a/packages/cli/src/nonInteractiveCliAgentSession.ts b/packages/cli/src/nonInteractiveCliAgentSession.ts index e0a532becf..e64b4a5686 100644 --- a/packages/cli/src/nonInteractiveCliAgentSession.ts +++ b/packages/cli/src/nonInteractiveCliAgentSession.ts @@ -193,6 +193,7 @@ export async function runNonInteractive({ let errorToHandle: unknown | undefined; let scheduler: Scheduler | undefined; + let session: LegacyAgentSession | undefined; let abortSession = () => {}; try { consolePatcher.patch(); @@ -296,7 +297,7 @@ export async function runNonInteractive({ } // Create LegacyAgentSession — owns the agentic loop - const session = new LegacyAgentSession({ + session = new LegacyAgentSession({ client: geminiClient, scheduler, config, @@ -305,7 +306,7 @@ export async function runNonInteractive({ // Wire Ctrl+C to session abort abortSession = () => { - void session.abort(); + void session?.abort(); }; abortController.signal.addEventListener('abort', abortSession); if (abortController.signal.aborted) { @@ -640,6 +641,7 @@ export async function runNonInteractive({ cleanupStdinCancellation(); abortController.signal.removeEventListener('abort', abortSession); + session?.dispose(); scheduler?.dispose(); consolePatcher.cleanup(); coreEvents.off(CoreEvent.UserFeedback, handleUserFeedback); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index d77c9adfb6..7cc9b4ec7b 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1180,6 +1180,10 @@ Logging in with Google... Restarting Gemini CLI to continue. [config, getPreferredEditor], ); + useEffect(() => () => { + streamAgent?.dispose?.(); + }, [streamAgent]); + const activeStream = streamAgent ? // eslint-disable-next-line react-hooks/rules-of-hooks useAgentStream({ diff --git a/packages/core/src/agent/agent-session.ts b/packages/core/src/agent/agent-session.ts index 6a4c295fc8..e07695532d 100644 --- a/packages/core/src/agent/agent-session.ts +++ b/packages/core/src/agent/agent-session.ts @@ -34,6 +34,10 @@ export class AgentSession implements AgentProtocol { return this._protocol.abort(); } + dispose(): void { + this._protocol.dispose?.(); + } + get events(): readonly AgentEvent[] { return this._protocol.events; } diff --git a/packages/core/src/agent/legacy-agent-session.ts b/packages/core/src/agent/legacy-agent-session.ts index b1ade30a24..587c107c69 100644 --- a/packages/core/src/agent/legacy-agent-session.ts +++ b/packages/core/src/agent/legacy-agent-session.ts @@ -71,6 +71,7 @@ export class LegacyAgentProtocol implements AgentProtocol { private _agentEndEmitted = false; private _activeStreamId?: string; private _abortController = new AbortController(); + private _disposalController = new AbortController(); private _nextStreamIdOverride?: string; private _lastToolStatuses = new Map(); @@ -106,10 +107,16 @@ export class LegacyAgentProtocol implements AgentProtocol { this._config.messageBus.subscribe( MessageBusType.TOOL_CALLS_UPDATE, this._handleToolCallsUpdate.bind(this), + { signal: this._disposalController.signal }, ); } } + dispose(): void { + this._disposalController.abort(); + void this.abort(); + } + get events(): readonly AgentEvent[] { return this._events; } diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts index 1f6990703b..9595135da3 100644 --- a/packages/core/src/agent/types.ts +++ b/packages/core/src/agent/types.ts @@ -37,6 +37,11 @@ export interface AgentProtocol extends Trajectory { */ abort(): Promise; + /** + * Disposes of the protocol, cleaning up any long-lived resources. + */ + dispose?(): void; + /** * AgentProtocol implements the Trajectory interface and can retrieve existing events. */ diff --git a/scripts/build_package.js b/scripts/build_package.js index 5523748a39..279e46fa94 100644 --- a/scripts/build_package.js +++ b/scripts/build_package.js @@ -18,7 +18,7 @@ // limitations under the License. import { execSync } from 'node:child_process'; -import { writeFileSync, existsSync, cpSync, rmSync } from 'node:fs'; +import { writeFileSync, existsSync, cpSync } from 'node:fs'; import { join, basename } from 'node:path'; if (!process.cwd().includes('packages')) { @@ -48,14 +48,7 @@ if (packageName === 'core') { const docsSource = join(process.cwd(), '..', '..', 'docs'); const docsTarget = join(process.cwd(), 'dist', 'docs'); if (existsSync(docsSource)) { - if (existsSync(docsTarget)) { - rmSync(docsTarget, { recursive: true, force: true }); - } - cpSync(docsSource, docsTarget, { - recursive: true, - dereference: true, - force: true, - }); + cpSync(docsSource, docsTarget, { recursive: true, dereference: true }); console.log('Copied documentation to dist/docs'); } }