feat: allow queuing messages during compression (#24071) (#26506)

Author: cocosheng-g

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();
+    });
+  });
 });

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,

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

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

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<typeof useMessageQueue>;
     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([]);
+      });
+    });
+  });
 });

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<string[]>([]);
 
@@ -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 {