kimi-code/packages/agent-core-v2/test/wireRecord/persistence.test.ts

import { randomBytes } from 'node:crypto';
import { mkdir, readFile, readdir, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'pathe';

import { afterEach, describe, expect, it } from 'vitest';

import { SyncDescriptor } from '#/_base/di/descriptors';
import { DisposableStore } from '#/_base/di/lifecycle';
import { TestInstantiationService } from '#/_base/di/test';
import { IAgentBlobService, AgentBlobServiceImpl } from '#/agent/blob';
import { IBlobStore } from '#/persistence/interface/blobStore';
import { BlobStoreService } from '#/persistence/backends/node-fs/blobStoreService';
import { IBootstrapService } from '#/app/bootstrap';
import { IHostFileSystem } from '#/os/interface/hostFileSystem';
import { HostFileSystem } from '#/os/backends/node-local/hostFsService';
import { AgentContextMemoryService } from '#/agent/contextMemory/contextMemoryService';
import { IAgentContextMemoryService, type ContextMessage } from '#/agent/contextMemory';
import { IAgentContextOpsService, AgentContextOpsService } from '#/agent/contextOps';
import { IAgentRecordService, AgentRecordService } from '#/agent/record';
import { IAgentScopeContext, makeAgentScopeContext } from '#/agent/scopeContext';
import {
  AppendLogStore,
  AGENT_WIRE_PROTOCOL_VERSION,
  IFileSystemStorageService,
  IAppendLogStore,
  IAgentWireRecordService,
  type PersistedWireRecord,
  AgentWireRecordService,
} from '#/index';
import { FileStorageService } from '#/persistence/backends/node-fs/fileStorageService';
import { InMemoryStorageService } from '#/persistence/backends/memory/inMemoryStorageService';
import { stubBootstrap } from '../bootstrap/stubs';

const cleanups: string[] = [];
const disposables: DisposableStore[] = [];
const SCOPE = 'wire-test';
const KEY = 'wire.jsonl';

afterEach(async () => {
  for (const store of disposables.splice(0)) {
    store.dispose();
  }
  for (const dir of cleanups.splice(0)) {
    await rm(dir, { recursive: true, force: true }).catch(() => {});
  }
});

async function makeDir(prefix: string): Promise<string> {
  const dir = join(tmpdir(), `${prefix}-${randomBytes(6).toString('hex')}`);
  await mkdir(dir, { recursive: true });
  cleanups.push(dir);
  return dir;
}

async function readLines(path: string): Promise<string[]> {
  const raw = await readFile(path, 'utf8');
  return raw.split('\n').filter((line) => line.length > 0);
}

function createAppendLogHarness(storage: IFileSystemStorageService): IAppendLogStore {
  const disposable = new DisposableStore();
  disposables.push(disposable);

  const ix = disposable.add(new TestInstantiationService());
  ix.stub(IFileSystemStorageService, storage);
  ix.set(IAppendLogStore, new SyncDescriptor(AppendLogStore));
  return ix.get(IAppendLogStore);
}

async function createFileAppendLogHarness(): Promise<{
  readonly dir: string;
  readonly log: IAppendLogStore;
}> {
  const dir = await makeDir('wire-jsonl-test');
  return {
    dir,
    log: createAppendLogHarness(new FileStorageService(dir)),
  };
}

async function collect<R>(log: IAppendLogStore, scope = SCOPE, key = KEY): Promise<R[]> {
  const records: R[] = [];
  for await (const record of log.read<R>(scope, key)) {
    records.push(record);
  }
  return records;
}

async function createWireHarness(): Promise<{
  readonly dir: string;
  readonly storage: IFileSystemStorageService;
  readonly wire: IAgentWireRecordService;
  readonly context: IAgentContextMemoryService;
  readonly contextOps: IAgentContextOpsService;
}> {
  const dir = await makeDir('wire-service-test');
  const disposable = new DisposableStore();
  disposables.push(disposable);

  const storage = new FileStorageService(dir);
  const ix = disposable.add(new TestInstantiationService());
  ix.stub(IFileSystemStorageService, storage);
  ix.set(IBlobStore, new SyncDescriptor(BlobStoreService));
  ix.stub(IBootstrapService, stubBootstrap(dir));
  ix.stub(IHostFileSystem, new HostFileSystem());
  ix.set(IAppendLogStore, new SyncDescriptor(AppendLogStore));
  ix.stub(IAgentScopeContext, makeAgentScopeContext({ agentId: 'test', agentScope: '' }));
  ix.set(IAgentBlobService, new SyncDescriptor(AgentBlobServiceImpl));

  ix.set(IAgentWireRecordService, new SyncDescriptor(AgentWireRecordService, [{}]));
  ix.set(IAgentRecordService, new SyncDescriptor(AgentRecordService));
  ix.set(IAgentContextMemoryService, new SyncDescriptor(AgentContextMemoryService));
  ix.set(IAgentContextOpsService, new SyncDescriptor(AgentContextOpsService));
  // Force-instantiate the context ops domain so its operation facets are
  // registered before any wire records are appended.
  ix.get(IAgentContextOpsService);

  return {
    dir,
    storage,
    wire: ix.get(IAgentWireRecordService),
    context: ix.get(IAgentContextMemoryService),
    contextOps: ix.get(IAgentContextOpsService),
  };
}

async function readPersistedWireRecords(
  storage: IFileSystemStorageService,
): Promise<readonly PersistedWireRecord[]> {
  const keys = await storage.list('wire');
  if (keys.length === 0) return [];
  expect(keys).toHaveLength(1);
  const log = createAppendLogHarness(storage);
  return collect<PersistedWireRecord>(log, 'wire', keys[0]);
}

describe('AppendLogStore file persistence', () => {
  it('writes only the appended record', async () => {
    const { dir, log } = await createFileAppendLogHarness();

    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'hello' }],
      origin: { kind: 'user' },
    });
    await log.close();

    const lines = await readLines(join(dir, SCOPE, KEY));
    expect(lines).toHaveLength(1);
    expect(JSON.parse(lines[0]!)['type']).toBe('turn.prompt');
  });

  it('appends to an existing file without injecting records', async () => {
    const dir = await makeDir('wire-jsonl-test');
    const first = createAppendLogHarness(new FileStorageService(dir));
    first.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'one' }],
      origin: { kind: 'user' },
    });
    await first.close();

    const second = createAppendLogHarness(new FileStorageService(dir));
    second.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'two' }],
      origin: { kind: 'user' },
    });
    await second.close();

    const lines = await readLines(join(dir, SCOPE, KEY));
    expect(lines).toHaveLength(2);
    expect(lines.map((line) => JSON.parse(line)['type'])).toEqual([
      'turn.prompt',
      'turn.prompt',
    ]);
  });

  it('returns appended metadata records from read() output', async () => {
    const { log } = await createFileAppendLogHarness();
    log.append(SCOPE, KEY, {
      type: 'metadata',
      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
      created_at: 1,
    });
    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'hi' }],
      origin: { kind: 'user' },
    });
    await log.close();

    const records = await collect<PersistedWireRecord>(log);
    expect(records).toHaveLength(2);
    expect(records[0]).toMatchObject({
      type: 'metadata',
      protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
    });
    expect(records[1]!.type).toBe('turn.prompt');
  });

  it('rewrites records from the beginning and then appends after them', async () => {
    const { dir, log } = await createFileAppendLogHarness();
    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'old' }],
      origin: { kind: 'user' },
    });
    await log.rewrite(SCOPE, KEY, [
      {
        type: 'metadata',
        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
        created_at: 1,
      },
      {
        type: 'turn.prompt',
        input: [{ type: 'text', text: 'new' }],
        origin: { kind: 'user' },
      },
    ]);
    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'later' }],
      origin: { kind: 'user' },
    });
    await log.flush();

    const lines = await readLines(join(dir, SCOPE, KEY));
    expect(lines.map((line) => JSON.parse(line)['type'])).toEqual([
      'metadata',
      'turn.prompt',
      'turn.prompt',
    ]);
    expect(JSON.parse(lines[1]!)['input'][0]['text']).toBe('new');
    expect(JSON.parse(lines[2]!)['input'][0]['text']).toBe('later');
  });

  it('rewrites already flushed records from the beginning', async () => {
    const { dir, log } = await createFileAppendLogHarness();
    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'old' }],
      origin: { kind: 'user' },
    });
    await log.flush();

    await log.rewrite(SCOPE, KEY, [
      {
        type: 'metadata',
        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
        created_at: 1,
      },
      {
        type: 'turn.prompt',
        input: [{ type: 'text', text: 'new' }],
        origin: { kind: 'user' },
      },
    ]);
    await log.flush();

    const lines = await readLines(join(dir, SCOPE, KEY));
    expect(lines.map((line) => JSON.parse(line)['type'])).toEqual([
      'metadata',
      'turn.prompt',
    ]);
    expect(JSON.parse(lines[1]!)['input'][0]['text']).toBe('new');
  });

  it('flushes pending records on close', async () => {
    const { dir, log } = await createFileAppendLogHarness();

    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'late' }],
      origin: { kind: 'user' },
    });
    await log.close();

    const lines = await readLines(join(dir, SCOPE, KEY));
    expect(lines).toHaveLength(1);
    expect(JSON.parse(lines[0]!)['type']).toBe('turn.prompt');
  });

  it('propagates write failures from flush', async () => {
    const dir = await makeDir('wire-jsonl-test');
    await mkdir(join(dir, SCOPE, KEY), { recursive: true });
    const log = createAppendLogHarness(new FileStorageService(dir));

    log.append(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'first' }],
      origin: { kind: 'user' },
    });

    await expect(log.flush()).rejects.toBeInstanceOf(Error);
  });
});

