From 87d252b29203833c78122790986a76de0a8038df Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Tue, 5 May 2026 11:45:00 -0400 Subject: [PATCH] feat: allow queuing messages during compression (#24071) --- packages/cli/src/ui/AppContainer.test.tsx | 46 +++++++++++- packages/cli/src/ui/AppContainer.tsx | 23 +++++- .../cli/src/ui/commands/compressCommand.ts | 75 ++++++++++--------- packages/cli/src/ui/hooks/useMessageQueue.ts | 4 + 4 files changed, 109 insertions(+), 39 deletions(-) diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 92a519856a..7786294b39 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -100,7 +100,7 @@ import { type LoadedSettings } from '../config/settings.js'; import { createMockSettings } from '../test-utils/settings.js'; import type { InitializationResult } from '../core/initializer.js'; import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js'; -import { StreamingState } from './types.js'; +import { StreamingState, MessageType } from './types.js'; import { UIStateContext, type UIState } from './contexts/UIStateContext.js'; import { UIActionsContext, @@ -3576,4 +3576,48 @@ describe('AppContainer State Management', () => { unmount(); }); }); + + describe('Compression Queuing', () => { + it('queues messages during compression instead of handling as steering hints', async () => { + const { checkPermissions } = await import( + './hooks/atCommandProcessor.js' + ); + vi.mocked(checkPermissions).mockResolvedValue([]); + + vi.spyOn(mockConfig, 'isModelSteeringEnabled').mockReturnValue(true); + + const actual = await vi.importActual('./hooks/useMessageQueue.js'); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const { useMessageQueue: realUseMessageQueue } = actual as any; + mockedUseMessageQueue.mockImplementation(realUseMessageQueue); + + // Start compression by mocking pendingHistoryItems to include a pending compression + mockedUseGeminiStream.mockImplementation(() => ({ + ...DEFAULT_GEMINI_STREAM_MOCK, + pendingHistoryItems: [ + { + type: MessageType.COMPRESSION, + compression: { + isPending: true, + originalTokenCount: null, + newTokenCount: null, + compressionStatus: null, + }, + }, + ], + })); + + const { unmount } = await act(async () => renderAppContainer()); + + // Submit a message + await act(async () => + capturedUIActions.handleFinalSubmit('follow up message'), + ); + + // Verify it was queued, not submitted as steering hint + expect(capturedUIState.messageQueue).toContain('follow up message'); + + unmount(); + }); + }); }); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index d8b1e1d277..70cfd29402 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1310,6 +1310,15 @@ Logging in with Google... Restarting Gemini CLI to continue. const { isMcpReady } = useMcpStatus(config); + const isCompressing = useMemo( + () => + pendingHistoryItems.some( + (item) => + item.type === MessageType.COMPRESSION && item.compression.isPending, + ), + [pendingHistoryItems], + ); + const { messageQueue, addMessage, @@ -1321,6 +1330,7 @@ Logging in with Google... Restarting Gemini CLI to continue. streamingState, submitQuery, isMcpReady, + isCompressing, }); cancelHandlerRef.current = useCallback( @@ -1415,7 +1425,10 @@ Logging in with Google... Restarting Gemini CLI to continue. } const isMcpOrConfigReady = isConfigInitialized && isMcpReady; - if ((isSlash && isConfigInitialized) || (isIdle && isMcpOrConfigReady)) { + if ( + (isSlash && isConfigInitialized) || + (isIdle && !isCompressing && isMcpOrConfigReady) + ) { if (!isSlash) { const permissions = await checkPermissions(submittedValue, config); if (permissions.length > 0) { @@ -1438,7 +1451,12 @@ Logging in with Google... Restarting Gemini CLI to continue. void submitQuery(submittedValue); } else { // Check messageQueue.length === 0 to only notify on the first queued item - if (isIdle && !isMcpOrConfigReady && messageQueue.length === 0) { + if ( + isIdle && + !isCompressing && + !isMcpOrConfigReady && + messageQueue.length === 0 + ) { coreEvents.emitFeedback( 'info', !isConfigInitialized @@ -1458,6 +1476,7 @@ Logging in with Google... Restarting Gemini CLI to continue. slashCommands, isMcpReady, streamingState, + isCompressing, messageQueue.length, pendingHistoryItems, config, diff --git a/packages/cli/src/ui/commands/compressCommand.ts b/packages/cli/src/ui/commands/compressCommand.ts index 6d53667010..81efc98d80 100644 --- a/packages/cli/src/ui/commands/compressCommand.ts +++ b/packages/cli/src/ui/commands/compressCommand.ts @@ -13,7 +13,7 @@ export const compressCommand: SlashCommand = { description: 'Compresses the context by replacing it with a summary', kind: CommandKind.BUILT_IN, autoExecute: true, - action: async (context) => { + action: (context) => { const { ui } = context; if (ui.pendingItem) { ui.addItem( @@ -36,48 +36,51 @@ export const compressCommand: SlashCommand = { }, }; - try { - ui.setPendingItem(pendingMessage); - const promptId = `compress-${Date.now()}`; - const compressed = - await context.services.agentContext?.geminiClient?.tryCompressChat( - promptId, - true, - ); - if (compressed) { - ui.addItem( - { - type: MessageType.COMPRESSION, - compression: { - isPending: false, - originalTokenCount: compressed.originalTokenCount, - newTokenCount: compressed.newTokenCount, - compressionStatus: compressed.compressionStatus, + ui.setPendingItem(pendingMessage); + + void (async () => { + try { + const promptId = `compress-${Date.now()}`; + const compressed = + await context.services.agentContext?.geminiClient?.tryCompressChat( + promptId, + true, + ); + if (compressed) { + ui.addItem( + { + type: MessageType.COMPRESSION, + compression: { + isPending: false, + originalTokenCount: compressed.originalTokenCount, + newTokenCount: compressed.newTokenCount, + compressionStatus: compressed.compressionStatus, + }, + } as HistoryItemCompression, + Date.now(), + ); + } else { + ui.addItem( + { + type: MessageType.ERROR, + text: 'Failed to compress chat history.', }, - } as HistoryItemCompression, - Date.now(), - ); - } else { + Date.now(), + ); + } + } catch (e) { ui.addItem( { type: MessageType.ERROR, - text: 'Failed to compress chat history.', + text: `Failed to compress chat history: ${ + e instanceof Error ? e.message : String(e) + }`, }, Date.now(), ); + } finally { + ui.setPendingItem(null); } - } 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/hooks/useMessageQueue.ts b/packages/cli/src/ui/hooks/useMessageQueue.ts index 93bb0ab7a9..f746273a16 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.ts +++ b/packages/cli/src/ui/hooks/useMessageQueue.ts @@ -12,6 +12,7 @@ export interface UseMessageQueueOptions { streamingState: StreamingState; submitQuery: (query: string) => void; isMcpReady: boolean; + isCompressing?: boolean; } export interface UseMessageQueueReturn { @@ -32,6 +33,7 @@ export function useMessageQueue({ streamingState, submitQuery, isMcpReady, + isCompressing = false, }: UseMessageQueueOptions): UseMessageQueueReturn { const [messageQueue, setMessageQueue] = useState([]); @@ -69,6 +71,7 @@ export function useMessageQueue({ if ( isConfigInitialized && streamingState === StreamingState.Idle && + !isCompressing && isMcpReady && messageQueue.length > 0 ) { @@ -84,6 +87,7 @@ export function useMessageQueue({ isMcpReady, messageQueue, submitQuery, + isCompressing, ]); return {