diff --git a/packages/git/src/queries.test.ts b/packages/git/src/queries.test.ts index e5ed66a661..b1e5055996 100644 --- a/packages/git/src/queries.test.ts +++ b/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); + }); +}); diff --git a/packages/git/src/queries.ts b/packages/git/src/queries.ts index 467092bc53..bdb4e3dfec 100644 --- a/packages/git/src/queries.ts +++ b/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 { + 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. diff --git a/packages/shared/src/analytics-events.ts b/packages/shared/src/analytics-events.ts index 62850de083..00c159689f 100644 --- a/packages/shared/src/analytics-events.ts +++ b/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; diff --git a/packages/workspace-server/src/db/repositories/workspace-repository.mock.ts b/packages/workspace-server/src/db/repositories/workspace-repository.mock.ts index 78b112ed33..25493aacde 100644 --- a/packages/workspace-server/src/db/repositories/workspace-repository.mock.ts +++ b/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: () => {}, diff --git a/packages/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index 0513282303..f7bc19f654 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/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({ diff --git a/packages/workspace-server/src/services/workspace/workspace.ts b/packages/workspace-server/src/services/workspace/workspace.ts index 2d9a100777..1661ec71d9 100644 --- a/packages/workspace-server/src/services/workspace/workspace.ts +++ b/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 private creatingWorkspaces = new Map>(); private branchWatcherInitialized = false; + /** Tasks with an in-flight staleness check, so reads don't stack git calls. */ + private staleLinkChecks = new Set(); private findTaskAssociation(taskId: string): TaskAssociation | null { const workspace = this.workspaceRepo.findByTaskId(taskId); @@ -384,7 +387,10 @@ export class WorkspaceService extends TypedEventEmitter 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 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 { + 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 { @@ -1080,7 +1123,7 @@ export class WorkspaceService extends TypedEventEmitter } 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 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,