fix(server): cleanup follows the branch an agent checked out

apps/server/src/orchestration-v2/ProviderTurnStartService.test.ts

@@ -295,6 +295,7 @@ function makeLocalCommandHarness(input: {
   readonly resolveKnowledge?: ProjectKnowledge.ProjectKnowledge["Service"]["resolveForTurn"];
   readonly readSentDeliveries?: ProjectKnowledge.ProjectKnowledge["Service"]["readSentDeliveries"];
   readonly worktreePath?: string;
+  readonly branch?: string;
   readonly worktreesDir?: string;
   readonly workspaceRoot?: string;
   readonly gitWorkflow?: GitWorkflow.GitWorkflowService["Service"];
@@ -396,7 +397,7 @@ function makeLocalCommandHarness(input: {
       id: threadId,
       activeProviderThreadId: providerThreadId,
       projectId: ProjectId.make("project-native-account-command"),
-      branch: input.worktreePath === undefined ? null : "feature/restore",
+      branch: input.worktreePath === undefined ? null : (input.branch ?? "feature/restore"),
       worktreePath: input.worktreePath ?? null,
       lineage: { parentThreadId: null, relationshipToParent: null, rootThreadId: threadId },
     } as OrchestrationV2ThreadProjection["thread"],
@@ -1317,6 +1318,53 @@ effectIt.effect(
     }).pipe(Effect.scoped, Effect.provide(layerRealGit)),
 );
 
+effectIt.effect("recreates a rebound thread on the branch its agent switched to", () =>
+  Effect.gen(function* () {
+    const fileSystem = yield* FileSystem.FileSystem;
+    const path = yield* Path.Path;
+    const gitWorkflow = yield* GitWorkflow.GitWorkflowService;
+    const driver = yield* GitVcsDriver.GitVcsDriver;
+    const root = yield* fileSystem.makeTempDirectoryScoped({ prefix: "t3-recreate-rebound-" });
+    const workspaceRoot = path.join(root, "project");
+    const worktreesDir = path.join(root, "worktrees");
+    const worktreePath = path.join(worktreesDir, "project", "restore");
+    yield* fileSystem.makeDirectory(workspaceRoot);
+    const git = (cwd: string, args: ReadonlyArray<string>) =>
+      driver
+        .execute({ operation: "test.worktreeRebound", cwd, args })
+        .pipe(Effect.map((result) => result.stdout.trim()));
+    yield* git(workspaceRoot, ["init", "-b", "main"]);
+    yield* git(workspaceRoot, ["config", "user.name", "Test"]);
+    yield* git(workspaceRoot, ["config", "user.email", "test@example.com"]);
+    yield* git(workspaceRoot, ["commit", "--allow-empty", "-m", "base"]);
+    yield* git(workspaceRoot, ["worktree", "add", "-b", "feature/restore", worktreePath]);
+    yield* git(worktreePath, ["switch", "-c", "agent/work"]);
+    yield* git(worktreePath, ["commit", "--allow-empty", "-m", "agent work"]);
+    const agentTip = yield* git(worktreePath, ["rev-parse", "HEAD"]);
+    yield* fileSystem.remove(worktreePath, { recursive: true });
+    const harness = makeLocalCommandHarness({
+      text: "Continue",
+      branch: "agent/work",
+      workspaceRoot,
+      worktreePath,
+      worktreesDir,
+      fileSystem,
+      gitWorkflow,
+      failReadsAfterRunning: true,
+    });
+    const open = harness.open.getMockImplementation()!;
+    harness.open.mockImplementation(() =>
+      Effect.gen(function* () {
+        expect(yield* git(worktreePath, ["branch", "--show-current"])).toBe("agent/work");
+        expect(yield* git(worktreePath, ["rev-parse", "HEAD"])).toBe(agentTip);
+        return yield* open();
+      }).pipe(Effect.orDie),
+    );
+    yield* harness.start;
+    expect(harness.open).toHaveBeenCalledOnce();
+  }).pipe(Effect.scoped, Effect.provide(layerRealGit)),
+);
+
 effectIt.effect(
   "preserves recreation interrupts and retries, then fails the final attempt before provider start",
   () =>

apps/server/src/orchestration-v2/RunFinalizationQueue.test.ts

@@ -23,6 +23,7 @@ import * as RunFinalization from "./RunFinalizationService.ts";
 import * as RuntimeRequests from "./RuntimeRequestService.ts";
 import * as ThreadManagement from "./ThreadManagementService.ts";
 import * as ThreadTitles from "./ThreadTitleRegenerationService.ts";
+import * as WorktreeBranches from "./WorktreeBranches.ts";
 
 it.effect.each([
   ["starts the next turn after a post-capture Git status failure", false],
@@ -73,6 +74,7 @@ it.effect.each([
       Layer.mergeAll(
         projections,
         observer,
+        Layer.mock(WorktreeBranches.WorktreeBranches)({}),
         Layer.mock(CheckpointCapture.CheckpointCaptureServiceV2)({
           execute: () =>
             Effect.gen(function* () {

apps/server/src/orchestration-v2/RunFinalizationService.test.ts

@@ -4,9 +4,12 @@ import * as NodeFSP from "node:fs/promises";
 import * as NodeOS from "node:os";
 import * as NodePath from "node:path";
 
+import * as NodeCrypto from "@effect/platform-node/NodeCrypto";
+import * as NodeServices from "@effect/platform-node/NodeServices";
 import { assert, it } from "@effect/vitest";
 import {
   CheckpointScopeId,
+  GitCommandError,
   MessageId,
   ProjectId,
   ProviderInstanceId,
@@ -18,22 +21,27 @@ import * as Deferred from "effect/Deferred";
 import * as DateTime from "effect/DateTime";
 import * as Effect from "effect/Effect";
 import * as Fiber from "effect/Fiber";
+import * as FileSystem from "effect/FileSystem";
 import * as Layer from "effect/Layer";
 import * as Stream from "effect/Stream";
 import * as Option from "effect/Option";
 import * as TestClock from "effect/testing/TestClock";
 
+import * as ServerConfig from "../config.ts";
 import * as CodeIntelligence from "../project/knowledge/CodeIntelligence.ts";
 import * as ModelEvaluations from "../modelEvaluations/ModelEvaluations.ts";
 import * as PullRequestService from "../pullRequest/PullRequestService.ts";
+import * as GitVcsDriver from "../vcs/GitVcsDriver.ts";
 import * as VcsStatusBroadcaster from "../vcs/VcsStatusBroadcaster.ts";
 import * as WorkspaceEntries from "../workspace/WorkspaceEntries.ts";
 import * as WorkspacePaths from "../workspace/WorkspacePaths.ts";
 import { withWorkspaceLease } from "../workspace/workspaceLease.ts";
 import * as CheckpointCapture from "./CheckpointCaptureService.ts";
 import * as ProjectionStore from "./ProjectionStore.ts";
+import * as Orchestrator from "./Orchestrator.ts";
 import * as ProjectStore from "./ProjectStore.ts";
 import * as RunFinalization from "./RunFinalizationService.ts";
+import * as WorktreeBranches from "./WorktreeBranches.ts";
 
 it.effect("coalesces completed turns without starving later project refreshes", () => {
   const projectA = ProjectId.make("project-turn-refresh-a");
@@ -76,7 +84,7 @@ it.effect("coalesces completed turns without starving later project refreshes",
   }).pipe(Effect.provide(layer));
 });
 
-it.effect.each(["workspace refresh", "evaluation enqueue"] as const)(
+it.effect.each(["workspace refresh", "evaluation enqueue", "branch sync"] as const)(
   "preserves completed checkpoint finalization when %s fails",
   (failure) => {
     const threadId = ThreadId.make("thread_refresh_failure");
@@ -120,6 +128,26 @@ it.effect.each(["workspace refresh", "evaluation enqueue"] as const)(
                 checkpointScopes: [{ id: scopeId, runId, kind: "root_run" as const, cwd: "/repo" }],
                 checkpoints: [],
               }),
+            getThreadShell: () =>
+              Effect.succeed({
+                id: threadId,
+                branch: "feature",
+                worktreePath: "/repo",
+              } as OrchestrationV2ThreadShell),
+          }),
+          Layer.mock(WorktreeBranches.WorktreeBranches)({
+            sync: () =>
+              Effect.gen(function* () {
+                captured.push("sync");
+                if (failure === "branch sync")
+                  return yield* new GitCommandError({
+                    operation: "test",
+                    command: "git",
+                    cwd: "/repo",
+                    detail: "Git index is locked.",
+                  });
+                return { _tag: "Bound" as const, branch: "feature" };
+              }),
           }),
           Layer.succeed(RunFinalization.RunFinalizationObserver, {
             refresh: () =>
@@ -151,7 +179,7 @@ it.effect.each(["workspace refresh", "evaluation enqueue"] as const)(
     return Effect.gen(function* () {
       const service = yield* RunFinalization.RunFinalizationService;
       yield* service.finalize({ threadId, runId, scopeId });
-      assert.deepEqual(captured, ["checkpoint", "evaluation", "refresh"]);
+      assert.deepEqual(captured, ["checkpoint", "evaluation", "sync", "refresh"]);
     }).pipe(Effect.provide(layer));
   },
 );
@@ -239,6 +267,146 @@ it.effect.each(
   }).pipe(Effect.provide(layer));
 });
 
+it.live.each(
+  (
+    [
+      {
+        label: "rebinds a thread that ended on a new branch before its pull request refresh",
+        checkout: ["switch", "-c", "agent/work"],
+        worktreePath: "worktree",
+        refName: "agent/work",
+        branch: "agent/work",
+        lookups: 1,
+      },
+      {
+        label: "leaves a thread that ended on a detached HEAD alone",
+        checkout: ["switch", "--detach"],
+        worktreePath: "worktree",
+        refName: null,
+        branch: "feature",
+        lookups: 0,
+      },
+      {
+        label: "leaves a thread that ended on the default branch alone",
+        checkout: ["switch", "main"],
+        worktreePath: "worktree",
+        refName: "main",
+        branch: "feature",
+        lookups: 0,
+      },
+      {
+        label: "leaves a thread working in the project root alone",
+        checkout: ["switch", "-c", "agent/work"],
+        worktreePath: null,
+        refName: "agent/work",
+        branch: "feature",
+        lookups: 0,
+      },
+    ] as const
+  ).map((scenario) => [scenario.label, scenario] as const),
+)("%s", ([, scenario]) =>
+  Effect.gen(function* () {
+    const threadId = ThreadId.make("thread-branch-sync");
+    const runId = RunId.make("completed-run");
+    const scopeId = CheckpointScopeId.make("scope-branch-sync");
+    const lookedUp = yield* Deferred.make<void>();
+    const fs = yield* FileSystem.FileSystem;
+    const git = yield* GitVcsDriver.GitVcsDriver;
+    const root = yield* fs.realPath(
+      yield* fs.makeTempDirectoryScoped({ prefix: "t3-branch-sync-" }),
+    );
+    const repo = `${root}/repo`;
+    const worktree = `${root}/feature`;
+    yield* fs.makeDirectory(repo);
+    const command = (cwd: string, args: ReadonlyArray<string>) =>
+      git.execute({ operation: "run-finalization-test", cwd, args });
+    yield* command(repo, ["init", "-b", "main"]);
+    yield* command(repo, ["config", "user.email", "finalize@example.test"]);
+    yield* command(repo, ["config", "user.name", "Finalize test"]);
+    yield* command(repo, ["commit", "--allow-empty", "-m", "initial"]);
+    yield* command(repo, ["worktree", "add", "-b", "feature", worktree]);
+    if (scenario.refName === "main") yield* command(repo, ["switch", "--detach"]);
+    yield* command(worktree, scenario.checkout);
+
+    let thread = {
+      id: threadId,
+      branch: "feature",
+      worktreePath: scenario.worktreePath === null ? null : worktree,
+      activeRunId: null,
+      archivedAt: null,
+      createdAt: DateTime.makeUnsafe(0),
+    } as OrchestrationV2ThreadShell;
+    let lookups = 0;
+    const projections = Layer.mock(ProjectionStore.ProjectionStoreV2)({
+      getCheckpointContext: () =>
+        Effect.succeed({
+          runs: [],
+          checkpointScopes: [{ id: scopeId, runId, kind: "root_run" as const, cwd: worktree }],
+          checkpoints: [],
+        }),
+      getThreadShell: () => Effect.sync(() => thread),
+    });
+    const layer = RunFinalization.layer.pipe(
+      Layer.provideMerge(RunFinalization.layerObserver),
+      Layer.provide(WorktreeBranches.layer),
+      Layer.provide(
+        Layer.mergeAll(
+          projections,
+          Layer.mock(ProjectStore.ProjectStoreV2)({
+            getShell: () => Effect.succeedSome({ workspaceRoot: repo }) as never,
+          }),
+          Layer.mock(ModelEvaluations.ModelEvaluations)({}),
+          Layer.mock(CheckpointCapture.CheckpointCaptureServiceV2)({ execute: () => Effect.void }),
+          Layer.mock(Orchestrator.OrchestratorV2)({
+            getShellSnapshot: () =>
+              Effect.sync(() => ({ threads: [thread], archivedThreads: [] })) as never,
+            dispatch: ((command: { readonly branch: string }) =>
+              Effect.sync(() => {
+                thread = { ...thread, branch: command.branch };
+              })) as never,
+          }),
+          Layer.mock(WorkspaceEntries.WorkspaceEntries)({ refresh: () => Effect.void }),
+          Layer.mock(CodeIntelligence.CodeIntelligence)({ refreshIfIndexed: () => Effect.void }),
+          Layer.mock(PullRequestService.PullRequestService)({
+            refreshAfterTurn: () => Effect.void,
+          }),
+          Layer.mock(VcsStatusBroadcaster.VcsStatusBroadcaster)({
+            refreshLocalStatus: () =>
+              Effect.succeed({
+                isRepo: true,
+                hasPrimaryRemote: false,
+                isDefaultRef: scenario.refName === "main",
+                refName: scenario.refName,
+                hasWorkingTreeChanges: false,
+                workingTree: { files: [], insertions: 0, deletions: 0 },
+              }),
+            refreshPullRequestStatus: () =>
+              Effect.suspend(() => {
+                lookups++;
+                return Deferred.succeed(lookedUp, undefined).pipe(Effect.as(null));
+              }),
+          }),
+        ),
+      ),
+    );
+    yield* Effect.gen(function* () {
+      const service = yield* RunFinalization.RunFinalizationService;
+      yield* service.finalize({ threadId, runId, scopeId });
+      if (scenario.lookups > 0) yield* Deferred.await(lookedUp);
+      assert.equal(thread.branch, scenario.branch);
+      assert.equal(lookups, scenario.lookups);
+    }).pipe(Effect.provide(layer));
+  }).pipe(
+    Effect.provide(
+      GitVcsDriver.layer.pipe(
+        Layer.provideMerge(ServerConfig.layerTest(process.cwd(), { prefix: "t3-branch-sync-" })),
+        Layer.provideMerge(Layer.mergeAll(NodeServices.layer, NodeCrypto.layer)),
+      ),
+    ),
+    Effect.scoped,
+  ),
+);
+
 it.effect("finalizes a run without waiting for the code index refresh its checkpoint started", () =>
   Effect.acquireUseRelease(
     Effect.promise(async () => {
@@ -263,6 +431,7 @@ it.effect("finalizes a run without waiting for the code index refresh its checkp
             Layer.provide(
               Layer.mergeAll(
                 Layer.mock(ModelEvaluations.ModelEvaluations)({}),
+                Layer.mock(WorktreeBranches.WorktreeBranches)({}),
                 Layer.mock(CheckpointCapture.CheckpointCaptureServiceV2)({
                   execute: () => Effect.void,
                 }),

apps/server/src/orchestration-v2/RunFinalizationService.ts

@@ -13,6 +13,7 @@ import * as WorkspaceEntries from "../workspace/WorkspaceEntries.ts";
 import * as ModelEvaluations from "../modelEvaluations/ModelEvaluations.ts";
 import * as CheckpointCapture from "./CheckpointCaptureService.ts";
 import * as ProjectionStore from "./ProjectionStore.ts";
+import * as WorktreeBranches from "./WorktreeBranches.ts";
 
 export class RunFinalizationError extends Schema.TaggedError<RunFinalizationError>()(
   "RunFinalizationError",
@@ -57,6 +58,7 @@ const make = Effect.gen(function* () {
   const projections = yield* ProjectionStore.ProjectionStoreV2;
   const observer = yield* RunFinalizationObserver;
   const modelEvaluations = yield* ModelEvaluations.ModelEvaluations;
+  const worktreeBranches = yield* WorktreeBranches.WorktreeBranches;
 
   const finalize: RunFinalizationService["Service"]["finalize"] = Effect.fn(
     "RunFinalizationService.finalize",
@@ -94,6 +96,20 @@ const make = Effect.gen(function* () {
         );
     }
     if (cwd !== undefined) {
+      yield* projections.getThreadShell(input.threadId).pipe(
+        Effect.flatMap((thread) =>
+          // An exact match excludes project-root threads and nested-repo checkpoint scopes.
+          thread?.worktreePath === cwd
+            ? worktreeBranches.sync({ worktreePath: cwd, runId: input.runId })
+            : Effect.void,
+        ),
+        Effect.catchCause((cause) =>
+          Effect.logWarning("Failed to sync worktree branch after checkpoint capture", {
+            ...input,
+            cause,
+          }),
+        ),
+      );
       // Retrying a display refresh after the checkpoint is durable blocks the next turn.
       yield* observer.refresh({ cwd, threadId: input.threadId, runId: input.runId }).pipe(
         Effect.catchTags({