From e5dc7095fdda00efbb38bd869b971d162b8db65a Mon Sep 17 00:00:00 2001 From: Taylor Ho Date: Mon, 13 Apr 2026 17:53:33 -1000 Subject: [PATCH] feat: wire retry and edit buttons on user messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add retryMessage(messageId) to useChat — truncates from the target message onward and resends the preceding user message. Works for both user and assistant messages. - Add editMessage(messageId) to useChat — truncates from the target user message onward and populates the chat input draft with the original text for revision. - Wire onRetryMessage and onEditMessage from ChatView through MessageTimeline to MessageBubble. - Remove role gate so retry icon shows on both user and assistant messages (edit remains user-only). --- ui/goose2/src/features/chat/hooks/useChat.ts | 104 +++++++++++++++--- .../src/features/chat/stores/chatStore.ts | 16 +++ ui/goose2/src/features/chat/ui/ChatInput.tsx | 25 ++++- ui/goose2/src/features/chat/ui/ChatView.tsx | 32 ++++++ .../src/features/chat/ui/MessageBubble.tsx | 2 +- .../src/features/chat/ui/MessageTimeline.tsx | 4 +- .../src/shared/i18n/locales/en/chat.json | 4 + 7 files changed, 164 insertions(+), 23 deletions(-) diff --git a/ui/goose2/src/features/chat/hooks/useChat.ts b/ui/goose2/src/features/chat/hooks/useChat.ts index 1150f1074a..b8ac850613 100644 --- a/ui/goose2/src/features/chat/hooks/useChat.ts +++ b/ui/goose2/src/features/chat/hooks/useChat.ts @@ -326,34 +326,98 @@ export function useChat( }); }, [getStreamingPersonaId, store, sessionId]); + /** Retry from a specific message — truncates everything from that message + * onward and re-sends the preceding user message. Works for both user + * messages (re-send that message) and assistant messages (re-send the + * user message that triggered it). */ + const retryMessage = useCallback( + async (messageId: string) => { + // Guard: don't truncate while the agent is still responding + if (chatState === "streaming" || chatState === "thinking") return; + + const sessionMessages = store.messagesBySession[sessionId] ?? []; + const targetIndex = sessionMessages.findIndex( + (m) => m.id === messageId, + ); + if (targetIndex === -1) return; + + const targetMessage = sessionMessages[targetIndex]; + + // Determine which user message to re-send + let userMessage: (typeof sessionMessages)[number] | undefined; + let truncateFromIndex: number; + + if (targetMessage.role === "user") { + userMessage = targetMessage; + truncateFromIndex = targetIndex; + } else { + // Find the preceding user message + const userIndex = findLastIndex( + sessionMessages.slice(0, targetIndex + 1), + (m) => m.role === "user", + ); + if (userIndex === -1) return; + userMessage = sessionMessages[userIndex]; + truncateFromIndex = userIndex; + } + + // Truncate from the user message onward (removes user msg + all responses) + store.setMessages(sessionId, sessionMessages.slice(0, truncateFromIndex)); + + const textContent = userMessage.content.find((c) => c.type === "text"); + if (textContent && "text" in textContent) { + const targetPersonaId = userMessage.metadata?.targetPersonaId; + const targetPersonaName = userMessage.metadata?.targetPersonaName; + await sendMessage( + textContent.text, + targetPersonaId + ? { id: targetPersonaId, name: targetPersonaName } + : undefined, + ); + } + }, + [sessionId, store, sendMessage, chatState], + ); + const retryLastMessage = useCallback(async () => { const sessionMessages = store.messagesBySession[sessionId] ?? []; - // Find the last user message const lastUserIndex = findLastIndex( sessionMessages, (m) => m.role === "user", ); if (lastUserIndex === -1) return; + await retryMessage(sessionMessages[lastUserIndex].id); + }, [sessionId, store, retryMessage]); - const lastUserMessage = sessionMessages[lastUserIndex]; + /** Enter edit mode for a user message — non-destructive. Populates the + * input draft with the original text and sets editing state. Truncation + * happens only when the user actually sends (handled in ChatView). */ + const editMessage = useCallback( + (messageId: string) => { + // Guard: don't enter edit mode while the agent is still responding + if (chatState === "streaming" || chatState === "thinking") return; - // Remove all messages after (and including) the last assistant response - const messagesToKeep = sessionMessages.slice(0, lastUserIndex); - store.setMessages(sessionId, messagesToKeep); + const sessionMessages = store.messagesBySession[sessionId] ?? []; + const target = sessionMessages.find((m) => m.id === messageId); + if (!target || target.role !== "user") return; - // Extract the text and resend - const textContent = lastUserMessage.content.find((c) => c.type === "text"); - if (textContent && "text" in textContent) { - const targetPersonaId = lastUserMessage.metadata?.targetPersonaId; - const targetPersonaName = lastUserMessage.metadata?.targetPersonaName; - await sendMessage( - textContent.text, - targetPersonaId - ? { id: targetPersonaId, name: targetPersonaName } - : undefined, - ); - } - }, [sessionId, store, sendMessage]); + // Extract the original text to pre-fill the input + const textContent = target.content.find((c) => c.type === "text"); + const originalText = + textContent && "text" in textContent ? textContent.text : ""; + + // Enter edit mode — history stays intact until send + store.setEditingMessageId(sessionId, messageId); + store.setDraft(sessionId, originalText); + }, + [sessionId, store, chatState], + ); + + /** Cancel edit mode — clears editing state and draft. */ + const cancelEdit = useCallback(() => { + store.setEditingMessageId(sessionId, null); + store.clearDraft(sessionId); + }, [sessionId, store]); const clearChat = useCallback(() => { abortRef.current?.abort(); @@ -375,6 +439,10 @@ export function useChat( stopGeneration, stopStreaming, retryLastMessage, + retryMessage, + editMessage, + cancelEdit, + editingMessageId: store.editingMessageIdBySession[sessionId] ?? null, clearChat, isStreaming, }; diff --git a/ui/goose2/src/features/chat/stores/chatStore.ts b/ui/goose2/src/features/chat/stores/chatStore.ts index f7781b8ef9..5d404cefef 100644 --- a/ui/goose2/src/features/chat/stores/chatStore.ts +++ b/ui/goose2/src/features/chat/stores/chatStore.ts @@ -72,6 +72,8 @@ interface ChatStoreState { isConnected: boolean; loadingSessionIds: Set; scrollTargetMessageBySession: Record; + /** Message ID currently being edited per session (non-destructive edit mode). */ + editingMessageIdBySession: Record; } interface ChatStoreActions { @@ -108,6 +110,7 @@ interface ChatStoreActions { dismissQueuedMessage: (sessionId: string) => void; setDraft: (sessionId: string, text: string) => void; clearDraft: (sessionId: string) => void; + setEditingMessageId: (sessionId: string, messageId: string | null) => void; setSessionLoading: (sessionId: string, loading: boolean) => void; setScrollTargetMessage: ( sessionId: string, @@ -126,6 +129,7 @@ export const useChatStore = create((set, get) => ({ sessionStateById: {}, queuedMessageBySession: {}, draftsBySession: loadCachedDrafts(), + editingMessageIdBySession: {}, activeSessionId: null, isConnected: false, loadingSessionIds: new Set(), @@ -430,6 +434,15 @@ export const useChatStore = create((set, get) => ({ persistDrafts(get().draftsBySession); }, + setEditingMessageId: (sessionId, messageId) => { + set((state) => ({ + editingMessageIdBySession: { + ...state.editingMessageIdBySession, + [sessionId]: messageId, + }, + })); + }, + // Session loading (replay) setSessionLoading: (sessionId, loading) => set((state) => { @@ -477,12 +490,15 @@ export const useChatStore = create((set, get) => ({ const { [sessionId]: ____, ...remainingDrafts } = state.draftsBySession; const { [sessionId]: _____, ...remainingTargets } = state.scrollTargetMessageBySession; + const { [sessionId]: ______, ...remainingEditing } = + state.editingMessageIdBySession; return { messagesBySession: rest, sessionStateById: remainingSessionState, queuedMessageBySession: remainingQueued, draftsBySession: remainingDrafts, scrollTargetMessageBySession: remainingTargets, + editingMessageIdBySession: remainingEditing, activeSessionId: state.activeSessionId === sessionId ? null : state.activeSessionId, }; diff --git a/ui/goose2/src/features/chat/ui/ChatInput.tsx b/ui/goose2/src/features/chat/ui/ChatInput.tsx index fdf2bea7d0..2b816ac5ff 100644 --- a/ui/goose2/src/features/chat/ui/ChatInput.tsx +++ b/ui/goose2/src/features/chat/ui/ChatInput.tsx @@ -1,6 +1,6 @@ import { useState, useRef, useCallback, useEffect, useMemo } from "react"; import { open } from "@tauri-apps/plugin-dialog"; -import { X } from "lucide-react"; +import { Pencil, X } from "lucide-react"; import { useTranslation } from "react-i18next"; import type { AcpProvider } from "@/shared/api/acp"; import type { Persona } from "@/shared/types/agents"; @@ -64,6 +64,10 @@ interface ChatInputProps { }) => void; contextTokens?: number; contextLimit?: number; + /** Message ID currently being edited (non-null = edit mode active). */ + editingMessageId?: string | null; + /** Cancel edit mode — clears editing state and restores input. */ + onCancelEdit?: () => void; } export function ChatInput({ @@ -94,6 +98,8 @@ export function ChatInput({ onCreateProject, contextTokens = 0, contextLimit = 0, + editingMessageId = null, + onCancelEdit, }: ChatInputProps) { const { t } = useTranslation("chat"); const [text, setTextRaw] = useState(initialValue); @@ -395,6 +401,23 @@ export function ChatInput({ )} + {editingMessageId && ( +
+ + + {t("edit.label")} + + +
+ )} +