diff --git a/packages/cli/src/ui/commands/compressCommand.test.ts b/packages/cli/src/ui/commands/compressCommand.test.ts index ecf3d85ccd..8e49850276 100644 --- a/packages/cli/src/ui/commands/compressCommand.test.ts +++ b/packages/cli/src/ui/commands/compressCommand.test.ts @@ -4,28 +4,37 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { - CompressionStatus, - type ChatCompressionInfo, - type GeminiClient, -} from '@google/gemini-cli-core'; +import * as Core 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'; +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = (await importOriginal()) as any; + return { + ...actual, + tokenLimit: vi.fn(), + }; +}); + describe('compressCommand', () => { let context: ReturnType; let mockTryCompressChat: ReturnType; beforeEach(() => { mockTryCompressChat = vi.fn(); + vi.mocked(Core.tokenLimit).mockReturnValue(1000); context = createMockCommandContext({ services: { agentContext: { + config: { + getModel: () => 'test-model', + getContextWindowCompressionThreshold: () => 0.2, + }, geminiClient: { tryCompressChat: mockTryCompressChat, - } as unknown as GeminiClient, + } as unknown as Core.GeminiClient, }, }, }); @@ -36,10 +45,10 @@ describe('compressCommand', () => { type: MessageType.COMPRESSION, compression: { isPending: true, - originalTokenCount: null, - newTokenCount: null, + beforePercentage: null, + afterPercentage: null, + threshold: null, compressionStatus: null, - model: 'test-model', }, }; await compressCommand.action!(context, ''); @@ -55,9 +64,9 @@ describe('compressCommand', () => { }); it('should set pending item, call tryCompressChat, and add result on success', async () => { - const compressedResult: ChatCompressionInfo = { + const compressedResult: Core.ChatCompressionInfo = { originalTokenCount: 200, - compressionStatus: CompressionStatus.COMPRESSED, + compressionStatus: Core.CompressionStatus.COMPRESSED, newTokenCount: 100, }; mockTryCompressChat.mockResolvedValue(compressedResult); @@ -69,9 +78,9 @@ describe('compressCommand', () => { compression: { isPending: true, compressionStatus: null, - originalTokenCount: null, - newTokenCount: null, - model: 'test-model', + beforePercentage: null, + afterPercentage: null, + threshold: null, }, }); @@ -85,10 +94,10 @@ describe('compressCommand', () => { type: MessageType.COMPRESSION, compression: { isPending: false, - compressionStatus: CompressionStatus.COMPRESSED, - originalTokenCount: 200, - newTokenCount: 100, - model: 'test-model', + compressionStatus: Core.CompressionStatus.COMPRESSED, + beforePercentage: 20, + afterPercentage: 10, + threshold: 20, }, }, expect.any(Number), diff --git a/packages/cli/src/ui/commands/compressCommand.ts b/packages/cli/src/ui/commands/compressCommand.ts index 29a813e4a7..d1ab1c9085 100644 --- a/packages/cli/src/ui/commands/compressCommand.ts +++ b/packages/cli/src/ui/commands/compressCommand.ts @@ -6,6 +6,7 @@ import { MessageType, type HistoryItemCompression } from '../types.js'; import { CommandKind, type SlashCommand } from './types.js'; +import { tokenLimit, CompressionStatus } from '@google/gemini-cli-core'; export const compressCommand: SlashCommand = { name: 'compress', @@ -14,7 +15,21 @@ export const compressCommand: SlashCommand = { kind: CommandKind.BUILT_IN, autoExecute: true, action: async (context) => { - const { ui } = context; + const { ui, services } = context; + const agentContext = services.agentContext; + if (!agentContext) { + ui.addItem( + { + type: MessageType.ERROR, + text: 'Agent context not found.', + }, + Date.now(), + ); + return; + } + + const config = agentContext.config; + if (ui.pendingItem) { ui.addItem( { @@ -30,31 +45,39 @@ export const compressCommand: SlashCommand = { type: MessageType.COMPRESSION, compression: { isPending: true, - originalTokenCount: null, - newTokenCount: null, + beforePercentage: null, + afterPercentage: null, + threshold: null, compressionStatus: null, - model: context.services.config?.getModel(), }, }; try { ui.setPendingItem(pendingMessage); const promptId = `compress-${Date.now()}`; - const compressed = - await context.services.agentContext?.geminiClient?.tryCompressChat( - promptId, - true, - ); + const compressed = await agentContext.geminiClient.tryCompressChat( + promptId, + true, + ); if (compressed) { + const limit = tokenLimit(config.getModel()); + const threshold = config.getContextWindowCompressionThreshold(); + const beforePercentage = Math.round( + (compressed.originalTokenCount / limit) * 100, + ); + const afterPercentage = Math.round( + (compressed.newTokenCount / limit) * 100, + ); + ui.addItem( { type: MessageType.COMPRESSION, compression: { isPending: false, - originalTokenCount: compressed.originalTokenCount, - newTokenCount: compressed.newTokenCount, + beforePercentage, + afterPercentage, + threshold, compressionStatus: compressed.compressionStatus, - model: context.services.config?.getModel(), }, } as HistoryItemCompression, Date.now(), diff --git a/packages/cli/src/ui/components/messages/CompressionMessage.test.tsx b/packages/cli/src/ui/components/messages/CompressionMessage.test.tsx index 064191e018..fbea908547 100644 --- a/packages/cli/src/ui/components/messages/CompressionMessage.test.tsx +++ b/packages/cli/src/ui/components/messages/CompressionMessage.test.tsx @@ -11,17 +11,22 @@ import { } from './CompressionMessage.js'; import { CompressionStatus } from '@google/gemini-cli-core'; import { type CompressionProps } from '../../types.js'; -import { describe, it, expect } from 'vitest'; +import { describe, it, expect, vi, afterEach } from 'vitest'; describe('', () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + const createCompressionProps = ( overrides: Partial = {}, ): CompressionDisplayProps => ({ compression: { isPending: false, - originalTokenCount: null, - newTokenCount: null, + beforePercentage: null, + afterPercentage: null, compressionStatus: CompressionStatus.COMPRESSED, + isManual: true, ...overrides, }, }); @@ -29,9 +34,10 @@ describe('', () => { describe('pending state', () => { it('renders pending message when compression is in progress', async () => { const props = createCompressionProps({ isPending: true }); - const { lastFrame, unmount } = await renderWithProviders( + const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( , ); + await waitUntilReady(); const output = lastFrame(); expect(output).toContain('Compressing chat history'); @@ -43,9 +49,10 @@ describe('', () => { it('renders success message when tokens are reduced', async () => { const props = createCompressionProps({ isPending: false, - originalTokenCount: 100, - newTokenCount: 50, + beforePercentage: 22, + afterPercentage: 6, compressionStatus: CompressionStatus.COMPRESSED, + thresholdPercentage: 50, }); const { lastFrame, unmount } = await renderWithProviders( , @@ -54,45 +61,16 @@ describe('', () => { expect(output).not.toContain('✦'); expect(output).toContain( - 'Context compressed from 100 tokens to 50 tokens. Change threshold in /settings.', + 'Context compressed (22% → 6%). Adjust threshold (50%) in /settings.', ); unmount(); }); - - it.each([ - { original: 50000, newTokens: 25000 }, // Large compression - { original: 700000, newTokens: 350000 }, // Very large compression - ])( - 'renders success message for large successful compression (from $original to $newTokens)', - async ({ original, newTokens }) => { - const props = createCompressionProps({ - isPending: false, - originalTokenCount: original, - newTokenCount: newTokens, - compressionStatus: CompressionStatus.COMPRESSED, - }); - const { lastFrame, unmount } = await renderWithProviders( - , - ); - const output = lastFrame(); - - expect(output).not.toContain('✦'); - expect(output).toContain( - `Context compressed from ${original.toLocaleString()} tokens to ${newTokens.toLocaleString()} tokens. Change threshold in /settings.`, - ); - expect(output).not.toContain('Skipping compression'); - expect(output).not.toContain('did not reduce size'); - unmount(); - }, - ); }); describe('skipped compression (tokens increased or same)', () => { it('renders skip message when compression would increase token count', async () => { const props = createCompressionProps({ isPending: false, - originalTokenCount: 50, - newTokenCount: 75, compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT, }); @@ -107,117 +85,6 @@ describe('', () => { ); unmount(); }); - - it('renders skip message when token counts are equal', async () => { - const props = createCompressionProps({ - isPending: false, - originalTokenCount: 50, - newTokenCount: 50, - compressionStatus: - CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT, - }); - const { lastFrame, unmount } = await renderWithProviders( - , - ); - const output = lastFrame(); - - expect(output).toContain( - 'Compression was not beneficial for this history size.', - ); - unmount(); - }); - }); - - describe('message content validation', () => { - it.each([ - { - original: 200, - newTokens: 80, - expected: - 'Context compressed from 200 tokens to 80 tokens. Change threshold in /settings.', - }, - { - original: 500, - newTokens: 150, - expected: - 'Context compressed from 500 tokens to 150 tokens. Change threshold in /settings.', - }, - { - original: 1500, - newTokens: 400, - expected: `Context compressed from ${(1500).toLocaleString()} tokens to 400 tokens. Change threshold in /settings.`, - }, - ])( - 'displays correct compression statistics (from $original to $newTokens)', - async ({ original, newTokens, expected }) => { - const props = createCompressionProps({ - isPending: false, - originalTokenCount: original, - newTokenCount: newTokens, - compressionStatus: CompressionStatus.COMPRESSED, - }); - const { lastFrame, unmount } = await renderWithProviders( - , - ); - const output = lastFrame(); - - expect(output).toContain(expected); - unmount(); - }, - ); - - it.each([ - { original: 50, newTokens: 60 }, // Increased - { original: 100, newTokens: 100 }, // Same - { original: 49999, newTokens: 50000 }, // Just under 50k threshold - ])( - 'shows skip message for small histories when new tokens >= original tokens ($original -> $newTokens)', - async ({ original, newTokens }) => { - const props = createCompressionProps({ - isPending: false, - originalTokenCount: original, - newTokenCount: newTokens, - compressionStatus: - CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT, - }); - const { lastFrame, unmount } = await renderWithProviders( - , - ); - const output = lastFrame(); - - expect(output).toContain( - 'Compression was not beneficial for this history size.', - ); - expect(output).not.toContain('compressed from'); - unmount(); - }, - ); - - it.each([ - { original: 50000, newTokens: 50100 }, // At 50k threshold - { original: 700000, newTokens: 710000 }, // Large history case - { original: 100000, newTokens: 100000 }, // Large history, same count - ])( - 'shows compression failure message for large histories when new tokens >= original tokens ($original -> $newTokens)', - async ({ original, newTokens }) => { - const props = createCompressionProps({ - isPending: false, - originalTokenCount: original, - newTokenCount: newTokens, - compressionStatus: - CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT, - }); - const { lastFrame, unmount } = await renderWithProviders( - , - ); - const output = lastFrame(); - - expect(output).toContain('compression did not reduce size'); - expect(output).not.toContain('compressed from'); - expect(output).not.toContain('Compression was not beneficial'); - unmount(); - }, - ); }); describe('failure states', () => { diff --git a/packages/cli/src/ui/components/messages/CompressionMessage.tsx b/packages/cli/src/ui/components/messages/CompressionMessage.tsx index cc800caadf..404d5149e3 100644 --- a/packages/cli/src/ui/components/messages/CompressionMessage.tsx +++ b/packages/cli/src/ui/components/messages/CompressionMessage.tsx @@ -9,7 +9,7 @@ import type { CompressionProps } from '../../types.js'; import { CliSpinner } from '../CliSpinner.js'; import { theme } from '../../semantic-colors.js'; import { SCREEN_READER_MODEL_PREFIX } from '../../textConstants.js'; -import { CompressionStatus, tokenLimit } from '@google/gemini-cli-core'; +import { CompressionStatus } from '@google/gemini-cli-core'; export interface CompressionDisplayProps { compression: CompressionProps; @@ -24,37 +24,22 @@ export function CompressionMessage({ }: CompressionDisplayProps): React.JSX.Element { const { isPending, - originalTokenCount, - newTokenCount, + beforePercentage, + afterPercentage, + threshold, compressionStatus, - model, } = compression; - const originalTokens = originalTokenCount ?? 0; - const newTokens = newTokenCount ?? 0; - const getCompressionText = () => { if (isPending) { return 'Compressing chat history'; } - const limit = model ? tokenLimit(model) : 0; - const formatPercent = (tokens: number) => - limit > 0 - ? `${Math.round((tokens / limit) * 100)}% (${tokens.toLocaleString()} tokens)` - : `${tokens.toLocaleString()} tokens`; - switch (compressionStatus) { case CompressionStatus.COMPRESSED: - return `Context compressed from ${formatPercent(originalTokens)} to ${formatPercent(newTokens)}. Change threshold in /settings.`; + return `Context compressed (${beforePercentage}% ➔ ${afterPercentage}%). Adjust threshold (${threshold}%) in /settings.`; case CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT: - // For smaller histories (< 50k tokens), compression overhead likely exceeds benefits - if (originalTokens < 50000) { - return 'Compression was not beneficial for this history size.'; - } - // For larger histories where compression should work but didn't, - // this suggests an issue with the compression process itself - return 'Chat history compression did not reduce size.'; + return 'Compression was not beneficial for this history size.'; case CompressionStatus.COMPRESSION_FAILED_TOKEN_COUNT_ERROR: return 'Could not compress chat history due to a token counting error.'; case CompressionStatus.COMPRESSION_FAILED_EMPTY_SUMMARY: diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 1d0711c466..08aa6b419d 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -332,6 +332,7 @@ describe('useGeminiStream', () => { getIdeMode: vi.fn(() => false), getEnableHooks: vi.fn(() => false), getShowContextWindowWarning: vi.fn(() => false), + getContextWindowCompressionThreshold: vi.fn(() => 0.2), } as unknown as Config; beforeEach(() => { @@ -2485,7 +2486,7 @@ describe('useGeminiStream', () => { remainingTokens: 70, shouldShow: true, expectedMessage: - 'Context window is 30% full. Message size (30 tokens) might exceed the limit.\nPlease try reducing the size of your message or use the /compress command to compress the chat history.', + 'Context 30% full. Message may exceed limit. Reduce size or /compress.', }, ])( 'should $name', @@ -2584,6 +2585,9 @@ describe('useGeminiStream', () => { it('should add informational messages when ChatCompressed event is received', async () => { vi.mocked(tokenLimit).mockReturnValue(10000); + vi.mocked( + mockConfig.getContextWindowCompressionThreshold, + ).mockReturnValue(0.2); // Setup mock to return a stream with ChatCompressed event mockSendMessageStream.mockReturnValue( (async function* () { @@ -2595,6 +2599,21 @@ describe('useGeminiStream', () => { compressionStatus: 'compressed', }, }; + yield { + type: ServerGeminiEventType.Content, + value: 'Response after compression', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { + finishReason: 'STOP', + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 20, + totalTokenCount: 30, + }, + }, + }; })(), ); @@ -2609,13 +2628,13 @@ describe('useGeminiStream', () => { await waitFor(() => { expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ - type: MessageType.COMPRESSION, + type: 'compression', compression: { isPending: false, - originalTokenCount: 1000, - newTokenCount: 500, + beforePercentage: 10, + afterPercentage: 5, + threshold: 20, compressionStatus: 'compressed', - model: 'gemini-2.5-pro', }, }), expect.any(Number), diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index c7a83cd3ae..a6f2f20de9 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -1143,15 +1143,28 @@ export const useGeminiStream = ( setPendingHistoryItem(null); } + const limit = tokenLimit(config.getModel()); + const beforePercentage = + eventValue?.originalTokenCount != null + ? Math.round((eventValue.originalTokenCount / limit) * 100) + : null; + const afterPercentage = + eventValue?.newTokenCount != null + ? Math.round((eventValue.newTokenCount / limit) * 100) + : null; + const threshold = Math.round( + config.getContextWindowCompressionThreshold() * 100, + ); + addItem( { type: 'compression', compression: { isPending: false, - originalTokenCount: eventValue?.originalTokenCount ?? null, - newTokenCount: eventValue?.newTokenCount ?? null, + beforePercentage, + afterPercentage, + threshold, compressionStatus: eventValue?.compressionStatus ?? null, - model: config.getModel(), }, timestamp: new Date(userMessageTimestamp), } as HistoryItemWithoutId, @@ -1185,7 +1198,7 @@ export const useGeminiStream = ( ((limit - remainingTokenCount) / limit) * 100, ); - const text = `Context window is ${usedPercentage}% full. Message size (${estimatedRequestTokenCount.toLocaleString()} tokens) might exceed the limit.\nPlease try reducing the size of your message or use the /compress command to compress the chat history.`; + const text = `Context ${usedPercentage}% full. Message may exceed limit. Reduce size or /compress.`; addItem({ type: 'info', diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 1fc928e9f4..d15666d838 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -138,10 +138,10 @@ export interface IndividualToolCallDisplay { export interface CompressionProps { isPending: boolean; - originalTokenCount: number | null; - newTokenCount: number | null; + beforePercentage: number | null; + afterPercentage: number | null; + threshold: number | null; compressionStatus: CompressionStatus | null; - model?: string; } /**