From c4f29baed01f47ae2a43c4f3aa7a2ed3e5d859a6 Mon Sep 17 00:00:00 2001
From: Bersabel Tadesse <brsabel.t@gmail.com>
Date: Mon, 3 Aug 2026 10:59:36 -0700
Subject: [PATCH] Keep Prompt Shaper running across new-thread navigation

---
 plugins/improve-prompt/app.test.tsx | 90 +++++++++++++++++++++++++----
 plugins/improve-prompt/app.tsx      | 12 ++--
 plugins/improve-prompt/dist/app.js  |  2 +-
 3 files changed, 87 insertions(+), 17 deletions(-)

diff --git a/plugins/improve-prompt/app.test.tsx b/plugins/improve-prompt/app.test.tsx
index b52819d..9348162 100644
--- a/plugins/improve-prompt/app.test.tsx
+++ b/plugins/improve-prompt/app.test.tsx
@@ -537,17 +537,6 @@ describe("Improve Prompt composer action", () => {
         queuedMessageId: "qmsg_2",
       },
     },
-    {
-      name: "new-thread project",
-      sourceScope: {
-        kind: "new-thread" as const,
-        projectId: "proj_1",
-      },
-      destinationScope: {
-        kind: "new-thread" as const,
-        projectId: "proj_2",
-      },
-    },
   ])(
     "cancels a pending request across a keyed $name unmount/remount",
     async ({ sourceScope, destinationScope }) => {
@@ -625,6 +614,85 @@ describe("Improve Prompt composer action", () => {
     },
   );
 
+  it("keeps a new-thread enhancement running while navigating away and back", async () => {
+    const result = deferred<{
+      requestId: string;
+      helperThreadId: string;
+      status: "complete";
+      enhancedPrompt: string;
+      assumptions: null;
+      createdAt: number;
+      completedAt: number;
+    }>();
+    const cancelEnhancement = vi.fn(() => ({ cancelled: true as const }));
+    configureAction({
+      text: "rough new-thread draft",
+      attachmentCount: 1,
+      scope: { kind: "new-thread", projectId: "proj_1" },
+      rpc: {
+        startEnhancement: () => ({
+          requestId: REQUEST_ID,
+          helperThreadId: "thr_helper",
+        }),
+        getEnhancement: () => result.promise,
+        cancelEnhancement,
+      },
+    });
+    const Action = await loadAction();
+    const source = mountAction(Action);
+
+    fireEvent.click(screen.getByRole("button", { name: "Improve prompt" }));
+    await waitFor(() => {
+      expect(actionSlot.inspection.composer.textEffect).toEqual({
+        className: "bb-improve-prompt-shimmer",
+      });
+      expect(window.sessionStorage.length).toBe(1);
+    });
+
+    source.lifecycle.unmount();
+    await driveComposerScope({ kind: "thread", threadId: "thr_other" });
+    await driveComposerText("other thread draft");
+    const destination = mountAction(Action);
+    await waitFor(() => {
+      expect(
+        screen.getByRole("button", { name: "Improve prompt" }),
+      ).not.toBeNull();
+      expect(actionSlot.inspection.composer.textEffect).toBeNull();
+    });
+    expect(cancelEnhancement).not.toHaveBeenCalled();
+    expect(window.sessionStorage.length).toBe(1);
+
+    destination.lifecycle.unmount();
+    await driveComposerScope({ kind: "new-thread", projectId: "proj_1" });
+    await driveComposerText("rough new-thread draft");
+    mountAction(Action);
+    await screen.findByRole("button", {
+      name: "Cancel prompt improvement",
+    });
+    expect(cancelEnhancement).not.toHaveBeenCalled();
+
+    await act(async () => {
+      result.resolve({
+        requestId: REQUEST_ID,
+        helperThreadId: "thr_helper",
+        status: "complete",
+        enhancedPrompt: "Enhanced after new-thread navigation.",
+        assumptions: null,
+        createdAt: 1,
+        completedAt: 2,
+      });
+      await result.promise;
+    });
+
+    await waitFor(() => {
+      expect(actionSlot.inspection.composer.text).toBe(
+        "Enhanced after new-thread navigation.",
+      );
+      expect(window.sessionStorage.length).toBe(0);
+    });
+    expect(actionSlot.inspection.composer.attachmentCount).toBe(1);
+  });
+
   it("replaces the latest edited draft and restores it through inline Undo", async () => {
     const result = deferred<{
       requestId: string;
diff --git a/plugins/improve-prompt/app.tsx b/plugins/improve-prompt/app.tsx
index 27f7db7..92ec202 100644
--- a/plugins/improve-prompt/app.tsx
+++ b/plugins/improve-prompt/app.tsx
@@ -375,13 +375,15 @@ function PromptShaperAction() {
       const detachedRequest = pendingRef.current;
       pendingRef.current = null;
 
-      // Thread requests intentionally survive keyed navigation so returning
-      // to the thread can recover and reconcile them. Every other composer
-      // scope is ephemeral: the host's full-scope key unmounts this action
-      // when its owner changes, so invalidate and cancel from this cleanup.
+      // Thread and new-thread requests intentionally survive navigation so
+      // returning to their durable scope can recover and reconcile them.
+      // Side-chat and queued-message scopes are ephemeral: the host's
+      // full-scope key unmounts this action when their owner changes, so
+      // invalidate and cancel those requests from this cleanup.
       if (
         detachedRequest === null ||
-        mountedComposerScopeKindRef.current === "thread"
+        mountedComposerScopeKindRef.current === "thread" ||
+        mountedComposerScopeKindRef.current === "new-thread"
       ) {
         return;
       }
diff --git a/plugins/improve-prompt/dist/app.js b/plugins/improve-prompt/dist/app.js
index 1f25acd..556951f 100644
--- a/plugins/improve-prompt/dist/app.js
+++ b/plugins/improve-prompt/dist/app.js
@@ -5210,7 +5210,7 @@ function PromptShaperAction() {
     return () => {
       const detachedRequest = pendingRef.current;
       pendingRef.current = null;
-      if (detachedRequest === null || mountedComposerScopeKindRef.current === "thread") {
+      if (detachedRequest === null || mountedComposerScopeKindRef.current === "thread" || mountedComposerScopeKindRef.current === "new-thread") {
         return;
       }
       const cancellationRequest = {