diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index 087f2b55b7..2e312e778b 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -44,6 +44,7 @@ import { SSE_KEEPALIVE_INTERVAL_MS, } from "./agent-server"; import { type JwtPayload, SANDBOX_CONNECTION_AUDIENCE } from "./jwt"; +import type { ExistingPrCheckoutResult } from "./pr-checkout"; const mockedClaudeSdk = vi.hoisted(() => { const createSuccessResult = () => ({ @@ -232,6 +233,13 @@ interface TestableServer { inboxReportUrl?: string | null, ): string; buildDetectedPrContext(prUrl: string): string; + buildExistingPrCheckoutPromise( + prUrl: string | null, + ): Promise | null; + logExistingPrCheckoutResult( + prUrl: string | null, + result: ExistingPrCheckoutResult, + ): void; buildSessionSystemPrompt( prUrl?: string | null, slackThreadUrl?: string | null, @@ -3871,6 +3879,78 @@ describe("AgentServer HTTP Mode", () => { delete process.env.POSTHOG_CODE_INTERACTION_ORIGIN; }); }); + + describe("buildExistingPrCheckoutPromise", () => { + const prUrl = "https://github.com/org/repo/pull/1"; + + afterEach(() => { + delete process.env.POSTHOG_CODE_INTERACTION_ORIGIN; + }); + + // Guards the gating condition: a review-first run (no auto-publish) must + // not silently check out a PR branch the prompt told the agent to leave + // alone. Regressing the guard to always-checkout would fail here. + it("does not check out when auto-publish is off", () => { + const s = createServer(); + const promise = ( + s as unknown as TestableServer + ).buildExistingPrCheckoutPromise(prUrl); + expect(promise).toBeNull(); + }); + + it("does not check out when there is no prUrl", () => { + process.env.POSTHOG_CODE_INTERACTION_ORIGIN = "slack"; + const s = createServer(); + const promise = ( + s as unknown as TestableServer + ).buildExistingPrCheckoutPromise(null); + expect(promise).toBeNull(); + }); + + it("does not check out when createPr is false, even on a Slack-origin run", () => { + process.env.POSTHOG_CODE_INTERACTION_ORIGIN = "slack"; + const s = createServer({ createPr: false }); + const promise = ( + s as unknown as TestableServer + ).buildExistingPrCheckoutPromise(prUrl); + expect(promise).toBeNull(); + }); + + it("does not check out when no repository is connected", () => { + process.env.POSTHOG_CODE_INTERACTION_ORIGIN = "slack"; + const s = createServer({ repositoryPath: undefined }); + const promise = ( + s as unknown as TestableServer + ).buildExistingPrCheckoutPromise(prUrl); + expect(promise).toBeNull(); + }); + + it("starts a checkout when auto-publish is on for a Slack-origin run", () => { + process.env.POSTHOG_CODE_INTERACTION_ORIGIN = "slack"; + const s = createServer(); + const promise = ( + s as unknown as TestableServer + ).buildExistingPrCheckoutPromise(prUrl); + expect(promise).toBeInstanceOf(Promise); + // Sanity: the promise resolves to a checkout result shape (it will fail + // against the synthetic URL with no real gh, which is fine — we only + // assert the promise was actually kicked off). + expect(typeof promise).toBe("object"); + }); + + // Guards the failure fallback: a transient gh failure must surface as a + // warn, never throw or abort startup. Regressing the failed branch to + // `throw` would fail here. + it("logs a warning for a failed checkout result without throwing", () => { + const s = createServer(); + expect(() => + (s as unknown as TestableServer).logExistingPrCheckoutResult(prUrl, { + status: "failed", + error: "gh unavailable", + }), + ).not.toThrow(); + }); + }); }); // Exercises getPendingUserPrompt directly (no HTTP server / git repo) so we can diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 208c643c4b..6f81afbc91 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -94,6 +94,10 @@ import { import { TaskRunEventStreamSender } from "./event-stream-sender"; import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt"; import { type McpRelayResponse, McpRelayServer } from "./mcp-relay-server"; +import { + checkoutExistingPullRequest, + type ExistingPrCheckoutResult, +} from "./pr-checkout"; import { resolveRtkSavings } from "./rtk-savings"; import { RunUsageAccumulator } from "./run-usage"; import { @@ -1505,28 +1509,55 @@ export class AgentServer { }; await this.waitForRepoReady(); - await this.installSkillBundleArtifacts( - payload.task_id, - payload.run_id, - this.getArtifactsById(preTaskRun?.artifacts, pendingUserArtifactIds), - ); - - const nativeResume = await this.prepareNativeResume( - payload, - posthogAPI, - preTaskRun, - runtimeAdapter, - sessionCwd, - initialPermissionMode, - ); + const existingPrCheckoutPromise = + this.buildExistingPrCheckoutPromise(prUrl); + // Overlap the best-effort PR checkout with the rest of session setup. The + // checkout promise is always awaited in `finally` so a throw from + // installSkillBundleArtifacts / prepareNativeResume / startMcpRelayServer + // can never abandon an in-flight `gh pr checkout` that would keep mutating + // the working tree after session start has been abandoned — the awaited + // settle (plus the checkout's own abort-on-return) cancels it. The overlap + // is safe despite both touching repositoryPath: skill bundles install under + // `.posthog/skills//...`, which is gitignored (untracked) in target + // repos, so `git checkout` — which only updates tracked files — cannot + // conflict with those writes or leave them associated with the wrong branch. + let nativeResume: { sessionId: string; warm: boolean } | null; let effectiveSessionMeta: typeof sessionMeta & { nativeGoal?: NonNullable; } = sessionMeta; + let sessionMcpServers: RemoteMcpServer[]; + try { + await this.installSkillBundleArtifacts( + payload.task_id, + payload.run_id, + this.getArtifactsById(preTaskRun?.artifacts, pendingUserArtifactIds), + ); - const sessionMcpServers = [ - ...(this.config.mcpServers ?? []), - ...(await this.startMcpRelayServer()), - ]; + nativeResume = await this.prepareNativeResume( + payload, + posthogAPI, + preTaskRun, + runtimeAdapter, + sessionCwd, + initialPermissionMode, + ); + + sessionMcpServers = [ + ...(this.config.mcpServers ?? []), + ...(await this.startMcpRelayServer()), + ]; + } finally { + // Always consume the checkout result — on the success path this is the + // intended await; on a throw it ensures the in-flight checkout settles + // (and aborts its children) instead of mutating the tree in the + // background. checkoutExistingPullRequest never rejects. + if (existingPrCheckoutPromise) { + this.logExistingPrCheckoutResult( + prUrl, + await existingPrCheckoutPromise, + ); + } + } let acpSessionId: string | null = null; if (nativeResume) { @@ -3143,6 +3174,54 @@ export class AgentServer { return `Continue working on the existing PR branch. If it is not already checked out, check it out with \`gh pr checkout ${prUrl}\`. Do not check it out again when it is already active.`; } + /** + * Fire-and-overlap: starts the best-effort PR-branch checkout so it runs + * concurrently with the rest of session setup, returning the promise (or + * null when there is nothing to check out). Only runs when auto-publishing, + * matching the system-prompt fallback's gate: a review-first run must not + * silently check out a branch the prompt told the agent to leave alone. + */ + private buildExistingPrCheckoutPromise( + prUrl: string | null, + ): Promise | null { + if (!prUrl || !this.config.repositoryPath) { + return null; + } + if (!this.shouldAutoPublishCloudChanges()) { + return null; + } + return checkoutExistingPullRequest({ + repositoryPath: this.config.repositoryPath, + prUrl, + }); + } + + /** + * Consume a pre-checkout result without throwing — a transient `gh` failure + * must fall back to the agent's own checkout (via the system-prompt + * instruction), never abort session start. + */ + private logExistingPrCheckoutResult( + prUrl: string | null, + result: ExistingPrCheckoutResult, + ): void { + if (result.status === "failed") { + this.logger.warn( + "Existing PR pre-checkout failed; agent will retry if needed", + { + prUrl, + error: result.error, + }, + ); + } else { + this.logger.debug("Existing PR branch prepared before session start", { + prUrl, + branch: result.branch, + alreadyActive: result.status === "already_active", + }); + } + } + private buildDetectedPrContext(prUrl: string): string { if (!this.shouldAutoPublishCloudChanges()) { return ( 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..73d3a725bd --- /dev/null +++ b/packages/agent/src/server/pr-checkout.test.ts @@ -0,0 +1,305 @@ +import { describe, expect, it, vi } from "vitest"; +import { checkoutExistingPullRequest } from "./pr-checkout"; + +const PR_URL = "https://github.com/PostHog/code/pull/1"; +const PR_BRANCH = "posthog-code/fix-checkout"; +const PR_HEAD_OID = "abc1234567890abcdef1234567890abcdef12345"; +// A SHA that differs from the PR head, used to force the "needs checkout" path. +const OTHER_OID = "0".repeat(40); + +type RunCommand = ( + executable: string, + args: string[], + cwd: string, + signal: AbortSignal, +) => Promise<{ stdout: string }>; + +/** + * Builds a runCommand mock that answers the four git/gh calls the checkout + * makes: `git branch --show-current`, `git rev-parse HEAD`, `gh pr view + * ... --json headRefName,headRefOid`, and `git remote get-url origin`. Records + * every call so assertions can pin exact args (including the prUrl and cwd) + * rather than just call counts. + */ +function buildRunCommand({ + currentBranch = PR_BRANCH, + currentHead = PR_HEAD_OID, + prBranch = PR_BRANCH, + prHeadOid = PR_HEAD_OID, + originUrl = "https://github.com/PostHog/code.git", + checkoutError, +}: { + currentBranch?: string; + currentHead?: string; + prBranch?: string; + prHeadOid?: string; + originUrl?: string; + checkoutError?: Error; +} = {}): { runCommand: RunCommand; calls: ReturnType } { + const calls = vi.fn(); + const runCommand: RunCommand = async (executable, args, cwd, _signal) => { + calls(executable, args, cwd); + if (executable === "git" && args[0] === "branch") { + return { stdout: `${currentBranch}\n` }; + } + if (executable === "git" && args[0] === "rev-parse") { + return { stdout: `${currentHead}\n` }; + } + if (executable === "git" && args[0] === "remote") { + return { stdout: `${originUrl}\n` }; + } + if (executable === "gh" && args[1] === "view") { + return { stdout: `${prBranch}\n${prHeadOid}\n` }; + } + if (executable === "gh" && args[1] === "checkout") { + if (checkoutError) { + throw checkoutError; + } + return { stdout: "" }; + } + return { stdout: "" }; + }; + return { runCommand, calls }; +} + +const ghCheckoutCall = (calls: ReturnType) => + calls.mock.calls.find( + ([executable, args]: string[]) => + executable === "gh" && args[0] === "pr" && args[1] === "checkout", + ); + +describe("checkoutExistingPullRequest", () => { + it.each([ + { + name: "skips checkout when attached to the PR branch at its head commit", + currentBranch: PR_BRANCH, + currentHead: PR_HEAD_OID, + expectedStatus: "already_active", + expectedCheckoutCalls: 0, + }, + { + name: "checks out when HEAD is on the PR branch but behind its head", + currentBranch: PR_BRANCH, + currentHead: OTHER_OID, + expectedStatus: "checked_out", + expectedCheckoutCalls: 1, + }, + { + name: "checks out when on another branch entirely", + currentBranch: "main", + currentHead: OTHER_OID, + expectedStatus: "checked_out", + expectedCheckoutCalls: 1, + }, + ])( + "$name", + async ({ + currentBranch, + currentHead, + expectedStatus, + expectedCheckoutCalls, + }) => { + const { runCommand, calls } = buildRunCommand({ + currentBranch, + currentHead, + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result.status).toBe(expectedStatus); + const checkoutCalls = calls.mock.calls.filter( + ([executable, args]: string[]) => + executable === "gh" && args[0] === "pr" && args[1] === "checkout", + ); + expect(checkoutCalls).toHaveLength(expectedCheckoutCalls); + + // Every git/gh call runs in the repository path — a regression that + // dropped (or changed) the cwd would not otherwise be caught, since the + // mock ignores it for dispatch. + for (const callCwd of calls.mock.calls.map( + ([, , cwd]: string[]) => cwd, + )) { + expect(callCwd).toBe("/tmp/repo"); + } + }, + ); + + it("passes the prUrl through to `gh pr checkout` (not just the call count)", async () => { + const { runCommand, calls } = buildRunCommand({ + currentBranch: "main", + currentHead: OTHER_OID, + }); + await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(ghCheckoutCall(calls)).toBeTruthy(); + // Pins the full args so a future refactor that drops the URL (or passes + // the wrong variable) fails instead of silently calling `gh pr checkout` + // with no argument. + expect(ghCheckoutCall(calls)?.[1]).toEqual(["pr", "checkout", PR_URL]); + }); + + it("checks out when the local branch shares the PR head's name but points at a different commit", async () => { + // Fork-PR edge case: a local branch named like the fork's head branch but + // tracking a different remote (different commit). Comparing only branch + // names would wrongly short-circuit as already_active. The SHA comparison + // must catch the mismatch and check out. + const { runCommand, calls } = buildRunCommand({ + currentBranch: PR_BRANCH, + currentHead: "deadbeef".repeat(5), + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result.status).toBe("checked_out"); + expect(ghCheckoutCall(calls)).toBeTruthy(); + }); + + it("checks out when HEAD is detached at the PR head commit (does not skip)", async () => { + // Detached HEAD at the PR head commit must NOT short-circuit as + // already_active: new commits would not attach to the PR branch. The + // branch-name check (empty on detached HEAD) forces a checkout. + const { runCommand, calls } = buildRunCommand({ + currentBranch: "", + currentHead: PR_HEAD_OID, + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result.status).toBe("checked_out"); + expect(ghCheckoutCall(calls)).toBeTruthy(); + }); + + it("fails when the PR belongs to a different repository than the workspace origin", async () => { + // `gh pr checkout ` treats a full URL as a repository override, so a + // foreign PR URL could pull an attacker-controlled branch in. The PR's + // owner/repo must match the workspace's origin. + const { runCommand, calls } = buildRunCommand({ + originUrl: "https://github.com/PostHog/posthog.git", + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result.status).toBe("failed"); + expect(result).toMatchObject({ + error: expect.stringContaining("not in the workspace repository"), + }); + expect(ghCheckoutCall(calls)).toBeUndefined(); + }); + + it("returns failed when the pull request head branch is empty", async () => { + const { runCommand } = buildRunCommand({ + prBranch: "", + prHeadOid: "", + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result).toEqual({ + status: "failed", + error: "Pull request head branch is empty", + }); + }); + + it("checks out when gh pr view omits the head OID (falls back from SHA match)", async () => { + // If headRefOid is unavailable, the SHA guard is skipped so checkout + // proceeds rather than assuming already_active. + const { runCommand, calls } = buildRunCommand({ + currentBranch: "main", + currentHead: OTHER_OID, + prHeadOid: "", + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result.status).toBe("checked_out"); + expect(ghCheckoutCall(calls)).toBeTruthy(); + }); + + it("returns a failure when gh pr view succeeds but gh pr checkout fails", async () => { + const { runCommand } = buildRunCommand({ + currentBranch: "main", + currentHead: OTHER_OID, + checkoutError: new Error("branch not found remotely"), + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result).toEqual({ + status: "failed", + error: "branch not found remotely", + }); + }); + + it("restores the original branch and HEAD after an interrupted checkout", async () => { + // A deadline-killed `gh pr checkout` can leave the working tree partially + // switched. The failure path must roll back to the pre-checkout state so + // the agent's fallback checkout starts clean. + const { runCommand, calls } = buildRunCommand({ + currentBranch: "main", + currentHead: OTHER_OID, + checkoutError: new Error("signal: aborted"), + }); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand, + }); + + expect(result.status).toBe("failed"); + const gitArgs = ( + calls.mock.calls as [string, string[], string, AbortSignal][] + ) + .filter(([executable]) => executable === "git") + .map(([, args]) => args.join(" ")); + expect(gitArgs).toContain("merge --abort"); + expect(gitArgs).toContain("checkout --force main"); + expect(gitArgs).toContain(`reset --hard ${OTHER_OID}`); + }); + + it("returns a failure so startup can fall back to agent checkout", async () => { + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: PR_URL, + runCommand: vi.fn().mockRejectedValue(new Error("gh unavailable")), + }); + + expect(result).toEqual({ status: "failed", error: "gh unavailable" }); + }); + + it("rejects a non-PR URL without invoking git or gh", async () => { + const runCommand = vi.fn(); + const result = await checkoutExistingPullRequest({ + repositoryPath: "/tmp/repo", + prUrl: "https://github.com/PostHog/code/blob/main/README.md", + runCommand, + }); + + expect(result.status).toBe("failed"); + expect(runCommand).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/agent/src/server/pr-checkout.ts b/packages/agent/src/server/pr-checkout.ts new file mode 100644 index 0000000000..044ad4d969 --- /dev/null +++ b/packages/agent/src/server/pr-checkout.ts @@ -0,0 +1,229 @@ +import { execFile } from "node:child_process"; +import { promisify } from "node:util"; +import { parseGithubUrl } from "@posthog/git/utils"; + +const execFileAsync = promisify(execFile); + +/** + * Overall budget for the pre-session PR checkout. This is a best-effort + * optimization: the agent still falls back to its own `gh pr checkout` via the + * system-prompt instruction when this pre-checkout fails or times out, so the + * deadline only bounds how long session start waits for it — it does not gate + * correctness. Per-command child timeouts (below) are larger so a single slow + * `gh` call is not cut short prematurely, but the total wall-clock spent here + * is capped so a hung `gh` (auth refresh, stalled fetch) can't stall startup. + */ +const PRE_CHECKOUT_DEADLINE_MS = 60_000; + +/** Per-child-process timeout; larger than the overall deadline on purpose. */ +const PER_COMMAND_TIMEOUT_MS = 120_000; +const PER_COMMAND_MAX_BUFFER = 10 * 1024 * 1024; + +type CommandResult = { stdout: string }; +type RunCommand = ( + executable: string, + args: string[], + cwd: string, + signal: AbortSignal, +) => 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, + signal: AbortSignal, +): Promise { + const { stdout } = await execFileAsync(executable, args, { + cwd, + encoding: "utf8", + timeout: PER_COMMAND_TIMEOUT_MS, + maxBuffer: PER_COMMAND_MAX_BUFFER, + // Killing the child on abort guarantees no in-flight checkout keeps + // mutating the working tree once we've stopped waiting for it. + signal, + }); + return { stdout }; +} + +/** + * The pull request the workspace should check out, parsed from prUrl. prUrl + * originates from the agent's own prior `task_run.output.pr_url`, which the + * codebase documents as user-writable, so we never hand an arbitrary value + * straight to `gh`. A non-PR value yields null and the agent's lazy checkout + * takes over. + */ +function parsePrUrl(prUrl: string): { + owner: string; + repo: string; + number: number; +} | null { + const parsed = parseGithubUrl(prUrl); + if (parsed?.kind !== "pr") { + return null; + } + return { owner: parsed.owner, repo: parsed.repo, number: parsed.number }; +} + +export async function checkoutExistingPullRequest({ + repositoryPath, + prUrl, + runCommand = defaultRunCommand, + deadlineMs = PRE_CHECKOUT_DEADLINE_MS, +}: { + repositoryPath: string; + prUrl: string; + runCommand?: RunCommand; + deadlineMs?: number; +}): Promise { + const pr = parsePrUrl(prUrl); + if (!pr) { + return { + status: "failed", + error: `Not a recognized pull request URL: ${prUrl}`, + }; + } + + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), deadlineMs); + + const withAbort = ( + executable: string, + args: string[], + cwd: string, + ): Promise => + runCommand(executable, args, cwd, controller.signal); + + try { + const [currentBranchResult, currentHeadResult, prInfoResult, originResult] = + await Promise.all([ + withAbort("git", ["branch", "--show-current"], repositoryPath), + withAbort("git", ["rev-parse", "HEAD"], repositoryPath), + withAbort( + "gh", + [ + "pr", + "view", + prUrl, + "--json", + "headRefName,headRefOid", + "--jq", + '.headRefName + "\\n" + .headRefOid', + ], + repositoryPath, + ), + withAbort("git", ["remote", "get-url", "origin"], repositoryPath), + ]); + const [prBranch, prHeadOid] = prInfoResult.stdout + .split("\n") + .map((line) => line.trim()); + const currentBranch = currentBranchResult.stdout.trim(); + const currentHead = currentHeadResult.stdout.trim(); + + if (!prBranch) { + return { status: "failed", error: "Pull request head branch is empty" }; + } + + // Reject PRs whose repository does not match the workspace's origin. `gh pr + // checkout ` treats a full URL as a repository override and fetches + // that PR's ref into the workspace, so a foreign PR URL (the prUrl field is + // user-writable) could pull an attacker-controlled branch in before the + // agent starts. Require the PR's owner/repo to match the connected repo. + const origin = parseGithubUrl(originResult.stdout.trim()); + if ( + !origin || + origin.owner.toLowerCase() !== pr.owner.toLowerCase() || + origin.repo.toLowerCase() !== pr.repo.toLowerCase() + ) { + return { + status: "failed", + error: `Pull request ${prUrl} is not in the workspace repository`, + }; + } + + // Only skip the checkout when HEAD is attached to the PR's branch AND at + // its head commit. Matching the SHA alone is not enough: a detached HEAD + // at the same commit would skip checkout, and new commits would then not + // attach to the PR branch. Matching the branch name alone is not enough + // either (a fork PR can share a name with a different-local-remote branch). + if (prHeadOid && currentBranch === prBranch && currentHead === prHeadOid) { + return { status: "already_active", branch: prBranch }; + } + + // originalHead/originalBranch capture the state before checkout so an + // interrupted `gh pr checkout` (e.g. the deadline aborting it mid-switch) + // can be rolled back instead of leaving a partially-switched working tree + // for the agent's fallback checkout to trip over. + const originalHead = currentHead; + const originalBranch = currentBranch; + try { + await withAbort("gh", ["pr", "checkout", prUrl], repositoryPath); + } catch (error) { + await restoreRepoState( + runCommand, + repositoryPath, + originalBranch, + originalHead, + ); + return { + status: "failed", + error: error instanceof Error ? error.message : String(error), + }; + } + return { status: "checked_out", branch: prBranch }; + } catch (error) { + return { + status: "failed", + error: error instanceof Error ? error.message : String(error), + }; + } finally { + clearTimeout(timeout); + // Ensure no child survives the deadline if we exit early (e.g. the second + // command failed before the first timed out). + controller.abort(); + } +} + +/** + * Best-effort restore of the working tree after an interrupted `gh pr + * checkout`. `gh pr checkout` may run `git fetch` + `git checkout`; if it is + * killed mid-checkout the index and working tree can be left partially + * switched. We abort any in-progress merge, return to the original branch (or + * commit, if HEAD was detached), and hard-reset tracked files to the original + * commit so the agent's fallback checkout starts from a clean, known state. + * Runs outside the checkout's deadline (per-command timeout still applies) and + * swallows its own errors so a failed restore never masks the original failure. + */ +async function restoreRepoState( + runCommand: RunCommand, + repositoryPath: string, + originalBranch: string, + originalHead: string, +): Promise { + // Fresh signal — the checkout's controller has already aborted. + const signal = new AbortController().signal; + const restore = (args: string[]): Promise => + runCommand("git", args, repositoryPath, signal); + try { + // Clear any merge state `gh pr checkout` may have left behind. + await restore(["merge", "--abort"]).catch(() => {}); + // Detached HEAD pre-checkout: return to the original commit. Otherwise + // return to the original branch (re-attaching HEAD to it). + await restore( + originalBranch + ? ["checkout", "--force", originalBranch] + : ["checkout", "--force", originalHead], + ).catch(() => {}); + // Discard partial index/working-tree changes from the interrupted + // checkout, restoring tracked files to the original commit. Untracked + // files (e.g. agent skill bundles under .posthog/) are left alone. + await restore(["reset", "--hard", originalHead]).catch(() => {}); + } catch { + // Best-effort: a restore failure must not mask the original checkout + // failure. The fallback `gh pr checkout` will surface any lingering state. + } +}