feat: add draft persistence for chat input during message editing

- Created DraftPersistenceProvider context to manage draft state
- Integrated draft saving when starting message edit
- Integrated draft restoration when canceling or saving edit
- Added automatic cleanup after restoration to prevent memory leaks
- Added comprehensive test coverage

webview-ui/src/components/chat/ChatView.tsx

@@ -60,6 +60,7 @@ import { QueuedMessages } from "./QueuedMessages"
 import DismissibleUpsell from "../common/DismissibleUpsell"
 import { useCloudUpsell } from "@src/hooks/useCloudUpsell"
 import { Cloud } from "lucide-react"
+import { DraftPersistenceProvider, useDraftPersistence } from "./hooks/useDraftPersistence"
 
 export interface ChatViewProps {
 	isHidden: boolean
@@ -75,11 +76,12 @@ export const MAX_IMAGES_PER_MESSAGE = 20 // This is the Anthropic limit.
 
 const isMac = navigator.platform.toUpperCase().indexOf("MAC") >= 0
 
-const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewProps> = (
+const ChatViewInner: React.ForwardRefRenderFunction<ChatViewRef, ChatViewProps> = (
 	{ isHidden, showAnnouncement, hideAnnouncement },
 	ref,
 ) => {
 	const isMountedRef = useRef(true)
+	const { saveCurrentDraft, restoreDraft } = useDraftPersistence()
 
 	const [audioBaseUri] = useState(() => {
 		const w = window as any
@@ -867,24 +869,36 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	useEvent("message", handleMessage)
 
 	// Begin editing from a row (WhatsApp-style overlay)
-	const handleBeginEdit = useCallback((message: ClineMessage) => {
-		setEditingOverlay({
-			ts: message.ts,
-			text: message.text || "",
-			images: message.images || [],
-		})
-		setInputValue(message.text || "")
-		setSelectedImages(message.images || [])
-		// Focus input when beginning edit
-		setTimeout(() => textAreaRef.current?.focus(), 0)
-	}, [])
+	const handleBeginEdit = useCallback(
+		(message: ClineMessage) => {
+			// Save the current draft before starting edit
+			saveCurrentDraft(inputValue)
+
+			setEditingOverlay({
+				ts: message.ts,
+				text: message.text || "",
+				images: message.images || [],
+			})
+			setInputValue(message.text || "")
+			setSelectedImages(message.images || [])
+			// Focus input when beginning edit
+			setTimeout(() => textAreaRef.current?.focus(), 0)
+		},
+		[inputValue, saveCurrentDraft],
+	)
 
 	const handleCancelEditOverlay = useCallback(() => {
 		setEditingOverlay(null)
-		setInputValue("")
+		// Restore the draft when canceling edit
+		const restoredDraft = restoreDraft()
+		if (restoredDraft !== null) {
+			setInputValue(restoredDraft)
+		} else {
+			setInputValue("")
+		}
 		setSelectedImages([])
 		setTimeout(() => textAreaRef.current?.focus(), 0)
-	}, [])
+	}, [restoreDraft])
 
 	const handleSubmitEdited = useCallback(() => {
 		if (!editingOverlay) return
@@ -895,9 +909,15 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			images: selectedImages,
 		})
 		setEditingOverlay(null)
-		setInputValue("")
+		// Restore the draft after saving edit
+		const restoredDraft = restoreDraft()
+		if (restoredDraft !== null) {
+			setInputValue(restoredDraft)
+		} else {
+			setInputValue("")
+		}
 		setSelectedImages([])
-	}, [editingOverlay, inputValue, selectedImages])
+	}, [editingOverlay, inputValue, selectedImages, restoreDraft])
 
 	// NOTE: the VSCode window needs to be focused for this to work.
 	useMount(() => textAreaRef.current?.focus())
@@ -2142,6 +2162,14 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	)
 }
 
