diff --git a/src/components/Chat/Chat.test.tsx b/src/components/Chat/Chat.test.tsx index 9d58e8e6..c9617fa8 100644 --- a/src/components/Chat/Chat.test.tsx +++ b/src/components/Chat/Chat.test.tsx @@ -10,11 +10,15 @@ const mockState = vi.hoisted(() => ({ handler: undefined as ((value: { content: string; images?: string[] }) => void) | undefined, history: [] as string[], + isActive: false, + restoreQueuedMessage: undefined as (() => string | undefined) | undefined, testInput: '', shouldReset: false, clear() { this.handler = undefined; this.history = []; + this.isActive = false; + this.restoreQueuedMessage = undefined; this.testInput = ''; this.shouldReset = true; }, @@ -217,6 +221,8 @@ vi.mock('./ChatInput', () => ({ history?: string[]; onSubmit?: (value: { content: string; images?: string[] }) => void; onInterrupt?: () => void; + onRestoreQueuedMessage?: () => string | undefined; + isActive?: boolean; isDisabled?: boolean; }) => { if (props.onSubmit) { @@ -224,6 +230,8 @@ vi.mock('./ChatInput', () => ({ } mockState.history = props.history ?? []; + mockState.isActive = props.isActive ?? false; + mockState.restoreQueuedMessage = props.onRestoreQueuedMessage; if (props.onInterrupt) { interruptState.handler = props.onInterrupt; @@ -404,6 +412,113 @@ describe('Chat', () => { expect(onMessagesChange).toHaveBeenCalled(); }); + it('returns undefined from restoreQueuedMessage when queue is empty', async () => { + const chat = ( + + ); + const { rerender } = renderWithTheme(chat); + rerender(chat); + await time.tick(); + + expect(mockState.restoreQueuedMessage?.()).toBeUndefined(); + }); + + it('shows queued messages and restores the latest message for editing', async () => { + let releaseStream: (() => void) | undefined; + let streamIndex = 0; + vi.mocked(ollama.streamChat).mockImplementation(async function* () { + streamIndex += 1; + if (streamIndex === 1) { + await new Promise((resolve) => { + releaseStream = resolve; + }); + } + yield { type: 'content', content: 'Done' }; + }); + + const chat = ( + + ); + const { lastFrame, rerender } = renderWithTheme(chat); + + mockState.handler?.({ content: 'first' }); + rerender(chat); + await time.tick(); + mockState.handler?.({ content: 'second' }); + mockState.handler?.({ content: 'third' }); + rerender(chat); + await time.tick(); + + expect(lastFrame()).toContain('Queued messages:'); + expect(lastFrame()).toContain('second'); + expect(lastFrame()).toContain('third'); + expect(mockState.restoreQueuedMessage?.()).toBe('third'); + rerender(chat); + await time.tick(); + expect(lastFrame()).not.toContain(' ↳ third'); + + releaseStream?.(); + await vi.waitFor(() => { + rerender(chat); + expect(ollama.streamChat).toHaveBeenCalledTimes(2); + expect(lastFrame()).not.toContain('Queued messages:'); + }); + }); + + it('does not queue a message with images submitted while a turn is active', async () => { + let releaseStream: (() => void) | undefined; + let streamIndex = 0; + vi.mocked(ollama.streamChat).mockImplementation(async function* () { + streamIndex += 1; + if (streamIndex === 1) { + await new Promise((resolve) => { + releaseStream = resolve; + }); + } + yield { type: 'content', content: 'Done' }; + }); + + const chat = ( + + ); + const { lastFrame, rerender } = renderWithTheme(chat); + + mockState.handler?.({ content: 'first' }); + rerender(chat); + await time.tick(); + + mockState.handler?.({ content: 'with image', images: ['/tmp/img.png'] }); + rerender(chat); + await time.tick(); + + expect(lastFrame()).not.toContain('Queued messages:'); + + releaseStream?.(); + await vi.waitFor(() => { + rerender(chat); + expect(ollama.streamChat).toHaveBeenCalledTimes(1); + expect(lastFrame()).not.toContain('Thinking'); + }); + }); + it('derives prompt history from user messages and excludes slash commands', async () => { renderWithTheme( { choosePlanMode(MODE.AUTO); await time.tick(); - }, 20_000); + }); it('executes an approved plan immediately in auto mode', async () => { const { streamChat } = ollama; diff --git a/src/components/Chat/Chat.tsx b/src/components/Chat/Chat.tsx index 90b5aa6e..cd501a5b 100644 --- a/src/components/Chat/Chat.tsx +++ b/src/components/Chat/Chat.tsx @@ -19,7 +19,7 @@ import { agents, memory, ollama, tools } from '@/utils'; import { ChatInput, type SubmittedInput } from './ChatInput'; import { MEMORY_COMMANDS } from './CommandMenu'; import { ChatActionType, InterruptReason } from './constants'; -import { useCompact, useRunTurn } from './hooks'; +import { useCompact, useMessageQueue, useRunTurn } from './hooks'; import { chatReducer, createInitialChatState } from './reducer'; import { ToolProgress } from './ToolProgress'; @@ -111,6 +111,7 @@ export function Chat({ toolProgress, } = state; const abortControllerRef = useRef(null); + const activeTurnRef = useRef(false); const persistedSnapshotRef = useRef(''); const [compactError, setCompactError] = useState(null); const compact = useCompact({ @@ -125,6 +126,7 @@ export function Chat({ }); useEffect(() => { + activeTurnRef.current = false; dispatch({ type: ChatActionType.ResetSession, messages: sessionMessages, @@ -190,6 +192,7 @@ export function Chat({ type: ChatActionType.SetLoading, isLoading: true, }); + activeTurnRef.current = true; // Add instruction to execute the plan const executeInstruction: ollama.Message = { @@ -202,7 +205,11 @@ export function Chat({ const executeMessages = [...planMessages, executeInstruction]; - await runTurn(executeMessages, selectedMode); + try { + await runTurn(executeMessages, selectedMode); + } finally { + activeTurnRef.current = false; + } }, [onModeChange, pendingPlan, runTurn], ); @@ -222,64 +229,104 @@ export function Chat({ type: ChatActionType.SetLoading, isLoading: true, }); + activeTurnRef.current = true; + + try { + switch (decision) { + case DECISION.APPROVE: { + dispatch({ + type: ChatActionType.SetStreamingMessage, + message: { role: ROLE.ASSISTANT, content: '' }, + }); + const result = await tools.executeToolCall(toolCall); + + const toolResultMessage: ollama.Message = { + role: ROLE.SYSTEM, + content: tools.formatToolResultContent( + toolCall.function.name, + result, + toolCall.function.arguments, + ), + toolResult: { + name: toolCall.function.name, + // v8 ignore next + ...(result.diff ? { diff: result.diff } : {}), + // v8 ignore next + ...(result.error ? { error: result.error } : {}), + }, + }; - switch (decision) { - case DECISION.APPROVE: { - dispatch({ - type: ChatActionType.SetStreamingMessage, - message: { role: ROLE.ASSISTANT, content: '' }, - }); - const result = await tools.executeToolCall(toolCall); - - const toolResultMessage: ollama.Message = { - role: ROLE.SYSTEM, - content: tools.formatToolResultContent( - toolCall.function.name, - result, - toolCall.function.arguments, - ), - toolResult: { - name: toolCall.function.name, - // v8 ignore next - ...(result.diff ? { diff: result.diff } : {}), - // v8 ignore next - ...(result.error ? { error: result.error } : {}), - }, - }; + const newMessages = [...approvedMessages, toolResultMessage]; + dispatch({ + type: ChatActionType.CommitMessages, + messages: newMessages, + }); - const newMessages = [...approvedMessages, toolResultMessage]; - dispatch({ - type: ChatActionType.CommitMessages, - messages: newMessages, - }); + await runTurn(newMessages, mode); + break; + } - await runTurn(newMessages, mode); - break; + case DECISION.REJECT: { + const toolResultMessage: ollama.Message = { + role: ROLE.SYSTEM, + content: tools.formatToolResultContent( + toolCall.function.name, + { + content: '', + error: 'Tool call rejected by user', + }, + toolCall.function.arguments, + ), + }; + dispatch({ + type: ChatActionType.ToolRejected, + messages: [...approvedMessages, toolResultMessage], + }); + break; + } } + } finally { + activeTurnRef.current = false; + } + }, + [mode, pendingToolCall, runTurn], + ); - case DECISION.REJECT: { - const toolResultMessage: ollama.Message = { - role: ROLE.SYSTEM, - content: tools.formatToolResultContent( - toolCall.function.name, - { - content: '', - error: 'Tool call rejected by user', - }, - toolCall.function.arguments, - ), - }; - dispatch({ - type: ChatActionType.ToolRejected, - messages: [...approvedMessages, toolResultMessage], - }); - break; + const runUserPrompt = useCallback( + async ({ content, images }: SubmittedInput) => { + const userMessage: ollama.Message = { + role: ROLE.USER, + content, + ...(images?.length ? { images } : {}), + }; + + const updatedMessages = [...messages, userMessage]; + dispatch({ + type: ChatActionType.StartTurn, + message: userMessage, + }); + activeTurnRef.current = true; + + try { + if (mode === MODE.PLAN) { + await runTurnReadOnly(updatedMessages); + } else { + await runTurn(updatedMessages); } + } finally { + activeTurnRef.current = false; } }, - [mode, pendingToolCall, runTurn], + [messages, mode, runTurn, runTurnReadOnly], ); + const { enqueueMessage, queuedMessages, restoreLatestMessage } = + useMessageQueue({ + isPaused: isLoading || !!pendingPlan || !!pendingToolCall, + onRunMessage: runUserPrompt, + resetKey: sessionId, + }); + const handleSubmit = useCallback( async ({ content, images }: SubmittedInput) => { const userContent = content.trim(); @@ -289,6 +336,18 @@ export function Chat({ return; } + if (activeTurnRef.current || isLoading) { + if ( + userContent && + !images?.length && + !userContent.startsWith('/') && + !userContent.startsWith('!') + ) { + enqueueMessage({ content: userContent }); + } + return; + } + if (userContent === '/compact' && !images?.length) { const error = await compact(); if (error) { @@ -338,46 +397,38 @@ export function Chat({ type: ChatActionType.StartTurn, message: { role: ROLE.USER, content: userContent }, }); + activeTurnRef.current = true; - const result = await tools.runShell(command); - const output = result.content.trim(); - const errorLine = result.error ? `\nError: ${result.error}` : ''; - dispatch({ - type: ChatActionType.AppendMessage, - message: { - role: ROLE.SYSTEM, - content: `$ ${command}${output ? `\n${output}` : ''}${errorLine}`, - }, - }); - dispatch({ - type: ChatActionType.SetLoading, - isLoading: false, - }); + try { + const result = await tools.runShell(command); + const output = result.content.trim(); + const errorLine = result.error ? `\nError: ${result.error}` : ''; + dispatch({ + type: ChatActionType.AppendMessage, + message: { + role: ROLE.SYSTEM, + content: `$ ${command}${output ? `\n${output}` : ''}${errorLine}`, + }, + }); + dispatch({ + type: ChatActionType.SetLoading, + isLoading: false, + }); + } finally { + activeTurnRef.current = false; + } return; } - const userMessage: ollama.Message = { - role: ROLE.USER, - content: userContent, - ...(images?.length ? { images } : {}), - }; - - const updatedMessages = [...messages, userMessage]; - dispatch({ - type: ChatActionType.StartTurn, - message: userMessage, - }); - - // Use plan mode stream if in plan mode, otherwise normal stream - if (mode === MODE.PLAN) { - await runTurnReadOnly(updatedMessages); - } else { - await runTurn(updatedMessages); - } + await runUserPrompt({ content: userContent, images }); }, - [compact, messages, mode, onCommand, runTurn, runTurnReadOnly], + [compact, enqueueMessage, isLoading, onCommand, runUserPrompt], ); + const handleRestoreQueuedMessage = useCallback(() => { + return restoreLatestMessage()?.content; + }, [restoreLatestMessage]); + return ( + + {queuedMessages.length > 0 && ( + + Queued messages: + {queuedMessages.map(({ content }, index) => ( + + + ↳ {content} + + + ))} + + )} + diff --git a/src/components/Chat/ChatInput.test.tsx b/src/components/Chat/ChatInput.test.tsx index 8f603558..8b6751b2 100644 --- a/src/components/Chat/ChatInput.test.tsx +++ b/src/components/Chat/ChatInput.test.tsx @@ -670,6 +670,30 @@ describe('ChatInput', () => { expect(lastFrame()).toContain('[image-1.png]'); }); + it('shows an error when staging a clipboard image via Ctrl+V while a turn is active', async () => { + const { lastFrame, stdin } = renderInput({ isActive: true }); + + stdin.write('\x16'); + await time.tick(); + + expect(clipboard.saveClipboardImage).not.toHaveBeenCalled(); + expect(lastFrame()).toContain("Images can't be queued."); + }); + + it('shows an error when pasting an image path while a turn is active', async () => { + const { lastFrame } = renderInput({ isActive: true }); + const inputProps = mockTextInput.mock.calls.at(-1)?.[0] as + { onChange?: (value: string) => void } | undefined; + + inputProps?.onChange?.( + `"${join(testDirectory, 'screen.png')}" explain this`, + ); + await time.tick(); + + expect(lastFrame()).toContain("Images can't be queued."); + expect(lastFrame()).not.toContain('[screen.png]'); + }); + it('hides the placeholder when an attachment is staged without typed text', async () => { const { lastFrame, stdin } = renderInput(); @@ -969,6 +993,64 @@ describe('ChatInput', () => { expect(onInterrupt).toHaveBeenCalledOnce(); }); + it('accepts and submits normal input while a turn is active', async () => { + const onSubmit = vi.fn(); + const { stdin } = renderInput({ isActive: true, onSubmit }); + stdin.write('follow up'); + await time.tick(); + stdin.write(KEY.ENTER); + await time.tick(); + expect(onSubmit).toHaveBeenCalledWith({ content: 'follow up' }); + }); + + it('keeps commands in the composer while a turn is active', async () => { + const onSubmit = vi.fn(); + const { lastFrame, stdin } = renderInput({ isActive: true, onSubmit }); + stdin.write('!pwd'); + await time.tick(); + stdin.write(KEY.ENTER); + await time.tick(); + expect(onSubmit).not.toHaveBeenCalled(); + expect(lastFrame()).toContain('!pwd'); + expect(lastFrame()).toContain("Commands and images can't be queued."); + }); + + it('interrupts an active turn on Ctrl+C even with a draft', async () => { + const onInterrupt = vi.fn(); + const { lastFrame, stdin } = renderInput({ isActive: true, onInterrupt }); + stdin.write('draft'); + await time.tick(); + stdin.write(KEY.CTRL_C); + await time.tick(); + expect(onInterrupt).toHaveBeenCalledOnce(); + expect(lastFrame()).toContain('draft'); + }); + + it('restores the latest queued message with Up on blank active input', async () => { + const onRestoreQueuedMessage = vi.fn(() => 'queued follow-up'); + const { lastFrame, stdin } = renderInput({ + isActive: true, + onRestoreQueuedMessage, + }); + stdin.write(KEY.UP); + await time.tick(); + expect(onRestoreQueuedMessage).toHaveBeenCalledOnce(); + expect(lastFrame()).toContain('queued follow-up'); + }); + + it('falls through to history navigation when onRestoreQueuedMessage returns undefined on blank active input', async () => { + const onRestoreQueuedMessage = vi.fn(() => undefined); + const { lastFrame, stdin } = renderInput({ + history: ['previous prompt'], + isActive: true, + onRestoreQueuedMessage, + }); + stdin.write(KEY.UP); + await time.tick(); + expect(onRestoreQueuedMessage).toHaveBeenCalledOnce(); + expect(lastFrame()).toContain('previous prompt'); + }); + it('calls onInterrupt on Esc when disabled', async () => { const onInterrupt = vi.fn(); const { stdin } = renderInput({ isDisabled: true, onInterrupt }); diff --git a/src/components/Chat/ChatInput.tsx b/src/components/Chat/ChatInput.tsx index 9a861843..e4ca7940 100644 --- a/src/components/Chat/ChatInput.tsx +++ b/src/components/Chat/ChatInput.tsx @@ -21,8 +21,10 @@ import { useHistorySearch } from './hooks'; interface Props { history: string[]; + isActive?: boolean; isDisabled?: boolean; onInterrupt?: () => void; + onRestoreQueuedMessage?: () => string | undefined; onSubmit: (value: SubmittedInput) => void; } @@ -60,8 +62,10 @@ function cleanupAttachments(attachments: Attachment[]) { export function ChatInput({ history: sessionHistory, + isActive = false, isDisabled = false, onInterrupt, + onRestoreQueuedMessage, onSubmit, }: Props) { const theme = useTheme(); @@ -155,6 +159,11 @@ export function ChatInput({ }, []); const attachClipboardImage = useCallback(() => { + if (isActive) { + setError("Images can't be queued."); + return; + } + try { const path = clipboard.saveClipboardImage( `image-${String(nextClipboardImageRef.current)}`, @@ -164,7 +173,7 @@ export function ChatInput({ } catch (error) { setError(error instanceof Error ? error.message : String(error)); } - }, [stageAttachments]); + }, [isActive, stageAttachments]); const handleSelectFileSuggestion = useCallback( (nextInput: FileSuggestionRef) => { @@ -213,6 +222,11 @@ export function ChatInput({ extractImageAttachments(nextInput); if (nextAttachments.length) { + if (isActive) { + setError("Images can't be queued."); + return; + } + stageAttachments(nextAttachments); setInput(remainingInput); setCursorPosition(remainingInput.length); @@ -227,7 +241,7 @@ export function ChatInput({ resetHistorySearch(); setError(null); }, - [input, resetHistorySearch, stageAttachments], + [input, isActive, resetHistorySearch, stageAttachments], ); const submitAndReset = useCallback( @@ -239,6 +253,16 @@ export function ChatInput({ return; } + if ( + isActive && + (imagePaths.length > 0 || + trimmedInput.startsWith('/') || + trimmedInput.startsWith('!')) + ) { + setError("Commands and images can't be queued."); + return; + } + onSubmit({ content: trimmedInput, ...(imagePaths.length ? { images: imagePaths } : {}), @@ -249,7 +273,7 @@ export function ChatInput({ resetInput(trimmedInput.startsWith('/')); fileSuggestionRef.current = null; }, - [attachments, onSubmit, resetInput], + [attachments, isActive, onSubmit, resetInput], ); const isMultilineInput = input.includes('\n'); @@ -362,6 +386,11 @@ export function ChatInput({ return; } + if (isActive && (isEscape || isCtrlC)) { + onInterrupt?.(); + return; + } + if (historySearch.isActive) { if (isEscape || isCtrlC) { cancelHistorySearch(); @@ -419,6 +448,16 @@ export function ChatInput({ } if (key.upArrow) { + if (isActive && !input && !hasAttachments) { + const queuedMessage = onRestoreQueuedMessage?.(); + if (queuedMessage !== undefined) { + setInput(queuedMessage); + setCursorPosition(queuedMessage.length); + setError(null); + return; + } + } + handleHistoryNavigation('up'); return; } diff --git a/src/components/Chat/hooks/index.ts b/src/components/Chat/hooks/index.ts index 1c9a844a..29240406 100644 --- a/src/components/Chat/hooks/index.ts +++ b/src/components/Chat/hooks/index.ts @@ -1,3 +1,4 @@ export { useCompact } from './useCompact'; export { useHistorySearch } from './useHistorySearch'; +export { useMessageQueue } from './useMessageQueue'; export { useRunTurn } from './useRunTurn'; diff --git a/src/components/Chat/hooks/useMessageQueue.ts b/src/components/Chat/hooks/useMessageQueue.ts new file mode 100644 index 00000000..bf65e136 --- /dev/null +++ b/src/components/Chat/hooks/useMessageQueue.ts @@ -0,0 +1,67 @@ +import { useCallback, useEffect, useRef, useState } from 'react'; + +import type { SubmittedInput } from '../ChatInput'; + +interface UseMessageQueueOptions { + isPaused: boolean; + onRunMessage: (message: SubmittedInput) => Promise; + resetKey: string; +} + +interface MessageQueue { + enqueueMessage: (message: SubmittedInput) => void; + queuedMessages: SubmittedInput[]; + restoreLatestMessage: () => SubmittedInput | undefined; +} + +export function useMessageQueue({ + isPaused, + onRunMessage, + resetKey, +}: UseMessageQueueOptions): MessageQueue { + const isDrainingRef = useRef(false); + const [isDraining, setIsDraining] = useState(false); + const [queuedMessages, setQueuedMessages] = useState([]); + + useEffect(() => { + isDrainingRef.current = false; + setIsDraining(false); + setQueuedMessages([]); + }, [resetKey]); + + useEffect(() => { + if ( + isPaused || + !queuedMessages.length || + isDraining || + isDrainingRef.current + ) { + return; + } + + const [nextMessage] = queuedMessages; + isDrainingRef.current = true; + setIsDraining(true); + setQueuedMessages((current) => current.slice(1)); + void onRunMessage(nextMessage).finally(() => { + isDrainingRef.current = false; + setIsDraining(false); + }); + }, [isDraining, isPaused, onRunMessage, queuedMessages]); + + const enqueueMessage = useCallback((message: SubmittedInput) => { + setQueuedMessages((current) => [...current, message]); + }, []); + + const restoreLatestMessage = useCallback(() => { + const latestMessage = queuedMessages.at(-1); + if (!latestMessage) { + return undefined; + } + + setQueuedMessages((current) => current.slice(0, -1)); + return latestMessage; + }, [queuedMessages]); + + return { enqueueMessage, queuedMessages, restoreLatestMessage }; +}