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

.changeset/agent-record-migrations.md

@@ -0,0 +1,6 @@
+---
+"@moonshot-ai/agent-core": patch
+"@moonshot-ai/kimi-code": patch
+---
+
+Add wire record migration handling during session replay.

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

@@ -1,11 +1,15 @@
 import type { Agent } from '..';
 import {
   AGENT_WIRE_PROTOCOL_VERSION,
-  type AgentRecord,
-  type AgentRecordPersistence,
-} from './types';
+  migrateWireRecord,
+  resolveWireMigrations,
+  type WireMigration,
+  type WireMigrationRecord,
+} from './migration';
+import type { AgentRecord, AgentRecordPersistence } from './types';
 
 export * from './types';
+export { AGENT_WIRE_PROTOCOL_VERSION } from './migration';
 export {
   FileSystemAgentRecordPersistence,
   InMemoryAgentRecordPersistence,
@@ -138,11 +142,37 @@ export class AgentRecords {
 
   async replay(): Promise<void> {
     if (!this.persistence) throw new Error('No persistence provided for AgentRecords');
+    let migrations: readonly WireMigration[] = [];
+    let hasMetadata = false;
+    let shouldRewrite = false;
+    const replayedRecords: AgentRecord[] = [];
     for await (const record of this.persistence.read()) {
-      if (!this.metadataInitialized) {
+      if (!hasMetadata) {
+        if (record.type !== 'metadata') {
+          throw new Error('AgentRecords replay expected metadata as the first record');
+        }
+        hasMetadata = true;
         this.metadataInitialized = true;
+        const readVersion = record.protocol_version;
+        migrations = resolveWireMigrations(readVersion);
+        shouldRewrite = readVersion !== AGENT_WIRE_PROTOCOL_VERSION;
       }
-      this.restore(record);
+      let migratedRecord = migrateWireRecord(
+        record as WireMigrationRecord,
+        migrations,
+      ) as AgentRecord;
+      if (migratedRecord.type === 'metadata') {
+        migratedRecord = {
+          ...migratedRecord,
+          protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+        };
+      }
+      replayedRecords.push(migratedRecord);
+      this.restore(migratedRecord);
+    }
+    if (shouldRewrite) {
+      this.persistence.rewrite(replayedRecords);
+      await this.persistence.flush();
     }
   }
 

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

@@ -0,0 +1,68 @@
+// Wire protocol versions currently support only the `number.number` format.
+export const AGENT_WIRE_PROTOCOL_VERSION = '1.0';
+
+export interface WireMigrationRecord {
+  readonly type: string;
+  [key: string]: unknown;
+}
+
+export interface WireMigration {
+  readonly sourceVersion: string;
+  readonly targetVersion: string;
+  migrateRecord(record: WireMigrationRecord): WireMigrationRecord;
+}
+
+const MIGRATIONS: readonly WireMigration[] = [];
+
+export function resolveWireMigrations(readVersion: string): readonly WireMigration[] {
+  if (compareWireVersions(readVersion, AGENT_WIRE_PROTOCOL_VERSION) === 0) {
+    return [];
+  }
+  if (compareWireVersions(readVersion, AGENT_WIRE_PROTOCOL_VERSION) > 0) {
+    throw new Error(
+      `Unsupported wire protocol version: ${readVersion} (current: ${AGENT_WIRE_PROTOCOL_VERSION})`,
+    );
+  }
+
+  const migrations: WireMigration[] = [];
+  let version = readVersion;
+  while (compareWireVersions(version, AGENT_WIRE_PROTOCOL_VERSION) < 0) {
+    const migration = findMigration(version);
+    if (migration === undefined) {
+      throw new Error(`Missing wire migration for version ${version}`);
+    }
+    migrations.push(migration);
+    version = migration.targetVersion;
+  }
+
+  return migrations;
+}
+
+export function migrateWireRecord(
+  record: WireMigrationRecord,
+  migrations: readonly WireMigration[],
+): WireMigrationRecord {
+  return migrations.reduce(
+    (current, migration) => migration.migrateRecord(current),
+    record,
+  );
+}
+
+function findMigration(sourceVersion: string): WireMigration | undefined {
+  for (const migration of MIGRATIONS) {
+    if (migration.sourceVersion === sourceVersion) return migration;
+  }
+}
+
+function compareWireVersions(a: string, b: string): number {
+  const partsA = a.split('.');
+  const partsB = b.split('.');
+  const maxLength = Math.max(partsA.length, partsB.length);
+
+  for (let i = 0; i < maxLength; i++) {
+    const diff = Number(partsA[i] ?? '0') - Number(partsB[i] ?? '0');
+    if (diff !== 0) return diff;
+  }
+
+  return 0;
+}

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

@@ -86,8 +86,6 @@ export type AgentRecordOf<K extends keyof AgentRecordEvents> = Extract<
   { readonly type: K }
 >;
 
-export const AGENT_WIRE_PROTOCOL_VERSION = '1.0';
-
 export interface AgentRecordPersistence {
   read(): AsyncIterable<AgentRecord>;
   append(input: AgentRecord): void;

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

@@ -1,4 +1,4 @@
-import { AGENT_WIRE_PROTOCOL_VERSION } from '../../agent/records/types';
+import { AGENT_WIRE_PROTOCOL_VERSION } from '../../agent/records';
 import type { SessionWireScan } from '#/session/export/wire-scan';
 import type { ExportSessionManifest, SessionSummary } from '#/rpc/core-api';
 

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

@@ -14,7 +14,10 @@ import {
 } from '../../../src/agent';
 import type { CompactionStrategy } from '../../../src/agent/compaction';
 import type { ApprovalResponse } from '../../../src/agent/permission';
-import { InMemoryAgentRecordPersistence } from '../../../src/agent/records';
+import {
+  AGENT_WIRE_PROTOCOL_VERSION,
+  InMemoryAgentRecordPersistence,
+} from '../../../src/agent/records';
 import type { KimiConfig } from '../../../src/config';
 import type { ExecutableToolResult } from '../../../src/loop';
 import type { Logger } from '../../../src/logging';
@@ -284,7 +287,9 @@ export class AgentTestContext {
       providerManager: this.agent.providerManager,
       generate: failOnResumeGenerate,
       compactionStrategy: this.options.compactionStrategy,
-      persistence: new InMemoryAgentRecordPersistence(this.recordHistory.map(cloneRecord)),
+      persistence: new InMemoryAgentRecordPersistence(
+        withMetadata(this.recordHistory.map(cloneRecord)),
+      ),
     });
 
     await resumed.agent.resume();
@@ -439,8 +444,12 @@ export class AgentTestContext {
   private wrapPersistence(persistence: AgentRecordPersistence): AgentRecordPersistence {
     return {
       read: () => this.readAndCapturePersistence(persistence),
-      append: (event) => persistence.append(event),
-      rewrite: (records) => persistence.rewrite(records),
+      append: (event) => {
+        persistence.append(event);
+      },
+      rewrite: (records) => {
+        persistence.rewrite(records);
+      },
       flush: () => persistence.flush(),
       close: () => persistence.close(),
     };
@@ -619,3 +628,15 @@ function buildSkillPrompt(content: string, args: string | undefined): string {
 function cloneRecord(event: AgentRecord): AgentRecord {
   return structuredClone(event);
 }
+
+function withMetadata(events: readonly AgentRecord[]): readonly AgentRecord[] {
+  if (events.length === 0 || events[0]?.type === 'metadata') return events;
+  return [
+    {
+      type: 'metadata',
+      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+      created_at: 1,
+    },
+    ...events,
+  ];
+}

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

@@ -0,0 +1,150 @@
+import { describe, expect, it } from 'vitest';
+
+import {
+  AGENT_WIRE_PROTOCOL_VERSION,
+  AgentRecords,
+  InMemoryAgentRecordPersistence,
+  type AgentRecord,
+} from '../../../src/agent/records';
+
+describe('AgentRecords persistence metadata', () => {
+  it('writes metadata before the first persisted record', async () => {
+    const persistence = new InMemoryAgentRecordPersistence();
+    const records = new AgentRecords(() => {}, persistence);
+
+    records.logRecord({
+      type: 'turn.prompt',
+      input: [{ type: 'text', text: 'hello' }],
+      origin: { kind: 'user' },
+    });
+    await records.flush();
+
+    expect(persistence.records).toHaveLength(2);
+    expect(persistence.records[0]).toMatchObject({
+      type: 'metadata',
+      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+    });
+    expect(persistence.records[1]?.type).toBe('turn.prompt');
+  });
+
+  it('does not write metadata when replaying an empty stream', async () => {
+    const persistence = new InMemoryAgentRecordPersistence();
+    const records = new AgentRecords(() => {}, persistence);
+
+    await records.replay();
+    records.logRecord({
+      type: 'turn.prompt',
+      input: [{ type: 'text', text: 'one' }],
+      origin: { kind: 'user' },
+    });
+    await records.flush();
+
+    expect(persistence.records.map((record) => record.type)).toEqual([
+      'metadata',
+      'turn.prompt',
+    ]);
+  });
+
+  it('rejects replaying a non-empty stream without metadata', async () => {
+    const persistence = new InMemoryAgentRecordPersistence([
+      {
+        type: 'turn.prompt',
+        input: [{ type: 'text', text: 'one' }],
+        origin: { kind: 'user' },
+      },
+    ]);
+    const records = new AgentRecords(() => {}, persistence);
+
+    await expect(records.replay()).rejects.toThrow(
+      'AgentRecords replay expected metadata as the first record',
+    );
+  });
+
+  it('does not duplicate metadata after replaying existing records', async () => {
+    const persistence = new InMemoryAgentRecordPersistence([
+      {
+        type: 'metadata',
+        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+        created_at: 1,
+      },
+      {
+        type: 'turn.prompt',
+        input: [{ type: 'text', text: 'one' }],
+        origin: { kind: 'user' },
+      },
+    ]);
+    const records = new AgentRecords(() => {}, persistence);
+
+    await records.replay();
+    records.logRecord({
+      type: 'turn.prompt',
+      input: [{ type: 'text', text: 'two' }],
+      origin: { kind: 'user' },
+    });
+    await records.flush();
+
+    expect(persistence.records.map((record) => record.type)).toEqual([
+      'metadata',
+      'turn.prompt',
+      'turn.prompt',
+    ]);
+    expect(persistence.records.filter((record) => record.type === 'metadata')).toHaveLength(1);
+  });
+
+  it('does not rewrite records that already use the current wire version', async () => {
+    const persistence = new RecordingInMemoryAgentRecordPersistence([
+      {
+        type: 'metadata',
+        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+        created_at: 1,
+      },
+      {
+        type: 'turn.prompt',
+        input: [{ type: 'text', text: 'one' }],
+        origin: { kind: 'user' },
+      },
+    ]);
+    const records = new AgentRecords(() => {}, persistence);
+
+    await records.replay();
+
+    expect(persistence.rewrites).toEqual([]);
+  });
+
+  it('rejects replaying records from a newer wire version', async () => {
+    const persistence = new InMemoryAgentRecordPersistence([
+      {
+        type: 'metadata',
+        protocol_version: '9.9',
+        created_at: 1,
+      },
+    ]);
+    const records = new AgentRecords(() => {}, persistence);
+
+    await expect(records.replay()).rejects.toThrow(
+      `Unsupported wire protocol version: 9.9 (current: ${AGENT_WIRE_PROTOCOL_VERSION})`,
+    );
+  });
+
+  it('rejects replaying records without a registered migration path', async () => {
+    const persistence = new InMemoryAgentRecordPersistence([
+      {
+        type: 'metadata',
+        protocol_version: '0.9',
+        created_at: 1,
+      },
+    ]);
+    const records = new AgentRecords(() => {}, persistence);
+
+    await expect(records.replay()).rejects.toThrow('Missing wire migration for version 0.9');
+  });
+});
+
+class RecordingInMemoryAgentRecordPersistence extends InMemoryAgentRecordPersistence {
+  readonly rewrites: AgentRecord[][] = [];
+
+  override rewrite(records: readonly AgentRecord[]): void {
+    this.rewrites.push([...records]);
+    super.rewrite(records);
+  }
+}

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

@@ -0,0 +1,35 @@
+import { describe, expect, it } from 'vitest';
+
+import {
+  migrateWireRecord,
+  type WireMigration,
+} from '../../../src/agent/records/migration';
+
+describe('wire record migrations', () => {
+  it('applies migrations in order', () => {
+    const migrations: WireMigration[] = [
+      {
+        sourceVersion: '0.8',
+        targetVersion: '0.9',
+        migrateRecord: (record) => ({
+          ...record,
+          first: true,
+        }),
+      },
+      {
+        sourceVersion: '0.9',
+        targetVersion: '1.0',
+        migrateRecord: (record) => ({
+          ...record,
+          second: record['first'] === true,
+        }),
+      },
+    ];
+
+    expect(migrateWireRecord({ type: 'metadata' }, migrations)).toEqual({
+      type: 'metadata',
+      first: true,
+      second: true,
+    });
+  });
+});

packages/agent-core/test/agent/records/persistence.test.ts

@@ -7,11 +7,10 @@ import { afterEach, describe, expect, it } from 'vitest';
 
 import {
   AGENT_WIRE_PROTOCOL_VERSION,
-  AgentRecords,
   FileSystemAgentRecordPersistence,
   InMemoryAgentRecordPersistence,
   type AgentRecord,
-} from '../../src/agent/records';
+} from '../../../src/agent/records';
 
 const cleanups: string[] = [];
 
@@ -249,82 +248,3 @@ describe('InMemoryAgentRecordPersistence', () => {
     expect(persistence.records).toEqual(records);
   });
 });
-
-describe('AgentRecords persistence metadata', () => {
-  it('writes metadata before the first persisted record', async () => {
-    const wirePath = await makeWirePath();
-    const persistence = new FileSystemAgentRecordPersistence(wirePath);
-    const records = new AgentRecords(() => {}, persistence);
-
-    records.logRecord({
-      type: 'turn.prompt',
-      input: [{ type: 'text', text: 'hello' }],
-      origin: { kind: 'user' },
-    });
-    await records.flush();
-
-    const lines = await readLines(wirePath);
-    expect(lines).toHaveLength(2);
-    const meta = JSON.parse(lines[0]!) as Record<string, unknown>;
-    expect(meta).toMatchObject({
-      type: 'metadata',
-      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
-    });
-    expect(typeof meta['created_at']).toBe('number');
-    expect(JSON.parse(lines[1]!)['type']).toBe('turn.prompt');
-  });
-
-  it('does not write metadata when replaying an empty stream', async () => {
-    const wirePath = await makeWirePath();
-    const persistence = new FileSystemAgentRecordPersistence(wirePath);
-    const records = new AgentRecords(() => {}, persistence);
-
-    await records.replay();
-    records.logRecord({
-      type: 'turn.prompt',
-      input: [{ type: 'text', text: 'one' }],
-      origin: { kind: 'user' },
-    });
-    await records.flush();
-
-    const lines = await readLines(wirePath);
-    expect(lines).toHaveLength(2);
-    expect(lines.map((line) => JSON.parse(line)['type'])).toEqual([
-      'metadata',
-      'turn.prompt',
-    ]);
-  });
-
-  it('does not duplicate metadata after replaying existing records', async () => {
-    const wirePath = await makeWirePath();
-    const persistence = new FileSystemAgentRecordPersistence(wirePath);
-    persistence.append({
-      type: 'metadata',
-      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
-      created_at: 1,
-    });
-    persistence.append({
-      type: 'turn.prompt',
-      input: [{ type: 'text', text: 'one' }],
-      origin: { kind: 'user' },
-    });
-    await persistence.flush();
-
-    const records = new AgentRecords(() => {}, persistence);
-    await records.replay();
-    records.logRecord({
-      type: 'turn.prompt',
-      input: [{ type: 'text', text: 'two' }],
-      origin: { kind: 'user' },
-    });
-    await records.flush();
-
-    const lines = await readLines(wirePath);
-    expect(lines.map((line) => JSON.parse(line)['type'])).toEqual([
-      'metadata',
-      'turn.prompt',
-      'turn.prompt',
-    ]);
-    expect(lines.filter((line) => JSON.parse(line)['type'] === 'metadata')).toHaveLength(1);
-  });
-});

packages/agent-core/test/agent/resume.test.ts

@@ -5,7 +5,10 @@ import { join } from 'node:path';
 import { describe, expect, it, vi } from 'vitest';
 
 import type { AgentRecord } from '../../src/agent';
-import { InMemoryAgentRecordPersistence } from '../../src/agent/records';
+import {
+  AGENT_WIRE_PROTOCOL_VERSION,
+  InMemoryAgentRecordPersistence,
+} from '../../src/agent/records';
 import { appendTaskOutput, writeTask } from '../../src/tools/background/persist';
 import { createFakeKaos } from '../tools/fixtures/fake-kaos';
 import { testAgent } from './harness/agent';
@@ -286,12 +289,28 @@ describe('Agent resume', () => {
 class RecordingAgentPersistence extends InMemoryAgentRecordPersistence {
   readonly appended: AgentRecord[] = [];
 
+  constructor(events: readonly AgentRecord[]) {
+    super(withMetadata(events));
+  }
+
   override append(input: AgentRecord): void {
     this.appended.push(input);
     super.append(input);
   }
 }
 
+function withMetadata(events: readonly AgentRecord[]): readonly AgentRecord[] {
+  if (events.length === 0 || events[0]?.type === 'metadata') return events;
+  return [
+    {
+      type: 'metadata',
+      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+      created_at: 1,
+    },
+    ...events,
+  ];
+}
+
 function resumeHistory(): AgentRecord[] {
   return [
     {

packages/agent-core/test/session/subagent-host.test.ts

@@ -7,6 +7,7 @@ import type { ToolCall } from '@moonshot-ai/kosong';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
 import type { Agent } from '../../src/agent';
+import { AGENT_WIRE_PROTOCOL_VERSION } from '../../src/agent/records';
 import type { ResolvedAgentProfile } from '../../src/profile';
 import type { SDKSessionRPC } from '../../src/rpc';
 import { Session } from '../../src/session';
@@ -1049,7 +1050,18 @@ function stat(kind: 'dir' | 'file') {
 
 async function writeWire(homedir: string, records: readonly Record<string, unknown>[]) {
   await mkdir(homedir, { recursive: true });
-  const text = records.map((record) => JSON.stringify(record)).join('\n');
+  const wireRecords =
+    records.length === 0
+      ? []
+      : [
+          {
+            type: 'metadata',
+            protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
+            created_at: 1,
+          },
+          ...records,
+        ];
+  const text = wireRecords.map((record) => JSON.stringify(record)).join('\n');
   await writeFile(join(homedir, 'wire.jsonl'), text.length === 0 ? '' : `${text}\n`, 'utf-8');
 }