From 01768d7759b81a0f1483c751206f6afcae6fc505 Mon Sep 17 00:00:00 2001 From: Jacob Richman Date: Fri, 30 May 2025 22:18:01 +0000 Subject: [PATCH] feat: add --show_memory_usage flag to display memory usage in status bar (#606) --- packages/cli/src/config/config.test.ts | 99 +++++++++++++++++-- packages/cli/src/config/config.ts | 7 ++ packages/cli/src/config/settings.ts | 1 + packages/cli/src/ui/App.tsx | 3 + packages/cli/src/ui/components/Footer.tsx | 5 + .../src/ui/components/MemoryUsageDisplay.tsx | 40 ++++++++ .../ui/hooks/slashCommandProcessor.test.ts | 35 ++++++- .../cli/src/ui/hooks/slashCommandProcessor.ts | 4 + packages/cli/src/ui/utils/formatters.ts | 16 +++ packages/server/src/config/config.ts | 7 ++ 10 files changed, 205 insertions(+), 12 deletions(-) create mode 100644 packages/cli/src/ui/components/MemoryUsageDisplay.tsx create mode 100644 packages/cli/src/ui/utils/formatters.ts diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index b24e4170..e443bbe5 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -7,16 +7,10 @@ // packages/cli/src/config/config.test.ts import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; -// import * as fsPromises from 'fs/promises'; -// import * as fsSync from 'fs'; import * as os from 'os'; -// import * as path from 'path'; // Unused, so removing -// import { readPackageUp } from 'read-package-up'; -// import { -// loadHierarchicalGeminiMemory, -// } from './config'; -// import { Settings } from './settings'; -// import * as ServerConfig from '@gemini-code/server'; +import { loadCliConfig } from './config.js'; +import { Settings } from './settings.js'; +import * as ServerConfig from '@gemini-code/server'; const MOCK_HOME_DIR = '/mock/home/user'; @@ -28,7 +22,92 @@ vi.mock('os', async (importOriginal) => { }; }); -// Further mocking of fs, read-package-up, etc. would go here if tests were active. +vi.mock('read-package-up', () => ({ + readPackageUp: vi.fn(() => + Promise.resolve({ packageJson: { version: 'test-version' } }), + ), +})); + +vi.mock('@gemini-code/server', async () => { + const actualServer = await vi.importActual( + '@gemini-code/server', + ); + return { + ...actualServer, + loadEnvironment: vi.fn(), + createServerConfig: vi.fn((params) => ({ + // Mock the config object and its methods + getApiKey: () => params.apiKey, + getModel: () => params.model, + getSandbox: () => params.sandbox, + getTargetDir: () => params.targetDir, + getDebugMode: () => params.debugMode, + getQuestion: () => params.question, + getFullContext: () => params.fullContext, + getCoreTools: () => params.coreTools, + getToolDiscoveryCommand: () => params.toolDiscoveryCommand, + getToolCallCommand: () => params.toolCallCommand, + getMcpServerCommand: () => params.mcpServerCommand, + getMcpServers: () => params.mcpServers, + getUserAgent: () => params.userAgent, + getUserMemory: () => params.userMemory, + getGeminiMdFileCount: () => params.geminiMdFileCount, + getVertexAI: () => params.vertexai, + getShowMemoryUsage: () => params.showMemoryUsage, // Added for the test + // Add any other methods that are called on the config object + setUserMemory: vi.fn(), + setGeminiMdFileCount: vi.fn(), + })), + loadServerHierarchicalMemory: vi.fn(() => + Promise.resolve({ memoryContent: '', fileCount: 0 }), + ), + }; +}); + +describe('loadCliConfig', () => { + const originalArgv = process.argv; + const originalEnv = { ...process.env }; + + beforeEach(() => { + vi.resetAllMocks(); + vi.mocked(os.homedir).mockReturnValue(MOCK_HOME_DIR); + process.env.GEMINI_API_KEY = 'test-api-key'; // Ensure API key is set for tests + }); + + afterEach(() => { + process.argv = originalArgv; + process.env = originalEnv; + vi.restoreAllMocks(); + }); + + it('should set showMemoryUsage to true when --memory flag is present', async () => { + process.argv = ['node', 'script.js', '--show_memory_usage']; + const settings: Settings = {}; + const config = await loadCliConfig(settings); + expect(config.getShowMemoryUsage()).toBe(true); + }); + + it('should set showMemoryUsage to false when --memory flag is not present', async () => { + process.argv = ['node', 'script.js']; + const settings: Settings = {}; + const config = await loadCliConfig(settings); + expect(config.getShowMemoryUsage()).toBe(false); + }); + + it('should set showMemoryUsage to false by default from settings if CLI flag is not present', async () => { + process.argv = ['node', 'script.js']; + const settings: Settings = { showMemoryUsage: false }; + const config = await loadCliConfig(settings); + expect(config.getShowMemoryUsage()).toBe(false); + }); + + it('should prioritize CLI flag over settings for showMemoryUsage (CLI true, settings false)', async () => { + process.argv = ['node', 'script.js', '--show_memory_usage']; + const settings: Settings = { showMemoryUsage: false }; + const config = await loadCliConfig(settings); + expect(config.getShowMemoryUsage()).toBe(true); + }); +}); describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => { beforeEach(() => { diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 14b02538..1808f545 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -35,6 +35,7 @@ interface CliArgs { debug: boolean | undefined; prompt: string | undefined; all_files: boolean | undefined; + show_memory_usage: boolean | undefined; } async function parseArguments(): Promise { @@ -67,6 +68,11 @@ async function parseArguments(): Promise { description: 'Include ALL files in context?', default: false, }) + .option('show_memory_usage', { + type: 'boolean', + description: 'Show memory usage in status bar', + default: false, + }) .help() .alias('h', 'help') .strict().argv; @@ -152,6 +158,7 @@ export async function loadCliConfig(settings: Settings): Promise { userMemory: memoryContent, geminiMdFileCount: fileCount, vertexai: useVertexAI, + showMemoryUsage: argv.show_memory_usage || false, }; return createServerConfig(configParams); diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index 25ec5cb4..efb083cc 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -27,6 +27,7 @@ export interface Settings { toolCallCommand?: string; mcpServerCommand?: string; mcpServers?: Record; + showMemoryUsage?: boolean; // Add other settings here. } diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index dcd4d490..7b1eb2cb 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -442,6 +442,9 @@ export const App = ({ corgiMode={corgiMode} errorCount={errorCount} showErrorDetails={showErrorDetails} + showMemoryUsage={ + config.getDebugMode() || config.getShowMemoryUsage() + } /> diff --git a/packages/cli/src/ui/components/Footer.tsx b/packages/cli/src/ui/components/Footer.tsx index 04a2f96f..22615779 100644 --- a/packages/cli/src/ui/components/Footer.tsx +++ b/packages/cli/src/ui/components/Footer.tsx @@ -9,6 +9,8 @@ import { Box, Text } from 'ink'; import { Colors } from '../colors.js'; import { shortenPath, tildeifyPath } from '@gemini-code/server'; import { ConsoleSummaryDisplay } from './ConsoleSummaryDisplay.js'; +import process from 'node:process'; +import { MemoryUsageDisplay } from './MemoryUsageDisplay.js'; interface FooterProps { model: string; @@ -20,6 +22,7 @@ interface FooterProps { corgiMode: boolean; errorCount: number; showErrorDetails: boolean; + showMemoryUsage?: boolean; } export const Footer: React.FC = ({ @@ -31,6 +34,7 @@ export const Footer: React.FC = ({ corgiMode, errorCount, showErrorDetails, + showMemoryUsage, }) => ( @@ -86,6 +90,7 @@ export const Footer: React.FC = ({ )} + {showMemoryUsage && } ); diff --git a/packages/cli/src/ui/components/MemoryUsageDisplay.tsx b/packages/cli/src/ui/components/MemoryUsageDisplay.tsx new file mode 100644 index 00000000..5d9a7c49 --- /dev/null +++ b/packages/cli/src/ui/components/MemoryUsageDisplay.tsx @@ -0,0 +1,40 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import React, { useEffect, useState } from 'react'; +import { Box, Text } from 'ink'; +import { Colors } from '../colors.js'; +import process from 'node:process'; +import { formatMemoryUsage } from '../utils/formatters.js'; + +export const MemoryUsageDisplay: React.FC = () => { + const [memoryUsage, setMemoryUsage] = useState(''); + const [memoryUsageColor, setMemoryUsageColor] = useState( + Colors.SubtleComment, + ); + + useEffect(() => { + const updateMemory = () => { + const usage = process.memoryUsage().rss; + setMemoryUsage(formatMemoryUsage(usage)); + setMemoryUsageColor( + usage >= 2 * 1024 * 1024 * 1024 + ? Colors.AccentRed + : Colors.SubtleComment, + ); + }; + const intervalId = setInterval(updateMemory, 2000); + updateMemory(); // Initial update + return () => clearInterval(intervalId); + }, []); + + return ( + + | + {memoryUsage} + + ); +}; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 1f049658..7696c40d 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -9,9 +9,37 @@ const { mockProcessExit } = vi.hoisted(() => ({ })); vi.mock('node:process', () => ({ + default: { + exit: mockProcessExit, + cwd: vi.fn(() => '/mock/cwd'), + get env() { + return process.env; + }, // Use a getter to ensure current process.env is used + platform: 'test-platform', + version: 'test-node-version', + memoryUsage: vi.fn(() => ({ + rss: 12345678, + heapTotal: 23456789, + heapUsed: 10234567, + external: 1234567, + arrayBuffers: 123456, + })), + }, + // Provide top-level exports as well for compatibility exit: mockProcessExit, cwd: vi.fn(() => '/mock/cwd'), - env: { ...process.env }, + get env() { + return process.env; + }, // Use a getter here too + platform: 'test-platform', + version: 'test-node-version', + memoryUsage: vi.fn(() => ({ + rss: 12345678, + heapTotal: 23456789, + heapUsed: 10234567, + external: 1234567, + arrayBuffers: 123456, + })), })); vi.mock('node:fs/promises', () => ({ @@ -227,7 +255,7 @@ describe('useSlashCommandProcessor', () => { seatbeltProfileVar?: string, ) => { const cliVersion = 'test-version'; - const osVersion = `${process.platform} ${process.version}`; + const osVersion = 'test-platform test-node-version'; let sandboxEnvStr = 'no sandbox'; if (sandboxEnvVar && sandboxEnvVar !== 'sandbox-exec') { sandboxEnvStr = sandboxEnvVar.replace(/^gemini-(?:code-)?/, ''); @@ -235,6 +263,8 @@ describe('useSlashCommandProcessor', () => { sandboxEnvStr = `sandbox-exec (${seatbeltProfileVar || 'unknown'})`; } const modelVersion = 'test-model'; + // Use the mocked memoryUsage value + const memoryUsage = '11.8 MB'; const diagnosticInfo = ` ## Describe the bug @@ -249,6 +279,7 @@ Add any other context about the problem here. * **Operating System:** ${osVersion} * **Sandbox Environment:** ${sandboxEnvStr} * **Model Version:** ${modelVersion} +* **Memory Usage:** ${memoryUsage} `; let url = 'https://github.com/google-gemini/gemini-cli/issues/new?template=bug_report.md'; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 5ec07c91..8f294380 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -7,11 +7,13 @@ import { useCallback, useMemo } from 'react'; import { type PartListUnion } from '@google/genai'; import open from 'open'; +import process from 'node:process'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; import { Config } from '@gemini-code/server'; import { Message, MessageType, HistoryItemWithoutId } from '../types.js'; import { createShowMemoryAction } from './useShowMemoryCommand.js'; import { GIT_COMMIT_INFO } from '../../generated/git-commit.js'; +import { formatMemoryUsage } from '../utils/formatters.js'; export interface SlashCommandActionReturn { shouldScheduleTool?: boolean; @@ -206,6 +208,7 @@ export const useSlashCommandProcessor = ( sandboxEnv = `sandbox-exec (${process.env.SEATBELT_PROFILE || 'unknown'})`; } const modelVersion = config?.getModel() || 'Unknown'; + const memoryUsage = formatMemoryUsage(process.memoryUsage().rss); const diagnosticInfo = ` ## Describe the bug @@ -220,6 +223,7 @@ Add any other context about the problem here. * **Operating System:** ${osVersion} * **Sandbox Environment:** ${sandboxEnv} * **Model Version:** ${modelVersion} +* **Memory Usage:** ${memoryUsage} `; let bugReportUrl = diff --git a/packages/cli/src/ui/utils/formatters.ts b/packages/cli/src/ui/utils/formatters.ts new file mode 100644 index 00000000..ab02160e --- /dev/null +++ b/packages/cli/src/ui/utils/formatters.ts @@ -0,0 +1,16 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export const formatMemoryUsage = (bytes: number): string => { + const gb = bytes / (1024 * 1024 * 1024); + if (bytes < 1024 * 1024) { + return `${(bytes / 1024).toFixed(1)} KB`; + } + if (bytes < 1024 * 1024 * 1024) { + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; + } + return `${gb.toFixed(2)} GB`; +}; diff --git a/packages/server/src/config/config.ts b/packages/server/src/config/config.ts index fd9ed954..9c03a5c1 100644 --- a/packages/server/src/config/config.ts +++ b/packages/server/src/config/config.ts @@ -54,6 +54,7 @@ export interface ConfigParameters { geminiMdFileCount?: number; alwaysSkipModificationConfirmation?: boolean; vertexai?: boolean; + showMemoryUsage?: boolean; } export class Config { @@ -75,6 +76,7 @@ export class Config { private geminiMdFileCount: number; private alwaysSkipModificationConfirmation: boolean; private readonly vertexai: boolean | undefined; + private readonly showMemoryUsage: boolean; constructor(params: ConfigParameters) { this.apiKey = params.apiKey; @@ -95,6 +97,7 @@ export class Config { this.alwaysSkipModificationConfirmation = params.alwaysSkipModificationConfirmation ?? false; this.vertexai = params.vertexai; + this.showMemoryUsage = params.showMemoryUsage ?? false; this.toolRegistry = createToolRegistry(this); } @@ -181,6 +184,10 @@ export class Config { getVertexAI(): boolean | undefined { return this.vertexai; } + + getShowMemoryUsage(): boolean { + return this.showMemoryUsage; + } } function findEnvFile(startDir: string): string | null {