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
2026-07-13 15:36:48 +02:00
parent 218097b7cc
commit 13d6ad598e
11 changed files with 462 additions and 79 deletions
@@ -1,7 +1,11 @@
 import { describe, expect, it } from "vitest";
 import type { UserComposerAttachment } from "@/attachments/types";
 import type { GitHubSearchItem } from "@getpaseo/protocol/messages";
-import { findCheckoutHintPrAttachment, syncPickerPrAttachment } from "./new-workspace-picker-state";
+import {
+  clearPickerPrAttachmentForTargetChange,
+  findCheckoutHintPrAttachment,
+  syncPickerPrAttachment,
+} from "./new-workspace-picker-state";

 function makePrItem(number: number, title: string, headRefName = "feature/x"): GitHubSearchItem {
  return {
@@ -19,8 +23,9 @@ function makePrItem(number: number, title: string, headRefName = "feature/x"): G

 function prAttachment(
  item: GitHubSearchItem,
+  owner?: "new-workspace-picker",
 ): Extract<UserComposerAttachment, { kind: "github_pr" }> {
-  return { kind: "github_pr", item };
+  return { kind: "github_pr", item, ...(owner ? { owner } : {}) };
 }

 function issueAttachment(number: number): UserComposerAttachment {
@@ -43,57 +48,88 @@ describe("syncPickerPrAttachment", () => {
    const pr = makePrItem(202, "Refactor picker");
    const result = syncPickerPrAttachment({
      attachments: [],
-      previousPickerPrNumber: null,
      item: { kind: "github-pr", item: pr },
    });
-    expect(result.attachedPrNumber).toBe(202);
-    expect(result.attachments).toEqual([prAttachment(pr)]);
+    expect(result).toEqual([prAttachment(pr, "new-workspace-picker")]);
  });

  it("selects a branch without modifying attachments when no previous picker PR", () => {
    const issue = issueAttachment(44);
    const result = syncPickerPrAttachment({
      attachments: [issue],
-      previousPickerPrNumber: null,
      item: { kind: "branch", name: "dev" },
    });
-    expect(result.attachedPrNumber).toBeNull();
-    expect(result.attachments).toEqual([issue]);
+    expect(result).toEqual([issue]);
  });

  it("replaces the previous picker PR when a different PR is selected", () => {
    const prA = makePrItem(202, "Refactor picker", "feature/picker");
    const prB = makePrItem(303, "Polish chip", "feature/chip");
    const result = syncPickerPrAttachment({
-      attachments: [prAttachment(prA)],
-      previousPickerPrNumber: 202,
+      attachments: [prAttachment(prA, "new-workspace-picker")],
      item: { kind: "github-pr", item: prB },
    });
-    expect(result.attachedPrNumber).toBe(303);
-    expect(result.attachments).toEqual([prAttachment(prB)]);
+    expect(result).toEqual([prAttachment(prB, "new-workspace-picker")]);
  });

  it("removes the previous picker PR and adds no new attachment when a branch is selected", () => {
    const pr = makePrItem(202, "Refactor picker");
    const issue = issueAttachment(44);
    const result = syncPickerPrAttachment({
-      attachments: [issue, prAttachment(pr)],
-      previousPickerPrNumber: 202,
+      attachments: [issue, prAttachment(pr, "new-workspace-picker")],
      item: { kind: "branch", name: "dev" },
    });
-    expect(result.attachedPrNumber).toBeNull();
-    expect(result.attachments).toEqual([issue]);
+    expect(result).toEqual([issue]);
  });

  it("does not duplicate a PR that was already manually attached by the user", () => {
    const pr = makePrItem(202, "Refactor picker");
    const result = syncPickerPrAttachment({
      attachments: [prAttachment(pr)],
-      previousPickerPrNumber: null,
      item: { kind: "github-pr", item: pr },
    });
-    expect(result.attachedPrNumber).toBeNull();
-    expect(result.attachments).toHaveLength(1);
+    expect(result).toEqual([prAttachment(pr)]);
+  });
+
+  it("clears a persisted picker selection without removing user-added attachments", () => {
+    const pickerPr = prAttachment(makePrItem(202, "Picker PR"), "new-workspace-picker");
+    const manuallyAttachedPr = prAttachment(makePrItem(303, "Manual PR"));
+    const issue = issueAttachment(44);
+
+    const result = syncPickerPrAttachment({
+      attachments: [issue, pickerPr, manuallyAttachedPr],
+      item: null,
+    });
+
+    expect(result).toEqual([issue, manuallyAttachedPr]);
+  });
+});
+
+describe("clearPickerPrAttachmentForTargetChange", () => {
+  it("keeps the picker selection when the target is reselected", () => {
+    const pickerPr = prAttachment(makePrItem(202, "Picker PR"), "new-workspace-picker");
+    const attachments = [pickerPr];
+
+    expect(
+      clearPickerPrAttachmentForTargetChange({
+        attachments,
+        currentTargetId: "server-a",
+        nextTargetId: "server-a",
+      }),
+    ).toBe(attachments);
+  });
+
+  it("clears only the picker-owned PR when the target changes", () => {
+    const pickerPr = prAttachment(makePrItem(202, "Picker PR"), "new-workspace-picker");
+    const manualPr = prAttachment(makePrItem(303, "Manual PR"));
+
+    expect(
+      clearPickerPrAttachmentForTargetChange({
+        attachments: [pickerPr, manualPr],
+        currentTargetId: "server-a",
+        nextTargetId: "server-b",
+      }),
+    ).toEqual([manualPr]);
  });
 });

@@ -1,37 +1,60 @@
-import type { UserComposerAttachment } from "@/attachments/types";
+import {
+  NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER,
+  type UserComposerAttachment,
+} from "@/attachments/types";
 import type { PickerItem } from "./new-workspace-picker-item";

-// The picker "owns" at most one PR attachment at a time. When the user selects
-// a different item the previously-owned PR is removed before the new one is added.
-// User-added attachments for other PRs/issues are left untouched.
+function isPickerOwnedPrAttachment(attachment: UserComposerAttachment): attachment is Extract<
+  UserComposerAttachment,
+  { kind: "github_pr" }
+> & {
+  owner: typeof NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER;
+} {
+  return (
+    attachment.kind === "github_pr" && attachment.owner === NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER
+  );
+}
+
+// Ownership lives on the attachment because drafts outlive this component.
+// The picker owns at most one PR; user-added PRs and issues remain untouched.
 export function syncPickerPrAttachment(input: {
  attachments: UserComposerAttachment[];
-  previousPickerPrNumber: number | null;
-  item: PickerItem;
-}): { attachments: UserComposerAttachment[]; attachedPrNumber: number | null } {
-  let nextAttachments = input.attachments;
-  let attachedPrNumber: number | null = null;
+  item: PickerItem | null;
+}): UserComposerAttachment[] {
+  const nextAttachments = input.attachments.filter(
+    (attachment) => !isPickerOwnedPrAttachment(attachment),
+  );

-  if (input.previousPickerPrNumber !== null) {
-    nextAttachments = nextAttachments.filter(
-      (attachment) =>
-        attachment.kind !== "github_pr" || attachment.item.number !== input.previousPickerPrNumber,
-    );
-  }
-
-  if (input.item.kind === "github-pr") {
+  if (input.item?.kind === "github-pr") {
    const selectedPr = input.item.item;
    const hasExistingPrAttachment = nextAttachments.some(
      (attachment) =>
        attachment.kind === "github_pr" && attachment.item.number === selectedPr.number,
    );
    if (!hasExistingPrAttachment) {
-      nextAttachments = [...nextAttachments, { kind: "github_pr", item: selectedPr }];
-      attachedPrNumber = selectedPr.number;
+      return [
+        ...nextAttachments,
+        {
+          kind: "github_pr",
+          item: selectedPr,
+          owner: NEW_WORKSPACE_PICKER_ATTACHMENT_OWNER,
+        },
+      ];
    }
  }

-  return { attachments: nextAttachments, attachedPrNumber };
+  return nextAttachments;
+}
+
+export function clearPickerPrAttachmentForTargetChange(input: {
+  attachments: UserComposerAttachment[];
+  currentTargetId: string;
+  nextTargetId: string;
+}): UserComposerAttachment[] {
+  if (input.currentTargetId === input.nextTargetId) {
+    return input.attachments;
+  }
+  return syncPickerPrAttachment({ attachments: input.attachments, item: null });
 }

 export function findCheckoutHintPrAttachment(input: {
@@ -50,7 +50,7 @@ import { navigateToWorkspace } from "@/stores/navigation-active-workspace-store"
 import { useLastWorkspaceSelection } from "@/stores/navigation-active-workspace-store";
 import { normalizeWorkspaceDescriptor, useSessionStore } from "@/stores/session-store";
 import { useWorkspace } from "@/stores/session-store-hooks";
-import { generateDraftId } from "@/stores/draft-keys";
+import { buildNewWorkspaceDraftKey, generateDraftId } from "@/stores/draft-keys";
 import { useDraftStore } from "@/stores/draft-store";
 import { useProjectPickerStore } from "@/stores/project-picker-store";
 import { isActiveCreateFlowForDraft, useCreateFlowStore } from "@/stores/create-flow-store";
@@ -92,7 +92,11 @@ import {
  type PickerCheckoutRequest,
  type PickerItem,
 } from "./new-workspace-picker-item";
-import { findCheckoutHintPrAttachment, syncPickerPrAttachment } from "./new-workspace-picker-state";
+import {
+  clearPickerPrAttachmentForTargetChange,
+  findCheckoutHintPrAttachment,
+  syncPickerPrAttachment,
+} from "./new-workspace-picker-state";
 import {
  resolveNewWorkspaceAutomaticServerId,
  resolveNewWorkspaceInitialServerId,
@@ -179,7 +183,6 @@ interface PickerOptionData {

 interface PickerSelection {
  item: PickerItem;
-  attachedPrNumber: number | null;
 }

 const BRANCH_OPTION_PREFIX = "branch:";
@@ -786,18 +789,6 @@ function getContentStyle(input: { isCompact: boolean; insetBottom: number }) {
  return [styles.content, styles.contentCentered];
 }

-function buildNewWorkspaceDraftKey(input: {
-  selectedServerId: string;
-  selectedSourceDirectory: string | null;
-  draftId?: string;
-}): string {
-  const explicitDraftId = input.draftId?.trim();
-  if (explicitDraftId) {
-    return `new-workspace:draft:${explicitDraftId}`;
-  }
-  return `new-workspace:${input.selectedServerId}:${input.selectedSourceDirectory ?? "choose-project"}`;
-}
-
 function getSelectedPickerItem(selection: PickerSelection | null): PickerItem | null {
  if (!selection) return null;
  return selection.item;
@@ -1682,11 +1673,7 @@ export function NewWorkspaceScreen({
  const projectIconDataByProjectKey = useProjectIconDataByProjectKey({
    projects: projectIconTargets,
  });
-  const draftKey = buildNewWorkspaceDraftKey({
-    selectedServerId,
-    selectedSourceDirectory,
-    draftId,
-  });
+  const draftKey = buildNewWorkspaceDraftKey(draftId);
  const forkDraftSetup = usePendingWorkspaceDraftSetup(draftId);
  const draftContextScopeKey = useDraftWorkspaceAttachmentScopeKey(draftId);
  const visibleDraftContextScopeKeys = useMemo(
@@ -1801,22 +1788,16 @@ export function NewWorkspaceScreen({
  }, [selectedItem]);
  const selectPickerItem = useCallback(
    (item: PickerItem) => {
-      const next = syncPickerPrAttachment({
+      const nextAttachments = syncPickerPrAttachment({
        attachments: chatDraft.attachments,
-        previousPickerPrNumber: manualPickerSelection?.attachedPrNumber ?? null,
        item,
      });

-      setManualPickerSelection({
-        item,
-        attachedPrNumber: next.attachedPrNumber,
-      });
-      if (next.attachments !== chatDraft.attachments) {
-        chatDraft.setAttachments(next.attachments);
-      }
+      setManualPickerSelection({ item });
+      chatDraft.setAttachments(nextAttachments);
      setPickerOpen(false);
    },
-    [chatDraft, manualPickerSelection?.attachedPrNumber],
+    [chatDraft],
  );

  const handleSelectOption = useCallback(
@@ -1828,6 +1809,20 @@ export function NewWorkspaceScreen({
    [itemById, selectPickerItem],
  );

+  const clearManualPickerSelectionForTargetChange = useCallback(
+    (currentTargetId: string, nextTargetId: string) => {
+      const nextAttachments = clearPickerPrAttachmentForTargetChange({
+        attachments: chatDraft.attachments,
+        currentTargetId,
+        nextTargetId,
+      });
+      if (nextAttachments === chatDraft.attachments) return;
+      chatDraft.setAttachments(nextAttachments);
+      setManualPickerSelection(null);
+    },
+    [chatDraft],
+  );
+
  const handleSelectProjectOption = useCallback(
    (id: string) => {
      // selectProjectOption enforces selectability (worktree-only when
@@ -1835,9 +1830,17 @@ export function NewWorkspaceScreen({
      // canCreateWorktree or non-git projects become unselectable.
      selectProjectOption(id);
      setProjectPickerOpen(false);
-      setManualPickerSelection(null);
+      clearManualPickerSelectionForTargetChange(selectedProjectOptionId, id);
    },
-    [selectProjectOption],
+    [clearManualPickerSelectionForTargetChange, selectProjectOption, selectedProjectOptionId],
+  );
+
+  const handleSelectWorkspaceHost = useCallback(
+    (id: string) => {
+      handleSelectHost(id);
+      clearManualPickerSelectionForTargetChange(selectedServerId, id);
+    },
+    [clearManualPickerSelectionForTargetChange, handleSelectHost, selectedServerId],
  );

  const handleAddProject = useCallback(() => {
@@ -2155,7 +2158,7 @@ export function NewWorkspaceScreen({
    host: {
      allHosts,
      selectedServerId,
-      onSelect: handleSelectHost,
+      onSelect: handleSelectWorkspaceHost,
      openState: hostPickerOpen,
      onOpenChange: handleHostPickerOpenChange,
      anchorRef: hostPickerAnchorRef,