NianToB/tests/unit/agent-system-documents.test.ts

import { mkdir, readFile, rm, writeFile } from 'node:fs/promises';
import { join } from 'node:path';
import { beforeEach, describe, expect, it, vi } from 'vitest';

const { listAgentsSnapshotMock, testHome, testRuntime } = vi.hoisted(() => {
  const suffix = Math.random().toString(36).slice(2);
  return {
    listAgentsSnapshotMock: vi.fn(),
    testHome: `/tmp/yinian-agent-system-documents-${suffix}`,
    testRuntime: `/tmp/yinian-agent-system-documents-runtime-${suffix}`,
  };
});

vi.mock('@electron/utils/agent-config', () => ({
  listAgentsSnapshot: listAgentsSnapshotMock,
}));

vi.mock('@electron/utils/paths', () => ({
  expandPath: (path: string) => path.replace(/^~(?=$|[\\/])/, testHome),
  getOpenClawDir: () => testRuntime,
}));

async function writeRuntimeTemplates(): Promise<void> {
  const templateDir = join(testRuntime, 'docs', 'reference', 'templates');
  await mkdir(templateDir, { recursive: true });
  await writeFile(join(templateDir, 'SOUL.md'), '# SOUL.md template\n', 'utf8');
  await writeFile(join(templateDir, 'IDENTITY.md'), '# IDENTITY.md template\n', 'utf8');
  await writeFile(join(templateDir, 'USER.md'), '# USER.md template\n', 'utf8');
  await writeFile(join(templateDir, 'AGENTS.md'), '# AGENTS.md template\n', 'utf8');
  await writeFile(join(templateDir, 'TOOLS.md'), '# TOOLS.md template\n', 'utf8');
  await writeFile(join(templateDir, 'HEARTBEAT.md'), '# HEARTBEAT.md template\n', 'utf8');
  await writeFile(join(templateDir, 'BOOT.md'), '# BOOT.md template\n', 'utf8');
}

function mockAgents(): void {
  listAgentsSnapshotMock.mockResolvedValue({
    agents: [
      {
        id: 'main',
        name: 'Main',
        isDefault: true,
        workspace: '~/.openclaw/workspace',
        agentDir: '~/.openclaw/agents/main/agent',
        modelDisplay: 'model',
        modelRef: null,
        overrideModelRef: null,
        inheritedModel: false,
        mainSessionKey: 'agent:main:main',
        channelTypes: [],
      },
      {
        id: 'coder',
        name: 'Coder',
        isDefault: false,
        workspace: '~/.openclaw/workspace-coder',
        agentDir: '~/.openclaw/agents/coder/agent',
        modelDisplay: 'model',
        modelRef: null,
        overrideModelRef: null,
        inheritedModel: false,
        mainSessionKey: 'agent:coder:main',
        channelTypes: [],
      },
    ],
    defaultAgentId: 'main',
    defaultModelRef: null,
    configuredChannelTypes: [],
    channelOwners: {},
    channelAccountOwners: {},
  });
}

describe('agent system documents', () => {
  beforeEach(async () => {
    vi.resetModules();
    vi.clearAllMocks();
    await rm(testHome, { recursive: true, force: true });
    await rm(testRuntime, { recursive: true, force: true });
    await writeRuntimeTemplates();
    mockAgents();
  });

  it('reads existing workspace docs and falls back to runtime templates for missing docs', async () => {
    const workspace = join(testHome, '.openclaw', 'workspace');
    await mkdir(workspace, { recursive: true });
    await writeFile(join(workspace, 'SOUL.md'), '# Custom soul\n', 'utf8');

    const { readAgentSystemDocuments } = await import('@electron/utils/agent-system-documents');
    const snapshot = await readAgentSystemDocuments();
    const soul = snapshot.documents.find((document) => document.kind === 'soul');
    const agent = snapshot.documents.find((document) => document.kind === 'agent');
    const heartbeat = snapshot.documents.find((document) => document.kind === 'heartbeat');

    expect(snapshot.selectedAgentId).toBe('main');
    expect(snapshot.paths.workspace).toBe(workspace);
    expect(snapshot.documents.map((document) => document.fileName)).toEqual([
      'SOUL.md',
      'IDENTITY.md',
      'USER.md',
      'AGENTS.md',
      'TOOLS.md',
      'HEARTBEAT.md',
      'BOOT.md',
    ]);
    expect(soul).toMatchObject({
      exists: true,
      source: 'workspace',
      content: '# Custom soul\n',
    });
    expect(agent).toMatchObject({
      exists: false,
      source: 'template',
      fileName: 'AGENTS.md',
      templateAvailable: true,
    });
    expect(agent?.content).toContain('# AGENTS.md');
    expect(heartbeat).toMatchObject({
      exists: false,
      source: 'template',
      fileName: 'HEARTBEAT.md',
      templateAvailable: true,
    });
  });

  it('saves a selected agent document into that agent workspace', async () => {
    const workspace = join(testHome, '.openclaw', 'workspace-coder');

    const { saveAgentSystemDocument } = await import('@electron/utils/agent-system-documents');
    const snapshot = await saveAgentSystemDocument('coder', 'tool', '# Tool notes\n');

    await expect(readFile(join(workspace, 'TOOLS.md'), 'utf8')).resolves.toBe('# Tool notes\n');
    expect(snapshot.selectedAgentId).toBe('coder');
    expect(snapshot.documents.find((document) => document.kind === 'tool')).toMatchObject({
      exists: true,
      source: 'workspace',
      content: '# Tool notes\n',
    });
  });

  it('restores a document from the OpenClaw runtime template', async () => {
    const workspace = join(testHome, '.openclaw', 'workspace');
    await mkdir(workspace, { recursive: true });
    await writeFile(join(workspace, 'AGENTS.md'), '# Old agent doc\n', 'utf8');

    const { resetAgentSystemDocument } = await import('@electron/utils/agent-system-documents');
    const snapshot = await resetAgentSystemDocument('main', 'agent');
    const restored = await readFile(join(workspace, 'AGENTS.md'), 'utf8');

    expect(restored).toContain('# AGENTS.md');
    expect(restored).not.toContain('# Old agent doc');
    expect(snapshot.documents.find((document) => document.kind === 'agent')).toMatchObject({
      exists: true,
      source: 'workspace',
    });
  });
});