mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-19 21:56:00 +00:00
fix(agent-core): cap foreground shell output to prevent OOM crash (#1285)
* fix(agent-core): cap foreground shell output to prevent OOM crash A foreground command that streams a very large or unbounded amount of output (e.g. `b3sum --length 18446744073709551615`) grew the live-output buffer until Node aborted with a JavaScript heap out-of-memory error (exit 134). The per-command output is now capped at 16 MiB: on breach the command is gracefully terminated (SIGTERM -> grace -> SIGKILL) and the result carries a message pointing at redirecting large output to a file. The per-task output ring buffer is also made O(1) per chunk (was O(n^2)), which previously starved the event loop and the foreground timeout. Background/detached tasks are exempt. * fix(agent-core): stop buffering output after the foreground cap trips After the 16 MiB foreground ceiling tripped, appendOutput still enqueued every subsequent chunk into the per-task disk write chain during the SIGTERM grace window. A producer that ignores SIGTERM could keep that chain — and the chunk strings each pending write retains — growing until SIGKILL, re-introducing the same out-of-memory risk the cap prevents. Once the cap has tripped, keep only the bounded in-memory ring buffer and stop feeding the disk write chain. Interrupt/timeout capture behaviour is unchanged (they do not set outputLimitTripped).
This commit is contained in:
parent
3ea84a56e4
commit
c434b4c3e6
4 changed files with 248 additions and 4 deletions
5
.changeset/cap-foreground-command-output.md
Normal file
5
.changeset/cap-foreground-command-output.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
---
|
||||
"@moonshot-ai/agent-core": patch
|
||||
---
|
||||
|
||||
Cap the output a single foreground shell command may stream so a runaway command can no longer crash the process. A command that produces a very large or unbounded amount of output (e.g. `b3sum --length 18446744073709551615`) previously grew the live-output buffer until Node aborted with a JavaScript heap out-of-memory error; it is now gracefully terminated once its output exceeds 16 MiB, and the result explains how to redirect large output to a file instead. The per-task output ring buffer is also maintained in O(1) per chunk rather than O(n²). Background tasks are unaffected.
|
||||
|
|
@ -58,8 +58,20 @@ interface ManagedTask {
|
|||
readonly taskId: string;
|
||||
readonly task: BackgroundTask;
|
||||
readonly outputChunks: string[];
|
||||
/**
|
||||
* Running total of characters currently held in `outputChunks`, maintained
|
||||
* incrementally so the ring-buffer cap stays O(1) per chunk instead of
|
||||
* re-summing every chunk (which was O(n²) over a command's lifetime).
|
||||
*/
|
||||
outputRingChars: number;
|
||||
/** Total UTF-8 bytes observed, including chunks dropped from the live ring buffer. */
|
||||
outputSizeBytes: number;
|
||||
/**
|
||||
* True once a foreground command has crossed `MAX_FOREGROUND_OUTPUT_BYTES`
|
||||
* and termination has been requested. One-shot guard so the ceiling fires
|
||||
* exactly once.
|
||||
*/
|
||||
outputLimitTripped: boolean;
|
||||
status: BackgroundTaskStatus;
|
||||
/** Normalized registration options. Current mutable state stays on ManagedTask. */
|
||||
readonly options: RegisterBackgroundTaskOptions;
|
||||
|
|
@ -110,6 +122,28 @@ interface ManagedTask {
|
|||
const MAX_OUTPUT_BYTES = 1024 * 1024; // 1 MiB
|
||||
const NOTIFICATION_FALLBACK_PREVIEW_BYTES = 3_000;
|
||||
|
||||
/**
|
||||
* Hard ceiling on the combined output a single *foreground* command may stream
|
||||
* before it is force-terminated (SIGTERM → grace → SIGKILL). It guards the
|
||||
* live-forward path — which has no memory bound of its own — from a runaway
|
||||
* command (e.g. `b3sum --length <huge>`) whose output would otherwise grow the
|
||||
* process heap until Node aborts with an out-of-memory crash.
|
||||
*
|
||||
* Detached (background) tasks are exempt: their output is ring-buffered and
|
||||
* spilled to disk, so it never accumulates unbounded in memory.
|
||||
*/
|
||||
const MAX_FOREGROUND_OUTPUT_BYTES = 16 * 1024 * 1024; // 16 MiB
|
||||
|
||||
/** Terminal `stopReason` recorded when a foreground command trips the output ceiling. */
|
||||
function foregroundOutputLimitReason(): string {
|
||||
const mib = Math.floor(MAX_FOREGROUND_OUTPUT_BYTES / (1024 * 1024));
|
||||
return (
|
||||
`Output limit exceeded: the command produced more than ${mib} MiB and was ` +
|
||||
'terminated. Redirect large output to a file (e.g. `command > out.txt`) and ' +
|
||||
'inspect it in slices instead.'
|
||||
);
|
||||
}
|
||||
|
||||
const SIGTERM_GRACE_MS = 5_000;
|
||||
const USER_INTERRUPT_REASON = 'Interrupted by user';
|
||||
|
||||
|
|
@ -281,7 +315,9 @@ export class BackgroundManager {
|
|||
taskId,
|
||||
task,
|
||||
outputChunks: [],
|
||||
outputRingChars: 0,
|
||||
outputSizeBytes: 0,
|
||||
outputLimitTripped: false,
|
||||
status: 'running',
|
||||
options: entryOptions,
|
||||
startedAt: Date.now(),
|
||||
|
|
@ -594,14 +630,39 @@ export class BackgroundManager {
|
|||
private appendOutput(entry: ManagedTask, chunk: string): void {
|
||||
entry.outputSizeBytes += Buffer.byteLength(chunk, 'utf-8');
|
||||
entry.outputChunks.push(chunk);
|
||||
// Enforce output cap: drop oldest chunks when over budget.
|
||||
let total = entry.outputChunks.reduce((s, c) => s + c.length, 0);
|
||||
while (total > MAX_OUTPUT_BYTES && entry.outputChunks.length > 1) {
|
||||
entry.outputRingChars += chunk.length;
|
||||
// Enforce the ring-buffer cap: drop oldest chunks when over budget. The
|
||||
// running total keeps this O(1) amortized per chunk; re-summing the whole
|
||||
// buffer on every chunk was O(n²) over a command's lifetime and could
|
||||
// starve the event loop (and so the foreground timeout) under a high-rate
|
||||
// output stream.
|
||||
while (entry.outputRingChars > MAX_OUTPUT_BYTES && entry.outputChunks.length > 1) {
|
||||
const removed = entry.outputChunks.shift();
|
||||
if (removed === undefined) break;
|
||||
total -= removed.length;
|
||||
entry.outputRingChars -= removed.length;
|
||||
}
|
||||
|
||||
// Foreground output ceiling: a single non-detached command must not grow
|
||||
// the (unbounded) live-forward buffer until the process runs out of memory.
|
||||
// Trip once, then request graceful termination through the shared stop path
|
||||
// (SIGTERM → grace → SIGKILL). Detached tasks are exempt — their output is
|
||||
// ring-buffered and spilled to disk, so it never accumulates in memory.
|
||||
if (
|
||||
!entry.outputLimitTripped &&
|
||||
!this.isDetached(entry) &&
|
||||
entry.outputSizeBytes > MAX_FOREGROUND_OUTPUT_BYTES
|
||||
) {
|
||||
entry.outputLimitTripped = true;
|
||||
void this.stop(entry.taskId, foregroundOutputLimitReason());
|
||||
}
|
||||
|
||||
// Once the cap has tripped the task is being terminated: keep only the
|
||||
// bounded in-memory ring buffer above and stop feeding the (unbounded) disk
|
||||
// write chain. A producer that ignores SIGTERM could otherwise keep the
|
||||
// chain — and the chunk strings each pending write retains — growing through
|
||||
// the grace window until SIGKILL, re-introducing the OOM this cap prevents.
|
||||
if (entry.outputLimitTripped) return;
|
||||
|
||||
if (this.persistence === undefined) return;
|
||||
|
||||
// Foreground tasks keep their full output in memory and only touch disk
|
||||
|
|
|
|||
|
|
@ -139,6 +139,12 @@ function observeProcessStream(
|
|||
const onData = (chunk: string): void => {
|
||||
if (chunk.length === 0) return;
|
||||
sink.appendOutput(chunk);
|
||||
// Once the manager has begun terminating the task — an output-limit trip
|
||||
// (see MAX_FOREGROUND_OUTPUT_BYTES), a user interrupt, or a timeout —
|
||||
// `appendOutput` above may synchronously abort the signal. Stop forwarding
|
||||
// live output from that point so the unbounded forward buffer cannot keep
|
||||
// growing while the process is being killed.
|
||||
if (sink.signal.aborted) return;
|
||||
onOutput?.(kind, chunk);
|
||||
};
|
||||
stream.on('data', onData);
|
||||
|
|
|
|||
172
packages/agent-core/test/agent/background/output-limit.test.ts
Normal file
172
packages/agent-core/test/agent/background/output-limit.test.ts
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
/**
|
||||
* Foreground output ceiling.
|
||||
*
|
||||
* A single non-detached command that streams more output than the per-command
|
||||
* limit must be force-terminated instead of growing the (unbounded)
|
||||
* live-forward buffer until the process runs out of memory. Detached
|
||||
* (background) tasks are exempt: their output is ring-buffered and spilled to
|
||||
* disk, so it never accumulates in memory.
|
||||
*/
|
||||
|
||||
import { mkdtempSync, rmSync } from 'node:fs';
|
||||
import { tmpdir } from 'node:os';
|
||||
import { Readable, type Writable } from 'node:stream';
|
||||
|
||||
import type { KaosProcess } from '@moonshot-ai/kaos';
|
||||
import { join } from 'pathe';
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { ProcessBackgroundTask } from '../../../src/agent/background';
|
||||
import { createBackgroundManager, waitForTerminal } from './helpers';
|
||||
|
||||
const MiB = 1024 * 1024;
|
||||
const LIMIT_BYTES = 16 * MiB;
|
||||
|
||||
/**
|
||||
* A process that streams `chunks` of stdout, then exits 0 on its own — unless
|
||||
* it is killed first, in which case `wait()` resolves with the signal's exit
|
||||
* code and the stream is destroyed (simulating the child dying on SIGTERM).
|
||||
*/
|
||||
function streamingProcess(chunks: string[]): {
|
||||
proc: KaosProcess;
|
||||
kill: ReturnType<typeof vi.fn>;
|
||||
} {
|
||||
const stdout = Readable.from(chunks);
|
||||
const stderr = Readable.from([]);
|
||||
let resolveWait!: (code: number) => void;
|
||||
const waitP = new Promise<number>((resolve) => {
|
||||
resolveWait = resolve;
|
||||
});
|
||||
stdout.on('end', () => {
|
||||
resolveWait(0);
|
||||
});
|
||||
const kill = vi.fn(async (signal: string) => {
|
||||
stdout.destroy();
|
||||
resolveWait(signal === 'SIGKILL' ? 137 : 143);
|
||||
});
|
||||
const proc = {
|
||||
stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable,
|
||||
stdout,
|
||||
stderr,
|
||||
pid: 4242,
|
||||
exitCode: null,
|
||||
wait: () => waitP,
|
||||
kill,
|
||||
dispose: vi.fn().mockResolvedValue(undefined),
|
||||
} as unknown as KaosProcess;
|
||||
return { proc, kill };
|
||||
}
|
||||
|
||||
/**
|
||||
* A process that keeps streaming all of `chunks` regardless of SIGTERM (only
|
||||
* SIGKILL stops it) — simulating a producer that ignores the graceful stop and
|
||||
* keeps writing through the SIGTERM grace window.
|
||||
*/
|
||||
function sigtermIgnoringProcess(chunks: string[]): { proc: KaosProcess; kill: ReturnType<typeof vi.fn> } {
|
||||
const stdout = Readable.from(chunks);
|
||||
const stderr = Readable.from([]);
|
||||
let resolveWait!: (code: number) => void;
|
||||
const waitP = new Promise<number>((resolve) => {
|
||||
resolveWait = resolve;
|
||||
});
|
||||
stdout.on('end', () => {
|
||||
resolveWait(0);
|
||||
});
|
||||
const kill = vi.fn(async (signal: string) => {
|
||||
if (signal === 'SIGKILL') {
|
||||
stdout.destroy();
|
||||
resolveWait(137);
|
||||
}
|
||||
// SIGTERM is intentionally ignored.
|
||||
});
|
||||
const proc = {
|
||||
stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable,
|
||||
stdout,
|
||||
stderr,
|
||||
pid: 4243,
|
||||
exitCode: null,
|
||||
wait: () => waitP,
|
||||
kill,
|
||||
dispose: vi.fn().mockResolvedValue(undefined),
|
||||
} as unknown as KaosProcess;
|
||||
return { proc, kill };
|
||||
}
|
||||
|
||||
describe('BackgroundManager — foreground output ceiling', () => {
|
||||
it('terminates a foreground command that exceeds the output limit and stops forwarding', async () => {
|
||||
const { manager } = createBackgroundManager();
|
||||
// 20 MiB total, well past the 16 MiB ceiling.
|
||||
const chunks = Array.from({ length: 20 }, () => 'x'.repeat(MiB));
|
||||
const { proc, kill } = streamingProcess(chunks);
|
||||
|
||||
let forwardedChars = 0;
|
||||
const onOutput = vi.fn((_kind: 'stdout' | 'stderr', text: string) => {
|
||||
forwardedChars += text.length;
|
||||
});
|
||||
|
||||
const taskId = manager.registerTask(
|
||||
new ProcessBackgroundTask(proc, 'b3sum --length 18446744073709551615', 'hash', onOutput),
|
||||
{ detached: false, signal: new AbortController().signal, timeoutMs: 60_000 },
|
||||
);
|
||||
|
||||
const info = await waitForTerminal(manager, taskId);
|
||||
|
||||
expect(info?.status).toBe('killed');
|
||||
expect(info?.stopReason ?? '').toMatch(/output limit/i);
|
||||
expect(kill).toHaveBeenCalledWith('SIGTERM');
|
||||
// The live-forward path is capped at the ceiling rather than draining the
|
||||
// full 20 MiB into the (unbounded) transcript/stderr buffer.
|
||||
expect(forwardedChars).toBeLessThanOrEqual(LIMIT_BYTES);
|
||||
});
|
||||
|
||||
it('does not terminate a detached (background) task for the same output', async () => {
|
||||
const { manager } = createBackgroundManager();
|
||||
const chunks = Array.from({ length: 20 }, () => 'x'.repeat(MiB));
|
||||
const { proc, kill } = streamingProcess(chunks);
|
||||
|
||||
const taskId = manager.registerTask(new ProcessBackgroundTask(proc, 'producer', 'bg'), {
|
||||
detached: true,
|
||||
timeoutMs: 60_000,
|
||||
});
|
||||
|
||||
const info = await waitForTerminal(manager, taskId);
|
||||
|
||||
expect(info?.status).toBe('completed');
|
||||
expect(kill).not.toHaveBeenCalledWith('SIGTERM');
|
||||
});
|
||||
|
||||
it('stops enqueuing output to disk once the foreground cap trips', async () => {
|
||||
const sessionDir = mkdtempSync(join(tmpdir(), 'bpm-limit-'));
|
||||
try {
|
||||
const { manager, persistence } = createBackgroundManager({ sessionDir });
|
||||
let persistedChars = 0;
|
||||
const spy = vi
|
||||
.spyOn(persistence!, 'appendTaskOutput')
|
||||
.mockImplementation(async (_id: string, chunk: string) => {
|
||||
persistedChars += chunk.length;
|
||||
});
|
||||
|
||||
// 20 MiB, and the producer ignores SIGTERM so it keeps writing through
|
||||
// the whole grace window.
|
||||
const chunks = Array.from({ length: 20 }, () => 'x'.repeat(MiB));
|
||||
const { proc } = sigtermIgnoringProcess(chunks);
|
||||
|
||||
const taskId = manager.registerTask(
|
||||
new ProcessBackgroundTask(proc, 'runaway', 'hash', () => {}),
|
||||
{ detached: false, signal: new AbortController().signal, timeoutMs: 60_000 },
|
||||
);
|
||||
|
||||
const info = await waitForTerminal(manager, taskId);
|
||||
|
||||
expect(info?.status).toBe('killed');
|
||||
// Before the fix every chunk of the 20 MiB is enqueued into the disk
|
||||
// write chain (retaining each string until its write drains); afterwards
|
||||
// enqueuing stops at the ceiling so the chain cannot grow unbounded.
|
||||
expect(persistedChars).toBeLessThanOrEqual(17 * MiB);
|
||||
|
||||
spy.mockRestore();
|
||||
} finally {
|
||||
rmSync(sessionDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue