summaryrefslogtreecommitdiff
path: root/packages/cli/src/ui/hooks
diff options
context:
space:
mode:
Diffstat (limited to 'packages/cli/src/ui/hooks')
-rw-r--r--packages/cli/src/ui/hooks/useMessageQueue.test.ts226
-rw-r--r--packages/cli/src/ui/hooks/useMessageQueue.ts69
2 files changed, 295 insertions, 0 deletions
diff --git a/packages/cli/src/ui/hooks/useMessageQueue.test.ts b/packages/cli/src/ui/hooks/useMessageQueue.test.ts
new file mode 100644
index 00000000..01e49afe
--- /dev/null
+++ b/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);
+ });
+});
diff --git a/packages/cli/src/ui/hooks/useMessageQueue.ts b/packages/cli/src/ui/hooks/useMessageQueue.ts
new file mode 100644
index 00000000..f7bbe1eb
--- /dev/null
+++ b/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,
+ };
+}