fix(app): retain checkout selection after workspace failure

Return the submitted checkout through the existing draft failure handoff. Keep client-owned receipt ordering and retries out of app presentation, and retain direct coverage of newer draft cleanup while removing the unused background-agent wrapper.

docs/architecture.md

@@ -364,8 +364,8 @@ so it cannot erase newer input.
 A failed creation goes to the existing owner of whatever the host did not create, using
 the receipt's resources and `failedStage`. Creation has no failure UI of its own:
 
-- No workspace: the app returns to New workspace with the submitted draft and shows the
-  error the way that screen shows its own submit errors.
+- No workspace: the app returns to New workspace with the submitted draft and checkout
+  selection, and shows the error the way that screen shows its own submit errors.
 - Workspace but no agent: the agent tab becomes an existing-workspace agent draft holding
   the prompt, in the abandoned create-flow state. Its error and Send retry are that draft's
   own; the retry creates the agent in the kept workspace.

packages/app/e2e/browser/creation-idempotency.spec.ts

@@ -1,3 +1,4 @@
+import { expectPickerSelected } from "../support/helpers/new-workspace";
 import {
   hideCreationLifecycle,
   holdWorkspaceProvisioning,
@@ -303,6 +304,39 @@ test("a failed workspace returns to New workspace with the prompt and its error"
   await page.screenshot({ path: testInfo.outputPath("failed-workspace-returned.png") });
 });
 
