feat: persist chat input text across VSCode restarts

- Add draft message storage key and methods to ClineProvider
- Add saveDraftMessage, getDraftMessage, clearDraftMessage WebviewMessage types
- Add draftMessage ExtensionMessage type
- Add webview message handlers for draft persistence
- Update ChatView.tsx with debounced save/restore logic (500ms debounce)
- Clear draft on task switch, message send, and chat reset
- Draft persists at workspace level (each workspace maintains its own draft)
- Add unit tests for draft message handlers
This commit is contained in:
Hannes Rudolph 2025-12-02 18:27:39 -07:00
parent 74d1ed7276
commit 50d40f260b
6 changed files with 264 additions and 1 deletions

View file

@ -145,6 +145,7 @@ export class ClineProvider
private recentTasksCache?: string[]
private pendingOperations: Map<string, PendingEditOperation> = new Map()
private static readonly PENDING_OPERATION_TIMEOUT_MS = 30000 // 30 seconds
private static readonly DRAFT_MESSAGE_KEY = "roo.draftMessage"
private cloudOrganizationsCache: CloudOrganizationMembership[] | null = null
private cloudOrganizationsCacheTimestamp: number | null = null
@ -550,6 +551,43 @@ export class ClineProvider
this.log(`[clearAllPendingEditOperations] Cleared all pending operations`)
}
// Draft Message Management
/**
* Save draft message to workspace state
* @param text - The draft message text
* @param images - Array of base64 image data URLs
*/
public async saveDraftMessage(text: string, images: string[]): Promise<void> {
// Don't save empty drafts
if (!text.trim() && images.length === 0) {
await this.clearDraftMessage()
return
}
const draft = {
text,
images,
timestamp: Date.now(),
}
await this.context.workspaceState.update(ClineProvider.DRAFT_MESSAGE_KEY, draft)
}
/**
* Get draft message from workspace state
* @returns The saved draft or undefined
*/
public getDraftMessage(): { text: string; images: string[]; timestamp: number } | undefined {
return this.context.workspaceState.get(ClineProvider.DRAFT_MESSAGE_KEY)
}
/**
* Clear draft message from workspace state
*/
public async clearDraftMessage(): Promise<void> {
await this.context.workspaceState.update(ClineProvider.DRAFT_MESSAGE_KEY, undefined)
}
/*
VSCode extensions use the disposable pattern to clean up resources when the sidebar/editor tab is closed by the user or system. This applies to event listening, commands, interacting with the UI, etc.
- https://vscode-docs.readthedocs.io/en/stable/extensions/patterns-and-principles/
@ -1622,6 +1660,9 @@ export class ClineProvider
}
async showTaskWithId(id: string) {
// Clear draft when switching tasks
await this.clearDraftMessage()
if (id !== this.getCurrentTask()?.taskId) {
// Non-current task.
const { historyItem } = await this.getTaskWithId(id)

View file

@ -0,0 +1,150 @@
// npx vitest run src/core/webview/__tests__/webviewMessageHandler.draftMessage.spec.ts
import { webviewMessageHandler } from "../webviewMessageHandler"
import type { ClineProvider } from "../ClineProvider"
vi.mock("vscode", () => ({
window: {
showInformationMessage: vi.fn(),
showErrorMessage: vi.fn(),
},
workspace: {
workspaceFolders: [{ uri: { fsPath: "/mock/workspace" } }],
},
}))
describe("webviewMessageHandler - Draft Message", () => {
let mockClineProvider: ClineProvider
beforeEach(() => {
vi.clearAllMocks()
mockClineProvider = {
saveDraftMessage: vi.fn().mockResolvedValue(undefined),
getDraftMessage: vi.fn(),
clearDraftMessage: vi.fn().mockResolvedValue(undefined),
postMessageToWebview: vi.fn(),
getState: vi.fn().mockResolvedValue({
apiConfiguration: {},
}),
contextProxy: {
context: {
extensionPath: "/mock/extension/path",
globalStorageUri: { fsPath: "/mock/global/storage" },
},
setValue: vi.fn(),
getValue: vi.fn(),
},
log: vi.fn(),
} as unknown as ClineProvider
})
describe("saveDraftMessage", () => {
it("should save draft message with text and images", async () => {
await webviewMessageHandler(mockClineProvider, {
type: "saveDraftMessage",
text: "Test draft message",
images: ["data:image/png;base64,abc123"],
})
expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("Test draft message", [
"data:image/png;base64,abc123",
])
})
it("should save draft message with only text", async () => {
await webviewMessageHandler(mockClineProvider, {
type: "saveDraftMessage",
text: "Text only draft",
})
expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("Text only draft", [])
})
it("should save draft message with only images", async () => {
await webviewMessageHandler(mockClineProvider, {
type: "saveDraftMessage",
images: ["data:image/png;base64,image1", "data:image/png;base64,image2"],
})
expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("", [
"data:image/png;base64,image1",
"data:image/png;base64,image2",
])
})
it("should handle empty text and images", async () => {
await webviewMessageHandler(mockClineProvider, {
type: "saveDraftMessage",
})
expect(mockClineProvider.saveDraftMessage).toHaveBeenCalledWith("", [])
})
})
describe("getDraftMessage", () => {
it("should return saved draft with text and images", async () => {
const mockDraft = {
text: "Saved draft",
images: ["data:image/png;base64,savedImage"],
timestamp: 1234567890,
}
;(mockClineProvider.getDraftMessage as ReturnType<typeof vi.fn>).mockReturnValue(mockDraft)
await webviewMessageHandler(mockClineProvider, {
type: "getDraftMessage",
})
expect(mockClineProvider.getDraftMessage).toHaveBeenCalled()
expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({
type: "draftMessage",
text: "Saved draft",
images: ["data:image/png;base64,savedImage"],
})
})
it("should handle when no draft is saved", async () => {
;(mockClineProvider.getDraftMessage as ReturnType<typeof vi.fn>).mockReturnValue(undefined)
await webviewMessageHandler(mockClineProvider, {
type: "getDraftMessage",
})
expect(mockClineProvider.getDraftMessage).toHaveBeenCalled()
expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({
type: "draftMessage",
text: undefined,
images: undefined,
})
})
it("should handle draft with only text", async () => {
const mockDraft = {
text: "Text only",
images: [],
timestamp: 1234567890,
}
;(mockClineProvider.getDraftMessage as ReturnType<typeof vi.fn>).mockReturnValue(mockDraft)
await webviewMessageHandler(mockClineProvider, {
type: "getDraftMessage",
})
expect(mockClineProvider.postMessageToWebview).toHaveBeenCalledWith({
type: "draftMessage",
text: "Text only",
images: [],
})
})
})
describe("clearDraftMessage", () => {
it("should clear the draft message", async () => {
await webviewMessageHandler(mockClineProvider, {
type: "clearDraftMessage",
})
expect(mockClineProvider.clearDraftMessage).toHaveBeenCalled()
})
})
})

View file

@ -3100,6 +3100,28 @@ export const webviewMessageHandler = async (
break
}
/**
* Draft Message Persistence
*/
case "saveDraftMessage":
await provider.saveDraftMessage(message.text || "", message.images || [])
break
case "getDraftMessage": {
const draft = provider.getDraftMessage()
await provider.postMessageToWebview({
type: "draftMessage",
text: draft?.text,
images: draft?.images,
})
break
}
case "clearDraftMessage":
await provider.clearDraftMessage()
break
default: {
// console.log(`Unhandled message type: ${message.type}`)
//

View file

@ -131,6 +131,7 @@ export interface ExtensionMessage {
| "interactionRequired"
| "browserSessionUpdate"
| "browserSessionNavigate"
| "draftMessage"
text?: string
payload?: any // Add a generic payload for now, can refine later
// Checkpoint warning message

View file

@ -175,6 +175,9 @@ export interface WebviewMessage {
| "browserPanelDidLaunch"
| "openDebugApiHistory"
| "openDebugUiHistory"
| "saveDraftMessage"
| "getDraftMessage"
| "clearDraftMessage"
text?: string
editedMessageContent?: string
tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "cloud"

View file

@ -137,6 +137,29 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
const [sendingDisabled, setSendingDisabled] = useState(false)
const [selectedImages, setSelectedImages] = useState<string[]>([])
// Debounced draft save to prevent excessive writes
const saveDraftDebounced = useMemo(
() =>
debounce((text: string, images: string[]) => {
vscode.postMessage({ type: "saveDraftMessage", text, images })
}, 500),
[],
)
// Cleanup debounce on unmount
useEffect(() => {
return () => {
if (saveDraftDebounced && typeof (saveDraftDebounced as any).clear === "function") {
;(saveDraftDebounced as any).clear()
}
}
}, [saveDraftDebounced])
// Request saved draft on mount
useEffect(() => {
vscode.postMessage({ type: "getDraftMessage" })
}, [])
// We need to hold on to the ask because useEffect > lastMessage will always
// let us know when an ask comes in and handle it, but by the time
// handleMessage is called, the last message might not be the ask anymore
@ -557,8 +580,20 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
// setPrimaryButtonText(undefined)
// setSecondaryButtonText(undefined)
disableAutoScrollRef.current = false
// Clear saved draft
vscode.postMessage({ type: "clearDraftMessage" })
}, [])
// Wrapper to save draft when input value changes
const handleInputValueChange = useCallback(
(value: string) => {
setInputValue(value)
saveDraftDebounced(value, selectedImages)
},
[selectedImages, saveDraftDebounced],
)
/**
* Handles sending messages to the extension
* @param text - The message text to send
@ -840,6 +875,17 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
case "interactionRequired":
playSound("notification")
break
case "draftMessage":
// Only restore draft if there's no active task input
if (messagesRef.current.length === 0 || clineAskRef.current === undefined) {
if (message.text) {
setInputValue(message.text)
}
if (message.images && message.images.length > 0) {
setSelectedImages(message.images)
}
}
break
}
// textAreaRef.current is not explicitly required here since React
// guarantees that ref will be stable across re-renders, and we're
@ -1568,7 +1614,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
<ChatTextArea
ref={textAreaRef}
inputValue={inputValue}
setInputValue={setInputValue}
setInputValue={handleInputValueChange}
sendingDisabled={sendingDisabled || isProfileDisabled}
selectApiConfigDisabled={sendingDisabled && clineAsk !== "api_req_failed"}
placeholderText={placeholderText}