mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-19 13:45:28 +00:00
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
This commit is contained in:
Luyu Cheng
GitHub
parent
cdaa80b778
commit
8440801de4
31 changed files with 1843 additions and 46 deletions
5
.changeset/wait-for-tool.md
Normal file
5
.changeset/wait-for-tool.md
Normal file
|
|
@ -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.
|
||||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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#运行时开关)。
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
13
packages/agent-core-v2/docs/wire-manifest.d.ts
vendored
13
packages/agent-core-v2/docs/wire-manifest.d.ts
vendored
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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: [],
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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}` : '',
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ const DEFAULT_APPROVE_TOOLS = new Set([
|
|||
'TodoList',
|
||||
'TaskList',
|
||||
'TaskOutput',
|
||||
'WaitFor',
|
||||
'CronList',
|
||||
'WebSearch',
|
||||
'FetchURL',
|
||||
|
|
|
|||
|
|
@ -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>;
|
||||
|
|
|
|||
|
|
@ -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>(),
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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}`;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
@ -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.
|
||||
|
|
@ -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<typeof WaitForInputSchema>;
|
||||
|
||||
export interface IWaitForTool extends AgentTool<WaitForInput> { readonly _serviceBrand: undefined }
|
||||
export const IWaitForTool = createDecorator<IWaitForTool>('waitForTool');
|
||||
|
|
@ -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<string, unknown> = 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<ExecutableToolResult> {
|
||||
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<AgentTaskInfo | undefined> {
|
||||
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<string> {
|
||||
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<string[]> {
|
||||
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),
|
||||
});
|
||||
|
|
@ -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<WaitForCompletedEvent>({
|
||||
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<ModelSwitchEvent>({
|
||||
owner: 'kimi-code',
|
||||
comment: 'The active model is bound or switched.',
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ const TOWER_WORKER_TOOLS = [
|
|||
'TaskOutput',
|
||||
'TaskStop',
|
||||
'TodoList',
|
||||
'WaitFor',
|
||||
'WebSearch',
|
||||
'FetchURL',
|
||||
'Write',
|
||||
|
|
|
|||
|
|
@ -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';
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -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<number>((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<number>((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<string>(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<void>,
|
||||
}),
|
||||
);
|
||||
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();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -18,12 +18,12 @@ import {
|
|||
type AgentTaskInfo,
|
||||
} from '#/agent/task/task';
|
||||
import { renderNotificationXml } from '#/agent/task/notificationXml';
|
||||
import { AgentTaskService } from '#/agent/task/taskService';
|
||||
import { AgentTaskService, taskNotificationDeliveryKey } from '#/agent/task/taskService';
|
||||
import { ProcessTask } from '#/agent/tools/os/bash/process-task';
|
||||
import type { IHostProcess } from '#/os/interface/hostProcess';
|
||||
import { IConfigRegistry, IConfigService } from '#/app/config/config';
|
||||
import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory';
|
||||
import type { ContextMessage } from '#/agent/contextMemory/types';
|
||||
import type { ContextMessage, TaskOrigin } from '#/agent/contextMemory/types';
|
||||
import { IAgentLoopService } from '#/agent/loop/loop';
|
||||
import { IAgentScopeContext, makeAgentScopeContext } from '#/agent/scopeContext/scopeContext';
|
||||
import { IAgentStateService } from '#/agent/state/agentState';
|
||||
|
|
@ -31,9 +31,13 @@ import { AgentStateService } from '#/agent/state/agentStateService';
|
|||
import { ISessionContext, makeSessionContext } from '#/session/sessionContext/sessionContext';
|
||||
import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore';
|
||||
import { IFileSystemStorageService } from '#/persistence/interface/storage';
|
||||
import { ITelemetryService } from '#/app/telemetry/telemetry';
|
||||
import { ITelemetryService, noopTelemetryService } from '#/app/telemetry/telemetry';
|
||||
import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
|
||||
import { SubagentTask } from '#/agent/tools/agent/subagent-task';
|
||||
import { type WaitForInput } from '#/agent/tools/task/task-wait/task-wait';
|
||||
import { WaitForTool } from '#/agent/tools/task/task-wait/taskWaitTool';
|
||||
import { IWireService } from '#/wire/wire';
|
||||
import { WireService } from '#/wire/wireService';
|
||||
import { IEventBus } from '#/app/event/eventBus';
|
||||
import { EventBusService } from '#/app/event/eventBusService';
|
||||
import { IAgentBlobService } from '#/agent/blob/agentBlobService';
|
||||
|
|
@ -41,11 +45,15 @@ import { ContextSpliced } from '#/agent/contextMemory/contextEvents';
|
|||
import { IEventDispatcher } from '#/state/eventDispatcher';
|
||||
import { EventDispatcherService } from '#/state/eventDispatcherService';
|
||||
import { ITaskService } from '#/app/task/task';
|
||||
import { IAppendLogStore } from '#/persistence/interface/appendLogStore';
|
||||
import { AppendLogStore } from '#/persistence/backends/node-fs/appendLogStore';
|
||||
import { InMemoryStorageService } from '#/persistence/backends/memory/inMemoryStorageService';
|
||||
|
||||
import { stubLog } from '../../_base/log/stubs';
|
||||
import { stubContextMemory } from '../contextMemory/stubs';
|
||||
import { stubLoopWithHooks } from '../loop/stubs';
|
||||
import { stubContextMemory, type StubContextMemory } from '../contextMemory/stubs';
|
||||
import { stubLoopWithHooks, type StubLoop } from '../loop/stubs';
|
||||
import { stubFlag } from '../../app/flag/stubs';
|
||||
import { executeTool } from '../../tools/fixtures/execute-tool';
|
||||
import type { TaskServiceTestManager } from './stubs';
|
||||
|
||||
function fakeProcessTask(): AgentTask {
|
||||
|
|
@ -251,6 +259,199 @@ describe('AgentTaskService', () => {
|
|||
expect(terminated?.['outputTail']).toBeUndefined();
|
||||
});
|
||||
|
||||
function stubLoop(): StubLoop {
|
||||
return ix.get(IAgentLoopService) as unknown as StubLoop;
|
||||
}
|
||||
|
||||
async function waitForCondition(condition: () => boolean): Promise<void> {
|
||||
for (let attempt = 0; attempt < 100; attempt++) {
|
||||
if (condition()) return;
|
||||
await new Promise((resolve) => setTimeout(resolve, 1));
|
||||
}
|
||||
}
|
||||
|
||||
it('enqueues a terminal notification for a finished detached task', async () => {
|
||||
const svc = ix.get(IAgentTaskService);
|
||||
const taskId = svc.registerTask(outputtingTask('done\n'));
|
||||
|
||||
await svc.wait(taskId, 1000);
|
||||
const loop = stubLoop();
|
||||
await waitForCondition(() => loop.hasPendingRequests());
|
||||
|
||||
expect(loop.hasPendingRequests()).toBe(true);
|
||||
});
|
||||
|
||||
it('markTasksDeliveredViaWait suppresses the automatic terminal notification', async () => {
|
||||
const svc = ix.get(IAgentTaskService);
|
||||
const taskId = svc.registerTask(outputtingTask('done\n'));
|
||||
svc.markTasksDeliveredViaWait([{ taskId, status: 'completed' }]);
|
||||
|
||||
await svc.wait(taskId, 1000);
|
||||
const loop = stubLoop();
|
||||
await waitForCondition(() => loop.hasPendingRequests());
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
expect(loop.hasPendingRequests()).toBe(false);
|
||||
expect(loop.launches).toEqual([]);
|
||||
|
||||
const deliveryKey = `${taskId}\0completed\0task:${taskId}:completed`;
|
||||
const states = ix.get(IAgentStateService);
|
||||
await waitForCondition(() => states.get(taskNotificationDeliveryKey).length > 0);
|
||||
expect(states.get(taskNotificationDeliveryKey)).toContain(deliveryKey);
|
||||
});
|
||||
|
||||
it('aborts an already-enqueued terminal notification when the task is marked delivered via wait', async () => {
|
||||
const svc = ix.get(IAgentTaskService);
|
||||
const taskId = svc.registerTask(outputtingTask('done\n'));
|
||||
|
||||
await svc.wait(taskId, 1000);
|
||||
const loop = stubLoop();
|
||||
await waitForCondition(() => loop.hasPendingRequests());
|
||||
expect(loop.hasPendingRequests()).toBe(true);
|
||||
|
||||
svc.markTasksDeliveredViaWait([{ taskId, status: 'completed' }]);
|
||||
|
||||
expect(loop.hasPendingRequests()).toBe(false);
|
||||
});
|
||||
|
||||
it('suppresses only the notification whose status was reported via wait', async () => {
|
||||
const svc = ix.get(IAgentTaskService);
|
||||
const taskId = svc.registerTask(outputtingTask('done\n'));
|
||||
svc.markTasksDeliveredViaWait([{ taskId, status: 'failed' }]);
|
||||
|
||||
await svc.wait(taskId, 1000);
|
||||
const loop = stubLoop();
|
||||
await waitForCondition(() => loop.hasPendingRequests());
|
||||
|
||||
expect(loop.hasPendingRequests()).toBe(true);
|
||||
});
|
||||
|
||||
it('keeps the automatic notification of tasks that were not reported via wait', async () => {
|
||||
const svc = ix.get(IAgentTaskService);
|
||||
const taskA = svc.registerTask(outputtingTask('a\n'));
|
||||
const taskB = svc.registerTask(outputtingTask('b\n'));
|
||||
svc.markTasksDeliveredViaWait([{ taskId: taskA, status: 'completed' }]);
|
||||
|
||||
await svc.wait(taskA, 1000);
|
||||
await svc.wait(taskB, 1000);
|
||||
const loop = stubLoop();
|
||||
await waitForCondition(() => loop.hasPendingRequests());
|
||||
|
||||
const context = ix.get(IAgentContextMemoryService) as StubContextMemory;
|
||||
loop.drainNextBatch(context);
|
||||
|
||||
const delivered = context.messages.filter((message) => message.origin?.kind === 'task');
|
||||
expect(delivered.map((message) => (message.origin as TaskOrigin).taskId)).toEqual([taskB]);
|
||||
});
|
||||
|
||||
function waitContext(toolCallId: string, args: WaitForInput) {
|
||||
return { turnId: 0, toolCallId, args, signal: new AbortController().signal };
|
||||
}
|
||||
|
||||
function waitResultString(result: { readonly output: string | readonly unknown[] }): string {
|
||||
expect(typeof result.output).toBe('string');
|
||||
return result.output as string;
|
||||
}
|
||||
|
||||
function pendingSubagentTask(agentId: string, description: string): {
|
||||
task: SubagentTask;
|
||||
settle: (value: { result: string }) => void;
|
||||
} {
|
||||
let settle!: (value: { result: string }) => void;
|
||||
const completion = new Promise<{ result: string }>((resolve) => {
|
||||
settle = resolve;
|
||||
});
|
||||
return {
|
||||
task: new SubagentTask(
|
||||
{ agentId, profileName: 'coder', completion },
|
||||
description,
|
||||
new AbortController(),
|
||||
),
|
||||
settle,
|
||||
};
|
||||
}
|
||||
|
||||
it('unwinds a nested wait chain leaf-first without deadlocking', async () => {
|
||||
const docs = mapBackedDocs();
|
||||
const bytes = new InMemoryStorageService();
|
||||
const mainSvc = buildAgentIx('main', docs, bytes).get(IAgentTaskService);
|
||||
const childSvc = buildAgentIx('child-1', docs, bytes).get(IAgentTaskService);
|
||||
const mainTool = new WaitForTool(mainSvc, noopTelemetryService, stubFlag(true));
|
||||
const childTool = new WaitForTool(childSvc, noopTelemetryService, stubFlag(true));
|
||||
|
||||
const leaf = pendingSubagentTask('agent-grandchild', 'leaf work');
|
||||
const taskC = childSvc.registerTask(leaf.task);
|
||||
await childSvc.suppressTerminalNotification(taskC);
|
||||
|
||||
const childWait = executeTool(
|
||||
childTool,
|
||||
waitContext('wait_child', { timeout: 30, task_id: taskC }),
|
||||
);
|
||||
const order: string[] = [];
|
||||
void childWait.then(() => {
|
||||
order.push('childWait');
|
||||
});
|
||||
const completionM = childWait.then(() => {
|
||||
order.push('taskM');
|
||||
return { result: 'parent done after child' };
|
||||
});
|
||||
const taskM = mainSvc.registerTask(
|
||||
new SubagentTask(
|
||||
{ agentId: 'agent-parent', profileName: 'coder', completion: completionM },
|
||||
'parent work',
|
||||
new AbortController(),
|
||||
),
|
||||
);
|
||||
const mainWait = executeTool(
|
||||
mainTool,
|
||||
waitContext('wait_main', { timeout: 30, task_id: taskM }),
|
||||
);
|
||||
void mainWait.then(() => {
|
||||
order.push('mainWait');
|
||||
});
|
||||
|
||||
leaf.settle({ result: 'leaf findings' });
|
||||
|
||||
const childResult = waitResultString(await childWait);
|
||||
const mainResult = waitResultString(await mainWait);
|
||||
expect(childResult).toContain('wait_status: completed');
|
||||
expect(childResult).toContain('leaf findings');
|
||||
expect(mainResult).toContain('wait_status: completed');
|
||||
expect(mainResult).toContain('parent done after child');
|
||||
expect(order).toEqual(['childWait', 'taskM', 'mainWait']);
|
||||
});
|
||||
|
||||
it('rejects waiting on a task owned by another agent, so a wait cycle cannot form', async () => {
|
||||
const docs = mapBackedDocs();
|
||||
const bytes = new InMemoryStorageService();
|
||||
const mainSvc = buildAgentIx('main', docs, bytes).get(IAgentTaskService);
|
||||
const childSvc = buildAgentIx('child-1', docs, bytes).get(IAgentTaskService);
|
||||
const mainTool = new WaitForTool(mainSvc, noopTelemetryService, stubFlag(true));
|
||||
const childTool = new WaitForTool(childSvc, noopTelemetryService, stubFlag(true));
|
||||
|
||||
const parent = pendingSubagentTask('agent-parent', 'parent work');
|
||||
const taskM = mainSvc.registerTask(parent.task);
|
||||
const leaf = pendingSubagentTask('agent-grandchild', 'leaf work');
|
||||
const taskC = childSvc.registerTask(leaf.task);
|
||||
|
||||
const childWaitingOnParent = await executeTool(
|
||||
childTool,
|
||||
waitContext('wait_cross_up', { timeout: 30, task_id: taskM }),
|
||||
);
|
||||
expect(childWaitingOnParent.isError).toBe(true);
|
||||
expect(waitResultString(childWaitingOnParent)).toContain(`Task not found: ${taskM}`);
|
||||
|
||||
const parentWaitingOnChild = await executeTool(
|
||||
mainTool,
|
||||
waitContext('wait_cross_down', { timeout: 30, task_id: taskC }),
|
||||
);
|
||||
expect(parentWaitingOnChild.isError).toBe(true);
|
||||
expect(waitResultString(parentWaitingOnChild)).toContain(`Task not found: ${taskC}`);
|
||||
|
||||
parent.settle({ result: 'parent done' });
|
||||
leaf.settle({ result: 'leaf done' });
|
||||
});
|
||||
|
||||
function stubTaskConfig(value: unknown): void {
|
||||
ix.stub(IConfigService, {
|
||||
get: ((domain: string) => (domain === 'task' ? value : undefined)) as IConfigService['get'],
|
||||
|
|
@ -530,6 +731,90 @@ describe('AgentTaskService', () => {
|
|||
return ix;
|
||||
}
|
||||
|
||||
function buildWiredAgentIx(
|
||||
agentId: string,
|
||||
docs: IAtomicDocumentStore,
|
||||
bytes: IFileSystemStorageService,
|
||||
context: StubContextMemory,
|
||||
): TestInstantiationService {
|
||||
const ix = disposables.add(new TestInstantiationService());
|
||||
ix.stub(ILogService, stubLog());
|
||||
ix.stub(IAgentConversationUndoParticipantRegistry, {
|
||||
register: () => toDisposable(() => {}),
|
||||
list: () => [],
|
||||
});
|
||||
ix.stub(IEventBus, disposables.add(new EventBusService()));
|
||||
ix.stub(IAgentContextInjectorService, {
|
||||
register: () => toDisposable(() => {}),
|
||||
});
|
||||
ix.stub(ITaskService, {
|
||||
run: () => {
|
||||
throw new Error('ITaskService.run is not used by this test');
|
||||
},
|
||||
defer: () => {
|
||||
throw new Error('ITaskService.defer is not used by this test');
|
||||
},
|
||||
});
|
||||
ix.stub(IAgentContextMemoryService, context);
|
||||
ix.stub(ITelemetryService, { track: () => {}, track2: () => {} });
|
||||
ix.stub(IAgentLoopService, stubLoopWithHooks());
|
||||
ix.stub(IConfigService, {
|
||||
get: (() => undefined) as IConfigService['get'],
|
||||
});
|
||||
ix.stub(
|
||||
ISessionContext,
|
||||
makeSessionContext({
|
||||
sessionId: 'test-session',
|
||||
workspaceId: 'test-ws',
|
||||
sessionDir: '/tmp/test-session',
|
||||
sessionScope: 'sessions/test-ws/test-session',
|
||||
cwd: '/tmp/test-session',
|
||||
}),
|
||||
);
|
||||
ix.stub(
|
||||
IAgentScopeContext,
|
||||
makeAgentScopeContext({
|
||||
agentId,
|
||||
agentScope: `sessions/test-ws/test-session/agents/${agentId}`,
|
||||
}),
|
||||
);
|
||||
ix.stub(IAtomicDocumentStore, docs);
|
||||
ix.stub(IFileSystemStorageService, bytes);
|
||||
ix.stub(IAgentBlobService, noopBlob);
|
||||
ix.set(IAppendLogStore, new SyncDescriptor(AppendLogStore));
|
||||
ix.set(IWireService, new SyncDescriptor(WireService));
|
||||
ix.set(IAgentStateService, new AgentStateService());
|
||||
ix.set(IEventDispatcher, new SyncDescriptor(EventDispatcherService));
|
||||
ix.set(IAgentTaskService, new SyncDescriptor(AgentTaskService));
|
||||
return ix;
|
||||
}
|
||||
|
||||
it('rebuilds wait-delivered keys on restore and skips their re-delivery', async () => {
|
||||
const docs = mapBackedDocs();
|
||||
const bytes = new InMemoryStorageService();
|
||||
|
||||
const one = buildWiredAgentIx('main', docs, bytes, stubContextMemory());
|
||||
const svc1 = one.get(IAgentTaskService);
|
||||
await one.get(IEventDispatcher).restore();
|
||||
|
||||
const taskA = svc1.registerTask(outputtingTask('a\n'));
|
||||
const taskB = svc1.registerTask(outputtingTask('b\n'));
|
||||
svc1.markTasksDeliveredViaWait([{ taskId: taskA, status: 'completed' }]);
|
||||
await svc1.wait(taskA, 1000);
|
||||
await svc1.wait(taskB, 1000);
|
||||
await one.get(IEventDispatcher).flush();
|
||||
|
||||
const context2 = stubContextMemory();
|
||||
const two = buildWiredAgentIx('main', docs, bytes, context2);
|
||||
two.get(IAgentTaskService);
|
||||
await two.get(IEventDispatcher).restore();
|
||||
|
||||
const keyA = `${taskA}\0completed\0task:${taskA}:completed`;
|
||||
expect(two.get(IAgentStateService).get(taskNotificationDeliveryKey)).toContain(keyA);
|
||||
const redelivered = context2.messages.filter((message) => message.origin?.kind === 'task');
|
||||
expect(redelivered.map((message) => (message.origin as TaskOrigin).taskId)).toEqual([taskB]);
|
||||
});
|
||||
|
||||
it('restore touches only the agent own task records', async () => {
|
||||
const docs = mapBackedDocs();
|
||||
const bytes = new InMemoryStorageService();
|
||||
|
|
|
|||
|
|
@ -1,27 +1,42 @@
|
|||
import { PassThrough, Readable, type Writable } from 'node:stream';
|
||||
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import type {
|
||||
AgentTask,
|
||||
AgentTaskInfo,
|
||||
AgentTaskOutputSnapshot,
|
||||
AgentTaskTrackOptions,
|
||||
ForegroundTaskReleaseReason,
|
||||
IAgentTaskEntry,
|
||||
import {
|
||||
IAgentTaskService,
|
||||
RegisterAgentTaskOptions,
|
||||
type AgentTask,
|
||||
type AgentTaskInfo,
|
||||
type AgentTaskOutputSnapshot,
|
||||
type AgentTaskTrackOptions,
|
||||
type AgentTaskWaitDelivery,
|
||||
type ForegroundTaskReleaseReason,
|
||||
type IAgentTaskEntry,
|
||||
type RegisterAgentTaskOptions,
|
||||
} from '#/agent/task/task';
|
||||
import { TERMINAL_STATUSES } from '#/agent/task/types';
|
||||
import { type AgentTaskStatus, TERMINAL_STATUSES } from '#/agent/task/types';
|
||||
import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
|
||||
import { TaskListInputSchema } from '#/agent/tools/task/task-list/task-list';
|
||||
import { TaskListTool } from '#/agent/tools/task/task-list/taskListTool';
|
||||
import { TaskOutputInputSchema } from '#/agent/tools/task/task-output/task-output';
|
||||
import { TaskOutputTool } from '#/agent/tools/task/task-output/taskOutputTool';
|
||||
import { TaskStopInputSchema } from '#/agent/tools/task/task-stop/task-stop';
|
||||
import { TaskStopTool } from '#/agent/tools/task/task-stop/taskStopTool';
|
||||
import { WaitForInputSchema } from '#/agent/tools/task/task-wait/task-wait';
|
||||
import { WaitForTool } from '#/agent/tools/task/task-wait/taskWaitTool';
|
||||
import { abortError } from '#/_base/utils/abort';
|
||||
import type { ITaskHandle } from '#/app/task/task';
|
||||
import type { IHostProcess } from '#/os/interface/hostProcess';
|
||||
import { compileToolArgsValidator, validateToolArgs } from '#/tool/args-validator';
|
||||
import type { ProcessTaskInfo } from '#/agent/tools/os/bash/process-task';
|
||||
import { ProcessTask, type ProcessTaskInfo } from '#/agent/tools/os/bash/process-task';
|
||||
import { SubagentTask } from '#/agent/tools/agent/subagent-task';
|
||||
import type { SubagentTaskInfo } from '#/agent/tools/agent/subagent-task';
|
||||
import { IWaitForTool } from '#/agent/tools/task/task-wait/task-wait';
|
||||
import { IAgentLoopService } from '#/agent/loop/loop';
|
||||
import { executeTool } from '../../../tools/fixtures/execute-tool';
|
||||
import { recordingTelemetry, type TelemetryRecord } from '../../../app/telemetry/stubs';
|
||||
import { stubFlag } from '../../../app/flag/stubs';
|
||||
import { agentService, createTestAgent, telemetryServices } from '../../../harness';
|
||||
import { stubLoopWithHooks } from '../../loop/stubs';
|
||||
|
||||
const signal = new AbortController().signal;
|
||||
|
||||
|
|
@ -99,6 +114,14 @@ class FakeTaskService implements IAgentTaskService {
|
|||
readonly stopCalls: Array<{ taskId: string; reason: string | undefined }> = [];
|
||||
readonly suppressCalls: string[] = [];
|
||||
readonly waitCalls: Array<{ taskId: string; timeoutMs: number | undefined }> = [];
|
||||
readonly waitDeliveries: Array<readonly AgentTaskWaitDelivery[]> = [];
|
||||
waitDelegate:
|
||||
| ((
|
||||
taskId: string,
|
||||
timeoutMs: number | undefined,
|
||||
signal: AbortSignal | undefined,
|
||||
) => Promise<AgentTaskInfo | undefined>)
|
||||
| undefined;
|
||||
|
||||
private readonly entries = new Map<string, FakeTaskEntry>();
|
||||
|
||||
|
|
@ -110,6 +133,16 @@ class FakeTaskService implements IAgentTaskService {
|
|||
return info.taskId;
|
||||
}
|
||||
|
||||
settle(taskId: string, status: AgentTaskStatus = 'completed'): void {
|
||||
const entry = this.entries.get(taskId);
|
||||
if (entry === undefined) return;
|
||||
entry.info = {
|
||||
...entry.info,
|
||||
status,
|
||||
endedAt: entry.info.endedAt ?? 1_700_000_002_000,
|
||||
} as AgentTaskInfo;
|
||||
}
|
||||
|
||||
track(_handle: ITaskHandle, _options: AgentTaskTrackOptions): IAgentTaskEntry {
|
||||
throw new Error('track is not implemented in FakeTaskService.');
|
||||
}
|
||||
|
|
@ -136,10 +169,13 @@ class FakeTaskService implements IAgentTaskService {
|
|||
|
||||
persistOutput(_taskId: string): void {}
|
||||
|
||||
readonly failSnapshotTaskIds = new Set<string>();
|
||||
|
||||
async getOutputSnapshot(
|
||||
taskId: string,
|
||||
_maxPreviewBytes: number,
|
||||
): Promise<AgentTaskOutputSnapshot> {
|
||||
if (this.failSnapshotTaskIds.has(taskId)) throw new Error('snapshot read failed');
|
||||
return this.entries.get(taskId)?.output ?? outputSnapshot();
|
||||
}
|
||||
|
||||
|
|
@ -159,6 +195,10 @@ class FakeTaskService implements IAgentTaskService {
|
|||
} as AgentTaskInfo;
|
||||
}
|
||||
|
||||
markTasksDeliveredViaWait(tasks: readonly AgentTaskWaitDelivery[]): void {
|
||||
this.waitDeliveries.push(tasks);
|
||||
}
|
||||
|
||||
detach(taskId: string): AgentTaskInfo | undefined {
|
||||
const entry = this.entries.get(taskId);
|
||||
if (entry === undefined) return undefined;
|
||||
|
|
@ -202,9 +242,12 @@ class FakeTaskService implements IAgentTaskService {
|
|||
async wait(
|
||||
taskId: string,
|
||||
timeoutMs?: number,
|
||||
_signal?: AbortSignal,
|
||||
signal?: AbortSignal,
|
||||
): Promise<AgentTaskInfo | undefined> {
|
||||
this.waitCalls.push({ taskId, timeoutMs });
|
||||
if (this.waitDelegate !== undefined) {
|
||||
return this.waitDelegate(taskId, timeoutMs, signal);
|
||||
}
|
||||
return this.entries.get(taskId)?.info;
|
||||
}
|
||||
|
||||
|
|
@ -702,3 +745,466 @@ describe('TaskStopTool', () => {
|
|||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WaitForTool', () => {
|
||||
function waitTelemetry(): { records: TelemetryRecord[]; telemetry: ReturnType<typeof recordingTelemetry> } {
|
||||
const records: TelemetryRecord[] = [];
|
||||
return { records, telemetry: recordingTelemetry(records) };
|
||||
}
|
||||
|
||||
function lastEvent(records: TelemetryRecord[]): TelemetryRecord | undefined {
|
||||
return records.findLast((record) => record.event === 'wait_for_completed');
|
||||
}
|
||||
|
||||
it('has name and accepts the current schema', () => {
|
||||
const tool = new WaitForTool(new FakeTaskService(), recordingTelemetry([]), stubFlag(true));
|
||||
|
||||
expect(tool.name).toBe('WaitFor');
|
||||
expect(WaitForInputSchema.safeParse({ timeout: 60 }).success).toBe(true);
|
||||
expect(WaitForInputSchema.safeParse({ timeout: 60, task_id: 'bash-1' }).success).toBe(true);
|
||||
expect(WaitForInputSchema.safeParse({ timeout: 600 }).success).toBe(true);
|
||||
expect(WaitForInputSchema.safeParse({}).success).toBe(false);
|
||||
expect(WaitForInputSchema.safeParse({ timeout: 0 }).success).toBe(false);
|
||||
expect(WaitForInputSchema.safeParse({ timeout: -5 }).success).toBe(false);
|
||||
expect(WaitForInputSchema.safeParse({ timeout: 601 }).success).toBe(false);
|
||||
expect(WaitForInputSchema.safeParse({ timeout: 1.5 }).success).toBe(false);
|
||||
expect(tool.parameters).toMatchObject({
|
||||
type: 'object',
|
||||
additionalProperties: false,
|
||||
required: ['timeout'],
|
||||
properties: {
|
||||
timeout: { type: 'integer' },
|
||||
task_id: { type: 'string' },
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('returns error and tracks task_not_found for an unknown task_id', async () => {
|
||||
const { records, telemetry } = waitTelemetry();
|
||||
const result = await executeTool(
|
||||
new WaitForTool(new FakeTaskService(), telemetry, stubFlag(true)),
|
||||
context('wait_unknown', { timeout: 10, task_id: 'bash-unknown0' }),
|
||||
);
|
||||
|
||||
expect(result.isError).toBe(true);
|
||||
expect(outputString(result)).toContain('Task not found: bash-unknown0');
|
||||
expect(lastEvent(records)?.properties).toMatchObject({
|
||||
outcome: 'task_not_found',
|
||||
timeout_ms: 10_000,
|
||||
has_task_id: true,
|
||||
extra_completed_count: 0,
|
||||
});
|
||||
});
|
||||
|
||||
it('returns immediately without waiting when no background tasks are running', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, recordingTelemetry([]), stubFlag(true)),
|
||||
context('wait_none', { timeout: 10 }),
|
||||
);
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: no_tasks');
|
||||
expect(output).toContain('No background tasks are running');
|
||||
expect(tasks.waitCalls).toEqual([]);
|
||||
expect(tasks.waitDeliveries).toEqual([]);
|
||||
});
|
||||
|
||||
it('returns a finished task immediately and marks it delivered via wait', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
const taskId = tasks.add(
|
||||
processTask({
|
||||
taskId: 'bash-done0002',
|
||||
status: 'completed',
|
||||
endedAt: 1_700_000_001_000,
|
||||
exitCode: 0,
|
||||
}),
|
||||
outputSnapshot('DONE-OUTPUT\n'),
|
||||
);
|
||||
|
||||
const { records, telemetry } = waitTelemetry();
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, telemetry, stubFlag(true)),
|
||||
context('wait_done', { timeout: 10, task_id: taskId }),
|
||||
);
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: completed');
|
||||
expect(output).toContain('status: completed');
|
||||
expect(output).toContain('[finished]');
|
||||
expect(output).toContain('[output]\nDONE-OUTPUT');
|
||||
expect(tasks.waitDeliveries).toEqual([[{ taskId, status: 'completed' }]]);
|
||||
expect(lastEvent(records)?.properties).toMatchObject({
|
||||
outcome: 'completed',
|
||||
has_task_id: true,
|
||||
extra_completed_count: 0,
|
||||
});
|
||||
});
|
||||
|
||||
it('reports tasks that finished during the wait and marks all of them delivered', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-wait001', description: 'main wait' }), outputSnapshot('WAITED-OUT\n'));
|
||||
tasks.add(processTask({ taskId: 'bash-extra001', description: 'side task' }));
|
||||
tasks.waitDelegate = async (taskId) => {
|
||||
tasks.settle('bash-wait001');
|
||||
tasks.settle('bash-extra001', 'failed');
|
||||
return tasks.getTask(taskId);
|
||||
};
|
||||
|
||||
const { records, telemetry } = waitTelemetry();
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, telemetry, stubFlag(true)),
|
||||
context('wait_extras', { timeout: 10, task_id: 'bash-wait001' }),
|
||||
);
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: completed');
|
||||
expect(output).toContain('[completed_during_wait]');
|
||||
expect(output).toContain('task_id: bash-extra001');
|
||||
expect(output).toContain('status: failed');
|
||||
expect(tasks.waitDeliveries).toEqual([
|
||||
[
|
||||
{ taskId: 'bash-wait001', status: 'completed' },
|
||||
{ taskId: 'bash-extra001', status: 'failed' },
|
||||
],
|
||||
]);
|
||||
expect(lastEvent(records)?.properties).toMatchObject({
|
||||
outcome: 'completed',
|
||||
extra_completed_count: 1,
|
||||
});
|
||||
});
|
||||
|
||||
it('waits for any running task when task_id is omitted', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-a1', description: 'task A' }), outputSnapshot('A-OUT\n'));
|
||||
tasks.add(processTask({ taskId: 'bash-b1', description: 'task B' }));
|
||||
tasks.waitDelegate = async (taskId) => {
|
||||
if (taskId === 'bash-a1') tasks.settle('bash-a1');
|
||||
return tasks.getTask(taskId);
|
||||
};
|
||||
|
||||
const { records, telemetry } = waitTelemetry();
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, telemetry, stubFlag(true)),
|
||||
context('wait_any', { timeout: 10 }),
|
||||
);
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: completed');
|
||||
expect(output).toContain('task_id: bash-a1');
|
||||
expect(output).toContain('[output]\nA-OUT');
|
||||
expect(output).toContain('[still_running]');
|
||||
expect(output).toContain('task_id: bash-b1');
|
||||
expect(tasks.waitCalls).toHaveLength(2);
|
||||
expect(tasks.waitDeliveries).toEqual([[{ taskId: 'bash-a1', status: 'completed' }]]);
|
||||
expect(lastEvent(records)?.properties).toMatchObject({
|
||||
outcome: 'completed',
|
||||
has_task_id: false,
|
||||
extra_completed_count: 0,
|
||||
});
|
||||
});
|
||||
|
||||
it('returns the still-running list on timeout without marking anything delivered', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-running9', description: 'slow task' }));
|
||||
|
||||
const { records, telemetry } = waitTelemetry();
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, telemetry, stubFlag(true)),
|
||||
context('wait_timeout', { timeout: 10, task_id: 'bash-running9' }),
|
||||
);
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: timed_out');
|
||||
expect(output).toContain('not an error');
|
||||
expect(output).toContain('[still_running]');
|
||||
expect(output).toContain('bash-running9');
|
||||
expect(tasks.waitDeliveries).toEqual([]);
|
||||
expect(lastEvent(records)?.properties).toMatchObject({
|
||||
outcome: 'timed_out',
|
||||
timeout_ms: 10_000,
|
||||
has_task_id: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('propagates an abort of the execution signal and tracks the aborted outcome', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-abort01' }));
|
||||
tasks.waitDelegate = (_taskId, _timeoutMs, waitSignal) =>
|
||||
new Promise<never>((_resolve, reject) => {
|
||||
waitSignal?.addEventListener('abort', () => reject(abortError()), { once: true });
|
||||
});
|
||||
|
||||
const { records, telemetry } = waitTelemetry();
|
||||
const controller = new AbortController();
|
||||
const pending = executeTool(
|
||||
new WaitForTool(tasks, telemetry, stubFlag(true)),
|
||||
context('wait_abort', { timeout: 600, task_id: 'bash-abort01' }, controller.signal),
|
||||
);
|
||||
controller.abort();
|
||||
|
||||
await expect(pending).rejects.toThrow('Aborted');
|
||||
expect(tasks.waitDeliveries).toEqual([]);
|
||||
expect(lastEvent(records)?.properties).toMatchObject({ outcome: 'aborted' });
|
||||
});
|
||||
|
||||
it('propagates an abort from a general wait and leaves tasks running', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-abort02' }));
|
||||
tasks.add(processTask({ taskId: 'bash-abort03' }));
|
||||
tasks.waitDelegate = (_taskId, _timeoutMs, waitSignal) =>
|
||||
new Promise<never>((_resolve, reject) => {
|
||||
waitSignal?.addEventListener('abort', () => reject(abortError()), { once: true });
|
||||
});
|
||||
|
||||
const controller = new AbortController();
|
||||
const pending = executeTool(
|
||||
new WaitForTool(tasks, recordingTelemetry([]), stubFlag(true)),
|
||||
context('wait_abort_any', { timeout: 600 }, controller.signal),
|
||||
);
|
||||
controller.abort();
|
||||
|
||||
await expect(pending).rejects.toThrow('Aborted');
|
||||
expect(tasks.getTask('bash-abort02')?.status).toBe('running');
|
||||
expect(tasks.getTask('bash-abort03')?.status).toBe('running');
|
||||
expect(tasks.waitDeliveries).toEqual([]);
|
||||
});
|
||||
|
||||
it('does not mark tasks delivered when formatting the result fails', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
const taskId = tasks.add(
|
||||
processTask({
|
||||
taskId: 'bash-fmtfail1',
|
||||
status: 'completed',
|
||||
endedAt: 1_700_000_001_000,
|
||||
exitCode: 0,
|
||||
}),
|
||||
);
|
||||
tasks.failSnapshotTaskIds.add(taskId);
|
||||
|
||||
await expect(
|
||||
executeTool(
|
||||
new WaitForTool(tasks, recordingTelemetry([]), stubFlag(true)),
|
||||
context('wait_fmt_fail', { timeout: 10, task_id: taskId }),
|
||||
),
|
||||
).rejects.toThrow('snapshot read failed');
|
||||
expect(tasks.waitDeliveries).toEqual([]);
|
||||
});
|
||||
|
||||
it('aborts the losing waits once the race resolves', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-win0001' }));
|
||||
tasks.add(processTask({ taskId: 'bash-lose001' }));
|
||||
const signals = new Map<string, AbortSignal>();
|
||||
tasks.waitDelegate = (taskId, _timeoutMs, waitSignal) => {
|
||||
signals.set(taskId, waitSignal!);
|
||||
if (taskId === 'bash-win0001') {
|
||||
tasks.settle('bash-win0001');
|
||||
return Promise.resolve(tasks.getTask(taskId));
|
||||
}
|
||||
return new Promise<AgentTaskInfo | undefined>(() => {});
|
||||
};
|
||||
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, recordingTelemetry([]), stubFlag(true)),
|
||||
context('wait_losers', { timeout: 600 }),
|
||||
);
|
||||
|
||||
expect(outputString(result)).toContain('wait_status: completed');
|
||||
expect(signals.get('bash-lose001')?.aborted).toBe(true);
|
||||
});
|
||||
|
||||
it('rejects execution when the wait_for flag is off', async () => {
|
||||
const tasks = new FakeTaskService();
|
||||
tasks.add(processTask({ taskId: 'bash-flagoff1' }));
|
||||
|
||||
const result = await executeTool(
|
||||
new WaitForTool(tasks, recordingTelemetry([]), stubFlag(false)),
|
||||
context('wait_flag_off', { timeout: 10, task_id: 'bash-flagoff1' }),
|
||||
);
|
||||
|
||||
expect(result.isError).toBe(true);
|
||||
expect(outputString(result)).toContain('wait_for experimental flag is off');
|
||||
expect(tasks.waitCalls).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('WaitForTool (harness)', () => {
|
||||
function immediateProcess(exitCode: number, stdoutText = ''): IHostProcess {
|
||||
return {
|
||||
_serviceBrand: undefined,
|
||||
stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable,
|
||||
stdout: Readable.from(stdoutText ? [stdoutText] : []),
|
||||
stderr: Readable.from([]),
|
||||
pid: 10000 + exitCode,
|
||||
exitCode,
|
||||
wait: vi.fn().mockResolvedValue(exitCode) as IHostProcess['wait'],
|
||||
kill: vi.fn().mockResolvedValue(undefined) as IHostProcess['kill'],
|
||||
dispose: vi.fn().mockResolvedValue(undefined) as IHostProcess['dispose'],
|
||||
};
|
||||
}
|
||||
|
||||
function controllableProcess(): {
|
||||
proc: IHostProcess;
|
||||
pushOutput: (text: string) => void;
|
||||
resolveWait: (code: number) => void;
|
||||
} {
|
||||
const stdout = new PassThrough();
|
||||
let resolveWait!: (code: number) => void;
|
||||
const waitPromise = new Promise<number>((resolve) => {
|
||||
resolveWait = resolve;
|
||||
});
|
||||
const proc = {
|
||||
_serviceBrand: undefined,
|
||||
stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable,
|
||||
stdout,
|
||||
stderr: Readable.from([]),
|
||||
pid: 10099,
|
||||
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;
|
||||
return {
|
||||
proc,
|
||||
pushOutput: (text) => {
|
||||
stdout.write(text);
|
||||
},
|
||||
resolveWait: (code) => {
|
||||
stdout.end();
|
||||
resolveWait(code);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function waitForTerminal(tasks: IAgentTaskService, taskId: string): Promise<void> {
|
||||
const deadline = Date.now() + 30_000;
|
||||
while (Date.now() <= deadline) {
|
||||
const info = await tasks.wait(taskId, 5);
|
||||
if (info !== undefined && TERMINAL_STATUSES.has(info.status)) return;
|
||||
await new Promise((resolve) => setTimeout(resolve, 1));
|
||||
}
|
||||
throw new Error(`Timed out waiting for task to terminate: ${taskId}`);
|
||||
}
|
||||
|
||||
it('waits for a real registered task end-to-end and suppresses its notification', async () => {
|
||||
const records: TelemetryRecord[] = [];
|
||||
const loop = stubLoopWithHooks();
|
||||
const ctx = createTestAgent(
|
||||
telemetryServices(recordingTelemetry(records)),
|
||||
agentService(IAgentLoopService, loop),
|
||||
);
|
||||
try {
|
||||
const tasks = ctx.get(IAgentTaskService);
|
||||
const tool = ctx.get(IAgentToolRegistryService).resolve('WaitFor');
|
||||
expect(tool).toBeDefined();
|
||||
|
||||
const slow = controllableProcess();
|
||||
const taskId = tasks.registerTask(new ProcessTask(slow.proc, 'echo done', 'wait target'));
|
||||
const pending = executeTool(tool!, context('wait_e2e', { timeout: 30, task_id: taskId }));
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
slow.pushOutput('DONE-OUTPUT\n');
|
||||
slow.resolveWait(0);
|
||||
const result = await pending;
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: completed');
|
||||
expect(output).toContain(`task_id: ${taskId}`);
|
||||
expect(output).toContain('[finished]');
|
||||
expect(output).toContain('[output]\nDONE-OUTPUT');
|
||||
expect(ctx.allEvents.some((event) => event.event === 'task.waitDelivered')).toBe(true);
|
||||
|
||||
expect(loop.hasPendingRequests()).toBe(false);
|
||||
loop.drainNextBatch(ctx.context);
|
||||
expect(ctx.context.get().some((message) => message.origin?.kind === 'task')).toBe(false);
|
||||
expect(ctx.allEvents.some((event) => event.event === 'task.notified')).toBe(false);
|
||||
expect(ctx.llmCalls).toHaveLength(0);
|
||||
expect(
|
||||
records.findLast((record) => record.event === 'wait_for_completed')?.properties,
|
||||
).toMatchObject({ outcome: 'completed', has_task_id: true, extra_completed_count: 0 });
|
||||
} finally {
|
||||
await ctx.dispose();
|
||||
}
|
||||
});
|
||||
|
||||
it('does not include tasks registered after the wait started', async () => {
|
||||
const ctx = createTestAgent();
|
||||
try {
|
||||
const tasks = ctx.get(IAgentTaskService);
|
||||
const tool = ctx.get(IAgentToolRegistryService).resolve('WaitFor');
|
||||
expect(tool).toBeDefined();
|
||||
|
||||
const slow = controllableProcess();
|
||||
const taskA = tasks.registerTask(new ProcessTask(slow.proc, 'sleep 30', 'slow'));
|
||||
const pending = executeTool(tool!, context('wait_race', { timeout: 30 }));
|
||||
|
||||
const late = controllableProcess();
|
||||
const taskB = tasks.registerTask(new ProcessTask(late.proc, 'echo b', 'late comer'));
|
||||
await tasks.suppressTerminalNotification(taskB);
|
||||
late.pushOutput('B-OUT\n');
|
||||
late.resolveWait(0);
|
||||
await waitForTerminal(tasks, taskB);
|
||||
|
||||
const race = await Promise.race([
|
||||
pending.then(() => 'resolved' as const),
|
||||
new Promise<'pending'>((resolve) => {
|
||||
setTimeout(() => resolve('pending'), 50);
|
||||
}),
|
||||
]);
|
||||
expect(race).toBe('pending');
|
||||
|
||||
slow.pushOutput('A-OUT\n');
|
||||
slow.resolveWait(0);
|
||||
const result = await pending;
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: completed');
|
||||
expect(output).toContain(`task_id: ${taskA}`);
|
||||
expect(output).not.toContain(taskB);
|
||||
expect(output).not.toContain('[completed_during_wait]');
|
||||
expect(ctx.allEvents.filter((event) => event.event === 'task.waitDelivered')).toHaveLength(1);
|
||||
} finally {
|
||||
await ctx.dispose();
|
||||
}
|
||||
});
|
||||
|
||||
it('returns from a wait on a task that never settles once the timeout elapses', async () => {
|
||||
const ctx = createTestAgent();
|
||||
try {
|
||||
const tasks = ctx.get(IAgentTaskService);
|
||||
const tool = ctx.get(IWaitForTool);
|
||||
const taskId = tasks.registerTask(
|
||||
new SubagentTask(
|
||||
{
|
||||
agentId: 'agent-hang',
|
||||
profileName: 'coder',
|
||||
completion: new Promise<{ result: string }>(() => {}),
|
||||
},
|
||||
'hung work',
|
||||
new AbortController(),
|
||||
),
|
||||
);
|
||||
|
||||
const result = await executeTool(tool, context('wait_hang', { timeout: 1, task_id: taskId }));
|
||||
const output = outputString(result);
|
||||
|
||||
expect(result.isError ?? false).toBe(false);
|
||||
expect(output).toContain('wait_status: timed_out');
|
||||
expect(output).toContain('[still_running]');
|
||||
expect(output).toContain(taskId);
|
||||
} finally {
|
||||
await ctx.dispose();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest';
|
|||
import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory';
|
||||
import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants';
|
||||
import { ContextApplyCompaction } from '#/agent/contextMemory/contextEvents';
|
||||
import type { TaskOrigin } from '#/agent/contextMemory/types';
|
||||
import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction';
|
||||
import { IAgentLoopService } from '#/agent/loop/loop';
|
||||
import { MessageStepRequest } from '#/agent/loop/stepRequest';
|
||||
|
|
@ -10,6 +11,8 @@ import { turnKey } from '#/agent/loop/turnOps';
|
|||
import { IAgentPlanService } from '#/features/plan/plan';
|
||||
import { planKey } from '#/features/plan/planOps';
|
||||
import { IAgentPromptService } from '#/agent/prompt/prompt';
|
||||
import { IAgentTaskService, type AgentTask } from '#/agent/task/task';
|
||||
import { taskNotificationDeliveryKey } from '#/agent/task/taskService';
|
||||
import { IAgentConversationUndoService } from '#/agent/undo/undo';
|
||||
import { ContextUndone } from '#/agent/undo/undoService';
|
||||
import { AgentStatusUpdated } from '#/agent/usage/usageEvents';
|
||||
|
|
@ -501,4 +504,41 @@ describe('AgentConversationUndoService', () => {
|
|||
expect(wireEvents).toContain('context.undo');
|
||||
expect(wireEvents).not.toContain('log.cut');
|
||||
});
|
||||
|
||||
it('re-delivers wait-reported task notifications after conversation undo', async () => {
|
||||
setup();
|
||||
const undo = ctx.get(IAgentConversationUndoService);
|
||||
const tasks = ctx.get(IAgentTaskService);
|
||||
ctx.appendTurnExchange('u1', 'a1');
|
||||
|
||||
const completingTask = (output: string): AgentTask => ({
|
||||
idPrefix: 'test',
|
||||
kind: 'process',
|
||||
description: 'fake process task',
|
||||
start: async (sink) => {
|
||||
sink.appendOutput(output);
|
||||
await sink.settle({ status: 'completed' });
|
||||
},
|
||||
toInfo: (base) => ({ ...base, kind: 'process', command: 'echo', pid: 0, exitCode: null }),
|
||||
});
|
||||
|
||||
const taskA = tasks.registerTask(completingTask('a\n'));
|
||||
const taskB = tasks.registerTask(completingTask('b\n'));
|
||||
tasks.markTasksDeliveredViaWait([
|
||||
{ taskId: taskA, status: 'completed' },
|
||||
{ taskId: taskB, status: 'completed' },
|
||||
]);
|
||||
await tasks.wait(taskA, 1000);
|
||||
await tasks.wait(taskB, 1000);
|
||||
|
||||
expect(ctx.context.get().some((message) => message.origin?.kind === 'task')).toBe(false);
|
||||
expect(ctx.agentState.get(taskNotificationDeliveryKey)).toHaveLength(2);
|
||||
|
||||
await undo.undo(1);
|
||||
|
||||
const redelivered = ctx.context.get().filter((message) => message.origin?.kind === 'task');
|
||||
expect(redelivered.map((message) => (message.origin as TaskOrigin).taskId).sort()).toEqual(
|
||||
[taskA, taskB].sort(),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -82,6 +82,7 @@ const V2_RECORD_TYPES: ReadonlySet<string> = new Set([
|
|||
'tower_mode.exit',
|
||||
'task.started',
|
||||
'task.terminated',
|
||||
'task.waitDelivered',
|
||||
'interaction.request',
|
||||
'interaction.resolved',
|
||||
'plan.revision',
|
||||
|
|
|
|||
|
|
@ -591,6 +591,9 @@ function createFakeTaskService(options: { maxRunningTasks?: number } = {}): {
|
|||
async suppressTerminalNotification(): Promise<void> {
|
||||
},
|
||||
|
||||
markTasksDeliveredViaWait(): void {
|
||||
},
|
||||
|
||||
detach(taskId: string): AgentTaskInfo | undefined {
|
||||
const entry = tasks.get(taskId);
|
||||
if (entry === undefined) return undefined;
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -414,7 +414,7 @@ describe('Agent resume', () => {
|
|||
expect(ctx.llmInputs()).toMatchInlineSnapshot(`
|
||||
call 1:
|
||||
system: <system-prompt>
|
||||
tools: Agent, AgentSwarm, AskUserQuestion, Bash, CreateGoal, Edit, EnterPlanMode, ExitPlanMode, FetchURL, GetGoal, Glob, Grep, Read, SetGoalBudget, Skill, TaskList, TaskOutput, TaskStop, TodoList, UpdateGoal, Write
|
||||
tools: Agent, AgentSwarm, AskUserQuestion, Bash, CreateGoal, Edit, EnterPlanMode, ExitPlanMode, FetchURL, GetGoal, Glob, Grep, Read, SetGoalBudget, Skill, TaskList, TaskOutput, TaskStop, TodoList, UpdateGoal, WaitFor, Write
|
||||
messages:
|
||||
user: text "Historical prompt before skill"
|
||||
assistant: [] calls call_resume_write:Write { "path": "result.txt" }, call_resume_skill:Skill { "skill": "review" }
|
||||
|
|
|
|||
|
|
@ -449,8 +449,9 @@ function projectResumedAgents(
|
|||
* the engines (the subagent/cron docs embed engine-specific facts), and
|
||||
* v1 additionally registers the `select_tools` meta tool v2 has no
|
||||
* counterpart for — both are engine design, not resume data. v2's default
|
||||
* profile also carries `TowerInit` (the tower-mode entry point); tower is
|
||||
* v2-only, so the tool is projected out of both rosters. A model-less
|
||||
* profile also carries `TowerInit` (the tower-mode entry point) and
|
||||
* `WaitFor` (the background-task wait primitive); both are v2-only, so the
|
||||
* tools are projected out of both rosters. A model-less
|
||||
* agent's roster is not compared at all (v1 initializes builtin tools
|
||||
* only on a profiled agent; v2 exposes them unbound).
|
||||
*/
|
||||
|
|
@ -469,6 +470,7 @@ function projectResumedAgent(agent: ResumedAgentState, home: HomePair): unknown
|
|||
projected['tools'] = tools
|
||||
.filter((tool) => tool['name'] !== 'select_tools')
|
||||
.filter((tool) => tool['name'] !== 'TowerInit')
|
||||
.filter((tool) => tool['name'] !== 'WaitFor')
|
||||
.map((tool) => ({ name: tool['name'], active: tool['active'], source: tool['source'] }))
|
||||
.toSorted((a, b) => String(a.name).localeCompare(String(b.name)));
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue