diff options
Diffstat (limited to 'packages/cli/src/ui/hooks')
| -rw-r--r-- | packages/cli/src/ui/hooks/slashCommandProcessor.test.ts | 39 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/slashCommandProcessor.ts | 68 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 10 |
3 files changed, 113 insertions, 4 deletions
diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index c2873bd6..73669651 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -62,6 +62,7 @@ import { getMCPServerStatus, MCPDiscoveryState, getMCPDiscoveryState, + GeminiClient, } from '@gemini-cli/core'; import { useSessionStats } from '../contexts/SessionContext.js'; @@ -100,6 +101,8 @@ describe('useSlashCommandProcessor', () => { let mockOpenEditorDialog: ReturnType<typeof vi.fn>; let mockPerformMemoryRefresh: ReturnType<typeof vi.fn>; let mockSetQuittingMessages: ReturnType<typeof vi.fn>; + let mockTryCompressChat: ReturnType<typeof vi.fn>; + let mockGeminiClient: GeminiClient; let mockConfig: Config; let mockCorgiMode: ReturnType<typeof vi.fn>; const mockUseSessionStats = useSessionStats as Mock; @@ -115,8 +118,13 @@ describe('useSlashCommandProcessor', () => { mockOpenEditorDialog = vi.fn(); mockPerformMemoryRefresh = vi.fn().mockResolvedValue(undefined); mockSetQuittingMessages = vi.fn(); + mockTryCompressChat = vi.fn(); + mockGeminiClient = { + tryCompressChat: mockTryCompressChat, + } as unknown as GeminiClient; mockConfig = { getDebugMode: vi.fn(() => false), + getGeminiClient: () => mockGeminiClient, getSandbox: vi.fn(() => 'test-sandbox'), getModel: vi.fn(() => 'test-model'), getProjectRoot: vi.fn(() => '/test/dir'), @@ -944,4 +952,35 @@ Add any other context about the problem here. expect(commandResult).toBe(true); }); }); + + describe('/compress command', () => { + it('should call tryCompressChat(true)', async () => { + const { handleSlashCommand } = getProcessor(); + mockTryCompressChat.mockImplementationOnce(async (force?: boolean) => { + // TODO: Check that we have a pending compression item in the history. + expect(force).toBe(true); + return { + originalTokenCount: 100, + newTokenCount: 50, + }; + }); + + await act(async () => { + handleSlashCommand('/compress'); + }); + expect(mockGeminiClient.tryCompressChat).toHaveBeenCalledWith(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 861d7bd9..97374e4f 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -9,6 +9,7 @@ import { type PartListUnion } from '@google/genai'; import open from 'open'; import process from 'node:process'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; +import { useStateAndRef } from './useStateAndRef.js'; import { Config, GitService, @@ -80,6 +81,13 @@ export const useSlashCommandProcessor = ( return new GitService(config.getProjectRoot()); }, [config]); + const pendingHistoryItems: HistoryItemWithoutId[] = []; + const [pendingCompressionItemRef, setPendingCompressionItem] = + useStateAndRef<HistoryItemWithoutId | null>(null); + if (pendingCompressionItemRef.current != null) { + pendingHistoryItems.push(pendingCompressionItemRef.current); + } + const addMessage = useCallback( (message: Message) => { // Convert Message to HistoryItemWithoutId @@ -105,6 +113,11 @@ export const useSlashCommandProcessor = ( stats: message.stats, duration: message.duration, }; + } else if (message.type === MessageType.COMPRESSION) { + historyItemContent = { + type: 'compression', + compression: message.compression, + }; } else { historyItemContent = { type: message.type as @@ -641,6 +654,57 @@ Add any other context about the problem here. }, 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, + }, + }); + try { + const compressed = await config! + .getGeminiClient()! + .tryCompressChat(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?.getCheckpointEnabled()) { @@ -767,6 +831,8 @@ Add any other context about the problem here. loadHistory, addItem, setQuittingMessages, + pendingCompressionItemRef, + setPendingCompressionItem, ]); const handleSlashCommand = useCallback( @@ -830,5 +896,5 @@ Add any other context about the problem here. [addItem, slashCommands, addMessage], ); - return { handleSlashCommand, slashCommands }; + return { handleSlashCommand, slashCommands, pendingHistoryItems }; }; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 920ec490..bff38a2b 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -13,6 +13,7 @@ import { ServerGeminiStreamEvent as GeminiEvent, ServerGeminiContentEvent as ContentEvent, ServerGeminiErrorEvent as ErrorEvent, + ServerGeminiChatCompressedEvent, getErrorMessage, isNodeError, MessageSenderType, @@ -368,11 +369,14 @@ export const useGeminiStream = ( ); const handleChatCompressionEvent = useCallback( - () => + (eventValue: ServerGeminiChatCompressedEvent['value']) => addItem( { type: 'info', - text: `IMPORTANT: this conversation approached the input token limit for ${config.getModel()}. We'll send a compressed context to the model for any future messages.`, + 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).`, }, Date.now(), ), @@ -406,7 +410,7 @@ export const useGeminiStream = ( handleErrorEvent(event.value, userMessageTimestamp); break; case ServerGeminiEventType.ChatCompressed: - handleChatCompressionEvent(); + handleChatCompressionEvent(event.value); break; case ServerGeminiEventType.UsageMetadata: addUsage(event.value); |