diff --git a/packages/cli/src/cli/cmd/run.ts b/packages/cli/src/cli/cmd/run.ts index 7e17273..3313b7e 100644 --- a/packages/cli/src/cli/cmd/run.ts +++ b/packages/cli/src/cli/cmd/run.ts @@ -35,6 +35,7 @@ import { SkillTool } from "../../tool/skill" import { BashTool } from "../../tool/bash" import { TodoWriteTool } from "../../tool/todo" import { Locale } from "../../util/locale" +import { Stdout } from "../stdout" type ToolProps = { input: Tool.InferParameters @@ -465,7 +466,7 @@ export const RunCommand = cmd({ function emit(type: string, data: Record) { if (args.format === "json") { - process.stdout.write(JSON.stringify({ type, timestamp: Date.now(), sessionID, ...data }) + EOL) + Stdout.write(JSON.stringify({ type, timestamp: Date.now(), sessionID, ...data }) + EOL) return true } return false diff --git a/packages/cli/src/cli/shutdown.ts b/packages/cli/src/cli/shutdown.ts new file mode 100644 index 0000000..bef93d3 --- /dev/null +++ b/packages/cli/src/cli/shutdown.ts @@ -0,0 +1,13 @@ +import { Log } from "../util/log" +import { Stdout } from "./stdout" + +export namespace Shutdown { + export async function flush() { + await Stdout.flush().catch((error) => { + Log.Default.error("stdout flush failed", { + error: error instanceof Error ? error.message : error, + }) + process.exitCode = 1 + }) + } +} diff --git a/packages/cli/src/cli/stdout.ts b/packages/cli/src/cli/stdout.ts new file mode 100644 index 0000000..a59a0af --- /dev/null +++ b/packages/cli/src/cli/stdout.ts @@ -0,0 +1,62 @@ +const state = { + bound: false, + closed: false, + pending: new Set>(), +} +let failure: unknown + +function epipe(error: unknown) { + return error instanceof Error && "code" in error && error.code === "EPIPE" +} + +function fail(error: unknown) { + if (epipe(error)) { + state.closed = true + return false + } + failure = error + return true +} + +function bind() { + if (state.bound) return + state.bound = true + process.stdout.on("error", fail) +} + +export namespace Stdout { + export function write(chunk: string) { + bind() + if (failure !== undefined) { + const pending = Promise.reject(failure) + pending.catch(() => {}) + return pending + } + if (state.closed || process.stdout.destroyed || process.stdout.writableEnded) return Promise.resolve() + const pending = new Promise((resolve, reject) => { + process.stdout.write(chunk, (error) => { + if (error && fail(error)) { + reject(error) + return + } + resolve() + }) + }) + state.pending.add(pending) + pending.catch(() => {}) + pending.then( + () => state.pending.delete(pending), + () => state.pending.delete(pending), + ) + return pending + } + + export async function flush() { + while (state.pending.size) await Promise.allSettled(state.pending) + if (failure !== undefined) throw failure + } + + export function isClosed() { + return state.closed + } +} diff --git a/packages/cli/src/headless.ts b/packages/cli/src/headless.ts index 71f2eff..f9140aa 100644 --- a/packages/cli/src/headless.ts +++ b/packages/cli/src/headless.ts @@ -19,6 +19,7 @@ import path from "path" import { Global } from "./global" import { JsonMigration } from "./storage/json-migration" import { Database } from "./storage/db" +import { Shutdown } from "./cli/shutdown" process.on("unhandledRejection", (e) => { Log.Default.error("rejection", { @@ -154,5 +155,6 @@ try { } process.exitCode = 1 } finally { + await Shutdown.flush() process.exit() } diff --git a/packages/cli/src/index.ts b/packages/cli/src/index.ts index 8a333a5..e2d4b9c 100644 --- a/packages/cli/src/index.ts +++ b/packages/cli/src/index.ts @@ -29,6 +29,7 @@ import path from "path" import { Global } from "./global" import { JsonMigration } from "./storage/json-migration" import { Database } from "./storage/db" +import { Shutdown } from "./cli/shutdown" process.on("unhandledRejection", (e) => { Log.Default.error("rejection", { @@ -212,5 +213,6 @@ try { // Most notably, some docker-container-based MCP servers don't handle such signals unless // run using `docker run --init`. // Explicitly exit to avoid any hanging subprocesses. + await Shutdown.flush() process.exit() } diff --git a/packages/cli/test/cli/fixture/stdout.ts b/packages/cli/test/cli/fixture/stdout.ts new file mode 100644 index 0000000..9bc7760 --- /dev/null +++ b/packages/cli/test/cli/fixture/stdout.ts @@ -0,0 +1,47 @@ +import { Stdout } from "../../../src/cli/stdout" +import { Shutdown } from "../../../src/cli/shutdown" + +if (process.argv.includes("--idle")) process.exit(process.stdout.listenerCount("error")) +if (process.argv.includes("--shutdown-error")) { + Stdout.write("") + await Stdout.flush() + process.stdout.emit("error", Object.assign(new Error("broken stdout"), { code: "EIO" })) + await Shutdown.flush() + const status = process.argv.find((arg) => arg.startsWith("--status="))?.slice("--status=".length) + if (status) await Bun.write(status, String(process.exitCode)) + process.exit() +} +if (process.argv.includes("--error")) { + Stdout.write("") + await Stdout.flush() + process.stdout.emit("error", Object.assign(new Error("broken stdout"), { code: "EIO" })) + const write = await Stdout.write("").then( + () => false, + (error) => error instanceof Error && error.message === "broken stdout", + ) + const flush = await Stdout.flush().then( + () => false, + (error) => error instanceof Error && error.message === "broken stdout", + ) + process.exit(write && flush ? 0 : 2) +} + +const records = Array.from({ length: 1024 }, (_, index) => + JSON.stringify({ + type: "data", + index, + value: "x".repeat(2048), + }), +) + +Array.from({ length: process.argv.includes("--epipe") ? 8 : 1 }).forEach(() => + records.forEach((record) => Stdout.write(record + "\n")), +) +Stdout.write(JSON.stringify({ type: "complete" }) + "\n") +await Stdout.flush() +if (process.argv.includes("--epipe")) { + const status = process.argv.find((arg) => arg.startsWith("--status="))?.slice("--status=".length) + if (status) await Bun.write(status, Stdout.isClosed() ? "closed" : "open") + process.exit(Stdout.isClosed() ? 0 : 2) +} +process.exit() diff --git a/packages/cli/test/cli/stdout.test.ts b/packages/cli/test/cli/stdout.test.ts new file mode 100644 index 0000000..afeffde --- /dev/null +++ b/packages/cli/test/cli/stdout.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, test } from "bun:test" +import { $ } from "bun" +import path from "path" +import { tmpdir } from "../fixture/fixture" + +describe("stdout", () => { + test("flushes NDJSON beyond pipe capacity before forced exit", async () => { + const child = Bun.spawn([process.execPath, path.join(import.meta.dir, "fixture", "stdout.ts")], { + cwd: path.join(import.meta.dir, "../.."), + stdout: "pipe", + stderr: "pipe", + }) + const output = new Response(child.stdout).text() + const error = new Response(child.stderr).text() + + expect(await child.exited).toBe(0) + expect(await error).toBe("") + + const lines = (await output) + .trim() + .split("\n") + .map((line) => JSON.parse(line)) + expect(lines).toHaveLength(1025) + expect(lines.at(-1)).toEqual({ type: "complete" }) + }) + + test("treats a closed output pipe as controlled", async () => { + await using tmp = await tmpdir() + const status = path.join(tmp.path, "status") + const child = await $`${process.execPath} ${path.join( + import.meta.dir, + "fixture", + "stdout.ts", + )} --epipe --status=${status} | ${process.execPath} -e ${"await Bun.stdin.stream().getReader().read(); process.exit()"}` + .cwd(path.join(import.meta.dir, "../..")) + .quiet() + .nothrow() + + expect(child.exitCode).toBe(0) + expect(child.stderr.toString()).toBe("") + expect(await Bun.file(status).text()).toBe("closed") + }) + + test("does not bind an error handler until output is written", async () => { + const child = Bun.spawn([process.execPath, path.join(import.meta.dir, "fixture", "stdout.ts"), "--idle"], { + cwd: path.join(import.meta.dir, "../.."), + stdout: "pipe", + stderr: "pipe", + }) + + expect(await child.exited).toBe(0) + expect(await new Response(child.stderr).text()).toBe("") + }) + + test("surfaces unexpected output errors from write and flush", async () => { + const child = Bun.spawn([process.execPath, path.join(import.meta.dir, "fixture", "stdout.ts"), "--error"], { + cwd: path.join(import.meta.dir, "../.."), + stdout: "pipe", + stderr: "pipe", + }) + + expect(await child.exited).toBe(0) + expect(await new Response(child.stderr).text()).toBe("") + }) + + test("flush failure sets exit status without blocking forced exit", async () => { + await using tmp = await tmpdir() + const status = path.join(tmp.path, "status") + const child = Bun.spawn( + [process.execPath, path.join(import.meta.dir, "fixture", "stdout.ts"), "--shutdown-error", `--status=${status}`], + { + cwd: path.join(import.meta.dir, "../.."), + stdout: "pipe", + stderr: "pipe", + }, + ) + + expect(await child.exited).toBe(1) + expect(await Bun.file(status).text()).toBe("1") + expect(await new Response(child.stderr).text()).toBe("") + }) +})