diff --git a/apps/mobile/src/state/queued-run-edit.test.ts b/apps/mobile/src/state/queued-run-edit.test.ts index 67e456369b890..fe50dbf497728 100644 --- a/apps/mobile/src/state/queued-run-edit.test.ts +++ b/apps/mobile/src/state/queued-run-edit.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vite-plus/test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test"; import { ChatAttachmentId, ComposerContextId, @@ -10,7 +10,34 @@ import { import { composerDraftEnvironmentId } from "../lib/composerAttachmentUploadQueue"; import { isQueuedEditDraftKey, queuedEditDraftKey } from "./queued-edit-draft-key"; -import { resolveQueuedEditPayload, type QueuedRunEdit } from "./queued-run-edit"; +import { + beginQueuedRunEdit, + endQueuedRunEdit, + getQueuedRunEdit, + recoverQueuedRunEdit, + resolveQueuedEditPayload, + type QueuedRunEdit, +} from "./queued-run-edit"; +import { + clearComposerDraft, + flushComposerDrafts, + getComposerDraftSnapshot, + setComposerDraftText, + waitForComposerDraftsLoaded, +} from "./use-composer-drafts"; + +vi.mock("expo-file-system", () => ({ + Directory: class { + create() {} + }, + File: class { + exists = false; + create() {} + write() {} + moveSync() {} + }, + Paths: { document: { uri: "file:///documents" } }, +})); const image = (id: string): ChatAttachment => ({ type: "image", @@ -104,3 +131,58 @@ describe("resolveQueuedEditPayload", () => { expect(payload.context).toBeUndefined(); }); }); + +describe("recoverQueuedRunEdit", () => { + const threadKey = "env:thread-recover"; + const runId = RunId.make("run-1"); + const editKey = queuedEditDraftKey(threadKey, runId); + + beforeEach(async () => { + await waitForComposerDraftsLoaded(); + }); + + afterEach(async () => { + endQueuedRunEdit(threadKey); + clearComposerDraft(threadKey); + await flushComposerDrafts(); + }); + + it("drops an untouched edit without restoring the sent message", async () => { + beginQueuedRunEdit(threadKey, edit([])); + + const recovery = recoverQueuedRunEdit(threadKey, runId); + await flushComposerDrafts(); + + expect(getComposerDraftSnapshot(threadKey).text).toBe(""); + expect(getComposerDraftSnapshot(editKey).text).toBe(""); + expect(getQueuedRunEdit(threadKey)).toBeNull(); + expect(recovery).toBe("clean"); + }); + + it("recovers a changed edit into an empty thread draft", async () => { + beginQueuedRunEdit(threadKey, edit([])); + setComposerDraftText(editKey, "revised"); + + const recovery = recoverQueuedRunEdit(threadKey, runId); + await flushComposerDrafts(); + + expect(getComposerDraftSnapshot(threadKey).text).toBe("revised"); + expect(getComposerDraftSnapshot(editKey).text).toBe(""); + expect(getQueuedRunEdit(threadKey)).toBeNull(); + expect(recovery).toBe("kept"); + }); + + it("preserves an occupied thread draft when discarding a changed edit", async () => { + setComposerDraftText(threadKey, "unsent draft"); + beginQueuedRunEdit(threadKey, edit([])); + setComposerDraftText(editKey, "revised"); + + const recovery = recoverQueuedRunEdit(threadKey, runId); + await flushComposerDrafts(); + + expect(getComposerDraftSnapshot(threadKey).text).toBe("unsent draft"); + expect(getComposerDraftSnapshot(editKey).text).toBe(""); + expect(getQueuedRunEdit(threadKey)).toBeNull(); + expect(recovery).toBe("discarded"); + }); +}); diff --git a/apps/mobile/src/state/queued-run-edit.ts b/apps/mobile/src/state/queued-run-edit.ts index 6eac91a5adfc6..b0c325d401aaf 100644 --- a/apps/mobile/src/state/queued-run-edit.ts +++ b/apps/mobile/src/state/queued-run-edit.ts @@ -19,6 +19,8 @@ import { appAtomRegistry } from "./atom-registry"; import { queuedEditDraftKey } from "./queued-edit-draft-key"; import { clearComposerDraft, + getComposerDraftSnapshot, + mergeComposerDraftContent, setComposerDraftContext, setComposerDraftText, } from "./use-composer-drafts"; @@ -93,6 +95,35 @@ export function endQueuedRunEdit( setQueuedRunEdit(threadKey, null); } +export function recoverQueuedRunEdit( + threadKey: string, + runId: RunId, +): "kept" | "discarded" | "clean" | null { + const edit = getQueuedRunEdit(threadKey); + if (edit === null || edit.runId !== runId) return null; + const editDraft = getComposerDraftSnapshot(queuedEditDraftKey(threadKey, runId)); + const dirty = editDraft.text !== edit.originalText || editDraft.attachments.length > 0; + if (!dirty) { + // Untouched text was already sent; restoring it would duplicate the message. + endQueuedRunEdit(threadKey); + return "clean"; + } + const threadDraft = getComposerDraftSnapshot(threadKey); + const keepable = + editDraft.text.trim().length > 0 && + threadDraft.text.trim().length === 0 && + threadDraft.attachments.length === 0; + if (keepable) { + void mergeComposerDraftContent(threadKey, { + text: editDraft.text, + attachments: editDraft.attachments, + ...(editDraft.context ? { context: editDraft.context } : {}), + }); + } + endQueuedRunEdit(threadKey, { deferAttachmentCleanup: keepable }); + return keepable ? "kept" : "discarded"; +} + /** * Attachments replace the message's list wholesale, so kept server attachments * are resent alongside the newly uploaded ones, in composer order. diff --git a/apps/mobile/src/state/use-thread-composer-state.ts b/apps/mobile/src/state/use-thread-composer-state.ts index 4129527668fd8..a5e2ff64959d3 100644 --- a/apps/mobile/src/state/use-thread-composer-state.ts +++ b/apps/mobile/src/state/use-thread-composer-state.ts @@ -97,6 +97,7 @@ import { environmentThreadDetails } from "./threads"; import { endQueuedRunEdit, getQueuedRunEdit, + recoverQueuedRunEdit, queuedEditDraftKey, removeQueuedRunEditAttachment, resolveQueuedEditPayload, @@ -383,27 +384,15 @@ export function useThreadComposerState() { (run) => run.id === editedRunId && run.status === "queued", ); if (stillQueued) return; - const editDraftKey = queuedEditDraftKey(selectedThreadKey, editedRunId); - const editDraft = getComposerDraftSnapshot(editDraftKey); - const threadDraft = getComposerDraftSnapshot(selectedThreadKey); - const keepable = - editDraft.text.trim().length > 0 && - threadDraft.text.trim().length === 0 && - threadDraft.attachments.length === 0; - if (keepable) { - void mergeComposerDraftContent(selectedThreadKey, { - text: editDraft.text, - attachments: editDraft.attachments, - ...(editDraft.context ? { context: editDraft.context } : {}), - }); + const recovery = recoverQueuedRunEdit(selectedThreadKey, editedRunId); + if (recovery === "kept" || recovery === "discarded") { + setThreadComposerError( + selectedThreadKey, + recovery === "kept" + ? "That message already started. Your edit is back in the composer." + : "That message already started, so the edit was discarded.", + ); } - endQueuedRunEdit(selectedThreadKey, { deferAttachmentCleanup: keepable }); - setThreadComposerError( - selectedThreadKey, - keepable - ? "That message already started. Your edit is back in the composer." - : "That message already started, so the edit was discarded.", - ); }, [editedRunId, selectedThreadKey, selectedThreadRuns]); const activeThreadBusy = threadRuntimeIsActive(selectedThreadRuntime);