feat(ui): implement message queuing during streaming responses (#6049)

Co-authored-by: Jacob Richman <jacob314@gmail.com>

packages/cli/src/ui/hooks/useMessageQueue.test.ts

@@ -0,0 +1,226 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { renderHook, act } from '@testing-library/react';
+import { useMessageQueue } from './useMessageQueue.js';
+import { StreamingState } from '../types.js';
+
+describe('useMessageQueue', () => {
+  let mockSubmitQuery: ReturnType<typeof vi.fn>;
+
+  beforeEach(() => {
+    mockSubmitQuery = vi.fn();
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.clearAllMocks();
+  });
+
+  it('should initialize with empty queue', () => {
+    const { result } = renderHook(() =>
+      useMessageQueue({
+        streamingState: StreamingState.Idle,
+        submitQuery: mockSubmitQuery,
+      }),
+    );
+
+    expect(result.current.messageQueue).toEqual([]);
+    expect(result.current.getQueuedMessagesText()).toBe('');
+  });
+
+  it('should add messages to queue', () => {
+    const { result } = renderHook(() =>
+      useMessageQueue({
+        streamingState: StreamingState.Responding,
+        submitQuery: mockSubmitQuery,
+      }),
+    );
+
+    act(() => {
+      result.current.addMessage('Test message 1');
+      result.current.addMessage('Test message 2');
+    });
+
+    expect(result.current.messageQueue).toEqual([
+      'Test message 1',
+      'Test message 2',
+    ]);
+  });
+
+  it('should filter out empty messages', () => {
+    const { result } = renderHook(() =>
+      useMessageQueue({
+        streamingState: StreamingState.Responding,
+        submitQuery: mockSubmitQuery,
+      }),
+    );
+
+    act(() => {
+      result.current.addMessage('Valid message');
+      result.current.addMessage('   '); // Only whitespace
+      result.current.addMessage(''); // Empty
+      result.current.addMessage('Another valid message');
+    });
+
+    expect(result.current.messageQueue).toEqual([
+      'Valid message',
+      'Another valid message',
+    ]);
+  });
+
+  it('should clear queue', () => {
+    const { result } = renderHook(() =>
+      useMessageQueue({
+        streamingState: StreamingState.Responding,
+        submitQuery: mockSubmitQuery,
+      }),
+    );
+
+    act(() => {
+      result.current.addMessage('Test message');
+    });
+
+    expect(result.current.messageQueue).toEqual(['Test message']);
+
+    act(() => {
+      result.current.clearQueue();
+    });
+
+    expect(result.current.messageQueue).toEqual([]);
+  });
+
+  it('should return queued messages as text with double newlines', () => {
+    const { result } = renderHook(() =>
+      useMessageQueue({
+        streamingState: StreamingState.Responding,
+        submitQuery: mockSubmitQuery,
+      }),
+    );
+
+    act(() => {
+      result.current.addMessage('Message 1');
+      result.current.addMessage('Message 2');
+      result.current.addMessage('Message 3');
+    });
+
+    expect(result.current.getQueuedMessagesText()).toBe(
+      'Message 1\n\nMessage 2\n\nMessage 3',
+    );
+  });
+
+  it('should auto-submit queued messages when transitioning to Idle', () => {
+    const { result, rerender } = renderHook(
+      ({ streamingState }) =>
+        useMessageQueue({
+          streamingState,
+          submitQuery: mockSubmitQuery,
+        }),
+      {
+        initialProps: { streamingState: StreamingState.Responding },
+      },
+    );
+
+    // Add some messages
+    act(() => {
+      result.current.addMessage('Message 1');
+      result.current.addMessage('Message 2');
+    });
+
+    expect(result.current.messageQueue).toEqual(['Message 1', 'Message 2']);
+
+    // Transition to Idle
+    rerender({ streamingState: StreamingState.Idle });
+
+    expect(mockSubmitQuery).toHaveBeenCalledWith('Message 1\n\nMessage 2');
+    expect(result.current.messageQueue).toEqual([]);
+  });
+
+  it('should not auto-submit when queue is empty', () => {
+    const { rerender } = renderHook(
+      ({ streamingState }) =>
+        useMessageQueue({
+          streamingState,
+          submitQuery: mockSubmitQuery,
+        }),
+      {
+        initialProps: { streamingState: StreamingState.Responding },
+      },
+    );
+
+    // Transition to Idle with empty queue
+    rerender({ streamingState: StreamingState.Idle });
+
+    expect(mockSubmitQuery).not.toHaveBeenCalled();
+  });
+
+  it('should not auto-submit when not transitioning to Idle', () => {
+    const { result, rerender } = renderHook(
+      ({ streamingState }) =>
+        useMessageQueue({
+          streamingState,
+          submitQuery: mockSubmitQuery,
+        }),
+      {
+        initialProps: { streamingState: StreamingState.Responding },
+      },
+    );
+
+    // Add messages
+    act(() => {
+      result.current.addMessage('Message 1');
+    });
+
+    // Transition to WaitingForConfirmation (not Idle)
+    rerender({ streamingState: StreamingState.WaitingForConfirmation });
+
+    expect(mockSubmitQuery).not.toHaveBeenCalled();
+    expect(result.current.messageQueue).toEqual(['Message 1']);
+  });
+
+  it('should handle multiple state transitions correctly', () => {
+    const { result, rerender } = renderHook(
+      ({ streamingState }) =>
+        useMessageQueue({
+          streamingState,
+          submitQuery: mockSubmitQuery,
+        }),
+      {
+        initialProps: { streamingState: StreamingState.Idle },
+      },
+    );
+
+    // Start responding
+    rerender({ streamingState: StreamingState.Responding });
+
+    // Add messages while responding
+    act(() => {
+      result.current.addMessage('First batch');
+    });
+
+    // Go back to idle - should submit
+    rerender({ streamingState: StreamingState.Idle });
+
+    expect(mockSubmitQuery).toHaveBeenCalledWith('First batch');
+    expect(result.current.messageQueue).toEqual([]);
+
+    // Start responding again
+    rerender({ streamingState: StreamingState.Responding });
+
+    // Add more messages
+    act(() => {
+      result.current.addMessage('Second batch');
+    });
+
+    // Go back to idle - should submit again
+    rerender({ streamingState: StreamingState.Idle });
+
+    expect(mockSubmitQuery).toHaveBeenCalledWith('Second batch');
+    expect(mockSubmitQuery).toHaveBeenCalledTimes(2);
+  });
+});

packages/cli/src/ui/hooks/useMessageQueue.ts

@@ -0,0 +1,69 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useCallback, useEffect, useState } from 'react';
+import { StreamingState } from '../types.js';
+
+export interface UseMessageQueueOptions {
+  streamingState: StreamingState;
+  submitQuery: (query: string) => void;
+}
+
+export interface UseMessageQueueReturn {
+  messageQueue: string[];
+  addMessage: (message: string) => void;
+  clearQueue: () => void;
+  getQueuedMessagesText: () => string;
+}
+
+/**
+ * Hook for managing message queuing during streaming responses.
+ * Allows users to queue messages while the AI is responding and automatically
+ * sends them when streaming completes.
+ */
+export function useMessageQueue({
+  streamingState,
+  submitQuery,
+}: UseMessageQueueOptions): UseMessageQueueReturn {
+  const [messageQueue, setMessageQueue] = useState<string[]>([]);
+
+  // Add a message to the queue
+  const addMessage = useCallback((message: string) => {
+    const trimmedMessage = message.trim();
+    if (trimmedMessage.length > 0) {
+      setMessageQueue((prev) => [...prev, trimmedMessage]);
+    }
+  }, []);
+
+  // Clear the entire queue
+  const clearQueue = useCallback(() => {
+    setMessageQueue([]);
+  }, []);
+
+  // Get all queued messages as a single text string
+  const getQueuedMessagesText = useCallback(() => {
+    if (messageQueue.length === 0) return '';
+    return messageQueue.join('\n\n');
+  }, [messageQueue]);
+
+  // Process queued messages when streaming becomes idle
+  useEffect(() => {
+    if (streamingState === StreamingState.Idle && messageQueue.length > 0) {
+      // Combine all messages with double newlines for clarity
+      const combinedMessage = messageQueue.join('\n\n');
+      // Clear the queue and submit
+      setMessageQueue([]);
+      submitQuery(combinedMessage);
+    }
+  }, [streamingState, messageQueue, submitQuery]);
+
+  return {
+    messageQueue,
+    addMessage,
+    clearQueue,
+    getQueuedMessagesText,
+  };
+}