zclaw_openfang/tests/desktop/agent-memory.test.ts

/**
 * Tests for Agent Memory System (Phase 1)
 *
 * Covers: MemoryManager, AgentIdentityManager, MemoryExtractor
 */

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import {
  MemoryManager,
  resetMemoryManager,
  type MemoryEntry,
} from '../../desktop/src/lib/agent-memory';
import {
  AgentIdentityManager,
  resetAgentIdentityManager,
} from '../../desktop/src/lib/agent-identity';
import {
  MemoryExtractor,
  resetMemoryExtractor,
} from '../../desktop/src/lib/memory-extractor';

// === Mock localStorage ===

const localStorageMock = (() => {
  let store: Record<string, string> = {};
  return {
    getItem: (key: string) => store[key] ?? null,
    setItem: (key: string, value: string) => { store[key] = value; },
    removeItem: (key: string) => { delete store[key]; },
    clear: () => { store = {}; },
  };
})();

vi.stubGlobal('localStorage', localStorageMock);

// =============================================
// MemoryManager Tests
// =============================================

describe('MemoryManager', () => {
  let mgr: MemoryManager;

  beforeEach(() => {
    localStorageMock.clear();
    resetMemoryManager();
    mgr = new MemoryManager();
  });

  describe('save', () => {
    it('saves a memory entry and assigns id/timestamps', async () => {
      const entry = await mgr.save({
        agentId: 'agent-1',
        content: '用户喜欢简洁回答',
        type: 'preference',
        importance: 7,
        source: 'auto',
        tags: ['communication'],
      });

      expect(entry.id).toMatch(/^mem_/);
      expect(entry.content).toBe('用户喜欢简洁回答');
      expect(entry.type).toBe('preference');
      expect(entry.importance).toBe(7);
      expect(entry.createdAt).toBeTruthy();
      expect(entry.accessCount).toBe(0);
    });

    it('deduplicates similar content for same agent+type', async () => {
      await mgr.save({
        agentId: 'agent-1',
        content: '用户的公司叫 ACME',
        type: 'fact',
        importance: 6,
        source: 'auto',
        tags: ['company'],
      });

      // Save very similar content
      const second = await mgr.save({
        agentId: 'agent-1',
        content: '用户的公司叫 ACME Corp',
        type: 'fact',
        importance: 8,
        source: 'auto',
        tags: ['company', 'name'],
      });

      // Should update existing entry, not create a new one
      const all = await mgr.getAll('agent-1');
      expect(all.length).toBe(1);
      expect(all[0].importance).toBe(8); // Takes higher importance
      expect(all[0].tags).toContain('name'); // Merged tags
    });

    it('does not deduplicate across different agents', async () => {
      await mgr.save({
        agentId: 'agent-1',
        content: '用户喜欢TypeScript',
        type: 'preference',
        importance: 6,
        source: 'auto',
        tags: [],
      });

      await mgr.save({
        agentId: 'agent-2',
        content: '用户喜欢TypeScript',
        type: 'preference',
        importance: 6,
        source: 'auto',
        tags: [],
      });

      const all1 = await mgr.getAll('agent-1');
      const all2 = await mgr.getAll('agent-2');
      expect(all1.length).toBe(1);
      expect(all2.length).toBe(1);
    });
  });

  describe('search', () => {
    beforeEach(async () => {
      await mgr.save({ agentId: 'a1', content: '用户的项目使用 React 和 TypeScript', type: 'fact', importance: 7, source: 'auto', tags: ['tech'] });
      await mgr.save({ agentId: 'a1', content: '用户喜欢简洁的代码风格', type: 'preference', importance: 6, source: 'auto', tags: ['coding'] });
      await mgr.save({ agentId: 'a1', content: '飞书API需要先验证token再调用', type: 'lesson', importance: 8, source: 'auto', tags: ['feishu', 'api'] });
      await mgr.save({ agentId: 'a2', content: '另一个 agent 的记忆', type: 'fact', importance: 5, source: 'auto', tags: [] });
    });

    it('finds relevant memories by keyword', async () => {
      const results = await mgr.search('TypeScript React', { agentId: 'a1' });
      expect(results.length).toBeGreaterThan(0);
      expect(results[0].content).toContain('TypeScript');
    });

    it('filters by agentId', async () => {
      const results = await mgr.search('记忆', { agentId: 'a1' });
      expect(results.every(r => r.agentId === 'a1')).toBe(true);
    });

    it('filters by type', async () => {
      const results = await mgr.search('代码', { agentId: 'a1', type: 'preference' });
      expect(results.every(r => r.type === 'preference')).toBe(true);
    });

    it('filters by minImportance', async () => {
      const results = await mgr.search('飞书 API token', { agentId: 'a1', minImportance: 7 });
      expect(results.every(r => r.importance >= 7)).toBe(true);
    });

    it('filters by tags', async () => {
      const results = await mgr.search('API', { agentId: 'a1', tags: ['feishu'] });
      expect(results.length).toBeGreaterThan(0);
      expect(results[0].tags).toContain('feishu');
    });

    it('returns empty for no match', async () => {
      const results = await mgr.search('量子物理', { agentId: 'a1' });
      expect(results).toHaveLength(0);
    });

    it('updates access metadata on search hit', async () => {
      const before = await mgr.getAll('a1');
      const hitEntry = before.find(e => e.content.includes('TypeScript'))!;
      const beforeCount = hitEntry.accessCount;

      await mgr.search('TypeScript', { agentId: 'a1' });

      const after = await mgr.get(hitEntry.id);
      expect(after!.accessCount).toBe(beforeCount + 1);
    });
  });

  describe('getAll', () => {
    it('returns entries for the requested agent', async () => {
      await mgr.save({ agentId: 'a1', content: 'memory A', type: 'fact', importance: 5, source: 'auto', tags: [] });
      await mgr.save({ agentId: 'a1', content: 'memory B', type: 'fact', importance: 5, source: 'auto', tags: [] });
      await mgr.save({ agentId: 'a2', content: 'other agent', type: 'fact', importance: 5, source: 'auto', tags: [] });

      const all = await mgr.getAll('a1');
      expect(all).toHaveLength(2);
      expect(all.every(e => e.agentId === 'a1')).toBe(true);
    });

    it('respects limit', async () => {
      for (let i = 0; i < 5; i++) {
        await mgr.save({ agentId: 'a1', content: `memory ${i}`, type: 'fact', importance: 5, source: 'auto', tags: [] });
      }

      const limited = await mgr.getAll('a1', { limit: 3 });
      expect(limited).toHaveLength(3);
    });
  });

  describe('forget', () => {
    it('removes a specific memory', async () => {
      const entry = await mgr.save({ agentId: 'a1', content: 'to forget', type: 'fact', importance: 5, source: 'auto', tags: [] });
      await mgr.forget(entry.id);
      expect(await mgr.get(entry.id)).toBeNull();
    });
  });

  describe('prune', () => {
    it('removes old low-importance entries', async () => {
      // Create an old low-importance entry
      const entry = await mgr.save({ agentId: 'a1', content: 'old unimportant', type: 'context', importance: 2, source: 'auto', tags: [] });
      // Manually set lastAccessedAt to 60 days ago
      const raw = JSON.parse(localStorageMock.getItem('zclaw-agent-memories')!);
      const idx = raw.findIndex((e: MemoryEntry) => e.id === entry.id);
      raw[idx].lastAccessedAt = new Date(Date.now() - 60 * 24 * 60 * 60 * 1000).toISOString();
      localStorageMock.setItem('zclaw-agent-memories', JSON.stringify(raw));

      // Reload and prune
      resetMemoryManager();
      const mgr2 = new MemoryManager();
      const pruned = await mgr2.prune({ maxAgeDays: 30, minImportance: 3, agentId: 'a1' });
      expect(pruned).toBe(1);
    });

    it('keeps high-importance entries even when old', async () => {
      const entry = await mgr.save({ agentId: 'a1', content: 'important old', type: 'fact', importance: 9, source: 'auto', tags: [] });
      const raw = JSON.parse(localStorageMock.getItem('zclaw-agent-memories')!);
      const idx = raw.findIndex((e: MemoryEntry) => e.id === entry.id);
      raw[idx].lastAccessedAt = new Date(Date.now() - 60 * 24 * 60 * 60 * 1000).toISOString();
      localStorageMock.setItem('zclaw-agent-memories', JSON.stringify(raw));

      resetMemoryManager();
      const mgr2 = new MemoryManager();
      const pruned = await mgr2.prune({ maxAgeDays: 30, minImportance: 3, agentId: 'a1' });
      expect(pruned).toBe(0);
    });
  });

  describe('exportToMarkdown', () => {
    it('produces formatted markdown with categories', async () => {
      await mgr.save({ agentId: 'a1', content: '公司叫 ACME', type: 'fact', importance: 7, source: 'auto', tags: ['company'] });
      await mgr.save({ agentId: 'a1', content: '喜欢简洁回答', type: 'preference', importance: 6, source: 'auto', tags: [] });

      const md = await mgr.exportToMarkdown('a1');
      expect(md).toContain('Agent Memory Export');
      expect(md).toContain('事实');
      expect(md).toContain('偏好');
      expect(md).toContain('ACME');
      expect(md).toContain('简洁');
    });

    it('returns empty notice for agent with no memories', async () => {
      const md = await mgr.exportToMarkdown('nonexistent');
      expect(md).toContain('No memories recorded');
    });
  });

  describe('stats', () => {
    it('returns correct statistics', async () => {
      await mgr.save({ agentId: 'a1', content: 'fact1', type: 'fact', importance: 5, source: 'auto', tags: [] });
      await mgr.save({ agentId: 'a1', content: 'pref1', type: 'preference', importance: 5, source: 'auto', tags: [] });
      await mgr.save({ agentId: 'a1', content: 'lesson1', type: 'lesson', importance: 5, source: 'auto', tags: [] });

      const stats = await mgr.stats('a1');
      expect(stats.totalEntries).toBe(3);
      expect(stats.byType['fact']).toBe(1);
      expect(stats.byType['preference']).toBe(1);
      expect(stats.byType['lesson']).toBe(1);
    });
  });

  describe('persistence', () => {
    it('survives manager recreation (simulates app restart)', async () => {
      await mgr.save({ agentId: 'a1', content: 'persistent memory', type: 'fact', importance: 8, source: 'auto', tags: [] });

      // Create a new manager (simulates restart)
      const mgr2 = new MemoryManager();
      const all = await mgr2.getAll('a1');
      expect(all.length).toBe(1);
      expect(all[0].content).toBe('persistent memory');
    });
  });
});