describe('AgentWireRecordService persistence', () => {
  it('persists raw system reminder args while applying the wrapped context message', async () => {
    const { storage, wire, context, contextOps } = await createWireHarness();

    contextOps.appendSystemReminder(' Remember this. ', {
      kind: 'injection',
      variant: 'host',
    });
    await wire.flush();

    expect(context.get()[0]?.content).toEqual([
      { type: 'text', text: '<system-reminder>\nRemember this.\n</system-reminder>' },
    ]);
    expect(context.get()[0]?.origin).toEqual({
      kind: 'injection',
      variant: 'host',
    });

    const records = await readPersistedWireRecords(storage);
    expect(records.map((record) => record.type)).toEqual([
      'metadata',
      'context.append_system_reminder',
    ]);
    const record = records[1] as {
      type: 'context.append_system_reminder';
      args: readonly [string, ContextMessage['origin']];
    };
    expect(record.args).toEqual(['Remember this.', { kind: 'injection', variant: 'host' }]);
  });

  it('restores string system reminder args with origin as wrapped context messages', async () => {
    const { wire, context } = await createWireHarness();

    await wire.restore([
      {
        type: 'metadata',
        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
        created_at: 1,
      },
      {
        type: 'context.append_system_reminder',
        args: ['Remember this.', { kind: 'injection', variant: 'host' }],
      },
    ]);

    const message = context.get()[0];
    expect(message?.content).toEqual([
      { type: 'text', text: '<system-reminder>\nRemember this.\n</system-reminder>' },
    ]);
    expect(message?.origin).toEqual({
      kind: 'injection',
      variant: 'host',
    });
  });

  it('offloads large content part data URIs to blobsDir during append', async () => {
    const { dir, storage, wire } = await createWireHarness();
    const payload = 'X'.repeat(5000);
    const dataUri = `data:image/png;base64,${payload}`;
    const message: ContextMessage = {
      role: 'user',
      content: [{ type: 'image_url', imageUrl: { url: dataUri } }],
      toolCalls: [],
      origin: { kind: 'user' },
    };

    wire.append({
      type: 'context.append',
      args: [message],
    });
    await wire.flush();

    const records = await readPersistedWireRecords(storage);
    expect(records.map((record) => record.type)).toEqual(['metadata', 'context.append']);
    const record = records[1] as { type: 'context.append'; args: readonly ContextMessage[] };
    const appendedMessages = record.args;
    const url = (
      appendedMessages[0]!.content[0] as unknown as { imageUrl: { url: string } }
    ).imageUrl.url;
    expect(url.startsWith('blobref:')).toBe(true);

    const blobFiles = await readdir(join(dir, 'blobs'));
    expect(blobFiles).toHaveLength(1);
    expect((await readFile(join(dir, 'blobs', blobFiles[0]!))).toString('base64')).toBe(payload);
  });
});

describe('wire record append-log persistence', () => {
  it('can be backed by the same IAppendLogStore contract as file persistence', async () => {
    const storage = new InMemoryStorageService();
    const log = createAppendLogHarness(storage);

    log.append<PersistedWireRecord>(SCOPE, KEY, {
      type: 'turn.prompt',
      input: [{ type: 'text', text: 'one' }],
      origin: { kind: 'user' },
    });
    await log.rewrite<PersistedWireRecord>(SCOPE, KEY, [
      {
        type: 'metadata',
        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
        created_at: 1,
      },
    ]);

    expect(await collect<PersistedWireRecord>(log)).toEqual([
      {
        type: 'metadata',
        protocol_version: AGENT_WIRE_PROTOCOL_VERSION,
        created_at: 1,
      },
    ]);
  });
});