From e4815a91c1aaeeb0941c0156fbe03dda0aa3b291 Mon Sep 17 00:00:00 2001 From: Adam Bowker Date: Mon, 6 Jul 2026 11:13:05 -0400 Subject: [PATCH 1/4] feat(workspace): auto-unlink linked branches that no longer exist MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A task's linkedBranch is never cleaned up when its branch is deleted (the usual end state after a PR merge), so returning to a finished task and asking a follow-up from main nags with a wrong-branch warning on every send — a link like that can never match the working tree again. Production data suggests this drives the ~59% "continue anyway" rate on the mismatch dialog. getWorkspace now self-heals: when the linked branch mismatches the current branch, a new anyBranchRefExists git query checks for the branch as a local head or remote-tracking ref on any remote (no network; tags don't count). If nothing remains, the link is dropped with source "auto" on the existing Branch unlinked analytics event. Being on the linked branch proves it exists, so only mismatched reads pay the extra ref lookup, and an in-flight guard keeps concurrent reads from stacking git calls. Also makes the mock workspace repository's updateLinkedBranch persist instead of being a no-op, so service tests can exercise link state. Generated-By: PostHog Code Task-Id: adec8438-4ca0-42f0-9d9d-948de60c3b23 --- packages/git/src/queries.test.ts | 38 ++++++++- packages/git/src/queries.ts | 28 ++++++ packages/shared/src/analytics-events.ts | 5 +- .../repositories/workspace-repository.mock.ts | 10 ++- .../src/services/workspace/workspace.test.ts | 85 +++++++++++++++++++ .../src/services/workspace/workspace.ts | 62 +++++++++++++- 6 files changed, 222 insertions(+), 6 deletions(-) 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..a5aac4ff4a 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,89 @@ 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("unlinks a mismatched branch that has no refs left", async () => { + seedLinkedTask("feat/gone"); + vi.mocked(anyBranchRefExists).mockResolvedValue(false); + const emitted = vi.fn(); + service.on(WorkspaceServiceEvent.LinkedBranchChanged, emitted); + + const workspace = await service.getWorkspace("t1"); + + expect(workspace?.linkedBranch).toBeNull(); + 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("keeps a mismatched branch that still has a ref", async () => { + seedLinkedTask("feat/alive"); + vi.mocked(anyBranchRefExists).mockResolvedValue(true); + + const workspace = await service.getWorkspace("t1"); + + expect(workspace?.linkedBranch).toBe("feat/alive"); + expect(mocks.analytics.track).not.toHaveBeenCalled(); + }); + + 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, From ed27d0a9d518d2e38fc542e402d00df0e2470b9e Mon Sep 17 00:00:00 2001 From: Adam Bowker Date: Mon, 6 Jul 2026 11:13:07 -0400 Subject: [PATCH 2/4] review: parameterise stale-link healing outcome tests Greptile review feedback: the unlink/keep cases share one shape (seed, mock anyBranchRefExists, read, assert linkedBranch), so they collapse into an it.each on refExists -> outcome; the unlink side effects (event, analytics, persisted null) move to their own focused test. Generated-By: PostHog Code Task-Id: adec8438-4ca0-42f0-9d9d-948de60c3b23 --- .../src/services/workspace/workspace.test.ts | 30 +++++++++++-------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/packages/workspace-server/src/services/workspace/workspace.test.ts b/packages/workspace-server/src/services/workspace/workspace.test.ts index a5aac4ff4a..f7bc19f654 100644 --- a/packages/workspace-server/src/services/workspace/workspace.test.ts +++ b/packages/workspace-server/src/services/workspace/workspace.test.ts @@ -296,15 +296,29 @@ describe("WorkspaceService", () => { vi.mocked(mocks.analytics.track).mockClear(); } - it("unlinks a mismatched branch that has no refs left", async () => { + 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); - const workspace = await service.getWorkspace("t1"); + await service.getWorkspace("t1"); - expect(workspace?.linkedBranch).toBeNull(); expect(emitted).toHaveBeenCalledWith({ taskId: "t1", branchName: null }); expect(mocks.analytics.track).toHaveBeenCalledWith( ANALYTICS_EVENTS.BRANCH_UNLINKED, @@ -313,16 +327,6 @@ describe("WorkspaceService", () => { expect(mocks.workspaceRepo.findByTaskId("t1")?.linkedBranch).toBeNull(); }); - it("keeps a mismatched branch that still has a ref", async () => { - seedLinkedTask("feat/alive"); - vi.mocked(anyBranchRefExists).mockResolvedValue(true); - - const workspace = await service.getWorkspace("t1"); - - expect(workspace?.linkedBranch).toBe("feat/alive"); - expect(mocks.analytics.track).not.toHaveBeenCalled(); - }); - it("skips the check when on the linked branch", async () => { seedLinkedTask("main", "main"); From b0577f3a89ebccca84a4b83af53644d4722ee593 Mon Sep 17 00:00:00 2001 From: Adam Bowker Date: Mon, 6 Jul 2026 11:13:47 -0400 Subject: [PATCH 3/4] feat(workspace): replace branch mismatch modal with non-blocking banner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wrong-branch AlertDialog intercepted the user's send at the worst moment (the Enter keypress) and blocked it behind a modal. Production data says its premise is wrong most of the time: of ~1,250 dialogs in 60 days, only ~22% chose "Switch branch" while ~59% chose "Continue anyway" — and dismissal was session-scoped, so the same users re-hit the dialog ~4x per day. Replace it with a slim banner above the composer, shown as soon as the mismatch exists (before typing) via a new composerNotice slot on SessionView. Sending is never intercepted; the agent works on the current branch either way, which is what the majority already chose. Actions: - Switch branch: checkout the linked branch, loading and errors inline. - Use current branch: re-link the task to the current branch via workspace.linkBranch — a durable resolution that replaces the old per-session "Continue anyway" nag. - Dismiss: hides for the session (re-arms when the branch changes), as before. Analytics keep the same event names; the action union becomes switch | relink | dismiss, and "shown" now fires once per banner appearance instead of once per intercepted send. The old dialog, its hook, and the beforeSubmit interception are deleted; the banner-ineligible cloud case is unchanged (branchName is null there). Generated-By: PostHog Code Task-Id: adec8438-4ca0-42f0-9d9d-948de60c3b23 --- ...g.test.ts => branchMismatchBanner.test.ts} | 72 +++-- ...matchDialog.ts => branchMismatchBanner.ts} | 34 ++- packages/shared/src/analytics-events.ts | 6 +- .../sessions/components/SessionView.tsx | 18 ++ .../task-detail/BranchMismatchBanner.test.tsx | 67 +++++ .../task-detail/BranchMismatchBanner.tsx | 75 +++++ .../task-detail/BranchMismatchDialog.tsx | 133 --------- .../task-detail/components/TaskLogsPanel.tsx | 18 +- .../workspace/useBranchMismatchBanner.test.ts | 217 ++++++++++++++ .../workspace/useBranchMismatchBanner.ts | 145 +++++++++ .../workspace/useBranchMismatchDialog.test.ts | 282 ------------------ .../workspace/useBranchMismatchDialog.ts | 143 --------- 12 files changed, 598 insertions(+), 612 deletions(-) rename packages/core/src/workspace/{branchMismatchDialog.test.ts => branchMismatchBanner.test.ts} (53%) rename packages/core/src/workspace/{branchMismatchDialog.ts => branchMismatchBanner.ts} (75%) create mode 100644 packages/ui/src/features/task-detail/BranchMismatchBanner.test.tsx create mode 100644 packages/ui/src/features/task-detail/BranchMismatchBanner.tsx delete mode 100644 packages/ui/src/features/task-detail/BranchMismatchDialog.tsx create mode 100644 packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts create mode 100644 packages/ui/src/features/workspace/useBranchMismatchBanner.ts delete mode 100644 packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts delete mode 100644 packages/ui/src/features/workspace/useBranchMismatchDialog.ts diff --git a/packages/core/src/workspace/branchMismatchDialog.test.ts b/packages/core/src/workspace/branchMismatchBanner.test.ts similarity index 53% rename from packages/core/src/workspace/branchMismatchDialog.test.ts rename to packages/core/src/workspace/branchMismatchBanner.test.ts index 632fc5d475..94f4137cff 100644 --- a/packages/core/src/workspace/branchMismatchDialog.test.ts +++ b/packages/core/src/workspace/branchMismatchBanner.test.ts @@ -3,9 +3,9 @@ import { describe, expect, it } from "vitest"; import { buildBranchMismatchAnalyticsEvent, buildCheckoutBranchRequest, - decideBeforeSubmit, - resolveSwitchErrorMessage, -} from "./branchMismatchDialog"; + buildLinkBranchRequest, + resolveBranchMismatchError, +} from "./branchMismatchBanner"; const context = { taskId: "task-1", @@ -14,16 +14,6 @@ const context = { hasUncommittedChanges: true, }; -describe("decideBeforeSubmit", () => { - it("allows submit when not warning", () => { - expect(decideBeforeSubmit(false)).toBe(true); - }); - - it("blocks submit when warning", () => { - expect(decideBeforeSubmit(true)).toBe(false); - }); -}); - describe("buildBranchMismatchAnalyticsEvent", () => { it("builds the warning-shown event", () => { expect(buildBranchMismatchAnalyticsEvent("shown", context)).toEqual({ @@ -37,21 +27,24 @@ describe("buildBranchMismatchAnalyticsEvent", () => { }); }); - it("builds the action event", () => { - expect(buildBranchMismatchAnalyticsEvent("switch", context)).toEqual({ - event: ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, - properties: { - task_id: "task-1", - action: "switch", - linked_branch: "feat/foo", - current_branch: "main", - }, - }); - }); + it.each(["switch", "relink", "dismiss"] as const)( + "builds the %s action event", + (action) => { + expect(buildBranchMismatchAnalyticsEvent(action, context)).toEqual({ + event: ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, + properties: { + task_id: "task-1", + action, + linked_branch: "feat/foo", + current_branch: "main", + }, + }); + }, + ); it("returns null without both branches", () => { expect( - buildBranchMismatchAnalyticsEvent("cancel", { + buildBranchMismatchAnalyticsEvent("dismiss", { ...context, linkedBranch: null, }), @@ -76,14 +69,29 @@ describe("buildCheckoutBranchRequest", () => { }); }); -describe("resolveSwitchErrorMessage", () => { - it("uses error message", () => { - expect(resolveSwitchErrorMessage(new Error("dirty worktree"))).toBe( - "dirty worktree", - ); +describe("buildLinkBranchRequest", () => { + it("builds the request", () => { + expect(buildLinkBranchRequest("task-1", "main")).toEqual({ + taskId: "task-1", + branchName: "main", + }); }); - it("falls back for non-errors", () => { - expect(resolveSwitchErrorMessage("oops")).toBe("Failed to switch branch"); + it("returns null without a current branch", () => { + expect(buildLinkBranchRequest("task-1", null)).toBeNull(); + }); +}); + +describe("resolveBranchMismatchError", () => { + it("uses the error message", () => { + expect( + resolveBranchMismatchError(new Error("dirty worktree"), "nope"), + ).toBe("dirty worktree"); + }); + + it.each([["oops"], [new Error("")]])("falls back for %s", (error) => { + expect(resolveBranchMismatchError(error, "Failed to switch branch")).toBe( + "Failed to switch branch", + ); }); }); diff --git a/packages/core/src/workspace/branchMismatchDialog.ts b/packages/core/src/workspace/branchMismatchBanner.ts similarity index 75% rename from packages/core/src/workspace/branchMismatchDialog.ts rename to packages/core/src/workspace/branchMismatchBanner.ts index d92d2388b3..9da8791f6a 100644 --- a/packages/core/src/workspace/branchMismatchDialog.ts +++ b/packages/core/src/workspace/branchMismatchBanner.ts @@ -4,10 +4,10 @@ import { type BranchMismatchWarningShownProperties, } from "@posthog/shared"; -export type BranchMismatchDialogAction = +export type BranchMismatchBannerAction = | "switch" - | "continue" - | "cancel" + | "relink" + | "dismiss" | "shown"; export interface BranchMismatchContext { @@ -22,6 +22,11 @@ export interface CheckoutBranchRequest { branchName: string; } +export interface LinkBranchRequest { + taskId: string; + branchName: string; +} + export type BranchMismatchAnalyticsEvent = | { event: typeof ANALYTICS_EVENTS.BRANCH_MISMATCH_WARNING_SHOWN; @@ -32,12 +37,8 @@ export type BranchMismatchAnalyticsEvent = properties: BranchMismatchActionProperties; }; -export function decideBeforeSubmit(shouldWarn: boolean): boolean { - return !shouldWarn; -} - export function buildBranchMismatchAnalyticsEvent( - action: BranchMismatchDialogAction, + action: BranchMismatchBannerAction, context: BranchMismatchContext, ): BranchMismatchAnalyticsEvent | null { const { taskId, linkedBranch, currentBranch, hasUncommittedChanges } = @@ -79,6 +80,19 @@ export function buildCheckoutBranchRequest( return { directoryPath: repoPath, branchName: linkedBranch }; } -export function resolveSwitchErrorMessage(error: unknown): string { - return error instanceof Error ? error.message : "Failed to switch branch"; +export function buildLinkBranchRequest( + taskId: string, + currentBranch: string | null, +): LinkBranchRequest | null { + if (!currentBranch) { + return null; + } + return { taskId, branchName: currentBranch }; +} + +export function resolveBranchMismatchError( + error: unknown, + fallback: string, +): string { + return error instanceof Error && error.message ? error.message : fallback; } diff --git a/packages/shared/src/analytics-events.ts b/packages/shared/src/analytics-events.ts index 00c159689f..869be6a153 100644 --- a/packages/shared/src/analytics-events.ts +++ b/packages/shared/src/analytics-events.ts @@ -315,8 +315,10 @@ export interface TourEventProperties { total_steps?: number; } -// Branch mismatch events -type BranchMismatchAction = "switch" | "continue" | "cancel"; +// Branch mismatch events. The banner replaced the blocking dialog: "relink" +// re-links the task to the current branch (was "continue"), "dismiss" hides +// the banner for the session (was "cancel"). +type BranchMismatchAction = "switch" | "relink" | "dismiss"; export interface BranchMismatchWarningShownProperties { task_id: string; diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx index 387b6e5804..5f00a4c56e 100644 --- a/packages/ui/src/features/sessions/components/SessionView.tsx +++ b/packages/ui/src/features/sessions/components/SessionView.tsx @@ -94,6 +94,8 @@ interface SessionViewProps { isActiveSession?: boolean; /** Hide the message input and permission UI — log-only view. */ hideInput?: boolean; + /** Ambient notice rendered above the composer (e.g. branch mismatch). */ + composerNotice?: React.ReactNode; } const DEFAULT_ERROR_MESSAGE = @@ -218,6 +220,7 @@ export function SessionView({ compact = false, isActiveSession = true, hideInput = false, + composerNotice, }: SessionViewProps) { const sessionService = useService(SESSION_SERVICE); useSessionEventsResidency(taskId); @@ -623,6 +626,21 @@ export function SessionView({ + {!hideInput && composerNotice && ( + + {composerNotice} + + )} + {hasError && !showInlineBanner ? ( [0]>, +) { + const handlers = { + onSwitch: vi.fn(), + onUseCurrentBranch: vi.fn(), + onDismiss: vi.fn(), + }; + render( + + + , + ); + return handlers; +} + +describe("BranchMismatchBanner", () => { + it("shows both branches", () => { + renderBanner(); + + expect(screen.getByText("feat/foo")).toBeInTheDocument(); + expect(screen.getByText("main")).toBeInTheDocument(); + }); + + it("wires each action to its handler", async () => { + const user = userEvent.setup(); + const { onSwitch, onUseCurrentBranch, onDismiss } = renderBanner(); + + await user.click(screen.getByRole("button", { name: "Switch branch" })); + await user.click( + screen.getByRole("button", { name: "Use current branch" }), + ); + await user.click(screen.getByRole("button", { name: "Dismiss" })); + + expect(onSwitch).toHaveBeenCalledTimes(1); + expect(onUseCurrentBranch).toHaveBeenCalledTimes(1); + expect(onDismiss).toHaveBeenCalledTimes(1); + }); + + it("shows the action error", () => { + renderBanner({ actionError: "dirty worktree" }); + + expect(screen.getByText("dirty worktree")).toBeInTheDocument(); + }); + + it("disables actions while a switch is in flight", () => { + renderBanner({ isSwitching: true }); + + expect( + screen.getByRole("button", { name: "Use current branch" }), + ).toBeDisabled(); + }); +}); diff --git a/packages/ui/src/features/task-detail/BranchMismatchBanner.tsx b/packages/ui/src/features/task-detail/BranchMismatchBanner.tsx new file mode 100644 index 0000000000..89ce2c6245 --- /dev/null +++ b/packages/ui/src/features/task-detail/BranchMismatchBanner.tsx @@ -0,0 +1,75 @@ +import { GitBranch, X } from "@phosphor-icons/react"; +import { Button, Code, Flex, IconButton, Text } from "@radix-ui/themes"; +import type { BranchMismatchBannerState } from "../workspace/useBranchMismatchBanner"; + +/** + * Slim, non-blocking strip above the composer for a task whose working tree + * is on a different branch than the task's linked branch. Sending is never + * intercepted — the agent works on the current branch either way. + */ +export function BranchMismatchBanner({ + linkedBranch, + currentBranch, + actionError, + isSwitching, + isRelinking, + onSwitch, + onUseCurrentBranch, + onDismiss, +}: BranchMismatchBannerState) { + const busy = isSwitching || isRelinking; + return ( + + + + + + Task is linked to {linkedBranch} but + you're on {currentBranch} — the agent + works on the current branch. + + + + + + + + + + + {actionError && ( + + {actionError} + + )} + + ); +} diff --git a/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx b/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx deleted file mode 100644 index 23c5cb5b78..0000000000 --- a/packages/ui/src/features/task-detail/BranchMismatchDialog.tsx +++ /dev/null @@ -1,133 +0,0 @@ -import { GitBranch, Warning } from "@phosphor-icons/react"; -import { - AlertDialog, - Button, - Callout, - Code, - Flex, - Text, -} from "@radix-ui/themes"; - -interface BranchMismatchDialogProps { - open: boolean; - linkedBranch: string; - currentBranch: string; - hasUncommittedChanges: boolean; - switchError: string | null; - onSwitch: () => void; - onContinue: () => void; - onCancel: () => void; - isSwitching?: boolean; -} - -function BranchLabel({ name }: { name: string }) { - return ( - - - - {name} - - - ); -} - -export function BranchMismatchDialog({ - open, - linkedBranch, - currentBranch, - hasUncommittedChanges, - switchError, - onSwitch, - onContinue, - onCancel, - isSwitching, -}: BranchMismatchDialogProps) { - return ( - { - if (!isOpen) onCancel(); - }} - > - - - - - Wrong branch - - - - This task is linked to a different branch than the one you're - currently on. The agent will make changes on the current branch. - - - - - Linked - - - - - - Current - - - - - - {hasUncommittedChanges && !switchError && ( - - - You have uncommitted changes on your current branch. If needed, - commit or stash them first. - - - )} - - {switchError && ( - - {switchError} - - )} - - - - - - - - - - - - - - - ); -} diff --git a/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx b/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx index b30345ce69..d713fdec29 100644 --- a/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx +++ b/packages/ui/src/features/task-detail/components/TaskLogsPanel.tsx @@ -14,10 +14,10 @@ import { useSessionConnection } from "../../sessions/hooks/useSessionConnection" import { useSessionViewState } from "../../sessions/hooks/useSessionViewState"; import { useRestoreTask } from "../../suspension/useRestoreTask"; import { useSuspendedTaskIds } from "../../suspension/useSuspendedTaskIds"; -import { useBranchMismatchDialog } from "../../workspace/useBranchMismatchDialog"; +import { useBranchMismatchBanner } from "../../workspace/useBranchMismatchBanner"; import { useWorkspaceLoaded } from "../../workspace/useWorkspace"; import { useCreateWorkspace } from "../../workspace/useWorkspaceMutations"; -import { BranchMismatchDialog } from "../BranchMismatchDialog"; +import { BranchMismatchBanner } from "../BranchMismatchBanner"; import { WorkspaceSetupPrompt } from "./WorkspaceSetupPrompt"; interface TaskLogsPanelProps { @@ -80,11 +80,7 @@ export function TaskLogsPanel({ taskId, task, hideInput }: TaskLogsPanelProps) { handleBashCommand, } = useSessionCallbacks({ taskId, task, session, repoPath }); - const { handleBeforeSubmit, dialogProps } = useBranchMismatchDialog({ - taskId, - repoPath, - onSendPrompt: handleSendPrompt, - }); + const branchMismatchBanner = useBranchMismatchBanner({ taskId, repoPath }); const slackThreadUrl = typeof task.latest_run?.state?.slack_thread_url === "string" @@ -166,7 +162,6 @@ export function TaskLogsPanel({ taskId, task, hideInput }: TaskLogsPanelProps) { isRestoring={isRestoring} isPromptPending={isPromptPending} promptStartedAt={promptStartedAt} - onBeforeSubmit={handleBeforeSubmit} onSendPrompt={handleSendPrompt} onBashCommand={isCloud ? undefined : handleBashCommand} onCancelPrompt={handleCancelPrompt} @@ -183,12 +178,15 @@ export function TaskLogsPanel({ taskId, task, hideInput }: TaskLogsPanelProps) { isCloud={isCloud} cloudStatus={cloudStatus} slackThreadUrl={slackThreadUrl} + composerNotice={ + branchMismatchBanner ? ( + + ) : undefined + } /> - - {dialogProps && } ); } diff --git a/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts b/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts new file mode 100644 index 0000000000..2a7d4dab7b --- /dev/null +++ b/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts @@ -0,0 +1,217 @@ +import { act, renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +let mockShouldWarn = false; +const mockDismissWarning = vi.fn(); + +const mockGuard = vi.hoisted(() => ({ + useBranchMismatchGuard: vi.fn( + (): { + shouldWarn: boolean; + linkedBranch: string | null; + currentBranch: string | null; + dismissWarning: () => void; + } => ({ + shouldWarn: mockShouldWarn, + linkedBranch: "feat/foo", + currentBranch: "main", + dismissWarning: mockDismissWarning, + }), + ), +})); +vi.mock("./useBranchMismatch", () => mockGuard); + +vi.mock("../git-interaction/useGitQueries", () => ({ + useGitQueries: () => ({ hasChanges: false }), +})); + +vi.mock("../git-interaction/gitCacheKeys", () => ({ + invalidateGitBranchQueries: vi.fn(), +})); + +// Tag each trpc mutation with a key so the useMutation mock can hand back a +// distinct mutate spy per mutation. +vi.mock("@posthog/host-router/react", () => ({ + useHostTRPC: () => ({ + git: { + checkoutBranch: { + mutationOptions: (opts: Record) => ({ + ...opts, + mutationKey: ["checkout"], + }), + }, + }, + workspace: { + linkBranch: { + mutationOptions: (opts: Record) => ({ + ...opts, + mutationKey: ["link"], + }), + }, + }, + }), +})); + +type CapturedMutation = { + onSuccess?: () => void; + onError?: (e: Error) => void; +}; +const captured: Record = {}; +const mutateSpies: Record> = {}; + +vi.mock("@tanstack/react-query", () => ({ + useMutation: (opts: { mutationKey?: string[] } & CapturedMutation) => { + const key = opts.mutationKey?.[0] ?? "unknown"; + captured[key] = opts; + mutateSpies[key] ??= vi.fn(); + return { mutate: mutateSpies[key], isPending: false }; + }, +})); + +vi.mock("../../shell/logger", () => ({ + logger: { scope: () => ({ error: vi.fn() }) }, +})); + +const mockTrack = vi.fn(); +vi.mock("../../shell/analytics", () => ({ + track: (...args: unknown[]) => mockTrack(...args), +})); + +import { ANALYTICS_EVENTS } from "@posthog/shared"; +import { useBranchMismatchBanner } from "./useBranchMismatchBanner"; + +function renderBanner(overrides?: { shouldWarn?: boolean }) { + mockShouldWarn = overrides?.shouldWarn ?? true; + mockGuard.useBranchMismatchGuard.mockReturnValue({ + shouldWarn: mockShouldWarn, + linkedBranch: "feat/foo", + currentBranch: "main", + dismissWarning: mockDismissWarning, + }); + + return renderHook(() => + useBranchMismatchBanner({ taskId: "task-1", repoPath: "/repo" }), + ); +} + +describe("useBranchMismatchBanner", () => { + beforeEach(() => { + vi.clearAllMocks(); + for (const key of Object.keys(captured)) delete captured[key]; + mockShouldWarn = false; + }); + + it("returns null and tracks nothing when there is no mismatch", () => { + const { result } = renderBanner({ shouldWarn: false }); + + expect(result.current).toBeNull(); + expect(mockTrack).not.toHaveBeenCalled(); + }); + + it("returns banner state and tracks shown once per appearance", () => { + const { result, rerender } = renderBanner(); + + expect(result.current).toMatchObject({ + linkedBranch: "feat/foo", + currentBranch: "main", + actionError: null, + }); + rerender(); + + expect(mockTrack).toHaveBeenCalledTimes(1); + expect(mockTrack).toHaveBeenCalledWith( + ANALYTICS_EVENTS.BRANCH_MISMATCH_WARNING_SHOWN, + { + task_id: "task-1", + linked_branch: "feat/foo", + current_branch: "main", + has_uncommitted_changes: false, + }, + ); + }); + + it("onSwitch checks out the linked branch and tracks switch", () => { + const { result } = renderBanner(); + + mockTrack.mockClear(); + act(() => { + result.current?.onSwitch(); + }); + + expect(mutateSpies.checkout).toHaveBeenCalledWith({ + directoryPath: "/repo", + branchName: "feat/foo", + }); + expect(mockTrack).toHaveBeenCalledWith( + ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, + { + task_id: "task-1", + action: "switch", + linked_branch: "feat/foo", + current_branch: "main", + }, + ); + + act(() => { + captured.checkout.onSuccess?.(); + }); + expect(mockDismissWarning).toHaveBeenCalled(); + }); + + it("onSwitch failure surfaces the error without dismissing", () => { + const { result } = renderBanner(); + + act(() => { + result.current?.onSwitch(); + }); + act(() => { + captured.checkout.onError?.(new Error("dirty worktree")); + }); + + expect(result.current?.actionError).toBe("dirty worktree"); + expect(mockDismissWarning).not.toHaveBeenCalled(); + }); + + it("onUseCurrentBranch re-links the task and tracks relink", () => { + const { result } = renderBanner(); + + mockTrack.mockClear(); + act(() => { + result.current?.onUseCurrentBranch(); + }); + + expect(mutateSpies.link).toHaveBeenCalledWith({ + taskId: "task-1", + branchName: "main", + }); + expect(mockTrack).toHaveBeenCalledWith( + ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, + { + task_id: "task-1", + action: "relink", + linked_branch: "feat/foo", + current_branch: "main", + }, + ); + }); + + it("onDismiss dismisses for the session and tracks dismiss", () => { + const { result } = renderBanner(); + + mockTrack.mockClear(); + act(() => { + result.current?.onDismiss(); + }); + + expect(mockDismissWarning).toHaveBeenCalled(); + expect(mockTrack).toHaveBeenCalledWith( + ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, + { + task_id: "task-1", + action: "dismiss", + linked_branch: "feat/foo", + current_branch: "main", + }, + ); + }); +}); diff --git a/packages/ui/src/features/workspace/useBranchMismatchBanner.ts b/packages/ui/src/features/workspace/useBranchMismatchBanner.ts new file mode 100644 index 0000000000..ff17a36f84 --- /dev/null +++ b/packages/ui/src/features/workspace/useBranchMismatchBanner.ts @@ -0,0 +1,145 @@ +import { + type BranchMismatchBannerAction, + buildBranchMismatchAnalyticsEvent, + buildCheckoutBranchRequest, + buildLinkBranchRequest, + resolveBranchMismatchError, +} from "@posthog/core/workspace/branchMismatchBanner"; +import { useHostTRPC } from "@posthog/host-router/react"; +import { useMutation } from "@tanstack/react-query"; +import { useCallback, useEffect, useRef, useState } from "react"; +import { track } from "../../shell/analytics"; +import { logger } from "../../shell/logger"; +import { invalidateGitBranchQueries } from "../git-interaction/gitCacheKeys"; +import { useGitQueries } from "../git-interaction/useGitQueries"; +import { useBranchMismatchGuard } from "./useBranchMismatch"; + +const log = logger.scope("branch-mismatch"); + +export interface BranchMismatchBannerState { + linkedBranch: string; + currentBranch: string; + actionError: string | null; + isSwitching: boolean; + isRelinking: boolean; + onSwitch: () => void; + onUseCurrentBranch: () => void; + onDismiss: () => void; +} + +interface UseBranchMismatchBannerOptions { + taskId: string; + repoPath: string | null; +} + +/** + * Ambient (non-blocking) surface for a task whose working tree is on a + * different branch than the one linked to the task. Unlike the old dialog it + * never intercepts sending: the user can switch to the linked branch, re-link + * the task to the current branch (a durable choice), or dismiss for the + * session. + */ +export function useBranchMismatchBanner({ + taskId, + repoPath, +}: UseBranchMismatchBannerOptions): BranchMismatchBannerState | null { + const { shouldWarn, linkedBranch, currentBranch, dismissWarning } = + useBranchMismatchGuard(taskId); + const { hasChanges: hasUncommittedChanges } = useGitQueries( + repoPath ?? undefined, + ); + const [actionError, setActionError] = useState(null); + + const emitAction = useCallback( + (action: BranchMismatchBannerAction) => { + const analytics = buildBranchMismatchAnalyticsEvent(action, { + taskId, + linkedBranch, + currentBranch, + hasUncommittedChanges, + }); + if (analytics) track(analytics.event, analytics.properties); + }, + [taskId, linkedBranch, currentBranch, hasUncommittedChanges], + ); + + // One "shown" per appearance, re-armed when the banner hides. + const shownRef = useRef(false); + useEffect(() => { + if (shouldWarn && !shownRef.current) { + shownRef.current = true; + emitAction("shown"); + } else if (!shouldWarn) { + shownRef.current = false; + setActionError(null); + } + }, [shouldWarn, emitAction]); + + const trpc = useHostTRPC(); + const { mutate: checkoutBranch, isPending: isSwitching } = useMutation( + trpc.git.checkoutBranch.mutationOptions({ + onSuccess: () => { + if (repoPath) invalidateGitBranchQueries(repoPath); + // The mismatch clears on its own once the branch queries refetch; + // dismissing hides the banner immediately instead of a beat later. + dismissWarning(); + setActionError(null); + }, + onError: (error) => { + log.error("Failed to switch branch", error); + setActionError( + resolveBranchMismatchError(error, "Failed to switch branch"), + ); + }, + }), + ); + + const { mutate: linkBranch, isPending: isRelinking } = useMutation( + trpc.workspace.linkBranch.mutationOptions({ + // The LinkedBranchChanged event invalidates the workspace query, which + // clears the mismatch and hides the banner. + onSuccess: () => setActionError(null), + onError: (error) => { + log.error("Failed to re-link branch", error); + setActionError( + resolveBranchMismatchError(error, "Failed to update the task branch"), + ); + }, + }), + ); + + const handleSwitch = useCallback(() => { + const request = buildCheckoutBranchRequest(repoPath, linkedBranch); + if (!request) return; + setActionError(null); + emitAction("switch"); + checkoutBranch(request); + }, [repoPath, linkedBranch, emitAction, checkoutBranch]); + + const handleUseCurrentBranch = useCallback(() => { + const request = buildLinkBranchRequest(taskId, currentBranch); + if (!request) return; + setActionError(null); + emitAction("relink"); + linkBranch(request); + }, [taskId, currentBranch, emitAction, linkBranch]); + + const handleDismiss = useCallback(() => { + emitAction("dismiss"); + setActionError(null); + dismissWarning(); + }, [emitAction, dismissWarning]); + + if (!shouldWarn || !linkedBranch || !currentBranch) return null; + + return { + linkedBranch, + currentBranch, + actionError, + isSwitching, + isRelinking, + onSwitch: handleSwitch, + onUseCurrentBranch: handleUseCurrentBranch, + onDismiss: handleDismiss, + }; +} diff --git a/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts b/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts deleted file mode 100644 index 47ce066f9a..0000000000 --- a/packages/ui/src/features/workspace/useBranchMismatchDialog.test.ts +++ /dev/null @@ -1,282 +0,0 @@ -import { act, renderHook } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -let mockShouldWarn = false; -const mockDismissWarning = vi.fn(); - -const mockGuard = vi.hoisted(() => ({ - useBranchMismatchGuard: vi.fn( - (): { - shouldWarn: boolean; - linkedBranch: string | null; - currentBranch: string | null; - dismissWarning: () => void; - } => ({ - shouldWarn: mockShouldWarn, - linkedBranch: "feat/foo", - currentBranch: "main", - dismissWarning: mockDismissWarning, - }), - ), -})); -vi.mock("./useBranchMismatch", () => mockGuard); - -vi.mock("../git-interaction/useGitQueries", () => ({ - useGitQueries: () => ({ hasChanges: false }), -})); - -vi.mock("../git-interaction/gitCacheKeys", () => ({ - invalidateGitBranchQueries: vi.fn(), -})); - -vi.mock("@posthog/host-router/react", () => ({ - useHostTRPC: () => ({ - git: { - checkoutBranch: { - mutationOptions: (opts: Record) => opts, - }, - }, - }), -})); - -let capturedMutationOptions: { - onSuccess?: () => void; - onError?: (e: Error) => void; -} = {}; -const mockMutate = vi.fn(); - -vi.mock("@tanstack/react-query", () => ({ - useMutation: (opts: Record) => { - capturedMutationOptions = opts as typeof capturedMutationOptions; - return { mutate: mockMutate, isPending: false }; - }, -})); - -vi.mock("../../shell/logger", () => ({ - logger: { scope: () => ({ error: vi.fn() }) }, -})); - -const mockTrack = vi.fn(); -vi.mock("../../shell/analytics", () => ({ - track: (...args: unknown[]) => mockTrack(...args), -})); - -import { ANALYTICS_EVENTS } from "@posthog/shared"; -import { useBranchMismatchDialog } from "./useBranchMismatchDialog"; - -function renderDialog(overrides?: { shouldWarn?: boolean }) { - mockShouldWarn = overrides?.shouldWarn ?? false; - mockGuard.useBranchMismatchGuard.mockReturnValue({ - shouldWarn: mockShouldWarn, - linkedBranch: "feat/foo", - currentBranch: "main", - dismissWarning: mockDismissWarning, - }); - - const onSendPrompt = vi.fn(); - const hook = renderHook(() => - useBranchMismatchDialog({ - taskId: "task-1", - repoPath: "/repo", - onSendPrompt, - }), - ); - return { ...hook, onSendPrompt }; -} - -describe("useBranchMismatchDialog", () => { - beforeEach(() => { - vi.clearAllMocks(); - capturedMutationOptions = {}; - mockShouldWarn = false; - }); - - describe("handleBeforeSubmit", () => { - it("returns true when shouldWarn is false", () => { - const { result } = renderDialog({ shouldWarn: false }); - const clearEditor = vi.fn(); - - const allowed = result.current.handleBeforeSubmit("hello", clearEditor); - - expect(allowed).toBe(true); - expect(result.current.dialogProps?.open).toBeFalsy(); - }); - - it("returns false and opens dialog when shouldWarn is true", () => { - const { result } = renderDialog({ shouldWarn: true }); - const clearEditor = vi.fn(); - - let allowed = true; - act(() => { - allowed = result.current.handleBeforeSubmit("hello", clearEditor); - }); - - expect(allowed).toBe(false); - expect(result.current.dialogProps?.open).toBe(true); - expect(clearEditor).not.toHaveBeenCalled(); - expect(mockTrack).toHaveBeenCalledWith( - ANALYTICS_EVENTS.BRANCH_MISMATCH_WARNING_SHOWN, - { - task_id: "task-1", - linked_branch: "feat/foo", - current_branch: "main", - has_uncommitted_changes: false, - }, - ); - }); - }); - - describe("handleContinue", () => { - it("sends the pending message and clears editor", () => { - const { result, onSendPrompt } = renderDialog({ shouldWarn: true }); - const clearEditor = vi.fn(); - - act(() => { - result.current.handleBeforeSubmit("hello", clearEditor); - }); - - mockTrack.mockClear(); - act(() => { - result.current.dialogProps?.onContinue(); - }); - - expect(onSendPrompt).toHaveBeenCalledWith("hello"); - expect(clearEditor).toHaveBeenCalled(); - expect(mockDismissWarning).toHaveBeenCalled(); - expect(result.current.dialogProps?.open).toBe(false); - expect(mockTrack).toHaveBeenCalledWith( - ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, - { - task_id: "task-1", - action: "continue", - linked_branch: "feat/foo", - current_branch: "main", - }, - ); - }); - }); - - describe("handleCancel", () => { - it("clears state without sending or clearing editor", () => { - const { result, onSendPrompt } = renderDialog({ shouldWarn: true }); - const clearEditor = vi.fn(); - - act(() => { - result.current.handleBeforeSubmit("hello", clearEditor); - }); - expect(result.current.dialogProps?.open).toBe(true); - - mockTrack.mockClear(); - act(() => { - result.current.dialogProps?.onCancel(); - }); - - expect(onSendPrompt).not.toHaveBeenCalled(); - expect(clearEditor).not.toHaveBeenCalled(); - expect(result.current.dialogProps?.open).toBe(false); - expect(mockTrack).toHaveBeenCalledWith( - ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, - { - task_id: "task-1", - action: "cancel", - linked_branch: "feat/foo", - current_branch: "main", - }, - ); - }); - }); - - describe("handleSwitch", () => { - it("calls checkoutBranch mutation and tracks switch action", () => { - const { result } = renderDialog({ shouldWarn: true }); - const clearEditor = vi.fn(); - - act(() => { - result.current.handleBeforeSubmit("hello", clearEditor); - }); - - mockTrack.mockClear(); - act(() => { - result.current.dialogProps?.onSwitch(); - }); - - expect(mockMutate).toHaveBeenCalledWith({ - directoryPath: "/repo", - branchName: "feat/foo", - }); - expect(mockTrack).toHaveBeenCalledWith( - ANALYTICS_EVENTS.BRANCH_MISMATCH_ACTION, - { - task_id: "task-1", - action: "switch", - linked_branch: "feat/foo", - current_branch: "main", - }, - ); - }); - - it("on success: sends pending message and clears editor", () => { - const { result, onSendPrompt } = renderDialog({ shouldWarn: true }); - const clearEditor = vi.fn(); - - act(() => { - result.current.handleBeforeSubmit("hello", clearEditor); - }); - - act(() => { - result.current.dialogProps?.onSwitch(); - }); - - // Simulate mutation success - act(() => { - capturedMutationOptions.onSuccess?.(); - }); - - expect(onSendPrompt).toHaveBeenCalledWith("hello"); - expect(clearEditor).toHaveBeenCalled(); - expect(mockDismissWarning).toHaveBeenCalled(); - }); - - it("on error: shows error without sending message", () => { - const { result, onSendPrompt } = renderDialog({ shouldWarn: true }); - const clearEditor = vi.fn(); - - act(() => { - result.current.handleBeforeSubmit("hello", clearEditor); - }); - - act(() => { - result.current.dialogProps?.onSwitch(); - }); - - act(() => { - capturedMutationOptions.onError?.(new Error("dirty worktree")); - }); - - expect(onSendPrompt).not.toHaveBeenCalled(); - expect(clearEditor).not.toHaveBeenCalled(); - expect(result.current.dialogProps?.switchError).toBe("dirty worktree"); - }); - }); - - describe("dialogProps", () => { - it("is null when no linked branch", () => { - mockGuard.useBranchMismatchGuard.mockReturnValue({ - shouldWarn: false, - linkedBranch: null, - currentBranch: "main", - dismissWarning: mockDismissWarning, - }); - - const { result } = renderHook(() => - useBranchMismatchDialog({ - taskId: "task-1", - repoPath: "/repo", - onSendPrompt: vi.fn(), - }), - ); - - expect(result.current.dialogProps).toBeNull(); - }); - }); -}); diff --git a/packages/ui/src/features/workspace/useBranchMismatchDialog.ts b/packages/ui/src/features/workspace/useBranchMismatchDialog.ts deleted file mode 100644 index bc76a9255b..0000000000 --- a/packages/ui/src/features/workspace/useBranchMismatchDialog.ts +++ /dev/null @@ -1,143 +0,0 @@ -import { - type BranchMismatchDialogAction, - buildBranchMismatchAnalyticsEvent, - buildCheckoutBranchRequest, - decideBeforeSubmit, - resolveSwitchErrorMessage, -} from "@posthog/core/workspace/branchMismatchDialog"; -import { useHostTRPC } from "@posthog/host-router/react"; -import { ANALYTICS_EVENTS } from "@posthog/shared"; -import { useMutation } from "@tanstack/react-query"; -import { useCallback, useRef, useState } from "react"; -import { track } from "../../shell/analytics"; -import { logger } from "../../shell/logger"; -import { invalidateGitBranchQueries } from "../git-interaction/gitCacheKeys"; -import { useGitQueries } from "../git-interaction/useGitQueries"; -import { useBranchMismatchGuard } from "./useBranchMismatch"; - -const log = logger.scope("branch-mismatch"); - -interface UseBranchMismatchDialogOptions { - taskId: string; - repoPath: string | null; - onSendPrompt: (text: string) => void; -} - -export function useBranchMismatchDialog({ - taskId, - repoPath, - onSendPrompt, -}: UseBranchMismatchDialogOptions) { - const { shouldWarn, linkedBranch, currentBranch, dismissWarning } = - useBranchMismatchGuard(taskId); - - // State drives dialog visibility (`open`), refs avoid stale closures in - // mutation callbacks (onSuccess / handleContinue) that capture at mount time. - const [pendingMessage, setPendingMessage] = useState(null); - const pendingMessageRef = useRef(null); - const pendingClearRef = useRef<(() => void) | null>(null); - const onSendPromptRef = useRef(onSendPrompt); - onSendPromptRef.current = onSendPrompt; - const [switchError, setSwitchError] = useState(null); - - const { hasChanges: hasUncommittedChanges } = useGitQueries( - repoPath ?? undefined, - ); - - const emitAction = useCallback( - (action: BranchMismatchDialogAction) => { - const analytics = buildBranchMismatchAnalyticsEvent(action, { - taskId, - linkedBranch, - currentBranch, - hasUncommittedChanges, - }); - if (!analytics) return; - if (analytics.event === ANALYTICS_EVENTS.BRANCH_MISMATCH_WARNING_SHOWN) { - track(analytics.event, analytics.properties); - } else { - track(analytics.event, analytics.properties); - } - }, - [taskId, linkedBranch, currentBranch, hasUncommittedChanges], - ); - - const trpc = useHostTRPC(); - const { mutate: checkoutBranch, isPending: isSwitching } = useMutation( - trpc.git.checkoutBranch.mutationOptions({ - onSuccess: () => { - if (repoPath) invalidateGitBranchQueries(repoPath); - dismissWarning(); - pendingClearRef.current?.(); - pendingClearRef.current = null; - const message = pendingMessageRef.current; - if (message) onSendPromptRef.current(message); - setPendingMessage(null); - pendingMessageRef.current = null; - }, - onError: (error) => { - log.error("Failed to switch branch", error); - setSwitchError(resolveSwitchErrorMessage(error)); - }, - }), - ); - - const handleBeforeSubmit = useCallback( - (text: string, clearEditor: () => void): boolean => { - if (!decideBeforeSubmit(shouldWarn)) { - setPendingMessage(text); - pendingMessageRef.current = text; - pendingClearRef.current = clearEditor; - emitAction("shown"); - return false; - } - return true; - }, - [shouldWarn, emitAction], - ); - - const handleSwitch = useCallback(() => { - const request = buildCheckoutBranchRequest(repoPath, linkedBranch); - if (!request) return; - setSwitchError(null); - emitAction("switch"); - checkoutBranch(request); - }, [linkedBranch, repoPath, emitAction, checkoutBranch]); - - const handleContinue = useCallback(() => { - emitAction("continue"); - dismissWarning(); - pendingClearRef.current?.(); - pendingClearRef.current = null; - const message = pendingMessageRef.current; - if (message) onSendPromptRef.current(message); - setPendingMessage(null); - pendingMessageRef.current = null; - setSwitchError(null); - }, [dismissWarning, emitAction]); - - const handleCancel = useCallback(() => { - emitAction("cancel"); - setPendingMessage(null); - pendingMessageRef.current = null; - pendingClearRef.current = null; - setSwitchError(null); - }, [emitAction]); - - const dialogProps = - linkedBranch && currentBranch - ? { - open: pendingMessage !== null, - linkedBranch, - currentBranch, - hasUncommittedChanges, - switchError, - onSwitch: handleSwitch, - onContinue: handleContinue, - onCancel: handleCancel, - isSwitching, - } - : null; - - return { handleBeforeSubmit, dialogProps }; -} From 363bf8d1fa3bd4093606ec1e285f0a5a45cfc86a Mon Sep 17 00:00:00 2001 From: Adam Bowker Date: Mon, 6 Jul 2026 11:32:40 -0400 Subject: [PATCH 4/4] review: dismiss banner on relink success, cover linkBranch callbacks Greptile review feedback: "Use current branch" success now calls dismissWarning() like the checkout path does, so the banner hides immediately instead of waiting on the LinkedBranchChanged round-trip (and can't stick around if that event is dropped). Adds the missing tests for the linkBranch success and error callbacks, mirroring the switch-path coverage. Generated-By: PostHog Code Task-Id: adec8438-4ca0-42f0-9d9d-948de60c3b23 --- .../workspace/useBranchMismatchBanner.test.ts | 19 +++++++++++++++++++ .../workspace/useBranchMismatchBanner.ts | 10 +++++++--- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts b/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts index 2a7d4dab7b..dedd99e8fb 100644 --- a/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts +++ b/packages/ui/src/features/workspace/useBranchMismatchBanner.test.ts @@ -193,6 +193,25 @@ describe("useBranchMismatchBanner", () => { current_branch: "main", }, ); + + act(() => { + captured.link.onSuccess?.(); + }); + expect(mockDismissWarning).toHaveBeenCalled(); + }); + + it("onUseCurrentBranch failure surfaces the error without dismissing", () => { + const { result } = renderBanner(); + + act(() => { + result.current?.onUseCurrentBranch(); + }); + act(() => { + captured.link.onError?.(new Error("no such task")); + }); + + expect(result.current?.actionError).toBe("no such task"); + expect(mockDismissWarning).not.toHaveBeenCalled(); }); it("onDismiss dismisses for the session and tracks dismiss", () => { diff --git a/packages/ui/src/features/workspace/useBranchMismatchBanner.ts b/packages/ui/src/features/workspace/useBranchMismatchBanner.ts index ff17a36f84..331e562097 100644 --- a/packages/ui/src/features/workspace/useBranchMismatchBanner.ts +++ b/packages/ui/src/features/workspace/useBranchMismatchBanner.ts @@ -96,9 +96,13 @@ export function useBranchMismatchBanner({ const { mutate: linkBranch, isPending: isRelinking } = useMutation( trpc.workspace.linkBranch.mutationOptions({ - // The LinkedBranchChanged event invalidates the workspace query, which - // clears the mismatch and hides the banner. - onSuccess: () => setActionError(null), + onSuccess: () => { + // The LinkedBranchChanged event refreshes the workspace query and + // clears the mismatch; dismissing hides the banner immediately + // instead of a beat later (or never, if the event is dropped). + dismissWarning(); + setActionError(null); + }, onError: (error) => { log.error("Failed to re-link branch", error); setActionError(