// =============================================
// AgentIdentityManager Tests
// =============================================

describe('AgentIdentityManager', () => {
  let idMgr: AgentIdentityManager;

  beforeEach(() => {
    localStorageMock.clear();
    resetAgentIdentityManager();
    idMgr = new AgentIdentityManager();
  });

  describe('getIdentity', () => {
    it('returns defaults for new agent', () => {
      const identity = idMgr.getIdentity('new-agent');
      expect(identity.soul).toContain('ZCLAW');
      expect(identity.instructions).toContain('指令');
      expect(identity.userProfile).toContain('用户画像');
    });

    it('returns same identity on second call', () => {
      const first = idMgr.getIdentity('test');
      const second = idMgr.getIdentity('test');
      expect(first.soul).toBe(second.soul);
    });
  });

  describe('buildSystemPrompt', () => {
    it('combines soul + instructions + memory context', () => {
      const prompt = idMgr.buildSystemPrompt('test', '## 相关记忆\n- 用户喜欢中文');
      expect(prompt).toContain('ZCLAW');
      expect(prompt).toContain('指令');
      expect(prompt).toContain('相关记忆');
    });

    it('excludes default user profile', () => {
      const prompt = idMgr.buildSystemPrompt('test');
      expect(prompt).not.toContain('尚未收集到');
    });

    it('includes updated user profile', () => {
      idMgr.updateUserProfile('test', '## 用户偏好\n- 喜欢TypeScript');
      const prompt = idMgr.buildSystemPrompt('test');
      expect(prompt).toContain('TypeScript');
    });
  });

  describe('updateUserProfile', () => {
    it('updates profile and creates snapshot', () => {
      idMgr.updateUserProfile('test', '新的用户画像');
      const identity = idMgr.getIdentity('test');
      expect(identity.userProfile).toBe('新的用户画像');

      const snapshots = idMgr.getSnapshots('test');
      expect(snapshots.length).toBeGreaterThan(0);
    });
  });

  describe('appendToUserProfile', () => {
    it('appends content to existing profile', () => {
      idMgr.appendToUserProfile('test', '- 新发现的偏好');
      const identity = idMgr.getIdentity('test');
      expect(identity.userProfile).toContain('新发现的偏好');
    });
  });

  describe('change proposals', () => {
    it('creates a pending proposal', () => {
      const proposal = idMgr.proposeChange('test', 'soul', '新的人格定义', '根据用户反馈调整');
      expect(proposal.status).toBe('pending');
      expect(proposal.suggestedContent).toBe('新的人格定义');

      const pending = idMgr.getPendingProposals('test');
      expect(pending).toHaveLength(1);
    });

    it('approves a proposal and updates identity', () => {
      const proposal = idMgr.proposeChange('test', 'soul', '新的 SOUL', '改进');
      idMgr.approveProposal(proposal.id);

      const identity = idMgr.getIdentity('test');
      expect(identity.soul).toBe('新的 SOUL');

      const pending = idMgr.getPendingProposals('test');
      expect(pending).toHaveLength(0);
    });

    it('rejects a proposal without changing identity', () => {
      const oldIdentity = idMgr.getIdentity('test');
      const proposal = idMgr.proposeChange('test', 'soul', '被拒绝的内容', '原因');
      idMgr.rejectProposal(proposal.id);

      const identity = idMgr.getIdentity('test');
      expect(identity.soul).toBe(oldIdentity.soul);
    });
  });

  describe('snapshots and rollback', () => {
    it('can rollback to a previous snapshot', () => {
      idMgr.updateUserProfile('test', 'version 1');
      idMgr.updateUserProfile('test', 'version 2');

      const snapshots = idMgr.getSnapshots('test');
      expect(snapshots.length).toBeGreaterThanOrEqual(2);

      // snapshots[0] = most recent (taken before 'version 2' was set, contains 'version 1')
      const snapshotWithV1 = snapshots[0];
      idMgr.restoreSnapshot('test', snapshotWithV1.id);

      const identity = idMgr.getIdentity('test');
      expect(identity.userProfile).toBe('version 1');
    });
  });

  describe('direct edit', () => {
    it('updates file and creates snapshot', () => {
      idMgr.updateFile('test', 'instructions', '# 自定义指令\n\n新指令内容');
      const identity = idMgr.getIdentity('test');
      expect(identity.instructions).toContain('自定义指令');
    });
  });

  describe('persistence', () => {
    it('survives recreation', () => {
      idMgr.updateUserProfile('test', '持久化测试');

      resetAgentIdentityManager();
      const mgr2 = new AgentIdentityManager();
      const identity = mgr2.getIdentity('test');
      expect(identity.userProfile).toBe('持久化测试');
    });
  });
});

