jcarr
/
gemini-cli
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

migrate compress command (#4271)

This commit is contained in:
Abhi 2025-07-15 21:59:16 -04:00 committed by GitHub
parent e88b9362dc
commit b72e3dfb43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 230 additions and 96 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
packages/cli/src/services/CommandService.test.ts
View File

@ -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,

2
packages/cli/src/services/CommandService.ts
View File

@ -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<SlashCommand[]> => [
@ -21,6 +22,7 @@ const loadBuiltInCommands = async (): Promise<SlashCommand[]> => [
authCommand,
chatCommand,
clearCommand,
compressCommand,
extensionsCommand,
helpCommand,
memoryCommand,

2
packages/cli/src/test-utils/mockCommandContext.ts
View File

@ -44,6 +44,8 @@ export const createMockCommandContext = (
addItem: vi.fn(),
clear: vi.fn(),
setDebugMessage: vi.fn(),
pendingItem: null,
setPendingItem: vi.fn(),
},
session: {
stats: {

129
packages/cli/src/ui/commands/compressCommand.test.ts Normal file
View File

@ -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<typeof createMockCommandContext>;
let mockTryCompressChat: ReturnType<typeof vi.fn>;
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);
});
});

77
packages/cli/src/ui/commands/compressCommand.ts Normal file
View File

@ -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);
}
},
};

14
packages/cli/src/ui/commands/types.ts
View File

@ -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: {

34
packages/cli/src/ui/hooks/slashCommandProcessor.test.ts
View File

@ -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),
);
});
});
});

60
packages/cli/src/ui/hooks/slashCommandProcessor.ts
View File

@ -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(