From 5dfbb739e5d1d44f9adb00cd46736e035ea98ab2 Mon Sep 17 00:00:00 2001 From: Anjaligarhwal Date: Tue, 5 May 2026 02:47:36 +0530 Subject: [PATCH] feat(cli): add /bug-memory command and auto-capture heap snapshot in /bug (#25639) --- .../src/services/BuiltinCommandLoader.test.ts | 3 + .../cli/src/services/BuiltinCommandLoader.ts | 2 + .../cli/src/ui/commands/bugCommand.test.ts | 125 +++++++++++++++++- packages/cli/src/ui/commands/bugCommand.ts | 53 ++++++++ .../src/ui/commands/bugMemoryCommand.test.ts | 121 +++++++++++++++++ .../cli/src/ui/commands/bugMemoryCommand.ts | 86 ++++++++++++ .../cli/src/ui/utils/memorySnapshot.test.ts | 84 ++++++++++++ packages/cli/src/ui/utils/memorySnapshot.ts | 30 +++++ 8 files changed, 503 insertions(+), 1 deletion(-) create mode 100644 packages/cli/src/ui/commands/bugMemoryCommand.test.ts create mode 100644 packages/cli/src/ui/commands/bugMemoryCommand.ts create mode 100644 packages/cli/src/ui/utils/memorySnapshot.test.ts create mode 100644 packages/cli/src/ui/utils/memorySnapshot.ts diff --git a/packages/cli/src/services/BuiltinCommandLoader.test.ts b/packages/cli/src/services/BuiltinCommandLoader.test.ts index d53273134c..aca91ab9d8 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.test.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.test.ts @@ -71,6 +71,9 @@ vi.mock('../ui/commands/agentsCommand.js', () => ({ agentsCommand: { name: 'agents' }, })); vi.mock('../ui/commands/bugCommand.js', () => ({ bugCommand: {} })); +vi.mock('../ui/commands/bugMemoryCommand.js', () => ({ + bugMemoryCommand: { name: 'bug-memory' }, +})); vi.mock('../ui/commands/chatCommand.js', () => ({ chatCommand: { name: 'chat', diff --git a/packages/cli/src/services/BuiltinCommandLoader.ts b/packages/cli/src/services/BuiltinCommandLoader.ts index 1c5288707c..5312d834e4 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.ts @@ -22,6 +22,7 @@ import { aboutCommand } from '../ui/commands/aboutCommand.js'; import { agentsCommand } from '../ui/commands/agentsCommand.js'; import { authCommand } from '../ui/commands/authCommand.js'; import { bugCommand } from '../ui/commands/bugCommand.js'; +import { bugMemoryCommand } from '../ui/commands/bugMemoryCommand.js'; import { chatCommand, debugCommand } from '../ui/commands/chatCommand.js'; import { clearCommand } from '../ui/commands/clearCommand.js'; import { commandsCommand } from '../ui/commands/commandsCommand.js'; @@ -123,6 +124,7 @@ export class BuiltinCommandLoader implements ICommandLoader { ...(this.config?.isAgentsEnabled() ? [agentsCommand] : []), authCommand, bugCommand, + bugMemoryCommand, { ...chatCommand, subCommands: chatResumeSubCommands, diff --git a/packages/cli/src/ui/commands/bugCommand.test.ts b/packages/cli/src/ui/commands/bugCommand.test.ts index f767805b01..a51c7af12c 100644 --- a/packages/cli/src/ui/commands/bugCommand.test.ts +++ b/packages/cli/src/ui/commands/bugCommand.test.ts @@ -12,10 +12,33 @@ import { createMockCommandContext } from '../../test-utils/mockCommandContext.js import { getVersion, type Config } from '@google/gemini-cli-core'; import { GIT_COMMIT_INFO } from '../../generated/git-commit.js'; import { formatBytes } from '../utils/formatters.js'; +import { MessageType } from '../types.js'; +import { captureHeapSnapshot } from '../utils/memorySnapshot.js'; + +const { memoryUsageMock } = vi.hoisted(() => ({ + memoryUsageMock: vi.fn(() => ({ + rss: 0, + heapTotal: 0, + heapUsed: 0, + external: 0, + arrayBuffers: 0, + })), +})); // Mock dependencies vi.mock('open'); vi.mock('../utils/formatters.js'); +vi.mock('../utils/memorySnapshot.js', () => ({ + captureHeapSnapshot: vi.fn(), + MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES: 2 * 1024 * 1024 * 1024, +})); +vi.mock('node:fs/promises', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + stat: vi.fn().mockResolvedValue({ size: 4096 }), + }; +}); vi.mock('../utils/historyExportUtils.js', async (importOriginal) => { const actual = await importOriginal(); @@ -53,7 +76,7 @@ vi.mock('node:process', () => ({ version: 'v20.0.0', // Keep other necessary process properties if needed by other parts of the code env: process.env, - memoryUsage: () => ({ rss: 0 }), + memoryUsage: memoryUsageMock, }, })); @@ -69,6 +92,13 @@ describe('bugCommand', () => { beforeEach(() => { vi.mocked(getVersion).mockResolvedValue('0.1.0'); vi.mocked(formatBytes).mockReturnValue('100 MB'); + memoryUsageMock.mockReturnValue({ + rss: 0, + heapTotal: 0, + heapUsed: 0, + external: 0, + arrayBuffers: 0, + }); vi.stubEnv('SANDBOX', 'gemini-test'); vi.useFakeTimers(); vi.setSystemTime(new Date('2024-01-01T00:00:00Z')); @@ -218,4 +248,97 @@ describe('bugCommand', () => { expect(open).toHaveBeenCalledWith(expectedUrl); }); + + const buildHighMemoryContext = (tempDir: string | undefined) => + createMockCommandContext({ + services: { + agentContext: { + config: { + getModel: () => 'gemini-pro', + getBugCommand: () => undefined, + getIdeMode: () => false, + getContentGeneratorConfig: () => ({ authType: 'oauth-personal' }), + storage: tempDir ? { getProjectTempDir: () => tempDir } : undefined, + getSessionId: vi.fn().mockReturnValue('test-session-id'), + } as unknown as Config, + geminiClient: { getChat: () => ({ getHistory: () => [] }) }, + }, + }, + }); + + it('captures a heap snapshot AFTER opening the bug URL when RSS exceeds 2 GB', async () => { + memoryUsageMock.mockReturnValue({ + rss: 3 * 1024 * 1024 * 1024, + heapTotal: 0, + heapUsed: 0, + external: 0, + arrayBuffers: 0, + }); + vi.mocked(captureHeapSnapshot).mockResolvedValueOnce(undefined); + + const tempDir = path.join('/tmp', 'gemini-test'); + const context = buildHighMemoryContext(tempDir); + + if (!bugCommand.action) throw new Error('Action is not defined'); + await bugCommand.action(context, 'A memory bug'); + + const now = new Date('2024-01-01T00:00:00Z').getTime(); + const expectedSnapshotPath = path.join( + tempDir, + `bug-memory-${now}.heapsnapshot`, + ); + expect(captureHeapSnapshot).toHaveBeenCalledWith(expectedSnapshotPath); + + const addItem = vi.mocked(context.ui.addItem); + const callOrder = addItem.mock.invocationCallOrder; + const openOrder = vi.mocked(open).mock.invocationCallOrder[0]; + // The URL message must precede the "capturing" message so the user sees + // the URL before the 20+ second snapshot starts. + expect(callOrder[0]).toBeLessThan(openOrder); + expect(callOrder[1]).toBeGreaterThan(openOrder); + expect(addItem.mock.calls[1][0].text).toContain('High memory usage'); + expect(addItem.mock.calls[2][0].text).toContain('Heap snapshot saved'); + expect(addItem.mock.calls[2][0].text).toContain(expectedSnapshotPath); + expect(addItem.mock.calls[2][0].type).toBe(MessageType.INFO); + }); + + it('skips auto-capture when RSS is below the 2 GB threshold', async () => { + memoryUsageMock.mockReturnValue({ + rss: 1 * 1024 * 1024 * 1024, + heapTotal: 0, + heapUsed: 0, + external: 0, + arrayBuffers: 0, + }); + const context = buildHighMemoryContext('/tmp/gemini-test'); + + if (!bugCommand.action) throw new Error('Action is not defined'); + await bugCommand.action(context, 'A light bug'); + + expect(captureHeapSnapshot).not.toHaveBeenCalled(); + }); + + it('reports an error if the auto-capture fails but does not throw', async () => { + memoryUsageMock.mockReturnValue({ + rss: 3 * 1024 * 1024 * 1024, + heapTotal: 0, + heapUsed: 0, + external: 0, + arrayBuffers: 0, + }); + vi.mocked(captureHeapSnapshot).mockRejectedValueOnce( + new Error('inspector failure'), + ); + const context = buildHighMemoryContext('/tmp/gemini-test'); + + if (!bugCommand.action) throw new Error('Action is not defined'); + await expect( + bugCommand.action(context, 'A memory bug'), + ).resolves.toBeUndefined(); + + const addItem = vi.mocked(context.ui.addItem).mock.calls; + const lastCall = addItem[addItem.length - 1][0]; + expect(lastCall.type).toBe(MessageType.ERROR); + expect(lastCall.text).toContain('inspector failure'); + }); }); diff --git a/packages/cli/src/ui/commands/bugCommand.ts b/packages/cli/src/ui/commands/bugCommand.ts index e146491dec..19bc7183d0 100644 --- a/packages/cli/src/ui/commands/bugCommand.ts +++ b/packages/cli/src/ui/commands/bugCommand.ts @@ -22,6 +22,11 @@ import { } from '@google/gemini-cli-core'; import { terminalCapabilityManager } from '../utils/terminalCapabilityManager.js'; import { exportHistoryToFile } from '../utils/historyExportUtils.js'; +import { + captureHeapSnapshot, + MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES, +} from '../utils/memorySnapshot.js'; +import { stat } from 'node:fs/promises'; import path from 'node:path'; export const bugCommand: SlashCommand = { @@ -129,6 +134,54 @@ export const bugCommand: SlashCommand = { Date.now(), ); } + + const rss = process.memoryUsage().rss; + const tempDir = config?.storage?.getProjectTempDir(); + if (rss >= MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES && tempDir) { + const snapshotPath = path.join( + tempDir, + `bug-memory-${Date.now()}.heapsnapshot`, + ); + context.ui.addItem( + { + type: MessageType.INFO, + text: `High memory usage detected (${formatBytes(rss)}). Capturing V8 heap snapshot to ${snapshotPath}.\nThis can take 20+ seconds and the CLI may be temporarily unresponsive; please do not exit.`, + }, + Date.now(), + ); + try { + const startedAt = Date.now(); + await captureHeapSnapshot(snapshotPath); + const durationMs = Date.now() - startedAt; + let sizeText = ''; + try { + const { size } = await stat(snapshotPath); + sizeText = ` (${formatBytes(size)})`; + } catch { + // Size reporting is best-effort; the snapshot itself was captured successfully. + } + context.ui.addItem( + { + type: MessageType.INFO, + text: `Heap snapshot saved${sizeText} in ${durationMs}ms:\n${snapshotPath}\n\nConsider attaching it to your bug report only if it does not contain sensitive information.`, + }, + Date.now(), + ); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : String(error); + debugLogger.error( + `Failed to capture heap snapshot for bug report: ${errorMessage}`, + ); + context.ui.addItem( + { + type: MessageType.ERROR, + text: `Failed to capture heap snapshot: ${errorMessage}`, + }, + Date.now(), + ); + } + } }, }; diff --git a/packages/cli/src/ui/commands/bugMemoryCommand.test.ts b/packages/cli/src/ui/commands/bugMemoryCommand.test.ts new file mode 100644 index 0000000000..8a93db9527 --- /dev/null +++ b/packages/cli/src/ui/commands/bugMemoryCommand.test.ts @@ -0,0 +1,121 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import path from 'node:path'; +import { bugMemoryCommand } from './bugMemoryCommand.js'; +import { captureHeapSnapshot } from '../utils/memorySnapshot.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { MessageType } from '../types.js'; +import type { Config } from '@google/gemini-cli-core'; + +vi.mock('../utils/memorySnapshot.js', () => ({ + captureHeapSnapshot: vi.fn(), + MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES: 2 * 1024 * 1024 * 1024, +})); + +vi.mock('node:fs/promises', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + stat: vi.fn().mockResolvedValue({ size: 1234 }), + }; +}); + +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + debugLogger: { + error: vi.fn(), + log: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + }, + }; +}); + +function makeContextWithTempDir(tempDir: string | undefined) { + return createMockCommandContext({ + services: { + agentContext: { + config: { + storage: tempDir ? { getProjectTempDir: () => tempDir } : undefined, + } as unknown as Config, + }, + }, + }); +} + +describe('bugMemoryCommand', () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2024-01-01T00:00:00Z')); + }); + + afterEach(() => { + vi.clearAllMocks(); + vi.useRealTimers(); + }); + + it('declares itself as a non-auto-executing built-in command', () => { + expect(bugMemoryCommand.name).toBe('bug-memory'); + expect(bugMemoryCommand.autoExecute).toBe(false); + expect(bugMemoryCommand.description).toBeTruthy(); + }); + + it('captures a heap snapshot and reports the file path', async () => { + const tempDir = path.join('/tmp', 'gemini-test'); + const context = makeContextWithTempDir(tempDir); + vi.mocked(captureHeapSnapshot).mockResolvedValueOnce(undefined); + + if (!bugMemoryCommand.action) throw new Error('Action missing'); + await bugMemoryCommand.action(context, ''); + + const expectedPath = path.join( + tempDir, + `bug-memory-${new Date('2024-01-01T00:00:00Z').getTime()}.heapsnapshot`, + ); + expect(captureHeapSnapshot).toHaveBeenCalledWith(expectedPath); + + const addItemCalls = vi.mocked(context.ui.addItem).mock.calls; + expect(addItemCalls).toHaveLength(2); + expect(addItemCalls[0][0]).toMatchObject({ type: MessageType.INFO }); + expect(addItemCalls[0][0].text).toContain(expectedPath); + expect(addItemCalls[1][0]).toMatchObject({ type: MessageType.INFO }); + expect(addItemCalls[1][0].text).toContain('Heap snapshot saved'); + expect(addItemCalls[1][0].text).toContain(expectedPath); + }); + + it('surfaces an error if capture fails', async () => { + const context = makeContextWithTempDir('/tmp/gemini-test'); + vi.mocked(captureHeapSnapshot).mockRejectedValueOnce( + new Error('inspector disconnected'), + ); + + if (!bugMemoryCommand.action) throw new Error('Action missing'); + await bugMemoryCommand.action(context, ''); + + const addItemCalls = vi.mocked(context.ui.addItem).mock.calls; + const lastCall = addItemCalls[addItemCalls.length - 1][0]; + expect(lastCall.type).toBe(MessageType.ERROR); + expect(lastCall.text).toContain('inspector disconnected'); + }); + + it('emits an error when no project temp directory is available', async () => { + const context = makeContextWithTempDir(undefined); + + if (!bugMemoryCommand.action) throw new Error('Action missing'); + await bugMemoryCommand.action(context, ''); + + expect(captureHeapSnapshot).not.toHaveBeenCalled(); + const addItemCalls = vi.mocked(context.ui.addItem).mock.calls; + expect(addItemCalls).toHaveLength(1); + expect(addItemCalls[0][0].type).toBe(MessageType.ERROR); + expect(addItemCalls[0][0].text).toContain('temp directory'); + }); +}); diff --git a/packages/cli/src/ui/commands/bugMemoryCommand.ts b/packages/cli/src/ui/commands/bugMemoryCommand.ts new file mode 100644 index 0000000000..cd43ce8902 --- /dev/null +++ b/packages/cli/src/ui/commands/bugMemoryCommand.ts @@ -0,0 +1,86 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { stat } from 'node:fs/promises'; +import path from 'node:path'; +import process from 'node:process'; +import { debugLogger } from '@google/gemini-cli-core'; +import { + type CommandContext, + type SlashCommand, + CommandKind, +} from './types.js'; +import { MessageType } from '../types.js'; +import { formatBytes } from '../utils/formatters.js'; +import { captureHeapSnapshot } from '../utils/memorySnapshot.js'; + +export const bugMemoryCommand: SlashCommand = { + name: 'bug-memory', + description: 'Capture a V8 heap snapshot to disk to attach to a bug report', + kind: CommandKind.BUILT_IN, + autoExecute: false, + action: async (context: CommandContext): Promise => { + const tempDir = + context.services.agentContext?.config?.storage?.getProjectTempDir(); + if (!tempDir) { + context.ui.addItem( + { + type: MessageType.ERROR, + text: 'Cannot capture heap snapshot: project temp directory is unavailable.', + }, + Date.now(), + ); + return; + } + + const filePath = path.join( + tempDir, + `bug-memory-${Date.now()}.heapsnapshot`, + ); + const rss = process.memoryUsage().rss; + + context.ui.addItem( + { + type: MessageType.INFO, + text: `Capturing V8 heap snapshot (current RSS: ${formatBytes(rss)}).\nThis can take 20+ seconds and the CLI may be temporarily unresponsive — please do not exit.\nDestination: ${filePath}`, + }, + Date.now(), + ); + + const startedAt = Date.now(); + try { + await captureHeapSnapshot(filePath); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + debugLogger.error(`Failed to capture heap snapshot: ${message}`); + context.ui.addItem( + { + type: MessageType.ERROR, + text: `Failed to capture heap snapshot: ${message}`, + }, + Date.now(), + ); + return; + } + + const durationMs = Date.now() - startedAt; + let sizeText = ''; + try { + const { size } = await stat(filePath); + sizeText = ` (${formatBytes(size)})`; + } catch { + // Size reporting is best-effort; the snapshot itself was captured successfully. + } + + context.ui.addItem( + { + type: MessageType.INFO, + text: `Heap snapshot saved${sizeText} in ${durationMs}ms:\n${filePath}\n\nLoad it in Chrome DevTools → Memory → "Load" to analyze. Attach it to your bug report only if it does not contain sensitive information.`, + }, + Date.now(), + ); + }, +}; diff --git a/packages/cli/src/ui/utils/memorySnapshot.test.ts b/packages/cli/src/ui/utils/memorySnapshot.test.ts new file mode 100644 index 0000000000..91fac95197 --- /dev/null +++ b/packages/cli/src/ui/utils/memorySnapshot.test.ts @@ -0,0 +1,84 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { Readable } from 'node:stream'; +import { + captureHeapSnapshot, + MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES, +} from './memorySnapshot.js'; + +const { mkdirMock, pipelineMock, getHeapSnapshotMock, createWriteStreamMock } = + vi.hoisted(() => ({ + mkdirMock: vi.fn(async () => undefined), + pipelineMock: vi.fn(async () => undefined), + getHeapSnapshotMock: vi.fn(), + createWriteStreamMock: vi.fn(), + })); + +vi.mock('node:fs/promises', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, mkdir: mkdirMock }; +}); + +vi.mock('node:fs', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, createWriteStream: createWriteStreamMock }; +}); + +vi.mock('node:v8', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, getHeapSnapshot: getHeapSnapshotMock }; +}); + +vi.mock('node:stream/promises', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, pipeline: pipelineMock }; +}); + +describe('captureHeapSnapshot', () => { + beforeEach(() => { + mkdirMock.mockClear(); + pipelineMock.mockClear(); + getHeapSnapshotMock.mockClear().mockReturnValue(Readable.from([])); + createWriteStreamMock + .mockClear() + .mockReturnValue({ write: vi.fn(), end: vi.fn() }); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it('exports the 2 GB auto-capture threshold', () => { + expect(MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES).toBe(2 * 1024 * 1024 * 1024); + }); + + it('creates the target directory and pipelines the V8 snapshot to disk', async () => { + const target = '/tmp/gemini-test/snapshot.heapsnapshot'; + + await captureHeapSnapshot(target); + + expect(mkdirMock).toHaveBeenCalledWith('/tmp/gemini-test', { + recursive: true, + }); + expect(getHeapSnapshotMock).toHaveBeenCalledTimes(1); + expect(createWriteStreamMock).toHaveBeenCalledWith(target); + expect(pipelineMock).toHaveBeenCalledTimes(1); + expect(pipelineMock).toHaveBeenCalledWith( + getHeapSnapshotMock.mock.results[0].value, + createWriteStreamMock.mock.results[0].value, + ); + }); + + it('propagates pipeline failures to the caller', async () => { + pipelineMock.mockRejectedValueOnce(new Error('write failed')); + + await expect( + captureHeapSnapshot('/tmp/gemini-test/fail.heapsnapshot'), + ).rejects.toThrow('write failed'); + }); +}); diff --git a/packages/cli/src/ui/utils/memorySnapshot.ts b/packages/cli/src/ui/utils/memorySnapshot.ts new file mode 100644 index 0000000000..746f3a5d0f --- /dev/null +++ b/packages/cli/src/ui/utils/memorySnapshot.ts @@ -0,0 +1,30 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { createWriteStream } from 'node:fs'; +import { mkdir } from 'node:fs/promises'; +import { dirname } from 'node:path'; +import { pipeline } from 'node:stream/promises'; +import { getHeapSnapshot } from 'node:v8'; + +/** + * RSS threshold at which `/bug` auto-captures a heap snapshot. + */ +export const MEMORY_SNAPSHOT_AUTO_THRESHOLD_BYTES = 2 * 1024 * 1024 * 1024; + +/** + * Capture a V8 heap snapshot from the current process and write it to disk. + * + * `v8.getHeapSnapshot()` returns a Readable stream whose producer is V8's + * internal snapshot generator. Piping it through `node:stream/promises`' + * `pipeline` propagates backpressure end-to-end, so even a multi-gigabyte + * heap is written without buffering the serialized snapshot in memory. + * Nothing is exposed over a debugger port. + */ +export async function captureHeapSnapshot(filePath: string): Promise { + await mkdir(dirname(filePath), { recursive: true }); + await pipeline(getHeapSnapshot(), createWriteStream(filePath)); +}