diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 569176573c..6ffd0f7420 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/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 { diff --git a/packages/agent/src/server/pr-checkout.test.ts b/packages/agent/src/server/pr-checkout.test.ts new file mode 100644 index 0000000000..1a0e2435aa --- /dev/null +++ b/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" }); + }); +}); diff --git a/packages/agent/src/server/pr-checkout.ts b/packages/agent/src/server/pr-checkout.ts new file mode 100644 index 0000000000..23dea0b96f --- /dev/null +++ b/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; + +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 { + 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 { + 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), + }; + } +}