From 8b06072b1f6e6ce3d7e73c192ee0c9e9ac8d86c6 Mon Sep 17 00:00:00 2001 From: star Date: Sat, 27 Jun 2026 19:47:33 +0800 Subject: [PATCH] feat(bash): add yield_time_ms and explicit accesses --- .changeset/fresh-tools-march.md | 5 + .../agent-core/src/agent/background/index.ts | 3 +- packages/agent-core/src/agent/tool/index.ts | 33 +++-- .../src/tools/builtin/shell/bash.md | 2 + .../src/tools/builtin/shell/bash.ts | 68 ++++++++- .../agent-core/test/agent/context.test.ts | 65 ++++++++- packages/agent-core/test/tools/bash.test.ts | 129 +++++++++++++++++- 7 files changed, 289 insertions(+), 16 deletions(-) create mode 100644 .changeset/fresh-tools-march.md diff --git a/.changeset/fresh-tools-march.md b/.changeset/fresh-tools-march.md new file mode 100644 index 000000000..2d155f7c1 --- /dev/null +++ b/.changeset/fresh-tools-march.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Let shell tool calls yield partial output before long-running commands finish. diff --git a/packages/agent-core/src/agent/background/index.ts b/packages/agent-core/src/agent/background/index.ts index 6f7aaed3e..f87da675b 100644 --- a/packages/agent-core/src/agent/background/index.ts +++ b/packages/agent-core/src/agent/background/index.ts @@ -439,10 +439,11 @@ export class BackgroundManager { return this.toInfo(entry); } - persistOutput(taskId: string): void { + async persistOutput(taskId: string): Promise { const entry = this.tasks.get(taskId); if (entry === undefined) return; this.startOutputPersist(entry); + await Promise.all([entry.outputWriteQueue, this.persistLive(entry)]); } /** Stop a running task. SIGTERM → 5s grace → SIGKILL. */ diff --git a/packages/agent-core/src/agent/tool/index.ts b/packages/agent-core/src/agent/tool/index.ts index f16577253..290a39439 100644 --- a/packages/agent-core/src/agent/tool/index.ts +++ b/packages/agent-core/src/agent/tool/index.ts @@ -4,7 +4,7 @@ import picomatch from 'picomatch'; import type { Agent } from '..'; import { makeErrorPayload } from '../../errors'; -import type { ExecutableTool, ToolUpdate } from '../../loop'; +import type { ExecutableTool, ExecutableToolContext, ToolUpdate } from '../../loop'; import { createMcpAuthTool } from '../../mcp/auth-tool'; import type { McpConnectionManager, McpServerEntry } from '../../mcp'; import { mcpResultToExecutableOutput } from '../../mcp/output'; @@ -113,14 +113,7 @@ export class ToolManager { const controller = new AbortController(); if (commandId !== undefined) this.shellCommandControllers.set(commandId, controller); try { - const execution = await bash.resolveExecution({ command, timeout: SHELL_FOREGROUND_TIMEOUT_S }); - if (!('execute' in execution)) { - const output = - typeof execution.output === 'string' ? execution.output : 'Command failed.'; - this.agent.context.appendBashOutput('', output); - return { stdout: '', stderr: output, isError: true }; - } - const result = await execution.execute({ + const toolContext: ExecutableToolContext = { turnId: '', toolCallId: 'shell-command', signal: controller.signal, @@ -140,7 +133,13 @@ export class ToolManager { this.agent.emitEvent({ type: 'shell.started', commandId, taskId }); } }, - }); + }; + const result = await (bash instanceof b.BashTool + ? bash.executeShellCommand( + { command, timeout: SHELL_FOREGROUND_TIMEOUT_S }, + toolContext, + ) + : this.executeShellCommandViaTool(bash, command, toolContext)); isError = result.isError === true; // Detached to background (ctrl+b): the BashTool returns the background @@ -179,6 +178,20 @@ export class ToolManager { return { stdout, stderr, isError }; } + private async executeShellCommandViaTool( + bash: BuiltinTool, + command: string, + context: ExecutableToolContext, + ) { + const execution = await bash.resolveExecution({ command, timeout: SHELL_FOREGROUND_TIMEOUT_S }); + if (!('execute' in execution)) { + const output = + typeof execution.output === 'string' ? execution.output : 'Command failed.'; + return { output, isError: true as const }; + } + return execution.execute(context); + } + cancelShellCommand(commandId: string): void { this.shellCommandControllers.get(commandId)?.abort(); } diff --git a/packages/agent-core/src/tools/builtin/shell/bash.md b/packages/agent-core/src/tools/builtin/shell/bash.md index 7a6b97dd3..1eee3fa1a 100644 --- a/packages/agent-core/src/tools/builtin/shell/bash.md +++ b/packages/agent-core/src/tools/builtin/shell/bash.md @@ -15,6 +15,8 @@ The stdout and stderr will be combined and returned as a string. The output may If `run_in_background=true`, the command will be started as a background task and this tool will return a task ID instead of waiting for command completion. When doing that, you must provide a short `description`. Background commands default to a {{ DEFAULT_BACKGROUND_TIMEOUT_S }}s timeout and `timeout` is capped at {{ MAX_BACKGROUND_TIMEOUT_S }}s; set `disable_timeout=true` only when the task should run without a timeout. You will be automatically notified when the task completes. After starting one, default to returning control to the user instead of immediately waiting on it. Use `TaskOutput` for a non-blocking status/output snapshot, and only set `block=true` when you explicitly want to wait for completion. Use `TaskStop` only if the task must be cancelled. If a human user wants to inspect background tasks themselves, point them to the `/tasks` command, which opens an interactive panel; it has no subcommands. +**yield_time_ms:** Wait before yielding output. Defaults to 10000 ms; effective range is 250-30000 ms. The command continues running during this wait. If the command finishes within the yield window, the full result is returned immediately. Otherwise, partial output is returned with a `task_id` you can use to poll or stop the command. + **Guidelines for safety and security:** - Each shell tool call will be executed in a fresh shell environment. The shell variables, current working directory changes, and the shell history is not preserved between calls. To run a command in a particular directory, pass the `cwd` argument (or use absolute paths) rather than relying on a `cd` from an earlier call. - The tool call will return after the command is finished. You shall not use this tool to execute an interactive command or a command that may run forever. For possibly long-running foreground commands, set the `timeout` argument in seconds. Foreground commands default to {{ DEFAULT_TIMEOUT_S }}s and allow up to {{ MAX_TIMEOUT_S }}s. diff --git a/packages/agent-core/src/tools/builtin/shell/bash.ts b/packages/agent-core/src/tools/builtin/shell/bash.ts index 8198a9506..ebbb8e36e 100644 --- a/packages/agent-core/src/tools/builtin/shell/bash.ts +++ b/packages/agent-core/src/tools/builtin/shell/bash.ts @@ -23,11 +23,18 @@ */ import type { Kaos, KaosProcess } from '@moonshot-ai/kaos'; +import { sleep } from '@antfu/utils'; import { z } from 'zod'; import { ProcessBackgroundTask, type BackgroundManager } from '../../../agent/background'; import type { BuiltinTool } from '../../../agent/tool'; -import type { ExecutableToolResult, ToolExecution, ToolUpdate } from '../../../loop/types'; +import { ToolAccesses } from '../../../loop/tool-access'; +import type { + ExecutableToolContext, + ExecutableToolResult, + ToolExecution, + ToolUpdate, +} from '../../../loop/types'; import { renderPrompt } from '../../../utils/render-prompt'; import { toInputJsonSchema } from '../../support/input-schema'; import { literalRulePattern, matchesGlobRuleSubject } from '../../support/rule-match'; @@ -78,6 +85,15 @@ export const BashInputSchema = z .describe( 'If true, do not apply a timeout to the command. Only applies when run_in_background is true.', ), + yield_time_ms: z + .number() + .int() + .min(250) + .max(30000) + .optional() + .describe( + 'Wait before yielding output. Defaults to 10000 ms; effective range is 250-30000 ms. The command continues running during this wait.', + ), }) .superRefine((val, ctx) => { if (val.timeout === undefined) return; @@ -147,6 +163,10 @@ function withoutBackgroundDescription(description: string): string { .replace( /\r?\n- Prefer `run_in_background=true`[\s\S]*?conversation to continue before the command finishes\./, '\n- Do not set `run_in_background=true`; background task management tools are not available.', + ) + .replace( + /\n\n\*\*yield_time_ms:\*\*[\s\S]*?Otherwise, partial output is returned with a `task_id` you can use to poll or stop the command\./, + '\n\n**yield_time_ms:** Background task tools are disabled for this agent, so foreground commands do not yield a `task_id`; they wait for completion, timeout, detach, or cancellation.', ); } @@ -176,6 +196,7 @@ export class BashTool implements BuiltinTool { resolveExecution(args: BashInput): ToolExecution { const preview = args.command.length > 50 ? `${args.command.slice(0, 50)}…` : args.command; return { + accesses: ToolAccesses.all(), description: args.run_in_background ? `Starting background: ${preview}` : `Running: ${preview}`, @@ -193,6 +214,19 @@ export class BashTool implements BuiltinTool { }; } + async executeShellCommand( + args: BashInput, + context: ExecutableToolContext, + ): Promise { + return this.execution( + args, + context.signal, + context.onUpdate, + context.onForegroundTaskStart, + { autoYield: false }, + ); + } + private spawn(effectiveCwd: string, command: string): Promise { const shellCwd = this.isWindowsBash ? windowsPathToPosixPath(effectiveCwd) : effectiveCwd; const shellArgs = [ @@ -225,6 +259,7 @@ export class BashTool implements BuiltinTool { signal: AbortSignal, onUpdate?: ((update: ToolUpdate) => void) | undefined, onForegroundTaskStart?: ((taskId: string) => void) | undefined, + options: { readonly autoYield?: boolean } = {}, ): Promise { const validationError = this.validateRunRequest(args, signal); if (validationError !== undefined) return validationError; @@ -262,7 +297,7 @@ export class BashTool implements BuiltinTool { onUpdate?.({ kind, text }); builder.write(text); if (!foregroundOutputPersisted && builder.truncated && foregroundTaskId !== undefined) { - this.backgroundManager.persistOutput(foregroundTaskId); + void this.backgroundManager.persistOutput(foregroundTaskId); foregroundOutputPersisted = true; } }; @@ -303,7 +338,34 @@ export class BashTool implements BuiltinTool { } try { - const release = await this.backgroundManager.waitForForegroundRelease(taskId); + const completionOrDetach = this.backgroundManager.waitForForegroundRelease(taskId); + // Only yield a foreground command when the model can actually manage the + // resulting task. Subagent profiles can expose Bash without TaskOutput or + // TaskStop, so returning a task_id there would strand the command result. + const release = this.allowBackground && options.autoYield !== false + ? await Promise.race([ + completionOrDetach, + sleep(args.yield_time_ms ?? 10_000).then(() => 'yielded' as const), + ]) + : await completionOrDetach; + if (release === 'yielded') { + // Command is still running after yield_time_ms. Return partial output + // with wall_time_seconds so the caller knows how long we waited. + await this.backgroundManager.persistOutput(taskId); + collectForegroundOutput = false; + const partialResult = builder.ok(''); + const wallSeconds = ((args.yield_time_ms ?? 10_000) / 1000).toFixed(1); + const partialOutput = partialResult.output; + return { + isError: false, + output: + (partialOutput.length > 0 ? partialOutput + '\n\n' : '') + + `Command still running after ${wallSeconds}s yield. ` + + `task_id: ${taskId}\n` + + `Use TaskOutput(task_id="${taskId}", block=false) to poll for more output, ` + + `or TaskStop(task_id="${taskId}") to cancel.`, + }; + } if (release === 'detached') { collectForegroundOutput = false; return this.backgroundStartedResult( diff --git a/packages/agent-core/test/agent/context.test.ts b/packages/agent-core/test/agent/context.test.ts index 580bda69c..ce6b7f4f6 100644 --- a/packages/agent-core/test/agent/context.test.ts +++ b/packages/agent-core/test/agent/context.test.ts @@ -1,4 +1,4 @@ -import { Readable, type Writable } from 'node:stream'; +import { PassThrough, Readable, type Writable } from 'node:stream'; import type { KaosProcess } from '@moonshot-ai/kaos'; import type { Message } from '@moonshot-ai/kosong'; @@ -133,6 +133,69 @@ describe('Agent context', () => { expect(textOf(ctx.agent.context.history[1]!)).toContain('hello'); }); + it('does not auto-yield manual shell commands after the Bash yield window', async () => { + vi.useFakeTimers(); + try { + const stdout = new PassThrough(); + const stderr = new PassThrough(); + let exitCode: number | null = null; + let resolveWait: (code: number) => void = () => {}; + const waitPromise = new Promise((resolve) => { + resolveWait = resolve; + }); + const proc: KaosProcess = { + stdin: { end: vi.fn(), write: vi.fn() } as unknown as Writable, + stdout, + stderr, + pid: 2, + get exitCode(): number | null { + return exitCode; + }, + wait: vi.fn(async () => waitPromise), + kill: vi.fn(async () => {}), + dispose: vi.fn(async () => { + stdout.destroy(); + stderr.destroy(); + }), + }; + const execWithEnv = vi.fn().mockResolvedValue(proc); + const kaos = createFakeKaos({ execWithEnv }); + const ctx = testAgent({ kaos }); + ctx.configure(); + + const running = ctx.agent.tools.runShellCommand('sleep 20'); + let settled = false; + void running.finally(() => { + settled = true; + }); + + await vi.waitFor(() => { + expect(execWithEnv).toHaveBeenCalledTimes(1); + }); + await vi.advanceTimersByTimeAsync(10_000); + + expect(settled).toBe(false); + + stdout.write('done\n'); + stdout.end(); + stderr.end(); + exitCode = 0; + resolveWait(0); + + const result = await running; + + expect(result).toMatchObject({ + stdout: 'done\n', + stderr: '', + isError: false, + }); + expect(result.backgrounded).toBeUndefined(); + expect(result.stdout).not.toContain('task_id:'); + } finally { + vi.useRealTimers(); + } + }); + it('surfaces the failure reason when a shell command fails with no output', async () => { const fakeProcess = (exitCode: number): KaosProcess => { const out = Readable.from([]); diff --git a/packages/agent-core/test/tools/bash.test.ts b/packages/agent-core/test/tools/bash.test.ts index 682aa26f2..102c04199 100644 --- a/packages/agent-core/test/tools/bash.test.ts +++ b/packages/agent-core/test/tools/bash.test.ts @@ -1,4 +1,4 @@ -import { mkdtempSync, readFileSync, rmSync } from 'node:fs'; +import { existsSync, mkdtempSync, readFileSync, rmSync } from 'node:fs'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { PassThrough, Readable, type Writable } from 'node:stream'; @@ -328,6 +328,18 @@ describe('BashTool', () => { } }); + it('does not expose sandbox permissions before execution supports them', () => { + const tool = bashTool(createFakeKaos({ osEnv: posixEnv }), '/workspace'); + const properties = (tool.parameters as { properties: Record }).properties; + + expect(properties).not.toHaveProperty('sandbox_permissions'); + expect(properties).not.toHaveProperty('additional_permissions'); + expect(properties).not.toHaveProperty('justification'); + expect(tool.description).not.toContain('sandbox_permissions'); + expect(tool.description).not.toContain('additional_permissions'); + expect(tool.description).not.toContain('require_escalated'); + }); + it('exposes a default timeout in the JSON Schema', () => { const tool = bashTool(createFakeKaos({ osEnv: posixEnv }), '/workspace'); const properties = (tool.parameters as { properties: Record }) @@ -669,6 +681,107 @@ describe('BashTool', () => { }); }); + it('does not yield a foreground task when background tools are disabled', async () => { + vi.useFakeTimers(); + try { + const { proc, finish } = pendingProcess(); + const manager = createBackgroundManager().manager; + const tool = bashTool( + createFakeKaos({ + execWithEnv: vi.fn().mockResolvedValue(proc), + osEnv: posixEnv, + }), + '/workspace', + manager, + { allowBackground: false }, + ); + + const running = executeTool( + tool, + context({ command: 'sleep 10', timeout: 60, yield_time_ms: 250 }), + ); + let settled = false; + void running.finally(() => { + settled = true; + }); + + await vi.waitFor(() => { + expect(manager.list(false)).toHaveLength(1); + }); + await vi.advanceTimersByTimeAsync(250); + + expect(settled).toBe(false); + const task = manager.list(false)[0]!; + expect(task).toMatchObject({ detached: false }); + + (proc.stdout as PassThrough).write('done\n'); + finish(); + const result = await running; + + expect(result).toMatchObject({ + isError: false, + message: 'Command executed successfully.', + }); + expect(result.output).toContain('done\n'); + expect(result.output).not.toContain('task_id:'); + expect(result.output).not.toContain('TaskOutput'); + expect(result.output).not.toContain('TaskStop'); + } finally { + vi.useRealTimers(); + } + }); + + it('persists complete output after yielding a foreground task', async () => { + vi.useFakeTimers(); + const sessionDir = mkdtempSync(join(tmpdir(), 'bash-yield-persist-')); + try { + const { proc, finish } = pendingProcess(); + const manager = createBackgroundManager({ sessionDir }).manager; + const tool = bashTool( + createFakeKaos({ + execWithEnv: vi.fn().mockResolvedValue(proc), + osEnv: posixEnv, + }), + '/workspace', + manager, + ); + + const running = executeTool( + tool, + context({ command: 'sleep 10', timeout: 60, yield_time_ms: 250 }), + ); + await vi.waitFor(() => { + expect(manager.list(false)).toHaveLength(1); + }); + const task = manager.list(false)[0]!; + + (proc.stdout as PassThrough).write('before yield\n'); + await vi.advanceTimersByTimeAsync(250); + const result = await running; + + expect(result.output).toContain(`task_id: ${task.taskId}`); + expect(existsSync(join(sessionDir, 'tasks', `${task.taskId}.json`))).toBe(true); + + (proc.stdout as PassThrough).write('after yield\n'); + await vi.waitFor(async () => { + const snapshot = await manager.getOutputSnapshot(task.taskId, 0); + expect(snapshot.fullOutputAvailable).toBe(true); + expect(snapshot.outputPath).toBeTruthy(); + }); + + const snapshot = await manager.getOutputSnapshot(task.taskId, 0); + expect(readFileSync(snapshot.outputPath!, 'utf8')).toBe('before yield\nafter yield\n'); + + finish(); + await expect(manager.wait(task.taskId)).resolves.toMatchObject({ + status: 'completed', + }); + } finally { + rmSync(sessionDir, { recursive: true, force: true }); + vi.useRealTimers(); + } + }); + it('keeps task metadata independent when noisy foreground output is capped before detach', async () => { const { proc, finish } = pendingProcess(); const manager = createBackgroundManager().manager; @@ -1301,4 +1414,18 @@ describe('BashTool prompt / runtime consistency', () => { // (`Command failed with exit code: N`), never via a system tag. expect(tool.description).not.toMatch(/exit code will be provided in a system tag/); }); + + it('documents yield_time_ms without task polling when background tools are disabled', () => { + const tool = bashTool( + createFakeKaos({ osEnv: posixEnv }), + '/workspace', + createBackgroundManager().manager, + { allowBackground: false }, + ); + + expect(tool.description).toContain('**yield_time_ms:**'); + expect(tool.description).toContain('foreground commands do not yield a `task_id`'); + expect(tool.description).not.toContain('TaskOutput'); + expect(tool.description).not.toContain('TaskStop'); + }); });