fix(web): preserve unsaved queued-message edits across navigation#23990

@@ -0,0 +1,124 @@
+import { EnvironmentId, MessageId, RunId, ThreadId } from "@t3tools/contracts";
+import { scopeThreadRef } from "@t3tools/client-runtime/environment";
+import { useRef } from "react";
+import { expect, it, vi } from "vite-plus/test";
+import { render } from "../../../test/render";
+import {
+  partializeComposerDraftStoreState,
+  useComposerDraftStore,
+  useComposerThreadDraft,
+} from "../../composerDraftStore";
+import { useQueuedRunEdit } from "./useQueuedRunEdit";
+
+const commands = vi.hoisted(() => ({
+  save: vi.fn(async (_input: unknown) => ({ _tag: "Success" })),
+}));
+vi.mock(import("../../state/use-atom-command"), () => ({
+  useAtomCommand: (() => commands.save) as never,
+}));
+vi.mock(import("../../assets/assetUrls"), () => ({ useAssetUrls: () => [] }));
+
+function Editor() {
+  const environmentId = EnvironmentId.make("environment:queue-edit");
+  const threadId = ThreadId.make("thread:queue-edit");
+  const threadRef = scopeThreadRef(environmentId, threadId);
+  const promptRef = useRef("");
+  const edit = useQueuedRunEdit({
+    environmentId,
+    routeThreadRef: threadRef,
+    threadDraftTarget: threadRef,
+    activeThreadId: threadId,
+    serverProjection: null,
+    promptRef,
+    composerRef: { current: null },
+  });
+  const draft = useComposerThreadDraft(edit.draftTarget);
+  return (
+    <>
+      <button
+        onClick={() =>
+          edit.beginEdit({
+            runId: RunId.make("run:queued"),
+            messageId: MessageId.make("message:queued"),
+            text: "Original message",
+            attachments: [
+              {
+                type: "file",
+                id: "file:report",
+                name: "report.pdf",
+                mimeType: "application/pdf",
+                sizeBytes: 6,
+              },
+            ],
+          })
+        }
+      >
+        Edit queued message
+      </button>
+      <textarea
+        aria-label="Message"
+        value={draft.prompt}
+        onChange={(event) =>
+          useComposerDraftStore.getState().setPrompt(edit.draftTarget, event.target.value)
+        }
+      />
+      {edit.editingAttachments?.map(({ attachment }) => (
+        <button key={attachment.id} onClick={() => edit.removeEditingAttachment(attachment.id)}>
+          Remove {attachment.name}
+        </button>
+      ))}
+      <button onClick={edit.cancelEdit}>Cancel edit</button>
+      <button onClick={() => edit.saveEdit(draft, () => undefined)}>Save edit</button>
+    </>
+  );
+}
+
+it.each(["Updated message", ""])(
+  "restores the unsaved edit %j and attachment removals after leaving the thread",
+  async (text) => {
+    commands.save.mockClear();
+    useComposerDraftStore.setState({ draftsByThreadKey: {}, draftThreadsByThreadKey: {} });
+    await using screen = await render(<Editor />, { router: false });
+    await screen.getByRole("textbox", { name: "Message" }).fill("Separate draft");
+    await screen.getByRole("button", { name: "Edit queued message" }).click();
+    await screen.getByRole("textbox", { name: "Message" }).fill(text);
+    await screen.getByRole("button", { name: "Remove report.pdf" }).click();
+    await screen.rerender(<div>Another thread</div>);
+    const persisted = JSON.parse(
+      JSON.stringify(partializeComposerDraftStoreState(useComposerDraftStore.getState())),
+    );
+    useComposerDraftStore.setState(
+      useComposerDraftStore.persist.getOptions().merge!(
+        persisted,
+        useComposerDraftStore.getInitialState(),
+      ),
+    );
+    await screen.rerender(<Editor />);
+    await screen.getByRole("button", { name: "Edit queued message" }).click();
+    await expect.element(screen.getByRole("textbox", { name: "Message" })).toHaveValue(text);
+    await expect
+      .element(screen.getByRole("button", { name: "Remove report.pdf" }))
+      .not.toBeInTheDocument();
+    if (text) {
+      await screen.getByRole("button", { name: "Save edit" }).click();
+      expect(commands.save).toHaveBeenCalledWith(
+        expect.objectContaining({
+          input: expect.objectContaining({
+            text,
+            edit: expect.objectContaining({ attachments: [] }),
+          }),
+        }),
+      );

@@ -219,17 +220,34 @@ export function useQueuedRunEdit(input: {
 
   const beginEdit = useStableCallback((request: EditQueuedRunRequest) => {
     if (activeThreadId === null) return;
-    if (editing !== null && editing.runId !== request.runId) clearEditDraft(editing.runId);
     const target = queuedEditDraftTarget(routeThreadRef, request.runId);
     const store = useComposerDraftStore.getState();
-    store.clearComposerContent(target);
-    store.setPrompt(target, request.text);
+    const savedDraft = store.getComposerDraft(target);

@@ -0,0 +1,138 @@
+import { EnvironmentId, MessageId, RunId, ThreadId } from "@t3tools/contracts";
+import { scopeThreadRef } from "@t3tools/client-runtime/environment";
+import { useRef } from "react";
+import { expect, it, vi } from "vite-plus/test";
+import { render } from "../../../test/render";
+import {
+  partializeComposerDraftStoreState,
+  useComposerDraftStore,
+  useComposerThreadDraft,
+} from "../../composerDraftStore";
+import { useQueuedRunEdit } from "./useQueuedRunEdit";
+
+const commands = vi.hoisted(() => ({
+  save: vi.fn(async (_input: unknown) => ({ _tag: "Success" })),
+}));
+vi.mock(import("../../state/use-atom-command"), () => ({
+  useAtomCommand: (() => commands.save) as never,
+}));
+vi.mock(import("../../assets/assetUrls"), () => ({ useAssetUrls: () => [] }));
+
+function Editor({ text = "Original message" }: { readonly text?: string }) {
+  const environmentId = EnvironmentId.make("environment:queue-edit");
+  const threadId = ThreadId.make("thread:queue-edit");
+  const threadRef = scopeThreadRef(environmentId, threadId);
+  const promptRef = useRef("");
+  const edit = useQueuedRunEdit({
+    environmentId,
+    routeThreadRef: threadRef,
+    threadDraftTarget: threadRef,
+    activeThreadId: threadId,
+    serverProjection: null,
+    promptRef,
+    composerRef: { current: null },
+  });
+  const draft = useComposerThreadDraft(edit.draftTarget);
+  return (
+    <>
+      <button
+        onClick={() =>
+          edit.beginEdit({
+            runId: RunId.make("run:queued"),
+            messageId: MessageId.make("message:queued"),
+            text,
+            attachments: [
+              {
+                type: "file",
+                id: "file:report",
+                name: "report.pdf",
+                mimeType: "application/pdf",
+                sizeBytes: 6,
+              },
+            ],
+          })
+        }
+      >
+        Edit queued message
+      </button>
+      <textarea
+        aria-label="Message"
+        value={draft.prompt}
+        onChange={(event) =>
+          useComposerDraftStore.getState().setPrompt(edit.draftTarget, event.target.value)
+        }
+      />
+      {edit.editingAttachments?.map(({ attachment }) => (
+        <button key={attachment.id} onClick={() => edit.removeEditingAttachment(attachment.id)}>
+          Remove {attachment.name}
+        </button>
+      ))}
+      <button onClick={edit.cancelEdit}>Cancel edit</button>
+      <button onClick={() => edit.saveEdit(draft, () => undefined)}>Save edit</button>
+    </>
+  );
+}
+
+it.each(["Updated message", ""])(
+  "restores the unsaved edit %j and attachment removals after leaving the thread",
+  async (text) => {
+    commands.save.mockClear();
+    useComposerDraftStore.setState({ draftsByThreadKey: {}, draftThreadsByThreadKey: {} });
+    await using screen = await render(<Editor />, { router: false });

@@ -254,34 +280,45 @@ export function useQueuedRunEdit(input: {
     );
   });
 
-  // Exit edit mode when the edited run leaves the queue (it started, or was
-  // cancelled from another client). A dirty edit moves into the thread's own
-  // draft when that draft is empty; otherwise it is dropped with a toast.
+  // Recover retained edits too: their run may have left the queue while this
+  // thread was unmounted. A dirty edit uses the same keep-or-discard path.
+  const runs = serverProjection?.runs;
   useEffect(() => {
-    if (editing === null || activeThreadId !== editing.threadId) return;
-    if (serverProjection === null) return;
-    const run = serverProjection.runs.find((candidate) => candidate.id === editing.runId);
-    if (run !== undefined && run.status === "queued") return;
-    const recovery = recoverQueuedMessageEdit({
-      editTarget: queuedEditDraftTarget(routeThreadRef, editing.runId),
-      threadTarget: threadDraftTarget,
-      originalText: editing.originalText,
-    });
-    if (recovery !== "clean") {
-      toastManager.add(
-        stackedThreadToast({
-          type: recovery === "kept" ? "info" : "warning",
-          title: "Queued message is no longer queued",
-          description:
-            recovery === "kept"
-              ? "Your unsaved edit was kept in the composer."
-              : "Your unsaved edit was discarded.",
-        }),
-      );
+    if (activeThreadId === null || runs === undefined) return;
+    const prefix = `queued-edit:${scopedThreadKey(routeThreadRef)}:`;
+    const queuedTargets = new Set(
+      runs
+        .filter((run) => run.status === "queued")
+        .map((run) => queuedEditDraftTarget(routeThreadRef, run.id)),
+    );
+    const drafts = useComposerDraftStore.getState().draftsByThreadKey;
+    for (const [key, draft] of Object.entries(drafts)) {
+      if (!key.startsWith(prefix) || draft.queuedMessageEdit === undefined) continue;
+      const editTarget = DraftId.make(key);
+      if (queuedTargets.has(editTarget)) continue;
+      const recovery = recoverQueuedMessageEdit({