From 2627c9067d5f67fc28038f2054bb1d87314a8277 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 1 Jul 2026 18:57:49 +0800 Subject: [PATCH 1/6] refactor: merge session subagent host service --- .../src/agent/background/agent-task.ts | 4 +- .../src/session/subagentHost/agentTool.ts | 2 +- .../defaultSessionSubagentHost.ts | 558 ----------------- .../src/session/subagentHost/index.ts | 2 +- .../src/session/subagentHost/subagentHost.ts | 20 +- .../subagentHost/subagentHostService.ts | 583 ++++++++++++++++-- .../agent-core-v2/test/background/ids.test.ts | 4 +- .../test/background/manager.test.ts | 4 +- .../test/background/rpc-events.test.ts | 4 +- packages/agent-core-v2/test/harness/agent.ts | 104 +++- .../test/subagentHost/agent-tool.test.ts | 17 +- ...ts => subagentHostService.default.test.ts} | 93 ++- .../subagentHost/subagentHostService.test.ts | 6 +- .../agent-core-v2/test/swarm/swarm.test.ts | 15 +- packages/agent-core-v2/test/tool/tool.test.ts | 16 +- packages/agent-core-v2/test/turn/turn.test.ts | 25 +- 16 files changed, 765 insertions(+), 692 deletions(-) delete mode 100644 packages/agent-core-v2/src/session/subagentHost/defaultSessionSubagentHost.ts rename packages/agent-core-v2/test/subagentHost/{defaultSessionSubagentHost.test.ts => subagentHostService.default.test.ts} (88%) diff --git a/packages/agent-core-v2/src/agent/background/agent-task.ts b/packages/agent-core-v2/src/agent/background/agent-task.ts index 6c308386b..327e37aff 100644 --- a/packages/agent-core-v2/src/agent/background/agent-task.ts +++ b/packages/agent-core-v2/src/agent/background/agent-task.ts @@ -17,7 +17,7 @@ export type SubagentHandle = { readonly completion: Promise; }; -export interface SessionSubagentHost { +interface SubagentDetachHost { markActiveChildDetached(agentId: string): void; } @@ -46,7 +46,7 @@ export class AgentBackgroundTask implements BackgroundTask { constructor( private readonly handle: SubagentHandle, readonly description: string, - private readonly subagentHost: Pick, + private readonly subagentHost: SubagentDetachHost, private readonly abortController: AbortController, ) { this.agentId = handle.agentId; diff --git a/packages/agent-core-v2/src/session/subagentHost/agentTool.ts b/packages/agent-core-v2/src/session/subagentHost/agentTool.ts index cff7ac029..9661183bb 100644 --- a/packages/agent-core-v2/src/session/subagentHost/agentTool.ts +++ b/packages/agent-core-v2/src/session/subagentHost/agentTool.ts @@ -2,7 +2,7 @@ * AgentTool — collaboration tool for spawning task subagents. * * Unlike the built-in tools (Read/Write/Edit/Bash/Grep/Glob), this is a - * "collaboration tool". It uses `ISessionSubagentHost` (injected via the constructor + * "collaboration tool". It uses the session subagent host (injected via the constructor * rather than through the runtime) to create in-process subagent loop instances. * * Foreground and background subagents both run through the background service. diff --git a/packages/agent-core-v2/src/session/subagentHost/defaultSessionSubagentHost.ts b/packages/agent-core-v2/src/session/subagentHost/defaultSessionSubagentHost.ts deleted file mode 100644 index 229142a66..000000000 --- a/packages/agent-core-v2/src/session/subagentHost/defaultSessionSubagentHost.ts +++ /dev/null @@ -1,558 +0,0 @@ -import { - APIProviderRateLimitError, - isProviderRateLimitError, - type TokenUsage, -} from '@moonshot-ai/kosong'; - -import { linkAbortSignal, userCancellationReason } from '#/_base/utils/abort'; -import { IAgentLifecycleService } from '#/session/agent-lifecycle'; -import type { IScopeHandle } from '#/_base/di/scope'; -import { - IAgentContextMemoryService, - type ContextMessage, - type PromptOrigin, -} from '#/agent/contextMemory'; -import { ErrorCodes, toKimiErrorPayload, type KimiErrorPayload } from '#/errors'; -import { IAgentEventSinkService } from '#/agent/eventSink'; -import { IAgentExternalHooksService } from '#/agent/externalHooks'; -import { isAbortError } from '#/agent/loop/errors'; -import { - DenyAllPermissionPolicyService, - IAgentPermissionPolicyService, -} from '#/agent/permissionPolicy'; -import { IAgentProfileService } from '#/agent/profile'; -import { ISessionMetadata } from '#/session/session-metadata'; -import { IAgentSystemReminderService } from '#/agent/systemReminder'; -import { ITelemetryService } from '#/app/telemetry'; -import { IAgentPromptService } from '#/agent/prompt'; -import { IAgentUsageService } from '#/agent/usage'; -import type { Turn } from '#/agent/turn'; - -import { DEFAULT_AGENT_SUBAGENT_PROFILES, EXPLORE_ROLE_ADDITIONAL } from './profiles'; -import { - resolveSwarmMaxConcurrency, - SubagentBatch, - type SubagentResult, - type SubagentSuspendedEvent, -} from './subagent-batch'; -import SUMMARY_CONTINUATION_PROMPT from './summary-continuation.md?raw'; -import { - type QueuedSubagentTask, - type RunSubagentOptions, - type SessionSubagentHost, - type SpawnSubagentOptions, - type SubagentHandle, -} from './subagentHost'; - -const SUBAGENT_PROMPT_ORIGIN: PromptOrigin = { kind: 'system_trigger', name: 'subagent' }; -const SUMMARY_MIN_LENGTH = 200; -const SUMMARY_CONTINUATION_ATTEMPTS = 1; -const HOOK_TEXT_PREVIEW_LENGTH = 500; -const TOOL_CALL_DISABLED_MESSAGE = - 'Tool calls are disabled for side questions. Answer with text only.'; -const SIDE_QUESTION_SYSTEM_REMINDER = ` -This is a side-channel conversation with the user. You should answer user questions directly based on what you already know. - -IMPORTANT: -- You are a separate, lightweight instance. -- The main agent continues independently; do not reference being interrupted. -- Do not call any tools. All tool calls are disabled and will be rejected. - Even though tool definitions are visible in this request, they exist only - for technical reasons (prompt cache). You must not use them. -- Respond only with text based on what you already know from the conversation - and this side-channel conversation. -- Follow-up turns may happen in this side-channel conversation. -- If you do not know the answer, say so directly. -`; - -export class DefaultSessionSubagentHost implements SessionSubagentHost { - private readonly activeChildren = new Map< - string, - { readonly controller: AbortController; runInBackground: boolean } - >(); - private readonly swarmItems = new Map(); - - constructor( - private readonly agents: IAgentLifecycleService, - private readonly ownerAgentId: string, - private readonly metadata?: ISessionMetadata, - ) {} - - getSwarmItem(agentId: string): string | undefined { - return this.swarmItems.get(agentId); - } - - async startBtw(): Promise { - const parent = await this.ensureParent(); - const child = await this.agents.create({ parentAgentId: this.ownerAgentId, type: 'sub' }); - - const parentProfile = parent.accessor.get(IAgentProfileService); - const childProfile = child.accessor.get(IAgentProfileService); - const parentData = parentProfile.data(); - // A side-question agent inherits the parent's model, thinking level, and - // system prompt so it answers from the same posture. - childProfile.update({ - modelAlias: parentData.modelAlias, - thinkingLevel: parentData.thinkingLevel, - systemPrompt: parentData.systemPrompt, - // Keep the parent's loop tools visible (prompt-cache parity) even though - // every call is denied below. - activeToolNames: parentData.activeToolNames, - }); - - // Project the parent's history into the child so it can answer from what - // the main agent already knows. - const parentMessages = parent.accessor.get(IAgentContextMemoryService)?.get(); - if (parentMessages !== undefined && parentMessages.length > 0) { - child.accessor.get(IAgentContextMemoryService)?.splice(0, 0, parentMessages); - } - - child.accessor - .get(IAgentSystemReminderService) - ?.appendSystemReminder(SIDE_QUESTION_SYSTEM_REMINDER.trim(), { - kind: 'system_trigger', - name: 'btw', - }); - - // Disable every tool call: side questions are answered with text only. - child.accessor - .get(IAgentPermissionPolicyService) - ?.registerPolicy(new DenyAllPermissionPolicyService(TOOL_CALL_DISABLED_MESSAGE)); - - return child.id; - } - - async spawn(options: SpawnSubagentOptions): Promise { - options.signal.throwIfAborted(); - const parent = await this.ensureParent(); - const child = await this.agents.create({ - parentAgentId: this.ownerAgentId, - cwd: parent.accessor.get(IAgentProfileService).data().cwd, - type: 'sub', - swarmItem: options.swarmItem, - }); - if (options.swarmItem !== undefined) this.swarmItems.set(child.id, options.swarmItem); - this.configureChild(parent, child, options.profileName); - this.emitSpawned(parent, child.id, options.profileName, options); - const completion = this.runWithActiveChild( - child, - options, - parent, - options.profileName, - (turnRef, controller) => this.runPromptTurn(child, parent, options, options.profileName, turnRef, controller), - ); - return { agentId: child.id, profileName: options.profileName, resumed: false, completion }; - } - - async resume(agentId: string, options: RunSubagentOptions): Promise { - options.signal.throwIfAborted(); - const parent = await this.ensureParent(); - const child = await this.requireChild(agentId); - const profileName = child.accessor.get(IAgentProfileService).data().profileName ?? 'subagent'; - this.emitSpawned(parent, child.id, profileName, options); - const completion = this.runWithActiveChild( - child, - options, - parent, - profileName, - (turnRef, controller) => this.runPromptTurn(child, parent, options, profileName, turnRef, controller), - ); - return { agentId, profileName, resumed: true, completion }; - } - - async retry(agentId: string, options: RunSubagentOptions): Promise { - options.signal.throwIfAborted(); - const parent = await this.ensureParent(); - const child = await this.requireChild(agentId); - const profileName = child.accessor.get(IAgentProfileService).data().profileName ?? 'subagent'; - this.emitSpawned(parent, child.id, profileName, options); - const completion = this.runWithActiveChild( - child, - options, - parent, - profileName, - (turnRef, controller) => this.runRetryTurn(child, parent, options, profileName, turnRef, controller), - ); - return { agentId, profileName, resumed: true, completion }; - } - - async getProfileName(agentId: string): Promise { - if (this.metadata !== undefined) { - const meta = (await this.metadata.read()).agents?.[agentId]; - if (meta?.type !== 'sub' || meta.parentAgentId !== this.ownerAgentId) return undefined; - } - const child = this.agents.getHandle(agentId); - if (child === undefined) return undefined; - return child.accessor.get(IAgentProfileService).data().profileName; - } - - markActiveChildDetached(agentId: string): void { - const child = this.activeChildren.get(agentId); - if (child !== undefined) child.runInBackground = true; - } - - async runQueued(tasks: readonly QueuedSubagentTask[]): Promise>> { - const maxConcurrency = resolveSwarmMaxConcurrency(); - return new SubagentBatch(this, tasks, { maxConcurrency }).run(); - } - - cancelAll(reason: unknown = userCancellationReason()): void { - // v2 tracks every subagent (including descendants spawned by subagents) in - // the single session-scoped host, so aborting the foreground children here - // cancels the whole tree — there is no per-agent host to recurse into. - for (const [, child] of this.activeChildren) { - if (child.runInBackground) continue; - child.controller.abort(reason); - } - } - - suspended(event: SubagentSuspendedEvent): void { - const parent = this.agents.getHandle(this.ownerAgentId); - parent?.accessor.get(IAgentEventSinkService)?.emit({ - type: 'subagent.suspended', - subagentId: event.agentId, - reason: event.reason, - }); - } - - private async ensureParent(): Promise { - const existing = this.agents.getHandle(this.ownerAgentId); - if (existing !== undefined) return existing; - if (this.ownerAgentId === 'main') return this.agents.createMain(); - throw new Error(`Parent agent "${this.ownerAgentId}" does not exist`); - } - - private async requireChild(agentId: string): Promise { - if (this.metadata !== undefined) { - const meta = (await this.metadata.read()).agents?.[agentId]; - if (meta === undefined) throw new Error(`Agent instance "${agentId}" does not exist`); - if (meta.type !== 'sub') throw new Error(`Agent instance "${agentId}" is not a subagent`); - if (meta.parentAgentId !== this.ownerAgentId) { - throw new Error(`Agent instance "${agentId}" does not belong to this parent agent`); - } - } - const child = this.agents.getHandle(agentId); - if (child === undefined) throw new Error(`Agent instance "${agentId}" does not exist`); - if (this.activeChildren.has(agentId)) { - throw new Error(`Agent instance "${agentId}" is already running`); - } - return child; - } - - private configureChild(parent: IScopeHandle, child: IScopeHandle, profileName: string): void { - const parentProfile = parent.accessor.get(IAgentProfileService); - const childProfile = child.accessor.get(IAgentProfileService); - const parentData = parentProfile.data(); - const profile = DEFAULT_AGENT_SUBAGENT_PROFILES[profileName]; - const activeToolNames = - profileName === 'coder' - ? (parentData.activeToolNames ?? profile?.tools) - : profile?.tools; - childProfile.update({ - cwd: parentData.cwd, - modelAlias: parentData.modelAlias, - thinkingLevel: parentData.thinkingLevel, - profileName, - // `explore` extends the parent (agent) prompt with its read-only - // exploration role, mirroring v1's `explore.yaml` (`extends: agent` + - // `roleAdditional`). Full profile resolution via `applyProfile` (AGENTS.md - // context assembly) and v1's `inheritUserTools` are not yet wired in v2, - // so the standard explore tool set is used directly. - systemPrompt: - profileName === 'explore' - ? `${parentData.systemPrompt}\n\n${EXPLORE_ROLE_ADDITIONAL}` - : parentData.systemPrompt, - activeToolNames, - }); - } - - private emitSpawned( - parent: IScopeHandle, - subagentId: string, - profileName: string, - options: RunSubagentOptions, - ): void { - parent.accessor.get(IAgentEventSinkService)?.emit({ - type: 'subagent.spawned', - subagentId, - subagentName: profileName, - parentToolCallId: options.parentToolCallId, - parentToolCallUuid: options.parentToolCallUuid, - parentAgentId: this.ownerAgentId, - description: options.description, - swarmIndex: options.swarmIndex, - runInBackground: options.runInBackground, - }); - parent.accessor.get(ITelemetryService)?.track('subagent_created', { - subagent_name: profileName, - run_in_background: options.runInBackground, - }); - } - - private emitStarted(parent: IScopeHandle, subagentId: string): void { - parent.accessor.get(IAgentEventSinkService)?.emit({ type: 'subagent.started', subagentId }); - } - - private emitCompleted( - parent: IScopeHandle, - subagentId: string, - resultSummary: string, - usage?: TokenUsage, - ): void { - parent.accessor.get(IAgentEventSinkService)?.emit({ - type: 'subagent.completed', - subagentId, - resultSummary, - usage, - }); - } - - private emitFailed( - parent: IScopeHandle, - subagentId: string, - error: unknown, - options: RunSubagentOptions, - ): void { - if (isAbortError(error)) return; - if (shouldSuppressQueuedAttemptFailureEvent(options, error)) return; - parent.accessor.get(IAgentEventSinkService)?.emit({ - type: 'subagent.failed', - subagentId, - error: errorMessage(error), - }); - } - - private async triggerSubagentStart( - parent: IScopeHandle, - profileName: string, - prompt: string, - signal: AbortSignal, - ): Promise { - await parent.accessor.get(IAgentExternalHooksService)?.triggerSubagentStart( - { - agentName: profileName, - prompt: prompt.slice(0, HOOK_TEXT_PREVIEW_LENGTH), - }, - signal, - ); - } - - private triggerSubagentStop(parent: IScopeHandle, profileName: string, result: string): void { - parent.accessor.get(IAgentExternalHooksService)?.triggerSubagentStop({ - agentName: profileName, - response: result.slice(0, HOOK_TEXT_PREVIEW_LENGTH), - }); - } - - private observeFirstRequest(turn: Turn, options: RunSubagentOptions): void { - if (options.onReady === undefined) return; - void turn.ready.then(() => options.onReady?.()).catch(() => {}); - } - - private async runWithActiveChild( - child: IScopeHandle, - options: RunSubagentOptions, - parent: IScopeHandle, - profileName: string, - run: ( - turn: { current?: Turn }, - controller: AbortController, - ) => Promise<{ result: string; usage?: TokenUsage }>, - ): Promise<{ result: string; usage?: TokenUsage }> { - const controller = new AbortController(); - this.activeChildren.set(child.id, { controller, runInBackground: options.runInBackground }); - const unlink = linkAbortSignal(options.signal, controller); - const turnRef: { current?: Turn } = {}; - this.emitStarted(parent, child.id); - try { - const result = await run(turnRef, controller); - this.emitCompleted(parent, child.id, result.result, result.usage); - this.triggerSubagentStop(parent, profileName, result.result); - return result; - } catch (error) { - this.emitFailed(parent, child.id, error, options); - throw error; - } finally { - unlink(); - if (controller.signal.aborted) { - turnRef.current?.abortController.abort(controller.signal.reason); - } - this.activeChildren.delete(child.id); - } - } - - private async runPromptTurn( - child: IScopeHandle, - parent: IScopeHandle, - options: RunSubagentOptions, - profileName: string, - turnRef: { current?: Turn }, - controller: AbortController, - ): Promise<{ result: string; usage?: TokenUsage }> { - options.signal.throwIfAborted(); - await this.triggerSubagentStart(parent, profileName, options.prompt, options.signal); - options.signal.throwIfAborted(); - - const turn = child.accessor.get(IAgentPromptService).prompt({ - role: 'user', - content: [{ type: 'text', text: options.prompt }], - toolCalls: [], - origin: SUBAGENT_PROMPT_ORIGIN, - }); - if (turn === undefined) { - throw new Error('Subagent turn could not be started'); - } - turnRef.current = turn; - this.observeFirstRequest(turn, options); - const result = await this.awaitTurn(turn, controller); - classifyTurnResult(result); - const summary = await this.completeSummary(child, controller, turnRef); - const usage = child.accessor.get(IAgentUsageService)?.status().total; - return { result: summary, usage }; - } - - private async runRetryTurn( - child: IScopeHandle, - parent: IScopeHandle, - options: RunSubagentOptions, - profileName: string, - turnRef: { current?: Turn }, - controller: AbortController, - ): Promise<{ result: string; usage?: TokenUsage }> { - options.signal.throwIfAborted(); - await this.triggerSubagentStart(parent, profileName, options.prompt, options.signal); - options.signal.throwIfAborted(); - - // Retry the existing turn in place (no new user message appended), mirroring - // v1's `child.turn.retry('agent-host')`. - const turn = child.accessor.get(IAgentPromptService).retry('agent-host'); - if (turn === undefined) { - throw new Error(`Agent instance "${child.id}" could not start a retry turn`); - } - turnRef.current = turn; - this.observeFirstRequest(turn, options); - const result = await this.awaitTurn(turn, controller); - classifyTurnResult(result); - const summary = await this.completeSummary(child, controller, turnRef); - const usage = child.accessor.get(IAgentUsageService)?.status().total; - return { result: summary, usage }; - } - - private async awaitTurn( - turn: Turn, - controller: AbortController, - ): Promise<{ reason: string; error?: unknown }> { - const onAbort = (): void => { - turn.abortController.abort(controller.signal.reason); - }; - controller.signal.addEventListener('abort', onAbort, { once: true }); - try { - return await Promise.race([turn.result, abortPromise(controller.signal)]); - } finally { - controller.signal.removeEventListener('abort', onAbort); - } - } - - private async completeSummary( - child: IScopeHandle, - controller: AbortController, - turnRef: { current?: Turn }, - ): Promise { - let summary = latestAssistantText(child.accessor.get(IAgentContextMemoryService).get()); - if (summary.trim().length >= SUMMARY_MIN_LENGTH) return summary; - - for (let attempt = 0; attempt < SUMMARY_CONTINUATION_ATTEMPTS; attempt++) { - const turn = child.accessor.get(IAgentPromptService).prompt({ - role: 'user', - content: [{ type: 'text', text: SUMMARY_CONTINUATION_PROMPT }], - toolCalls: [], - origin: SUBAGENT_PROMPT_ORIGIN, - }); - if (turn === undefined) break; - turnRef.current = turn; - const result = await this.awaitTurn(turn, controller); - if (result.reason !== 'completed') break; - const continued = latestAssistantText(child.accessor.get(IAgentContextMemoryService).get()); - if (continued.trim().length > 0) summary = continued; - if (summary.trim().length >= SUMMARY_MIN_LENGTH) break; - } - return summary; - } -} - -/** - * Map a finished subagent turn to the v1 error taxonomy: - * - `filtered` → provider safety policy block - * - provider rate limit → `APIProviderRateLimitError` (so the swarm batch - * requeues the attempt via `isProviderRateLimitError`) - * - `cancelled` → user cancellation reason - * - * `max_tokens` is intentionally not classified here: v2's turn result collapses - * every non-aborted/non-filtered stop into `completed`, so the subagent host - * cannot observe a max_tokens stop. See the migration notes for this deliberate - * drop. - */ -function classifyTurnResult(result: { reason: string; error?: unknown }): void { - if (result.reason === 'filtered') { - throw new Error('Subagent turn blocked by provider safety policy'); - } - if (result.reason === 'failed') { - const error = result.error; - if (isProviderRateLimitError(error)) throw error; - const payload = toKimiErrorPayload(error); - if (payload.code === ErrorCodes.PROVIDER_RATE_LIMIT) { - throw providerRateLimitErrorFromPayload(payload); - } - throw error instanceof Error ? error : new Error(String(error ?? 'Subagent turn failed')); - } - if (result.reason === 'cancelled') { - throw userCancellationReason(); - } -} - -function shouldSuppressQueuedAttemptFailureEvent( - options: RunSubagentOptions, - error: unknown, -): boolean { - if (options.suppressRateLimitFailureEvent !== true) return false; - if (isProviderRateLimitError(error)) return true; - return isAbortError(error) || options.signal.aborted; -} - -function providerRateLimitErrorFromPayload(error: KimiErrorPayload): APIProviderRateLimitError { - const requestId = - typeof error.details?.['requestId'] === 'string' ? error.details['requestId'] : null; - return new APIProviderRateLimitError(error.message, requestId); -} - -function abortPromise(signal: AbortSignal): Promise { - if (signal.aborted) { - return Promise.reject(signal.reason ?? userCancellationReason()); - } - return new Promise((_resolve, reject) => { - signal.addEventListener('abort', () => reject(signal.reason ?? userCancellationReason()), { - once: true, - }); - }); -} - -function errorMessage(error: unknown): string { - return error instanceof Error ? error.message : String(error); -} - -function latestAssistantText(messages: readonly ContextMessage[]): string { - for (let i = messages.length - 1; i >= 0; i--) { - const message = messages[i]!; - if (message.role !== 'assistant') continue; - return contentText(message.content); - } - return ''; -} - -function contentText(content: ContextMessage['content']): string { - if (typeof content === 'string') return content; - return content - .filter((part): part is Extract<(typeof content)[number], { type: 'text' }> => part.type === 'text') - .map((part) => part.text) - .join(''); -} diff --git a/packages/agent-core-v2/src/session/subagentHost/index.ts b/packages/agent-core-v2/src/session/subagentHost/index.ts index 2fa673ea0..ca9a738f7 100644 --- a/packages/agent-core-v2/src/session/subagentHost/index.ts +++ b/packages/agent-core-v2/src/session/subagentHost/index.ts @@ -4,7 +4,7 @@ export * from './subagentHost'; export * from './subagentHostService'; -export * from './defaultSessionSubagentHost'; export * from './profiles'; export { AgentTool, AgentToolInputSchema, AgentToolOutputSchema } from './agentTool'; export type { AgentToolInput, AgentToolOutput } from './agentTool'; +export type { QueuedSubagentTask } from './subagent-batch'; diff --git a/packages/agent-core-v2/src/session/subagentHost/subagentHost.ts b/packages/agent-core-v2/src/session/subagentHost/subagentHost.ts index d8a01667f..f77d4dc95 100644 --- a/packages/agent-core-v2/src/session/subagentHost/subagentHost.ts +++ b/packages/agent-core-v2/src/session/subagentHost/subagentHost.ts @@ -36,9 +36,11 @@ export type SubagentHandle = { }>; }; -export interface SessionSubagentHost { +export interface ISessionSubagentHost { + readonly _serviceBrand: undefined; getSwarmItem(agentId: string): string | undefined; startBtw(): Promise; + generateAgentsMd(): Promise; spawn(options: SpawnSubagentOptions): Promise; resume(agentId: string, options: RunSubagentOptions): Promise; retry(agentId: string, options: RunSubagentOptions): Promise; @@ -52,21 +54,5 @@ export interface SessionSubagentHost { } export type QueuedSubagentRunResult = SubagentResult; -export type { QueuedSubagentTask }; - -export interface ISessionSubagentHost { - readonly _serviceBrand: undefined; - getSwarmItem(agentId: string): string | undefined; - startBtw(): Promise; - generateAgentsMd(): Promise; - spawn(options: SpawnSubagentOptions): Promise; - resume(agentId: string, options: RunSubagentOptions): Promise; - getProfileName(agentId: string): Promise; - markActiveChildDetached(agentId: string): void; - runQueued(tasks: readonly QueuedSubagentTask[]): Promise>>; - cancelAll(reason?: unknown): void; - suspended(event: SubagentSuspendedEvent): void; -} - export const ISessionSubagentHost = createDecorator('sessionSubagentHost'); diff --git a/packages/agent-core-v2/src/session/subagentHost/subagentHostService.ts b/packages/agent-core-v2/src/session/subagentHost/subagentHostService.ts index 8d4e185b5..27e598a72 100644 --- a/packages/agent-core-v2/src/session/subagentHost/subagentHostService.ts +++ b/packages/agent-core-v2/src/session/subagentHost/subagentHostService.ts @@ -1,39 +1,94 @@ -import type { - QueuedSubagentRunResult, - QueuedSubagentTask, - SessionSubagentHost, - SpawnSubagentOptions, - RunSubagentOptions, - SubagentHandle, -} from './subagentHost'; -import type { SubagentSuspendedEvent } from './subagent-batch'; import { - ISessionSubagentHost, -} from './subagentHost'; + APIProviderRateLimitError, + isProviderRateLimitError, + type TokenUsage, +} from '@moonshot-ai/kosong'; + import { Disposable } from '#/_base/di'; import { InstantiationType } from '#/_base/di/extensions'; -import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; -import { IAgentLifecycleService } from '#/session/agent-lifecycle'; -import { IAgentBackgroundService } from '#/agent/background'; -import { ILogService } from '#/app/log'; -import { IAgentProfileService } from '#/agent/profile'; +import { + LifecycleScope, + registerScopedService, + type IScopeHandle, +} from '#/_base/di/scope'; +import { linkAbortSignal, userCancellationReason } from '#/_base/utils/abort'; import { IKaos } from '#/app/kaos'; +import { ILogService } from '#/app/log'; +import { ITelemetryService } from '#/app/telemetry'; +import { IAgentBackgroundService } from '#/agent/background'; +import { + IAgentContextMemoryService, + type ContextMessage, + type PromptOrigin, +} from '#/agent/contextMemory'; +import { IAgentEventSinkService } from '#/agent/eventSink'; +import { IAgentExternalHooksService } from '#/agent/externalHooks'; +import { isAbortError } from '#/agent/loop/errors'; +import { + DenyAllPermissionPolicyService, + IAgentPermissionPolicyService, +} from '#/agent/permissionPolicy'; +import { IAgentProfileService } from '#/agent/profile'; +import { IAgentPromptService } from '#/agent/prompt'; +import { IAgentSystemReminderService } from '#/agent/systemReminder'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry'; +import type { Turn } from '#/agent/turn'; +import { IAgentUsageService } from '#/agent/usage'; +import { ErrorCodes, toKimiErrorPayload, type KimiErrorPayload } from '#/errors'; +import { IAgentLifecycleService } from '#/session/agent-lifecycle'; import { ISessionProcessRunner } from '#/session/process'; import { ISessionMetadata } from '#/session/session-metadata'; -import { IAgentToolRegistryService } from '#/agent/toolRegistry'; import { AgentTool } from './agentTool'; -import { DefaultSessionSubagentHost } from './defaultSessionSubagentHost'; -import { DEFAULT_AGENT_SUBAGENT_PROFILES } from './profiles'; +import { DEFAULT_AGENT_SUBAGENT_PROFILES, EXPLORE_ROLE_ADDITIONAL } from './profiles'; +import { + resolveSwarmMaxConcurrency, + SubagentBatch, + type QueuedSubagentTask, + type SubagentResult, + type SubagentSuspendedEvent, +} from './subagent-batch'; +import SUMMARY_CONTINUATION_PROMPT from './summary-continuation.md?raw'; +import { + ISessionSubagentHost, + type RunSubagentOptions, + type SpawnSubagentOptions, + type SubagentHandle, +} from './subagentHost'; + +const SUBAGENT_PROMPT_ORIGIN: PromptOrigin = { kind: 'system_trigger', name: 'subagent' }; +const SUMMARY_MIN_LENGTH = 200; +const SUMMARY_CONTINUATION_ATTEMPTS = 1; +const HOOK_TEXT_PREVIEW_LENGTH = 500; +const TOOL_CALL_DISABLED_MESSAGE = + 'Tool calls are disabled for side questions. Answer with text only.'; +const SIDE_QUESTION_SYSTEM_REMINDER = ` +This is a side-channel conversation with the user. You should answer user questions directly based on what you already know. + +IMPORTANT: +- You are a separate, lightweight instance. +- The main agent continues independently; do not reference being interrupted. +- Do not call any tools. All tool calls are disabled and will be rejected. + Even though tool definitions are visible in this request, they exist only + for technical reasons (prompt cache). You must not use them. +- Respond only with text based on what you already know from the conversation + and this side-channel conversation. +- Follow-up turns may happen in this side-channel conversation. +- If you do not know the answer, say so directly. +`; export class SessionSubagentHostService extends Disposable implements ISessionSubagentHost { declare readonly _serviceBrand: undefined; - private readonly host: SessionSubagentHost; + private readonly activeChildren = new Map< + string, + { readonly controller: AbortController; runInBackground: boolean } + >(); + private readonly swarmItems = new Map(); + private readonly ownerAgentId = 'main'; constructor( - subagentHost: SessionSubagentHost | undefined, - @IAgentLifecycleService agents: IAgentLifecycleService, - @ISessionMetadata metadata: ISessionMetadata, + @IAgentLifecycleService private readonly agents: IAgentLifecycleService, + @ISessionMetadata private readonly metadata: ISessionMetadata | undefined, @IAgentToolRegistryService toolRegistry: IAgentToolRegistryService, @IAgentBackgroundService background: IAgentBackgroundService, @IAgentProfileService profile: IAgentProfileService, @@ -42,8 +97,6 @@ export class SessionSubagentHostService extends Disposable implements ISessionSu @ILogService log?: ILogService, ) { super(); - this.host = subagentHost ?? new DefaultSessionSubagentHost(agents, 'main', metadata); - this._register( toolRegistry.register( new AgentTool(this, background, DEFAULT_AGENT_SUBAGENT_PROFILES, { @@ -60,15 +113,51 @@ export class SessionSubagentHostService extends Disposable implements ISessionSu } getSwarmItem(agentId: string): string | undefined { - return this.host?.getSwarmItem(agentId); + return this.swarmItems.get(agentId); } - startBtw(): Promise { - return this.host.startBtw(); + async startBtw(): Promise { + const parent = await this.ensureParent(); + const child = await this.agents.create({ parentAgentId: this.ownerAgentId, type: 'sub' }); + + const parentProfile = parent.accessor.get(IAgentProfileService); + const childProfile = child.accessor.get(IAgentProfileService); + const parentData = parentProfile.data(); + // A side-question agent inherits the parent's model, thinking level, and + // system prompt so it answers from the same posture. + childProfile.update({ + modelAlias: parentData.modelAlias, + thinkingLevel: parentData.thinkingLevel, + systemPrompt: parentData.systemPrompt, + // Keep the parent's loop tools visible (prompt-cache parity) even though + // every call is denied below. + activeToolNames: parentData.activeToolNames, + }); + + // Project the parent's history into the child so it can answer from what + // the main agent already knows. + const parentMessages = parent.accessor.get(IAgentContextMemoryService)?.get(); + if (parentMessages !== undefined && parentMessages.length > 0) { + child.accessor.get(IAgentContextMemoryService)?.splice(0, 0, parentMessages); + } + + child.accessor + .get(IAgentSystemReminderService) + ?.appendSystemReminder(SIDE_QUESTION_SYSTEM_REMINDER.trim(), { + kind: 'system_trigger', + name: 'btw', + }); + + // Disable every tool call: side questions are answered with text only. + child.accessor + .get(IAgentPermissionPolicyService) + ?.registerPolicy(new DenyAllPermissionPolicyService(TOOL_CALL_DISABLED_MESSAGE)); + + return child.id; } async generateAgentsMd(): Promise { - const handle = await this.host.spawn({ + const handle = await this.spawn({ profileName: 'coder', parentToolCallId: 'generate-agents-md', prompt: 'Initialize AGENTS.md for this workspace.', @@ -79,39 +168,439 @@ export class SessionSubagentHostService extends Disposable implements ISessionSu await handle.completion; } - spawn(options: SpawnSubagentOptions): Promise { - return this.host.spawn(options); + async spawn(options: SpawnSubagentOptions): Promise { + options.signal.throwIfAborted(); + const parent = await this.ensureParent(); + const child = await this.agents.create({ + parentAgentId: this.ownerAgentId, + cwd: parent.accessor.get(IAgentProfileService).data().cwd, + type: 'sub', + swarmItem: options.swarmItem, + }); + if (options.swarmItem !== undefined) this.swarmItems.set(child.id, options.swarmItem); + this.configureChild(parent, child, options.profileName); + this.emitSpawned(parent, child.id, options.profileName, options); + const completion = this.runWithActiveChild( + child, + options, + parent, + options.profileName, + (turnRef, controller) => this.runPromptTurn(child, parent, options, options.profileName, turnRef, controller), + ); + return { agentId: child.id, profileName: options.profileName, resumed: false, completion }; } - resume(agentId: string, options: RunSubagentOptions): Promise { - return this.host.resume(agentId, options); + async resume(agentId: string, options: RunSubagentOptions): Promise { + options.signal.throwIfAborted(); + const parent = await this.ensureParent(); + const child = await this.requireChild(agentId); + const profileName = child.accessor.get(IAgentProfileService).data().profileName ?? 'subagent'; + this.emitSpawned(parent, child.id, profileName, options); + const completion = this.runWithActiveChild( + child, + options, + parent, + profileName, + (turnRef, controller) => this.runPromptTurn(child, parent, options, profileName, turnRef, controller), + ); + return { agentId, profileName, resumed: true, completion }; } - getProfileName(agentId: string): Promise { - return this.host.getProfileName(agentId); + async retry(agentId: string, options: RunSubagentOptions): Promise { + options.signal.throwIfAborted(); + const parent = await this.ensureParent(); + const child = await this.requireChild(agentId); + const profileName = child.accessor.get(IAgentProfileService).data().profileName ?? 'subagent'; + this.emitSpawned(parent, child.id, profileName, options); + const completion = this.runWithActiveChild( + child, + options, + parent, + profileName, + (turnRef, controller) => this.runRetryTurn(child, parent, options, profileName, turnRef, controller), + ); + return { agentId, profileName, resumed: true, completion }; + } + + async getProfileName(agentId: string): Promise { + if (this.metadata !== undefined) { + const meta = (await this.metadata.read()).agents?.[agentId]; + if (meta?.type !== 'sub' || meta.parentAgentId !== this.ownerAgentId) return undefined; + } + const child = this.agents.getHandle(agentId); + if (child === undefined) return undefined; + return child.accessor.get(IAgentProfileService).data().profileName; } markActiveChildDetached(agentId: string): void { - this.host.markActiveChildDetached(agentId); + const child = this.activeChildren.get(agentId); + if (child !== undefined) child.runInBackground = true; } - cancelAll(reason?: unknown): void { - this.host.cancelAll(reason); + async runQueued(tasks: readonly QueuedSubagentTask[]): Promise>> { + const maxConcurrency = resolveSwarmMaxConcurrency(); + return new SubagentBatch(this, tasks, { maxConcurrency }).run(); + } + + cancelAll(reason: unknown = userCancellationReason()): void { + // v2 tracks every subagent (including descendants spawned by subagents) in + // the single session-scoped host, so aborting the foreground children here + // cancels the whole tree — there is no per-agent host to recurse into. + for (const [, child] of this.activeChildren) { + if (child.runInBackground) continue; + child.controller.abort(reason); + } } suspended(event: SubagentSuspendedEvent): void { - this.host.suspended(event); + const parent = this.agents.getHandle(this.ownerAgentId); + parent?.accessor.get(IAgentEventSinkService)?.emit({ + type: 'subagent.suspended', + subagentId: event.agentId, + reason: event.reason, + }); } - runQueued( - tasks: readonly QueuedSubagentTask[], - ): Promise>> { - const subagentHost = this.host; - if (subagentHost === undefined) { - throw new Error('Subagent host is not configured.'); - } - return subagentHost.runQueued(tasks); + private async ensureParent(): Promise { + const existing = this.agents.getHandle(this.ownerAgentId); + if (existing !== undefined) return existing; + if (this.ownerAgentId === 'main') return this.agents.createMain(); + throw new Error(`Parent agent "${this.ownerAgentId}" does not exist`); } + + private async requireChild(agentId: string): Promise { + if (this.metadata !== undefined) { + const meta = (await this.metadata.read()).agents?.[agentId]; + if (meta === undefined) throw new Error(`Agent instance "${agentId}" does not exist`); + if (meta.type !== 'sub') throw new Error(`Agent instance "${agentId}" is not a subagent`); + if (meta.parentAgentId !== this.ownerAgentId) { + throw new Error(`Agent instance "${agentId}" does not belong to this parent agent`); + } + } + const child = this.agents.getHandle(agentId); + if (child === undefined) throw new Error(`Agent instance "${agentId}" does not exist`); + if (this.activeChildren.has(agentId)) { + throw new Error(`Agent instance "${agentId}" is already running`); + } + return child; + } + + private configureChild(parent: IScopeHandle, child: IScopeHandle, profileName: string): void { + const parentProfile = parent.accessor.get(IAgentProfileService); + const childProfile = child.accessor.get(IAgentProfileService); + const parentData = parentProfile.data(); + const profile = DEFAULT_AGENT_SUBAGENT_PROFILES[profileName]; + const activeToolNames = + profileName === 'coder' + ? (parentData.activeToolNames ?? profile?.tools) + : profile?.tools; + childProfile.update({ + cwd: parentData.cwd, + modelAlias: parentData.modelAlias, + thinkingLevel: parentData.thinkingLevel, + profileName, + // `explore` extends the parent (agent) prompt with its read-only + // exploration role, mirroring v1's `explore.yaml` (`extends: agent` + + // `roleAdditional`). Full profile resolution via `applyProfile` (AGENTS.md + // context assembly) and v1's `inheritUserTools` are not yet wired in v2, + // so the standard explore tool set is used directly. + systemPrompt: + profileName === 'explore' + ? `${parentData.systemPrompt}\n\n${EXPLORE_ROLE_ADDITIONAL}` + : parentData.systemPrompt, + activeToolNames, + }); + } + + private emitSpawned( + parent: IScopeHandle, + subagentId: string, + profileName: string, + options: RunSubagentOptions, + ): void { + parent.accessor.get(IAgentEventSinkService)?.emit({ + type: 'subagent.spawned', + subagentId, + subagentName: profileName, + parentToolCallId: options.parentToolCallId, + parentToolCallUuid: options.parentToolCallUuid, + parentAgentId: this.ownerAgentId, + description: options.description, + swarmIndex: options.swarmIndex, + runInBackground: options.runInBackground, + }); + parent.accessor.get(ITelemetryService)?.track('subagent_created', { + subagent_name: profileName, + run_in_background: options.runInBackground, + }); + } + + private emitStarted(parent: IScopeHandle, subagentId: string): void { + parent.accessor.get(IAgentEventSinkService)?.emit({ type: 'subagent.started', subagentId }); + } + + private emitCompleted( + parent: IScopeHandle, + subagentId: string, + resultSummary: string, + usage?: TokenUsage, + ): void { + parent.accessor.get(IAgentEventSinkService)?.emit({ + type: 'subagent.completed', + subagentId, + resultSummary, + usage, + }); + } + + private emitFailed( + parent: IScopeHandle, + subagentId: string, + error: unknown, + options: RunSubagentOptions, + ): void { + if (isAbortError(error)) return; + if (shouldSuppressQueuedAttemptFailureEvent(options, error)) return; + parent.accessor.get(IAgentEventSinkService)?.emit({ + type: 'subagent.failed', + subagentId, + error: errorMessage(error), + }); + } + + private async triggerSubagentStart( + parent: IScopeHandle, + profileName: string, + prompt: string, + signal: AbortSignal, + ): Promise { + await parent.accessor.get(IAgentExternalHooksService)?.triggerSubagentStart( + { + agentName: profileName, + prompt: prompt.slice(0, HOOK_TEXT_PREVIEW_LENGTH), + }, + signal, + ); + } + + private triggerSubagentStop(parent: IScopeHandle, profileName: string, result: string): void { + parent.accessor.get(IAgentExternalHooksService)?.triggerSubagentStop({ + agentName: profileName, + response: result.slice(0, HOOK_TEXT_PREVIEW_LENGTH), + }); + } + + private observeFirstRequest(turn: Turn, options: RunSubagentOptions): void { + if (options.onReady === undefined) return; + void turn.ready.then(() => options.onReady?.()).catch(() => {}); + } + + private async runWithActiveChild( + child: IScopeHandle, + options: RunSubagentOptions, + parent: IScopeHandle, + profileName: string, + run: ( + turn: { current?: Turn }, + controller: AbortController, + ) => Promise<{ result: string; usage?: TokenUsage }>, + ): Promise<{ result: string; usage?: TokenUsage }> { + const controller = new AbortController(); + this.activeChildren.set(child.id, { controller, runInBackground: options.runInBackground }); + const unlink = linkAbortSignal(options.signal, controller); + const turnRef: { current?: Turn } = {}; + this.emitStarted(parent, child.id); + try { + const result = await run(turnRef, controller); + this.emitCompleted(parent, child.id, result.result, result.usage); + this.triggerSubagentStop(parent, profileName, result.result); + return result; + } catch (error) { + this.emitFailed(parent, child.id, error, options); + throw error; + } finally { + unlink(); + if (controller.signal.aborted) { + turnRef.current?.abortController.abort(controller.signal.reason); + } + this.activeChildren.delete(child.id); + } + } + + private async runPromptTurn( + child: IScopeHandle, + parent: IScopeHandle, + options: RunSubagentOptions, + profileName: string, + turnRef: { current?: Turn }, + controller: AbortController, + ): Promise<{ result: string; usage?: TokenUsage }> { + options.signal.throwIfAborted(); + await this.triggerSubagentStart(parent, profileName, options.prompt, options.signal); + options.signal.throwIfAborted(); + + const turn = child.accessor.get(IAgentPromptService).prompt({ + role: 'user', + content: [{ type: 'text', text: options.prompt }], + toolCalls: [], + origin: SUBAGENT_PROMPT_ORIGIN, + }); + if (turn === undefined) { + throw new Error('Subagent turn could not be started'); + } + turnRef.current = turn; + this.observeFirstRequest(turn, options); + const result = await this.awaitTurn(turn, controller); + classifyTurnResult(result); + const summary = await this.completeSummary(child, controller, turnRef); + const usage = child.accessor.get(IAgentUsageService)?.status().total; + return { result: summary, usage }; + } + + private async runRetryTurn( + child: IScopeHandle, + parent: IScopeHandle, + options: RunSubagentOptions, + profileName: string, + turnRef: { current?: Turn }, + controller: AbortController, + ): Promise<{ result: string; usage?: TokenUsage }> { + options.signal.throwIfAborted(); + await this.triggerSubagentStart(parent, profileName, options.prompt, options.signal); + options.signal.throwIfAborted(); + + // Retry the existing turn in place (no new user message appended), mirroring + // v1's `child.turn.retry('agent-host')`. + const turn = child.accessor.get(IAgentPromptService).retry('agent-host'); + if (turn === undefined) { + throw new Error(`Agent instance "${child.id}" could not start a retry turn`); + } + turnRef.current = turn; + this.observeFirstRequest(turn, options); + const result = await this.awaitTurn(turn, controller); + classifyTurnResult(result); + const summary = await this.completeSummary(child, controller, turnRef); + const usage = child.accessor.get(IAgentUsageService)?.status().total; + return { result: summary, usage }; + } + + private async awaitTurn( + turn: Turn, + controller: AbortController, + ): Promise<{ reason: string; error?: unknown }> { + const onAbort = (): void => { + turn.abortController.abort(controller.signal.reason); + }; + controller.signal.addEventListener('abort', onAbort, { once: true }); + try { + return await Promise.race([turn.result, abortPromise(controller.signal)]); + } finally { + controller.signal.removeEventListener('abort', onAbort); + } + } + + private async completeSummary( + child: IScopeHandle, + controller: AbortController, + turnRef: { current?: Turn }, + ): Promise { + let summary = latestAssistantText(child.accessor.get(IAgentContextMemoryService).get()); + if (summary.trim().length >= SUMMARY_MIN_LENGTH) return summary; + + for (let attempt = 0; attempt < SUMMARY_CONTINUATION_ATTEMPTS; attempt++) { + const turn = child.accessor.get(IAgentPromptService).prompt({ + role: 'user', + content: [{ type: 'text', text: SUMMARY_CONTINUATION_PROMPT }], + toolCalls: [], + origin: SUBAGENT_PROMPT_ORIGIN, + }); + if (turn === undefined) break; + turnRef.current = turn; + const result = await this.awaitTurn(turn, controller); + if (result.reason !== 'completed') break; + const continued = latestAssistantText(child.accessor.get(IAgentContextMemoryService).get()); + if (continued.trim().length > 0) summary = continued; + if (summary.trim().length >= SUMMARY_MIN_LENGTH) break; + } + return summary; + } +} + +/** + * Map a finished subagent turn to the v1 error taxonomy: + * - `filtered` → provider safety policy block + * - provider rate limit → `APIProviderRateLimitError` (so the swarm batch + * requeues the attempt via `isProviderRateLimitError`) + * - `cancelled` → user cancellation reason + * + * `max_tokens` is intentionally not classified here: v2's turn result collapses + * every non-aborted/non-filtered stop into `completed`, so the subagent host + * cannot observe a max_tokens stop. See the migration notes for this deliberate + * drop. + */ +function classifyTurnResult(result: { reason: string; error?: unknown }): void { + if (result.reason === 'filtered') { + throw new Error('Subagent turn blocked by provider safety policy'); + } + if (result.reason === 'failed') { + const error = result.error; + if (isProviderRateLimitError(error)) throw error; + const payload = toKimiErrorPayload(error); + if (payload.code === ErrorCodes.PROVIDER_RATE_LIMIT) { + throw providerRateLimitErrorFromPayload(payload); + } + throw error instanceof Error ? error : new Error(String(error ?? 'Subagent turn failed')); + } + if (result.reason === 'cancelled') { + throw userCancellationReason(); + } +} + +function shouldSuppressQueuedAttemptFailureEvent( + options: RunSubagentOptions, + error: unknown, +): boolean { + if (options.suppressRateLimitFailureEvent !== true) return false; + if (isProviderRateLimitError(error)) return true; + return isAbortError(error) || options.signal.aborted; +} + +function providerRateLimitErrorFromPayload(error: KimiErrorPayload): APIProviderRateLimitError { + const requestId = + typeof error.details?.['requestId'] === 'string' ? error.details['requestId'] : null; + return new APIProviderRateLimitError(error.message, requestId); +} + +function abortPromise(signal: AbortSignal): Promise { + if (signal.aborted) { + return Promise.reject(signal.reason ?? userCancellationReason()); + } + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => reject(signal.reason ?? userCancellationReason()), { + once: true, + }); + }); +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function latestAssistantText(messages: readonly ContextMessage[]): string { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]!; + if (message.role !== 'assistant') continue; + return contentText(message.content); + } + return ''; +} + +function contentText(content: ContextMessage['content']): string { + if (typeof content === 'string') return content; + return content + .filter((part): part is Extract<(typeof content)[number], { type: 'text' }> => part.type === 'text') + .map((part) => part.text) + .join(''); } registerScopedService( diff --git a/packages/agent-core-v2/test/background/ids.test.ts b/packages/agent-core-v2/test/background/ids.test.ts index 7eb83fc2a..9cb4f0999 100644 --- a/packages/agent-core-v2/test/background/ids.test.ts +++ b/packages/agent-core-v2/test/background/ids.test.ts @@ -9,7 +9,7 @@ import { IAgentBackgroundService, ProcessBackgroundTask, } from '#/agent/background'; -import type { SessionSubagentHost, SubagentHandle } from '#/session/subagentHost'; +import type { ISessionSubagentHost, SubagentHandle } from '#/session/subagentHost'; import { createTestAgent, type TestAgentContext } from '../harness'; import { createBackgroundTaskPersistence } from './stubs'; @@ -36,7 +36,7 @@ function agentTask( handle, description, { markActiveChildDetached: vi.fn() } as unknown as Pick< - SessionSubagentHost, + ISessionSubagentHost, 'markActiveChildDetached' >, new AbortController(), diff --git a/packages/agent-core-v2/test/background/manager.test.ts b/packages/agent-core-v2/test/background/manager.test.ts index a862f08ac..c7a9a8690 100644 --- a/packages/agent-core-v2/test/background/manager.test.ts +++ b/packages/agent-core-v2/test/background/manager.test.ts @@ -17,7 +17,7 @@ import { ProcessBackgroundTask, type BackgroundTaskInfo, } from '#/agent/background'; -import type { SessionSubagentHost, SubagentHandle } from '#/session/subagentHost'; +import type { ISessionSubagentHost, SubagentHandle } from '#/session/subagentHost'; import { isUserCancellation, userCancellationReason } from '#/_base/utils/abort'; import { configServices, @@ -79,7 +79,7 @@ function agentTask( options: { readonly agentId?: string; readonly subagentType?: string; - readonly subagentHost?: Pick; + readonly subagentHost?: Pick; readonly abortController?: AbortController; readonly timeoutMs?: number; } = {}, diff --git a/packages/agent-core-v2/test/background/rpc-events.test.ts b/packages/agent-core-v2/test/background/rpc-events.test.ts index 484bc7cce..dc6c99cc1 100644 --- a/packages/agent-core-v2/test/background/rpc-events.test.ts +++ b/packages/agent-core-v2/test/background/rpc-events.test.ts @@ -21,7 +21,7 @@ import { IAgentContextMemoryService } from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; import type { HookEngine } from '#/agent/externalHooks/engine'; import { IAgentPromptService } from '#/agent/prompt'; -import type { SessionSubagentHost, SubagentHandle } from '#/session/subagentHost'; +import type { ISessionSubagentHost, SubagentHandle } from '#/session/subagentHost'; import { configServices, createTestAgent, @@ -82,7 +82,7 @@ function agentTask( options: { readonly agentId?: string; readonly subagentType?: string; - readonly subagentHost?: Pick; + readonly subagentHost?: Pick; readonly abortController?: AbortController; readonly timeoutMs?: number; } = {}, diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index 94df4d9cd..6f14c1373 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -57,6 +57,7 @@ import { IAgentShellToolsService, IStorageService, ISessionSubagentHost, + ISessionProcessRunner, IAgentSwarmService, ITelemetryService, ISessionTerminalBackend, @@ -80,13 +81,15 @@ import { bootstrapSeed, createAppScope, resolveBootstrapOptions, + type QueuedSubagentRunResult, + type QueuedSubagentTask, type IDisposable, type Scope, type ScopeSeed, type ServiceIdentifier, } from '#/index'; import { Event } from '#/_base/event'; -import { toDisposable } from '#/_base/di'; +import { Disposable, toDisposable } from '#/_base/di'; import type { PromisifyMethods } from '#/_base/utils/types'; import type { ApprovalResponse } from '#/session/approval'; import type { BackgroundTaskInfo } from '#/agent/background'; @@ -131,7 +134,11 @@ import { ISessionSkillCatalog } from '#/session/sessionSkillCatalog/skillCatalog import { AgentSkillService } from '#/agent/skill/skillService'; import { ModelSkillTool } from '#/agent/skill/tools/modelSkill'; import type { SkillCatalog } from '#/app/globalSkillCatalog/types'; -import { SessionSubagentHostService, type SessionSubagentHost } from '#/session/subagentHost'; +import { + AgentTool, + DEFAULT_AGENT_SUBAGENT_PROFILES, + SessionSubagentHostService, +} from '#/session/subagentHost'; import type { ExecutableToolOutput as ToolOutput, ToolResult } from '#/agent/tool'; import type { PersistedWireRecord, @@ -541,8 +548,91 @@ function createSessionSkillCatalog(catalog: SkillCatalog): ISessionSkillCatalog }; } -export function subagentHostServices(host: SessionSubagentHost): TestAgentServiceOverride { - return agentService(ISessionSubagentHost, new SyncDescriptor(SessionSubagentHostService, [host])); +class TestSessionSubagentHostService extends Disposable implements ISessionSubagentHost { + declare readonly _serviceBrand: undefined; + + constructor( + private readonly host: ISessionSubagentHost, + @IAgentToolRegistryService toolRegistry: IAgentToolRegistryService, + @IAgentBackgroundService background: IAgentBackgroundService, + @IAgentProfileService profile: IAgentProfileService, + @IKaos kaos: IKaos, + @ISessionProcessRunner runner: ISessionProcessRunner, + @ILogService log?: ILogService, + ) { + super(); + this._register( + toolRegistry.register( + new AgentTool(this, background, DEFAULT_AGENT_SUBAGENT_PROFILES, { + log, + gitContext: { cwd: kaos.cwd, runner }, + canRunInBackground: () => { + return profile.isToolActive('TaskList') && + profile.isToolActive('TaskOutput') && + profile.isToolActive('TaskStop'); + }, + }), + ), + ); + } + + getSwarmItem(agentId: string): string | undefined { + return this.host.getSwarmItem(agentId); + } + + startBtw(): Promise { + return this.host.startBtw(); + } + + generateAgentsMd(): Promise { + return this.host.generateAgentsMd(); + } + + spawn( + options: Parameters[0], + ): ReturnType { + return this.host.spawn(options); + } + + resume( + agentId: string, + options: Parameters[1], + ): ReturnType { + return this.host.resume(agentId, options); + } + + retry( + agentId: string, + options: Parameters[1], + ): ReturnType { + return this.host.retry(agentId, options); + } + + getProfileName(agentId: string): Promise { + return this.host.getProfileName(agentId); + } + + markActiveChildDetached(agentId: string): void { + this.host.markActiveChildDetached(agentId); + } + + runQueued( + tasks: readonly QueuedSubagentTask[], + ): Promise>> { + return this.host.runQueued(tasks); + } + + cancelAll(reason?: unknown): void { + this.host.cancelAll(reason); + } + + suspended(event: Parameters[0]): void { + this.host.suspended(event); + } +} + +export function subagentHostServices(host: ISessionSubagentHost): TestAgentServiceOverride { + return agentService(ISessionSubagentHost, new SyncDescriptor(TestSessionSubagentHostService, [host])); } export function goalServices(options: GoalServiceOptions): TestAgentServiceOverride { @@ -939,7 +1029,7 @@ export class AgentTestContext { reg.defineDescriptor(IAgentUserToolService, new SyncDescriptor(AgentUserToolService)); reg.defineDescriptor( ISessionSubagentHost, - new SyncDescriptor(SessionSubagentHostService, [unavailableSubagentHost()]), + new SyncDescriptor(TestSessionSubagentHostService, [unavailableSubagentHost()]), ); }, ], this.serviceOverrides, 'agent'), @@ -1773,13 +1863,15 @@ function createTerminalBackend(): ISessionTerminalBackend { }; } -function unavailableSubagentHost(): SessionSubagentHost { +function unavailableSubagentHost(): ISessionSubagentHost { const fail = async (): Promise => { throw new Error('Subagent host is not configured in this test.'); }; return { + _serviceBrand: undefined, getSwarmItem: () => undefined, startBtw: fail, + generateAgentsMd: fail, spawn: fail, resume: fail, retry: fail, diff --git a/packages/agent-core-v2/test/subagentHost/agent-tool.test.ts b/packages/agent-core-v2/test/subagentHost/agent-tool.test.ts index 89fa11922..d28ff4f49 100644 --- a/packages/agent-core-v2/test/subagentHost/agent-tool.test.ts +++ b/packages/agent-core-v2/test/subagentHost/agent-tool.test.ts @@ -9,7 +9,6 @@ import { AgentToolInputSchema, DEFAULT_SUBAGENT_TIMEOUT_MS, type ISessionSubagentHost, - type SessionSubagentHost, } from '#/session/subagentHost'; import type { AgentToolSubagentMap } from '#/session/subagentHost/agentTool'; import { ToolAccesses } from '#/agent/tool'; @@ -69,7 +68,7 @@ describe('AgentTool direct contract', () => { canRunInBackground, log, }: { - readonly host?: SessionSubagentHost; + readonly host?: ISessionSubagentHost; readonly maxRunningTasks?: number; readonly subagents?: AgentToolSubagentMap; readonly canRunInBackground?: () => boolean; @@ -77,7 +76,7 @@ describe('AgentTool direct contract', () => { } = {}): { readonly ctx: TestAgentContext; readonly background: IAgentBackgroundService; - readonly host: SessionSubagentHost; + readonly host: ISessionSubagentHost; readonly tool: AgentTool; } { const ctx = @@ -739,7 +738,7 @@ describe('Agent tool service runtime', () => { describe('with a resolving subagent host', () => { let ctx: TestAgentContext; - let subagentHost: SessionSubagentHost; + let subagentHost: ISessionSubagentHost; let profile: IAgentProfileService; let tools: IAgentToolRegistryService; @@ -857,7 +856,7 @@ describe('Agent tool service runtime', () => { describe('with a non-resuming subagent host', () => { let ctx: TestAgentContext; - let subagentHost: SessionSubagentHost; + let subagentHost: ISessionSubagentHost; let profile: IAgentProfileService; let tools: IAgentToolRegistryService; @@ -902,11 +901,13 @@ describe('Agent tool service runtime', () => { }); function createSubagentHost( - overrides: Partial = {}, -): SessionSubagentHost { - const host: SessionSubagentHost = { + overrides: Partial = {}, +): ISessionSubagentHost { + const host: ISessionSubagentHost = { + _serviceBrand: undefined, getSwarmItem: vi.fn(), startBtw: vi.fn().mockResolvedValue('btw-url'), + generateAgentsMd: vi.fn().mockResolvedValue(undefined), spawn: vi.fn(), resume: vi.fn(), retry: vi.fn(), diff --git a/packages/agent-core-v2/test/subagentHost/defaultSessionSubagentHost.test.ts b/packages/agent-core-v2/test/subagentHost/subagentHostService.default.test.ts similarity index 88% rename from packages/agent-core-v2/test/subagentHost/defaultSessionSubagentHost.test.ts rename to packages/agent-core-v2/test/subagentHost/subagentHostService.default.test.ts index 73e905103..349e5b0bc 100644 --- a/packages/agent-core-v2/test/subagentHost/defaultSessionSubagentHost.test.ts +++ b/packages/agent-core-v2/test/subagentHost/subagentHostService.default.test.ts @@ -1,9 +1,14 @@ -import { describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { APIProviderRateLimitError } from '@moonshot-ai/kosong'; -import { IAgentLifecycleService } from '#/session/agent-lifecycle'; +import { SyncDescriptor } from '#/_base/di/descriptors'; +import { DisposableStore } from '#/_base/di/lifecycle'; +import { TestInstantiationService } from '#/_base/di/test'; import type { IScopeHandle } from '#/_base/di/scope'; +import { IKaos } from '#/app/kaos'; +import { ITelemetryService } from '#/app/telemetry'; +import { IAgentBackgroundService } from '#/agent/background'; import { IAgentContextMemoryService } from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; import { IAgentExternalHooksService } from '#/agent/externalHooks'; @@ -14,9 +19,12 @@ import { import { IAgentProfileService } from '#/agent/profile'; import { IAgentPromptService } from '#/agent/prompt'; import { IAgentSystemReminderService } from '#/agent/systemReminder'; -import { ITelemetryService } from '#/app/telemetry'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry'; import { IAgentUsageService } from '#/agent/usage'; -import { DefaultSessionSubagentHost } from '#/session/subagentHost/defaultSessionSubagentHost'; +import { IAgentLifecycleService } from '#/session/agent-lifecycle'; +import { ISessionProcessRunner } from '#/session/process'; +import { ISessionMetadata } from '#/session/session-metadata'; +import { ISessionSubagentHost, SessionSubagentHostService } from '#/session/subagentHost'; const CHILD_SUMMARY = 'child summary '.repeat(20); @@ -157,12 +165,39 @@ function makeAgents(parent: IScopeHandle, children: Record }; } -describe('DefaultSessionSubagentHost', () => { +describe('SessionSubagentHostService built-in host behavior', () => { + let disposables: DisposableStore; + let ix: TestInstantiationService; + + beforeEach(() => { + disposables = new DisposableStore(); + ix = disposables.add(new TestInstantiationService()); + }); + + afterEach(() => { + disposables.dispose(); + }); + + function createHost( + agents: IAgentLifecycleService, + metadata?: Partial, + ): SessionSubagentHostService { + ix.stub(IAgentLifecycleService, agents); + if (metadata !== undefined) ix.stub(ISessionMetadata, metadata); + ix.stub(IAgentToolRegistryService, { register: vi.fn(() => ({ dispose: () => {} })) }); + ix.stub(IAgentBackgroundService, {}); + ix.stub(IAgentProfileService, { isToolActive: vi.fn().mockReturnValue(false) }); + ix.stub(IKaos, { cwd: '/repo' }); + ix.stub(ISessionProcessRunner, { exec: vi.fn() }); + ix.set(ISessionSubagentHost, new SyncDescriptor(SessionSubagentHostService)); + return ix.get(ISessionSubagentHost) as SessionSubagentHostService; + } + it('aborts a running subagent when the caller signal aborts', async () => { const parent = fakeScope('main'); const child = fakeScope('child', { result: new Promise(() => {}) }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const controller = new AbortController(); const handle = await host.spawn({ @@ -183,7 +218,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main', { events }); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'explore', @@ -211,7 +246,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main', { initialText: 'short' }); const child = fakeScope('child', { initialText: 'short' }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'coder', @@ -231,7 +266,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); await host.spawn({ profileName: 'coder', @@ -266,11 +301,7 @@ describe('DefaultSessionSubagentHost', () => { }, }), }; - const host = new DefaultSessionSubagentHost( - agents as unknown as IAgentLifecycleService, - 'main', - metadata as never, - ); + const host = createHost(agents as unknown as IAgentLifecycleService, metadata as never); await expect( host.resume('child', { @@ -293,7 +324,7 @@ describe('DefaultSessionSubagentHost', () => { result: Promise.resolve({ reason: 'failed', error: new Error('boom') }), }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'coder', @@ -317,7 +348,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main', { result: new Promise(() => {}), events }); const child = fakeScope('child', { result: new Promise(() => {}) }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const controller = new AbortController(); const handle = await host.spawn({ @@ -347,7 +378,7 @@ describe('DefaultSessionSubagentHost', () => { createMain: vi.fn(), create: vi.fn(), }; - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.resume('child', { parentToolCallId: 'call_agent', @@ -377,7 +408,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main', { parentMessages }); const child = fakeScope('btw-child'); const agents = makeAgents(parent, { 'btw-child': child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); await expect(host.startBtw()).resolves.toBe('btw-child'); expect(agents.create).toHaveBeenCalledWith({ parentAgentId: 'main', type: 'sub' }); @@ -415,7 +446,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const results = await host.runQueued([ { @@ -445,7 +476,7 @@ describe('DefaultSessionSubagentHost', () => { pending.push(d); return fakeScope(`child-${pending.length}`, { result: d.promise }); }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const tasks = Array.from({ length: 4 }, (_, index) => ({ kind: 'spawn' as const, @@ -488,7 +519,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main', { result: new Promise(() => {}) }); const child = fakeScope('child', { result: new Promise(() => {}) }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); await host.spawn({ profileName: 'coder', @@ -507,7 +538,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'explore', @@ -534,7 +565,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'explore', @@ -561,7 +592,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'coder', @@ -585,7 +616,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child', { ready: ready.promise }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const onReady = vi.fn(); const handle = await host.spawn({ @@ -613,7 +644,7 @@ describe('DefaultSessionSubagentHost', () => { createMain: vi.fn(), create: vi.fn(), }; - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.retry('child', { parentToolCallId: 'call_agent', @@ -633,7 +664,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child', { result: Promise.resolve({ reason: 'filtered' }) }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'coder', @@ -656,7 +687,7 @@ describe('DefaultSessionSubagentHost', () => { }), }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'coder', @@ -674,7 +705,7 @@ describe('DefaultSessionSubagentHost', () => { const events: unknown[] = []; const parent = fakeScope('main', { events }); const agents = makeAgents(parent, {}); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); host.suspended({ task: { @@ -703,7 +734,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child', { result: new Promise(() => {}) }); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'coder', @@ -730,7 +761,7 @@ describe('DefaultSessionSubagentHost', () => { const parent = fakeScope('main'); const child = fakeScope('child'); const agents = makeAgents(parent, { child }); - const host = new DefaultSessionSubagentHost(agents as unknown as IAgentLifecycleService, 'main'); + const host = createHost(agents as unknown as IAgentLifecycleService); const handle = await host.spawn({ profileName: 'explore', diff --git a/packages/agent-core-v2/test/subagentHost/subagentHostService.test.ts b/packages/agent-core-v2/test/subagentHost/subagentHostService.test.ts index 333ddde76..6cddd2de4 100644 --- a/packages/agent-core-v2/test/subagentHost/subagentHostService.test.ts +++ b/packages/agent-core-v2/test/subagentHost/subagentHostService.test.ts @@ -44,7 +44,7 @@ describe('SessionSubagentHostService DI wiring', () => { afterEach(() => disposables.dispose()); - it('constructs a DefaultSessionSubagentHost when no host override is provided', async () => { + it('uses the built-in host behavior when no host override is provided', async () => { const register = vi.fn(() => ({ dispose: () => {} })); const parent = fakeProfileScope('main'); const child = fakeProfileScope('child'); @@ -56,7 +56,7 @@ describe('SessionSubagentHostService DI wiring', () => { ix.stub(ISessionMetadata, { ready: Promise.resolve(), read: vi.fn().mockResolvedValue({ agents: {} }), - onDidChange: () => ({ dispose: () => {} }), + onDidChangeMetadata: () => ({ dispose: () => {} }), update: vi.fn(), setTitle: vi.fn(), setArchived: vi.fn(), @@ -68,7 +68,7 @@ describe('SessionSubagentHostService DI wiring', () => { ix.stub(IKaos, { cwd: '/repo' }); ix.stub(ISessionProcessRunner, { exec: vi.fn() }); ix.stub(ILogService, { warn: vi.fn(), info: vi.fn(), debug: vi.fn(), error: vi.fn() }); - ix.set(ISessionSubagentHost, new SyncDescriptor(SessionSubagentHostService, [undefined])); + ix.set(ISessionSubagentHost, new SyncDescriptor(SessionSubagentHostService)); const service = ix.get(ISessionSubagentHost); diff --git a/packages/agent-core-v2/test/swarm/swarm.test.ts b/packages/agent-core-v2/test/swarm/swarm.test.ts index 071bd342a..997ee41c5 100644 --- a/packages/agent-core-v2/test/swarm/swarm.test.ts +++ b/packages/agent-core-v2/test/swarm/swarm.test.ts @@ -41,12 +41,21 @@ function mockSubagentHost({ readonly getSwarmItem?: (agentId: string) => string | undefined; // eslint-disable-next-line @typescript-eslint/no-explicit-any readonly runQueued?: (...args: any[]) => any; -} = {}) { +} = {}): ISessionSubagentHost { return { + _serviceBrand: undefined, getSwarmItem: vi.fn(getSwarmItem), + startBtw: vi.fn(), + generateAgentsMd: vi.fn(), + spawn: vi.fn(), + resume: vi.fn(), + retry: vi.fn(), + getProfileName: vi.fn(), + markActiveChildDetached: vi.fn(), runQueued, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any; + cancelAll: vi.fn(), + suspended: vi.fn(), + }; } function mockSwarmMode() { diff --git a/packages/agent-core-v2/test/tool/tool.test.ts b/packages/agent-core-v2/test/tool/tool.test.ts index 5f990b54c..359bde0a5 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -7,7 +7,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { IAgentContextMemoryService } from '#/agent/contextMemory'; import { HookEngine } from '#/agent/externalHooks/engine'; import { IAgentProfileService } from '#/agent/profile'; -import type { SessionSubagentHost } from '#/session/subagentHost'; +import type { ISessionSubagentHost } from '#/session/subagentHost'; import { IAgentToolRegistryService } from '#/agent/toolRegistry'; import { createFakeKaos } from '../tools/fixtures/fake-kaos'; import { @@ -220,7 +220,7 @@ describe('Agent tools', () => { }); describe('foreground Agent tool recovery', () => { - let subagentHost: SessionSubagentHost; + let subagentHost: ISessionSubagentHost; beforeEach(() => { const completion = Promise.reject( @@ -228,6 +228,10 @@ describe('Agent tools', () => { ); void completion.catch(() => undefined); subagentHost = { + _serviceBrand: undefined, + getSwarmItem: vi.fn(), + startBtw: vi.fn(), + generateAgentsMd: vi.fn(), spawn: vi.fn().mockResolvedValue({ agentId: 'agent-child', profileName: 'coder', @@ -235,7 +239,13 @@ describe('Agent tools', () => { completion, }), resume: vi.fn(), - } as unknown as SessionSubagentHost; + retry: vi.fn(), + getProfileName: vi.fn(), + markActiveChildDetached: vi.fn(), + runQueued: vi.fn(), + cancelAll: vi.fn(), + suspended: vi.fn(), + }; ctx = createTestAgent(subagentHostServices(subagentHost)); profile = ctx.get(IAgentProfileService); profile.update({ activeToolNames: ['Agent'] }); diff --git a/packages/agent-core-v2/test/turn/turn.test.ts b/packages/agent-core-v2/test/turn/turn.test.ts index 4b0d7d990..587d50dae 100644 --- a/packages/agent-core-v2/test/turn/turn.test.ts +++ b/packages/agent-core-v2/test/turn/turn.test.ts @@ -35,7 +35,7 @@ import { IAgentToolRegistryService } from '#/agent/toolRegistry'; import type { QueuedSubagentRunResult, QueuedSubagentTask, - SessionSubagentHost, + ISessionSubagentHost, } from '#/session/subagentHost'; import { recordingTelemetry, type TelemetryRecord } from '../telemetry/stubs'; import { createFakeKaos } from '../tools/fixtures/fake-kaos'; @@ -457,7 +457,7 @@ describe('Agent turn flow', () => { })); }); const subagentHost = mockSubagentHost({ - runQueued: runQueued as unknown as SessionSubagentHost['runQueued'], + runQueued: runQueued as unknown as ISessionSubagentHost['runQueued'], }); const ctx = testAgent(subagentHostServices(subagentHost)); ctx.configure({ tools: ['AgentSwarm'] }); @@ -1995,11 +1995,24 @@ function agentSwarmCall(): ToolCall { }; } -function mockSubagentHost>( +function mockSubagentHost>( host: T, -): T & SessionSubagentHost { - return { spawn: vi.fn(), resume: vi.fn(), runQueued: vi.fn(), ...host } as unknown as T & - SessionSubagentHost; +): T & ISessionSubagentHost { + return { + _serviceBrand: undefined, + getSwarmItem: vi.fn(), + startBtw: vi.fn(), + generateAgentsMd: vi.fn(), + spawn: vi.fn(), + resume: vi.fn(), + retry: vi.fn(), + getProfileName: vi.fn(), + markActiveChildDetached: vi.fn(), + runQueued: vi.fn(), + cancelAll: vi.fn(), + suspended: vi.fn(), + ...host, + } as T & ISessionSubagentHost; } interface ApiErrorTelemetryCase { From 186ed6eaa567c3eb00d59aaf6f3d9496c575a8a9 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 1 Jul 2026 21:29:51 +0800 Subject: [PATCH 2/6] feat: implement goal core workflow --- GOAL.md | 231 ++++++ .../agent-core-v2/docs/di-scope-domains.puml | 6 + packages/agent-core-v2/src/agent/goal/goal.ts | 2 +- .../src/agent/goal/goalService.ts | 233 +++++- .../src/agent/goal/tools/update-goal.ts | 36 +- packages/agent-core-v2/test/goal/goal.test.ts | 772 +++++++++++------- packages/server-v2/src/transport/errors.ts | 6 + packages/server-v2/test/rpc.test.ts | 87 ++ 8 files changed, 1019 insertions(+), 354 deletions(-) create mode 100644 GOAL.md diff --git a/GOAL.md b/GOAL.md new file mode 100644 index 000000000..c0fdc2a36 --- /dev/null +++ b/GOAL.md @@ -0,0 +1,231 @@ +# Goal 功能拆分 + +本文把 agent-core 中 goal mode 的能力拆成三部分: + +1. 核心工作流:没有它就不能运行 goal。 +2. 统计 / token 数限制:让 goal 可度量、可限额、可审计。 +3. 用户交互相关:让用户可以安全启动、理解、控制和恢复 goal。 + +## 1. 核心工作流 + +核心工作流是 goal mode 的运行骨架。它负责创建结构化目标、维护状态机、把普通 turn 串成自治多轮执行,并让模型用机器可读状态结束或停放目标。 + +### 目标状态 + +同一个 main agent 同时最多只有一个当前 goal。goal 不是普通聊天文本,而是 runtime 持有的结构化状态,至少包含目标、可选完成标准、当前状态、停止原因和运行统计。 + +状态分为四类: + +- `active`:正在被 goal driver 推进。只有这个状态会自动运行下一轮。 +- `paused`:暂停但保留目标。通常来自用户暂停、中断、进程恢复后降级、provider 或 runtime 错误。可以恢复。 +- `blocked`:目标遇到真实阻塞但保留目标。通常来自模型判断需要外部输入、目标无法按当前表述完成、预算达到、prompt hook 阻止。可以恢复。 +- `complete`:瞬时完成状态。runtime 发出完成事件后立即清除 goal,不长期持久化。 + +没有 `cancelled` 状态。取消就是清除 goal,并提醒模型忽略之前关于该目标的 active reminder。 + +### 创建和替换 + +创建 goal 时,runtime 需要校验目标不能为空、不能过长。已有 active、paused 或 blocked goal 时,默认拒绝创建新 goal,防止静默覆盖。只有用户或调用方明确要求替换时,才先清除旧 goal,再创建新 goal。 + +新 goal 创建后进入 `active`,写入持久记录,并发出 goal 更新事件。 + +### 多轮驱动 + +goal driver 的职责是把一个 active goal 推进成连续的普通 turn: + +- turn 开始时如果 goal 已经是 `active`,进入 goal driver。 +- 普通 turn 中如果模型创建了 goal,或把 paused/blocked goal 恢复成 active,当前 turn 结束后 goal driver 接管继续执行。 +- driver 每次只运行一个普通 turn。 +- 每个 turn 结束后读取 goal 状态。 +- goal 仍是 `active` 时,runtime 自动追加 continuation prompt 并启动下一轮。 +- goal 变成 `paused`、`blocked` 或被清除时,driver 停止。 + +模型如果不调用状态更新工具,且 goal 仍是 active,runtime 会继续下一轮。模型不能只靠自然语言说“完成了”来结束 goal,必须给出结构化状态信号。 + +### Goal 注入 + +每个 goal turn 的边界,runtime 会把当前 goal 状态注入上下文。注入内容包括: + +- 当前正在 goal mode。 +- 目标和完成标准是什么。 +- 目标文本是用户提供的数据,不能覆盖 system/developer 指令、工具 schema、权限规则或 host 控制。 +- 当前状态和进度。 +- 模型应该做简短自审,然后推进一个连贯工作切片。 +- 简单、已完成、不可能、不安全、矛盾的目标,应在同一轮内直接标记 complete 或 blocked。 +- 只有全部要求完成、验证通过、没有下一步有用动作时,才能标记 complete。 +- 外部条件或用户输入阻塞时,应标记 blocked。 +- 不要只做了计划、总结、第一版或部分结果就标记 complete。 + +goal 注入只在 turn / continuation 边界做,不在每个 model step 都做,避免上下文重复膨胀,也有利于 prompt cache。 + +paused 和 blocked goal 的注入更轻: + +- paused:提醒模型目标存在但当前不应自治推进,除非用户明确要求继续。 +- blocked:提醒模型目标被阻塞且当前不自治推进,除非用户要求处理或恢复。 + +### Continuation prompt + +当 goal 仍是 active,runtime 会追加一个系统触发输入,含义相当于“继续朝当前 active goal 工作”。它不只是简单续跑,还要求模型每轮重新判断: + +- 是否已经完成。 +- 是否遇到真实阻塞。 +- 是否应该只推进一个合理切片后继续下一轮。 +- 是否应该避免发散或启动无关工作。 +- 除非真实阻塞,否则不要向用户要输入。 + +### 完成、阻塞和暂停 + +模型通过结构化状态更新控制 goal 生命周期: + +- `complete`:目标已满足,runtime 发出完成事件并清除 goal。 +- `blocked`:遇到真实阻塞,runtime 保留 goal 并停止自治推进。 +- `paused`:暂时放下 goal,runtime 保留 goal 并停止自治推进。 +- `active`:恢复 paused 或 blocked goal。 + +状态更新工具的输入应保持窄,只表达机器状态。完成总结或阻塞原因由模型随后给用户说明。 + +当模型标记 complete 后,runtime 应再给模型一次收尾机会,生成简短最终回复,说明 goal 已完成、主要做了什么、跑了什么验证。 + +当模型标记 blocked 后,runtime 应再给模型一次收尾机会,说明具体阻塞、需要什么输入或变化才能继续。 + +如果当前 turn 已经没有 step 预算,不应为了收尾总结强行再跑一步,避免把“没法写总结”变成 turn 失败。 + +### 错误停车 + +goal mode 把技术运行失败视为可恢复停车: + +- 用户中断当前 turn:goal 变 paused。 +- provider rate limit:goal 变 paused。 +- provider 连接错误、认证错误、API 错误:goal 变 paused。 +- 模型配置错误:goal 变 paused。 +- runtime 异常:goal 变 paused。 +- provider safety filter:goal 变 paused。 + +业务、规则或外部条件阻塞则变 blocked: + +- prompt hook 阻止目标。 +- 模型判断无法继续。 +- 预算达到。 +- 需要用户或外部系统提供新条件。 + +### 持久化和恢复 + +goal 的创建、更新、完成、阻塞、清除应写入可恢复记录。session 恢复时,runtime 用记录重建 goal。 + +恢复时如果发现 goal 原来是 active,不应自动继续跑,而是降级为 paused。因为旧进程中的 active turn 不可能还活着,自动继续会造成重启后偷偷消耗资源。 + +paused 和 blocked 原样保留。complete 理论上不长期存在,因为完成后会清除。 + +fork session 时不继承源 session 的 goal,并提醒模型不要继续源 session 的旧目标。 + +## 2. 统计 / token 数限制 + +这一部分让 goal 可度量、可限额、可审计。没有它,goal 仍然可以运行,但不可控。 + +### 运行统计 + +goal 统计包括: + +- continuation turn 数。 +- token 数。 +- active wall-clock 时间。 + +统计只在 goal 是 `active` 时增长。paused 和 blocked 期间不继续计数。 + +turn 统计在每个 goal turn 准备运行时增加,因此模型在某一轮里标记 complete 时,这一轮也计入最终统计。 + +token 统计在 model step 结束后累计。没有 active goal 时,不记入 goal。token 统计应以静默更新为主,不应每一步都刷 UI。 + +时间统计只计算 active pursuit 时间。进入 active 时开启计时区间,离开 active 时折算进累计时间;pause/resume 会形成新的 active 区间。 + +### 预算 + +goal 预算包括: + +- turn budget。 +- token budget。 +- wall-clock budget。 + +默认没有预算。只有用户明确给出硬限制时才设置,例如“最多 20 轮”“不超过 500k token”“30 分钟内”。模糊表达如“尽快”“别花太久”不能设置预算,模型也不能自行发明预算。 + +时间预算需要合理范围。过短或过长应拒绝。turn 和 token 预算应规范化为正整数。 + +### 预算硬停 + +预算检查应发生在 goal turn 开始前和结束后。token budget 还应在 model step 后触发停止,避免超额后继续下一步。 + +一旦达到预算,runtime 应直接把 goal 标记为 blocked,原因是配置预算已达到。这个 blocked 仍可恢复,但如果预算不变,恢复后可能立刻再次 blocked。 + +### 预算引导和最终统计 + +当预算未接近时,模型提示应鼓励稳定推进。当任一预算达到 75% 以上时,提示应转为收敛,避免启动新的可选工作。 + +complete 和 blocked 的最终回复提示应包含 worked turns、elapsed time、tokens used 等统计信息。UI 事件也应带当前 snapshot 和变化类型。 + +telemetry 可以记录 goal 创建、预算设置、continuation、状态变化、清除等事件,但不应包含目标文本、停止原因等敏感内容。 + +## 3. 用户交互相关 + +这一部分让用户可以安全启动、理解、控制和恢复 goal。没有它,runtime 仍可能运行,但交互体验和安全边界不足。 + +### 生命周期控制 + +用户可以直接控制 goal: + +- 创建。 +- 查看。 +- 暂停。 +- 恢复。 +- 取消。 + +这些操作可以不经过模型 turn。pause 把 active goal 变 paused;resume 把 paused 或 blocked goal 变 active;cancel 直接清除当前 goal。 + +resume 会清除旧停止原因,表示开始新的尝试。paused/blocked goal 不会因为用户发普通消息就自动继续。 + +### 模型发起 goal 的确认 + +模型可以代表用户创建 goal,但只有在用户明确要求启动 goal、自治工作,或宿主 goal-intake 提示要求时才应该这样做。普通请求不能被模型擅自升级成 goal。 + +模型发起 CreateGoal 时,非 auto 权限模式下应触发用户确认。确认菜单允许用户选择本次 goal 的运行权限模式。用户拒绝则 goal 不创建。 + +`GetGoal`、`SetGoalBudget`、`UpdateGoal` 只改 goal runtime 状态,默认可以更容易批准。真正写文件、跑 shell、访问敏感路径等仍走普通权限系统。 + +### 暂停、阻塞和取消后的提示 + +paused goal 的上下文提示应说明目标存在但当前不应继续做,除非用户明确要求继续。 + +blocked goal 的上下文提示应说明目标被阻塞且当前不自治推进,可以在用户要求时帮助解阻,否则正常处理当前请求。 + +cancel 后应追加提醒,让模型忽略旧 goal 的 active reminder,避免旧上下文诱导模型继续已经取消的目标。 + +### 完成和阻塞的用户回复 + +complete 后,goal 被清除,模型应给用户一条简短完成总结,说明完成了什么、做了什么验证。 + +blocked 后,goal 保留,模型应给用户一条简短阻塞说明,说明具体阻塞和继续所需输入、权限、外部条件或变更。 + +### Tool 暴露和隔离 + +goal 工具只给 main agent。subagent 不应直接创建、恢复、结束主 goal。 + +没有 goal 时,模型不应看到 `UpdateGoal` 和 `SetGoalBudget`。有 goal 时才暴露这些控制工具。 + +goal ID 不应暴露给模型,因为它只是 runtime/UI 内部标识,没有用户语义。 + +### 辅助写 goal + +`write-goal` 类能力用于帮助用户把粗糙意图整理成适合 goal mode 的完成契约。好的 goal 应明确: + +- end state:什么条件必须变成真。 +- proof:用什么可观察证据证明完成。 +- boundaries:工作范围和禁止触碰的内容。 +- loop:如何迭代推进。 +- stop rule:什么情况下停止并报告,而不是强行继续。 + +预算是 opt-in,不应默认加入,也不应把 turn cap 写进目标文本。 + +### UI 和会话语义 + +goal 创建、暂停、恢复、阻塞、完成、清除都应发出 goal updated 事件。lifecycle 变化和 completion 变化应区分。completion 是一次终局事件,然后 snapshot 变 null。blocked/paused 保留 snapshot,UI 可以继续展示可恢复 goal。 + +session 恢复时,active goal 会变 paused,避免重启后自动继续。fork session 时不继承 goal,并提醒模型不要继续源 session 的目标。 diff --git a/packages/agent-core-v2/docs/di-scope-domains.puml b/packages/agent-core-v2/docs/di-scope-domains.puml index 2927ece5a..14916f4fa 100644 --- a/packages/agent-core-v2/docs/di-scope-domains.puml +++ b/packages/agent-core-v2/docs/di-scope-domains.puml @@ -243,6 +243,10 @@ goal --> systemReminder #34495E goal --> replayBuilder #34495E goal --> telemetry #34495E goal --> contextInjector #34495E +goal --> contextMemory #34495E +goal --> turn #34495E +goal --> toolRegistry #34495E +goal --> permissionMode #34495E skill --> prompt #34495E skill --> eventSink #34495E skill --> wireRecord #34495E @@ -349,6 +353,8 @@ fullCompaction ..> turn #16A085 : hooks.onContextOverflow permissionRules ..> wireRecord #16A085 : permission.rules.add / record_approval_result skill ..> wireRecord #16A085 : skill.activate goal ..> wireRecord #16A085 : goal.create/update/clear +goal ..> turn #16A085 : turn lifecycle hooks +goal ..> eventSink #16A085 : hook.result / goal.updated microCompaction ..> wireRecord #16A085 : micro_compaction.apply / full_compaction.complete swarm ..> wireRecord #16A085 : swarm_mode.enter/exit swarm ..> turn #16A085 : hooks.onEnded diff --git a/packages/agent-core-v2/src/agent/goal/goal.ts b/packages/agent-core-v2/src/agent/goal/goal.ts index 7fe24a13e..ae9d7be04 100644 --- a/packages/agent-core-v2/src/agent/goal/goal.ts +++ b/packages/agent-core-v2/src/agent/goal/goal.ts @@ -17,7 +17,7 @@ export interface IAgentGoalService { createGoal(input: CreateGoalInput, actor?: GoalActor): Promise; pauseGoal(input?: GoalReasonInput, actor?: GoalActor): Promise; resumeGoal(input?: GoalReasonInput, actor?: GoalActor): Promise; - cancelGoal(actor?: GoalActor): Promise; + cancelGoal(input?: GoalReasonInput, actor?: GoalActor): Promise; setBudgetLimits( input: { readonly budgetLimits: GoalBudgetLimits }, actor?: GoalActor, diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index 2f18aefc7..2180e2a12 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -1,17 +1,42 @@ -import { - randomUUID } from 'node:crypto'; +/** + * `goal` domain (L4) - `IAgentGoalService` implementation. + * + * Owns the per-agent goal lifecycle; persists records through `wireRecord`, + * broadcasts through `eventSink`, injects reminders through `contextInjector`, + * drives continuation turns through `turn`, updates context through + * `contextMemory`, writes system reminders through `systemReminder`, registers + * model tools through `toolRegistry`, and reports telemetry through + * `telemetry`. Bound at Agent scope. + */ + +import { randomUUID } from 'node:crypto'; import { InstantiationType } from '#/_base/di/extensions'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; +import { Disposable } from "#/_base/di"; import { - Disposable, -} from "#/_base/di"; -import { ErrorCodes, KimiError } from "#/errors"; + ErrorCodes, + KimiError, + toKimiErrorPayload, + type KimiErrorPayload, +} from "#/errors"; import { IAgentContextInjectorService } from '#/agent/contextInjector'; +import { + ensureMessageId, + IAgentContextMemoryService, + type ContextMessage, + type PromptOrigin, +} from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; import { IAgentPermissionModeService } from '#/agent/permissionMode'; import { IAgentReplayBuilderService } from '#/agent/replayBuilder'; import { IAgentSystemReminderService } from '#/agent/systemReminder'; +import { + IAgentTurnService, + type Turn, + type TurnEndedContext, + type TurnStepContext, +} from '#/agent/turn'; import type { TelemetryProperties } from '#/app/telemetry'; import { ITelemetryService } from '#/app/telemetry'; import { IAgentToolRegistryService } from '#/agent/toolRegistry'; @@ -38,6 +63,10 @@ import type { } from './types'; import { CreateGoalTool } from '#/agent/goal/tools/create-goal'; import { GetGoalTool } from '#/agent/goal/tools/get-goal'; +import { + buildGoalBlockedReasonPrompt, + buildGoalCompletionSummaryPrompt, +} from '#/agent/goal/tools/outcome-prompts'; import { SetGoalBudgetTool } from '#/agent/goal/tools/set-goal-budget'; import { UpdateGoalTool } from '#/agent/goal/tools/update-goal'; @@ -76,6 +105,36 @@ const GOAL_FORK_CLEARED_REMINDER = [ 'Handle requests normally unless the user starts a new goal.', ].join(' '); +const GOAL_CONTINUATION_ORIGIN: PromptOrigin = { + kind: 'system_trigger', + name: 'goal_continuation', +}; +const GOAL_COMPLETION_REMINDER_NAME = 'goal_completion_summary'; +const GOAL_BLOCKED_REMINDER_NAME = 'goal_blocked_reason'; +const GOAL_RATE_LIMIT_PAUSE_REASON = 'Paused after provider rate limit'; +const GOAL_PROVIDER_CONNECTION_PAUSE_PREFIX = 'Paused after provider connection error'; +const GOAL_PROVIDER_AUTH_PAUSE_PREFIX = 'Paused after provider authentication error'; +const GOAL_PROVIDER_API_PAUSE_PREFIX = 'Paused after provider API error'; +const GOAL_MODEL_CONFIG_PAUSE_PREFIX = 'Paused after model configuration error'; +const GOAL_RUNTIME_PAUSE_PREFIX = 'Paused after runtime error'; +const GOAL_PROVIDER_FILTERED_PAUSE_REASON = 'Paused after provider safety policy block'; +const LLM_NOT_SET_MESSAGE = 'LLM not set, send "/login" to login'; + +const GOAL_CONTINUATION_PROMPT = [ + 'Continue working toward the active goal.', + 'Keep the self-audit brief. 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, weigh the objective and any completion criteria', + 'against the work done so far. Goal mode is iterative: do one coherent slice of work, then', + 'reassess. Call UpdateGoal with `complete` only when all required work is done, any stated', + 'validation has passed, and there is no useful next action. Do not mark complete after only', + 'producing a plan, summary, first pass, or partial result. If an external condition or required', + 'user input prevents progress, or the objective cannot be completed as stated, call UpdateGoal', + 'with `blocked`. Otherwise keep going - use the existing conversation context and your tools,', + 'and do not ask the user for input unless a real blocker prevents progress.', +].join(' '); + export interface GoalServiceOptions { readonly enabled?: boolean | (() => boolean); readonly injection?: GoalInjectionOptions; @@ -98,6 +157,10 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { declare readonly _serviceBrand: undefined; private state: GoalState | undefined; + private readonly goalDrivenTurns = new Set(); + private readonly countedGoalTurns = new Set(); + private readonly goalOutcomeContinuationTurns = new Set(); + private readonly promptHookBlockedTurns = new Set(); constructor( private readonly options: GoalServiceOptions = {}, @@ -106,7 +169,9 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { @IAgentSystemReminderService private readonly reminders: IAgentSystemReminderService, @IAgentReplayBuilderService private readonly replayBuilder: IAgentReplayBuilderService, @ITelemetryService private readonly telemetry: ITelemetryService, - @IAgentContextInjectorService private readonly dynamicInjector: IAgentContextInjectorService, + @IAgentContextInjectorService dynamicInjector: IAgentContextInjectorService, + @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, + @IAgentTurnService private readonly turnService: IAgentTurnService, @IAgentToolRegistryService toolRegistry: IAgentToolRegistryService, @IAgentPermissionModeService private readonly permissionMode: IAgentPermissionModeService, ) { @@ -146,11 +211,42 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { this.normalizeAfterReplay(); }), ); + this._register( + turnService.hooks.onLaunched.register('goal-track-launched-turn', (ctx, next) => { + this.handleTurnLaunched(ctx.turn); + return next(); + }), + ); + this._register( + turnService.hooks.beforeStep.register('goal-count-turn', async (ctx, next) => { + await this.handleBeforeStep(ctx); + await next(); + }), + ); + this._register( + turnService.hooks.afterStep.register('goal-outcome-continuation', async (ctx, next) => { + await next(); + this.handleAfterStep(ctx); + }), + ); + this._register( + turnService.hooks.onEnded.register('goal-drive-continuation', async (ctx, next) => { + await next(); + await this.handleTurnEnded(ctx); + }), + ); + this._register( + events.on((event) => { + if (event.type === 'hook.result' && event.blocked === true) { + this.promptHookBlockedTurns.add(event.turnId); + } + }), + ); this._register(toolRegistry.register(new CreateGoalTool(this, this.permissionMode))); this._register(toolRegistry.register(new GetGoalTool(this))); this._register(toolRegistry.register(new SetGoalBudgetTool(this))); - this._register(toolRegistry.register(new UpdateGoalTool(this, this.reminders))); + this._register(toolRegistry.register(new UpdateGoalTool(this))); } get enabled(): boolean { @@ -289,7 +385,10 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { return this.toSnapshot(state); } - async cancelGoal(actor: GoalActor = 'user'): Promise { + async cancelGoal( + _input: GoalReasonInput = {}, + actor: GoalActor = 'user', + ): Promise { const state = this.requireState(); const snapshot = this.toSnapshot(state); this.clearInternal(actor); @@ -314,7 +413,14 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { change: { kind: 'lifecycle', status: 'blocked', reason: input.reason, actor }, }); this.appendStatusUpdate(state, actor, input.reason); - return this.toSnapshot(state); + const snapshot = this.toSnapshot(state); + if (actor === 'model') { + this.reminders.appendSystemReminder(buildGoalBlockedReasonPrompt(snapshot), { + kind: 'system_trigger', + name: GOAL_BLOCKED_REMINDER_NAME, + }); + } + return snapshot; } async markComplete( @@ -334,6 +440,12 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { stats: this.statsOf(state), actor, }); + if (actor === 'model') { + this.reminders.appendSystemReminder(buildGoalCompletionSummaryPrompt(snapshot), { + kind: 'system_trigger', + name: GOAL_COMPLETION_REMINDER_NAME, + }); + } this.clearInternal(actor); return snapshot; } @@ -361,6 +473,66 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { return this.toSnapshot(state); } + private handleTurnLaunched(turn: Turn): void { + if (this.state?.status === 'active') this.goalDrivenTurns.add(turn.id); + this.goalOutcomeContinuationTurns.delete(turn.id); + this.promptHookBlockedTurns.delete(turn.id); + } + + private async handleBeforeStep(ctx: TurnStepContext): Promise { + if (!this.goalDrivenTurns.has(ctx.turn.id)) return; + if (this.countedGoalTurns.has(ctx.turn.id)) return; + this.countedGoalTurns.add(ctx.turn.id); + await this.incrementTurn(); + } + + private handleAfterStep(ctx: TurnStepContext): void { + if (this.goalOutcomeContinuationTurns.has(ctx.turn.id)) return; + if (!isGoalOutcomeReminder(this.context.get().at(-1))) return; + this.goalOutcomeContinuationTurns.add(ctx.turn.id); + ctx.continueTurn = true; + } + + private async handleTurnEnded(ctx: TurnEndedContext): Promise { + this.goalDrivenTurns.delete(ctx.turn.id); + this.countedGoalTurns.delete(ctx.turn.id); + this.goalOutcomeContinuationTurns.delete(ctx.turn.id); + + const blockedByPromptHook = this.promptHookBlockedTurns.delete(ctx.turn.id); + if (blockedByPromptHook) { + await this.markBlocked({ reason: 'Blocked by UserPromptSubmit hook' }); + return; + } + + if (ctx.result.reason === 'cancelled') { + await this.pauseOnInterrupt({ reason: 'Paused after interruption' }); + return; + } + if (ctx.result.reason === 'failed') { + await this.pauseActiveGoal({ reason: goalFailurePauseReason(ctx.result.error) }); + return; + } + if (ctx.result.reason === 'filtered') { + await this.pauseActiveGoal({ reason: GOAL_PROVIDER_FILTERED_PAUSE_REASON }); + return; + } + + if (this.state?.status !== 'active') return; + if (this.turnService.getActiveTurn() !== undefined) return; + this.launchContinuationTurn(); + } + + private launchContinuationTurn(): void { + const message = ensureMessageId({ + role: 'user', + content: [{ type: 'text', text: GOAL_CONTINUATION_PROMPT }], + toolCalls: [], + origin: GOAL_CONTINUATION_ORIGIN, + }); + this.context.splice(this.context.get().length, 0, [message]); + this.turnService.launch(GOAL_CONTINUATION_ORIGIN, message.id); + } + private normalizeAfterReplay(): void { const state = this.state; if (state === undefined) return; @@ -593,10 +765,51 @@ function normalizeCompletionCriterion(value: string | undefined): string | undef return trimmed?.length ? trimmed : undefined; } +function isGoalOutcomeReminder(message: ContextMessage | undefined): boolean { + if (message?.origin?.kind !== 'system_trigger') return false; + return ( + message.origin.name === GOAL_COMPLETION_REMINDER_NAME || + message.origin.name === GOAL_BLOCKED_REMINDER_NAME + ); +} + +function goalFailurePauseReason(error: unknown): string { + const payload = normalizeGoalErrorPayload(error); + switch (payload.code) { + case ErrorCodes.PROVIDER_RATE_LIMIT: + return GOAL_RATE_LIMIT_PAUSE_REASON; + case ErrorCodes.PROVIDER_CONNECTION_ERROR: + return pauseReasonWithMessage(GOAL_PROVIDER_CONNECTION_PAUSE_PREFIX, payload.message); + case ErrorCodes.PROVIDER_AUTH_ERROR: + return pauseReasonWithMessage(GOAL_PROVIDER_AUTH_PAUSE_PREFIX, payload.message); + case ErrorCodes.PROVIDER_API_ERROR: + return pauseReasonWithMessage(GOAL_PROVIDER_API_PAUSE_PREFIX, payload.message); + case ErrorCodes.MODEL_NOT_CONFIGURED: + return pauseReasonWithMessage(GOAL_MODEL_CONFIG_PAUSE_PREFIX, LLM_NOT_SET_MESSAGE); + case ErrorCodes.MODEL_CONFIG_INVALID: + return pauseReasonWithMessage(GOAL_MODEL_CONFIG_PAUSE_PREFIX, payload.message); + default: + return pauseReasonWithMessage(GOAL_RUNTIME_PAUSE_PREFIX, payload.message); + } +} + +function normalizeGoalErrorPayload(error: unknown): KimiErrorPayload { + const payload = toKimiErrorPayload(error); + if (payload.code === ErrorCodes.MODEL_NOT_CONFIGURED) { + return { ...payload, message: LLM_NOT_SET_MESSAGE }; + } + return payload; +} + +function pauseReasonWithMessage(prefix: string, message: string | undefined): string { + const trimmed = message?.trim(); + return trimmed === undefined || trimmed.length === 0 ? prefix : `${prefix}: ${trimmed}`; +} + registerScopedService( LifecycleScope.Agent, IAgentGoalService, AgentGoalService, - InstantiationType.Delayed, + InstantiationType.Eager, 'goal', ); diff --git a/packages/agent-core-v2/src/agent/goal/tools/update-goal.ts b/packages/agent-core-v2/src/agent/goal/tools/update-goal.ts index 96ceebbe7..28a5d9bae 100644 --- a/packages/agent-core-v2/src/agent/goal/tools/update-goal.ts +++ b/packages/agent-core-v2/src/agent/goal/tools/update-goal.ts @@ -6,25 +6,17 @@ * * The argument is intentionally just a status enum — no reason or evidence. The * model explains itself in its own reply; the status is the machine-readable - * signal. The tool is only offered to the model while a goal exists. + * signal. */ import { z } from 'zod'; import { toInputJsonSchema } from '#/_base/tools/support/input-schema'; -import type { IAgentSystemReminderService } from '#/agent/systemReminder'; import type { BuiltinTool, ToolExecution } from '#/agent/tool'; import type { IAgentGoalService } from '#/agent/goal/goal'; -import { - buildGoalBlockedReasonPrompt, - buildGoalCompletionSummaryPrompt, -} from './outcome-prompts'; import DESCRIPTION from './update-goal.md?raw'; -const GOAL_COMPLETION_REMINDER_NAME = 'goal_completion_summary'; -const GOAL_BLOCKED_REMINDER_NAME = 'goal_blocked_reason'; - export const UpdateGoalToolInputSchema = z .object({ status: z @@ -40,10 +32,7 @@ export class UpdateGoalTool implements BuiltinTool { readonly description: string = DESCRIPTION; readonly parameters: Record = toInputJsonSchema(UpdateGoalToolInputSchema); - constructor( - private readonly goal: IAgentGoalService, - private readonly reminders: IAgentSystemReminderService, - ) {} + constructor(private readonly goal: IAgentGoalService) {} resolveExecution(args: UpdateGoalToolInput): ToolExecution { return { @@ -56,28 +45,11 @@ export class UpdateGoalTool implements BuiltinTool { return { output: 'Goal resumed.' }; } if (args.status === 'complete') { - const completed = await this.goal.markComplete({}, 'model'); - // `complete` is transient: markComplete announces then clears the - // record. Store the summary request as a system reminder, so the next - // provider request ends with a user message after the UpdateGoal tool - // result. Anthropic-compatible providers reject trailing assistant - // messages as unsupported prefill. - if (completed !== null) { - this.reminders.appendSystemReminder(buildGoalCompletionSummaryPrompt(completed), { - kind: 'system_trigger', - name: GOAL_COMPLETION_REMINDER_NAME, - }); - } + await this.goal.markComplete({}, 'model'); return { output: 'Goal marked complete.', stopTurn: true }; } if (args.status === 'blocked') { - const blocked = await this.goal.markBlocked({}, 'model'); - if (blocked !== null) { - this.reminders.appendSystemReminder(buildGoalBlockedReasonPrompt(blocked), { - kind: 'system_trigger', - name: GOAL_BLOCKED_REMINDER_NAME, - }); - } + await this.goal.markBlocked({}, 'model'); return { output: 'Goal marked blocked.', stopTurn: true }; } await this.goal.pauseGoal({}, 'model'); diff --git a/packages/agent-core-v2/test/goal/goal.test.ts b/packages/agent-core-v2/test/goal/goal.test.ts index 8915e9f4f..60cb89115 100644 --- a/packages/agent-core-v2/test/goal/goal.test.ts +++ b/packages/agent-core-v2/test/goal/goal.test.ts @@ -5,15 +5,18 @@ import { IAgentContextMemoryService } from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; import { IAgentGoalService, type AgentGoalService } from '#/agent/goal'; import { IAgentReplayBuilderService } from '#/agent/replayBuilder'; +import { IAgentTurnService, type Turn, type TurnResult } from '#/agent/turn'; import type { PersistedWireRecord, WireRecord } from '#/agent/wireRecord'; import { recordingTelemetry, type TelemetryRecord } from '../telemetry/stubs'; import { InMemoryWireRecordPersistence, + agentService, createTestAgent, telemetryServices, wireRecordPersistenceServices, type TestAgentContext, } from '../harness'; +import { stubTurn, type StubTurn } from '../turn/stubs'; type GoalServiceTestManager = IAgentGoalService & AgentGoalService; type GoalRecord = Extract; @@ -35,6 +38,30 @@ async function restoreGoalRecords( await ctx.restore(records as readonly PersistedWireRecord[]); } +function makeTurn(id: number): Turn { + return { + id, + abortController: new AbortController(), + ready: Promise.resolve(), + result: Promise.resolve({ reason: 'completed' }), + }; +} + +async function runGoalStep(turnService: IAgentTurnService, turn: Turn): Promise { + const step = { turn, continueTurn: false }; + await turnService.hooks.beforeStep.run(step); + await turnService.hooks.afterStep.run(step); + return step.continueTurn; +} + +async function endTurn( + turnService: IAgentTurnService, + turn: Turn, + result: TurnResult = { reason: 'completed' }, +): Promise { + await turnService.hooks.onEnded.run({ turn, result }); +} + describe('AgentGoalService', () => { let ctx: TestAgentContext; let context: IAgentContextMemoryService; @@ -70,350 +97,473 @@ describe('AgentGoalService', () => { } }); -describe('AgentGoalService creation', () => { - it('creates a goal and exposes it through getGoal', async () => { - const snapshot = await goals.createGoal({ objective: 'Ship feature X' }); + describe('AgentGoalService creation', () => { + it('creates a goal and exposes it through getGoal', async () => { + const snapshot = await goals.createGoal({ objective: 'Ship feature X' }); - expect(snapshot.objective).toBe('Ship feature X'); - expect(snapshot.status).toBe('active'); - expect(goals.getGoal().goal?.goalId).toBe(snapshot.goalId); - }); - - it('stores a completion criterion when provided', async () => { - const snapshot = await goals.createGoal({ - objective: 'Ship feature X', - completionCriterion: ' tests pass ', + expect(snapshot.objective).toBe('Ship feature X'); + expect(snapshot.status).toBe('active'); + expect(goals.getGoal().goal?.goalId).toBe(snapshot.goalId); }); - expect(snapshot.completionCriterion).toBe('tests pass'); - expect(goals.getGoal().goal?.completionCriterion).toBe('tests pass'); - }); + it('stores a completion criterion when provided', async () => { + const snapshot = await goals.createGoal({ + objective: 'Ship feature X', + completionCriterion: ' tests pass ', + }); - it('sets no default work caps when none is provided', async () => { - const snapshot = await goals.createGoal({ objective: 'Do work' }); - - expect(snapshot.budget.turnBudget).toBeNull(); - expect(snapshot.budget.tokenBudget).toBeNull(); - expect(snapshot.budget.wallClockBudgetMs).toBeNull(); - expect(snapshot.budget.overBudget).toBe(false); - }); - - it('rejects empty and too-long objectives', async () => { - await expect(goals.createGoal({ objective: ' ' })).rejects.toMatchObject({ - code: ErrorCodes.GOAL_OBJECTIVE_EMPTY, + expect(snapshot.completionCriterion).toBe('tests pass'); + expect(goals.getGoal().goal?.completionCriterion).toBe('tests pass'); }); - await expect(goals.createGoal({ objective: 'x'.repeat(4001) })).rejects.toMatchObject({ - code: ErrorCodes.GOAL_OBJECTIVE_TOO_LONG, + + it('sets no default work caps when none is provided', async () => { + const snapshot = await goals.createGoal({ objective: 'Do work' }); + + expect(snapshot.budget.turnBudget).toBeNull(); + expect(snapshot.budget.tokenBudget).toBeNull(); + expect(snapshot.budget.wallClockBudgetMs).toBeNull(); + expect(snapshot.budget.overBudget).toBe(false); + }); + + it('rejects empty and too-long objectives', async () => { + await expect(goals.createGoal({ objective: ' ' })).rejects.toMatchObject({ + code: ErrorCodes.GOAL_OBJECTIVE_EMPTY, + }); + await expect(goals.createGoal({ objective: 'x'.repeat(4001) })).rejects.toMatchObject({ + code: ErrorCodes.GOAL_OBJECTIVE_TOO_LONG, + }); + }); + + it('rejects duplicate active, paused, and blocked goals without replace', async () => { + await goals.createGoal({ objective: 'first' }); + await expect(goals.createGoal({ objective: 'second' })).rejects.toMatchObject({ + code: ErrorCodes.GOAL_ALREADY_EXISTS, + }); + await goals.pauseGoal(); + await expect(goals.createGoal({ objective: 'second' })).rejects.toMatchObject({ + code: ErrorCodes.GOAL_ALREADY_EXISTS, + }); + await goals.resumeGoal(); + await goals.markBlocked({ reason: 'stuck' }); + await expect(goals.createGoal({ objective: 'second' })).rejects.toMatchObject({ + code: ErrorCodes.GOAL_ALREADY_EXISTS, + }); + }); + + it('replaces an existing goal when replace is set', async () => { + const first = await goals.createGoal({ objective: 'first' }); + const second = await goals.createGoal({ objective: 'second', replace: true }); + await ctx.wireRecord.flush(); + + expect(second.goalId).not.toBe(first.goalId); + expect(goals.getGoal().goal?.objective).toBe('second'); + expect(goalRecords(records).map((record) => record.type)).toEqual([ + 'goal.create', + 'goal.clear', + 'goal.create', + ]); + }); + + it('cancels with dispatcher-style empty input', async () => { + await goals.createGoal({ objective: 'work' }); + const removed = await goals.cancelGoal({}); + expect(removed.status).toBe('active'); + expect(goals.getGoal().goal).toBeNull(); }); }); - it('rejects duplicate active, paused, and blocked goals without replace', async () => { - await goals.createGoal({ objective: 'first' }); - await expect(goals.createGoal({ objective: 'second' })).rejects.toMatchObject({ - code: ErrorCodes.GOAL_ALREADY_EXISTS, + describe('AgentGoalService lifecycle', () => { + it('emits typed lifecycle and completion changes', async () => { + await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); + expect(events.at(-1)?.change).toBeUndefined(); + + await goals.pauseGoal(); + expect(events.at(-1)?.change).toMatchObject({ kind: 'lifecycle', status: 'paused' }); + + await goals.resumeGoal(); + expect(events.at(-1)?.change).toMatchObject({ kind: 'lifecycle', status: 'active' }); + + await goals.markComplete({ reason: 'done' }, 'model'); + const completion = events.find((event) => event.change?.kind === 'completion')?.change; + expect(completion).toMatchObject({ kind: 'completion', status: 'complete', reason: 'done' }); + expect(goals.getGoal().goal).toBeNull(); + expect(events.at(-1)?.snapshot).toBeNull(); }); - await goals.pauseGoal(); - await expect(goals.createGoal({ objective: 'second' })).rejects.toMatchObject({ - code: ErrorCodes.GOAL_ALREADY_EXISTS, + + it('keeps blocked goals resumable', async () => { + await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); + const blocked = await goals.markBlocked({ reason: 'need creds' }); + expect(blocked?.status).toBe('blocked'); + expect(blocked?.terminalReason).toBe('need creds'); + + const resumed = await goals.resumeGoal(); + expect(resumed.status).toBe('active'); + expect(resumed.terminalReason).toBeUndefined(); }); - await goals.resumeGoal(); - await goals.markBlocked({ reason: 'stuck' }); - await expect(goals.createGoal({ objective: 'second' })).rejects.toMatchObject({ - code: ErrorCodes.GOAL_ALREADY_EXISTS, + + it('pauseOnInterrupt parks active goals and no-ops for stopped goals', async () => { + await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); + const paused = await goals.pauseOnInterrupt({ reason: 'Paused after interruption' }); + expect(paused?.status).toBe('paused'); + expect(paused?.terminalReason).toBe('Paused after interruption'); + + expect(await goals.pauseOnInterrupt({ reason: 'again' })).toBeNull(); + expect(goals.getGoal().goal?.status).toBe('paused'); + }); + + it('cancelGoal discards the goal and throws when missing', async () => { + await goals.createGoal({ objective: 'work' }); + const removed = await goals.cancelGoal(); + expect(removed.status).toBe('active'); + expect(goals.getGoal()).toEqual({ goal: null }); + const reminder = context.get().at(-1); + expect(reminder?.origin).toEqual({ kind: 'system_trigger', name: 'goal_cancelled' }); + expect(JSON.stringify(reminder?.content)).toContain('Ignore earlier active-goal reminders'); + await expect(goals.cancelGoal()).rejects.toMatchObject({ code: ErrorCodes.GOAL_NOT_FOUND }); }); }); - it('replaces an existing goal when replace is set', async () => { - const first = await goals.createGoal({ objective: 'first' }); - const second = await goals.createGoal({ objective: 'second', replace: true }); - await ctx.wireRecord.flush(); + describe('AgentGoalService accounting and budgets', () => { + it('counts tokens and turns only while active', async () => { + await goals.createGoal({ objective: 'work' }); + await goals.recordTokenUsage(30); + await goals.incrementTurn(); + expect(goals.getGoal().goal).toMatchObject({ tokensUsed: 30, turnsUsed: 1 }); - expect(second.goalId).not.toBe(first.goalId); - expect(goals.getGoal().goal?.objective).toBe('second'); - expect(goalRecords(records).map((record) => record.type)).toEqual([ - 'goal.create', - 'goal.clear', - 'goal.create', - ]); - }); -}); - -describe('AgentGoalService lifecycle', () => { - it('emits typed lifecycle and completion changes', async () => { - await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); - expect(events.at(-1)?.change).toBeUndefined(); - - await goals.pauseGoal(); - expect(events.at(-1)?.change).toMatchObject({ kind: 'lifecycle', status: 'paused' }); - - await goals.resumeGoal(); - expect(events.at(-1)?.change).toMatchObject({ kind: 'lifecycle', status: 'active' }); - - await goals.markComplete({ reason: 'done' }, 'model'); - const completion = events.find((event) => event.change?.kind === 'completion')?.change; - expect(completion).toMatchObject({ kind: 'completion', status: 'complete', reason: 'done' }); - expect(goals.getGoal().goal).toBeNull(); - expect(events.at(-1)?.snapshot).toBeNull(); - }); - - it('keeps blocked goals resumable', async () => { - await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); - const blocked = await goals.markBlocked({ reason: 'need creds' }); - expect(blocked?.status).toBe('blocked'); - expect(blocked?.terminalReason).toBe('need creds'); - - const resumed = await goals.resumeGoal(); - expect(resumed.status).toBe('active'); - expect(resumed.terminalReason).toBeUndefined(); - }); - - it('pauseOnInterrupt parks active goals and no-ops for stopped goals', async () => { - await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); - const paused = await goals.pauseOnInterrupt({ reason: 'Paused after interruption' }); - expect(paused?.status).toBe('paused'); - expect(paused?.terminalReason).toBe('Paused after interruption'); - - expect(await goals.pauseOnInterrupt({ reason: 'again' })).toBeNull(); - expect(goals.getGoal().goal?.status).toBe('paused'); - }); - - it('cancelGoal discards the goal and throws when missing', async () => { - await goals.createGoal({ objective: 'work' }); - const removed = await goals.cancelGoal(); - expect(removed.status).toBe('active'); - expect(goals.getGoal()).toEqual({ goal: null }); - const reminder = context.get().at(-1); - expect(reminder?.origin).toEqual({ kind: 'system_trigger', name: 'goal_cancelled' }); - expect(JSON.stringify(reminder?.content)).toContain('Ignore earlier active-goal reminders'); - await expect(goals.cancelGoal()).rejects.toMatchObject({ code: ErrorCodes.GOAL_NOT_FOUND }); - }); -}); - -describe('AgentGoalService accounting and budgets', () => { - it('counts tokens and turns only while active', async () => { - await goals.createGoal({ objective: 'work' }); - await goals.recordTokenUsage(30); - await goals.incrementTurn(); - expect(goals.getGoal().goal).toMatchObject({ tokensUsed: 30, turnsUsed: 1 }); - - await goals.pauseGoal(); - await goals.recordTokenUsage(12); - await goals.incrementTurn(); - expect(goals.getGoal().goal).toMatchObject({ tokensUsed: 30, turnsUsed: 1 }); - }); - - it('sets budget limits through SetGoalBudget-style updates', async () => { - await goals.createGoal({ objective: 'work' }); - const snapshot = await goals.setBudgetLimits({ - budgetLimits: { tokenBudget: 100, turnBudget: 2, wallClockBudgetMs: 1000 }, - }, 'model'); - - expect(snapshot.budget.tokenBudget).toBe(100); - expect(snapshot.budget.turnBudget).toBe(2); - expect(snapshot.budget.wallClockBudgetMs).toBe(1000); - }); - - it('tracks telemetry without goal text', async () => { - await goals.createGoal({ objective: 'private objective', replace: true }); - await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 100 } }, 'model'); - await goals.incrementTurn(); - await goals.pauseGoal({ reason: 'private pause reason' }); - await goals.resumeGoal(); - await goals.markComplete({ reason: 'private completion reason' }, 'model'); - - expect(telemetry.map((record) => record.event)).toEqual([ - 'goal_created', - 'goal_budget_set', - 'goal_continued', - 'goal_status_changed', - 'goal_status_changed', - 'goal_status_changed', - 'goal_cleared', - ]); - expect(telemetry[0]?.properties).toEqual({ actor: 'user', replace: true }); - expect(telemetry[1]?.properties).toMatchObject({ actor: 'model', has_token_budget: true }); - expect(telemetry[3]?.properties).toMatchObject({ status: 'paused', actor: 'user' }); - expect(JSON.stringify(telemetry)).not.toContain('private objective'); - expect(JSON.stringify(telemetry)).not.toContain('private pause reason'); - expect(JSON.stringify(telemetry)).not.toContain('private completion reason'); - }); -}); - -describe('AgentGoalService records', () => { - it('records only replay-relevant create/update/clear fields', async () => { - await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); - await goals.recordTokenUsage(5); - await goals.incrementTurn(); - await goals.setBudgetLimits({ budgetLimits: { turnBudget: 2 } }, 'model'); - await goals.markBlocked({ reason: 'stuck' }); - await goals.cancelGoal(); - await ctx.wireRecord.flush(); - - const recordsWithoutMetadata = goalRecords(records); - expect(recordsWithoutMetadata).toEqual([ - expect.objectContaining({ - type: 'goal.create', - goalId: expect.any(String), - objective: 'work', - completionCriterion: 'tests pass', - }), - expect.objectContaining({ type: 'goal.update', tokensUsed: 5 }), - expect.objectContaining({ type: 'goal.update', turnsUsed: 1 }), - expect.objectContaining({ - type: 'goal.update', - budgetLimits: { turnBudget: 2 }, - }), - expect.objectContaining({ - type: 'goal.update', - status: 'blocked', - reason: 'stuck', - actor: 'runtime', - }), - expect.objectContaining({ type: 'goal.clear' }), - ]); - expect(recordsWithoutMetadata[0]).not.toHaveProperty('actor'); - expect(recordsWithoutMetadata[0]).not.toHaveProperty('budgetLimits'); - expect(recordsWithoutMetadata[1]).not.toHaveProperty('goalId'); - expect(recordsWithoutMetadata[1]).not.toHaveProperty('status'); - expect(recordsWithoutMetadata.at(-1)).not.toHaveProperty('goalId'); - expect(recordsWithoutMetadata.at(-1)).not.toHaveProperty('reason'); - }); - - it('restores state from patch records', async () => { - await restoreGoalRecords(ctx, goals, [ - { - type: 'goal.create', - goalId: 'g1', - objective: 'work', - completionCriterion: 'tests pass', - time: Date.parse('2026-01-01T00:00:00.000Z'), - }, - { type: 'goal.update', tokensUsed: 5 }, - { type: 'goal.update', turnsUsed: 1 }, - { type: 'goal.update', budgetLimits: { turnBudget: 2 } }, - { type: 'goal.update', status: 'blocked', reason: 'stuck' }, - ]); - - expect(goals.getGoal().goal).toMatchObject({ - objective: 'work', - completionCriterion: 'tests pass', - status: 'blocked', - terminalReason: 'stuck', - tokensUsed: 5, - turnsUsed: 1, + await goals.pauseGoal(); + await goals.recordTokenUsage(12); + await goals.incrementTurn(); + expect(goals.getGoal().goal).toMatchObject({ tokensUsed: 30, turnsUsed: 1 }); + }); + + it('sets budget limits through SetGoalBudget-style updates', async () => { + await goals.createGoal({ objective: 'work' }); + const snapshot = await goals.setBudgetLimits({ + budgetLimits: { tokenBudget: 100, turnBudget: 2, wallClockBudgetMs: 1000 }, + }, 'model'); + + expect(snapshot.budget.tokenBudget).toBe(100); + expect(snapshot.budget.turnBudget).toBe(2); + expect(snapshot.budget.wallClockBudgetMs).toBe(1000); + }); + + it('tracks telemetry without goal text', async () => { + await goals.createGoal({ objective: 'private objective', replace: true }); + await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 100 } }, 'model'); + await goals.incrementTurn(); + await goals.pauseGoal({ reason: 'private pause reason' }); + await goals.resumeGoal(); + await goals.markComplete({ reason: 'private completion reason' }, 'model'); + + expect(telemetry.map((record) => record.event)).toEqual([ + 'goal_created', + 'goal_budget_set', + 'goal_continued', + 'goal_status_changed', + 'goal_status_changed', + 'goal_status_changed', + 'goal_cleared', + ]); + expect(telemetry[0]?.properties).toEqual({ actor: 'user', replace: true }); + expect(telemetry[1]?.properties).toMatchObject({ actor: 'model', has_token_budget: true }); + expect(telemetry[3]?.properties).toMatchObject({ status: 'paused', actor: 'user' }); + expect(JSON.stringify(telemetry)).not.toContain('private objective'); + expect(JSON.stringify(telemetry)).not.toContain('private pause reason'); + expect(JSON.stringify(telemetry)).not.toContain('private completion reason'); }); - expect(goals.getGoal().goal?.budget.turnBudget).toBe(2); }); - it('projects restored goal status changes into replay records', async () => { - await restoreGoalRecords(ctx, goals, [ - { - type: 'goal.create', - goalId: 'g1', - objective: 'work', - completionCriterion: 'tests pass', - time: Date.parse('2026-01-01T00:00:00.000Z'), - }, - { type: 'goal.update', tokensUsed: 5 }, - { type: 'goal.update', turnsUsed: 1 }, - { - type: 'goal.update', - status: 'paused', - reason: 'break', - actor: 'runtime', - }, - { type: 'goal.update', status: 'active', actor: 'user' }, - { - type: 'goal.update', - status: 'complete', - reason: 'done', - actor: 'model', - }, - ]); + describe('AgentGoalService records', () => { + it('records only replay-relevant create/update/clear fields', async () => { + await goals.createGoal({ objective: 'work', completionCriterion: 'tests pass' }); + await goals.recordTokenUsage(5); + await goals.incrementTurn(); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 2 } }, 'model'); + await goals.markBlocked({ reason: 'stuck' }); + await goals.cancelGoal(); + await ctx.wireRecord.flush(); - expect(replayBuilder.buildResult()).toEqual([ - expect.objectContaining({ - type: 'goal_updated', - snapshot: expect.objectContaining({ objective: 'work', status: 'active' }), - change: { kind: 'created' }, - }), - expect.objectContaining({ - type: 'goal_updated', - snapshot: expect.objectContaining({ status: 'paused', terminalReason: 'break' }), - change: { kind: 'lifecycle', status: 'paused', reason: 'break', actor: 'runtime' }, - }), - expect.objectContaining({ - type: 'goal_updated', - snapshot: expect.objectContaining({ status: 'active' }), - change: { kind: 'lifecycle', status: 'active', reason: undefined, actor: 'user' }, - }), - expect.objectContaining({ - type: 'goal_updated', - snapshot: expect.objectContaining({ - status: 'complete', - terminalReason: 'done', - turnsUsed: 1, - tokensUsed: 5, + const recordsWithoutMetadata = goalRecords(records); + expect(recordsWithoutMetadata).toEqual([ + expect.objectContaining({ + type: 'goal.create', + goalId: expect.any(String), + objective: 'work', + completionCriterion: 'tests pass', }), - change: { - kind: 'completion', + expect.objectContaining({ type: 'goal.update', tokensUsed: 5 }), + expect.objectContaining({ type: 'goal.update', turnsUsed: 1 }), + expect.objectContaining({ + type: 'goal.update', + budgetLimits: { turnBudget: 2 }, + }), + expect.objectContaining({ + type: 'goal.update', + status: 'blocked', + reason: 'stuck', + actor: 'runtime', + }), + expect.objectContaining({ type: 'goal.clear' }), + ]); + expect(recordsWithoutMetadata[0]).not.toHaveProperty('actor'); + expect(recordsWithoutMetadata[0]).not.toHaveProperty('budgetLimits'); + expect(recordsWithoutMetadata[1]).not.toHaveProperty('goalId'); + expect(recordsWithoutMetadata[1]).not.toHaveProperty('status'); + expect(recordsWithoutMetadata.at(-1)).not.toHaveProperty('goalId'); + expect(recordsWithoutMetadata.at(-1)).not.toHaveProperty('reason'); + }); + + it('restores state from patch records', async () => { + await restoreGoalRecords(ctx, goals, [ + { + type: 'goal.create', + goalId: 'g1', + objective: 'work', + completionCriterion: 'tests pass', + time: Date.parse('2026-01-01T00:00:00.000Z'), + }, + { type: 'goal.update', tokensUsed: 5 }, + { type: 'goal.update', turnsUsed: 1 }, + { type: 'goal.update', budgetLimits: { turnBudget: 2 } }, + { type: 'goal.update', status: 'blocked', reason: 'stuck' }, + ]); + + expect(goals.getGoal().goal).toMatchObject({ + objective: 'work', + completionCriterion: 'tests pass', + status: 'blocked', + terminalReason: 'stuck', + tokensUsed: 5, + turnsUsed: 1, + }); + expect(goals.getGoal().goal?.budget.turnBudget).toBe(2); + }); + + it('projects restored goal status changes into replay records', async () => { + await restoreGoalRecords(ctx, goals, [ + { + type: 'goal.create', + goalId: 'g1', + objective: 'work', + completionCriterion: 'tests pass', + time: Date.parse('2026-01-01T00:00:00.000Z'), + }, + { type: 'goal.update', tokensUsed: 5 }, + { type: 'goal.update', turnsUsed: 1 }, + { + type: 'goal.update', + status: 'paused', + reason: 'break', + actor: 'runtime', + }, + { type: 'goal.update', status: 'active', actor: 'user' }, + { + type: 'goal.update', status: 'complete', reason: 'done', - stats: { turnsUsed: 1, tokensUsed: 5, wallClockMs: 0 }, actor: 'model', }, - }), + ]); + + expect(replayBuilder.buildResult()).toEqual([ + expect.objectContaining({ + type: 'goal_updated', + snapshot: expect.objectContaining({ objective: 'work', status: 'active' }), + change: { kind: 'created' }, + }), + expect.objectContaining({ + type: 'goal_updated', + snapshot: expect.objectContaining({ status: 'paused', terminalReason: 'break' }), + change: { kind: 'lifecycle', status: 'paused', reason: 'break', actor: 'runtime' }, + }), + expect.objectContaining({ + type: 'goal_updated', + snapshot: expect.objectContaining({ status: 'active' }), + change: { kind: 'lifecycle', status: 'active', reason: undefined, actor: 'user' }, + }), + expect.objectContaining({ + type: 'goal_updated', + snapshot: expect.objectContaining({ + status: 'complete', + terminalReason: 'done', + turnsUsed: 1, + tokensUsed: 5, + }), + change: { + kind: 'completion', + status: 'complete', + reason: 'done', + stats: { turnsUsed: 1, tokensUsed: 5, wallClockMs: 0 }, + actor: 'model', + }, + }), + ]); + }); + + it('keeps resume-normalization pauses in core replay records', async () => { + await restoreGoalRecords(ctx, goals, [ + { + type: 'goal.create', + goalId: 'g1', + objective: 'work', + time: Date.parse('2026-01-01T00:00:00.000Z'), + }, + { + type: 'goal.update', + status: 'paused', + reason: 'Paused after agent resume', + }, + ]); + + expect(replayBuilder.buildResult().at(-1)).toMatchObject({ + type: 'goal_updated', + snapshot: { status: 'paused', terminalReason: 'Paused after agent resume' }, + change: { + kind: 'lifecycle', + status: 'paused', + reason: 'Paused after agent resume', + actor: undefined, + }, + }); + }); + + it('normalizes active replayed goals to paused', async () => { + records.length = 0; + await restoreGoalRecords(ctx, goals, [ + { + type: 'goal.create', + goalId: 'g1', + objective: 'resume me', + }, + ]); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'paused', + terminalReason: 'Paused after agent resume', + }); + expect(goalRecords(records)).toEqual([ + expect.objectContaining({ + type: 'goal.update', + status: 'paused', + reason: 'Paused after agent resume', + }), + ]); + }); + }); +}); + +describe('AgentGoalService core workflow hooks', () => { + let ctx: TestAgentContext | undefined; + let context: IAgentContextMemoryService; + let goals: IAgentGoalService; + let turnService: StubTurn; + let eventSink: IAgentEventSinkService; + + beforeEach(() => { + turnService = stubTurn(); + turnService.hooks.beforeStep.register('turn-before-step-event', (_ctx, next) => next()); + ctx = createTestAgent(agentService(IAgentTurnService, turnService)); + context = ctx.get(IAgentContextMemoryService); + goals = ctx.get(IAgentGoalService); + eventSink = ctx.get(IAgentEventSinkService); + }); + + afterEach(async () => { + await ctx?.dispose(); + }); + + it('counts an active goal turn and launches the next continuation', async () => { + await goals.createGoal({ objective: 'finish the task' }); + + const turn = makeTurn(1); + await turnService.hooks.onLaunched.run({ turn }); + await runGoalStep(turnService, turn); + await endTurn(turnService, turn); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'active', + turnsUsed: 1, + }); + expect(turnService.launches).toEqual([ + { kind: 'system_trigger', name: 'goal_continuation' }, + ]); + expect(context.get().at(-1)?.origin).toEqual({ + kind: 'system_trigger', + name: 'goal_continuation', + }); + expect(JSON.stringify(context.get().at(-1)?.content)).toContain('Continue working toward'); + }); + + it('continues after creating a goal mid-turn without counting the starter turn', async () => { + const turn = makeTurn(2); + await turnService.hooks.onLaunched.run({ turn }); + await runGoalStep(turnService, turn); + + await goals.createGoal({ objective: 'finish the task' }, 'model'); + await endTurn(turnService, turn); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'active', + turnsUsed: 0, + }); + expect(turnService.launches).toEqual([ + { kind: 'system_trigger', name: 'goal_continuation' }, ]); }); - it('keeps resume-normalization pauses in core replay records', async () => { - await restoreGoalRecords(ctx, goals, [ - { - type: 'goal.create', - goalId: 'g1', - objective: 'work', - time: Date.parse('2026-01-01T00:00:00.000Z'), - }, - { - type: 'goal.update', - status: 'paused', - reason: 'Paused after agent resume', - }, - ]); + it('requests one final outcome turn after model completion', async () => { + await goals.createGoal({ objective: 'finish the task' }); - expect(replayBuilder.buildResult().at(-1)).toMatchObject({ - type: 'goal_updated', - snapshot: { status: 'paused', terminalReason: 'Paused after agent resume' }, - change: { - kind: 'lifecycle', - status: 'paused', - reason: 'Paused after agent resume', - actor: undefined, - }, + const turn = makeTurn(3); + await turnService.hooks.onLaunched.run({ turn }); + const step = { turn, continueTurn: false }; + await turnService.hooks.beforeStep.run(step); + + await goals.markComplete({}, 'model'); + await turnService.hooks.afterStep.run(step); + await endTurn(turnService, turn); + + expect(step.continueTurn).toBe(true); + expect(goals.getGoal().goal).toBeNull(); + expect(turnService.launches).toEqual([]); + expect(context.get().at(-1)?.origin).toEqual({ + kind: 'system_trigger', + name: 'goal_completion_summary', }); }); - it('normalizes active replayed goals to paused', async () => { - records.length = 0; - await restoreGoalRecords(ctx, goals, [ - { - type: 'goal.create', - goalId: 'g1', - objective: 'resume me', - }, - ]); + it('pauses active goals after failed turns', async () => { + await goals.createGoal({ objective: 'finish the task' }); + + const turn = makeTurn(4); + await turnService.hooks.onLaunched.run({ turn }); + await endTurn(turnService, turn, { reason: 'failed', error: new Error('boom') }); expect(goals.getGoal().goal).toMatchObject({ status: 'paused', - terminalReason: 'Paused after agent resume', + terminalReason: 'Paused after runtime error: boom', }); - expect(goalRecords(records)).toEqual([ - expect.objectContaining({ - type: 'goal.update', - status: 'paused', - reason: 'Paused after agent resume', - }), - ]); + expect(turnService.launches).toEqual([]); + }); + + it('blocks active goals when the user prompt hook blocks the turn', async () => { + await goals.createGoal({ objective: 'finish the task' }); + + const turn = makeTurn(5); + await turnService.hooks.onLaunched.run({ turn }); + eventSink.emit({ + type: 'hook.result', + turnId: turn.id, + hookEvent: 'UserPromptSubmit', + content: 'blocked', + blocked: true, + }); + await endTurn(turnService, turn); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'blocked', + terminalReason: 'Blocked by UserPromptSubmit hook', + }); + expect(turnService.launches).toEqual([]); }); }); -}); diff --git a/packages/server-v2/src/transport/errors.ts b/packages/server-v2/src/transport/errors.ts index dd53a847c..9d0331b97 100644 --- a/packages/server-v2/src/transport/errors.ts +++ b/packages/server-v2/src/transport/errors.ts @@ -34,6 +34,12 @@ const KIMI_TO_PROTOCOL: Record = { [ErrorCodes.PROMPT_NOT_FOUND]: ErrorCode.PROMPT_NOT_FOUND, [ErrorCodes.SESSION_BUSY]: ErrorCode.SESSION_BUSY, [ErrorCodes.PROMPT_ALREADY_COMPLETED]: ErrorCode.PROMPT_ALREADY_COMPLETED, + [ErrorCodes.GOAL_ALREADY_EXISTS]: ErrorCode.GOAL_ALREADY_EXISTS, + [ErrorCodes.GOAL_NOT_FOUND]: ErrorCode.GOAL_NOT_FOUND, + [ErrorCodes.GOAL_STATUS_INVALID]: ErrorCode.GOAL_STATUS_INVALID, + [ErrorCodes.GOAL_NOT_RESUMABLE]: ErrorCode.GOAL_NOT_RESUMABLE, + [ErrorCodes.GOAL_OBJECTIVE_EMPTY]: ErrorCode.GOAL_OBJECTIVE_EMPTY, + [ErrorCodes.GOAL_OBJECTIVE_TOO_LONG]: ErrorCode.GOAL_OBJECTIVE_TOO_LONG, }; /** diff --git a/packages/server-v2/test/rpc.test.ts b/packages/server-v2/test/rpc.test.ts index ab7c82baa..0dc87d0aa 100644 --- a/packages/server-v2/test/rpc.test.ts +++ b/packages/server-v2/test/rpc.test.ts @@ -29,6 +29,22 @@ interface SessionMetaWire { archived: boolean; } +interface GoalSnapshotWire { + goalId: string; + objective: string; + completionCriterion?: string; + status: 'active' | 'paused' | 'blocked' | 'complete'; + turnsUsed: number; + tokensUsed: number; + wallClockMs: number; + budget: unknown; + terminalReason?: string; +} + +interface GoalToolResultWire { + goal: GoalSnapshotWire | null; +} + describe('server-v2 /api/v2 RPC', () => { let server: RunningServer | undefined; let home: string | undefined; @@ -224,6 +240,77 @@ describe('server-v2 /api/v2 RPC', () => { expect(body.data.isError).not.toBe(true); }); + it('controls goals through goal:* RPC', async () => { + const id = await createSession(home as string); + await createMainAgent(id); + + const created = await call( + 'POST', + `/api/v2/session/${id}/agent/main/goal:create`, + { objective: 'finish the migration' }, + ); + expect(created.body.code).toBe(0); + expect(created.body.data).toMatchObject({ + objective: 'finish the migration', + status: 'active', + }); + + const read = await call( + 'GET', + `/api/v2/session/${id}/agent/main/goal:get`, + ); + expect(read.body.code).toBe(0); + expect(read.body.data.goal).toMatchObject({ + objective: 'finish the migration', + status: 'active', + }); + + const paused = await call( + 'POST', + `/api/v2/session/${id}/agent/main/goal:pause`, + {}, + ); + expect(paused.body.data.status).toBe('paused'); + + const resumed = await call( + 'POST', + `/api/v2/session/${id}/agent/main/goal:resume`, + {}, + ); + expect(resumed.body.data.status).toBe('active'); + + const cancelled = await call( + 'POST', + `/api/v2/session/${id}/agent/main/goal:cancel`, + {}, + ); + expect(cancelled.body.code).toBe(0); + expect(cancelled.body.data.status).toBe('active'); + + const afterCancel = await call( + 'GET', + `/api/v2/session/${id}/agent/main/goal:get`, + ); + expect(afterCancel.body.data.goal).toBeNull(); + }); + + it('maps goal errors through RPC envelopes', async () => { + const id = await createSession(home as string); + await createMainAgent(id); + + await call( + 'POST', + `/api/v2/session/${id}/agent/main/goal:create`, + { objective: 'first' }, + ); + const duplicate = await call( + 'POST', + `/api/v2/session/${id}/agent/main/goal:create`, + { objective: 'second' }, + ); + expect(duplicate.body.code).toBe(40913); + }); + it('lists and installs plugins through plugins:* RPC', async () => { const pluginRoot = await mkdtemp(join(tmpdir(), 'server-v2-plugin-source-')); try { From a78c1707abd74c240baaa0dbbad5d33ab070aafe Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 1 Jul 2026 21:40:41 +0800 Subject: [PATCH 3/6] feat: enforce goal budgets --- .../src/agent/goal/goalService.ts | 68 +++++++++++++- packages/agent-core-v2/test/goal/goal.test.ts | 91 ++++++++++++++++++- .../agent-core-v2/test/goal/injection.test.ts | 7 +- 3 files changed, 159 insertions(+), 7 deletions(-) diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index 2180e2a12..bb8310e01 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -118,6 +118,7 @@ const GOAL_PROVIDER_API_PAUSE_PREFIX = 'Paused after provider API error'; const GOAL_MODEL_CONFIG_PAUSE_PREFIX = 'Paused after model configuration error'; const GOAL_RUNTIME_PAUSE_PREFIX = 'Paused after runtime error'; const GOAL_PROVIDER_FILTERED_PAUSE_REASON = 'Paused after provider safety policy block'; +const GOAL_BUDGET_BLOCK_PREFIX = 'Blocked after goal budget reached'; const LLM_NOT_SET_MESSAGE = 'LLM not set, send "/login" to login'; const GOAL_CONTINUATION_PROMPT = [ @@ -161,6 +162,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { private readonly countedGoalTurns = new Set(); private readonly goalOutcomeContinuationTurns = new Set(); private readonly promptHookBlockedTurns = new Set(); + private readonly observedUsageTotalsByTurn = new Map(); constructor( private readonly options: GoalServiceOptions = {}, @@ -239,6 +241,10 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { events.on((event) => { if (event.type === 'hook.result' && event.blocked === true) { this.promptHookBlockedTurns.add(event.turnId); + return; + } + if (event.type === 'agent.status.updated' && event.usage?.currentTurn !== undefined) { + this.handleUsageStatusUpdated(event.usage.currentTurn); } }), ); @@ -382,7 +388,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { actor, ...budgetTelemetryProperties(input.budgetLimits), }); - return this.toSnapshot(state); + return this.blockIfBudgetReached(state) ?? this.toSnapshot(state); } async cancelGoal( @@ -455,12 +461,16 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { } async recordTokenUsage(tokenDelta: number): Promise { + return this.accountTokenUsage(tokenDelta); + } + + private accountTokenUsage(tokenDelta: number): GoalSnapshot | null { const state = this.state; if (state === undefined || state.status !== 'active') return null; state.tokensUsed += Math.max(0, tokenDelta); this.persistState(state, { silent: true }); this.appendGoalUpdate({ tokensUsed: state.tokensUsed }); - return this.toSnapshot(state); + return this.blockIfBudgetReached(state) ?? this.toSnapshot(state); } async incrementTurn(): Promise { @@ -470,13 +480,14 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { this.persistState(state); this.appendGoalUpdate({ turnsUsed: state.turnsUsed }); this.telemetry.track('goal_continued', { turns_used: state.turnsUsed }); - return this.toSnapshot(state); + return this.blockIfBudgetReached(state) ?? this.toSnapshot(state); } private handleTurnLaunched(turn: Turn): void { if (this.state?.status === 'active') this.goalDrivenTurns.add(turn.id); this.goalOutcomeContinuationTurns.delete(turn.id); this.promptHookBlockedTurns.delete(turn.id); + this.observedUsageTotalsByTurn.delete(turn.id); } private async handleBeforeStep(ctx: TurnStepContext): Promise { @@ -497,6 +508,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { this.goalDrivenTurns.delete(ctx.turn.id); this.countedGoalTurns.delete(ctx.turn.id); this.goalOutcomeContinuationTurns.delete(ctx.turn.id); + this.observedUsageTotalsByTurn.delete(ctx.turn.id); const blockedByPromptHook = this.promptHookBlockedTurns.delete(ctx.turn.id); if (blockedByPromptHook) { @@ -518,10 +530,22 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { } if (this.state?.status !== 'active') return; + if (this.blockIfBudgetReached(this.state) !== null) return; if (this.turnService.getActiveTurn() !== undefined) return; this.launchContinuationTurn(); } + private handleUsageStatusUpdated(usage: TokenUsageLike): void { + const turn = this.turnService.getActiveTurn(); + if (turn === undefined || !this.goalDrivenTurns.has(turn.id)) return; + const total = tokenUsageTotal(usage); + const previous = this.observedUsageTotalsByTurn.get(turn.id) ?? 0; + this.observedUsageTotalsByTurn.set(turn.id, total); + const delta = total - previous; + if (delta <= 0) return; + this.accountTokenUsage(delta); + } + private launchContinuationTurn(): void { const message = ensureMessageId({ role: 'user', @@ -714,6 +738,26 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { terminalReason: state.terminalReason, }; } + + private blockIfBudgetReached(state: GoalState): GoalSnapshot | null { + if (state.status !== 'active') return null; + const reason = goalBudgetBlockReason(this.toSnapshot(state).budget); + if (reason === undefined) return null; + this.applyStatus(state, 'blocked'); + state.terminalReason = reason; + this.persistState(state, { + change: { kind: 'lifecycle', status: 'blocked', reason, actor: 'runtime' }, + }); + this.appendStatusUpdate(state, 'runtime', reason); + return this.toSnapshot(state); + } +} + +interface TokenUsageLike { + readonly inputCacheRead: number; + readonly inputCacheCreation: number; + readonly inputOther: number; + readonly output: number; } function liveWallClockMs(state: GoalState, now: number = Date.now()): number { @@ -752,6 +796,20 @@ function computeBudgetReport( }; } +function goalBudgetBlockReason(budget: GoalBudgetReport): string | undefined { + const reached: string[] = []; + if (budget.turnBudgetReached) { + reached.push(`turn budget ${budget.turnBudget ?? ''}`.trim()); + } + if (budget.tokenBudgetReached) { + reached.push(`token budget ${budget.tokenBudget ?? ''}`.trim()); + } + if (budget.wallClockBudgetReached) { + reached.push(`wall-clock budget ${budget.wallClockBudgetMs ?? ''}ms`.trim()); + } + return reached.length === 0 ? undefined : `${GOAL_BUDGET_BLOCK_PREFIX}: ${reached.join(', ')}`; +} + function budgetTelemetryProperties(limits: GoalBudgetLimits): TelemetryProperties { return { has_token_budget: limits.tokenBudget !== undefined, @@ -760,6 +818,10 @@ function budgetTelemetryProperties(limits: GoalBudgetLimits): TelemetryPropertie }; } +function tokenUsageTotal(usage: TokenUsageLike): number { + return usage.inputCacheRead + usage.inputCacheCreation + usage.inputOther + usage.output; +} + function normalizeCompletionCriterion(value: string | undefined): string | undefined { const trimmed = value?.trim(); return trimmed?.length ? trimmed : undefined; diff --git a/packages/agent-core-v2/test/goal/goal.test.ts b/packages/agent-core-v2/test/goal/goal.test.ts index 60cb89115..ac2d7a699 100644 --- a/packages/agent-core-v2/test/goal/goal.test.ts +++ b/packages/agent-core-v2/test/goal/goal.test.ts @@ -247,6 +247,41 @@ describe('AgentGoalService', () => { expect(snapshot.budget.wallClockBudgetMs).toBe(1000); }); + it('blocks when a token budget is reached', async () => { + await goals.createGoal({ objective: 'work' }); + await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 10 } }, 'model'); + + const snapshot = await goals.recordTokenUsage(10); + + expect(snapshot).toMatchObject({ + status: 'blocked', + tokensUsed: 10, + terminalReason: 'Blocked after goal budget reached: token budget 10', + }); + expect(goals.getGoal().goal).toMatchObject({ + status: 'blocked', + budget: { + tokenBudgetReached: true, + overBudget: true, + }, + }); + }); + + it('blocks when a newly set budget is already exhausted', async () => { + await goals.createGoal({ objective: 'work' }); + await goals.incrementTurn(); + + const snapshot = await goals.setBudgetLimits( + { budgetLimits: { turnBudget: 1 } }, + 'model', + ); + + expect(snapshot).toMatchObject({ + status: 'blocked', + terminalReason: 'Blocked after goal budget reached: turn budget 1', + }); + }); + it('tracks telemetry without goal text', async () => { await goals.createGoal({ objective: 'private objective', replace: true }); await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 100 } }, 'model'); @@ -460,7 +495,7 @@ describe('AgentGoalService core workflow hooks', () => { let eventSink: IAgentEventSinkService; beforeEach(() => { - turnService = stubTurn(); + turnService = stubTurn({ hasActiveTurn: true }); turnService.hooks.beforeStep.register('turn-before-step-event', (_ctx, next) => next()); ctx = createTestAgent(agentService(IAgentTurnService, turnService)); context = ctx.get(IAgentContextMemoryService); @@ -494,6 +529,60 @@ describe('AgentGoalService core workflow hooks', () => { expect(JSON.stringify(context.get().at(-1)?.content)).toContain('Continue working toward'); }); + it('blocks at the turn budget instead of launching a continuation', async () => { + await goals.createGoal({ objective: 'finish the task' }); + await goals.setBudgetLimits({ budgetLimits: { turnBudget: 1 } }, 'model'); + + const turn = makeTurn(11); + await turnService.hooks.onLaunched.run({ turn }); + await runGoalStep(turnService, turn); + await endTurn(turnService, turn); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'blocked', + turnsUsed: 1, + terminalReason: 'Blocked after goal budget reached: turn budget 1', + }); + expect(turnService.launches).toEqual([]); + }); + + it('accounts live usage status deltas for active goal turns', async () => { + await goals.createGoal({ objective: 'finish the task' }); + await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 7 } }, 'model'); + + const turn = turnService.launch({ kind: 'user' }); + await turnService.hooks.onLaunched.run({ turn }); + + eventSink.emit({ + type: 'agent.status.updated', + usage: { + currentTurn: { + inputCacheRead: 0, + inputCacheCreation: 0, + inputOther: 4, + output: 0, + }, + }, + }); + eventSink.emit({ + type: 'agent.status.updated', + usage: { + currentTurn: { + inputCacheRead: 0, + inputCacheCreation: 0, + inputOther: 4, + output: 3, + }, + }, + }); + + expect(goals.getGoal().goal).toMatchObject({ + status: 'blocked', + tokensUsed: 7, + terminalReason: 'Blocked after goal budget reached: token budget 7', + }); + }); + it('continues after creating a goal mid-turn without counting the starter turn', async () => { const turn = makeTurn(2); await turnService.hooks.onLaunched.run({ turn }); diff --git a/packages/agent-core-v2/test/goal/injection.test.ts b/packages/agent-core-v2/test/goal/injection.test.ts index 4321da8fd..715365443 100644 --- a/packages/agent-core-v2/test/goal/injection.test.ts +++ b/packages/agent-core-v2/test/goal/injection.test.ts @@ -174,15 +174,16 @@ describe('GoalInjection content', () => { expect(text).toContain('avoid starting new discretionary work'); }); - it('has no separate over-budget guidance (the runtime auto-blocks instead)', async () => { + it('shows a blocked note once a budget is reached', async () => { const text = (await readGoalReminder(async (goals) => { await goals.createGoal({ objective: 'work' }); await goals.setBudgetLimits({ budgetLimits: { turnBudget: 2 } }, 'model'); await goals.incrementTurn(); await goals.incrementTurn(); // 2/2 = 100% }))!; - expect(text).not.toContain('report the best terminal state'); - expect(text).toContain('nearing a budget'); + expect(text).toContain('currently blocked'); + expect(text).toContain('Blocked after goal budget reached: turn budget 2'); + expect(text).not.toContain('Budget guidance'); }); it('tells the model to call UpdateGoal to finish', async () => { From 50d9660f56c2dff9c3f1341a906763ad760e5762 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 1 Jul 2026 22:02:35 +0800 Subject: [PATCH 4/6] chore: rename tool dedupe domain directory --- packages/agent-core-v2/scripts/check-domain-layers.mjs | 2 +- .../src/agent/{tooldedup => toolDedupe}/index.ts | 0 .../src/agent/{tooldedup => toolDedupe}/toolDedupe.ts | 0 .../src/agent/{tooldedup => toolDedupe}/toolDedupeService.ts | 0 packages/agent-core-v2/src/index.ts | 2 +- packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts | 4 ++-- 6 files changed, 4 insertions(+), 4 deletions(-) rename packages/agent-core-v2/src/agent/{tooldedup => toolDedupe}/index.ts (100%) rename packages/agent-core-v2/src/agent/{tooldedup => toolDedupe}/toolDedupe.ts (100%) rename packages/agent-core-v2/src/agent/{tooldedup => toolDedupe}/toolDedupeService.ts (100%) diff --git a/packages/agent-core-v2/scripts/check-domain-layers.mjs b/packages/agent-core-v2/scripts/check-domain-layers.mjs index 83a1389e6..b63bd6e63 100644 --- a/packages/agent-core-v2/scripts/check-domain-layers.mjs +++ b/packages/agent-core-v2/scripts/check-domain-layers.mjs @@ -102,7 +102,7 @@ const DOMAIN_LAYER = new Map([ ['swarm', 4], ['scopeContext', 4], ['usage', 4], - ['tooldedup', 4], + ['toolDedupe', 4], ['contextMemory', 4], ['contextInjector', 4], ['systemReminder', 4], diff --git a/packages/agent-core-v2/src/agent/tooldedup/index.ts b/packages/agent-core-v2/src/agent/toolDedupe/index.ts similarity index 100% rename from packages/agent-core-v2/src/agent/tooldedup/index.ts rename to packages/agent-core-v2/src/agent/toolDedupe/index.ts diff --git a/packages/agent-core-v2/src/agent/tooldedup/toolDedupe.ts b/packages/agent-core-v2/src/agent/toolDedupe/toolDedupe.ts similarity index 100% rename from packages/agent-core-v2/src/agent/tooldedup/toolDedupe.ts rename to packages/agent-core-v2/src/agent/toolDedupe/toolDedupe.ts diff --git a/packages/agent-core-v2/src/agent/tooldedup/toolDedupeService.ts b/packages/agent-core-v2/src/agent/toolDedupe/toolDedupeService.ts similarity index 100% rename from packages/agent-core-v2/src/agent/tooldedup/toolDedupeService.ts rename to packages/agent-core-v2/src/agent/toolDedupe/toolDedupeService.ts diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index b27edeb4f..ed91ddfc0 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -36,7 +36,7 @@ export * from '#/agent/plan'; export * from '#/agent/goal'; export * from '#/agent/swarm'; export * from '#/agent/usage'; -export * from '#/agent/tooldedup'; +export * from '#/agent/toolDedupe'; export * from '#/agent/background'; import '#/agent/cron'; diff --git a/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts b/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts index 2474cd29b..68f6806fa 100644 --- a/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts +++ b/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts @@ -7,8 +7,8 @@ import { IAgentToolDedupeService, AgentToolDedupeService, __testing as toolDedupTesting, -} from '#/agent/tooldedup'; -import type { ToolDedupResult } from '#/agent/tooldedup'; +} from '#/agent/toolDedupe'; +import type { ToolDedupResult } from '#/agent/toolDedupe'; import { IAgentToolExecutorService } from '#/agent/toolExecutor'; import { IAgentTurnService } from '#/agent/turn'; import { recordingTelemetry, type TelemetryRecord } from '../telemetry/stubs'; From 44834808f69332576b309bc917842a6a04a466ba Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 1 Jul 2026 22:15:22 +0800 Subject: [PATCH 5/6] fix(agent-core-v2): fix tests --- .../src/agent/agentTool/agentToolService.ts | 8 +- .../src/agent/goal/goalService.ts | 43 ++--- packages/agent-core-v2/src/agent/loop/loop.ts | 9 +- .../src/agent/loop/loopService.ts | 29 ++-- .../agent-core-v2/src/agent/rpc/rpcService.ts | 8 +- .../agent-core-v2/src/agent/swarm/index.ts | 1 + packages/agent-core-v2/src/agent/turn/turn.ts | 11 ++ .../src/agent/turn/turnService.ts | 3 + packages/agent-core-v2/src/index.ts | 3 +- .../test/agentTool/agentToolService.test.ts | 4 +- .../test/background/rpc-events.test.ts | 48 ++++-- .../test/fullCompaction/full.test.ts | 47 +++--- packages/agent-core-v2/test/goal/goal.test.ts | 70 +++++--- packages/agent-core-v2/test/harness/agent.ts | 150 ++++++++++-------- packages/agent-core-v2/test/harness/index.ts | 2 + .../promptLegacy/promptLegacyService.test.ts | 1 + .../test/tools/fixtures/fake-exec.ts | 37 +++++ packages/agent-core-v2/test/turn/stubs.ts | 1 + packages/agent-core-v2/test/turn/turn.test.ts | 4 +- 19 files changed, 307 insertions(+), 172 deletions(-) create mode 100644 packages/agent-core-v2/test/tools/fixtures/fake-exec.ts diff --git a/packages/agent-core-v2/src/agent/agentTool/agentToolService.ts b/packages/agent-core-v2/src/agent/agentTool/agentToolService.ts index dae20a419..87e9b1d4b 100644 --- a/packages/agent-core-v2/src/agent/agentTool/agentToolService.ts +++ b/packages/agent-core-v2/src/agent/agentTool/agentToolService.ts @@ -9,7 +9,7 @@ * identity through `scopeContext`, creates child agents through * `agent-lifecycle`, reads the parent check through `session-metadata`, gates * background execution through the agent `profile`, and gathers git context - * through `execContext` (cwd) + `process` (runner). + * through `kaos` (cwd) + `process` (runner). */ import { Disposable } from '#/_base/di'; @@ -19,7 +19,7 @@ import { IAgentBackgroundService } from '#/agent/background'; import { IAgentProfileService } from '#/agent/profile'; import { IAgentScopeContext } from '#/agent/scopeContext'; import { IAgentToolRegistryService } from '#/agent/toolRegistry'; -import { IExecContext } from '#/session/execContext'; +import { IKaos } from '#/app/kaos'; import { ILogService } from '#/app/log'; import { IAgentLifecycleService } from '#/session/agent-lifecycle'; import { ISessionProcessRunner } from '#/session/process'; @@ -40,7 +40,7 @@ export class AgentToolService extends Disposable implements IAgentToolService { @IAgentToolRegistryService toolRegistry: IAgentToolRegistryService, @IAgentBackgroundService background: IAgentBackgroundService, @IAgentProfileService profile: IAgentProfileService, - @IExecContext execCtx: IExecContext, + @IKaos kaos: IKaos, @ISessionProcessRunner processRunner: ISessionProcessRunner, @ILogService log?: ILogService, ) { @@ -53,7 +53,7 @@ export class AgentToolService extends Disposable implements IAgentToolService { metadata, background, profile, - cwd: execCtx.cwd, + cwd: kaos.cwd, processRunner, log, runOverride: runner, diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index bb8310e01..a6f376c35 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -10,6 +10,7 @@ */ import { randomUUID } from 'node:crypto'; +import type { TokenUsage } from '@moonshot-ai/kosong'; import { InstantiationType } from '#/_base/di/extensions'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; @@ -36,6 +37,7 @@ import { type Turn, type TurnEndedContext, type TurnStepContext, + type TurnStepUsageContext, } from '#/agent/turn'; import type { TelemetryProperties } from '#/app/telemetry'; import { ITelemetryService } from '#/app/telemetry'; @@ -162,7 +164,6 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { private readonly countedGoalTurns = new Set(); private readonly goalOutcomeContinuationTurns = new Set(); private readonly promptHookBlockedTurns = new Set(); - private readonly observedUsageTotalsByTurn = new Map(); constructor( private readonly options: GoalServiceOptions = {}, @@ -225,6 +226,12 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { await next(); }), ); + this._register( + turnService.hooks.onStepUsage.register('goal-record-step-usage', async (ctx, next) => { + this.handleStepUsage(ctx); + await next(); + }), + ); this._register( turnService.hooks.afterStep.register('goal-outcome-continuation', async (ctx, next) => { await next(); @@ -241,10 +248,6 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { events.on((event) => { if (event.type === 'hook.result' && event.blocked === true) { this.promptHookBlockedTurns.add(event.turnId); - return; - } - if (event.type === 'agent.status.updated' && event.usage?.currentTurn !== undefined) { - this.handleUsageStatusUpdated(event.usage.currentTurn); } }), ); @@ -487,7 +490,6 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { if (this.state?.status === 'active') this.goalDrivenTurns.add(turn.id); this.goalOutcomeContinuationTurns.delete(turn.id); this.promptHookBlockedTurns.delete(turn.id); - this.observedUsageTotalsByTurn.delete(turn.id); } private async handleBeforeStep(ctx: TurnStepContext): Promise { @@ -497,6 +499,14 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { await this.incrementTurn(); } + private handleStepUsage(ctx: TurnStepUsageContext): void { + if (!this.goalDrivenTurns.has(ctx.turn.id)) return; + const snapshot = this.accountTokenUsage(tokenUsageTotal(ctx.usage)); + if (snapshot?.budget.overBudget === true) { + ctx.stopTurn = true; + } + } + private handleAfterStep(ctx: TurnStepContext): void { if (this.goalOutcomeContinuationTurns.has(ctx.turn.id)) return; if (!isGoalOutcomeReminder(this.context.get().at(-1))) return; @@ -508,7 +518,6 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { this.goalDrivenTurns.delete(ctx.turn.id); this.countedGoalTurns.delete(ctx.turn.id); this.goalOutcomeContinuationTurns.delete(ctx.turn.id); - this.observedUsageTotalsByTurn.delete(ctx.turn.id); const blockedByPromptHook = this.promptHookBlockedTurns.delete(ctx.turn.id); if (blockedByPromptHook) { @@ -535,17 +544,6 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { this.launchContinuationTurn(); } - private handleUsageStatusUpdated(usage: TokenUsageLike): void { - const turn = this.turnService.getActiveTurn(); - if (turn === undefined || !this.goalDrivenTurns.has(turn.id)) return; - const total = tokenUsageTotal(usage); - const previous = this.observedUsageTotalsByTurn.get(turn.id) ?? 0; - this.observedUsageTotalsByTurn.set(turn.id, total); - const delta = total - previous; - if (delta <= 0) return; - this.accountTokenUsage(delta); - } - private launchContinuationTurn(): void { const message = ensureMessageId({ role: 'user', @@ -753,13 +751,6 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { } } -interface TokenUsageLike { - readonly inputCacheRead: number; - readonly inputCacheCreation: number; - readonly inputOther: number; - readonly output: number; -} - function liveWallClockMs(state: GoalState, now: number = Date.now()): number { if (state.status === 'active' && state.wallClockResumedAt !== undefined) { return state.wallClockMs + Math.max(0, now - state.wallClockResumedAt); @@ -818,7 +809,7 @@ function budgetTelemetryProperties(limits: GoalBudgetLimits): TelemetryPropertie }; } -function tokenUsageTotal(usage: TokenUsageLike): number { +function tokenUsageTotal(usage: TokenUsage): number { return usage.inputCacheRead + usage.inputCacheCreation + usage.inputOther + usage.output; } diff --git a/packages/agent-core-v2/src/agent/loop/loop.ts b/packages/agent-core-v2/src/agent/loop/loop.ts index c760cf490..f69fa1d01 100644 --- a/packages/agent-core-v2/src/agent/loop/loop.ts +++ b/packages/agent-core-v2/src/agent/loop/loop.ts @@ -1,9 +1,16 @@ import { createDecorator } from "#/_base/di"; import type { HookSlot } from '#/hooks'; -import type { Turn, TurnContextOverflowContext, TurnResult, TurnStepContext } from '#/agent/turn'; +import type { + Turn, + TurnContextOverflowContext, + TurnResult, + TurnStepContext, + TurnStepUsageContext, +} from '#/agent/turn'; export interface LoopRunHooks { readonly beforeStep: HookSlot; + readonly onStepUsage: HookSlot; readonly afterStep: HookSlot; readonly onContextOverflow: HookSlot; } diff --git a/packages/agent-core-v2/src/agent/loop/loopService.ts b/packages/agent-core-v2/src/agent/loop/loopService.ts index 10e736a12..21f1772fa 100644 --- a/packages/agent-core-v2/src/agent/loop/loopService.ts +++ b/packages/agent-core-v2/src/agent/loop/loopService.ts @@ -151,17 +151,28 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { toolExecutor: this.toolExecutor, maxSteps: this.config.get(LOOP_CONTROL_SECTION)?.maxStepsPerTurn, maxRetryAttempts: this.config.get(LOOP_CONTROL_SECTION)?.maxRetriesPerStep, - recordStepUsage: (usage, context) => { + recordStepUsage: async (usage, context) => { const tokens = tokenUsageTotal(usage); - if (tokens <= 0) return; - if (context.toolCallCount > 0) { - this.pendingMeasurements.set(context.stepUuid, { - tokens, - remainingToolCalls: context.toolCallCount, - }); - } else { - this.contextSize.measured(this.measurementLength(context.stepUuid), tokens); + if (tokens > 0) { + if (context.toolCallCount > 0) { + this.pendingMeasurements.set(context.stepUuid, { + tokens, + remainingToolCalls: context.toolCallCount, + }); + } else { + this.contextSize.measured(this.measurementLength(context.stepUuid), tokens); + } } + const usageContext = { + turn, + usage, + stepNumber: context.stepNumber, + stepUuid: context.stepUuid, + toolCallCount: context.toolCallCount, + stopTurn: false, + }; + await hooks?.onStepUsage.run(usageContext); + return usageContext.stopTurn ? { stopTurn: true } : undefined; }, }); if (result.stopReason === 'aborted') { diff --git a/packages/agent-core-v2/src/agent/rpc/rpcService.ts b/packages/agent-core-v2/src/agent/rpc/rpcService.ts index 66573f492..8990d65e6 100644 --- a/packages/agent-core-v2/src/agent/rpc/rpcService.ts +++ b/packages/agent-core-v2/src/agent/rpc/rpcService.ts @@ -14,6 +14,7 @@ import { userCancellationReason } from '#/_base/utils/abort'; import { IAgentPermissionGate } from '#/agent/permissionGate'; import { IAgentPermissionModeService } from '#/agent/permissionMode/permissionMode'; import { IAgentPlanService } from '#/agent/plan'; +import { IKaos } from '#/app/kaos'; import { expandCommandArguments, IPluginService } from '#/app/plugin'; import { IAgentProfileService } from '#/agent/profile'; import { IAgentPromptService } from '#/agent/prompt'; @@ -21,8 +22,6 @@ import { IAgentQuestionToolsService } from '#/agent/questionTools'; import { ISessionMetadata, type SessionMetaPatch } from '#/session/session-metadata'; import { BashTool, IAgentShellToolsService } from '#/agent/shellTools'; import { IAgentSkillService } from '#/agent/skill'; -import { IHostEnvironment } from '#/app/hostEnvironment'; -import { IExecContext } from '#/session/execContext'; import { ISessionProcessRunner } from '#/session/process'; import { IAgentToolService } from '#/agent/agentTool'; import { @@ -109,8 +108,7 @@ export class AgentRPCService implements IAgentRPCService { @IAgentFileToolsService private readonly fileTools: IAgentFileToolsService, @IAgentShellToolsService private readonly shellTools: IAgentShellToolsService, @ISessionProcessRunner private readonly processRunner: ISessionProcessRunner, - @IHostEnvironment private readonly env: IHostEnvironment, - @IExecContext private readonly ctx: IExecContext, + @IKaos private readonly kaos: IKaos, @IAgentBackgroundService private readonly background: IAgentBackgroundService, @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, @IAgentContextSizeService private readonly contextSize: IAgentContextSizeService, @@ -139,7 +137,7 @@ export class AgentRPCService implements IAgentRPCService { private ensureBashTool() { const existing = this.toolRegistry.resolve('Bash'); if (existing !== undefined) return existing; - const bash = new BashTool(this.processRunner, this.env, this.ctx, this.background); + const bash = new BashTool(this.processRunner, this.kaos, this.background); this.toolRegistry.register(bash); return bash; } diff --git a/packages/agent-core-v2/src/agent/swarm/index.ts b/packages/agent-core-v2/src/agent/swarm/index.ts index 831d46c42..e5ecb2dbd 100644 --- a/packages/agent-core-v2/src/agent/swarm/index.ts +++ b/packages/agent-core-v2/src/agent/swarm/index.ts @@ -6,3 +6,4 @@ export * from './swarm'; export * from './swarmService'; +export * from './subagentBatch'; diff --git a/packages/agent-core-v2/src/agent/turn/turn.ts b/packages/agent-core-v2/src/agent/turn/turn.ts index 04ba16d29..430224998 100644 --- a/packages/agent-core-v2/src/agent/turn/turn.ts +++ b/packages/agent-core-v2/src/agent/turn/turn.ts @@ -1,4 +1,5 @@ import { createDecorator } from "#/_base/di"; +import type { TokenUsage } from '@moonshot-ai/kosong'; import type { ContextMessage, PromptOrigin } from '#/agent/contextMemory'; import type { Hooks } from '#/hooks'; @@ -21,6 +22,15 @@ export interface TurnStepContext { continueTurn: boolean; } +export interface TurnStepUsageContext { + readonly turn: Turn; + readonly usage: TokenUsage; + readonly stepNumber: number; + readonly stepUuid: string; + readonly toolCallCount: number; + stopTurn: boolean; +} + export interface TurnContextOverflowContext { readonly turn: Turn; readonly error: unknown; @@ -54,6 +64,7 @@ export interface IAgentTurnService { onLaunched: { turn: Turn }; onEnded: TurnEndedContext; beforeStep: TurnStepContext; + onStepUsage: TurnStepUsageContext; afterStep: TurnStepContext; onContextOverflow: TurnContextOverflowContext; }>; diff --git a/packages/agent-core-v2/src/agent/turn/turnService.ts b/packages/agent-core-v2/src/agent/turn/turnService.ts index 526055c91..3840d3775 100644 --- a/packages/agent-core-v2/src/agent/turn/turnService.ts +++ b/packages/agent-core-v2/src/agent/turn/turnService.ts @@ -16,6 +16,7 @@ import type { TurnEndedContext, TurnResult, TurnStepContext, + TurnStepUsageContext, } from './turn'; import { IAgentTurnService } from './turn'; @@ -44,6 +45,7 @@ export class AgentTurnService implements IAgentTurnService { onLaunched: new OrderedHookSlot<{ turn: Turn }>(), onEnded: new OrderedHookSlot(), beforeStep: new OrderedHookSlot(), + onStepUsage: new OrderedHookSlot(), afterStep: new OrderedHookSlot(), onContextOverflow: new OrderedHookSlot(), }; @@ -133,6 +135,7 @@ export class AgentTurnService implements IAgentTurnService { } result = await this.loop.runTurn(turn, { beforeStep: this.hooks.beforeStep, + onStepUsage: this.hooks.onStepUsage, afterStep: this.hooks.afterStep, onContextOverflow: this.hooks.onContextOverflow, }); diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index ed91ddfc0..346690143 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -59,8 +59,7 @@ export * from '#/app/gateway'; export * from '#/session/workspaceContext'; export * from '#/app/workspaceRegistry'; export * from '#/app/hostFolderBrowser'; -export * from '#/app/hostEnvironment'; -export * from '#/session/execContext'; +export * from '#/app/kaos'; export * from '#/session/agentFs'; export * from '#/session/process'; export * from '#/session/terminal'; diff --git a/packages/agent-core-v2/test/agentTool/agentToolService.test.ts b/packages/agent-core-v2/test/agentTool/agentToolService.test.ts index 78425c8a8..140d289f6 100644 --- a/packages/agent-core-v2/test/agentTool/agentToolService.test.ts +++ b/packages/agent-core-v2/test/agentTool/agentToolService.test.ts @@ -5,12 +5,12 @@ import { DisposableStore } from '#/_base/di/lifecycle'; import { TestInstantiationService } from '#/_base/di/test'; import { IAgentBackgroundService } from '#/agent/background'; import { IAgentLifecycleService } from '#/session/agent-lifecycle'; +import { IKaos } from '#/app/kaos'; import { ILogService } from '#/app/log'; import { IAgentProfileService } from '#/agent/profile'; import { IAgentScopeContext } from '#/agent/scopeContext'; import { AgentToolService, IAgentToolService } from '#/agent/agentTool'; import { IAgentToolRegistryService } from '#/agent/toolRegistry'; -import { IExecContext, createExecContext } from '#/session/execContext'; import { ISessionMetadata } from '#/session/session-metadata'; import { ISessionProcessRunner } from '#/session/process'; @@ -35,7 +35,7 @@ describe('AgentToolService DI wiring', () => { ix.stub(IAgentToolRegistryService, { register }); ix.stub(IAgentBackgroundService, {}); ix.stub(IAgentProfileService, { isToolActive: vi.fn().mockReturnValue(false) }); - ix.stub(IExecContext, createExecContext('/repo')); + ix.stub(IKaos, { cwd: '/repo' }); ix.stub(ISessionProcessRunner, { exec: vi.fn() }); ix.stub(ILogService, { warn: vi.fn(), info: vi.fn(), debug: vi.fn(), error: vi.fn() }); ix.set(IAgentToolService, new SyncDescriptor(AgentToolService, [undefined])); diff --git a/packages/agent-core-v2/test/background/rpc-events.test.ts b/packages/agent-core-v2/test/background/rpc-events.test.ts index 1065211d9..60cd26ba4 100644 --- a/packages/agent-core-v2/test/background/rpc-events.test.ts +++ b/packages/agent-core-v2/test/background/rpc-events.test.ts @@ -23,6 +23,7 @@ import type { HookEngine } from '#/agent/externalHooks/engine'; import { IAgentPromptService } from '#/agent/prompt'; import type { SubagentDetachHandle } from '#/agent/background'; import type { SubagentHandle } from '#/agent/agentTool'; +import { ISessionMetadata } from '#/session/session-metadata'; import { configServices, createTestAgent, @@ -238,6 +239,17 @@ function createBackgroundManager(options: { }; } +async function cleanupSessionDir( + sessionDir: string, + fixture?: BackgroundServiceFixture, +): Promise { + if (fixture !== undefined) { + await fixture.ctx.get(ISessionMetadata).ready; + await fixture.ctx.dispose(); + } + await rm(sessionDir, { recursive: true, force: true }); +} + function firstAppendedContextMessage(agent: FakeBackgroundAgent): TestContextMessage { const call = agent.context.appendUserMessage.mock.calls[0] as unknown as [ number, @@ -367,6 +379,7 @@ describe('BackgroundManager — event emission', () => { it('emits background.task.terminated when a restored task is marked lost', async () => { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-agent-reconcile-')); + let fixture: BackgroundServiceFixture | undefined; try { const persistence = createBackgroundTaskPersistence(sessionDir); await persistence.writeTask( @@ -379,7 +392,8 @@ describe('BackgroundManager — event emission', () => { status: 'running', }), ); - const { agent, manager } = createBackgroundManager({ sessionDir }); + fixture = createBackgroundManager({ sessionDir }); + const { agent, manager } = fixture; await manager.loadFromDisk(); await manager.reconcile(); @@ -392,7 +406,7 @@ describe('BackgroundManager — event emission', () => { }), }); } finally { - await rm(sessionDir, { recursive: true, force: true }); + await cleanupSessionDir(sessionDir, fixture); } }); }); @@ -477,11 +491,13 @@ describe('BackgroundManager — notification delivery', () => { it('replays restored terminal agent task notifications when undelivered', async () => { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-agent-replay-')); + let fixture: BackgroundServiceFixture | undefined; try { const persistence = createBackgroundTaskPersistence(sessionDir); await persistence.writeTask(persistedAgent()); await persistence.appendTaskOutput('agent-done0000', 'restored subagent summary'); - const { agent, manager } = createBackgroundManager({ sessionDir }); + fixture = createBackgroundManager({ sessionDir }); + const { agent, manager } = fixture; await manager.loadFromDisk(); await manager.reconcile(); @@ -503,17 +519,19 @@ describe('BackgroundManager — notification delivery', () => { expect(text).toContain(' { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-bash-replay-')); + let fixture: BackgroundServiceFixture | undefined; try { const persistence = createBackgroundTaskPersistence(sessionDir); await persistence.writeTask(persistedProcess()); await persistence.appendTaskOutput('bash-done0000', 'restored shell output'); - const { agent, manager } = createBackgroundManager({ sessionDir }); + fixture = createBackgroundManager({ sessionDir }); + const { agent, manager } = fixture; await manager.loadFromDisk(); await manager.reconcile(); @@ -535,19 +553,21 @@ describe('BackgroundManager — notification delivery', () => { expect(text).toContain(' { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-bash-tail-')); + let fixture: BackgroundServiceFixture | undefined; try { const taskId = 'bash-large000'; const largeOutput = `early-output-marker\n${'x'.repeat(8_000)}\nfinal output line`; const persistence = createBackgroundTaskPersistence(sessionDir); await persistence.writeTask(persistedProcess({ taskId })); await persistence.appendTaskOutput(taskId, largeOutput); - const { agent, manager } = createBackgroundManager({ sessionDir }); + fixture = createBackgroundManager({ sessionDir }); + const { agent, manager } = fixture; await manager.loadFromDisk(); await manager.reconcile(); @@ -562,12 +582,13 @@ describe('BackgroundManager — notification delivery', () => { expect(text).not.toContain('final output line'); expect(text).not.toContain('early-output-marker'); } finally { - await rm(sessionDir, { recursive: true, force: true }); + await cleanupSessionDir(sessionDir, fixture); } }); it('does not replay restored notifications already marked delivered', async () => { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-agent-replay-')); + let fixture: BackgroundServiceFixture | undefined; try { const origin = { kind: 'background_task', @@ -578,7 +599,8 @@ describe('BackgroundManager — notification delivery', () => { const persistence = createBackgroundTaskPersistence(sessionDir); await persistence.writeTask(persistedAgent({ taskId: 'agent-seen0000' })); await persistence.appendTaskOutput('agent-seen0000', 'already delivered summary'); - const { agent, ctx, manager } = createBackgroundManager({ sessionDir }); + fixture = createBackgroundManager({ sessionDir }); + const { agent, ctx, manager } = fixture; const context = ctx.get(IAgentContextMemoryService); context.splice(context.get().length, 0, [ { @@ -597,12 +619,13 @@ describe('BackgroundManager — notification delivery', () => { expect(agent.turn.steer).not.toHaveBeenCalled(); expect(agent.context.appendUserMessage).not.toHaveBeenCalled(); } finally { - await rm(sessionDir, { recursive: true, force: true }); + await cleanupSessionDir(sessionDir, fixture); } }); it('does not double-notify newly lost restored agent tasks', async () => { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-agent-lost-')); + let fixture: BackgroundServiceFixture | undefined; try { const persistence = createBackgroundTaskPersistence(sessionDir); await persistence.writeTask( @@ -613,7 +636,8 @@ describe('BackgroundManager — notification delivery', () => { status: 'running', }), ); - const { agent, manager } = createBackgroundManager({ sessionDir }); + fixture = createBackgroundManager({ sessionDir }); + const { agent, manager } = fixture; await manager.loadFromDisk(); await manager.reconcile(); @@ -634,7 +658,7 @@ describe('BackgroundManager — notification delivery', () => { 'Background agent lost', ); } finally { - await rm(sessionDir, { recursive: true, force: true }); + await cleanupSessionDir(sessionDir, fixture); } }); diff --git a/packages/agent-core-v2/test/fullCompaction/full.test.ts b/packages/agent-core-v2/test/fullCompaction/full.test.ts index c1313e38a..15903dfd5 100644 --- a/packages/agent-core-v2/test/fullCompaction/full.test.ts +++ b/packages/agent-core-v2/test/fullCompaction/full.test.ts @@ -869,35 +869,43 @@ describe('FullCompaction', () => { it('names truncated compaction responses when retries are exhausted', async () => { vi.useFakeTimers(); + const firstAttemptFinished = deferred(); let attempts = 0; const generate: GenerateFn = async () => { attempts += 1; + if (attempts === 1) { + firstAttemptFinished.resolve(); + } return { ...textResult('Partial summary.'), finishReason: 'truncated', rawFinishReason: 'length', }; }; - const ctx = testAgent({ generate, fullCompaction: { compactionStrategy: alwaysCompactOnce } }); - ctx.configure(); + const ctx = testAgent({ generate }); + ctx.configure({ + provider: CATALOGUED_PROVIDER, + modelCapabilities: CATALOGUED_MODEL_CAPABILITIES, + }); + ctx.appendExchange(1, 'old user one', 'old assistant one', 20); + ctx.appendExchange(2, 'recent user two', 'recent assistant two', 80); + const failed = ctx.once('error'); - await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Trigger truncated auto compaction' }] }); + await ctx.rpc.beginCompaction({}); + await firstAttemptFinished.promise; await vi.advanceTimersByTimeAsync(60_000); - const events = await ctx.untilTurnEnd(); + await failed; expect(attempts).toBe(5); - expect(events).toContainEqual( + expect(ctx.newEvents()).toContainEqual( expect.objectContaining({ - event: 'turn.ended', - args: { - turnId: 0, - reason: 'failed', - error: expect.objectContaining({ - code: 'compaction.failed', - message: - 'CompactionTruncatedError: Compaction response was truncated before producing a complete summary.', - }), - }, + event: 'error', + args: expect.objectContaining({ + code: 'compaction.failed', + message: + 'CompactionTruncatedError: Compaction response was truncated before producing a complete summary.', + name: 'KimiError', + }), }), ); await ctx.expectResumeMatches(); @@ -1905,12 +1913,13 @@ describe('FullCompaction', () => { provider: CATALOGUED_PROVIDER, modelCapabilities: CATALOGUED_MODEL_CAPABILITIES, }); - // Set maxOutputSize on the harness's internal kimiConfig — the - // compaction path reads it via resolveModelContext().maxOutputSize. + // Set maxOutputSize on the harness's internal kimiConfig. Keep it below + // the Kimi model context window so provider-side context clipping does not + // hide whether compaction passed this configured value through. const models = (ctx as unknown as MutableKimiConfig).kimiConfig.models; models![CATALOGUED_PROVIDER.model] = { ...models![CATALOGUED_PROVIDER.model]!, - maxOutputSize: 384000, + maxOutputSize: 64_000, }; ctx.appendExchange(1, 'old user one', 'old assistant one', 20); ctx.newEvents(); @@ -1919,7 +1928,7 @@ describe('FullCompaction', () => { await ctx.untilTurnEnd(); expect(callCount).toBe(3); - expect(compactionMaxCompletionTokens).toEqual([384000]); + expect(compactionMaxCompletionTokens).toEqual([64_000]); }); it('uses default 128k hardCap when maxOutputSize is not configured', async () => { diff --git a/packages/agent-core-v2/test/goal/goal.test.ts b/packages/agent-core-v2/test/goal/goal.test.ts index ac2d7a699..2069c5323 100644 --- a/packages/agent-core-v2/test/goal/goal.test.ts +++ b/packages/agent-core-v2/test/goal/goal.test.ts @@ -1,4 +1,5 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import type { TokenUsage } from '@moonshot-ai/kosong'; import { ErrorCodes } from '#/errors'; import { IAgentContextMemoryService } from '#/agent/contextMemory'; @@ -54,6 +55,23 @@ async function runGoalStep(turnService: IAgentTurnService, turn: Turn): Promise< return step.continueTurn; } +async function recordStepUsage( + turnService: IAgentTurnService, + turn: Turn, + usage: TokenUsage, +): Promise { + const usageContext = { + turn, + usage, + stepNumber: 1, + stepUuid: 'step-1', + toolCallCount: 0, + stopTurn: false, + }; + await turnService.hooks.onStepUsage.run(usageContext); + return usageContext.stopTurn; +} + async function endTurn( turnService: IAgentTurnService, turn: Turn, @@ -546,35 +564,25 @@ describe('AgentGoalService core workflow hooks', () => { expect(turnService.launches).toEqual([]); }); - it('accounts live usage status deltas for active goal turns', async () => { + it('accounts step usage through the turn usage hook for active goal turns', async () => { await goals.createGoal({ objective: 'finish the task' }); await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 7 } }, 'model'); const turn = turnService.launch({ kind: 'user' }); await turnService.hooks.onLaunched.run({ turn }); - eventSink.emit({ - type: 'agent.status.updated', - usage: { - currentTurn: { - inputCacheRead: 0, - inputCacheCreation: 0, - inputOther: 4, - output: 0, - }, - }, - }); - eventSink.emit({ - type: 'agent.status.updated', - usage: { - currentTurn: { - inputCacheRead: 0, - inputCacheCreation: 0, - inputOther: 4, - output: 3, - }, - }, - }); + expect(await recordStepUsage(turnService, turn, { + inputCacheRead: 0, + inputCacheCreation: 0, + inputOther: 4, + output: 0, + })).toBe(false); + expect(await recordStepUsage(turnService, turn, { + inputCacheRead: 0, + inputCacheCreation: 0, + inputOther: 0, + output: 3, + })).toBe(true); expect(goals.getGoal().goal).toMatchObject({ status: 'blocked', @@ -583,6 +591,22 @@ describe('AgentGoalService core workflow hooks', () => { }); }); + it('ignores step usage for non-goal turns', async () => { + await goals.createGoal({ objective: 'finish the task' }); + + const turn = makeTurn(99); + expect(await recordStepUsage(turnService, turn, { + inputCacheRead: 0, + inputCacheCreation: 0, + inputOther: 10, + output: 5, + })).toBe(false); + expect(goals.getGoal().goal).toMatchObject({ + status: 'active', + tokensUsed: 0, + }); + }); + it('continues after creating a goal mid-turn without counting the starter turn', async () => { const turn = makeTurn(2); await turnService.hooks.onLaunched.run({ turn }); diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index de8edb800..9baf5e315 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -3,6 +3,7 @@ import { isAbsolute, relative, resolve } from 'node:path'; import { Readable, type Writable } from 'node:stream'; import { createControlledPromise } from '@antfu/utils'; +import type { Kaos } from '@moonshot-ai/kaos'; import { isToolCall, isToolCallPart, @@ -32,6 +33,7 @@ import { IAppendLogStore, IAppendLogStorage, ISessionApprovalService, + ISessionMetadata, IAtomicDocumentStorage, IAgentBackgroundService, IBlobStorage, @@ -45,7 +47,7 @@ import { IAgentExternalHooksService, IAgentFileToolsService, IAgentFullCompactionService, - IHostEnvironment, + IKaos, IAgentLLMRequesterService, ILogService, IAgentMcpService, @@ -91,7 +93,6 @@ import { type ServiceIdentifier, type AgentToolRunOverride, } from '#/index'; -import { IExecContext, createExecContext, execContextSeed } from '#/session/execContext'; import type { IProcess } from '#/session/process'; import type { AgentSwarmToolHost } from '#/agent/swarm/tools/agent-swarm'; import { Event } from '#/_base/event'; @@ -148,7 +149,8 @@ import type { WireRecordRestoreResult, } from '#/agent/wireRecord'; import type { PathAccessOperation } from '#/session/workspaceContext'; -import { createFakeAgentFs, createFakeHostEnvironment, createFakeProcessRunner } from '../tools/fixtures/fake-exec'; +import { createFakeAgentFs, createFakeProcessRunner } from '../tools/fixtures/fake-exec'; +import { createFakeKaos } from '../tools/fixtures/fake-kaos'; import { createScriptedGenerate } from './scripted-generate'; import { @@ -159,8 +161,7 @@ import { } from './snapshots'; import { recordAgentEvents, type RecordedEventEntry } from '../snapshot/events'; -const TEST_HOST_ENVIRONMENT: IHostEnvironment = createFakeHostEnvironment(); -const TEST_HOME_DIR: string = TEST_HOST_ENVIRONMENT.homeDir; +const TEST_HOME_DIR = '/home/test'; const MOCK_PROVIDER = { type: 'kimi', @@ -401,69 +402,80 @@ function defineServiceValue( } } +type KaosOverride = + | IKaos + | Kaos + | { readonly cwd?: string; readonly envLayers?: readonly Record[] }; + /** - * Session-scope override for the execution-environment atoms - * (`IHostEnvironment` / `IExecContext` / `ISessionAgentFileSystem` / - * `ISessionProcessRunner`). Replaces the v1 `kaosServices(kaos)` helper — - * tests now pass just the atoms they care about. + * Session-scope override for the execution environment and derived atoms. */ export interface ExecEnvOverride { - readonly hostEnvironment?: IHostEnvironment | Partial; - readonly execContext?: IExecContext | { readonly cwd?: string; readonly envLayers?: readonly Record[] }; + readonly kaos?: KaosOverride; + readonly execContext?: { readonly cwd?: string; readonly envLayers?: readonly Record[] }; readonly agentFs?: ISessionAgentFileSystem | Partial; readonly processRunner?: ISessionProcessRunner | Partial; } /** - * Register a fake execution-environment atom set for a test session. Any - * unspecified atom keeps the harness default (fake host env, `/workspace` - * exec ctx, throwing fs/runner). + * Register a fake execution-environment set for a test session. Any + * unspecified atom keeps the harness default fake `IKaos` and the real + * services backed by it. */ export function execEnvServices(override: ExecEnvOverride = {}): TestAgentServiceOverride { - return [ - override.hostEnvironment !== undefined - ? appService(IHostEnvironment, resolveHostEnvironmentOverride(override.hostEnvironment)) - : appServices(() => {}), - sessionServices((reg) => { - if (override.execContext !== undefined) { - const ctx = resolveExecContextOverride(override.execContext); - for (const [id, value] of execContextSeed(ctx)) { - reg.defineInstance(id as ServiceIdentifier, value); - } - } - if (override.agentFs !== undefined) { - reg.defineInstance(ISessionAgentFileSystem, resolveAgentFsOverride(override.agentFs)); - } - if (override.processRunner !== undefined) { - reg.defineInstance(ISessionProcessRunner, resolveProcessRunnerOverride(override.processRunner)); - } - reg.defineDescriptor(ISessionWorkspaceContext, new SyncDescriptor(SessionWorkspaceContextService)); - }), - ]; + return sessionServices((reg) => { + const kaosOverride = override.kaos ?? override.execContext; + if (kaosOverride !== undefined) { + reg.defineInstance(IKaos, resolveKaosOverride(kaosOverride)); + } + if (override.agentFs !== undefined) { + reg.defineInstance(ISessionAgentFileSystem, resolveAgentFsOverride(override.agentFs)); + } + if (override.processRunner !== undefined) { + reg.defineInstance(ISessionProcessRunner, resolveProcessRunnerOverride(override.processRunner)); + } + reg.defineDescriptor(ISessionWorkspaceContext, new SyncDescriptor(SessionWorkspaceContextService)); + }); } -function resolveHostEnvironmentOverride( - input: IHostEnvironment | Partial, -): IHostEnvironment { - // A full `IHostEnvironment` is an object literal (constructed via - // `createFakeHostEnvironment` or the real service). Anything that is not a - // plain-object literal (no prototype tricks) is treated as a full instance - // and passed through. Plain object literals with partial fields fall - // through to the fake factory. Since both `createFakeHostEnvironment` and - // real services produce plain-shape objects, we discriminate on the - // presence of `osKind` — a full env always has it. - if (typeof (input as IHostEnvironment).osKind === 'string') { - return input as IHostEnvironment; - } - return createFakeHostEnvironment(input as Partial); +export function kaosServices(input: IKaos | Kaos): TestAgentServiceOverride { + return sessionService(IKaos, resolveKaosOverride(input)); } -function resolveExecContextOverride( - input: IExecContext | { readonly cwd?: string; readonly envLayers?: readonly Record[] }, -): IExecContext { - if ('withCwd' in input && 'withEnv' in input) return input as IExecContext; - const partial = input as { readonly cwd?: string; readonly envLayers?: readonly Record[] }; - return createExecContext(partial.cwd ?? '/workspace', partial.envLayers ?? []); +function resolveKaosOverride(input: KaosOverride): IKaos { + if (isIKaos(input)) return input; + if (isKaosBackend(input)) return wrapKaos(input); + return wrapKaos(createFakeKaos(undefined, input.envLayers).withCwd(input.cwd ?? '/workspace')); +} + +function isIKaos(input: KaosOverride): input is IKaos { + return typeof (input as IKaos).backend === 'object' && typeof (input as IKaos).getcwd === 'function'; +} + +function isKaosBackend(input: KaosOverride): input is Kaos { + return typeof (input as Kaos).execWithEnv === 'function' && typeof (input as Kaos).getcwd === 'function'; +} + +function wrapKaos(backend: Kaos): IKaos { + return { + _serviceBrand: undefined, + get name() { + return backend.name; + }, + get cwd() { + return backend.getcwd(); + }, + get osEnv() { + return backend.osEnv; + }, + backend, + pathClass: () => backend.pathClass(), + normpath: (path) => backend.normpath(path), + gethome: () => backend.gethome(), + getcwd: () => backend.getcwd(), + withCwd: (cwd) => wrapKaos(backend.withCwd(cwd)), + withEnv: (env) => wrapKaos(backend.withEnv(env)), + }; } function resolveAgentFsOverride( @@ -690,6 +702,13 @@ export function createCommandRunner(stdout: string, exitCode = 0): ISessionProce }); } +export function createCommandKaos(stdout: string, exitCode = 0): Kaos { + const runner = createCommandRunner(stdout, exitCode); + return createFakeKaos({ + execWithEnv: async (args, env) => runner.exec(args, { env }) as unknown as ReturnType, + }); +} + export function testAgent(...inputs: readonly TestAgentInput[]): AgentTestContext { return createTestAgent(...inputs); } @@ -967,12 +986,6 @@ export class AgentTestContext { })) { reg.defineInstance(id, value); } - // Fake `IHostEnvironment` — a real `HostEnvironmentService` would kick - // off an async probe (spawn `sh --version` on Windows, `os.homedir()`, - // etc.) at App-scope construction; the harness stubs it with a - // deterministic Linux/bash snapshot instead. Tests can override with - // `execEnvServices({ hostEnvironment: … })`. - reg.defineInstance(IHostEnvironment, TEST_HOST_ENVIRONMENT); // In-memory Storage-layer backend. The `InMemoryStorageService` is no // longer auto-registered, so the harness seeds it here to keep a // workable default for storage-backed services. Tests that need durable @@ -1023,14 +1036,10 @@ export class AgentTestContext { reg.defineInstance(ISessionInteractionService, this.createInteractionService()); reg.defineInstance(ISessionApprovalService, this.createApprovalService()); reg.defineInstance(ISessionQuestionService, this.createQuestionService()); - // Seed the session `IExecContext` (was `IKaos` in the old harness). - for (const [id, value] of execContextSeed(createExecContext(this.cwd))) { - reg.defineInstance(id as ServiceIdentifier, value); - } + reg.defineInstance(IKaos, wrapKaos(createFakeKaos().withCwd(this.cwd))); // Note: `ISessionAgentFileSystem` and `ISessionProcessRunner` are - // auto-registered by their service files (backed by `IExecContext` - // and Node fs/spawn). Tests that need a fake override them via - // `execEnvServices({ agentFs: … })` / `execEnvServices({ processRunner: … })`. + // auto-registered by their service files and backed by `IKaos`. + // Tests that need a fake override them via `execEnvServices`. reg.defineInstance(ISessionTerminalBackend, createTerminalBackend()); reg.defineDescriptor(ISessionWorkspaceContext, new SyncDescriptor(SessionWorkspaceContextService)); reg.defineDescriptor(ISessionModelResolver, new SyncDescriptor(ConfigBackedModelResolver, [{}])); @@ -1502,6 +1511,7 @@ export class AgentTestContext { } async expectResumeMatches(): Promise { + await this.waitForSessionMetadata(); await this.drainWirePersistence(); const profile = this.get(IAgentProfileService); const configSnapshot = structuredClone(this.get(IConfigService).getAll() as KimiConfig); @@ -1518,14 +1528,20 @@ export class AgentTestContext { try { await resumed.restorePersisted(); + await resumed.waitForSessionMetadata(); // oxlint-disable-next-line jest/no-standalone-expect expect(resumeStateSnapshot(resumed)).toEqual(resumeStateSnapshot(this)); } finally { + await resumed.waitForSessionMetadata(); await resumed.dispose(); } } + private async waitForSessionMetadata(): Promise { + await this.session.accessor.get(ISessionMetadata).ready; + } + private async drainWirePersistence(): Promise { for (let i = 0; i < 5; i += 1) { await Promise.resolve(); diff --git a/packages/agent-core-v2/test/harness/index.ts b/packages/agent-core-v2/test/harness/index.ts index 772d36128..8977ad19f 100644 --- a/packages/agent-core-v2/test/harness/index.ts +++ b/packages/agent-core-v2/test/harness/index.ts @@ -4,6 +4,7 @@ export { agentServices, backgroundServices, configServices, + createCommandKaos, createCommandRunner, createTestAgent, appService, @@ -15,6 +16,7 @@ export { goalServices, homeDirServices, InMemoryWireRecordPersistence, + kaosServices, llmGenerateServices, logServices, mcpServices, diff --git a/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts b/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts index 556708388..606d5539c 100644 --- a/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts +++ b/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts @@ -86,6 +86,7 @@ function createHarness(): Harness { onLaunched: { run: async () => {} }, onEnded: { run: async () => {} }, beforeStep: { run: async () => {} }, + onStepUsage: { run: async () => {} }, afterStep: { run: async () => {} }, }, } as unknown as IAgentTurnService; diff --git a/packages/agent-core-v2/test/tools/fixtures/fake-exec.ts b/packages/agent-core-v2/test/tools/fixtures/fake-exec.ts new file mode 100644 index 000000000..c719ee47d --- /dev/null +++ b/packages/agent-core-v2/test/tools/fixtures/fake-exec.ts @@ -0,0 +1,37 @@ +import type { ISessionAgentFileSystem } from '#/session/agentFs'; +import type { ISessionProcessRunner } from '#/session/process'; + +function notImplemented(name: string): never { + throw new Error(`${name} not implemented - override it in the test`); +} + +export function createFakeProcessRunner( + overrides: Partial = {}, +): ISessionProcessRunner { + return { + _serviceBrand: undefined, + exec: () => notImplemented('FakeProcessRunner.exec'), + ...overrides, + }; +} + +export function createFakeAgentFs( + overrides: Partial = {}, +): ISessionAgentFileSystem { + const cwd = overrides.cwd ?? '/workspace'; + const fs: ISessionAgentFileSystem = { + _serviceBrand: undefined, + cwd, + readText: () => notImplemented('FakeAgentFs.readText'), + writeText: () => notImplemented('FakeAgentFs.writeText'), + readBytes: () => notImplemented('FakeAgentFs.readBytes'), + readLines: () => notImplemented('FakeAgentFs.readLines'), + writeBytes: () => notImplemented('FakeAgentFs.writeBytes'), + stat: () => notImplemented('FakeAgentFs.stat'), + readdir: () => notImplemented('FakeAgentFs.readdir'), + glob: () => notImplemented('FakeAgentFs.glob'), + mkdir: () => notImplemented('FakeAgentFs.mkdir'), + withCwd: (nextCwd) => createFakeAgentFs({ ...overrides, cwd: nextCwd }), + }; + return { ...fs, ...overrides }; +} diff --git a/packages/agent-core-v2/test/turn/stubs.ts b/packages/agent-core-v2/test/turn/stubs.ts index c3fcaebff..3d6596341 100644 --- a/packages/agent-core-v2/test/turn/stubs.ts +++ b/packages/agent-core-v2/test/turn/stubs.ts @@ -44,6 +44,7 @@ function makeHooks(): IAgentTurnService['hooks'] { 'onLaunched', 'onEnded', 'beforeStep', + 'onStepUsage', 'afterStep', 'onContextOverflow', ]) as IAgentTurnService['hooks']; diff --git a/packages/agent-core-v2/test/turn/turn.test.ts b/packages/agent-core-v2/test/turn/turn.test.ts index ebbc967b2..01ed27854 100644 --- a/packages/agent-core-v2/test/turn/turn.test.ts +++ b/packages/agent-core-v2/test/turn/turn.test.ts @@ -18,10 +18,10 @@ import { describe, expect, it, vi } from 'vitest'; import { abortError, abortable } from '#/_base/utils/abort'; import { ISessionAgentFileSystem } from '#/session/agentFs'; import type { ContextMessage } from '#/agent/contextMemory'; +import { IKaos } from '#/app/kaos'; import { IOAuthService } from '#/app/auth'; import { ErrorCodes, KimiError } from '#/errors'; import { HookEngine } from '#/agent/externalHooks/engine'; -import { IHostEnvironment } from '#/app/hostEnvironment'; import type { ILogger as Logger, LogPayload } from '#/app/log'; import { IAgentMcpService } from '#/agent/mcp'; import { McpConnectionManager } from '#/agent/mcp/connection-manager'; @@ -1723,7 +1723,7 @@ describe('Agent turn flow', () => { }); const registration = registerMediaTools(ctx.get(IAgentToolRegistryService), { fs: ctx.get(ISessionAgentFileSystem), - env: ctx.get(IHostEnvironment), + kaos: ctx.get(IKaos), workspace: { workspaceDir: '/workspace', additionalDirs: [] }, capabilities: mediaCapabilities(), videoUploader, From b8ef84f5e5caf4c419d46276a94d5c334698a7f8 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Wed, 1 Jul 2026 22:35:25 +0800 Subject: [PATCH 6/6] refactor: move step hooks to loop service --- .../agent-core-v2/docs/di-scope-domains.puml | 11 +++-- .../agent-core-v2/docs/di-scope-domains.svg | 2 +- .../contextInjector/contextInjectorService.ts | 4 +- .../fullCompaction/fullCompactionService.ts | 12 +++--- .../src/agent/goal/goalService.ts | 14 ++++--- packages/agent-core-v2/src/agent/loop/loop.ts | 17 ++++---- .../src/agent/loop/loopService.ts | 29 ++++++++++---- .../microCompaction/microCompactionService.ts | 8 ++-- .../src/agent/prompt/promptService.ts | 6 ++- .../src/agent/toolDedupe/toolDedupeService.ts | 10 ++--- packages/agent-core-v2/src/agent/turn/turn.ts | 4 -- .../src/agent/turn/turnService.ts | 26 +++++------- .../test/contextInjector/manager.test.ts | 5 ++- packages/agent-core-v2/test/goal/goal.test.ts | 40 +++++++++++-------- .../promptLegacy/promptLegacyService.test.ts | 3 -- .../test/toolDedup/tool-dedup.test.ts | 4 +- packages/agent-core-v2/test/turn/stubs.ts | 25 ++++++++++-- 17 files changed, 129 insertions(+), 91 deletions(-) diff --git a/packages/agent-core-v2/docs/di-scope-domains.puml b/packages/agent-core-v2/docs/di-scope-domains.puml index a021d1006..a513ad57d 100644 --- a/packages/agent-core-v2/docs/di-scope-domains.puml +++ b/packages/agent-core-v2/docs/di-scope-domains.puml @@ -170,6 +170,7 @@ contextMemory --> wireRecord #34495E contextMemory --> replayBuilder #34495E contextInjector --> contextMemory #34495E contextInjector --> turn #34495E +contextInjector --> loop #34495E contextInjector --> systemReminder #34495E contextSize --> contextMemory #34495E contextSize --> eventSink #34495E @@ -185,6 +186,7 @@ profile --> modelProvider #34495E profile --> chatProvider #34495E prompt --> contextMemory #34495E prompt --> turn #34495E +prompt --> loop #34495E prompt --> wireRecord #34495E prompt --> eventSink #34495E turn --> loop #34495E @@ -217,7 +219,7 @@ llmRequestLog --> log #34495E toolExecutor --> toolRegistry #34495E toolStore --> wireRecord #34495E toolDedup --> telemetry #34495E -toolDedup --> turn #34495E +toolDedup --> loop #34495E toolDedup --> toolExecutor #34495E permissionGate --> permissionMode #34495E permissionGate --> permissionRules #34495E @@ -249,6 +251,7 @@ goal --> telemetry #34495E goal --> contextInjector #34495E goal --> contextMemory #34495E goal --> turn #34495E +goal --> loop #34495E goal --> toolRegistry #34495E goal --> permissionMode #34495E skill --> prompt #34495E @@ -299,13 +302,14 @@ fullCompaction --> eventSink #34495E fullCompaction --> replayBuilder #34495E fullCompaction --> externalHooks #34495E fullCompaction --> turn #34495E +fullCompaction --> loop #34495E microCompaction --> contextMemory #34495E microCompaction --> contextSize #34495E microCompaction --> wireRecord #34495E microCompaction --> flag #34495E microCompaction --> profile #34495E microCompaction --> telemetry #34495E -microCompaction --> turn #34495E +microCompaction --> loop #34495E externalHooks --> toolExecutor #34495E externalHooks --> config #34495E todoList --> contextMemory #34495E @@ -353,11 +357,12 @@ profile ..> wireRecord #16A085 : config.update / tools.set_active_tools contextMemory ..> wireRecord #16A085 : context.splice cron ..> wireRecord #16A085 : cron.add / delete / cursor fullCompaction ..> wireRecord #16A085 : full_compaction.begin/cancel/complete -fullCompaction ..> turn #16A085 : hooks.onContextOverflow +fullCompaction ..> loop #16A085 : hooks.onContextOverflow permissionRules ..> wireRecord #16A085 : permission.rules.add / record_approval_result skill ..> wireRecord #16A085 : skill.activate goal ..> wireRecord #16A085 : goal.create/update/clear goal ..> turn #16A085 : turn lifecycle hooks +goal ..> loop #16A085 : step/usage hooks goal ..> eventSink #16A085 : hook.result / goal.updated microCompaction ..> wireRecord #16A085 : micro_compaction.apply / full_compaction.complete swarm ..> wireRecord #16A085 : swarm_mode.enter/exit diff --git a/packages/agent-core-v2/docs/di-scope-domains.svg b/packages/agent-core-v2/docs/di-scope-domains.svg index 35d6e9c84..6a3845bca 100644 --- a/packages/agent-core-v2/docs/di-scope-domains.svg +++ b/packages/agent-core-v2/docs/di-scope-domains.svg @@ -1 +1 @@ -App scope (process-wide)Session scope (per session)Agent scope (per agent)bootstrapAppIBootstrapServicelogAppILogServiceILogWriterServicetelemetryAppITelemetryServiceeventAppIEventServicestorageAppIStorageServiceIAppendLogStoreIAtomicDocumentStoreIAtomicTomlDocumentStorefilestoreAppIFileStoregatewayAppIRestGatewayIWSGatewayIWSBroadcastServicesession-lifecycleAppISessionLifecycleServicesession-indexAppISessionIndexhostFsAppIHostFileSystemworkspaceRegistryAppIWorkspaceRegistryhostFolderBrowserAppIHostFolderBrowserkaosAppIKaosFactoryauthAppIOAuthServiceIAuthSummaryServiceproviderAppIProviderServiceflagAppIFlagServiceIFlagRegistryconfigAppIConfigRegistryIConfigServicechatProviderAppIChatProviderFactorymodelAppIModelServicemodelCatalogAppIModelCatalogServicesessionSessionISessionServicesession-contextSessionISessionContext (seed)session-metadataSessionISessionMetadatasession-activitySessionISessionActivityagent-lifecycleSessionIAgentLifecycleServiceinteractionSessionIInteractionServiceworkspaceContextSessionIWorkspaceContextsessionLogSessionISessionLogServiceILogWriterServicekaosSessionIKaos (seed)agentFsSessionIAgentFileSystemIFsServiceapprovalSessionIApprovalServicequestionSessionIQuestionServicesubagentHostSessionISubagentHostprocessSessionIProcessRunnerIProcessterminalSessionITerminalServiceITerminalBackendmodelProviderSessionIModelProvider (seed)wireRecordAgentIWireRecord (event hub)eventSinkAgentIEventSinkblobStoreAgentIBlobStoreServicecontextMemoryAgentIContextMemorycontextProjectorAgentIContextProjectorcontextInjectorAgentIContextInjectorcontextSizeAgentIContextSizeServicesystemReminderAgentISystemReminderServicereplayBuilderAgentIReplayBuilderServiceprofileAgentIProfileServicepromptAgentIPromptServiceturnAgentITurnServiceloopAgentILoopServicellmRequesterAgentILLMRequesterllmRequestLogAgentILLMRequestLogServicetoolRegistryAgentIToolRegistrytoolExecutorAgentIToolExecutortoolStoreAgentIToolStoreServicetoolDedupAgentIToolDedupepermissionAgentIPermissionGatepermissionModeAgentIPermissionModeServicepermissionPolicyAgentIPermissionPolicyServicepermissionRulesAgentIPermissionRulesServiceplanAgentIPlanServicegoalAgentIGoalServiceskillAgentIAgentSkillServiceuserToolAgentIUserToolServicebackgroundAgentIBackgroundServicecronAgentICronServiceswarmAgentISwarmServicemcpAgentIMcpServicefullCompactionAgentIFullCompactionmicroCompactionAgentIMicroCompactionServiceexternalHooksAgentIExternalHooksServicetodoListAgentITodoListServiceusageAgentIUsageServicerpcAgentIAgentRPCServicefileToolsAgentIFileToolsServiceshellToolsAgentIShellToolsServiceonDidChangecontext.splicehooks.onSplicedcontext_size.measuredconfig.update / tools.set_active_toolshooks.onSplicedhooks.onResumeEndedtools.update_storepermission.set_modepermission.rules.add / record_approval_resultplan_mode.enter/cancel/exitgoal.create/update/clearskill.activatetools.register_/unregister_user_toolbackground.task.started/terminatedhooks.onSplicedcron.add / delete / cursorswarm_mode.enter/exitfull_compaction.begin/cancel/completehooks.onContextOverflowmicro_compaction.apply / full_compaction.completeusage.recordsubscribehooks.onEndedScope = node color      App (process-wide)      Session (per session)      Agent (per agent)Edgessolid: DI injection (ctor @IX)dashed: event-driven (subscribe/emit)direction: consumer ---> providerNotesGenerated from `node scripts/dep-graph.mjs` output;`_base` / seed / options deps are omitted. \ No newline at end of file +App scope (process-wide)Session scope (per session)Agent scope (per agent)bootstrapAppIBootstrapServicelogAppILogServiceILogWriterServicetelemetryAppITelemetryServiceeventAppIEventServicestorageAppIStorageServiceIAppendLogStoreIAtomicDocumentStoreIAtomicTomlDocumentStorefilestoreAppIFileStoregatewayAppIRestGatewayIWSGatewaysession-lifecycleAppISessionLifecycleServicesession-indexAppISessionIndexhostFsAppIHostFileSystemworkspaceRegistryAppIWorkspaceRegistryhostFolderBrowserAppIHostFolderBrowserkaosAppIKaosFactoryauthAppIOAuthServiceIAuthSummaryServiceproviderAppIProviderServiceflagAppIFlagServiceIFlagRegistryconfigAppIConfigRegistryIConfigServicechatProviderAppIChatProviderFactorymodelAppIModelServicemodelCatalogAppIModelCatalogServiceglobalSkillCatalogAppIGlobalSkillCatalogISkillCatalogStoresessionSessionISessionServicesession-contextSessionISessionContext (seed)session-metadataSessionISessionMetadatasession-activitySessionISessionActivityagent-lifecycleSessionIAgentLifecycleServiceinteractionSessionIInteractionServiceworkspaceContextSessionIWorkspaceContextsessionLogSessionISessionLogServiceILogWriterServicesessionSkillCatalogSessionISessionSkillCatalogkaosSessionIKaos (seed)agentFsSessionIAgentFileSystemIFsServiceapprovalSessionIApprovalServicequestionSessionIQuestionServiceprocessSessionIProcessRunnerIProcessterminalSessionITerminalServiceITerminalBackendmodelProviderSessionIModelProvider (seed)wireRecordAgentIAgentWireRecordService (event hub)eventSinkAgentIAgentEventSinkServiceblobStoreAgentIAgentBlobStoreServicecontextMemoryAgentIAgentContextMemoryServicecontextProjectorAgentIAgentContextProjectorServicecontextInjectorAgentIAgentContextInjectorServicecontextSizeAgentIAgentContextSizeServicesystemReminderAgentIAgentSystemReminderServicereplayBuilderAgentIAgentReplayBuilderServiceprofileAgentIAgentProfileServicepromptAgentIAgentPromptServiceturnAgentIAgentTurnServiceloopAgentIAgentLoopServicellmRequesterAgentIAgentLLMRequesterServicellmRequestLogAgentIAgentLLMRequestLogServicetoolRegistryAgentIAgentToolRegistryServicetoolExecutorAgentIAgentToolExecutorServicetoolStoreAgentIAgentToolStoreServicetoolDedupAgentIAgentToolDedupeServicepermissionGateAgentIAgentPermissionGatepermissionModeAgentIAgentPermissionModeServicepermissionPolicyAgentIAgentPermissionPolicyServicepermissionRulesAgentIAgentPermissionRulesServiceplanAgentIAgentPlanServicegoalAgentIAgentGoalServiceskillAgentIAgentSkillServicequestionToolsAgentIAgentQuestionToolsServiceuserToolAgentIAgentUserToolServicebackgroundAgentIAgentBackgroundServicecronAgentIAgentCronServiceswarmAgentIAgentSwarmServicemcpAgentIAgentMcpServicefullCompactionAgentIAgentFullCompactionServicemicroCompactionAgentIAgentMicroCompactionServiceexternalHooksAgentIAgentExternalHooksServicetodoListAgentIAgentTodoListServiceusageAgentIAgentUsageServicerpcAgentIAgentRPCServicefileToolsAgentIAgentFileToolsServiceshellToolsAgentIAgentShellToolsServiceagentToolAgentIAgentToolServicescopeContextAgentIAgentScopeContext (seed)onDidChangeConfigurationcontext.splicehooks.onSplicedcontext_size.measuredconfig.update / tools.set_active_toolshooks.onSplicedhooks.onResumeEndedtools.update_storepermission.set_modepermission.rules.add / record_approval_resultplan_mode.enter/cancel/exitgoal.create/update/clearhook.result / goal.updatedturn lifecycle hooksstep/usage hooksskill.activatetools.register_/unregister_user_toolbackground.task.started/terminatedhooks.onSplicedcron.add / delete / cursorswarm_mode.enter/exitfull_compaction.begin/cancel/completehooks.onContextOverflowmicro_compaction.apply / full_compaction.completeusage.recordhooks.onEndedScope = node color      App (process-wide)      Session (per session)      Agent (per agent)Edgessolid: DI injection (ctor @IX)dashed: event-driven (subscribe/emit)direction: consumer ---> providerNotesGenerated from `node scripts/dep-graph.mjs` output;`_base` / seed / options deps are omitted. \ No newline at end of file diff --git a/packages/agent-core-v2/src/agent/contextInjector/contextInjectorService.ts b/packages/agent-core-v2/src/agent/contextInjector/contextInjectorService.ts index 5cea5d766..ce84d4622 100644 --- a/packages/agent-core-v2/src/agent/contextInjector/contextInjectorService.ts +++ b/packages/agent-core-v2/src/agent/contextInjector/contextInjectorService.ts @@ -6,6 +6,7 @@ import { InstantiationType } from '#/_base/di/extensions'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; import { IAgentContextMemoryService } from '#/agent/contextMemory'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentSystemReminderService } from '#/agent/systemReminder'; import { IAgentTurnService } from '#/agent/turn'; import type { ContextMessage } from '#/agent/contextMemory'; @@ -32,11 +33,12 @@ export class AgentContextInjectorService extends Disposable implements IAgentCon constructor( @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, @IAgentTurnService turnService: IAgentTurnService, + @IAgentLoopService loopService: IAgentLoopService, @IAgentSystemReminderService private readonly reminders: IAgentSystemReminderService, ) { super(); this._register( - turnService.hooks.beforeStep.register('context-injector', async (_ctx, next) => { + loopService.hooks.beforeStep.register('context-injector', async (_ctx, next) => { await next(); await this.inject(); }), diff --git a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts index 0aef4791b..270d52047 100644 --- a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts +++ b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts @@ -21,6 +21,7 @@ import { IAgentContextSizeService } from '#/agent/contextSize'; import { IAgentEventSinkService } from '#/agent/eventSink'; import { IAgentExternalHooksService } from '#/agent/externalHooks'; import { IAgentLLMRequesterService, type LLMEvent } from '#/agent/llmRequester'; +import { IAgentLoopService } from '#/agent/loop'; import { isAbortError } from '#/agent/loop/errors'; import { retryBackoffDelays, sleepForRetry } from '#/agent/loop/retry'; import { IAgentProfileService } from '#/agent/profile'; @@ -109,6 +110,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull @IAgentReplayBuilderService private readonly replayBuilder: IAgentReplayBuilderService, @IAgentExternalHooksService private readonly externalHooks: IAgentExternalHooksService, @IAgentTurnService turnService: IAgentTurnService, + @IAgentLoopService loopService: IAgentLoopService, ) { super(); this.strategy = @@ -121,19 +123,19 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull }), ); this._register( - turnService.hooks.beforeStep.register('full-compaction', async (ctx, next) => { + loopService.hooks.beforeStep.register('full-compaction', async (ctx, next) => { await this.beforeStep(ctx.turn.abortController.signal, ctx.turn.id); await next(); }), ); this._register( - turnService.hooks.afterStep.register('full-compaction', async (_ctx, next) => { + loopService.hooks.afterStep.register('full-compaction', async (_ctx, next) => { await this.afterStep(); await next(); }), ); this._register( - turnService.hooks.onContextOverflow.register('full-compaction', async (ctx, next) => { + loopService.hooks.onContextOverflow.register('full-compaction', async (ctx, next) => { await this.onContextOverflow(ctx, next); }), ); @@ -592,8 +594,8 @@ function isTodoItem(value: unknown): value is TodoItem { export { AgentFullCompactionService as FullCompaction }; -// Construct eagerly (not delayed): the service registers turn-lifecycle hooks -// (onLaunched / beforeStep / afterStep) in its constructor that drive auto +// Construct eagerly (not delayed): the service registers turn and loop hooks +// (onLaunched / beforeStep / afterStep / onContextOverflow) that drive auto // compaction. With delayed instantiation the eager `accessor.get(IAgentFullCompactionService)` // only realizes a proxy, so the hooks would not register until the first RPC — // after turns have already run without the auto-compaction gate. diff --git a/packages/agent-core-v2/src/agent/goal/goalService.ts b/packages/agent-core-v2/src/agent/goal/goalService.ts index a6f376c35..fe6d2b7af 100644 --- a/packages/agent-core-v2/src/agent/goal/goalService.ts +++ b/packages/agent-core-v2/src/agent/goal/goalService.ts @@ -3,9 +3,9 @@ * * Owns the per-agent goal lifecycle; persists records through `wireRecord`, * broadcasts through `eventSink`, injects reminders through `contextInjector`, - * drives continuation turns through `turn`, updates context through - * `contextMemory`, writes system reminders through `systemReminder`, registers - * model tools through `toolRegistry`, and reports telemetry through + * drives continuation turns through `turn`, participates in steps through + * `loop`, updates context through `contextMemory`, writes system reminders + * through `systemReminder`, registers model tools through `toolRegistry`, and reports telemetry through * `telemetry`. Bound at Agent scope. */ @@ -29,6 +29,7 @@ import { type PromptOrigin, } from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentPermissionModeService } from '#/agent/permissionMode'; import { IAgentReplayBuilderService } from '#/agent/replayBuilder'; import { IAgentSystemReminderService } from '#/agent/systemReminder'; @@ -175,6 +176,7 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { @IAgentContextInjectorService dynamicInjector: IAgentContextInjectorService, @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, @IAgentTurnService private readonly turnService: IAgentTurnService, + @IAgentLoopService loopService: IAgentLoopService, @IAgentToolRegistryService toolRegistry: IAgentToolRegistryService, @IAgentPermissionModeService private readonly permissionMode: IAgentPermissionModeService, ) { @@ -221,19 +223,19 @@ export class AgentGoalService extends Disposable implements IAgentGoalService { }), ); this._register( - turnService.hooks.beforeStep.register('goal-count-turn', async (ctx, next) => { + loopService.hooks.beforeStep.register('goal-count-turn', async (ctx, next) => { await this.handleBeforeStep(ctx); await next(); }), ); this._register( - turnService.hooks.onStepUsage.register('goal-record-step-usage', async (ctx, next) => { + loopService.hooks.onStepUsage.register('goal-record-step-usage', async (ctx, next) => { this.handleStepUsage(ctx); await next(); }), ); this._register( - turnService.hooks.afterStep.register('goal-outcome-continuation', async (ctx, next) => { + loopService.hooks.afterStep.register('goal-outcome-continuation', async (ctx, next) => { await next(); this.handleAfterStep(ctx); }), diff --git a/packages/agent-core-v2/src/agent/loop/loop.ts b/packages/agent-core-v2/src/agent/loop/loop.ts index f69fa1d01..1dfd6d756 100644 --- a/packages/agent-core-v2/src/agent/loop/loop.ts +++ b/packages/agent-core-v2/src/agent/loop/loop.ts @@ -1,5 +1,5 @@ import { createDecorator } from "#/_base/di"; -import type { HookSlot } from '#/hooks'; +import type { Hooks } from '#/hooks'; import type { Turn, TurnContextOverflowContext, @@ -8,16 +8,15 @@ import type { TurnStepUsageContext, } from '#/agent/turn'; -export interface LoopRunHooks { - readonly beforeStep: HookSlot; - readonly onStepUsage: HookSlot; - readonly afterStep: HookSlot; - readonly onContextOverflow: HookSlot; -} - export interface IAgentLoopService { readonly _serviceBrand: undefined; - runTurn(turn: Turn, hooks?: LoopRunHooks): Promise; + readonly hooks: Hooks<{ + beforeStep: TurnStepContext; + onStepUsage: TurnStepUsageContext; + afterStep: TurnStepContext; + onContextOverflow: TurnContextOverflowContext; + }>; + runTurn(turn: Turn): Promise; } export const IAgentLoopService = createDecorator('agentLoopService'); diff --git a/packages/agent-core-v2/src/agent/loop/loopService.ts b/packages/agent-core-v2/src/agent/loop/loopService.ts index 21f1772fa..15bf76cff 100644 --- a/packages/agent-core-v2/src/agent/loop/loopService.ts +++ b/packages/agent-core-v2/src/agent/loop/loopService.ts @@ -43,6 +43,7 @@ import { IAgentContextSizeService } from '#/agent/contextSize'; import { IAgentEventSinkService } from '#/agent/eventSink'; import { IAgentExternalHooksService } from '#/agent/externalHooks'; import { IAgentLLMRequesterService } from '#/agent/llmRequester'; +import { OrderedHookSlot } from '#/hooks'; import { ILogService } from '#/app/log'; import { IAgentProfileService } from '#/agent/profile'; import { IConfigRegistry, IConfigService } from '#/app/config'; @@ -57,7 +58,7 @@ import type { LoopRecordedEvent, } from './events'; import type { LLM, LLMChatParams, LLMChatResponse } from './llm'; -import { IAgentLoopService, type LoopRunHooks } from './loop'; +import { IAgentLoopService } from './loop'; import { LOOP_CONTROL_SECTION, LoopControlSchema, @@ -82,6 +83,13 @@ type TelemetryProperties = Record; export class AgentLoopService extends Disposable implements IAgentLoopService { declare readonly _serviceBrand: undefined; + readonly hooks: IAgentLoopService['hooks'] = { + beforeStep: new OrderedHookSlot(), + onStepUsage: new OrderedHookSlot(), + afterStep: new OrderedHookSlot(), + onContextOverflow: new OrderedHookSlot(), + }; + private readonly openSteps = new Map(); private readonly toolCallStartedAt = new Map(); private readonly toolCallDupType = new Map(); @@ -108,6 +116,11 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { @ILogService private readonly log: ILogService, ) { super(); + this._register( + this.hooks.beforeStep.register('turn-before-step-event', async (_ctx, next) => { + await next(); + }), + ); configRegistry.registerSection(LOOP_CONTROL_SECTION, LoopControlSchema, { fromToml: loopControlFromToml, toToml: loopControlToToml, @@ -127,11 +140,11 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { ); } - async runTurn(turn: Turn, hooks: LoopRunHooks | undefined): Promise { + async runTurn(turn: Turn): Promise { const startedAt = Date.now(); this.protocolTurnId = turn.id; const llm = this.createLLM(turn.id); - const loopHooks = this.loopHooks(turn, hooks); + const loopHooks = this.loopHooks(turn); try { // Preflight the model configuration before any step begins. Legacy reads // `config.model` at the top of its step loop, so a missing model fails the @@ -171,7 +184,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { toolCallCount: context.toolCallCount, stopTurn: false, }; - await hooks?.onStepUsage.run(usageContext); + await this.hooks.onStepUsage.run(usageContext); return usageContext.stopTurn ? { stopTurn: true } : undefined; }, }); @@ -185,7 +198,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { } catch (error) { if (isContextOverflowError(error)) { const context = { turn, error, handled: false }; - await hooks?.onContextOverflow.run(context); + await this.hooks.onContextOverflow.run(context); if (context.handled) continue; } throw error; @@ -579,17 +592,17 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { }); } - private loopHooks(turn: Turn, hooks: LoopRunHooks | undefined): LoopHooks { + private loopHooks(turn: Turn): LoopHooks { let continueAfterStop = false; let stopHookContinuationUsed = false; return { beforeStep: async () => { - await hooks?.beforeStep.run({ turn, continueTurn: false }); + await this.hooks.beforeStep.run({ turn, continueTurn: false }); return undefined; }, afterStep: async (context) => { const turnContext = { turn, continueTurn: false }; - await hooks?.afterStep.run(turnContext); + await this.hooks.afterStep.run(turnContext); if (context.stopReason !== 'tool_use' && turnContext.continueTurn) { continueAfterStop = true; } diff --git a/packages/agent-core-v2/src/agent/microCompaction/microCompactionService.ts b/packages/agent-core-v2/src/agent/microCompaction/microCompactionService.ts index a6d68232c..42a359e69 100644 --- a/packages/agent-core-v2/src/agent/microCompaction/microCompactionService.ts +++ b/packages/agent-core-v2/src/agent/microCompaction/microCompactionService.ts @@ -4,7 +4,7 @@ * Tracks cache-miss compaction cutoffs over `contextMemory`, sizes context via * `contextSize`, resolves model capacity through `profile`, persists cutoffs * through `wireRecord`, gates behavior through `flag`, emits telemetry, and - * participates in `turn` hooks. Bound at Agent scope. + * participates in `loop` hooks. Bound at Agent scope. */ import type { ContentPart } from '@moonshot-ai/kosong'; @@ -23,9 +23,9 @@ import { IConfigRegistry, IConfigService } from '#/app/config'; import { IAgentContextMemoryService } from '#/agent/contextMemory'; import { IAgentContextSizeService } from '#/agent/contextSize'; import { IFlagService } from '#/app/flag'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentProfileService } from '#/agent/profile'; import { ITelemetryService } from '#/app/telemetry'; -import { IAgentTurnService } from '#/agent/turn'; import type { ContextMessage } from '#/agent/contextMemory'; import { IAgentWireRecordService } from '#/agent/wireRecord'; import { @@ -71,7 +71,7 @@ export class AgentMicroCompactionService @IFlagService private readonly flags: IFlagService, @IAgentProfileService private readonly profile: IAgentProfileService, @ITelemetryService private readonly telemetry: ITelemetryService, - @IAgentTurnService turn: IAgentTurnService, + @IAgentLoopService loop: IAgentLoopService, @IConfigRegistry configRegistry: IConfigRegistry, @IConfigService private readonly config: IConfigService, ) { @@ -86,7 +86,7 @@ export class AgentMicroCompactionService }), ); this._register( - turn.hooks.beforeStep.register( + loop.hooks.beforeStep.register( 'micro-compaction', async (_ctx, next) => { this.detect(); diff --git a/packages/agent-core-v2/src/agent/prompt/promptService.ts b/packages/agent-core-v2/src/agent/prompt/promptService.ts index e68ad3be5..1fd086bb6 100644 --- a/packages/agent-core-v2/src/agent/prompt/promptService.ts +++ b/packages/agent-core-v2/src/agent/prompt/promptService.ts @@ -9,6 +9,7 @@ import { type ContextMessage, } from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentTurnService, type Turn } from '#/agent/turn'; import { IAgentWireRecordService } from '#/agent/wireRecord'; import { IAgentPromptService } from './prompt'; @@ -21,14 +22,15 @@ export class AgentPromptService implements IAgentPromptService { constructor( @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, @IAgentTurnService private readonly turnService: IAgentTurnService, + @IAgentLoopService loopService: IAgentLoopService, @IAgentWireRecordService private readonly wireRecord: IAgentWireRecordService, @IAgentEventSinkService private readonly events: IAgentEventSinkService, ) { - turnService.hooks.beforeStep.register('prompt-service-steer-before-step', async (_ctx, next) => { + loopService.hooks.beforeStep.register('prompt-service-steer-before-step', async (_ctx, next) => { this.flushSteerQueue(); await next(); }); - turnService.hooks.afterStep.register('prompt-service-steer', async (ctx, next) => { + loopService.hooks.afterStep.register('prompt-service-steer', async (ctx, next) => { await next(); if (this.flushSteerQueue()) { ctx.continueTurn = true; diff --git a/packages/agent-core-v2/src/agent/toolDedupe/toolDedupeService.ts b/packages/agent-core-v2/src/agent/toolDedupe/toolDedupeService.ts index e5399d3e7..84d1fe62b 100644 --- a/packages/agent-core-v2/src/agent/toolDedupe/toolDedupeService.ts +++ b/packages/agent-core-v2/src/agent/toolDedupe/toolDedupeService.ts @@ -1,7 +1,7 @@ /** * `toolDedup` domain (L4) — `IAgentToolDedupeService` implementation. * - * Self-wiring plugin: its constructor registers `turn` beforeStep/afterStep + * Self-wiring plugin: its constructor registers `loop` beforeStep/afterStep * hooks and `toolExecutor` onWillExecuteTool/onDidExecuteTool hooks to drive * same-step suppression and cross-step repeat reminders, and reports repeat * telemetry through `telemetry`. Constructed eagerly at Agent scope so the @@ -13,8 +13,8 @@ import { Disposable } from '#/_base/di/lifecycle'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; import { canonicalTelemetryArgs } from '#/_base/utils/canonical-args'; import { ITelemetryService } from '#/app/telemetry/telemetry'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentToolExecutorService } from '#/agent/toolExecutor'; -import { IAgentTurnService } from '#/agent/turn'; import type { ContentPart } from '@moonshot-ai/kosong'; import { IAgentToolDedupeService, type ToolDedupResult } from './toolDedupe'; @@ -109,15 +109,15 @@ export class AgentToolDedupeService extends Disposable implements IAgentToolDedu constructor( @ITelemetryService private readonly telemetry: ITelemetryService, - @IAgentTurnService turn: IAgentTurnService, + @IAgentLoopService loop: IAgentLoopService, @IAgentToolExecutorService toolExecutor: IAgentToolExecutorService, ) { super(); - turn.hooks.beforeStep.register('toolDedup', async (_ctx, next) => { + loop.hooks.beforeStep.register('toolDedup', async (_ctx, next) => { this.beginStep(); await next(); }); - turn.hooks.afterStep.register('toolDedup', async (_ctx, next) => { + loop.hooks.afterStep.register('toolDedup', async (_ctx, next) => { this.endStep(); await next(); }); diff --git a/packages/agent-core-v2/src/agent/turn/turn.ts b/packages/agent-core-v2/src/agent/turn/turn.ts index 430224998..344098e5c 100644 --- a/packages/agent-core-v2/src/agent/turn/turn.ts +++ b/packages/agent-core-v2/src/agent/turn/turn.ts @@ -63,10 +63,6 @@ export interface IAgentTurnService { readonly hooks: Hooks<{ onLaunched: { turn: Turn }; onEnded: TurnEndedContext; - beforeStep: TurnStepContext; - onStepUsage: TurnStepUsageContext; - afterStep: TurnStepContext; - onContextOverflow: TurnContextOverflowContext; }>; } diff --git a/packages/agent-core-v2/src/agent/turn/turnService.ts b/packages/agent-core-v2/src/agent/turn/turnService.ts index 3840d3775..8b5d0f88b 100644 --- a/packages/agent-core-v2/src/agent/turn/turnService.ts +++ b/packages/agent-core-v2/src/agent/turn/turnService.ts @@ -12,11 +12,8 @@ import { ITelemetryService } from '#/app/telemetry'; import { IAgentWireRecordService } from '#/agent/wireRecord'; import type { Turn, - TurnContextOverflowContext, TurnEndedContext, TurnResult, - TurnStepContext, - TurnStepUsageContext, } from './turn'; import { IAgentTurnService } from './turn'; @@ -44,10 +41,6 @@ export class AgentTurnService implements IAgentTurnService { readonly hooks = { onLaunched: new OrderedHookSlot<{ turn: Turn }>(), onEnded: new OrderedHookSlot(), - beforeStep: new OrderedHookSlot(), - onStepUsage: new OrderedHookSlot(), - afterStep: new OrderedHookSlot(), - onContextOverflow: new OrderedHookSlot(), }; constructor( @@ -61,10 +54,14 @@ export class AgentTurnService implements IAgentTurnService { wireRecord.register('turn.launch', (record) => { this.restoreLaunch(record.turnId); }); - this.hooks.beforeStep.register('turn-before-step-event', async (ctx, next) => { - await next(); - this.resolveReady(ctx.turn); - }); + this.loop.hooks.beforeStep.register( + 'turn-ready-before-step', + async (ctx, next) => { + await next(); + this.resolveReady(ctx.turn); + }, + { before: 'turn-before-step-event' }, + ); this.events.on((event) => { if (event.type === 'agent.status.updated' && event.planMode !== undefined) { this.planModeActive = event.planMode; @@ -133,12 +130,7 @@ export class AgentTurnService implements IAgentTurnService { result = promptHookResult; return result; } - result = await this.loop.runTurn(turn, { - beforeStep: this.hooks.beforeStep, - onStepUsage: this.hooks.onStepUsage, - afterStep: this.hooks.afterStep, - onContextOverflow: this.hooks.onContextOverflow, - }); + result = await this.loop.runTurn(turn); return result; } catch (error) { if (turn.abortController.signal.aborted) { diff --git a/packages/agent-core-v2/test/contextInjector/manager.test.ts b/packages/agent-core-v2/test/contextInjector/manager.test.ts index 4e45ec83e..402043f5a 100644 --- a/packages/agent-core-v2/test/contextInjector/manager.test.ts +++ b/packages/agent-core-v2/test/contextInjector/manager.test.ts @@ -8,6 +8,7 @@ import { import { IAgentContextInjectorService } from '#/agent/contextInjector'; import { AgentContextInjectorService } from '#/agent/contextInjector/contextInjectorService'; import { IAgentContextMemoryService, type ContextMessage } from '#/agent/contextMemory'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentProfileService } from '#/agent/profile'; import { IAgentSystemReminderService } from '#/agent/systemReminder'; import { AgentSystemReminderService } from '#/agent/systemReminder/systemReminderService'; @@ -17,7 +18,7 @@ import { IAgentToolRegistryService } from '#/agent/toolRegistry'; import { IAgentToolStoreService } from '#/agent/toolStore'; import { IAgentTurnService } from '#/agent/turn'; import { registerContextMemoryServices } from '../contextMemory/stubs'; -import { stubTurnWithHooks } from '../turn/stubs'; +import { stubLoopWithHooks, stubTurnWithHooks } from '../turn/stubs'; type InjectableContextInjector = IAgentContextInjectorService & { inject(): Promise; @@ -66,6 +67,7 @@ describe('AgentContextInjectorService', () => { base: [registerContextMemoryServices], strict: true, additionalServices: (reg) => { + reg.defineInstance(IAgentLoopService, stubLoopWithHooks()); reg.defineInstance(IAgentTurnService, stubTurnWithHooks()); reg.define(IAgentSystemReminderService, AgentSystemReminderService); reg.define(IAgentContextInjectorService, AgentContextInjectorService); @@ -218,6 +220,7 @@ describe('AgentContextInjectorService registration', () => { base: [registerContextMemoryServices], strict: true, additionalServices: (reg) => { + reg.defineInstance(IAgentLoopService, stubLoopWithHooks()); reg.defineInstance(IAgentTurnService, stubTurnWithHooks()); reg.define(IAgentSystemReminderService, AgentSystemReminderService); reg.define(IAgentContextInjectorService, AgentContextInjectorService); diff --git a/packages/agent-core-v2/test/goal/goal.test.ts b/packages/agent-core-v2/test/goal/goal.test.ts index 2069c5323..1e3f637bc 100644 --- a/packages/agent-core-v2/test/goal/goal.test.ts +++ b/packages/agent-core-v2/test/goal/goal.test.ts @@ -5,6 +5,7 @@ import { ErrorCodes } from '#/errors'; import { IAgentContextMemoryService } from '#/agent/contextMemory'; import { IAgentEventSinkService } from '#/agent/eventSink'; import { IAgentGoalService, type AgentGoalService } from '#/agent/goal'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentReplayBuilderService } from '#/agent/replayBuilder'; import { IAgentTurnService, type Turn, type TurnResult } from '#/agent/turn'; import type { PersistedWireRecord, WireRecord } from '#/agent/wireRecord'; @@ -17,7 +18,7 @@ import { wireRecordPersistenceServices, type TestAgentContext, } from '../harness'; -import { stubTurn, type StubTurn } from '../turn/stubs'; +import { stubLoopWithHooks, stubTurn, type StubTurn } from '../turn/stubs'; type GoalServiceTestManager = IAgentGoalService & AgentGoalService; type GoalRecord = Extract; @@ -48,15 +49,15 @@ function makeTurn(id: number): Turn { }; } -async function runGoalStep(turnService: IAgentTurnService, turn: Turn): Promise { +async function runGoalStep(loopService: IAgentLoopService, turn: Turn): Promise { const step = { turn, continueTurn: false }; - await turnService.hooks.beforeStep.run(step); - await turnService.hooks.afterStep.run(step); + await loopService.hooks.beforeStep.run(step); + await loopService.hooks.afterStep.run(step); return step.continueTurn; } async function recordStepUsage( - turnService: IAgentTurnService, + loopService: IAgentLoopService, turn: Turn, usage: TokenUsage, ): Promise { @@ -68,7 +69,7 @@ async function recordStepUsage( toolCallCount: 0, stopTurn: false, }; - await turnService.hooks.onStepUsage.run(usageContext); + await loopService.hooks.onStepUsage.run(usageContext); return usageContext.stopTurn; } @@ -510,12 +511,17 @@ describe('AgentGoalService core workflow hooks', () => { let context: IAgentContextMemoryService; let goals: IAgentGoalService; let turnService: StubTurn; + let loopService: IAgentLoopService; let eventSink: IAgentEventSinkService; beforeEach(() => { turnService = stubTurn({ hasActiveTurn: true }); - turnService.hooks.beforeStep.register('turn-before-step-event', (_ctx, next) => next()); - ctx = createTestAgent(agentService(IAgentTurnService, turnService)); + loopService = stubLoopWithHooks(); + loopService.hooks.beforeStep.register('turn-before-step-event', (_ctx, next) => next()); + ctx = createTestAgent( + agentService(IAgentTurnService, turnService), + agentService(IAgentLoopService, loopService), + ); context = ctx.get(IAgentContextMemoryService); goals = ctx.get(IAgentGoalService); eventSink = ctx.get(IAgentEventSinkService); @@ -530,7 +536,7 @@ describe('AgentGoalService core workflow hooks', () => { const turn = makeTurn(1); await turnService.hooks.onLaunched.run({ turn }); - await runGoalStep(turnService, turn); + await runGoalStep(loopService, turn); await endTurn(turnService, turn); expect(goals.getGoal().goal).toMatchObject({ @@ -553,7 +559,7 @@ describe('AgentGoalService core workflow hooks', () => { const turn = makeTurn(11); await turnService.hooks.onLaunched.run({ turn }); - await runGoalStep(turnService, turn); + await runGoalStep(loopService, turn); await endTurn(turnService, turn); expect(goals.getGoal().goal).toMatchObject({ @@ -564,20 +570,20 @@ describe('AgentGoalService core workflow hooks', () => { expect(turnService.launches).toEqual([]); }); - it('accounts step usage through the turn usage hook for active goal turns', async () => { + it('accounts step usage through the loop usage hook for active goal turns', async () => { await goals.createGoal({ objective: 'finish the task' }); await goals.setBudgetLimits({ budgetLimits: { tokenBudget: 7 } }, 'model'); const turn = turnService.launch({ kind: 'user' }); await turnService.hooks.onLaunched.run({ turn }); - expect(await recordStepUsage(turnService, turn, { + expect(await recordStepUsage(loopService, turn, { inputCacheRead: 0, inputCacheCreation: 0, inputOther: 4, output: 0, })).toBe(false); - expect(await recordStepUsage(turnService, turn, { + expect(await recordStepUsage(loopService, turn, { inputCacheRead: 0, inputCacheCreation: 0, inputOther: 0, @@ -595,7 +601,7 @@ describe('AgentGoalService core workflow hooks', () => { await goals.createGoal({ objective: 'finish the task' }); const turn = makeTurn(99); - expect(await recordStepUsage(turnService, turn, { + expect(await recordStepUsage(loopService, turn, { inputCacheRead: 0, inputCacheCreation: 0, inputOther: 10, @@ -610,7 +616,7 @@ describe('AgentGoalService core workflow hooks', () => { it('continues after creating a goal mid-turn without counting the starter turn', async () => { const turn = makeTurn(2); await turnService.hooks.onLaunched.run({ turn }); - await runGoalStep(turnService, turn); + await runGoalStep(loopService, turn); await goals.createGoal({ objective: 'finish the task' }, 'model'); await endTurn(turnService, turn); @@ -630,10 +636,10 @@ describe('AgentGoalService core workflow hooks', () => { const turn = makeTurn(3); await turnService.hooks.onLaunched.run({ turn }); const step = { turn, continueTurn: false }; - await turnService.hooks.beforeStep.run(step); + await loopService.hooks.beforeStep.run(step); await goals.markComplete({}, 'model'); - await turnService.hooks.afterStep.run(step); + await loopService.hooks.afterStep.run(step); await endTurn(turnService, turn); expect(step.continueTurn).toBe(true); diff --git a/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts b/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts index 606d5539c..a344cc8fb 100644 --- a/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts +++ b/packages/agent-core-v2/test/promptLegacy/promptLegacyService.test.ts @@ -85,9 +85,6 @@ function createHarness(): Harness { hooks: { onLaunched: { run: async () => {} }, onEnded: { run: async () => {} }, - beforeStep: { run: async () => {} }, - onStepUsage: { run: async () => {} }, - afterStep: { run: async () => {} }, }, } as unknown as IAgentTurnService; diff --git a/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts b/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts index 68f6806fa..c387c324a 100644 --- a/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts +++ b/packages/agent-core-v2/test/toolDedup/tool-dedup.test.ts @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; import { ITelemetryService } from '#/app/telemetry'; +import { IAgentLoopService } from '#/agent/loop'; import { IAgentToolDedupeService, AgentToolDedupeService, @@ -12,7 +13,7 @@ import type { ToolDedupResult } from '#/agent/toolDedupe'; import { IAgentToolExecutorService } from '#/agent/toolExecutor'; import { IAgentTurnService } from '#/agent/turn'; import { recordingTelemetry, type TelemetryRecord } from '../telemetry/stubs'; -import { stubToolExecutor, stubTurnWithHooks } from '../turn/stubs'; +import { stubLoopWithHooks, stubToolExecutor, stubTurnWithHooks } from '../turn/stubs'; const { REMINDER_TEXT_1, REMINDER_TEXT_3, makeReminderText2 } = toolDedupTesting; @@ -30,6 +31,7 @@ function createDeduper(telemetry = recordingTelemetry(telemetryEvents)): ToolDed const ix = createServices(disposables, { additionalServices: (reg) => { reg.defineInstance(ITelemetryService, telemetry); + reg.defineInstance(IAgentLoopService, stubLoopWithHooks()); reg.defineInstance(IAgentTurnService, stubTurnWithHooks()); reg.defineInstance(IAgentToolExecutorService, stubToolExecutor()); reg.define(IAgentToolDedupeService, AgentToolDedupeService); diff --git a/packages/agent-core-v2/test/turn/stubs.ts b/packages/agent-core-v2/test/turn/stubs.ts index 3d6596341..cc0341829 100644 --- a/packages/agent-core-v2/test/turn/stubs.ts +++ b/packages/agent-core-v2/test/turn/stubs.ts @@ -7,6 +7,7 @@ import { createHooks } from '#/hooks'; import type { PromptOrigin } from '#/agent/contextMemory'; +import type { IAgentLoopService } from '#/agent/loop'; import type { IAgentTurnService, Turn } from '#/agent/turn'; import type { IAgentToolExecutorService } from '#/agent/toolExecutor'; @@ -43,11 +44,16 @@ function makeHooks(): IAgentTurnService['hooks'] { return createHooks([ 'onLaunched', 'onEnded', + ]) as IAgentTurnService['hooks']; +} + +function makeLoopHooks(): IAgentLoopService['hooks'] { + return createHooks([ 'beforeStep', 'onStepUsage', 'afterStep', 'onContextOverflow', - ]) as IAgentTurnService['hooks']; + ]) as IAgentLoopService['hooks']; } /** A configurable `IAgentTurnService` stub backed by real `OrderedHookSlot`s. */ @@ -75,9 +81,9 @@ export function stubTurn(options: StubTurnOptions = {}): StubTurn { /** * An `IAgentTurnService` stub backed by real `OrderedHookSlot`s. Use when the system - * under test registers turn-lifecycle hooks (`onLaunched` / `beforeStep` / - * `afterStep`) in its constructor, or when a test needs to drive those hooks - * directly. `launch` returns a minimal {@link Turn}; `getActiveTurn` is a no-op. + * under test registers turn-lifecycle hooks (`onLaunched` / `onEnded`) in its + * constructor, or when a test needs to drive those hooks directly. `launch` + * returns a minimal {@link Turn}; `getActiveTurn` is a no-op. */ export function stubTurnWithHooks(): IAgentTurnService { const turn = makeTurn(0); @@ -90,6 +96,17 @@ export function stubTurnWithHooks(): IAgentTurnService { }; } +/** An `IAgentLoopService` stub backed by real loop lifecycle hook slots. */ +export function stubLoopWithHooks(): IAgentLoopService { + const hooks = makeLoopHooks(); + hooks.beforeStep.register('turn-before-step-event', (_ctx, next) => next()); + return { + _serviceBrand: undefined, + hooks, + runTurn: async () => ({ reason: 'completed' }), + }; +} + /** * An `IAgentToolExecutorService` stub whose tool-execution hooks (`onWillExecuteTool` / * `onDidExecuteTool`) are real `OrderedHookSlot`s, so services that register