diff --git a/src/main/services/agent-workspace-tools.ts b/src/main/services/agent-workspace-tools.ts index 3ee80479..949e25f0 100644 --- a/src/main/services/agent-workspace-tools.ts +++ b/src/main/services/agent-workspace-tools.ts @@ -1,15 +1,13 @@ import { readdirSync, readFileSync, statSync } from "node:fs"; -import { spawn } from "node:child_process"; import { join } from "node:path"; import { resolveWorkspacePath } from "./files/workspace-path"; +import { runCommand } from "./workspace-command"; import { optionalString, requireString } from "./agent-tool-args"; import type { AgentToolContext } from "./agent-tools"; import type { ToolCall, ToolDefinition } from "./provider-tool-turn"; import { PAGE_TOOL_PROPERTIES, pageText, readPageArgs, type PageRequest } from "./text-paging"; const LIST_FILES_LIMIT = 400; -const COMMAND_TIMEOUT_MS = 120_000; -const COMMAND_OUTPUT_CHAR_LIMIT = 16_000; /** Dependency and VCS internals: huge, machine-managed, never what an agent is looking for. */ const LIST_SKIPPED_DIRECTORIES = new Set([".git", "node_modules", ".pnpm-store", "out", "dist"]); @@ -115,35 +113,3 @@ function readFile(workspaceRoot: string, relativePath: string, page: PageRequest if (!stats.isFile()) throw new Error(`${relativePosix} is not a file.`); return pageText(relativePosix, readFileSync(absolute, "utf8"), page); } - -function runCommand(workspaceRoot: string, command: string): Promise { - return new Promise((resolvePromise) => { - const child = spawn(command, { shell: true, cwd: workspaceRoot, windowsHide: true }); - let output = ""; - let killedByTimeout = false; - const timer = setTimeout(() => { - killedByTimeout = true; - child.kill(); - }, COMMAND_TIMEOUT_MS); - const collect = (chunk: Buffer): void => { - output += chunk.toString("utf8"); - }; - child.stdout.on("data", collect); - child.stderr.on("data", collect); - child.on("error", (error) => { - clearTimeout(timer); - resolvePromise(`ERROR: ${error.message}`); - }); - child.on("close", (code) => { - clearTimeout(timer); - const trimmed = output.length > COMMAND_OUTPUT_CHAR_LIMIT - ? `${output.slice(0, COMMAND_OUTPUT_CHAR_LIMIT / 2)}\n[... output truncated ...]\n${output.slice(-COMMAND_OUTPUT_CHAR_LIMIT / 2)}` - : output; - if (killedByTimeout) { - resolvePromise(`ERROR: command timed out after ${COMMAND_TIMEOUT_MS / 1000}s and was killed. Partial output:\n${trimmed}`); - return; - } - resolvePromise(`exit code ${code ?? 0}\n${trimmed.trim() || "(no output)"}`); - }); - }); -} diff --git a/src/main/services/workspace-command.ts b/src/main/services/workspace-command.ts new file mode 100644 index 00000000..41a90db0 --- /dev/null +++ b/src/main/services/workspace-command.ts @@ -0,0 +1,71 @@ +import { spawn } from "node:child_process"; + +const COMMAND_TIMEOUT_MS = 120_000; +const KEEP_BYTES = 8_000; + +/** Keeps the first and last KEEP_BYTES of a stream and counts what it drops. */ +export class BoundedOutput { + private head: Buffer = Buffer.alloc(0); + private tail: Buffer = Buffer.alloc(0); + private dropped = 0; + + push(chunk: Buffer): void { + let rest = chunk; + if (this.head.length < KEEP_BYTES) { + const room = KEEP_BYTES - this.head.length; + this.head = ownCopy(Buffer.concat([this.head, rest.subarray(0, room)])); + rest = rest.subarray(room); + } + if (rest.length === 0) return; + const joined = Buffer.concat([this.tail, rest]); + const overflow = Math.max(0, joined.length - KEEP_BYTES); + this.dropped += overflow; + this.tail = ownCopy(joined.subarray(overflow)); + } + + /** Bytes of memory the kept output occupies. */ + retainedBytes(): number { + return this.head.buffer.byteLength + this.tail.buffer.byteLength; + } + + text(): string { + if (this.dropped === 0) return Buffer.concat([this.head, this.tail]).toString("utf8"); + const gap = `\n[... ${this.dropped} bytes of output dropped ...]\n`; + return `${this.head.toString("utf8")}${gap}${this.tail.toString("utf8")}`; + } +} + +/** A standalone allocation, so a small view never pins a large chunk. */ +function ownCopy(view: Buffer): Buffer { + const copy = Buffer.allocUnsafeSlow(view.length); + view.copy(copy); + return copy; +} + +export function runCommand(workspaceRoot: string, command: string): Promise { + return new Promise((resolvePromise) => { + const child = spawn(command, { shell: true, cwd: workspaceRoot, windowsHide: true }); + const output = new BoundedOutput(); + let killedByTimeout = false; + const timer = setTimeout(() => { + killedByTimeout = true; + child.kill(); + }, COMMAND_TIMEOUT_MS); + const collect = (chunk: Buffer): void => output.push(chunk); + child.stdout.on("data", collect); + child.stderr.on("data", collect); + child.on("error", (error) => { + clearTimeout(timer); + resolvePromise(`ERROR: ${error.message}`); + }); + child.on("close", (code) => { + clearTimeout(timer); + const text = output.text(); + if (killedByTimeout) { + resolvePromise(`ERROR: command timed out after ${COMMAND_TIMEOUT_MS / 1000}s and was killed. Partial output:\n${text}`); + return; + } + resolvePromise(`exit code ${code ?? 0}\n${text.trim() || "(no output)"}`); + }); + }); +} diff --git a/tests/unit/workspace-command.test.ts b/tests/unit/workspace-command.test.ts new file mode 100644 index 00000000..02798614 --- /dev/null +++ b/tests/unit/workspace-command.test.ts @@ -0,0 +1,40 @@ +import { tmpdir } from "node:os"; +import { describe, expect, test } from "vitest"; +import { BoundedOutput, runCommand } from "../../src/main/services/workspace-command"; + +describe("bounded command output", () => { + test("holds at most head and tail while counting what it drops", () => { + const output = new BoundedOutput(); + output.push(Buffer.from("START-")); + output.push(Buffer.alloc(10_000_000, 120)); + expect(output.retainedBytes()).toBeLessThanOrEqual(16_000); + for (let index = 0; index < 1_000; index += 1) output.push(Buffer.alloc(10_000, 120)); + expect(output.retainedBytes()).toBeLessThanOrEqual(16_000); + output.push(Buffer.from("-END")); + const text = output.text(); + expect(text.startsWith("START-x")).toBe(true); + expect(text.endsWith("x-END")).toBe(true); + expect(text).toContain(`[... ${6 + 20_000_000 + 4 - 16_000} bytes of output dropped ...]`); + }); + + test("output that fits decodes as one piece, even across the head boundary", () => { + const output = new BoundedOutput(); + output.push(Buffer.from(`${"a".repeat(7_999)}邮END`, "utf8")); + expect(output.text()).toBe(`${"a".repeat(7_999)}邮END`); + }); + + test("small output comes back whole", () => { + const output = new BoundedOutput(); + output.push(Buffer.from("hello ")); + output.push(Buffer.from("world")); + expect(output.text()).toBe("hello world"); + }); + + test("a command printing megabytes returns its start, end, and the dropped count", async () => { + const result = await runCommand(tmpdir(), "node -e \"process.stdout.write('BEGIN' + 'a'.repeat(5000000) + 'FINISH')\""); + expect(result.startsWith("exit code 0\nBEGINaaa")).toBe(true); + expect(result.endsWith("aaaFINISH")).toBe(true); + expect(result).toContain("bytes of output dropped"); + expect(result.length).toBeLessThan(16_200); + }); +});