From 911e7c3fcfc8a005b1b8d90388260d1a4032f76f Mon Sep 17 00:00:00 2001 From: _Kerman Date: Fri, 12 Jun 2026 11:10:52 +0800 Subject: [PATCH] fix: replay compaction records on resume (#617) --- .changeset/replay-compaction-records.md | 6 ++ .../src/tui/controllers/session-replay.ts | 28 +++++++ apps/kimi-code/src/tui/kimi-tui.ts | 6 +- apps/kimi-code/src/tui/types.ts | 1 + .../kimi-code/test/tui/message-replay.test.ts | 55 ++++++++++++ .../agent-core/src/agent/compaction/full.ts | 7 ++ .../agent-core/src/agent/context/index.ts | 20 +++-- packages/agent-core/src/agent/replay/index.ts | 12 +++ packages/agent-core/src/rpc/resumed.ts | 2 + packages/agent-core/test/agent/resume.test.ts | 83 +++++++++++++++++++ 10 files changed, 213 insertions(+), 7 deletions(-) create mode 100644 .changeset/replay-compaction-records.md diff --git a/.changeset/replay-compaction-records.md b/.changeset/replay-compaction-records.md new file mode 100644 index 000000000..3bc6c9de3 --- /dev/null +++ b/.changeset/replay-compaction-records.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kimi-code": patch +--- + +Show completed and cancelled compaction records correctly when resuming a session. diff --git a/apps/kimi-code/src/tui/controllers/session-replay.ts b/apps/kimi-code/src/tui/controllers/session-replay.ts index b7d761346..4a13fe373 100644 --- a/apps/kimi-code/src/tui/controllers/session-replay.ts +++ b/apps/kimi-code/src/tui/controllers/session-replay.ts @@ -45,6 +45,7 @@ import type { SessionEventHandler } from './session-event-handler'; import type { TUIState } from '../tui-state'; type GoalReplayRecord = Extract; +type CompactionReplayRecord = Extract; type GoalReplayLifecycleChange = GoalChange & { readonly kind: 'lifecycle' }; export interface SessionReplayHost { @@ -178,6 +179,9 @@ export class SessionReplayRenderer { case 'message': this.renderMessage(context, record.message); return; + case 'compaction': + this.renderCompaction(context, record); + return; case 'goal_updated': this.renderGoalReplayRecord(context, record); return; @@ -373,6 +377,30 @@ export class SessionReplayRenderer { }); } + private renderCompaction(context: ReplayRenderContext, record: CompactionReplayRecord): void { + this.flushAssistant(context); + if (record.result === undefined) return; + if (record.result === 'cancelled') { + this.host.appendTranscriptEntry({ + ...replayEntry(context, 'status', 'Compaction cancelled', 'plain'), + compactionData: { + result: 'cancelled', + instruction: record.instruction, + }, + }); + return; + } + + this.host.appendTranscriptEntry({ + ...replayEntry(context, 'status', 'Compaction complete', 'plain'), + compactionData: { + tokensBefore: record.result.tokensBefore, + tokensAfter: record.result.tokensAfter, + instruction: record.instruction, + }, + }); + } + private renderGoalReplayRecord(context: ReplayRenderContext, record: GoalReplayRecord): void { this.flushAssistant(context); const { change } = record; diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 2eb00d2c6..7a43affc5 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -1308,7 +1308,11 @@ export class KimiTUI { if (entry.compactionData !== undefined) { const data = entry.compactionData; const block = new CompactionComponent(this.state.ui, data.instruction); - block.markDone(data.tokensBefore, data.tokensAfter); + if (data.result === 'cancelled') { + block.markCanceled(); + } else { + block.markDone(data.tokensBefore, data.tokensAfter); + } return block; } diff --git a/apps/kimi-code/src/tui/types.ts b/apps/kimi-code/src/tui/types.ts index 836db9311..bbf047073 100644 --- a/apps/kimi-code/src/tui/types.ts +++ b/apps/kimi-code/src/tui/types.ts @@ -104,6 +104,7 @@ export interface BackgroundAgentStatusData { } export interface CompactionTranscriptData { + readonly result?: 'cancelled'; readonly tokensBefore?: number; readonly tokensAfter?: number; readonly instruction?: string; diff --git a/apps/kimi-code/test/tui/message-replay.test.ts b/apps/kimi-code/test/tui/message-replay.test.ts index 8a02b505b..51a5ac4bc 100644 --- a/apps/kimi-code/test/tui/message-replay.test.ts +++ b/apps/kimi-code/test/tui/message-replay.test.ts @@ -990,6 +990,61 @@ describe('KimiTUI resume message replay', () => { expect(transcript).toContain('hook response 2'); }); + it('renders replayed compaction records as completed compaction blocks', async () => { + const driver = await replayIntoDriver([ + message('user', [{ type: 'text', text: 'prompt before compaction' }]), + { + type: 'compaction', + result: { + summary: 'Compacted transcript summary.', + compactedCount: 4, + tokensBefore: 120, + tokensAfter: 24, + }, + instruction: 'preserve implementation notes', + }, + message('user', [{ type: 'text', text: 'prompt after compaction' }]), + ]); + + const compactionEntry = driver.state.transcriptEntries.find( + (entry) => entry.compactionData !== undefined, + ); + expect(compactionEntry?.compactionData).toEqual({ + tokensBefore: 120, + tokensAfter: 24, + instruction: 'preserve implementation notes', + }); + const transcript = stripAnsi(driver.state.transcriptContainer.render(120).join('\n')); + expect(transcript).toContain('Compaction complete'); + expect(transcript).toContain('120 → 24 tokens'); + expect(transcript).toContain('preserve implementation notes'); + expect(transcript).not.toContain('Compacted transcript summary.'); + }); + + it('renders replayed cancelled compaction records as cancelled compaction blocks', async () => { + const driver = await replayIntoDriver([ + message('user', [{ type: 'text', text: 'prompt before cancellation' }]), + { + type: 'compaction', + result: 'cancelled', + instruction: 'preserve implementation notes', + }, + message('user', [{ type: 'text', text: 'prompt after cancellation' }]), + ]); + + const compactionEntry = driver.state.transcriptEntries.find( + (entry) => entry.compactionData !== undefined, + ); + expect(compactionEntry?.compactionData).toEqual({ + result: 'cancelled', + instruction: 'preserve implementation notes', + }); + const transcript = stripAnsi(driver.state.transcriptContainer.render(120).join('\n')); + expect(transcript).toContain('Compaction cancelled'); + expect(transcript).toContain('preserve implementation notes'); + expect(transcript).not.toContain('Compaction complete'); + }); + it('renders plan permission and approval replay notices', async () => { const driver = await replayIntoDriver([ { time: REPLAY_TIME, type: 'plan_updated', enabled: true }, diff --git a/packages/agent-core/src/agent/compaction/full.ts b/packages/agent-core/src/agent/compaction/full.ts index 3c52043a9..2b8a624c6 100644 --- a/packages/agent-core/src/agent/compaction/full.ts +++ b/packages/agent-core/src/agent/compaction/full.ts @@ -100,6 +100,10 @@ export class FullCompaction { } if (this.compactionCountInTurn > this.strategy.maxCompactionPerTurn) return; if (this.agent.records.restoring) { + this.agent.replayBuilder.push({ + type: 'compaction', + instruction: data.instruction, + }); return; } const compactedCount = this.strategy.computeCompactCount(this.agent.context.history, data.source); @@ -139,6 +143,9 @@ export class FullCompaction { } private markCanceled(): void { + this.agent.replayBuilder.patchLast('compaction', { + result: 'cancelled', + }); if (!this.compacting) return; this.agent.records.logRecord({ type: 'full_compaction.cancel', diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index 19ca54f07..4746eeb36 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -146,26 +146,34 @@ export class ContextMemory { } } - applyCompaction(summary: CompactionResult): void { + applyCompaction(result: CompactionResult): void { this.agent.records.logRecord({ type: 'context.apply_compaction', - ...summary, + ...result, + }); + this.agent.replayBuilder.patchLast('compaction', { + result: { + summary: result.summary, + compactedCount: result.compactedCount, + tokensBefore: result.tokensBefore, + tokensAfter: result.tokensAfter, + }, }); this._history = [ { role: 'assistant', - content: [{ type: 'text', text: summary.summary }], + content: [{ type: 'text', text: result.summary }], toolCalls: [], origin: { kind: 'compaction_summary' }, }, - ...this._history.slice(summary.compactedCount), + ...this._history.slice(result.compactedCount), ]; this.openSteps.clear(); this.flushDeferredMessagesIfToolExchangeClosed(); - this._tokenCount = summary.tokensAfter; + this._tokenCount = result.tokensAfter; this.tokenCountCoveredMessageCount = this._history.length; this.agent.microCompaction.reset(); - this.agent.injection.onContextCompacted(summary.compactedCount); + this.agent.injection.onContextCompacted(result.compactedCount); this.agent.emitStatusUpdated(); } diff --git a/packages/agent-core/src/agent/replay/index.ts b/packages/agent-core/src/agent/replay/index.ts index c5f851869..c74747928 100644 --- a/packages/agent-core/src/agent/replay/index.ts +++ b/packages/agent-core/src/agent/replay/index.ts @@ -17,6 +17,18 @@ export class ReplayBuilder { } } + patchLast( + type: T, + patch: Partial>, + ): void { + if (this.agent.records.restoring) { + const last = this.records.at(-1); + if (last && last.type === type) { + Object.assign(last, patch); + } + } + } + removeLastMessages(removedMessages: ReadonlySet): void { if (removedMessages.size === 0) return; for (let i = this.records.length - 1; i >= 0; i--) { diff --git a/packages/agent-core/src/rpc/resumed.ts b/packages/agent-core/src/rpc/resumed.ts index 4fc635331..e90f2c353 100644 --- a/packages/agent-core/src/rpc/resumed.ts +++ b/packages/agent-core/src/rpc/resumed.ts @@ -1,5 +1,6 @@ import type { AgentType } from '#/agent'; import type { BackgroundTaskInfo } from '#/agent/background'; +import type { CompactionResult } from '#/agent/compaction'; import type { AgentConfigData, AgentConfigUpdateData } from '#/agent/config'; import type { AgentContextData, ContextMessage } from '#/agent/context'; import type { GoalChange, GoalSnapshot } from '#/agent/goal'; @@ -16,6 +17,7 @@ import type { SessionMeta } from '#/session'; export type AgentReplayRecordPayload = | { type: 'message'; message: ContextMessage } + | { type: 'compaction'; result?: CompactionResult | 'cancelled'; instruction?: string } | { type: 'goal_updated'; snapshot: GoalSnapshot; diff --git a/packages/agent-core/test/agent/resume.test.ts b/packages/agent-core/test/agent/resume.test.ts index 83df5fbc2..39b1e42cd 100644 --- a/packages/agent-core/test/agent/resume.test.ts +++ b/packages/agent-core/test/agent/resume.test.ts @@ -224,6 +224,89 @@ describe('Agent resume', () => { } }); + it('projects restored compactions into replay records', async () => { + const persistence = new RecordingAgentPersistence([ + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'Historical prompt before compaction' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'full_compaction.begin', + source: 'manual', + instruction: 'preserve implementation notes', + }, + { + type: 'full_compaction.complete', + }, + { + type: 'context.apply_compaction', + summary: 'Compacted implementation notes.', + compactedCount: 1, + tokensBefore: 120, + tokensAfter: 24, + }, + ]); + const ctx = testAgent({ persistence }); + + await ctx.agent.resume(); + + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'assistant', + content: [{ type: 'text', text: 'Compacted implementation notes.' }], + origin: { kind: 'compaction_summary' }, + }), + ]); + expect(ctx.agent.replayBuilder.buildResult()).toEqual([ + expect.objectContaining({ + type: 'message', + message: expect.objectContaining({ + role: 'user', + content: [{ type: 'text', text: 'Historical prompt before compaction' }], + }), + }), + { + type: 'compaction', + result: { + summary: 'Compacted implementation notes.', + compactedCount: 1, + tokensBefore: 120, + tokensAfter: 24, + }, + instruction: 'preserve implementation notes', + }, + ]); + }); + + it('projects restored cancelled compactions into replay records', async () => { + const persistence = new RecordingAgentPersistence([ + { + type: 'full_compaction.begin', + source: 'manual', + instruction: 'preserve implementation notes', + }, + { + type: 'full_compaction.cancel', + }, + ]); + const ctx = testAgent({ persistence }); + + await ctx.agent.resume(); + + expect(ctx.agent.replayBuilder.buildResult()).toEqual([ + { + type: 'compaction', + result: 'cancelled', + instruction: 'preserve implementation notes', + }, + ]); + }); + it('persists undelivered restored background notifications during resume', async () => { const persistence = new RecordingAgentPersistence([ {