diff --git a/desktop/src/components/chat/ChatInput.test.tsx b/desktop/src/components/chat/ChatInput.test.tsx index 7bddf9e3..67c989b5 100644 --- a/desktop/src/components/chat/ChatInput.test.tsx +++ b/desktop/src/components/chat/ChatInput.test.tsx @@ -452,22 +452,27 @@ describe('ChatInput file mentions', () => { attachments: [], }) expect(screen.queryByTestId('pending-user-message')).not.toBeInTheDocument() - expect(useChatStore.getState().sessions[sessionId]?.messages.at(-1)).toMatchObject({ - type: 'assistant_text', - content: 'working', - }) + expect(useChatStore.getState().sessions[sessionId]?.messages).toMatchObject([ + { type: 'assistant_text', content: 'workingstill answering' }, + { type: 'user_text', content: 'please adjust the current direction' }, + ]) act(() => { + useChatStore.getState().handleServerMessage(sessionId, { + type: 'tool_result', + toolUseId: 'tool-1', + content: 'tool finished', + isError: false, + }) useChatStore.getState().handleServerMessage(sessionId, { type: 'user_message_replay', content: 'please adjust the current direction', }) }) - expect(useChatStore.getState().sessions[sessionId]?.messages.at(-1)).toMatchObject({ - type: 'user_text', - content: 'please adjust the current direction', - }) + const guidedMessages = useChatStore.getState().sessions[sessionId]?.messages + .filter((message) => message.type === 'user_text' && message.content === 'please adjust the current direction') + expect(guidedMessages).toHaveLength(1) }) it('edits and deletes queued prompts without sending them', async () => { diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index 1685ef63..321c07a6 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -1372,9 +1372,8 @@ export const useChatStore = create((set, get) => ({ const queuedMessage = (session?.queuedUserMessages ?? []).find((message) => message.id === messageId) if (!session || !queuedMessage) return - get().removeQueuedUserMessage(sessionId, messageId) - if (session.chatState === 'idle') { + get().removeQueuedUserMessage(sessionId, messageId) get().sendMessage( sessionId, queuedMessage.content, @@ -1387,6 +1386,22 @@ export const useChatStore = create((set, get) => ({ return } + const now = Date.now() + set((state) => ({ + sessions: updateSessionIn(state.sessions, sessionId, (currentSession) => { + const pendingText = `${currentSession.streamingText}${consumePendingDelta(sessionId)}` + const baseMessages = pendingText.trim() + ? appendAssistantTextMessage(currentSession.messages, pendingText, now) + : currentSession.messages + return { + messages: appendOptimisticQueuedUserMessage(baseMessages, queuedMessage, now), + queuedUserMessages: (currentSession.queuedUserMessages ?? []) + .filter((message) => message.id !== messageId), + ...(pendingText.trim() ? { streamingText: '' } : {}), + } + }), + })) + wsManager.send(sessionId, { type: 'user_message', content: queuedMessage.content, @@ -2740,6 +2755,19 @@ function appendReplayedUserMessage( if (!displayContent) return messages const modelContent = parsed.modelContent ?? content.trim() + const optimisticIndex = findOptimisticQueuedUserMessageIndex(messages, modelContent) + if (optimisticIndex >= 0) { + const optimisticMessage = messages[optimisticIndex] + if (optimisticMessage?.type === 'user_text') { + const { optimisticQueued: _optimisticQueued, ...confirmedMessage } = optimisticMessage + return [ + ...messages.slice(0, optimisticIndex), + confirmedMessage, + ...messages.slice(optimisticIndex + 1), + ] + } + } + const last = messages[messages.length - 1] if ( last?.type === 'user_text' && @@ -2761,6 +2789,63 @@ function appendReplayedUserMessage( ] } +function appendOptimisticQueuedUserMessage( + messages: UIMessage[], + message: QueuedUserMessage, + timestamp: number, +): UIMessage[] { + const displayContent = message.displayContent.trim() + const modelContent = message.content.trim() + const attachments = mapQueuedDisplayAttachments(message.displayAttachments) + if (!displayContent && !attachments) return messages + + return [ + ...messages, + { + id: nextId(), + type: 'user_text', + content: displayContent, + ...(modelContent && modelContent !== displayContent ? { modelContent } : {}), + ...(attachments ? { attachments } : {}), + timestamp, + optimisticQueued: true, + }, + ] +} + +function mapQueuedDisplayAttachments(attachments?: AttachmentRef[]): UIAttachment[] | undefined { + if (!attachments?.length) return undefined + return attachments.map((attachment) => ({ + type: attachment.type, + name: attachment.name || attachment.path || attachment.mimeType || attachment.type, + path: attachment.path, + data: attachment.data, + mimeType: attachment.mimeType, + isDirectory: attachment.isDirectory, + lineStart: attachment.lineStart, + lineEnd: attachment.lineEnd, + note: attachment.note, + quote: attachment.quote, + })) +} + +function findOptimisticQueuedUserMessageIndex( + messages: UIMessage[], + modelContent: string, +): number { + for (let index = messages.length - 1; index >= 0; index -= 1) { + const message = messages[index] + if ( + message?.type === 'user_text' && + message.optimisticQueued && + (message.modelContent ?? message.content).trim() === modelContent + ) { + return index + } + } + return -1 +} + function replaceQueuedMessageDisplayContent( message: QueuedUserMessage, nextDisplayContent: string, diff --git a/desktop/src/types/chat.ts b/desktop/src/types/chat.ts index a4e7c614..15a17875 100644 --- a/desktop/src/types/chat.ts +++ b/desktop/src/types/chat.ts @@ -265,7 +265,7 @@ export type TaskSummaryItem = { } export type UIMessage = - | { id: string; type: 'user_text'; content: string; modelContent?: string; transcriptMessageId?: string; timestamp: number; attachments?: UIAttachment[]; pending?: boolean } + | { id: string; type: 'user_text'; content: string; modelContent?: string; transcriptMessageId?: string; timestamp: number; attachments?: UIAttachment[]; pending?: boolean; optimisticQueued?: boolean } | { id: string; type: 'assistant_text'; content: string; transcriptMessageId?: string; timestamp: number; model?: string } | { id: string; type: 'thinking'; content: string; timestamp: number } | {