Preserve pull request status across branch checkouts

apps/mobile/src/state/use-thread-pr.ts

@@ -15,6 +15,10 @@ export {
  * (environmentId, cwd) by the atom family, so many rows on the same worktree
  * or project root share one stream — and virtualization means only visible
  * rows subscribe at all.
+ *
+ * Keyed by the thread's recorded branch, not the current checkout: the
+ * streamed status is the fast path while the checkout matches, and a
+ * branch-keyed lookup keeps the badge alive once the checkout moves away.
  */
 export function useThreadPr(
   thread: EnvironmentThreadShell,
@@ -29,13 +33,25 @@ export function useThreadPr(
         })
       : null,
   );
+  const branchPr = useEnvironmentQuery(
+    thread.branch !== null && cwd !== null
+      ? vcsEnvironment.branchPr({
+          environmentId: thread.environmentId,
+          input: { cwd, branch: thread.branch },
+        })
+      : null,
+  );
 
   const status = gitStatus.data;
-  if (status === null || thread.branch === null || status.refName !== thread.branch) {
-    return null;
+  // When the status describes the thread's own branch it is authoritative
+  // about that branch's PR — including when it reports none, so this must not
+  // fall through to a warm branch lookup that would resurrect a stale badge.
+  if (status !== null && thread.branch !== null && status.refName === thread.branch) {
+    return status.pr ? presentThreadPr(status.pr, status.sourceControlProvider) : null;
   }
-  if (!status.pr) {
+  const pr = branchPr.data?.pr;
+  if (!pr) {
     return null;
   }
-  return presentThreadPr(status.pr, status.sourceControlProvider);
+  return presentThreadPr(pr, status?.sourceControlProvider);
 }

apps/server/src/auth/RpcAuthorization.ts

@@ -64,6 +64,7 @@ export const RPC_REQUIRED_SCOPES = {
   [WS_METHODS.subscribeVcsStatus]: AuthOrchestrationReadScope,
   [WS_METHODS.subscribeResourceTelemetry]: AuthOrchestrationReadScope,
   [WS_METHODS.vcsRefreshStatus]: AuthOrchestrationReadScope,
+  [WS_METHODS.vcsBranchPr]: AuthOrchestrationReadScope,
   [WS_METHODS.vcsPull]: AuthOrchestrationOperateScope,
   [WS_METHODS.gitRunStackedAction]: AuthOrchestrationOperateScope,
   [WS_METHODS.gitResolvePullRequest]: AuthOrchestrationOperateScope,

apps/server/src/git/GitManager.ts

@@ -23,6 +23,8 @@ import {
   GitRunStackedActionInput,
   GitRunStackedActionResult,
   GitStackedAction,
+  type VcsBranchPrInput,
+  type VcsBranchPrResult,
   VcsStatusInput,
   type VcsStatusLocalResult,
   type VcsStatusRemoteResult,
@@ -87,6 +89,9 @@ export class GitManager extends Context.Service<
       input: VcsStatusInput,
       options?: GitVcsDriver.GitRemoteStatusOptions,
     ) => Effect.Effect<VcsStatusRemoteResult | null, GitManagerServiceError>;
+    readonly branchPr: (
+      input: VcsBranchPrInput,
+    ) => Effect.Effect<VcsBranchPrResult, GitManagerServiceError>;
     readonly invalidateLocalStatus: (cwd: string) => Effect.Effect<void, never>;
     readonly invalidateRemoteStatus: (cwd: string) => Effect.Effect<void, never>;
     readonly invalidateStatus: (cwd: string) => Effect.Effect<void, never>;
@@ -1684,6 +1689,43 @@ export const make = Effect.gen(function* () {
     });
     return mergeGitStatusParts(local, remote);
   });
