diff --git a/.agents/skills/agent-core-dev/orient.md b/.agents/skills/agent-core-dev/orient.md index 3047c43d45..05d1b0ed7e 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`, `rewind` | | L7 | boundary / edge | `gateway`, `rpc`, `approval`, `question`, `*Legacy` | ## File-header comment convention diff --git a/.changeset/fix-undo-injections.md b/.changeset/fix-undo-injections.md new file mode 100644 index 0000000000..9f94819dda --- /dev/null +++ b/.changeset/fix-undo-injections.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix undo leaving injected context from removed turns in the conversation history. 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 af3ccf11cc..9b6e84ea5b 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/scripts/check-domain-layers.mjs b/packages/agent-core-v2/scripts/check-domain-layers.mjs index d02848869c..1ae6e900b8 100644 --- a/packages/agent-core-v2/scripts/check-domain-layers.mjs +++ b/packages/agent-core-v2/scripts/check-domain-layers.mjs @@ -214,6 +214,11 @@ const DOMAIN_LAYER = new Map([ ['sessionExport', 6], ['interaction', 6], ['sessionMetadata', 6], + // `undo` owns the undo pipeline (quiesce → context.undo → reconcile): it + // coordinates L4 agent domains (loop / prompt / contextMemory / + // fullCompaction), L5 task delivery, and `sessionMetadata`, so it sits in + // L6 beside the other cross-agent coordinators. + ['undo', 6], ['sessionActivity', 6], ['session', 6], ['terminal', 6], diff --git a/packages/agent-core-v2/src/agent/contextMemory/contextMemoryService.ts b/packages/agent-core-v2/src/agent/contextMemory/contextMemoryService.ts index 53485e99e4..94dc6cefff 100644 --- a/packages/agent-core-v2/src/agent/contextMemory/contextMemoryService.ts +++ b/packages/agent-core-v2/src/agent/contextMemory/contextMemoryService.ts @@ -1,24 +1,12 @@ /** * `contextMemory` domain (L4) — `IAgentContextMemoryService` implementation. * - * Owns the per-agent conversation history in the wire `ContextModel` - * (`ContextMessage[]`): reads through `wire.getModel`, writes through the - * v1 wire Ops (`append` / `appendLoopEvent` / `clear` / `undo` / - * `applyCompaction`). - * As the sole live mutation gateway for the history, it also cascades a - * (non-persisted) `context_size.measured` Op alongside every mutation that - * changes the measured prefix — `clear` resets it, `applyCompaction` adopts - * `tokensAfter`, and `undo` rebases it (to an estimate when the measured - * aggregate is truncated); `append` leaves the measured prefix untouched since - * new messages are the unmeasured tail (see `contextSizeService`). - * Splice-shaped mutations publish `context.spliced` from the live path only - * (replay rebuilds the Model silently and never invokes these methods), so - * existing subscribers observe the same change regardless of which Op was - * persisted. Messages - * are persisted without local ids — the on-disk record matches v1's field set - * and public message ids are derived from the transcript index. Blob - * dehydrate/rehydrate is declared on `ContextModel.blobs`. Bound at - * Agent scope. + * Owns per-agent conversation history through `wire`, maintains measurements + * with `contextSize`, and broadcasts live mutations through `event`. Every + * splice-shaped mutation (`clear` / `applyCompaction` / `undo`) publishes + * `context.spliced` from the live path only — replay rebuilds silently — and + * `undo` additionally rebases the measured token prefix to an estimate when + * the cut truncates it. Bound at Agent scope. */ import { Disposable } from '#/_base/di/lifecycle'; diff --git a/packages/agent-core-v2/src/agent/contextMemory/contextOps.ts b/packages/agent-core-v2/src/agent/contextMemory/contextOps.ts index 8446efadd0..b8c6a83b20 100644 --- a/packages/agent-core-v2/src/agent/contextMemory/contextOps.ts +++ b/packages/agent-core-v2/src/agent/contextMemory/contextOps.ts @@ -22,6 +22,11 @@ * `popSwarmModeReminder`) so the pop replays from the `swarm_mode.exit` record * itself, exactly like v1's restore-time `popMatchedMessage`. * + * `context.undo` counts conversation ticks with the single `isUndoAnchor` + * predicate (`./conversationTime`) — the same definition the checkpoint + * protocol pushes with, so anchor counting and checkpoint pushing can never + * drift apart. + * * Blob handling is declared as a `ModelBlobCodec` on `ContextModel.blobs`: * - `dehydrate(record, transform)`: at dispatch time, traverses message content * in `context.append_message` and `context.append_loop_event` records, @@ -43,6 +48,7 @@ import { createCompactionSummaryMessage, type ContextCompactionShapeInput, } from './compactionHandoff'; +import { isUndoAnchor, isValidUndoCount } from './conversationTime'; import { foldAppendMessage, foldLoopEvent, @@ -304,7 +310,7 @@ export function computeUndoCut(state: readonly ContextMessage[], count: number): stoppedAtCompaction = true; break; } - if (isRealUserPrompt(message)) { + if (isUndoAnchor(message)) { remaining--; removedCount++; cutIndex = i; @@ -353,21 +359,13 @@ export function formatUndoUnavailableMessage( } 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..c52324b165 100644 --- a/packages/agent-core-v2/src/agent/contextMemory/contextTranscript.ts +++ b/packages/agent-core-v2/src/agent/contextMemory/contextTranscript.ts @@ -1,33 +1,9 @@ /** - * `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 message and transcript consumers with full pre-compaction history, + * folded-context length, and stable turn identity while preserving live + * undo/clear semantics. Scope-agnostic. */ import { type ContentPart, type ToolCall } from '#/kosong/contract/message'; @@ -36,11 +12,11 @@ import type { WireRecord } from '#/wire/record'; import { COMPACT_USER_MESSAGE_MAX_TOKENS, collectCompactableUserMessages, - isRealUserInput, selectRecentUserMessages, } from './compactionHandoff'; +import { isUndoAnchor } from './conversationTime'; import type { LoopRecordedEvent } from './loopEventFold'; -import type { ContextMessage } from './types'; +import type { ContextMessage, PromptOrigin } from './types'; import { isVacuousContentPart } from './vacuousContent'; const TOOL_INTERRUPTED_ON_RESUME_OUTPUT = @@ -49,9 +25,19 @@ const TOOL_INTERRUPTED_ON_RESUME_OUTPUT = export interface ContextTranscript { readonly entries: readonly ContextMessage[]; readonly times: readonly (number | undefined)[]; + readonly turnIds: readonly (number | undefined)[]; + readonly turns: readonly ContextTranscriptTurn[]; + readonly stableTurnIds: boolean; readonly foldedLength: number; } +export interface ContextTranscriptTurn { + readonly turnId: number; + readonly input: readonly ContentPart[]; + readonly origin: PromptOrigin; + readonly time?: number; +} + export interface ContextTranscriptReducer { add(record: WireRecord): void; result(): ContextTranscript; @@ -70,8 +56,32 @@ interface MutableMessage { interface MutableEntry { message: MutableMessage; time?: number; + turnId?: number; + opensTurn: boolean; + order: number; +} + +interface MutableTurn extends ContextTranscriptTurn { + readonly order: number; +} + +interface PendingSteerBase { + readonly input: readonly ContentPart[]; + readonly origin: PromptOrigin; + readonly time?: number; + readonly order: number; + readonly expectsMessage: boolean; } +type PendingSteer = + | (PendingSteerBase & { readonly state: 'recorded' }) + | (PendingSteerBase & { readonly state: 'message-bound'; readonly entry: MutableEntry }) + | (PendingSteerBase & { + readonly state: 'turn-bound'; + readonly turnId: number; + readonly opensTurn: boolean; + }); + export function reduceContextTranscript(records: Iterable): ContextTranscript { const reducer = createContextTranscriptReducer(); for (const record of records) reducer.add(record); @@ -86,6 +96,15 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { const pendingToolResultIds = new Set(); let deferred: MutableEntry[] = []; let lastOpenStepUuid: string | undefined; + let turns: MutableTurn[] = []; + let nextTurnId = 0; + let currentTurnId: number | undefined; + const cancelledTurnIds = new Set(); + let pendingPromptAnchor: { readonly turnId: number } | undefined; + let pendingSteers: PendingSteer[] = []; + let stableTurnIds = true; + let recordOrder = 0; + const toolTurnIds = new Map(); const push = (...entries: MutableEntry[]): void => { transcript.push(...entries); @@ -109,16 +128,25 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { isError: true, }, time, + turnId: toolTurnIds.get(toolCallId), + opensTurn: false, + order: recordOrder, }); pendingToolResultIds.delete(toolCallId); + toolTurnIds.delete(toolCallId); } flushDeferredIfToolExchangeClosed(); }; - const resetOpenState = (): void => { + const resetOpenState = (clearPendingInput = true): void => { openSteps.clear(); pendingToolResultIds.clear(); deferred = []; lastOpenStepUuid = undefined; + toolTurnIds.clear(); + if (clearPendingInput) { + pendingPromptAnchor = undefined; + pendingSteers = []; + } }; const settleStep = (uuid: string): void => { const entry = openSteps.get(uuid); @@ -133,6 +161,21 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { }; const applyLoopEvent = (event: LoopRecordedEvent, time: number | undefined): void => { + const observedTurnId = readLoopTurnId(event); + if (observedTurnId !== undefined) { + resolvePendingSteers(observedTurnId); + currentTurnId = observedTurnId; + nextTurnId = Math.max(nextTurnId, observedTurnId + 1); + if (!turns.some((turn) => turn.turnId === observedTurnId)) { + turns.push({ + turnId: observedTurnId, + input: [], + origin: { kind: 'retry' }, + time, + order: recordOrder, + }); + } + } switch (event.type) { case 'step.begin': { closePendingToolResults(time); @@ -140,6 +183,9 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { const entry: MutableEntry = { message: { role: 'assistant', content: [], toolCalls: [] }, time, + turnId: observedTurnId ?? currentTurnId, + opensTurn: false, + order: recordOrder, }; push(entry); openSteps.set(event.uuid, entry); @@ -168,6 +214,7 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { }; openStep.message.toolCalls.push(call); pendingToolResultIds.add(event.toolCallId); + toolTurnIds.set(event.toolCallId, openStep.turnId); return; } case 'tool.result': { @@ -181,8 +228,12 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { isError: event.result.isError, }, time, + turnId: toolTurnIds.get(event.toolCallId), + opensTurn: false, + order: recordOrder, }); pendingToolResultIds.delete(event.toolCallId); + toolTurnIds.delete(event.toolCallId); flushDeferredIfToolExchangeClosed(); return; } @@ -192,24 +243,153 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { const applyUndo = (count: number): void => { if (count <= 0) return; let removedUserCount = 0; + let cutIndex = -1; + let cutEntry: MutableEntry | undefined; for (let i = transcript.length - 1; i >= clearFloor; i--) { - const message = transcript[i]!.message; - if (message.origin?.kind === 'injection') continue; - if (message.origin?.kind === 'compaction_summary') break; - transcript.splice(i, 1); - foldedLength = Math.max(0, foldedLength - 1); - if (isRealUserInput(message)) { - removedUserCount++; - if (removedUserCount >= count) break; + const entry = transcript[i]!; + const message = entry.message; + if (message.origin?.kind === 'compaction_summary') { + cutIndex = i + 1; + break; + } + if (!isUndoAnchor(message)) continue; + removedUserCount++; + if (removedUserCount >= count) { + cutIndex = i; + cutEntry = entry; + break; } } + if (cutIndex >= 0) { + const removedEntries = transcript.length - cutIndex; + transcript.splice(cutIndex); + foldedLength = Math.max(0, foldedLength - removedEntries); + } + if (cutEntry !== undefined) { + turns = turns.filter( + (turn) => + turn.order < cutEntry.order && + (!cutEntry.opensTurn || turn.turnId !== cutEntry.turnId), + ); + } + currentTurnId = turns.at(-1)?.turnId; resetOpenState(); }; const add = (record: WireRecord): void => { + recordOrder += 1; switch (record.type) { + case 'turn.prompt': { + if (pendingSteers.length > 0) stableTurnIds = false; + const input = readTurnInput(record); + const origin = readTurnOrigin(record); + while (cancelledTurnIds.delete(nextTurnId)) nextTurnId += 1; + const turnId = nextTurnId; + nextTurnId += 1; + currentTurnId = turnId; + turns.push({ turnId, input, origin, time: record.time, order: recordOrder }); + pendingSteers = []; + pendingPromptAnchor = + input.length > 0 && + isUndoAnchor({ role: 'user', content: [...input], toolCalls: [], origin }) + ? { turnId } + : undefined; + break; + } + case 'turn.steer': { + const input = readTurnInput(record); + const origin = readTurnOrigin(record); + pendingSteers.push({ + state: 'recorded', + input, + origin, + time: record.time, + order: recordOrder, + expectsMessage: input.length > 0, + }); + break; + } + case 'turn.cancel': { + const turnId = readTurnCancelId(record); + const target = readTurnCancelTarget(record); + if (target === 'queued') { + if (turnId !== undefined && turnId >= nextTurnId) { + cancelledTurnIds.add(turnId); + } + break; + } + if (target === undefined) { + stableTurnIds = false; + pendingSteers = []; + pendingPromptAnchor = undefined; + break; + } + if (pendingSteers.length === 0) { + if (turnId !== undefined && turnId >= nextTurnId) cancelledTurnIds.add(turnId); + pendingPromptAnchor = undefined; + break; + } + if ( + turnId === undefined || + (turnId !== currentTurnId && turnId !== nextTurnId) + ) { + stableTurnIds = false; + pendingSteers = []; + } else { + resolvePendingSteers(turnId, true); + currentTurnId = turnId; + nextTurnId = Math.max(nextTurnId, turnId + 1); + if (!turns.some((turn) => turn.turnId === turnId)) { + turns.push({ + turnId, + input: [], + origin: { kind: 'retry' }, + time: record.time, + order: recordOrder, + }); + } + } + pendingPromptAnchor = undefined; + break; + } case 'context.append_message': { - const entry = toMutableEntry(record['message'] as ContextMessage, record.time); + const message = record['message'] as ContextMessage; + const entry = toMutableEntry( + message, + record.time, + currentTurnId, + false, + recordOrder, + ); + if (isUndoAnchor(message) && pendingPromptAnchor !== undefined) { + entry.turnId = pendingPromptAnchor.turnId; + entry.opensTurn = true; + pendingPromptAnchor = undefined; + const redundantSteer = pendingSteers.findIndex( + (steer) => + steer.state !== 'message-bound' && + steer.expectsMessage && + hasSameOriginKind(steer.origin, message.origin), + ); + if (redundantSteer >= 0) pendingSteers.splice(redundantSteer, 1); + } else { + const steerIndex = pendingSteers.findIndex( + (candidate) => + candidate.state !== 'message-bound' && + candidate.expectsMessage && + hasSameOriginKind(candidate.origin, message.origin), + ); + const steer = pendingSteers[steerIndex]; + if (steer !== undefined) { + if (steer.state === 'turn-bound') { + entry.turnId = steer.turnId; + entry.opensTurn = steer.opensTurn; + pendingSteers.splice(steerIndex, 1); + } else { + pendingSteers[steerIndex] = { ...steer, state: 'message-bound', entry }; + } + } + } if (pendingToolResultIds.size > 0) deferred.push(entry); else push(entry); break; @@ -226,9 +406,12 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { origin: { kind: 'compaction_summary' }, }, time: record.time, + turnId: currentTurnId, + opensTurn: false, + order: recordOrder, }); foldedLength = recoverFoldedLength(record, transcript, clearFloor, foldedLength); - resetOpenState(); + resetOpenState(false); break; } case 'context.undo': @@ -249,12 +432,58 @@ export function createContextTranscriptReducer(): ContextTranscriptReducer { result: () => ({ entries: transcript.map((e) => e.message), times: transcript.map((e) => e.time), + turnIds: transcript.map((e) => e.turnId), + turns: turns + .toSorted((a, b) => a.turnId - b.turnId) + .map(({ order: _order, ...turn }) => turn), + stableTurnIds, foldedLength, }), }; + + function resolvePendingSteers(observedTurnId: number, settleUnbound = false): void { + const unresolved = pendingSteers.filter( + (steer) => steer.state !== 'turn-bound' && steer.expectsMessage, + ); + const opensNewTurn = + unresolved.length > 0 && + (currentTurnId === undefined || observedTurnId !== currentTurnId); + if (opensNewTurn) { + const opener = unresolved[0]!; + if (!turns.some((turn) => turn.turnId === observedTurnId)) { + turns.push({ + turnId: observedTurnId, + input: opener.input, + origin: opener.origin, + time: opener.time, + order: opener.order, + }); + } + } + let unresolvedIndex = 0; + pendingSteers = pendingSteers.flatMap((steer) => { + if (steer.state === 'turn-bound') return settleUnbound ? [] : [steer]; + if (!steer.expectsMessage) return []; + const opensTurn = opensNewTurn && unresolvedIndex === 0; + unresolvedIndex += 1; + if (steer.state === 'message-bound') { + steer.entry.turnId = observedTurnId; + steer.entry.opensTurn = opensTurn; + return []; + } + if (settleUnbound) return []; + return [{ ...steer, state: 'turn-bound', turnId: observedTurnId, opensTurn }]; + }); + } } -function toMutableEntry(message: ContextMessage, time: number | undefined): MutableEntry { +function toMutableEntry( + message: ContextMessage, + time: number | undefined, + turnId: number | undefined, + opensTurn: boolean, + order: number, +): MutableEntry { return { message: { ...(message.id !== undefined ? { id: message.id } : {}), @@ -266,9 +495,50 @@ function toMutableEntry(message: ContextMessage, time: number | undefined): Muta ...(message.origin !== undefined ? { origin: message.origin } : {}), }, time, + turnId, + opensTurn, + order, }; } +function readLoopTurnId(event: LoopRecordedEvent): number | undefined { + if (!('turnId' in event) || event.turnId === undefined) return undefined; + const value = Number.parseInt(event.turnId, 10); + return Number.isInteger(value) && value >= 0 ? value : undefined; +} + +function readTurnInput(record: WireRecord): readonly ContentPart[] { + const input = record['input']; + return Array.isArray(input) ? (input as readonly ContentPart[]) : []; +} + +function readTurnOrigin(record: WireRecord): PromptOrigin { + const origin = record['origin']; + if (origin !== null && typeof origin === 'object' && 'kind' in origin) { + return origin as PromptOrigin; + } + return { kind: 'user' }; +} + +function readTurnCancelId(record: WireRecord): number | undefined { + const turnId = record['turnId']; + return typeof turnId === 'number' && Number.isInteger(turnId) && turnId >= 0 + ? turnId + : undefined; +} + +function readTurnCancelTarget(record: WireRecord): 'active' | 'queued' | undefined { + const target = record['target']; + return target === 'active' || target === 'queued' ? target : undefined; +} + +function hasSameOriginKind( + left: PromptOrigin, + right: ContextMessage['origin'], +): boolean { + return left.kind === (right?.kind ?? 'user'); +} + function recoverFoldedLength( record: WireRecord, transcript: readonly MutableEntry[], 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..5314780eb7 --- /dev/null +++ b/packages/agent-core-v2/src/agent/contextMemory/conversationTime.ts @@ -0,0 +1,73 @@ +/** + * `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 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/conversationUndoReconciliation.ts b/packages/agent-core-v2/src/agent/contextMemory/conversationUndoReconciliation.ts new file mode 100644 index 0000000000..0e9e2bf25b --- /dev/null +++ b/packages/agent-core-v2/src/agent/contextMemory/conversationUndoReconciliation.ts @@ -0,0 +1,66 @@ +/** + * `contextMemory` domain (L4) — Agent-scoped post-undo reconciliation registry. + * + * Hosts state-repair and derived-projection participants for the undo + * coordinator. Bound at Agent scope. + */ + +import { InstantiationType } from '#/_base/di/extensions'; +import { createDecorator } from '#/_base/di/instantiation'; +import { Disposable, toDisposable, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; + +export interface AgentConversationUndoReconciliationParticipant { + readonly id: string; + readonly phase?: AgentConversationUndoReconciliationPhase; + reconcileAfterUndo(): Promise; +} + +export type AgentConversationUndoReconciliationPhase = 'state' | 'projection'; + +export interface IAgentConversationUndoReconciliationRegistry { + readonly _serviceBrand: undefined; + + register(participant: AgentConversationUndoReconciliationParticipant): IDisposable; + list(): readonly AgentConversationUndoReconciliationParticipant[]; +} + +export const IAgentConversationUndoReconciliationRegistry = + createDecorator( + 'agentConversationUndoReconciliationRegistry', + ); + +class AgentConversationUndoReconciliationRegistry + extends Disposable + implements IAgentConversationUndoReconciliationRegistry +{ + declare readonly _serviceBrand: undefined; + + private readonly participants = new Map(); + + register(participant: AgentConversationUndoReconciliationParticipant): IDisposable { + if (this.participants.has(participant.id)) { + throw new Error( + `Conversation undo reconciliation 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 AgentConversationUndoReconciliationParticipant[] { + return [...this.participants.values()]; + } +} + +registerScopedService( + LifecycleScope.Agent, + IAgentConversationUndoReconciliationRegistry, + AgentConversationUndoReconciliationRegistry, + InstantiationType.Eager, + 'conversationUndoReconciliation', +); diff --git a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts index e89c49a6ab..a179ed8690 100644 --- a/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts +++ b/packages/agent-core-v2/src/agent/fullCompaction/fullCompactionService.ts @@ -326,8 +326,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 3b16df8ca6..c894edd505 100644 --- a/packages/agent-core-v2/src/agent/loop/loopService.ts +++ b/packages/agent-core-v2/src/agent/loop/loopService.ts @@ -101,9 +101,11 @@ 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 nextReservedTurnId: number | undefined; private readonly settleWaiters: Array<() => void> = []; + private quiescenceDepth = 0; private activeRequestTrace: LLMRequestTrace | undefined; private lastRequestTraceId: string | undefined; private disposing = false; @@ -131,6 +133,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(); } @@ -141,6 +147,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': @@ -163,10 +181,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 { @@ -199,10 +213,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; } @@ -212,7 +250,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'; @@ -226,12 +264,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) => { @@ -240,7 +283,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(); @@ -296,6 +343,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) { @@ -344,7 +399,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(); @@ -475,6 +530,7 @@ export class AgentLoopService extends Disposable implements IAgentLoopService { if (step.state === 'queued' || step.state === 'running') step.cancel(reason); } this.activeTurnJob = undefined; + this.maybeSettle(); } registerLoopErrorHandler( @@ -1013,6 +1069,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 60e3dbd4f2..fce78c891d 100644 --- a/packages/agent-core-v2/src/agent/plan/planOps.ts +++ b/packages/agent-core-v2/src/agent/plan/planOps.ts @@ -1,38 +1,33 @@ /** - * `plan` domain (L4) — wire Model (`PlanModel`) and the `plan_mode.enter` - * (`planModeEnter`) / `plan_mode.cancel` (`planModeCancel`) / `plan_mode.exit` - * (`planModeExit`) Ops that mirror the plan-mode lifecycle into a persisted, - * replayable `{ active, id }` state. + * `plan` domain (L4) — persists plan-mode conversation state. * - * The Model holds the persistent, replayable fields — whether plan mode is - * active and the plan id. The persisted records carry 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`. - * Each `apply` returns the same reference on a no-op (re-entering the same - * plan, or cancelling/exiting while already inactive) so the wire's - * reference-equality gate stays quiet. The side effects — `telemetryContext` - * mode, plan-directory/file fs I/O, and the `agent.status.updated` planMode - * slice — are NOT part of `apply`: they run after `wire.dispatch` on the live - * path, and `wire.replay` rebuilds the Model silently from the persisted - * `plan_mode.*` records (seeded by `sessionLifecycle`). The legacy - * `toReplay: plan_updated` projection is dropped (inert — nothing reads it). - * Consumed by the Agent-scope `planService`. + * Keeps plan mode aligned with conversation undo through `contextMemory`'s + * checkpoint protocol and exposes the state consumed by the Agent-scope plan + * service. */ import { z } from 'zod'; -import { defineModel } from '#/wire/model'; +import { + defineCheckpointedModel, + type Checkpointed, +} from '#/agent/contextMemory/conversationTime'; export interface PlanState { readonly active: boolean; readonly id?: string; } -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 : { active: true, id: p.id }), + apply: (s, p) => + s.current.active && s.current.id === p.id + ? s + : { ...s, current: { active: true, id: p.id } }, toEvent: () => ({ type: 'agent.status.updated' as const, planMode: true }), }); @@ -46,12 +41,14 @@ 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 }), + apply: (s) => + s.current.active ? { ...s, current: { active: false } } : 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 }), + apply: (s) => + s.current.active ? { ...s, current: { active: false } } : s, toEvent: () => ({ type: 'agent.status.updated' as const, planMode: false }), }); diff --git a/packages/agent-core-v2/src/agent/plan/planService.ts b/packages/agent-core-v2/src/agent/plan/planService.ts index 9c63cec358..7ca6ef3488 100644 --- a/packages/agent-core-v2/src/agent/plan/planService.ts +++ b/packages/agent-core-v2/src/agent/plan/planService.ts @@ -18,6 +18,7 @@ import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory' import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { PlanModeInjection } from '#/agent/plan/injection/planModeInjection'; import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; +import { IEventBus } from '#/app/event/eventBus'; import { IAgentTelemetryContextService } from '#/app/telemetry/agentTelemetryContext'; import { IHostFileSystem } from '#/os/interface/hostFileSystem'; import { ISessionContext } from '#/session/sessionContext/sessionContext'; @@ -42,6 +43,7 @@ export class AgentPlanService extends Disposable implements IAgentPlanService { @IHostFileSystem private readonly hostFs: IHostFileSystem, @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, @@ -54,24 +56,31 @@ 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)); } 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 { @@ -118,7 +127,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 f3bf39b49b..a6af045452 100644 --- a/packages/agent-core-v2/src/agent/prompt/promptService.ts +++ b/packages/agent-core-v2/src/agent/prompt/promptService.ts @@ -14,7 +14,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'; @@ -157,13 +156,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); 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 cfa198b40a..ae0e878f51 100644 --- a/packages/agent-core-v2/src/agent/rpc/core-api.ts +++ b/packages/agent-core-v2/src/agent/rpc/core-api.ts @@ -313,7 +313,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 cd27ffbea8..58570e8a62 100644 --- a/packages/agent-core-v2/src/agent/rpc/rpcService.ts +++ b/packages/agent-core-v2/src/agent/rpc/rpcService.ts @@ -19,6 +19,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'; @@ -64,6 +65,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, @@ -127,10 +130,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 ff59ca0bf1..b161bbab2d 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -1,30 +1,12 @@ /** * `task` domain (L5) — `AgentTaskService` implementation. * - * Owns the agent's registry of running and restored tasks: - * registers and drives tasks to completion, retains a bounded output ring, - * persists task state and output through task persistence rooted at the - * agent's own scope (v1's per-agent `/agents//tasks/` - * layout), lets only the main agent read through the previous v2 - * session-level task root without writing back to it, reads - * limits through `config`, records lifecycle and broadcasts through `wire` - * (`task.started` / `task.terminated` Ops into `TaskModel`, plus the matching - * signals), restores ghosts through a single `wire.hooks.onDidRestore` hook - * (wire replay -> disk load -> reconcile, in that order), delivers live - * terminal notifications by enqueueing `TaskNotificationStepRequest`s onto - * `loop` with `activeOrNewTurn` admission (mid-turn ones fold into the active turn's - * following step; idle ones launch a fresh turn themselves, matching v1's - * `turn.steer`, so the model consumes the notification without waiting for - * the user), silently appends restored notifications through `contextMemory`, - * re-surfaces active tasks through `contextInjector` after compaction, and - * requests every owned task to stop on session close (`stopAllOnExit` — v1's - * `stopBackgroundTasksOnExit`) with configurable SIGTERM grace and SIGKILL - * escalation. `keepAliveOnExit` skips task-manager teardown so independently - * living external work such as processes can continue; Session-scoped agents - * remain governed by the Session lifecycle. Scope disposal paths that bypass - * graceful close synchronously cancel/abort work and immediately attempt a - * best-effort force-stop to reduce the risk of surviving child processes. - * Bound at Agent scope. + * Owns task lifecycle, retained output, persistence, restoration, and terminal + * notification delivery. Coordinates App-level task handles, `loop`, + * `contextMemory`, `contextInjector`, `config`, Agent and Session identity, + * atomic-document and file storage, `eventBus`, `wire`, and `telemetry`; + * notification delivery follows conversation undo through the checkpoint and + * reconciliation contracts. Bound at Agent scope. */ import { randomBytes } from 'node:crypto'; @@ -36,12 +18,15 @@ import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; import type { ContentPart } from '#/kosong/contract/message'; import { Disposable } from '#/_base/di/lifecycle'; +import { ILogService } from '#/_base/log/log'; import { abortable, userCancellationReason, } 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 { IAgentConversationUndoReconciliationRegistry } from '#/agent/contextMemory/conversationUndoReconciliation'; import type { ContextMessage, TaskOrigin } from '#/agent/contextMemory/types'; import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { IAgentLoopService } from '#/agent/loop/loop'; @@ -61,7 +46,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, @@ -107,17 +91,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]; }, }, ); @@ -199,7 +181,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, @@ -207,6 +192,10 @@ export class TaskNotificationStepRequest extends MessageStepRequest { admission: 'activeOrNewTurn', }); } + + override onWillMaterialize(): void { + this.onWillDeliver?.(); + } } export class AgentTaskService extends Disposable implements IAgentTaskService { @@ -214,9 +203,12 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { private readonly tasks = new Map(); private readonly ghosts = new Map(); + private readonly buildingNotificationKeys = new Set(); private readonly scheduledNotificationKeys = new Set(); private readonly deliveredNotificationKeys = new Set(); + private readonly pendingNotificationRequests = new Map(); private readonly persistence: AgentTaskPersistence; + private notificationRestoreQueue: Promise = Promise.resolve(); private activeTaskReminderPending = false; constructor( @@ -232,6 +224,9 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { @IEventBus private readonly eventBus: IEventBus, @IAgentContextInjectorService injector: IAgentContextInjectorService, @IAgentLoopService private readonly loop: IAgentLoopService, + @IAgentConversationUndoReconciliationRegistry + conversationUndoReconciliation: IAgentConversationUndoReconciliationRegistry, + @ILogService private readonly log: ILogService, ) { super(); const fallbackRoot = @@ -245,9 +240,15 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { byteStore, fallbackRoot, ); + this._register( + conversationUndoReconciliation.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(); @@ -470,6 +471,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; @@ -976,7 +992,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); } @@ -1001,17 +1019,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); @@ -1042,35 +1085,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 { @@ -1093,7 +1152,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 39dbc581f4..e798e9eb8f 100644 --- a/packages/agent-core-v2/src/agent/toolSelect/toolSelectService.ts +++ b/packages/agent-core-v2/src/agent/toolSelect/toolSelectService.ts @@ -66,14 +66,19 @@ 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(); }), ); } + 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..47a12af8bb --- /dev/null +++ b/packages/agent-core-v2/src/agent/undo/undoService.ts @@ -0,0 +1,249 @@ +/** + * `undo` domain (L6) — `IAgentConversationUndoService` implementation. + * + * Owns idle conversation undo coordination and restored observable state. + * Coordinates `contextMemory`, conversation reconciliation, `fullCompaction`, + * `loop`, `prompt`, Agent and Session identity, `sessionMetadata`, `event`, + * `eventBus`, `telemetry`, and `wire`. Bound at Agent scope. + */ + +import { InstantiationType } from '#/_base/di/extensions'; +import { Disposable, type IDisposable } from '#/_base/di/lifecycle'; +import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; +import { ILogService } from '#/_base/log/log'; +import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; +import { + IAgentConversationUndoReconciliationRegistry, + type AgentConversationUndoReconciliationPhase, +} from '#/agent/contextMemory/conversationUndoReconciliation'; +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, + @IAgentConversationUndoReconciliationRegistry + private readonly participants: IAgentConversationUndoReconciliationRegistry, + @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()); + 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('state'); + await this.flushAfterCommit('state reconciliation'); + await this.reconcileParticipants('projection'); + await this.flushAfterCommit('projection 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(): number { + let depth = Number.POSITIVE_INFINITY; + for (const def of CHECKPOINTED_MODELS) { + const state = this.wire.getModel(def) as Checkpointed; + depth = Math.min(depth, state.checkpoints.length); + } + return depth; + } + + 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 = this.checkpointDepth(); + if (depth >= turns) return; + throw new Error2( + ErrorCodes.SESSION_UNDO_UNAVAILABLE, + formatUndoUnavailableMessage({ + ok: false, + reason: 'compaction_boundary', + requested: turns, + undoable: depth, + }), + { + details: { + reason: 'compaction_boundary', + requestedCount: turns, + undoableCount: depth, + }, + }, + ); + } + + private async reconcileParticipants( + phase: AgentConversationUndoReconciliationPhase, + ): Promise { + const participants = this.participants + .list() + .filter((participant) => (participant.phase ?? 'state') === phase); + 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, + InstantiationType.Eager, + 'undo', +); diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 6f833e0b32..17ca327b43 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -386,6 +386,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/conversationUndoReconciliation'; +export * from '#/agent/contextMemory/conversationTime'; export * from '#/agent/contextMemory/loopEventFold'; export * from '#/agent/contextMemory/messageId'; export * from '#/agent/contextMemory/messageProjection'; @@ -454,6 +456,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 a7dcd94fb4..e55cfb831b 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'; @@ -25,6 +17,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'; @@ -42,6 +35,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, @@ -73,7 +67,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 { @@ -93,7 +87,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 { @@ -102,6 +98,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 { @@ -130,9 +138,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..0111cbd06f 100644 --- a/packages/agent-core-v2/test/agent/contextMemory/context.test.ts +++ b/packages/agent-core-v2/test/agent/contextMemory/context.test.ts @@ -579,12 +579,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 +592,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 +629,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 +652,7 @@ describe('Agent context', () => { }, ); - ctx.undoHistory(1); + await ctx.undoHistory(1); expect(context.get()).toEqual([ expect.objectContaining({ 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..f601f5ffe5 100644 --- a/packages/agent-core-v2/test/agent/contextMemory/contextTranscript.test.ts +++ b/packages/agent-core-v2/test/agent/contextMemory/contextTranscript.test.ts @@ -67,6 +67,22 @@ function undo(count: number): WireRecord { return { type: 'context.undo', count }; } +function promptTurn(input: string, origin: PromptOrigin): WireRecord { + return { type: 'turn.prompt', input: [{ type: 'text', text: input }], origin }; +} + +function steerTurn(input: string): WireRecord { + return { + type: 'turn.steer', + input: [{ type: 'text', text: input }], + origin: { kind: 'user' }, + }; +} + +function emptySteerTurn(): WireRecord { + return { type: 'turn.steer', input: [], origin: { kind: 'user' } }; +} + function texts(result: ContextTranscript): string[] { return result.entries.map((m) => m.content.map((p) => (p.type === 'text' ? p.text : `[${p.type}]`)).join(''), @@ -172,6 +188,398 @@ describe('reduceContextTranscript', () => { expect(texts(result)).toEqual(['message A', 'reply A']); }); + it('removes a post-anchor injection when undo cuts its turn', () => { + const result = reduceContextTranscript([ + promptTurn('undo me', { kind: 'user' }), + appendMessage(userMessage('undo me', { kind: 'user' })), + appendMessage(assistantMessage('undone answer')), + appendMessage(userMessage('internal reminder', { kind: 'injection', variant: 'todo' })), + undo(1), + promptTurn('keep me', { kind: 'user' }), + appendMessage(userMessage('keep me', { kind: 'user' })), + appendMessage(assistantMessage('kept answer')), + ]); + + expect(texts(result)).toEqual(['keep me', 'kept answer']); + expect(result.turnIds).toEqual([1, 1]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([1]); + }); + + it('keeps stable turn ids when undo removes a steer after a hidden retry turn', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ + type: 'content.part', + stepUuid: 's0', + turnId: '0', + part: { type: 'text', text: 'a0' }, + }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + promptTurn('', { kind: 'retry' }), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ + type: 'content.part', + stepUuid: 's1', + turnId: '1', + part: { type: 'text', text: 'retry answer' }, + }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + promptTurn('u2', { kind: 'user' }), + appendMessage(userMessage('u2', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's2', turnId: '2' }), + loopEvent({ + type: 'content.part', + stepUuid: 's2', + turnId: '2', + part: { type: 'text', text: 'before steer' }, + }), + loopEvent({ type: 'step.end', uuid: 's2', turnId: '2' }), + steerTurn('remove me'), + appendMessage(userMessage('remove me', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's3', turnId: '2' }), + loopEvent({ + type: 'content.part', + stepUuid: 's3', + turnId: '2', + part: { type: 'text', text: 'after steer' }, + }), + loopEvent({ type: 'step.end', uuid: 's3', turnId: '2' }), + undo(1), + ]); + + expect(texts(result)).toEqual(['u0', 'a0', 'retry answer', 'u2', 'before steer']); + expect(result.turnIds).toEqual([0, 0, 1, 2, 2]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0, 1, 2]); + }); + + it('promotes a legacy idle steer to the new turn reported by its loop events', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ + type: 'content.part', + stepUuid: 's0', + turnId: '0', + part: { type: 'text', text: 'a0' }, + }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('u1'), + appendMessage(userMessage('u1', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ + type: 'content.part', + stepUuid: 's1', + turnId: '1', + part: { type: 'text', text: 'a1' }, + }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + ]); + + expect(texts(result)).toEqual(['u0', 'a0', 'u1', 'a1']); + expect(result.turnIds).toEqual([0, 0, 1, 1]); + expect(result.turns).toMatchObject([ + { turnId: 0, input: [{ type: 'text', text: 'u0' }], origin: { kind: 'user' } }, + { turnId: 1, input: [{ type: 'text', text: 'u1' }], origin: { kind: 'user' } }, + ]); + }); + + it('removes the promoted legacy idle-steer turn when its anchor is undone', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ + type: 'content.part', + stepUuid: 's0', + turnId: '0', + part: { type: 'text', text: 'a0' }, + }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('u1'), + appendMessage(userMessage('u1', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ + type: 'content.part', + stepUuid: 's1', + turnId: '1', + part: { type: 'text', text: 'a1' }, + }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + undo(1), + ]); + + expect(texts(result)).toEqual(['u0', 'a0']); + expect(result.turnIds).toEqual([0, 0]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0]); + }); + + it.each(['messages-first', 'event-first'] as const)( + 'assigns multiple same-origin legacy steers FIFO when %s', + (order) => { + const steers = [steerTurn('u1'), steerTurn('u2')]; + const messages = [ + appendMessage(userMessage('u1', { kind: 'user' })), + appendMessage(userMessage('u2', { kind: 'user' })), + ]; + const stepBegin = loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }); + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + ...steers, + ...(order === 'messages-first' ? [...messages, stepBegin] : [stepBegin, ...messages]), + loopEvent({ + type: 'content.part', + stepUuid: 's1', + turnId: '1', + part: { type: 'text', text: 'a1' }, + }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + ]); + + const idsByText = new Map(texts(result).map((text, index) => [text, result.turnIds[index]])); + expect(idsByText.get('u1')).toBe(1); + expect(idsByText.get('u2')).toBe(1); + expect(idsByText.get('a1')).toBe(1); + expect(result.turns[1]).toMatchObject({ + turnId: 1, + input: [{ type: 'text', text: 'u1' }], + }); + }, + ); + + it('does not retain an empty steer after observing its turn', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + emptySteerTurn(), + loopEvent({ type: 'step.begin', uuid: 'empty', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 'empty', turnId: '0' }), + steerTurn('u1'), + appendMessage(userMessage('u1', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + ]); + + expect(result.turnIds).toEqual([0, 1]); + expect(result.turns).toMatchObject([ + { turnId: 0, input: [{ type: 'text', text: 'u0' }] }, + { turnId: 1, input: [{ type: 'text', text: 'u1' }] }, + ]); + expect(result.stableTurnIds).toBe(true); + }); + + it('uses the message-producing steer as the opener when an empty steer precedes it', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + emptySteerTurn(), + steerTurn('u1'), + appendMessage(userMessage('u1', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + promptTurn('u2', { kind: 'user' }), + appendMessage(userMessage('u2', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's2', turnId: '2' }), + loopEvent({ type: 'step.end', uuid: 's2', turnId: '2' }), + ]); + + expect(result.turnIds).toEqual([0, 1, 2]); + expect(result.turns).toMatchObject([ + { turnId: 0, input: [{ type: 'text', text: 'u0' }] }, + { turnId: 1, input: [{ type: 'text', text: 'u1' }] }, + { turnId: 2, input: [{ type: 'text', text: 'u2' }] }, + ]); + expect(result.stableTurnIds).toBe(true); + }); + + it('removes the message-producing steer turn on undo when an empty steer preceded it', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + emptySteerTurn(), + steerTurn('u1'), + appendMessage(userMessage('u1', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + undo(1), + ]); + + expect(texts(result)).toEqual(['u0']); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0]); + }); + + it('keeps restored task notifications off removed turn ids after undo', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + ...assistantStep('s0', 'a0'), + promptTurn('u1', { kind: 'user' }), + appendMessage(userMessage('u1', { kind: 'user' })), + ...assistantStep('s1', 'a1'), + undo(1), + appendMessage( + userMessage('task done', { + kind: 'task', + taskId: 'bash-001', + status: 'completed', + notificationId: 'task:bash-001:completed', + }), + ), + ]); + + expect(texts(result)).toEqual(['u0', 'a0', 'task done']); + expect(result.turnIds).toEqual([0, 0, 0]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0]); + expect(result.stableTurnIds).toBe(true); + }); + + it('uses an active cancel turn id to resolve an idle steer before its first loop event', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('cancelled'), + appendMessage(userMessage('cancelled', { kind: 'user' })), + { type: 'turn.cancel', turnId: 1, target: 'active' }, + promptTurn('u2', { kind: 'user' }), + appendMessage(userMessage('u2', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's2', turnId: '2' }), + loopEvent({ type: 'step.end', uuid: 's2', turnId: '2' }), + ]); + + expect(result.turnIds).toEqual([0, 1, 2]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0, 1, 2]); + expect(result.turns[1]).toMatchObject({ + input: [{ type: 'text', text: 'cancelled' }], + origin: { kind: 'user' }, + }); + expect(result.stableTurnIds).toBe(true); + }); + + it('marks an id-less legacy cancel as ambiguous instead of exposing partial stable ids', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('cancelled'), + appendMessage(userMessage('cancelled', { kind: 'user' })), + { type: 'turn.cancel' }, + promptTurn('next', { kind: 'user' }), + appendMessage(userMessage('next', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's2', turnId: '2' }), + loopEvent({ type: 'step.end', uuid: 's2', turnId: '2' }), + ]); + + expect(result.stableTurnIds).toBe(false); + }); + + it('does not create a phantom turn for an orphan legacy cancel id', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + { type: 'turn.cancel', turnId: 99 }, + promptTurn('u1', { kind: 'user' }), + appendMessage(userMessage('u1', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's1', turnId: '1' }), + loopEvent({ type: 'step.end', uuid: 's1', turnId: '1' }), + ]); + + expect(result.turnIds).toEqual([0, 1]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0, 1]); + expect(result.stableTurnIds).toBe(false); + }); + + it('preserves a queued-turn id gap after its cancel record', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + { type: 'turn.cancel', turnId: 1, target: 'queued' }, + promptTurn('u2', { kind: 'user' }), + appendMessage(userMessage('u2', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's2', turnId: '2' }), + loopEvent({ type: 'step.end', uuid: 's2', turnId: '2' }), + ]); + + expect(result.turnIds).toEqual([0, 2]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0, 2]); + expect(result.stableTurnIds).toBe(true); + }); + + it('keeps an active-turn steer separate from a concurrently cancelled queued turn', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('steered into active'), + appendMessage(userMessage('steered into active', { kind: 'user' })), + { type: 'turn.cancel', turnId: 1, target: 'queued' }, + loopEvent({ type: 'step.begin', uuid: 's0-next', turnId: '0' }), + loopEvent({ + type: 'content.part', + stepUuid: 's0-next', + turnId: '0', + part: { type: 'text', text: 'active answer' }, + }), + loopEvent({ type: 'step.end', uuid: 's0-next', turnId: '0' }), + promptTurn('u2', { kind: 'user' }), + appendMessage(userMessage('u2', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's2', turnId: '2' }), + loopEvent({ type: 'step.end', uuid: 's2', turnId: '2' }), + ]); + + expect(result.turnIds).toEqual([0, 0, 0, 2]); + expect(result.turns.map((turn) => turn.turnId)).toEqual([0, 2]); + expect(result.stableTurnIds).toBe(true); + }); + + it('marks a legacy cancel ambiguous when a pending steer could belong to another turn', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('pending'), + appendMessage(userMessage('pending', { kind: 'user' })), + { type: 'turn.cancel', turnId: 1 }, + ]); + + expect(result.turns.map((turn) => turn.turnId)).toEqual([0]); + expect(result.stableTurnIds).toBe(false); + }); + + it('marks a cancel id ambiguous when it cannot resolve a pending steer', () => { + const result = reduceContextTranscript([ + promptTurn('u0', { kind: 'user' }), + appendMessage(userMessage('u0', { kind: 'user' })), + loopEvent({ type: 'step.begin', uuid: 's0', turnId: '0' }), + loopEvent({ type: 'step.end', uuid: 's0', turnId: '0' }), + steerTurn('pending'), + appendMessage(userMessage('pending', { kind: 'user' })), + { type: 'turn.cancel', turnId: 99 }, + ]); + + expect(result.turns.map((turn) => turn.turnId)).toEqual([0]); + expect(result.stableTurnIds).toBe(false); + }); + 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 3fe4f6b895..2f8310bd62 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 6191e48afc..9217c88fe3 100644 --- a/packages/agent-core-v2/test/agent/loop/stubs.ts +++ b/packages/agent-core-v2/test/agent/loop/stubs.ts @@ -69,6 +69,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 ca7fa17992..1599260a6d 100644 --- a/packages/agent-core-v2/test/agent/plan/plan.test.ts +++ b/packages/agent-core-v2/test/agent/plan/plan.test.ts @@ -752,11 +752,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 3c6e1c6109..f1aa376854 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 { @@ -59,20 +60,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([ @@ -93,11 +94,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([ @@ -121,6 +122,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(); @@ -136,7 +156,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', }); @@ -152,6 +172,6 @@ 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); }); }); 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 9c5dbfa54f..888ce76df1 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 37b468cd05..2ce7587320 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 f873dfeb97..6c9a40d651 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 { IAgentConversationUndoReconciliationRegistry } from '#/agent/contextMemory/conversationUndoReconciliation'; import { IAgentContextInjectorService, type ContextInjectionContext, @@ -45,6 +47,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'; @@ -87,6 +90,11 @@ describe('AgentTaskService', () => { ix = disposables.add(new TestInstantiationService()); eventBus = disposables.add(new EventBusService()); injectionProviders = new Map(); + ix.stub(ILogService, stubLog()); + ix.stub(IAgentConversationUndoReconciliationRegistry, { + register: () => toDisposable(() => {}), + list: () => [], + }); ix.stub(IWireService, stubWireService()); ix.stub(IEventBus, eventBus); ix.stub(IAgentContextInjectorService, { @@ -393,6 +401,11 @@ describe('AgentTaskService', () => { captureRestoreHook?: (hook: RestoreHook) => void, ): TestInstantiationService { const ix = disposables.add(new TestInstantiationService()); + ix.stub(ILogService, stubLog()); + ix.stub(IAgentConversationUndoReconciliationRegistry, { + 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 db6e07250f..e287fc38bd 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'; @@ -185,7 +186,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 cf63a8ecef..b86471ef04 100644 --- a/packages/agent-core-v2/test/agent/toolSelect/toolSelectService.test.ts +++ b/packages/agent-core-v2/test/agent/toolSelect/toolSelectService.test.ts @@ -223,6 +223,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..9159b17092 --- /dev/null +++ b/packages/agent-core-v2/test/agent/undo/undo.test.ts @@ -0,0 +1,499 @@ +/** + * 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 { IAgentConversationUndoReconciliationRegistry } from '#/agent/contextMemory/conversationUndoReconciliation'; +import { contextApplyCompaction } from '#/agent/contextMemory/contextOps'; +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 { 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('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 rebuilding projections', 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(IAgentConversationUndoReconciliationRegistry); + participants.register({ + id: 'test.state', + reconcileAfterUndo: async () => { + order.push('state'); + }, + }); + participants.register({ + id: 'test.projection', + phase: 'projection', + reconcileAfterUndo: async () => { + order.push('projection'); + }, + }); + 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', + 'projection', + 'flush', + 'context.undone', + ]); + } finally { + subscription.dispose(); + flush.mockRestore(); + } + }); + + it.each([ + [1, []], + [2, ['state']], + [3, ['state', 'projection']], + ] 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(IAgentConversationUndoReconciliationRegistry); + participants.register({ + id: 'test.flush-failure-state', + reconcileAfterUndo: async () => { + reconciled.push('state'); + }, + }); + participants.register({ + id: 'test.flush-failure-projection', + phase: 'projection', + reconcileAfterUndo: async () => { + reconciled.push('projection'); + }, + }); + 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 projection 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(IAgentConversationUndoReconciliationRegistry).register({ + id: 'test.serial-projection', + phase: 'projection', + 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 43bcbcabbb..9b48a97cb5 100644 --- a/packages/agent-core-v2/test/app/externalHooksRunner/integration.test.ts +++ b/packages/agent-core-v2/test/app/externalHooksRunner/integration.test.ts @@ -106,7 +106,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 d27b01575f..480d9bc7bc 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'; @@ -145,6 +146,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 } from '#/kosong/model/model'; import { IModelCatalog, type Model } from '#/kosong/model/catalog'; import { ModelCatalog } from '#/kosong/model/catalogService'; @@ -347,6 +349,7 @@ interface ResumeStateSnapshot { readonly context: { readonly history: readonly ContextMessage[]; }; + readonly checkpointedModels: Readonly>; readonly permission: Omit, 'rules'>; readonly usage: Omit, 'currentTurn'>; } @@ -1383,6 +1386,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' }, @@ -1413,9 +1428,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 { @@ -1502,6 +1517,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); } @@ -2105,6 +2130,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 1142d1dcbf..ca1721780d 100644 --- a/packages/agent-core-v2/test/index.test.ts +++ b/packages/agent-core-v2/test/index.test.ts @@ -147,7 +147,9 @@ 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' }, + ]); }); }); 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 f12bec39d3..3529f874fa 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -231,6 +231,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..4d177f9b93 100644 --- a/packages/agent-core-v2/test/wire/resume.test.ts +++ b/packages/agent-core-v2/test/wire/resume.test.ts @@ -4,9 +4,14 @@ import { join } from 'pathe'; import { describe, expect, it, vi } from 'vitest'; +import { + resetUnexpectedErrorHandler, + setUnexpectedErrorHandler, +} from '#/_base/errors/unexpectedError'; import { WIRE_PROTOCOL_VERSION, IAgentGoalService, + reduceContextTranscript, type WireRecord, type PromptOrigin, } from '#/index'; @@ -170,6 +175,46 @@ 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 }); + const transcript = reduceContextTranscript(persistence.records); + expect(transcript.turnIds).toEqual([0, 2, 2]); + expect(transcript.stableTurnIds).toBe(true); + }); + it('projects restored pending tool results before later user messages', async () => { const persistence = new RecordingAgentPersistence([ resumeConfigRecord(), @@ -772,6 +817,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 e47a0f7637..d635ff6bed 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` @@ -75,7 +75,7 @@ import { ErrorCodes, IAgentContextMemoryService, IAgentProfileService, - IAgentPromptService, + IAgentConversationUndoService, IAgentFullCompactionService, IAgentLifecycleService, IAgentRPCService, @@ -689,12 +689,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([ diff --git a/packages/kap-server/src/services/transcript/coreBinding.ts b/packages/kap-server/src/services/transcript/coreBinding.ts index a708ee3ac5..2f0c6f07f0 100644 --- a/packages/kap-server/src/services/transcript/coreBinding.ts +++ b/packages/kap-server/src/services/transcript/coreBinding.ts @@ -21,9 +21,12 @@ import { IAgentLifecycleService, IAgentActivityView, + IAgentConversationUndoReconciliationRegistry, + IAgentContextMemoryService, IEventBus, ISessionMetadata, ISessionInteractionService, + ISessionTodoService, MAIN_AGENT_ID, type AgentMeta, type IDisposable, @@ -31,7 +34,14 @@ import { type Interaction, type ISessionScopeHandle, } from '@moonshot-ai/agent-core-v2'; -import type { AgentDescriptor, TranscriptChangeEvent, TranscriptStore } from '@moonshot-ai/transcript'; +import type { + AgentTranscriptSnapshot, + AgentDescriptor, + TranscriptItem, + TranscriptChangeEvent, + TranscriptStore, +} from '@moonshot-ai/transcript'; +import { groupMessagesIntoSnapshot } from '@moonshot-ai/transcript'; import { AgentTranscriptProjector, type ProjectorInteraction } from './coreEventMap'; @@ -72,9 +82,12 @@ export function bindSessionTranscript( * state-style/idempotent, so replaying a locally-unchanged upsert is safe. */ onOps?: (event: TranscriptChangeEvent) => void, + onUndo?: (agentId: string) => void, + loadPostUndoSnapshot?: (agentId: string) => Promise, ): TranscriptBinding { const agents = session.accessor.get(IAgentLifecycleService); const interactions = session.accessor.get(ISessionInteractionService); + const todos = session.accessor.get(ISessionTodoService); const disposables: IDisposable[] = []; /** Per-agent subscriptions (bus listeners capturing the agent's projector) — disposed with the agent. */ const agentDisposables = new Map(); @@ -94,6 +107,8 @@ export function bindSessionTranscript( const earlyResolves = new Map(); /** Agents whose pendings may announce live (their backfill+seed has run). */ const seededAgents = new Set(); + /** Projection outcomes waiting for their matching `context.undone` event. */ + const undoProjections = new Map(); let seededAll = false; const isSeeded = (agentId: string): boolean => seededAll || seededAgents.has(agentId); @@ -165,9 +180,70 @@ export function bindSessionTranscript( // Agent-scoped, so this sees only this agent's events. The subscription is // tracked per agent and disposed with it — the listener captures the // projector, so a dead agent must not keep projecting into the store. - const bus = handle.accessor.get(IEventBus); - const busD = bus.subscribe((event) => applyOps(handle.id, projector.map(event))); const list = agentDisposables.get(handle.id) ?? []; + if (loadPostUndoSnapshot !== undefined) { + const reconciliation = handle.accessor.get(IAgentConversationUndoReconciliationRegistry); + list.push( + reconciliation.register({ + id: 'transcript.projection', + phase: 'projection', + reconcileAfterUndo: async () => { + onUndo?.(handle.id); + let projected = false; + try { + const rebuilt = await loadPostUndoSnapshot(handle.id); + if (rebuilt !== undefined) { + const transcript = store.ensureAgent(handle.id); + applyOps(handle.id, [ + { + op: 'reset', + agentId: handle.id, + snapshot: conversationResetSnapshot(rebuilt, transcript.snapshot()), + }, + ]); + projected = true; + } + } catch (error) { + logger?.warn( + { + sessionId: store.sessionId, + agentId: handle.id, + err: error instanceof Error ? error.message : error, + }, + 'transcript: post-undo projection failed, falling back to live context', + ); + } finally { + const outcomes = undoProjections.get(handle.id) ?? []; + outcomes.push(projected); + undoProjections.set(handle.id, outcomes); + } + }, + }), + ); + } + const bus = handle.accessor.get(IEventBus); + const busD = bus.subscribe((event) => { + if (event.type === 'context.undone') { + const outcomes = undoProjections.get(handle.id); + const projected = outcomes?.shift(); + if (outcomes?.length === 0) undoProjections.delete(handle.id); + if (projected === true) return; + if (projected === undefined) onUndo?.(handle.id); + const transcript = store.ensureAgent(handle.id); + const rebuilt = groupMessagesIntoSnapshot( + handle.accessor.get(IAgentContextMemoryService).get(), + ); + applyOps(handle.id, [ + { + op: 'reset', + agentId: handle.id, + snapshot: conversationResetSnapshot(rebuilt, transcript.snapshot()), + }, + ]); + return; + } + applyOps(handle.id, projector.map(event)); + }); list.push(busD); agentDisposables.set(handle.id, list); }; @@ -235,10 +311,23 @@ export function bindSessionTranscript( agentDisposables.delete(agentId); subscribedAgents.delete(agentId); projectors.delete(agentId); + undoProjections.delete(agentId); store.markDisposed(agentId, new Date().toISOString()); }), ); + const projectTodos = (items: ReturnType): void => { + applyOps(MAIN_AGENT_ID, [ + { + op: 'todo.upsert', + todo: { todoId: 'todo', items, updatedAt: new Date().toISOString() }, + }, + ]); + }; + disposables.push(todos.onDidChange(projectTodos)); + const initialTodos = todos.getTodos(); + if (initialTodos.length > 0) projectTodos(initialTodos); + // Interactions already pending at bind time are REGISTERED without frames // (ownership must exist immediately so a resolve arriving before the // deferred seed still routes); their frames land in seedPendingInteractions, @@ -330,10 +419,60 @@ export function bindSessionTranscript( knownInteractions.clear(); unseeded.clear(); earlyResolves.clear(); + undoProjections.clear(); }, }; } +function conversationResetSnapshot( + rebuilt: AgentTranscriptSnapshot, + current: AgentTranscriptSnapshot, +): AgentTranscriptSnapshot { + const survivingToolCallIds = new Set(); + for (const item of rebuilt.items) { + if (item.kind !== 'turn') continue; + for (const step of item.steps) { + for (const frame of step.frames) { + if (frame.kind === 'tool') survivingToolCallIds.add(frame.toolCallId); + } + } + } + const attachments = new Map( + rebuilt.attachments.map((attachment) => [attachment.attachmentId, attachment]), + ); + for (const attachment of current.attachments) { + attachments.set(attachment.attachmentId, attachment); + } + return { + ...rebuilt, + items: preserveTaskRefs(rebuilt.items, current.items), + tasks: current.tasks, + interactions: current.interactions.filter((interaction) => + survivingToolCallIds.has(interaction.toolCallId), + ), + attachments: [...attachments.values()], + todos: current.todos, + meta: current.meta, + }; +} + +function preserveTaskRefs( + rebuilt: readonly TranscriptItem[], + current: readonly TranscriptItem[], +): readonly TranscriptItem[] { + const ids = new Set(rebuilt.map(transcriptItemId)); + const taskRefs = current.filter( + (item) => item.kind === 'taskref' && !ids.has(item.refId), + ); + return taskRefs.length === 0 ? rebuilt : [...rebuilt, ...taskRefs]; +} + +function transcriptItemId(item: TranscriptItem): string { + if (item.kind === 'turn') return item.turnId; + if (item.kind === 'marker') return item.markerId; + return item.refId; +} + export function descriptorFromMeta(agentId: string, meta: AgentMeta | undefined): AgentDescriptor { const parentFromLabels = meta?.labels?.['parentAgentId']; const swarmItem = meta?.labels?.['swarmItem'] ?? meta?.swarmItem; diff --git a/packages/kap-server/src/services/transcript/coreEventMap.ts b/packages/kap-server/src/services/transcript/coreEventMap.ts index de7fbf6f49..c29fd94efd 100644 --- a/packages/kap-server/src/services/transcript/coreEventMap.ts +++ b/packages/kap-server/src/services/transcript/coreEventMap.ts @@ -14,9 +14,8 @@ * - `tool.call.delta` / `tool.progress` are NOT projected in v1 (argument and * progress streaming are dropped; `tool.result` carries the terminal * state). Known limitation. - * - `context.spliced` (undo/clear) is projected as a bare 'undo' marker with - * the raw payload — no `items.remove` reconstruction in v1. Known - * limitation. + * - `context.spliced` projects the visible undo marker. The session binding + * handles `context.undone` as a full conversation projection reset. * - `error` / `warning` become `marker.upsert{ marker: 'notice' }` and never * enter a step. * - No `swarm.*` / `plan.*` mode-transition events exist on the v2 bus today; @@ -201,9 +200,9 @@ export class AgentTranscriptProjector { }), ]; case 'context.spliced': - // Known limitation: undo/clear projects as a bare 'undo' marker (raw - // payload attached); no `items.remove` reconstruction in v1. return [this.markerOp('undo', restOf(event))]; + case 'context.undone': + return []; case 'error': return [this.noticeOp('error', event.message, restOf(event))]; case 'warning': diff --git a/packages/kap-server/src/services/transcript/transcriptService.ts b/packages/kap-server/src/services/transcript/transcriptService.ts index 51c7af06d1..5ffb430d95 100644 --- a/packages/kap-server/src/services/transcript/transcriptService.ts +++ b/packages/kap-server/src/services/transcript/transcriptService.ts @@ -101,6 +101,7 @@ export class TranscriptService { private readonly opsListeners = new Map void>>(); /** Debounced post-turn heals: `${sessionId}:${agentId}` → pending ordinals + timer. */ private readonly healTimers = new Map; timer: NodeJS.Timeout }>(); + private readonly projectionGenerations = new Map(); constructor(private readonly deps: TranscriptServiceDeps) { // Live entries must not outlive their session: once it closes or archives, @@ -130,8 +131,13 @@ export class TranscriptService { const store = new TranscriptStore(sessionId); let binding: TranscriptBinding; try { - binding = bindSessionTranscript(store, session, this.deps.logger, (event) => - this.handleLiveOps(sessionId, event), + binding = bindSessionTranscript( + store, + session, + this.deps.logger, + (event) => this.handleLiveOps(sessionId, event), + (agentId) => this.invalidateAgentProjection(sessionId, agentId), + (agentId) => this.readColdSnapshot(sessionId, agentId), ); } catch (error) { // The session's core scope can be disposed mid-bind during shutdown @@ -221,6 +227,7 @@ export class TranscriptService { * without colliding. */ private async backfillAgent(sessionId: string, store: TranscriptStore, agentId: string): Promise { + const generation = this.projectionGeneration(sessionId, agentId); let snapshot: AgentTranscriptSnapshot | undefined; try { snapshot = await this.readColdSnapshot(sessionId, agentId); @@ -231,7 +238,10 @@ export class TranscriptService { ); } // The entry may have been dropped (session closed) while reading from disk. - if (this.live.get(sessionId)?.store !== store) return; + if ( + this.live.get(sessionId)?.store !== store || + this.projectionGeneration(sessionId, agentId) !== generation + ) return; const transcript = store.ensureAgent(agentId); if (snapshot !== undefined) { // Turns merge live-first (`healTurnOps`): ops the projector landed @@ -341,6 +351,20 @@ export class TranscriptService { this.healTimers.set(key, { ordinals, timer }); } + private projectionGeneration(sessionId: string, agentId: string): number { + return this.projectionGenerations.get(`${sessionId}:${agentId}`) ?? 0; + } + + private invalidateAgentProjection(sessionId: string, agentId: string): void { + const key = `${sessionId}:${agentId}`; + this.projectionGenerations.set(key, this.projectionGeneration(sessionId, agentId) + 1); + const heal = this.healTimers.get(key); + if (heal !== undefined) { + clearTimeout(heal.timer); + this.healTimers.delete(key); + } + } + /** * A backfill rebuilds every turn as 'completed' — the cold grouping cannot * see in-flight work. When the agent's loop is actually mid-turn, re-assert @@ -394,6 +418,7 @@ export class TranscriptService { ): Promise { const entry = this.live.get(sessionId); if (entry === undefined) return; + const generation = this.projectionGeneration(sessionId, agentId); let snapshot: AgentTranscriptSnapshot | undefined; try { snapshot = await this.readColdSnapshot(sessionId, agentId); @@ -405,7 +430,11 @@ export class TranscriptService { return; } // The entry may have been dropped (session closed) while reading from disk. - if (snapshot === undefined || this.live.get(sessionId)?.store !== entry.store) return; + if ( + snapshot === undefined || + this.live.get(sessionId)?.store !== entry.store || + this.projectionGeneration(sessionId, agentId) !== generation + ) return; const transcript = entry.store.getAgent(agentId); if (transcript === undefined) return; const ops: TranscriptOperation[] = []; @@ -480,8 +509,13 @@ export class TranscriptService { } throw error; } - const messages = [...reduceContextTranscript(records).entries]; - return groupMessagesIntoSnapshot(messages); + const transcript = reduceContextTranscript(records); + return groupMessagesIntoSnapshot( + transcript.entries, + transcript.stableTurnIds + ? { turnIds: transcript.turnIds, turns: transcript.turns } + : undefined, + ); } /** Dispose the live store + binding for a session (session closed / server shutdown). */ @@ -493,6 +527,9 @@ export class TranscriptService { this.healTimers.delete(key); } } + for (const key of this.projectionGenerations.keys()) { + if (key.startsWith(`${sessionId}:`)) this.projectionGenerations.delete(key); + } const entry = this.live.get(sessionId); if (entry === undefined) return; this.live.delete(sessionId); diff --git a/packages/kap-server/test/services/transcript.test.ts b/packages/kap-server/test/services/transcript.test.ts index 1b6f551ec4..072bb58c28 100644 --- a/packages/kap-server/test/services/transcript.test.ts +++ b/packages/kap-server/test/services/transcript.test.ts @@ -12,20 +12,27 @@ import { join } from 'node:path'; import { IAgentLifecycleService, + IAgentConversationUndoReconciliationRegistry, + IAgentContextMemoryService, IAgentLoopService, IEventBus, ISessionIndex, ISessionInteractionService, ISessionLifecycleService, ISessionMetadata, + ISessionTodoService, SessionInteractionService, + reduceContextTranscript, type DomainEvent, + type ContextMessage, type ISessionScopeHandle, type Scope, + type WireRecord, } from '@moonshot-ai/agent-core-v2'; import { AgentTranscript, TranscriptStore, + groupMessagesIntoSnapshot, type AgentTranscriptSnapshot, type AppendOp, type FrameUpsertOp, @@ -34,7 +41,7 @@ import { type TranscriptTask, type TranscriptTurn, } from '@moonshot-ai/transcript'; -import { describe, expect, it } from 'vitest'; +import { describe, expect, it, vi } from 'vitest'; import { bindSessionTranscript } from '../../src/services/transcript/coreBinding'; import { AgentTranscriptProjector } from '../../src/services/transcript/coreEventMap'; @@ -692,6 +699,7 @@ describe('AgentTranscriptProjector', () => { feed(ev({ type: 'compaction.started', trigger: 'auto' })); feed(ev({ type: 'compaction.completed', result: { kept: 3 } })); feed(ev({ type: 'context.spliced', start: 1, deleteCount: 2, messages: [] })); + feed(ev({ type: 'context.undone', turns: 1 })); const markers = tx .getItems() @@ -1010,9 +1018,41 @@ describe('bindSessionTranscript', () => { interface FakeAgentHandle { readonly id: string; readonly bus: FakeBus; + reconcile(phase: 'state' | 'projection'): Promise; readonly accessor: { get: (token: unknown) => unknown }; } + class FakeReconciliationRegistry { + private readonly participants = new Map< + string, + { + readonly phase?: 'state' | 'projection'; + reconcileAfterUndo(): Promise; + } + >(); + register(participant: { + readonly id: string; + readonly phase?: 'state' | 'projection'; + reconcileAfterUndo(): Promise; + }): { dispose: () => void } { + this.participants.set(participant.id, participant); + return { + dispose: () => { + if (this.participants.get(participant.id) === participant) { + this.participants.delete(participant.id); + } + }, + }; + } + async reconcile(phase: 'state' | 'projection'): Promise { + await Promise.all( + [...this.participants.values()] + .filter((participant) => (participant.phase ?? 'state') === phase) + .map((participant) => participant.reconcileAfterUndo()), + ); + } + } + class FakeAgents { private readonly handles = new Map(); private readonly createHandlers = new Set<(handle: FakeAgentHandle) => void>(); @@ -1031,17 +1071,26 @@ describe('bindSessionTranscript', () => { this.disposeHandlers.add(cb); return { dispose: () => this.disposeHandlers.delete(cb) }; } - add(id: string, opts?: { loopStatus?: unknown }): FakeAgentHandle { + add( + id: string, + opts?: { loopStatus?: unknown; history?: readonly ContextMessage[] }, + ): FakeAgentHandle { const bus = new FakeBus(); + const reconciliation = new FakeReconciliationRegistry(); const handle: FakeAgentHandle = { id, bus, + reconcile: (phase) => reconciliation.reconcile(phase), accessor: { get: (token: unknown) => { if (token === IEventBus) return bus; + if (token === IAgentConversationUndoReconciliationRegistry) return reconciliation; if (token === IAgentLoopService) { return { status: () => opts?.loopStatus ?? { state: 'idle' } }; } + if (token === IAgentContextMemoryService) { + return { get: () => opts?.history ?? [] }; + } return undefined; }, }, @@ -1056,9 +1105,31 @@ describe('bindSessionTranscript', () => { } } + class FakeTodos { + private readonly handlers = new Set<(items: readonly { title: string; status: 'pending' | 'in_progress' | 'done' }[]) => void>(); + private current: readonly { title: string; status: 'pending' | 'in_progress' | 'done' }[] = []; + readonly onDidChange = ( + handler: (items: readonly { title: string; status: 'pending' | 'in_progress' | 'done' }[]) => void, + ): { dispose: () => void } => { + this.handlers.add(handler); + return { dispose: () => this.handlers.delete(handler) }; + }; + getTodos(): readonly { title: string; status: 'pending' | 'in_progress' | 'done' }[] { + return this.current; + } + setTodos(items: readonly { title: string; status: 'pending' | 'in_progress' | 'done' }[]): void { + this.current = items; + for (const handler of this.handlers) handler(items); + } + clear(): void { + this.setTodos([]); + } + } + function fakeSession( interactions: SessionInteractionService, agents?: FakeAgents, + todos = new FakeTodos(), ): ISessionScopeHandle { return { accessor: { @@ -1073,6 +1144,7 @@ describe('bindSessionTranscript', () => { ); } if (token === ISessionInteractionService) return interactions; + if (token === ISessionTodoService) return todos; if (token === ISessionMetadata) return { read: async () => ({ agents: {} }) }; return undefined; }, @@ -1080,6 +1152,543 @@ describe('bindSessionTranscript', () => { } as unknown as ISessionScopeHandle; } + it('resets conversation items from the authoritative post-undo context', () => { + const interactions = new SessionInteractionService(); + const agents = new FakeAgents(); + const main = agents.add('main', { + history: [ + { + role: 'user', + content: [{ type: 'text', text: 'kept' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'assistant', + content: [{ type: 'text', text: 'answer' }], + toolCalls: [], + }, + ], + }); + const store = new TranscriptStore('s1'); + store.ensureAgent('main').apply([ + { + op: 'turn.upsert', + turn: { + kind: 'turn', + turnId: 't0', + ordinal: 0, + state: 'completed', + origin: { kind: 'user' }, + }, + }, + ]); + const ops: TranscriptOperation[] = []; + const binding = bindSessionTranscript( + store, + fakeSession(interactions, agents), + undefined, + (event) => ops.push(...event.ops), + ); + const emitTurn = (turnId: number, origin: Record): void => { + main.bus.emit(ev({ type: 'turn.started', turnId, origin })); + main.bus.emit(ev({ type: 'turn.ended', turnId, reason: 'completed' })); + }; + + emitTurn(1, { kind: 'skill_activation', skillName: 'review', trigger: 'user-slash' }); + emitTurn(2, { kind: 'shell_command' }); + emitTurn(3, { kind: 'plugin_command', pluginId: 'p', commandName: 'c', trigger: 'user-slash' }); + emitTurn(4, { kind: 'task', taskId: 'task-1' }); + main.bus.emit( + ev({ + type: 'task.started', + info: { + taskId: 'task-1', + kind: 'process', + description: 'build', + status: 'running', + detached: true, + startedAt: 1_700_000_000_000, + endedAt: null, + }, + }), + ); + main.bus.emit(ev({ type: 'context.spliced', start: 1, deleteCount: 4, messages: [] })); + main.bus.emit(ev({ type: 'context.undone', turns: 2 })); + + expect( + store + .getAgent('main') + ?.getItems() + .filter((item): item is TranscriptTurn => item.kind === 'turn') + .map((turn) => turn.turnId), + ).toEqual(['t0']); + expect(store.getAgent('main')?.getTask('task-1')).toMatchObject({ state: 'running' }); + expect(store.getAgent('main')?.getItems().filter((item) => item.kind === 'marker')).toEqual([]); + expect(ops.filter((op) => op.op === 'reset')).toHaveLength(1); + expect(ops.filter((op) => op.op === 'items.remove')).toEqual([]); + binding.dispose(); + }); + + it('does not infer the undo cut from the displayed turn count', () => { + const interactions = new SessionInteractionService(); + const agents = new FakeAgents(); + const main = agents.add('main', { + history: [ + { + role: 'user', + content: [{ type: 'text', text: 'kept' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + ], + }); + const store = new TranscriptStore('s1'); + store.ensureAgent('main').apply([ + { + op: 'turn.upsert', + turn: { + kind: 'turn', + turnId: 't0', + ordinal: 0, + state: 'completed', + origin: { kind: 'user' }, + }, + }, + ]); + const ops: TranscriptOperation[] = []; + const binding = bindSessionTranscript( + store, + fakeSession(interactions, agents), + undefined, + (event) => ops.push(...event.ops), + ); + main.bus.emit(ev({ type: 'turn.started', turnId: 1, origin: { kind: 'user' } })); + main.bus.emit(ev({ type: 'turn.ended', turnId: 1, reason: 'completed' })); + + main.bus.emit(ev({ type: 'context.spliced', start: 0, deleteCount: 3, messages: [] })); + main.bus.emit(ev({ type: 'context.undone', turns: 5 })); + + expect( + store + .getAgent('main') + ?.getItems() + .filter((item): item is TranscriptTurn => item.kind === 'turn') + .map((turn) => turn.turnId), + ).toEqual(['t0']); + expect(ops.filter((op) => op.op === 'reset')).toHaveLength(1); + expect(ops.filter((op) => op.op === 'items.remove')).toEqual([]); + binding.dispose(); + }); + + it('rebuilds the post-undo projection from full journal history after compaction', async () => { + const interactions = new SessionInteractionService(); + const agents = new FakeAgents(); + const foldedHistory: ContextMessage[] = [ + { + role: 'user', + content: [{ type: 'text', text: 'before compaction' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'user', + content: [{ type: 'text', text: 'summary' }], + toolCalls: [], + origin: { kind: 'compaction_summary' }, + }, + { + role: 'user', + content: [{ type: 'text', text: 'surviving turn' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'assistant', + content: [{ type: 'text', text: 'surviving answer' }], + toolCalls: [], + }, + ]; + const main = agents.add('main', { history: foldedHistory }); + const fullHistory: ContextMessage[] = [ + foldedHistory[0]!, + { + role: 'assistant', + content: [{ type: 'text', text: 'checking old state' }], + toolCalls: [ + { + type: 'function', + id: 'call-old', + name: 'Read', + arguments: '{"path":"/old"}', + }, + ], + }, + { + role: 'tool', + content: [{ type: 'text', text: 'old tool result' }], + toolCalls: [], + toolCallId: 'call-old', + }, + { + role: 'assistant', + content: [{ type: 'text', text: 'old answer' }], + toolCalls: [], + }, + ...foldedHistory.slice(1), + ]; + const doomedHistory: ContextMessage[] = [ + ...fullHistory, + { + role: 'user', + content: [{ type: 'text', text: 'undone turn' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'assistant', + content: [{ type: 'text', text: 'undone answer' }], + toolCalls: [], + }, + ]; + const store = new TranscriptStore('s1'); + store.ensureAgent('main').apply([ + { + op: 'reset', + agentId: 'main', + snapshot: groupMessagesIntoSnapshot(doomedHistory), + }, + ]); + const ops: TranscriptOperation[] = []; + let invalidations = 0; + const binding = bindSessionTranscript( + store, + fakeSession(interactions, agents), + undefined, + (event) => ops.push(...event.ops), + () => { + invalidations += 1; + }, + async () => groupMessagesIntoSnapshot(fullHistory), + ); + main.bus.emit( + ev({ + type: 'task.started', + info: { + taskId: 'task-1', + kind: 'process', + description: 'build', + status: 'running', + detached: true, + startedAt: 1_700_000_000_000, + endedAt: null, + }, + }), + ); + + await main.reconcile('projection'); + main.bus.emit(ev({ type: 'context.undone', turns: 1 })); + + const transcript = store.getAgent('main'); + const items = transcript?.getItems() ?? []; + const turns = items.filter((item): item is TranscriptTurn => item.kind === 'turn'); + expect(turns.map((turn) => turn.prompt)).toEqual([ + 'before compaction', + 'surviving turn', + ]); + expect( + turns[0]?.steps.flatMap((step) => step.frames).find((frame) => frame.kind === 'tool'), + ).toMatchObject({ output: 'old tool result' }); + expect(items.filter((item) => item.kind === 'marker')).toHaveLength(1); + expect(items.some((item) => item.kind === 'turn' && item.prompt === 'undone turn')).toBe(false); + expect(transcript?.getTask('task-1')).toMatchObject({ state: 'running' }); + expect(ops.filter((op) => op.op === 'reset')).toHaveLength(1); + expect(invalidations).toBe(1); + binding.dispose(); + }); + + it('keeps only interactions anchored to surviving tool calls after undo', async () => { + const interactions = new SessionInteractionService(); + const agents = new FakeAgents(); + const main = agents.add('main'); + const survivingHistory: ContextMessage[] = [ + { + role: 'user', + content: [{ type: 'text', text: 'kept prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'assistant', + content: [], + toolCalls: [ + { + type: 'function', + id: 'call-kept', + name: 'Read', + arguments: '{"path":"/kept"}', + }, + ], + }, + { + role: 'tool', + content: [{ type: 'text', text: 'kept result' }], + toolCalls: [], + toolCallId: 'call-kept', + }, + { + role: 'assistant', + content: [{ type: 'text', text: 'kept answer' }], + toolCalls: [], + }, + ]; + const doomedHistory: ContextMessage[] = [ + ...survivingHistory, + { + role: 'user', + content: [{ type: 'text', text: 'undone prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'assistant', + content: [], + toolCalls: [ + { + type: 'function', + id: 'call-undone-resolved', + name: 'Write', + arguments: '{"path":"/undone"}', + }, + { + type: 'function', + id: 'call-undone-pending', + name: 'Bash', + arguments: '{"command":"false"}', + }, + ], + }, + ]; + const store = new TranscriptStore('s1'); + store.ensureAgent('main').apply([ + { + op: 'reset', + agentId: 'main', + snapshot: groupMessagesIntoSnapshot(doomedHistory), + }, + { + op: 'interaction.upsert', + interaction: { + interactionId: 'approval-kept', + interactionKind: 'approval', + toolCallId: 'call-kept', + state: 'approved', + response: { decision: 'approved' }, + }, + }, + { + op: 'interaction.upsert', + interaction: { + interactionId: 'approval-undone', + interactionKind: 'approval', + toolCallId: 'call-undone-resolved', + state: 'approved', + response: { decision: 'approved' }, + }, + }, + { + op: 'interaction.upsert', + interaction: { + interactionId: 'question-undone', + interactionKind: 'question', + toolCallId: 'call-undone-pending', + state: 'pending', + request: { questions: [{ question: 'Continue?', options: [] }] }, + }, + }, + ]); + const binding = bindSessionTranscript( + store, + fakeSession(interactions, agents), + undefined, + undefined, + undefined, + async () => groupMessagesIntoSnapshot(survivingHistory), + ); + + await main.reconcile('projection'); + main.bus.emit(ev({ type: 'context.undone', turns: 1 })); + + const transcript = store.getAgent('main'); + expect([...transcript!.getInteractions().keys()]).toEqual(['approval-kept']); + expect(transcript?.getInteraction('approval-kept')).toMatchObject({ + toolCallId: 'call-kept', + state: 'approved', + }); + expect(transcript?.getInteraction('approval-undone')).toBeUndefined(); + expect(transcript?.listPendingInteractions()).toEqual([]); + binding.dispose(); + }); + + it('preserves surviving turn ids and world-time task refs across journal projection', async () => { + const interactions = new SessionInteractionService(); + const agents = new FakeAgents(); + const main = agents.add('main'); + const survivingHistory: ContextMessage[] = [ + { + role: 'user', + content: [{ type: 'text', text: 'kept prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + { + role: 'assistant', + content: [{ type: 'text', text: 'kept answer' }], + toolCalls: [], + }, + ]; + const store = new TranscriptStore('s1'); + const binding = bindSessionTranscript( + store, + fakeSession(interactions, agents), + undefined, + undefined, + undefined, + async () => + groupMessagesIntoSnapshot(survivingHistory, { + turnIds: [0, 0], + turns: [ + { turnId: 0, input: survivingHistory[0]!.content, origin: { kind: 'user' } }, + { turnId: 1, input: [], origin: { kind: 'retry' } }, + ], + }), + ); + const emitTurn = (turnId: number, origin: Record): void => { + main.bus.emit(ev({ type: 'turn.started', turnId, origin })); + main.bus.emit(ev({ type: 'turn.ended', turnId, reason: 'completed' })); + }; + + emitTurn(0, { kind: 'user' }); + emitTurn(1, { kind: 'retry' }); + emitTurn(2, { kind: 'user' }); + main.bus.emit( + ev({ + type: 'task.started', + info: { + taskId: 'task-world-time', + kind: 'process', + description: 'background build', + status: 'running', + detached: true, + startedAt: 1_700_000_000_000, + endedAt: null, + }, + }), + ); + + await main.reconcile('projection'); + main.bus.emit(ev({ type: 'context.undone', turns: 1 })); + + const items = store.getAgent('main')?.getItems() ?? []; + expect( + items + .filter((item): item is TranscriptTurn => item.kind === 'turn') + .map((turn) => turn.turnId), + ).toEqual(['t0', 't1']); + expect(items).toContainEqual( + expect.objectContaining({ kind: 'taskref', taskId: 'task-world-time' }), + ); + expect(store.getAgent('main')?.getTask('task-world-time')).toMatchObject({ + state: 'running', + }); + binding.dispose(); + }); + + it('omits a phantom turn when undo is followed by a restored task notification', () => { + const records: WireRecord[] = [ + { + type: 'turn.prompt', + input: [{ type: 'text', text: 'kept prompt' }], + origin: { kind: 'user' }, + }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'kept prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'turn.prompt', + input: [{ type: 'text', text: 'undone prompt' }], + origin: { kind: 'user' }, + }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'undone prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { type: 'context.undo', count: 1 }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'background task completed' }], + toolCalls: [], + origin: { + kind: 'task', + taskId: 'bash-001', + status: 'completed', + notificationId: 'task:bash-001:completed', + }, + }, + }, + ]; + const transcript = reduceContextTranscript(records); + + const snapshot = groupMessagesIntoSnapshot(transcript.entries, { + turnIds: transcript.turnIds, + turns: transcript.turns, + }); + + expect( + snapshot.items + .filter((item): item is TranscriptTurn => item.kind === 'turn') + .map((turn) => ({ turnId: turn.turnId, prompt: turn.prompt })), + ).toEqual([{ turnId: 't0', prompt: 'kept prompt' }]); + }); + + it('projects session Todo changes into the main live transcript', () => { + const todos = new FakeTodos(); + todos.setTodos([{ title: 'pre-bind', status: 'pending' }]); + const store = new TranscriptStore('s1'); + const binding = bindSessionTranscript( + store, + fakeSession(new SessionInteractionService(), undefined, todos), + ); + + expect(store.getAgent('main')?.getTodo('todo')?.items).toEqual([ + { title: 'pre-bind', status: 'pending' }, + ]); + + todos.setTodos([{ title: 'kept', status: 'pending' }]); + expect(store.getAgent('main')?.getTodo('todo')?.items).toEqual([ + { title: 'kept', status: 'pending' }, + ]); + + binding.dispose(); + todos.setTodos([{ title: 'ignored after dispose', status: 'done' }]); + expect(store.getAgent('main')?.getTodo('todo')?.items).toEqual([ + { title: 'kept', status: 'pending' }, + ]); + }); + it('registers pre-bind pendings without frames and replays an early resolve at seed time', () => { const interactions = new SessionInteractionService(); interactions.enqueue({ @@ -1477,6 +2086,62 @@ describe('bindSessionTranscript', () => { } }); + it('drops a stale backfill that completes after an undo reset', async () => { + const home = await seedWireHome(); + try { + const history: ContextMessage[] = [ + { + role: 'user', + content: [{ type: 'text', text: 'kept after undo' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + ]; + const agents = new FakeAgents(); + const main = agents.add('main', { history }); + const service = new TranscriptService({ + homeDir: home, + core: fakeCoreWithAgents(new SessionInteractionService(), agents), + }); + let markStarted!: () => void; + let release!: () => void; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const blocked = new Promise((resolve) => { + release = resolve; + }); + vi.spyOn(service, 'readColdSnapshot').mockImplementation(async () => { + markStarted(); + await blocked; + return groupMessagesIntoSnapshot([ + { + role: 'user', + content: [{ type: 'text', text: 'stale before undo' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + ]); + }); + + const store = service.forSessionLive('s1'); + await started; + main.bus.emit(ev({ type: 'context.undone', turns: 1 })); + release(); + await service.whenReady('s1'); + + const prompts = store + ?.getAgent('main') + ?.getItems() + .filter((item): item is TranscriptTurn => item.kind === 'turn') + .map((turn) => turn.prompt); + expect(prompts).toEqual(['kept after undo']); + service.dropSession('s1'); + } finally { + await rm(home, { recursive: true, force: true }); + } + }); + it('re-asserts running when the backfill rebuilds the live turn completed', async () => { const home = await seedWireHome(); try { diff --git a/packages/kap-server/test/sessionEventBroadcaster.test.ts b/packages/kap-server/test/sessionEventBroadcaster.test.ts index 54720a8279..4c9b3682f5 100644 --- a/packages/kap-server/test/sessionEventBroadcaster.test.ts +++ b/packages/kap-server/test/sessionEventBroadcaster.test.ts @@ -10,6 +10,7 @@ import type { IScopeHandle, Scope } from '@moonshot-ai/agent-core-v2'; import { ContextSizeModel, IAgentActivityView, + IAgentConversationUndoReconciliationRegistry, IAgentContextSizeService, IAgentLifecycleService, IAgentProfileService, @@ -18,6 +19,7 @@ import { IEventService, ISessionInteractionService, ISessionLifecycleService, + ISessionTodoService, IWireService, ISessionMetadata, SessionInteractionService, @@ -78,6 +80,9 @@ class FakeAgentHandle { private readonly services = new Map(); constructor(readonly id: string) { this.services.set(IEventBus, this.bus); + this.services.set(IAgentConversationUndoReconciliationRegistry, { + register: () => ({ dispose: () => {} }), + }); this.accessor = { get: (token: unknown) => this.services.get(token), }; @@ -191,6 +196,13 @@ function makeCore( if (t === ISessionInteractionService) return lifecycle.interactions; // Minimal metadata read for the transcript binding's descriptor pass. if (t === ISessionMetadata) return { read: async () => ({ agents: metaAgents }) }; + // Minimal todo facade for the transcript binding's change projection. + if (t === ISessionTodoService) { + return { + onDidChange: () => ({ dispose: () => {} }), + getTodos: () => [], + }; + } return undefined; }, }; diff --git a/packages/kap-server/test/sessions.test.ts b/packages/kap-server/test/sessions.test.ts index 69cda99b17..5bcc8a9318 100644 --- a/packages/kap-server/test/sessions.test.ts +++ b/packages/kap-server/test/sessions.test.ts @@ -614,8 +614,8 @@ 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('rejects an unsupported action suffix (40001)', async () => { diff --git a/packages/transcript/src/granularity/filterOps.ts b/packages/transcript/src/granularity/filterOps.ts index 4c1ca10860..3703caf83b 100644 --- a/packages/transcript/src/granularity/filterOps.ts +++ b/packages/transcript/src/granularity/filterOps.ts @@ -22,7 +22,16 @@ export function filterOpsForGrade( ): TranscriptOperation[] { const rank = GRADE_RANK[grade]; if (rank === 0) return []; - return ops.filter((op) => admits(grade, op)); + const filtered: TranscriptOperation[] = []; + for (const op of ops) { + if (!admits(grade, op)) continue; + filtered.push( + op.op === 'reset' + ? { ...op, snapshot: redactSnapshotForGrade(grade, op.snapshot) } + : op, + ); + } + return filtered; } function admits(grade: TranscriptGrade, op: TranscriptOperation): boolean { diff --git a/packages/transcript/src/history/groupTurns.ts b/packages/transcript/src/history/groupTurns.ts index 4b4424fd8b..ca8817fa45 100644 --- a/packages/transcript/src/history/groupTurns.ts +++ b/packages/transcript/src/history/groupTurns.ts @@ -10,12 +10,9 @@ * - streamed-vs-persisted duplication is assumed already resolved upstream; * - interaction frames do not appear (approvals are not persisted as * context messages); - * - persisted messages carry no turn ids, so turn ordinals are assigned by - * grouping — **0-based, matching the engine's live turn numbering** — and - * can drift from the engine's ids when hidden origins (e.g. retries) make - * the engine consume an ordinal that grouping cannot see. Alignment is - * what makes a rebuilt slice safe to merge into a live store (backfill): - * the engine's next turn continues at `t` without colliding. + * - when the source supplies journal-recovered turn identity, it remains + * authoritative across hidden retries and same-turn steers; callers with + * only flat legacy messages use 0-based best-effort grouping. * * The input type is structural so the engine's `ContextMessage` is directly * assignable without a dependency from this package onto the engine. @@ -60,6 +57,17 @@ export interface HistoryMessage { readonly origin?: { readonly kind: string }; } +export interface HistoryTurnSeed { + readonly turnId: number; + readonly input: readonly HistoryContentPart[]; + readonly origin: { readonly kind: string }; +} + +export interface HistoryGroupingOptions { + readonly turnIds: readonly (number | undefined)[]; + readonly turns: readonly HistoryTurnSeed[]; +} + interface TurnDraft { turnId: string; ordinal: number; @@ -99,7 +107,11 @@ const FALLBACK_ORIGIN: TurnOrigin = { kind: 'other' }; export function groupMessagesIntoSnapshot( messages: readonly HistoryMessage[], + options?: HistoryGroupingOptions, ): AgentTranscriptSnapshot { + if (options !== undefined && options.turns.length > 0) { + return groupMessagesByStableTurn(messages, options); + } const items: TranscriptItem[] = []; const attachments: TranscriptAttachment[] = []; let turn: TurnDraft | undefined; @@ -256,8 +268,165 @@ export function groupMessagesIntoSnapshot( return { items, tasks: [], interactions: [], attachments, todos: [], meta: {} }; } +function groupMessagesByStableTurn( + messages: readonly HistoryMessage[], + options: HistoryGroupingOptions, +): AgentTranscriptSnapshot { + const attachments: TranscriptAttachment[] = []; + const drafts = new Map(); + const markers = new Map(); + let markerCount = 0; + let lastTurnId: number | undefined; + + const ensureDraft = (turnId: number, message?: HistoryMessage): TurnDraft => { + let draft = drafts.get(turnId); + if (draft !== undefined) return draft; + draft = { + turnId: `t${turnId}`, + ordinal: turnId, + origin: message === undefined ? FALLBACK_ORIGIN : mapOrigin(message), + prompt: message === undefined || isHiddenPrompt(message) ? undefined : textOf(message), + steps: [], + }; + drafts.set(turnId, draft); + return draft; + }; + + for (const seed of options.turns) { + const message: HistoryMessage = { + role: 'user', + content: seed.input, + origin: seed.origin, + }; + ensureDraft(seed.turnId, message); + } + + const collectAttachments = (message: HistoryMessage): string[] | undefined => { + const ids: string[] = []; + for (const part of message.content ?? []) { + if (part.type === 'image' || part.type === 'video' || part.type === 'audio') { + if (!('source' in part) || part.source === undefined) continue; + const source = part.source as HistoryMediaSource; + const entity: TranscriptAttachment = { + attachmentId: `att_${attachments.length + 1}`, + mediaType: source.kind === 'base64' ? source.media_type : `${part.type}/*`, + source: + source.kind === 'url' + ? { kind: 'url', url: source.url } + : source.kind === 'file' + ? { kind: 'file', fileId: source.file_id } + : undefined, + }; + attachments.push(entity); + ids.push(entity.attachmentId); + } else if (part.type === 'file' && 'file_id' in part) { + const entity: TranscriptAttachment = { + attachmentId: `att_${attachments.length + 1}`, + mediaType: part.media_type as string, + name: part.name as string, + size: part.size as number, + source: { kind: 'file', fileId: part.file_id as string }, + }; + attachments.push(entity); + ids.push(entity.attachmentId); + } + } + return ids.length > 0 ? ids : undefined; + }; + + for (const [index, message] of messages.entries()) { + if (message.role === 'system') continue; + const assignedTurnId = options.turnIds[index] ?? lastTurnId; + if (assignedTurnId !== undefined) lastTurnId = assignedTurnId; + const draft = assignedTurnId === undefined ? undefined : ensureDraft(assignedTurnId); + const originKind = message.origin?.kind; + + if (message.role === 'user') { + const markerKey = originKind === undefined ? undefined : MARKER_USER_ORIGINS[originKind]; + if (markerKey !== undefined) { + const list = markers.get(assignedTurnId) ?? []; + markerCount += 1; + list.push({ + kind: 'marker', + markerId: `m${markerCount}`, + marker: markerKey, + payload: { text: textOf(message), origin: message.origin }, + }); + markers.set(assignedTurnId, list); + } + if (draft !== undefined && draft.attachmentIds === undefined) { + draft.attachmentIds = collectAttachments(message); + } + continue; + } + + if (message.role === 'assistant' && draft !== undefined) { + const stepOrdinal = draft.steps.length + 1; + const step: StepDraft = { + stepId: `${draft.turnId}.${stepOrdinal}`, + ordinal: stepOrdinal, + frames: [], + }; + draft.steps.push(step); + let frameCount = 0; + const nextFrameId = (): string => `${step.stepId}.f${++frameCount}`; + for (const part of message.content ?? []) { + if (part.type === 'text' && 'text' in part && typeof part.text === 'string' && part.text.length > 0) { + step.frames.push({ kind: 'text', frameId: nextFrameId(), role: 'assistant', text: part.text }); + } else if (part.type === 'think' && 'think' in part && typeof part.think === 'string' && part.think.length > 0) { + step.frames.push({ kind: 'thinking', frameId: nextFrameId(), text: part.think }); + } + } + for (const call of message.toolCalls ?? []) { + step.frames.push({ + kind: 'tool', + frameId: `${step.stepId}.${call.id}`, + toolCallId: call.id, + name: call.name, + state: 'running', + input: parseArguments(call.arguments), + }); + } + continue; + } + + if (message.role === 'tool' && draft !== undefined) { + const frame = currentTurnToolFrame(draft, message.toolCallId); + if (frame?.kind !== 'tool') continue; + const output = textOf(message); + replaceToolFrame(draft, message.toolCallId!, { + ...frame, + state: message.isError ? 'error' : 'done', + output, + error: message.isError ? output : undefined, + }); + } + } + + const items: TranscriptItem[] = [...(markers.get(undefined) ?? [])]; + const emittedMarkerOrdinals = new Set(); + for (const draft of [...drafts.values()].toSorted((a, b) => a.ordinal - b.ordinal)) { + items.push(draftToTurnItem(draft), ...(markers.get(draft.ordinal) ?? [])); + emittedMarkerOrdinals.add(draft.ordinal); + } + for (const [ordinal, entries] of [...markers.entries()].toSorted(([a], [b]) => { + if (a === undefined) return -1; + if (b === undefined) return 1; + return a - b; + })) { + if (ordinal === undefined || emittedMarkerOrdinals.has(ordinal)) continue; + items.push(...entries); + } + return { items, tasks: [], interactions: [], attachments, todos: [], meta: {} }; +} + // ---------------------------------------------------------------- helpers +function isHiddenPrompt(message: HistoryMessage): boolean { + const kind = message.origin?.kind; + return kind !== undefined && HIDDEN_USER_ORIGINS.has(kind); +} + /** Whether a hidden-origin user message opened its own engine turn. */ function opensOwnTurn(message: HistoryMessage): boolean { const origin = message.origin as { kind?: unknown; name?: unknown } | undefined; diff --git a/packages/transcript/test/layers.test.ts b/packages/transcript/test/layers.test.ts index 605845939f..548a378d7a 100644 --- a/packages/transcript/test/layers.test.ts +++ b/packages/transcript/test/layers.test.ts @@ -138,6 +138,15 @@ describe('granularity', () => { expect(turnGrade.items[1]?.kind).toBe('marker'); expect(redactSnapshotForGrade('block', snapshot)).toBe(snapshot); expect(redactSnapshotForGrade('delta', snapshot)).toBe(snapshot); + + const reset: TranscriptOperation = { op: 'reset', agentId: 'main', snapshot }; + const filteredReset = filterOpsForGrade('turn', [reset])[0]; + expect(filteredReset?.op).toBe('reset'); + expect( + filteredReset?.op === 'reset' && + filteredReset.snapshot.items[0]?.kind === 'turn' && + filteredReset.snapshot.items[0].steps, + ).toEqual([]); }); }); @@ -504,4 +513,61 @@ describe('groupMessagesIntoSnapshot (cold path)', () => { if (taskTurn?.kind !== 'turn') throw new Error('expected turn'); expect(taskTurn.origin).toMatchObject({ kind: 'task', taskId: 'b83rhswvs' }); }); + + it('uses persisted turn ids so retry turns and same-turn steers do not renumber history', () => { + const messages = [ + { role: 'user', content: [{ type: 'text' as const, text: 'u0' }], origin: { kind: 'user' } }, + { role: 'assistant', content: [{ type: 'text' as const, text: 'a0' }] }, + { role: 'assistant', content: [{ type: 'text' as const, text: 'retry answer' }] }, + { role: 'user', content: [{ type: 'text' as const, text: 'u2' }], origin: { kind: 'user' } }, + { role: 'assistant', content: [{ type: 'text' as const, text: 'before steer' }] }, + ]; + const snapshot = groupMessagesIntoSnapshot(messages, { + turnIds: [0, 0, 1, 2, 2], + turns: [ + { turnId: 0, input: messages[0]!.content, origin: { kind: 'user' } }, + { turnId: 1, input: [], origin: { kind: 'retry' } }, + { turnId: 2, input: messages[3]!.content, origin: { kind: 'user' } }, + ], + }); + + const turns = snapshot.items.filter((item) => item.kind === 'turn'); + expect(turns.map((turn) => turn.turnId)).toEqual(['t0', 't1', 't2']); + expect(turns[1]?.prompt).toBeUndefined(); + expect(turns[1]?.steps[0]?.frames[0]).toMatchObject({ + kind: 'text', + text: 'retry answer', + }); + expect(turns[2]?.steps).toHaveLength(1); + expect(turns[2]?.steps[0]?.frames[0]).toMatchObject({ + kind: 'text', + text: 'before steer', + }); + }); + + it('projects a legacy idle steer as its own stable turn', () => { + const messages = [ + { role: 'user', content: [{ type: 'text' as const, text: 'u0' }], origin: { kind: 'user' } }, + { role: 'assistant', content: [{ type: 'text' as const, text: 'a0' }] }, + { role: 'user', content: [{ type: 'text' as const, text: 'u1' }], origin: { kind: 'user' } }, + { role: 'assistant', content: [{ type: 'text' as const, text: 'a1' }] }, + ]; + const snapshot = groupMessagesIntoSnapshot(messages, { + turnIds: [0, 0, 1, 1], + turns: [ + { turnId: 0, input: messages[0]!.content, origin: { kind: 'user' } }, + { turnId: 1, input: messages[2]!.content, origin: { kind: 'user' } }, + ], + }); + + const turns = snapshot.items.filter((item) => item.kind === 'turn'); + expect(turns.map((turn) => [turn.turnId, turn.prompt])).toEqual([ + ['t0', 'u0'], + ['t1', 'u1'], + ]); + expect(turns.map((turn) => turn.steps[0]?.frames[0])).toMatchObject([ + { kind: 'text', text: 'a0' }, + { kind: 'text', text: 'a1' }, + ]); + }); });