Files
shockbot/mcp/checkout.ts
T

522 lines
20 KiB
TypeScript

import { createHash } from "node:crypto";
import { statSync, unlinkSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import { type } from "arktype";
import { log } from "../utils/cli.ts";
import { countLines, createDiffCoverageState } from "../utils/diffCoverage.ts";
import { $git, $gitFetchWithDeepen } from "../utils/gitAuth.ts";
import type { Gitea } from "../utils/gitea.ts";
import type { ChangedFileWithPatch } from "../utils/gitea.ts";
import { executeLifecycleHook } from "../utils/lifecycle.ts";
import { computeIncrementalDiff } from "../utils/rangeDiff.ts";
import { retry } from "../utils/retry.ts";
import { $ } from "../utils/shell.ts";
import { rejectIfLeadingDash } from "./git.ts";
import { commentableLinesForFile } from "./review.ts";
import type { ToolContext } from "./server.ts";
import { execute, tool } from "./shared.ts";
export type FormatFilesResult = { content: string; toc: string };
export type FetchAndFormatPrDiffResult = FormatFilesResult & {
files: DiffFile[];
};
export type DiffFile = {
filename?: string | undefined;
patch?: string | undefined;
};
export function formatFilesWithLineNumbers(files: DiffFile[]): FormatFilesResult {
const output: string[] = [];
const tocEntries: Array<{ filename: string; startLine: number; endLine: number }> = [];
const tocHeaderSize = 1 + files.length + 2;
let currentLine = tocHeaderSize + 1;
for (const file of files) {
const filename = file.filename ?? "(unknown)";
const fileStartLine = currentLine;
output.push(`diff --git a/${filename} b/${filename}`);
output.push(`--- a/${filename}`);
output.push(`+++ b/${filename}`);
currentLine += 3;
if (!file.patch) {
output.push("(binary file or no changes)");
output.push("");
currentLine += 2;
tocEntries.push({ filename, startLine: fileStartLine, endLine: currentLine - 1 });
continue;
}
const lines = file.patch.split("\n");
let oldLine = 0;
let newLine = 0;
for (const line of lines) {
const hunkMatch = line.match(/^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@/);
if (hunkMatch) {
oldLine = parseInt(hunkMatch[1], 10);
newLine = parseInt(hunkMatch[2], 10);
output.push(line);
currentLine++;
continue;
}
const changeType = line[0] || " ";
const code = line.slice(1);
if (changeType === "-") {
output.push(`| ${padNum(oldLine)} | | - | ${code}`);
oldLine++;
} else if (changeType === "+") {
output.push(`| | ${padNum(newLine)} | + | ${code}`);
newLine++;
} else if (changeType === " " || changeType === "\\") {
if (changeType === "\\") {
output.push(line);
} else {
output.push(`| ${padNum(oldLine)} | ${padNum(newLine)} | | ${code}`);
oldLine++;
newLine++;
}
} else {
output.push(line);
}
currentLine++;
}
output.push("");
currentLine++;
tocEntries.push({ filename, startLine: fileStartLine, endLine: currentLine - 1 });
}
const tocLines = [`## Files (${files.length})`];
for (const entry of tocEntries) {
const anchor = createHash("sha256").update(entry.filename).digest("hex");
tocLines.push(`- ${entry.filename} → lines ${entry.startLine}-${entry.endLine} · diff-${anchor}`);
}
tocLines.push("", "---", "");
const toc = tocLines.join("\n");
return { content: toc + output.join("\n"), toc };
}
function padNum(n: number): string {
return n.toString().padStart(4, " ");
}
export const CheckoutPr = type({
pull_number: type.number.describe("the pull request number to checkout"),
});
export type CheckoutPrResult = {
success: true;
number: number;
title: string;
body: string | null;
base: string;
localBranch: string;
remoteBranch: string;
isFork: boolean;
maintainerCanModify: boolean;
url: string;
headRepo: string;
diffPath: string;
incrementalDiffPath?: string | undefined;
toc: string;
commitCount: number;
commitLog: string;
commitLogTruncated: boolean;
commitLogUnavailable: boolean;
hookWarning?: string | undefined;
instructions: string;
};
export async function fetchAndFormatPrDiff(
ctx: ToolContext,
pullNumber: number
): Promise<FetchAndFormatPrDiffResult> {
const raw = await ctx.gitea.paginate(ctx.gitea.rest.repository.repoGetPullRequestFiles, {
owner: ctx.repo.owner,
repo: ctx.repo.name,
index: pullNumber,
});
// ChangedFile from the SDK omits `patch`; Gitea does return it, so we map here.
const files: DiffFile[] = raw.map((f) => ({
filename: f.filename,
patch: (f as unknown as { patch?: string }).patch,
}));
return { ...formatFilesWithLineNumbers(files), files };
}
import { captureInitialHead } from "../utils/setup.ts";
export type PrData = {
number: number;
headSha: string;
headRef: string;
headRepoFullName: string;
baseRef: string;
baseRepoFullName: string;
maintainerCanModify: boolean;
};
const STALE_LOCK_AGE_MS = 30_000;
const PULL_REF_RETRY_DELAYS_MS = [2_000, 5_000, 10_000];
const PULL_REF_MISSING_PATTERN = /couldn't find remote ref pull\/\d+\/head/i;
const GIT_LOCK_PATHS = [".git/shallow.lock", ".git/index.lock", ".git/objects/maintenance.lock"] as const;
function cleanupStaleGitLocks(): void {
const now = Date.now();
for (const relPath of GIT_LOCK_PATHS) {
let mtimeMs: number;
try { mtimeMs = statSync(relPath).mtimeMs; } catch { continue; }
if (now - mtimeMs < STALE_LOCK_AGE_MS) continue;
try { unlinkSync(relPath); log.warning(`» removed stale ${relPath}`); } catch {}
}
}
type CheckoutPrBranchParams = {
gitToken: string;
owner: string;
name: string;
gitea: Gitea;
toolState: import("../toolState.ts").ToolState;
shell: import("../external.ts").ShellPermission;
postCheckoutScript: string | null;
beforeSha?: string | undefined;
};
async function abortIfPullRequestMoved(args: { gitea: Gitea; owner: string; repo: string; pr: PrData }): Promise<void> {
try {
const data = (await args.gitea.rest.repository.repoGetPullRequest({ owner: args.owner, repo: args.repo, index: args.pr.number })).data;
if (data.state !== "open" || data.head?.sha !== args.pr.headSha) {
throw new Error(`PR #${args.pr.number} is no longer in the state it was at dispatch. Aborting.`);
}
} catch (e) {
if (e instanceof Error && e.message.includes("no longer")) throw e;
// API error — lenient, don't abort
}
}
type CreateTempBranchParams = { gitea: Gitea; owner: string; repo: string; branchName: string; sha: string };
async function createTempBranch(params: CreateTempBranchParams) {
await params.gitea.rest.repository.repoCreateBranch({
owner: params.owner,
repo: params.repo,
body: { new_branch_name: params.branchName, old_ref_name: params.sha },
});
return {
async [Symbol.asyncDispose]() {
try {
await params.gitea.rest.repository.repoDeleteBranch({ owner: params.owner, repo: params.repo, branch: params.branchName });
log.debug(`» deleted temp branch ${params.branchName}`);
} catch (e) {
log.debug(`» failed to delete temp branch ${params.branchName}: ${e instanceof Error ? e.message : String(e)}`);
}
},
};
}
async function ensureBeforeShaReachable(params: {
sha: string; gitea: Gitea; owner: string; repo: string; gitToken: string; isShallow: boolean;
}): Promise<boolean> {
try { $("git", ["cat-file", "-t", params.sha], { log: false }); return true; } catch {}
const branchName = `shockbot/tmp/${params.sha.slice(0, 12)}`;
try {
await using _ref = await createTempBranch({ gitea: params.gitea, owner: params.owner, repo: params.repo, branchName, sha: params.sha });
await $gitFetchWithDeepen(
["--no-tags", ...(params.isShallow ? ["--depth=1"] : []), "origin", branchName],
{ token: params.gitToken },
`before_sha temp branch ${branchName}`
);
return true;
} catch (e) {
log.debug(`» failed to fetch before_sha: ${e instanceof Error ? e.message : String(e)}`);
return false;
}
}
export async function checkoutPrBranch(
pr: PrData,
params: CheckoutPrBranchParams
): Promise<{ hookWarning?: string | undefined }> {
const { gitea, owner, name, gitToken, toolState, beforeSha } = params;
log.info(`» checking out PR #${pr.number}...`);
rejectIfLeadingDash(pr.baseRef, "PR base ref");
rejectIfLeadingDash(pr.headRef, "PR head ref");
cleanupStaleGitLocks();
const isFork = pr.headRepoFullName !== pr.baseRepoFullName;
const localBranch = `pr-${pr.number}`;
const isShallow = $("git", ["rev-parse", "--is-shallow-repository"], { log: false }).trim() === "true";
toolState.checkoutSha = $("git", ["rev-parse", "HEAD"], { log: false }).trim();
const alreadyOnBranch = toolState.checkoutSha === pr.headSha;
log.debug(`» fetching base branch (${pr.baseRef})...`);
await $gitFetchWithDeepen(["--no-tags", "origin", pr.baseRef], { token: gitToken }, `base branch ${pr.baseRef}`);
if (!alreadyOnBranch) {
$("git", ["checkout", "-B", pr.baseRef, `origin/${pr.baseRef}`], { log: false });
log.debug(`» fetching PR #${pr.number} (${localBranch})...`);
await retry(
async () => {
try {
await $gitFetchWithDeepen(
["--no-tags", "origin", `+pull/${pr.number}/head:${localBranch}`],
{ token: gitToken },
`PR #${pr.number}`
);
} catch (e) {
const msg = e instanceof Error ? e.message : String(e);
if (PULL_REF_MISSING_PATTERN.test(msg)) {
await abortIfPullRequestMoved({ gitea, owner, repo: name, pr });
}
throw e;
}
},
{
delaysMs: PULL_REF_RETRY_DELAYS_MS,
label: `pull/${pr.number}/head fetch`,
shouldRetry: (e) => PULL_REF_MISSING_PATTERN.test(e instanceof Error ? e.message : String(e)),
}
);
$("git", ["checkout", localBranch], { log: false });
toolState.checkoutSha = $("git", ["rev-parse", "HEAD"], { log: false }).trim();
}
const beforeShaReachable = beforeSha
? await ensureBeforeShaReachable({ sha: beforeSha, gitea, owner, repo: name, gitToken, isShallow })
: false;
if (isShallow) {
let deepenDepth = 0;
try {
const [prComp, beforeComp] = await Promise.all([
gitea.rest.repository.repoCompareDiff({ owner, repo: name, basehead: `${pr.baseRef}...${toolState.checkoutSha}` }),
beforeSha && beforeShaReachable
? gitea.rest.repository.repoCompareDiff({ owner, repo: name, basehead: `${pr.baseRef}...${beforeSha}` })
: undefined,
]);
const prTotal = prComp.data.total_commits ?? 0;
const beforeTotal = beforeComp?.data.total_commits ?? 0;
deepenDepth = Math.max(prTotal, beforeTotal) + 10;
log.debug(`» compare: PR=${prTotal}, before=${beforeTotal}, deepen=${deepenDepth}`);
} catch {
deepenDepth = 1000;
log.debug(`» compare API failed, falling back to --deepen=${deepenDepth}`);
}
if (deepenDepth) {
await $git("fetch", [`--deepen=${deepenDepth}`, "--no-tags", "origin"], { token: gitToken });
}
}
if (isFork) {
const remoteName = `pr-${pr.number}`;
const giteaUrl = (process.env.GITEA_URL ?? "https://git.shockvpn.com").replace(/\/$/, "");
const forkUrl = `${giteaUrl}/${pr.headRepoFullName}.git`;
try { $("git", ["remote", "add", remoteName, forkUrl], { log: false }); }
catch { $("git", ["remote", "set-url", remoteName, forkUrl], { log: false }); }
$("git", ["config", `branch.${localBranch}.pushRemote`, remoteName], { log: false });
$("git", ["config", `branch.${localBranch}.merge`, `refs/heads/${pr.headRef}`], { log: false });
if (!pr.maintainerCanModify) log.warning(`» fork PR has maintainer_can_modify=false — push will likely fail.`);
toolState.pushUrl = forkUrl;
} else {
$("git", ["config", `branch.${localBranch}.pushRemote`, "origin"], { log: false });
$("git", ["config", `branch.${localBranch}.merge`, `refs/heads/${pr.headRef}`], { log: false });
}
toolState.issueNumber = pr.number;
toolState.pushDest = { remoteName: isFork ? `pr-${pr.number}` : "origin", remoteBranch: pr.headRef, localBranch };
const postCheckoutHook = await executeLifecycleHook({
event: "post-checkout",
script: params.postCheckoutScript,
normalizeWorkingTreeAfter: true,
});
return { hookWarning: postCheckoutHook.warning };
}
const inFlightCheckouts = new Map<number, Promise<CheckoutPrResult>>();
type InitialHead = NonNullable<ToolContext["toolState"]["initialHead"]>;
function headsEqual(a: InitialHead, b: InitialHead): boolean {
if (a.kind === "branch" && b.kind === "branch") return a.name === b.name;
if (a.kind === "detached" && b.kind === "detached") return a.sha === b.sha;
return false;
}
function describeHead(h: InitialHead): string {
return h.kind === "branch" ? `branch \`${h.name}\`` : `detached HEAD \`${h.sha}\``;
}
export function CheckoutPrTool(ctx: ToolContext) {
const runCheckout = async (pull_number: number): Promise<CheckoutPrResult> => {
const prResult = await ctx.gitea.rest.repository.repoGetPullRequest({
owner: ctx.repo.owner,
repo: ctx.repo.name,
index: pull_number,
});
const prData = prResult.data;
const headRepo = prData.head?.repo;
if (!headRepo) throw new Error(`PR #${pull_number} source repository was deleted`);
const pr: PrData = {
number: pull_number,
headSha: prData.head?.sha ?? "",
headRef: prData.head?.ref ?? "",
headRepoFullName: headRepo.full_name ?? "",
baseRef: prData.base?.ref ?? "",
baseRepoFullName: prData.base?.repo?.full_name ?? "",
maintainerCanModify: prData.allow_maintainer_edit ?? false,
};
const checkoutResult = await checkoutPrBranch(pr, {
gitea: ctx.gitea,
owner: ctx.repo.owner,
name: ctx.repo.name,
gitToken: ctx.gitToken,
toolState: ctx.toolState,
shell: ctx.payload.shell,
postCheckoutScript: ctx.postCheckoutScript,
beforeSha: ctx.toolState.beforeSha,
});
const tempDir = process.env.SHOCKBOT_TEMP_DIR;
if (!tempDir) throw new Error("SHOCKBOT_TEMP_DIR not set");
const headShort = ctx.toolState.checkoutSha!.slice(0, 7);
let incrementalDiffPath: string | undefined;
if (ctx.toolState.beforeSha && ctx.toolState.checkoutSha) {
const beforeShort = ctx.toolState.beforeSha.slice(0, 7);
const incremental = computeIncrementalDiff({
baseBranch: pr.baseRef,
beforeSha: ctx.toolState.beforeSha,
headSha: ctx.toolState.checkoutSha,
});
if (incremental) {
incrementalDiffPath = join(tempDir, `pr-${pull_number}-${beforeShort}-${headShort}-incremental.diff`);
writeFileSync(incrementalDiffPath, incremental);
log.info(`» incremental diff computed → ${incrementalDiffPath}`);
}
}
const formatResult = await fetchAndFormatPrDiff(ctx, pull_number);
const diffPath = join(tempDir, `pr-${pull_number}-${headShort}.diff`);
writeFileSync(diffPath, formatResult.content);
log.debug(`wrote diff to ${diffPath} (${formatResult.content.length} bytes)`);
ctx.toolState.diffCoverage = createDiffCoverageState({
diffPath,
totalLines: countLines({ content: formatResult.content }),
toc: formatResult.toc,
previous: ctx.toolState.diffCoverage,
});
const cached = new Map<string, ReturnType<typeof commentableLinesForFile>>();
for (const file of formatResult.files) {
if (file.filename) cached.set(file.filename, commentableLinesForFile(file.patch));
}
ctx.toolState.commentableLinesByFile = cached;
ctx.toolState.commentableLinesPullNumber = pull_number;
ctx.toolState.commentableLinesCheckoutSha = ctx.toolState.checkoutSha;
const COMMIT_LOG_MAX = 200;
const baseRange = `origin/${pr.baseRef}..HEAD`;
let commitCount = 0;
let commitLog = "";
let commitLogUnavailable = false;
try {
commitCount = parseInt($("git", ["rev-list", "--count", baseRange], { log: false }).trim() || "0", 10);
commitLog = $("git", ["log", "--oneline", `--max-count=${COMMIT_LOG_MAX}`, baseRange], { log: false });
} catch {
commitLogUnavailable = true;
}
return {
success: true,
number: prData.number!,
title: prData.title ?? "",
body: prData.body ?? null,
base: pr.baseRef,
localBranch: `pr-${pull_number}`,
remoteBranch: `refs/heads/${pr.headRef}`,
isFork: pr.headRepoFullName !== pr.baseRepoFullName,
maintainerCanModify: pr.maintainerCanModify,
url: prData.html_url ?? "",
headRepo: pr.headRepoFullName,
diffPath,
incrementalDiffPath,
toc: formatResult.toc,
commitCount,
commitLog,
commitLogTruncated: commitCount > COMMIT_LOG_MAX,
commitLogUnavailable,
hookWarning: checkoutResult.hookWarning,
instructions:
`the diff file at diffPath contains a table of contents (TOC) listing every changed file with its line range. ` +
`use the TOC line ranges as your checklist and read specific files from the diff. ` +
`for example, if the TOC says "src/foo.ts → lines 5-42", read lines 5-42 from diffPath. ` +
`review files selectively based on relevance. ` +
`to inspect the PR's changed files, use diffPath — do NOT run git diff commands. ` +
(incrementalDiffPath
? ` IMPORTANT: read incrementalDiffPath FIRST to understand what changed since last review, then use diffPath for full context.`
: "") +
(checkoutResult.hookWarning ? ` HOOK WARNING: the post-checkout hook reported a non-fatal failure.` : "") +
(commitLogUnavailable ? ` NOTE: commit metadata is partial (shallow fetch).` : commitCount > COMMIT_LOG_MAX ? ` NOTE: commitLog was capped at ${COMMIT_LOG_MAX} entries.` : ""),
} satisfies CheckoutPrResult;
};
return tool({
name: "checkout_pr",
timeoutMs: 600_000,
description:
"Checkout a pull request branch locally. Returns diffPath pointing to the formatted diff file. " +
"Example: `checkout_pr({ pull_number: 1234 })`. Large repos can take several minutes.",
parameters: CheckoutPr,
execute: execute(async ({ pull_number }) => {
const inFlight = inFlightCheckouts.get(pull_number);
if (inFlight) {
log.info(`» checkout_pr({pull_number:${pull_number}}) already in flight — sharing result`);
return inFlight;
}
const dirty = $("git", ["status", "--porcelain"], { log: false }).trim();
if (dirty) {
throw new Error(
`cannot checkout PR #${pull_number} while the working tree has uncommitted changes. ` +
`commit or discard with \`git restore --staged --worktree .\` / \`git clean -fd\` before retrying.\n${dirty}`
);
}
const initialHead = ctx.toolState.initialHead;
if (initialHead) {
const currentHead = captureInitialHead(process.cwd());
const targetBranch = `pr-${pull_number}`;
const onTarget = currentHead.kind === "branch" && currentHead.name === targetBranch;
const onInitial = headsEqual(currentHead, initialHead);
if (!onTarget && !onInitial) {
const recoverCmd = initialHead.kind === "branch" ? `git checkout ${initialHead.name}` : `git checkout ${initialHead.sha}`;
throw new Error(
`cannot checkout PR #${pull_number} from ${describeHead(currentHead)}. ` +
`recover with \`${recoverCmd}\` first.`
);
}
}
const promise = runCheckout(pull_number);
inFlightCheckouts.set(pull_number, promise);
try { return await promise; }
finally { inFlightCheckouts.delete(pull_number); }
}),
});
}