diff --git a/ui/goose2/scripts/check-file-sizes.mjs b/ui/goose2/scripts/check-file-sizes.mjs index 0c45b763c29e..3dc02998a18b 100644 --- a/ui/goose2/scripts/check-file-sizes.mjs +++ b/ui/goose2/scripts/check-file-sizes.mjs @@ -11,9 +11,29 @@ const EXCEPTIONS = { "Drag-and-drop handlers plus activeProjectId highlight for draft-in-project sessions.", }, "src/features/chat/ui/ChatView.tsx": { - limit: 535, + limit: 600, justification: - "ACP prewarm guards, project-aware working dir selection, working context sync, and chat bootstrapping still live together here.", + "ACP prewarm guards, project-aware working dir selection, working context sync, chat bootstrapping, and inline edit/retry orchestration still live together here.", + }, + "src/features/chat/hooks/__tests__/useChat.test.ts": { + limit: 840, + justification: + "Edit/retry integration tests cover attachment preservation, persona round-tripping, truncation, and draft lifecycle in one cohesive suite.", + }, + "src/features/chat/stores/chatStore.ts": { + limit: 520, + justification: + "Edit-mode draft state, per-session editing message tracking, and cleanup logic added for inline edit/retry.", + }, + "src/features/chat/ui/MessageBubble.tsx": { + limit: 610, + justification: + "Inline edit textarea, save/cancel controls, and attachment re-display for edit/retry colocated with message rendering.", + }, + "src/features/chat/ui/__tests__/MessageBubble.test.tsx": { + limit: 560, + justification: + "Edit/retry button visibility, inline edit save/cancel, and attachment preservation tests added to existing bubble suite.", }, "src/features/chat/ui/__tests__/ContextPanel.test.tsx": { limit: 550, diff --git a/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts b/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts index 141826952fa6..f59ae2b5daf1 100644 --- a/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts +++ b/ui/goose2/src/features/chat/hooks/__tests__/useChat.test.ts @@ -59,6 +59,8 @@ describe("useChat", () => { sessionStateById: {}, activeSessionId: null, isConnected: true, + editingMessageIdBySession: {}, + draftsBySession: {}, }); useChatSessionStore.setState({ sessions: [], @@ -435,4 +437,399 @@ describe("useChat", () => { }, ]); }); + + // --------------------------------------------------------------------------- + // retryMessage + // --------------------------------------------------------------------------- + + describe("retryMessage", () => { + function seedConversation( + sessionId: string, + msgs: Array<{ + id: string; + role: "user" | "assistant"; + text: string; + personaId?: string; + personaName?: string; + }>, + ) { + const messages: Message[] = msgs.map((m) => ({ + id: m.id, + role: m.role, + created: Date.now(), + content: [{ type: "text", text: m.text }], + metadata: { + userVisible: true, + agentVisible: true, + ...(m.personaId + ? { targetPersonaId: m.personaId, targetPersonaName: m.personaName } + : {}), + }, + })); + useChatStore.getState().setMessages(sessionId, messages); + } + + it("truncates from target user message and re-sends its text", async () => { + mockAcpSendMessage.mockResolvedValue(undefined); + + seedConversation("session-1", [ + { id: "u1", role: "user", text: "First question" }, + { id: "a1", role: "assistant", text: "First answer" }, + { id: "u2", role: "user", text: "Second question" }, + { id: "a2", role: "assistant", text: "Second answer" }, + ]); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("u2"); + }); + + // History should be truncated to just the first exchange, then a new user + // message appended by sendMessage. + const messages = useChatStore.getState().messagesBySession["session-1"]; + // u1, a1 survive the truncation; sendMessage adds a new user message + expect(messages.length).toBeGreaterThanOrEqual(3); + expect(messages[0].id).toBe("u1"); + expect(messages[1].id).toBe("a1"); + // The third message is the re-sent user message (new id) + expect(messages[2].role).toBe("user"); + expect(messages[2].content[0]).toEqual({ + type: "text", + text: "Second question", + }); + + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "goose", + "Second question", + expect.objectContaining({ personaId: undefined }), + ); + }); + + it("preserves persona when retrying a persona-targeted message", async () => { + mockAcpSendMessage.mockResolvedValue(undefined); + + seedConversation("session-1", [ + { + id: "u1", + role: "user", + text: "Hello persona", + personaId: "persona-a", + personaName: "Persona A", + }, + { id: "a1", role: "assistant", text: "Hi there" }, + ]); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("u1"); + }); + + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "goose", + "Hello persona", + expect.objectContaining({ + personaId: "persona-a", + personaName: "Persona A", + }), + ); + }); + + it("handles retrying an assistant message by finding the preceding user message", async () => { + mockAcpSendMessage.mockResolvedValue(undefined); + + seedConversation("session-1", [ + { id: "u1", role: "user", text: "My question" }, + { id: "a1", role: "assistant", text: "My answer" }, + ]); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("a1"); + }); + + // Truncation should start from u1 (the preceding user message) + const messages = useChatStore.getState().messagesBySession["session-1"]; + // Only the re-sent user message should remain (u1 and a1 were truncated) + expect(messages[0].role).toBe("user"); + expect(messages[0].content[0]).toEqual({ + type: "text", + text: "My question", + }); + expect(mockAcpSendMessage).toHaveBeenCalledTimes(1); + }); + + it("guards against retrying while streaming", async () => { + seedConversation("session-1", [ + { id: "u1", role: "user", text: "Hello" }, + { id: "a1", role: "assistant", text: "World" }, + ]); + useChatStore.getState().setChatState("session-1", "streaming"); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("u1"); + }); + + // Messages should be untouched — retry was blocked + const messages = useChatStore.getState().messagesBySession["session-1"]; + expect(messages).toHaveLength(2); + expect(mockAcpSendMessage).not.toHaveBeenCalled(); + }); + + it("guards against retrying while thinking", async () => { + seedConversation("session-1", [ + { id: "u1", role: "user", text: "Hello" }, + ]); + useChatStore.getState().setChatState("session-1", "thinking"); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("u1"); + }); + + const messages = useChatStore.getState().messagesBySession["session-1"]; + expect(messages).toHaveLength(1); + expect(mockAcpSendMessage).not.toHaveBeenCalled(); + }); + }); + + // --------------------------------------------------------------------------- + // editMessage + // --------------------------------------------------------------------------- + + describe("editMessage", () => { + it("sets editing state for a user message", () => { + const userMsg: Message = { + id: "u1", + role: "user", + created: Date.now(), + content: [{ type: "text", text: "Hello" }], + metadata: { userVisible: true, agentVisible: true }, + }; + useChatStore.getState().setMessages("session-1", [userMsg]); + + const { result } = renderHook(() => useChat("session-1")); + + act(() => { + result.current.editMessage("u1"); + }); + + expect(result.current.editingMessageId).toBe("u1"); + }); + + it("guards against entering edit mode while streaming", () => { + const userMsg: Message = { + id: "u1", + role: "user", + created: Date.now(), + content: [{ type: "text", text: "Hello" }], + metadata: { userVisible: true, agentVisible: true }, + }; + useChatStore.getState().setMessages("session-1", [userMsg]); + useChatStore.getState().setChatState("session-1", "streaming"); + + const { result } = renderHook(() => useChat("session-1")); + + act(() => { + result.current.editMessage("u1"); + }); + + expect(result.current.editingMessageId).toBeNull(); + }); + + it("refuses to edit an assistant message", () => { + const assistantMsg: Message = { + id: "a1", + role: "assistant", + created: Date.now(), + content: [{ type: "text", text: "Hello" }], + metadata: { userVisible: true, agentVisible: true }, + }; + useChatStore.getState().setMessages("session-1", [assistantMsg]); + + const { result } = renderHook(() => useChat("session-1")); + + act(() => { + result.current.editMessage("a1"); + }); + + expect(result.current.editingMessageId).toBeNull(); + }); + }); + + // --------------------------------------------------------------------------- + // cancelEdit + // --------------------------------------------------------------------------- + + // --------------------------------------------------------------------------- + // retryMessage — attachment & image-only edge cases + // --------------------------------------------------------------------------- + + describe("retryMessage edge cases", () => { + it("does not truncate history when retrying an image-only message (no text)", async () => { + // Seed a user message that has only image content, no text + const imageOnlyMsg: Message = { + id: "img-1", + role: "user", + created: Date.now(), + content: [ + { + type: "image", + source: { + type: "base64", + mediaType: "image/png", + data: "iVBORw0KGgo=", + }, + }, + ], + metadata: { userVisible: true, agentVisible: true }, + }; + const assistantReply: Message = { + id: "a1", + role: "assistant", + created: Date.now(), + content: [{ type: "text", text: "I see the image" }], + metadata: { userVisible: true, agentVisible: true }, + }; + useChatStore + .getState() + .setMessages("session-1", [imageOnlyMsg, assistantReply]); + mockAcpSendMessage.mockResolvedValue(undefined); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("img-1"); + }); + + // The image-only message should be retried (not silently dropped). + // sendMessage should have been called with the image attachment. + expect(mockAcpSendMessage).toHaveBeenCalledTimes(1); + // The re-sent message should carry the image + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "goose", + expect.any(String), + expect.objectContaining({ + images: [[expect.any(String), "image/png"]], + }), + ); + }); + + it("preserves file attachments when retrying a message", async () => { + mockAcpSendMessage.mockResolvedValue(undefined); + + const msgWithAttachments: Message = { + id: "u1", + role: "user", + created: Date.now(), + content: [{ type: "text", text: "Check this file" }], + metadata: { + userVisible: true, + agentVisible: true, + attachments: [ + { type: "file", name: "report.pdf", path: "/tmp/report.pdf" }, + { type: "directory", name: "src", path: "/tmp/src" }, + ], + }, + }; + useChatStore.getState().setMessages("session-1", [msgWithAttachments]); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("u1"); + }); + + // The re-sent user message should carry the original attachments + const messages = useChatStore.getState().messagesBySession["session-1"]; + const reSent = messages.find((m) => m.role === "user"); + expect(reSent?.metadata?.attachments).toEqual([ + expect.objectContaining({ + type: "file", + name: "report.pdf", + path: "/tmp/report.pdf", + }), + expect.objectContaining({ + type: "directory", + name: "src", + path: "/tmp/src", + }), + ]); + }); + + it("preserves image content blocks when retrying a message with text and images", async () => { + mockAcpSendMessage.mockResolvedValue(undefined); + + const msgWithImage: Message = { + id: "u1", + role: "user", + created: Date.now(), + content: [ + { type: "text", text: "Look at this" }, + { + type: "image", + source: { + type: "base64", + mediaType: "image/jpeg", + data: "base64data", + }, + }, + ], + metadata: { userVisible: true, agentVisible: true }, + }; + useChatStore.getState().setMessages("session-1", [msgWithImage]); + + const { result } = renderHook(() => useChat("session-1")); + + await act(async () => { + await result.current.retryMessage("u1"); + }); + + expect(mockAcpSendMessage).toHaveBeenCalledWith( + "session-1", + "goose", + expect.stringContaining("Look at this"), + expect.objectContaining({ + images: [["base64data", "image/jpeg"]], + }), + ); + }); + }); + + describe("cancelEdit", () => { + it("clears editing state but preserves compose draft", () => { + const userMsg: Message = { + id: "u1", + role: "user", + created: Date.now(), + content: [{ type: "text", text: "Hello" }], + metadata: { userVisible: true, agentVisible: true }, + }; + useChatStore.getState().setMessages("session-1", [userMsg]); + useChatStore.getState().setEditingMessageId("session-1", "u1"); + useChatStore.getState().setDraft("session-1", "unsent compose text"); + + const { result } = renderHook(() => useChat("session-1")); + + expect(result.current.editingMessageId).toBe("u1"); + + act(() => { + result.current.cancelEdit(); + }); + + expect(result.current.editingMessageId).toBeNull(); + // cancelEdit must NOT wipe the compose draft — only editing state + expect(useChatStore.getState().draftsBySession["session-1"]).toBe( + "unsent compose text", + ); + }); + }); }); diff --git a/ui/goose2/src/features/chat/hooks/useChat.ts b/ui/goose2/src/features/chat/hooks/useChat.ts index cd45ba8a0acc..8d4485922295 100644 --- a/ui/goose2/src/features/chat/hooks/useChat.ts +++ b/ui/goose2/src/features/chat/hooks/useChat.ts @@ -23,6 +23,7 @@ import { buildAcpImages, buildAttachmentPromptPreamble, buildMessageAttachments, + rebuildAttachmentDrafts, } from "../lib/attachments"; function getErrorMessage(error: unknown): string { @@ -326,34 +327,99 @@ export function useChat( }); }, [getStreamingPersonaId, store, sessionId]); - 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; + /** 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; + } + + // Extract re-sendable content BEFORE truncating so we never lose data + const textContent = userMessage.content.find((c) => c.type === "text"); + const text = textContent && "text" in textContent ? textContent.text : ""; + + // Reconstruct attachment drafts from stored metadata + image content + const attachmentDrafts = rebuildAttachmentDrafts(userMessage); + const hasContent = text.trim() || attachmentDrafts.length > 0; - const lastUserMessage = sessionMessages[lastUserIndex]; + // Bail if there's nothing to re-send — don't truncate + if (!hasContent) return; - // Remove all messages after (and including) the last assistant response - const messagesToKeep = sessionMessages.slice(0, lastUserIndex); - store.setMessages(sessionId, messagesToKeep); + // Truncate from the user message onward (removes user msg + all responses) + store.setMessages(sessionId, sessionMessages.slice(0, truncateFromIndex)); - // 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; + const targetPersonaId = userMessage.metadata?.targetPersonaId; + const targetPersonaName = userMessage.metadata?.targetPersonaName; await sendMessage( - textContent.text, + text, targetPersonaId ? { id: targetPersonaId, name: targetPersonaName } : undefined, + attachmentDrafts.length > 0 ? attachmentDrafts : undefined, ); - } - }, [sessionId, store, sendMessage]); + }, + [sessionId, store, sendMessage, chatState], + ); + + const retryLastMessage = useCallback(async () => { + const sessionMessages = store.messagesBySession[sessionId] ?? []; + const lastUserIndex = findLastIndex( + sessionMessages, + (m) => m.role === "user", + ); + if (lastUserIndex === -1) return; + await retryMessage(sessionMessages[lastUserIndex].id); + }, [sessionId, store, retryMessage]); + + /** 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; + + const sessionMessages = store.messagesBySession[sessionId] ?? []; + const target = sessionMessages.find((m) => m.id === messageId); + if (!target || target.role !== "user") return; + + // Enter edit mode — history stays intact until send. + // The inline edit reads text directly from the message, no draft needed. + store.setEditingMessageId(sessionId, messageId); + }, + [sessionId, store, chatState], + ); + + /** Cancel edit mode — clears only the editing state, not the compose draft. */ + const cancelEdit = useCallback(() => { + store.setEditingMessageId(sessionId, null); + }, [sessionId, store]); const clearChat = useCallback(() => { abortRef.current?.abort(); @@ -375,6 +441,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/lib/__tests__/attachments.test.ts b/ui/goose2/src/features/chat/lib/__tests__/attachments.test.ts new file mode 100644 index 000000000000..9e5583eaa40d --- /dev/null +++ b/ui/goose2/src/features/chat/lib/__tests__/attachments.test.ts @@ -0,0 +1,150 @@ +import { describe, it, expect } from "vitest"; +import { rebuildAttachmentDrafts } from "../attachments"; +import type { Message } from "@/shared/types/messages"; + +describe("rebuildAttachmentDrafts", () => { + it("skips image metadata entries when content blocks already provide base64", () => { + const msg: Message = { + id: "m1", + role: "user", + created: Date.now(), + content: [ + { + type: "image", + source: { + type: "base64", + mediaType: "image/png", + data: "iVBORw0KGgo=", + }, + }, + ], + metadata: { + userVisible: true, + agentVisible: true, + attachments: [ + { + type: "file", + name: "screenshot.png", + path: "/tmp/screenshot.png", + mimeType: "image/png", + }, + ], + }, + }; + + const drafts = rebuildAttachmentDrafts(msg); + + // Should produce exactly 1 image draft from the content block, + // NOT a second file draft from the metadata entry + expect(drafts).toHaveLength(1); + expect(drafts[0].kind).toBe("image"); + if (drafts[0].kind === "image") { + expect(drafts[0].base64).toBe("iVBORw0KGgo="); + } + }); + + it("preserves non-image file metadata alongside image content blocks", () => { + const msg: Message = { + id: "m2", + role: "user", + created: Date.now(), + content: [ + { + type: "image", + source: { + type: "base64", + mediaType: "image/jpeg", + data: "/9j/4AAQ=", + }, + }, + { type: "text", text: "See attached" }, + ], + metadata: { + userVisible: true, + agentVisible: true, + attachments: [ + { + type: "file", + name: "photo.jpg", + path: "/tmp/photo.jpg", + mimeType: "image/jpeg", + }, + { + type: "file", + name: "data.csv", + path: "/tmp/data.csv", + mimeType: "text/csv", + }, + { + type: "directory", + name: "src", + path: "/project/src", + }, + ], + }, + }; + + const drafts = rebuildAttachmentDrafts(msg); + + // 1 image from content block + 1 csv file + 1 directory = 3 + // photo.jpg metadata entry skipped (image already from content block) + expect(drafts).toHaveLength(3); + expect(drafts.map((d) => d.kind)).toEqual(["image", "file", "directory"]); + expect(drafts[1].name).toBe("data.csv"); + expect(drafts[2].name).toBe("src"); + }); + + it("includes image metadata entries when no content blocks exist", () => { + const msg: Message = { + id: "m3", + role: "user", + created: Date.now(), + content: [{ type: "text", text: "file attached" }], + metadata: { + userVisible: true, + agentVisible: true, + attachments: [ + { + type: "file", + name: "photo.jpg", + path: "/tmp/photo.jpg", + mimeType: "image/jpeg", + }, + ], + }, + }; + + const drafts = rebuildAttachmentDrafts(msg); + + // No image content blocks → metadata image entry should be included as file + expect(drafts).toHaveLength(1); + expect(drafts[0].kind).toBe("file"); + expect(drafts[0].name).toBe("photo.jpg"); + }); + + it("preserves pathless browser-uploaded file attachments", () => { + const msg: Message = { + id: "m4", + role: "user", + created: Date.now(), + content: [{ type: "text", text: "see attached" }], + metadata: { + userVisible: true, + agentVisible: true, + attachments: [ + { + type: "file", + name: "report.pdf", + mimeType: "application/pdf", + }, + ], + }, + }; + + const drafts = rebuildAttachmentDrafts(msg); + + expect(drafts).toHaveLength(1); + expect(drafts[0].kind).toBe("file"); + expect(drafts[0].name).toBe("report.pdf"); + }); +}); diff --git a/ui/goose2/src/features/chat/lib/attachments.ts b/ui/goose2/src/features/chat/lib/attachments.ts index a69a38361c29..0358e94800e9 100644 --- a/ui/goose2/src/features/chat/lib/attachments.ts +++ b/ui/goose2/src/features/chat/lib/attachments.ts @@ -1,5 +1,6 @@ import type { ChatAttachmentDraft, + Message, MessageAttachment, } from "@/shared/types/messages"; @@ -66,3 +67,63 @@ export function buildAcpImages( return images.length > 0 ? images : undefined; } + +/** + * Reconstruct ChatAttachmentDraft[] from a stored user message so retry/edit + * can forward attachments that were present on the original send. Image + * content blocks carry base64 data; file/directory attachments are rebuilt + * from metadata.attachments. + */ +export function rebuildAttachmentDrafts( + message: Message, +): ChatAttachmentDraft[] { + const drafts: ChatAttachmentDraft[] = []; + + // Rebuild image drafts from ImageContent blocks embedded in the message + for (const block of message.content) { + if (block.type === "image" && block.source.type === "base64") { + drafts.push({ + id: crypto.randomUUID(), + kind: "image", + name: "image", + mimeType: block.source.mediaType, + base64: block.source.data, + previewUrl: `data:${block.source.mediaType};base64,${block.source.data}`, + }); + } + } + + // Rebuild file/directory drafts from metadata.attachments, skipping image + // entries when we already reconstructed image drafts from content blocks. + // An uploaded image is stored as both a base64 content block and a metadata + // file entry — including both would duplicate the image in the re-send. + // + // TODO: This heuristic assumes all uploaded images produce both a content + // block and a metadata entry. If a future code path stores an image only in + // metadata (without a corresponding content block), this broad skip will + // silently drop it. If that changes, switch to per-image matching by + // name or content hash instead of the blanket `hasImageDrafts` flag. + const hasImageDrafts = drafts.some((d) => d.kind === "image"); + for (const att of message.metadata?.attachments ?? []) { + if (att.type === "directory" && att.path) { + drafts.push({ + id: crypto.randomUUID(), + kind: "directory", + name: att.name, + path: att.path, + }); + } else if (att.type === "file") { + // Skip image file entries when content blocks already provide the base64 + if (hasImageDrafts && att.mimeType?.startsWith("image/")) continue; + drafts.push({ + id: crypto.randomUUID(), + kind: "file", + name: att.name, + ...(att.path ? { path: att.path } : {}), + mimeType: att.mimeType, + }); + } + } + + return drafts; +} diff --git a/ui/goose2/src/features/chat/stores/__tests__/chatStore.test.ts b/ui/goose2/src/features/chat/stores/__tests__/chatStore.test.ts index 72873531b155..c64c1d2cabfc 100644 --- a/ui/goose2/src/features/chat/stores/__tests__/chatStore.test.ts +++ b/ui/goose2/src/features/chat/stores/__tests__/chatStore.test.ts @@ -159,6 +159,29 @@ describe("chatStore", () => { expect(store.activeSessionId).toBeNull(); }); + it("sets and clears editingMessageId per session", () => { + const store = useChatStore.getState(); + + store.setEditingMessageId("s1", "msg-1"); + expect(useChatStore.getState().editingMessageIdBySession.s1).toBe("msg-1"); + expect( + useChatStore.getState().editingMessageIdBySession.s2, + ).toBeUndefined(); + + store.setEditingMessageId("s1", null); + expect(useChatStore.getState().editingMessageIdBySession.s1).toBeNull(); + }); + + it("cleans up editingMessageId on session cleanup", () => { + const store = useChatStore.getState(); + + store.setEditingMessageId("s1", "msg-1"); + store.cleanupSession("s1"); + expect( + useChatStore.getState().editingMessageIdBySession.s1, + ).toBeUndefined(); + }); + it("stores and clears scroll targets per session", () => { const store = useChatStore.getState(); diff --git a/ui/goose2/src/features/chat/stores/chatStore.ts b/ui/goose2/src/features/chat/stores/chatStore.ts index f7781b8ef9bb..5d404cefefe4 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 fdf2bea7d051..3d12f44946ca 100644 --- a/ui/goose2/src/features/chat/ui/ChatInput.tsx +++ b/ui/goose2/src/features/chat/ui/ChatInput.tsx @@ -1,4 +1,11 @@ -import { useState, useRef, useCallback, useEffect, useMemo } from "react"; +import { + useState, + useRef, + useCallback, + useEffect, + useLayoutEffect, + useMemo, +} from "react"; import { open } from "@tauri-apps/plugin-dialog"; import { X } from "lucide-react"; import { useTranslation } from "react-i18next"; @@ -107,6 +114,28 @@ export function ChatInput({ const [isCompact, setIsCompact] = useState(false); const textareaRef = useRef(null); const containerRef = useRef(null); + + // Re-sync local text when the parent changes initialValue (e.g. cancel edit → re-edit). + // Guard avoids redundant state updates on every keystroke when initialValue + // round-trips through the store during active editing. + useEffect(() => { + const ta = textareaRef.current; + if (ta && ta.value !== initialValue) { + setTextRaw(initialValue); + } + }, [initialValue]); + + // Recalculate textarea height after React commits new text to the DOM. + // useLayoutEffect fires synchronously after DOM mutation but before paint, + // so scrollHeight already reflects the current value. + // biome-ignore lint/correctness/useExhaustiveDependencies: text triggers height recalc via scrollHeight + useLayoutEffect(() => { + const ta = textareaRef.current; + if (ta) { + ta.style.height = "auto"; + ta.style.height = `${Math.min(ta.scrollHeight, 200)}px`; + } + }, [text]); const { attachments, addBrowserFiles, @@ -228,9 +257,7 @@ export function ChatInput({ setText(value); const cursorPosition = event.target.selectionStart ?? value.length; detectMention(value, cursorPosition); - const textarea = event.target; - textarea.style.height = "auto"; - textarea.style.height = `${Math.min(textarea.scrollHeight, 200)}px`; + // Height recalc handled by useLayoutEffect([text]) }; const handlePaste = useCallback( diff --git a/ui/goose2/src/features/chat/ui/ChatView.tsx b/ui/goose2/src/features/chat/ui/ChatView.tsx index 0478343dae04..84028e1a0efe 100644 --- a/ui/goose2/src/features/chat/ui/ChatView.tsx +++ b/ui/goose2/src/features/chat/ui/ChatView.tsx @@ -1,4 +1,5 @@ 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"; @@ -8,6 +9,7 @@ import { LoadingGoose } from "./LoadingGoose"; import { ChatLoadingSkeleton } from "./ChatLoadingSkeleton"; import { useChat } from "../hooks/useChat"; import { useMessageQueue } from "../hooks/useMessageQueue"; +import { rebuildAttachmentDrafts } from "../lib/attachments"; import { useChatStore } from "../stores/chatStore"; import { useAgentStore } from "@/features/agents/stores/agentStore"; import { useProviderSelection } from "@/features/agents/hooks/useProviderSelection"; @@ -337,6 +339,10 @@ export function ChatView({ tokenState, sendMessage, stopStreaming, + retryMessage, + editMessage, + cancelEdit, + editingMessageId, streamingMessageId, } = useChat( activeSessionId, @@ -351,8 +357,12 @@ export function ChatView({ (s.messagesBySession[activeSessionId]?.length ?? 0) === 0, ); - const deferredSend = useRef<{ + // 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); @@ -379,7 +389,7 @@ export function ChatView({ } handlePersonaChange(personaId); // Defer the send until after persona state updates - deferredSend.current = { text, attachments }; + pendingSend.current = { reason: "persona", text, attachments }; return; } // Queue if agent is busy and no message already queued @@ -402,10 +412,55 @@ export function ChatView({ ], ); + /** Save an inline edit: truncate history from the edited message onward, then send the new text. */ + const handleSaveEdit = useCallback( + (messageId: string, text: string) => { + if (chatState !== "idle") { + 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); + + // 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"); + }); + + sendMessage( + text, + targetPersonaId + ? { id: targetPersonaId, name: targetPersonaName } + : undefined, + originalAttachments.length > 0 ? originalAttachments : undefined, + ); + }, + [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 (deferredSend.current && selectedPersona) { - const { text, attachments } = deferredSend.current; - deferredSend.current = null; + if ( + pendingSend.current && + pendingSend.current.reason === "persona" && + selectedPersona + ) { + const { text, attachments } = pendingSend.current; + pendingSend.current = null; sendMessage(text, undefined, attachments); } }, [sendMessage, selectedPersona]); @@ -465,6 +520,11 @@ export function ChatView({ scrollTargetMessageId={scrollTarget?.messageId ?? null} scrollTargetQuery={scrollTarget?.query ?? null} onScrollTargetHandled={handleScrollTargetHandled} + onRetryMessage={retryMessage} + onEditMessage={editMessage} + editingMessageId={editingMessageId} + onSaveEdit={handleSaveEdit} + onCancelEdit={cancelEdit} /> )} diff --git a/ui/goose2/src/features/chat/ui/MessageBubble.tsx b/ui/goose2/src/features/chat/ui/MessageBubble.tsx index 6363f98c0a6f..6580156724de 100644 --- a/ui/goose2/src/features/chat/ui/MessageBubble.tsx +++ b/ui/goose2/src/features/chat/ui/MessageBubble.tsx @@ -1,11 +1,17 @@ -import { useState, memo } from "react"; +import { + useState, + useRef, + useEffect, + useLayoutEffect, + useCallback, + memo, +} from "react"; import { useTranslation } from "react-i18next"; import { Copy, Check, RotateCcw, Pencil, - User, FileText, FolderClosed, } from "lucide-react"; @@ -30,6 +36,12 @@ import { ReasoningTrigger, ReasoningContent, } from "@/shared/ui/ai-elements/reasoning"; +import { Button } from "@/shared/ui/button"; +import { + HoverCard, + HoverCardTrigger, + HoverCardContent, +} from "@/shared/ui/hover-card"; import { ToolChainCards, type ToolChainItem } from "./ToolChainCards"; import { ClickableImage } from "./ClickableImage"; import { useArtifactLinkHandler } from "@/features/chat/hooks/useArtifactLinkHandler"; @@ -89,6 +101,9 @@ interface MessageBubbleProps { onCopy?: () => void; onRetryMessage?: (messageId: string) => void; onEditMessage?: (messageId: string) => void; + isEditing?: boolean; + onSaveEdit?: (messageId: string, text: string) => void; + onCancelEdit?: () => void; } interface ContentSection { @@ -306,6 +321,9 @@ export const MessageBubble = memo(function MessageBubble({ isStreaming, onRetryMessage, onEditMessage, + isEditing = false, + onSaveEdit, + onCancelEdit, }: MessageBubbleProps) { const { t } = useTranslation(["chat", "common"]); const { formatDate } = useLocaleFormatting(); @@ -323,6 +341,55 @@ export const MessageBubble = memo(function MessageBubble({ .map((c) => c.text) .join("\n"); + // Inline edit state — only allocated for user messages + const canEdit = role === "user"; + const [editText, setEditText] = useState(canEdit ? textContent : ""); + const editTextareaRef = useRef(null); + + // Reset edit text when entering edit mode + useEffect(() => { + if (canEdit && isEditing) { + setEditText(textContent); + } + }, [canEdit, isEditing, textContent]); + + // Auto-resize on every text change + // biome-ignore lint/correctness/useExhaustiveDependencies: editText triggers height recalc via scrollHeight + useLayoutEffect(() => { + const ta = editTextareaRef.current; + if (canEdit && isEditing && ta) { + ta.style.height = "auto"; + ta.style.height = `${Math.min(ta.scrollHeight, 300)}px`; + } + }, [canEdit, isEditing, editText]); + + // Focus only when entering edit mode, not on every keystroke + useEffect(() => { + if (canEdit && isEditing) { + editTextareaRef.current?.focus(); + } + }, [canEdit, isEditing]); + + const handleEditSave = useCallback(() => { + const trimmed = editText.trim(); + if (trimmed && onSaveEdit) { + onSaveEdit(message.id, trimmed); + } + }, [editText, message.id, onSaveEdit]); + + const handleEditKeyDown = useCallback( + (e: React.KeyboardEvent) => { + if (e.key === "Escape") { + e.preventDefault(); + onCancelEdit?.(); + } else if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleEditSave(); + } + }, + [handleEditSave, onCancelEdit], + ); + if (role === "system") { return (
@@ -360,22 +427,25 @@ export const MessageBubble = memo(function MessageBubble({ return (
- {isUser ? ( -
- -
- ) : null} -
{showAssistantIdentity ? ( @@ -403,76 +473,124 @@ export const MessageBubble = memo(function MessageBubble({
) : null} - {/* biome-ignore lint/a11y/useKeyWithClickEvents: delegated link handler */} - {/* biome-ignore lint/a11y/noStaticElementInteractions: delegated link handler */} -
- {messageAttachments.length > 0 && ( -
- {messageAttachments.map((attachment) => ( - - ))} + {isEditing ? ( +
+