diff --git a/ui/goose2/src/features/chat/hooks/useChat.ts b/ui/goose2/src/features/chat/hooks/useChat.ts index b5cbf521a2..d4dbdd6765 100644 --- a/ui/goose2/src/features/chat/hooks/useChat.ts +++ b/ui/goose2/src/features/chat/hooks/useChat.ts @@ -129,6 +129,7 @@ export function useChat( text: string, overridePersona?: { id: string; name?: string }, attachments?: ChatAttachmentDraft[], + truncateBeforeMessageId?: string, ) => { const images = buildAcpImages(attachments); const hasAttachments = (attachments?.length ?? 0) > 0; @@ -242,6 +243,7 @@ export function useChat( images: images?.map( (img) => [img.base64, img.mimeType] as [string, string], ), + truncateBeforeMessageId, }); store.setChatState(sessionId, "idle"); @@ -371,7 +373,9 @@ export function useChat( // Bail if there's nothing to re-send — don't truncate if (!hasContent) return; - // Truncate from the user message onward (removes user msg + all responses) + // Truncate locally for immediate UI feedback; the backend truncates + // via _meta.truncate_before_message_id on the PromptRequest. + const truncateMessageId = userMessage.id; store.setMessages(sessionId, sessionMessages.slice(0, truncateFromIndex)); const targetPersonaId = userMessage.metadata?.targetPersonaId; @@ -382,6 +386,7 @@ export function useChat( ? { id: targetPersonaId, name: targetPersonaName } : undefined, attachmentDrafts.length > 0 ? attachmentDrafts : undefined, + truncateMessageId, ); }, [sessionId, store, sendMessage, chatState], diff --git a/ui/goose2/src/features/chat/ui/ChatView.tsx b/ui/goose2/src/features/chat/ui/ChatView.tsx index 84028e1a0e..39538b69f4 100644 --- a/ui/goose2/src/features/chat/ui/ChatView.tsx +++ b/ui/goose2/src/features/chat/ui/ChatView.tsx @@ -1,5 +1,4 @@ import { useState, useEffect, useRef, useCallback, useMemo } from "react"; -import { flushSync } from "react-dom"; import { useTranslation } from "react-i18next"; import { AnimatePresence } from "motion/react"; import { MessageTimeline } from "./MessageTimeline"; @@ -360,9 +359,7 @@ export function ChatView({ // Unified deferred send — a single ref + effect drains any pending send // that had to wait for a state update (persona switch or edit truncation). const pendingSend = useRef<{ - reason: "persona" | "edit"; text: string; - persona?: { id: string; name?: string }; attachments?: ChatAttachmentDraft[]; } | null>(null); const queue = useMessageQueue(activeSessionId, chatState, sendMessage); @@ -389,7 +386,7 @@ export function ChatView({ } handlePersonaChange(personaId); // Defer the send until after persona state updates - pendingSend.current = { reason: "persona", text, attachments }; + pendingSend.current = { text, attachments }; return; } // Queue if agent is busy and no message already queued @@ -413,52 +410,77 @@ export function ChatView({ ); /** Save an inline edit: truncate history from the edited message onward, then send the new text. */ + /** Save an inline edit: truncate locally for instant UI feedback, then + * send with truncation metadata so the backend also rewinds both the + * display log and LLM context. + * + * When the user edits while streaming, stopStreaming() sets chatState + * to "idle" in the store, but sendMessage's closure still captures the + * stale "streaming" value and silently bails. We defer the send into a + * ref so the next render — which carries the fresh "idle" chatState and + * a new sendMessage — picks it up via the useEffect below. */ + const pendingEditSend = useRef<{ + text: string; + truncateMessageId: string; + persona?: { id: string; name?: string }; + attachments?: ChatAttachmentDraft[]; + } | null>(null); const handleSaveEdit = useCallback( (messageId: string, text: string) => { - if (chatState !== "idle") { + const wasIdle = chatState === "idle"; + if (!wasIdle) { stopStreaming(); } const store = useChatStore.getState(); const allMessages = store.messagesBySession[activeSessionId] ?? []; const editIndex = allMessages.findIndex((m) => m.id === messageId); if (editIndex === -1) { - // Message vanished — bail to avoid appending to untruncated history store.setEditingMessageId(activeSessionId, null); return; } - // Preserve persona and attachments from the original message const originalMessage = allMessages[editIndex]; const targetPersonaId = originalMessage.metadata?.targetPersonaId; const targetPersonaName = originalMessage.metadata?.targetPersonaName; const originalAttachments = rebuildAttachmentDrafts(originalMessage); + const persona = targetPersonaId + ? { id: targetPersonaId, name: targetPersonaName } + : undefined; + const attachments = + originalAttachments.length > 0 ? originalAttachments : undefined; - // Use flushSync to commit state updates synchronously so we can call - // sendMessage immediately — no timing gap, no race condition. - flushSync(() => { - store.setMessages(activeSessionId, allMessages.slice(0, editIndex)); - store.setEditingMessageId(activeSessionId, null); - store.setChatState(activeSessionId, "idle"); - }); + // Local truncation for immediate UI feedback. + store.setMessages(activeSessionId, allMessages.slice(0, editIndex)); + store.setEditingMessageId(activeSessionId, null); + store.setChatState(activeSessionId, "idle"); - sendMessage( - text, - targetPersonaId - ? { id: targetPersonaId, name: targetPersonaName } - : undefined, - originalAttachments.length > 0 ? originalAttachments : undefined, - ); + if (wasIdle) { + // chatState was already "idle" — sendMessage's closure is fresh. + sendMessage(text, persona, attachments, messageId); + } else { + // Defer until React re-renders with fresh chatState / sendMessage. + pendingEditSend.current = { + text, + truncateMessageId: messageId, + persona, + attachments, + }; + } }, [activeSessionId, chatState, stopStreaming, sendMessage], ); - // Single drain effect for deferred sends (persona switch path only now — - // edit path uses flushSync and calls sendMessage synchronously). useEffect(() => { - if ( - pendingSend.current && - pendingSend.current.reason === "persona" && - selectedPersona - ) { + if (pendingEditSend.current && chatState === "idle") { + const { text, truncateMessageId, persona, attachments } = + pendingEditSend.current; + pendingEditSend.current = null; + sendMessage(text, persona, attachments, truncateMessageId); + } + }, [chatState, sendMessage]); + + // Drain deferred sends (persona switch only). + useEffect(() => { + if (pendingSend.current && selectedPersona) { const { text, attachments } = pendingSend.current; pendingSend.current = null; sendMessage(text, undefined, attachments); diff --git a/ui/goose2/src/shared/api/acp.ts b/ui/goose2/src/shared/api/acp.ts index 8bfd8914fd..be1b7c7ee9 100644 --- a/ui/goose2/src/shared/api/acp.ts +++ b/ui/goose2/src/shared/api/acp.ts @@ -18,6 +18,10 @@ export interface AcpSendMessageOptions { personaName?: string; /** Image attachments as [base64Data, mimeType] pairs. */ images?: [string, string][]; + /** When set, the backend truncates this message and all subsequent messages + * from both the display log and LLM context before appending the new prompt. + * Used by edit and retry to rewind the conversation. */ + truncateBeforeMessageId?: string; } export interface AcpPrepareSessionOptions { @@ -36,7 +40,7 @@ export async function acpSendMessage( prompt: string, options: AcpSendMessageOptions = {}, ): Promise { - const { systemPrompt, personaId, images } = options; + const { systemPrompt, personaId, images, truncateBeforeMessageId } = options; const gooseSessionId = sessionTracker.getGooseSessionId(sessionId, personaId); if (!gooseSessionId) { @@ -55,10 +59,14 @@ export async function acpSendMessage( } } + const meta = truncateBeforeMessageId + ? { truncate_before_message_id: truncateBeforeMessageId } + : undefined; + const messageId = crypto.randomUUID(); setActiveMessageId(gooseSessionId, messageId); - await directAcp.prompt(gooseSessionId, content); + await directAcp.prompt(gooseSessionId, content, meta); clearActiveMessageId(gooseSessionId); } diff --git a/ui/goose2/src/shared/api/acpApi.ts b/ui/goose2/src/shared/api/acpApi.ts index b992c1be05..fb382997cb 100644 --- a/ui/goose2/src/shared/api/acpApi.ts +++ b/ui/goose2/src/shared/api/acpApi.ts @@ -135,7 +135,8 @@ export async function loadSession( export async function prompt( sessionId: string, content: ContentBlock[], + meta?: Record, ): Promise { const client = await getClient(); - return client.prompt({ sessionId, prompt: content }); + return client.prompt({ sessionId, prompt: content, _meta: meta }); }