feat(workspace): auto-unlink linked branches that no longer exist

packages/git/src/queries.test.ts

@@ -1,9 +1,10 @@
 import { mkdtemp, rm, unlink, writeFile } from "node:fs/promises";
-import { tmpdir } from "node:os";
+import { devNull, tmpdir } from "node:os";
 import path from "node:path";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import { createGitClient } from "./client";
 import {
+  anyBranchRefExists,
   type ChangedFileInfo,
   computeDiffStatsFromFiles,
   detectDefaultBranch,
@@ -505,3 +506,38 @@ describe("remoteBranchExists", () => {
     expect(await remoteBranchExists(repoDir, "main")).toBe(false);
   });
 });
+
+describe("anyBranchRefExists", () => {
+  let repoDir: string;
+
+  // Builds refs via plumbing (commit-tree + update-ref) so the fixture also
+  // works in sandboxes where `git commit` is unavailable.
+  beforeEach(async () => {
+    repoDir = await mkdtemp(path.join(tmpdir(), "posthog-code-refs-"));
+    const git = createGitClient(repoDir);
+    await git.init(["--initial-branch", "main"]);
+    await git.addConfig("user.name", "Test");
+    await git.addConfig("user.email", "test@example.com");
+    const tree = (
+      await git.raw(["hash-object", "-w", "-t", "tree", devNull])
+    ).trim();
+    const sha = (await git.raw(["commit-tree", tree, "-m", "seed"])).trim();
+    await git.raw(["update-ref", "refs/heads/feat/local", sha]);
+    await git.raw(["update-ref", "refs/remotes/upstream/feat/remote", sha]);
+    await git.raw(["update-ref", "refs/tags/feat/tag-only", sha]);
+  });
+
+  afterEach(async () => {
+    await rm(repoDir, { recursive: true, force: true });
+  });
+
+  it.each([
+    { branch: "feat/local", expected: true },
+    { branch: "feat/remote", expected: true },
+    { branch: "feat/gone", expected: false },
+    // A tag with the name does not resurrect a deleted branch.
+    { branch: "feat/tag-only", expected: false },
+  ])("returns $expected for '$branch'", async ({ branch, expected }) => {
+    expect(await anyBranchRefExists(repoDir, branch)).toBe(expected);
+  });
+});

packages/git/src/queries.ts

@@ -224,6 +224,34 @@ export async function branchExists(
   );
 }
 