+  // PR lookup for a branch that is not necessarily checked out. Threads keep
+  // a recorded branch, and their PR badge must survive the checkout moving
+  // elsewhere — so this derives the branch's upstream from git config instead
+  // of HEAD and funnels into the same cached lookup as status.
+  const branchPr: GitManager["Service"]["branchPr"] = Effect.fn("branchPr")(function* (input) {
+    const cwd = yield* normalizeStatusCacheKey(input.cwd);
+    const branch = input.branch;
+    const [remoteName, mergeRef, originHead] = yield* Effect.all(
+      [
+        readConfigValueNullable(cwd, `branch.${branch}.remote`),
+        readConfigValueNullable(cwd, `branch.${branch}.merge`),
+        gitCore
+          .execute({
+            operation: "GitManager.branchPr.defaultBranch",
+            cwd,
+            args: ["symbolic-ref", "--short", "refs/remotes/origin/HEAD"],
+            allowNonZeroExit: true,
+          })
+          .pipe(Effect.orElseSucceed(() => null)),
+      ],
+      { concurrency: "unbounded" },
+    );
+    const upstreamRef =
+      remoteName !== null && mergeRef !== null && mergeRef.startsWith("refs/heads/")
+        ? `${remoteName}/${mergeRef.slice("refs/heads/".length)}`
+        : null;
+    const originHeadBranch =
+      originHead !== null && originHead.exitCode === 0
+        ? (originHead.stdout.trim().split("/").slice(1).join("/") ?? "")
+        : "";
+    const isDefaultBranch =
+      originHeadBranch.length > 0
+        ? branch === originHeadBranch
+        : branch === "main" || branch === "master";
+    const pr = yield* lookupStatusPr(cwd, { branch, upstreamRef, isDefaultBranch });
+    return { branch, pr } satisfies VcsBranchPrResult;
+  });
   const invalidateLocalStatus: GitManager["Service"]["invalidateLocalStatus"] = Effect.fn(
     "invalidateLocalStatus",
   )(function* (cwd) {
@@ -2136,6 +2178,7 @@ export const make = Effect.gen(function* () {
   return GitManager.of({
     localStatus,
     remoteStatus,
+    branchPr,
     status,
     invalidateLocalStatus,
     invalidateRemoteStatus,

apps/server/src/git/GitWorkflowService.ts

@@ -22,6 +22,8 @@ import {
   type GitResolvePullRequestResult,
   type GitRunStackedActionInput,
   type GitRunStackedActionResult,
+  type VcsBranchPrInput,
+  type VcsBranchPrResult,
   type VcsStatusInput,
   type VcsStatusLocalResult,
   type VcsStatusRemoteResult,
@@ -45,6 +47,9 @@ export class GitWorkflowService extends Context.Service<
       input: VcsStatusInput,
       options?: GitVcsDriver.GitRemoteStatusOptions,
     ) => Effect.Effect<VcsStatusRemoteResult | null, GitManagerServiceError>;
+    readonly branchPr: (
+      input: VcsBranchPrInput,
+    ) => Effect.Effect<VcsBranchPrResult, GitManagerServiceError>;
     readonly invalidateLocalStatus: (cwd: string) => Effect.Effect<void, never>;
     readonly invalidateRemoteStatus: (cwd: string) => Effect.Effect<void, never>;
     readonly invalidateStatus: (cwd: string) => Effect.Effect<void, never>;
@@ -270,6 +275,14 @@ export const make = Effect.gen(function* () {
           isGitRepository ? gitManager.remoteStatus(input, options) : Effect.succeed(null),
         ),
       ),
+    branchPr: (input) =>
+      detectGitRepositoryForStatus("GitWorkflowService.branchPr", input.cwd).pipe(
+        Effect.flatMap((isGitRepository) =>
+          isGitRepository
+            ? gitManager.branchPr(input)
+            : Effect.succeed({ branch: input.branch, pr: null }),
+        ),
+      ),
     invalidateLocalStatus: gitManager.invalidateLocalStatus,
     invalidateRemoteStatus: gitManager.invalidateRemoteStatus,
     invalidateStatus: gitManager.invalidateStatus,

apps/server/src/vcs/VcsStatusBroadcaster.test.ts

@@ -6,6 +6,7 @@ import * as DateTime from "effect/DateTime";
 import * as Duration from "effect/Duration";
 import * as Effect from "effect/Effect";
 import * as Exit from "effect/Exit";
+import * as Fiber from "effect/Fiber";
 import * as FileSystem from "effect/FileSystem";
 import * as Layer from "effect/Layer";
 import * as Logger from "effect/Logger";
@@ -208,7 +209,7 @@ describe("VcsStatusBroadcaster", () => {
     }).pipe(Effect.provide(makeTestLayer(state)));
   });
 
-  it.effect("keeps the cached snapshot unchanged when a refresh branch fails", () => {
+  it.effect("keeps local progress and recovers the remote half after a failed refresh", () => {
     const state = {
       currentLocalStatus: baseLocalStatus,
       currentRemoteStatus: baseRemoteStatus,
@@ -273,14 +274,104 @@ describe("VcsStatusBroadcaster", () => {
       state.failRemoteStatus = true;
 
       const refreshExit = yield* broadcaster.refreshStatus("/repo").pipe(Effect.exit);
-      const cached = yield* broadcaster.getStatus({ cwd: "/repo" });
-
       assert.isTrue(Exit.isFailure(refreshExit));
-      assert.deepStrictEqual(cached, baseStatus);
+
+      // The local half landed even though the refresh failed. The remote half
+      // was computed for the previous checkout, so the refName change dropped
+      // it rather than pairing the old branch's data with the new branch; it
+      // repopulates on the next successful read.
+      state.failRemoteStatus = false;
+      const cached = yield* broadcaster.getStatus({ cwd: "/repo" });
+      assert.deepStrictEqual(cached, {
+        ...state.currentLocalStatus,
+        ...state.currentRemoteStatus,
+      });
     }).pipe(Effect.provide(testLayer));
   });
 
-  it.effect("refreshes only the cached local snapshot when requested", () => {
+  it.effect("returns the new checkout when one lands mid-refresh", () => {
+    const state = {
+      currentLocalStatus: baseLocalStatus,
+      currentRemoteStatus: baseRemoteStatus,
+      localStatusCalls: 0,
+      remoteStatusCalls: 0,
+      localInvalidationCalls: 0,
+      remoteInvalidationCalls: 0,
+    };
+
+    return Effect.gen(function* () {
+      const remoteEntered = yield* Deferred.make<void>();
+      const releaseRemote = yield* Deferred.make<void>();
+      // Armed only around the refresh under test, so neither the priming read
+      // nor the follow-up refresh that the checkout change kicks off blocks.
+      const gate = { armed: false };
+      const testLayer = VcsStatusBroadcaster.layer.pipe(
+        Layer.provideMerge(NodeServices.layer),
+        Layer.provide(makeBackgroundPolicyLayer(() => true)),
+        Layer.provide(
+          Layer.mock(GitWorkflowService.GitWorkflowService)({
+            localStatus: () =>
+              Effect.sync(() => {
+                state.localStatusCalls += 1;
+                return state.currentLocalStatus;
+              }),
+            remoteStatus: () =>
+              Effect.suspend(() => {
+                state.remoteStatusCalls += 1;
+                if (!gate.armed) {
+                  return Effect.succeed(state.currentRemoteStatus);
+                }
+                gate.armed = false;
+                return Deferred.succeed(remoteEntered, undefined).pipe(
+                  Effect.andThen(Deferred.await(releaseRemote)),
+                  Effect.as(state.currentRemoteStatus),
+                );
+              }),
+            invalidateLocalStatus: () =>
+              Effect.sync(() => {
+                state.localInvalidationCalls += 1;
+              }),
+            invalidateRemoteStatus: () =>
+              Effect.sync(() => {
+                state.remoteInvalidationCalls += 1;
+              }),
+            invalidateStatus: () =>
+              Effect.sync(() => {
+                state.localInvalidationCalls += 1;
+                state.remoteInvalidationCalls += 1;
+              }),
+          }),
+        ),
+      );
+
+      return yield* Effect.gen(function* () {
+        const broadcaster = yield* VcsStatusBroadcaster.VcsStatusBroadcaster;
+        yield* broadcaster.getStatus({ cwd: "/repo" });
+
+        gate.armed = true;
+        const refreshFiber = yield* Effect.forkChild(broadcaster.refreshStatus("/repo"), {
+          startImmediately: true,
+        });
+        yield* Deferred.await(remoteEntered);
+
+        // A checkout lands while the slow remote read is still in flight.
+        state.currentLocalStatus = {
+          ...baseLocalStatus,
+          refName: "feature/checked-out-mid-refresh",
+        };
+        yield* broadcaster.refreshLocalStatus("/repo");
+
+        yield* Deferred.succeed(releaseRemote, undefined);
+        const refreshed = yield* Fiber.join(refreshFiber);
+
+        // The refresh must not hand back the branch that was checked out when
+        // it started; its own remote result was discarded as superseded.
+        assert.equal(refreshed.refName, "feature/checked-out-mid-refresh");
+      }).pipe(Effect.provide(testLayer));
+    });
+  });
+
+  it.effect("refreshes only the cached local snapshot when the checkout is unchanged", () => {
     const state = {
       currentLocalStatus: baseLocalStatus,
       currentRemoteStatus: baseRemoteStatus,
@@ -296,7 +387,6 @@ describe("VcsStatusBroadcaster", () => {
 
       state.currentLocalStatus = {
         ...baseLocalStatus,
-        refName: "feature/local-only-refresh",
         hasWorkingTreeChanges: true,
       };
 
@@ -316,6 +406,53 @@ describe("VcsStatusBroadcaster", () => {
     }).pipe(Effect.provide(makeTestLayer(state)));
   });
 
+  it.effect("repopulates the remote half when a local refresh changes the checkout", () => {
+    const state = {
+      currentLocalStatus: baseLocalStatus,
+      currentRemoteStatus: baseRemoteStatus,
+      localStatusCalls: 0,
+      remoteStatusCalls: 0,
+      localInvalidationCalls: 0,
+      remoteInvalidationCalls: 0,
+    };
+
+    return Effect.gen(function* () {
+      const broadcaster = yield* VcsStatusBroadcaster.VcsStatusBroadcaster;
+      const initial = yield* broadcaster.getStatus({ cwd: "/repo" });
+
+      state.currentLocalStatus = {
+        ...baseLocalStatus,
+        refName: "feature/local-only-refresh",
+        hasWorkingTreeChanges: true,
+      };
+      state.currentRemoteStatus = {
+        ...baseRemoteStatus,
+        aheadCount: 4,
+      };
+
+      const refreshedLocal = yield* broadcaster.refreshLocalStatus("/repo");
+
+      // The refName change drops the stale remote half and forks a remote
+      // refresh for the new branch; let that fiber settle.
+      for (let i = 0; i < 100 && state.remoteStatusCalls < 2; i++) {
+        yield* Effect.yieldNow;
+      }
+
+      const cached = yield* broadcaster.getStatus({ cwd: "/repo" });
+
+      assert.deepStrictEqual(initial, baseStatus);
+      assert.deepStrictEqual(refreshedLocal, state.currentLocalStatus);
+      assert.deepStrictEqual(cached, {
+        ...state.currentLocalStatus,
+        ...state.currentRemoteStatus,
+      });
+      assert.equal(state.localStatusCalls, 2);
+      assert.equal(state.remoteStatusCalls, 2);
+      assert.equal(state.localInvalidationCalls, 1);
+      assert.equal(state.remoteInvalidationCalls, 1);
+    }).pipe(Effect.provide(makeTestLayer(state)));
+  });
+
   it.effect("normalizes symlinked CWDs before cache lookup and workflow calls", () => {
     const seenCwds: string[] = [];
     const state = {