Skip to content

Commit 2004aed

Browse files
kermanxkermanx
authored
feat(agent-core): add agent record migrations (#22)
* feat(agent-core): add agent record migrations * test(agent-core): include wire metadata in subagent fixtures * docs(agent-core): document wire version format
1 parent 0da6073 commit 2004aed

11 files changed

Lines changed: 354 additions & 95 deletions

File tree

‎.changeset/agent-record-migrations.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
"@moonshot-ai/agent-core": patch
3+
"@moonshot-ai/kimi-code": patch
4+
---
5+
6+
Add wire record migration handling during session replay.

‎packages/agent-core/src/agent/records/index.ts‎

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
import type { Agent } from '..';
22
import {
33
AGENT_WIRE_PROTOCOL_VERSION,
4-
type AgentRecord,
5-
type AgentRecordPersistence,
6-
} from './types';
4+
migrateWireRecord,
5+
resolveWireMigrations,
6+
type WireMigration,
7+
type WireMigrationRecord,
8+
} from './migration';
9+
import type { AgentRecord, AgentRecordPersistence } from './types';
710

811
export * from './types';
12+
export { AGENT_WIRE_PROTOCOL_VERSION } from './migration';
913
export {
1014
FileSystemAgentRecordPersistence,
1115
InMemoryAgentRecordPersistence,
@@ -138,11 +142,37 @@ export class AgentRecords {
138142

139143
async replay(): Promise<void> {
140144
if (!this.persistence) throw new Error('No persistence provided for AgentRecords');
145+
let migrations: readonly WireMigration[] = [];
146+
let hasMetadata = false;
147+
let shouldRewrite = false;
148+
const replayedRecords: AgentRecord[] = [];
141149
for await (const record of this.persistence.read()) {
142-
if (!this.metadataInitialized) {
150+
if (!hasMetadata) {
151+
if (record.type !== 'metadata') {
152+
throw new Error('AgentRecords replay expected metadata as the first record');
153+
}
154+
hasMetadata = true;
143155
this.metadataInitialized = true;
156+
const readVersion = record.protocol_version;
157+
migrations = resolveWireMigrations(readVersion);
158+
shouldRewrite = readVersion !== AGENT_WIRE_PROTOCOL_VERSION;
144159
}
145-
this.restore(record);
160+
let migratedRecord = migrateWireRecord(
161+
record as WireMigrationRecord,
162+
migrations,
163+
) as AgentRecord;
164+
if (migratedRecord.type === 'metadata') {
165+
migratedRecord = {
166+
...migratedRecord,
167+
protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
168+
};
169+
}
170+
replayedRecords.push(migratedRecord);
171+
this.restore(migratedRecord);
172+
}
173+
if (shouldRewrite) {
174+
this.persistence.rewrite(replayedRecords);
175+
await this.persistence.flush();
146176
}
147177
}
148178

‎packages/agent-core/src/agent/records/migration/index.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
// Wire protocol versions currently support only the `number.number` format.
2+
export const AGENT_WIRE_PROTOCOL_VERSION = '1.0';
3+
4+
export interface WireMigrationRecord {
5+
readonly type: string;
6+
[key: string]: unknown;
7+
}
8+
9+
export interface WireMigration {
10+
readonly sourceVersion: string;
11+
readonly targetVersion: string;
12+
migrateRecord(record: WireMigrationRecord): WireMigrationRecord;
13+
}
14+
15+
const MIGRATIONS: readonly WireMigration[] = [];
16+
17+
export function resolveWireMigrations(readVersion: string): readonly WireMigration[] {
18+
if (compareWireVersions(readVersion, AGENT_WIRE_PROTOCOL_VERSION) === 0) {
19+
return [];
20+
}
21+
if (compareWireVersions(readVersion, AGENT_WIRE_PROTOCOL_VERSION) > 0) {
22+
throw new Error(
23+
`Unsupported wire protocol version: ${readVersion} (current: ${AGENT_WIRE_PROTOCOL_VERSION})`,
24+
);
25+
}
26+
27+
const migrations: WireMigration[] = [];
28+
let version = readVersion;
29+
while (compareWireVersions(version, AGENT_WIRE_PROTOCOL_VERSION) < 0) {
30+
const migration = findMigration(version);
31+
if (migration === undefined) {
32+
throw new Error(`Missing wire migration for version ${version}`);
33+
}
34+
migrations.push(migration);
35+
version = migration.targetVersion;
36+
}
37+
38+
return migrations;
39+
}
40+
41+
export function migrateWireRecord(
42+
record: WireMigrationRecord,
43+
migrations: readonly WireMigration[],
44+
): WireMigrationRecord {
45+
return migrations.reduce(
46+
(current, migration) => migration.migrateRecord(current),
47+
record,
48+
);
49+
}
50+
51+
function findMigration(sourceVersion: string): WireMigration | undefined {
52+
for (const migration of MIGRATIONS) {
53+
if (migration.sourceVersion === sourceVersion) return migration;
54+
}
55+
}
56+
57+
function compareWireVersions(a: string, b: string): number {
58+
const partsA = a.split('.');
59+
const partsB = b.split('.');
60+
const maxLength = Math.max(partsA.length, partsB.length);
61+
62+
for (let i = 0; i < maxLength; i++) {
63+
const diff = Number(partsA[i] ?? '0') - Number(partsB[i] ?? '0');
64+
if (diff !== 0) return diff;
65+
}
66+
67+
return 0;
68+
}

‎packages/agent-core/src/agent/records/types.ts‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,6 @@ export type AgentRecordOf<K extends keyof AgentRecordEvents> = Extract<
8686
{ readonly type: K }
8787
>;
8888

89-
export const AGENT_WIRE_PROTOCOL_VERSION = '1.0';
90-
9189
export interface AgentRecordPersistence {
9290
read(): AsyncIterable<AgentRecord>;
9391
append(input: AgentRecord): void;

‎packages/agent-core/src/session/export/manifest.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { AGENT_WIRE_PROTOCOL_VERSION } from '../../agent/records/types';
1+
import { AGENT_WIRE_PROTOCOL_VERSION } from '../../agent/records';
22
import type { SessionWireScan } from '#/session/export/wire-scan';
33
import type { ExportSessionManifest, SessionSummary } from '#/rpc/core-api';
44

‎packages/agent-core/test/agent/harness/agent.ts‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ import {
1414
} from '../../../src/agent';
1515
import type { CompactionStrategy } from '../../../src/agent/compaction';
1616
import type { ApprovalResponse } from '../../../src/agent/permission';
17-
import { InMemoryAgentRecordPersistence } from '../../../src/agent/records';
17+
import {
18+
AGENT_WIRE_PROTOCOL_VERSION,
19+
InMemoryAgentRecordPersistence,
20+
} from '../../../src/agent/records';
1821
import type { KimiConfig } from '../../../src/config';
1922
import type { ExecutableToolResult } from '../../../src/loop';
2023
import type { Logger } from '../../../src/logging';
@@ -284,7 +287,9 @@ export class AgentTestContext {
284287
providerManager: this.agent.providerManager,
285288
generate: failOnResumeGenerate,
286289
compactionStrategy: this.options.compactionStrategy,
287-
persistence: new InMemoryAgentRecordPersistence(this.recordHistory.map(cloneRecord)),
290+
persistence: new InMemoryAgentRecordPersistence(
291+
withMetadata(this.recordHistory.map(cloneRecord)),
292+
),
288293
});
289294

290295
await resumed.agent.resume();
@@ -439,8 +444,12 @@ export class AgentTestContext {
439444
private wrapPersistence(persistence: AgentRecordPersistence): AgentRecordPersistence {
440445
return {
441446
read: () => this.readAndCapturePersistence(persistence),
442-
append: (event) => persistence.append(event),
443-
rewrite: (records) => persistence.rewrite(records),
447+
append: (event) => {
448+
persistence.append(event);
449+
},
450+
rewrite: (records) => {
451+
persistence.rewrite(records);
452+
},
444453
flush: () => persistence.flush(),
445454
close: () => persistence.close(),
446455
};
@@ -619,3 +628,15 @@ function buildSkillPrompt(content: string, args: string | undefined): string {
619628
function cloneRecord(event: AgentRecord): AgentRecord {
620629
return structuredClone(event);
621630
}
631+
632+
function withMetadata(events: readonly AgentRecord[]): readonly AgentRecord[] {
633+
if (events.length === 0 || events[0]?.type === 'metadata') return events;
634+
return [
635+
{
636+
type: 'metadata',
637+
protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
638+
created_at: 1,
639+
},
640+
...events,
641+
];
642+
}

‎packages/agent-core/test/agent/records/index.test.ts‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import { describe, expect, it } from 'vitest';
2+
3+
import {
4+
AGENT_WIRE_PROTOCOL_VERSION,
5+
AgentRecords,
6+
InMemoryAgentRecordPersistence,
7+
type AgentRecord,
8+
} from '../../../src/agent/records';
9+
10+
describe('AgentRecords persistence metadata', () => {
11+
it('writes metadata before the first persisted record', async () => {
12+
const persistence = new InMemoryAgentRecordPersistence();
13+
const records = new AgentRecords(() => {}, persistence);
14+
15+
records.logRecord({
16+
type: 'turn.prompt',
17+
input: [{ type: 'text', text: 'hello' }],
18+
origin: { kind: 'user' },
19+
});
20+
await records.flush();
21+
22+
expect(persistence.records).toHaveLength(2);
23+
expect(persistence.records[0]).toMatchObject({
24+
type: 'metadata',
25+
protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
26+
});
27+
expect(persistence.records[1]?.type).toBe('turn.prompt');
28+
});
29+
30+
it('does not write metadata when replaying an empty stream', async () => {
31+
const persistence = new InMemoryAgentRecordPersistence();
32+
const records = new AgentRecords(() => {}, persistence);
33+
34+
await records.replay();
35+
records.logRecord({
36+
type: 'turn.prompt',
37+
input: [{ type: 'text', text: 'one' }],
38+
origin: { kind: 'user' },
39+
});
40+
await records.flush();
41+
42+
expect(persistence.records.map((record) => record.type)).toEqual([
43+
'metadata',
44+
'turn.prompt',
45+
]);
46+
});
47+
48+
it('rejects replaying a non-empty stream without metadata', async () => {
49+
const persistence = new InMemoryAgentRecordPersistence([
50+
{
51+
type: 'turn.prompt',
52+
input: [{ type: 'text', text: 'one' }],
53+
origin: { kind: 'user' },
54+
},
55+
]);
56+
const records = new AgentRecords(() => {}, persistence);
57+
58+
await expect(records.replay()).rejects.toThrow(
59+
'AgentRecords replay expected metadata as the first record',
60+
);
61+
});
62+
63+
it('does not duplicate metadata after replaying existing records', async () => {
64+
const persistence = new InMemoryAgentRecordPersistence([
65+
{
66+
type: 'metadata',
67+
protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
68+
created_at: 1,
69+
},
70+
{
71+
type: 'turn.prompt',
72+
input: [{ type: 'text', text: 'one' }],
73+
origin: { kind: 'user' },
74+
},
75+
]);
76+
const records = new AgentRecords(() => {}, persistence);
77+
78+
await records.replay();
79+
records.logRecord({
80+
type: 'turn.prompt',
81+
input: [{ type: 'text', text: 'two' }],
82+
origin: { kind: 'user' },
83+
});
84+
await records.flush();
85+
86+
expect(persistence.records.map((record) => record.type)).toEqual([
87+
'metadata',
88+
'turn.prompt',
89+
'turn.prompt',
90+
]);
91+
expect(persistence.records.filter((record) => record.type === 'metadata')).toHaveLength(1);
92+
});
93+
94+
it('does not rewrite records that already use the current wire version', async () => {
95+
const persistence = new RecordingInMemoryAgentRecordPersistence([
96+
{
97+
type: 'metadata',
98+
protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
99+
created_at: 1,
100+
},
101+
{
102+
type: 'turn.prompt',
103+
input: [{ type: 'text', text: 'one' }],
104+
origin: { kind: 'user' },
105+
},
106+
]);
107+
const records = new AgentRecords(() => {}, persistence);
108+
109+
await records.replay();
110+
111+
expect(persistence.rewrites).toEqual([]);
112+
});
113+
114+
it('rejects replaying records from a newer wire version', async () => {
115+
const persistence = new InMemoryAgentRecordPersistence([
116+
{
117+
type: 'metadata',
118+
protocol_version: '9.9',
119+
created_at: 1,
120+
},
121+
]);
122+
const records = new AgentRecords(() => {}, persistence);
123+
124+
await expect(records.replay()).rejects.toThrow(
125+
`Unsupported wire protocol version: 9.9 (current: ${AGENT_WIRE_PROTOCOL_VERSION})`,
126+
);
127+
});
128+
129+
it('rejects replaying records without a registered migration path', async () => {
130+
const persistence = new InMemoryAgentRecordPersistence([
131+
{
132+
type: 'metadata',
133+
protocol_version: '0.9',
134+
created_at: 1,
135+
},
136+
]);
137+
const records = new AgentRecords(() => {}, persistence);
138+
139+
await expect(records.replay()).rejects.toThrow('Missing wire migration for version 0.9');
140+
});
141+
});
142+
143+
class RecordingInMemoryAgentRecordPersistence extends InMemoryAgentRecordPersistence {
144+
readonly rewrites: AgentRecord[][] = [];
145+
146+
override rewrite(records: readonly AgentRecord[]): void {
147+
this.rewrites.push([...records]);
148+
super.rewrite(records);
149+
}
150+
}

‎packages/agent-core/test/agent/records/migration.test.ts‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import { describe, expect, it } from 'vitest';
2+
3+
import {
4+
migrateWireRecord,
5+
type WireMigration,
6+
} from '../../../src/agent/records/migration';
7+
8+
describe('wire record migrations', () => {
9+
it('applies migrations in order', () => {
10+
const migrations: WireMigration[] = [
11+
{
12+
sourceVersion: '0.8',
13+
targetVersion: '0.9',
14+
migrateRecord: (record) => ({
15+
...record,
16+
first: true,
17+
}),
18+
},
19+
{
20+
sourceVersion: '0.9',
21+
targetVersion: '1.0',
22+
migrateRecord: (record) => ({
23+
...record,
24+
second: record['first'] === true,
25+
}),
26+
},
27+
];
28+
29+
expect(migrateWireRecord({ type: 'metadata' }, migrations)).toEqual({
30+
type: 'metadata',
31+
first: true,
32+
second: true,
33+
});
34+
});
35+
});

0 commit comments

Comments
 (0)