Skip to content

Commit da55146

Browse files
vincentkocvincentkoc
committed
fix(agent-sessions): fail oversized exec output
1 parent 2252cf6 commit da55146

2 files changed

Lines changed: 237 additions & 22 deletions

File tree

‎src/agents/sessions/exec.test.ts‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
import { EventEmitter } from "node:events";
2+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
3+
4+
const { spawnMock, waitForChildProcessMock } = vi.hoisted(() => ({
5+
spawnMock: vi.fn(),
6+
waitForChildProcessMock: vi.fn(),
7+
}));
8+
9+
vi.mock("node:child_process", () => ({
10+
spawn: spawnMock,
11+
}));
12+
13+
vi.mock("../utils/child-process.js", () => ({
14+
waitForChildProcess: waitForChildProcessMock,
15+
}));
16+
17+
type StubChild = EventEmitter & {
18+
kill: ReturnType<typeof vi.fn>;
19+
stderr: EventEmitter;
20+
stdout: EventEmitter;
21+
};
22+
23+
function createStubChild(): StubChild {
24+
const child = new EventEmitter() as StubChild;
25+
child.stdout = new EventEmitter();
26+
child.stderr = new EventEmitter();
27+
child.kill = vi.fn();
28+
return child;
29+
}
30+
31+
function createDeferred<T>() {
32+
let resolve!: (value: T) => void;
33+
let reject!: (error: Error) => void;
34+
const promise = new Promise<T>((res, rej) => {
35+
resolve = res;
36+
reject = rej;
37+
});
38+
return { promise, resolve, reject };
39+
}
40+
41+
describe("execCommand", () => {
42+
beforeEach(() => {
43+
spawnMock.mockReset();
44+
waitForChildProcessMock.mockReset();
45+
vi.useRealTimers();
46+
});
47+
48+
afterEach(() => {
49+
vi.useRealTimers();
50+
vi.restoreAllMocks();
51+
});
52+
53+
it("bounds retained stdout and stderr independently", async () => {
54+
const child = createStubChild();
55+
const wait = createDeferred<number | null>();
56+
spawnMock.mockReturnValue(child);
57+
waitForChildProcessMock.mockReturnValue(wait.promise);
58+
const { execCommand } = await import("./exec.js");
59+
60+
const resultPromise = execCommand("cmd", [], "/tmp", { maxOutputChars: 256 });
61+
child.stdout.emit("data", Buffer.from(`${"a".repeat(300)}stdout-tail`));
62+
child.stderr.emit("data", Buffer.from(`${"b".repeat(300)}stderr-tail`));
63+
wait.resolve(0);
64+
65+
const result = await resultPromise;
66+
expect(result.code).toBe(0);
67+
expect(result.stdout.length).toBeLessThanOrEqual(256);
68+
expect(result.stderr.length).toBeLessThanOrEqual(256);
69+
expect(result.stdout.endsWith("stdout-tail")).toBe(true);
70+
expect(result.stderr.endsWith("stderr-tail")).toBe(true);
71+
expect(result.stdoutTruncatedChars).toBeGreaterThan(0);
72+
expect(result.stderrTruncatedChars).toBeGreaterThan(0);
73+
});
74+
75+
it("honors caller-supplied small output caps", async () => {
76+
const child = createStubChild();
77+
const wait = createDeferred<number | null>();
78+
spawnMock.mockReturnValue(child);
79+
waitForChildProcessMock.mockReturnValue(wait.promise);
80+
const { execCommand } = await import("./exec.js");
81+
82+
const resultPromise = execCommand("cmd", [], "/tmp", { maxOutputChars: 3 });
83+
child.stdout.emit("data", Buffer.from("abcdef"));
84+
wait.resolve(0);
85+
86+
const result = await resultPromise;
87+
expect(result.code).toBe(0);
88+
expect(result.stdout).toBe("def");
89+
expect(result.stdoutTruncatedChars).toBe(3);
90+
});
91+
92+
it("fails instead of silently truncating default exec output", async () => {
93+
const child = createStubChild();
94+
const wait = createDeferred<number | null>();
95+
spawnMock.mockReturnValue(child);
96+
waitForChildProcessMock.mockReturnValue(wait.promise);
97+
const { execCommand } = await import("./exec.js");
98+
99+
const resultPromise = execCommand("cmd", [], "/tmp");
100+
child.stdout.emit("data", Buffer.from("x".repeat(16 * 1024 * 1024 + 1)));
101+
wait.resolve(0);
102+
103+
const result = await resultPromise;
104+
expect(child.kill).toHaveBeenCalledWith("SIGTERM");
105+
expect(result.code).toBe(1);
106+
expect(result.killed).toBe(true);
107+
expect(result.outputLimitExceeded).toBe("stdout");
108+
expect(result.stdout.length).toBe(16 * 1024 * 1024);
109+
expect(result.stdoutTruncatedChars).toBe(1);
110+
expect(result.stderr).toContain("exec stdout exceeded output limit");
111+
});
112+
113+
it("escalates timed-out commands to SIGKILL after the grace period", async () => {
114+
vi.useFakeTimers();
115+
const child = createStubChild();
116+
const wait = createDeferred<number | null>();
117+
spawnMock.mockReturnValue(child);
118+
waitForChildProcessMock.mockReturnValue(wait.promise);
119+
const { execCommand } = await import("./exec.js");
120+
121+
const resultPromise = execCommand("cmd", [], "/tmp", { timeout: 10 });
122+
await vi.advanceTimersByTimeAsync(10);
123+
expect(child.kill).toHaveBeenCalledWith("SIGTERM");
124+
expect(child.kill).not.toHaveBeenCalledWith("SIGKILL");
125+
126+
await vi.advanceTimersByTimeAsync(4_999);
127+
expect(child.kill).not.toHaveBeenCalledWith("SIGKILL");
128+
await vi.advanceTimersByTimeAsync(1);
129+
expect(child.kill).toHaveBeenCalledWith("SIGKILL");
130+
131+
wait.resolve(null);
132+
const result = await resultPromise;
133+
expect(result.killed).toBe(true);
134+
});
135+
});

