From 1bb8a3501027a72bcef22a7cbc86a6e8eee6e65c Mon Sep 17 00:00:00 2001
From: Adam Bowker <adam.b@posthog.com>
Date: Thu, 30 Apr 2026 14:31:47 -0700
Subject: [PATCH] feat(code): harden branch linking during PR creation flow

---
 apps/code/src/main/services/git/service.test.ts       |  7 ++++---
 apps/code/src/main/services/git/service.ts            | 11 +++++++++++
 .../git-interaction/hooks/useGitInteraction.ts        | 11 -----------
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/apps/code/src/main/services/git/service.test.ts b/apps/code/src/main/services/git/service.test.ts
index ed891e798..3abc28e15 100644
--- a/apps/code/src/main/services/git/service.test.ts
+++ b/apps/code/src/main/services/git/service.test.ts
@@ -24,6 +24,7 @@ vi.mock("../../utils/logger.js", () => ({
 }));
 
 import type { LlmGatewayService } from "../llm-gateway/service";
+import type { WorkspaceService } from "../workspace/service";
 import { GitService } from "./service";
 
 describe("GitService.getPrChangedFiles", () => {
@@ -31,7 +32,7 @@ describe("GitService.getPrChangedFiles", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    service = new GitService({} as LlmGatewayService);
+    service = new GitService({} as LlmGatewayService, {} as WorkspaceService);
   });
 
   it("flattens paginated GH API results and maps file statuses", async () => {
@@ -139,7 +140,7 @@ describe("GitService.getGhAuthToken", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    service = new GitService({} as LlmGatewayService);
+    service = new GitService({} as LlmGatewayService, {} as WorkspaceService);
   });
 
   it("returns the authenticated GitHub CLI token", async () => {
@@ -197,7 +198,7 @@ describe("GitService.getPrUrlForBranch", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    service = new GitService({} as LlmGatewayService);
+    service = new GitService({} as LlmGatewayService, {} as WorkspaceService);
   });
 
   it("returns the PR URL for a branch via gh pr list", async () => {
diff --git a/apps/code/src/main/services/git/service.ts b/apps/code/src/main/services/git/service.ts
index a40eeaa8e..19d8da28b 100644
--- a/apps/code/src/main/services/git/service.ts
+++ b/apps/code/src/main/services/git/service.ts
@@ -41,6 +41,7 @@ import { MAIN_TOKENS } from "../../di/tokens";
 import { logger } from "../../utils/logger";
 import { TypedEventEmitter } from "../../utils/typed-event-emitter";
 import type { LlmGatewayService } from "../llm-gateway/service";
+import type { WorkspaceService } from "../workspace/service";
 import { CreatePrSaga } from "./create-pr-saga";
 import type {
   ChangedFile,
@@ -117,6 +118,8 @@ export class GitService extends TypedEventEmitter<GitServiceEvents> {
   constructor(
     @inject(MAIN_TOKENS.LlmGatewayService)
     private readonly llmGateway: LlmGatewayService,
+    @inject(MAIN_TOKENS.WorkspaceService)
+    private readonly workspaceService: WorkspaceService,
   ) {
     super();
   }
@@ -626,6 +629,14 @@ export class GitService extends TypedEventEmitter<GitServiceEvents> {
       includePrStatus: true,
     });
 
+    if (input.taskId) {
+      const linkedBranch =
+        input.branchName ?? (await getCurrentBranch(directoryPath));
+      if (linkedBranch) {
+        this.workspaceService.linkBranch(input.taskId, linkedBranch, "user");
+      }
+    }
+
     emitProgress(
       "complete",
       "Pull request created",
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 3581f3ffb..0b4c8cc98 100644
--- a/apps/code/src/renderer/features/git-interaction/hooks/useGitInteraction.ts
+++ b/apps/code/src/renderer/features/git-interaction/hooks/useGitInteraction.ts
@@ -282,17 +282,6 @@ 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) {