From c28805598a4f62fb513c0b9b79dead4d2da41a5d Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Tue, 21 Jul 2026 08:59:13 +0100 Subject: [PATCH] perf(agent): overlap existing PR checkout with startup Resolve and prepare an existing pull request branch while the remaining cloud session setup continues. Skip redundant checkout when the target branch is already active and preserve the agent fallback if GitHub CLI preparation fails. Generated-By: PostHog Code Task-Id: 51deba7c-20d3-46e2-b2b5-ca35fb1670d1 --- packages/agent/src/server/agent-server.ts | 29 ++++++++ packages/agent/src/server/pr-checkout.test.ts | 61 +++++++++++++++++ packages/agent/src/server/pr-checkout.ts | 68 +++++++++++++++++++ 3 files changed, 158 insertions(+) create mode 100644 packages/agent/src/server/pr-checkout.test.ts create mode 100644 packages/agent/src/server/pr-checkout.ts diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 208c643c4b..7951fe2681 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -94,6 +94,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 { @@ -1505,6 +1506,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, @@ -1528,6 +1538,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), + }; + } +}