diff --git a/.agents/skills/agent-core-dev/orient.md b/.agents/skills/agent-core-dev/orient.md index 402ff9304b..54f74f0083 100644 --- a/.agents/skills/agent-core-dev/orient.md +++ b/.agents/skills/agent-core-dev/orient.md @@ -64,7 +64,7 @@ So a Session-scoped service is not "L1" — e.g. `session` is Session-scoped but | L3 | registries & capabilities | `tool`, `toolRegistry`, `permission*`, `flag`, `skill`, `plugin` | | L4 | agent behaviour | `turn`, `loop`, `prompt`, `profile`, `contextMemory`, `goal`, `plan`, `swarm` | | L5 | async lifecycle | `background`, `mcp`, `cron`, `agentTool` | -| L6 | coordination | `session`, `agentLifecycle`, `sessionMetadata`, `interaction`, `terminal` | +| L6 | coordination | `session`, `agentLifecycle`, `sessionMetadata`, `interaction`, `terminal`, `undo` | | L7 | boundary / edge | `gateway`, `rpc`, `approval`, `question`, `*Legacy` | ## File-header comment convention diff --git a/.changeset/undo-rewind-consistency.md b/.changeset/undo-rewind-consistency.md new file mode 100644 index 0000000000..0a751fb58d --- /dev/null +++ b/.changeset/undo-rewind-consistency.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix /undo to restore conversation history, todo lists, plan mode, and task notifications consistently. diff --git a/apps/kimi-web/src/components/chat/ChatPane.vue b/apps/kimi-web/src/components/chat/ChatPane.vue index 5886e0ef90..79ad6c5d19 100644 --- a/apps/kimi-web/src/components/chat/ChatPane.vue +++ b/apps/kimi-web/src/components/chat/ChatPane.vue @@ -274,7 +274,7 @@ function onQueueDragEnd(): void { } // Id of the most recent user turn — the only one offered an "edit & resend" -// affordance (undo only rewinds the latest exchange). +// affordance (undo only removes the latest exchange). const lastUserTurnId = computed(() => { for (let i = props.turns.length - 1; i >= 0; i--) { if (props.turns[i]!.role === 'user') return props.turns[i]!.id; @@ -314,10 +314,10 @@ function compactionDividerLabel(turn: ChatTurn): string { // Per-turn copy button state (keyed by turn id) const copiedTurn = ref(null); -// Undo in-flight guard (keyed by turn id) — set while the server rewinds the +// Undo in-flight guard (keyed by turn id) — set while the server undoes the // turn so a second undo can't fire until the first one settles. const undoingTurnId = ref(null); -// Fallback that releases the undoing state if the server rewind never removes +// Fallback that releases the undoing state if the server undo never removes // the turn (e.g. the undo failed). Without it the guard in confirmEditMessage // would block any further undo. let undoFallbackTimer: ReturnType | null = null; @@ -339,7 +339,7 @@ function confirmEditMessage(turn: ChatTurn): void { if (undoingTurnId.value !== null) return; undoingTurnId.value = turn.id; emit('editMessage', { text: turn.text, attachments: turn.attachments }); - // Fallback: if the server rewind never removes the turn (e.g. it failed), + // Fallback: if the server undo never removes the turn (e.g. it failed), // release the guard so the user can retry. undoFallbackTimer = setTimeout(() => { undoFallbackTimer = null; @@ -347,7 +347,7 @@ function confirmEditMessage(turn: ChatTurn): void { }, UNDO_FALLBACK_MS); } -// Release the undoing guard once the server rewind has actually removed the turn +// Release the undoing guard once the server undo has actually removed the turn // from the list (post-render, so the element is already gone). watch( () => props.turns, diff --git a/apps/kimi-web/src/components/chat/ConversationPane.vue b/apps/kimi-web/src/components/chat/ConversationPane.vue index 1b6b1a9e79..0464beac9b 100644 --- a/apps/kimi-web/src/components/chat/ConversationPane.vue +++ b/apps/kimi-web/src/components/chat/ConversationPane.vue @@ -828,7 +828,7 @@ watch(scrollKey, async (next, prev) => { } await nextTick(); if (following.value || hasUserActionFollowLock()) { - // A rewind (undo / compaction) shortens the transcript — glide to the new + // An undo or compaction shortens the transcript — glide to the new // bottom smoothly; growth (new turns / streaming) snaps instantly so the // follow keeps up with the tail. scrollToBottom(next.length < prev.length); @@ -928,9 +928,9 @@ function handleComposerSubmit(payload: { text: string; attachments: PromptAttach emit('submit', payload); } -// Undo ("edit & resend") rewinds the transcript asynchronously — the server +// Undo ("edit & resend") shortens the transcript asynchronously — the server // round-trip in App.vue's handleEditMessage truncates the turns after this emit -// returns. Scrolling here would target the pre-rewind bottom and fight the +// returns. Scrolling here would target the pre-undo bottom and fight the // bubble-exit animation, so we only arm the follow state; the scrollKey watcher // smooth-scrolls once the truncated turns actually land. function handleEditMessage(payload: { diff --git a/docs/en/reference/slash-commands.md b/docs/en/reference/slash-commands.md index 75c8033564..d3ab8bf4cb 100644 --- a/docs/en/reference/slash-commands.md +++ b/docs/en/reference/slash-commands.md @@ -32,7 +32,7 @@ Some commands are only available in the idle state. Executing these commands whi | `/fork` | — | Fork a new session from the current one, preserving the full conversation history | No | | `/title []` | `/rename` | Without arguments, display the current session title; with an argument, set a new title (max 200 characters) | Yes | | `/compact []` | — | Compact the current conversation context to free up token usage; an optional custom instruction can hint to the model what to preserve | No | -| `/undo []` | — | Undo recent prompts from the active context. Without a count, opens a selector; with a count, undoes that many prompts. Prompts before the last compaction cannot be undone | No | +| `/undo []` | — | Undo recent prompts from the active context. Without a count, opens a selector; with a count, undoes that many prompts. Prompts before the last compaction cannot be undone. Undoing also rolls back the todo list and plan mode state produced by those prompts (code changes are not reverted) | No | | `/reload` | — | Reload the current session and apply the latest `config.toml` settings (providers, models, etc.) and `tui.toml` UI preferences, without restarting the CLI | No | | `/reload-tui` | — | Reload only the `tui.toml` UI preferences (theme, editor, notifications, etc.) without rebuilding the session | Yes | | `/init` | — | Analyze the current codebase and generate `AGENTS.md` | No | diff --git a/docs/zh/reference/slash-commands.md b/docs/zh/reference/slash-commands.md index ec4d2b79dc..a1882cf87e 100644 --- a/docs/zh/reference/slash-commands.md +++ b/docs/zh/reference/slash-commands.md @@ -32,7 +32,7 @@ | `/fork` | — | 基于当前会话 fork 一份新会话，保留完整对话历史 | 否 | | `/title []` | `/rename` | 不带参数时显示当前会话标题；带参数时设置为新标题（最长 200 字符） | 是 | | `/compact []` | — | 压缩当前对话上下文，释放 token 占用；可附带自定义指令，提示模型压缩时保留哪些信息 | 否 | -| `/undo []` | — | 从当前上下文撤销最近的提示词。不带数量时打开选择器；带数量时撤销对应条数。最后一次上下文压缩之前的提示词不能撤销 | 否 | +| `/undo []` | — | 从当前上下文撤销最近的提示词。不带数量时打开选择器；带数量时撤销对应条数。最后一次上下文压缩之前的提示词不能撤销。撤销会一并回滚这些提示词产生的 todo 列表和计划模式状态（不回滚代码改动） | 否 | | `/init` | — | 分析当前代码库并生成 `AGENTS.md` | 否 | | `/export-md []` | `/export` | 将当前会话导出为 Markdown 文件 | 否 | | `/export-debug-zip` | — | 将当前会话导出为调试用 ZIP 压缩包（与 [`kimi export`](./kimi-command.md#kimi-export) 行为一致） | 否 | diff --git a/packages/agent-core-v2/AGENTS.md b/packages/agent-core-v2/AGENTS.md index b63c3df196..8646b066a2 100644 --- a/packages/agent-core-v2/AGENTS.md +++ b/packages/agent-core-v2/AGENTS.md @@ -53,6 +53,10 @@ Business domains **do not implement persistence themselves** — they depend on Business code must not `import 'node:fs'`, write SQL, hand-roll append-logs / atomic writes, or hold file handles. Generic Stores are named by **access pattern** (`IAppendLogStore`, `IAtomicDocumentStore`); only domain-unique Stores are named after the domain (`ISessionIndex`). See `.agents/skills/agent-core-dev/persistence.md` for the full layering rules and decision tree. +## Conversation undo + +`context.undo` is the only persisted undo fact. `contextMemory/conversationTime.ts` owns the conversation clock (`isUndoAnchor` — the single tick predicate used by `computeUndoCut`, the checkpoint reducers, and the transcript reducer) and the checkpoint protocol. A wire Model whose state must follow conversation undo (todo, plan, task-notification delivery, …) **MUST** be defined with `defineCheckpointedModel` — never hand-roll the push/clear/restore reducers — which also registers it into `CHECKPOINTED_MODELS` for the undo pipeline's pre-cut depth check. World-time state (turn counters, task registries, revision counters) must stay outside checkpointed Models. + ## Docs Per-domain references live in `docs/`. diff --git a/packages/agent-core-v2/docs/state-manifest.d.ts b/packages/agent-core-v2/docs/state-manifest.d.ts index e26afd4492..6c6e8ae80c 100644 --- a/packages/agent-core-v2/docs/state-manifest.d.ts +++ b/packages/agent-core-v2/docs/state-manifest.d.ts @@ -653,6 +653,7 @@ export interface SessionStateSnapshot { } | /* InjectionOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'injection'; readonly variant: string; + readonly ownerPromptId?: string; } | /* ShellCommandOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'shell_command'; readonly phase: 'input' | 'output'; @@ -728,6 +729,7 @@ export interface AgentStateSnapshot { } | /* InjectionOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'injection'; readonly variant: string; + readonly ownerPromptId?: string; } | /* ShellCommandOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'shell_command'; readonly phase: 'input' | 'output'; @@ -851,6 +853,7 @@ export interface AgentStateSnapshot { } | /* InjectionOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'injection'; readonly variant: string; + readonly ownerPromptId?: string; } | /* ShellCommandOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'shell_command'; readonly phase: 'input' | 'output'; @@ -906,6 +909,7 @@ export interface AgentStateSnapshot { } | /* InjectionOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'injection'; readonly variant: string; + readonly ownerPromptId?: string; } | /* ShellCommandOrigin — packages/agent-core-v2/src/agent/contextMemory/types.ts */ { readonly kind: 'shell_command'; readonly phase: 'input' | 'output'; @@ -1002,7 +1006,7 @@ export interface AgentStateSnapshot { 'llmRequester.lastConfigLogSignature': string | undefined; 'llmRequester.mediaDegradedTurns': Set; 'llmRequester.mediaStrippedTurns': Map; 'llmRequester.turnConfigs': Map 0 && + isPromptOwnedInjection(state[cutIndex - 1]!, message) + ) { + cutIndex--; + } } } return { cutIndex, removedCount, stoppedAtCompaction }; @@ -317,7 +333,11 @@ export function isFullyUndoable(cut: UndoCut, count: number): boolean { return cut.cutIndex >= 0 && cut.removedCount >= count; } -export type UndoUnavailableReason = 'empty' | 'compaction_boundary' | 'insufficient'; +export type UndoUnavailableReason = + | 'empty' + | 'compaction_boundary' + | 'insufficient' + | 'checkpoint_lost'; export type UndoPrecheck = | { readonly ok: true } @@ -349,25 +369,19 @@ export function formatUndoUnavailableMessage( return 'Nothing to undo: would cross a compaction boundary'; case 'insufficient': return `Nothing to undo: only ${precheck.undoable} of ${precheck.requested} requested turn(s) available`; + case 'checkpoint_lost': + return 'Nothing to undo: conversation state checkpoints are incomplete'; } } export const contextUndo = ContextModel.defineOp('context.undo', { - schema: z.object({ count: z.number() }), + schema: z.object({ + count: z.number().int().positive().max(Number.MAX_SAFE_INTEGER), + }), apply: (state, p) => { - if (p.count <= 0 || state.length === 0) return state; + if (!isValidUndoCount(p.count) || state.length === 0) return state; const cut = computeUndoCut(state, p.count); if (!isFullyUndoable(cut, p.count)) return state; return resetFold(state.slice(0, cut.cutIndex)) as ContextMessage[]; }, }); - -function isRealUserPrompt(message: ContextMessage): boolean { - if (message.role !== 'user') return false; - const origin = message.origin; - if (origin === undefined || origin.kind === 'user') return true; - return ( - (origin.kind === 'skill_activation' || origin.kind === 'plugin_command') && - origin.trigger === 'user-slash' - ); -} diff --git a/packages/agent-core-v2/src/agent/contextMemory/contextTranscript.ts b/packages/agent-core-v2/src/agent/contextMemory/contextTranscript.ts index c719356be6..6279b19bf2 100644 --- a/packages/agent-core-v2/src/agent/contextMemory/contextTranscript.ts +++ b/packages/agent-core-v2/src/agent/contextMemory/contextTranscript.ts @@ -1,33 +1,8 @@ /** - * `contextMemory` transcript reducer — rebuilds the FULL message history of an - * agent from its `context.*` wire records for UI display (snapshot / messages). + * `contextMemory` domain (L4) — rebuilds display history from the wire journal. * - * The live `ContextModel` (`ContextMemoryService`) rewrites the model-facing - * context on `context.apply_compaction` into `[...keptUserMessages, - * compaction_summary]`, so reading the live context after a compaction loses - * everything before the fold. The wire log keeps every record, though, so this - * reducer re-reduces the `context.*` records with the same semantics as the - * live `ContextMemory` restore, EXCEPT that `context.apply_compaction` KEEPS - * the full history and appends a user-role summary marker — the same view the - * v1 transcript / TUI shows after resume. `foldedLength` tracks what the live - * (folded) `context.history.length` would be, so a caller can detect and - * append an unflushed live tail. - * - * Mirrors v1 `reduceWireRecords` - * (`packages/agent-core/src/services/message/transcript.ts`): - * - `context.append_message` → append (deferred while a tool exchange is open) - * - `context.append_loop_event` → step.begin/content.part/tool.call mutate the - * open assistant; tool.result appends a tool - * message with the raw output; settling a step - * (at its end or at the next begin) drops an - * output-free assistant, mirroring the live fold - * - `context.apply_compaction` → keep the full history, append the user-role - * summary marker, recover `foldedLength` from - * the recorded kept-count fields - * - `context.undo` → remove tail messages (skip injections, stop - * at compaction summaries / clear floor) - * - `context.clear` → keep prior transcript entries but reset the - * folded view + * Supplies transcript consumers with full pre-compaction history and folded + * context length while preserving undo/clear semantics. Scope-agnostic. */ import { type ContentPart, type ToolCall } from '#/kosong/contract/message'; @@ -36,9 +11,9 @@ import type { WireRecord } from '#/wire/record'; import { COMPACT_USER_MESSAGE_MAX_TOKENS, collectCompactableUserMessages, - isRealUserInput, selectRecentUserMessages, } from './compactionHandoff'; +import { isPromptOwnedInjection, isUndoAnchor } from './conversationTime'; import type { LoopRecordedEvent } from './loopEventFold'; import type { ContextMessage } from './types'; import { isVacuousContentPart } from './vacuousContent'; @@ -198,9 +173,19 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { if (message.origin?.kind === 'compaction_summary') break; transcript.splice(i, 1); foldedLength = Math.max(0, foldedLength - 1); - if (isRealUserInput(message)) { + if (isUndoAnchor(message)) { removedUserCount++; - if (removedUserCount >= count) break; + if (removedUserCount >= count) { + while ( + i > clearFloor && + isPromptOwnedInjection(transcript[i - 1]!.message, message) + ) { + transcript.splice(i - 1, 1); + i--; + foldedLength = Math.max(0, foldedLength - 1); + } + break; + } } } resetOpenState(); diff --git a/packages/agent-core-v2/src/agent/contextMemory/conversationTime.ts b/packages/agent-core-v2/src/agent/contextMemory/conversationTime.ts new file mode 100644 index 0000000000..b0196a91a4 --- /dev/null +++ b/packages/agent-core-v2/src/agent/contextMemory/conversationTime.ts @@ -0,0 +1,85 @@ +/** + * `contextMemory` domain (L4) — shared conversation clock and checkpointed + * wire-Model factory. + * + * Defines the undo anchor vocabulary and registers conversation-time Models + * for undo validation. Scope-agnostic. + */ + +import { defineModel, type ModelDef } from '#/wire/model'; + +import type { ContextMessage } from './types'; + +export function isUndoAnchor(message: ContextMessage): boolean { + if (message.role !== 'user') return false; + const origin = message.origin; + if (origin === undefined || origin.kind === 'user') return true; + return ( + (origin.kind === 'skill_activation' || origin.kind === 'plugin_command') && + origin.trigger === 'user-slash' + ); +} + +export function isPromptOwnedInjection( + message: ContextMessage, + prompt: ContextMessage, +): boolean { + const origin = message.origin; + return ( + origin?.kind === 'injection' && + origin.ownerPromptId !== undefined && + origin.ownerPromptId === prompt.id + ); +} + +export function isValidUndoCount(count: number): boolean { + return Number.isSafeInteger(count) && count > 0; +} + +export interface Checkpointed { + readonly current: T; + readonly checkpoints: readonly T[]; +} + +export const CHECKPOINTED_MODELS: ModelDef>[] = []; + +export interface CheckpointModelOptions { + readonly onAppendMessage?: (current: T, message: ContextMessage) => T; +} + +export function defineCheckpointedModel( + name: string, + initial: () => T, + opts?: CheckpointModelOptions, +): ModelDef> { + const def = defineModel>( + name, + () => ({ current: initial(), checkpoints: [] }), + { + reducers: { + 'context.append_message': (state, { message }) => { + if (isUndoAnchor(message)) { + return { ...state, checkpoints: [...state.checkpoints, state.current] }; + } + if (opts?.onAppendMessage === undefined) return state; + const current = opts.onAppendMessage(state.current, message); + return current === state.current ? state : { ...state, current }; + }, + 'context.apply_compaction': (state) => + state.checkpoints.length === 0 ? state : { ...state, checkpoints: [] }, + 'context.clear': (state) => + state.checkpoints.length === 0 ? state : { ...state, checkpoints: [] }, + 'context.undo': (state, { count }) => { + if (!isValidUndoCount(count) || state.checkpoints.length < count) return state; + const checkpointIndex = state.checkpoints.length - count; + return { + current: state.checkpoints[checkpointIndex]!, + checkpoints: state.checkpoints.slice(0, checkpointIndex), + }; + }, + }, + }, + ); + CHECKPOINTED_MODELS.push(def as ModelDef>); + return def; +} diff --git a/packages/agent-core-v2/src/agent/contextMemory/conversationUndoParticipants.ts b/packages/agent-core-v2/src/agent/contextMemory/conversationUndoParticipants.ts new file mode 100644 index 0000000000..d65a6739c4 --- /dev/null +++ b/packages/agent-core-v2/src/agent/contextMemory/conversationUndoParticipants.ts @@ -0,0 +1,62 @@ +/** + * `contextMemory` domain (L4) — Agent-scoped post-undo reconciliation registry. + * + * Hosts state-repair participants for the undo coordinator. Bound at Agent + * scope. + */ + +import { createDecorator } from '#/_base/di/instantiation'; +import { Disposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, ScopeActivation, registerScopedService } from '#/_base/di/scope'; + +export interface AgentConversationUndoParticipant { + readonly id: string; + reconcileAfterUndo(): Promise; +} + +export interface IAgentConversationUndoParticipantRegistry { + readonly _serviceBrand: undefined; + + register(participant: AgentConversationUndoParticipant): IDisposable; + list(): readonly AgentConversationUndoParticipant[]; +} + +export const IAgentConversationUndoParticipantRegistry = + createDecorator( + 'agentConversationUndoParticipantRegistry', + ); + +class AgentConversationUndoParticipantRegistry + extends Disposable + implements IAgentConversationUndoParticipantRegistry +{ + declare readonly _serviceBrand: undefined; + + private readonly participants = new Map(); + + register(participant: AgentConversationUndoParticipant): IDisposable { + if (this.participants.has(participant.id)) { + throw new Error( + `Conversation undo participant "${participant.id}" is already registered`, + ); + } + this.participants.set(participant.id, participant); + return toDisposable(() => { + if (this.participants.get(participant.id) === participant) { + this.participants.delete(participant.id); + } + }); + } + + list(): readonly AgentConversationUndoParticipant[] { + return [...this.participants.values()]; + } +} + +registerScopedService( + LifecycleScope.Agent, + IAgentConversationUndoParticipantRegistry, + AgentConversationUndoParticipantRegistry, + ScopeActivation.OnScopeCreated, + 'conversationUndoParticipants', +); diff --git a/packages/agent-core-v2/src/agent/contextMemory/types.ts b/packages/agent-core-v2/src/agent/contextMemory/types.ts index dd73e87090..24736b5547 100644 --- a/packages/agent-core-v2/src/agent/contextMemory/types.ts +++ b/packages/agent-core-v2/src/agent/contextMemory/types.ts @@ -33,6 +33,7 @@ export interface PluginCommandOrigin { export interface InjectionOrigin { readonly kind: 'injection'; readonly variant: string; + readonly ownerPromptId?: string; } export interface ShellCommandOrigin { diff --git a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts index 1f4116e57d..85f3c51fa7 100644 --- a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts +++ b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts @@ -404,8 +404,7 @@ export class AgentFullCompactionService extends Disposable implements IAgentFull }; } - /** Scope disposal is the fallback abort path; normal agent teardown aborts - * and awaits the exposed task before disposing the scope. */ + /** Scope disposal is the abort path: tear down any in-flight compaction. */ override dispose(): void { if (this._compacting !== null && !this._compacting.abortController.signal.aborted) { this._compacting.abortController.abort(); diff --git a/packages/agent-core-v2/src/agent/loop/loop.ts b/packages/agent-core-v2/src/agent/loop/loop.ts index d20a659a16..241a1e25a9 100644 --- a/packages/agent-core-v2/src/agent/loop/loop.ts +++ b/packages/agent-core-v2/src/agent/loop/loop.ts @@ -146,6 +146,8 @@ export interface IAgentLoopService { cancel(turnId?: number, reason?: unknown): boolean; + tryAcquireQuiescence(): IDisposable | undefined; + /** Resolves once no turn is active and none are queued — the disposal drain * awaited by `agentLifecycle.remove`. */ settled(): Promise; diff --git a/packages/agent-core-v2/src/agent/loop/loopService.ts b/packages/agent-core-v2/src/agent/loop/loopService.ts index 25bc361355..69a67fd771 100644 --- a/packages/agent-core-v2/src/agent/loop/loopService.ts +++ b/packages/agent-core-v2/src/agent/loop/loopService.ts @@ -119,8 +119,10 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { private readonly pendingAssignments = new Map>>(); private readonly errorHandlers: LoopErrorHandler[] = []; private readonly pendingTurns: TurnJob[] = []; + private readonly heldAdmissions: HeldAdmission[] = []; private activeTurnJob: TurnJob | undefined; private readonly settleWaiters: Array<() => void> = []; + private quiescenceDepth = 0; private activeRequestTrace: LLMRequestTrace | undefined; constructor( @@ -174,6 +176,10 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { request.abort(); this.rejectAssignment(request, reason); } + for (const { request } of this.heldAdmissions.splice(0)) { + request.abort(); + this.rejectAssignment(request, reason); + } this.maybeSettle(); super.dispose(); } @@ -184,6 +190,18 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { void assignment.catch(() => undefined); this.pendingAssignments.set(request, assignment); + if (this.quiescenceDepth > 0) { + this.heldAdmissions.push({ request, options }); + } else { + this.admit(request, options); + } + return { + assigned: assignment, + abort: (reason) => this.abortRequest(request, reason), + }; + } + + private admit(request: StepRequest, options?: StepEnqueueOptions): void { const active = this.activeTurnJob; switch (request.admission) { case 'newTurn': @@ -206,10 +224,6 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { this.assignStep(active, request, options); break; } - return { - assigned: assignment, - abort: (reason) => this.abortRequest(request, reason), - }; } private createAndQueueTurn(request: StepRequest): void { @@ -242,10 +256,34 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { ); } + tryAcquireQuiescence(): IDisposable | undefined { + if (this.disposing) throw abortError('Agent loop disposed'); + if (this.activeTurnJob !== undefined || this.hasPendingRequests()) return undefined; + this.quiescenceDepth += 1; + return toDisposable(() => this.releaseQuiescence()); + } + + private releaseQuiescence(): void { + if (this.quiescenceDepth === 0) return; + this.quiescenceDepth -= 1; + if (this.quiescenceDepth > 0 || this.disposing) return; + this.pumpTurns(); + for (const admission of this.heldAdmissions.splice(0)) { + if (admission.request.aborted) continue; + try { + this.admit(admission.request, admission.options); + } catch (error) { + admission.request.abort(); + this.rejectAssignment(admission.request, error); + } + } + this.pumpTurns(); + } + private cancelActiveTurn(turnId: number | undefined, cancellation: unknown): boolean { const job = this.activeTurnJob; if (job === undefined || (turnId !== undefined && job.turn.id !== turnId)) return false; - this.wire.dispatch(cancelTurn({ turnId })); + this.wire.dispatch(cancelTurn({ turnId: job.turn.id, target: 'active' })); job.controller.abort(cancellation); return true; } @@ -255,7 +293,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { if (index < 0) return false; const [job] = this.pendingTurns.splice(index, 1); if (job === undefined || job.turn.state !== 'queued') return false; - this.wire.dispatch(cancelTurn({ turnId })); + this.wire.dispatch(cancelTurn({ turnId, target: 'queued' })); for (const step of job.steps.values()) step.cancel(cancellation); job.controller.abort(cancellation); job.turn.state = 'cancelled'; @@ -269,12 +307,17 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { return ( this.activeTurnJob?.queue.hasPendingRequests() === true || this.standaloneStepQueue.hasPendingRequests() || - this.pendingTurns.length > 0 + this.pendingTurns.length > 0 || + this.heldAdmissions.some(({ request }) => !request.aborted) ); } settled(): Promise { - if (this.activeTurnJob === undefined && this.pendingTurns.length === 0) { + if ( + this.activeTurnJob === undefined && + this.pendingTurns.length === 0 && + this.heldAdmissions.length === 0 + ) { return Promise.resolve(); } return new Promise((resolve) => { @@ -283,7 +326,11 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { } private maybeSettle(): void { - if (this.activeTurnJob !== undefined || this.pendingTurns.length > 0) return; + if ( + this.activeTurnJob !== undefined || + this.pendingTurns.length > 0 || + this.heldAdmissions.length > 0 + ) return; if (this.settleWaiters.length === 0) return; const waiters = this.settleWaiters.splice(0); for (const resolve of waiters) resolve(); @@ -339,6 +386,14 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { } private abortRequest(request: StepRequest, reason?: unknown): boolean { + const heldIndex = this.heldAdmissions.findIndex((entry) => entry.request === request); + if (heldIndex >= 0) { + this.heldAdmissions.splice(heldIndex, 1); + if (!request.abort()) return false; + this.rejectAssignment(request, reason ?? userCancellationReason()); + this.maybeSettle(); + return true; + } for (const job of [this.activeTurnJob, ...this.pendingTurns]) { if (job === undefined) continue; if (job.turn.state === 'queued' && job.request === request) { @@ -387,7 +442,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { } private pumpTurns(): void { - if (this.disposing || this.activeTurnJob !== undefined) return; + if (this.disposing || this.quiescenceDepth > 0 || this.activeTurnJob !== undefined) return; const job = this.pendingTurns.shift(); if (job === undefined) { this.maybeSettle(); @@ -518,6 +573,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { if (step.state === 'queued' || step.state === 'running') step.cancel(reason); } this.activeTurnJob = undefined; + this.maybeSettle(); } registerLoopErrorHandler( @@ -1056,6 +1112,11 @@ interface TurnJob { readonly turn: MutableTurn; } +interface HeldAdmission { + readonly request: StepRequest; + readonly options?: StepEnqueueOptions; +} + interface LoopRuntime { readonly turnId: number; readonly turnSignal: AbortSignal; diff --git a/packages/agent-core-v2/src/agent/loop/turnOps.ts b/packages/agent-core-v2/src/agent/loop/turnOps.ts index f76ee74cf6..0a6714e92c 100644 --- a/packages/agent-core-v2/src/agent/loop/turnOps.ts +++ b/packages/agent-core-v2/src/agent/loop/turnOps.ts @@ -1,22 +1,9 @@ /** - * `loop` domain (L4) — wire Model (`TurnModel`) and the Ops that bookkeep the - * agent's turn lifecycle on the wire. + * `loop` domain (L4) — persists and restores monotonically increasing turn + * identity. * - * Declares the next turn id as a wire Model (initial `0`). The persisted - * `turn.prompt` record carries exactly v1's field set (`{ input, origin }` — - * no `turnId`), and `apply` mirrors v1's `restorePrompt()`: every record - * advances the counter by one, so the counter is restored by counting - * turn starts. Every turn is started by `loopService.enqueue` admitting a - * request that creates a new Turn, which dispatches one - * `turn.prompt` per start. As a belt-and-suspenders for v1-written logs whose - * internally-driven turns (goal continuations) have no `turn.prompt` record, - * `TurnModel` also registers a cross-model reducer on - * `context.append_loop_event` that raises the counter past any `turnId` - * observed in a replayed loop event — the v1 `observeRestoredTurnId` - * semantics. The `turn.started` / `turn.ended` / `error` signals are not part - * of this Op set and remain on their existing path (published by the loop - * service around a run). Consumed by the Agent-scope `loopService` (which - * reads the next turn id on admission). + * Owns the next available turn id, including cancelled queued reservations and + * legacy loop-event observations. Consumed by the Agent-scope `loopService`. */ import { z } from 'zod'; @@ -27,22 +14,27 @@ import type { PromptOrigin } from '#/agent/contextMemory/types'; export interface TurnModelState { readonly nextTurnId: number; + readonly cancelledTurnIds: readonly number[]; } -export const TurnModel = defineModel('turn', () => ({ nextTurnId: 0 }), { - reducers: { - 'context.append_loop_event': (state, { event }) => { - if (event.type === 'tool.result' || event.turnId === undefined) { - return state; - } +export const TurnModel = defineModel( + 'turn', + () => ({ nextTurnId: 0, cancelledTurnIds: [] }), + { + reducers: { + 'context.append_loop_event': (state, { event }) => { + if (event.type === 'tool.result' || event.turnId === undefined) { + return state; + } - const turnId = Number.parseInt(event.turnId, 10); - return Number.isInteger(turnId) && turnId >= state.nextTurnId - ? { nextTurnId: turnId + 1 } - : state; + const turnId = Number.parseInt(event.turnId, 10); + return Number.isInteger(turnId) && turnId >= state.nextTurnId + ? advanceTurnClock(state, turnId + 1) + : state; + }, }, }, -}); +); const turnInputShape = { input: z.custom(), @@ -59,7 +51,7 @@ declare module '#/wire/types' { export const promptTurn = TurnModel.defineOp('turn.prompt', { schema: z.object(turnInputShape), - apply: (s) => ({ nextTurnId: s.nextTurnId + 1 }), + apply: (s) => advanceTurnClock(s, s.nextTurnId + 1), }); export const steerTurn = TurnModel.defineOp('turn.steer', { @@ -68,6 +60,27 @@ export const steerTurn = TurnModel.defineOp('turn.steer', { }); export const cancelTurn = TurnModel.defineOp('turn.cancel', { - schema: z.object({ turnId: z.number().optional() }), - apply: (s) => s, + schema: z.object({ + turnId: z.number().optional(), + target: z.enum(['active', 'queued']).optional(), + }), + apply: (s, { turnId, target }) => { + if (target === undefined || turnId === undefined || turnId < s.nextTurnId) return s; + return advanceTurnClock(s, s.nextTurnId, [...s.cancelledTurnIds, turnId]); + }, }); + +function advanceTurnClock( + state: TurnModelState, + nextTurnId: number, + cancelledTurnIds: readonly number[] = state.cancelledTurnIds, +): TurnModelState { + const pendingCancellations = new Set( + cancelledTurnIds.filter((turnId) => turnId >= nextTurnId), + ); + while (pendingCancellations.delete(nextTurnId)) nextTurnId += 1; + return { + nextTurnId, + cancelledTurnIds: [...pendingCancellations].toSorted((a, b) => a - b), + }; +} diff --git a/packages/agent-core-v2/src/agent/plan/planOps.ts b/packages/agent-core-v2/src/agent/plan/planOps.ts index 8bccc920b0..1781af0184 100644 --- a/packages/agent-core-v2/src/agent/plan/planOps.ts +++ b/packages/agent-core-v2/src/agent/plan/planOps.ts @@ -7,13 +7,14 @@ * reference-only fact. * * The Model holds the persistent, replayable fields — whether plan mode is - * active, the plan id, and the last recorded revision version per plan id. - * The lifecycle records keep exactly v1's field set (`{ id }`); the plan file - * path is NOT persisted — it is derived from the id at read time - * (`planService.planFilePathFor`), matching v1's `restoreEnter`. Plan - * content is recorded separately: every ExitPlanMode submit snapshots the - * plan file into blob storage and persists a `plan.revision` record carrying - * only the reference (`{ id, version, path, sha256, bytes }`, `path` + * active, the plan id, and the last recorded revision version per plan id — + * wrapped in `contextMemory`'s checkpoint protocol so plan mode stays aligned + * with conversation undo. The lifecycle records keep exactly v1's field set + * (`{ id }`); the plan file path is NOT persisted — it is derived from the id + * at read time (`planService.planFilePathFor`), matching v1's `restoreEnter`. + * Plan content is recorded separately: every ExitPlanMode submit snapshots + * the plan file into blob storage and persists a `plan.revision` record + * carrying only the reference (`{ id, version, path, sha256, bytes }`, `path` * homeDir-relative) — never the content. `revisionCount` tracks the latest * version per plan id so `recordRevision` can mint the next version * replay-consistently; it is kept across enter/exit so a re-entered plan id @@ -34,7 +35,10 @@ import { z } from 'zod'; -import { defineModel } from '#/wire/model'; +import { + defineCheckpointedModel, + type Checkpointed, +} from '#/agent/contextMemory/conversationTime'; export interface PlanState { readonly active: boolean; @@ -42,14 +46,16 @@ export interface PlanState { readonly revisionCount?: Readonly>; } -export const PlanModel = defineModel('plan', () => ({ active: false })); +export type PlanModelState = Checkpointed; + +export const PlanModel = defineCheckpointedModel('plan', (): PlanState => ({ active: false })); export const planModeEnter = PlanModel.defineOp('plan_mode.enter', { schema: z.object({ id: z.string() }), apply: (s, p) => - s.active && s.id === p.id + s.current.active && s.current.id === p.id ? s - : { active: true, id: p.id, revisionCount: s.revisionCount }, + : { ...s, current: { active: true, id: p.id, revisionCount: s.current.revisionCount } }, toEvent: () => ({ type: 'agent.status.updated' as const, planMode: true }), }); @@ -64,13 +70,19 @@ declare module '#/wire/types' { export const planModeCancel = PlanModel.defineOp('plan_mode.cancel', { schema: z.object({ id: z.string().optional() }), - apply: (s) => (s.active === false ? s : { active: false, revisionCount: s.revisionCount }), + apply: (s) => + s.current.active + ? { ...s, current: { active: false, revisionCount: s.current.revisionCount } } + : s, toEvent: () => ({ type: 'agent.status.updated' as const, planMode: false }), }); export const planModeExit = PlanModel.defineOp('plan_mode.exit', { schema: z.object({ id: z.string().optional() }), - apply: (s) => (s.active === false ? s : { active: false, revisionCount: s.revisionCount }), + apply: (s) => + s.current.active + ? { ...s, current: { active: false, revisionCount: s.current.revisionCount } } + : s, toEvent: () => ({ type: 'agent.status.updated' as const, planMode: false }), }); @@ -98,7 +110,10 @@ export const planRevision = PlanModel.defineOp('plan.revision', { }), apply: (s, p) => ({ ...s, - revisionCount: { ...s.revisionCount, [p.id]: p.version }, + current: { + ...s.current, + revisionCount: { ...s.current.revisionCount, [p.id]: p.version }, + }, }), toEvent: (p) => ({ type: 'plan.revision' as const, diff --git a/packages/agent-core-v2/src/agent/plan/planService.ts b/packages/agent-core-v2/src/agent/plan/planService.ts index a81cc6e3a2..c6690ed7c4 100644 --- a/packages/agent-core-v2/src/agent/plan/planService.ts +++ b/packages/agent-core-v2/src/agent/plan/planService.ts @@ -42,6 +42,7 @@ import type { BeforeToolExecuteEvent, ResolvedToolExecutionHookContext, } from '#/agent/toolExecutor/toolHooks'; +import { IEventBus } from '#/app/event/eventBus'; import { IAgentTelemetryContextService } from '#/app/telemetry/agentTelemetryContext'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import { IHostFileSystem } from '#/os/interface/hostFileSystem'; @@ -74,6 +75,7 @@ export class AgentPlanService extends Disposable implements IAgentPlanService { @IBlobStore private readonly blobs: IBlobStore, @IAgentContextInjectorService dynamicInjector: IAgentContextInjectorService, @IAgentTelemetryContextService private readonly telemetryContext: IAgentTelemetryContextService, + @IEventBus eventBus: IEventBus, @IWireService private readonly wire: IWireService, @ISessionContext private readonly sessionCtx: ISessionContext, @IAgentScopeContext private readonly agentCtx: IAgentScopeContext, @@ -93,6 +95,15 @@ export class AgentPlanService extends Disposable implements IAgentPlanService { await next(); }), ); + this._register( + eventBus.subscribe('context.undone', () => { + this.restoreTelemetryMode(); + eventBus.publish({ + type: 'agent.status.updated', + planMode: this.isActive, + }); + }), + ); this._register(new PlanModeInjection(dynamicInjector, this, this.context, states)); this._register(this.registerPlanGuard(toolExecutor)); @@ -152,19 +163,17 @@ export class AgentPlanService extends Disposable implements IAgentPlanService { } private get isActive(): boolean { - return this.wire.getModel(PlanModel).active; + return this.wire.getModel(PlanModel).current.active; } private currentPlanFilePath(): PlanFilePath { - const state = this.wire.getModel(PlanModel); + const state = this.wire.getModel(PlanModel).current; if (!state.active || state.id === undefined) return null; return this.planFilePathFor(state.id); } private restoreTelemetryMode(): void { - if (this.isActive) { - this.telemetryContext.set({ mode: 'plan' }); - } + this.telemetryContext.set({ mode: this.isActive ? 'plan' : 'agent' }); } private createPlanId(): string { @@ -211,7 +220,7 @@ export class AgentPlanService extends Disposable implements IAgentPlanService { } async recordRevision(): Promise { - const state = this.wire.getModel(PlanModel); + const state = this.wire.getModel(PlanModel).current; if (!state.active || state.id === undefined) return; const id = state.id; const content = await this.hostFs.readText(this.planFilePathFor(id)); @@ -232,7 +241,7 @@ export class AgentPlanService extends Disposable implements IAgentPlanService { } async status(): Promise { - const state = this.wire.getModel(PlanModel); + const state = this.wire.getModel(PlanModel).current; if (!state.active || state.id === undefined) return null; const path = this.planFilePathFor(state.id); let content = ''; diff --git a/packages/agent-core-v2/src/agent/prompt/prompt.ts b/packages/agent-core-v2/src/agent/prompt/prompt.ts index 8f9110a246..d5045dd025 100644 --- a/packages/agent-core-v2/src/agent/prompt/prompt.ts +++ b/packages/agent-core-v2/src/agent/prompt/prompt.ts @@ -55,7 +55,6 @@ export interface IAgentPromptService { abort(promptId: string, reason?: Error): boolean; inject(message: ContextMessage): Promise; retry(): Promise; - undo(count: number): number; clear(): void; readonly hooks: Hooks<{ onBeforeSubmitPrompt: PromptSubmitContext }>; } diff --git a/packages/agent-core-v2/src/agent/prompt/promptMetadataText.ts b/packages/agent-core-v2/src/agent/prompt/promptMetadataText.ts new file mode 100644 index 0000000000..e31a469823 --- /dev/null +++ b/packages/agent-core-v2/src/agent/prompt/promptMetadataText.ts @@ -0,0 +1,66 @@ +/** + * `prompt` domain (L4) — safe, displayable metadata text derived from prompts. + * + * Shared by prompt submission and undo projection so `lastPrompt` uses one + * normalization, redaction, and length limit, with image captions supplied by + * the `media` domain. + */ + +import type { ContentPart } from '#/kosong/contract/message'; +import { extractImageCompressionCaptions } from '#/agent/media/image-compress'; + +const MAX_TITLE_LENGTH = 200; +const MAX_LAST_PROMPT_LENGTH = 4000; + +export function titleFromPromptMetadataText(text: string): string { + return text.slice(0, MAX_TITLE_LENGTH); +} + +export function promptMetadataTextFromContentParts( + parts: readonly ContentPart[], +): string | undefined { + const texts: string[] = []; + for (const part of parts) { + const text = promptPartText(part); + if (text !== undefined) texts.push(text); + } + return promptMetadataTextFromText(texts.join('\n')); +} + +export function promptMetadataTextFromText(text: string): string | undefined { + const sanitized = text + .replaceAll( + /-----BEGIN [^-]*PRIVATE KEY-----[\s\S]*?-----END [^-]*PRIVATE KEY-----/gi, + '[redacted]', + ) + .replaceAll(/\b(authorization)\s*:\s*bearer\s+\S+/gi, '$1: Bearer [redacted]') + .replaceAll( + /\b(api[_-]?key|token|secret|password|passwd|pwd)\b\s*[:=]\s*(?:"[^"]*"|'[^']*'|\S+)/gi, + '$1=[redacted]', + ) + .replaceAll(/\bsk-[A-Za-z0-9_-]{12,}\b/g, '[redacted]') + .replaceAll(/\b[A-Za-z0-9][A-Za-z0-9+/=_-]{39,}\b/g, '[redacted]') + .replaceAll(/\p{Cc}+/gu, ' ') + .replaceAll(/\s+/g, ' ') + .trim(); + + if (sanitized.length === 0) return undefined; + return sanitized.slice(0, MAX_LAST_PROMPT_LENGTH); +} + +function promptPartText(part: ContentPart): string | undefined { + switch (part.type) { + case 'text': { + const { text } = extractImageCompressionCaptions(part.text); + return text.trim().length === 0 ? undefined : text; + } + case 'image_url': + return '[image]'; + case 'audio_url': + return '[audio]'; + case 'video_url': + return '[video]'; + case 'think': + return undefined; + } +} diff --git a/packages/agent-core-v2/src/agent/prompt/promptService.ts b/packages/agent-core-v2/src/agent/prompt/promptService.ts index 6d93a7f5b2..aeed006505 100644 --- a/packages/agent-core-v2/src/agent/prompt/promptService.ts +++ b/packages/agent-core-v2/src/agent/prompt/promptService.ts @@ -19,7 +19,6 @@ import { extractImageCompressionCaptions } from '#/agent/media/image-compress'; import { userCancellationReason } from '#/_base/utils/abort'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import { newMessageId } from '#/agent/contextMemory/messageId'; -import { formatUndoUnavailableMessage, precheckUndo } from '#/agent/contextMemory/contextOps'; import { USER_PROMPT_ORIGIN, type ContextMessage } from '#/agent/contextMemory/types'; import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction'; import { IAgentLoopService, type Turn, type TurnResult } from '#/agent/loop/loop'; @@ -174,13 +173,6 @@ export class AgentPromptService implements IAgentPromptService { async retry(): Promise { return (await this.loop.enqueue(new RetryStepRequest()).assigned).turn; } - undo(count: number): number { - if (count <= 0) return 0; - const check = precheckUndo(this.context.get(), count); - if (!check.ok) throw new Error2(ErrorCodes.SESSION_UNDO_UNAVAILABLE, formatUndoUnavailableMessage(check), { details: { reason: check.reason, requestedCount: count, undoableCount: check.undoable } }); - return this.context.undo(count).removedCount; - } - clear(): void { for (const item of this.pending.slice()) this.abort(item.id); if (this.active !== undefined) this.abort(this.active.id); @@ -246,8 +238,15 @@ export class AgentPromptService implements IAgentPromptService { return { message: captions.length === 0 ? message : { ...message, content: parts }, captions }; } private appendPrompt(message: ContextMessage, captions: readonly string[]): void { - for (const caption of captions) this.reminders.appendSystemReminder(caption, { kind: 'injection', variant: 'image_compression' }); - if (message.content.length > 0) this.context.append(message); + const ownerPromptId = message.id ?? newMessageId(); + for (const caption of captions) { + this.reminders.appendSystemReminder(caption, { + kind: 'injection', + variant: 'image_compression', + ownerPromptId, + }); + } + if (message.content.length > 0) this.context.append({ ...message, id: ownerPromptId }); } private async deliverToolResult(ctx: ToolDidExecuteContext): Promise { const delivery = ctx.result.delivery; if (delivery === undefined) return; diff --git a/packages/agent-core-v2/src/agent/prompt/promptStepRequests.ts b/packages/agent-core-v2/src/agent/prompt/promptStepRequests.ts index c4ec2ee3d8..40e204c13a 100644 --- a/packages/agent-core-v2/src/agent/prompt/promptStepRequests.ts +++ b/packages/agent-core-v2/src/agent/prompt/promptStepRequests.ts @@ -16,12 +16,14 @@ */ import { USER_PROMPT_ORIGIN, type ContextMessage } from '#/agent/contextMemory/types'; +import { newMessageId } from '#/agent/contextMemory/messageId'; import { StepRequest, type StepRequestOptions, type TurnSeed } from '#/agent/loop/stepRequest'; import { gateImageFormatParts } from '#/agent/media/image-compress'; import type { IAgentSystemReminderService } from '#/agent/systemReminder/systemReminder'; abstract class UserMessageStepRequest extends StepRequest { protected readonly message: ContextMessage; + private readonly ownerPromptId: string; constructor( message: ContextMessage, @@ -30,7 +32,12 @@ abstract class UserMessageStepRequest extends StepRequest { options?: StepRequestOptions, ) { super(options); - this.message = { ...message, content: gateImageFormatParts(message.content) }; + this.ownerPromptId = message.id ?? newMessageId(); + this.message = { + ...message, + id: this.ownerPromptId, + content: gateImageFormatParts(message.content), + }; } override get turnSeed(): TurnSeed { @@ -42,6 +49,7 @@ abstract class UserMessageStepRequest extends StepRequest { this.reminders.appendSystemReminder(caption, { kind: 'injection', variant: 'image_compression', + ownerPromptId: this.ownerPromptId, }); } } diff --git a/packages/agent-core-v2/src/agent/rpc/core-api.ts b/packages/agent-core-v2/src/agent/rpc/core-api.ts index 8b4261d3c5..0b6668471e 100644 --- a/packages/agent-core-v2/src/agent/rpc/core-api.ts +++ b/packages/agent-core-v2/src/agent/rpc/core-api.ts @@ -314,7 +314,7 @@ export interface AgentAPI { prompt: (payload: PromptPayload) => PromptLaunchResult | undefined; steer: (payload: SteerPayload) => PromptLaunchResult | undefined; cancel: (payload: CancelPayload) => void; - undoHistory: (payload: UndoHistoryPayload) => number; + undoHistory: (payload: UndoHistoryPayload) => Promise; setPermission: (payload: SetPermissionPayload) => void; cancelCompaction: (payload: EmptyPayload) => void; activateSkill: (payload: ActivateSkillPayload) => void; diff --git a/packages/agent-core-v2/src/agent/rpc/prompt-metadata.ts b/packages/agent-core-v2/src/agent/rpc/prompt-metadata.ts index 3e15261ced..0bf80ce6a0 100644 --- a/packages/agent-core-v2/src/agent/rpc/prompt-metadata.ts +++ b/packages/agent-core-v2/src/agent/rpc/prompt-metadata.ts @@ -7,11 +7,14 @@ * prompt adapter so both surfaces keep the same easy-title behavior. */ -import type { ContentPart } from '#/kosong/contract/message'; import type { IEventService } from '#/app/event/event'; import type { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata'; -import { extractImageCompressionCaptions } from '#/agent/media/image-compress'; +import { + promptMetadataTextFromContentParts, + promptMetadataTextFromText, + titleFromPromptMetadataText, +} from '#/agent/prompt/promptMetadataText'; import type { ActivatePluginCommandPayload, @@ -19,33 +22,16 @@ import type { PromptPayload, } from './core-api'; -const MAX_TITLE_LENGTH = 200; -const MAX_LAST_PROMPT_LENGTH = 4000; - -export function titleFromPromptMetadataText(text: string): string { - return text.slice(0, MAX_TITLE_LENGTH); -} +export { promptMetadataTextFromContentParts, titleFromPromptMetadataText }; export function promptMetadataTextFromPayload(payload: PromptPayload): string | undefined { return promptMetadataTextFromContentParts(payload.input); } -export function promptMetadataTextFromContentParts( - parts: readonly ContentPart[], -): string | undefined { - const texts: string[] = []; - for (const part of parts) { - const text = promptPartText(part); - if (text !== undefined) texts.push(text); - } - return sanitizeAndTruncatePromptText(texts.join('\n'), MAX_LAST_PROMPT_LENGTH); -} - export function promptMetadataTextFromSkill(payload: ActivateSkillPayload): string | undefined { const args = payload.args?.trim(); - return sanitizeAndTruncatePromptText( + return promptMetadataTextFromText( args === undefined || args.length === 0 ? `/${payload.name}` : `/${payload.name} ${args}`, - MAX_LAST_PROMPT_LENGTH, ); } @@ -54,9 +40,8 @@ export function promptMetadataTextFromPluginCommand( ): string | undefined { const args = payload.args?.trim(); const command = `/${payload.pluginId}:${payload.commandName}`; - return sanitizeAndTruncatePromptText( + return promptMetadataTextFromText( args === undefined || args.length === 0 ? command : `${command} ${args}`, - MAX_LAST_PROMPT_LENGTH, ); } @@ -98,41 +83,3 @@ export async function applyPromptMetadataUpdate( }, }); } - -function promptPartText(part: ContentPart): string | undefined { - switch (part.type) { - case 'text': { - const { text } = extractImageCompressionCaptions(part.text); - return text.trim().length === 0 ? undefined : text; - } - case 'image_url': - return '[image]'; - case 'audio_url': - return '[audio]'; - case 'video_url': - return '[video]'; - case 'think': - return undefined; - } -} - -function sanitizeAndTruncatePromptText(text: string, maxLength: number): string | undefined { - const sanitized = text - .replaceAll( - /-----BEGIN [^-]*PRIVATE KEY-----[\s\S]*?-----END [^-]*PRIVATE KEY-----/gi, - '[redacted]', - ) - .replaceAll(/\b(authorization)\s*:\s*bearer\s+\S+/gi, '$1: Bearer [redacted]') - .replaceAll( - /\b(api[_-]?key|token|secret|password|passwd|pwd)\b\s*[:=]\s*(?:"[^"]*"|'[^']*'|\S+)/gi, - '$1=[redacted]', - ) - .replaceAll(/\bsk-[A-Za-z0-9_-]{12,}\b/g, '[redacted]') - .replaceAll(/\b[A-Za-z0-9][A-Za-z0-9+/=_-]{39,}\b/g, '[redacted]') - .replaceAll(/\p{Cc}+/gu, ' ') - .replaceAll(/\s+/g, ' ') - .trim(); - - if (sanitized.length === 0) return undefined; - return sanitized.slice(0, maxLength); -} diff --git a/packages/agent-core-v2/src/agent/rpc/rpcService.ts b/packages/agent-core-v2/src/agent/rpc/rpcService.ts index 75b78cab5f..bb7ffef1d4 100644 --- a/packages/agent-core-v2/src/agent/rpc/rpcService.ts +++ b/packages/agent-core-v2/src/agent/rpc/rpcService.ts @@ -18,6 +18,7 @@ import { IPluginService } from '#/app/plugin/plugin'; import { ProfileError } from '#/agent/profile/profile'; import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy'; import { IAgentPromptService } from '#/agent/prompt/prompt'; +import { IAgentConversationUndoService } from '#/agent/undo/undo'; import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata'; import { ISessionContext } from '#/session/sessionContext/sessionContext'; import { IAgentSkillService } from '#/agent/skill/skill'; @@ -63,6 +64,8 @@ export class AgentRPCService implements IAgentRPCService { constructor( @IAgentPromptService private readonly promptService: IAgentPromptService, + @IAgentConversationUndoService + private readonly conversationUndo: IAgentConversationUndoService, @IAgentLoopService private readonly loop: IAgentLoopService, @IAgentToolPolicyService private readonly toolPolicy: IAgentToolPolicyService, @IAgentPermissionModeService private readonly permissionMode: IAgentPermissionModeService, @@ -126,10 +129,8 @@ export class AgentRPCService implements IAgentRPCService { this.loop.cancel(turnId); } - undoHistory(payload: UndoHistoryPayload): number { - const undone = this.promptService.undo(payload.count); - this.telemetry.track2('conversation_undo', { count: payload.count }); - return undone; + async undoHistory(payload: UndoHistoryPayload): Promise { + return this.conversationUndo.undo(payload.count); } setPermission(payload: SetPermissionPayload): void { diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index d7a2fc5d8e..cc454d3dec 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -31,8 +31,11 @@ * into `agentState` (`IAgentStateService`) and read/written through it; the * live `tasks` registry stays a plain field because a `ManagedTask` holds * resources (promise chains, an `AbortController`, task handles) that must - * not be snapshotted, as does the `persistence` construction-time helper. - * Bound at Agent scope. + * not be snapshotted, as do the `persistence` construction-time helper and + * the notification delivery machinery (`buildingNotificationKeys`, + * `pendingNotificationRequests`, `notificationRestoreQueue`). + * Notification delivery follows conversation undo through the checkpoint and + * reconciliation contracts. Bound at Agent scope. */ import { randomBytes } from 'node:crypto'; @@ -43,6 +46,7 @@ import { LifecycleScope, ScopeActivation, registerScopedService } from '#/_base/ import type { ContentPart } from '#/kosong/contract/message'; import { Disposable } from '#/_base/di/lifecycle'; +import { ILogService } from '#/_base/log/log'; import { defineState } from '#/_base/state/stateRegistry'; import { abortable, @@ -50,6 +54,8 @@ import { } from '#/_base/utils/abort'; import { escapeXml, escapeXmlAttr } from '#/_base/utils/xml-escape'; import { IEventBus } from '#/app/event/eventBus'; +import { defineCheckpointedModel } from '#/agent/contextMemory/conversationTime'; +import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants'; import type { ContextMessage, TaskOrigin } from '#/agent/contextMemory/types'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { IAgentLoopService } from '#/agent/loop/loop'; @@ -70,7 +76,6 @@ import { ISessionContext } from '#/session/sessionContext/sessionContext'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; import { IFileSystemStorageService } from '#/persistence/interface/storage'; import { ITelemetryService } from '#/app/telemetry/telemetry'; -import { defineModel } from '#/wire/model'; import { IWireService } from '#/wire/wire'; import { IAgentTaskService, @@ -116,17 +121,15 @@ interface AgentTaskNotificationBuildContext { readonly notification: AgentTaskNotification; } -const TaskNotificationDeliveryModel = defineModel( +const TaskNotificationDeliveryModel = defineCheckpointedModel( 'task.notificationDelivery', - () => [], + (): readonly string[] => [], { - reducers: { - 'context.append_message': (state, payload: { message?: unknown }) => { - const origin = taskOriginFromMessage(payload.message); - if (origin === undefined) return state; - const key = notificationKey(origin); - return state.includes(key) ? state : [...state, key]; - }, + onAppendMessage: (current, message) => { + const origin = taskOriginFromMessage(message); + if (origin === undefined) return current; + const key = notificationKey(origin); + return current.includes(key) ? current : [...current, key]; }, }, ); @@ -210,7 +213,10 @@ declare module '#/app/event/eventBus' { } export class TaskNotificationStepRequest extends MessageStepRequest { - constructor(message: ContextMessage) { + constructor( + message: ContextMessage, + private readonly onWillDeliver?: () => void, + ) { super(message, { kind: 'task_notification', mergeable: true, @@ -218,6 +224,10 @@ export class TaskNotificationStepRequest extends MessageStepRequest { admission: 'activeOrNewTurn', }); } + + override onWillMaterialize(): void { + this.onWillDeliver?.(); + } } export const taskGhostsKey = defineState>( @@ -241,7 +251,10 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { declare readonly _serviceBrand: undefined; private readonly tasks = new Map(); + private readonly buildingNotificationKeys = new Set(); + private readonly pendingNotificationRequests = new Map(); private readonly persistence: AgentTaskPersistence; + private notificationRestoreQueue: Promise = Promise.resolve(); constructor( @ITelemetryService private readonly telemetry: ITelemetryService, @@ -256,6 +269,9 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { @IEventBus private readonly eventBus: IEventBus, @IAgentContextInjectorService injector: IAgentContextInjectorService, @IAgentLoopService private readonly loop: IAgentLoopService, + @IAgentConversationUndoParticipantRegistry + undoParticipants: IAgentConversationUndoParticipantRegistry, + @ILogService private readonly log: ILogService, @IAgentStateService private readonly states: IAgentStateService, ) { super(); @@ -274,9 +290,15 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { byteStore, fallbackRoot, ); + this._register( + undoParticipants.register({ + id: 'task.notificationDelivery', + reconcileAfterUndo: () => this.reconcileNotificationDeliveryAfterUndo(), + }), + ); this._register( this.wire.hooks.onDidRestore.register('task', async (_ctx, next) => { - for (const key of this.wire.getModel(TaskNotificationDeliveryModel)) { + for (const key of this.wire.getModel(TaskNotificationDeliveryModel).current) { this.deliveredNotificationKeys.add(key); } await this.restoreAfterReplay(); @@ -519,6 +541,21 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { return result; } + private async reconcileNotificationDeliveryAfterUndo(): Promise { + const restoredKeys = new Set(this.wire.getModel(TaskNotificationDeliveryModel).current); + for (const [key, request] of this.pendingNotificationRequests) { + if (request.aborted) this.clearPendingNotification(key, request); + } + this.deliveredNotificationKeys.clear(); + for (const key of restoredKeys) this.deliveredNotificationKeys.add(key); + for (const key of this.scheduledNotificationKeys) { + if (restoredKeys.has(key) || !this.pendingNotificationRequests.has(key)) { + this.scheduledNotificationKeys.delete(key); + } + } + await this.restoreAgentTaskNotifications(); + } + persistOutput(taskId: string): void { const entry = this.tasks.get(taskId); if (entry === undefined) return; @@ -1025,7 +1062,9 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { if (!this.isDetached(entry)) return; entry.terminalFired = true; const info = this.toInfo(entry); - void this.notifyAgentTask(info).catch(() => { }); + void this.notifyAgentTask(info).catch((error) => { + this.log.error('task notification delivery failed', { taskId: info.taskId, error }); + }); this.recordTaskTerminated(info, this.retainedOutputTail(entry)); } @@ -1057,17 +1096,42 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { private async notifyAgentTask(info: AgentTaskInfo): Promise { const context = await this.buildAgentTaskNotificationContext(info); if (context === undefined) return; - const request = new TaskNotificationStepRequest({ - role: 'user', - content: [...context.content], - toolCalls: [], - origin: context.origin, - }); - this.loop.enqueue(request); - this.fireNotificationHook(context.notification); + const key = notificationKey(context.origin); + const request = new TaskNotificationStepRequest( + { + role: 'user', + content: [...context.content], + toolCalls: [], + origin: context.origin, + }, + () => this.fireNotificationHook(context.notification), + ); + this.pendingNotificationRequests.set(key, request); + try { + const receipt = this.loop.enqueue(request); + void receipt.assigned + .then(({ step }) => step.result) + .then( + () => { + if (request.aborted) this.clearPendingNotification(key, request); + }, + () => this.clearPendingNotification(key, request), + ); + } catch (error) { + this.clearPendingNotification(key, request); + throw error; + } } - private async restoreAgentTaskNotifications(): Promise { + private restoreAgentTaskNotifications(): Promise { + const restore = this.notificationRestoreQueue.then(() => + this.restoreAgentTaskNotificationsNow(), + ); + this.notificationRestoreQueue = restore.catch(() => {}); + return restore; + } + + private async restoreAgentTaskNotificationsNow(): Promise { for (const info of this.list(false)) { if (!isAgentTaskTerminal(info.status)) continue; await this.restoreAgentTaskNotification(info); @@ -1098,35 +1162,51 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { notificationId: `task:${info.taskId}:${info.status}`, }; const key = notificationKey(origin); + if (this.buildingNotificationKeys.has(key)) return undefined; if (this.scheduledNotificationKeys.has(key)) return undefined; if (this.deliveredNotificationKeys.has(key)) return undefined; if (this.hasDeliveredNotification(key)) return undefined; - this.scheduledNotificationKeys.add(key); - - let output = await this.getOutputSnapshot(info.taskId, 0); - if (!output.fullOutputAvailable) { - output = await this.getOutputSnapshot(info.taskId, NOTIFICATION_FALLBACK_PREVIEW_BYTES); + this.buildingNotificationKeys.add(key); + try { + let output = emptyOutputSnapshot(); + try { + output = await this.getOutputSnapshot(info.taskId, 0); + if (!output.fullOutputAvailable) { + output = await this.getOutputSnapshot(info.taskId, NOTIFICATION_FALLBACK_PREVIEW_BYTES); + } + } catch (error) { + this.log.error('task notification output read failed; delivering without output', { + taskId: info.taskId, + error, + }); + } + if (this.isTerminalNotificationSuppressed(info.taskId)) return undefined; + if (this.scheduledNotificationKeys.has(key)) return undefined; + if (this.deliveredNotificationKeys.has(key)) return undefined; + if (this.hasDeliveredNotification(key)) return undefined; + this.scheduledNotificationKeys.add(key); + const notification: AgentTaskNotification = { + id: origin.notificationId, + category: 'task', + type: `task.${info.status}`, + source_kind: 'background_task', + source_id: info.taskId, + agent_id: info.kind === 'agent' ? info.agentId : undefined, + title: `Background ${info.kind} ${info.status}`, + severity: info.status === 'completed' ? 'info' : 'warning', + body: buildAgentTaskNotificationBody(info), + children: agentTaskNotificationChildren(output), + }; + const content = [ + { + type: 'text', + text: renderNotificationXml(notification), + }, + ] as const; + return { content, origin, notification }; + } finally { + this.buildingNotificationKeys.delete(key); } - if (this.isTerminalNotificationSuppressed(info.taskId)) return undefined; - const notification: AgentTaskNotification = { - id: origin.notificationId, - category: 'task', - type: `task.${info.status}`, - source_kind: 'background_task', - source_id: info.taskId, - agent_id: info.kind === 'agent' ? info.agentId : undefined, - title: `Background ${info.kind} ${info.status}`, - severity: info.status === 'completed' ? 'info' : 'warning', - body: buildAgentTaskNotificationBody(info), - children: agentTaskNotificationChildren(output), - }; - const content = [ - { - type: 'text', - text: renderNotificationXml(notification), - }, - ] as const; - return { content, origin, notification }; } private fireNotificationHook(notification: AgentTaskNotification): void { @@ -1149,7 +1229,18 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { } private markDeliveredNotification(origin: TaskNotificationOrigin): void { - this.deliveredNotificationKeys.add(notificationKey(origin)); + const key = notificationKey(origin); + this.scheduledNotificationKeys.delete(key); + this.pendingNotificationRequests.delete(key); + this.deliveredNotificationKeys.add(key); + } + + private clearPendingNotification(key: string, request: TaskNotificationStepRequest): void { + if (this.pendingNotificationRequests.get(key) !== request) return; + this.pendingNotificationRequests.delete(key); + if (!this.deliveredNotificationKeys.has(key) && !this.hasDeliveredNotification(key)) { + this.scheduledNotificationKeys.delete(key); + } } private hasDeliveredNotification(key: string): boolean { diff --git a/packages/agent-core-v2/src/agent/toolSelect/toolSelectService.ts b/packages/agent-core-v2/src/agent/toolSelect/toolSelectService.ts index e6422c3ccb..ce22e42b26 100644 --- a/packages/agent-core-v2/src/agent/toolSelect/toolSelectService.ts +++ b/packages/agent-core-v2/src/agent/toolSelect/toolSelectService.ts @@ -75,10 +75,7 @@ export class AgentToolSelectService extends Disposable implements IAgentToolSele this._register( eventBus.subscribe('context.spliced', (splice) => { if (splice.deleteCount === 0 || this.pendingLoaded.size === 0) return; - const landed = collectLoadedDynamicToolNames(this.context.get()); - for (const name of this.pendingLoaded) { - if (!landed.has(name)) this.pendingLoaded.delete(name); - } + this.dropPendingLoadedNotLanded(); }), ); } @@ -87,6 +84,14 @@ export class AgentToolSelectService extends Disposable implements IAgentToolSele return this.states.get(toolSelectPendingLoadedKey); } + private dropPendingLoadedNotLanded(): void { + if (this.pendingLoaded.size === 0) return; + const landed = collectLoadedDynamicToolNames(this.context.get()); + for (const name of this.pendingLoaded) { + if (!landed.has(name)) this.pendingLoaded.delete(name); + } + } + enabled(): boolean { const capabilities = this.profile.getModelCapabilities(); return ( diff --git a/packages/agent-core-v2/src/agent/undo/undo.ts b/packages/agent-core-v2/src/agent/undo/undo.ts new file mode 100644 index 0000000000..9ccec70307 --- /dev/null +++ b/packages/agent-core-v2/src/agent/undo/undo.ts @@ -0,0 +1,23 @@ +/** + * `undo` domain (L6) — Agent-scoped conversation undo contract. + * + * Defines the availability and idle-only execution surface shared by every + * undo entry point. Bound at Agent scope. + */ + +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; + +export interface UndoAvailability { + readonly maxTurns: number; + readonly stoppedAtCompaction: boolean; +} + +export interface IAgentConversationUndoService { + readonly _serviceBrand: undefined; + + availability(): UndoAvailability; + undo(turns: number): Promise; +} + +export const IAgentConversationUndoService: ServiceIdentifier = + createDecorator('agentConversationUndoService'); diff --git a/packages/agent-core-v2/src/agent/undo/undoService.ts b/packages/agent-core-v2/src/agent/undo/undoService.ts new file mode 100644 index 0000000000..9441dec57c --- /dev/null +++ b/packages/agent-core-v2/src/agent/undo/undoService.ts @@ -0,0 +1,248 @@ +/** + * `undo` domain (L6) — `IAgentConversationUndoService` implementation. + * + * Owns idle conversation undo coordination and restored observable state. + * Coordinates `contextMemory`, undo participants, `fullCompaction`, + * `loop`, `prompt`, Agent and Session identity, `sessionMetadata`, `event`, + * `eventBus`, `telemetry`, and `wire`. Bound at Agent scope. + */ + +import { Disposable, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { ILogService } from '#/_base/log/log'; +import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; +import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants'; +import { + computeUndoCut, + formatUndoUnavailableMessage, + precheckUndo, +} from '#/agent/contextMemory/contextOps'; +import { + CHECKPOINTED_MODELS, + isUndoAnchor, + isValidUndoCount, + type Checkpointed, +} from '#/agent/contextMemory/conversationTime'; +import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction'; +import { IAgentLoopService } from '#/agent/loop/loop'; +import { IAgentPromptService } from '#/agent/prompt/prompt'; +import { promptMetadataTextFromContentParts } from '#/agent/prompt/promptMetadataText'; +import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; +import { IEventService } from '#/app/event/event'; +import { IEventBus } from '#/app/event/eventBus'; +import { ITelemetryService } from '#/app/telemetry/telemetry'; +import { ErrorCodes, Error2 } from '#/errors'; +import { MAIN_AGENT_ID } from '#/session/agentLifecycle/agentLifecycle'; +import { ISessionContext } from '#/session/sessionContext/sessionContext'; +import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata'; +import { IWireService } from '#/wire/wire'; + +import { IAgentConversationUndoService, type UndoAvailability } from './undo'; + +declare module '#/app/event/eventBus' { + interface DomainEventMap { + 'context.undone': { turns: number }; + } +} + +export class AgentConversationUndoService + extends Disposable + implements IAgentConversationUndoService +{ + declare readonly _serviceBrand: undefined; + + private undoQueue: Promise = Promise.resolve(); + + constructor( + @IAgentLoopService private readonly loop: IAgentLoopService, + @IAgentFullCompactionService private readonly fullCompaction: IAgentFullCompactionService, + @IAgentPromptService private readonly prompt: IAgentPromptService, + @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, + @IAgentConversationUndoParticipantRegistry + private readonly participants: IAgentConversationUndoParticipantRegistry, + @IAgentScopeContext private readonly agentCtx: IAgentScopeContext, + @ISessionContext private readonly session: ISessionContext, + @ISessionMetadata private readonly metadata: ISessionMetadata, + @IEventService private readonly eventService: IEventService, + @IEventBus private readonly eventBus: IEventBus, + @ITelemetryService private readonly telemetry: ITelemetryService, + @IWireService private readonly wire: IWireService, + @ILogService private readonly log: ILogService, + ) { + super(); + } + + availability(): UndoAvailability { + const cut = computeUndoCut(this.context.get(), Number.MAX_SAFE_INTEGER); + const maxTurns = Math.min(cut.removedCount, this.checkpointDepth().depth); + return { + maxTurns, + stoppedAtCompaction: cut.stoppedAtCompaction || maxTurns < cut.removedCount, + }; + } + + async undo(turns: number): Promise { + if (!isValidUndoCount(turns)) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + 'Undo count must be a positive safe integer', + { details: { field: 'count' } }, + ); + } + const run = this.undoQueue.then(() => this.undoNow(turns)); + this.undoQueue = run.then( + () => undefined, + () => undefined, + ); + return run; + } + + private async undoNow(turns: number): Promise { + let quiescence: IDisposable | undefined; + try { + quiescence = this.loop.tryAcquireQuiescence(); + if (quiescence === undefined) { + throw this.busyError('loop'); + } + if (this.fullCompaction.compacting !== null) { + throw this.busyError('compaction'); + } + this.assertUndoAvailable(turns); + this.context.undo(turns); + await this.flushAfterCommit('context cut'); + await this.reconcileParticipants(); + await this.flushAfterCommit('state reconciliation'); + await this.reconcileLastPromptSafely(); + this.telemetry.track2('conversation_undo', { count: turns }); + this.eventBus.publish({ type: 'context.undone', turns }); + return turns; + } finally { + quiescence?.dispose(); + } + } + + private checkpointDepth(): { depth: number; model: string } { + let depth = Number.POSITIVE_INFINITY; + let model = ''; + for (const def of CHECKPOINTED_MODELS) { + const state = this.wire.getModel(def) as Checkpointed; + if (state.checkpoints.length < depth) { + depth = state.checkpoints.length; + model = def.name; + } + } + return { depth, model }; + } + + private busyError(reason: 'loop' | 'compaction'): Error2 { + const message = reason === 'loop' + ? 'Cannot undo while a turn is active or queued. Wait for it to finish, then retry.' + : 'Cannot undo while conversation compaction is running. Wait for it to finish, then retry.'; + return new Error2(ErrorCodes.SESSION_BUSY, message, { details: { reason } }); + } + + private assertUndoAvailable(turns: number): void { + const check = precheckUndo(this.context.get(), turns); + if (!check.ok) { + throw new Error2( + ErrorCodes.SESSION_UNDO_UNAVAILABLE, + formatUndoUnavailableMessage(check), + { + details: { + reason: check.reason, + requestedCount: check.requested, + undoableCount: check.undoable, + }, + }, + ); + } + const { depth, model } = this.checkpointDepth(); + if (depth >= turns) return; + // A compaction explains missing checkpoints (they are cleared at the + // boundary); without one, a checkpointed model failed to track an anchor. + const fullCut = computeUndoCut(this.context.get(), Number.MAX_SAFE_INTEGER); + const reason = fullCut.stoppedAtCompaction ? 'compaction_boundary' : 'checkpoint_lost'; + throw new Error2( + ErrorCodes.SESSION_UNDO_UNAVAILABLE, + formatUndoUnavailableMessage({ + ok: false, + reason, + requested: turns, + undoable: depth, + }), + { + details: { + reason, + requestedCount: turns, + undoableCount: depth, + model, + }, + }, + ); + } + + private async reconcileParticipants(): Promise { + const participants = this.participants.list(); + const results = await Promise.allSettled( + participants.map((participant) => participant.reconcileAfterUndo()), + ); + results.forEach((result, index) => { + if (result.status === 'fulfilled') return; + this.log.error('undo participant reconciliation failed', { + participantId: participants[index]?.id, + error: result.reason, + }); + }); + } + + private async reconcileLastPromptSafely(): Promise { + try { + await this.reconcileLastPrompt(); + } catch (error) { + this.log.error('undo lastPrompt reconciliation failed', { error }); + } + } + + private async flushAfterCommit(stage: string): Promise { + try { + await this.wire.flush(); + } catch (error) { + this.log.error('undo wire flush failed after in-memory commit', { stage, error }); + throw error; + } + } + + private async reconcileLastPrompt(): Promise { + if (this.agentCtx.agentId !== MAIN_AGENT_ID) return; + const pending = this.prompt.list().pending.at(-1); + let lastPrompt = pending === undefined + ? undefined + : promptMetadataTextFromContentParts(pending.message.content); + if (lastPrompt === undefined) { + const history = this.context.get(); + for (let i = history.length - 1; i >= 0; i--) { + const message = history[i]!; + if (!isUndoAnchor(message)) continue; + lastPrompt = promptMetadataTextFromContentParts(message.content); + if (lastPrompt !== undefined) break; + } + } + await this.metadata.update({ lastPrompt }); + this.eventService.publish({ + type: 'session.meta.updated', + payload: { + agentId: MAIN_AGENT_ID, + sessionId: this.session.sessionId, + patch: { lastPrompt }, + }, + }); + } +} + +registerScopedService( + LifecycleScope.Agent, + IAgentConversationUndoService, + AgentConversationUndoService, + ScopeActivation.OnScopeCreated, + 'undo', +); diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 023fe5aba8..c78cd2f213 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -453,6 +453,8 @@ export * from '#/agent/contextMemory/contextMemory'; export * from '#/agent/contextMemory/contextMemoryService'; export * from '#/agent/contextMemory/contextOps'; export * from '#/agent/contextMemory/compactionHandoff'; +export * from '#/agent/contextMemory/conversationUndoParticipants'; +export * from '#/agent/contextMemory/conversationTime'; export * from '#/agent/contextMemory/loopEventFold'; export * from '#/agent/contextMemory/messageId'; export * from '#/agent/contextMemory/messageProjection'; @@ -520,6 +522,8 @@ import '#/app/messageLegacy/errors'; export * from '#/app/messageLegacy/messageLegacy'; export * from '#/app/messageLegacy/messageLegacyService'; export * from '#/agent/replayBuilder/types'; +export * from '#/agent/undo/undo'; +export * from '#/agent/undo/undoService'; export * from '#/agent/shellCommand/shellCommand'; export * from '#/agent/shellCommand/shellCommandService'; export * from '#/agent/rpc/rpc'; diff --git a/packages/agent-core-v2/src/session/todo/sessionTodoService.ts b/packages/agent-core-v2/src/session/todo/sessionTodoService.ts index 0752bfae06..d05a5a523c 100644 --- a/packages/agent-core-v2/src/session/todo/sessionTodoService.ts +++ b/packages/agent-core-v2/src/session/todo/sessionTodoService.ts @@ -1,20 +1,12 @@ /** * `todo` domain (L4) — `ISessionTodoService` implementation. * - * Holds the session's shared todo list as a stateless facade over the main - * agent's `TodoModel`: `getTodos` reads `wire.getModel(TodoModel)` live, and - * every mutation only dispatches a `tools.update_store` Op to the main agent's - * wire (the single source of truth and replayable timeline), then emits - * `onDidChange` from the rebuilt Model. The service keeps no list copy of its - * own, so the live view and the post-replay view can never drift. Binds the - * `TodoListTool` and the stale-todo reminder into every agent (`onDidCreate`), - * borrowing each agent's services through its `IAgentScopeHandle.accessor`. - * Per-agent bindings are disposed when the agent is disposed. Bound at + * Provides session-wide todo access through the main agent's `wire`, binds + * todo capabilities into each agent, and publishes changes through its typed + * event. The main agent's wire owns the replayable state (including the + * undo-checkpointed `TodoModel`); this facade keeps no list copy of its own + * and there is deliberately no second session-level wire aggregate. Bound at * Session scope. - * - * The session owns the todo facade and tool bindings, while the main Agent wire - * owns the replayable state. This is an explicit cross-scope orchestration - * boundary: there is no second session-level wire aggregate or journal. */ import { Disposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle'; @@ -29,6 +21,7 @@ import { Emitter } from '#/_base/event'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy'; +import { IEventBus } from '#/app/event/eventBus'; import { IAgentLifecycleService } from '#/session/agentLifecycle/agentLifecycle'; import { IWireService } from '#/wire/wire'; @@ -46,6 +39,7 @@ export class SessionTodoService extends Disposable implements ISessionTodoServic readonly onDidChange = this.onDidChangeEmitter.event; private readonly agentBindings = new Map(); + private lastKnownTodos: readonly TodoItem[] = []; constructor( @IAgentLifecycleService private readonly agentLifecycle: IAgentLifecycleService, @@ -77,7 +71,7 @@ export class SessionTodoService extends Disposable implements ISessionTodoServic getTodos(): readonly TodoItem[] { const main = this.agentLifecycle.get(MAIN_AGENT_ID); if (main === undefined) return []; - return main.accessor.get(IWireService).getModel(TodoModel); + return main.accessor.get(IWireService).getModel(TodoModel).current; } setTodos(todos: readonly TodoItem[]): void { @@ -97,7 +91,9 @@ export class SessionTodoService extends Disposable implements ISessionTodoServic if (main === undefined) return; const wire = main.accessor.get(IWireService); wire.dispatch(todoSet({ key: 'todo', value: todos })); - this.onDidChangeEmitter.fire(wire.getModel(TodoModel)); + const current = wire.getModel(TodoModel).current; + this.lastKnownTodos = current; + this.onDidChangeEmitter.fire(current); } private bindAgent(handle: IAgentScopeHandle): void { @@ -106,6 +102,18 @@ export class SessionTodoService extends Disposable implements ISessionTodoServic handle.id, injector.register(TODO_LIST_REMINDER_VARIANT, () => this.staleReminder(handle)), ); + if (handle.id !== MAIN_AGENT_ID) return; + + this.lastKnownTodos = handle.accessor.get(IWireService).getModel(TodoModel).current; + this.trackAgentBinding( + handle.id, + handle.accessor.get(IEventBus).subscribe('context.undone', () => { + const current = handle.accessor.get(IWireService).getModel(TodoModel).current; + if (todoItemsEqual(current, this.lastKnownTodos)) return; + this.lastKnownTodos = current; + this.onDidChangeEmitter.fire(current); + }), + ); } private staleReminder(handle: IAgentScopeHandle): string | undefined { @@ -134,9 +142,17 @@ export class SessionTodoService extends Disposable implements ISessionTodoServic disposable.dispose(); } this.agentBindings.delete(agentId); + if (agentId === MAIN_AGENT_ID) this.lastKnownTodos = []; } } +function todoItemsEqual(a: readonly TodoItem[], b: readonly TodoItem[]): boolean { + return ( + a.length === b.length && + a.every((item, index) => item.title === b[index]?.title && item.status === b[index]?.status) + ); +} + registerScopedService( LifecycleScope.Session, ISessionTodoService, diff --git a/packages/agent-core-v2/src/session/todo/todoOps.ts b/packages/agent-core-v2/src/session/todo/todoOps.ts index 29a931aa4b..cf036d029d 100644 --- a/packages/agent-core-v2/src/session/todo/todoOps.ts +++ b/packages/agent-core-v2/src/session/todo/todoOps.ts @@ -1,31 +1,23 @@ /** - * `todo` domain (L4) — wire Model (`TodoModel`) and the `tools.update_store` - * Op (`todoSet`) for the session's shared todo list. + * `todo` domain (L4) — persists the session's shared todo document. * - * Declares the todo list as `readonly TodoItem[]` (initial `[]`). The - * persisted record is v1's `tools.update_store` (`{ key: 'todo', value }`), so - * the on-disk vocabulary stays exactly v1's and `wire.replay` — of both v2 and - * v1 sessions — rebuilds the Model from the shared append log. `apply` is the - * single log→model boundary: it ignores non-`todo` keys and sanitizes the - * value through `readTodoItems`, so every consumer (`getTodos`, the tool - * render, the stale reminder, the compaction summary) can trust the Model - * without re-validating. Consumed cross-scope by the Session-scope - * `SessionTodoService`: it dispatches to the MAIN agent's wire (the single - * source of truth and replayable timeline), and `getTodos` reads the rebuilt - * Model back from that same wire after restore. The Ops register into the - * global `OP_REGISTRY` at import time, so they are in place before the main - * agent restores. + * Validates todo state through the local item contract, keeps it aligned with + * conversation undo through `contextMemory`, and serves the Session-scope todo + * facade from the main agent's wire. */ import { z } from 'zod'; -import { defineModel } from '#/wire/model'; +import { + defineCheckpointedModel, + type Checkpointed, +} from '#/agent/contextMemory/conversationTime'; import { readTodoItems, type TodoItem } from './todoItem'; -export type TodoModelState = readonly TodoItem[]; +export type TodoModelState = Checkpointed; -export const TodoModel = defineModel('todo', () => []); +export const TodoModel = defineCheckpointedModel('todo', (): readonly TodoItem[] => []); declare module '#/wire/types' { interface PersistedOpMap { @@ -35,5 +27,6 @@ declare module '#/wire/types' { export const todoSet = TodoModel.defineOp('tools.update_store', { schema: z.object({ key: z.string(), value: z.unknown() }), - apply: (s, p) => (p.key === 'todo' ? readTodoItems(p.value) : s), + apply: (s, p) => + p.key === 'todo' ? { ...s, current: readTodoItems(p.value) } : s, }); diff --git a/packages/agent-core-v2/test/agent/contextMemory/context.test.ts b/packages/agent-core-v2/test/agent/contextMemory/context.test.ts index 2189dac79b..3178f59ebe 100644 --- a/packages/agent-core-v2/test/agent/contextMemory/context.test.ts +++ b/packages/agent-core-v2/test/agent/contextMemory/context.test.ts @@ -2,6 +2,7 @@ import type { Message } from '#/kosong/contract/message'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { estimateTokensForMessages } from '#/kosong/contract/tokens'; +import { buildImageCompressionCaption } from '#/agent/media/image-compress'; import type { ContextMessage } from '#/agent/contextMemory/types'; import { IWireService } from '#/wire/wire'; import { @@ -579,12 +580,12 @@ describe('Agent context', () => { expect(contextSize.get()).toEqual({ size: 0, measured: 0, estimated: 0 }); }); - it('rebases the measured prefix to an estimate when undo truncates it', () => { - ctx.appendAssistantTextWithUsage(1, 'a1', 1_000); - ctx.appendAssistantTextWithUsage(2, 'a2', 2_000); + it('rebases the measured prefix to an estimate when undo truncates it', async () => { + ctx.appendTurnExchange('u1', 'a1', 1_000); + ctx.appendTurnExchange('u2', 'a2', 2_000); expect(contextSize.get().measured).toBe(2_000); - ctx.undoHistory(1); + await ctx.undoHistory(1); const surviving = context.get(); expect(surviving.map((m) => m.role)).toEqual(['user', 'assistant']); @@ -592,20 +593,20 @@ describe('Agent context', () => { expect(contextSize.get()).toEqual({ size: estimate, measured: estimate, estimated: 0 }); }); - it('keeps the measured prefix when undo removes only the unmeasured tail', () => { - ctx.appendAssistantTextWithUsage(1, 'a1', 1_000); - ctx.appendUserMessage([{ type: 'text', text: 'unmeasured follow up' }]); + it('keeps the measured prefix when undo removes only the unmeasured tail', async () => { + ctx.appendTurnExchange('u1', 'a1', 1_000); + ctx.appendTurnExchange('u2', 'a2'); expect(contextSize.get().measured).toBe(1_000); - ctx.undoHistory(1); + await ctx.undoHistory(1); expect(context.get().map((m) => m.role)).toEqual(['user', 'assistant']); expect(contextSize.get()).toEqual({ size: 1_000, measured: 1_000, estimated: 0 }); }); - it('undo only counts real user prompts, skipping task notifications', () => { - ctx.appendAssistantText(1, 'first response'); - ctx.appendAssistantText(2, 'second response'); + it('undo only counts real user prompts, skipping task notifications', async () => { + ctx.appendTurnExchange('u1', 'first response'); + ctx.appendTurnExchange('u2', 'second response'); context.append( { @@ -629,14 +630,14 @@ describe('Agent context', () => { 'user', ]); - ctx.undoHistory(1); + await ctx.undoHistory(1); expect(context.get().map((m) => m.role)).toEqual(['user', 'assistant']); }); - it('removes injection messages inside the undone turn', () => { - context.append(userMessage('earlier question', { kind: 'user' })); - context.append(userMessage('do the work', { kind: 'user' })); + it('removes injection messages inside the undone turn', async () => { + ctx.appendUserTurn('earlier question'); + ctx.appendUserTurn('do the work'); context.append( userMessage('Plan mode is active', { kind: 'injection', @@ -652,7 +653,7 @@ describe('Agent context', () => { }, ); - ctx.undoHistory(1); + await ctx.undoHistory(1); expect(context.get()).toEqual([ expect.objectContaining({ @@ -663,6 +664,30 @@ describe('Agent context', () => { ]); }); + it('removes a pre-anchor image compression reminder when undoing its prompt', async () => { + profile.update({ activeToolNames: [] }); + const caption = buildImageCompressionCaption({ + original: { width: 3264, height: 666, byteLength: 344 * 1024, mimeType: 'image/png' }, + final: { width: 2000, height: 408, byteLength: 282 * 1024, mimeType: 'image/png' }, + originalPath: '/tmp/originals/shot.png', + }); + ctx.mockNextResponse({ type: 'text', text: 'done' }); + await ctx.rpc.prompt({ + input: [{ type: 'text', text: `inspect this image ${caption}` }], + }); + await ctx.untilTurnEnd(); + + expect(context.get()).toMatchObject([ + { origin: { kind: 'injection', variant: 'image_compression' } }, + { origin: { kind: 'user' } }, + { role: 'assistant' }, + ]); + + await ctx.undoHistory(1); + + expect(context.get()).toEqual([]); + }); + describe('notification projection', () => { it('does not merge a cron-fire envelope into an adjacent user message', () => { const cronEnvelope = diff --git a/packages/agent-core-v2/test/agent/contextMemory/contextTranscript.test.ts b/packages/agent-core-v2/test/agent/contextMemory/contextTranscript.test.ts index 710c136b28..c016a62bdb 100644 --- a/packages/agent-core-v2/test/agent/contextMemory/contextTranscript.test.ts +++ b/packages/agent-core-v2/test/agent/contextMemory/contextTranscript.test.ts @@ -172,6 +172,25 @@ describe('reduceContextTranscript', () => { expect(texts(result)).toEqual(['message A', 'reply A']); }); + it('removes a pre-anchor image compression reminder owned by the undone prompt', () => { + const result = reduceContextTranscript([ + appendMessage( + userMessage('compressed image', { + kind: 'injection', + variant: 'image_compression', + ownerPromptId: 'prompt-1', + }), + ), + appendMessage({ ...userMessage('undo me', { kind: 'user' }), id: 'prompt-1' }), + appendMessage(assistantMessage('undone answer')), + undo(1), + appendMessage(userMessage('keep me', { kind: 'user' })), + appendMessage(assistantMessage('kept answer')), + ]); + + expect(texts(result)).toEqual(['keep me', 'kept answer']); + }); + it('undo stops at a compaction summary', () => { const result = reduceContextTranscript([ appendMessage(userMessage('old')), diff --git a/packages/agent-core-v2/test/agent/contextMemory/undoPrecheck.test.ts b/packages/agent-core-v2/test/agent/contextMemory/undoPrecheck.test.ts index 61b29e4b01..acbe72909f 100644 --- a/packages/agent-core-v2/test/agent/contextMemory/undoPrecheck.test.ts +++ b/packages/agent-core-v2/test/agent/contextMemory/undoPrecheck.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it } from 'vitest'; -import { precheckUndo } from '#/agent/contextMemory/contextOps'; +import { + computeUndoCut, + contextUndo, + isFullyUndoable, +} from '#/agent/contextMemory/contextOps'; import type { ContextMessage } from '#/agent/contextMemory/types'; function text(value: string): { type: 'text'; text: string } { @@ -40,63 +44,82 @@ function compaction(): ContextMessage { const USER_ORIGIN: ContextMessage['origin'] = { kind: 'user' }; -describe('precheckUndo', () => { - it('returns ok when enough real user prompts exist', () => { - expect(precheckUndo([user(USER_ORIGIN), assistant()], 1)).toEqual({ ok: true }); +describe('computeUndoCut', () => { + it('finds the cut for the last real user prompt', () => { + const cut = computeUndoCut([user(USER_ORIGIN), assistant()], 1); + expect(cut).toEqual({ cutIndex: 0, removedCount: 1, stoppedAtCompaction: false }); + expect(isFullyUndoable(cut, 1)).toBe(true); }); it('skips trailing non-user messages while scanning', () => { - expect(precheckUndo([user(USER_ORIGIN), assistant(), assistant()], 1)).toEqual({ ok: true }); + const cut = computeUndoCut([user(USER_ORIGIN), assistant(), assistant()], 1); + expect(cut.cutIndex).toBe(0); + expect(isFullyUndoable(cut, 1)).toBe(true); }); it('treats a user message without origin as a real prompt (legacy)', () => { - expect(precheckUndo([user(), assistant()], 1)).toEqual({ ok: true }); + const cut = computeUndoCut([user(), assistant()], 1); + expect(cut.cutIndex).toBe(0); + expect(isFullyUndoable(cut, 1)).toBe(true); }); - it('returns empty when the history has no real user prompt', () => { - expect(precheckUndo([], 1)).toEqual({ - ok: false, - reason: 'empty', - requested: 1, - undoable: 0, - }); + it('finds nothing when the history has no real user prompt', () => { + const cut = computeUndoCut([], 1); + expect(cut).toEqual({ cutIndex: -1, removedCount: 0, stoppedAtCompaction: false }); + expect(isFullyUndoable(cut, 1)).toBe(false); }); - it('returns empty when only injections are present', () => { - expect(precheckUndo([injection(), assistant()], 1)).toEqual({ - ok: false, - reason: 'empty', - requested: 1, - undoable: 0, - }); + it('skips injections without counting them', () => { + const cut = computeUndoCut([injection(), assistant()], 1); + expect(cut.cutIndex).toBe(-1); + expect(isFullyUndoable(cut, 1)).toBe(false); }); - it('returns insufficient when some but fewer than count prompts exist', () => { + it('counts fewer prompts than requested as not fully undoable', () => { const history = [user(USER_ORIGIN), assistant(), user(USER_ORIGIN), assistant()]; - expect(precheckUndo(history, 3)).toEqual({ - ok: false, - reason: 'insufficient', - requested: 3, - undoable: 2, - }); + const cut = computeUndoCut(history, 3); + expect(cut.removedCount).toBe(2); + expect(isFullyUndoable(cut, 3)).toBe(false); }); - it('returns compaction_boundary when a summary is hit before count is met', () => { - expect(precheckUndo([user(USER_ORIGIN), compaction(), assistant()], 1)).toEqual({ - ok: false, - reason: 'compaction_boundary', - requested: 1, - undoable: 0, - }); + it('stops at a compaction summary', () => { + const cut = computeUndoCut([user(USER_ORIGIN), compaction(), assistant()], 1); + expect(cut).toEqual({ cutIndex: -1, removedCount: 0, stoppedAtCompaction: true }); + expect(isFullyUndoable(cut, 1)).toBe(false); }); - it('reports compaction_boundary over insufficient when the boundary stops the scan', () => { + it('stops at a compaction summary even after counting some prompts', () => { const history = [user(USER_ORIGIN), compaction(), user(USER_ORIGIN), assistant()]; - expect(precheckUndo(history, 2)).toEqual({ - ok: false, - reason: 'compaction_boundary', - requested: 2, - undoable: 1, - }); + const cut = computeUndoCut(history, 2); + expect(cut.removedCount).toBe(1); + expect(cut.stoppedAtCompaction).toBe(true); + expect(isFullyUndoable(cut, 2)).toBe(false); }); }); + +describe('contextUndo op', () => { + it('slices the history at the cut point, dropping post-cut injections too', () => { + const state = [ + user(USER_ORIGIN), + assistant(), + user(USER_ORIGIN), + injection(), + assistant(), + ]; + const next = contextUndo.apply(state, { count: 1 }); + expect(next).toEqual([user(USER_ORIGIN), assistant()]); + }); + + it('returns the same reference when not fully undoable', () => { + const state = [user(USER_ORIGIN), compaction(), assistant()]; + expect(contextUndo.apply(state, { count: 1 })).toBe(state); + }); + + it.each([0, 0.5, Number.MAX_SAFE_INTEGER + 1])( + 'returns the same reference for invalid count %s', + (count) => { + const state = [user(USER_ORIGIN), assistant()]; + expect(contextUndo.apply(state, { count })).toBe(state); + }, + ); +}); diff --git a/packages/agent-core-v2/test/agent/loop/loop.test.ts b/packages/agent-core-v2/test/agent/loop/loop.test.ts index 67d0e2393d..02a7a24767 100644 --- a/packages/agent-core-v2/test/agent/loop/loop.test.ts +++ b/packages/agent-core-v2/test/agent/loop/loop.test.ts @@ -529,6 +529,65 @@ describe('Agent loop', () => { expect(ctx.llmCalls).toHaveLength(3); }); + it('refuses a quiescence lease while a turn is active without cancelling it', async () => { + let started!: () => void; + const activeStarted = new Promise((resolve) => { + started = resolve; + }); + let release!: () => void; + const canFinish = new Promise((resolve) => { + release = resolve; + }); + const hook = loop.hooks.onWillBeginStep.register('test-quiescence', async (_hookCtx, next) => { + started(); + await canFinish; + await next(); + }); + + const active = (await loop.enqueue(nextTurnMessage('active')).assigned).turn; + await activeStarted; + + expect(loop.tryAcquireQuiescence()).toBeUndefined(); + expect(active.signal.aborted).toBe(false); + + hook.dispose(); + ctx.mockNextResponse({ type: 'text', text: 'completed normally' }); + release(); + await expect(active.result).resolves.toMatchObject({ type: 'completed' }); + }); + + it('holds new admissions until an idle quiescence lease is released', async () => { + const lease = loop.tryAcquireQuiescence(); + expect(lease).toBeDefined(); + const held = loop.enqueue(nextTurnMessage('held')); + let assigned = false; + void held.assigned.then(() => { + assigned = true; + }); + + await Promise.resolve(); + expect(assigned).toBe(false); + expect(loop.status()).toMatchObject({ state: 'idle', hasPendingRequests: true }); + + ctx.mockNextResponse({ type: 'text', text: 'after undo' }); + lease?.dispose(); + const resumed = (await held.assigned).turn; + await expect(resumed.result).resolves.toMatchObject({ type: 'completed' }); + }); + + it('can abort an admission while quiescence holds it', async () => { + const lease = loop.tryAcquireQuiescence(); + expect(lease).toBeDefined(); + const held = loop.enqueue(nextTurnMessage('held')); + + expect(held.abort()).toBe(true); + await expect(held.assigned).rejects.toBeDefined(); + expect(loop.hasPendingRequests()).toBe(false); + + lease?.dispose(); + expect(loop.status().state).toBe('idle'); + }); + it('cancels a running step without cancelling its turn and continues the next step', async () => { let releaseRunning!: () => void; const running = new Promise((resolve) => { diff --git a/packages/agent-core-v2/test/agent/loop/stubs.ts b/packages/agent-core-v2/test/agent/loop/stubs.ts index 8ea6a4dc2f..8a14046f23 100644 --- a/packages/agent-core-v2/test/agent/loop/stubs.ts +++ b/packages/agent-core-v2/test/agent/loop/stubs.ts @@ -72,6 +72,7 @@ export function stubLoopWithHooks(options: StubLoopOptions = {}): StubLoop { async run() { return { type: 'completed', steps: 0, truncated: false }; }, status() { return { state: active !== undefined ? 'running' : 'idle', activeTurnId: active?.id, pendingTurnIds: [], hasPendingRequests: queue.hasPendingRequests() }; }, cancel(turnId, reason) { cancels.push({ turnId, reason }); if (active === undefined || (turnId !== undefined && active.id !== turnId)) return false; active.cancel(reason); return true; }, + tryAcquireQuiescence: () => toDisposable(() => {}), hasPendingRequests: () => queue.hasPendingRequests(), registerLoopErrorHandler: errorHandlers.register, settled: () => Promise.resolve(), drainNextBatch(context) { const batch = queue.takeNextBatch(); if (!batch) return undefined; materialize(batch.driver, context); for (const r of batch.merged) materialize(r, context); return batch; }, diff --git a/packages/agent-core-v2/test/agent/plan/plan.test.ts b/packages/agent-core-v2/test/agent/plan/plan.test.ts index da6623abf9..8977ccdc6e 100644 --- a/packages/agent-core-v2/test/agent/plan/plan.test.ts +++ b/packages/agent-core-v2/test/agent/plan/plan.test.ts @@ -947,11 +947,11 @@ describe('Plan service', () => { it('keeps the preserved injection index aligned after undo removes earlier messages', async () => { await plan.enter('test-plan', false); - ctx.appendUserMessage([{ type: 'text', text: 'draft the plan' }]); + ctx.appendUserTurn('draft the plan'); await injectDynamic(); ctx.appendAssistantTurn(1, 'Plan drafted.'); - ctx.undoHistory(1); + await ctx.undoHistory(1); ctx.appendUserMessage([{ type: 'text', text: 'new plan request' }]); await injectDynamic(); diff --git a/packages/agent-core-v2/test/agent/plan/planOps.test.ts b/packages/agent-core-v2/test/agent/plan/planOps.test.ts index c8b8900dca..f3fe745f85 100644 --- a/packages/agent-core-v2/test/agent/plan/planOps.test.ts +++ b/packages/agent-core-v2/test/agent/plan/planOps.test.ts @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; import { DisposableStore } from '#/_base/di/lifecycle'; import { TestInstantiationService } from '#/_base/di/test'; +import { contextAppendMessage, contextUndo } from '#/agent/contextMemory/contextOps'; import { IEventBus } from '#/app/event/eventBus'; import { EventBusService } from '#/app/event/eventBusService'; import { @@ -60,20 +61,20 @@ async function readRecords(key = KEY): Promise { describe('plan ops (wire-backed)', () => { it('enter/cancel/exit drive active state and persist flat records', async () => { - expect(wire.getModel(PlanModel).active).toBe(false); + expect(wire.getModel(PlanModel).current.active).toBe(false); wire.dispatch(planModeEnter({ id: 'p1' })); - expect(wire.getModel(PlanModel)).toEqual({ + expect(wire.getModel(PlanModel).current).toEqual({ active: true, id: 'p1', }); wire.dispatch(planModeCancel({ id: 'p1' })); - expect(wire.getModel(PlanModel)).toEqual({ active: false }); + expect(wire.getModel(PlanModel).current).toEqual({ active: false }); wire.dispatch(planModeEnter({ id: 'p2' })); wire.dispatch(planModeExit({})); - expect(wire.getModel(PlanModel).active).toBe(false); + expect(wire.getModel(PlanModel).current.active).toBe(false); const records = await readRecords(); expect(records.map((record) => record.type)).toEqual([ @@ -94,11 +95,11 @@ describe('plan ops (wire-backed)', () => { it('cancel and exit both deactivate plan mode but emit distinct record types', async () => { wire.dispatch(planModeEnter({ id: 'p1' })); wire.dispatch(planModeCancel({ id: 'p1' })); - expect(wire.getModel(PlanModel)).toEqual({ active: false }); + expect(wire.getModel(PlanModel).current).toEqual({ active: false }); wire.dispatch(planModeEnter({ id: 'p2' })); wire.dispatch(planModeExit({ id: 'p2' })); - expect(wire.getModel(PlanModel)).toEqual({ active: false }); + expect(wire.getModel(PlanModel).current).toEqual({ active: false }); const records = await readRecords(); expect(records.map((record) => record.type)).toEqual([ @@ -122,6 +123,25 @@ describe('plan ops (wire-backed)', () => { expect(wire.getModel(PlanModel)).toBe(active); }); + it('ignores an invalid undo count without corrupting checkpoint state', () => { + wire.dispatch( + contextAppendMessage({ + message: { + role: 'user', + content: [{ type: 'text', text: 'keep me' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }), + ); + const checkpointed = wire.getModel(PlanModel); + + wire.dispatch(contextUndo({ count: 0.5 })); + + expect(wire.getModel(PlanModel)).toBe(checkpointed); + expect(wire.getModel(PlanModel).current).toEqual({ active: false }); + }); + it('replay rebuilds active state silently', async () => { wire.dispatch(planModeEnter({ id: 'p1' })); const records = await readRecords(); @@ -137,7 +157,7 @@ describe('plan ops (wire-backed)', () => { testWireScope(SCOPE, 'plan-replay'), records, ); - expect(host.wire.getModel(PlanModel)).toEqual({ + expect(host.wire.getModel(PlanModel).current).toEqual({ active: true, id: 'p1', }); @@ -153,7 +173,7 @@ describe('plan ops (wire-backed)', () => { { type: 'plan_mode.cancel', id: 'p1' }, ], ); - expect(cancelled.wire.getModel(PlanModel).active).toBe(false); + expect(cancelled.wire.getModel(PlanModel).current.active).toBe(false); }); it('plan.revision persists a flat reference record and advances the per-id counter', async () => { @@ -167,7 +187,7 @@ describe('plan ops (wire-backed)', () => { bytes: 12, }), ); - expect(wire.getModel(PlanModel)).toEqual({ + expect(wire.getModel(PlanModel).current).toEqual({ active: true, id: 'p1', revisionCount: { p1: 1 }, @@ -182,7 +202,7 @@ describe('plan ops (wire-backed)', () => { bytes: 20, }), ); - expect(wire.getModel(PlanModel).revisionCount).toEqual({ p1: 2 }); + expect(wire.getModel(PlanModel).current.revisionCount).toEqual({ p1: 2 }); const records = await readRecords(); expect(records.map((record) => record.type)).toEqual([ @@ -222,7 +242,7 @@ describe('plan ops (wire-backed)', () => { }), ); host.wire.dispatch(planModeExit({})); - expect(host.wire.getModel(PlanModel)).toEqual({ + expect(host.wire.getModel(PlanModel).current).toEqual({ active: false, revisionCount: { p1: 1 }, }); @@ -230,7 +250,7 @@ describe('plan ops (wire-backed)', () => { // Re-entering the same plan id continues the counter instead of // restarting it, so later revisions never overwrite earlier blobs. host.wire.dispatch(planModeEnter({ id: 'p1' })); - expect(host.wire.getModel(PlanModel).revisionCount).toEqual({ p1: 1 }); + expect(host.wire.getModel(PlanModel).current.revisionCount).toEqual({ p1: 1 }); expect( emissions.filter((e) => (e as { type: string }).type === 'plan.revision'), @@ -279,7 +299,7 @@ describe('plan ops (wire-backed)', () => { testWireScope(SCOPE, 'plan-revision-replay'), records, ); - expect(host.wire.getModel(PlanModel)).toEqual({ + expect(host.wire.getModel(PlanModel).current).toEqual({ active: true, id: 'p1', revisionCount: { p1: 2 }, diff --git a/packages/agent-core-v2/test/agent/rpc/undoHistory.test.ts b/packages/agent-core-v2/test/agent/rpc/undoHistory.test.ts index a607af735d..3c586f7b2b 100644 --- a/packages/agent-core-v2/test/agent/rpc/undoHistory.test.ts +++ b/packages/agent-core-v2/test/agent/rpc/undoHistory.test.ts @@ -1,5 +1,7 @@ import { afterEach, describe, expect, it } from 'vitest'; +import { ErrorCodes } from '#/errors'; + import { createTestAgent, telemetryServices, @@ -22,7 +24,7 @@ describe('undoHistory RPC', () => { it('tracks conversation_undo after undoing history', async () => { records = []; ctx = createTestAgent(telemetryServices(recordingTelemetry(records))); - ctx.appendUserMessage([{ type: 'text', text: 'undo me' }]); + ctx.appendUserTurn('undo me'); const undone = await ctx.rpc.undoHistory({ count: 1 }); @@ -32,4 +34,19 @@ describe('undoHistory RPC', () => { properties: { agent_id: 'main', count: 1 }, }); }); + + it('rejects a fractional count without changing persisted history', async () => { + records = []; + ctx = createTestAgent(telemetryServices(recordingTelemetry(records))); + ctx.appendUserTurn('keep me'); + const history = ctx.context.get(); + + await expect(ctx.rpc.undoHistory({ count: 0.5 })).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + details: { field: 'count' }, + }); + + expect(ctx.context.get()).toBe(history); + expect(records).not.toContainEqual(expect.objectContaining({ event: 'conversation_undo' })); + }); }); diff --git a/packages/agent-core-v2/test/agent/skill/skill.test.ts b/packages/agent-core-v2/test/agent/skill/skill.test.ts index 1aeee70249..a46bd364d9 100644 --- a/packages/agent-core-v2/test/agent/skill/skill.test.ts +++ b/packages/agent-core-v2/test/agent/skill/skill.test.ts @@ -69,7 +69,6 @@ describe('AgentSkillService', () => { reg.definePartialInstance(IAgentPromptService, { enqueue: ({ message }: { message: ContextMessage }) => { prompted.push(message); return Promise.resolve({ launched: Promise.resolve(fakeTurn()) } as never); }, retry: () => Promise.resolve(undefined), - undo: () => 0, clear: () => {}, }); registerTestAgentWireServices(reg, 'wire/skill-test'); @@ -162,7 +161,6 @@ describe('SkillTool', () => { reg.definePartialInstance(IAgentPromptService, { enqueue: ({ message }: { message: ContextMessage }) => { prompted.push(message); return Promise.resolve({ launched: Promise.resolve(fakeTurn()) } as never); }, retry: () => Promise.resolve(undefined), - undo: () => 0, clear: () => {}, }); registerTestAgentWireServices(reg, 'wire/skill-test'); diff --git a/packages/agent-core-v2/test/agent/task/rpc-events.test.ts b/packages/agent-core-v2/test/agent/task/rpc-events.test.ts index 21c876ad1a..628d9bd7cb 100644 --- a/packages/agent-core-v2/test/agent/task/rpc-events.test.ts +++ b/packages/agent-core-v2/test/agent/task/rpc-events.test.ts @@ -25,6 +25,9 @@ import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory' import { IEventBus } from '#/app/event/eventBus'; import type { IExternalHooksRunnerService } from '#/app/externalHooksRunner/externalHooksRunner'; import { IAgentLoopService } from '#/agent/loop/loop'; +import { MessageStepRequest } from '#/agent/loop/stepRequest'; +import { IAgentConversationUndoService } from '#/agent/undo/undo'; +import { ErrorCodes } from '#/errors'; import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata'; import { configServices, @@ -714,6 +717,109 @@ describe('AgentTaskService — notification delivery', () => { } }); + it('re-delivers a terminal task notification removed by undo when output is unavailable', async () => { + const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-agent-undo-')); + let fixture: TaskServiceFixture | undefined; + try { + const persistence = createAgentTaskPersistence(sessionDir); + await persistence.writeTask(persistedAgent()); + await persistence.appendTaskOutput('agent-done0000', 'restored subagent summary'); + fixture = createAgentTaskService({ sessionDir }); + const { agent, ctx, manager } = fixture; + ctx.appendUserTurn('start the background task'); + agent.context.appendUserMessage.mockClear(); + + await manager.loadFromDisk(); + await manager.reconcile(); + await vi.waitFor(() => { + expect(agent.context.appendUserMessage).toHaveBeenCalledTimes(1); + }); + vi.spyOn(manager, 'getOutputSnapshot').mockRejectedValueOnce( + new Error('output unavailable'), + ); + + await ctx.get(IAgentConversationUndoService).undo(1); + + expect(agent.context.appendUserMessage).toHaveBeenCalledTimes(2); + expect(ctx.context.get().some((message) => message.origin?.kind === 'user')).toBe(false); + expect( + ctx.context.get().filter((message) => message.origin?.kind === 'task'), + ).toHaveLength(1); + } finally { + await cleanupSessionDir(sessionDir, fixture); + } + }); + + it('preserves a queued notification when undo rejects an active turn', async () => { + const fixture = createAgentTaskService(); + const { ctx, manager } = fixture; + const loop = ctx.get(IAgentLoopService); + let markStarted!: () => void; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + let release!: () => void; + const canFinish = new Promise((resolve) => { + release = resolve; + }); + const hook = loop.hooks.onWillBeginStep.register('test-notification-undo', async (_hookCtx, next) => { + markStarted(); + await canFinish; + await next(); + }); + + try { + ctx.appendTurnExchange('kept prompt', 'kept answer'); + const active = ( + await loop.enqueue( + new MessageStepRequest( + { + role: 'user', + content: [{ type: 'text', text: 'remove me' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { admission: 'newTurn' }, + ), + ).assigned + ).turn; + await started; + const taskId = registerProcess(manager, immediateProcess(0, 'done'), 'echo done', 'done'); + await vi.waitFor(() => { + expect(manager.getTask(taskId)?.status).toBe('completed'); + expect(loop.hasPendingRequests()).toBe(true); + }); + expect(notifiedCount(ctx)).toBe(0); + + await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_BUSY, + details: { reason: 'loop' }, + }); + expect(active.signal.aborted).toBe(false); + expect( + ctx.context.get().filter((message) => message.origin?.kind === 'task'), + ).toEqual([]); + + ctx.mockNextResponse({ type: 'text', text: 'notification acknowledged' }); + ctx.mockNextResponse({ type: 'text', text: 'turn completed' }); + release(); + await expect(active.result).resolves.toMatchObject({ type: 'completed' }); + expect( + ctx.context.get().filter((message) => message.origin?.kind === 'task'), + ).toEqual([ + expect.objectContaining({ + origin: expect.objectContaining({ taskId, status: 'completed' }), + }), + ]); + expect(notifiedCount(ctx)).toBe(1); + } finally { + release(); + hook.dispose(); + await ctx.get(ISessionMetadata).ready; + await ctx.dispose(); + } + }); + it('does not double-notify newly lost restored agent tasks', async () => { const sessionDir = await mkdtemp(join(tmpdir(), 'kimi-bg-agent-lost-')); let fixture: TaskServiceFixture | undefined; diff --git a/packages/agent-core-v2/test/agent/task/taskService.test.ts b/packages/agent-core-v2/test/agent/task/taskService.test.ts index 1b2e8f49bd..d7e3f49319 100644 --- a/packages/agent-core-v2/test/agent/task/taskService.test.ts +++ b/packages/agent-core-v2/test/agent/task/taskService.test.ts @@ -13,7 +13,9 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; import { DisposableStore, toDisposable } from '#/_base/di/lifecycle'; +import { ILogService } from '#/_base/log/log'; import { TestInstantiationService } from '#/_base/di/test'; +import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants'; import { IAgentContextInjectorService, type ContextInjectionContext, @@ -47,6 +49,7 @@ import { EventBusService } from '#/app/event/eventBusService'; import { ITaskService } from '#/app/task/task'; import { InMemoryStorageService } from '#/persistence/backends/memory/inMemoryStorageService'; +import { stubLog } from '../../_base/log/stubs'; import { stubContextMemory } from '../contextMemory/stubs'; import { stubLoopWithHooks } from '../loop/stubs'; import type { TaskServiceTestManager } from './stubs'; @@ -89,6 +92,11 @@ describe('AgentTaskService', () => { ix = disposables.add(new TestInstantiationService()); eventBus = disposables.add(new EventBusService()); injectionProviders = new Map(); + ix.stub(ILogService, stubLog()); + ix.stub(IAgentConversationUndoParticipantRegistry, { + register: () => toDisposable(() => {}), + list: () => [], + }); ix.stub(IWireService, stubWireService()); ix.stub(IEventBus, eventBus); ix.stub(IAgentContextInjectorService, { @@ -461,6 +469,11 @@ describe('AgentTaskService', () => { captureRestoreHook?: (hook: RestoreHook) => void, ): TestInstantiationService { const ix = disposables.add(new TestInstantiationService()); + ix.stub(ILogService, stubLog()); + ix.stub(IAgentConversationUndoParticipantRegistry, { + register: () => toDisposable(() => {}), + list: () => [], + }); ix.stub(IWireService, stubWireService(captureRestoreHook)); ix.stub(IEventBus, disposables.add(new EventBusService())); ix.stub(IAgentContextInjectorService, { diff --git a/packages/agent-core-v2/test/agent/toolSelect/toolSelect.e2e.test.ts b/packages/agent-core-v2/test/agent/toolSelect/toolSelect.e2e.test.ts index a74ac91146..69f151edb0 100644 --- a/packages/agent-core-v2/test/agent/toolSelect/toolSelect.e2e.test.ts +++ b/packages/agent-core-v2/test/agent/toolSelect/toolSelect.e2e.test.ts @@ -21,6 +21,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; +import { IAgentConversationUndoService } from '#/agent/undo/undo'; import type { ContextMessage } from '#/agent/contextMemory/types'; import type { ExecutableTool, ToolExecution } from '#/tool/toolContract'; import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor'; @@ -235,7 +236,7 @@ describe('progressive tool disclosure end-to-end', () => { await ctx.rpc.prompt({ input: [{ type: 'text', text: 'load alpha' }] }); await ctx.untilTurnEnd(); - ctx.get(IAgentContextMemoryService).undo(1); + await ctx.get(IAgentConversationUndoService).undo(1); const afterUndo = ctx.get(IAgentContextMemoryService).get(); expect(afterUndo.some((message) => message.tools?.some((tool) => tool.name === MCP_ALPHA))).toBe( false, diff --git a/packages/agent-core-v2/test/agent/toolSelect/toolSelectService.test.ts b/packages/agent-core-v2/test/agent/toolSelect/toolSelectService.test.ts index 881d9c1fbf..0258f7b2ef 100644 --- a/packages/agent-core-v2/test/agent/toolSelect/toolSelectService.test.ts +++ b/packages/agent-core-v2/test/agent/toolSelect/toolSelectService.test.ts @@ -230,6 +230,10 @@ class FakeLoopService implements IAgentLoopService { throw new Error('unused in this suite'); } + tryAcquireQuiescence(): IDisposable | undefined { + return toDisposable(() => {}); + } + hasPendingRequests(): boolean { return false; } diff --git a/packages/agent-core-v2/test/agent/undo/undo.test.ts b/packages/agent-core-v2/test/agent/undo/undo.test.ts new file mode 100644 index 0000000000..170c3b63c6 --- /dev/null +++ b/packages/agent-core-v2/test/agent/undo/undo.test.ts @@ -0,0 +1,509 @@ +/** + * Scenario: undo validation and restoration across conversation-scoped models. + * Responsibility: AgentConversationUndoService commits one undo and publishes + * restored observable state. + * Wiring: full TestAgentContext with real wire models and event bus. + * Run: pnpm --filter @moonshot-ai/agent-core-v2 test -- test/agent/undo/undo.test.ts + */ + +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; +import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants'; +import { contextApplyCompaction } from '#/agent/contextMemory/contextOps'; +import { + CHECKPOINTED_MODELS, + type Checkpointed, +} from '#/agent/contextMemory/conversationTime'; +import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction'; +import { IAgentLoopService } from '#/agent/loop/loop'; +import { MessageStepRequest } from '#/agent/loop/stepRequest'; +import { TurnModel } from '#/agent/loop/turnOps'; +import { IAgentPlanService } from '#/agent/plan/plan'; +import { PlanModel } from '#/agent/plan/planOps'; +import { IAgentPromptService } from '#/agent/prompt/prompt'; +import { IAgentConversationUndoService } from '#/agent/undo/undo'; +import { IEventBus } from '#/app/event/eventBus'; +import { IAgentTelemetryContextService } from '#/app/telemetry/agentTelemetryContext'; +import { ErrorCodes } from '#/errors'; +import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata'; +import { TodoModel, todoSet } from '#/session/todo/todoOps'; +import { defineModel } from '#/wire/model'; +import { IWireService } from '#/wire/wire'; + +import { createTestAgent, telemetryServices, type TestAgentContext } from '../../harness'; +import { recordingTelemetry, type TelemetryRecord } from '../../app/telemetry/stubs'; + +describe('AgentConversationUndoService', () => { + let ctx: TestAgentContext; + let records: TelemetryRecord[]; + + afterEach(async () => { + try { + await ctx.expectResumeMatches(); + } finally { + await ctx.dispose(); + } + }); + + function setup() { + records = []; + ctx = createTestAgent(telemetryServices(recordingTelemetry(records))); + ctx.get(IAgentContextMemoryService); + return ctx; + } + + it('exposes availability from context history', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + expect(undo.availability()).toEqual({ maxTurns: 0, stoppedAtCompaction: false }); + + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + expect(undo.availability()).toEqual({ maxTurns: 2, stoppedAtCompaction: false }); + }); + + it('rejects undo with structured reasons', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + + await expect(undo.undo(1)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_UNDO_UNAVAILABLE, + details: { reason: 'empty', requestedCount: 1, undoableCount: 0 }, + }); + + ctx.appendTurnExchange('u1', 'a1'); + await expect(undo.undo(2)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_UNDO_UNAVAILABLE, + details: { reason: 'insufficient', requestedCount: 2, undoableCount: 1 }, + }); + }); + + it.each([ + 0, + -1, + 0.5, + Number.MAX_SAFE_INTEGER + 1, + Number.POSITIVE_INFINITY, + Number.NaN, + ])('rejects invalid undo count %s without mutating history', async (count) => { + setup(); + ctx.appendTurnExchange('u1', 'a1'); + const history = ctx.context.get(); + + await expect(ctx.get(IAgentConversationUndoService).undo(count)).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + details: { field: 'count' }, + }); + + expect(ctx.context.get()).toBe(history); + }); + + it('returns session.busy for an active turn without cancelling it', async () => { + setup(); + const loop = ctx.get(IAgentLoopService); + let started!: () => void; + let release!: () => void; + const didStart = new Promise((resolve) => { + started = resolve; + }); + const canFinish = new Promise((resolve) => { + release = resolve; + }); + const hook = loop.hooks.onWillBeginStep.register('test-invalid-undo', async (_hookCtx, next) => { + started(); + await canFinish; + await next(); + }); + ctx.mockNextResponse({ type: 'text', text: 'system result' }); + const turn = ( + await loop.enqueue( + new MessageStepRequest( + { + role: 'user', + content: [{ type: 'text', text: 'system work' }], + toolCalls: [], + origin: { kind: 'system_trigger', name: 'test' }, + }, + { admission: 'newTurn' }, + ), + ).assigned + ).turn; + await didStart; + const history = ctx.context.get(); + + await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_BUSY, + details: { reason: 'loop' }, + }); + expect(turn.signal.aborted).toBe(false); + expect(loop.status().state).toBe('running'); + expect(ctx.context.get()).toBe(history); + + hook.dispose(); + release(); + await expect(turn.result).resolves.toMatchObject({ type: 'completed' }); + }); + + it('returns session.busy for active compaction without cancelling it', async () => { + setup(); + ctx.appendTurnExchange('u1', 'a1'); + const history = ctx.context.get(); + const compaction = ctx.get(IAgentFullCompactionService); + const abortController = new AbortController(); + const active = vi.spyOn(compaction, 'compacting', 'get').mockReturnValue({ + abortController, + promise: new Promise(() => {}), + trigger: 'manual', + tokenCount: 2, + }); + + try { + await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_BUSY, + details: { reason: 'compaction' }, + }); + expect(abortController.signal.aborted).toBe(false); + expect(ctx.context.get()).toBe(history); + } finally { + active.mockRestore(); + } + }); + + it('refuses to cross a compaction boundary', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + ctx.appendTurnExchange('u1', 'a1'); + ctx.get(IAgentContextMemoryService).applyCompaction({ + summary: 'summary of u1', + compactedCount: 2, + tokensBefore: 100, + tokensAfter: 10, + }); + ctx.appendTurnExchange('u2', 'a2'); + + expect(undo.availability()).toEqual({ maxTurns: 1, stoppedAtCompaction: true }); + await expect(undo.undo(2)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_UNDO_UNAVAILABLE, + details: { reason: 'compaction_boundary', requestedCount: 2, undoableCount: 1 }, + }); + + await undo.undo(1); + const history = ctx.context.get(); + expect(history.map((m) => m.role)).toEqual(['user', 'user']); + expect(history[1]?.origin?.kind).toBe('compaction_summary'); + }); + + it('refuses loudly when a legacy compaction leaves anchors without checkpoints', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + const wire = ctx.get(IWireService); + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + wire.dispatch(contextApplyCompaction({ summary: 'legacy summary', compactedCount: 2 })); + expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'user', 'assistant']); + + await expect(undo.undo(1)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_UNDO_UNAVAILABLE, + details: { reason: 'compaction_boundary', requestedCount: 1, undoableCount: 0 }, + }); + expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'user', 'assistant']); + }); + + it('attributes a checkpoint depth failure to the limiting model', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + ctx.appendTurnExchange('u1', 'a1'); + // A checkpointed model that never tracks anchors (no reducers) drags the + // depth to 0 without any compaction in history. + const defective = defineModel>('testDefective', () => ({ + current: null, + checkpoints: [], + })); + CHECKPOINTED_MODELS.push(defective); + + try { + await expect(undo.undo(1)).rejects.toMatchObject({ + code: ErrorCodes.SESSION_UNDO_UNAVAILABLE, + details: { + reason: 'checkpoint_lost', + requestedCount: 1, + undoableCount: 0, + model: 'testDefective', + }, + }); + } finally { + CHECKPOINTED_MODELS.splice(CHECKPOINTED_MODELS.indexOf(defective), 1); + } + expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'assistant']); + }); + + it('restores todos to their pre-turn value', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + const wire = ctx.get(IWireService); + ctx.appendTurnExchange('u1', 'a1'); + wire.dispatch(todoSet({ key: 'todo', value: [{ title: 'kept', status: 'pending' }] })); + ctx.appendTurnExchange('u2', 'a2'); + wire.dispatch(todoSet({ key: 'todo', value: [{ title: 'doomed', status: 'pending' }] })); + + await undo.undo(1); + + expect(wire.getModel(TodoModel).current).toEqual([{ title: 'kept', status: 'pending' }]); + }); + + it('restores plan mode and its telemetry mirror to their pre-turn value', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + const wire = ctx.get(IWireService); + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + await ctx.get(IAgentPlanService).enter('plan-x', false); + const restoredModes: boolean[] = []; + const subscription = ctx.get(IEventBus).subscribe('agent.status.updated', (event) => { + if (event.planMode !== undefined) restoredModes.push(event.planMode); + }); + + try { + await undo.undo(1); + + expect(wire.getModel(PlanModel).current.active).toBe(false); + expect(ctx.get(IAgentTelemetryContextService).get().mode).toBe('agent'); + expect(restoredModes).toEqual([false]); + } finally { + subscription.dispose(); + } + }); + + it('does not roll back world-time turn bookkeeping', async () => { + setup(); + const undo = ctx.get(IAgentConversationUndoService); + const wire = ctx.get(IWireService); + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + expect(wire.getModel(TurnModel).nextTurnId).toBe(2); + + await undo.undo(1); + + expect(wire.getModel(TurnModel).nextTurnId).toBe(2); + }); + + it('flushes state reconciliation before publishing undo', async () => { + setup(); + const wire = ctx.get(IWireService); + const order: string[] = []; + const flush = vi.spyOn(wire, 'flush'); + const originalFlush = flush.getMockImplementation(); + flush.mockImplementation(async () => { + order.push('flush'); + await originalFlush?.(); + }); + const participants = ctx.get(IAgentConversationUndoParticipantRegistry); + participants.register({ + id: 'test.state', + reconcileAfterUndo: async () => { + order.push('state'); + }, + }); + const subscription = ctx.get(IEventBus).subscribe('context.undone', () => { + order.push('context.undone'); + }); + ctx.appendTurnExchange('u1', 'a1'); + + try { + await ctx.get(IAgentConversationUndoService).undo(1); + + expect(order).toEqual(['flush', 'state', 'flush', 'context.undone']); + } finally { + subscription.dispose(); + flush.mockRestore(); + } + }); + + it.each([ + [1, []], + [2, ['state']], + ] as const)( + 'rejects the committed undo when post-cut flush %i fails', + async (failureCall, expectedReconciled) => { + setup(); + const wire = ctx.get(IWireService); + const originalFlush = wire.flush.bind(wire); + let flushCalls = 0; + const storageError = new Error('storage unavailable'); + const flush = vi.spyOn(wire, 'flush').mockImplementation(async () => { + flushCalls += 1; + if (flushCalls === failureCall) throw storageError; + await originalFlush(); + }); + const reconciled: string[] = []; + const participants = ctx.get(IAgentConversationUndoParticipantRegistry); + participants.register({ + id: 'test.flush-failure-state', + reconcileAfterUndo: async () => { + reconciled.push('state'); + }, + }); + const undone: number[] = []; + const subscription = ctx.get(IEventBus).subscribe('context.undone', ({ turns }) => { + undone.push(turns); + }); + ctx.appendTurnExchange('u1', 'a1'); + + try { + await expect(ctx.get(IAgentConversationUndoService).undo(1)).rejects.toBe(storageError); + expect(ctx.context.get()).toEqual([]); + expect(reconciled).toEqual(expectedReconciled); + expect(undone).toEqual([]); + expect(records.filter((record) => record.event === 'conversation_undo')).toEqual([]); + } finally { + subscription.dispose(); + flush.mockRestore(); + } + }, + ); + + it('serializes concurrent undos through state reconciliation', async () => { + setup(); + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + let releaseFirst!: () => void; + const firstBlocked = new Promise((resolve) => { + releaseFirst = resolve; + }); + let markFirstStarted!: () => void; + const firstStarted = new Promise((resolve) => { + markFirstStarted = resolve; + }); + let calls = 0; + let active = 0; + let maxActive = 0; + ctx.get(IAgentConversationUndoParticipantRegistry).register({ + id: 'test.serial-state', + reconcileAfterUndo: async () => { + calls += 1; + active += 1; + maxActive = Math.max(maxActive, active); + if (calls === 1) { + markFirstStarted(); + await firstBlocked; + } + active -= 1; + }, + }); + + const first = ctx.get(IAgentConversationUndoService).undo(1); + await firstStarted; + const second = ctx.get(IAgentConversationUndoService).undo(1); + await Promise.resolve(); + + expect(calls).toBe(1); + expect(ctx.context.get().map((message) => message.role)).toEqual(['user', 'assistant']); + releaseFirst(); + await Promise.all([first, second]); + + expect(calls).toBe(2); + expect(maxActive).toBe(1); + expect(ctx.context.get()).toEqual([]); + }); + + it('publishes context.undone and tracks conversation_undo', async () => { + setup(); + ctx.get(IAgentConversationUndoService); + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + + await ctx.rpc.undoHistory({ count: 1 }); + + expect(records).toContainEqual({ + event: 'conversation_undo', + properties: { agent_id: 'main', count: 1 }, + }); + expect(ctx.context.get().map((m) => m.role)).toEqual(['user', 'assistant']); + }); + + it('clears lastPrompt when undo removes the only prompt', async () => { + setup(); + const metadata = ctx.get(ISessionMetadata); + await metadata.ready; + await metadata.update({ lastPrompt: 'u1' }); + ctx.appendTurnExchange('u1', 'a1'); + + await ctx.get(IAgentConversationUndoService).undo(1); + + await expect(metadata.read()).resolves.toMatchObject({ lastPrompt: undefined }); + }); + + it('uses the newest pending prompt as lastPrompt after undo', async () => { + setup(); + const metadata = ctx.get(ISessionMetadata); + await metadata.ready; + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + const list = vi.spyOn(ctx.get(IAgentPromptService), 'list').mockReturnValue({ + active: undefined, + pending: [ + { + id: 'queued', + userMessageId: 'queued', + createdAt: new Date(0).toISOString(), + state: 'pending', + message: { + role: 'user', + content: [{ type: 'text', text: 'queued prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + ], + }); + + try { + await ctx.get(IAgentConversationUndoService).undo(1); + await expect(metadata.read()).resolves.toMatchObject({ lastPrompt: 'queued prompt' }); + } finally { + list.mockRestore(); + } + }); + + it('treats metadata reconciliation failure as non-fatal after committing undo', async () => { + setup(); + ctx.appendTurnExchange('u1', 'a1'); + ctx.appendTurnExchange('u2', 'a2'); + const update = vi.spyOn(ctx.get(ISessionMetadata), 'update').mockRejectedValueOnce( + new Error('metadata write failed'), + ); + const undone: number[] = []; + const subscription = ctx.get(IEventBus).subscribe('context.undone', ({ turns }) => { + undone.push(turns); + }); + + try { + await expect(ctx.get(IAgentConversationUndoService).undo(1)).resolves.toBe(1); + + expect(ctx.context.get().map((message) => message.role)).toEqual(['user', 'assistant']); + expect(undone).toEqual([1]); + expect(records).toContainEqual({ + event: 'conversation_undo', + properties: { agent_id: 'main', count: 1 }, + }); + } finally { + subscription.dispose(); + update.mockRestore(); + } + }); + + it('persists context.undo without introducing a wire-level cut record', async () => { + setup(); + ctx.appendTurnExchange('u1', 'a1'); + + await ctx.get(IAgentConversationUndoService).undo(1); + await ctx.get(IWireService).flush(); + + const wireEvents = ctx.allEvents + .filter((event) => event.type === '[wire]') + .map((event) => event.event); + expect(wireEvents).toContain('context.undo'); + expect(wireEvents).not.toContain('log.cut'); + }); +}); diff --git a/packages/agent-core-v2/test/app/externalHooksRunner/integration.test.ts b/packages/agent-core-v2/test/app/externalHooksRunner/integration.test.ts index 4f66c01450..f447eed466 100644 --- a/packages/agent-core-v2/test/app/externalHooksRunner/integration.test.ts +++ b/packages/agent-core-v2/test/app/externalHooksRunner/integration.test.ts @@ -107,7 +107,7 @@ function stubContextMemory(): IAgentContextMemoryService & { undo: (count) => { const cut = computeUndoCut(messages, count); if (cut.cutIndex >= 0 && cut.removedCount >= count) { - messages.splice(cut.cutIndex, messages.length - cut.cutIndex); + messages.splice(cut.cutIndex); } return cut; }, diff --git a/packages/agent-core-v2/test/app/gateway/gateway.test.ts b/packages/agent-core-v2/test/app/gateway/gateway.test.ts index d6d4b908eb..e2194cd5ed 100644 --- a/packages/agent-core-v2/test/app/gateway/gateway.test.ts +++ b/packages/agent-core-v2/test/app/gateway/gateway.test.ts @@ -56,7 +56,6 @@ describe('RestGateway', () => { abort: () => true, inject: () => Promise.resolve(undefined), retry: () => Promise.resolve(undefined), - undo: () => 0, clear: () => {}, hooks: createHooks(['onBeforeSubmitPrompt']) as IAgentPromptService['hooks'], }; diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index 4e0b8a3f3a..7c9d868ca0 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -16,6 +16,7 @@ import { IAgentBlobService } from '#/agent/blob/agentBlobService'; import { AgentBlobServiceImpl } from '#/agent/blob/agentBlobServiceImpl'; import { IHostEnvironment } from '#/os/interface/hostEnvironment'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; +import { CHECKPOINTED_MODELS, type Checkpointed } from '#/agent/contextMemory/conversationTime'; import type { ContextMessage } from '#/agent/contextMemory/types'; import { ISessionCronService } from '#/session/cron/sessionCronService'; import { SessionCronServiceImpl } from '#/session/cron/sessionCronServiceImpl'; @@ -144,6 +145,7 @@ import { import { IEventBus } from '#/app/event/eventBus'; import { IWireService } from '#/wire/wire'; import { WireService } from '#/wire/wireService'; +import { promptTurn } from '#/agent/loop/turnOps'; import { IModelService, type ModelsSection } from '#/kosong/model/model'; import { DEFAULT_MODEL_SECTION, @@ -358,6 +360,7 @@ interface ResumeStateSnapshot { readonly context: { readonly history: readonly ContextMessage[]; }; + readonly checkpointedModels: Readonly>; readonly permission: Omit, 'rules'>; readonly usage: Omit, 'currentTurn'>; } @@ -1460,6 +1463,18 @@ export class AgentTestContext { }); } + appendUserTurn(text: string): void { + this.get(IWireService).dispatch( + promptTurn({ input: [{ type: 'text', text }], origin: { kind: 'user' } }), + ); + this.appendMessage({ + role: 'user', + content: [{ type: 'text', text }], + toolCalls: [], + origin: { kind: 'user' }, + }); + } + appendSystemReminder( content: string, origin: ContextMessage['origin'] = { kind: 'injection', variant: 'system-reminder' }, @@ -1490,9 +1505,9 @@ export class AgentTestContext { this.get(IAgentPromptService).clear(); } - undoHistory(count: number): number { + async undoHistory(count: number): Promise { const rpcMethods = this.get(IAgentRPCService); - return rpcMethods.undoHistory({ count }) as unknown as number; + return rpcMethods.undoHistory({ count }); } newEvents(): EventSnapshot { @@ -1579,6 +1594,16 @@ export class AgentTestContext { this.coverUsage(tokenTotal); } + appendTurnExchange(userText: string, assistantText: string, tokenTotal?: number): void { + this.appendUserTurn(userText); + this.appendAssistantMessage({ + role: 'assistant', + content: [{ type: 'text', text: assistantText }], + toolCalls: [], + }); + this.coverUsage(tokenTotal); + } + appendAssistantText(step: number, text: string): void { this.appendAssistantTextWithUsage(step, text); } @@ -2182,6 +2207,12 @@ function resumeStateSnapshot(ctx: AgentTestContext): ResumeStateSnapshot { return { config: configStateSnapshot(ctx), context: resumeContextSnapshot(ctx), + checkpointedModels: Object.fromEntries( + CHECKPOINTED_MODELS.map((model) => [ + model.name, + (ctx.get(IWireService).getModel(model) as Checkpointed).current, + ]), + ), permission: permissionData, usage: usageStatus, }; diff --git a/packages/agent-core-v2/test/index.test.ts b/packages/agent-core-v2/test/index.test.ts index a8131c9f25..db7cd48bad 100644 --- a/packages/agent-core-v2/test/index.test.ts +++ b/packages/agent-core-v2/test/index.test.ts @@ -2,6 +2,7 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { WIRE_PROTOCOL_VERSION, + CHECKPOINTED_MODELS, IAgentContextMemoryService, IAgentContextSizeService, IAgentGoalService, @@ -23,6 +24,7 @@ import { IAppendLogStore } from '#/persistence/interface/appendLogStore'; import { IFileSystemStorageService } from '#/persistence/interface/storage'; import { todoSet, TodoModel } from '#/session/todo/todoOps'; import { OP_REGISTRY } from '#/wire/op'; +import { MODEL_CROSS_REDUCERS } from '#/wire/model'; import { IWireService } from '#/wire/wire'; import { AGENT_WIRE_RECORD_KEY } from '#/wire/record'; import { registerTestAgentWire, restoreTestAgentWire } from './wire/stubs'; @@ -162,7 +164,42 @@ describe('v1 wire vocabulary', () => { await restoreTestAgentWire(fresh, log2, SCOPE, records); - expect(fresh.getModel(TodoModel)).toEqual([{ title: 'restore me', status: 'in_progress' }]); + expect(fresh.getModel(TodoModel).current).toEqual([ + { title: 'restore me', status: 'in_progress' }, + ]); + }); +}); + +describe('conversation-time checkpoint registration', () => { + // Models that react to context.* records but deliberately stay on world time + // (ephemeral notice state that must not travel through undo) are exempt. + // Registering a new context-reacting model without `defineCheckpointedModel` + // fails this test — add the name here only with a justification. + const CHECKPOINT_EXEMPT_MODELS: ReadonlySet = new Set([ + // goalForkNotice is one-shot reminder bookkeeping, not conversation state. + 'goalForkNotice', + ]); + const CONTEXT_OPS = [ + 'context.append_message', + 'context.apply_compaction', + 'context.clear', + 'context.undo', + ]; + + it('registers every context-reacting model as checkpointed or explicitly exempt', () => { + const violations: string[] = []; + let entries = 0; + for (const opType of CONTEXT_OPS) { + for (const entry of MODEL_CROSS_REDUCERS.get(opType) ?? []) { + entries += 1; + if (CHECKPOINTED_MODELS.includes(entry.model)) continue; + if (CHECKPOINT_EXEMPT_MODELS.has(entry.model.name)) continue; + violations.push(`${entry.model.name} (on ${opType})`); + } + } + // Guard against a vacuous pass when module loading changes. + expect(entries).toBeGreaterThan(0); + expect(violations).toEqual([]); }); }); diff --git a/packages/agent-core-v2/test/session/todo/sessionTodo.test.ts b/packages/agent-core-v2/test/session/todo/sessionTodo.test.ts index c970b352ee..47c7158550 100644 --- a/packages/agent-core-v2/test/session/todo/sessionTodo.test.ts +++ b/packages/agent-core-v2/test/session/todo/sessionTodo.test.ts @@ -1,3 +1,10 @@ +/** + * Scenario: session-shared Todo state, including undo restoration. + * Responsibility: SessionTodoService exposes the main wire state and emits observable changes. + * Wiring: lightweight lifecycle/agent fakes with real event-bus behavior. + * Run: pnpm --filter @moonshot-ai/agent-core-v2 test -- test/session/todo/sessionTodo.test.ts + */ + import { describe, expect, it } from 'vitest'; import type { ServiceIdentifier, ServicesAccessor } from '#/_base/di/instantiation'; @@ -8,7 +15,10 @@ import { Emitter } from '#/_base/event'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import { IAgentProfileService } from '#/agent/profile/profile'; +import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy'; import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; +import { IEventBus } from '#/app/event/eventBus'; +import { EventBusService } from '#/app/event/eventBusService'; import { createHooks } from '#/hooks'; import { IAgentLifecycleService } from '#/session/agentLifecycle/agentLifecycle'; import { ISessionTodoService } from '#/session/todo/sessionTodo'; @@ -27,6 +37,7 @@ interface FakeAgent { readonly registeredTools: string[]; readonly registeredVariants: string[]; readonly appended: RecordedTodoSet[]; + readonly eventBus: EventBusService; readonly restore: (records: readonly WireRecord[]) => Promise; } @@ -34,6 +45,7 @@ function makeFakeAgent(agentId: string): FakeAgent { const registeredTools: string[] = []; const registeredVariants: string[] = []; const appended: RecordedTodoSet[] = []; + const eventBus = new EventBusService(); let todoState: readonly TodoItem[] = []; @@ -97,7 +109,7 @@ function makeFakeAgent(agentId: string): FakeAgent { }, restore: async () => {}, flush: async () => {}, - getModel: () => todoState, + getModel: () => ({ current: todoState, checkpoints: [] }), subscribe: () => toDisposable(() => {}), } as unknown as IWireService; @@ -108,6 +120,8 @@ function makeFakeAgent(agentId: string): FakeAgent { if (id === IInstantiationService) return instantiationStub as unknown as T; if (id === IAgentContextMemoryService) return memoryStub as unknown as T; if (id === IAgentProfileService) return profileStub as unknown as T; + if (id === IAgentToolPolicyService) return profileStub as unknown as T; + if (id === IEventBus) return eventBus as unknown as T; if (id === IWireService) return wireStub as unknown as T; throw new Error(`unexpected service request in fake agent: ${String(id)}`); }, @@ -125,6 +139,7 @@ function makeFakeAgent(agentId: string): FakeAgent { registeredTools, registeredVariants, appended, + eventBus, restore, }; } @@ -205,6 +220,24 @@ describe('SessionTodoService', () => { ]); }); + it('fires the restored list once when undo changes the main wire state', async () => { + const main = makeFakeAgent('main'); + const lifecycle = makeLifecycleStub([main.handle]); + const service = new SessionTodoService(lifecycle.service); + service.setTodos([{ title: 'doomed', status: 'in_progress' }]); + + const seen: Array = []; + const subscription = service.onDidChange((todos) => seen.push(todos)); + await main.restore([ + { type: 'tools.update_store', key: 'todo', value: [{ title: 'kept', status: 'pending' }] }, + ]); + main.eventBus.publish({ type: 'context.undone', turns: 1 }); + main.eventBus.publish({ type: 'context.undone', turns: 1 }); + subscription.dispose(); + + expect(seen).toEqual([[{ title: 'kept', status: 'pending' }]]); + }); + it('appends a tools.update_store record to the main agent wire on setTodos', () => { const main = makeFakeAgent('main'); const lifecycle = makeLifecycleStub([main.handle]); diff --git a/packages/agent-core-v2/test/tool/tool.test.ts b/packages/agent-core-v2/test/tool/tool.test.ts index f0e3adbfb7..fa98629ac1 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -294,6 +294,13 @@ function createAgentLifecycleStub(options: AgentLifecycleStubOptions = {}): Agen unregister: () => {}, } as never; } + if (serviceId === IEventBus) { + return { + _serviceBrand: undefined, + publish: () => {}, + subscribe: () => noopDisposable(), + } as never; + } if (serviceId === IWireService) { return { _serviceBrand: undefined, diff --git a/packages/agent-core-v2/test/wire/resume.test.ts b/packages/agent-core-v2/test/wire/resume.test.ts index 92afce4d32..337a69c17c 100644 --- a/packages/agent-core-v2/test/wire/resume.test.ts +++ b/packages/agent-core-v2/test/wire/resume.test.ts @@ -4,6 +4,10 @@ import { join } from 'pathe'; import { describe, expect, it, vi } from 'vitest'; +import { + resetUnexpectedErrorHandler, + setUnexpectedErrorHandler, +} from '#/_base/errors/unexpectedError'; import { WIRE_PROTOCOL_VERSION, IAgentGoalService, @@ -170,6 +174,43 @@ describe('Agent resume', () => { }); }); + it('restores a cancelled queued-turn gap before allocating the next turn', async () => { + const persistence = new RecordingAgentPersistence([ + resumeConfigRecord(), + { + type: 'turn.prompt', + input: [{ type: 'text', text: 'Historical prompt' }], + origin: { kind: 'user' }, + }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'Historical prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'context.append_loop_event', + event: { type: 'step.begin', uuid: 'historical-step', turnId: '0' }, + }, + { + type: 'context.append_loop_event', + event: { type: 'step.end', uuid: 'historical-step', turnId: '0' }, + }, + { type: 'turn.cancel', turnId: 1, target: 'queued' }, + ] as WireRecord[]); + const ctx = testAgent({ persistence, autoConfigure: false }); + + await ctx.restorePersisted(); + ctx.mockNextResponse({ type: 'text', text: 'Fresh response.' }); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Fresh prompt' }] }); + await ctx.untilTurnEnd(); + + expect(findRpcEvent(ctx.allEvents, 'turn.started')?.args).toMatchObject({ turnId: 2 }); + }); + it('projects restored pending tool results before later user messages', async () => { const persistence = new RecordingAgentPersistence([ resumeConfigRecord(), @@ -772,6 +813,47 @@ describe('Agent resume', () => { expect(ctx.context.get()[0]?.role).toBe('user'); expect(ctx.context.get()[1]?.role).toBe('assistant'); }); + + it('skips a fractional undo record on resume without corrupting checkpointed state', async () => { + const unexpected: unknown[] = []; + setUnexpectedErrorHandler((error) => unexpected.push(error)); + const persistence = new RecordingAgentPersistence([ + { + type: 'metadata', + protocol_version: '1.4', + created_at: 1, + }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'keep me' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { type: 'context.undo', count: 0.5 }, + ] as unknown as WireRecord[]); + const ctx = testAgent({ persistence, autoConfigure: false }); + + try { + await ctx.restorePersisted(); + + expect(ctx.context.get()).toHaveLength(1); + await expect(ctx.get(IAgentPlanService).status()).resolves.toBeNull(); + expect(unexpected).toHaveLength(1); + expect(unexpected[0]).toMatchObject({ + code: 'wire.unknown_record', + details: { type: 'context.undo', index: 1 }, + }); + } finally { + try { + await ctx.dispose(); + } finally { + resetUnexpectedErrorHandler(); + } + } + }); }); class RecordingAgentPersistence extends InMemoryWireRecordPersistence { diff --git a/packages/kap-server/src/routes/sessions.ts b/packages/kap-server/src/routes/sessions.ts index 2886902f94..1623004e09 100644 --- a/packages/kap-server/src/routes/sessions.ts +++ b/packages/kap-server/src/routes/sessions.ts @@ -21,7 +21,7 @@ * the native v2 services directly (`ISessionLifecycleService.fork` / `archive` / `restore`, * `IAgentFullCompactionService.begin`, `IAgentRPCService.cancel`); there is no * v1-only projection to centralize, so no adapter is involved. `undo` likewise - * calls `IAgentPromptService.undo` directly (it now throws + * calls `IAgentConversationUndoService.undo` directly (it throws * `session.undo_unavailable` with a structured reason) and only borrows * `ISessionLegacyService.status` for the cross-domain status rollup. The * `/sessions/{id}/children` endpoints call `ISessionLifecycleService.createChild` @@ -77,7 +77,7 @@ import { ErrorCodes, IAgentContextMemoryService, IAgentProfileService, - IAgentPromptService, + IAgentConversationUndoService, IAgentFullCompactionService, IAgentRPCService, IAuthSummaryService, @@ -691,12 +691,11 @@ export function registerSessionsRoutes(app: SessionRouteHost, core: Scope): void if (parsed.action === 'undo') { const body = undoSessionRequestSchema.parse(req.body); const agent = await resolveMainAgent(core, parsed.id); - // `prompt.undo` throws `session.undo_unavailable` (with a structured - // `reason`) when the history cannot satisfy `count`; it is a no-op - // until the precheck passes, so the post-undo read below always sees - // the cut applied. The status rollup stays in the legacy adapter - // (cross-domain) and is reused here verbatim. - agent.accessor.get(IAgentPromptService).undo(body.count); + // The conversation undo service throws `session.undo_unavailable` (with a + // structured `reason`) when fewer than `count` turns may be cut; + // it quiesces the loop/compaction first, so the post-undo read + // below always sees the cut applied. + await agent.accessor.get(IAgentConversationUndoService).undo(body.count); const history = agent.accessor.get(IAgentContextMemoryService).get(); requestLog(req)?.info({ session_id: parsed.id, action: 'undo' }, 'session action completed'); const [summary, status] = await Promise.all([ @@ -1219,6 +1218,7 @@ function sendMappedError( reply.send(errEnvelope(ErrorCode.SESSION_NOT_FOUND, err.message, requestId, err.stack)); return; case 'session.fork_active_turn': + case ErrorCodes.SESSION_BUSY: reply.send(errEnvelope(ErrorCode.SESSION_BUSY, err.message, requestId, err.stack)); return; case 'compaction.unable': diff --git a/packages/kap-server/test/sessions.test.ts b/packages/kap-server/test/sessions.test.ts index 69cda99b17..1e50124d03 100644 --- a/packages/kap-server/test/sessions.test.ts +++ b/packages/kap-server/test/sessions.test.ts @@ -1,7 +1,7 @@ /** * Scenario: v1-compatible session routes, including blocked-goal Web resume. * Responsibilities: verify HTTP envelopes, persisted reads, and session actions. - * Wiring: real kap-server; goal-resume cases observe the agent event stream. + * Wiring: real kap-server; route errors stub the agent service contract. * Run: `pnpm --filter @moonshot-ai/kap-server exec vitest run test/sessions.test.ts`. */ import { randomBytes } from 'node:crypto'; @@ -10,11 +10,14 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { inflateRawSync } from 'node:zlib'; -import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { + Error2, + ErrorCodes, IBootstrapService, type DomainEvent, + IAgentConversationUndoService, IAgentGoalService, IAgentLifecycleService, IEventBus, @@ -614,8 +617,35 @@ describe('server-v2 /api/v1/sessions', () => { expect(res.body.code).toBe(40911); expect(res.body.msg).toMatch(/nothing to undo/i); // The thrown Error2's stack is surfaced so operators can locate the - // source — the precheck/throw now lives in the native prompt service. - expect(res.body.stack).toEqual(expect.stringContaining('promptService')); + // source — the precheck/throw now lives in the undo service. + expect(res.body.stack).toEqual(expect.stringContaining('undoService')); + }); + + it('returns 40901 when :undo reports a busy session', async () => { + const created = await postJson('/api/v1/sessions', { + metadata: { cwd: home as string }, + }); + const session = (server as RunningServer).core.accessor + .get(ISessionLifecycleService) + .get(created.body.data.id); + if (session === undefined) throw new Error('expected live session'); + const agent = await session.accessor + .get(IAgentLifecycleService) + .create({ agentId: MAIN_AGENT_ID }); + const undo = vi + .spyOn(agent.accessor.get(IAgentConversationUndoService), 'undo') + .mockRejectedValue(new Error2(ErrorCodes.SESSION_BUSY, 'session is busy')); + + try { + const response = await postJson( + `/api/v1/sessions/${created.body.data.id}:undo`, + { count: 1 }, + ); + + expect(response.body.code).toBe(40901); + } finally { + undo.mockRestore(); + } }); it('rejects an unsupported action suffix (40001)', async () => {