-
Notifications
You must be signed in to change notification settings - Fork 354
feat(bash): add yield_time_ms and explicit accesses #1140
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| --- | ||
| "@moonshot-ai/kimi-code": patch | ||
| --- | ||
|
|
||
| Let shell tool calls yield partial output before long-running commands finish. |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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<BashInput> { | |
| 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<BashInput> { | |
| }; | ||
| } | ||
|
|
||
| async executeShellCommand( | ||
| args: BashInput, | ||
| context: ExecutableToolContext, | ||
| ): Promise<ExecutableToolResult> { | ||
| return this.execution( | ||
| args, | ||
| context.signal, | ||
| context.onUpdate, | ||
| context.onForegroundTaskStart, | ||
| { autoYield: false }, | ||
| ); | ||
| } | ||
|
|
||
| private spawn(effectiveCwd: string, command: string): Promise<KaosProcess> { | ||
| const shellCwd = this.isWindowsBash ? windowsPathToPosixPath(effectiveCwd) : effectiveCwd; | ||
| const shellArgs = [ | ||
|
|
@@ -225,6 +259,7 @@ export class BashTool implements BuiltinTool<BashInput> { | |
| signal: AbortSignal, | ||
| onUpdate?: ((update: ToolUpdate) => void) | undefined, | ||
| onForegroundTaskStart?: ((taskId: string) => void) | undefined, | ||
| options: { readonly autoYield?: boolean } = {}, | ||
| ): Promise<ExecutableToolResult> { | ||
| const validationError = this.validateRunRequest(args, signal); | ||
| if (validationError !== undefined) return validationError; | ||
|
|
@@ -262,7 +297,7 @@ export class BashTool implements BuiltinTool<BashInput> { | |
| 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<BashInput> { | |
| } | ||
|
|
||
| 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; | ||
|
Comment on lines
+351
to
+355
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a foreground Bash call reaches Useful? React with 👍 / 👎. |
||
| const partialResult = builder.ok(''); | ||
| const wallSeconds = ((args.yield_time_ms ?? 10_000) / 1000).toFixed(1); | ||
| const partialOutput = partialResult.output; | ||
| return { | ||
| isError: false, | ||
|
Comment on lines
+359
to
+360
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When this return path resolves, the shell process is still running, but the loop releases the call's Useful? React with 👍 / 👎. |
||
| output: | ||
| (partialOutput.length > 0 ? partialOutput + '\n\n' : '') + | ||
| `<system>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.</system>`, | ||
| }; | ||
| } | ||
| if (release === 'detached') { | ||
| collectForegroundOutput = false; | ||
| return this.backgroundStartedResult( | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When this BashTool is used by
ToolManager.runShellCommandfor the TUI!path, that caller resolves Bash with only{ command, timeout: SHELL_FOREGROUND_TIMEOUT_S }and only treats results starting withtask_id:as backgrounded. With this new 10s default, any manual shell command that runs longer than 10 seconds returns through the yielded branch, but its handle is embedded inresult.outputas a<system>...task_id...message whilerunShellCommandreports only the streamed stdout/stderr and drops the handle, leaving the process running without a visible way for the user to poll or cancel it. Please make yielding opt-in for this path or teach the shell-command caller to handle yielded results.Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed in 7e275e1: manual
!shell commands now call Bash through an internal no-auto-yield execution path, so they wait for completion, timeout, cancellation, or explicit ctrl+b detach instead of returning a hidden yieldedtask_id. Added a regression test inpackages/agent-core/test/agent/context.test.ts.