diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 35f99cff67..44535d8041 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -93,33 +93,47 @@ describe("useConversationEditingState", () => { setMobileViewport(false); }); - const renderEditing = () => { + const renderEditing = (...args: [] | [string | undefined]) => { const onSend = vi.fn().mockResolvedValue(undefined); const onDeleteQueuedMessage = vi.fn().mockResolvedValue(undefined); const chatInputRef = createRef(); + const inputValueRef = { current: "" }; + // createRef returns { current: null }, but we need it initialized + // to "" so the hook sees a string. + (inputValueRef as { current: string }).current = ""; + + const resolvedChatID = args.length === 0 ? chatID : args[0]; const hook = renderHook(() => useConversationEditingState({ - chatID, + chatID: resolvedChatID, onSend, onDeleteQueuedMessage, chatInputRef, + inputValueRef, }), ); - return { ...hook, onSend }; + return { ...hook, onSend, inputValueRef }; }; it("persists and removes drafts via handleContentChange", () => { const { result, unmount } = renderEditing(); act(() => { - result.current.handleContentChange("work in progress"); + result.current.handleContentChange( + "work in progress", + "work in progress", + false, + ); }); expect(localStorage.getItem(expectedKey)).toBe("work in progress"); act(() => { - result.current.handleContentChange(""); + // Even though the serialized state is non-empty (Lexical always + // produces a JSON object), the draft is removed when the plain + // text content is empty. + result.current.handleContentChange("", '{"root":{"children":[]}}', false); }); expect(localStorage.getItem(expectedKey)).toBeNull(); @@ -128,49 +142,69 @@ describe("useConversationEditingState", () => { it("loads edit text into the composer and restores the prior draft on cancel without refocusing", () => { const { result, unmount } = renderEditing(); - const mockInput = createMockChatInputHandle("work in progress"); - result.current.chatInputRef.current = mockInput.handle; + + // Simulate the user typing a draft via handleContentChange. + act(() => { + result.current.handleContentChange( + "work in progress", + "work in progress", + false, + ); + }); + + const remountKeyBefore = result.current.remountKey; act(() => { result.current.handleEditUserMessage(7, "edited message"); }); expect(result.current.editingMessageId).toBe(7); - expect(mockInput.setValue).toHaveBeenCalledWith("edited message"); - expect(mockInput.focus).toHaveBeenCalledTimes(1); + expect(result.current.editorInitialValue).toBe("edited message"); + expect(result.current.remountKey).toBe(remountKeyBefore + 1); + + const remountKeyAfterEdit = result.current.remountKey; act(() => { result.current.handleCancelHistoryEdit(); }); expect(result.current.editingMessageId).toBeNull(); - expect(mockInput.setValue).toHaveBeenLastCalledWith("work in progress"); - expect(mockInput.currentValue.value).toBe("work in progress"); - expect(mockInput.focus).toHaveBeenCalledTimes(1); + expect(result.current.editorInitialValue).toBe("work in progress"); + expect(result.current.remountKey).toBe(remountKeyAfterEdit + 1); unmount(); }); it("loads queue edit text into the composer and restores the prior draft on cancel without refocusing", () => { const { result, unmount } = renderEditing(); - const mockInput = createMockChatInputHandle("work in progress"); - result.current.chatInputRef.current = mockInput.handle; + + // Simulate the user typing a draft via handleContentChange. + act(() => { + result.current.handleContentChange( + "work in progress", + "work in progress", + false, + ); + }); + + const remountKeyBefore = result.current.remountKey; act(() => { result.current.handleStartQueueEdit(9, "queued message", []); }); expect(result.current.editingQueuedMessageID).toBe(9); - expect(mockInput.setValue).toHaveBeenCalledWith("queued message"); - expect(mockInput.focus).toHaveBeenCalledTimes(1); + expect(result.current.editorInitialValue).toBe("queued message"); + expect(result.current.remountKey).toBe(remountKeyBefore + 1); + + const remountKeyAfterEdit = result.current.remountKey; act(() => { result.current.handleCancelQueueEdit(); }); expect(result.current.editingQueuedMessageID).toBeNull(); - expect(mockInput.setValue).toHaveBeenLastCalledWith("work in progress"); - expect(mockInput.currentValue.value).toBe("work in progress"); - expect(mockInput.focus).toHaveBeenCalledTimes(1); + expect(result.current.editorInitialValue).toBe("work in progress"); + expect(result.current.remountKey).toBe(remountKeyAfterEdit + 1); unmount(); }); @@ -180,6 +214,10 @@ describe("useConversationEditingState", () => { const mockInput = createMockChatInputHandle("draft before edit"); result.current.chatInputRef.current = mockInput.handle; + // Edit/cancel now drive the editor via editorInitialValue + + // remountKey, so focus is never called on the mock during + // edit and cancel flows. handleSendFromInput is the only + // path that calls focus and it skips on mobile viewports. act(() => { result.current.handleEditUserMessage(7, "edited message"); }); @@ -205,8 +243,6 @@ describe("useConversationEditingState", () => { it("falls back to the persisted draft when history edit starts before hydration", () => { localStorage.setItem(expectedKey, "persisted draft"); const { result, unmount } = renderEditing(); - const mockInput = createMockChatInputHandle(""); - result.current.chatInputRef.current = mockInput.handle; act(() => { result.current.handleEditUserMessage(7, "edited message"); @@ -216,16 +252,15 @@ describe("useConversationEditingState", () => { result.current.handleCancelHistoryEdit(); }); - expect(mockInput.setValue).toHaveBeenLastCalledWith("persisted draft"); - expect(mockInput.currentValue.value).toBe("persisted draft"); + // The hook reads the persisted draft from localStorage when + // inputValueRef hasn't been updated by handleContentChange yet. + expect(result.current.editorInitialValue).toBe("persisted draft"); unmount(); }); it("falls back to the persisted draft when queue edit starts before hydration", () => { localStorage.setItem(expectedKey, "persisted draft"); const { result, unmount } = renderEditing(); - const mockInput = createMockChatInputHandle(""); - result.current.chatInputRef.current = mockInput.handle; act(() => { result.current.handleStartQueueEdit(9, "queued message", []); @@ -235,16 +270,19 @@ describe("useConversationEditingState", () => { result.current.handleCancelQueueEdit(); }); - expect(mockInput.setValue).toHaveBeenLastCalledWith("persisted draft"); - expect(mockInput.currentValue.value).toBe("persisted draft"); + expect(result.current.editorInitialValue).toBe("persisted draft"); unmount(); }); it("prefers the live editor value over stale persisted draft state", () => { localStorage.setItem(expectedKey, "stale persisted draft"); const { result, unmount } = renderEditing(); - const mockInput = createMockChatInputHandle("live draft"); - result.current.chatInputRef.current = mockInput.handle; + + // Simulate the editor emitting a content change, which updates + // inputValueRef to the live value. + act(() => { + result.current.handleContentChange("live draft", "live draft", false); + }); act(() => { result.current.handleEditUserMessage(7, "edited message"); @@ -254,31 +292,38 @@ describe("useConversationEditingState", () => { result.current.handleCancelHistoryEdit(); }); - expect(mockInput.setValue).toHaveBeenLastCalledWith("live draft"); - expect(mockInput.currentValue.value).toBe("live draft"); + expect(result.current.editorInitialValue).toBe("live draft"); unmount(); }); - it("can load the same edit text again after send without relying on a remount", async () => { + it("can load the same edit text again after send", async () => { const { result, onSend, unmount } = renderEditing(); const mockInput = createMockChatInputHandle(); result.current.chatInputRef.current = mockInput.handle; + const remountKeyBefore = result.current.remountKey; + act(() => { result.current.handleEditUserMessage(7, "hello"); }); + expect(result.current.remountKey).toBe(remountKeyBefore + 1); + await act(async () => { await result.current.handleSendFromInput("hello"); }); + const remountKeyAfterSend = result.current.remountKey; + act(() => { result.current.handleEditUserMessage(7, "hello"); }); + // remountKey increments each time an edit is loaded, even for + // the same text, so the editor is forced to reinitialize. + expect(result.current.remountKey).toBe(remountKeyAfterSend + 1); + expect(result.current.editorInitialValue).toBe("hello"); expect(onSend).toHaveBeenCalledWith("hello", undefined, 7); - expect(mockInput.setValue).toHaveBeenNthCalledWith(1, "hello"); - expect(mockInput.setValue).toHaveBeenNthCalledWith(2, "hello"); unmount(); }); @@ -298,4 +343,311 @@ describe("useConversationEditingState", () => { expect(localStorage.getItem(expectedKey)).toBeNull(); unmount(); }); + + it("does not write a draft key when chatID is undefined", () => { + const { result, unmount } = renderEditing(undefined); + + act(() => { + result.current.handleContentChange("should not persist", "{}", false); + }); + + // The ref is still updated even without persistence. + expect(result.current.inputValueRef.current).toBe("should not persist"); + // No draft for "undefined" chatID should appear. + expect( + localStorage.getItem(`${draftInputStorageKeyPrefix}undefined`), + ).toBeNull(); + unmount(); + }); + + it("calls focus on the input ref after a successful send", async () => { + const { result, onSend, unmount } = renderEditing(); + + // Attach a mock ChatMessageInputRef to the chatInputRef + const mockFocus = vi.fn(); + const mockClear = vi.fn(); + const mockInputRef = { + focus: mockFocus, + clear: mockClear, + setValue: vi.fn(), + insertText: vi.fn(), + getValue: vi.fn().mockReturnValue(""), + addFileReference: vi.fn(), + getContentParts: vi.fn().mockReturnValue([]), + }; // The hook exposes chatInputRef – assign the mock to it. + result.current.chatInputRef.current = mockInputRef; + + await act(async () => { + result.current.handleSendFromInput("hello"); + await vi.waitFor(() => { + expect(onSend).toHaveBeenCalledWith("hello", undefined, undefined); + }); + }); + + expect(mockClear).toHaveBeenCalled(); + expect(mockFocus).toHaveBeenCalled(); + unmount(); + }); + + it("initializes with the correct draft for each chatID", () => { + const chatA = "chat-aaa"; + const chatB = "chat-bbb"; + localStorage.setItem(`${draftInputStorageKeyPrefix}${chatA}`, "draft A"); + localStorage.setItem(`${draftInputStorageKeyPrefix}${chatB}`, "draft B"); + + // Each chatID should initialize with its own draft — this is + // what the key={agentId} wrapper guarantees at the component + // level (a new chatID means a full remount). + const hookA = renderEditing(chatA); + expect(hookA.result.current.editorInitialValue).toBe("draft A"); + hookA.unmount(); + + const hookB = renderEditing(chatB); + expect(hookB.result.current.editorInitialValue).toBe("draft B"); + hookB.unmount(); + }); + + it("clears the draft from localStorage on successful send", async () => { + localStorage.setItem(expectedKey, "draft to clear"); + + const { result, unmount } = renderEditing(); + + expect(localStorage.getItem(expectedKey)).toBe("draft to clear"); + + await act(async () => { + result.current.handleSendFromInput("hello"); + await vi.waitFor(() => { + expect(localStorage.getItem(expectedKey)).toBeNull(); + }); + }); + unmount(); + }); + + it("persists serialized editor state when provided", () => { + const { result, unmount } = renderEditing(); + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [ + { text: "review this" }, + { + type: "file-reference", + version: 1, + fileName: "main.go", + startLine: 1, + endLine: 10, + content: "code", + }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + + act(() => { + result.current.handleContentChange("review this", editorState, true); + }); + + // The serialized editor state should be stored, not the plain text. + expect(localStorage.getItem(expectedKey)).toBe(editorState); + expect(result.current.inputValueRef.current).toBe("review this"); + unmount(); + }); + + it("restores editorInitialState from a Lexical JSON draft", () => { + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [{ text: "hello" }], + type: "paragraph", + }, + ], + type: "root", + }, + }); + localStorage.setItem(expectedKey, editorState); + + const { result, unmount } = renderEditing(); + + expect(result.current.initialEditorState).toBe(editorState); + expect(result.current.editorInitialValue).toBe("hello"); + unmount(); + }); + + it("falls back to plain text for legacy drafts", () => { + localStorage.setItem(expectedKey, "legacy plain text"); + + const { result, unmount } = renderEditing(); + + expect(result.current.initialEditorState).toBeUndefined(); + expect(result.current.editorInitialValue).toBe("legacy plain text"); + unmount(); + }); + + it("persists file-reference-only drafts (no text content)", () => { + const { result, unmount } = renderEditing(); + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [ + { + type: "file-reference", + version: 1, + fileName: "main.go", + startLine: 1, + endLine: 10, + content: "code", + }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + + act(() => { + // Empty text but hasFileReferences=true should still persist. + result.current.handleContentChange("", editorState, true); + }); + + expect(localStorage.getItem(expectedKey)).toBe(editorState); + unmount(); + }); + + it("removes draft for whitespace-only content without file references", () => { + localStorage.setItem(expectedKey, "old draft"); + const { result, unmount } = renderEditing(); + + act(() => { + result.current.handleContentChange(" ", '{"root":{}}', false); + }); + + expect(localStorage.getItem(expectedKey)).toBeNull(); + unmount(); + }); + + it("preserves serialized editor state across history edit then cancel", () => { + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [ + { text: "my draft", type: "text" }, + { + type: "file-reference", + version: 1, + fileName: "main.go", + startLine: 1, + endLine: 10, + content: "code", + }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + localStorage.setItem(expectedKey, editorState); + + const { result, unmount } = renderEditing(); + + expect(result.current.initialEditorState).toBe(editorState); + expect(result.current.editorInitialValue).toBe("my draft"); + + // Simulate typing so localStorage reflects the current draft. + act(() => { + result.current.handleContentChange("my draft", editorState, true); + }); + + // Start editing a history message. + act(() => { + result.current.handleEditUserMessage(42, "old message text"); + }); + + expect(result.current.editingMessageId).toBe(42); + expect(result.current.initialEditorState).toBeUndefined(); + expect(result.current.editorInitialValue).toBe("old message text"); + + // Cancel — should restore both plain text and serialized state. + act(() => { + result.current.handleCancelHistoryEdit(); + }); + + expect(result.current.editingMessageId).toBeNull(); + expect(result.current.initialEditorState).toBe(editorState); + expect(result.current.editorInitialValue).toBe("my draft"); + unmount(); + }); + + it("preserves serialized editor state across queue edit then cancel", () => { + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [{ text: "queued draft", type: "text" }], + type: "paragraph", + }, + ], + type: "root", + }, + }); + localStorage.setItem(expectedKey, editorState); + + const { result, unmount } = renderEditing(); + + act(() => { + result.current.handleContentChange("queued draft", editorState, false); + }); + + act(() => { + result.current.handleStartQueueEdit(99, "queued msg", []); + }); + + expect(result.current.editingQueuedMessageID).toBe(99); + expect(result.current.initialEditorState).toBeUndefined(); + + act(() => { + result.current.handleCancelQueueEdit(); + }); + + expect(result.current.editingQueuedMessageID).toBeNull(); + expect(result.current.initialEditorState).toBe(editorState); + expect(result.current.editorInitialValue).toBe("queued draft"); + unmount(); + }); + + it("returns undefined initialEditorState after edit then cancel with plain-text draft", () => { + localStorage.setItem(expectedKey, "plain text draft"); + + const { result, unmount } = renderEditing(); + + expect(result.current.initialEditorState).toBeUndefined(); + + act(() => { + result.current.handleContentChange( + "plain text draft", + "plain text draft", + false, + ); + }); + + act(() => { + result.current.handleEditUserMessage(1, "editing"); + }); + + act(() => { + result.current.handleCancelHistoryEdit(); + }); + + expect(result.current.initialEditorState).toBeUndefined(); + expect(result.current.editorInitialValue).toBe("plain text draft"); + unmount(); + }); }); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 83dfad0682..6485388c5d 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -1,4 +1,5 @@ -import { type FC, useEffect, useRef, useState } from "react"; +import { type FC, useEffect, useLayoutEffect, useRef, useState } from "react"; + import { useInfiniteQuery, useMutation, @@ -61,6 +62,7 @@ import { saveMCPSelection, } from "./components/MCPServerPicker"; import { useGitWatcher } from "./hooks/useGitWatcher"; +import { type ParsedDraft, parseStoredDraft } from "./utils/draftStorage"; import { getModelOptionsFromConfigs, getModelSelectorPlaceholder, @@ -81,15 +83,20 @@ const lastModelConfigIDStorageKey = "agents.last-model-config-id"; /** @internal Exported for testing. */ export const draftInputStorageKeyPrefix = "agents.draft-input."; -/** @internal Exported for testing. */ +/** + * Read the persisted plain-text draft for a given chat ID. + * Returns the text portion of the draft (stripping Lexical JSON + * wrapper if present) for backward compatibility. + */ export function getPersistedDraftInputValue( chatID: string | undefined, ): string { if (typeof window === "undefined" || !chatID) { return ""; } - - return localStorage.getItem(`${draftInputStorageKeyPrefix}${chatID}`) ?? ""; + return parseStoredDraft( + localStorage.getItem(`${draftInputStorageKeyPrefix}${chatID}`), + ).text; } /** @internal Exported for testing. */ @@ -102,41 +109,45 @@ export function useConversationEditingState(deps: { ) => Promise; onDeleteQueuedMessage: (id: number) => Promise; chatInputRef: React.RefObject; + inputValueRef: React.RefObject; }) { - const { chatID, onSend, onDeleteQueuedMessage, chatInputRef } = deps; + const { chatID, onSend, onDeleteQueuedMessage, chatInputRef, inputValueRef } = + deps; const draftStorageKey = chatID ? `${draftInputStorageKeyPrefix}${chatID}` : null; - const getDraftBeforeEdit = () => { - const currentInputValue = chatInputRef.current?.getValue() ?? ""; - if (currentInputValue) { - return currentInputValue; + const [{ editorInitialValue, initialEditorState }, setDraftState] = useState( + () => { + if (typeof window === "undefined" || !draftStorageKey) { + return { editorInitialValue: "", initialEditorState: undefined }; + } + const draft = parseStoredDraft(localStorage.getItem(draftStorageKey)); + return { + editorInitialValue: draft.text, + initialEditorState: draft.editorState, + }; + }, + ); + + // Monotonic counter to force LexicalComposer remount. + const [remountKey, setRemountKey] = useState(0); + + // Sync the ref with the initial draft value so callers that + // read inputValueRef.current see the persisted draft. Uses a + // layout effect so the value is available before paint. + const initialSyncDone = useRef(false); + useLayoutEffect(() => { + if (!initialSyncDone.current && editorInitialValue) { + initialSyncDone.current = true; + (inputValueRef as React.MutableRefObject).current = + editorInitialValue; } - // The editor seeds its initial value after paint, so the live editor can - // still be empty while the persisted draft already exists. - if (typeof window === "undefined" || !draftStorageKey) { - return ""; - } - return localStorage.getItem(draftStorageKey) ?? ""; - }; - const replaceInputValue = (content: string) => { - chatInputRef.current?.setValue(content); - }; - const focusInputIfDesktop = () => { - if (!isMobileViewport()) { - chatInputRef.current?.focus(); - } - }; - const replaceInputValueAndFocus = (content: string) => { - replaceInputValue(content); - focusInputIfDesktop(); - }; + }, [editorInitialValue, inputValueRef]); // -- History editing state -- const [editingMessageId, setEditingMessageId] = useState(null); - const [draftBeforeHistoryEdit, setDraftBeforeHistoryEdit] = useState< - string | null - >(null); + const [draftBeforeHistoryEdit, setDraftBeforeHistoryEdit] = + useState(null); const [editingFileBlocks, setEditingFileBlocks] = useState< readonly ChatMessagePart[] >([]); @@ -146,49 +157,85 @@ export function useConversationEditingState(deps: { text: string, fileBlocks?: readonly ChatMessagePart[], ) => { - setDraftBeforeHistoryEdit((prev) => - editingMessageId !== null ? prev : getDraftBeforeEdit(), - ); + if (editingMessageId === null) { + // Read the current serialized editor state from localStorage + // (kept up-to-date by handleContentChange) rather than from + // the stale initialEditorState React state. + const currentEditorState = draftStorageKey + ? parseStoredDraft(localStorage.getItem(draftStorageKey)).editorState + : undefined; + setDraftBeforeHistoryEdit({ + text: inputValueRef.current, + editorState: currentEditorState, + }); + } setEditingMessageId(messageId); - replaceInputValueAndFocus(text); + setDraftState({ + editorInitialValue: text, + initialEditorState: undefined, + }); + setRemountKey((k) => k + 1); + inputValueRef.current = text; setEditingFileBlocks(fileBlocks ?? []); }; const handleCancelHistoryEdit = () => { - const draft = draftBeforeHistoryEdit ?? ""; + const savedText = draftBeforeHistoryEdit?.text ?? ""; + const savedState = draftBeforeHistoryEdit?.editorState; + setDraftState({ + editorInitialValue: savedText, + initialEditorState: savedState, + }); + setRemountKey((k) => k + 1); + inputValueRef.current = savedText; setEditingMessageId(null); setDraftBeforeHistoryEdit(null); setEditingFileBlocks([]); - replaceInputValue(draft); }; // -- Queue editing state -- const [editingQueuedMessageID, setEditingQueuedMessageID] = useState< number | null >(null); - const [draftBeforeQueueEdit, setDraftBeforeQueueEdit] = useState< - string | null - >(null); + const [draftBeforeQueueEdit, setDraftBeforeQueueEdit] = + useState(null); const handleStartQueueEdit = ( id: number, text: string, fileBlocks: readonly ChatMessagePart[], ) => { - setDraftBeforeQueueEdit((prev) => - editingQueuedMessageID === null ? getDraftBeforeEdit() : prev, - ); + if (editingQueuedMessageID === null) { + const currentEditorState = draftStorageKey + ? parseStoredDraft(localStorage.getItem(draftStorageKey)).editorState + : undefined; + setDraftBeforeQueueEdit({ + text: inputValueRef.current, + editorState: currentEditorState, + }); + } setEditingQueuedMessageID(id); - replaceInputValueAndFocus(text); + setDraftState({ + editorInitialValue: text, + initialEditorState: undefined, + }); + setRemountKey((k) => k + 1); + inputValueRef.current = text; setEditingFileBlocks(fileBlocks); }; const handleCancelQueueEdit = () => { - const draft = draftBeforeQueueEdit ?? ""; + const savedText = draftBeforeQueueEdit?.text ?? ""; + const savedState = draftBeforeQueueEdit?.editorState; + setDraftState({ + editorInitialValue: savedText, + initialEditorState: savedState, + }); + setRemountKey((k) => k + 1); + inputValueRef.current = savedText; setEditingQueuedMessageID(null); setDraftBeforeQueueEdit(null); setEditingFileBlocks([]); - replaceInputValue(draft); }; // Wraps the parent onSend to clear local input/editing state @@ -201,7 +248,10 @@ export function useConversationEditingState(deps: { await onSend(message, fileIds, editedMessageID); // Clear input and editing state on success. chatInputRef.current?.clear(); - focusInputIfDesktop(); + if (!isMobileViewport()) { + chatInputRef.current?.focus(); + } + inputValueRef.current = ""; if (draftStorageKey) { localStorage.removeItem(draftStorageKey); } @@ -218,10 +268,29 @@ export function useConversationEditingState(deps: { } }; - const handleContentChange = (content: string) => { + const handleContentChange = ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => { + inputValueRef.current = content; + + // Don't overwrite the persisted draft while editing a + // history or queued message — the original draft (possibly + // containing file-reference chips) is saved in React state + // and should survive a cancel. + if (editingMessageId !== null || editingQueuedMessageID !== null) { + return; + } + if (draftStorageKey) { - if (content) { - localStorage.setItem(draftStorageKey, content); + const shouldPersist = content.trim() || hasFileReferences; + if (shouldPersist) { + try { + localStorage.setItem(draftStorageKey, serializedEditorState); + } catch { + // QuotaExceededError — silently discard the draft. + } } else { localStorage.removeItem(draftStorageKey); } @@ -229,7 +298,11 @@ export function useConversationEditingState(deps: { }; return { + inputValueRef, chatInputRef, + editorInitialValue, + initialEditorState, + remountKey, editingMessageId, editingFileBlocks, handleEditUserMessage, @@ -360,10 +433,12 @@ const AgentChatPage: FC = () => { >(null); const scrollToBottomRef = useRef<(() => void) | null>(null); const chatInputRef = useRef(null); - // Read once on mount — agentId is stable because KeyedAgentChatPage - // remounts the entire component when the route param changes. - const [initialInputValue] = useState(() => - getPersistedDraftInputValue(agentId), + const inputValueRef = useRef( + agentId + ? parseStoredDraft( + localStorage.getItem(`${draftInputStorageKeyPrefix}${agentId}`), + ).text + : "", ); // Right panel open/closed state is owned here so the loading @@ -610,15 +685,13 @@ const AgentChatPage: FC = () => { const handleCommit = (repoRoot: string) => { const commitPrompt = `Commit and push the working changes in ${repoRoot}. If there are unstaged files, commit them too.`; - const current = chatInputRef.current?.getValue() ?? ""; + const current = inputValueRef.current; if (current.includes(commitPrompt)) { return; } const prefix = current.trim() ? "\n\n" : ""; chatInputRef.current?.insertText(prefix + commitPrompt); - if (!isMobileViewport()) { - chatInputRef.current?.focus(); - } + chatInputRef.current?.focus(); }; // Prefer the explicit PR number from the API, and only fall back to URL @@ -868,6 +941,7 @@ const AgentChatPage: FC = () => { onSend: handleSend, onDeleteQueuedMessage: handleDeleteQueuedMessage, chatInputRef, + inputValueRef, }); const chatTitle = chatQuery.data?.title; @@ -1034,10 +1108,7 @@ const AgentChatPage: FC = () => { persistedError={persistedError} isArchived={isArchived} hasWorkspace={Boolean(workspaceId)} - workspaceAgent={workspaceAgent} - workspace={workspace} store={store} - initialInputValue={initialInputValue} editing={editing} pendingEditMessageId={pendingEditMessageId} effectiveSelectedModel={effectiveSelectedModel} diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index 546dbe73ee..67092e452f 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -59,6 +59,9 @@ const buildEditing = ( overrides: Partial["editing"]> = {}, ) => ({ chatInputRef: { current: null }, + editorInitialValue: "", + initialEditorState: undefined, + remountKey: 0, editingMessageId: null as number | null, editingFileBlocks: [] as readonly ChatMessagePart[], handleEditUserMessage: fn(), @@ -112,7 +115,6 @@ const StoryAgentChatPageView: FC = ({ editing, ...overrides }) => { hasWorkspace: true, store: createChatStore(), pendingEditMessageId: null as number | null, - initialInputValue: "", effectiveSelectedModel: defaultModelConfigID, setSelectedModel: fn(), modelOptions: defaultModelOptions, @@ -434,8 +436,8 @@ export const EditingMessage: Story = { store={buildStoreWithMessages(editingMessages)} editing={{ editingMessageId: 3, + editorInitialValue: "Now tell me a joke", }} - initialInputValue="Now tell me a joke" /> ), }; @@ -448,8 +450,8 @@ export const EditingSaving: Story = { store={buildStoreWithMessages(editingMessages)} editing={{ editingMessageId: 3, + editorInitialValue: "Now tell me a better joke", }} - initialInputValue="Now tell me a better joke" pendingEditMessageId={3} isSubmissionPending /> diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index d43eacccca..24b6a8fdd0 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -32,6 +32,9 @@ type ChatStoreHandle = ReturnType["store"]; interface EditingState { chatInputRef: RefObject; + editorInitialValue: string; + initialEditorState: string | undefined; + remountKey: number; editingMessageId: number | null; editingFileBlocks: readonly ChatMessagePart[]; handleEditUserMessage: ( @@ -48,7 +51,11 @@ interface EditingState { ) => void; handleCancelQueueEdit: () => void; handleSendFromInput: (message: string, fileIds?: string[]) => void; - handleContentChange: (content: string) => void; + handleContentChange: ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => void; } interface AgentChatPageViewProps { @@ -69,9 +76,6 @@ interface AgentChatPageViewProps { editing: EditingState; pendingEditMessageId: number | null; - // Input configuration. - initialInputValue: string; - // Model/input configuration. effectiveSelectedModel: string; setSelectedModel: (model: string) => void; @@ -158,7 +162,6 @@ export const AgentChatPageView: FC = ({ store, editing, pendingEditMessageId, - initialInputValue, effectiveSelectedModel, setSelectedModel, modelOptions, @@ -343,7 +346,9 @@ export const AgentChatPageView: FC = ({ modelSelectorPlaceholder={modelSelectorPlaceholder} isModelCatalogLoading={isModelCatalogLoading} inputRef={editing.chatInputRef} - initialInputValue={initialInputValue} + initialValue={editing.editorInitialValue} + initialEditorState={editing.initialEditorState} + remountKey={editing.remountKey} onContentChange={editing.handleContentChange} editingQueuedMessageID={editing.editingQueuedMessageID} onStartQueueEdit={editing.handleStartQueueEdit} diff --git a/site/src/pages/AgentsPage/AgentsPage.test.ts b/site/src/pages/AgentsPage/AgentsPage.test.ts index 2c365bf717..7202db414b 100644 --- a/site/src/pages/AgentsPage/AgentsPage.test.ts +++ b/site/src/pages/AgentsPage/AgentsPage.test.ts @@ -38,7 +38,11 @@ describe("useEmptyStateDraft", () => { const { result, unmount } = renderDraft(); act(() => { - result.current.handleContentChange("work in progress"); + result.current.handleContentChange( + "work in progress", + "work in progress", + false, + ); }); expect(localStorage.getItem(emptyInputStorageKey)).toBe("work in progress"); @@ -51,7 +55,10 @@ describe("useEmptyStateDraft", () => { const { result, unmount } = renderDraft(); act(() => { - result.current.handleContentChange(""); + // Even though the serialized state is non-empty (Lexical always + // produces a JSON object), the draft is removed when the plain + // text content is empty. + result.current.handleContentChange("", '{"root":{"children":[]}}', false); }); expect(localStorage.getItem(emptyInputStorageKey)).toBeNull(); @@ -86,7 +93,7 @@ describe("useEmptyStateDraft", () => { // the re-render with the old content. Without the sentRef // guard this would re-persist the draft. act(() => { - result.current.handleContentChange("fix the bug"); + result.current.handleContentChange("fix the bug", "fix the bug", false); }); expect(localStorage.getItem(emptyInputStorageKey)).toBeNull(); @@ -98,7 +105,7 @@ describe("useEmptyStateDraft", () => { const { result, unmount } = renderDraft(); act(() => { - result.current.handleContentChange("original"); + result.current.handleContentChange("original", "original", false); }); expect(localStorage.getItem(emptyInputStorageKey)).toBe("original"); @@ -107,7 +114,11 @@ describe("useEmptyStateDraft", () => { }); act(() => { - result.current.handleContentChange("totally new content"); + result.current.handleContentChange( + "totally new content", + "totally new content", + false, + ); }); expect(localStorage.getItem(emptyInputStorageKey)).toBeNull(); unmount(); @@ -133,7 +144,7 @@ describe("useEmptyStateDraft", () => { const { result, unmount } = renderDraft(); act(() => { - result.current.handleContentChange("attempt one"); + result.current.handleContentChange("attempt one", "attempt one", false); }); expect(localStorage.getItem(emptyInputStorageKey)).toBe("attempt one"); @@ -148,7 +159,7 @@ describe("useEmptyStateDraft", () => { }); act(() => { - result.current.handleContentChange("attempt two"); + result.current.handleContentChange("attempt two", "attempt two", false); }); expect(localStorage.getItem(emptyInputStorageKey)).toBe("attempt two"); unmount(); @@ -169,6 +180,113 @@ describe("useEmptyStateDraft", () => { expect(localStorage.getItem(emptyInputStorageKey)).toBeNull(); unmount(); }); + + it("persists serialized editor state when provided", () => { + const { result, unmount } = renderDraft(); + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [ + { text: "review this" }, + { + type: "file-reference", + version: 1, + fileName: "main.go", + startLine: 1, + endLine: 10, + content: "code", + }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + + act(() => { + result.current.handleContentChange("review this", editorState, true); + }); + + expect(localStorage.getItem(emptyInputStorageKey)).toBe(editorState); + expect(result.current.getCurrentContent()).toBe("review this"); + unmount(); + }); + + it("restores initialEditorState from a Lexical JSON draft", () => { + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [{ text: "hello" }], + type: "paragraph", + }, + ], + type: "root", + }, + }); + localStorage.setItem(emptyInputStorageKey, editorState); + + const { result, unmount } = renderDraft(); + + expect(result.current.initialEditorState).toBe(editorState); + expect(result.current.initialInputValue).toBe("hello"); + unmount(); + }); + + it("falls back to plain text for legacy drafts", () => { + localStorage.setItem(emptyInputStorageKey, "legacy plain text"); + + const { result, unmount } = renderDraft(); + + expect(result.current.initialEditorState).toBeUndefined(); + expect(result.current.initialInputValue).toBe("legacy plain text"); + unmount(); + }); + + it("persists file-reference-only drafts (no text content)", () => { + const { result, unmount } = renderDraft(); + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [ + { + type: "file-reference", + version: 1, + fileName: "main.go", + startLine: 1, + endLine: 10, + content: "code", + }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + + act(() => { + result.current.handleContentChange("", editorState, true); + }); + + expect(localStorage.getItem(emptyInputStorageKey)).toBe(editorState); + unmount(); + }); + + it("removes draft for whitespace-only content without file references", () => { + localStorage.setItem(emptyInputStorageKey, "old draft"); + const { result, unmount } = renderDraft(); + + act(() => { + result.current.handleContentChange(" ", '{"root":{}}', false); + }); + + expect(localStorage.getItem(emptyInputStorageKey)).toBeNull(); + unmount(); + }); }); describe("useFileAttachments persistence", () => { diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx index 7f73672b52..98244feed4 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.stories.tsx @@ -2,12 +2,12 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import { useEffect, useRef } from "react"; import { expect, fn, userEvent, waitFor, within } from "storybook/test"; import type * as TypesGen from "#/api/typesGenerated"; -import type { ChatMessageInputRef } from "#/components/ChatMessageInput/ChatMessageInput"; import { AgentChatInput, type AgentContextUsage, type UploadState, } from "./AgentChatInput"; +import type { ChatMessageInputRef } from "./ChatMessageInput/ChatMessageInput"; const defaultModelConfigID = "model-config-1"; diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index c89d630340..a2685d2eff 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -24,10 +24,6 @@ import type * as TypesGen from "#/api/typesGenerated"; import type { ChatMessagePart, ChatQueuedMessage } from "#/api/typesGenerated"; import { Alert } from "#/components/Alert/Alert"; import { Button } from "#/components/Button/Button"; -import { - ChatMessageInput, - type ChatMessageInputRef, -} from "#/components/ChatMessageInput/ChatMessageInput"; import { Command, CommandEmpty, @@ -55,17 +51,21 @@ import { formatProviderLabel } from "../utils/modelOptions"; import type { UploadState } from "./AttachmentPreview"; import { AttachmentPreview } from "./AttachmentPreview"; import { ModelSelector, type ModelSelectorOption } from "./ChatElements"; +import { + ChatMessageInput, + type ChatMessageInputRef, +} from "./ChatMessageInput/ChatMessageInput"; import type { AgentContextUsage } from "./ContextUsageIndicator"; import { ContextUsageIndicator } from "./ContextUsageIndicator"; import { ImageLightbox } from "./ImageLightbox"; import { QueuedMessagesList } from "./QueuedMessagesList"; import { TextPreviewDialog } from "./TextPreviewDialog"; -export type { ChatMessageInputRef } from "#/components/ChatMessageInput/ChatMessageInput"; export { ImageThumbnail, type UploadState, } from "./AttachmentPreview"; +export type { ChatMessageInputRef } from "./ChatMessageInput/ChatMessageInput"; export type { AgentContextUsage } from "./ContextUsageIndicator"; interface AgentChatInputProps { @@ -77,8 +77,17 @@ interface AgentChatInputProps { inputRef?: React.Ref; // Initial text to seed the editor on first mount only. initialValue?: string; - // Called on every text change inside the editor. - onContentChange?: (content: string) => void; + // Serialized Lexical editor state for restoring drafts with + // file-reference chips. Takes precedence over initialValue. + initialEditorState?: string; + // Monotonic counter to force editor remount. + remountKey?: number; + // Called on every content change inside the editor. + onContentChange?: ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => void; // Model selector. selectedModel: string; onModelChange: (value: string) => void; @@ -201,6 +210,8 @@ export const AgentChatInput: FC = ({ isLoading, inputRef, initialValue, + initialEditorState, + remountKey, onContentChange, selectedModel, onModelChange, @@ -428,11 +439,15 @@ export const AgentChatInput: FC = ({ countInvisibleCharacters(initialValue ?? ""), ); - const handleContentChange = (content: string, hasRefs: boolean) => { + const handleContentChange = ( + content: string, + serializedEditorState: string, + hasRefs: boolean, + ) => { setHasContent(Boolean(content.trim())); setHasFileReferences(hasRefs); setInvisibleCharCount(countInvisibleCharacters(content)); - onContentChange?.(content); + onContentChange?.(content, serializedEditorState, hasRefs); }; // Re-focus the editor after a send completes (isLoading goes @@ -648,6 +663,8 @@ export const AgentChatInput: FC = ({ className="min-h-[60px] sm:min-h-24 w-full resize-none bg-transparent px-3 py-2 font-sans text-[15px] leading-6 text-content-primary placeholder:text-content-secondary disabled:cursor-not-allowed disabled:opacity-70" placeholder={placeholder} initialValue={initialValue} + initialEditorState={initialEditorState} + remountKey={remountKey} onChange={handleContentChange} onKeyDown={handleEditorKeyDown} onEnter={handleSubmit} diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx index 89cbc90776..bafbdd2e82 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx @@ -9,6 +9,7 @@ import { Button } from "#/components/Button/Button"; import { useDashboard } from "#/modules/dashboard/useDashboard"; import { docs } from "#/utils/docs"; import { useFileAttachments } from "../hooks/useFileAttachments"; +import { parseStoredDraft } from "../utils/draftStorage"; import { getModelSelectorPlaceholder, hasConfiguredModelsInCatalog, @@ -53,17 +54,30 @@ export type CreateChatOptions = { * @internal Exported for testing. */ export function useEmptyStateDraft() { - const [initialInputValue] = useState(() => { - return localStorage.getItem(emptyInputStorageKey) ?? ""; + const [{ initialInputValue, initialEditorState }] = useState(() => { + const draft = parseStoredDraft(localStorage.getItem(emptyInputStorageKey)); + return { + initialInputValue: draft.text, + initialEditorState: draft.editorState, + }; }); const inputValueRef = useRef(initialInputValue); const sentRef = useRef(false); - const handleContentChange = (content: string) => { + const handleContentChange = ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => { inputValueRef.current = content; if (!sentRef.current) { - if (content) { - localStorage.setItem(emptyInputStorageKey, content); + const shouldPersist = content.trim() || hasFileReferences; + if (shouldPersist) { + try { + localStorage.setItem(emptyInputStorageKey, serializedEditorState); + } catch { + // QuotaExceededError — silently discard the draft. + } } else { localStorage.removeItem(emptyInputStorageKey); } @@ -85,6 +99,7 @@ export function useEmptyStateDraft() { return { initialInputValue, + initialEditorState, getCurrentContent, handleContentChange, submitDraft, @@ -128,8 +143,13 @@ export const AgentCreateForm: FC = ({ isWorkspacesLoading, }) => { const { organizations } = useDashboard(); - const { initialInputValue, handleContentChange, submitDraft, resetDraft } = - useEmptyStateDraft(); + const { + initialInputValue, + initialEditorState, + handleContentChange, + submitDraft, + resetDraft, + } = useEmptyStateDraft(); const [initialLastModelConfigID] = useState(() => { return localStorage.getItem(lastModelConfigIDStorageKey) ?? ""; }); @@ -320,6 +340,7 @@ export const AgentCreateForm: FC = ({ isDisabled={isCreating || isForbidden} isLoading={isCreating} initialValue={initialInputValue} + initialEditorState={initialEditorState} onContentChange={handleContentChange} selectedModel={selectedModel} onModelChange={handleModelChange} diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx index 1303edf3af..5e681106f7 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx @@ -10,7 +10,6 @@ import { } from "react"; import type { UrlTransform } from "streamdown"; import type * as TypesGen from "#/api/typesGenerated"; -import { FileReferenceChip } from "#/components/ChatMessageInput/FileReferenceNode"; import { CopyButton } from "#/components/CopyButton/CopyButton"; import { Spinner } from "#/components/Spinner/Spinner"; import { @@ -34,6 +33,7 @@ import { Tool, } from "../ChatElements"; import { WebSearchSources } from "../ChatElements/tools"; +import { FileReferenceChip } from "../ChatMessageInput/FileReferenceNode"; import { ImageLightbox } from "../ImageLightbox"; import { TextPreviewDialog } from "../TextPreviewDialog"; import { getEditableUserMessagePayload } from "./messageParsing"; diff --git a/site/src/components/ChatMessageInput/ChatMessageInput.test.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.test.tsx similarity index 100% rename from site/src/components/ChatMessageInput/ChatMessageInput.test.tsx rename to site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.test.tsx diff --git a/site/src/components/ChatMessageInput/ChatMessageInput.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx similarity index 50% rename from site/src/components/ChatMessageInput/ChatMessageInput.tsx rename to site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx index 2d18d7cb89..33d28bd9b6 100644 --- a/site/src/components/ChatMessageInput/ChatMessageInput.tsx +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx @@ -11,6 +11,8 @@ import { $createTextNode, $getRoot, $getSelection, + $insertNodes, + $isParagraphNode, $isRangeSelection, COMMAND_PRIORITY_HIGH, FORMAT_ELEMENT_COMMAND, @@ -19,16 +21,12 @@ import { KEY_ENTER_COMMAND, type LexicalEditor, PASTE_COMMAND, - type ParagraphNode, } from "lexical"; import { type FC, - memo, - useCallback, useEffect, useImperativeHandle, useLayoutEffect, - useMemo, useRef, } from "react"; import { cn } from "#/utils/cn"; @@ -46,7 +44,7 @@ import { // Blocks Cmd+B/I/U and element formatting shortcuts so the editor // stays plain-text only. -const DisableFormattingPlugin: FC = memo(function DisableFormattingPlugin() { +const DisableFormattingPlugin: FC = function DisableFormattingPlugin() { const [editor] = useLexicalComposerContext(); useEffect(() => { @@ -65,7 +63,7 @@ const DisableFormattingPlugin: FC = memo(function DisableFormattingPlugin() { }, [editor]); return null; -}); +}; function insertPlainTextIntoEditor(editor: LexicalEditor, text: string) { editor.update(() => { @@ -76,17 +74,15 @@ function insertPlainTextIntoEditor(editor: LexicalEditor, text: string) { } const root = $getRoot(); const lastChild = root.getLastChild(); - if (lastChild) { - if (lastChild.getType() === "paragraph") { - const paragraph = lastChild as ParagraphNode; - const textNode = $createTextNode(text); - paragraph.append(textNode); - textNode.selectEnd(); - } else { - const textNode = $createTextNode(text); - lastChild.insertAfter(textNode); - textNode.selectEnd(); - } + + if (lastChild && $isParagraphNode(lastChild)) { + const textNode = $createTextNode(text); + lastChild.append(textNode); + textNode.selectEnd(); + } else if (lastChild) { + const textNode = $createTextNode(text); + lastChild.insertAfter(textNode); + textNode.selectEnd(); } else { const paragraph = $createParagraphNode(); const textNode = $createTextNode(text); @@ -97,6 +93,8 @@ function insertPlainTextIntoEditor(editor: LexicalEditor, text: string) { }); } +// Replaces the entire editor content with the given plain text. +// Used by the imperative setValue() method. function replacePlainTextInEditor(editor: LexicalEditor, text: string) { editor.update(() => { const root = $getRoot(); @@ -129,7 +127,7 @@ function replacePlainTextInEditor(editor: LexicalEditor, text: string) { const PasteSanitizationPlugin: FC<{ onFilePaste?: (file: File) => void; allowTextAttachmentPaste?: boolean; -}> = memo(function PasteSanitizationPlugin({ +}> = function PasteSanitizationPlugin({ onFilePaste, allowTextAttachmentPaste = true, }) { @@ -255,40 +253,44 @@ const PasteSanitizationPlugin: FC<{ }, [allowTextAttachmentPaste, editor, onFilePaste]); return null; -}); +}; // Handles Enter key behavior: plain Enter submits via the onEnter // callback, Shift+Enter inserts a newline. -const EnterKeyPlugin: FC<{ onEnter?: () => void }> = memo( - function EnterKeyPlugin({ onEnter }) { - const [editor] = useLexicalComposerContext(); +const EnterKeyPlugin: FC<{ onEnter?: () => void }> = function EnterKeyPlugin({ + onEnter, +}) { + const [editor] = useLexicalComposerContext(); - useEffect(() => { - return editor.registerCommand( - KEY_ENTER_COMMAND, - (event: KeyboardEvent | null) => { - if (event?.shiftKey) { - return false; - } - if (onEnter) { - event?.preventDefault(); - onEnter(); - } - return true; - }, - COMMAND_PRIORITY_HIGH, - ); - }, [editor, onEnter]); + useEffect(() => { + return editor.registerCommand( + KEY_ENTER_COMMAND, + (event: KeyboardEvent | null) => { + if (event?.shiftKey) { + return false; + } + if (onEnter) { + event?.preventDefault(); + onEnter(); + } + return true; + }, + COMMAND_PRIORITY_HIGH, + ); + }, [editor, onEnter]); - return null; - }, -); + return null; +}; // Fires the onChange callback with the editor's plain-text content // on every update. const ContentChangePlugin: FC<{ - onChange?: (content: string, hasFileReferences: boolean) => void; -}> = memo(function ContentChangePlugin({ onChange }) { + onChange?: ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => void; +}> = function ContentChangePlugin({ onChange }) { const [editor] = useLexicalComposerContext(); useEffect(() => { @@ -299,29 +301,79 @@ const ContentChangePlugin: FC<{ const root = $getRoot(); const content = root.getTextContent(); let hasRefs = false; + for (const child of root.getChildren()) { - if (child.getType() !== "paragraph") continue; - for (const node of (child as ParagraphNode).getChildren()) { + if (!$isParagraphNode(child)) continue; + + for (const node of child.getChildren()) { if (node instanceof FileReferenceNode) { hasRefs = true; break; } } + if (hasRefs) break; } - onChange(content, hasRefs); + const serialized = JSON.stringify(editorState.toJSON()); + onChange(content, serialized, hasRefs); }); }); }, [editor, onChange]); return null; -}); +}; + +// Seeds the editor with an initial value on first mount. When +// initialEditorState is provided (a serialized Lexical JSON string), +// it restores the full editor state including file-reference chips. +// Falls back to plain-text seeding via initialValue. +const ValueSyncPlugin: FC<{ + initialValue?: string; + initialEditorState?: string; +}> = function ValueSyncPlugin({ initialValue, initialEditorState }) { + const [editor] = useLexicalComposerContext(); + const hasInitialized = useRef(false); + + useEffect(() => { + if (hasInitialized.current) { + return; + } + hasInitialized.current = true; + + // Prefer restoring the full serialized editor state + // (preserves file-reference chips and node positions). + if (initialEditorState) { + try { + const parsed = editor.parseEditorState(initialEditorState); + editor.setEditorState(parsed); + return; + } catch { + // Malformed state — fall through to plain-text path. + } + } + + if (initialValue === undefined || initialValue === "") { + return; + } + + editor.update(() => { + const root = $getRoot(); + root.clear(); + const paragraph = $createParagraphNode(); + const textNode = $createTextNode(initialValue); + paragraph.append(textNode); + root.append(paragraph); + }); + }, [editor, initialValue, initialEditorState]); + + return null; +}; // Exposes the LexicalEditor instance to the parent via a callback // so it can be stored in a ref for imperative access. const InsertTextPlugin: FC<{ onEditorReady: (editor: LexicalEditor) => void; -}> = memo(function InsertTextPlugin({ onEditorReady }) { +}> = function InsertTextPlugin({ onEditorReady }) { const [editor] = useLexicalComposerContext(); useEffect(() => { @@ -329,7 +381,7 @@ const InsertTextPlugin: FC<{ }, [editor, onEditorReady]); return null; -}); +}; /** * Structured data for a file reference extracted from the editor. @@ -352,10 +404,17 @@ type EditorContentPart = readonly reference: FileReferenceData; }; +// Mutable variant used internally while building the parts +// array so we can append to the last text segment without +// casting away readonly. +type MutableTextPart = { type: "text"; text: string }; +type MutableFileRefPart = { + type: "file-reference"; + reference: FileReferenceData; +}; +type MutableContentPart = MutableTextPart | MutableFileRefPart; + export interface ChatMessageInputRef { - /** - * Replace the editor's plain-text content in a single Lexical update. - */ setValue: (text: string) => void; insertText: (text: string) => void; clear: () => void; @@ -378,7 +437,19 @@ interface ChatMessageInputProps extends Omit, "onChange" | "role" | "ref"> { placeholder?: string; initialValue?: string; - onChange?: (content: string, hasFileReferences: boolean) => void; + /** + * Serialized Lexical editor state JSON. When provided, the editor + * restores the full state (including file-reference chips) instead + * of using initialValue as plain text. + */ + initialEditorState?: string; + onChange?: ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => void; + /** Monotonic counter to force editor remount. */ + remountKey?: number; rows?: number; onEnter?: () => void; onFilePaste?: (file: File) => void; @@ -391,7 +462,7 @@ interface ChatMessageInputProps // Keeps the Lexical editor's editable state in sync with the // disabled prop so that the underlying contentEditable element // becomes truly non-interactive when the input is disabled. -const EditableStatePlugin: FC<{ disabled: boolean }> = memo( +const EditableStatePlugin: FC<{ disabled: boolean }> = function EditableStatePlugin({ disabled }) { const [editor] = useLexicalComposerContext(); @@ -400,241 +471,250 @@ const EditableStatePlugin: FC<{ disabled: boolean }> = memo( }, [editor, disabled]); return null; - }, -); + }; -const ChatMessageInput = memo( - ({ - className, - placeholder, - initialValue, - onChange, - rows, - onEnter, - onFilePaste, - allowTextAttachmentPaste, - disabled, - autoFocus, - "aria-label": ariaLabel, - ref, - ...props - }: ChatMessageInputProps & { ref?: React.Ref }) => { - const initialConfig = useMemo( - () => ({ - namespace: "ChatMessageInput", - theme: { - paragraph: "m-0", - inlineDecorator: "mx-1", - }, - onError: (error: Error) => console.error("Lexical error:", error), - nodes: [FileReferenceNode], - editable: !disabled, - }), - [disabled], - ); - const style = useMemo( - () => ({ - minHeight: rows ? `${rows * 1.5}rem` : undefined, - }), - [rows], - ); +const ChatMessageInput = ({ + className, + placeholder, + initialValue, + initialEditorState, + onChange, + remountKey, + rows, + onEnter, + onFilePaste, + allowTextAttachmentPaste, + disabled, + autoFocus, + "aria-label": ariaLabel, + ref, + ...props +}: ChatMessageInputProps & { ref?: React.Ref }) => { + const initialConfig = { + namespace: "ChatMessageInput", + theme: { + paragraph: "m-0", + inlineDecorator: "mx-1", + }, + onError: (error: Error) => console.error("Lexical error:", error), + nodes: [FileReferenceNode], + editable: !disabled, + }; + const style = { + minHeight: rows ? `${rows * 1.5}rem` : undefined, + }; - const editorRef = useRef(null); - const lastKnownValueRef = useRef(initialValue ?? ""); - const pendingReplacementRef = useRef(null); + const editorRef = useRef(null); + // Tracks the last known text content so getValue() can return + // a useful value before the Lexical editor hydrates. + const lastKnownValueRef = useRef(initialValue ?? ""); + // Queues a setValue call made before the editor ref is ready. + const pendingReplacementRef = useRef(null); - const replaceValueOrQueue = useCallback((text: string) => { - lastKnownValueRef.current = text; - const editor = editorRef.current; - if (!editor) { - pendingReplacementRef.current = text; - return; - } + const handleEditorReady = (editor: LexicalEditor) => { + editorRef.current = editor; + // Flush any queued setValue that arrived before the editor + // was ready (e.g. useLayoutEffect in a parent). + const pending = pendingReplacementRef.current; + if (pending !== null) { pendingReplacementRef.current = null; - replacePlainTextInEditor(editor, text); - }, []); + replacePlainTextInEditor(editor, pending); + } + }; - const handleEditorReady = useCallback((editor: LexicalEditor) => { - editorRef.current = editor; - const pendingReplacement = pendingReplacementRef.current; - if (pendingReplacement !== null) { + useImperativeHandle( + ref, + () => ({ + setValue: (text: string) => { + lastKnownValueRef.current = text; + const editor = editorRef.current; + if (!editor) { + pendingReplacementRef.current = text; + return; + } pendingReplacementRef.current = null; - replacePlainTextInEditor(editor, pendingReplacement); - return; - } - const initialText = lastKnownValueRef.current; - if (initialText) { - replacePlainTextInEditor(editor, initialText); - } - }, []); - - const handleContentChange = useCallback( - (content: string, hasFileReferences: boolean) => { - lastKnownValueRef.current = content; - onChange?.(content, hasFileReferences); + replacePlainTextInEditor(editor, text); }, - [onChange], - ); + insertText: (text: string) => { + const editor = editorRef.current; + if (!editor) return; - useImperativeHandle( - ref, - () => ({ - setValue: (text: string) => { - replaceValueOrQueue(text); - }, - insertText: (text: string) => { - const editor = editorRef.current; - if (!editor) return; - - insertPlainTextIntoEditor(editor, text); - }, - clear: () => { - replaceValueOrQueue(""); - }, - focus: () => { - const editor = editorRef.current; - if (!editor) return; - editor.focus(() => { - editor.update(() => { - const root = $getRoot(); - const last = root.getLastChild(); - if (!last) { - const paragraph = $createParagraphNode(); - root.append(paragraph); - paragraph.select(); - return; - } - last.selectEnd(); - }); - }); - }, - getValue: () => { - const editor = editorRef.current; - if (!editor) { - return lastKnownValueRef.current; + editor.update(() => { + const selection = $getSelection(); + if ($isRangeSelection(selection)) { + const textNode = $createTextNode(text); + $insertNodes([textNode]); + textNode.selectEnd(); + } else { + insertPlainTextIntoEditor(editor, text); } - let content = ""; - editor.getEditorState().read(() => { - content = $getRoot().getTextContent(); - }); - return content; - }, - addFileReference: (ref: FileReferenceData) => { - const editor = editorRef.current; - if (!editor) return; + }); + }, + clear: () => { + const editor = editorRef.current; + if (!editor) return; + editor.update(() => { + const root = $getRoot(); + root.clear(); + const paragraph = $createParagraphNode(); + root.append(paragraph); + paragraph.select(); + }); + }, + focus: () => { + const editor = editorRef.current; + if (!editor) return; + editor.focus(() => { editor.update(() => { const root = $getRoot(); - let paragraph = root.getFirstChild(); - if (!paragraph || paragraph.getType() !== "paragraph") { - paragraph = $createParagraphNode(); + const last = root.getLastChild(); + if (!last) { + const paragraph = $createParagraphNode(); root.append(paragraph); + paragraph.select(); + return; } - const chipNode = $createFileReferenceNode( - ref.fileName, - ref.startLine, - ref.endLine, - ref.content, - ); - (paragraph as ParagraphNode).append(chipNode); - chipNode.selectNext(); + last.selectEnd(); }); - }, - getContentParts: () => { - const editor = editorRef.current; - if (!editor) return []; - const parts: EditorContentPart[] = []; - editor.getEditorState().read(() => { - const paragraphs = $getRoot().getChildren(); - for (let i = 0; i < paragraphs.length; i++) { - const para = paragraphs[i]; - if (para.getType() !== "paragraph") continue; - // Separate paragraphs with a newline in the - // preceding text part, just like getTextContent(). - if (i > 0) { - const last = parts[parts.length - 1]; - if (last?.type === "text") { - (last as { text: string }).text += "\n"; - } else { - parts.push({ type: "text", text: "\n" }); - } - } - for (const node of (para as ParagraphNode).getChildren()) { - if (node instanceof FileReferenceNode) { - parts.push({ - type: "file-reference", - reference: { - fileName: node.__fileName, - startLine: node.__startLine, - endLine: node.__endLine, - content: node.__content, - }, - }); - } else { - // Text node (or any other inline) — - // merge into the last text part. - const t = node.getTextContent(); - if (!t) continue; - const last = parts[parts.length - 1]; - if (last?.type === "text") { - (last as { text: string }).text += t; - } else { - parts.push({ type: "text", text: t }); - } - } - } - } - }); - return parts; - }, - }), - [replaceValueOrQueue], - ); + }); + }, + getValue: () => { + const editor = editorRef.current; + if (!editor) { + return lastKnownValueRef.current; + } + let content = ""; + editor.getEditorState().read(() => { + content = $getRoot().getTextContent(); + }); + return content; + }, + addFileReference: (ref: FileReferenceData) => { + const editor = editorRef.current; + if (!editor) return; - return ( - -
*]:col-start-1 [&>*]:row-start-1", - disabled && "cursor-not-allowed opacity-50", - className, - )} - style={style} - {...props} - > - + editor.update(() => { + const root = $getRoot(); + const firstChild = root.getFirstChild(); + const paragraph = $isParagraphNode(firstChild) + ? firstChild + : $createParagraphNode(); + + if (!$isParagraphNode(firstChild)) { + root.append(paragraph); + } + + const chipNode = $createFileReferenceNode( + ref.fileName, + ref.startLine, + ref.endLine, + ref.content, + ); + paragraph.append(chipNode); + chipNode.selectNext(); + }); + }, + getContentParts: () => { + const editor = editorRef.current; + if (!editor) return []; + + const parts: MutableContentPart[] = []; + + const appendText = (str: string) => { + const last = parts[parts.length - 1]; + if (last?.type === "text") { + last.text += str; + } else { + parts.push({ type: "text", text: str }); + } + }; + + editor.getEditorState().read(() => { + const paragraphs = $getRoot().getChildren(); + + for (let i = 0; i < paragraphs.length; i++) { + const para = paragraphs[i]; + if (!$isParagraphNode(para)) continue; + + // Separate paragraphs with a newline in the + // preceding text part, just like getTextContent(). + if (i > 0) { + appendText("\n"); } - placeholder={ -
- {placeholder} -
+ + for (const node of para.getChildren()) { + if (node instanceof FileReferenceNode) { + parts.push({ + type: "file-reference", + reference: { + fileName: node.__fileName, + startLine: node.__startLine, + endLine: node.__endLine, + content: node.__content, + }, + }); + } else { + const t = node.getTextContent(); + if (t) appendText(t); + } } - ErrorBoundary={LexicalErrorBoundary} - /> - - - - - - - - {autoFocus && } -
-
- ); - }, -); + } + }); + + return parts as EditorContentPart[]; + }, + }), + [], + ); + + return ( + +
*]:col-start-1 [&>*]:row-start-1", + disabled && "cursor-not-allowed opacity-50", + className, + )} + style={style} + {...props} + > + + } + placeholder={ +
+ {placeholder} +
+ } + ErrorBoundary={LexicalErrorBoundary} + /> + + + + + + + + + {autoFocus && } +
+
+ ); +}; ChatMessageInput.displayName = "ChatMessageInput"; export { ChatMessageInput }; diff --git a/site/src/components/ChatMessageInput/FileReferenceNode.stories.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/FileReferenceNode.stories.tsx similarity index 100% rename from site/src/components/ChatMessageInput/FileReferenceNode.stories.tsx rename to site/src/pages/AgentsPage/components/ChatMessageInput/FileReferenceNode.stories.tsx diff --git a/site/src/components/ChatMessageInput/FileReferenceNode.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/FileReferenceNode.tsx similarity index 100% rename from site/src/components/ChatMessageInput/FileReferenceNode.tsx rename to site/src/pages/AgentsPage/components/ChatMessageInput/FileReferenceNode.tsx diff --git a/site/src/components/ChatMessageInput/pasteHelpers.test.ts b/site/src/pages/AgentsPage/components/ChatMessageInput/pasteHelpers.test.ts similarity index 100% rename from site/src/components/ChatMessageInput/pasteHelpers.test.ts rename to site/src/pages/AgentsPage/components/ChatMessageInput/pasteHelpers.test.ts diff --git a/site/src/components/ChatMessageInput/pasteHelpers.ts b/site/src/pages/AgentsPage/components/ChatMessageInput/pasteHelpers.ts similarity index 100% rename from site/src/components/ChatMessageInput/pasteHelpers.ts rename to site/src/pages/AgentsPage/components/ChatMessageInput/pasteHelpers.ts diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 34dc4955c1..f3e945946f 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -126,8 +126,14 @@ interface ChatPageInputProps { // Imperative editor handle plus the one-time initial draft, // owned by the conversation component. inputRef?: React.Ref; - initialInputValue?: string; - onContentChange?: (content: string) => void; + initialValue?: string; + initialEditorState?: string; + remountKey?: number; + onContentChange?: ( + content: string, + serializedEditorState: string, + hasFileReferences: boolean, + ) => void; editingQueuedMessageID: number | null; onStartQueueEdit: ( id: number, @@ -170,7 +176,9 @@ export const ChatPageInput: FC = ({ modelSelectorPlaceholder, isModelCatalogLoading = false, inputRef, - initialInputValue, + initialValue, + initialEditorState, + remountKey, onContentChange, editingQueuedMessageID, onStartQueueEdit, @@ -315,7 +323,9 @@ export const ChatPageInput: FC = ({ previewUrls={previewUrls} textContents={textContents} inputRef={inputRef} - initialValue={initialInputValue} + initialValue={initialValue} + initialEditorState={initialEditorState} + remountKey={remountKey} onContentChange={onContentChange} queuedMessages={queuedMessages} onDeleteQueuedMessage={onDeleteQueuedMessage} diff --git a/site/src/pages/AgentsPage/utils/draftStorage.test.ts b/site/src/pages/AgentsPage/utils/draftStorage.test.ts new file mode 100644 index 0000000000..e794060ca0 --- /dev/null +++ b/site/src/pages/AgentsPage/utils/draftStorage.test.ts @@ -0,0 +1,167 @@ +import { describe, expect, it } from "vitest"; +import { parseStoredDraft } from "./draftStorage"; + +describe("parseStoredDraft", () => { + it("returns empty text and no editorState for null input", () => { + const result = parseStoredDraft(null); + expect(result.text).toBe(""); + expect(result.editorState).toBeUndefined(); + }); + + it("returns empty text and no editorState for empty string", () => { + const result = parseStoredDraft(""); + expect(result.text).toBe(""); + expect(result.editorState).toBeUndefined(); + }); + + it("treats a plain text string as a legacy draft", () => { + const result = parseStoredDraft("hello world"); + expect(result.text).toBe("hello world"); + expect(result.editorState).toBeUndefined(); + }); + + it("treats valid JSON without a root key as a legacy draft", () => { + const raw = JSON.stringify({ foo: "bar" }); + const result = parseStoredDraft(raw); + expect(result.text).toBe(raw); + expect(result.editorState).toBeUndefined(); + }); + + it("treats JSON with a root key but no type as a legacy draft", () => { + const raw = JSON.stringify({ root: true }); + const result = parseStoredDraft(raw); + expect(result.text).toBe(raw); + expect(result.editorState).toBeUndefined(); + }); + + it("detects Lexical editor state JSON (has root key)", () => { + const state = JSON.stringify({ + root: { + children: [ + { + children: [{ text: "hello", type: "text" }], + type: "paragraph", + }, + ], + type: "root", + }, + }); + const result = parseStoredDraft(state); + expect(result.editorState).toBe(state); + expect(result.text).toBe("hello"); + }); + + it("extracts text from multiple paragraphs joined by newlines", () => { + const state = JSON.stringify({ + root: { + children: [ + { + children: [{ text: "line one", type: "text" }], + type: "paragraph", + }, + { + children: [{ text: "line two", type: "text" }], + type: "paragraph", + }, + ], + type: "root", + }, + }); + const result = parseStoredDraft(state); + expect(result.text).toBe("line one\n\nline two"); + }); + + it("skips non-text nodes (file-reference chips)", () => { + const state = JSON.stringify({ + root: { + children: [ + { + children: [ + { text: "review ", type: "text" }, + { + type: "file-reference", + fileName: "main.go", + startLine: 1, + endLine: 10, + content: "code", + }, + { text: " please", type: "text" }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + const result = parseStoredDraft(state); + expect(result.text).toBe("review please"); + }); + + it("handles empty root children", () => { + const state = JSON.stringify({ + root: { children: [], type: "root" }, + }); + const result = parseStoredDraft(state); + expect(result.text).toBe(""); + expect(result.editorState).toBe(state); + }); + + it("handles paragraphs with no children", () => { + const state = JSON.stringify({ + root: { + children: [{ type: "paragraph" }], + type: "root", + }, + }); + const result = parseStoredDraft(state); + expect(result.text).toBe(""); + expect(result.editorState).toBe(state); + }); + + it("extracts text from deeply nested structures", () => { + const state = JSON.stringify({ + root: { + children: [ + { + children: [ + { + children: [{ text: "nested", type: "text" }], + type: "listitem", + }, + ], + type: "list", + }, + ], + type: "root", + }, + }); + const result = parseStoredDraft(state); + expect(result.text).toBe("nested"); + }); + + it("extracts linebreak nodes as newline characters", () => { + const state = JSON.stringify({ + root: { + children: [ + { + children: [ + { text: "before", type: "text" }, + { type: "linebreak", version: 1 }, + { text: "after", type: "text" }, + ], + type: "paragraph", + }, + ], + type: "root", + }, + }); + const result = parseStoredDraft(state); + expect(result.text).toBe("before\nafter"); + }); + + it("handles malformed JSON gracefully", () => { + const result = parseStoredDraft("{not valid json"); + expect(result.text).toBe("{not valid json"); + expect(result.editorState).toBeUndefined(); + }); +}); diff --git a/site/src/pages/AgentsPage/utils/draftStorage.ts b/site/src/pages/AgentsPage/utils/draftStorage.ts new file mode 100644 index 0000000000..1b8b3ece20 --- /dev/null +++ b/site/src/pages/AgentsPage/utils/draftStorage.ts @@ -0,0 +1,83 @@ +/** + * Utilities for persisting and restoring chat input drafts. + * + * Drafts are stored in localStorage. The current format stores the + * serialized Lexical editor state JSON so that file-reference chips + * survive navigation. Legacy drafts (plain-text strings) are detected + * and handled transparently on read. + */ + +export interface ParsedDraft { + /** Plain text content for inputValueRef / send-button checks. */ + text: string; + /** + * The raw Lexical serialized editor state JSON string, if the + * stored draft was in the structured format. `undefined` for + * legacy plain-text drafts. + */ + editorState: string | undefined; +} + +/** + * Read a draft from localStorage and determine whether it is a + * Lexical editor state (JSON with a `root` key) or a legacy + * plain-text string. + */ +export function parseStoredDraft(raw: string | null): ParsedDraft { + if (!raw) { + return { text: "", editorState: undefined }; + } + try { + const parsed = JSON.parse(raw); + if (parsed?.root?.type === "root") { + return { text: extractPlainText(parsed), editorState: raw }; + } + } catch { + // Not JSON — treat as legacy plain-text draft. + } + return { text: raw, editorState: undefined }; +} + +/** + * Recursively walk a serialized Lexical node tree and extract the + * concatenated plain text. Mirrors `$getRoot().getTextContent()` + * without needing a live editor instance. + */ +function extractTextFromNode(node: { + text?: string; + children?: Array>; + type?: string; +}): string { + // Leaf text node. + if (typeof node.text === "string") { + return node.text; + } + // LineBreakNode serializes as { type: "linebreak" } with no + // text or children. Lexical's getTextContent() returns "\n". + if (node.type === "linebreak") { + return "\n"; + } + // FileReferenceNode and other non-text leaves contribute + // nothing to plain text, matching getTextContent() behavior. + if (!node.children) { + return ""; + } + const childTexts = node.children.map((child) => + extractTextFromNode(child as typeof node), + ); + // Join root-level children (paragraphs) with double + // newlines, matching Lexical's getTextContent() behavior. + if (node.type === "root") { + return childTexts.join("\n\n"); + } + return childTexts.join(""); +} + +function extractPlainText(state: { root?: Record }): string { + if (!state.root) { + return ""; + } + return extractTextFromNode( + state.root as Parameters[0], + ); +}