From 6a48f9445666b13255f42b978d0ea818011bfb7e Mon Sep 17 00:00:00 2001 From: fanhuanjie Date: Sat, 22 Aug 2026 10:32:23 +0800 Subject: [PATCH] fix(tool): enforce max_output_tokens in command translator - Add max_output_tokens parameter to bash and exec_command tools - Truncate output by token budget (bytes*4 approximation) instead of line limits - Append tool storage path after merge conflict annotations - Include head+tail truncation message with original token count - Update tool-script description and declarations - Add tests for token-based truncation and exec_command passthrough --- packages/opencode/src/tool/bash.ts | 70 ++++++++++++++----- packages/opencode/src/tool/tool-script.ts | 14 +++- packages/opencode/src/tool/tool-script.txt | 2 +- packages/opencode/test/tool/bash.test.ts | 27 +++++++ .../opencode/test/tool/tool-script.test.ts | 28 +++++--- 5 files changed, 111 insertions(+), 30 deletions(-) diff --git a/packages/opencode/src/tool/bash.ts b/packages/opencode/src/tool/bash.ts index 41b5917db..604858048 100644 --- a/packages/opencode/src/tool/bash.ts +++ b/packages/opencode/src/tool/bash.ts @@ -144,6 +144,14 @@ const GIT_DESTRUCTIVE = new Map>([ const Parameters = z.object({ command: z.string().describe("The command to execute"), timeout: z.number().describe("Optional timeout in milliseconds").optional(), + max_output_tokens: z + .number() + .int() + .positive() + .describe( + "Maximum approximate tokens returned inline. Full output is saved to tool storage when this limit is exceeded.", + ) + .optional(), workdir: z .string() .describe( @@ -353,6 +361,14 @@ function head(text: string, maxLines: number, maxBytes: number): string { return out.join("\n") } +function headBytes(text: string, maxBytes: number) { + const buf = Buffer.from(text, "utf-8") + if (buf.length <= maxBytes) return text + let end = maxBytes + while (end > 0 && (buf[end] & 0xc0) === 0x80) end-- + return buf.subarray(0, end).toString("utf-8") +} + function tail(text: string, maxLines: number, maxBytes: number) { const lines = text.split("\n") if (lines.length <= maxLines && Buffer.byteLength(text, "utf-8") <= maxBytes) { @@ -687,17 +703,20 @@ export const BashTool = Tool.define( cwd: string env: NodeJS.ProcessEnv timeout: number + maxOutputTokens?: number description: string }, ctx: Tool.Context, ) { - const bytes = Truncate.MAX_BYTES - const lines = Truncate.MAX_LINES + const bytes = input.maxOutputTokens ? input.maxOutputTokens * 4 : Truncate.MAX_BYTES + const lines = input.maxOutputTokens ? Number.MAX_SAFE_INTEGER : Truncate.MAX_LINES const keep = bytes * 2 let full = "" + let first = "" let last = "" const list: Chunk[] = [] let used = 0 + let total = 0 let file = "" let sink: ReturnType | undefined let cut = false @@ -718,6 +737,10 @@ export const BashTool = Tool.define( yield* Effect.forkScoped( Stream.runForEach(Stream.decodeText(handle.all), (chunk) => { const size = Buffer.byteLength(chunk, "utf-8") + total += size + if (input.maxOutputTokens && Buffer.byteLength(first, "utf-8") < Math.floor(bytes / 2)) { + first = headBytes(first + chunk, Math.floor(bytes / 2)) + } list.push({ text: chunk, size }) used += size while (used > keep && list.length > 1) { @@ -845,24 +868,30 @@ export const BashTool = Tool.define( if (!output) output = "(no output)" if (cut && file) { - // Check if tail contains error patterns — if so, prepend head for context - const tailScan = end.text.length > 2048 ? end.text.slice(-2048) : end.text - const hasErrors = ERROR_PATTERN.test(tailScan) - if (hasErrors) { - let fileContent: string | undefined - try { - fileContent = readFileSync(file, "utf-8") - } catch { - fileContent = undefined - } - if (fileContent) { - const headText = head(fileContent, HEAD_LINES, HEAD_BYTES) - output = `...output truncated (head+tail shown due to errors)...\n\nFull output saved to: ${file}\n\n${headText}\n\n...middle omitted...\n\n${end.text}` + if (input.maxOutputTokens) { + const suffix = tail(raw, Number.MAX_SAFE_INTEGER, bytes - Buffer.byteLength(first, "utf-8")).text + const shown = Buffer.byteLength(first, "utf-8") + Buffer.byteLength(suffix, "utf-8") + output = `Warning: truncated output (original token count: ${Math.ceil(total / 4)})\n\n${first}\n…${Math.ceil((total - shown) / 4)} tokens truncated…\n${suffix}` + } else { + // Check if tail contains error patterns — if so, prepend head for context + const tailScan = end.text.length > 2048 ? end.text.slice(-2048) : end.text + const hasErrors = ERROR_PATTERN.test(tailScan) + if (hasErrors) { + let fileContent: string | undefined + try { + fileContent = readFileSync(file, "utf-8") + } catch { + fileContent = undefined + } + if (fileContent) { + const headText = head(fileContent, HEAD_LINES, HEAD_BYTES) + output = `...output truncated (head+tail shown due to errors)...\n\nFull output saved to: ${file}\n\n${headText}\n\n...middle omitted...\n\n${end.text}` + } else { + output = `...output truncated...\n\nFull output saved to: ${file}\n\n` + output + } } else { output = `...output truncated...\n\nFull output saved to: ${file}\n\n` + output } - } else { - output = `...output truncated...\n\nFull output saved to: ${file}\n\n` + output } } @@ -874,14 +903,16 @@ export const BashTool = Tool.define( // unmerged paths, the result itself carries the rule (the conflict belongs // to the branch's owner) and the two literal commands that follow it — // because the model reads a tool result before its next tool call, and does - // not re-read a system prompt assembled requests ago. Appended LAST so it is - // the final thing in the result, and never blocking: see the module header. + // not re-read a system prompt assembled requests ago. Appended after command + // output and never blocking; a tool-storage pointer may follow it so a + // truncated result always ends with the address of its complete output. output += yield* MergeConflict.annotate({ git: gitSvc, cwd: input.cwd, command: input.command, output, }) + if (cut && file && input.maxOutputTokens) output += `\n\nFull output saved to: ${file}` if (sink) { const stream = sink yield* Effect.promise( @@ -1008,6 +1039,7 @@ export const BashTool = Tool.define( cwd, env: yield* shellEnv(ctx, cwd), timeout, + maxOutputTokens: params.max_output_tokens, description: params.description, }, ctx, diff --git a/packages/opencode/src/tool/tool-script.ts b/packages/opencode/src/tool/tool-script.ts index 461264946..ddd8cfc3f 100644 --- a/packages/opencode/src/tool/tool-script.ts +++ b/packages/opencode/src/tool/tool-script.ts @@ -30,6 +30,7 @@ const MAX_CODE_BYTES = 128 * 1024 const MAX_FILE_BYTES = 10 * 1024 * 1024 const TRACE_TAIL_ENTRIES = 20 const EXEC_COMMAND_DEFAULT_YIELD_TIME_MS = 10_000 +const EXEC_COMMAND_DEFAULT_MAX_OUTPUT_TOKENS = 10_000 const ExecCommandParameters = z.object({ cmd: z.string().describe("Shell command to execute."), @@ -38,7 +39,15 @@ const ExecCommandParameters = z.object({ .int() .min(1) .optional() - .describe(`Wait budget in milliseconds before the command is terminated. Defaults to ${EXEC_COMMAND_DEFAULT_YIELD_TIME_MS} ms.`), + .describe( + `Wait budget in milliseconds before the command is terminated. Defaults to ${EXEC_COMMAND_DEFAULT_YIELD_TIME_MS} ms.`, + ), + max_output_tokens: z + .number() + .int() + .positive() + .optional() + .describe(`Output token budget. Defaults to ${EXEC_COMMAND_DEFAULT_MAX_OUTPUT_TOKENS} tokens.`), workdir: z .string() .optional() @@ -46,13 +55,14 @@ const ExecCommandParameters = z.object({ }) const EXEC_COMMAND_DESCRIPTION = - "Runs a shell command through the permission-gated bash executor. `cmd` is required; `yield_time_ms` is optional and defaults to 10000 ms." + "Runs a shell command through the permission-gated bash executor. `yield_time_ms` and `max_output_tokens` default to 10000. Output exceeding the token budget is saved to tool storage." function execCommandArgs(args: unknown) { const input = ExecCommandParameters.parse(args) return { command: input.cmd, timeout: input.yield_time_ms ?? EXEC_COMMAND_DEFAULT_YIELD_TIME_MS, + max_output_tokens: input.max_output_tokens ?? EXEC_COMMAND_DEFAULT_MAX_OUTPUT_TOKENS, workdir: input.workdir, description: input.cmd.length > 80 ? `${input.cmd.slice(0, 77)}...` : input.cmd, } diff --git a/packages/opencode/src/tool/tool-script.txt b/packages/opencode/src/tool/tool-script.txt index 05db35a94..615ef45b3 100644 --- a/packages/opencode/src/tool/tool-script.txt +++ b/packages/opencode/src/tool/tool-script.txt @@ -36,7 +36,7 @@ The script environment provides JavaScript built-ins plus `tools`, `files`, and - Call built-in tools, and other tool IDs that are valid JavaScript identifiers, with `await tools.name(input)`. - Search `ALL_TOOLS` names and descriptions for the capability you need, then call the exact catalog name with `await tools[name](input)`; the same permission checks as direct calls apply. Do not infer availability or MCP status from a name prefix. -- Use `await tools.exec_command({ cmd, yield_time_ms?, workdir? })` for shell commands. `yield_time_ms` defaults to 10000 ms. The runtime maps it to the `bash` tool, preserving its permissions, execution, timeout, and truncation behavior; commands still running when the wait budget expires are terminated. +- Use `await tools.exec_command({ cmd, yield_time_ms?, max_output_tokens?, workdir? })` for shell commands. `yield_time_ms` and `max_output_tokens` default to 10000. Output exceeding the token budget is truncated inline and saved to tool storage, with its full-output path appended to the result. The runtime maps the call to the `bash` tool, preserving its permissions and execution behavior; commands still running when the wait budget expires are terminated. - MCP results carry parsed `structuredContent` in the `structured` field when the server provides it. - Use `Promise.all` or `Promise.allSettled` only for independent calls; at most 8 run concurrently. - Return a small JSON-serializable aggregate. Circular values, BigInt, and throwing getters fail execution. diff --git a/packages/opencode/test/tool/bash.test.ts b/packages/opencode/test/tool/bash.test.ts index 726c2cd07..f78b980f8 100644 --- a/packages/opencode/test/tool/bash.test.ts +++ b/packages/opencode/test/tool/bash.test.ts @@ -1518,6 +1518,33 @@ describe("tool.bash truncation", () => { }) }) + test("limits output by approximate tokens and appends the tool storage path", async () => { + await Instance.provide({ + directory: projectRoot, + fn: async () => { + const bash = await initBash() + const result = await Effect.runPromise( + bash.execute( + { + command: fill("bytes", 5000), + description: "Generate output exceeding token limit", + max_output_tokens: 100, + }, + ctx, + ), + ) + mustTruncate(result) + + const filepath = (result.metadata as { outputPath?: string }).outputPath + expect(filepath).toBeTruthy() + expect(result.output).toContain("Warning: truncated output (original token count: 1250)") + expect(result.output).toContain("tokens truncated…") + expect(result.output.endsWith(`Full output saved to: ${filepath}`)).toBe(true) + expect(await Filesystem.readText(filepath!)).toBe("a".repeat(5000)) + }, + }) + }) + test("does not truncate small output", async () => { await Instance.provide({ directory: projectRoot, diff --git a/packages/opencode/test/tool/tool-script.test.ts b/packages/opencode/test/tool/tool-script.test.ts index c85242cb8..b76a11e87 100644 --- a/packages/opencode/test/tool/tool-script.test.ts +++ b/packages/opencode/test/tool/tool-script.test.ts @@ -332,10 +332,11 @@ describe("exec", () => { }) test("exec_command maps to bash while direct bash remains backward compatible", async () => { - const seen: Array<{ command: string; timeout: number; description: string }> = [] + const seen: Array<{ command: string; timeout: number; max_output_tokens?: number; description: string }> = [] const parameters = z.object({ command: z.string(), timeout: z.number(), + max_output_tokens: z.number().optional(), workdir: z.string().optional(), description: z.string(), }) @@ -344,14 +345,19 @@ describe("exec", () => { description: "fake bash", parameters, execute: (args) => { - seen.push({ command: args.command, timeout: args.timeout, description: args.description }) + seen.push({ + command: args.command, + timeout: args.timeout, + max_output_tokens: args.max_output_tokens, + description: args.description, + }) return Effect.succeed({ title: args.description, output: `ran:${args.command}`, metadata: {} }) }, } const result = await runToolScript( `return await Promise.all([ tools.bash({ command: "direct", timeout: 25000, description: "direct bash" }), - tools.exec_command({ cmd: "alias", yield_time_ms: 15000 }), + tools.exec_command({ cmd: "alias", yield_time_ms: 15000, max_output_tokens: 25000 }), ])`, [bash], ) @@ -360,15 +366,16 @@ describe("exec", () => { expect(result.output).toContain("ran:direct") expect(result.output).toContain("ran:alias") expect(seen).toEqual(expect.arrayContaining([ - { command: "direct", timeout: 25000, description: "direct bash" }, - { command: "alias", timeout: 15000, description: "alias" }, + { command: "direct", timeout: 25000, max_output_tokens: undefined, description: "direct bash" }, + { command: "alias", timeout: 15000, max_output_tokens: 25000, description: "alias" }, ])) }) - test("exec_command defaults yield_time_ms to 10000 ms", async () => { + test("exec_command defaults yield_time_ms and max_output_tokens to 10000", async () => { const parameters = z.object({ command: z.string(), timeout: z.number(), + max_output_tokens: z.number(), workdir: z.string().optional(), description: z.string(), }) @@ -377,12 +384,16 @@ describe("exec", () => { description: "fake bash", parameters, execute: (args) => - Effect.succeed({ title: args.description, output: String(args.timeout), metadata: {} }), + Effect.succeed({ + title: args.description, + output: `${args.timeout}:${args.max_output_tokens}`, + metadata: {}, + }), } const result = await runToolScript(`return await tools.exec_command({ cmd: "echo ok" })`, [bash]) expect(result.metadata.status).toBe("completed") - expect(result.output).toContain('"output": "10000"') + expect(result.output).toContain('"output": "10000:10000"') }) test("lists exec_command instead of bash in the code-mode catalog", async () => { @@ -692,6 +703,7 @@ describe("renderToolScriptDeclarations", () => { expect(text).toContain("Alias for bash") expect(text).toContain("cmd: string") expect(text).toContain("yield_time_ms?: number") + expect(text).toContain("max_output_tokens?: number") expect(text).not.toContain("command: string") expect(text).not.toContain("timeout?: number") expect(text).not.toContain("interactive?: boolean")