Skip to content

Commit 48216d3

Browse files
daniel-lxsdaniel-lxs
committed
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
1 parent 08b8365 commit 48216d3

File tree

3 files changed

+273
-17
lines changed

3 files changed

+273
-17
lines changed

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

Lines changed: 45 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ import { QueuedMessages } from "./QueuedMessages"
6060
import DismissibleUpsell from "../common/DismissibleUpsell"
6161
import { useCloudUpsell } from "@src/hooks/useCloudUpsell"
6262
import { Cloud } from "lucide-react"
63+
import { DraftPersistenceProvider, useDraftPersistence } from "./hooks/useDraftPersistence"
6364

6465
export interface ChatViewProps {
6566
isHidden: boolean
@@ -75,11 +76,12 @@ export const MAX_IMAGES_PER_MESSAGE = 20 // This is the Anthropic limit.
7576

7677
const isMac = navigator.platform.toUpperCase().indexOf("MAC") >= 0
7778

78-
const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewProps> = (
79+
const ChatViewInner: React.ForwardRefRenderFunction<ChatViewRef, ChatViewProps> = (
7980
{ isHidden, showAnnouncement, hideAnnouncement },
8081
ref,
8182
) => {
8283
const isMountedRef = useRef(true)
84+
const { saveCurrentDraft, restoreDraft } = useDraftPersistence()
8385

8486
const [audioBaseUri] = useState(() => {
8587
const w = window as any
@@ -867,24 +869,36 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
867869
useEvent("message", handleMessage)
868870

869871
// Begin editing from a row (WhatsApp-style overlay)
870-
const handleBeginEdit = useCallback((message: ClineMessage) => {
871-
setEditingOverlay({
872-
ts: message.ts,
873-
text: message.text || "",
874-
images: message.images || [],
875-
})
876-
setInputValue(message.text || "")
877-
setSelectedImages(message.images || [])
878-
// Focus input when beginning edit
879-
setTimeout(() => textAreaRef.current?.focus(), 0)
880-
}, [])
872+
const handleBeginEdit = useCallback(
873+
(message: ClineMessage) => {
874+
// Save the current draft before starting edit
875+
saveCurrentDraft(inputValue)
876+
877+
setEditingOverlay({
878+
ts: message.ts,
879+
text: message.text || "",
880+
images: message.images || [],
881+
})
882+
setInputValue(message.text || "")
883+
setSelectedImages(message.images || [])
884+
// Focus input when beginning edit
885+
setTimeout(() => textAreaRef.current?.focus(), 0)
886+
},
887+
[inputValue, saveCurrentDraft],
888+
)
881889

882890
const handleCancelEditOverlay = useCallback(() => {
883891
setEditingOverlay(null)
884-
setInputValue("")
892+
// Restore the draft when canceling edit
893+
const restoredDraft = restoreDraft()
894+
if (restoredDraft !== null) {
895+
setInputValue(restoredDraft)
896+
} else {
897+
setInputValue("")
898+
}
885899
setSelectedImages([])
886900
setTimeout(() => textAreaRef.current?.focus(), 0)
887-
}, [])
901+
}, [restoreDraft])
888902

889903
const handleSubmitEdited = useCallback(() => {
890904
if (!editingOverlay) return
@@ -895,9 +909,15 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
895909
images: selectedImages,
896910
})
897911
setEditingOverlay(null)
898-
setInputValue("")
912+
// Restore the draft after saving edit
913+
const restoredDraft = restoreDraft()
914+
if (restoredDraft !== null) {
915+
setInputValue(restoredDraft)
916+
} else {
917+
setInputValue("")
918+
}
899919
setSelectedImages([])
900-
}, [editingOverlay, inputValue, selectedImages])
920+
}, [editingOverlay, inputValue, selectedImages, restoreDraft])
901921

902922
// NOTE: the VSCode window needs to be focused for this to work.
903923
useMount(() => textAreaRef.current?.focus())
@@ -2142,6 +2162,14 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
21422162
)
21432163
}
21442164

