mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-20 16:57:46 +00:00
Add /compress command to force a compression of the context (#986)
Related to https://b.corp.google.com/issues/423605555 - I figured this might be a simpler solution to start with, while still also being useful on its own even if we do implement that.
This commit is contained in:
@@ -170,7 +170,11 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
|
||||
}
|
||||
}, [config, addItem]);
|
||||
|
||||
const { handleSlashCommand, slashCommands } = useSlashCommandProcessor(
|
||||
const {
|
||||
handleSlashCommand,
|
||||
slashCommands,
|
||||
pendingHistoryItems: pendingSlashCommandHistoryItems,
|
||||
} = useSlashCommandProcessor(
|
||||
config,
|
||||
history,
|
||||
addItem,
|
||||
@@ -186,6 +190,7 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
|
||||
showToolDescriptions,
|
||||
setQuittingMessages,
|
||||
);
|
||||
const pendingHistoryItems = [...pendingSlashCommandHistoryItems];
|
||||
|
||||
const { rows: terminalHeight, columns: terminalWidth } = useTerminalSize();
|
||||
const { stdin, setRawMode } = useStdin();
|
||||
@@ -286,18 +291,23 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
|
||||
return editorType as EditorType;
|
||||
}, [settings, openEditorDialog]);
|
||||
|
||||
const { streamingState, submitQuery, initError, pendingHistoryItems } =
|
||||
useGeminiStream(
|
||||
config.getGeminiClient(),
|
||||
history,
|
||||
addItem,
|
||||
setShowHelp,
|
||||
config,
|
||||
setDebugMessage,
|
||||
handleSlashCommand,
|
||||
shellModeActive,
|
||||
getPreferredEditor,
|
||||
);
|
||||
const {
|
||||
streamingState,
|
||||
submitQuery,
|
||||
initError,
|
||||
pendingHistoryItems: pendingGeminiHistoryItems,
|
||||
} = useGeminiStream(
|
||||
config.getGeminiClient(),
|
||||
history,
|
||||
addItem,
|
||||
setShowHelp,
|
||||
config,
|
||||
setDebugMessage,
|
||||
handleSlashCommand,
|
||||
shellModeActive,
|
||||
getPreferredEditor,
|
||||
);
|
||||
pendingHistoryItems.push(...pendingGeminiHistoryItems);
|
||||
const { elapsedTime, currentLoadingPhrase } =
|
||||
useLoadingIndicator(streamingState);
|
||||
const showAutoAcceptIndicator = useAutoAcceptIndicator({ config });
|
||||
|
||||
@@ -13,6 +13,7 @@ import { InfoMessage } from './messages/InfoMessage.js';
|
||||
import { ErrorMessage } from './messages/ErrorMessage.js';
|
||||
import { ToolGroupMessage } from './messages/ToolGroupMessage.js';
|
||||
import { GeminiMessageContent } from './messages/GeminiMessageContent.js';
|
||||
import { CompressionMessage } from './messages/CompressionMessage.js';
|
||||
import { Box } from 'ink';
|
||||
import { AboutBox } from './AboutBox.js';
|
||||
import { StatsDisplay } from './StatsDisplay.js';
|
||||
@@ -81,5 +82,8 @@ export const HistoryItemDisplay: React.FC<HistoryItemDisplayProps> = ({
|
||||
isFocused={isFocused}
|
||||
/>
|
||||
)}
|
||||
{item.type === 'compression' && (
|
||||
<CompressionMessage compression={item.compression} />
|
||||
)}
|
||||
</Box>
|
||||
);
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import React from 'react';
|
||||
import { Box, Text } from 'ink';
|
||||
import { CompressionProps } from '../../types.js';
|
||||
import Spinner from 'ink-spinner';
|
||||
import { Colors } from '../../colors.js';
|
||||
|
||||
export interface CompressionDisplayProps {
|
||||
compression: CompressionProps;
|
||||
}
|
||||
|
||||
/*
|
||||
* Compression messages appear when the /compress command is ran, and show a loading spinner
|
||||
* while compression is in progress, followed up by some compression stats.
|
||||
*/
|
||||
export const CompressionMessage: React.FC<CompressionDisplayProps> = ({
|
||||
compression,
|
||||
}) => {
|
||||
const text = compression.isPending
|
||||
? 'Compressing chat history'
|
||||
: `Chat history compressed from ${compression.originalTokenCount} to ${compression.newTokenCount} tokens.`;
|
||||
|
||||
return (
|
||||
<Box flexDirection="row">
|
||||
<Box marginRight={1}>
|
||||
{compression.isPending ? (
|
||||
<Spinner type="dots" />
|
||||
) : (
|
||||
<Text color={Colors.AccentPurple}>✦</Text>
|
||||
)}
|
||||
</Box>
|
||||
<Box>
|
||||
<Text
|
||||
color={
|
||||
compression.isPending ? Colors.AccentPurple : Colors.AccentGreen
|
||||
}
|
||||
>
|
||||
{text}
|
||||
</Text>
|
||||
</Box>
|
||||
</Box>
|
||||
);
|
||||
};
|
||||
@@ -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),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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 };
|
||||
};
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -53,6 +53,12 @@ export interface IndividualToolCallDisplay {
|
||||
renderOutputAsMarkdown?: boolean;
|
||||
}
|
||||
|
||||
export interface CompressionProps {
|
||||
isPending: boolean;
|
||||
originalTokenCount?: number;
|
||||
newTokenCount?: number;
|
||||
}
|
||||
|
||||
export interface HistoryItemBase {
|
||||
text?: string; // Text content for user/gemini/info/error messages
|
||||
}
|
||||
@@ -113,6 +119,11 @@ export type HistoryItemUserShell = HistoryItemBase & {
|
||||
text: string;
|
||||
};
|
||||
|
||||
export type HistoryItemCompression = HistoryItemBase & {
|
||||
type: 'compression';
|
||||
compression: CompressionProps;
|
||||
};
|
||||
|
||||
// Using Omit<HistoryItem, 'id'> seems to have some issues with typescript's
|
||||
// type inference e.g. historyItem.type === 'tool_group' isn't auto-inferring that
|
||||
// 'tools' in historyItem.
|
||||
@@ -127,7 +138,8 @@ export type HistoryItemWithoutId =
|
||||
| HistoryItemAbout
|
||||
| HistoryItemToolGroup
|
||||
| HistoryItemStats
|
||||
| HistoryItemQuit;
|
||||
| HistoryItemQuit
|
||||
| HistoryItemCompression;
|
||||
|
||||
export type HistoryItem = HistoryItemWithoutId & { id: number };
|
||||
|
||||
@@ -140,6 +152,7 @@ export enum MessageType {
|
||||
STATS = 'stats',
|
||||
QUIT = 'quit',
|
||||
GEMINI = 'gemini',
|
||||
COMPRESSION = 'compression',
|
||||
}
|
||||
|
||||
// Simplified message structure for internal feedback
|
||||
@@ -172,6 +185,11 @@ export type Message =
|
||||
stats: CumulativeStats;
|
||||
duration: string;
|
||||
content?: string;
|
||||
}
|
||||
| {
|
||||
type: MessageType.COMPRESSION;
|
||||
compression: CompressionProps;
|
||||
timestamp: Date;
|
||||
};
|
||||
|
||||
export interface ConsoleMessageItem {
|
||||
|
||||
Reference in New Issue
Block a user