‎src/agents/sessions/exec.ts‎

Lines changed: 102 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@
55
import { spawn } from "node:child_process";
66
import { waitForChildProcess } from "../utils/child-process.js";
77

8+
const DEFAULT_OUTPUT_LIMIT_CHARS = 16 * 1024 * 1024;
9+
const FORCE_KILL_GRACE_MS = 5000;
10+
811
/**
912
* Options for executing shell commands.
1013
*/
@@ -15,6 +18,8 @@ export interface ExecOptions {
1518
timeout?: number;
1619
/** Working directory */
1720
cwd?: string;
21+
/** Optional maximum retained stdout/stderr characters per stream. */
22+
maxOutputChars?: number;
1823
}
1924

2025
/**
@@ -23,10 +28,46 @@ export interface ExecOptions {
2328
export interface ExecResult {
2429
stdout: string;
2530
stderr: string;
31+
stdoutTruncatedChars?: number;
32+
stderrTruncatedChars?: number;
33+
outputLimitExceeded?: "stdout" | "stderr";
2634
code: number;
2735
killed: boolean;
2836
}
2937

38+
type OutputCapture = {
39+
text: string;
40+
truncatedChars: number;
41+
};
42+
43+
function clampMaxOutputChars(value: number | undefined): number {
44+
if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
45+
return DEFAULT_OUTPUT_LIMIT_CHARS;
46+
}
47+
return Math.max(1, Math.floor(value));
48+
}
49+
50+
function appendCapturedOutput(
51+
current: OutputCapture,
52+
chunk: Buffer | string,
53+
maxOutputChars: number,
54+
truncateTail: boolean,
55+
): OutputCapture {
56+
const text = String(chunk);
57+
const combined = `${current.text}${text}`;
58+
const overflowChars = Math.max(0, combined.length - maxOutputChars);
59+
if (overflowChars === 0) {
60+
return {
61+
text: combined,
62+
truncatedChars: current.truncatedChars,
63+
};
64+
}
65+
return {
66+
text: truncateTail ? combined.slice(overflowChars) : combined.slice(0, maxOutputChars),
67+
truncatedChars: current.truncatedChars + overflowChars,
68+
};
69+
}
70+
3071
/**
3172
* Execute a shell command and return stdout/stderr/code.
3273
* Supports timeout and abort signal.
@@ -44,21 +85,64 @@ export async function execCommand(
4485
stdio: ["ignore", "pipe", "pipe"],
4586
});
4687

47-
let stdout = "";
48-
let stderr = "";
88+
let stdout: OutputCapture = { text: "", truncatedChars: 0 };
89+
let stderr: OutputCapture = { text: "", truncatedChars: 0 };
4990
let killed = false;
5091
let timeoutId: NodeJS.Timeout | undefined;
92+
let forceKillTimer: NodeJS.Timeout | undefined;
93+
let settled = false;
94+
const maxOutputChars = clampMaxOutputChars(options?.maxOutputChars);
95+
const truncateOutput = options?.maxOutputChars !== undefined;
96+
let outputLimitExceeded: "stdout" | "stderr" | undefined;
97+
const markOutputLimitExceeded = (stream: "stdout" | "stderr") => {
98+
if (!truncateOutput && !outputLimitExceeded) {
99+
outputLimitExceeded = stream;
100+
killProcess();
101+
}
102+
};
103+
const finish = (code: number) => {
104+
if (settled) {
105+
return;
106+
}
107+
settled = true;
108+
if (timeoutId) {
109+
clearTimeout(timeoutId);
110+
}
111+
if (forceKillTimer) {
112+
clearTimeout(forceKillTimer);
113+
}
114+
if (options?.signal) {
115+
options.signal.removeEventListener("abort", killProcess);
116+
}
117+
if (outputLimitExceeded) {
118+
stderr = appendCapturedOutput(
119+
stderr,
120+
`${stderr.text ? "\n" : ""}exec ${outputLimitExceeded} exceeded output limit ${maxOutputChars} chars`,
121+
maxOutputChars,
122+
true,
123+
);
124+
}
125+
resolve({
126+
stdout: stdout.text,
127+
stderr: stderr.text,
128+
stdoutTruncatedChars: stdout.truncatedChars || undefined,
129+
stderrTruncatedChars: stderr.truncatedChars || undefined,
130+
outputLimitExceeded,
131+
code: outputLimitExceeded ? 1 : code,
132+
killed,
133+
});
134+
};
51135

52136
const killProcess = () => {
53137
if (!killed) {
54138
killed = true;
55139
proc.kill("SIGTERM");
56-
// Force kill after 5 seconds if SIGTERM doesn't work
57-
setTimeout(() => {
58-
if (!proc.killed) {
140+
forceKillTimer = setTimeout(() => {
141+
if (!settled) {
59142
proc.kill("SIGKILL");
60143
}
61-
}, 5000);
144+
}, FORCE_KILL_GRACE_MS);
145+
forceKillTimer.unref?.();
62146
}
63147
};
64148

@@ -79,33 +163,29 @@ export async function execCommand(
79163
}
80164

81165
proc.stdout?.on("data", (data) => {
82-
stdout += data.toString();
166+
const before = stdout.truncatedChars;
167+
stdout = appendCapturedOutput(stdout, data, maxOutputChars, truncateOutput);
168+
if (stdout.truncatedChars > before) {
169+
markOutputLimitExceeded("stdout");
170+
}
83171
});
84172

85173
proc.stderr?.on("data", (data) => {
86-
stderr += data.toString();
174+
const before = stderr.truncatedChars;
175+
stderr = appendCapturedOutput(stderr, data, maxOutputChars, truncateOutput);
176+
if (stderr.truncatedChars > before) {
177+
markOutputLimitExceeded("stderr");
178+
}
87179
});
88180

89181
// Wait for process termination without hanging on inherited stdio handles
90182
// held open by detached descendants.
91183
waitForChildProcess(proc)
92184
.then((code) => {
93-
if (timeoutId) {
94-
clearTimeout(timeoutId);
95-
}
96-
if (options?.signal) {
97-
options.signal.removeEventListener("abort", killProcess);
98-
}
99-
resolve({ stdout, stderr, code: code ?? 0, killed });
185+
finish(code ?? 0);
100186
})
101187
.catch(() => {
102-
if (timeoutId) {
103-
clearTimeout(timeoutId);
104-
}
105-
if (options?.signal) {
106-
options.signal.removeEventListener("abort", killProcess);
107-
}
108-
resolve({ stdout, stderr, code: 1, killed });
188+
finish(1);
109189
});
110190
});
111191
}

0 commit comments

Comments
 (0)