Keep New Workspace prompts when switching projects or hosts (#2036)

* fix(app): preserve New Workspace prompts across target changes

The New Workspace draft was keyed by the selected host and project, so changing either target replaced the visible prompt with another draft scope. Treat the screen as one persistent draft surface and migrate the newest active legacy draft.

* test(app): express draft persistence as user actions

* fix(app): clear stale picker PR context

* fix(app): persist picker attachment ownership

* fix(app): keep workspace target context safe

packages/app/src/stores/draft-store/migration.test.ts

@@ -1,11 +1,60 @@
 import { describe, expect, it } from "vitest";
-import type { ComposerAttachment } from "@/attachments/types";
+import type { ComposerAttachment, UserComposerAttachment } from "@/attachments/types";
 import { migratePersistedState, type MigrateLegacyImages } from "./migration";
-import { isAttachmentMetadata } from "./state";
+import { isAttachmentMetadata, type DraftRecord } from "./state";
 
 const passThroughMigrateLegacyImages: MigrateLegacyImages = async (images) =>
   images.filter(isAttachmentMetadata);
 
+function activeDraft(
+  text: string,
+  updatedAt: number,
+  attachments: UserComposerAttachment[] = [],
+): DraftRecord {
+  return {
+    input: { text, attachments },
+    lifecycle: "active",
+    updatedAt,
+    version: 1,
+  };
+}
+
+function githubIssueAttachment(
+  number: number,
+): Extract<UserComposerAttachment, { kind: "github_issue" }> {
+  return {
+    kind: "github_issue",
+    item: {
+      kind: "issue",
+      number,
+      title: `Review item ${number}`,
+      url: `https://example.com/issues/${number}`,
+      state: "open",
+      body: null,
+      labels: [],
+    },
+  };
+}
+
+function githubPrAttachment(
+  number: number,
+): Extract<UserComposerAttachment, { kind: "github_pr" }> {
+  return {
+    kind: "github_pr",
+    item: {
+      kind: "pr",
+      number,
+      title: `Review item ${number}`,
+      url: `https://example.com/pulls/${number}`,
+      state: "open",
+      body: null,
+      labels: [],
+      baseRefName: "main",
+      headRefName: "feature/legacy",
+    },
+  };
+}
+
 function workspaceReviewAttachment(): Extract<ComposerAttachment, { kind: "review" }> {
   return {
     kind: "review",
@@ -47,6 +96,57 @@ function workspaceReviewAttachment(): Extract<ComposerAttachment, { kind: "revie
 }
 
 describe("draft-store migration", () => {
+  it("promotes the newest legacy New Workspace draft into the singleton surface", async () => {
+    const forkDraft = activeDraft("fork context", 1700000000003);
+    const agentDraft = activeDraft("agent prompt", 1700000000004);
+
+    const migrated = await migratePersistedState(
+      {
+        drafts: {
+          "new-workspace:server-a:/project/older": activeDraft(
+            "older new workspace prompt",
+            1700000000001,
+          ),
+          "new-workspace:server-b:/project/newer": activeDraft(
+            "newer new workspace prompt",
+            1700000000002,
+          ),
+          "new-workspace:draft:fork-1": forkDraft,
+          "agent:server-a:agent-1": agentDraft,
+        },
+        createModalDraft: null,
+      },
+      { migrateLegacyImages: passThroughMigrateLegacyImages, nowMs: 1700000000005 },
+    );
+
+    expect(migrated.drafts).toEqual({
+      "new-workspace": activeDraft("newer new workspace prompt", 1700000000002),
+      "new-workspace:draft:fork-1": forkDraft,
+      "agent:server-a:agent-1": agentDraft,
+    });
+  });
+
+  it("drops unowned checkout PR context when promoting a scoped New Workspace draft", async () => {
+    const issue = githubIssueAttachment(101);
+    const migrated = await migratePersistedState(
+      {
+        drafts: {
+          "new-workspace:server-a:/project/a": activeDraft("keep the prompt", 2, [
+            issue,
+            githubPrAttachment(202),
+          ]),
+        },
+        createModalDraft: null,
+      },
+      { migrateLegacyImages: passThroughMigrateLegacyImages, nowMs: 3 },
+    );
+
+    expect(migrated.drafts["new-workspace"]?.input).toEqual({
+      text: "keep the prompt",
+      attachments: [issue],
+    });
+  });
+
   it("normalizes legacy image metadata into image attachments and strips persisted preview URLs", async () => {
     const migrated = await migratePersistedState(
       {

packages/app/src/stores/draft-store/migration.ts

@@ -1,4 +1,5 @@
 import type { AttachmentMetadata, UserComposerAttachment } from "@/attachments/types";
+import { isLegacyNewWorkspaceDraftKey, NEW_WORKSPACE_DRAFT_KEY } from "@/stores/draft-keys";
 import {
   isAttachmentMetadata,
   isLegacyDraftImage,
@@ -98,6 +99,46 @@ async function buildMigratedDraftRecord(
   };
 }
 
+function migrateNewWorkspaceDraftKeys(
+  drafts: Record<string, DraftRecord>,
+): Record<string, DraftRecord> {
+  const legacyEntries = Object.entries(drafts).filter(([draftKey]) =>
+    isLegacyNewWorkspaceDraftKey(draftKey),
+  );
+  if (legacyEntries.length === 0) {
+    return drafts;
+  }
+
+  const nextDrafts = { ...drafts };
+  for (const [draftKey] of legacyEntries) {
+    delete nextDrafts[draftKey];
+  }
+
+  if (nextDrafts[NEW_WORKSPACE_DRAFT_KEY]) {
+    return nextDrafts;
+  }
+
+  const newestActiveDraft = legacyEntries
+    .map(([, draft]) => draft)
+    .filter((draft) => draft.lifecycle === "active")
+    .sort((left, right) => right.updatedAt - left.updatedAt)[0];
+  if (newestActiveDraft) {
+    // Legacy scoped drafts did not record whether a PR came from the Base
+    // picker. That checkout context is unsafe to carry onto a global surface.
+    nextDrafts[NEW_WORKSPACE_DRAFT_KEY] = {
+      ...newestActiveDraft,
+      input: {
+        ...newestActiveDraft.input,
+        attachments: newestActiveDraft.input.attachments.filter(
+          (attachment) => attachment.kind !== "github_pr",
+        ),
+      },
+    };
+  }
+
+  return nextDrafts;
+}
+
 export async function migratePersistedState(
   state: unknown,
   ports: { migrateLegacyImages: MigrateLegacyImages; nowMs: number },
@@ -129,7 +170,8 @@ export async function migratePersistedState(
   }
 
   return {
-    drafts: nextDrafts,
+    // COMPAT(newWorkspaceDraftSingleton): migrated in v0.1.108; remove after 2027-01-13.
+    drafts: migrateNewWorkspaceDraftKeys(nextDrafts),
     createModalDraft,
   };
 }

packages/app/src/stores/draft-store/state.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import { applyClearDraftRecord, pruneFinalizedDraftRecords } from "./state";
+import { applyClearDraftRecord, pruneFinalizedDraftRecords, toDraftInputIfReady } from "./state";
 
 describe("draft-store lifecycle", () => {
   it("prunes finalized tombstones after TTL", () => {
@@ -69,3 +69,35 @@ describe("draft-store lifecycle", () => {
     });
   });
 });
+
+describe("draft-store normalization", () => {
+  it("preserves New Workspace picker ownership when hydrating a draft", () => {
+    const pickerAttachment = {
+      kind: "github_pr" as const,
+      owner: "new-workspace-picker" as const,
+      item: {
+        kind: "pr" as const,
+        number: 202,
+        title: "Persist picker ownership",
+        url: "https://example.com/pull/202",
+        state: "open" as const,
+        body: null,
+        labels: [],
+        baseRefName: "main",
+        headRefName: "feature/picker-ownership",
+      },
+    };
+
+    expect(
+      toDraftInputIfReady({
+        input: { text: "Keep this prompt", attachments: [pickerAttachment] },
+        lifecycle: "active",
+        updatedAt: 1,
+        version: 1,
+      }),
+    ).toEqual({
+      text: "Keep this prompt",
+      attachments: [pickerAttachment],
+    });
+  });
+});

packages/app/src/stores/draft-store/state.ts

@@ -1,7 +1,11 @@
-import type { AttachmentMetadata, UserComposerAttachment } from "@/attachments/types";
+import {
+  NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER,
+  type AttachmentMetadata,
+  type UserComposerAttachment,
+} from "@/attachments/types";
 import { GitHubSearchItemSchema } from "@getpaseo/protocol/messages";
 
-export const DRAFT_STORE_VERSION = 4;
+export const DRAFT_STORE_VERSION = 5;
 export const FINALIZED_DRAFT_TTL_MS = 5 * 60 * 1000;
 
 export interface LegacyDraftImage {
@@ -82,6 +86,13 @@ export function isUserComposerAttachment(value: unknown): value is UserComposerA
   if (record.kind !== "github_issue" && record.kind !== "github_pr") {
     return false;
   }
+  if (
+    record.kind === "github_pr" &&
+    record.owner !== undefined &&
+    record.owner !== NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER
+  ) {
+    return false;
+  }
   return GitHubSearchItemSchema.safeParse(record.item).success;
 }
 
@@ -94,6 +105,15 @@ export function normalizeComposerAttachment(
       metadata: normalizeAttachmentMetadata(attachment.metadata),
     };
   }
+  if (attachment.kind === "github_pr") {
+    return {
+      kind: "github_pr",
+      item: attachment.item,
+      ...(attachment.owner === NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER
+        ? { owner: attachment.owner }
+        : {}),
+    };
+  }
   return attachment;
 }