From e80d7cc08303c78813508a9ea028f06fd3121e06 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Tue, 5 May 2026 13:52:08 -0400 Subject: [PATCH] feat: allow queuing messages during compression (#24071) (#26506) --- packages/cli/src/ui/AppContainer.test.tsx | 63 +++++++++++++++- packages/cli/src/ui/AppContainer.tsx | 23 +++++- .../src/ui/commands/compressCommand.test.ts | 5 ++ .../cli/src/ui/commands/compressCommand.ts | 73 ++++++++++--------- .../cli/src/ui/hooks/useMessageQueue.test.tsx | 49 +++++++++++++ packages/cli/src/ui/hooks/useMessageQueue.ts | 4 + 6 files changed, 179 insertions(+), 38 deletions(-) diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 92a519856a..ea9e0629d1 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,65 @@ describe('AppContainer State Management', () => { unmount(); }); }); + + describe('Compression Queuing', () => { + beforeEach(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, + }, + }, + ], + })); + }); + + it('queues messages during compression instead of handling as steering hints', async () => { + const { unmount } = await act(async () => renderAppContainer()); + + // Verify state isolation + expect(capturedUIState.streamingState).toBe(StreamingState.Idle); + + // 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(); + }); + + it('executes slash commands immediately during compression', async () => { + const { unmount } = await act(async () => renderAppContainer()); + + // Submit a slash command + await act(async () => capturedUIActions.handleFinalSubmit('/help')); + + // Verify it was NOT queued + expect(capturedUIState.messageQueue).not.toContain('/help'); + + unmount(); + }); + }); }); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index d8b1e1d277..5c7d4176bc 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) || + (!isCompressing && isIdle && 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.test.ts b/packages/cli/src/ui/commands/compressCommand.test.ts index fd60b54354..c91c42b8c4 100644 --- a/packages/cli/src/ui/commands/compressCommand.test.ts +++ b/packages/cli/src/ui/commands/compressCommand.test.ts @@ -42,6 +42,7 @@ describe('compressCommand', () => { }, }; await compressCommand.action!(context, ''); + await new Promise((r) => setTimeout(r, 0)); expect(context.ui.addItem).toHaveBeenCalledWith( expect.objectContaining({ type: MessageType.ERROR, @@ -62,6 +63,7 @@ describe('compressCommand', () => { mockTryCompressChat.mockResolvedValue(compressedResult); await compressCommand.action!(context, ''); + await new Promise((r) => setTimeout(r, 0)); expect(context.ui.setPendingItem).toHaveBeenNthCalledWith(1, { type: MessageType.COMPRESSION, @@ -98,6 +100,7 @@ describe('compressCommand', () => { mockTryCompressChat.mockResolvedValue(null); await compressCommand.action!(context, ''); + await new Promise((r) => setTimeout(r, 0)); expect(context.ui.addItem).toHaveBeenCalledWith( expect.objectContaining({ @@ -114,6 +117,7 @@ describe('compressCommand', () => { mockTryCompressChat.mockRejectedValue(error); await compressCommand.action!(context, ''); + await new Promise((r) => setTimeout(r, 0)); expect(context.ui.addItem).toHaveBeenCalledWith( expect.objectContaining({ @@ -128,6 +132,7 @@ describe('compressCommand', () => { it('should clear the pending item in a finally block', async () => { mockTryCompressChat.mockRejectedValue(new Error('some error')); await compressCommand.action!(context, ''); + await new Promise((r) => setTimeout(r, 0)); expect(context.ui.setPendingItem).toHaveBeenCalledWith(null); }); diff --git a/packages/cli/src/ui/commands/compressCommand.ts b/packages/cli/src/ui/commands/compressCommand.ts index 6d53667010..37ffc4930a 100644 --- a/packages/cli/src/ui/commands/compressCommand.ts +++ b/packages/cli/src/ui/commands/compressCommand.ts @@ -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.test.tsx b/packages/cli/src/ui/hooks/useMessageQueue.test.tsx index da6eea233c..cda2f97f39 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.test.tsx +++ b/packages/cli/src/ui/hooks/useMessageQueue.test.tsx @@ -29,6 +29,7 @@ describe('useMessageQueue', () => { streamingState: StreamingState; submitQuery: (query: string) => void; isMcpReady: boolean; + isCompressing?: boolean; }) => { let hookResult: ReturnType; function TestComponent(props: typeof initialProps) { @@ -402,4 +403,52 @@ describe('useMessageQueue', () => { expect(result.current.messageQueue).toEqual([]); }); }); + + describe('isCompressing logic', () => { + it('should not auto-submit when isCompressing is true, even if streamingState is Idle', async () => { + const { result } = await renderMessageQueueHook({ + isConfigInitialized: true, + streamingState: StreamingState.Idle, + submitQuery: mockSubmitQuery, + isMcpReady: true, + isCompressing: true, + }); + + // Add messages + act(() => { + result.current.addMessage('Compression message'); + }); + + expect(mockSubmitQuery).not.toHaveBeenCalled(); + expect(result.current.messageQueue).toEqual(['Compression message']); + }); + + it('should auto-submit queued messages when isCompressing becomes false', async () => { + const { result, rerender } = await renderMessageQueueHook({ + isConfigInitialized: true, + streamingState: StreamingState.Idle, + submitQuery: mockSubmitQuery, + isMcpReady: true, + isCompressing: true, + }); + + // Add messages + act(() => { + result.current.addMessage('Pending compression message 1'); + result.current.addMessage('Pending compression message 2'); + }); + + expect(mockSubmitQuery).not.toHaveBeenCalled(); + + // Transition isCompressing to false + rerender({ isCompressing: false }); + + await waitFor(() => { + expect(mockSubmitQuery).toHaveBeenCalledWith( + 'Pending compression message 1\n\nPending compression message 2', + ); + expect(result.current.messageQueue).toEqual([]); + }); + }); + }); }); 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 {