diff --git a/ui/src/components/IssueChatThread.test.tsx b/ui/src/components/IssueChatThread.test.tsx index 38a07c97c48..62e1e686d11 100644 --- a/ui/src/components/IssueChatThread.test.tsx +++ b/ui/src/components/IssueChatThread.test.tsx @@ -3405,8 +3405,13 @@ describe("IssueChatThread", () => { ); let root = createRoot(container); + const storagePrototype = Object.getPrototypeOf(sessionStorage) as Storage; + const originalStorageWrite = storagePrototype.setItem; const recoveryWrite = scenario === "unavailable recovery storage" - ? vi.spyOn(Object.getPrototypeOf(sessionStorage), "setItem").mockImplementation(() => { throw new Error("Storage full"); }) + ? vi.spyOn(storagePrototype, "setItem").mockImplementation(function (this: Storage, key: string, value: string) { + if (this === sessionStorage) throw new Error("Storage full"); + return originalStorageWrite.call(this, key, value); + }) : null; try { await act(async () => root.render(element(false))); @@ -3476,6 +3481,31 @@ describe("IssueChatThread", () => { } finally { await act(async () => root.unmount()); } }); + it("returns to the shared draft after finishing a recovery and leaving the task", async () => { + const a = "finished-recovery-a"; + const b = "finished-recovery-b"; + localStorage.setItem(a, "The other tab's preserved draft"); + preserveDraftInTab(a, "My recovered message", []); + const onAdd = vi.fn().mockResolvedValue(undefined); + const root = createRoot(container); + const element = (draftKey: string) => ( + + ); + const editor = () => container.querySelector('textarea[aria-label="Issue chat editor"]')!; + try { + await act(async () => root.render(element(a))); + expect(editor().value).toBe("My recovered message"); + await act(async () => (Array.from(container.querySelectorAll("button")).find(button => button.textContent === "Send") as HTMLButtonElement).click()); + expect(onAdd.mock.calls[0]![0]).toBe("My recovered message"); + expect(editor().value).toBe(""); + await act(async () => root.render(element(b))); + await act(async () => root.render(element(a))); + expect(editor().value).toBe("The other tab's preserved draft"); + expect(container.textContent).not.toContain("This draft is kept separately"); + } finally { await act(async () => root.unmount()); } + }); + it("stores and restores the composer draft per issue key", () => { vi.useFakeTimers(); const root = createRoot(container); diff --git a/ui/src/components/IssueChatThread.tsx b/ui/src/components/IssueChatThread.tsx index 6db7961bd87..1b7f8418e72 100644 --- a/ui/src/components/IssueChatThread.tsx +++ b/ui/src/components/IssueChatThread.tsx @@ -4682,6 +4682,9 @@ const IssueChatComposer = forwardRef< const [newRecovery, setDraftRecovery] = useState<{ sourceKey: string; key: string; persisted: boolean } | null>(null); const draftRecovery = newRecovery?.sourceKey === sharedDraftKey ? newRecovery : restoredRecovery; const draftKey = draftRecovery?.key ?? sharedDraftKey; + // Keep the active buffer through send completion. Re-entering a task may + // return to its shared draft after an empty recovery has been retired. + useEffect(() => setDraftRecovery(restoredRecovery), [sharedDraftKey, restoredRecovery]); const retiredDraftKeyRef = useRef(undefined); // Initialize before StrictMode's mount cleanup can flush an empty value over // the stored draft. The effect below handles subsequent task-key changes. diff --git a/ui/src/lib/composer-draft.test.ts b/ui/src/lib/composer-draft.test.ts index 26f83e16802..e03cdfb7240 100644 --- a/ui/src/lib/composer-draft.test.ts +++ b/ui/src/lib/composer-draft.test.ts @@ -45,6 +45,28 @@ describe("task draft upload receipts", () => { expect(loadDraftSubmission(key)?.attemptId).toBe(id); }); + it("retires a finished recovery but keeps or restores its mapping for unsent work", () => { + sessionStorage.clear(); + saveDraft(key, "Shared draft stays available"); + const fork = preserveDraftInTab(key, "Recovered message", []); + saveDraftSubmission(fork.key, { attemptId: id, reviewed: false }); + expect(settleDraftSubmission(fork.key, id)).toBe(true); + expect(loadDraftRecoveryKey(key)).toBeNull(); + expect(loadDraft(key)).toBe("Shared draft stays available"); + + saveDraft(fork.key, "Next recovered draft"); + expect(loadDraftRecoveryKey(key)).toBe(fork.key); + saveDraftSubmission(fork.key, { attemptId: id, reviewed: false, nextDraftOffset: 0, submittedAttachmentIds: [] }); + expect(settleDraftSubmission(fork.key, id)).toBe(true); + expect(loadDraftRecoveryKey(key)).toBe(fork.key); + expect(loadDraft(fork.key)).toBe("Next recovered draft"); + clearDraft(fork.key); + expect(loadDraftRecoveryKey(key)).toBeNull(); + saveDraftAttachments(fork.key, [receipt]); + expect(loadDraftRecoveryKey(key)).toBe(fork.key); + expect(loadDraftAttachments(fork.key)).toEqual([receipt]); + }); + it("settles only submitted text and attachments while preserving the next draft", () => { const nextId = "aaf8228f-0be7-45ae-a104-6fbe0af6f1d3"; const nextReceipt = { ...receipt, attachmentId: nextId, contentPath: `/api/attachments/${nextId}/content` }; diff --git a/ui/src/lib/composer-draft.ts b/ui/src/lib/composer-draft.ts index 35c82b56ce5..2ca07566def 100644 --- a/ui/src/lib/composer-draft.ts +++ b/ui/src/lib/composer-draft.ts @@ -51,6 +51,17 @@ export function preserveDraftInTab(draftKey: string, body: string, attachments: } } +function syncDraftRecoveryMarker(draftKey: string) { + const prefix = "paperclip:agent-chat-draft:recovered:"; + if (!draftKey.startsWith(prefix)) return; + const marker = `${draftKey}:recovery:v1`; + if (loadDraft(draftKey).trim() || loadDraftAttachments(draftKey).length || loadDraftSubmission(draftKey)) { + sessionStorage.setItem(marker, draftKey.slice(prefix.length)); + } else { + sessionStorage.removeItem(marker); + } +} + function mayWriteDraft(draftKey: string, attemptId?: string) { const pending = loadDraftSubmission(draftKey); return !pending || pending.attemptId === attemptId; @@ -64,6 +75,7 @@ export function saveDraft(draftKey: string, value: string, attemptId?: string) { } else { draftStorage(draftKey).removeItem(draftKey); } + syncDraftRecoveryMarker(draftKey); } catch { // Ignore browser storage failures. } @@ -75,6 +87,7 @@ export function clearDraft(draftKey: string, attemptId?: string) { draftStorage(draftKey).removeItem(draftKey); draftStorage(draftKey).removeItem(`${draftKey}:attachments:v1`); draftStorage(draftKey).removeItem(`${draftKey}:submission:v1`); + syncDraftRecoveryMarker(draftKey); } catch { // Ignore browser storage failures. } @@ -127,6 +140,7 @@ export function saveDraftSubmission( `${draftKey}:submission:v1`, JSON.stringify({ version: 1, draftKey, ...submission }), ); + syncDraftRecoveryMarker(draftKey); } catch { /* The composer also retains the fence in memory. */ } @@ -134,8 +148,10 @@ export function saveDraftSubmission( export function clearDraftSubmission(draftKey: string, attemptId: string) { try { - if (loadDraftSubmission(draftKey)?.attemptId === attemptId) + if (loadDraftSubmission(draftKey)?.attemptId === attemptId) { draftStorage(draftKey).removeItem(`${draftKey}:submission:v1`); + syncDraftRecoveryMarker(draftKey); + } } catch { /* Disabled browser storage is supported in memory. */ } @@ -238,6 +254,7 @@ export function saveDraftAttachments(draftKey: string, attachments: unknown, att JSON.stringify({ version: 1, draftKey, attachments: selected }), ); else draftStorage(draftKey).removeItem(`${draftKey}:attachments:v1`); + syncDraftRecoveryMarker(draftKey); } catch { /* Disabled/full browser storage must not break the composer. */ }