From e59c872b3dea9dee8206f990a761e3bfee3a1194 Mon Sep 17 00:00:00 2001 From: Jacob MacDonald Date: Tue, 17 Jun 2025 08:44:54 -0700 Subject: [PATCH] code review followup for compress command (#1097) Followup to https://github.com/google-gemini/gemini-cli/pull/986 --- .../messages/CompressionMessage.tsx | 3 +- .../ui/hooks/slashCommandProcessor.test.ts | 29 ++++++++++++++----- .../cli/src/ui/hooks/slashCommandProcessor.ts | 2 ++ packages/cli/src/ui/hooks/useGeminiStream.ts | 3 +- packages/cli/src/ui/types.ts | 4 +-- packages/core/src/core/client.ts | 10 ++++--- packages/core/src/core/turn.ts | 6 ++-- 7 files changed, 39 insertions(+), 18 deletions(-) diff --git a/packages/cli/src/ui/components/messages/CompressionMessage.tsx b/packages/cli/src/ui/components/messages/CompressionMessage.tsx index aaa56149..d1080668 100644 --- a/packages/cli/src/ui/components/messages/CompressionMessage.tsx +++ b/packages/cli/src/ui/components/messages/CompressionMessage.tsx @@ -23,7 +23,8 @@ export const CompressionMessage: React.FC = ({ }) => { const text = compression.isPending ? 'Compressing chat history' - : `Chat history compressed from ${compression.originalTokenCount} to ${compression.newTokenCount} tokens.`; + : `Chat history compressed from ${compression.originalTokenCount ?? 'unknown'}` + + ` to ${compression.newTokenCount ?? 'unknown'} tokens.`; return ( diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 68dc3ea9..2eec29ca 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -159,8 +159,8 @@ describe('useSlashCommandProcessor', () => { process.env = { ...globalThis.process.env }; }); - const getProcessor = (showToolDescriptions: boolean = false) => { - const { result } = renderHook(() => + const getProcessorHook = (showToolDescriptions: boolean = false) => + renderHook(() => useSlashCommandProcessor( mockConfig, [], @@ -178,8 +178,9 @@ describe('useSlashCommandProcessor', () => { mockSetQuittingMessages, ), ); - return result.current; - }; + + const getProcessor = (showToolDescriptions: boolean = false) => + getProcessorHook(showToolDescriptions).result.current; describe('/memory add', () => { it('should return tool scheduling info on valid input', async () => { @@ -1132,10 +1133,20 @@ Add any other context about the problem here. describe('/compress command', () => { it('should call tryCompressChat(true)', async () => { - const { handleSlashCommand } = getProcessor(); + const hook = getProcessorHook(); mockTryCompressChat.mockImplementationOnce(async (force?: boolean) => { - // TODO: Check that we have a pending compression item in the history. expect(force).toBe(true); + await act(async () => { + hook.rerender(); + }); + expect(hook.result.current.pendingHistoryItems).toContainEqual({ + type: MessageType.COMPRESSION, + compression: { + isPending: true, + originalTokenCount: null, + newTokenCount: null, + }, + }); return { originalTokenCount: 100, newTokenCount: 50, @@ -1143,8 +1154,12 @@ Add any other context about the problem here. }); await act(async () => { - handleSlashCommand('/compress'); + hook.result.current.handleSlashCommand('/compress'); }); + await act(async () => { + hook.rerender(); + }); + expect(hook.result.current.pendingHistoryItems).toEqual([]); expect(mockGeminiClient.tryCompressChat).toHaveBeenCalledWith(true); expect(mockAddItem).toHaveBeenNthCalledWith( 2, diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 916af3e7..8303ea55 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -760,6 +760,8 @@ Add any other context about the problem here. type: MessageType.COMPRESSION, compression: { isPending: true, + originalTokenCount: null, + newTokenCount: null, }, }); try { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index d2153edf..6d92af0d 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -379,7 +379,8 @@ export const useGeminiStream = ( text: `IMPORTANT: This conversation approached the input token limit for ${config.getModel()}. ` + `A compressed context will be sent for future messages (compressed from: ` + - `${eventValue.originalTokenCount} to ${eventValue.newTokenCount} tokens).`, + `${eventValue?.originalTokenCount ?? 'unknown'} to ` + + `${eventValue?.newTokenCount ?? 'unknown'} tokens).`, }, Date.now(), ), diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 3c0ec616..5cf372b4 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -55,8 +55,8 @@ export interface IndividualToolCallDisplay { export interface CompressionProps { isPending: boolean; - originalTokenCount?: number; - newTokenCount?: number; + originalTokenCount: number | null; + newTokenCount: number | null; } export interface HistoryItemBase { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 94c598bf..9cc8f328 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -475,9 +475,11 @@ export class GeminiClient { await cg.countTokens({ model: this.model, contents: newHistory }) ).totalTokens; - return { - originalTokenCount, - newTokenCount, - }; + return originalTokenCount && newTokenCount + ? { + originalTokenCount, + newTokenCount, + } + : null; } } diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 4c0a297e..5a8b57e0 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -110,13 +110,13 @@ export type ServerGeminiErrorEvent = { }; export interface ChatCompressionInfo { - originalTokenCount: number | undefined; - newTokenCount: number | undefined; + originalTokenCount: number; + newTokenCount: number; } export type ServerGeminiChatCompressedEvent = { type: GeminiEventType.ChatCompressed; - value: ChatCompressionInfo; + value: ChatCompressionInfo | null; }; export type ServerGeminiUsageMetadataEvent = {