+/**
+ * True when the branch exists as a local branch or as a remote-tracking
+ * ref on any remote. Unlike `branchExists`, a tag or raw commit-ish with
+ * the same name does not count, and nothing reaches the network — a
+ * remote branch only counts once it has been fetched.
+ */
+export async function anyBranchRefExists(
+  baseDir: string,
+  branchName: string,
+  options?: CreateGitClientOptions,
+): Promise<boolean> {
+  const manager = getGitOperationManager();
+  return manager.executeRead(
+    baseDir,
+    async (git) => {
+      const refs = await git.raw([
+        "for-each-ref",
+        "--count=1",
+        "--format=%(refname)",
+        `refs/heads/${branchName}`,
+        `refs/remotes/*/${branchName}`,
+      ]);
+      return refs.trim().length > 0;
+    },
+    { signal: options?.abortSignal },
+  );
+}
+
 /**
  * Checks whether a branch exists on the remote without fetching it.
  * Uses `git ls-remote --heads`, which is read-only and reaches the remote.

packages/shared/src/analytics-events.ts

@@ -166,8 +166,9 @@ export interface AgentFileActivityProperties {
   branch_name: string | null;
 }
 
-// Branch link events
-type BranchLinkSource = "agent" | "user" | "unknown";
+// Branch link events. "auto" marks self-healing unlinks of branches that no
+// longer exist anywhere (e.g. deleted after a PR merge).
+type BranchLinkSource = "agent" | "user" | "auto" | "unknown";
 
 export interface BranchLinkedProperties {
   task_id: string;

packages/workspace-server/src/db/repositories/workspace-repository.mock.ts

@@ -109,7 +109,15 @@ export function createMockWorkspaceRepository(): MockWorkspaceRepository {
         workspaces.delete(id);
       }
     },
-    updateLinkedBranch: () => {},
+    updateLinkedBranch: (taskId, linkedBranch) => {
+      const w = findLiveByTaskId(taskId);
+      if (!w) return;
+      workspaces.set(w.id, {
+        ...w,
+        linkedBranch,
+        updatedAt: new Date().toISOString(),
+      });
+    },
     updatePinnedAt: () => {},
     updateLastViewedAt: () => {},
     updateLastActivityAt: () => {},

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

@@ -3,6 +3,7 @@ import * as os from "node:os";
 import * as path from "node:path";
 import type { RootLogger } from "@posthog/di/logger";
 import {
+  anyBranchRefExists,
   branchExists,
   getCurrentBranch,
   getDefaultBranch,
@@ -35,6 +36,7 @@ vi.mock("@posthog/git/queries", async (importOriginal) => {
     getDefaultBranch: vi.fn(),
     getCurrentBranch: vi.fn(),
     branchExists: vi.fn(),
+    anyBranchRefExists: vi.fn(),
     remoteBranchExists: vi.fn(),
     hasTrackedFiles: vi.fn(),
   };
@@ -258,6 +260,93 @@ describe("WorkspaceService", () => {
     });
   });
 
+  describe("getWorkspace (stale linked branch healing)", () => {
+    const tempDirs: string[] = [];
+
+    beforeEach(() => {
+      vi.mocked(anyBranchRefExists).mockReset();
+    });
+
+    afterEach(() => {
+      for (const dir of tempDirs.splice(0)) {
+        fs.rmSync(dir, { recursive: true, force: true });
+      }
+    });
+
+    /** A fake worktree checkout whose HEAD points at `branch`. */
+    function mkWorktreeOn(branch: string): string {
+      const dir = fs.mkdtempSync(path.join(os.tmpdir(), "stale-link-"));
+      tempDirs.push(dir);
+      fs.mkdirSync(path.join(dir, ".git"));
+      fs.writeFileSync(
+        path.join(dir, ".git", "HEAD"),
+        `ref: refs/heads/${branch}\n`,
+      );
+      return dir;
+    }
+
+    function seedLinkedTask(linkedBranch: string, currentBranch = "main") {
+      seedWorktreeTask(mocks, {
+        taskId: "t1",
+        repoPath: "/code/myrepo",
+        name: "wt",
+        worktreePath: mkWorktreeOn(currentBranch),
+      });
+      service.linkBranch("t1", linkedBranch, "user");
+      vi.mocked(mocks.analytics.track).mockClear();
+    }
+
+    it.each([
+      { branch: "feat/gone", refExists: false, expected: null },
+      { branch: "feat/alive", refExists: true, expected: "feat/alive" },
+    ])(
+      "linkedBranch is $expected when refExists=$refExists",
+      async ({ branch, refExists, expected }) => {
+        seedLinkedTask(branch);
+        vi.mocked(anyBranchRefExists).mockResolvedValue(refExists);
+
+        const workspace = await service.getWorkspace("t1");
+
+        expect(workspace?.linkedBranch).toBe(expected);
+      },
+    );
+
+    it("emits, tracks, and persists the unlink when refs are gone", async () => {
+      seedLinkedTask("feat/gone");
+      vi.mocked(anyBranchRefExists).mockResolvedValue(false);
+      const emitted = vi.fn();
+      service.on(WorkspaceServiceEvent.LinkedBranchChanged, emitted);
+
+      await service.getWorkspace("t1");
+
+      expect(emitted).toHaveBeenCalledWith({ taskId: "t1", branchName: null });
+      expect(mocks.analytics.track).toHaveBeenCalledWith(
+        ANALYTICS_EVENTS.BRANCH_UNLINKED,
+        expect.objectContaining({ task_id: "t1", source: "auto" }),
+      );
+      expect(mocks.workspaceRepo.findByTaskId("t1")?.linkedBranch).toBeNull();
+    });
+
+    it("skips the check when on the linked branch", async () => {
+      seedLinkedTask("main", "main");
+
+      const workspace = await service.getWorkspace("t1");
+
+      expect(workspace?.linkedBranch).toBe("main");
+      expect(anyBranchRefExists).not.toHaveBeenCalled();
+    });
+
+    it("keeps the link when the staleness check fails", async () => {
+      seedLinkedTask("feat/unknown");
+      vi.mocked(anyBranchRefExists).mockRejectedValue(new Error("git broke"));
+
+      const workspace = await service.getWorkspace("t1");
+
+      expect(workspace?.linkedBranch).toBe("feat/unknown");
+      expect(mocks.analytics.track).not.toHaveBeenCalled();
+    });
+  });
+
   describe("getWorkspace (cloud mode)", () => {
     it("projects a cloud workspace without touching git or fs", async () => {
       mocks.workspaceRepo.create({

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

@@ -7,6 +7,7 @@ import {
 } from "@posthog/di/logger";
 import { createGitClient } from "@posthog/git/client";
 import {
+  anyBranchRefExists,
   branchExists,
   getCurrentBranch,
   getDefaultBranch,
@@ -157,6 +158,8 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
 
   private creatingWorkspaces = new Map<string, Promise<WorkspaceInfo>>();
   private branchWatcherInitialized = false;
+  /** Tasks with an in-flight staleness check, so reads don't stack git calls. */
+  private staleLinkChecks = new Set<string>();
 
   private findTaskAssociation(taskId: string): TaskAssociation | null {
     const workspace = this.workspaceRepo.findByTaskId(taskId);
@@ -384,7 +387,10 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
     this.log.info("Linked branch to task", { taskId, branchName, source });
   }
 
-  public unlinkBranch(taskId: string, source?: "agent" | "user"): void {
+  public unlinkBranch(
+    taskId: string,
+    source?: "agent" | "user" | "auto",
+  ): void {
     this.workspaceRepo.updateLinkedBranch(taskId, null);
     this.emit(WorkspaceServiceEvent.LinkedBranchChanged, {
       taskId,
@@ -397,6 +403,43 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
     this.log.info("Unlinked branch from task", { taskId, source });
   }
 
+  /**
+   * Drops a linked branch whose branch no longer exists anywhere — no local
+   * branch and no remote-tracking ref (the usual end state after a PR merge
+   * plus branch deletion). Such a link can never match the working tree
+   * again, so keeping it only produces wrong-branch warnings on every send.
+   * Returns true when the link was removed.
+   */
+  private async unlinkStaleLinkedBranch(
+    taskId: string,
+    linkedBranch: string,
+    repoPath: string,
+  ): Promise<boolean> {
+    if (this.staleLinkChecks.has(taskId)) return false;
+    this.staleLinkChecks.add(taskId);
+    try {
+      if (await anyBranchRefExists(repoPath, linkedBranch)) return false;
+      // Re-read: the link may have changed while the git check ran.
+      const row = this.workspaceRepo.findByTaskId(taskId);
+      if ((row?.linkedBranch ?? null) !== linkedBranch) return false;
+      this.log.info("Linked branch has no remaining refs, unlinking", {
+        taskId,
+        linkedBranch,
+      });
+      this.unlinkBranch(taskId, "auto");
+      return true;
+    } catch (error) {
+      this.log.warn("Failed to check linked branch staleness", {
+        taskId,
+        linkedBranch,
+        error,
+      });
+      return false;
+    } finally {
+      this.staleLinkChecks.delete(taskId);
+    }
+  }
+
   private getLocalWorktreePathIfExists(
     mainRepoPath: string,
   ): Promise<string | null> {
@@ -1080,7 +1123,7 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
     }
 
     const dbRow = this.workspaceRepo.findByTaskId(taskId);
-    const linkedBranch = dbRow?.linkedBranch ?? null;
+    let linkedBranch = dbRow?.linkedBranch ?? null;
 
     if (assoc.mode === "cloud") {
       return {
@@ -1116,6 +1159,21 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
       branchName = await getBranchFromPath(branchPath);
     }
 
+    // Only a mismatched link can be stale: being on the linked branch proves
+    // it exists, and without a current branch nothing warns anyway.
+    if (
+      linkedBranch &&
+      branchName &&
+      linkedBranch !== branchName &&
+      (await this.unlinkStaleLinkedBranch(
+        taskId,
+        linkedBranch,
+        worktreePath ?? folderPath,
+      ))
+    ) {
+      linkedBranch = null;
+    }
+
     return {
       taskId,
       folderId: assoc.folderId,