From 8440801de47ddae29224430048e1228b80cde370 Mon Sep 17 00:00:00 2001 From: Luyu Cheng <2239547+chengluyu@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:26:59 +0800 Subject: [PATCH] feat(agent-core-v2): add the `WaitFor` tool for waiting on background tasks (#3060) * feat(agent-core-v2): add the WaitFor tool for waiting on background tasks * fix(agent-core-v2): mark WaitFor deliveries only after formatting succeeds * fix(agent-core-v2): cancel losing waits once the WaitFor race resolves * test(node-sdk): project WaitFor out of the v1-v2 resume parity roster * fix(agent-core-v2): gate WaitFor goal guidance behind the wait_for flag * fix(agent-core-v2): gate WaitFor goal guidance on actual tool availability * fix(agent-core-v2): enforce the wait_for flag at WaitFor execution time * fix(agent-core-v2): consult the live tool policy in the WaitFor availability check --- .changeset/wait-for-tool.md | 5 + docs/en/reference/tools.md | 5 +- docs/zh/reference/tools.md | 5 +- .../agent-core-v2/docs/state-manifest.d.ts | 2 +- .../agent-core-v2/docs/wire-manifest.d.ts | 13 +- .../src/agent/goal/goalService.ts | 23 +- .../goal/injection/goal-active-reminder.md | 2 +- .../src/agent/goal/injection/goalInjection.ts | 11 +- .../policies/default-tool-approve.ts | 1 + packages/agent-core-v2/src/agent/task/task.ts | 6 + .../agent-core-v2/src/agent/task/taskOps.ts | 9 + .../src/agent/task/taskService.ts | 35 +- .../src/agent/tools/task/task-wait/flag.ts | 16 + .../agent/tools/task/task-wait/task-wait.md | 16 + .../agent/tools/task/task-wait/task-wait.ts | 29 + .../tools/task/task-wait/taskWaitTool.ts | 274 +++++++++ .../agent-core-v2/src/app/telemetry/events.ts | 20 + .../src/features/tower/workerProfile.ts | 1 + packages/agent-core-v2/src/index.ts | 2 + .../agentLifecycle/profile/profiles.ts | 2 + .../fullCompaction/fullCompaction.test.ts | 14 +- .../test/agent/goal/goal.test.ts | 505 +++++++++++++++++ .../test/agent/loop/loop.test.ts | 4 +- .../test/agent/task/taskService.test.ts | 295 +++++++++- .../test/agent/task/tools/task-tools.test.ts | 528 +++++++++++++++++- .../test/agent/undo/undo.test.ts | 40 ++ packages/agent-core-v2/test/index.test.ts | 1 + .../os/backends/node-local/tools/bash.test.ts | 3 + packages/agent-core-v2/test/tool/tool.test.ts | 14 +- .../agent-core-v2/test/wire/resume.test.ts | 2 +- packages/node-sdk/test/v1-v2-parity.test.ts | 6 +- 31 files changed, 1843 insertions(+), 46 deletions(-) create mode 100644 .changeset/wait-for-tool.md create mode 100644 packages/agent-core-v2/src/agent/tools/task/task-wait/flag.ts create mode 100644 packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.md create mode 100644 packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.ts create mode 100644 packages/agent-core-v2/src/agent/tools/task/task-wait/taskWaitTool.ts diff --git a/.changeset/wait-for-tool.md b/.changeset/wait-for-tool.md new file mode 100644 index 000000000..f195b6500 --- /dev/null +++ b/.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. diff --git a/docs/en/reference/tools.md b/docs/en/reference/tools.md index cf16c9200..d3d67e050 100644 --- a/docs/en/reference/tools.md +++ b/docs/en/reference/tools.md @@ -99,13 +99,14 @@ 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). @@ -113,6 +114,8 @@ Background task tools manage tasks started via `Bash`, `Agent`, or `AskUserQuest **`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). diff --git a/docs/zh/reference/tools.md b/docs/zh/reference/tools.md index 83fbdb094..eead29c92 100644 --- a/docs/zh/reference/tools.md +++ b/docs/zh/reference/tools.md @@ -99,13 +99,14 @@ 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）。 @@ -113,6 +114,8 @@ Plan 模式是一种受约束的工作状态：进入后 `Write` 与 `Edit` 只 **`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#运行时开关)。 diff --git a/packages/agent-core-v2/docs/state-manifest.d.ts b/packages/agent-core-v2/docs/state-manifest.d.ts index 9cefad33c..3400b04ee 100644 --- a/packages/agent-core-v2/docs/state-manifest.d.ts +++ b/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; // src/agent/tokenCounting/tokenCountingOps.ts diff --git a/packages/agent-core-v2/docs/wire-manifest.d.ts b/packages/agent-core-v2/docs/wire-manifest.d.ts index a3ff3e2e9..7c944d95e 100644 --- a/packages/agent-core-v2/docs/wire-manifest.d.ts +++ b/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; diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index ef356fcf0..3d336f1c6 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -13,7 +13,7 @@ import { isPlainRecord } from '#/_base/utils/canonical-args'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { ContextAppendMessage } from '#/agent/contextMemory/contextEvents'; import type { ContextMessage, PromptOrigin } from '#/agent/contextMemory/types'; -import { GoalInjection } from '#/agent/goal/injection/goalInjection'; +import { GoalInjection, GOAL_WAIT_FOR_GUIDANCE } from '#/agent/goal/injection/goalInjection'; import { IAgentLoopService, type AfterStepContext, @@ -31,11 +31,14 @@ import { IAgentPermissionModeService } from '#/agent/permissionMode/permissionMo import type { PermissionMode } from '#/agent/permissionPolicy/types'; import { IAgentToolApprovalService } from '#/agent/toolApproval/toolApproval'; import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; +import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy'; import type { BeforeToolExecuteEvent } from '#/agent/toolExecutor/toolHooks'; import { IAgentUsageService, type UsageRecordedContext } from '#/agent/usage/usage'; import type { GoalBudgetProperties } from '#/app/telemetry/events'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import { IConfigService } from '#/app/config/config'; +import { IFlagService } from '#/app/flag/flag'; import { ErrorCodes, Error2, @@ -47,6 +50,7 @@ import { IEventDispatcher } from '#/state/eventDispatcher'; import { defineState } from '#/state/state'; import { IAgentGoalService, type GoalReasonInput, type ResumeGoalInput } from './goal'; +import { WAIT_FOR_FLAG_ID } from '#/agent/tools/task/task-wait/flag'; import { IGoalDeadlineScheduler } from './goalDeadlineScheduler'; import { GoalClear, @@ -263,10 +267,13 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { @IAgentContextInjectorService injector: IAgentContextInjectorService, @IAgentLoopService private readonly loopService: IAgentLoopService, @IAgentToolExecutorService toolExecutor: IAgentToolExecutorService, + @IAgentToolRegistryService private readonly toolRegistry: IAgentToolRegistryService, + @IAgentToolPolicyService private readonly toolPolicy: IAgentToolPolicyService, @IAgentToolApprovalService private readonly toolApproval: IAgentToolApprovalService, @IAgentPermissionModeService private readonly permissionMode: IAgentPermissionModeService, @IAgentUsageService usageService: IAgentUsageService, @IConfigService private readonly config: IConfigService, + @IFlagService private readonly flags: IFlagService, @IGoalDeadlineScheduler private readonly deadlineScheduler: IGoalDeadlineScheduler, @IAgentScopeContext private readonly agentContext: IAgentScopeContext, @IAgentStateService private readonly states: IAgentStateService, @@ -291,6 +298,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { new GoalInjection( { getGoal: () => this.getGoal().goal, + isWaitForEnabled: () => this.isWaitForAvailable(), }, injector, ), @@ -895,15 +903,26 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { } catch {} } + private isWaitForAvailable(): boolean { + return ( + this.flags.enabled(WAIT_FOR_FLAG_ID) && + this.toolRegistry.resolve('WaitFor') !== undefined && + this.toolPolicy.isToolActive('WaitFor') + ); + } + private launchContinuationTurn(goalId: string, stepCapped = false): void { if (!this.isActiveGoal(goalId)) return; if (this.pendingContinuation !== undefined) return; + const prompt = stepCapped ? GOAL_STEP_CAP_CONTINUATION_PROMPT : GOAL_CONTINUATION_PROMPT; const message: ContextMessage = { role: 'user', content: [ { type: 'text', - text: stepCapped ? GOAL_STEP_CAP_CONTINUATION_PROMPT : GOAL_CONTINUATION_PROMPT, + text: this.isWaitForAvailable() + ? `${prompt} ${GOAL_WAIT_FOR_GUIDANCE}` + : prompt, }, ], toolCalls: [], diff --git a/packages/agent-core-v2/src/agent/goal/injection/goal-active-reminder.md b/packages/agent-core-v2/src/agent/goal/injection/goal-active-reminder.md index 527367f56..a15375571 100644 --- a/packages/agent-core-v2/src/agent/goal/injection/goal-active-reminder.md +++ b/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.${wait_for_guidance} diff --git a/packages/agent-core-v2/src/agent/goal/injection/goalInjection.ts b/packages/agent-core-v2/src/agent/goal/injection/goalInjection.ts index 6b6d979e3..e6e6567b5 100644 --- a/packages/agent-core-v2/src/agent/goal/injection/goalInjection.ts +++ b/packages/agent-core-v2/src/agent/goal/injection/goalInjection.ts @@ -8,8 +8,12 @@ import GOAL_PAUSED_REMINDER from './goal-paused-reminder.md?raw'; export interface GoalInjectionOptions { readonly getGoal: () => GoalSnapshot | null; + readonly isWaitForEnabled?: () => boolean; } +export const GOAL_WAIT_FOR_GUIDANCE = + '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.'; + export class GoalInjection extends Service { constructor( private readonly options: GoalInjectionOptions, @@ -24,7 +28,9 @@ export class GoalInjection extends Service { private reminder(): string | undefined { const goal = this.options.getGoal(); if (goal === null) return undefined; - if (goal.status === 'active') return buildGoalReminder(goal); + if (goal.status === 'active') { + return buildGoalReminder(goal, this.options.isWaitForEnabled?.() === true); + } if (goal.status === 'blocked') return buildBlockedNote(goal); if (goal.status === 'paused') return buildPausedNote(goal); return undefined; @@ -52,7 +58,7 @@ function buildPausedNote(goal: GoalSnapshot): string { }); } -function buildGoalReminder(goal: GoalSnapshot): string { +function buildGoalReminder(goal: GoalSnapshot, waitForEnabled: boolean): string { const budgets = formatBudgets(goal); return renderPrompt(GOAL_ACTIVE_REMINDER, { objective: escapeUntrustedText(goal.objective), @@ -61,6 +67,7 @@ function buildGoalReminder(goal: GoalSnapshot): string { progress: `${goal.turnsUsed} continuation turns, ${goal.tokensUsed} tokens, ${formatElapsed(goal.wallClockMs)} elapsed`, budgets_block: budgets.length > 0 ? `Budgets: ${budgets}.\n` : '', budget_guidance: isNearingBudget(goal) ? BUDGET_GUIDANCE_NEARING : BUDGET_GUIDANCE_WITHIN, + wait_for_guidance: waitForEnabled ? ` ${GOAL_WAIT_FOR_GUIDANCE}` : '', }); } diff --git a/packages/agent-core-v2/src/agent/permissionPolicy/policies/default-tool-approve.ts b/packages/agent-core-v2/src/agent/permissionPolicy/policies/default-tool-approve.ts index 4867b0e00..a2b79a22b 100644 --- a/packages/agent-core-v2/src/agent/permissionPolicy/policies/default-tool-approve.ts +++ b/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', diff --git a/packages/agent-core-v2/src/agent/task/task.ts b/packages/agent-core-v2/src/agent/task/task.ts index dde02bee3..61e6c008c 100644 --- a/packages/agent-core-v2/src/agent/task/task.ts +++ b/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; readOutput(taskId: string, tail?: number): Promise; suppressTerminalNotification(taskId: string): Promise; + markTasksDeliveredViaWait(tasks: readonly AgentTaskWaitDelivery[]): void; detach(taskId: string): AgentTaskInfo | undefined; stop(taskId: string, reason?: string): Promise; stopByUser(taskId: string): Promise; diff --git a/packages/agent-core-v2/src/agent/task/taskOps.ts b/packages/agent-core-v2/src/agent/task/taskOps.ts index b8bf3e8e9..5e1c45141 100644 --- a/packages/agent-core-v2/src/agent/task/taskOps.ts +++ b/packages/agent-core-v2/src/agent/task/taskOps.ts @@ -47,6 +47,15 @@ export class TaskNotified extends Event2 { } export interface TaskNotified extends AgentTaskNotificationContext {} +const taskWaitDeliveredSchema = z.object({ keys: z.array(z.string()) }); + +export class TaskWaitDelivered extends Event2> { + static override readonly type = 'task.waitDelivered'; + static override readonly durable = true; + static override readonly schema = taskWaitDeliveredSchema; +} +export interface TaskWaitDelivered extends z.infer {} + export const taskKey = defineState('task', (): TaskModelState => new Map()).replayable({ schema: z.custom(), }) diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index 1f2d908e1..d64d69998 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/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; @@ -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}`; } diff --git a/packages/agent-core-v2/src/agent/tools/task/task-wait/flag.ts b/packages/agent-core-v2/src/agent/tools/task/task-wait/flag.ts new file mode 100644 index 000000000..dd42774b7 --- /dev/null +++ b/packages/agent-core-v2/src/agent/tools/task/task-wait/flag.ts @@ -0,0 +1,16 @@ +import { type FlagDefinitionInput, registerFlagDefinition } from '#/app/flag/flagRegistry'; + +export const WAIT_FOR_FLAG_ID = 'wait_for'; +export const WAIT_FOR_FLAG_ENV = 'KIMI_CODE_EXPERIMENTAL_WAIT_FOR'; + +export const waitForFlag: FlagDefinitionInput = { + id: WAIT_FOR_FLAG_ID, + title: 'WaitFor tool', + description: + 'Give the model the WaitFor tool so it can wait for background tasks inside the current turn instead of ending the turn and being re-invoked.', + env: WAIT_FOR_FLAG_ENV, + default: true, + surface: 'core', +}; + +registerFlagDefinition(waitForFlag); diff --git a/packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.md b/packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.md new file mode 100644 index 000000000..30ebbc8fa --- /dev/null +++ b/packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.md @@ -0,0 +1,16 @@ +Wait for background tasks to finish without ending the current turn. + +Use this when your next step depends on the result of a running background task (a sub-agent, a background bash command, or a background AskUserQuestion). The call suspends inside the current turn until the task finishes or the timeout elapses, then returns the outcome so you can keep working in the same turn. While waiting, no LLM requests are made. + +Guidelines: + +- Do not call WaitFor right after dispatching work whose result you do not need yet — finished background tasks notify you automatically. WaitFor is for the moment you genuinely cannot proceed without a result. +- `timeout` is required, in seconds, capped at 600. To wait longer, call WaitFor again; waking up periodically also lets you re-evaluate the situation. +- A timeout is not an error: the result lists the tasks that are still running, and you decide whether to wait again or do other work meanwhile. +- Without `task_id`, the wait ends as soon as any background task that was running at call time finishes. Tasks started during the wait are not covered by it; their completion arrives via the usual automatic notification. +- With `task_id`, the wait ends when that task finishes. An unknown `task_id` is an error; a task that has already finished returns immediately. +- When no background tasks are running, WaitFor returns immediately without waiting. +- When the wait ends because a task finished, the result also lists other tasks that finished during the wait window, so failures surface with context. +- Waiting has no side effects on the waited tasks: WaitFor never stops a task, and interrupting the wait (for example, a user interruption) leaves every task running. +- A finished task's result is delivered exactly once: tasks reported by WaitFor do not also produce an automatic completion notification. +- You can only wait for background tasks started by this agent; task IDs belonging to other agents are unknown here. diff --git a/packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.ts b/packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.ts new file mode 100644 index 000000000..69b001414 --- /dev/null +++ b/packages/agent-core-v2/src/agent/tools/task/task-wait/task-wait.ts @@ -0,0 +1,29 @@ +import { z } from 'zod'; + +import { createDecorator } from '#/_base/di/instantiation'; +import { type AgentTool } from '#/tool/toolContract'; +import { DEFAULT_BACKGROUND_TIMEOUT_S } from '#/agent/tools/os/bash/bash'; + +export const WAIT_FOR_MAX_TIMEOUT_S = DEFAULT_BACKGROUND_TIMEOUT_S; + +export const WaitForInputSchema = z.object({ + timeout: z + .number() + .int() + .positive() + .max(WAIT_FOR_MAX_TIMEOUT_S) + .describe( + `Maximum time to wait, in seconds (1-${String(WAIT_FOR_MAX_TIMEOUT_S)}). A timeout is not an error: the tool returns the tasks that are still running, and you can call it again to keep waiting.`, + ), + task_id: z + .string() + .optional() + .describe( + 'The background task ID to wait for. When omitted, the wait ends as soon as any background task that was running at call time finishes.', + ), +}); + +export type WaitForInput = z.infer; + +export interface IWaitForTool extends AgentTool { readonly _serviceBrand: undefined } +export const IWaitForTool = createDecorator('waitForTool'); diff --git a/packages/agent-core-v2/src/agent/tools/task/task-wait/taskWaitTool.ts b/packages/agent-core-v2/src/agent/tools/task/task-wait/taskWaitTool.ts new file mode 100644 index 000000000..4c2733b75 --- /dev/null +++ b/packages/agent-core-v2/src/agent/tools/task/task-wait/taskWaitTool.ts @@ -0,0 +1,274 @@ +import { toInputJsonSchema } from '#/tool/input-schema'; +import { matchesGlobRuleSubject } from '#/tool/rule-match'; +import { + type ExecutableToolContext, + type ExecutableToolResult, + type ToolExecution, +} from '#/tool/toolContract'; +import { registerAgentToolService } from '#/agent/toolRegistry/toolContribution'; + +import { IAgentTaskService } from '#/agent/task/task'; +import type { AgentTaskInfo, AgentTaskOutputSnapshot } from '#/agent/task/task'; +import { TERMINAL_STATUSES } from '#/agent/task/types'; +import { formatPlainObject } from '#/agent/task/tools/format'; +import { formatTaskList } from '#/agent/tools/task/task-list/taskListTool'; +import { IFlagService } from '#/app/flag/flag'; +import { ITelemetryService } from '#/app/telemetry/telemetry'; +import { abortError, linkAbortSignal } from '#/_base/utils/abort'; +import { WAIT_FOR_FLAG_ID } from './flag'; +import { IWaitForTool, WaitForInputSchema, type WaitForInput } from './task-wait'; +import WAIT_FOR_DESCRIPTION from './task-wait.md?raw'; + +const OUTPUT_PREVIEW_BYTES = 32 * 1024; + +const PAGING_HINT_LINES = 300; + +type WaitForOutcome = 'completed' | 'timed_out' | 'task_not_found' | 'aborted'; + +function terminalReason(info: AgentTaskInfo): 'timed_out' | 'stopped' | 'failed' | undefined { + if (info.status === 'timed_out') return 'timed_out'; + if (info.status === 'killed' && info.stopReason !== undefined) return 'stopped'; + if (info.status === 'failed' && info.stopReason !== undefined) return 'failed'; + return undefined; +} + +function fullOutputHint(output: AgentTaskOutputSnapshot): string | undefined { + if (!output.fullOutputAvailable || output.outputPath === undefined) return undefined; + if (output.truncated) { + return ( + `Only the last ${String(OUTPUT_PREVIEW_BYTES)} bytes are shown above. ` + + 'Use the Read tool with the output_path to page through the full log ' + + `(parameters: path, line_offset, n_lines; read about ${String(PAGING_HINT_LINES)} ` + + 'lines per page).' + ); + } + return ( + 'The preview above is the complete output. Use the Read tool with the output_path ' + + 'if you need to re-read the full log later ' + + `(parameters: path, line_offset, n_lines; read about ${String(PAGING_HINT_LINES)} ` + + 'lines per page).' + ); +} + +export class WaitForTool implements IWaitForTool { + declare readonly _serviceBrand: undefined; + readonly name = 'WaitFor' as const; + readonly description: string = WAIT_FOR_DESCRIPTION; + readonly parameters: Record = toInputJsonSchema(WaitForInputSchema); + + constructor( + @IAgentTaskService private readonly tasks: IAgentTaskService, + @ITelemetryService private readonly telemetry: ITelemetryService, + @IFlagService private readonly flags: IFlagService, + ) {} + + resolveExecution(args: WaitForInput): ToolExecution { + return { + description: + args.task_id === undefined + ? `Waiting up to ${String(args.timeout)}s for any background task` + : `Waiting up to ${String(args.timeout)}s for task ${args.task_id}`, + approvalRule: this.name, + matchesRule: (ruleArgs) => matchesGlobRuleSubject(ruleArgs, args.task_id ?? 'any'), + execute: (ctx) => this.execute(args, ctx), + }; + } + + private async execute( + args: WaitForInput, + ctx: ExecutableToolContext, + ): Promise { + if (!this.flags.enabled(WAIT_FOR_FLAG_ID)) { + return { + isError: true, + output: 'WaitFor is disabled: the wait_for experimental flag is off.', + }; + } + const startedAt = Date.now(); + const timeoutMs = args.timeout * 1000; + const runningAtStart = this.tasks.list(true); + + if (args.task_id === undefined) { + if (runningAtStart.length === 0) { + this.track(args, startedAt, timeoutMs, 'completed', 0); + return { + output: [ + formatPlainObject({ waitStatus: 'no_tasks', waitedMs: 0, timeoutMs }), + 'No background tasks are running, so there is nothing to wait for. Finished tasks report back via automatic notification.', + ].join('\n\n'), + isError: false, + }; + } + } else if (this.tasks.getTask(args.task_id) === undefined) { + this.track(args, startedAt, timeoutMs, 'task_not_found', 0); + return { isError: true, output: `Task not found: ${args.task_id}` }; + } + + let waited: AgentTaskInfo | undefined; + try { + waited = + args.task_id === undefined + ? await this.waitAny(runningAtStart, timeoutMs, ctx.signal) + : await this.tasks.wait(args.task_id, timeoutMs, ctx.signal); + } catch (error) { + this.track(args, startedAt, timeoutMs, 'aborted', 0); + throw error; + } + + if (waited === undefined) { + this.track(args, startedAt, timeoutMs, 'task_not_found', 0); + return { isError: true, output: `Task not found: ${args.task_id ?? ''}` }; + } + + if (!TERMINAL_STATUSES.has(waited.status)) { + this.track(args, startedAt, timeoutMs, 'timed_out', 0); + return { output: this.formatTimeout(args, startedAt, timeoutMs), isError: false }; + } + + const extras = this.collectExtras(runningAtStart, waited.taskId); + const output = await this.formatCompleted(waited, extras, startedAt, timeoutMs); + this.tasks.markTasksDeliveredViaWait( + [waited, ...extras].map((info) => ({ taskId: info.taskId, status: info.status })), + ); + this.track(args, startedAt, timeoutMs, 'completed', extras.length); + return { output, isError: false }; + } + + private async waitAny( + running: readonly AgentTaskInfo[], + timeoutMs: number, + signal: AbortSignal, + ): Promise { + const controller = new AbortController(); + const unlink = linkAbortSignal(signal, controller); + try { + const outcomes = running.map((task) => + this.tasks.wait(task.taskId, timeoutMs, controller.signal).then( + (info) => ({ info, error: undefined }), + (error: unknown) => ({ + info: undefined, + error: error instanceof Error ? error : new Error(String(error)), + }), + ), + ); + const first = await Promise.race(outcomes); + if (first.error !== undefined) throw first.error; + return first.info; + } finally { + unlink(); + controller.abort(abortError()); + } + } + + private collectExtras( + runningAtStart: readonly AgentTaskInfo[], + finishedTaskId: string, + ): AgentTaskInfo[] { + const extras: AgentTaskInfo[] = []; + for (const task of runningAtStart) { + if (task.taskId === finishedTaskId) continue; + const current = this.tasks.getTask(task.taskId); + if (current !== undefined && TERMINAL_STATUSES.has(current.status)) extras.push(current); + } + return extras; + } + + private formatTimeout(args: WaitForInput, startedAt: number, timeoutMs: number): string { + const lines = [ + formatPlainObject({ + waitStatus: 'timed_out', + taskId: args.task_id, + waitedMs: Date.now() - startedAt, + timeoutMs, + }), + 'The wait ended before the task finished — a timeout is not an error. Call WaitFor again to keep waiting, or continue with other work; completion also arrives via automatic notification.', + ]; + const running = this.tasks.list(true); + if (running.length > 0) { + lines.push('', '[still_running]', formatTaskList(running, true)); + } + return lines.join('\n'); + } + + private async formatCompleted( + finished: AgentTaskInfo, + extras: readonly AgentTaskInfo[], + startedAt: number, + timeoutMs: number, + ): Promise { + const lines = [ + formatPlainObject({ + waitStatus: 'completed', + taskId: finished.taskId, + waitedMs: Date.now() - startedAt, + timeoutMs, + }), + '', + '[finished]', + ...(await this.formatFinishedTask(finished)), + ]; + if (extras.length > 0) { + lines.push( + '', + '[completed_during_wait]', + extras.map((extra) => formatPlainObject(extra)).join('\n---\n'), + 'Use TaskOutput with one of the task_id values above to read the full output.', + ); + } + const running = this.tasks.list(true); + if (running.length > 0) { + lines.push('', '[still_running]', formatTaskList(running, true)); + } + return lines.join('\n'); + } + + private async formatFinishedTask(info: AgentTaskInfo): Promise { + const output = await this.tasks.getOutputSnapshot(info.taskId, OUTPUT_PREVIEW_BYTES); + const lines = [ + formatPlainObject({ + ...info, + outputPath: output.outputPath, + terminalReason: terminalReason(info), + outputSizeBytes: output.outputSizeBytes, + outputPreviewBytes: output.previewBytes, + outputTruncated: output.truncated, + fullOutputAvailable: output.fullOutputAvailable, + fullOutputTool: + output.fullOutputAvailable && output.outputPath !== undefined ? 'Read' : undefined, + fullOutputHint: fullOutputHint(output), + }), + '', + ]; + if (output.truncated) { + lines.push( + output.fullOutputAvailable && output.outputPath !== undefined + ? `[Truncated. Full output: ${output.outputPath}]` + : '[Truncated. No persisted full log is available for this task.]', + ); + } + lines.push('[output]', output.preview || '[no output available]'); + return lines; + } + + private track( + args: WaitForInput, + startedAt: number, + timeoutMs: number, + outcome: WaitForOutcome, + extraCompletedCount: number, + ): void { + this.telemetry.track2('wait_for_completed', { + outcome, + timeout_ms: timeoutMs, + waited_ms: Date.now() - startedAt, + has_task_id: args.task_id !== undefined, + extra_completed_count: extraCompletedCount, + }); + } +} + +registerAgentToolService(IWaitForTool, WaitForTool, { + name: 'WaitFor', + domain: 'agentTask', + when: (accessor) => accessor.get(IFlagService).enabled(WAIT_FOR_FLAG_ID), +}); diff --git a/packages/agent-core-v2/src/app/telemetry/events.ts b/packages/agent-core-v2/src/app/telemetry/events.ts index 61ba9c1db..ac4153cfa 100644 --- a/packages/agent-core-v2/src/app/telemetry/events.ts +++ b/packages/agent-core-v2/src/app/telemetry/events.ts @@ -235,6 +235,14 @@ export interface BackgroundTaskCompletedEvent { status: 'running' | 'completed' | 'failed' | 'timed_out' | 'killed' | 'lost'; } +export interface WaitForCompletedEvent { + outcome: 'completed' | 'timed_out' | 'task_not_found' | 'aborted'; + timeout_ms: number; + waited_ms: number; + has_task_id: boolean; + extra_completed_count: number; +} + export interface ModelSwitchEvent { model: string; } @@ -682,6 +690,18 @@ export const telemetryEventDefinitions = { status: 'Terminal task status', }, }), + wait_for_completed: defineAgentTelemetryEvent({ + owner: 'kimi-code', + comment: 'A WaitFor tool call returns.', + properties: { + outcome: + 'How the wait ended: the waited task finished, the wait timed out, the task id was unknown, or the wait was aborted', + timeout_ms: 'Timeout argument in milliseconds', + waited_ms: 'Actual wall-clock wait time in milliseconds', + has_task_id: 'Whether a specific task id was given', + extra_completed_count: 'Number of additional tasks that finished within the wait window', + }, + }), model_switch: defineAgentTelemetryEvent({ owner: 'kimi-code', comment: 'The active model is bound or switched.', diff --git a/packages/agent-core-v2/src/features/tower/workerProfile.ts b/packages/agent-core-v2/src/features/tower/workerProfile.ts index 996b1d0f0..d5e35bb69 100644 --- a/packages/agent-core-v2/src/features/tower/workerProfile.ts +++ b/packages/agent-core-v2/src/features/tower/workerProfile.ts @@ -36,6 +36,7 @@ const TOWER_WORKER_TOOLS = [ 'TaskOutput', 'TaskStop', 'TodoList', + 'WaitFor', 'WebSearch', 'FetchURL', 'Write', diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 19f3729ba..0995dc5b2 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -398,6 +398,8 @@ export * from '#/agent/tools/task/task-output/task-output'; import '#/agent/tools/task/task-output/taskOutputTool'; export * from '#/agent/tools/task/task-stop/task-stop'; import '#/agent/tools/task/task-stop/taskStopTool'; +export * from '#/agent/tools/task/task-wait/task-wait'; +import '#/agent/tools/task/task-wait/taskWaitTool'; export * from '#/agent/task/task'; export * from '#/agent/task/taskOps'; export * from '#/agent/task/taskService'; diff --git a/packages/agent-core-v2/src/session/agentLifecycle/profile/profiles.ts b/packages/agent-core-v2/src/session/agentLifecycle/profile/profiles.ts index cc9ff3ac0..77f26482f 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/profile/profiles.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/profile/profiles.ts @@ -19,6 +19,7 @@ const AGENT_TOOLS = [ 'TaskList', 'TaskOutput', 'TaskStop', + 'WaitFor', 'CronCreate', 'CronList', 'CronDelete', @@ -57,6 +58,7 @@ const CODER_TOOLS = [ 'TaskOutput', 'TaskStop', 'TodoList', + 'WaitFor', 'WebSearch', 'FetchURL', 'Write', diff --git a/packages/agent-core-v2/test/agent/fullCompaction/fullCompaction.test.ts b/packages/agent-core-v2/test/agent/fullCompaction/fullCompaction.test.ts index 348d58926..2d18c2029 100644 --- a/packages/agent-core-v2/test/agent/fullCompaction/fullCompaction.test.ts +++ b/packages/agent-core-v2/test/agent/fullCompaction/fullCompaction.test.ts @@ -291,7 +291,7 @@ describe('FullCompaction', () => { properties: expect.objectContaining({ agent_id: 'main', source: 'manual', - tokens_before: 3_297, + tokens_before: 3_302, tokens_after: expect.any(Number), duration_ms: expect.any(Number), compacted_count: 6, @@ -570,7 +570,7 @@ describe('FullCompaction', () => { session_id: 'test-session', cwd: dir, trigger: 'auto', - token_count: 3_297, + token_count: 3_302, }); expect(post).toMatchObject({ hook_event_name: 'PostCompact', @@ -656,7 +656,7 @@ describe('FullCompaction', () => { event: 'compaction_finished', properties: expect.objectContaining({ source: 'manual', - tokens_before: 14_363, + tokens_before: 15_004, retry_count: 1, trace_id: 'trace-compact-1', }), @@ -1039,7 +1039,7 @@ describe('FullCompaction', () => { properties: expect.objectContaining({ agent_id: 'main', source: 'manual', - tokens_before: 14_363, + tokens_before: 15_004, duration_ms: expect.any(Number), round: 1, retry_count: 0, @@ -1264,7 +1264,7 @@ describe('FullCompaction', () => { event: 'compaction_failed', properties: expect.objectContaining({ source: 'manual', - tokens_before: 14_363, + tokens_before: 15_004, duration_ms: expect.any(Number), retry_count: 4, error_type: 'APIConnectionError', @@ -1637,8 +1637,8 @@ describe('FullCompaction', () => { event: 'compaction_finished', properties: expect.objectContaining({ source: 'auto', - tokens_before: 3_304, - tokens_after: 3_288, + tokens_before: 3_309, + tokens_after: 3_293, compacted_count: 7, retry_count: 0, }), diff --git a/packages/agent-core-v2/test/agent/goal/goal.test.ts b/packages/agent-core-v2/test/agent/goal/goal.test.ts index e3423f6ef..b56c78e86 100644 --- a/packages/agent-core-v2/test/agent/goal/goal.test.ts +++ b/packages/agent-core-v2/test/agent/goal/goal.test.ts @@ -1,6 +1,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { PassThrough, Readable, type Writable } from 'node:stream'; import { isUserCancellation } from '#/_base/utils/abort'; +import { Event } from '#/_base/event'; import { TurnEnded } from '#/agent/loop/turnOps'; import { TurnStarted } from '#/agent/loop/turnEvents'; @@ -11,6 +13,10 @@ import { IAgentGoalService } from '#/agent/goal/goal'; import { IGoalDeadlineScheduler } from '#/agent/goal/goalDeadlineScheduler'; import { type AgentGoalService } from '#/agent/goal/goalService'; import { GoalUpdated } from '#/agent/goal/goalOps'; +import { IAgentTaskService } from '#/agent/task/task'; +import { ProcessTask } from '#/agent/tools/os/bash/process-task'; +import { SubagentTask } from '#/agent/tools/agent/subagent-task'; +import type { IHostProcess, IHostProcessService } from '#/os/interface/hostProcess'; import { UpdateGoalToolInputSchema } from '#/agent/tools/goal/update-goal/update-goal'; import { UpdateGoalTool } from '#/agent/tools/goal/update-goal/updateGoalTool'; import { @@ -48,7 +54,9 @@ import { appService, agentService, createTestAgent as createHarnessTestAgent, + execEnvServices, permissionModeServices, + sessionService, telemetryServices, wireRecordPersistenceServices, type TestAgentContext, @@ -56,6 +64,10 @@ import { type TestAgentServiceOverride, } from '../../harness'; import { recordingTelemetry, type TelemetryRecord } from '../../app/telemetry/stubs'; +import { stubFlag } from '../../app/flag/stubs'; +import { IFlagService } from '#/app/flag/flag'; +import { ISessionToolPolicyGate } from '#/session/sessionToolPolicyGate/sessionToolPolicyGate'; +import { ISessionToolPolicy } from '#/session/sessionToolPolicy/sessionToolPolicy'; import { stubLoopWithHooks, type StubLoop } from '../loop/stubs'; import { stubToolExecutorEvents, type ToolExecutorEventStubs } from '../toolExecutor/stubs'; import { stubAgentSwarm } from './stubs'; @@ -1294,6 +1306,7 @@ describe('AgentGoalService core workflow hooks', () => { name: 'goal_continuation', }); expect(JSON.stringify(context.get().at(-1)?.content)).toContain('Continue working toward'); + expect(JSON.stringify(context.get().at(-1)?.content)).toContain('WaitFor'); }); it('blocks the next continuation only after the final allowed turn ends', async () => { @@ -2360,3 +2373,495 @@ describe('AgentGoalService fork boundaries', () => { expect(context.get()).toEqual([]); }); }); + +describe('AgentGoalService WaitFor regression', () => { + it('does not launch a goal continuation while WaitFor is pending, and the continuation prompt mentions WaitFor', async () => { + const ctx = createTestAgent(); + try { + ctx.configure({ tools: ['WaitFor', 'UpdateGoal'] }); + const tasks = ctx.get(IAgentTaskService); + + const stdout = new PassThrough(); + let resolveWait!: (code: number) => void; + const waitPromise = new Promise((resolve) => { + resolveWait = resolve; + }); + const proc = { + _serviceBrand: undefined, + stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable, + stdout, + stderr: Readable.from([]), + pid: 10098, + exitCode: null, + wait: vi.fn(() => waitPromise) as IHostProcess['wait'], + kill: vi.fn(async () => { + stdout.destroy(); + resolveWait(143); + }) as IHostProcess['kill'], + dispose: vi.fn().mockResolvedValue(undefined) as IHostProcess['dispose'], + } as IHostProcess; + tasks.registerTask(new ProcessTask(proc, 'sleep 30', 'bg work')); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + + const continuationTurnIds: number[] = []; + ctx.get(IEventBus).subscribe(TurnStarted, (event) => { + if (event.origin.kind === 'system_trigger' && event.origin.name === 'goal_continuation') { + continuationTurnIds.push(event.turnId); + } + }); + + ctx.mockNextResponse({ + type: 'function', + id: 'wait_1', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 30 }), + }); + ctx.mockNextResponse({ type: 'text', text: 'slice done' }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(1)); + await new Promise((resolve) => setTimeout(resolve, 20)); + + expect(continuationTurnIds).toEqual([]); + + stdout.end(); + resolveWait(0); + + await vi.waitFor(() => expect(continuationTurnIds).toHaveLength(1)); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(4)); + const continuationHistory = JSON.stringify(ctx.llmCalls[2]?.history); + expect(continuationHistory).toContain('Continue working toward the active goal'); + expect(continuationHistory).toContain('WaitFor'); + } finally { + await ctx.dispose(); + } + }); +}); + +describe('AgentGoalService WaitFor background scenarios', () => { + function controllableSpawn(): { + spawn: IHostProcessService['spawn']; + pushOutput: (text: string) => void; + finish: (code: number) => void; + } { + const stdout = new PassThrough(); + let resolveWait!: (code: number) => void; + const waitPromise = new Promise((resolve) => { + resolveWait = resolve; + }); + const proc: IHostProcess = { + _serviceBrand: undefined, + stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable, + stdout, + stderr: Readable.from([]), + pid: 10097, + exitCode: null, + wait: vi.fn(() => waitPromise) as IHostProcess['wait'], + kill: vi.fn(async () => { + stdout.destroy(); + resolveWait(143); + }) as IHostProcess['kill'], + dispose: vi.fn().mockResolvedValue(undefined) as IHostProcess['dispose'], + }; + return { + spawn: vi.fn(async () => proc), + pushOutput: (text) => { + stdout.write(text); + }, + finish: (code) => { + stdout.end(); + resolveWait(code); + }, + }; + } + + function watchTurns(ctx: TestAgentContext): { + continuationTurnIds: number[]; + endedReasons: string[]; + } { + const continuationTurnIds: number[] = []; + const endedReasons: string[] = []; + const eventBus = ctx.get(IEventBus); + eventBus.subscribe(TurnStarted, (event) => { + if (event.origin.kind === 'system_trigger' && event.origin.name === 'goal_continuation') { + continuationTurnIds.push(event.turnId); + } + }); + eventBus.subscribe(TurnEnded, (event) => { + endedReasons.push(event.reason); + }); + return { continuationTurnIds, endedReasons }; + } + + it('dispatches a background bash task, waits for it, and completes the goal in one turn', async () => { + const sh = controllableSpawn(); + const ctx = createTestAgent( + execEnvServices({ processRunner: { spawn: sh.spawn } }), + permissionModeServices('yolo'), + ); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + const { continuationTurnIds, endedReasons } = watchTurns(ctx); + + ctx.mockNextResponse({ + type: 'function', + id: 'bash_1', + name: 'Bash', + arguments: JSON.stringify({ command: 'sleep 30', run_in_background: true, description: 'bg sleep' }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'wait_1', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 30 }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(2)); + await new Promise((resolve) => setTimeout(resolve, 20)); + + expect(continuationTurnIds).toEqual([]); + + sh.pushOutput('BG-OUTPUT\n'); + sh.finish(0); + + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(4)); + expect(continuationTurnIds).toEqual([]); + const history = JSON.stringify(ctx.llmCalls[2]?.history); + expect(history).toContain('wait_status: completed'); + expect(history).toContain('BG-OUTPUT'); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + expect(endedReasons).toEqual(['completed']); + } finally { + await ctx.dispose(); + } + }); + + it('waits for a dispatched background subagent and completes the goal in one turn', async () => { + const ctx = createTestAgent(); + try { + ctx.configure(); + const tasks = ctx.get(IAgentTaskService); + let settle!: (value: { result: string }) => void; + const completion = new Promise<{ result: string }>((resolve) => { + settle = resolve; + }); + tasks.registerTask( + new SubagentTask( + { agentId: 'agent-child', profileName: 'coder', completion }, + 'investigate flaky test', + new AbortController(), + ), + ); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + const { continuationTurnIds, endedReasons } = watchTurns(ctx); + + ctx.mockNextResponse({ + type: 'function', + id: 'wait_1', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 30 }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(1)); + await new Promise((resolve) => setTimeout(resolve, 20)); + + expect(continuationTurnIds).toEqual([]); + + settle({ result: 'SUBAGENT-FINDINGS: the test is order-dependent' }); + + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(3)); + expect(continuationTurnIds).toEqual([]); + const history = JSON.stringify(ctx.llmCalls[1]?.history); + expect(history).toContain('wait_status: completed'); + expect(history).toContain('kind: agent'); + expect(history).toContain('SUBAGENT-FINDINGS'); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + expect(endedReasons).toEqual(['completed']); + } finally { + await ctx.dispose(); + } + }); + + it('waits again after a WaitFor timeout and still completes the goal without continuations', async () => { + const sh = controllableSpawn(); + const ctx = createTestAgent( + execEnvServices({ processRunner: { spawn: sh.spawn } }), + permissionModeServices('yolo'), + ); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + const { continuationTurnIds, endedReasons } = watchTurns(ctx); + + ctx.mockNextResponse({ + type: 'function', + id: 'bash_1', + name: 'Bash', + arguments: JSON.stringify({ command: 'sleep 30', run_in_background: true, description: 'bg sleep' }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'wait_1', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 1 }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'wait_2', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 30 }), + }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(3), { timeout: 5000 }); + + expect(continuationTurnIds).toEqual([]); + const timedOutHistory = JSON.stringify(ctx.llmCalls[2]?.history); + expect(timedOutHistory).toContain('wait_status: timed_out'); + expect(timedOutHistory).toContain('[still_running]'); + + sh.pushOutput('BG-OUTPUT\n'); + sh.finish(0); + + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(5)); + expect(continuationTurnIds).toEqual([]); + const completedHistory = JSON.stringify(ctx.llmCalls[3]?.history); + expect(completedHistory).toContain('wait_status: completed'); + expect(completedHistory).toContain('BG-OUTPUT'); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + expect(endedReasons).toEqual(['completed']); + } finally { + await ctx.dispose(); + } + }); + + it('runs a ten-turn goal chain with WaitFor in a continuation turn', async () => { + const sh = controllableSpawn(); + const ctx = createTestAgent( + execEnvServices({ processRunner: { spawn: sh.spawn } }), + permissionModeServices('yolo'), + ); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + const { continuationTurnIds, endedReasons } = watchTurns(ctx); + + ctx.mockNextResponse({ + type: 'function', + id: 'bash_1', + name: 'Bash', + arguments: JSON.stringify({ command: 'sleep 30', run_in_background: true, description: 'bg sleep' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'slice 1 done' }); + ctx.mockNextResponse({ + type: 'function', + id: 'wait_1', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 30 }), + }); + for (let round = 2; round <= 9; round++) { + ctx.mockNextResponse({ type: 'text', text: `slice ${String(round)} done` }); + } + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(3)); + + sh.pushOutput('BG-OUTPUT\n'); + sh.finish(0); + + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(13), { timeout: 5000 }); + + expect(continuationTurnIds).toHaveLength(9); + expect(endedReasons).toEqual(Array(10).fill('completed')); + const waitResultHistory = JSON.stringify(ctx.llmCalls[3]?.history); + expect(waitResultHistory).toContain('wait_status: completed'); + expect(waitResultHistory).toContain('BG-OUTPUT'); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + } finally { + await ctx.dispose(); + } + }); +}); + +describe('AgentGoalService WaitFor guidance gating', () => { + it('shows the WaitFor guidance in the active-goal reminder when the flag is on', async () => { + const ctx = createTestAgent(); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + + ctx.mockNextResponse({ type: 'text', text: 'slice done' }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(3)); + + expect(JSON.stringify(ctx.llmCalls[0])).toContain('re-invoked again and again'); + } finally { + await ctx.dispose(); + } + }); + + it('hides WaitFor from the reminder, the continuation prompt, and the tools when the flag is off', async () => { + const ctx = createTestAgent(appService(IFlagService, stubFlag(false))); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + + ctx.mockNextResponse({ type: 'text', text: 'slice done' }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(3)); + + const allCalls = JSON.stringify(ctx.llmCalls); + expect(allCalls).not.toContain('re-invoked again and again'); + for (const call of ctx.llmCalls) { + expect(call.tools.map((tool) => tool.name)).not.toContain('WaitFor'); + } + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + } finally { + await ctx.dispose(); + } + }); + + it('hides WaitFor guidance when a tool policy disables WaitFor even though the flag is on', async () => { + const ctx = createTestAgent( + sessionService(ISessionToolPolicyGate, { + _serviceBrand: undefined, + disabledTools: ['WaitFor'], + onDidChange: Event.None as Event, + }), + ); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + + ctx.mockNextResponse({ type: 'text', text: 'slice done' }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(3)); + + const allCalls = JSON.stringify(ctx.llmCalls); + expect(allCalls).not.toContain('WaitFor'); + expect(allCalls).not.toContain('re-invoked again and again'); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + } finally { + await ctx.dispose(); + } + }); + + it('hides WaitFor guidance once the session tool policy disables it mid-goal', async () => { + const ctx = createTestAgent(); + try { + ctx.configure({ tools: ['WaitFor', 'UpdateGoal'] }); + const tasks = ctx.get(IAgentTaskService); + let settle!: (value: { result: string }) => void; + const completion = new Promise<{ result: string }>((resolve) => { + settle = resolve; + }); + tasks.registerTask( + new SubagentTask( + { agentId: 'agent-child', profileName: 'coder', completion }, + 'bg work', + new AbortController(), + ), + ); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + + ctx.mockNextResponse({ + type: 'function', + id: 'wait_1', + name: 'WaitFor', + arguments: JSON.stringify({ timeout: 30 }), + }); + ctx.mockNextResponse({ type: 'text', text: 'slice done' }); + ctx.mockNextResponse({ + type: 'function', + id: 'ug_1', + name: 'UpdateGoal', + arguments: JSON.stringify({ status: 'complete' }), + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(1)); + expect(JSON.stringify(ctx.llmCalls[0])).toContain('re-invoked again and again'); + + await ctx.get(ISessionToolPolicy).setDisabledTools(['WaitFor']); + settle({ result: 'bg result' }); + + await vi.waitFor(() => expect(ctx.llmCalls).toHaveLength(4)); + const continuationCall = ctx.llmCalls[2]!; + const continuationPrompt = continuationCall.history.find((message) => + JSON.stringify(message).includes('Continue working toward the active goal'), + ); + expect(continuationPrompt).toBeDefined(); + expect(JSON.stringify(continuationPrompt)).not.toContain('re-invoked again and again'); + const freshReminder = continuationCall.history.at(-1); + expect(JSON.stringify(freshReminder)).toContain('active goal'); + expect(JSON.stringify(freshReminder)).not.toContain('re-invoked again and again'); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + } finally { + await ctx.dispose(); + } + }); +}); diff --git a/packages/agent-core-v2/test/agent/loop/loop.test.ts b/packages/agent-core-v2/test/agent/loop/loop.test.ts index feec8518e..bf78addf3 100644 --- a/packages/agent-core-v2/test/agent/loop/loop.test.ts +++ b/packages/agent-core-v2/test/agent/loop/loop.test.ts @@ -134,8 +134,8 @@ describe('Agent loop', () => { [emit] turn.step.started { "time": "