fix: replay compaction records on resume (#617)

This commit is contained in:
_Kerman 2026-06-12 11:10:52 +08:00 committed by GitHub
parent dcf30754d0
commit 911e7c3fcf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 213 additions and 7 deletions

View file

@ -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.

View file

@ -45,6 +45,7 @@ import type { SessionEventHandler } from './session-event-handler';
import type { TUIState } from '../tui-state';
type GoalReplayRecord = Extract<AgentReplayRecord, { type: 'goal_updated' }>;
type CompactionReplayRecord = Extract<AgentReplayRecord, { type: 'compaction' }>;
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;

View file

@ -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;
}

View file

@ -104,6 +104,7 @@ export interface BackgroundAgentStatusData {
}
export interface CompactionTranscriptData {
readonly result?: 'cancelled';
readonly tokensBefore?: number;
readonly tokensAfter?: number;
readonly instruction?: string;

View file

@ -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 },

View file

@ -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',

View file

@ -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();
}

View file

@ -17,6 +17,18 @@ export class ReplayBuilder {
}
}
patchLast<T extends AgentReplayRecord['type']>(
type: T,
patch: Partial<Extract<AgentReplayRecord, { type: T }>>,
): void {
if (this.agent.records.restoring) {
const last = this.records.at(-1);
if (last && last.type === type) {
Object.assign(last, patch);
}
}
}
removeLastMessages(removedMessages: ReadonlySet<ContextMessage>): void {
if (removedMessages.size === 0) return;
for (let i = this.records.length - 1; i >= 0; i--) {

View file

@ -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;

View file

@ -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([
{