From 93318a3e233340e105a138cfdfd1424c34897e1f Mon Sep 17 00:00:00 2001 From: Phil Haack Date: Fri, 12 Jun 2026 18:31:21 -0700 Subject: [PATCH 1/2] feat(worktree): reuse an existing worktree for the branch When starting a worktree task on a branch that already has a worktree checked out, confirm with the user and reuse that worktree for the task instead of silently creating a second, detached worktree at the same commit. - checkWorktreeBranch now reports an existingWorktreePath when a PostHog-managed worktree (under the worktree base path) is already on the branch. - The task-creation preflight prompts on that and passes reuseExistingWorktree through to the server, which registers the task against the existing worktree rather than running `git worktree add`. Stacked on the remote-only-branch checkout change; base that PR first. Generated-By: PostHog Code Task-Id: 1c8ffc23-3673-4b14-b5b7-50637fca8fb2 --- .../core/src/task-detail/taskCreationHost.ts | 1 + .../core/src/task-detail/taskCreationSaga.ts | 1 + packages/core/src/task-detail/taskInput.ts | 2 + packages/shared/src/task-creation-domain.ts | 3 + .../components/ExistingWorktreeDialog.tsx | 58 ++++++++++++++++ .../task-detail/hooks/useTaskCreation.ts | 22 ++++-- .../existingWorktreeConfirmStore.test.ts | 67 +++++++++++++++++++ .../stores/existingWorktreeConfirmStore.ts | 47 +++++++++++++ packages/ui/src/router/routes/__root.tsx | 4 ++ .../src/services/workspace/schemas.ts | 6 ++ .../src/services/workspace/workspace.ts | 63 ++++++++++++++--- 11 files changed, 258 insertions(+), 16 deletions(-) create mode 100644 packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx create mode 100644 packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.test.ts create mode 100644 packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.ts diff --git a/packages/core/src/task-detail/taskCreationHost.ts b/packages/core/src/task-detail/taskCreationHost.ts index 7329f050b..cc815bafe 100644 --- a/packages/core/src/task-detail/taskCreationHost.ts +++ b/packages/core/src/task-detail/taskCreationHost.ts @@ -16,6 +16,7 @@ export interface CreateWorkspaceArgs { mode: WorkspaceMode; branch?: string; allowRemoteBranchCheckout?: boolean; + reuseExistingWorktree?: boolean; } export interface CreatedWorkspaceInfo { diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index 8c55ed59c..f9561655f 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -100,6 +100,7 @@ export class TaskCreationSaga extends Saga< mode: workspaceMode, branch: branch ?? undefined, allowRemoteBranchCheckout: input.allowRemoteBranchCheckout, + reuseExistingWorktree: input.reuseExistingWorktree, }); }, rollback: async () => { diff --git a/packages/core/src/task-detail/taskInput.ts b/packages/core/src/task-detail/taskInput.ts index 3f3ed7f35..d412a4200 100644 --- a/packages/core/src/task-detail/taskInput.ts +++ b/packages/core/src/task-detail/taskInput.ts @@ -10,6 +10,7 @@ export interface PrepareTaskInputOptions { workspaceMode: WorkspaceMode; branch?: string | null; allowRemoteBranchCheckout?: boolean; + reuseExistingWorktree?: boolean; executionMode?: ExecutionMode; adapter?: "claude" | "codex"; model?: string; @@ -41,6 +42,7 @@ export function prepareTaskInput( workspaceMode: options.workspaceMode, branch: options.branch, allowRemoteBranchCheckout: options.allowRemoteBranchCheckout, + reuseExistingWorktree: options.reuseExistingWorktree, executionMode: options.executionMode, adapter: options.adapter, model: options.model, diff --git a/packages/shared/src/task-creation-domain.ts b/packages/shared/src/task-creation-domain.ts index 8bb489da2..fff65f947 100644 --- a/packages/shared/src/task-creation-domain.ts +++ b/packages/shared/src/task-creation-domain.ts @@ -23,6 +23,9 @@ export interface TaskCreationInput { // When the branch exists only on the remote, opt in to fetching and checking // it out locally into the worktree (set after the user confirms). allowRemoteBranchCheckout?: boolean; + // When a worktree is already checked out on the branch, opt in to reusing it + // for this task instead of creating a new one (set after the user confirms). + reuseExistingWorktree?: boolean; githubIntegrationId?: number; githubUserIntegrationId?: string; executionMode?: ExecutionMode; diff --git a/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx b/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx new file mode 100644 index 000000000..4e184660a --- /dev/null +++ b/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx @@ -0,0 +1,58 @@ +import { FolderOpen } from "@phosphor-icons/react"; +import { AlertDialog, Button, Code, Flex } from "@radix-ui/themes"; +import { useExistingWorktreeConfirmStore } from "../stores/existingWorktreeConfirmStore"; + +/** + * Globally-mounted confirmation shown when a user starts a worktree task on a + * branch that already has a worktree checked out. Confirming reuses that + * worktree for the task instead of creating a new one. + */ +export function ExistingWorktreeDialog() { + const isOpen = useExistingWorktreeConfirmStore((s) => s.isOpen); + const branch = useExistingWorktreeConfirmStore((s) => s.branch); + const worktreePath = useExistingWorktreeConfirmStore((s) => s.worktreePath); + const accept = useExistingWorktreeConfirmStore((s) => s.accept); + const cancel = useExistingWorktreeConfirmStore((s) => s.cancel); + + return ( + { + if (!open) cancel(); + }} + > + + + + + Worktree already exists + + + + A worktree is already checked out on{" "} + {branch ? {branch} : "this branch"} + {worktreePath ? ( + <> + {" "} + at {worktreePath} + + ) : null} + . Continue and use that worktree for this task? + + + + + + + + + + + + + ); +} diff --git a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts index 980d4165b..367c892c0 100644 --- a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts +++ b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts @@ -41,6 +41,7 @@ import { useSettingsStore } from "../../settings/settingsStore"; import { useCreateTask } from "../../tasks/useTaskCrudMutations"; import { useTourStore } from "../../tour/tourStore"; import { createFirstTaskTour } from "../../tour/tours/createFirstTaskTour"; +import { useExistingWorktreeConfirmStore } from "../stores/existingWorktreeConfirmStore"; import { useRemoteBranchConfirmStore } from "../stores/remoteBranchConfirmStore"; const log = logger.scope("task-creation"); @@ -184,18 +185,28 @@ export function useTaskCreation({ return false; } - // If the chosen worktree branch only exists on the remote, confirm before - // fetching and checking it out locally. Done before the pending view so - // the dialog (and a cancel) don't leave a half-started task on screen. + // Confirm a couple of worktree branch situations before starting the + // task. Done before the pending view so a dialog (and a cancel) don't + // leave a half-started task on screen. An existing worktree takes + // priority: a branch with one already checked out also exists locally. let allowRemoteBranchCheckout = false; + let reuseExistingWorktree = false; if (workspaceMode === "worktree" && branch && selectedDirectory) { try { - const { status } = + const { status, existingWorktreePath } = await hostClient.workspace.checkWorktreeBranch.query({ mainRepoPath: selectedDirectory, branch, }); - if (status === "remote-only") { + if (existingWorktreePath) { + const confirmed = await useExistingWorktreeConfirmStore + .getState() + .confirm(branch, existingWorktreePath); + if (!confirmed) { + return false; + } + reuseExistingWorktree = true; + } else if (status === "remote-only") { const confirmed = await useRemoteBranchConfirmStore .getState() .confirm(branch); @@ -250,6 +261,7 @@ export function useTaskCreation({ workspaceMode, branch, allowRemoteBranchCheckout, + reuseExistingWorktree, executionMode, adapter, model, diff --git a/packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.test.ts b/packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.test.ts new file mode 100644 index 000000000..3227b490d --- /dev/null +++ b/packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.test.ts @@ -0,0 +1,67 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { useExistingWorktreeConfirmStore } from "./existingWorktreeConfirmStore"; + +describe("existingWorktreeConfirmStore", () => { + beforeEach(() => { + useExistingWorktreeConfirmStore.setState({ + isOpen: false, + branch: null, + worktreePath: null, + resolve: null, + }); + }); + + it("starts closed", () => { + const state = useExistingWorktreeConfirmStore.getState(); + expect(state.isOpen).toBe(false); + expect(state.branch).toBeNull(); + expect(state.worktreePath).toBeNull(); + }); + + it("confirm opens the dialog with the branch and path", () => { + void useExistingWorktreeConfirmStore + .getState() + .confirm("feature/x", "/wt/feature-x"); + const state = useExistingWorktreeConfirmStore.getState(); + expect(state.isOpen).toBe(true); + expect(state.branch).toBe("feature/x"); + expect(state.worktreePath).toBe("/wt/feature-x"); + }); + + it("accept resolves the pending promise with true and closes", async () => { + const promise = useExistingWorktreeConfirmStore + .getState() + .confirm("feature/x", "/wt/feature-x"); + useExistingWorktreeConfirmStore.getState().accept(); + await expect(promise).resolves.toBe(true); + const state = useExistingWorktreeConfirmStore.getState(); + expect(state.isOpen).toBe(false); + expect(state.branch).toBeNull(); + expect(state.worktreePath).toBeNull(); + expect(state.resolve).toBeNull(); + }); + + it("cancel resolves the pending promise with false and closes", async () => { + const promise = useExistingWorktreeConfirmStore + .getState() + .confirm("feature/x", "/wt/feature-x"); + useExistingWorktreeConfirmStore.getState().cancel(); + await expect(promise).resolves.toBe(false); + expect(useExistingWorktreeConfirmStore.getState().isOpen).toBe(false); + }); + + it("opening a second dialog resolves the first as cancelled", async () => { + const first = useExistingWorktreeConfirmStore + .getState() + .confirm("first", "/wt/first"); + const second = useExistingWorktreeConfirmStore + .getState() + .confirm("second", "/wt/second"); + await expect(first).resolves.toBe(false); + expect(useExistingWorktreeConfirmStore.getState().worktreePath).toBe( + "/wt/second", + ); + useExistingWorktreeConfirmStore.getState().accept(); + await expect(second).resolves.toBe(true); + }); +}); diff --git a/packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.ts b/packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.ts new file mode 100644 index 000000000..1a99ee706 --- /dev/null +++ b/packages/ui/src/features/task-detail/stores/existingWorktreeConfirmStore.ts @@ -0,0 +1,47 @@ +import { create } from "zustand"; + +interface ExistingWorktreeConfirmState { + isOpen: boolean; + branch: string | null; + worktreePath: string | null; + resolve: ((confirmed: boolean) => void) | null; +} + +interface ExistingWorktreeConfirmActions { + /** + * Opens the confirmation dialog for a branch that already has a worktree and + * resolves to the user's choice. `true` reuses the existing worktree for the + * task; `false` cancels. + */ + confirm: (branch: string, worktreePath: string) => Promise; + accept: () => void; + cancel: () => void; +} + +type ExistingWorktreeConfirmStore = ExistingWorktreeConfirmState & + ExistingWorktreeConfirmActions; + +export const useExistingWorktreeConfirmStore = + create()((set, get) => ({ + isOpen: false, + branch: null, + worktreePath: null, + resolve: null, + + confirm: (branch, worktreePath) => + new Promise((resolve) => { + // Resolve any dialog already waiting before replacing it. + get().resolve?.(false); + set({ isOpen: true, branch, worktreePath, resolve }); + }), + + accept: () => { + get().resolve?.(true); + set({ isOpen: false, branch: null, worktreePath: null, resolve: null }); + }, + + cancel: () => { + get().resolve?.(false); + set({ isOpen: false, branch: null, worktreePath: null, resolve: null }); + }, + })); diff --git a/packages/ui/src/router/routes/__root.tsx b/packages/ui/src/router/routes/__root.tsx index 81f8fa75e..1f89db15b 100644 --- a/packages/ui/src/router/routes/__root.tsx +++ b/packages/ui/src/router/routes/__root.tsx @@ -23,6 +23,7 @@ import { ProjectSwitcher } from "@posthog/ui/features/sidebar/components/Project import { SidebarNavSection } from "@posthog/ui/features/sidebar/components/SidebarNavSection"; import { useSidebarData } from "@posthog/ui/features/sidebar/useSidebarData"; import { useVisualTaskOrder } from "@posthog/ui/features/sidebar/useVisualTaskOrder"; +import { ExistingWorktreeDialog } from "@posthog/ui/features/task-detail/components/ExistingWorktreeDialog"; import { RemoteBranchCheckoutDialog } from "@posthog/ui/features/task-detail/components/RemoteBranchCheckoutDialog"; import { useTasks } from "@posthog/ui/features/tasks/useTasks"; import { TourOverlay } from "@posthog/ui/features/tour/components/TourOverlay"; @@ -252,6 +253,7 @@ function RootLayout() { /> {billingEnabled && } + {import.meta.env.DEV && ( @@ -276,6 +278,7 @@ function RootLayout() { /> {billingEnabled && } + {import.meta.env.DEV && ( @@ -319,6 +322,7 @@ function RootLayout() { {billingEnabled && } + {import.meta.env.DEV && ( diff --git a/packages/workspace-server/src/services/workspace/schemas.ts b/packages/workspace-server/src/services/workspace/schemas.ts index 973c4628e..b1a93bed2 100644 --- a/packages/workspace-server/src/services/workspace/schemas.ts +++ b/packages/workspace-server/src/services/workspace/schemas.ts @@ -26,6 +26,9 @@ export const createWorkspaceInput = z // When set, a worktree branch that exists only on the remote is fetched and // checked out locally instead of failing. Gated behind a user confirmation. allowRemoteBranchCheckout: z.boolean().optional(), + // When set, an existing worktree already checked out on the branch is reused + // for the task instead of creating a new one. Gated behind a confirmation. + reuseExistingWorktree: z.boolean().optional(), }) .refine( (data) => @@ -46,6 +49,9 @@ export const checkWorktreeBranchOutput = z.object({ // "local": branch exists locally. "remote-only": exists on the remote but not // locally. "missing": found neither locally nor on the remote. status: z.enum(["trunk", "local", "remote-only", "missing"]), + // Path of an existing worktree already checked out on this branch, if any. + // Set independently of `status`; the renderer offers to reuse it. + existingWorktreePath: z.string().nullable(), }); export const reconcileCloudWorkspacesInput = z.object({ diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index 610238351..a8c231a53 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -430,28 +430,59 @@ export class WorkspaceService extends TypedEventEmitter ): Promise { const { mainRepoPath, branch } = options; - const defaultBranch = await getDefaultBranch(mainRepoPath, { - abortSignal: signal, - }).catch(() => - getCurrentBranch(mainRepoPath, { abortSignal: signal }).then( - (b) => b ?? "main", + const [existingWorktree, defaultBranch] = await Promise.all([ + this.findExistingWorktreeForBranch(mainRepoPath, branch), + getDefaultBranch(mainRepoPath, { abortSignal: signal }).catch(() => + getCurrentBranch(mainRepoPath, { abortSignal: signal }).then( + (b) => b ?? "main", + ), ), - ); + ]); + const existingWorktreePath = existingWorktree?.worktreePath ?? null; if (branch === defaultBranch) { - return { status: "trunk" }; + return { status: "trunk", existingWorktreePath }; } if (await branchExists(mainRepoPath, branch, { abortSignal: signal })) { - return { status: "local" }; + return { status: "local", existingWorktreePath }; } if ( await remoteBranchExists(mainRepoPath, branch, { abortSignal: signal }) ) { - return { status: "remote-only" }; + return { status: "remote-only", existingWorktreePath }; } - return { status: "missing" }; + return { status: "missing", existingWorktreePath }; + } + + /** + * Finds a PostHog-managed worktree (under the worktree base path) already + * checked out on `branch`, returning it as a `WorktreeInfo` ready to reuse, or + * null when none exists. Only base-path worktrees are considered because the + * task<->worktree association re-derives paths from the base path and name. + */ + private async findExistingWorktreeForBranch( + mainRepoPath: string, + branch: string, + ): Promise { + const worktreeBasePath = this.workspaceSettings.getWorktreeLocation(); + const twigWorktrees = await listTwigWorktrees( + mainRepoPath, + worktreeBasePath, + ); + const match = twigWorktrees.find((wt) => wt.branch === branch); + if (!match) return null; + + // baseBranch/createdAt are unknown for an already-existing worktree; mirror + // WorktreeManager.listWorktrees() and leave them empty rather than fabricate. + return { + worktreePath: match.worktreePath, + worktreeName: path.basename(path.dirname(match.worktreePath)), + branchName: branch, + baseBranch: "", + createdAt: "", + }; } async createWorkspace(options: CreateWorkspaceInput): Promise { @@ -485,6 +516,7 @@ export class WorkspaceService extends TypedEventEmitter branch, useExistingBranch, allowRemoteBranchCheckout, + reuseExistingWorktree, } = options; const existingWorkspace = await this.getWorkspaceInfo(taskId); @@ -582,7 +614,16 @@ export class WorkspaceService extends TypedEventEmitter this.provisioning.emitOutput(taskId, data); }; - if (isTrunkSelected) { + const existingWorktree = reuseExistingWorktree + ? await this.findExistingWorktreeForBranch(mainRepoPath, selectedBranch) + : null; + + if (existingWorktree) { + this.log.info( + `Reusing existing worktree for branch ${selectedBranch}: ${existingWorktree.worktreePath}`, + ); + worktree = existingWorktree; + } else if (isTrunkSelected) { this.log.info( `Trunk branch selected (${defaultBranch}), creating detached worktree`, ); From f73ce30526081be49643345f91818fa3b61ceb91 Mon Sep 17 00:00:00 2001 From: Phil Haack Date: Tue, 16 Jun 2026 09:16:40 -0700 Subject: [PATCH 2/2] Fix worktree reuse: block occupied worktrees, add tests, warn on delete --- .../components/ExistingWorktreeDialog.tsx | 19 ++- .../task-detail/hooks/useTaskCreation.ts | 23 ++- .../src/services/workspace/schemas.ts | 10 +- .../src/services/workspace/workspace.test.ts | 155 +++++++++++++++++- .../src/services/workspace/workspace.ts | 52 +++++- 5 files changed, 243 insertions(+), 16 deletions(-) diff --git a/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx b/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx index 4e184660a..cfd005e46 100644 --- a/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx +++ b/packages/ui/src/features/task-detail/components/ExistingWorktreeDialog.tsx @@ -1,5 +1,5 @@ -import { FolderOpen } from "@phosphor-icons/react"; -import { AlertDialog, Button, Code, Flex } from "@radix-ui/themes"; +import { FolderOpen, Warning } from "@phosphor-icons/react"; +import { AlertDialog, Button, Code, Flex, Text } from "@radix-ui/themes"; import { useExistingWorktreeConfirmStore } from "../stores/existingWorktreeConfirmStore"; /** @@ -40,9 +40,22 @@ export function ExistingWorktreeDialog() { . Continue and use that worktree for this task? + + + + Deleting this task later removes the worktree and any uncommitted + work in it, even though it existed beforehand. + + + - diff --git a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts index 367c892c0..f78602307 100644 --- a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts +++ b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts @@ -39,6 +39,7 @@ import { useTaskInputHistoryStore } from "../../message-editor/taskInputHistoryS import type { EditorHandle } from "../../message-editor/types"; import { useSettingsStore } from "../../settings/settingsStore"; import { useCreateTask } from "../../tasks/useTaskCrudMutations"; +import { useTasks } from "../../tasks/useTasks"; import { useTourStore } from "../../tour/tourStore"; import { createFirstTaskTour } from "../../tour/tours/createFirstTaskTour"; import { useExistingWorktreeConfirmStore } from "../stores/existingWorktreeConfirmStore"; @@ -166,6 +167,8 @@ export function useTaskCreation({ ); const { invalidateTasks } = useCreateTask(); const { isOnline } = useConnectivity(); + // Used to name the task occupying a branch's worktree when reuse is blocked. + const { data: tasks } = useTasks(); const hasRequiredPath = workspaceMode === "cloud" ? !!selectedRepository : !!selectedDirectory; @@ -187,17 +190,30 @@ export function useTaskCreation({ // Confirm a couple of worktree branch situations before starting the // task. Done before the pending view so a dialog (and a cancel) don't - // leave a half-started task on screen. An existing worktree takes - // priority: a branch with one already checked out also exists locally. + // leave a half-started task on screen. Reusing an existing worktree takes + // priority over checking out a remote branch. let allowRemoteBranchCheckout = false; let reuseExistingWorktree = false; if (workspaceMode === "worktree" && branch && selectedDirectory) { try { - const { status, existingWorktreePath } = + const { status, existingWorktreePath, existingWorktreeTaskId } = await hostClient.workspace.checkWorktreeBranch.query({ mainRepoPath: selectedDirectory, branch, }); + if (existingWorktreeTaskId) { + // The branch's worktree already belongs to another task. Don't + // create a duplicate; point the user at the task using it. + const occupant = tasks?.find( + (t) => t.id === existingWorktreeTaskId, + ); + toast.error("Worktree already in use", { + description: occupant + ? `${branch} already has a worktree used by "${occupant.title}". Open that task to keep working there.` + : `${branch} already has a worktree used by another task.`, + }); + return false; + } if (existingWorktreePath) { const confirmed = await useExistingWorktreeConfirmStore .getState() @@ -365,6 +381,7 @@ export function useTaskCreation({ onTaskCreated, hostClient, taskService, + tasks, ], ); diff --git a/packages/workspace-server/src/services/workspace/schemas.ts b/packages/workspace-server/src/services/workspace/schemas.ts index b1a93bed2..340967b04 100644 --- a/packages/workspace-server/src/services/workspace/schemas.ts +++ b/packages/workspace-server/src/services/workspace/schemas.ts @@ -49,9 +49,15 @@ export const checkWorktreeBranchOutput = z.object({ // "local": branch exists locally. "remote-only": exists on the remote but not // locally. "missing": found neither locally nor on the remote. status: z.enum(["trunk", "local", "remote-only", "missing"]), - // Path of an existing worktree already checked out on this branch, if any. - // Set independently of `status`; the renderer offers to reuse it. + // Path of an *unused* worktree already checked out on this branch, if any. + // Set only when no task is associated with it; the renderer offers to reuse + // it. Null when there is no managed worktree, or when one exists but is + // already taken by a task (see existingWorktreeTaskId). existingWorktreePath: z.string().nullable(), + // Id of the task already using the managed worktree on this branch, if any. + // When set, the renderer blocks reuse and points the user at that task. + // Mutually exclusive with existingWorktreePath. + existingWorktreeTaskId: z.string().nullable(), }); export const reconcileCloudWorkspacesInput = z.object({ diff --git a/packages/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 8161b1018..ff7ea3617 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -6,6 +6,7 @@ import { branchExists, getCurrentBranch, getDefaultBranch, + hasTrackedFiles, remoteBranchExists, } from "@posthog/git/queries"; import type { IAnalytics } from "@posthog/platform/analytics"; @@ -17,12 +18,14 @@ import { createMockWorkspaceRepository } from "../../db/repositories/workspace-r import { createMockWorktreeRepository } from "../../db/repositories/worktree-repository.mock"; import type { ProcessTrackingService } from "../process-tracking/process-tracking"; import type { SuspensionService } from "../suspension/suspension"; +import { listTwigWorktrees } from "../worktree-query/worktree-query"; import type { WorkspaceAgent, WorkspaceFileWatcher, WorkspaceFocus, WorkspaceProvisioning, } from "./ports"; +import type { CreateWorkspaceInput } from "./schemas"; import { WorkspaceService, WorkspaceServiceEvent } from "./workspace"; vi.mock("@posthog/git/queries", async (importOriginal) => { @@ -33,6 +36,7 @@ vi.mock("@posthog/git/queries", async (importOriginal) => { getCurrentBranch: vi.fn(), branchExists: vi.fn(), remoteBranchExists: vi.fn(), + hasTrackedFiles: vi.fn(), }; }); @@ -44,6 +48,7 @@ vi.mock("../worktree-query/worktree-query", async (importOriginal) => { return { ...actual, deleteWorktree: vi.fn(async () => {}), + listTwigWorktrees: vi.fn(), }; }); @@ -277,6 +282,7 @@ describe("WorkspaceService", () => { vi.mocked(getCurrentBranch).mockResolvedValue("main"); vi.mocked(branchExists).mockResolvedValue(false); vi.mocked(remoteBranchExists).mockResolvedValue(false); + vi.mocked(listTwigWorktrees).mockResolvedValue([]); }); it.each([ @@ -297,17 +303,162 @@ describe("WorkspaceService", () => { expect( await service.checkWorktreeBranch({ mainRepoPath, branch }), - ).toEqual({ status }); + ).toEqual({ + status, + existingWorktreePath: null, + existingWorktreeTaskId: null, + }); }, ); + it("offers an unused worktree on the branch for reuse", async () => { + vi.mocked(branchExists).mockResolvedValue(true); + vi.mocked(listTwigWorktrees).mockResolvedValue([ + { + worktreePath: "/tmp/worktrees/feature-x/repo", + head: "abc123", + branch: "feature/x", + }, + ]); + + expect( + await service.checkWorktreeBranch({ + mainRepoPath, + branch: "feature/x", + }), + ).toEqual({ + status: "local", + existingWorktreePath: "/tmp/worktrees/feature-x/repo", + existingWorktreeTaskId: null, + }); + }); + + it("reports the occupying task instead of offering reuse when the worktree is taken", async () => { + vi.mocked(branchExists).mockResolvedValue(true); + vi.mocked(listTwigWorktrees).mockResolvedValue([ + { + worktreePath: "/tmp/worktrees/feature-x/repo", + head: "abc123", + branch: "feature/x", + }, + ]); + // Associate a task with that worktree path so getWorktreeTasks finds it. + // deriveWorktreePath (new layout) reconstructs //, so + // name "feature-x" + repo "repo" resolves to the path above. + const folder = mocks.repositoryRepo.create({ path: mainRepoPath }); + const occupantWorkspace = mocks.workspaceRepo.create({ + taskId: "occupant-task", + repositoryId: folder.id, + mode: "worktree", + }); + mocks.worktreeRepo.create({ + workspaceId: occupantWorkspace.id, + name: "feature-x", + path: "/tmp/worktrees/feature-x/repo", + }); + + expect( + await service.checkWorktreeBranch({ + mainRepoPath, + branch: "feature/x", + }), + ).toEqual({ + status: "local", + existingWorktreePath: null, + existingWorktreeTaskId: "occupant-task", + }); + }); + it("falls back to the current branch as trunk when getDefaultBranch fails", async () => { vi.mocked(getDefaultBranch).mockRejectedValue(new Error("no remote")); vi.mocked(getCurrentBranch).mockResolvedValue("develop"); expect( await service.checkWorktreeBranch({ mainRepoPath, branch: "develop" }), - ).toEqual({ status: "trunk" }); + ).toEqual({ + status: "trunk", + existingWorktreePath: null, + existingWorktreeTaskId: null, + }); + }); + }); + + describe("createWorkspace (worktree reuse)", () => { + const mainRepoPath = "/tmp/repo"; + + beforeEach(() => { + vi.mocked(getDefaultBranch).mockResolvedValue("main"); + vi.mocked(getCurrentBranch).mockResolvedValue("main"); + // The reuse success path checks whether the worktree has files; pretend it + // does so the empty-workspace warning branch (and its fs reads) is skipped. + vi.mocked(hasTrackedFiles).mockResolvedValue(true); + }); + + function reuseInput(taskId: string): CreateWorkspaceInput { + return { + taskId, + mainRepoPath, + folderId: "folder-1", + folderPath: mainRepoPath, + mode: "worktree", + branch: "feature/x", + reuseExistingWorktree: true, + }; + } + + it("reuses an unused worktree and stores its layout-aware name (legacy layout)", async () => { + // Legacy layout is //, so the name is the final segment + // ("feature-x"), not the parent dir. No task owns it, so reuse proceeds and + // the recovered name is persisted via worktreeRepo.create. + vi.mocked(listTwigWorktrees).mockResolvedValue([ + { + worktreePath: "/tmp/worktrees/repo/feature-x", + head: "abc123", + branch: "feature/x", + }, + ]); + const createWorktree = vi.spyOn(mocks.worktreeRepo, "create"); + + const workspace = await service.createWorkspace(reuseInput("new-task")); + + expect(workspace.worktree?.worktreeName).toBe("feature-x"); + expect(workspace.worktree?.worktreePath).toBe( + "/tmp/worktrees/repo/feature-x", + ); + expect(createWorktree).toHaveBeenCalledWith( + expect.objectContaining({ + name: "feature-x", + path: "/tmp/worktrees/repo/feature-x", + }), + ); + }); + + it("fails the create step when the worktree was claimed between preflight and create", async () => { + vi.mocked(listTwigWorktrees).mockResolvedValue([ + { + worktreePath: "/tmp/worktrees/feature-x/repo", + head: "abc123", + branch: "feature/x", + }, + ]); + // Associate another task with that worktree path so the re-check's + // getWorktreeTasks finds an occupant (same fixture as the checkWorktreeBranch + // occupied case: new layout // round-trips to the path). + const folder = mocks.repositoryRepo.create({ path: mainRepoPath }); + const occupantWorkspace = mocks.workspaceRepo.create({ + taskId: "occupant-task", + repositoryId: folder.id, + mode: "worktree", + }); + mocks.worktreeRepo.create({ + workspaceId: occupantWorkspace.id, + name: "feature-x", + path: "/tmp/worktrees/feature-x/repo", + }); + + await expect( + service.createWorkspace(reuseInput("new-task")), + ).rejects.toThrow(/already used by task occupant-task/); }); }); diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index a8c231a53..caa9ea231 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/packages/workspace-server/src/services/workspace/workspace.ts @@ -438,22 +438,41 @@ export class WorkspaceService extends TypedEventEmitter ), ), ]); - const existingWorktreePath = existingWorktree?.worktreePath ?? null; + // Reuse is only offered for an *unused* worktree. If a task already holds + // the worktree on this branch, report that task instead so the renderer can + // block the duplicate and point the user at it. + let worktree: { + existingWorktreePath: string | null; + existingWorktreeTaskId: string | null; + } = { existingWorktreePath: null, existingWorktreeTaskId: null }; + if (existingWorktree) { + const [occupant] = this.getWorktreeTasks(existingWorktree.worktreePath); + worktree = occupant + ? { + existingWorktreePath: null, + existingWorktreeTaskId: occupant.taskId, + } + : { + existingWorktreePath: existingWorktree.worktreePath, + existingWorktreeTaskId: null, + }; + } + if (branch === defaultBranch) { - return { status: "trunk", existingWorktreePath }; + return { status: "trunk", ...worktree }; } if (await branchExists(mainRepoPath, branch, { abortSignal: signal })) { - return { status: "local", existingWorktreePath }; + return { status: "local", ...worktree }; } if ( await remoteBranchExists(mainRepoPath, branch, { abortSignal: signal }) ) { - return { status: "remote-only", existingWorktreePath }; + return { status: "remote-only", ...worktree }; } - return { status: "missing", existingWorktreePath }; + return { status: "missing", ...worktree }; } /** @@ -474,11 +493,22 @@ export class WorkspaceService extends TypedEventEmitter const match = twigWorktrees.find((wt) => wt.branch === branch); if (!match) return null; + // Recover the worktree name from its path, layout-aware so the stored name + // round-trips through deriveWorktreePath. New layout is `//` + // (name is the parent dir); legacy is `//` (name is the + // final segment). Distinguish by whether the final segment is the repo. + const repoName = path.basename(mainRepoPath); + const finalSegment = path.basename(match.worktreePath); + const worktreeName = + finalSegment === repoName + ? path.basename(path.dirname(match.worktreePath)) + : finalSegment; + // baseBranch/createdAt are unknown for an already-existing worktree; mirror // WorktreeManager.listWorktrees() and leave them empty rather than fabricate. return { worktreePath: match.worktreePath, - worktreeName: path.basename(path.dirname(match.worktreePath)), + worktreeName, branchName: branch, baseBranch: "", createdAt: "", @@ -618,7 +648,17 @@ export class WorkspaceService extends TypedEventEmitter ? await this.findExistingWorktreeForBranch(mainRepoPath, selectedBranch) : null; + // Reuse only an unused worktree. The renderer already gates on this, but + // re-check here so a lost race (the worktree got claimed between preflight + // and now) fails the saga step rather than sharing one worktree across two + // tasks. if (existingWorktree) { + const [occupant] = this.getWorktreeTasks(existingWorktree.worktreePath); + if (occupant) { + throw new Error( + `Worktree at ${existingWorktree.worktreePath} is already used by task ${occupant.taskId}`, + ); + } this.log.info( `Reusing existing worktree for branch ${selectedBranch}: ${existingWorktree.worktreePath}`, );