// =============================================
// MemoryExtractor Tests
// =============================================

describe('MemoryExtractor', () => {
  let extractor: MemoryExtractor;

  beforeEach(() => {
    localStorageMock.clear();
    resetMemoryManager();
    resetAgentIdentityManager();
    resetMemoryExtractor();
    extractor = new MemoryExtractor();
  });

  describe('extractFromConversation', () => {
    it('skips extraction for short conversations', async () => {
      const result = await extractor.extractFromConversation(
        [
          { role: 'user', content: '你好' },
          { role: 'assistant', content: '你好！' },
        ],
        'agent-1'
      );
      expect(result.saved).toBe(0);
    });

    it('extracts facts from user messages', async () => {
      const result = await extractor.extractFromConversation(
        [
          { role: 'user', content: '我的公司叫字节跳动' },
          { role: 'assistant', content: '好的，了解了。' },
          { role: 'user', content: '我在做一个AI助手项目' },
          { role: 'assistant', content: '听起来很有趣！' },
          { role: 'user', content: '帮我看看这个bug' },
          { role: 'assistant', content: '好的，让我看看。' },
        ],
        'agent-1'
      );

      expect(result.items.length).toBeGreaterThan(0);
      const facts = result.items.filter(i => i.type === 'fact');
      expect(facts.length).toBeGreaterThan(0);
    });

    it('extracts preferences from user messages', async () => {
      const result = await extractor.extractFromConversation(
        [
          { role: 'user', content: '请用中文回复我' },
          { role: 'assistant', content: '好的，我会用中文。' },
          { role: 'user', content: '我喜欢简洁的代码风格' },
          { role: 'assistant', content: '明白了。' },
          { role: 'user', content: '以后都用TypeScript写' },
          { role: 'assistant', content: '好的。' },
        ],
        'agent-1'
      );

      const prefs = result.items.filter(i => i.type === 'preference');
      expect(prefs.length).toBeGreaterThan(0);
    });

    it('saves extracted memories to MemoryManager', async () => {
      await extractor.extractFromConversation(
        [
          { role: 'user', content: '我的项目叫ZCLAW' },
          { role: 'assistant', content: '好的。' },
          { role: 'user', content: '我喜欢简洁回答' },
          { role: 'assistant', content: '了解。' },
          { role: 'user', content: '继续' },
          { role: 'assistant', content: '好的。' },
        ],
        'agent-1'
      );

      const memMgr = new MemoryManager();
      const all = await memMgr.getAll('agent-1');
      expect(all.length).toBeGreaterThan(0);
    });

    it('respects extraction cooldown', async () => {
      const msgs = [
        { role: 'user', content: '我的公司叫字节跳动' },
        { role: 'assistant', content: '好的。' },
        { role: 'user', content: '帮我写代码' },
        { role: 'assistant', content: '好的。' },
        { role: 'user', content: '继续' },
        { role: 'assistant', content: '好的。' },
      ];

      const r1 = await extractor.extractFromConversation(msgs, 'agent-1');
      const r2 = await extractor.extractFromConversation(msgs, 'agent-1');

      // Second call should be blocked by cooldown
      expect(r2.saved).toBe(0);
      // First call should have extracted something
      expect(r1.items.length).toBeGreaterThanOrEqual(0); // May or may not match patterns
    });
  });

  describe('buildExtractionPrompt', () => {
    it('produces a valid LLM prompt', () => {
      const prompt = extractor.buildExtractionPrompt([
        { role: 'user', content: '帮我配置飞书' },
        { role: 'assistant', content: '好的，需要app_id和app_secret' },
      ]);

      expect(prompt).toContain('请从以下对话中提取');
      expect(prompt).toContain('用户');
      expect(prompt).toContain('助手');
      expect(prompt).toContain('飞书');
    });
  });

  describe('parseExtractionResponse', () => {
    it('parses valid JSON response', () => {
      const response = `[
        {"content": "用户喜欢中文", "type": "preference", "importance": 7, "tags": ["lang"]},
        {"content": "公司叫ACME", "type": "fact", "importance": 6, "tags": ["company"]}
      ]`;

      const items = extractor.parseExtractionResponse(response);
      expect(items).toHaveLength(2);
      expect(items[0].type).toBe('preference');
      expect(items[1].type).toBe('fact');
    });

    it('handles wrapped JSON in markdown', () => {
      const response = "```json\n[{\"content\": \"test\", \"type\": \"fact\", \"importance\": 5, \"tags\": []}]\n```";
      const items = extractor.parseExtractionResponse(response);
      expect(items).toHaveLength(1);
    });

    it('returns empty for invalid response', () => {
      const items = extractor.parseExtractionResponse('no json here');
      expect(items).toHaveLength(0);
    });

    it('clamps importance to 1-10', () => {
      const response = '[{"content": "test", "type": "fact", "importance": 15, "tags": []}]';
      const items = extractor.parseExtractionResponse(response);
      expect(items[0].importance).toBe(10);
    });
  });
});