trunk-merge/pr-3684/9fcb83ff-3102-4a4f-ab61-0a19e076582f

packages/agent/src/server/agent-server.ts

@@ -102,6 +102,7 @@ import {
 import { TaskRunEventStreamSender } from "./event-stream-sender";
 import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt";
 import { type McpRelayResponse, McpRelayServer } from "./mcp-relay-server";
+import { checkoutExistingPullRequest } from "./pr-checkout";
 import { resolveRtkSavings } from "./rtk-savings";
 import { RunUsageAccumulator } from "./run-usage";
 import {
@@ -1625,6 +1626,15 @@ export class AgentServer {
     };
 
     await this.waitForRepoReady();
+    const existingPrCheckoutPromise =
+      prUrl &&
+      this.config.repositoryPath &&
+      this.shouldAutoPublishCloudChanges()
+        ? checkoutExistingPullRequest({
+            repositoryPath: this.config.repositoryPath,
+            prUrl,
+          })
+        : null;
     await this.installSkillBundleArtifacts(
       payload.task_id,
       payload.run_id,
@@ -1648,6 +1658,25 @@ export class AgentServer {
       ...(await this.startMcpRelayServer()),
     ];
 
+    if (existingPrCheckoutPromise) {
+      const checkoutResult = await existingPrCheckoutPromise;
+      if (checkoutResult.status === "failed") {
+        this.logger.warn(
+          "Existing PR pre-checkout failed; agent will retry if needed",
+          {
+            prUrl,
+            error: checkoutResult.error,
+          },
+        );
+      } else {
+        this.logger.debug("Existing PR branch prepared before session start", {
+          prUrl,
+          branch: checkoutResult.branch,
+          alreadyActive: checkoutResult.status === "already_active",
+        });
+      }
+    }
+
     let acpSessionId: string | null = null;
     if (nativeResume) {
       try {

packages/agent/src/server/pr-checkout.test.ts

@@ -0,0 +1,61 @@
+import { describe, expect, it, vi } from "vitest";
+import { checkoutExistingPullRequest } from "./pr-checkout";
+
+describe("checkoutExistingPullRequest", () => {
+  it.each([
+    {
+      name: "skips checkout when the pull request branch is already active",
+      currentBranch: "posthog-code/fix-checkout",
+      expectedStatus: "already_active",
+      expectedCheckoutCalls: 0,
+    },
+    {
+      name: "checks out the pull request when another branch is active",
+      currentBranch: "main",
+      expectedStatus: "checked_out",
+      expectedCheckoutCalls: 1,
+    },
+  ])(
+    "$name",
+    async ({ currentBranch, expectedStatus, expectedCheckoutCalls }) => {
+      const runCommand = vi.fn(
+        async (
+          executable: string,
+          args: string[],
+        ): Promise<{ stdout: string }> => {
+          if (executable === "git") {
+            return { stdout: `${currentBranch}\n` };
+          }
+          if (args[1] === "view") {
+            return { stdout: "posthog-code/fix-checkout\n" };
+          }
+          return { stdout: "" };
+        },
+      );
+
+      const result = await checkoutExistingPullRequest({
+        repositoryPath: "/tmp/repo",
+        prUrl: "https://github.com/PostHog/code/pull/1",
+        runCommand,
+      });
+
+      expect(result.status).toBe(expectedStatus);
+      expect(
+        runCommand.mock.calls.filter(
+          ([executable, args]) =>
+            executable === "gh" && args[0] === "pr" && args[1] === "checkout",
+        ),
+      ).toHaveLength(expectedCheckoutCalls);
+    },
+  );
+
+  it("returns a failure so startup can fall back to agent checkout", async () => {
+    const result = await checkoutExistingPullRequest({
+      repositoryPath: "/tmp/repo",
+      prUrl: "https://github.com/PostHog/code/pull/1",
+      runCommand: vi.fn().mockRejectedValue(new Error("gh unavailable")),
+    });
+
+    expect(result).toEqual({ status: "failed", error: "gh unavailable" });
+  });
+});

packages/agent/src/server/pr-checkout.ts

@@ -0,0 +1,68 @@
+import { execFile } from "node:child_process";
+import { promisify } from "node:util";
+
+const execFileAsync = promisify(execFile);
+
+type CommandResult = { stdout: string };
+type RunCommand = (
+  executable: string,
+  args: string[],
+  cwd: string,
+) => Promise<CommandResult>;
+
+export type ExistingPrCheckoutResult =
+  | { status: "already_active"; branch: string }
+  | { status: "checked_out"; branch: string }
+  | { status: "failed"; error: string };
+
+async function defaultRunCommand(
+  executable: string,
+  args: string[],
+  cwd: string,
+): Promise<CommandResult> {
+  const { stdout } = await execFileAsync(executable, args, {
+    cwd,
+    encoding: "utf8",
+    timeout: 120_000,
+    maxBuffer: 10 * 1024 * 1024,
+  });
+  return { stdout };
+}
+
+export async function checkoutExistingPullRequest({
+  repositoryPath,
+  prUrl,
+  runCommand = defaultRunCommand,
+}: {
+  repositoryPath: string;
+  prUrl: string;
+  runCommand?: RunCommand;
+}): Promise<ExistingPrCheckoutResult> {
+  try {
+    const [currentBranchResult, prBranchResult] = await Promise.all([
+      runCommand("git", ["branch", "--show-current"], repositoryPath),
+      runCommand(
+        "gh",
+        ["pr", "view", prUrl, "--json", "headRefName", "--jq", ".headRefName"],
+        repositoryPath,
+      ),
+    ]);
+    const currentBranch = currentBranchResult.stdout.trim();
+    const prBranch = prBranchResult.stdout.trim();
+
+    if (!prBranch) {
+      return { status: "failed", error: "Pull request head branch is empty" };
+    }
+    if (currentBranch === prBranch) {
+      return { status: "already_active", branch: prBranch };
+    }
+
+    await runCommand("gh", ["pr", "checkout", prUrl], repositoryPath);
+    return { status: "checked_out", branch: prBranch };
+  } catch (error) {
+    return {
+      status: "failed",
+      error: error instanceof Error ? error.message : String(error),
+    };
+  }
+}