fix(site/src/pages/AgentsPage): persist file-reference chips across chat navigation (#23854)

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<ChatMessageInputRef>();
+		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();
+	});
 });

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<void>;
 	onDeleteQueuedMessage: (id: number) => Promise<void>;
 	chatInputRef: React.RefObject<ChatMessageInputRef | null>;
+	inputValueRef: React.RefObject<string>;
 }) {
-	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<string>).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<number | null>(null);
-	const [draftBeforeHistoryEdit, setDraftBeforeHistoryEdit] = useState<
-		string | null
-	>(null);
+	const [draftBeforeHistoryEdit, setDraftBeforeHistoryEdit] =
+		useState<ParsedDraft | null>(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<ParsedDraft | null>(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<ChatMessageInputRef | null>(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}

site/src/pages/AgentsPage/AgentChatPageView.stories.tsx

@@ -59,6 +59,9 @@ const buildEditing = (
 	overrides: Partial<ComponentProps<typeof AgentChatPageView>["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<StoryProps> = ({ 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
 		/>

site/src/pages/AgentsPage/AgentChatPageView.tsx

@@ -32,6 +32,9 @@ type ChatStoreHandle = ReturnType<typeof useChatStore>["store"];
 
 interface EditingState {
 	chatInputRef: RefObject<ChatMessageInputRef | null>;
+	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<AgentChatPageViewProps> = ({
 	store,
 	editing,
 	pendingEditMessageId,
-	initialInputValue,
 	effectiveSelectedModel,
 	setSelectedModel,
 	modelOptions,
@@ -343,7 +346,9 @@ export const AgentChatPageView: FC<AgentChatPageViewProps> = ({
 							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}

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", () => {

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";
 

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<ChatMessageInputRef>;
 	// 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<AgentChatInputProps> = ({
 	isLoading,
 	inputRef,
 	initialValue,
+	initialEditorState,
+	remountKey,
 	onContentChange,
 	selectedModel,
 	onModelChange,
@@ -428,11 +439,15 @@ export const AgentChatInput: FC<AgentChatInputProps> = ({
 		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<AgentChatInputProps> = ({
 					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}

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<AgentCreateFormProps> = ({
 	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<AgentCreateFormProps> = ({
 					isDisabled={isCreating || isForbidden}
 					isLoading={isCreating}
 					initialValue={initialInputValue}
+					initialEditorState={initialEditorState}
 					onContentChange={handleContentChange}
 					selectedModel={selectedModel}
 					onModelChange={handleModelChange}

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";

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<React.ComponentProps<"div">, "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<ChatMessageInputRef> }) => {
-		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<ChatMessageInputRef> }) => {
+	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<LexicalEditor | null>(null);
-		const lastKnownValueRef = useRef(initialValue ?? "");
-		const pendingReplacementRef = useRef<string | null>(null);
+	const editorRef = useRef<LexicalEditor | null>(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<string | null>(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 (
-			<LexicalComposer initialConfig={initialConfig}>
-				<div
-					className={cn(
-						"grid w-full rounded-md bg-transparent text-base placeholder:text-content-secondary focus-visible:outline-none whitespace-pre-wrap break-words [&>*]:col-start-1 [&>*]:row-start-1",
-						disabled && "cursor-not-allowed opacity-50",
-						className,
-					)}
-					style={style}
-					{...props}
-				>
-					<RichTextPlugin
-						contentEditable={
-							<ContentEditable
-								className="outline-none w-full whitespace-pre-wrap overflow-y-auto max-h-[50vh] [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent] [&_p]:leading-normal [&_p:first-child]:mt-0 [&_p:last-child]:mb-0 py-px"
-								data-testid="chat-message-input"
-								style={{ minHeight: "inherit" }}
-								aria-label={ariaLabel}
-								aria-disabled={disabled}
-							/>
+				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={
-							<div className="pointer-events-none text-content-secondary [&_p]:leading-normal">
-								{placeholder}
-							</div>
+
+						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}
-					/>
-					<HistoryPlugin />
-					<DisableFormattingPlugin />
-					<PasteSanitizationPlugin
-						onFilePaste={onFilePaste}
-						allowTextAttachmentPaste={allowTextAttachmentPaste}
-					/>
-					<EnterKeyPlugin onEnter={disabled ? undefined : onEnter} />
-					<ContentChangePlugin onChange={handleContentChange} />
-					<InsertTextPlugin onEditorReady={handleEditorReady} />
-					<EditableStatePlugin disabled={!!disabled} />
-					{autoFocus && <AutoFocusPlugin />}
-				</div>
-			</LexicalComposer>
-		);
-	},
-);
+					}
+				});
+
+				return parts as EditorContentPart[];
+			},
+		}),
+		[],
+	);
+
+	return (
+		<LexicalComposer initialConfig={initialConfig} key={remountKey}>
+			<div
+				className={cn(
+					"grid w-full rounded-md bg-transparent text-base placeholder:text-content-secondary focus-visible:outline-none whitespace-pre-wrap break-words [&>*]:col-start-1 [&>*]:row-start-1",
+					disabled && "cursor-not-allowed opacity-50",
+					className,
+				)}
+				style={style}
+				{...props}
+			>
+				<RichTextPlugin
+					contentEditable={
+						<ContentEditable
+							className="outline-none w-full whitespace-pre-wrap overflow-y-auto max-h-[50vh] [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent] [&_p]:leading-normal [&_p:first-child]:mt-0 [&_p:last-child]:mb-0 py-px"
+							data-testid="chat-message-input"
+							style={{ minHeight: "inherit" }}
+							aria-label={ariaLabel}
+							aria-disabled={disabled}
+						/>
+					}
+					placeholder={
+						<div className="pointer-events-none text-content-secondary [&_p]:leading-normal">
+							{placeholder}
+						</div>
+					}
+					ErrorBoundary={LexicalErrorBoundary}
+				/>
+				<HistoryPlugin />
+				<DisableFormattingPlugin />
+				<PasteSanitizationPlugin
+					onFilePaste={onFilePaste}
+					allowTextAttachmentPaste={allowTextAttachmentPaste}
+				/>
+				<EnterKeyPlugin onEnter={disabled ? undefined : onEnter} />
+				<ContentChangePlugin onChange={onChange} />
+				<ValueSyncPlugin
+					initialValue={initialValue}
+					initialEditorState={initialEditorState}
+				/>
+				<InsertTextPlugin onEditorReady={handleEditorReady} />
+				<EditableStatePlugin disabled={!!disabled} />
+				{autoFocus && <AutoFocusPlugin />}
+			</div>
+		</LexicalComposer>
+	);
+};
 ChatMessageInput.displayName = "ChatMessageInput";
 
 export { ChatMessageInput };

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<ChatMessageInputRef>;
-	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<ChatPageInputProps> = ({
 	modelSelectorPlaceholder,
 	isModelCatalogLoading = false,
 	inputRef,
-	initialInputValue,
+	initialValue,
+	initialEditorState,
+	remountKey,
 	onContentChange,
 	editingQueuedMessageID,
 	onStartQueueEdit,
@@ -315,7 +323,9 @@ export const ChatPageInput: FC<ChatPageInputProps> = ({
 			previewUrls={previewUrls}
 			textContents={textContents}
 			inputRef={inputRef}
-			initialValue={initialInputValue}
+			initialValue={initialValue}
+			initialEditorState={initialEditorState}
+			remountKey={remountKey}
 			onContentChange={onContentChange}
 			queuedMessages={queuedMessages}
 			onDeleteQueuedMessage={onDeleteQueuedMessage}

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();
+	});
+});

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<Record<string, unknown>>;
+	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, unknown> }): string {
+	if (!state.root) {
+		return "";
+	}
+	return extractTextFromNode(
+		state.root as Parameters<typeof extractTextFromNode>[0],
+	);
+}