diff --git a/.gitignore b/.gitignore index a06c1377e..33722b2c8 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ dist/ packages/*/dist/ packages/*/dist-chrome/ packages/*/dist-firefox/ +*.cpuprofile # Environment .env diff --git a/package.json b/package.json index 958d51bfa..6b7b03075 100644 --- a/package.json +++ b/package.json @@ -17,6 +17,8 @@ "dev:tsc": "concurrently --names \"ai,web-ui\" --prefix-colors \"cyan,green\" \"cd packages/ai && npm run dev:tsc\" \"cd packages/web-ui && npm run dev:tsc\"", "check": "biome check --write --error-on-warnings . && tsgo --noEmit && npm run check:browser-smoke && cd packages/web-ui && npm run check", "check:browser-smoke": "node scripts/check-browser-smoke.mjs", + "profile:tui": "node scripts/profile-coding-agent-node.mjs --mode tui", + "profile:rpc": "node scripts/profile-coding-agent-node.mjs --mode rpc", "test": "npm run test --workspaces --if-present", "version:patch": "npm version patch -ws --no-git-tag-version && node scripts/sync-versions.js && shx rm -rf node_modules packages/*/node_modules package-lock.json && npm install", "version:minor": "npm version minor -ws --no-git-tag-version && node scripts/sync-versions.js && shx rm -rf node_modules packages/*/node_modules package-lock.json && npm install", diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index 7384799c1..2812d0d8e 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -754,6 +754,11 @@ export async function main(args: string[]) { stdinContent, ); const isInteractive = !parsed.print && parsed.mode === undefined; + const startupBenchmark = isTruthyEnvFlag(process.env.PI_STARTUP_BENCHMARK); + if (startupBenchmark && !isInteractive) { + console.error(chalk.red("Error: PI_STARTUP_BENCHMARK only supports interactive mode")); + process.exit(1); + } const mode = parsed.mode || "text"; initTheme(settingsManager.getTheme(), isInteractive); @@ -848,8 +853,7 @@ export async function main(args: string[]) { console.log(chalk.dim(`Model scope: ${modelList} ${chalk.gray("(Ctrl+P to cycle)")}`)); } - printTimings(); - const mode = new InteractiveMode(session, { + const interactiveMode = new InteractiveMode(session, { migratedProviders, modelFallbackMessage, initialMessage, @@ -857,7 +861,21 @@ export async function main(args: string[]) { initialMessages: parsed.messages, verbose: parsed.verbose, }); - await mode.run(); + if (startupBenchmark) { + await interactiveMode.init(); + interactiveMode.stop(); + stopThemeWatcher(); + if (process.stdout.writableLength > 0) { + await new Promise((resolve) => process.stdout.once("drain", resolve)); + } + if (process.stderr.writableLength > 0) { + await new Promise((resolve) => process.stderr.once("drain", resolve)); + } + return; + } + + printTimings(); + await interactiveMode.run(); } else { await runPrintMode(session, { mode, diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 4c9b7cf7a..7181e189b 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -596,9 +596,7 @@ export async function runRpcMode(session: AgentSession): Promise { */ let detachInput = () => {}; - async function checkShutdownRequested(): Promise { - if (!shutdownRequested) return; - + async function shutdown(): Promise { const currentRunner = session.extensionRunner; if (currentRunner?.hasHandlers("session_shutdown")) { await currentRunner.emit({ type: "session_shutdown" }); @@ -609,6 +607,11 @@ export async function runRpcMode(session: AgentSession): Promise { process.exit(0); } + async function checkShutdownRequested(): Promise { + if (!shutdownRequested) return; + await shutdown(); + } + const handleInputLine = async (line: string) => { try { const parsed = JSON.parse(line); @@ -636,9 +639,20 @@ export async function runRpcMode(session: AgentSession): Promise { } }; - detachInput = attachJsonlLineReader(process.stdin, (line) => { - void handleInputLine(line); - }); + const onInputEnd = () => { + void shutdown(); + }; + process.stdin.on("end", onInputEnd); + + detachInput = (() => { + const detachJsonl = attachJsonlLineReader(process.stdin, (line) => { + void handleInputLine(line); + }); + return () => { + detachJsonl(); + process.stdin.off("end", onInputEnd); + }; + })(); // Keep process alive forever return new Promise(() => {}); diff --git a/scripts/profile-coding-agent-node.mjs b/scripts/profile-coding-agent-node.mjs new file mode 100644 index 000000000..9127f0659 --- /dev/null +++ b/scripts/profile-coding-agent-node.mjs @@ -0,0 +1,550 @@ +import { 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 tui or rpc (default: tui) + --runs Number of measured runs (default: 1) + --warmup Number of warmup runs before measurements (default: 0) + --profile-dir CPU profile output directory + Default: profiles-node for Node, profiles-bun for Bun + --label Profile name prefix (default: -startup) + --runtime node, bun, or auto (default: auto) + --agent-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) + --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, + }; + + 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 === "--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, + }; +} + +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/coding-agent...\n"); + const startedAt = performance.now(); + const child = spawn("npm", ["run", "build"], { + cwd: packageDir, + 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) { + const benchmarkArgs = ["--no-session"]; + if (mode === "rpc") { + benchmarkArgs.push("--mode", "rpc"); + } + + if (runtime === "bun") { + return { + executable: "bun", + args: [ + "--cpu-prof", + `--cpu-prof-dir=${profileDir}`, + `--cpu-prof-name=${profileName}`, + srcCliPath, + ...benchmarkArgs, + ], + }; + } + + return { + executable: process.execPath, + args: [ + "--cpu-prof", + `--cpu-prof-dir=${profileDir}`, + `--cpu-prof-name=${profileName}`, + distCliPath, + ...benchmarkArgs, + ], + }; +} + +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); + 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 = join(profileDir, profileName); + if (!existsSync(profilePath)) { + throw new Error(`CPU profile was not written: ${profilePath}`); + } + + return { elapsedMs, profilePath }; + } 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); + 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 = join(profileDir, profileName); + if (!existsSync(profilePath)) { + throw new Error(`CPU profile was not written: ${profilePath}`); + } + + return { elapsedMs: readyElapsedMs, profilePath }; + } 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 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`); + process.stdout.write(` selected profile: ${toDisplayPath(maxElapsedRun.profilePath)}\n`); + process.stdout.write(` profiles dir: ${toDisplayPath(profileDir)}\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`); + process.stdout.write(` selected profile: ${toDisplayPath(maxElapsedRun.profilePath)}\n`); + process.stdout.write(` profiles dir: ${toDisplayPath(profileDir)}\n`); +} + +main().catch((error) => { + const message = error instanceof Error ? error.message : String(error); + console.error(message); + process.exit(1); +});