From aebd78e07993b66fb974c08f92a7cac9b7833cb4 Mon Sep 17 00:00:00 2001
From: Roo Code <roomote@roocode.com>
Date: Mon, 22 Sep 2025 23:37:05 +0000
Subject: [PATCH] feat: implement chat draft persistence

- Add chatDraft field to GlobalState type in global-settings.ts
- Implement draft saving/loading in ChatView component with 500ms debounce
- Add message handlers for draft persistence in webviewMessageHandler
- Add new WebviewMessage and ExtensionMessage types for draft operations
- Clear draft when task starts to avoid interference
- Persist both text and images in draft

Fixes #8236
---
 packages/types/src/global-settings.ts       |  8 +++++
 src/core/webview/webviewMessageHandler.ts   | 22 +++++++++++++
 src/shared/ExtensionMessage.ts              |  1 +
 src/shared/WebviewMessage.ts                |  3 ++
 webview-ui/src/components/chat/ChatView.tsx | 36 +++++++++++++++++++++
 5 files changed, 70 insertions(+)

diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts
index 7e79855f7e..15b596c7b1 100644
--- a/packages/types/src/global-settings.ts
+++ b/packages/types/src/global-settings.ts
@@ -43,6 +43,14 @@ export const globalSettingsSchema = z.object({
 	taskHistory: z.array(historyItemSchema).optional(),
 	dismissedUpsells: z.array(z.string()).optional(),
 
+	// Chat draft persistence
+	chatDraft: z
+		.object({
+			text: z.string(),
+			images: z.array(z.string()),
+		})
+		.optional(),
+
 	// Image generation settings (experimental) - flattened for simplicity
 	openRouterImageApiKey: z.string().optional(),
 	openRouterImageGenerationSelectedModel: z.string().optional(),
diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts
index 6274694da2..1ffeff7cc5 100644
--- a/src/core/webview/webviewMessageHandler.ts
+++ b/src/core/webview/webviewMessageHandler.ts
@@ -3058,5 +3058,27 @@ export const webviewMessageHandler = async (
 			})
 			break
 		}
+		case "persistDraft": {
+			// Save the draft to global state
+			await updateGlobalState("chatDraft", { text: message.text || "", images: message.images || [] })
+			break
+		}
+		case "clearPersistedDraft": {
+			// Clear the persisted draft
+			await updateGlobalState("chatDraft", undefined)
+			break
+		}
+		case "requestPersistedDraft": {
+			// Send the persisted draft to the webview
+			const draft = getGlobalState("chatDraft")
+			if (draft) {
+				await provider.postMessageToWebview({
+					type: "persistedDraft",
+					text: draft.text,
+					images: draft.images,
+				})
+			}
+			break
+		}
 	}
 }
diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts
index 7772e8d817..5c2f30b92a 100644
--- a/src/shared/ExtensionMessage.ts
+++ b/src/shared/ExtensionMessage.ts
@@ -124,6 +124,7 @@ export interface ExtensionMessage {
 		| "commands"
 		| "insertTextIntoTextarea"
 		| "dismissedUpsells"
+		| "persistedDraft"
 	text?: string
 	payload?: any // Add a generic payload for now, can refine later
 	action?:
diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts
index 541b445fb2..be3eb29159 100644
--- a/src/shared/WebviewMessage.ts
+++ b/src/shared/WebviewMessage.ts
@@ -227,6 +227,9 @@ export interface WebviewMessage {
 		| "editQueuedMessage"
 		| "dismissUpsell"
 		| "getDismissedUpsells"
+		| "persistDraft"
+		| "clearPersistedDraft"
+		| "requestPersistedDraft"
 	text?: string
 	editedMessageContent?: string
 	tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "cloud"
diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx
index d358c68f1c..bfd5acbe55 100644
--- a/webview-ui/src/components/chat/ChatView.tsx
+++ b/webview-ui/src/components/chat/ChatView.tsx
@@ -171,6 +171,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 	// Has to be after api_req_finished are all reduced into api_req_started messages.
 	const apiMetrics = useMemo(() => getApiMetrics(modifiedMessages), [modifiedMessages])
 
+	// Initialize input value from persisted draft if available
 	const [inputValue, setInputValue] = useState("")
 	const inputValueRef = useRef(inputValue)
 	const textAreaRef = useRef<HTMLTextAreaElement>(null)
@@ -224,6 +225,28 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		inputValueRef.current = inputValue
 	}, [inputValue])
 
+	// Load persisted draft on mount
+	useEffect(() => {
+		// Request the persisted draft from the extension
+		vscode.postMessage({ type: "requestPersistedDraft" })
+	}, [])
+
+	// Save draft when input changes (debounced)
+	useEffect(() => {
+		// Don't save empty drafts or when there's an active task
+		if (!task && (inputValue.trim() || selectedImages.length > 0)) {
+			const timer = setTimeout(() => {
+				vscode.postMessage({
+					type: "persistDraft",
+					text: inputValue,
+					images: selectedImages,
+				})
+			}, 500) // Debounce for 500ms
+
+			return () => clearTimeout(timer)
+		}
+	}, [inputValue, selectedImages, task])
+
 	useEffect(() => {
 		isMountedRef.current = true
 		return () => {
@@ -582,6 +605,9 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 		// setPrimaryButtonText(undefined)
 		// setSecondaryButtonText(undefined)
 		disableAutoScrollRef.current = false
+
+		// Clear the persisted draft when chat is reset (task started)
+		vscode.postMessage({ type: "clearPersistedDraft" })
 	}, [])
 
 	/**
@@ -793,6 +819,15 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 							break
 					}
 					break
+				case "persistedDraft":
+					// Restore the persisted draft if we don't have an active task
+					if (!task && message.text !== undefined) {
+						setInputValue(message.text)
+						if (message.images && message.images.length > 0) {
+							setSelectedImages(message.images)
+						}
+					}
+					break
 				case "selectedImages":
 					// Only handle selectedImages if it's not for editing context
 					// When context is "edit", ChatRow will handle the images
@@ -845,6 +880,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
 			handleSetChatBoxMessage,
 			handlePrimaryButtonClick,
 			handleSecondaryButtonClick,
+			task,
 		],
 	)