diff --git a/.changeset/feat-undo-command.md b/.changeset/feat-undo-command.md new file mode 100644 index 00000000..b43085f5 --- /dev/null +++ b/.changeset/feat-undo-command.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/agent-core": minor +"@moonshot-ai/kimi-code": minor +--- + +Add `/undo` slash command to withdraw the last prompt from conversation history, and keep replay records in sync when a prompt is undone. diff --git a/apps/kimi-code/src/tui/commands/dispatch.ts b/apps/kimi-code/src/tui/commands/dispatch.ts index 3bd878b0..930dcbb9 100644 --- a/apps/kimi-code/src/tui/commands/dispatch.ts +++ b/apps/kimi-code/src/tui/commands/dispatch.ts @@ -42,6 +42,7 @@ import { handleInitCommand, handleTitleCommand, } from './session'; +import { handleUndoCommand } from './undo'; // --------------------------------------------------------------------------- // Re-exports — keep existing consumers working @@ -78,6 +79,7 @@ export { handleInitCommand, handleTitleCommand, } from './session'; +export { handleUndoCommand } from './undo'; // --------------------------------------------------------------------------- // Host interface @@ -279,6 +281,9 @@ async function handleBuiltInSlashCommand( case 'logout': await handleLogoutCommand(host); return; + case 'undo': + await handleUndoCommand(host, args); + return; default: host.showError(`Unknown slash command: /${String(name)}`); return; diff --git a/apps/kimi-code/src/tui/commands/index.ts b/apps/kimi-code/src/tui/commands/index.ts index 60178b26..261f70ae 100644 --- a/apps/kimi-code/src/tui/commands/index.ts +++ b/apps/kimi-code/src/tui/commands/index.ts @@ -34,6 +34,7 @@ export { handleInitCommand, handleTitleCommand, } from './session'; +export { handleUndoCommand } from './undo'; export { promptApiKey, promptCatalogProviderSelection, diff --git a/apps/kimi-code/src/tui/commands/registry.ts b/apps/kimi-code/src/tui/commands/registry.ts index faf76b57..c2c99cf2 100644 --- a/apps/kimi-code/src/tui/commands/registry.ts +++ b/apps/kimi-code/src/tui/commands/registry.ts @@ -127,6 +127,13 @@ export const BUILTIN_SLASH_COMMANDS = [ priority: 60, availability: 'always', }, + { + name: 'undo', + aliases: [], + description: 'Withdraw the last prompt from the transcript', + priority: 80, + availability: 'idle-only', + }, { name: 'editor', aliases: [], diff --git a/apps/kimi-code/src/tui/commands/undo.ts b/apps/kimi-code/src/tui/commands/undo.ts new file mode 100644 index 00000000..752fd27a --- /dev/null +++ b/apps/kimi-code/src/tui/commands/undo.ts @@ -0,0 +1,188 @@ +import type { Component } from '@earendil-works/pi-tui'; + +import { WelcomeComponent } from '../components/chrome/welcome'; +import { AgentGroupComponent } from '../components/messages/agent-group'; +import { AssistantMessageComponent } from '../components/messages/assistant-message'; +import { BackgroundAgentStatusComponent } from '../components/messages/background-agent-status'; +import { CronMessageComponent } from '../components/messages/cron-message'; +import { ReadGroupComponent } from '../components/messages/read-group'; +import { SkillActivationComponent } from '../components/messages/skill-activation'; +import { ThinkingComponent } from '../components/messages/thinking'; +import { ToolCallComponent } from '../components/messages/tool-call'; +import { UserMessageComponent } from '../components/messages/user-message'; +import { NO_ACTIVE_SESSION_MESSAGE } from '../constant/kimi-tui'; +import type { TranscriptEntry } from '../types'; +import { formatErrorMessage } from '../utils/event-payload'; +import { getTranscriptComponentEntry } from '../utils/transcript-component-metadata'; +import type { SlashCommandHost } from './dispatch'; + +// --------------------------------------------------------------------------- +// Undo command +// --------------------------------------------------------------------------- + +export async function handleUndoCommand( + host: SlashCommandHost, + args: string = '', +): Promise { + if (host.state.appState.streamingPhase !== 'idle') { + host.showError('Cannot undo while streaming — press Esc or Ctrl-C first.'); + return; + } + + const count = parseUndoCount(args); + if (count === undefined) { + host.showError('Usage: /undo [count], where count is a positive integer.'); + return; + } + + const session = host.session; + if (session === undefined) { + host.showError(NO_ACTIVE_SESSION_MESSAGE); + return; + } + + const entries = host.state.transcriptEntries; + const lastUserIndex = findUndoAnchorEntryIndex(entries, count); + if (lastUserIndex === undefined) { + host.showError('Nothing to undo.'); + return; + } + + try { + await session.undoHistory(count); + } catch (error) { + const message = formatErrorMessage(error); + host.showError(`Failed to undo: ${message}`); + return; + } + + const children = host.state.transcriptContainer.children; + const lastUserComponentIndex = findUndoAnchorComponentIndex(children, count); + if (lastUserComponentIndex !== undefined) { + removeUndoContextComponents(children, lastUserComponentIndex); + host.state.transcriptContainer.invalidate(); + } + + const preservedEntries = entries.slice(lastUserIndex).filter( + (entry) => !isUndoContextEntry(entry), + ); + entries.splice(lastUserIndex, entries.length - lastUserIndex, ...preservedEntries); + + if (entries.length === 0) { + renderWelcome(host); + } + + host.state.ui.requestRender(); +} + +function parseUndoCount(args: string): number | undefined { + const value = args.trim(); + if (value.length === 0) return 1; + if (!/^[1-9]\d*$/.test(value)) return undefined; + const count = Number(value); + return Number.isSafeInteger(count) ? count : undefined; +} + +function isUndoAnchorEntry(entry: TranscriptEntry): boolean { + return ( + entry.kind === 'user' || + (entry.kind === 'skill_activation' && entry.skillTrigger === 'user-slash') + ); +} + +function findUndoAnchorEntryIndex( + entries: readonly TranscriptEntry[], + count: number, +): number | undefined { + let found = 0; + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry !== undefined && isUndoAnchorEntry(entry)) { + found++; + if (found === count) return i; + } + } + return undefined; +} + +function isUndoContextEntry(entry: TranscriptEntry): boolean { + switch (entry.kind) { + case 'user': + case 'assistant': + case 'tool_call': + case 'thinking': + case 'skill_activation': + case 'cron': + return true; + case 'status': + return entry.turnId !== undefined; + case 'welcome': + return false; + } +} + +function findUndoAnchorComponentIndex( + children: readonly Component[], + count: number, +): number | undefined { + let found = 0; + for (let i = children.length - 1; i >= 0; i--) { + const child = children[i]; + if (child !== undefined && isUndoAnchorComponent(child)) { + found++; + if (found === count) return i; + } + } + return undefined; +} + +function removeUndoContextComponents( + children: Component[], + startIndex: number, +): void { + for (let i = children.length - 1; i >= startIndex; i--) { + const child = children[i]; + if (child !== undefined && isUndoContextComponent(child)) { + children.splice(i, 1); + } + } +} + +function isUndoAnchorComponent(child: Component): boolean { + return ( + child instanceof UserMessageComponent || + (child instanceof SkillActivationComponent && child.trigger === 'user-slash') + ); +} + +function isUndoContextComponent(child: Component): boolean { + const entry = getTranscriptComponentEntry(child); + if (entry !== undefined) { + return isUndoContextEntry(entry); + } + + return ( + child instanceof UserMessageComponent || + child instanceof AssistantMessageComponent || + child instanceof ThinkingComponent || + child instanceof ToolCallComponent || + child instanceof AgentGroupComponent || + child instanceof ReadGroupComponent || + child instanceof SkillActivationComponent || + child instanceof BackgroundAgentStatusComponent || + child instanceof CronMessageComponent + ); +} + +function renderWelcome(host: SlashCommandHost): void { + if ( + host.state.transcriptContainer.children.some( + (child) => child instanceof WelcomeComponent, + ) + ) { + return; + } + host.state.transcriptContainer.addChild( + new WelcomeComponent(host.state.appState, host.state.theme.colors), + ); +} diff --git a/apps/kimi-code/src/tui/components/messages/skill-activation.ts b/apps/kimi-code/src/tui/components/messages/skill-activation.ts index 7f9e05c9..4526328c 100644 --- a/apps/kimi-code/src/tui/components/messages/skill-activation.ts +++ b/apps/kimi-code/src/tui/components/messages/skill-activation.ts @@ -16,11 +16,17 @@ import { Container, Text, Spacer } from '@earendil-works/pi-tui'; import chalk from 'chalk'; import type { ColorPalette } from '#/tui/theme/colors'; +import type { SkillActivationTrigger } from '#/tui/types'; const ARGS_PREVIEW_MAX = 200; export class SkillActivationComponent extends Container { - constructor(name: string, args: string | undefined, colors: ColorPalette) { + constructor( + name: string, + args: string | undefined, + colors: ColorPalette, + readonly trigger?: SkillActivationTrigger, + ) { super(); this.addChild(new Spacer(1)); const head = diff --git a/apps/kimi-code/src/tui/controllers/session-event-handler.ts b/apps/kimi-code/src/tui/controllers/session-event-handler.ts index 27b13a13..0a961c84 100644 --- a/apps/kimi-code/src/tui/controllers/session-event-handler.ts +++ b/apps/kimi-code/src/tui/controllers/session-event-handler.ts @@ -662,6 +662,7 @@ export class SessionEventHandler { skillActivationId: event.activationId, skillName: event.skillName, skillArgs: event.skillArgs, + skillTrigger: event.trigger, }); } diff --git a/apps/kimi-code/src/tui/controllers/session-replay.ts b/apps/kimi-code/src/tui/controllers/session-replay.ts index 27b76ceb..fa1af75f 100644 --- a/apps/kimi-code/src/tui/controllers/session-replay.ts +++ b/apps/kimi-code/src/tui/controllers/session-replay.ts @@ -345,6 +345,7 @@ export class SessionReplayRenderer { skillActivationId: skill.activationId, skillName: skill.skillName, skillArgs: skill.skillArgs, + skillTrigger: skill.trigger, }); } diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 9702ed78..589ef5a1 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -120,6 +120,7 @@ import { installTerminalFocusTracking } from './utils/terminal-focus'; import { notifyTerminalOnce } from './utils/terminal-notification'; import { installTerminalThemeTracking } from './utils/terminal-theme'; import { detectTmuxKeyboardWarning } from './utils/tmux-keyboard'; +import { markTranscriptComponent } from './utils/transcript-component-metadata'; import { nextTranscriptId } from './utils/transcript-id'; export type { TUIState } from './tui-state'; @@ -1177,6 +1178,7 @@ export class KimiTUI { entry.skillName ?? entry.content, entry.skillArgs, this.state.theme.colors, + entry.skillTrigger, ); case 'cron': return new CronMessageComponent( @@ -1241,6 +1243,7 @@ export class KimiTUI { this.state.transcriptEntries.push(entry); const component = this.createTranscriptComponent(entry); if (component) { + markTranscriptComponent(component, entry); this.state.transcriptContainer.addChild(component); this.state.ui.requestRender(); } @@ -1267,12 +1270,20 @@ export class KimiTUI { this.appendTranscriptEntry({ id: nextTranscriptId(), kind: 'status', + turnId: request.turnId === undefined ? undefined : String(request.turnId), renderMode: 'notice', content: parts.join(''), }); } private renderWelcome(): void { + if ( + this.state.transcriptContainer.children.some( + (child) => child instanceof WelcomeComponent, + ) + ) { + return; + } const welcome = new WelcomeComponent(this.state.appState, this.state.theme.colors); this.state.transcriptContainer.addChild(welcome); } diff --git a/apps/kimi-code/src/tui/types.ts b/apps/kimi-code/src/tui/types.ts index 5a1a8ee4..cc5526c7 100644 --- a/apps/kimi-code/src/tui/types.ts +++ b/apps/kimi-code/src/tui/types.ts @@ -114,6 +114,8 @@ export type TranscriptEntryKind = | 'skill_activation' | 'cron'; +export type SkillActivationTrigger = 'user-slash' | 'model-tool' | 'nested-skill'; + export interface TranscriptEntry { id: string; kind: TranscriptEntryKind; @@ -130,6 +132,7 @@ export interface TranscriptEntry { skillActivationId?: string; skillName?: string; skillArgs?: string; + skillTrigger?: SkillActivationTrigger; } export type LivePaneMode = diff --git a/apps/kimi-code/src/tui/utils/message-replay.ts b/apps/kimi-code/src/tui/utils/message-replay.ts index 8b83186d..3f1638eb 100644 --- a/apps/kimi-code/src/tui/utils/message-replay.ts +++ b/apps/kimi-code/src/tui/utils/message-replay.ts @@ -11,6 +11,7 @@ import type { import type { AppState, BackgroundAgentMetadata, + SkillActivationTrigger, ToolCallBlockData, TranscriptEntry, } from '#/tui/types'; @@ -38,6 +39,7 @@ export interface SkillActivationProjection { readonly activationId: string; readonly skillName: string; readonly skillArgs?: string; + readonly trigger: SkillActivationTrigger; } export interface ReplayBackgroundProjection { @@ -203,6 +205,7 @@ export function skillActivationFromOrigin( activationId: origin.activationId, skillName: origin.skillName, skillArgs: origin.skillArgs, + trigger: origin.trigger, }; } diff --git a/apps/kimi-code/src/tui/utils/transcript-component-metadata.ts b/apps/kimi-code/src/tui/utils/transcript-component-metadata.ts new file mode 100644 index 00000000..12151958 --- /dev/null +++ b/apps/kimi-code/src/tui/utils/transcript-component-metadata.ts @@ -0,0 +1,15 @@ +import type { Component } from '@earendil-works/pi-tui'; + +import type { TranscriptEntry } from '../types'; + +const componentEntries = new WeakMap(); + +export function markTranscriptComponent(component: Component, entry: TranscriptEntry): void { + componentEntries.set(component, entry); +} + +export function getTranscriptComponentEntry( + component: Component, +): TranscriptEntry | undefined { + return componentEntries.get(component); +} diff --git a/apps/kimi-code/test/tui/commands/registry.test.ts b/apps/kimi-code/test/tui/commands/registry.test.ts index 74737fb5..c481aa2c 100644 --- a/apps/kimi-code/test/tui/commands/registry.test.ts +++ b/apps/kimi-code/test/tui/commands/registry.test.ts @@ -97,6 +97,7 @@ describe('built-in slash command registry', () => { 'status', 'theme', 'title', + 'undo', 'usage', 'version', 'yolo', diff --git a/apps/kimi-code/test/tui/commands/resolve.test.ts b/apps/kimi-code/test/tui/commands/resolve.test.ts index 07381c0b..c839cba9 100644 --- a/apps/kimi-code/test/tui/commands/resolve.test.ts +++ b/apps/kimi-code/test/tui/commands/resolve.test.ts @@ -63,6 +63,11 @@ describe('resolveSlashCommandInput', () => { commandName: 'resume', reason: 'streaming', }); + expect(resolve('/undo', { isStreaming: true })).toEqual({ + kind: 'blocked', + commandName: 'undo', + reason: 'streaming', + }); }); it('blocks model and session pickers while compacting', () => { diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 58c6f15a..f23e475c 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -12,6 +12,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; import { ApprovalPanelComponent } from '#/tui/components/dialogs/approval-panel'; import { KIMI_CODE_PLUGIN_MARKETPLACE_URL } from '#/constant/app'; +import { WelcomeComponent } from '#/tui/components/chrome/welcome'; import { ModelSelectorComponent } from '#/tui/components/dialogs/model-selector'; import { PluginMcpSelectorComponent, @@ -110,6 +111,7 @@ function makeSession(overrides: Record = {}) { prompt: vi.fn(async () => {}), steer: vi.fn(async () => {}), init: vi.fn(async () => {}), + undoHistory: vi.fn(async () => {}), cancel: vi.fn(async () => {}), cancelCompaction: vi.fn(async () => {}), getStatus: vi.fn(async () => ({ @@ -685,6 +687,288 @@ describe('KimiTUI message flow', () => { ]); }); + it('keeps the transcript intact when undo RPC fails', async () => { + const session = makeSession({ + undoHistory: vi.fn(async () => { + throw new Error('core rpc unavailable'); + }), + }); + const { driver } = await makeDriver(session); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain( + 'Error: Failed to undo: core rpc unavailable', + ); + }); + + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('hello'); + }); + + it('does not duplicate welcome after undoing the only turn', async () => { + const { driver } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(driver.state.transcriptEntries).toEqual([]); + }); + + expect( + driver.state.transcriptContainer.children.filter( + (child) => child instanceof WelcomeComponent, + ), + ).toHaveLength(1); + }); + + it('keeps command notices that are not part of the undone context', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + driver.handleUserInput('/auto on'); + + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain('Auto mode: ON'); + }); + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).not.toContain('hello'); + expect(transcript).toContain('Auto mode: ON'); + expect(driver.state.appState.permissionMode).toBe('auto'); + }); + + it('removes turn-scoped background status entries and restores welcome', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + driver.sessionEventHandler.handleEvent( + { + type: 'background.task.started', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + info: { + taskId: 'bash-bg123456', + command: 'npm test', + description: 'Run tests in background', + status: 'running', + pid: 1234, + exitCode: null, + startedAt: Date.now(), + endedAt: null, + }, + } as Event, + () => {}, + ); + + await vi.waitFor(() => { + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('bash task started in background'); + expect(transcript).toContain('Run tests in background'); + }); + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + + const transcript = stripSgr(renderTranscript(driver)); + expect(driver.state.transcriptEntries).toEqual([]); + expect(transcript).not.toContain('hello'); + expect(transcript).not.toContain('bash task started in background'); + expect(transcript).not.toContain('Run tests in background'); + expect( + driver.state.transcriptContainer.children.filter( + (child) => child instanceof WelcomeComponent, + ), + ).toHaveLength(1); + }); + + it('removes approval notices from undone turns', async () => { + const { driver, session } = await makeDriver(); + const approvalHandler = vi.mocked(session.setApprovalHandler).mock.calls[0]?.[0] as + | ((request: ApprovalRequest) => Promise) + | undefined; + if (approvalHandler === undefined) throw new Error('expected approval handler'); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + const response = approvalHandler({ + turnId: 1, + toolCallId: 'call_bash', + toolName: 'Bash', + action: 'Run shell command', + display: { + kind: 'generic', + summary: 'Run shell command', + detail: { command: 'echo ok', description: 'Run a shell command' }, + }, + }); + + await vi.waitFor(() => { + expect(driver.state.editorContainer.children[0]).toBeInstanceOf(ApprovalPanelComponent); + }); + (driver.state.editorContainer.children[0] as ApprovalPanelComponent).handleInput('1'); + await expect(response).resolves.toMatchObject({ decision: 'approved' }); + + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain('Approved: Run shell command'); + }); + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).not.toContain('hello'); + expect(transcript).not.toContain('Approved: Run shell command'); + }); + + it('undoes multiple turns when a count is provided', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('first'); + driver.state.appState.streamingPhase = 'idle'; + driver.handleUserInput('second'); + driver.state.appState.streamingPhase = 'idle'; + driver.handleUserInput('third'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo 2'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(2); + }); + + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'first', + }), + ]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('first'); + expect(transcript).not.toContain('second'); + expect(transcript).not.toContain('third'); + }); + + it('rejects invalid undo counts without changing context', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo 0'); + + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain( + 'Error: Usage: /undo [count], where count is a positive integer.', + ); + }); + + expect(session.undoHistory).not.toHaveBeenCalled(); + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + }); + + it('undoes from the real user turn when the last skill activation came from the model', async () => { + const { driver } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.sessionEventHandler.handleEvent( + { + type: 'skill.activated', + agentId: 'main', + activationId: 'act-model', + skillName: 'review', + trigger: 'model-tool', + } as Event, + () => {}, + ); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(driver.state.transcriptEntries).toEqual([]); + }); + + expect(driver.state.transcriptEntries).toEqual([]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).not.toContain('hello'); + expect(transcript).not.toContain('review'); + }); + + it('keeps user-slash skill activations as undo anchors', async () => { + const { driver } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.sessionEventHandler.handleEvent( + { + type: 'skill.activated', + agentId: 'main', + activationId: 'act-user', + skillName: 'review', + trigger: 'user-slash', + } as Event, + () => {}, + ); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + }); + + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('hello'); + expect(transcript).not.toContain('review'); + }); + it('sends pasted image placeholders as image content parts', async () => { const { driver, session } = await makeDriver(); const imageStore = (driver as unknown as { imageStore: ImageAttachmentStore }).imageStore; diff --git a/packages/agent-core/src/agent/compaction/micro.ts b/packages/agent-core/src/agent/compaction/micro.ts index 938f8769..2dbfa5a5 100644 --- a/packages/agent-core/src/agent/compaction/micro.ts +++ b/packages/agent-core/src/agent/compaction/micro.ts @@ -32,8 +32,8 @@ export class MicroCompaction { this.config = { ...DEFAULT_CONFIG, ...config }; } - reset(): void { - this.cutoff = 0; + reset(maxCutoff = 0): void { + this.cutoff = Math.min(this.cutoff, maxCutoff); } apply(cutoff: number): void { diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index 4b3810f9..386dccd9 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -1,6 +1,7 @@ import { createToolMessage, type ContentPart, type Message } from '@moonshot-ai/kosong'; import type { Agent } from '..'; +import { ErrorCodes, KimiError } from '../../errors'; import type { ExecutableToolResult, LoopRecordedEvent } from '../../loop'; import { estimateTokensForMessages } from '../../utils/tokens'; import type { CompactionResult } from '../compaction'; @@ -71,6 +72,58 @@ export class ContextMemory { this.agent.emitStatusUpdated(); } + undo(count: number): void { + if (count <= 0) return; + if (this._history.length === 0) return; + + this.agent.records.logRecord({ type: 'context.undo', count }); + + let removedUserCount = 0; + const removedMessages = new Set(); + let stoppedAtBoundary = false; + for (let i = this._history.length - 1; i >= 0; i--) { + const message = this._history[i]; + if (message === undefined) continue; + if (message.origin?.kind === 'injection') continue; + if (message.origin?.kind === 'compaction_summary') { + stoppedAtBoundary = true; + break; + } + + removedMessages.add(message); + this._history.splice(i, 1); + this.agent.injection.onContextMessageRemoved(i); + + if (i < this.tokenCountCoveredMessageCount) { + this.tokenCountCoveredMessageCount--; + this._tokenCount -= estimateTokensForMessages([message]); + } + + if (isRealUserPrompt(message)) { + removedUserCount++; + if (removedUserCount >= count) break; + } + } + + this.agent.replayBuilder.removeLastMessages(removedMessages); + + this.openSteps.clear(); + this.pendingToolResultIds.clear(); + this.deferredMessages = []; + this.agent.microCompaction.reset(this._history.length); + this.agent.emitStatusUpdated(); + + if ( + !this.agent.records.restoring && + (stoppedAtBoundary || removedUserCount < count) + ) { + throw new KimiError( + ErrorCodes.REQUEST_INVALID, + 'Nothing to undo in the active context.', + ); + } + } + applyCompaction(summary: CompactionResult): void { this.agent.records.logRecord({ type: 'context.apply_compaction', @@ -263,3 +316,13 @@ function toolResultOutputForModel(result: ExecutableToolResult): string | Conten function isEmptyOutputText(output: string): boolean { return output.length === 0 || output.trim() === TOOL_OUTPUT_EMPTY_TEXT; } + +function isRealUserPrompt(message: ContextMessage): boolean { + if (message.role !== 'user') return false; + const origin = message.origin; + if (origin === undefined || origin.kind === 'user') return true; + if (origin.kind === 'skill_activation') { + return origin.trigger === 'user-slash'; + } + return false; +} diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index d38e0c32..a3738789 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/packages/agent-core/src/agent/index.ts @@ -298,6 +298,9 @@ export class Agent { } this.turn.cancel(payload.turnId); }, + undoHistory: (payload) => { + this.context.undo(payload.count); + }, setThinking: (payload) => { const wasEnabled = this.config.thinkingLevel !== 'off'; this.config.update({ thinkingLevel: payload.level }); diff --git a/packages/agent-core/src/agent/injection/injector.ts b/packages/agent-core/src/agent/injection/injector.ts index 1084e32e..504e412d 100644 --- a/packages/agent-core/src/agent/injection/injector.ts +++ b/packages/agent-core/src/agent/injection/injector.ts @@ -16,6 +16,15 @@ export abstract class DynamicInjector { } } + onContextMessageRemoved(index: number): void { + if (this.injectedAt === null) return; + if (index < this.injectedAt) { + this.injectedAt--; + } else if (index === this.injectedAt) { + this.injectedAt = null; + } + } + async inject(): Promise { const injection = await this.getInjection(); if (injection) { diff --git a/packages/agent-core/src/agent/injection/manager.ts b/packages/agent-core/src/agent/injection/manager.ts index edda42c4..1de6aa92 100644 --- a/packages/agent-core/src/agent/injection/manager.ts +++ b/packages/agent-core/src/agent/injection/manager.ts @@ -36,4 +36,10 @@ export class InjectionManager { } } } + + onContextMessageRemoved(index: number): void { + for (const injector of this.injectors) { + injector.onContextMessageRemoved(index); + } + } } diff --git a/packages/agent-core/src/agent/records/index.ts b/packages/agent-core/src/agent/records/index.ts index cf79270f..c257c816 100644 --- a/packages/agent-core/src/agent/records/index.ts +++ b/packages/agent-core/src/agent/records/index.ts @@ -82,6 +82,9 @@ function restoreAgentRecord(agent: Agent, input: AgentRecord): void { case 'context.apply_compaction': agent.context.applyCompaction(input); return; + case 'context.undo': + agent.context.undo(input.count); + return; case 'tools.register_user_tool': agent.tools.registerUserTool(input); return; diff --git a/packages/agent-core/src/agent/records/types.ts b/packages/agent-core/src/agent/records/types.ts index c8acefb7..0e78f545 100644 --- a/packages/agent-core/src/agent/records/types.ts +++ b/packages/agent-core/src/agent/records/types.ts @@ -70,6 +70,7 @@ export interface AgentRecordEvents { 'context.append_loop_event': { event: LoopRecordedEvent }; 'context.clear': {}; 'context.apply_compaction': CompactionResult; + 'context.undo': { count: number }; 'tools.update_store': ToolStoreUpdate; } diff --git a/packages/agent-core/src/agent/replay/index.ts b/packages/agent-core/src/agent/replay/index.ts index 0196e622..acb76005 100644 --- a/packages/agent-core/src/agent/replay/index.ts +++ b/packages/agent-core/src/agent/replay/index.ts @@ -1,5 +1,6 @@ import type { Agent } from '..'; import type { AgentReplayRecord } from '../..'; +import type { ContextMessage } from '../context'; export class ReplayBuilder { protected readonly records: AgentReplayRecord[] = []; @@ -12,6 +13,16 @@ export class ReplayBuilder { } } + removeLastMessages(removedMessages: ReadonlySet): void { + if (removedMessages.size === 0) return; + for (let i = this.records.length - 1; i >= 0; i--) { + const record = this.records[i]!; + if (record.type === 'message' && removedMessages.has(record.message)) { + this.records.splice(i, 1); + } + } + } + buildResult(): readonly AgentReplayRecord[] { return this.records; } diff --git a/packages/agent-core/src/rpc/core-api.ts b/packages/agent-core/src/rpc/core-api.ts index 504e9a30..ea6be5b9 100644 --- a/packages/agent-core/src/rpc/core-api.ts +++ b/packages/agent-core/src/rpc/core-api.ts @@ -154,6 +154,9 @@ export interface CancelPlanPayload { export interface BeginCompactionPayload { readonly instruction?: string; } +export interface UndoHistoryPayload { + readonly count: number; +} export interface RegisterToolPayload { readonly name: string; readonly description: string; @@ -265,6 +268,7 @@ export interface AgentAPI { prompt: (payload: PromptPayload) => void; steer: (payload: SteerPayload) => void; cancel: (payload: CancelPayload) => void; + undoHistory: (payload: UndoHistoryPayload) => void; setThinking: (payload: SetThinkingPayload) => void; setPermission: (payload: SetPermissionPayload) => void; setModel: (payload: SetModelPayload) => SetModelResult; diff --git a/packages/agent-core/src/rpc/core-impl.ts b/packages/agent-core/src/rpc/core-impl.ts index 26e0f7aa..24d153bf 100644 --- a/packages/agent-core/src/rpc/core-impl.ts +++ b/packages/agent-core/src/rpc/core-impl.ts @@ -84,6 +84,7 @@ import type { SkillSummary, SteerPayload, StopBackgroundPayload, + UndoHistoryPayload, UnregisterToolPayload, UpdateSessionMetadataPayload, } from './core-api'; @@ -430,6 +431,10 @@ export class KimiCore implements PromisableMethods { return this.sessionApi(sessionId).cancel(payload); } + undoHistory({ sessionId, ...payload }: SessionAgentPayload) { + return this.sessionApi(sessionId).undoHistory(payload); + } + async setModel({ sessionId, ...payload diff --git a/packages/agent-core/src/session/rpc.ts b/packages/agent-core/src/session/rpc.ts index be5eac82..883827f5 100644 --- a/packages/agent-core/src/session/rpc.ts +++ b/packages/agent-core/src/session/rpc.ts @@ -23,6 +23,7 @@ import type { SkillSummary, SteerPayload, StopBackgroundPayload, + UndoHistoryPayload, UnregisterToolPayload, UpdateSessionMetadataPayload, } from '#/rpc'; @@ -103,6 +104,10 @@ export class SessionAPIImpl implements PromisableMethods { return this.getAgent(agentId).cancel(payload); } + undoHistory({ agentId, ...payload }: AgentScopedPayload) { + return this.getAgent(agentId).undoHistory(payload); + } + setModel({ agentId, ...payload }: AgentScopedPayload) { return this.getAgent(agentId).setModel(payload); } diff --git a/packages/agent-core/test/agent/compaction/micro.test.ts b/packages/agent-core/test/agent/compaction/micro.test.ts index c0737954..7975f6a0 100644 --- a/packages/agent-core/test/agent/compaction/micro.test.ts +++ b/packages/agent-core/test/agent/compaction/micro.test.ts @@ -405,6 +405,39 @@ describe('MicroCompaction', () => { ]); }); + it('clamps cutoff when undo shortens the context', () => { + vi.useFakeTimers(); + const ctx = testAgent({ + microCompaction: { + keepRecentMessages: 2, + minContentTokens: 1, + cacheMissedThresholdMs: 60 * MINUTE, + minContextUsageRatio: 0, + }, + }); + + vi.setSystemTime(0); + appendMicroToolExchange(ctx, 1, { output: 'result one' }); + appendMicroToolExchange(ctx, 2, { output: 'result two' }); + appendMicroToolExchange(ctx, 3, { output: 'result three' }); + + vi.setSystemTime(61 * MINUTE); + ctx.agent.microCompaction.detect(); + expect(toolTexts(ctx.agent.context.messages)).toEqual([ + DEFAULT_MARKER, + DEFAULT_MARKER, + 'result three', + ]); + + ctx.agent.context.undo(2); + appendMicroToolExchange(ctx, 4, { output: 'result four' }); + + expect(toolTexts(ctx.agent.context.messages)).toEqual([ + DEFAULT_MARKER, + 'result four', + ]); + }); + it('tracks telemetry when a cache miss advances the micro-compaction cutoff', () => { vi.useFakeTimers(); const records: TelemetryRecord[] = []; diff --git a/packages/agent-core/test/agent/context.test.ts b/packages/agent-core/test/agent/context.test.ts index 5bf02e9f..610b97ca 100644 --- a/packages/agent-core/test/agent/context.test.ts +++ b/packages/agent-core/test/agent/context.test.ts @@ -507,6 +507,149 @@ describe('Agent context', () => { ); }); + it('undo only counts real user prompts, skipping background notifications', () => { + const ctx = testAgent(); + ctx.configure(); + + ctx.appendAssistantText(1, 'first response'); + ctx.appendAssistantText(2, 'second response'); + + // Append a background task notification (role: 'user' but not a real prompt) + ctx.agent.context.appendMessage({ + role: 'user', + content: [{ type: 'text', text: 'background task completed' }], + toolCalls: [], + origin: { + kind: 'background_task', + taskId: 'bash-001', + status: 'completed', + notificationId: 'task:bash-001:completed', + }, + }); + + expect(ctx.agent.context.history.map((m) => m.role)).toEqual([ + 'user', + 'assistant', + 'user', + 'assistant', + 'user', + ]); + + ctx.agent.context.undo(1); + + // Should remove the background notification, the second assistant, and the second user prompt + expect(ctx.agent.context.history.map((m) => m.role)).toEqual(['user', 'assistant']); + }); + + it('stops at compaction summary and records the requested undo count', () => { + const ctx = testAgent(); + ctx.configure(); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'old user message' }]); + ctx.agent.context.applyCompaction({ + summary: 'summary of compacted context', + compactedCount: 1, + tokensBefore: 100, + tokensAfter: 20, + }); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'recent user message' }]); + ctx.agent.context.appendMessage({ + role: 'assistant', + content: [{ type: 'text', text: 'recent answer' }], + toolCalls: [], + }); + ctx.newEvents(); + + expect(() => { + ctx.agent.context.undo(2); + }).toThrow('Nothing to undo in the active context.'); + + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'assistant', + origin: { kind: 'compaction_summary' }, + content: [{ type: 'text', text: 'summary of compacted context' }], + }), + ]); + expect(ctx.newEvents()).toContainEqual( + expect.objectContaining({ + type: '[wire]', + event: 'context.undo', + args: expect.objectContaining({ count: 2 }), + }), + ); + }); + + it('does not throw while restoring an undo that stops at compaction summary', () => { + const ctx = testAgent(); + ctx.configure(); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'old user message' }]); + ctx.agent.context.applyCompaction({ + summary: 'summary of compacted context', + compactedCount: 1, + tokensBefore: 100, + tokensAfter: 20, + }); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'recent user message' }]); + ctx.agent.context.appendMessage({ + role: 'assistant', + content: [{ type: 'text', text: 'recent answer' }], + toolCalls: [], + }); + + expect(() => { + ctx.agent.records.restore({ type: 'context.undo', count: 2 }); + }).not.toThrow(); + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'assistant', + origin: { kind: 'compaction_summary' }, + content: [{ type: 'text', text: 'summary of compacted context' }], + }), + ]); + }); + + it('preserves injection messages when undo removes the surrounding turn', () => { + const ctx = testAgent(); + ctx.configure(); + + ctx.dispatch({ + type: 'context.append_message', + message: userMessage('do the work', { kind: 'user' }), + }); + ctx.dispatch({ + type: 'context.append_message', + message: userMessage('Plan mode is active', { + kind: 'injection', + variant: 'plan_mode', + }), + }); + ctx.dispatch({ + type: 'context.append_message', + message: { + role: 'assistant', + content: [{ type: 'text', text: 'work done' }], + toolCalls: [], + }, + }); + + ctx.agent.context.undo(1); + + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'user', + origin: { kind: 'injection', variant: 'plan_mode' }, + }), + ]); + expect(ctx.agent.replayBuilder.buildResult()).toEqual([ + expect.objectContaining({ + type: 'message', + message: expect.objectContaining({ + origin: { kind: 'injection', variant: 'plan_mode' }, + }), + }), + ]); + }); + }); describe('Agent context notification projection', () => { diff --git a/packages/agent-core/test/agent/plan.test.ts b/packages/agent-core/test/agent/plan.test.ts index 8c32d48d..2f652584 100644 --- a/packages/agent-core/test/agent/plan.test.ts +++ b/packages/agent-core/test/agent/plan.test.ts @@ -579,6 +579,22 @@ describe('plan mode injection cadence', () => { expect(ctx.agent.context.history).toHaveLength(afterExit); await ctx.expectResumeMatches(); }); + + it('keeps the preserved injection index aligned after undo removes earlier messages', async () => { + const ctx = testAgent(); + ctx.configure(); + await ctx.agent.planMode.enter('test-plan', false); + + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'draft the plan' }]); + await ctx.agent.injection.inject(); + ctx.appendAssistantTurn(1, 'Plan drafted.'); + + ctx.agent.context.undo(1); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'new plan request' }]); + await ctx.agent.injection.inject(); + + expect(lastUserText(ctx.agent.context.history)).toContain('Plan mode is active'); + }); }); function delay(ms: number): Promise { diff --git a/packages/agent-core/test/agent/resume.test.ts b/packages/agent-core/test/agent/resume.test.ts index dafbc864..2d72b39e 100644 --- a/packages/agent-core/test/agent/resume.test.ts +++ b/packages/agent-core/test/agent/resume.test.ts @@ -323,6 +323,106 @@ describe('Agent resume', () => { }), }); }); + + it('removes replay messages matching undone history', async () => { + const persistence = new RecordingAgentPersistence([ + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'first prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.begin', + uuid: 'step-1', + turnId: '0', + step: 1, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'content.part', + uuid: 'part-1', + turnId: '0', + step: 1, + stepUuid: 'step-1', + part: { type: 'text', text: 'first response' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.end', + uuid: 'step-1', + turnId: '0', + step: 1, + }, + }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'second prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.begin', + uuid: 'step-2', + turnId: '1', + step: 1, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'content.part', + uuid: 'part-2', + turnId: '1', + step: 1, + stepUuid: 'step-2', + part: { type: 'text', text: 'second response' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.end', + uuid: 'step-2', + turnId: '1', + step: 1, + }, + }, + { type: 'context.undo', count: 1 }, + ]); + const ctx = testAgent({ persistence }); + + await ctx.agent.resume(); + + expect(ctx.agent.context.history).toHaveLength(2); + expect(ctx.agent.context.history[0]?.role).toBe('user'); + expect(ctx.agent.context.history[1]?.role).toBe('assistant'); + + const replay = ctx.agent.replayBuilder.buildResult(); + expect(replay).toHaveLength(2); + expect(replay[0]).toMatchObject({ + type: 'message', + message: expect.objectContaining({ role: 'user', content: [{ type: 'text', text: 'first prompt' }] }), + }); + expect(replay[1]).toMatchObject({ + type: 'message', + message: expect.objectContaining({ role: 'assistant', content: [{ type: 'text', text: 'first response' }] }), + }); + }); }); class RecordingAgentPersistence extends InMemoryAgentRecordPersistence { diff --git a/packages/node-sdk/src/rpc.ts b/packages/node-sdk/src/rpc.ts index 7346e5a5..d66ff679 100644 --- a/packages/node-sdk/src/rpc.ts +++ b/packages/node-sdk/src/rpc.ts @@ -318,6 +318,15 @@ export class SDKRpcClient { }); } + async undoHistory(input: SessionIdRpcInput & { count: number }): Promise { + const rpc = await this.getRpc(); + return rpc.undoHistory({ + sessionId: input.sessionId, + agentId: this.interactiveAgentId, + count: input.count, + }); + } + async getContext(input: SessionIdRpcInput): Promise { const rpc = await this.getRpc(); return rpc.getContext({ diff --git a/packages/node-sdk/src/session.ts b/packages/node-sdk/src/session.ts index 6dc395ef..24d8f286 100644 --- a/packages/node-sdk/src/session.ts +++ b/packages/node-sdk/src/session.ts @@ -166,6 +166,11 @@ export class Session { await this.rpc.cancelCompaction({ sessionId: this.id }); } + async undoHistory(count: number = 1): Promise { + this.ensureOpen(); + await this.rpc.undoHistory({ sessionId: this.id, count }); + } + async getContext(): Promise { this.ensureOpen(); return this.rpc.getContext({ sessionId: this.id });