+test("a failed workspace preserves its non-default branch through resubmission", async ({
+  creation,
+  provisioningFailure,
+  page,
+}) => {
+  const prompt = "Keep the selected worktree base when provisioning fails.";
+  const provisioning = await holdWorkspaceProvisioning(page);
+  await creation.openWorkspaceForm("worktree");
+  await creation.selectWorktreeBranch("recovery-base");
+  await provisioningFailure.arm();
+  await creation.submitPrompt(prompt, "Create");
+  const first = await provisioning.waitForRequest();
+  expect(first.source).toMatchObject({
+    kind: "worktree",
+    action: "branch-off",
+    refName: "refs/heads/recovery-base",
+  });
+  provisioning.release();
+  await creation.expectReturnedToNewWorkspace(prompt, "Provisioning outcome is uncertain for test");
+
+  await expectPickerSelected(page, "recovery-base");
+  await creation.resubmitWorkspace();
+  const second = await provisioning.waitForRequests(2);
+  expect(second.source).toMatchObject({
+    kind: "worktree",
+    action: "branch-off",
+    refName: "refs/heads/recovery-base",
+  });
+  provisioning.release();
+  await creation.expectOneCreatedWorkspace();
+  await creation.expectAgentCount(1);
+});
+
 test("a failed workspace keeps its queued follow-up and unsent text through resubmission", async ({
   creation,
   provisioningFailure,

packages/app/e2e/support/helpers/creation.ts

@@ -10,6 +10,8 @@ import { createAgentTabFromMenu } from "./workspace-tabs";
 import {
   connectNewWorkspaceDaemonClient,
   openNewWorkspaceComposer,
+  openStartingRefPicker,
+  startingRefRow,
   selectWorkspaceIsolation,
   loadSessionMessageReaders,
 } from "./new-workspace";
@@ -190,7 +192,10 @@ export async function createCreationScenario(page: Page) {
   const requests = await observeCreationRequests(page);
   const project = await seedWorkspace({
     repoPrefix: "creation-idempotency-",
-    repo: { files: [{ path: "nested/context.txt", content: "Nested agent working directory" }] },
+    repo: {
+      files: [{ path: "nested/context.txt", content: "Nested agent working directory" }],
+      branches: ["recovery-base"],
+    },
   });
   let workspaceId = project.workspaceId;
   const evictionWorkspaceIds = new Set<string>();
@@ -269,6 +274,10 @@ export async function createCreationScenario(page: Page) {
       await openNewWorkspaceComposer(page, project);
       await selectWorkspaceIsolation(page, isolation);
     },
+    async selectWorktreeBranch(name: string) {
+      await openStartingRefPicker(page);
+      await startingRefRow(page, `${name}, local branch`).click();
+    },
     async openAgentDraft() {
       await gotoWorkspace(page, project.workspaceId);
       await createAgentTabFromMenu(page);

packages/app/src/screens/new-workspace-picker-state.test.ts

@@ -163,6 +163,25 @@ describe("clearPickerPrAttachmentForTargetChange", () => {
 });
 
 describe("reducePickerSelection", () => {
+  it("restores the submitted branch without replacing a newer choice", () => {
+    const submitted = {
+      kind: "branch" as const,
+      name: "recovery-base",
+      refName: "refs/heads/recovery-base",
+      accessibilityLabel: "recovery-base, local branch",
+    };
+    const restoration = { type: "draft-restored" as const, item: submitted };
+    expect(reducePickerSelection(initialPickerSelectionState, restoration).selectedItem).toEqual(
+      submitted,
+    );
+
+    const newer = reducePickerSelection(initialPickerSelectionState, {
+      type: "picker-selected",
+      item: { ...submitted, name: "newer-base", refName: "refs/heads/newer-base" },
+    });
+    expect(reducePickerSelection(newer, restoration)).toBe(newer);
+  });
+
   it("selects a PR that was newly detected and added", () => {
     const item = { kind: "github-pr" as const, item: makePrItem(101, "A") };
     const detected = reducePickerSelection(initialPickerSelectionState, { type: "pr-detected" });

packages/app/src/screens/new-workspace-picker-state.ts

@@ -13,6 +13,7 @@ export type PickerSelectionEvent =
   | { type: "pr-detected" }
   | { type: "pr-added"; item: Extract<PickerItem, { kind: "github-pr" }> }
   | { type: "picker-selected"; item: PickerItem }
+  | { type: "draft-restored"; item: PickerItem | null }
   | { type: "target-changed" };
 
 export const initialPickerSelectionState: PickerSelectionState = {
@@ -33,6 +34,8 @@ export function reducePickerSelection(
         : state;
     case "picker-selected":
       return { selectedItem: event.item, allowAutoPrSelection: false };
+    case "draft-restored":
+      return state.selectedItem ? state : { selectedItem: event.item, allowAutoPrSelection: false };
     case "target-changed":
       return initialPickerSelectionState;
   }

packages/app/src/screens/new-workspace-screen.tsx

@@ -1,7 +1,7 @@
 import { getHostRuntimeStore } from "@/runtime/host-runtime";
 import type { CreateWorkspaceRequestOptions } from "@getpaseo/client/internal/daemon-client";
 import type { AgentSnapshotPayload, CreationSnapshot } from "@getpaseo/protocol/messages";
-import { submitWorkspaceChat, takeWorkspaceCreationError } from "@/workspace-creation";
+import { submitWorkspaceChat, takeWorkspaceCreationFailure } from "@/workspace-creation";
 import { useCallback, useEffect, useMemo, useReducer, useRef, useState } from "react";
 import type { ReactNode } from "react";
 import { useTranslation } from "react-i18next";
@@ -631,13 +631,6 @@ export function NewWorkspaceScreen({
     projects: projectIconTargets,
   });
   const draftKey = buildNewWorkspaceDraftKey(draftId);
-  // A chat whose workspace could not be created returns here with its draft.
-  useEffect(() => {
-    const message = takeWorkspaceCreationError(draftKey);
-    if (!message) return;
-    setErrorMessage(message);
-    toast.error(message);
-  }, [draftKey, toast]);
   const forkDraftSetup = usePendingWorkspaceDraftSetup(draftId);
   const draftContextScopeKey = useDraftWorkspaceAttachmentScopeKey(draftId);
   const visibleDraftContextScopeKeys = useMemo(
@@ -660,6 +653,14 @@ export function NewWorkspaceScreen({
   );
   const selectedItem = pickerSelection.selectedItem;
 
+  useEffect(() => {
+    const failure = takeWorkspaceCreationFailure(draftKey);
+    if (!failure) return;
+    dispatchPickerSelection({ type: "draft-restored", item: failure.selectedItem });
+    setErrorMessage(failure.error);
+    toast.error(failure.error);
+  }, [draftKey, toast]);
+
   const handleForgeChangeRequestDetected = useCallback(() => {
     dispatchPickerSelection({ type: "pr-detected" });
   }, []);

packages/app/src/screens/new-workspace/background-handoff.test.ts

@@ -1,12 +1,8 @@
 import { beforeEach, describe, expect, it } from "vitest";
-import { useCreateFlowStore } from "@/stores/create-flow-store";
 import { useDraftStore } from "@/stores/draft-store";
 import { useWorkspaceAttachmentsStore } from "@/attachments/workspace-attachments-store";
 import { useWorkspaceDraftSubmissionStore } from "@/stores/workspace-draft-submission-store";
-import {
-  captureWorkspaceDraftCleanup,
-  createWorkspaceAgentInBackground,
-} from "./background-handoff";
+import { captureWorkspaceDraftCleanup } from "./background-handoff";
 
 const DRAFT_ID = "draft_1";
 const DRAFT_KEY = "new-workspace";
@@ -32,144 +28,96 @@ function captureCleanup() {
   });
 }
 
-function run(createAgent: () => Promise<unknown>, clearConsumedDraft = captureCleanup()) {
-  return createWorkspaceAgentInBackground({ createAgent, clearConsumedDraft });
+function setSetup(cwd: string) {
+  useWorkspaceDraftSubmissionStore.getState().setDraftSetup({
+    draftId: DRAFT_ID,
+    setup: {
+      cwd,
+      provider: "codex",
+      modeId: null,
+      model: null,
+      thinkingOptionId: null,
+      featureValues: {},
+    },
+  });
+}
+
+function addContext() {
+  useWorkspaceAttachmentsStore.getState().addWorkspaceAttachment({
+    scopeKey: SCOPE_KEY,
+    attachment: {
+      kind: "chat_history",
+      id: "fork-context",
+      attachment: {
+        type: "text",
+        mimeType: "text/plain",
+        contextKind: "chat_history",
+        title: "Chat history",
+        text: "Preserve this context",
+      },
+      source: { serverId: "local", agentId: "agent-1" },
+    },
+  });
 }
 
 beforeEach(() => {
   useWorkspaceDraftSubmissionStore.setState({ pendingByDraftId: {}, setupByDraftId: {} });
-  useCreateFlowStore.setState({ pendingByDraftId: {} });
   useWorkspaceAttachmentsStore.setState({ attachmentsByScope: { [SCOPE_KEY]: [] } });
   saveDraft("do the thing");
 });
 
-describe("createWorkspaceAgentInBackground", () => {
-  it("preserves context added to a newer draft without editing its text", async () => {
-    await run(async () => {
-      useWorkspaceAttachmentsStore.getState().addWorkspaceAttachment({
-        scopeKey: SCOPE_KEY,
-        attachment: {
-          kind: "chat_history",
-          id: "fork-context",
-          attachment: {
-            type: "text",
-            mimeType: "text/plain",
-            contextKind: "chat_history",
-            title: "Chat history",
-            text: "Preserve this context",
-          },
-          source: { serverId: "local", agentId: "agent-1" },
-        },
-      });
-      useWorkspaceDraftSubmissionStore.getState().setDraftSetup({
-        draftId: DRAFT_ID,
-        setup: {
-          cwd: "/work/new-repo",
-          provider: "codex",
-          modeId: null,
-          model: null,
-          thinkingOptionId: null,
-          featureValues: {},
-        },
-      });
-    });
-    expect(useWorkspaceAttachmentsStore.getState().attachmentsByScope[SCOPE_KEY]).toHaveLength(1);
-    expect(useWorkspaceDraftSubmissionStore.getState().setupByDraftId[DRAFT_ID]?.setup.cwd).toBe(
-      "/work/new-repo",
-    );
-  });
+describe("captureWorkspaceDraftCleanup", () => {
+  it("clears the consumed draft setup, attachments and input", () => {
+    setSetup("/work/repo");
+    addContext();
+    const clearConsumedDraft = captureCleanup();
 
-  it("creates one agent and clears the consumed draft setup, attachments and input", async () => {
-    let calls = 0;
-    useWorkspaceDraftSubmissionStore.getState().setDraftSetup({
-      draftId: DRAFT_ID,
-      setup: {
-        cwd: "/work/repo",
-        provider: "claude",
-        modeId: null,
-        model: null,
-        thinkingOptionId: null,
-        featureValues: {},
-      },
-    });
+    clearConsumedDraft();
 
-    await run(async () => {
-      calls += 1;
-    });
-
-    expect(calls).toBe(1);
     expect(useWorkspaceDraftSubmissionStore.getState().setupByDraftId).toEqual({});
     expect(useWorkspaceAttachmentsStore.getState().attachmentsByScope[SCOPE_KEY]).toBeUndefined();
     expect(draftText()).toBe("");
   });
 
-  it("preserves context changed while workspace creation was pending", async () => {
+  it("preserves newer setup without requiring a text edit", () => {
+    setSetup("/work/repo");
     const clearConsumedDraft = captureCleanup();
-    useWorkspaceAttachmentsStore.getState().addWorkspaceAttachment({
-      scopeKey: SCOPE_KEY,
-      attachment: {
-        kind: "chat_history",
-        id: "fork-context",
-        attachment: {
-          type: "text",
-          mimeType: "text/plain",
-          contextKind: "chat_history",
-          title: "Chat history",
-          text: "Preserve this context",
-        },
-        source: { serverId: "local", agentId: "agent-1" },
-      },
-    });
-    await run(async () => undefined, clearConsumedDraft);
+    setSetup("/work/new-repo");
+
+    clearConsumedDraft();
+
+    expect(useWorkspaceDraftSubmissionStore.getState().setupByDraftId[DRAFT_ID]?.setup.cwd).toBe(
+      "/work/new-repo",
+    );
+    expect(draftText()).toBe("do the thing");
+  });
+
+  it("preserves context changed while workspace creation was pending", () => {
+    const clearConsumedDraft = captureCleanup();
+    addContext();
+
+    clearConsumedDraft();
+
     expect(useWorkspaceAttachmentsStore.getState().attachmentsByScope[SCOPE_KEY]).toHaveLength(1);
     expect(draftText()).toBe("do the thing");
   });
 
-  it("leaves no pending submission behind, so no draft tab can create a second agent", async () => {
-    await run(async () => undefined);
+  it("keeps a prompt typed into the shared draft while creation was in flight", () => {
+    const clearConsumedDraft = captureCleanup();
+    saveDraft("a different idea");
 
-    expect(useWorkspaceDraftSubmissionStore.getState().pendingByDraftId).toEqual({});
-    expect(useCreateFlowStore.getState().pendingByDraftId).toEqual({});
-  });
-
-  it("rethrows a failed creation and keeps the prompt for retry", async () => {
-    const failure = new Error("daemon exploded");
-
-    await expect(
-      run(async () => {
-        throw failure;
-      }),
-    ).rejects.toThrow(failure);
-
-    expect(draftText()).toBe("do the thing");
-  });
-
-  it("keeps a prompt typed into the shared draft while creation was in flight", async () => {
-    await run(async () => {
-      saveDraft("a different idea");
-    });
+    clearConsumedDraft();
 
     expect(draftText()).toBe("a different idea");
   });
-  it("keeps a newer draft written while workspace creation was pending", async () => {
-    useWorkspaceAttachmentsStore.getState().addWorkspaceAttachment({
-      scopeKey: SCOPE_KEY,
-      attachment: {
-        kind: "chat_history",
-        id: "fork-context",
-        attachment: {
-          type: "text",
-          mimeType: "text/plain",
-          contextKind: "chat_history",
-          title: "Chat history",
-          text: "Preserve this context",
-        },
-        source: { serverId: "local", agentId: "agent-1" },
-      },
-    });
+
+  it("keeps unchanged context when a newer draft owns it", () => {
+    addContext();
     const clearConsumedDraft = captureCleanup();
     saveDraft("a newer workspace idea");
-    await run(async () => undefined, clearConsumedDraft);
+
+    clearConsumedDraft();
+
     expect(draftText()).toBe("a newer workspace idea");
     expect(useWorkspaceAttachmentsStore.getState().attachmentsByScope[SCOPE_KEY]).toHaveLength(1);
   });

packages/app/src/screens/new-workspace/background-handoff.ts

@@ -34,21 +34,3 @@ export function captureWorkspaceDraftCleanup(input: WorkspaceDraftCleanupInput):
     input.clearDraft("sent");
   };
 }
-
-/**
- * The workspace draft tab normally issues create_agent once it mounts on the destination screen.
- * When the user leaves the New workspace screen before creation resolves, that tab never mounts,
- * so this path issues the request itself. The agent tab opens on its own when the workspace is
- * next visited — `reconcileTabs` auto-opens agents.
- *
- * The caller must not write a pending entry to the draft-submission or create-flow stores on this
- * path: those drive the draft tab's auto-submit, and the daemon does not dedupe create_agent by
- * clientMessageId, so a leftover entry would create a second agent when the workspace is opened.
- */
-export async function createWorkspaceAgentInBackground(input: {
-  clearConsumedDraft: () => void;
-  createAgent: () => Promise<unknown>;
-}): Promise<void> {
-  await input.createAgent();
-  input.clearConsumedDraft();
-}

packages/app/src/workspace-creation/index.ts

@@ -6,4 +6,4 @@ export {
   isAgentReadyForPrompts,
   useIsAgentReadyForPrompts,
 } from "./runtime";
-export { submitWorkspaceChat, takeWorkspaceCreationError } from "./submit";
+export { submitWorkspaceChat, takeWorkspaceCreationFailure } from "./submit";

packages/app/src/workspace-creation/internal/model.test.ts

@@ -208,7 +208,7 @@ describe("creation lifecycle ownership", () => {
     expect(lifecycle).toEqual(["submitted:preparing", "complete:complete", "adopted:complete"]);
   });
 
-  it("does not regress receipt revisions or release overlays on empty directory snapshots", async () => {
+  it("keeps presentation stable across receipt updates and empty directory snapshots", async () => {
     const response = deferred<{}>();
     const observed =
       deferred<
@@ -234,11 +234,16 @@ describe("creation lifecycle ownership", () => {
       phase: "workspace_ready" as const,
       error: null,
     };
-    onEvent(event);
+    const pendingSnapshot = model.getSnapshot();
+    let publications = 0;
+    const stop = model.subscribe(() => publications++);
     onEvent({ ...event, revision: 1, phase: "accepted" });
+    onEvent(event);
+    expect(model.getSnapshot()).toBe(pendingSnapshot);
+    expect(publications).toBe(0);
+    stop();
     model.observeResources("host", new Set(), new Set());
     expect(model.getWorkspace("host", attempt.workspaceId)).toMatchObject({
-      revision: 3,
       workspaceObserved: false,
       agentObserved: false,
     });

packages/app/src/workspace-creation/model/index.ts

@@ -34,7 +34,6 @@ export interface CreationAttempt extends Omit<Submission, SubmissionHandlers> {
   agentId: string;
   clientMessageId: string;
   phase: "preparing" | "creating" | "complete";
-  revision: number;
   workspaceObserved: boolean;
   agentObserved: boolean;
 }
@@ -45,8 +44,6 @@ interface Result {
 }
 interface Operation extends Pick<Submission, "onComplete" | "onWorkspaceFailed" | "onAgentFailed"> {
   attempt: CreationAttempt;
-  prepare: Submission["prepare"] | null;
-  request?: CreateWorkspaceRequestOptions;
   /** The host's latest account of this creation. */
   creation?: CreationSnapshot;
   running?: Promise<void>;
@@ -84,7 +81,7 @@ export function createWorkspaceCreationModel(dependencies: {
   const listeners = new Set<() => void>();
   let snapshot: readonly CreationAttempt[] = [];
   const operations = new Map<string, Operation>();
-  // Current submissions remain owned until resolved or dismissed. Adopted creations
+  // Current submissions remain owned until settled and adopted. Adopted creations
   // only need a bounded receipt to absorb duplicate commands from retained forms.
   const receipts = new Map<string, CreationAttempt>();
   const intentKey = (serverId: string, key: string) =>
@@ -119,11 +116,9 @@ export function createWorkspaceCreationModel(dependencies: {
     patch(operation, { phase: "complete" });
     dependencies.onComplete?.(operation.attempt);
     const onComplete = operation.onComplete;
-    operation.prepare = null;
     operation.onComplete = undefined;
     operation.onWorkspaceFailed = undefined;
     operation.onAgentFailed = undefined;
-    operation.request = undefined;
     onComplete?.();
     retire(operation);
   }
@@ -141,30 +136,26 @@ export function createWorkspaceCreationModel(dependencies: {
       operation.onAgentFailed?.(operation.attempt, error, operation.creation?.workspace);
     else operation.onWorkspaceFailed?.(operation.attempt, error);
   }
-  async function run(operation: NonNullable<ReturnType<typeof operations.get>>) {
+  async function run(operation: Operation, prepare: Submission["prepare"]) {
     try {
-      if (!operation.request) {
-        if (!operation.prepare) throw new Error("Completed creation cannot be submitted again");
-        const prepared = await operation.prepare();
-        operation.request = captureCreationInput({
-          ...prepared,
-          workspaceId: operation.attempt.workspaceId,
-          idempotencyKey: operation.attempt.operationKey,
-          agent: {
-            ...prepared.agent,
-            agentId: operation.attempt.agentId,
-            clientMessageId: operation.attempt.clientMessageId,
-          },
-        });
-      }
+      const prepared = await prepare();
+      const request = captureCreationInput({
+        ...prepared,
+        workspaceId: operation.attempt.workspaceId,
+        idempotencyKey: operation.attempt.operationKey,
+        agent: {
+          ...prepared.agent,
+          agentId: operation.attempt.agentId,
+          clientMessageId: operation.attempt.clientMessageId,
+        },
+      });
       patch(operation, { phase: "creating" });
       const result = await dependencies.execute(
         {
-          ...captureCreationInput(operation.request),
+          ...request,
+          // CreationClient owns receipt ordering; progress alone does not change presentation.
           onEvent: (event) => {
-            if (event.revision <= operation.attempt.revision) return;
             operation.creation = event;
-            patch(operation, { revision: event.revision });
           },
         },
         operation.attempt.serverId,
@@ -181,8 +172,8 @@ export function createWorkspaceCreationModel(dependencies: {
       operation.running = undefined;
     }
   }
-  function start(operation: NonNullable<ReturnType<typeof operations.get>>) {
-    operation.running = Promise.resolve().then(() => run(operation));
+  function start(operation: Operation, prepare: Submission["prepare"]) {
+    operation.running = Promise.resolve().then(() => run(operation, prepare));
   }
   return {
     subscribe(listener: () => void) {
@@ -204,12 +195,10 @@ export function createWorkspaceCreationModel(dependencies: {
         key,
         ...dependencies.identities(),
         phase: "preparing",
-        revision: -1,
         workspaceObserved: false,
         agentObserved: false,
       };
       const operation: Operation = {
-        prepare,
         onComplete,
         onWorkspaceFailed,
         onAgentFailed,
@@ -218,7 +207,7 @@ export function createWorkspaceCreationModel(dependencies: {
       operations.set(key, operation);
       dependencies.onSubmitted?.(operation.attempt);
       publish();
-      start(operation);
+      start(operation, prepare);
       return operation.attempt;
     },
     observeResources(

packages/app/src/workspace-creation/submit.ts

@@ -64,14 +64,18 @@ interface SubmitWorkspaceChatInput extends DraftCleanupInput {
   persistPreferences: () => Promise<unknown>;
 }
 
-// Failed creations return to New workspace, whose error lives in screen state.
-const workspaceErrors = new Map<string, string>();
+interface WorkspaceCreationFailure {
+  error: string;
+  selectedItem: PickerItem | null;
+}
 
-/** The error of a creation that returned to the New workspace draft with this key, once. */
-export function takeWorkspaceCreationError(draftKey: string): string | null {
-  const error = workspaceErrors.get(draftKey) ?? null;
-  workspaceErrors.delete(draftKey);
-  return error;
+// The returned draft takes back its form selection and ordinary submit error together.
+const workspaceFailures = new Map<string, WorkspaceCreationFailure>();
+
+export function takeWorkspaceCreationFailure(draftKey: string): WorkspaceCreationFailure | null {
+  const failure = workspaceFailures.get(draftKey) ?? null;
+  workspaceFailures.delete(draftKey);
+  return failure;
 }
 
 /**
@@ -213,7 +217,7 @@ export function submitWorkspaceChat(input: SubmitWorkspaceChatInput) {
     projectName: project.projectName,
     text: payload.text,
     onWorkspaceFailed: (attempt, error) => {
-      workspaceErrors.set(input.draftKey, error);
+      workspaceFailures.set(input.draftKey, { error, selectedItem });
       if (input.returnDraftId) restoreSubmittedContext(input.returnDraftId, payload);
       // The New workspace composer owns the follow-ups until the form is submitted again.
       transferComposerFollowUps(serverId, attempt.agentId, input.draftKey);