From 95b560b3ccd45db7d05657e2e79b82dfc43f8d88 Mon Sep 17 00:00:00 2001 From: chengluyu <2239547+chengluyu@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:22:52 +0800 Subject: [PATCH] fix(agent-core-v2): enforce goal deadlines --- .changeset/enforce-goal-deadlines.md | 5 + .../src/agent/goal/goalDeadlineScheduler.ts | 19 ++ .../goal/goalDeadlineSchedulerService.ts | 40 +++ .../src/agent/goal/goalService.ts | 113 ++++++- packages/agent-core-v2/src/index.ts | 2 + .../test/agent/goal/goal.test.ts | 281 ++++++++++++++++-- .../test/agent/goal/goalOps.test.ts | 9 + 7 files changed, 429 insertions(+), 40 deletions(-) create mode 100644 .changeset/enforce-goal-deadlines.md create mode 100644 packages/agent-core-v2/src/agent/goal/goalDeadlineScheduler.ts create mode 100644 packages/agent-core-v2/src/agent/goal/goalDeadlineSchedulerService.ts diff --git a/.changeset/enforce-goal-deadlines.md b/.changeset/enforce-goal-deadlines.md new file mode 100644 index 0000000000..c2b997e071 --- /dev/null +++ b/.changeset/enforce-goal-deadlines.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Enforce goal wall-clock budgets while model or tool work is still running. diff --git a/packages/agent-core-v2/src/agent/goal/goalDeadlineScheduler.ts b/packages/agent-core-v2/src/agent/goal/goalDeadlineScheduler.ts new file mode 100644 index 0000000000..4591c2b1a9 --- /dev/null +++ b/packages/agent-core-v2/src/agent/goal/goalDeadlineScheduler.ts @@ -0,0 +1,19 @@ +/** + * `goal` domain (L4) — wall-clock deadline scheduling contract. + * + * Defines the App-scoped `IGoalDeadlineScheduler` used by per-agent goal + * services to measure active time and arm hard wall-clock budget deadlines. + */ + +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; +import type { IDisposable } from '#/_base/di/lifecycle'; + +export interface IGoalDeadlineScheduler { + readonly _serviceBrand: undefined; + + now(): number; + schedule(delayMs: number, callback: () => void): IDisposable; +} + +export const IGoalDeadlineScheduler: ServiceIdentifier = + createDecorator('goalDeadlineScheduler'); diff --git a/packages/agent-core-v2/src/agent/goal/goalDeadlineSchedulerService.ts b/packages/agent-core-v2/src/agent/goal/goalDeadlineSchedulerService.ts new file mode 100644 index 0000000000..156fc26b73 --- /dev/null +++ b/packages/agent-core-v2/src/agent/goal/goalDeadlineSchedulerService.ts @@ -0,0 +1,40 @@ +/** + * `goal` domain (L4) — `IGoalDeadlineScheduler` implementation. + * + * Measures monotonic elapsed time and schedules disposable one-shot deadlines + * with the host timer API. Bound at App scope. + */ + +import { InstantiationType } from '#/_base/di/extensions'; +import { toDisposable, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; + +import { IGoalDeadlineScheduler } from './goalDeadlineScheduler'; + +export class GoalDeadlineSchedulerService implements IGoalDeadlineScheduler { + declare readonly _serviceBrand: undefined; + + now(): number { + return Number(process.hrtime.bigint() / 1_000_000n); + } + + schedule(delayMs: number, callback: () => void): IDisposable { + let timeout: ReturnType | undefined = setTimeout(() => { + timeout = undefined; + callback(); + }, Math.max(0, delayMs)); + timeout.unref?.(); + return toDisposable(() => { + if (timeout !== undefined) clearTimeout(timeout); + timeout = undefined; + }); + } +} + +registerScopedService( + LifecycleScope.App, + IGoalDeadlineScheduler, + GoalDeadlineSchedulerService, + InstantiationType.Delayed, + 'goal', +); diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index ab4794a8aa..73aec20a9b 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -1,5 +1,5 @@ /** - * `goal` domain (L4) - `IAgentGoalService` implementation. + * `goal` domain (L4) — `IAgentGoalService` implementation. * * Owns the per-agent goal lifecycle; persists the goal in the `wire` * `GoalModel` (`GoalState | null`) through the `goal.create` / `goal.update` / @@ -20,15 +20,17 @@ * loop pops it), accounts live * turn usage through `usage`, writes system reminders through * `systemReminder`, registers model tools through `toolRegistry`, and reports - * telemetry through `telemetry`. Bound at Agent scope. + * telemetry through `telemetry`. Measures time and arms hard deadlines through + * `goal`'s App-scoped deadline scheduler. Bound at Agent scope. */ import { randomUUID } from 'node:crypto'; import type { TurnEndedEvent, TurnStartedEvent } from '@moonshot-ai/protocol'; -import { Disposable } from '#/_base/di/lifecycle'; +import { Disposable, MutableDisposable, type IDisposable } from '#/_base/di/lifecycle'; import { InstantiationType } from '#/_base/di/extensions'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; +import { abortError } from '#/_base/utils/abort'; import { isPlainRecord } from '#/_base/utils/canonical-args'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import type { ContextMessage, PromptOrigin } from '#/agent/contextMemory/types'; @@ -60,6 +62,7 @@ import { defineModel } from '#/wire/model'; import { IEventBus } from '#/app/event/eventBus'; import { IAgentGoalService, type GoalReasonInput, type ResumeGoalInput } from './goal'; +import { IGoalDeadlineScheduler } from './goalDeadlineScheduler'; import { clearGoal, createGoal, GoalModel, updateGoal, type GoalState } from './goalOps'; import type { CreateGoalInput, @@ -205,6 +208,8 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { private readonly pendingContinuationGoals = new Map(); private readonly goalTurnTargets = new Map(); private readonly exhaustedTurnBudgetGoals = new Map(); + private readonly wallClockDeadline = this._register(new MutableDisposable()); + private liveWallClockStartedAt?: number; private pendingContinuation?: PendingContinuation; constructor( @@ -217,6 +222,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { @IAgentToolExecutorService toolExecutor: IAgentToolExecutorService, @IAgentUsageService usageService: IAgentUsageService, @IConfigService private readonly config: IConfigService, + @IGoalDeadlineScheduler private readonly deadlineScheduler: IGoalDeadlineScheduler, ) { super(); this._register( @@ -321,8 +327,10 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { wallClockResumedAt, }), ); + this.liveWallClockStartedAt = this.deadlineScheduler.now(); this.adoptStarterTurn(actor); const state = this.requireState(); + this.refreshWallClockDeadline(state); this.emitGoalUpdated(this.toSnapshot(state)); this.telemetry.track2('goal_created', { actor, replace: input.replace === true }); return this.toSnapshot(state); @@ -413,12 +421,18 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { actor, ...budgetTelemetryProperties(input.budgetLimits), }); - return this.blockIfBudgetReached(next) ?? this.toSnapshot(next); + const blocked = this.blockIfBudgetReached(next); + if (blocked !== null) return blocked; + this.refreshWallClockDeadline(next); + return this.toSnapshot(next); } async cancelGoal(_input: GoalReasonInput = {}, actor: GoalActor = 'user'): Promise { const state = this.requireState(); const snapshot = this.toSnapshot(state); + if (state.status === 'active' && this.liveTurnId !== undefined) { + this.loopService.cancel(this.liveTurnId); + } this.clearInternal(actor); if (actor === 'user') { this.reminders.appendSystemReminder(GOAL_CANCELLED_REMINDER, { @@ -750,17 +764,32 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { return this.goalTurnTargets.get(turnId) ?? this.goalDrivenTurns.get(turnId); } - private cancelPendingContinuation(preserveLiveContinuation = false): void { + private cancelPendingContinuation( + preserveLiveContinuation = false, + reason?: unknown, + ): void { const pending = this.pendingContinuation; if (preserveLiveContinuation && pending?.turnId === this.liveTurnId) return; this.pendingContinuation = undefined; - if (pending !== undefined && !pending.receipt.abort() && pending.turnId !== undefined) { - this.loopService.cancel(pending.turnId); + const aborted = + reason === undefined ? pending?.receipt.abort() : pending?.receipt.abort(reason); + if ( + pending !== undefined && + !aborted && + pending.turnId !== undefined + ) { + if (reason === undefined) { + this.loopService.cancel(pending.turnId); + } else { + this.loopService.cancel(pending.turnId, reason); + } } } private normalizeAfterReplay(): void { this.appendForkClearedReminder(); + this.wallClockDeadline.clear(); + this.liveWallClockStartedAt = undefined; const state = this.goalState; if (state === null) return; if (state.status === 'complete') { @@ -795,6 +824,8 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { ): void { if (this.goalState === null) return; this.cancelPendingContinuation(opts.preserveLiveContinuation === true); + this.wallClockDeadline.clear(); + this.liveWallClockStartedAt = undefined; this.wire.dispatch(clearGoal({})); if (opts.emit !== false) this.emitGoalUpdated(null); if (opts.track !== false) this.telemetry.track2('goal_cleared', { actor }); @@ -805,18 +836,29 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { status: GoalStatus, reason: string | undefined, actor: GoalActor, - opts: { readonly preserveLiveContinuation?: boolean } = {}, + opts: { + readonly preserveLiveContinuation?: boolean; + readonly cancellationReason?: unknown; + } = {}, ): GoalSnapshot { const wallClockMs = this.settleWallClock(state); const wallClockResumedAt = status === 'active' ? Date.now() : undefined; - if (status !== 'active' && state.status === 'active') { - this.cancelPendingContinuation(opts.preserveLiveContinuation === true); + if (status === 'active') { + this.liveWallClockStartedAt = this.deadlineScheduler.now(); + } else if (state.status === 'active') { + this.cancelPendingContinuation( + opts.preserveLiveContinuation === true, + opts.cancellationReason, + ); + this.wallClockDeadline.clear(); + this.liveWallClockStartedAt = undefined; } this.wire.dispatch( updateGoal({ status, reason, wallClockMs, wallClockResumedAt, actor }), ); const next = this.requireState(); if (status === 'active') this.adoptStarterTurn(actor); + if (status === 'active') this.refreshWallClockDeadline(next); this.emitGoalUpdated(this.toSnapshot(next), { kind: 'lifecycle', status, reason, actor }); this.trackStatusChanged(next, actor); return this.toSnapshot(next); @@ -846,6 +888,12 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { } private settleWallClock(state: GoalState): number { + if (state.status === 'active' && this.liveWallClockStartedAt !== undefined) { + return ( + state.wallClockMs + + Math.max(0, this.deadlineScheduler.now() - this.liveWallClockStartedAt) + ); + } if (state.status === 'active' && state.wallClockResumedAt !== undefined) { return state.wallClockMs + Math.max(0, Date.now() - state.wallClockResumedAt); } @@ -853,6 +901,12 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { } private liveWallClockMs(state: GoalState): number { + if (state.status === 'active' && this.liveWallClockStartedAt !== undefined) { + return ( + state.wallClockMs + + Math.max(0, this.deadlineScheduler.now() - this.liveWallClockStartedAt) + ); + } if (state.status === 'active' && state.wallClockResumedAt !== undefined) { return state.wallClockMs + Math.max(0, Date.now() - state.wallClockResumedAt); } @@ -890,6 +944,45 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { preserveLiveContinuation: true, }); } + + private refreshWallClockDeadline(state: GoalState): void { + this.wallClockDeadline.clear(); + const budgetMs = state.budgetLimits.wallClockBudgetMs; + if ( + state.status !== 'active' || + budgetMs === undefined || + this.liveWallClockStartedAt === undefined + ) { + return; + } + const remainingMs = Math.max(0, budgetMs - this.liveWallClockMs(state)); + this.wallClockDeadline.value = this.deadlineScheduler.schedule(remainingMs, () => { + this.handleWallClockDeadline(); + }); + } + + private handleWallClockDeadline(): void { + this.wallClockDeadline.clear(); + const state = this.goalState; + if (state === null || state.status !== 'active') return; + const budgetMs = state.budgetLimits.wallClockBudgetMs; + if (budgetMs === undefined) return; + if (this.liveWallClockMs(state) < budgetMs) { + this.refreshWallClockDeadline(state); + return; + } + const reason = goalBudgetBlockReason(this.toSnapshot(state).budget); + if (reason === undefined) return; + const cancellation = abortError(reason); + const liveTurnId = this.liveTurnId; + const pendingTurnId = this.pendingContinuation?.turnId; + this.applyLifecycle(state, 'blocked', reason, 'runtime', { + cancellationReason: cancellation, + }); + if (liveTurnId !== undefined && liveTurnId !== pendingTurnId) { + this.loopService.cancel(liveTurnId, cancellation); + } + } } function computeBudgetReport(state: GoalState, wallClockMs: number): GoalBudgetReport { diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 7c819c0181..6774c4ebb6 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -182,6 +182,8 @@ import '#/agent/goal/tools/create-goal'; import '#/agent/goal/tools/get-goal'; import '#/agent/goal/tools/set-goal-budget'; import '#/agent/goal/tools/update-goal'; +export * from '#/agent/goal/goalDeadlineScheduler'; +import '#/agent/goal/goalDeadlineSchedulerService'; export * from '#/agent/goal/goal'; export * from '#/agent/goal/goalService'; export * from '#/agent/goal/types'; diff --git a/packages/agent-core-v2/test/agent/goal/goal.test.ts b/packages/agent-core-v2/test/agent/goal/goal.test.ts index 3e9aa33e0a..7688f2292e 100644 --- a/packages/agent-core-v2/test/agent/goal/goal.test.ts +++ b/packages/agent-core-v2/test/agent/goal/goal.test.ts @@ -8,9 +8,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type { TurnEndedEvent } from '@moonshot-ai/protocol'; +import type { IDisposable } from '#/_base/di/lifecycle'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import { USER_PROMPT_ORIGIN } from '#/agent/contextMemory/types'; import { IAgentGoalService } from '#/agent/goal/goal'; +import { IGoalDeadlineScheduler } from '#/agent/goal/goalDeadlineScheduler'; import { type AgentGoalService } from '#/agent/goal/goalService'; import { UpdateGoalTool, UpdateGoalToolInputSchema } from '#/agent/goal/tools/update-goal'; import { IAgentLoopService, type AfterStepContext, type EnqueueReceipt, type Step, type Turn } from '#/agent/loop/loop'; @@ -20,6 +22,7 @@ import { type ToolExecutionResult, } from '#/agent/toolExecutor/toolExecutor'; import type { ToolBeforeExecuteContext } from '#/agent/toolExecutor/toolHooks'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; import { IAgentUsageService } from '#/agent/usage/usage'; import type { WireRecord } from '#/wire/record'; import { type DomainEvent, IEventBus } from '#/app/event/eventBus'; @@ -27,9 +30,11 @@ import { APIConnectionError, APIStatusError } from '#/app/llmProtocol/errors'; import type { ToolCall } from '#/app/llmProtocol/message'; import type { TokenUsage } from '#/app/llmProtocol/usage'; import { ErrorCodes, Error2, errorInfo, toKimiErrorPayload } from '#/errors'; +import type { ExecutableTool } from '#/tool/toolContract'; import { InMemoryWireRecordPersistence, + appService, agentService, createTestAgent, permissionModeServices, @@ -51,6 +56,103 @@ type TurnEndedInput = { readonly error?: unknown; }; +interface ManualDeadline { + readonly dueAt: number; + readonly callback: () => void; + cancelled: boolean; +} + +class ManualGoalDeadlineScheduler implements IGoalDeadlineScheduler { + declare readonly _serviceBrand: undefined; + + private currentTime = 0; + private readonly deadlines = new Set(); + + now(): number { + return this.currentTime; + } + + schedule(delayMs: number, callback: () => void): IDisposable { + const deadline: ManualDeadline = { + dueAt: this.currentTime + Math.max(0, delayMs), + callback, + cancelled: false, + }; + this.deadlines.add(deadline); + return { + dispose: () => { + deadline.cancelled = true; + this.deadlines.delete(deadline); + }, + }; + } + + advanceBy(deltaMs: number): void { + this.currentTime += deltaMs; + while (true) { + const due = [...this.deadlines] + .filter((deadline) => !deadline.cancelled && deadline.dueAt <= this.currentTime) + .toSorted((left, right) => left.dueAt - right.dueAt)[0]; + if (due === undefined) return; + this.deadlines.delete(due); + due.callback(); + } + } +} + +function deferred(): { readonly promise: Promise; readonly resolve: () => void } { + let resolve!: () => void; + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise; + }); + return { promise, resolve }; +} + +function waitForAbort(signal: AbortSignal): Promise { + if (signal.aborted) return Promise.reject(signal.reason); + return new Promise((_resolve, reject) => { + signal.addEventListener( + 'abort', + () => { + reject(signal.reason); + }, + { once: true }, + ); + }); +} + +function blockingGenerate(): { + readonly generate: NonNullable; + readonly started: Promise; + readonly signal: () => AbortSignal; +} { + const started = deferred(); + let activeSignal: AbortSignal | undefined; + const generate: NonNullable = async ( + _chat, + _systemPrompt, + _tools, + _history, + _callbacks, + options, + ) => { + const signal = options?.signal; + if (signal === undefined) throw new Error('Expected an LLM abort signal'); + options?.onRequestStart?.(); + activeSignal = signal; + started.resolve(); + return waitForAbort(signal); + }; + return { + generate, + started: started.promise, + signal: () => { + if (activeSignal === undefined) throw new Error('LLM request has not started'); + return activeSignal; + }, + }; +} + const zeroUsage: TokenUsage = { inputCacheRead: 0, inputCacheCreation: 0, @@ -589,10 +691,13 @@ describe('AgentGoalService core workflow hooks', () => { let toolExecutor: IAgentToolExecutorService; let usageService: IAgentUsageService; let eventBus: IEventBus; + let clock: ManualGoalDeadlineScheduler; beforeEach(() => { loopService = stubLoopWithHooks(); + clock = new ManualGoalDeadlineScheduler(); ctx = createTestAgent( + appService(IGoalDeadlineScheduler, clock), agentService(IAgentLoopService, loopService), permissionModeServices('auto'), ); @@ -933,40 +1038,32 @@ describe('AgentGoalService core workflow hooks', () => { it.each(['turn', 'token', 'wall-clock'] as const)( 'keeps a goal blocked when its %s budget is exhausted before resume', async (budget) => { - if (budget === 'wall-clock') { - vi.useFakeTimers(); - vi.setSystemTime(new Date('2026-01-01T00:00:00Z')); + await goals.createGoal({ objective: 'finish the task' }); + if (budget === 'turn') { + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); + } else if (budget === 'token') { + await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 1 } }, 'model'); + } else { + await goals.setBudgetLimits({ budgetLimits: { wallClockBudgetMs: 1 } }, 'model'); + clock.advanceBy(1); } - try { - await goals.createGoal({ objective: 'finish the task' }); - if (budget === 'turn') { - await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); - } else if (budget === 'token') { - await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 1 } }, 'model'); - } else { - await goals.setBudgetLimits({ budgetLimits: { wallClockBudgetMs: 1 } }, 'model'); - vi.advanceTimersByTime(1); - } - const turn = makeTurn(101); - eventBus.publish({ type: 'turn.started', turnId: turn.id, origin: USER_PROMPT_ORIGIN }); - if (budget === 'token') { - recordStepUsage(usageService, goals, turn, { ...zeroUsage, output: 1 }); - } else { - await runGoalStep(loopService, turn); - } - endTurn(eventBus, turn); - expect(loopService.status()).toMatchObject({ state: 'idle', hasPendingRequests: false }); + const turn = makeTurn(101); + eventBus.publish({ type: 'turn.started', turnId: turn.id, origin: USER_PROMPT_ORIGIN }); + if (budget === 'token') { + recordStepUsage(usageService, goals, turn, { ...zeroUsage, output: 1 }); + } else { + await runGoalStep(loopService, turn); + } + endTurn(eventBus, turn); + expect(loopService.status()).toMatchObject({ state: 'idle', hasPendingRequests: false }); - const resumed = await goals.resumeGoal({ continueIfBlocked: true }); + const resumed = await goals.resumeGoal({ continueIfBlocked: true }); - expect(resumed.status).toBe('blocked'); - expect(resumed.budget.overBudget).toBe(true); - expect(resumed.terminalReason).toMatch(/^Blocked after goal budget reached:/); - expect(loopService.launches).toEqual([]); - } finally { - if (budget === 'wall-clock') vi.useRealTimers(); - } + expect(resumed.status).toBe('blocked'); + expect(resumed.budget.overBudget).toBe(true); + expect(resumed.terminalReason).toMatch(/^Blocked after goal budget reached:/); + expect(loopService.launches).toEqual([]); }, ); @@ -1489,6 +1586,130 @@ describe('goal pause classification on provider errors', () => { }); }); +describe('AgentGoalService hard wall-clock deadline', () => { + it('aborts an in-flight LLM request when the wall-clock budget expires', async () => { + const clock = new ManualGoalDeadlineScheduler(); + const llm = blockingGenerate(); + const ctx = createTestAgent(appService(IGoalDeadlineScheduler, clock), { + generate: llm.generate, + }); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + await ctx + .get(IAgentGoalService) + .setBudgetLimits({ budgetLimits: { wallClockBudgetMs: 1_000 } }, 'user'); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await llm.started; + clock.advanceBy(1_000); + + expect(llm.signal().aborted).toBe(true); + const events = await ctx.untilTurnEnd(); + expect(events).toContainEqual( + expect.objectContaining({ + event: 'turn.ended', + args: expect.objectContaining({ reason: 'cancelled' }), + }), + ); + expect((await ctx.rpc.getGoal({})).goal).toMatchObject({ + status: 'blocked', + wallClockMs: 1_000, + budget: { wallClockBudgetReached: true }, + terminalReason: 'Blocked after goal budget reached: wall-clock budget 1000ms', + }); + } finally { + await ctx.dispose(); + } + }); + + it('aborts an in-flight tool execution when the wall-clock budget expires', async () => { + const clock = new ManualGoalDeadlineScheduler(); + const toolStarted = deferred(); + let toolSignal: AbortSignal | undefined; + const tool: ExecutableTool = { + name: 'SlowWork', + description: 'Wait for cancellation.', + parameters: { type: 'object', properties: {}, additionalProperties: false }, + resolveExecution: () => ({ + approvalRule: 'SlowWork', + accesses: [], + execute: async ({ signal }) => { + toolSignal = signal; + toolStarted.resolve(); + return waitForAbort(signal); + }, + }), + }; + const ctx = createTestAgent( + appService(IGoalDeadlineScheduler, clock), + permissionModeServices('yolo'), + ); + try { + ctx.get(IAgentToolRegistryService).register(tool); + ctx.configure({ tools: ['SlowWork'] }); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + await ctx + .get(IAgentGoalService) + .setBudgetLimits({ budgetLimits: { wallClockBudgetMs: 1_000 } }, 'user'); + ctx.mockNextResponse({ + type: 'function', + id: 'slow_work', + name: 'SlowWork', + arguments: '{}', + }); + + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await toolStarted.promise; + clock.advanceBy(1_000); + + expect(toolSignal?.aborted).toBe(true); + const events = await ctx.untilTurnEnd(); + expect(events).toContainEqual( + expect.objectContaining({ + event: 'turn.ended', + args: expect.objectContaining({ reason: 'cancelled' }), + }), + ); + expect((await ctx.rpc.getGoal({})).goal).toMatchObject({ + status: 'blocked', + budget: { wallClockBudgetReached: true }, + }); + } finally { + await ctx.dispose(); + } + }); + + it('keeps user cancellation authoritative when it precedes the wall-clock deadline', async () => { + const clock = new ManualGoalDeadlineScheduler(); + const llm = blockingGenerate(); + const ctx = createTestAgent(appService(IGoalDeadlineScheduler, clock), { + generate: llm.generate, + }); + try { + ctx.configure(); + await ctx.rpc.createGoal({ objective: 'finish bounded work' }); + await ctx + .get(IAgentGoalService) + .setBudgetLimits({ budgetLimits: { wallClockBudgetMs: 1_000 } }, 'user'); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'start work' }] }); + await llm.started; + + await ctx.rpc.cancelGoal({}); + expect(llm.signal()).toMatchObject({ + aborted: true, + reason: expect.objectContaining({ userCancelled: true }), + }); + clock.advanceBy(1_000); + + await ctx.untilTurnEnd(); + expect((await ctx.rpc.getGoal({})).goal).toBeNull(); + } finally { + await ctx.dispose(); + } + }); +}); + describe('AgentGoalService mid-turn budget stop', () => { it('grants one tool-free grace step when a token budget is reached mid-turn', async () => { const ctx = createTestAgent(); diff --git a/packages/agent-core-v2/test/agent/goal/goalOps.test.ts b/packages/agent-core-v2/test/agent/goal/goalOps.test.ts index 98b43c8de2..b55c4aa053 100644 --- a/packages/agent-core-v2/test/agent/goal/goalOps.test.ts +++ b/packages/agent-core-v2/test/agent/goal/goalOps.test.ts @@ -1,3 +1,9 @@ +/** + * Scenario: wire-backed goal lifecycle persistence and replay. + * Responsibilities: verify goal Ops, live events, and replay normalization through the service contract. + * Wiring: real goal/wire/event/deadline services with non-persistence collaborators stubbed. + * Run: `pnpm --filter @moonshot-ai/agent-core-v2 exec vitest run test/agent/goal/goalOps.test.ts`. + */ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; @@ -10,6 +16,8 @@ import { IConfigService } from '#/app/config/config'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import { IAgentGoalService } from '#/agent/goal/goal'; +import { IGoalDeadlineScheduler } from '#/agent/goal/goalDeadlineScheduler'; +import { GoalDeadlineSchedulerService } from '#/agent/goal/goalDeadlineSchedulerService'; import { AgentGoalService } from '#/agent/goal/goalService'; import { GoalModel } from '#/agent/goal/goalOps'; import { IAgentLoopService } from '#/agent/loop/loop'; @@ -114,6 +122,7 @@ function buildHost(key: string): { ix.stub(ITelemetryService, createTelemetryStub()); ix.stub(IAgentToolExecutorService, createToolExecutorStub()); ix.stub(IConfigService, createConfigStub()); + ix.set(IGoalDeadlineScheduler, new SyncDescriptor(GoalDeadlineSchedulerService)); ix.set(IAgentGoalService, new SyncDescriptor(AgentGoalService)); const wire = registerTestAgentWire(ix, testWireScope(SCOPE, key), { log: ix.get(IAppendLogStore),