-const ChatView = forwardRef(ChatViewComponent)
+const ChatViewWithRef = forwardRef(ChatViewInner)
+
+const ChatView: React.FC<ChatViewProps> = (props) => {
+	return (
+		<DraftPersistenceProvider>
+			<ChatViewWithRef {...props} />
+		</DraftPersistenceProvider>
+	)
+}
 
 export default ChatView

webview-ui/src/components/chat/__tests__/DraftPersistence.spec.tsx

@@ -0,0 +1,180 @@
+import { describe, it, expect, vi, beforeEach } from "vitest"
+import { render, screen, fireEvent, waitFor } from "@testing-library/react"
+import React from "react"
+import { DraftPersistenceProvider, useDraftPersistence } from "../hooks/useDraftPersistence"
+
+// Test component to interact with the draft persistence context
+const TestComponent = () => {
+	const { savedDraft, saveCurrentDraft, restoreDraft, clearDraft } = useDraftPersistence()
+	const [localDraft, setLocalDraft] = React.useState("")
+
+	return (
+		<div>
+			<input type="text" value={localDraft} onChange={(e) => setLocalDraft(e.target.value)} data-testid="input" />
+			<button onClick={() => saveCurrentDraft(localDraft)} data-testid="save">
+				Save Draft
+			</button>
+			<button
+				onClick={() => {
+					const draft = restoreDraft()
+					if (draft) setLocalDraft(draft)
+				}}
+				data-testid="restore">
+				Restore Draft
+			</button>
+			<button onClick={() => clearDraft()} data-testid="clear">
+				Clear Draft
+			</button>
+			<div data-testid="saved-draft">{savedDraft || "No draft"}</div>
+		</div>
+	)
+}
+
+describe("DraftPersistence", () => {
+	beforeEach(() => {
+		vi.clearAllMocks()
+	})
+
+	it("should save and restore a draft", async () => {
+		render(
+			<DraftPersistenceProvider>
+				<TestComponent />
+			</DraftPersistenceProvider>,
+		)
+
+		const input = screen.getByTestId("input")
+		const saveButton = screen.getByTestId("save")
+		const restoreButton = screen.getByTestId("restore")
+		const savedDraftDisplay = screen.getByTestId("saved-draft")
+
+		// Initially no draft
+		expect(savedDraftDisplay.textContent).toBe("No draft")
+
+		// Type some text
+		fireEvent.change(input, { target: { value: "My draft text" } })
+		expect(input).toHaveValue("My draft text")
+
+		// Save the draft
+		fireEvent.click(saveButton)
+		await waitFor(() => {
+			expect(savedDraftDisplay.textContent).toBe("My draft text")
+		})
+
+		// Clear the input
+		fireEvent.change(input, { target: { value: "" } })
+		expect(input).toHaveValue("")
+
+		// Restore the draft
+		fireEvent.click(restoreButton)
+		expect(input).toHaveValue("My draft text")
+
+		// After restoring, the saved draft should be cleared
+		await waitFor(() => {
+			expect(savedDraftDisplay.textContent).toBe("No draft")
+		})
+	})
+
+	it("should clear a draft", async () => {
+		render(
+			<DraftPersistenceProvider>
+				<TestComponent />
+			</DraftPersistenceProvider>,
+		)
+
+		const input = screen.getByTestId("input")
+		const saveButton = screen.getByTestId("save")
+		const clearButton = screen.getByTestId("clear")
+		const savedDraftDisplay = screen.getByTestId("saved-draft")
+
+		// Save a draft
+		fireEvent.change(input, { target: { value: "Draft to clear" } })
+		fireEvent.click(saveButton)
+		await waitFor(() => {
+			expect(savedDraftDisplay.textContent).toBe("Draft to clear")
+		})
+
+		// Clear the draft
+		fireEvent.click(clearButton)
+		await waitFor(() => {
+			expect(savedDraftDisplay.textContent).toBe("No draft")
+		})
+	})
+
+	it("should handle multiple save operations", async () => {
+		render(
+			<DraftPersistenceProvider>
+				<TestComponent />
+			</DraftPersistenceProvider>,
+		)
+
+		const input = screen.getByTestId("input")
+		const saveButton = screen.getByTestId("save")
+		const savedDraftDisplay = screen.getByTestId("saved-draft")
+
+		// Save first draft
+		fireEvent.change(input, { target: { value: "First draft" } })
+		fireEvent.click(saveButton)
+		await waitFor(() => {
+			expect(savedDraftDisplay.textContent).toBe("First draft")
+		})
+
+		// Save second draft (overwrites first)
+		fireEvent.change(input, { target: { value: "Second draft" } })
+		fireEvent.click(saveButton)
+		await waitFor(() => {
+			expect(savedDraftDisplay.textContent).toBe("Second draft")
+		})
+	})
+
+	it("should return null when restoring with no saved draft", () => {
+		render(
+			<DraftPersistenceProvider>
+				<TestComponent />
+			</DraftPersistenceProvider>,
+		)
+
+		const input = screen.getByTestId("input")
+		const restoreButton = screen.getByTestId("restore")
+
+		// Try to restore when no draft is saved
+		fireEvent.click(restoreButton)
+
+		// Input should remain empty
+		expect(input).toHaveValue("")
+	})
+
+	it("should provide no-op implementation when context is not available", () => {
+		// Component using the hook outside of provider
+		const ComponentWithoutProvider = () => {
+			const { savedDraft, saveCurrentDraft, restoreDraft, clearDraft } = useDraftPersistence()
+
+			return (
+				<div>
+					<div data-testid="saved">{savedDraft || "null"}</div>
+					<button onClick={() => saveCurrentDraft("test")} data-testid="save">
+						Save
+					</button>
+					<button onClick={() => restoreDraft()} data-testid="restore">
+						Restore
+					</button>
+					<button onClick={() => clearDraft()} data-testid="clear">
+						Clear
+					</button>
+				</div>
+			)
+		}
+
+		render(<ComponentWithoutProvider />)
+
+		const saved = screen.getByTestId("saved")
+		expect(saved.textContent).toBe("null")
+
+		// These should not throw errors even without provider
+		fireEvent.click(screen.getByTestId("save"))
+		fireEvent.click(screen.getByTestId("restore"))
+		fireEvent.click(screen.getByTestId("clear"))
+
+		// State should remain unchanged
+		expect(saved.textContent).toBe("null")
+	})
+})

