diff --git a/App/memmy-agent/src/core/agent-runtime/tools/exec-session.ts b/App/memmy-agent/src/core/agent-runtime/tools/exec-session.ts index 96724ad3c..acfb887e1 100644 --- a/App/memmy-agent/src/core/agent-runtime/tools/exec-session.ts +++ b/App/memmy-agent/src/core/agent-runtime/tools/exec-session.ts @@ -11,6 +11,8 @@ export const DEFAULT_MAX_OUTPUT_CHARS = 10_000; export const MAX_OUTPUT_CHARS = 50_000; const STDIO_DRAIN_GRACE_MS = 1000; const PROCESS_TERMINATE_GRACE_MS = 1000; +// Batch progress output without turning every stream chunk into a tool return. +const OUTPUT_COALESCE_MS = 1000; export const EXEC_LIFECYCLE_ERROR = "Error: Command process exited but descendants kept stdout/stderr open; the managed process tree was terminated."; @@ -41,10 +43,6 @@ export type ExecSessionInfo = { returncode: number | null; }; -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - function processExited(child: ChildProcess): boolean { return child.exitCode != null || child.signalCode != null; } @@ -169,18 +167,17 @@ export async function drainExitedProcessStreams({ streamsClosed, forceCloseStreams, options, + abortSignal, }: { child: ChildProcess; closePromise: Promise; streamsClosed: () => boolean; forceCloseStreams: () => void; options?: ManagedProcessOptions; + abortSignal?: AbortSignal | null; }): Promise { if (streamsClosed()) return null; - const closedDuringDrain = await Promise.race([ - closePromise.then(() => true), - sleep(STDIO_DRAIN_GRACE_MS).then(() => streamsClosed()), - ]); + const closedDuringDrain = await waitForDrain(closePromise, streamsClosed, abortSignal); if (closedDuringDrain || streamsClosed()) return null; await terminateManagedProcessTree(child, options); if (!streamsClosed()) forceCloseStreams(); @@ -201,20 +198,27 @@ function throwIfAborted(signal?: AbortSignal | null): void { if (signal?.aborted) throw createToolAbortError(); } -function abortableSleep(ms: number, signal?: AbortSignal | null): Promise { - if (ms <= 0) return Promise.resolve(); +function waitForDrain( + closePromise: Promise, + streamsClosed: () => boolean, + signal?: AbortSignal | null, +): Promise { throwIfAborted(signal); - if (!signal) return sleep(ms); return new Promise((resolve, reject) => { - const timer = setTimeout(() => { - signal.removeEventListener("abort", onAbort); - resolve(); - }, ms); - const onAbort = () => { + let settled = false; + const finish = (closed: boolean, error?: Error) => { + if (settled) return; + settled = true; clearTimeout(timer); - reject(createToolAbortError()); + signal?.removeEventListener("abort", onAbort); + if (error) reject(error); + else resolve(closed); }; - signal.addEventListener("abort", onAbort, { once: true }); + const onAbort = () => finish(false, createToolAbortError()); + const timer = setTimeout(() => finish(streamsClosed()), STDIO_DRAIN_GRACE_MS); + signal?.addEventListener("abort", onAbort, { once: true }); + void closePromise.then(() => finish(true)); + if (signal?.aborted) onAbort(); }); } @@ -241,6 +245,7 @@ export class ExecSession { deadlineMs: number; lastAccess = Date.now(); private chunks: string[] = []; + private outputStartedAt: number | null = null; private timedOut = false; private stderrStarted = false; private readonly stdoutDecoder: CommandOutputDecoder; @@ -251,6 +256,7 @@ export class ExecSession { private lifecycleError: string | null = null; private readonly closePromise: Promise; private resolveClose!: () => void; + private readonly outputWaiters = new Set<() => void>(); constructor({ command, @@ -296,7 +302,7 @@ export class ExecSession { args.push("-c", command); this.process = spawnProcess(resolvedShell, args, { cwd, env, detached: true }); } - this.process.on("error", (error) => this.chunks.push(`Error executing command: ${error.message}`)); + this.process.on("error", (error) => this.appendStdout(`Error executing command: ${error.message}`)); this.process.stdout.on("data", (chunk: Buffer) => { this.stdoutDecoder.push(chunk); this.appendStdout(this.stdoutDecoder.read()); @@ -348,7 +354,7 @@ export class ExecSession { ): Promise { this.lastAccess = Date.now(); throwIfAborted(abortSignal); - if (yieldTimeMs > 0 && !this.isDone()) await abortableSleep(Math.min(yieldTimeMs, MAX_YIELD_MS), abortSignal); + if (yieldTimeMs > 0) await this.waitForActivity(Math.min(yieldTimeMs, MAX_YIELD_MS), abortSignal); throwIfAborted(abortSignal); if (!this.isDone() && Date.now() >= this.deadlineMs) { this.timedOut = true; @@ -360,11 +366,13 @@ export class ExecSession { closePromise: this.closePromise, streamsClosed: () => this.streamsClosed, forceCloseStreams: () => this.forceCloseStreams(), + abortSignal, }); } throwIfAborted(abortSignal); const raw = this.chunks.join(""); this.chunks = []; + this.outputStartedAt = null; const [output, truncated] = truncateOutput(raw, maxOutputChars); const elapsed = Math.max(0, (Date.now() - this.startedAt) / 1000); const done = this.streamsClosed; @@ -381,14 +389,66 @@ export class ExecSession { }; } + private waitForActivity(ms: number, signal?: AbortSignal | null): Promise { + throwIfAborted(signal); + if (this.isDone() || this.streamsClosed) return Promise.resolve(); + return new Promise((resolve, reject) => { + let settled = false; + let outputTimer: ReturnType | undefined; + const finish = (error?: Error) => { + if (settled) return; + settled = true; + clearTimeout(waitTimer); + clearTimeout(outputTimer); + this.outputWaiters.delete(onOutput); + this.process.removeListener("exit", onDone); + this.process.removeListener("close", onDone); + signal?.removeEventListener("abort", onAbort); + if (error) reject(error); + else resolve(); + }; + const onDone = () => finish(); + const onAbort = () => finish(createToolAbortError()); + const onOutput = () => { + if (this.isDone() || this.streamsClosed) { + finish(); + } else if (this.chunks.length > 0 && outputTimer === undefined) { + // Anchor at the first readable output, never debounce continuous logs. + const remaining = Math.max(0, (this.outputStartedAt ?? Date.now()) + OUTPUT_COALESCE_MS - Date.now()); + if (remaining === 0) finish(); + else outputTimer = setTimeout(onDone, remaining); + } + }; + const waitMs = Math.max(0, Math.min(ms, this.deadlineMs - Date.now())); + const waitTimer = setTimeout(onDone, waitMs); + this.process.once("exit", onDone); + this.process.once("close", onDone); + this.outputWaiters.add(onOutput); + signal?.addEventListener("abort", onAbort, { once: true }); + // Recheck after subscribing so completion, cancellation and buffered + // output cannot fall between checking state and installing listeners. + if (signal?.aborted) onAbort(); + else onOutput(); + }); + } + + private notifyOutputWaiters(): void { + for (const wake of this.outputWaiters) wake(); + } + private appendStdout(text: string): void { - if (text) this.chunks.push(text); + if (!text) return; + this.chunks.push(text); + this.outputStartedAt ??= Date.now(); + this.notifyOutputWaiters(); } private appendStderr(text: string): void { if (!text) return; this.chunks.push(this.stderrStarted ? text : `STDERR:\n${text}`); + this.outputStartedAt ??= Date.now(); this.stderrStarted = true; + this.notifyOutputWaiters(); } private finishStdoutOnce(): void { @@ -412,6 +472,7 @@ export class ExecSession { if (this.streamsClosed) return; this.streamsClosed = true; this.resolveClose(); + this.notifyOutputWaiters(); } } diff --git a/App/memmy-agent/tests/core/agent-runtime/tools/exec-session-wait.test.ts b/App/memmy-agent/tests/core/agent-runtime/tools/exec-session-wait.test.ts new file mode 100644 index 000000000..620a9bb65 --- /dev/null +++ b/App/memmy-agent/tests/core/agent-runtime/tools/exec-session-wait.test.ts @@ -0,0 +1,276 @@ +import { EventEmitter } from "node:events"; +import { PassThrough } from "node:stream"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { ExecSession, type SessionPoll } from "../../../../src/core/agent-runtime/tools/exec-session.js"; + +const children: any[] = []; + +function childProcess() { + const child = new EventEmitter() as any; + child.stdin = new PassThrough(); + child.stdout = new PassThrough(); + child.stderr = new PassThrough(); + child.exitCode = null; + child.signalCode = null; + child.kill = vi.fn(() => { + child.signalCode = "SIGTERM"; + child.stdout.destroy(); + child.stderr.destroy(); + child.emit("exit", null, "SIGTERM"); + close(child); + return true; + }); + children.push(child); + return child; +} + +function close(child: any) { + child.stdout.end(); + child.stderr.end(); + child.emit("close", child.exitCode, child.signalCode); +} + +function exit(child: any, code = 0) { + child.exitCode = code; + child.emit("exit", code, null); + close(child); +} + +function sessionFor(child: any, timeoutS?: number) { + return new ExecSession({ command: "controlled child", timeoutS, spawnProcess: (() => child) as any }); +} + +function observe(pending: Promise) { + const state: { settled: boolean; poll?: SessionPoll; error?: Error } = { settled: false }; + void pending.then((poll) => { state.settled = true; state.poll = poll; }, (error) => { state.settled = true; state.error = error; }); + return state; +} + +beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(0); +}); + +afterEach(async () => { + for (const child of children.splice(0)) exit(child); + await vi.runOnlyPendingTimersAsync(); + vi.clearAllTimers(); + vi.useRealTimers(); +}); + +describe("exec session waiting", () => { + it("returns when a silent process completes instead of waiting 30 seconds", async () => { + const child = childProcess(); + const state = observe(sessionFor(child).poll(30_000, 10_000)); + await vi.advanceTimersByTimeAsync(100); + exit(child); + await vi.advanceTimersByTimeAsync(0); + expect(state.poll).toMatchObject({ done: true, exitCode: 0, output: "", elapsedS: 0.1 }); + expect(vi.getTimerCount()).toBe(0); + expect(child.listenerCount("exit")).toBe(0); + }); + + it("wakes on signal exit with a null exit code", async () => { + const child = childProcess(); + const state = observe(sessionFor(child).poll(30_000, 10_000)); + await vi.advanceTimersByTimeAsync(50); + child.kill(); + await vi.advanceTimersByTimeAsync(0); + expect(state.poll).toMatchObject({ done: true, exitCode: null, elapsedS: 0.05 }); + expect(vi.getTimerCount()).toBe(0); + }); + + it("returns immediately if the process has already completed", async () => { + const child = childProcess(); + const session = sessionFor(child); + child.stdout.write("finished\n"); + exit(child); + expect(await session.poll(30_000, 10_000)).toMatchObject({ done: true, output: "finished\n", elapsedS: 0 }); + expect(vi.getTimerCount()).toBe(0); + }); + + it("coalesces stdout and stderr for one second from the first readable output", async () => { + const child = childProcess(); + const session = sessionFor(child); + const state = observe(session.poll(30_000, 10_000)); + await vi.advanceTimersByTimeAsync(100); + child.stdout.write("first\n"); + await vi.advanceTimersByTimeAsync(500); + child.stderr.write("second\n"); + await vi.advanceTimersByTimeAsync(499); + expect(state.settled).toBe(false); + await vi.advanceTimersByTimeAsync(1); + expect(state.poll).toMatchObject({ done: false, output: "first\nSTDERR:\nsecond\n", elapsedS: 1.1 }); + expect(vi.getTimerCount()).toBe(0); + expect(child.listenerCount("exit")).toBe(0); + expect(child.listenerCount("close")).toBe(1); + child.stdout.write("third\n"); + expect((await session.poll(0, 10_000)).output).toBe("third\n"); + expect((await session.poll(0, 10_000)).output).toBe(""); + }); + + it("also coalesces output that was buffered before polling", async () => { + const child = childProcess(); + const session = sessionFor(child); + child.stdout.write("buffered\n"); + const state = observe(session.poll(30_000, 10_000)); + await vi.advanceTimersByTimeAsync(999); + expect(state.settled).toBe(false); + await vi.advanceTimersByTimeAsync(1); + expect(state.poll).toMatchObject({ done: false, output: "buffered\n", elapsedS: 1 }); + }); + + it("returns stderr-only output without waiting the full yield", async () => { + const child = childProcess(); + const session = sessionFor(child); + const state = observe(session.poll(30_000, 10_000)); + child.stderr.write("diagnostic\n"); + await vi.advanceTimersByTimeAsync(1000); + expect(state.poll).toMatchObject({ done: false, output: "STDERR:\ndiagnostic\n" }); + }); + + it("does not restart an expired output window when a later poll begins", async () => { + const child = childProcess(); + const session = sessionFor(child); + child.stdout.write("ready-before-poll\n"); + await vi.advanceTimersByTimeAsync(1200); + expect(await session.poll(30_000, 10_000)).toMatchObject({ done: false, output: "ready-before-poll\n", elapsedS: 1.2 }); + expect(vi.getTimerCount()).toBe(0); + }); + + it("does not restart the coalescing window for continuous small chunks", async () => { + const child = childProcess(); + const state = observe(sessionFor(child).poll(30_000, 10_000)); + for (let i = 1; i <= 80; i++) { + await vi.advanceTimersByTimeAsync(10); + child.stdout.write(`chunk-${i}\n`); + expect(state.settled).toBe(false); + } + await vi.advanceTimersByTimeAsync(209); + expect(state.settled).toBe(false); + await vi.advanceTimersByTimeAsync(1); + expect(state.poll?.output).toBe(Array.from({ length: 80 }, (_, i) => `chunk-${i + 1}\n`).join("")); + expect(state.poll?.elapsedS).toBe(1.01); + expect(vi.getTimerCount()).toBe(0); + }); + + it("respects a shorter yield instead of extending it for coalescing", async () => { + const child = childProcess(); + const state = observe(sessionFor(child).poll(200, 10_000)); + await vi.advanceTimersByTimeAsync(100); + child.stdout.write("progress\n"); + await vi.advanceTimersByTimeAsync(100); + expect(state.poll).toMatchObject({ done: false, output: "progress\n", elapsedS: 0.2 }); + expect(vi.getTimerCount()).toBe(0); + }); + + it("returns at the yield deadline when no output arrives", async () => { + const state = observe(sessionFor(childProcess()).poll(1000, 10_000)); + await vi.advanceTimersByTimeAsync(999); + expect(state.settled).toBe(false); + await vi.advanceTimersByTimeAsync(1); + expect(state.poll).toMatchObject({ done: false, output: "", elapsedS: 1 }); + }); + + it("keeps zero-yield reads immediate and preserves the output limit", async () => { + const child = childProcess(); + const session = sessionFor(child); + child.stdout.write("a".repeat(2000)); + expect(await session.poll(0, 1000)).toMatchObject({ done: false, truncatedChars: 1000, elapsedS: 0 }); + expect((await session.poll(0, 1000)).output).toBe(""); + expect(vi.getTimerCount()).toBe(0); + }); + + it("bypasses coalescing on completion and waits for the last stream bytes", async () => { + const child = childProcess(); + const session = sessionFor(child); + const state = observe(session.poll(30_000, 10_000)); + child.stdout.write("head|"); + await vi.advanceTimersByTimeAsync(100); + child.exitCode = 0; + child.emit("exit", 0, null); + await vi.advanceTimersByTimeAsync(0); + expect(state.settled).toBe(false); + await vi.advanceTimersByTimeAsync(50); + child.stdout.write("tail\n"); + child.stderr.write("final-error\n"); + close(child); + await vi.advanceTimersByTimeAsync(0); + expect(state.poll).toMatchObject({ done: true, output: "head|tail\nSTDERR:\nfinal-error\n", elapsedS: 0.15 }); + expect((await session.poll(0, 10_000)).output).toBe(""); + expect(vi.getTimerCount()).toBe(0); + }); + + it("cancels during coalescing and preserves unread output for a subsequent poll", async () => { + const child = childProcess(); + const session = sessionFor(child); + const controller = new AbortController(); + const state = observe(session.poll(30_000, 10_000, {}, controller.signal)); + child.stdout.write("keep-this\n"); + await vi.advanceTimersByTimeAsync(50); + controller.abort(); + await vi.advanceTimersByTimeAsync(0); + expect(state.error?.name).toBe("AbortError"); + expect(child.kill).not.toHaveBeenCalled(); + expect(vi.getTimerCount()).toBe(0); + expect(child.listenerCount("exit")).toBe(0); + expect((await session.poll(0, 10_000)).output).toBe("keep-this\n"); + }); + + it("cancels while draining an exited process without consuming its final output", async () => { + const child = childProcess(); + const session = sessionFor(child); + const controller = new AbortController(); + child.exitCode = 0; + child.emit("exit", 0, null); + const state = observe(session.poll(30_000, 10_000, {}, controller.signal)); + await vi.advanceTimersByTimeAsync(50); + controller.abort(); + await vi.advanceTimersByTimeAsync(0); + expect(state.error?.name).toBe("AbortError"); + expect(vi.getTimerCount()).toBe(0); + child.stdout.write("late-tail\n"); + close(child); + expect(await session.poll(0, 10_000)).toMatchObject({ done: true, output: "late-tail\n" }); + }); + + it("rejects an already-cancelled poll without installing wait listeners", async () => { + const child = childProcess(); + const controller = new AbortController(); + controller.abort(); + await expect(sessionFor(child).poll(30_000, 10_000, {}, controller.signal)).rejects.toMatchObject({ name: "AbortError" }); + expect(vi.getTimerCount()).toBe(0); + expect(child.listenerCount("exit")).toBe(0); + }); + + it("honors the command timeout before a longer yield expires", async () => { + const child = childProcess(); + const state = observe(sessionFor(child, 0.1).poll(30_000, 10_000)); + await vi.advanceTimersByTimeAsync(100); + expect(state.poll).toMatchObject({ done: true, timedOut: true, elapsedS: 0.1 }); + expect(child.kill).toHaveBeenCalledTimes(1); + expect(vi.getTimerCount()).toBe(0); + }); + + it("wakes when a spawn failure closes streams even without an exit code", async () => { + const child = childProcess(); + const state = observe(sessionFor(child).poll(30_000, 10_000)); + child.emit("error", new Error("spawn failed")); + close(child); + await vi.advanceTimersByTimeAsync(0); + expect(state.poll).toMatchObject({ done: true, exitCode: null, output: "Error executing command: spawn failed" }); + expect(vi.getTimerCount()).toBe(0); + }); + + it("settles once when completion and the yield deadline coincide", async () => { + const child = childProcess(); + const session = sessionFor(child); + setTimeout(() => exit(child), 1000); + const settled = vi.fn(); + const pending = session.poll(1000, 10_000).then((poll) => { settled(); return poll; }); + await vi.advanceTimersByTimeAsync(1000); + expect(await pending).toMatchObject({ done: true, exitCode: 0 }); + expect(settled).toHaveBeenCalledTimes(1); + expect(vi.getTimerCount()).toBe(0); + }); +});