Use stored worktree path when creating task terminals (#3074)

Author: haacked

packages/workspace-server/src/services/shell/shell.test.ts

@@ -1,4 +1,10 @@
-import { describe, expect, it, vi } from "vitest";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import path from "node:path";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { createMockRepositoryRepository } from "../../db/repositories/repository-repository.mock";
+import { createMockWorkspaceRepository } from "../../db/repositories/workspace-repository.mock";
+import { createMockWorktreeRepository } from "../../db/repositories/worktree-repository.mock";
 import { ShellEvent } from "./schemas";
 
 const mockPty = vi.hoisted(() => ({
@@ -7,6 +13,13 @@ const mockPty = vi.hoisted(() => ({
 
 vi.mock("node-pty", () => mockPty);
 
+const mockGitQueries = vi.hoisted(() => ({
+  getCurrentBranch: vi.fn(async () => "feature-branch"),
+  getDefaultBranch: vi.fn(async () => "main"),
+}));
+
+vi.mock("@posthog/git/queries", () => mockGitQueries);
+
 import { ShellService } from "./shell";
 
 function createMockPtyProcess() {
@@ -21,7 +34,11 @@ function createMockPtyProcess() {
   };
 }
 
-function createService() {
+function createService(overrides?: {
+  repositoryRepo?: unknown;
+  workspaceRepo?: unknown;
+  worktreeRepo?: unknown;
+}) {
   const processTracking = {
     register: vi.fn(),
     unregister: vi.fn(),
@@ -37,9 +54,9 @@ function createService() {
   };
   const service = new ShellService(
     processTracking as never,
-    {} as never,
-    {} as never,
-    {} as never,
+    (overrides?.repositoryRepo ?? {}) as never,
+    (overrides?.workspaceRepo ?? {}) as never,
+    (overrides?.worktreeRepo ?? {}) as never,
     { getWorktreeLocation: vi.fn(() => "/tmp/worktrees") } as never,
     logger as never,
   );
@@ -68,3 +85,72 @@ describe("ShellService.destroy", () => {
     expect(() => service.destroy("nonexistent")).not.toThrow();
   });
 });
+
+describe("ShellService.createSession workspace env", () => {
+  function createWorktreeTaskService(worktreePath: string) {
+    const repositoryRepo = createMockRepositoryRepository();
+    const workspaceRepo = createMockWorkspaceRepository();
+    const worktreeRepo = createMockWorktreeRepository();
+    const repo = repositoryRepo.create({ path: "/repos/code" });
+    const workspace = workspaceRepo.create({
+      taskId: "task-1",
+      repositoryId: repo.id,
+      mode: "worktree",
+    });
+    worktreeRepo.create({
+      workspaceId: workspace.id,
+      name: "spawn-tasks",
+      path: worktreePath,
+    });
+    return createService({ repositoryRepo, workspaceRepo, worktreeRepo });
+  }
+
+  function spawnedEnv(): Record<string, string | undefined> {
+    return mockPty.spawn.mock.calls[0][2].env;
+  }
+
+  beforeEach(() => {
+    mockPty.spawn.mockReset();
+    mockPty.spawn.mockReturnValue(createMockPtyProcess());
+    mockGitQueries.getCurrentBranch.mockResolvedValue("feature-branch");
+    mockGitQueries.getDefaultBranch.mockResolvedValue("main");
+  });
+
+  it("uses the stored worktree path when it exists on disk", async () => {
+    const tempDir = mkdtempSync(path.join(tmpdir(), "shell-test-"));
+    try {
+      const { service } = createWorktreeTaskService(tempDir);
+
+      await service.createSession({ sessionId: "session-1", taskId: "task-1" });
+
+      expect(spawnedEnv().POSTHOG_CODE_WORKSPACE_PATH).toBe(tempDir);
+      expect(mockGitQueries.getCurrentBranch).toHaveBeenCalledWith(tempDir);
+    } finally {
+      rmSync(tempDir, { recursive: true, force: true });
+    }
+  });
+
+  it("falls back to the derived path when the stored path is missing", async () => {
+    const { service } = createWorktreeTaskService("/does/not/exist");
+
+    await service.createSession({ sessionId: "session-1", taskId: "task-1" });
+
+    const derivedPath = path.join("/tmp/worktrees", "spawn-tasks", "code");
+    expect(spawnedEnv().POSTHOG_CODE_WORKSPACE_PATH).toBe(derivedPath);
+    expect(mockGitQueries.getCurrentBranch).toHaveBeenCalledWith(derivedPath);
+  });
+
+  it("still creates the shell when env construction fails", async () => {
+    mockGitQueries.getDefaultBranch.mockRejectedValue(
+      new Error("Cannot use simple-git on a directory that does not exist"),
+    );
+    const { service } = createWorktreeTaskService("/does/not/exist");
+
+    await expect(
+      service.createSession({ sessionId: "session-1", taskId: "task-1" }),
+    ).resolves.toBeDefined();
+
+    expect(mockPty.spawn).toHaveBeenCalledTimes(1);
+    expect(spawnedEnv().POSTHOG_CODE_WORKSPACE_PATH).toBeUndefined();
+  });
+});

packages/workspace-server/src/services/shell/shell.ts

@@ -424,16 +424,29 @@ export class ShellService extends TypedEventEmitter<ShellEvents> {
       const worktree = this.worktreeRepo.findByWorkspaceId(workspace.id);
       if (worktree) {
         worktreeName = worktree.name;
-        worktreePath = await this.deriveWorktreePath(repo.path, worktreeName);
+        // The stored path is authoritative — reused worktrees can live
+        // outside the managed worktree directory. Only derive when the
+        // stored path is gone (e.g. stale row after a location move).
+        worktreePath = existsSync(worktree.path)
+          ? worktree.path
+          : this.deriveWorktreePath(repo.path, worktreeName);
       }
     }
 
-    return buildWorkspaceEnv({
-      taskId,
-      folderPath: repo.path,
-      worktreePath,
-      worktreeName,
-      mode: workspace.mode,
-    });
+    try {
+      return await buildWorkspaceEnv({
+        taskId,
+        folderPath: repo.path,
+        worktreePath,
+        worktreeName,
+        mode: workspace.mode,
+      });
+    } catch (error) {
+      this.log.warn(
+        `Failed to build workspace env for task ${taskId}, starting shell without it`,
+        error,
+      );
+      return undefined;
+    }
   }
 }