From b72e3dfb43acc3faf8099fa758a4283b83f32ff6 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Tue, 15 Jul 2025 21:59:16 -0400 Subject: [PATCH] migrate compress command (#4271) --- .../cli/src/services/CommandService.test.ts | 8 +- packages/cli/src/services/CommandService.ts | 2 + .../cli/src/test-utils/mockCommandContext.ts | 2 + .../src/ui/commands/compressCommand.test.ts | 129 ++++++++++++++++++ .../cli/src/ui/commands/compressCommand.ts | 77 +++++++++++ packages/cli/src/ui/commands/types.ts | 14 +- .../ui/hooks/slashCommandProcessor.test.ts | 34 ----- .../cli/src/ui/hooks/slashCommandProcessor.ts | 60 +------- 8 files changed, 230 insertions(+), 96 deletions(-) create mode 100644 packages/cli/src/ui/commands/compressCommand.test.ts create mode 100644 packages/cli/src/ui/commands/compressCommand.ts diff --git a/packages/cli/src/services/CommandService.test.ts b/packages/cli/src/services/CommandService.test.ts index b1f6e496..1ee78e8c 100644 --- a/packages/cli/src/services/CommandService.test.ts +++ b/packages/cli/src/services/CommandService.test.ts @@ -16,6 +16,7 @@ import { themeCommand } from '../ui/commands/themeCommand.js'; import { statsCommand } from '../ui/commands/statsCommand.js'; import { privacyCommand } from '../ui/commands/privacyCommand.js'; import { aboutCommand } from '../ui/commands/aboutCommand.js'; +import { compressCommand } from '../ui/commands/compressCommand.js'; import { extensionsCommand } from '../ui/commands/extensionsCommand.js'; // Mock the command modules to isolate the service from the command implementations. @@ -43,12 +44,15 @@ vi.mock('../ui/commands/statsCommand.js', () => ({ vi.mock('../ui/commands/aboutCommand.js', () => ({ aboutCommand: { name: 'about', description: 'Mock About' }, })); +vi.mock('../ui/commands/compressCommand.js', () => ({ + compressCommand: { name: 'compress', description: 'Mock Compress' }, +})); vi.mock('../ui/commands/extensionsCommand.js', () => ({ extensionsCommand: { name: 'extensions', description: 'Mock Extensions' }, })); describe('CommandService', () => { - const subCommandLen = 10; + const subCommandLen = 11; describe('when using default production loader', () => { let commandService: CommandService; @@ -85,6 +89,7 @@ describe('CommandService', () => { expect(commandNames).toContain('stats'); expect(commandNames).toContain('privacy'); expect(commandNames).toContain('about'); + expect(commandNames).toContain('compress'); expect(commandNames).toContain('extensions'); }); @@ -116,6 +121,7 @@ describe('CommandService', () => { authCommand, chatCommand, clearCommand, + compressCommand, extensionsCommand, helpCommand, memoryCommand, diff --git a/packages/cli/src/services/CommandService.ts b/packages/cli/src/services/CommandService.ts index 50f2c63a..6c81cd0c 100644 --- a/packages/cli/src/services/CommandService.ts +++ b/packages/cli/src/services/CommandService.ts @@ -14,6 +14,7 @@ import { chatCommand } from '../ui/commands/chatCommand.js'; import { statsCommand } from '../ui/commands/statsCommand.js'; import { privacyCommand } from '../ui/commands/privacyCommand.js'; import { aboutCommand } from '../ui/commands/aboutCommand.js'; +import { compressCommand } from '../ui/commands/compressCommand.js'; import { extensionsCommand } from '../ui/commands/extensionsCommand.js'; const loadBuiltInCommands = async (): Promise => [ @@ -21,6 +22,7 @@ const loadBuiltInCommands = async (): Promise => [ authCommand, chatCommand, clearCommand, + compressCommand, extensionsCommand, helpCommand, memoryCommand, diff --git a/packages/cli/src/test-utils/mockCommandContext.ts b/packages/cli/src/test-utils/mockCommandContext.ts index bf7d814d..88da4a32 100644 --- a/packages/cli/src/test-utils/mockCommandContext.ts +++ b/packages/cli/src/test-utils/mockCommandContext.ts @@ -44,6 +44,8 @@ export const createMockCommandContext = ( addItem: vi.fn(), clear: vi.fn(), setDebugMessage: vi.fn(), + pendingItem: null, + setPendingItem: vi.fn(), }, session: { stats: { diff --git a/packages/cli/src/ui/commands/compressCommand.test.ts b/packages/cli/src/ui/commands/compressCommand.test.ts new file mode 100644 index 00000000..95a8bda3 --- /dev/null +++ b/packages/cli/src/ui/commands/compressCommand.test.ts @@ -0,0 +1,129 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { GeminiClient } from '@google/gemini-cli-core'; +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { compressCommand } from './compressCommand.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { MessageType } from '../types.js'; + +describe('compressCommand', () => { + let context: ReturnType; + let mockTryCompressChat: ReturnType; + + beforeEach(() => { + mockTryCompressChat = vi.fn(); + context = createMockCommandContext({ + services: { + config: { + getGeminiClient: () => + ({ + tryCompressChat: mockTryCompressChat, + }) as unknown as GeminiClient, + }, + }, + }); + }); + + it('should do nothing if a compression is already pending', async () => { + context.ui.pendingItem = { + type: MessageType.COMPRESSION, + compression: { + isPending: true, + originalTokenCount: null, + newTokenCount: null, + }, + }; + await compressCommand.action!(context, ''); + expect(context.ui.addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: MessageType.ERROR, + text: 'Already compressing, wait for previous request to complete', + }), + expect.any(Number), + ); + expect(context.ui.setPendingItem).not.toHaveBeenCalled(); + expect(mockTryCompressChat).not.toHaveBeenCalled(); + }); + + it('should set pending item, call tryCompressChat, and add result on success', async () => { + const compressedResult = { + originalTokenCount: 200, + newTokenCount: 100, + }; + mockTryCompressChat.mockResolvedValue(compressedResult); + + await compressCommand.action!(context, ''); + + expect(context.ui.setPendingItem).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + type: MessageType.COMPRESSION, + compression: { + isPending: true, + originalTokenCount: null, + newTokenCount: null, + }, + }), + ); + + expect(mockTryCompressChat).toHaveBeenCalledWith( + expect.stringMatching(/^compress-\d+$/), + true, + ); + + expect(context.ui.addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: MessageType.COMPRESSION, + compression: { + isPending: false, + originalTokenCount: 200, + newTokenCount: 100, + }, + }), + expect.any(Number), + ); + + expect(context.ui.setPendingItem).toHaveBeenNthCalledWith(2, null); + }); + + it('should add an error message if tryCompressChat returns falsy', async () => { + mockTryCompressChat.mockResolvedValue(null); + + await compressCommand.action!(context, ''); + + expect(context.ui.addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: MessageType.ERROR, + text: 'Failed to compress chat history.', + }), + expect.any(Number), + ); + expect(context.ui.setPendingItem).toHaveBeenCalledWith(null); + }); + + it('should add an error message if tryCompressChat throws', async () => { + const error = new Error('Compression failed'); + mockTryCompressChat.mockRejectedValue(error); + + await compressCommand.action!(context, ''); + + expect(context.ui.addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: MessageType.ERROR, + text: `Failed to compress chat history: ${error.message}`, + }), + expect.any(Number), + ); + expect(context.ui.setPendingItem).toHaveBeenCalledWith(null); + }); + + it('should clear the pending item in a finally block', async () => { + mockTryCompressChat.mockRejectedValue(new Error('some error')); + await compressCommand.action!(context, ''); + expect(context.ui.setPendingItem).toHaveBeenCalledWith(null); + }); +}); diff --git a/packages/cli/src/ui/commands/compressCommand.ts b/packages/cli/src/ui/commands/compressCommand.ts new file mode 100644 index 00000000..c3dfdf37 --- /dev/null +++ b/packages/cli/src/ui/commands/compressCommand.ts @@ -0,0 +1,77 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { HistoryItemCompression, MessageType } from '../types.js'; +import { SlashCommand } from './types.js'; + +export const compressCommand: SlashCommand = { + name: 'compress', + altName: 'summarize', + description: 'Compresses the context by replacing it with a summary.', + action: async (context) => { + const { ui } = context; + if (ui.pendingItem) { + ui.addItem( + { + type: MessageType.ERROR, + text: 'Already compressing, wait for previous request to complete', + }, + Date.now(), + ); + return; + } + + const pendingMessage: HistoryItemCompression = { + type: MessageType.COMPRESSION, + compression: { + isPending: true, + originalTokenCount: null, + newTokenCount: null, + }, + }; + + try { + ui.setPendingItem(pendingMessage); + const promptId = `compress-${Date.now()}`; + const compressed = await context.services.config + ?.getGeminiClient() + ?.tryCompressChat(promptId, true); + if (compressed) { + ui.addItem( + { + type: MessageType.COMPRESSION, + compression: { + isPending: false, + originalTokenCount: compressed.originalTokenCount, + newTokenCount: compressed.newTokenCount, + }, + } as HistoryItemCompression, + Date.now(), + ); + } else { + ui.addItem( + { + type: MessageType.ERROR, + text: 'Failed to compress chat history.', + }, + Date.now(), + ); + } + } catch (e) { + ui.addItem( + { + type: MessageType.ERROR, + text: `Failed to compress chat history: ${ + e instanceof Error ? e.message : String(e) + }`, + }, + Date.now(), + ); + } finally { + ui.setPendingItem(null); + } + }, +}; diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index 27db2be2..85a85abe 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -23,11 +23,6 @@ export interface CommandContext { }; // UI state and history management ui: { - // TODO - As more commands are add some additions may be needed or reworked using this new context. - // Ex. - // history: HistoryItem[]; - // pendingHistoryItems: HistoryItemWithoutId[]; - /** Adds a new item to the history display. */ addItem: UseHistoryManagerReturn['addItem']; /** Clears all history items and the console screen. */ @@ -36,6 +31,15 @@ export interface CommandContext { * Sets the transient debug message displayed in the application footer in debug mode. */ setDebugMessage: (message: string) => void; + /** The currently pending history item, if any. */ + pendingItem: HistoryItemWithoutId | null; + /** + * Sets a pending item in the history, which is useful for indicating + * that a long-running operation is in progress. + * + * @param item The history item to display as pending, or `null` to clear. + */ + setPendingItem: (item: HistoryItemWithoutId | null) => void; }; // Session-specific data session: { diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index f39795c0..20c8d7fe 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -1204,38 +1204,4 @@ describe('useSlashCommandProcessor', () => { expect(commandResult).toEqual({ type: 'handled' }); }); }); - - describe('/compress command', () => { - it('should call tryCompressChat(true)', async () => { - const hook = getProcessorHook(); - mockTryCompressChat.mockResolvedValue({ - originalTokenCount: 100, - newTokenCount: 50, - }); - - await act(async () => { - hook.result.current.handleSlashCommand('/compress'); - }); - await act(async () => { - hook.rerender(); - }); - expect(hook.result.current.pendingHistoryItems).toEqual([]); - expect(mockGeminiClient.tryCompressChat).toHaveBeenCalledWith( - 'Prompt Id not set', - true, - ); - expect(mockAddItem).toHaveBeenNthCalledWith( - 2, - expect.objectContaining({ - type: MessageType.COMPRESSION, - compression: { - isPending: false, - originalTokenCount: 100, - newTokenCount: 50, - }, - }), - expect.any(Number), - ); - }); - }); }); diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 181c4980..67dbfcdd 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -169,6 +169,8 @@ export const useSlashCommandProcessor = ( refreshStatic(); }, setDebugMessage: onDebugMessage, + pendingItem: pendingCompressionItemRef.current, + setPendingItem: setPendingCompressionItem, }, session: { stats: session.stats, @@ -184,6 +186,8 @@ export const useSlashCommandProcessor = ( refreshStatic, session.stats, onDebugMessage, + pendingCompressionItemRef, + setPendingCompressionItem, ], ); @@ -599,60 +603,6 @@ export const useSlashCommandProcessor = ( }, 100); }, }, - { - name: 'compress', - altName: 'summarize', - description: 'Compresses the context by replacing it with a summary.', - action: async (_mainCommand, _subCommand, _args) => { - if (pendingCompressionItemRef.current !== null) { - addMessage({ - type: MessageType.ERROR, - content: - 'Already compressing, wait for previous request to complete', - timestamp: new Date(), - }); - return; - } - setPendingCompressionItem({ - type: MessageType.COMPRESSION, - compression: { - isPending: true, - originalTokenCount: null, - newTokenCount: null, - }, - }); - try { - const compressed = await config! - .getGeminiClient()! - // TODO: Set Prompt id for CompressChat from SlashCommandProcessor. - .tryCompressChat('Prompt Id not set', true); - if (compressed) { - addMessage({ - type: MessageType.COMPRESSION, - compression: { - isPending: false, - originalTokenCount: compressed.originalTokenCount, - newTokenCount: compressed.newTokenCount, - }, - timestamp: new Date(), - }); - } else { - addMessage({ - type: MessageType.ERROR, - content: 'Failed to compress chat history.', - timestamp: new Date(), - }); - } - } catch (e) { - addMessage({ - type: MessageType.ERROR, - content: `Failed to compress chat history: ${e instanceof Error ? e.message : String(e)}`, - timestamp: new Date(), - }); - } - setPendingCompressionItem(null); - }, - }, ]; if (config?.getCheckpointingEnabled()) { @@ -786,8 +736,6 @@ export const useSlashCommandProcessor = ( gitService, loadHistory, setQuittingMessages, - pendingCompressionItemRef, - setPendingCompressionItem, ]); const handleSlashCommand = useCallback(