feat(agent-core-v2): add the WaitFor tool for waiting on background tasks

.changeset/wait-for-tool.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": minor
+---
+
+Add the WaitFor tool: the agent can now wait for a background task (sub-agent, background bash, or background question) inside the current turn — with an optional task ID and a required timeout of up to 600 seconds — instead of ending the turn and being re-invoked.

docs/en/reference/tools.md

@@ -99,20 +99,23 @@ Collaboration tools handle inter-Agent coordination, user interaction, and Skill
 
 ## Background Tasks
 
-Background task tools manage tasks started via `Bash`, `Agent`, or `AskUserQuestion`. When a task reaches a terminal state, its status and saved output path are automatically delivered back to the Agent; use `TaskOutput` to check progress early.
+Background task tools manage tasks started via `Bash`, `Agent`, or `AskUserQuestion`. When a task reaches a terminal state, its status and saved output path are automatically delivered back to the Agent; use `TaskOutput` to check progress early, or `WaitFor` to wait for a result inside the current turn.
 
 | Tool | Default Approval | Description |
 | --- | --- | --- |
 | `TaskList` | Auto-allow | List background tasks |
 | `TaskOutput` | Auto-allow | View the output of a background task |
 | `TaskStop` | Requires approval | Stop a running background task |
+| `WaitFor` | Auto-allow | Wait for background tasks to finish |
 
 **`TaskList`** returns the list of background tasks. Optional parameters: `active_only` (defaults to true; lists only running tasks) and `limit` (defaults to 20; range 1–100).
 
 **`TaskOutput`** returns the status and output of a task given its `task_id`. The inline preview includes at most the most recent 32 KB of content; the full log is saved to disk, and the tool also returns an `output_path` with a suggestion to use `Read` for paginated access. The call is always non-blocking — it returns the current snapshot immediately, and task completion is delivered via automatic notification.
 
 **`TaskStop`** accepts a `task_id` and optional `reason` (defaults to `Stopped by TaskStop`). Safe to call on tasks that are already in a terminal state.
 
