diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index f0f9cf19fa..85559aa2da 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -47,6 +47,7 @@ import { validateAuthMethod } from '../config/auth.js'; import { loadHierarchicalGeminiMemory } from '../config/config.js'; import process from 'node:process'; import { useHistory } from './hooks/useHistoryManager.js'; +import { useMemoryMonitor } from './hooks/useMemoryMonitor.js'; import { useThemeCommand } from './hooks/useThemeCommand.js'; import { useAuthCommand } from './auth/useAuth.js'; import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js'; @@ -123,6 +124,7 @@ const SHELL_HEIGHT_PADDING = 10; export const AppContainer = (props: AppContainerProps) => { const { settings, config, initializationResult } = props; const historyManager = useHistory(); + useMemoryMonitor(historyManager); const [corgiMode, setCorgiMode] = useState(false); const [debugMessage, setDebugMessage] = useState(''); const [quittingMessages, setQuittingMessages] = useState< diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.tsx index b14443e3c6..8748301968 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.tsx @@ -14,6 +14,7 @@ import { ErrorMessage } from './messages/ErrorMessage.js'; import { ToolGroupMessage } from './messages/ToolGroupMessage.js'; import { GeminiMessageContent } from './messages/GeminiMessageContent.js'; import { CompressionMessage } from './messages/CompressionMessage.js'; +import { WarningMessage } from './messages/WarningMessage.js'; import { Box } from 'ink'; import { AboutBox } from './AboutBox.js'; import { StatsDisplay } from './StatsDisplay.js'; @@ -66,6 +67,7 @@ export const HistoryItemDisplay: React.FC = ({ /> )} {item.type === 'info' && } + {item.type === 'warning' && } {item.type === 'error' && } {item.type === 'about' && ( = ({ text }) => { + const prefix = '⚠ '; + const prefixWidth = 3; + + return ( + + + {prefix} + + + + + + + + ); +}; diff --git a/packages/cli/src/ui/hooks/useMemoryMonitor.test.ts b/packages/cli/src/ui/hooks/useMemoryMonitor.test.ts new file mode 100644 index 0000000000..3250a33833 --- /dev/null +++ b/packages/cli/src/ui/hooks/useMemoryMonitor.test.ts @@ -0,0 +1,71 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { renderHook } from '@testing-library/react'; +import { vi } from 'vitest'; +import { + useMemoryMonitor, + MEMORY_CHECK_INTERVAL, + MEMORY_WARNING_THRESHOLD, +} from './useMemoryMonitor.js'; +import process from 'node:process'; +import { MessageType } from '../types.js'; + +describe('useMemoryMonitor', () => { + const memoryUsageSpy = vi.spyOn(process, 'memoryUsage'); + const addItem = vi.fn(); + + beforeEach(() => { + vi.useFakeTimers(); + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('should not warn when memory usage is below threshold', () => { + memoryUsageSpy.mockReturnValue({ + rss: MEMORY_WARNING_THRESHOLD / 2, + } as NodeJS.MemoryUsage); + renderHook(() => useMemoryMonitor({ addItem })); + vi.advanceTimersByTime(10000); + expect(addItem).not.toHaveBeenCalled(); + }); + + it('should warn when memory usage is above threshold', () => { + memoryUsageSpy.mockReturnValue({ + rss: MEMORY_WARNING_THRESHOLD * 1.5, + } as NodeJS.MemoryUsage); + renderHook(() => useMemoryMonitor({ addItem })); + vi.advanceTimersByTime(MEMORY_CHECK_INTERVAL); + expect(addItem).toHaveBeenCalledTimes(1); + expect(addItem).toHaveBeenCalledWith( + { + type: MessageType.WARNING, + text: 'High memory usage detected: 10.50 GB. If you experience a crash, please file a bug report by running `/bug`', + }, + expect.any(Number), + ); + }); + + it('should only warn once', () => { + memoryUsageSpy.mockReturnValue({ + rss: MEMORY_WARNING_THRESHOLD * 1.5, + } as NodeJS.MemoryUsage); + const { rerender } = renderHook(() => useMemoryMonitor({ addItem })); + vi.advanceTimersByTime(MEMORY_CHECK_INTERVAL); + expect(addItem).toHaveBeenCalledTimes(1); + + // Rerender and advance timers, should not warn again + memoryUsageSpy.mockReturnValue({ + rss: MEMORY_WARNING_THRESHOLD * 1.5, + } as NodeJS.MemoryUsage); + rerender(); + vi.advanceTimersByTime(MEMORY_CHECK_INTERVAL); + expect(addItem).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/cli/src/ui/hooks/useMemoryMonitor.ts b/packages/cli/src/ui/hooks/useMemoryMonitor.ts new file mode 100644 index 0000000000..7573eb0c2c --- /dev/null +++ b/packages/cli/src/ui/hooks/useMemoryMonitor.ts @@ -0,0 +1,41 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useEffect } from 'react'; +import process from 'node:process'; +import { type HistoryItemWithoutId, MessageType } from '../types.js'; + +export const MEMORY_WARNING_THRESHOLD = 7 * 1024 * 1024 * 1024; // 7GB in bytes +export const MEMORY_CHECK_INTERVAL = 60 * 1000; // one minute + +interface MemoryMonitorOptions { + addItem: (item: HistoryItemWithoutId, timestamp: number) => void; +} + +export const useMemoryMonitor = ({ addItem }: MemoryMonitorOptions) => { + useEffect(() => { + const intervalId = setInterval(() => { + const usage = process.memoryUsage().rss; + if (usage > MEMORY_WARNING_THRESHOLD) { + addItem( + { + type: MessageType.WARNING, + text: + `High memory usage detected: ${( + usage / + (1024 * 1024 * 1024) + ).toFixed(2)} GB. ` + + 'If you experience a crash, please file a bug report by running `/bug`', + }, + Date.now(), + ); + clearInterval(intervalId); + } + }, MEMORY_CHECK_INTERVAL); + + return () => clearInterval(intervalId); + }, [addItem]); +}; diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 3ab892b2cf..36a9bbb0fd 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -106,6 +106,11 @@ export type HistoryItemError = HistoryItemBase & { text: string; }; +export type HistoryItemWarning = HistoryItemBase & { + type: 'warning'; + text: string; +}; + export type HistoryItemAbout = HistoryItemBase & { type: 'about'; cliVersion: string; @@ -170,6 +175,7 @@ export type HistoryItemWithoutId = | HistoryItemGeminiContent | HistoryItemInfo | HistoryItemError + | HistoryItemWarning | HistoryItemAbout | HistoryItemHelp | HistoryItemToolGroup @@ -186,6 +192,7 @@ export type HistoryItem = HistoryItemWithoutId & { id: number }; export enum MessageType { INFO = 'info', ERROR = 'error', + WARNING = 'warning', USER = 'user', ABOUT = 'about', HELP = 'help', diff --git a/packages/cli/tsconfig.json b/packages/cli/tsconfig.json index 44424d671d..368b7d57e6 100644 --- a/packages/cli/tsconfig.json +++ b/packages/cli/tsconfig.json @@ -66,6 +66,7 @@ "src/ui/components/shared/vim-buffer-actions.test.ts", "src/ui/components/StatsDisplay.test.tsx", "src/ui/components/ToolStatsDisplay.test.tsx", + "src/ui/components/WarningMessage.test.tsx", "src/ui/contexts/SessionContext.test.tsx", "src/ui/hooks/slashCommandProcessor.test.ts", "src/ui/hooks/useAtCompletion.test.ts",