From ab59a633049273cb8721ef168292d3c3b028ab68 Mon Sep 17 00:00:00 2001
From: Adam Bowker <adam.b@posthog.com>
Date: Thu, 9 Apr 2026 10:54:27 -0700
Subject: [PATCH] feat(code): enable branch linking for local tasks

---
 .../src/main/services/workspace/service.ts    | 35 +++++++++++++++++++
 .../hooks/useGitInteraction.ts                | 17 +++++++++
 2 files changed, 52 insertions(+)

diff --git a/apps/code/src/main/services/workspace/service.ts b/apps/code/src/main/services/workspace/service.ts
index dd64795b4..8ad5bafb5 100644
--- a/apps/code/src/main/services/workspace/service.ts
+++ b/apps/code/src/main/services/workspace/service.ts
@@ -22,6 +22,7 @@ import { MAIN_TOKENS } from "../../di/tokens";
 import { logger } from "../../utils/logger";
 import { TypedEventEmitter } from "../../utils/typed-event-emitter";
 import { deriveWorktreePath } from "../../utils/worktree-helpers";
+import { AgentServiceEvent } from "../agent/schemas";
 import type { AgentService } from "../agent/service";
 import { FileWatcherEvent } from "../file-watcher/schemas";
 import type { FileWatcherService } from "../file-watcher/service";
@@ -237,6 +238,11 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
       this.handleFocusBranchRenamed.bind(this),
     );
 
+    this.agentService.on(
+      AgentServiceEvent.AgentFileActivity,
+      this.handleAgentFileActivity.bind(this),
+    );
+
     log.info("Branch watcher initialized");
   }
 
@@ -310,6 +316,35 @@ export class WorkspaceService extends TypedEventEmitter<WorkspaceServiceEvents>
     }
   }
 
+  private async handleAgentFileActivity({
+    taskId,
+    branchName,
+  }: {
+    taskId: string;
+    branchName: string | null;
+  }): Promise<void> {
+    if (!branchName) return;
+
+    const dbRow = this.workspaceRepo.findByTaskId(taskId);
+    if (!dbRow || dbRow.mode !== "local") return;
+    if (!dbRow.repositoryId) return;
+
+    const folderPath = this.getFolderPath(dbRow.repositoryId);
+    if (!folderPath) return;
+
+    try {
+      const defaultBranch = await getDefaultBranch(folderPath);
+      if (branchName === defaultBranch) return;
+    } catch {
+      // If we can't determine the default branch, still allow linking
+    }
+
+    const currentLinked = dbRow.linkedBranch ?? null;
+    if (currentLinked === branchName) return;
+
+    this.linkBranch(taskId, branchName);
+  }
+
   private updateAssociationBranchName(
     _taskId: string,
     _branchName: string,
diff --git a/apps/code/src/renderer/features/git-interaction/hooks/useGitInteraction.ts b/apps/code/src/renderer/features/git-interaction/hooks/useGitInteraction.ts
index acfffeb77..10b84473c 100644
--- a/apps/code/src/renderer/features/git-interaction/hooks/useGitInteraction.ts
+++ b/apps/code/src/renderer/features/git-interaction/hooks/useGitInteraction.ts
@@ -276,6 +276,17 @@ export function useGitInteraction(
       }
       if (store.createPrNeedsBranch) {
         invalidateGitBranchQueries(repoPath);
+        trpcClient.workspace.linkBranch
+          .mutate({ taskId, branchName: store.branchName.trim() })
+          .catch((err) =>
+            log.warn("Failed to link branch to task", { taskId, err }),
+          );
+      } else if (git.currentBranch) {
+        trpcClient.workspace.linkBranch
+          .mutate({ taskId, branchName: git.currentBranch })
+          .catch((err) =>
+            log.warn("Failed to link branch to task", { taskId, err }),
+          );
       }
 
       if (result.prUrl) {
@@ -529,6 +540,12 @@ export function useGitInteraction(
       trackGitAction(taskId, "branch-here", true);
       await queryClient.invalidateQueries(trpc.workspace.getAll.pathFilter());
 
+      trpcClient.workspace.linkBranch
+        .mutate({ taskId, branchName: store.branchName.trim() })
+        .catch((err) =>
+          log.warn("Failed to link branch to task", { taskId, err }),
+        );
+
       modal.closeBranch();
     } catch (error) {
       log.error("Failed to create branch", error);