Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
19 kB · 639 lines
JavaScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640import { existsSync, mkdirSync, mkdtempSync, rmSync } from "node:fs";import { spawn } from "node:child_process";import { tmpdir } from "node:os";import { dirname, join, relative, resolve } from "node:path";import { performance } from "node:perf_hooks";import { fileURLToPath } from "node:url";
const __dirname = dirname(fileURLToPath(import.meta.url));const repoRoot = resolve(__dirname, "..");const packageDir = join(repoRoot, "packages", "coding-agent");const distCliPath = join(packageDir, "dist", "cli.js");const srcCliPath = join(packageDir, "src", "cli.ts");const defaultNodeProfileDir = join(repoRoot, "profiles-node");const defaultBunProfileDir = join(repoRoot, "profiles-bun");const agentDirEnvName = "PI_CODING_AGENT_DIR";const startupBenchmarkEnvName = "PI_STARTUP_BENCHMARK";
function printHelp() { console.log(`Usage: node scripts/profile-coding-agent-node.mjs [options]
Profiles coding-agent startup with the runtime selected below:- npm run profile:tui -> builds packages/coding-agent and profiles TUI startup with Node- npm run profile:rpc -> builds packages/coding-agent and profiles RPC startup with Node- bun run profile:tui -> profiles TUI startup from src/cli.ts directly with Bun- bun run profile:rpc -> profiles RPC startup from src/cli.ts directly with Bun
Options: --mode <name> tui or rpc (default: tui) --runs <n> Number of measured runs (default: 1) --warmup <n> Number of warmup runs before measurements (default: 0) --profile-dir <dir> CPU profile output directory Default: profiles-node for Node, profiles-bun for Bun --label <name> Profile name prefix (default: <mode>-startup) --runtime <name> node, bun, or auto (default: auto) --agent-dir <dir> Use a specific PI_CODING_AGENT_DIR for the benchmark run --isolated-agent-dir Use a fresh temporary agent dir instead of the normal one --no-offline Do not force PI_OFFLINE=1 / PI_SKIP_VERSION_CHECK=1 --skip-build Reuse the current dist/cli.js without rebuilding first (Node only) --cpu-profile Write CPU profiles for benchmark runs --help Show this help
Notes: - By default the benchmark uses your normal configured agent dir, so global models/auth/settings work. - TUI mode measures startup until the interactive UI reaches first usable state. - RPC mode measures startup until a real get_state request receives a response, then closes stdin to exit cleanly. - CPU profiles are kept in the selected profile directory for later analysis.`);}
function parseIntegerFlag(value, name) { const parsed = Number.parseInt(value, 10); if (!Number.isFinite(parsed) || parsed < 0) { throw new Error(`Invalid ${name}: ${value}`); } return parsed;}
function parseRuntime(value) { if (value === "auto" || value === "node" || value === "bun") { return value; } throw new Error(`Invalid --runtime: ${value}`);}
function parseMode(value) { if (value === "tui" || value === "rpc") { return value; } throw new Error(`Invalid --mode: ${value}`);}
function parseArgs(argv) { const options = { mode: "tui", runs: 1, warmup: 0, profileDir: undefined, label: undefined, offline: true, build: true, runtime: "auto", agentDir: undefined, isolatedAgentDir: false, cpuProfile: false, };
for (let index = 0; index < argv.length; index++) { const arg = argv[index];
if (arg === "--help" || arg === "-h") { options.help = true; continue; }
if (arg === "--no-offline") { options.offline = false; continue; }
if (arg === "--isolated-agent-dir") { options.isolatedAgentDir = true; continue; }
if (arg === "--skip-build") { options.build = false; continue; }
if (arg === "--cpu-profile") { options.cpuProfile = true; continue; }
if ( (arg === "--mode" || arg === "--runs" || arg === "--warmup" || arg === "--profile-dir" || arg === "--label" || arg === "--runtime" || arg === "--agent-dir") && index + 1 >= argv.length ) { throw new Error(`Missing value for ${arg}`); }
if (arg === "--mode") { options.mode = parseMode(argv[++index]); continue; }
if (arg === "--runs") { options.runs = parseIntegerFlag(argv[++index], "--runs"); continue; }
if (arg === "--warmup") { options.warmup = parseIntegerFlag(argv[++index], "--warmup"); continue; }
if (arg === "--profile-dir") { options.profileDir = resolve(argv[++index]); continue; }
if (arg === "--label") { options.label = argv[++index]; continue; }
if (arg === "--runtime") { options.runtime = parseRuntime(argv[++index]); continue; }
if (arg === "--agent-dir") { options.agentDir = resolve(argv[++index]); continue; }
throw new Error(`Unknown option: ${arg}`); }
return options;}
function detectRuntimeFromPackageManager() { const userAgent = process.env.npm_config_user_agent ?? ""; return userAgent.startsWith("bun/") ? "bun" : "node";}
function resolveRuntime(requestedRuntime) { if (requestedRuntime === "auto") { return detectRuntimeFromPackageManager(); } return requestedRuntime;}
function resolveProfileDir(runtime, requestedProfileDir) { if (requestedProfileDir) { return requestedProfileDir; } return runtime === "bun" ? defaultBunProfileDir : defaultNodeProfileDir;}
function resolveLabel(mode, requestedLabel) { return requestedLabel ?? `${mode}-startup`;}
function formatMs(value) { return `${value.toFixed(1)}ms`;}
function toDisplayPath(path) { const relativePath = relative(repoRoot, path); if (relativePath !== "" && !relativePath.startsWith("..")) { return relativePath.replaceAll("\\", "/"); } return path;}
function summarize(values) { const sorted = [...values].sort((a, b) => a - b); const total = sorted.reduce((sum, value) => sum + value, 0); const middle = Math.floor(sorted.length / 2); const median = sorted.length % 2 === 0 ? (sorted[middle - 1] + sorted[middle]) / 2 : sorted[middle]; return { min: sorted[0], max: sorted[sorted.length - 1], avg: total / sorted.length, median, };}
function parseStartupTimings(stderr) { const lines = stderr.split(/\r?\n/); const timings = new Map(); let inBlock = false;
for (const line of lines) { if (line.includes("--- Startup Timings ---")) { inBlock = true; continue; } if (!inBlock) { continue; } if (line.includes("------------------------")) { break; } const match = line.match(/^\s+([^:]+):\s+(\d+)ms$/); if (!match) { continue; } timings.set(match[1], Number.parseInt(match[2], 10)); }
return timings;}
function summarizeTimingMaps(runs) { const valuesByLabel = new Map(); for (const run of runs) { for (const [label, value] of run.timings.entries()) { const values = valuesByLabel.get(label); if (values) { values.push(value); } else { valuesByLabel.set(label, [value]); } } }
const summaries = new Map(); for (const [label, values] of valuesByLabel.entries()) { summaries.set(label, summarize(values)); } return summaries;}
function toMetricName(label) { return `${label.replaceAll(/[^a-zA-Z0-9]+/g, "_").replaceAll(/^_+|_+$/g, "")}_ms`;}
async function waitForExit(child, errorPrefix) { return await new Promise((resolve, reject) => { child.once("error", reject); child.once("exit", (code, signal) => { if (signal) { reject(new Error(`${errorPrefix} exited from signal ${signal}`)); return; } resolve(code ?? 0); }); });}
async function runBuild() { process.stdout.write("Building packages/tui, packages/ai, packages/agent, and packages/coding-agent...\n"); const startedAt = performance.now(); const child = spawn( "npm", [ "run", "build", "--workspace", "packages/tui", "--workspace", "packages/ai", "--workspace", "packages/agent", "--workspace", "packages/coding-agent", ], { cwd: repoRoot, env: process.env, stdio: ["ignore", "pipe", "pipe"], shell: process.platform === "win32", }, );
let stdout = ""; let stderr = ""; child.stdout.setEncoding("utf8"); child.stdout.on("data", (chunk) => { stdout += chunk; }); child.stderr.setEncoding("utf8"); child.stderr.on("data", (chunk) => { stderr += chunk; });
const exitCode = await waitForExit(child, "Build"); if (exitCode !== 0) { if (stdout.trim()) { process.stdout.write(`${stdout}${stdout.endsWith("\n") ? "" : "\n"}`); } if (stderr.trim()) { process.stderr.write(`${stderr}${stderr.endsWith("\n") ? "" : "\n"}`); } throw new Error(`Build failed with exit code ${exitCode}`); }
process.stdout.write(`Build completed in ${formatMs(performance.now() - startedAt)}\n`);}
function getRuntimeCommand(runtime, mode, profileDir, profileName, cpuProfile) { const benchmarkArgs = ["--no-session"]; if (mode === "rpc") { benchmarkArgs.push("--mode", "rpc"); }
if (runtime === "bun") { const args = []; if (cpuProfile) { args.push("--cpu-prof", `--cpu-prof-dir=${profileDir}`, `--cpu-prof-name=${profileName}`); } args.push(srcCliPath, ...benchmarkArgs); return { executable: "bun", args, }; }
const args = []; if (cpuProfile) { args.push("--cpu-prof", `--cpu-prof-dir=${profileDir}`, `--cpu-prof-name=${profileName}`); } args.push(distCliPath, ...benchmarkArgs); return { executable: process.execPath, args, };}
function createBenchmarkEnv(options, isolatedAgentDir) { const env = { ...process.env }; if (options.agentDir) { env[agentDirEnvName] = options.agentDir; } else if (isolatedAgentDir) { env[agentDirEnvName] = isolatedAgentDir; } if (options.mode === "tui") { env[startupBenchmarkEnvName] = "1"; } if (options.offline) { env.PI_OFFLINE = "1"; env.PI_SKIP_VERSION_CHECK = "1"; } return env;}
async function runTuiBenchmarkRun({ runtime, runIndex, measuredIndex, options, profileDir }) { const runNumber = runIndex + 1; const suffix = String(runNumber).padStart(3, "0"); const profileName = `${options.label}-${suffix}.cpuprofile`; const tempRoot = options.isolatedAgentDir ? mkdtempSync(join(tmpdir(), "pi-startup-benchmark-")) : undefined; const isolatedAgentDir = tempRoot ? join(tempRoot, "agent") : undefined; if (isolatedAgentDir) { mkdirSync(isolatedAgentDir, { recursive: true }); }
const command = getRuntimeCommand(runtime, "tui", profileDir, profileName, options.cpuProfile); const child = spawn(command.executable, command.args, { cwd: packageDir, env: createBenchmarkEnv(options, isolatedAgentDir), stdio: ["inherit", "ignore", "pipe"], shell: process.platform === "win32" && runtime === "bun", });
let stderr = ""; child.stderr.setEncoding("utf8"); child.stderr.on("data", (chunk) => { stderr += chunk; });
const startedAt = performance.now(); const exitCode = await waitForExit(child, `Benchmark ${measuredIndex === undefined ? `warmup ${runNumber}` : `run ${measuredIndex}`}`); const elapsedMs = performance.now() - startedAt;
try { if (exitCode !== 0) { throw new Error(stderr.trim() || `Benchmark child exited with code ${exitCode}`); }
const profilePath = options.cpuProfile ? join(profileDir, profileName) : undefined; if (profilePath && !existsSync(profilePath)) { throw new Error(`CPU profile was not written: ${profilePath}`); }
return { elapsedMs, profilePath, timings: parseStartupTimings(stderr) }; } finally { if (tempRoot) { rmSync(tempRoot, { recursive: true, force: true }); } }}
function splitJsonLines(buffer, onLine) { let remaining = buffer; while (true) { const newlineIndex = remaining.indexOf("\n"); if (newlineIndex === -1) { return remaining; } const line = remaining.slice(0, newlineIndex); onLine(line.endsWith("\r") ? line.slice(0, -1) : line); remaining = remaining.slice(newlineIndex + 1); }}
async function runRpcBenchmarkRun({ runtime, runIndex, measuredIndex, options, profileDir }) { const runNumber = runIndex + 1; const suffix = String(runNumber).padStart(3, "0"); const profileName = `${options.label}-${suffix}.cpuprofile`; const tempRoot = options.isolatedAgentDir ? mkdtempSync(join(tmpdir(), "pi-startup-benchmark-")) : undefined; const isolatedAgentDir = tempRoot ? join(tempRoot, "agent") : undefined; if (isolatedAgentDir) { mkdirSync(isolatedAgentDir, { recursive: true }); }
const command = getRuntimeCommand(runtime, "rpc", profileDir, profileName, options.cpuProfile); const child = spawn(command.executable, command.args, { cwd: packageDir, env: createBenchmarkEnv(options, isolatedAgentDir), stdio: ["pipe", "pipe", "pipe"], shell: process.platform === "win32" && runtime === "bun", });
let stdoutBuffer = ""; let stderr = ""; let readyElapsedMs; let responseError; const requestId = `startup-benchmark-${runNumber}`; const startedAt = performance.now();
child.stdout.setEncoding("utf8"); child.stdout.on("data", (chunk) => { stdoutBuffer = splitJsonLines(stdoutBuffer + chunk, (line) => { if (line.trim() === "") { return; } let parsed; try { parsed = JSON.parse(line); } catch (error) { responseError = error instanceof Error ? error.message : String(error); return; }
if (parsed?.type !== "response" || parsed.id !== requestId || parsed.command !== "get_state") { return; }
if (parsed.success !== true) { responseError = typeof parsed.error === "string" ? parsed.error : "get_state failed"; return; }
if (readyElapsedMs === undefined) { readyElapsedMs = performance.now() - startedAt; child.stdin.end(); } }); });
child.stderr.setEncoding("utf8"); child.stderr.on("data", (chunk) => { stderr += chunk; });
child.stdin.setDefaultEncoding("utf8"); child.stdin.write(`${JSON.stringify({ id: requestId, type: "get_state" })}\n`);
const exitCode = await waitForExit(child, `Benchmark ${measuredIndex === undefined ? `warmup ${runNumber}` : `run ${measuredIndex}`}`);
try { if (responseError) { throw new Error(responseError); } if (readyElapsedMs === undefined) { throw new Error(stderr.trim() || "RPC benchmark did not receive get_state response"); } if (exitCode !== 0) { throw new Error(stderr.trim() || `Benchmark child exited with code ${exitCode}`); }
const profilePath = options.cpuProfile ? join(profileDir, profileName) : undefined; if (profilePath && !existsSync(profilePath)) { throw new Error(`CPU profile was not written: ${profilePath}`); }
return { elapsedMs: readyElapsedMs, profilePath, timings: parseStartupTimings(stderr) }; } finally { if (tempRoot) { rmSync(tempRoot, { recursive: true, force: true }); } }}
async function runBenchmarkRun(params) { if (params.options.mode === "rpc") { return await runRpcBenchmarkRun(params); } return await runTuiBenchmarkRun(params);}
async function main() { const options = parseArgs(process.argv.slice(2)); if (options.help) { printHelp(); return; }
if (options.agentDir && options.isolatedAgentDir) { throw new Error("--agent-dir and --isolated-agent-dir cannot be combined"); }
if (options.mode === "tui" && (!process.stdin.isTTY || !process.stdout.isTTY)) { throw new Error("TUI benchmark must be run from an interactive terminal."); }
const runtime = resolveRuntime(options.runtime); options.label = resolveLabel(options.mode, options.label); const profileDir = resolveProfileDir(runtime, options.profileDir);
if (runtime === "node" && options.build) { await runBuild(); } if (runtime === "bun") { process.stdout.write( `Using Bun runtime with ${options.mode === "rpc" ? "packages/coding-agent/src/cli.ts --mode rpc" : "packages/coding-agent/src/cli.ts"}\n`, ); }
const entryPath = runtime === "bun" ? srcCliPath : distCliPath; if (!existsSync(entryPath)) { throw new Error(`CLI entrypoint not found: ${entryPath}`); }
mkdirSync(profileDir, { recursive: true });
const measuredRuns = []; const totalRuns = options.warmup + options.runs; for (let runIndex = 0; runIndex < totalRuns; runIndex++) { const measuredIndex = runIndex >= options.warmup ? runIndex - options.warmup + 1 : undefined; const result = await runBenchmarkRun({ runtime, runIndex, measuredIndex, options, profileDir, });
process.stdout.write( `[${measuredIndex === undefined ? `warmup ${runIndex + 1}` : `run ${measuredIndex}`}] elapsed=${formatMs(result.elapsedMs)}\n`, );
if (measuredIndex !== undefined) { measuredRuns.push(result); } }
if (measuredRuns.length === 0) { process.stdout.write("\nNo measured runs requested.\n"); return; }
const elapsedSummary = summarize(measuredRuns.map((run) => run.elapsedMs)); const timingSummaries = summarizeTimingMaps(measuredRuns); const maxElapsedRun = measuredRuns.reduce((slowest, run) => (run.elapsedMs > slowest.elapsedMs ? run : slowest)); if (measuredRuns.length === 1) { process.stdout.write("\nResult\n"); process.stdout.write(` runtime: ${runtime}\n`); process.stdout.write(` mode: ${options.mode}\n`); process.stdout.write(` elapsed: ${formatMs(measuredRuns[0].elapsedMs)}\n`); for (const [label, summary] of timingSummaries.entries()) { process.stdout.write(` ${label}: ${formatMs(summary.median)}\n`); } if (options.cpuProfile && maxElapsedRun.profilePath) { process.stdout.write(` selected profile: ${toDisplayPath(maxElapsedRun.profilePath)}\n`); process.stdout.write(` profiles dir: ${toDisplayPath(profileDir)}\n`); } process.stdout.write(`METRIC startup_time_ms=${measuredRuns[0].elapsedMs.toFixed(1)}\n`); for (const [label, summary] of timingSummaries.entries()) { process.stdout.write(`METRIC ${toMetricName(label)}=${summary.median.toFixed(1)}\n`); } return; }
process.stdout.write("\nSummary\n"); process.stdout.write(` runtime: ${runtime}\n`); process.stdout.write(` mode: ${options.mode}\n`); process.stdout.write(` elapsed min: ${formatMs(elapsedSummary.min)}\n`); process.stdout.write(` elapsed median: ${formatMs(elapsedSummary.median)}\n`); process.stdout.write(` elapsed avg: ${formatMs(elapsedSummary.avg)}\n`); process.stdout.write(` elapsed max: ${formatMs(elapsedSummary.max)}\n`); for (const [label, summary] of timingSummaries.entries()) { process.stdout.write(` ${label} median: ${formatMs(summary.median)}\n`); } if (options.cpuProfile && maxElapsedRun.profilePath) { process.stdout.write(` selected profile: ${toDisplayPath(maxElapsedRun.profilePath)}\n`); process.stdout.write(` profiles dir: ${toDisplayPath(profileDir)}\n`); } process.stdout.write(`METRIC startup_time_ms=${elapsedSummary.median.toFixed(1)}\n`); for (const [label, summary] of timingSummaries.entries()) { process.stdout.write(`METRIC ${toMetricName(label)}=${summary.median.toFixed(1)}\n`); }}
main().catch((error) => { const message = error instanceof Error ? error.message : String(error); console.error(message); process.exit(1);});