2145-
const ChatView = forwardRef(ChatViewComponent)
2165+
const ChatViewWithRef = forwardRef(ChatViewInner)
2166+
2167+
const ChatView: React.FC<ChatViewProps> = (props) => {
2168+
return (
2169+
<DraftPersistenceProvider>
2170+
<ChatViewWithRef {...props} />
2171+
</DraftPersistenceProvider>
2172+
)
2173+
}
21462174

21472175
export default ChatView

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

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import { render, screen, fireEvent, waitFor } from "@testing-library/react"
3+
import React from "react"
4+
import { DraftPersistenceProvider, useDraftPersistence } from "../hooks/useDraftPersistence"
5+
6+
// Test component to interact with the draft persistence context
7+
const TestComponent = () => {
8+
const { savedDraft, saveCurrentDraft, restoreDraft, clearDraft } = useDraftPersistence()
9+
const [localDraft, setLocalDraft] = React.useState("")
10+
11+
return (
12+
<div>
13+
<input type="text" value={localDraft} onChange={(e) => setLocalDraft(e.target.value)} data-testid="input" />
14+
<button onClick={() => saveCurrentDraft(localDraft)} data-testid="save">
15+
Save Draft
16+
</button>
17+
<button
18+
onClick={() => {
19+
const draft = restoreDraft()
20+
if (draft) setLocalDraft(draft)
21+
}}
22+
data-testid="restore">
23+
Restore Draft
24+
</button>
25+
<button onClick={() => clearDraft()} data-testid="clear">
26+
Clear Draft
27+
</button>
28+
<div data-testid="saved-draft">{savedDraft || "No draft"}</div>
29+
</div>
30+
)
31+
}
32+
33+
describe("DraftPersistence", () => {
34+
beforeEach(() => {
35+
vi.clearAllMocks()
36+
})
37+
38+
it("should save and restore a draft", async () => {
39+
render(
40+
<DraftPersistenceProvider>
41+
<TestComponent />
42+
</DraftPersistenceProvider>,
43+
)
44+
45+
const input = screen.getByTestId("input")
46+
const saveButton = screen.getByTestId("save")
47+
const restoreButton = screen.getByTestId("restore")
48+
const savedDraftDisplay = screen.getByTestId("saved-draft")
49+
50+
// Initially no draft
51+
expect(savedDraftDisplay.textContent).toBe("No draft")
52+
53+
// Type some text
54+
fireEvent.change(input, { target: { value: "My draft text" } })
55+
expect(input).toHaveValue("My draft text")
56+
57+
// Save the draft
58+
fireEvent.click(saveButton)
59+
await waitFor(() => {
60+
expect(savedDraftDisplay.textContent).toBe("My draft text")
61+
})
62+
63+
// Clear the input
64+
fireEvent.change(input, { target: { value: "" } })
65+
expect(input).toHaveValue("")
66+
67+
// Restore the draft
68+
fireEvent.click(restoreButton)
69+
expect(input).toHaveValue("My draft text")
70+
71+
// After restoring, the saved draft should be cleared
72+
await waitFor(() => {
73+
expect(savedDraftDisplay.textContent).toBe("No draft")
74+
})
75+
})
76+
77+
it("should clear a draft", async () => {
78+
render(
79+
<DraftPersistenceProvider>
80+
<TestComponent />
81+
</DraftPersistenceProvider>,
82+
)
83+
84+
const input = screen.getByTestId("input")
85+
const saveButton = screen.getByTestId("save")
86+
const clearButton = screen.getByTestId("clear")
87+
const savedDraftDisplay = screen.getByTestId("saved-draft")
88+
89+
// Save a draft
90+
fireEvent.change(input, { target: { value: "Draft to clear" } })
91+
fireEvent.click(saveButton)
92+
await waitFor(() => {
93+
expect(savedDraftDisplay.textContent).toBe("Draft to clear")
94+
})
95+
96+
// Clear the draft
97+
fireEvent.click(clearButton)
98+
await waitFor(() => {
99+
expect(savedDraftDisplay.textContent).toBe("No draft")
100+
})
101+
})
102+
103+
it("should handle multiple save operations", async () => {
104+
render(
105+
<DraftPersistenceProvider>
106+
<TestComponent />
107+
</DraftPersistenceProvider>,
108+
)
109+
110+
const input = screen.getByTestId("input")
111+
const saveButton = screen.getByTestId("save")
112+
const savedDraftDisplay = screen.getByTestId("saved-draft")
113+
114+
// Save first draft
115+
fireEvent.change(input, { target: { value: "First draft" } })
116+
fireEvent.click(saveButton)
117+
await waitFor(() => {
118+
expect(savedDraftDisplay.textContent).toBe("First draft")
119+
})
120+
121+
// Save second draft (overwrites first)
122+
fireEvent.change(input, { target: { value: "Second draft" } })
123+
fireEvent.click(saveButton)
124+
await waitFor(() => {
125+
expect(savedDraftDisplay.textContent).toBe("Second draft")
126+
})
127+
})
128+
129+
it("should return null when restoring with no saved draft", () => {
130+
render(
131+
<DraftPersistenceProvider>
132+
<TestComponent />
133+
</DraftPersistenceProvider>,
134+
)
135+
136+
const input = screen.getByTestId("input")
137+
const restoreButton = screen.getByTestId("restore")
138+
139+
// Try to restore when no draft is saved
140+
fireEvent.click(restoreButton)
141+
142+
// Input should remain empty
143+
expect(input).toHaveValue("")
144+
})
145+
146+
it("should provide no-op implementation when context is not available", () => {
147+
// Component using the hook outside of provider
148+
const ComponentWithoutProvider = () => {
149+
const { savedDraft, saveCurrentDraft, restoreDraft, clearDraft } = useDraftPersistence()
150+
151+
return (
152+
<div>
153+
<div data-testid="saved">{savedDraft || "null"}</div>
154+
<button onClick={() => saveCurrentDraft("test")} data-testid="save">
155+
Save
156+
</button>
157+
<button onClick={() => restoreDraft()} data-testid="restore">
158+
Restore
159+
</button>
160+
<button onClick={() => clearDraft()} data-testid="clear">
161+
Clear
162+
</button>
163+
</div>
164+
)
165+
}
166+
167+
render(<ComponentWithoutProvider />)
168+
169+
const saved = screen.getByTestId("saved")
170+
expect(saved.textContent).toBe("null")
171+
172+
// These should not throw errors even without provider
173+
fireEvent.click(screen.getByTestId("save"))
174+
fireEvent.click(screen.getByTestId("restore"))
175+
fireEvent.click(screen.getByTestId("clear"))
176+
177+
// State should remain unchanged
178+
expect(saved.textContent).toBe("null")
179+
})
180+
})

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import React, { createContext, useContext, useState, useCallback, ReactNode } from "react"
2+
3+
interface DraftPersistenceContextType {
4+
savedDraft: string | null
5+
saveCurrentDraft: (draft: string) => void
6+
restoreDraft: () => string | null
7+
clearDraft: () => void
8+
}
9+
10+
const DraftPersistenceContext = createContext<DraftPersistenceContextType | undefined>(undefined)
11+
12+
export const DraftPersistenceProvider: React.FC<{ children: ReactNode }> = ({ children }) => {
13+
const [savedDraft, setSavedDraft] = useState<string | null>(null)
14+
15+
const saveCurrentDraft = useCallback((draft: string) => {
16+
setSavedDraft(draft)
17+
}, [])
18+
19+
const restoreDraft = useCallback(() => {
20+
const draft = savedDraft
21+
setSavedDraft(null) // Clear after restoring
22+
return draft
23+
}, [savedDraft])
24+
25+
const clearDraft = useCallback(() => {
26+
setSavedDraft(null)
27+
}, [])
28+
29+
return (
30+
<DraftPersistenceContext.Provider value={{ savedDraft, saveCurrentDraft, restoreDraft, clearDraft }}>
31+
{children}
32+
</DraftPersistenceContext.Provider>
33+
)
34+
}
35+
36+
export const useDraftPersistence = () => {
37+
const context = useContext(DraftPersistenceContext)
38+
if (!context) {
39+
// Return a no-op implementation if context is not available
40+
return {
41+
savedDraft: null,
42+
saveCurrentDraft: () => {},
43+
restoreDraft: () => null,
44+
clearDraft: () => {},
45+
}
46+
}
47+
return context
48+
}

0 commit comments

Comments
 (0)