webview-ui/src/components/chat/hooks/useDraftPersistence.tsx

@@ -0,0 +1,48 @@
+import React, { createContext, useContext, useState, useCallback, ReactNode } from "react"
+
+interface DraftPersistenceContextType {
+	savedDraft: string | null
+	saveCurrentDraft: (draft: string) => void
+	restoreDraft: () => string | null
+	clearDraft: () => void
+}
+
+const DraftPersistenceContext = createContext<DraftPersistenceContextType | undefined>(undefined)
+
+export const DraftPersistenceProvider: React.FC<{ children: ReactNode }> = ({ children }) => {
+	const [savedDraft, setSavedDraft] = useState<string | null>(null)
+
+	const saveCurrentDraft = useCallback((draft: string) => {
+		setSavedDraft(draft)
+	}, [])
+
+	const restoreDraft = useCallback(() => {
+		const draft = savedDraft
+		setSavedDraft(null) // Clear after restoring
+		return draft
+	}, [savedDraft])
+
+	const clearDraft = useCallback(() => {
+		setSavedDraft(null)
+	}, [])
+
+	return (
+		<DraftPersistenceContext.Provider value={{ savedDraft, saveCurrentDraft, restoreDraft, clearDraft }}>
+			{children}
+		</DraftPersistenceContext.Provider>
+	)
+}
+
+export const useDraftPersistence = () => {
+	const context = useContext(DraftPersistenceContext)
+	if (!context) {
+		// Return a no-op implementation if context is not available
+		return {
+			savedDraft: null,
+			saveCurrentDraft: () => {},
+			restoreDraft: () => null,
+			clearDraft: () => {},
+		}
+	}
+	return context
+}