diff options
Diffstat (limited to 'packages/cli/src/ui/hooks')
| -rw-r--r-- | packages/cli/src/ui/hooks/useMessageQueue.test.ts | 226 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useMessageQueue.ts | 69 |
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, + }; +} |