+**`WaitFor`** suspends the current turn until a background task finishes or the timeout elapses. Parameters: `timeout` (required, in seconds, max 600) and optional `task_id`. Without `task_id`, the wait ends as soon as any background task that was running at call time finishes; when no background tasks are running, it returns immediately. A timeout is not an error — the result lists the tasks still running, and the Agent can wait again or do other work meanwhile. A task whose result was reported by `WaitFor` does not also produce an automatic completion notification.
+
 ## Scheduled Tasks
 
 Scheduled task tools allow the Agent to re-inject a prompt into the current session at a future time — either as a one-time reminder or as a recurring cron-triggered task (periodic checks, daily reports, deployment monitoring, etc.). Schedules are bound to the session and remain active when you resume it with `kimi --session`, but are not carried into a brand-new session. A single session can hold at most 50 active scheduled tasks. Set `KIMI_DISABLE_CRON=1` to disable them entirely; see [Environment Variables](../configuration/env-vars.md#runtime-switches).

docs/zh/reference/tools.md

@@ -99,20 +99,23 @@ Plan 模式是一种受约束的工作状态：进入后 `Write` 与 `Edit` 只
 
 ## 后台任务
 
-后台任务工具用于管理通过 `Bash`、`Agent` 或 `AskUserQuestion` 启动的后台任务。任务进入终止状态时会自动把状态和已保存的输出路径送回 Agent；如需提前检查进度，使用 `TaskOutput`。
+后台任务工具用于管理通过 `Bash`、`Agent` 或 `AskUserQuestion` 启动的后台任务。任务进入终止状态时会自动把状态和已保存的输出路径送回 Agent；如需提前检查进度，使用 `TaskOutput`；如果下一步必须等待某个任务的结果，使用 `WaitFor` 在当前轮次内等待。
 
 | 工具 | 默认审批 | 说明 |
 | --- | --- | --- |
 | `TaskList` | 自动放行 | 列出后台任务 |
 | `TaskOutput` | 自动放行 | 查看后台任务的输出 |
 | `TaskStop` | 需审批 | 停止正在运行的后台任务 |
+| `WaitFor` | 自动放行 | 等待后台任务结束 |
 
 **`TaskList`** 返回后台任务列表。可选参数 `active_only`（默认 true，仅列出运行中的任务）和 `limit`（默认 20，取值范围 1–100）。
 
 **`TaskOutput`** 根据 `task_id` 返回任务状态与输出。内联预览最多包含最近 32 KB 的内容；完整日志保存在磁盘上，工具会一并返回 `output_path` 并提示通过 `Read` 分页读取。该调用始终是非阻塞的——立即返回当前快照，任务完成会通过自动通知送达。
 
 **`TaskStop`** 接受 `task_id` 和可选的 `reason`（默认 `Stopped by TaskStop`）。对已处于终止状态的任务也能安全调用。
 
+**`WaitFor`** 把当前轮次挂起，直到后台任务结束或超时。参数：`timeout`（必填，单位秒，上限 600）和可选的 `task_id`。不传 `task_id` 时，调用时刻运行中的任意一个后台任务结束即返回；当前没有运行中的后台任务时立即返回。超时不是错误——结果会列出仍在运行的任务，Agent 可以再次等待，也可以先处理其他工作。已通过 `WaitFor` 汇报结果的任务不会再推送自动完成通知。
+
 ## 定时任务
 
 定时任务工具允许 Agent 把一段 prompt 在未来某个时间重新注入到当前会话——既可以是一次性提醒，也可以是按 cron 周期触发的任务（定期巡检、每日报表、部署监控等）。计划绑定到会话，用 `kimi --session` 恢复会话后仍然有效，但不会带入全新的会话。单个会话最多保留 50 个生效中的定时任务。设置 `KIMI_DISABLE_CRON=1` 可整体禁用，详见[环境变量](../configuration/env-vars.md#运行时开关)。

packages/agent-core-v2/docs/state-manifest.d.ts

@@ -1516,7 +1516,7 @@ export interface AgentStateSnapshot {
     readonly terminalNotificationSuppressed?: boolean;
     readonly timeoutMs?: number;
   }>;
-  // replayable · durable · undoable — folds: ContextAppendMessage
+  // replayable · durable · undoable — folds: ContextAppendMessage, TaskWaitDelivered
   'task.notificationDelivery': readonly string[];
   'task.scheduledNotificationKeys': Set<string>;
   // src/agent/tokenCounting/tokenCountingOps.ts

packages/agent-core-v2/docs/wire-manifest.d.ts

@@ -24,7 +24,7 @@
 // cross-reducers), blobs (the folding states whose blob codec offloads inline
 // media to blob storage), owner (the source file declaring the class).
 
-// Index (48 record types)
+// Index (49 record types)
 //   config.update                      profile                                                               src/agent/profile/profileOps.ts
 //   context.append_loop_event          contextMemory, turn                                                   src/agent/contextMemory/contextEvents.ts
 //   context.append_message             contextMemory, goalForkNotice, plan, task.notificationDelivery, todo  src/agent/contextMemory/contextEvents.ts
@@ -58,6 +58,7 @@
 //   swarm_mode.exit                    contextMemory, swarm                                                  src/features/swarm/swarmOps.ts
 //   task.started                       task                                                                  src/agent/task/taskOps.ts
 //   task.terminated                    task                                                                  src/agent/task/taskOps.ts
+//   task.waitDelivered                 task.notificationDelivery                                             src/agent/task/taskOps.ts
 //   token_counting.measured            tokenCounting                                                         src/agent/tokenCounting/tokenCountingOps.ts
 //   token_counting.rebased             tokenCounting                                                         src/agent/tokenCounting/tokenCountingOps.ts
 //   token_counting.truncated           tokenCounting                                                         src/agent/tokenCounting/tokenCountingOps.ts
@@ -495,6 +496,15 @@ interface TaskTerminatedPayload {
   outputTail?: string;
 }
 
+/**
+ * states: task.notificationDelivery
+ * owner: src/agent/task/taskOps.ts
+ */
+interface TaskWaitDeliveredPayload {
+  _name: 'task.waitDelivered';
+  keys: string[];
+}
+
 /**
  * states: tokenCounting
  * owner: src/agent/tokenCounting/tokenCountingOps.ts
@@ -732,6 +742,7 @@ interface WirePayloadMap {
   "swarm_mode.exit": SwarmModeExitPayload;
   "task.started": TaskStartedPayload;
   "task.terminated": TaskTerminatedPayload;
+  "task.waitDelivered": TaskWaitDeliveredPayload;
   "token_counting.measured": TokenCountingMeasuredPayload;
   "token_counting.rebased": TokenCountingRebasedPayload;
   "token_counting.truncated": TokenCountingTruncatedPayload;

packages/agent-core-v2/src/agent/goal/goalService.ts

@@ -144,7 +144,11 @@ const GOAL_CONTINUATION_PROMPT = [
   'validation, would benefit from clarification, or needs more goal turns. Once the 3-turn',
   'threshold is met and you cannot make meaningful progress without user input or an',
   'external-state change, call UpdateGoal with `blocked`; do not keep reporting the blocker while',
-  'leaving the goal active. Do not ask the user for input unless a real blocker prevents progress.',
+  'leaving the goal active. If you are waiting for background sub-agents or bash tasks to',
+  'finish, call WaitFor to wait for them inside this turn instead of ending the turn; ending',
+  'the turn just gets you re-invoked again and again. You can also use the waiting time to do',
+  'useful parallel work. Either way, make sure every goal turn is productive. Do not ask the',
+  'user for input unless a real blocker prevents progress.',
 ].join(' ');
 
 const GOAL_STEP_CAP_CONTINUATION_PROMPT = [

packages/agent-core-v2/src/agent/goal/injection/goal-active-reminder.md

@@ -11,4 +11,4 @@ ${budgets_block}${budget_guidance}
 
 Before doing any goal work, check the objective and latest request for a clear hard budget limit. If one is present and the current goal does not already record that limit, call SetGoalBudget first. Do not invent budgets. If a requested budget is not reasonable, do not set it; tell the user it is not reasonable.
 
-Goal mode is iterative. Keep the self-audit brief each turn. Do not explore unrelated interpretations once the goal can be decided. If the objective is simple, already answered, impossible, unsafe, or contradictory, do not run another goal turn. Explain briefly if useful, then call UpdateGoal with `complete` or `blocked` in the same turn. Otherwise, choose one bounded, useful slice of work toward the objective. Do not try to finish a broad goal in one turn unless the whole goal is genuinely small. Most goal turns should not call UpdateGoal: after completing a useful slice, if material work remains, end the turn normally without calling UpdateGoal so the runtime can continue the goal in the next turn. Call UpdateGoal with `complete` only when all required work is done, any stated validation has passed, and there is no useful next action. Completion audit: before calling `complete`, verify the current state against the actual objective and every explicit requirement. Treat weak or indirect evidence as not complete. Do not mark complete after only producing a plan, summary, first pass, or partial result. Do not mark complete merely because a budget is nearly exhausted or you want to stop. Blocked audit: do not call UpdateGoal with `blocked` the first time you hit a blocker. Use `blocked` only for a genuine impasse: an external condition, required user input, missing credentials or permissions, or a persistent technical failure. For those non-terminal blockers, the same blocking condition must repeat for at least 3 consecutive goal turns before you call `blocked`, counting the original/user-triggered turn and automatic continuations. If a previously blocked goal is resumed, treat the resumed run as a fresh blocked audit. Exception: if the objective itself is impossible, unsafe, or contradictory, call UpdateGoal with `blocked` in the same turn; do not run more goal turns just to satisfy the audit. Do not use `blocked` because the work is large, hard, slow, uncertain, incomplete, still needs validation, would benefit from clarification, or needs more goal turns. Once the 3-turn threshold is met and you cannot make meaningful progress without user input or an external-state change, call UpdateGoal with `blocked`; do not keep reporting the blocker while leaving the goal active.
+Goal mode is iterative. Keep the self-audit brief each turn. Do not explore unrelated interpretations once the goal can be decided. If the objective is simple, already answered, impossible, unsafe, or contradictory, do not run another goal turn. Explain briefly if useful, then call UpdateGoal with `complete` or `blocked` in the same turn. Otherwise, choose one bounded, useful slice of work toward the objective. Do not try to finish a broad goal in one turn unless the whole goal is genuinely small. Most goal turns should not call UpdateGoal: after completing a useful slice, if material work remains, end the turn normally without calling UpdateGoal so the runtime can continue the goal in the next turn. Call UpdateGoal with `complete` only when all required work is done, any stated validation has passed, and there is no useful next action. Completion audit: before calling `complete`, verify the current state against the actual objective and every explicit requirement. Treat weak or indirect evidence as not complete. Do not mark complete after only producing a plan, summary, first pass, or partial result. Do not mark complete merely because a budget is nearly exhausted or you want to stop. Blocked audit: do not call UpdateGoal with `blocked` the first time you hit a blocker. Use `blocked` only for a genuine impasse: an external condition, required user input, missing credentials or permissions, or a persistent technical failure. For those non-terminal blockers, the same blocking condition must repeat for at least 3 consecutive goal turns before you call `blocked`, counting the original/user-triggered turn and automatic continuations. If a previously blocked goal is resumed, treat the resumed run as a fresh blocked audit. Exception: if the objective itself is impossible, unsafe, or contradictory, call UpdateGoal with `blocked` in the same turn; do not run more goal turns just to satisfy the audit. Do not use `blocked` because the work is large, hard, slow, uncertain, incomplete, still needs validation, would benefit from clarification, or needs more goal turns. Once the 3-turn threshold is met and you cannot make meaningful progress without user input or an external-state change, call UpdateGoal with `blocked`; do not keep reporting the blocker while leaving the goal active. If you are waiting for background sub-agents or bash tasks to finish, call WaitFor to wait for them inside this turn instead of ending the turn; ending the turn just gets you re-invoked again and again. You can also use the waiting time to do useful parallel work. Either way, make sure every goal turn is productive.

packages/agent-core-v2/src/agent/permissionPolicy/policies/default-tool-approve.ts

@@ -13,6 +13,7 @@ const DEFAULT_APPROVE_TOOLS = new Set([
   'TodoList',
   'TaskList',
   'TaskOutput',
+  'WaitFor',
   'CronList',
   'WebSearch',
   'FetchURL',

packages/agent-core-v2/src/agent/task/task.ts

@@ -65,6 +65,11 @@ export interface AgentTaskNotificationContext {
   readonly sourceId: string;
 }
 
+export interface AgentTaskWaitDelivery {
+  readonly taskId: string;
+  readonly status: AgentTaskStatus;
+}
+
 export interface IAgentTaskService {
   readonly _serviceBrand: undefined;
 
@@ -79,6 +84,7 @@ export interface IAgentTaskService {
   ): Promise<AgentTaskOutputSnapshot>;
   readOutput(taskId: string, tail?: number): Promise<string>;
   suppressTerminalNotification(taskId: string): Promise<void>;
+  markTasksDeliveredViaWait(tasks: readonly AgentTaskWaitDelivery[]): void;
   detach(taskId: string): AgentTaskInfo | undefined;
   stop(taskId: string, reason?: string): Promise<AgentTaskInfo | undefined>;
   stopByUser(taskId: string): Promise<AgentTaskInfo | undefined>;

packages/agent-core-v2/src/agent/task/taskOps.ts

@@ -47,6 +47,15 @@ export class TaskNotified extends Event2<AgentTaskNotificationContext> {
 }
 export interface TaskNotified extends AgentTaskNotificationContext {}
 
+const taskWaitDeliveredSchema = z.object({ keys: z.array(z.string()) });
+
+export class TaskWaitDelivered extends Event2<z.infer<typeof taskWaitDeliveredSchema>> {
+  static override readonly type = 'task.waitDelivered';
+  static override readonly durable = true;
+  static override readonly schema = taskWaitDeliveredSchema;
+}
+export interface TaskWaitDelivered extends z.infer<typeof taskWaitDeliveredSchema> {}
+
 export const taskKey = defineState('task', (): TaskModelState => new Map()).replayable({
   schema: z.custom<TaskModelState>(),
 })

packages/agent-core-v2/src/agent/task/taskService.ts

@@ -52,16 +52,18 @@ import {
   type AgentTaskOutputSnapshot,
   type AgentTaskStatus,
   type AgentTaskTrackOptions,
+  type AgentTaskWaitDelivery,
   type ForegroundTaskReleaseReason,
   type IAgentTaskEntry,
   type RegisterAgentTaskOptions,
 } from './task';
 import { resolveAgentTaskConfig } from './configSection';
 import { AgentTaskPersistence } from './persist';
-import { taskKey, TaskNotified, TaskStarted, TaskTerminated } from './taskOps';
+import { taskKey, TaskNotified, TaskStarted, TaskTerminated, TaskWaitDelivered } from './taskOps';
 import { formatTaskList } from '#/agent/tools/task/task-list/taskListTool';
 import '#/agent/tools/task/task-output/taskOutputTool';
 import '#/agent/tools/task/task-stop/taskStopTool';
+import '#/agent/tools/task/task-wait/taskWaitTool';
 
 interface ForegroundRelease {
   readonly promise: Promise<ForegroundTaskReleaseReason>;
@@ -100,6 +102,13 @@ export const taskNotificationDeliveryKey = defineState(
     if (!s.includes(key)) {
       s.push(key);
     }
+  })
+  .on(TaskWaitDelivered, (s, e) => {
+    for (const key of e.keys) {
+      if (!s.includes(key)) {
+        s.push(key);
+      }
+    }
   });
 
 interface ManagedTask {
@@ -604,6 +613,23 @@ export class AgentTaskService extends Disposable implements IAgentTaskService {
     if (ghost !== undefined) return;
   }
 
+  markTasksDeliveredViaWait(tasks: readonly AgentTaskWaitDelivery[]): void {
+    if (tasks.length === 0) return;
+    const keys: string[] = [];
+    for (const { taskId, status } of tasks) {
+      const origin: TaskNotificationOrigin = {
+        taskId,
+        status,
+        notificationId: taskNotificationId(taskId, status),
+      };
+      const key = notificationKey(origin);
+      this.pendingNotificationRequests.get(key)?.abort();
+      this.markDeliveredNotification(origin);
+      keys.push(key);
+    }
+    void this.dispatcher.dispatch(new TaskWaitDelivered({ keys }));
+  }
+
   detach(taskId: string): AgentTaskInfo | undefined {
     const entry = this.tasks.get(taskId);
     if (entry === undefined) return this.ghosts.get(taskId);
@@ -1063,6 +1089,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService {
     const context = await this.buildAgentTaskNotificationContext(info);
     if (context === undefined) return;
     const key = notificationKey(context.origin);
+    if (this.deliveredNotificationKeys.has(key)) return;
     const request = new TaskNotificationStepRequest(
       {
         role: 'user',
@@ -1125,7 +1152,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService {
       kind: 'task',
       taskId: info.taskId,
       status: info.status,
-      notificationId: `task:${info.taskId}:${info.status}`,
+      notificationId: taskNotificationId(info.taskId, info.status),
     };
     const key = notificationKey(origin);
     if (this.buildingNotificationKeys.has(key)) return undefined;
@@ -1345,6 +1372,10 @@ function isTaskOrigin(origin: unknown): origin is TaskNotificationOrigin {
   );
 }
 
+function taskNotificationId(taskId: string, status: string): string {
+  return `task:${taskId}:${status}`;
+}
+
 function notificationKey(origin: TaskNotificationOrigin): string {
   return `${origin.taskId}\0${origin.status}\0${origin.notificationId}`;
 }