/** * Tests for Context Compactor (Phase 2) * * Covers: token estimation, threshold checking, memory flush, compaction */ import { describe, it, expect, beforeEach, vi } from 'vitest'; import { ContextCompactor, resetContextCompactor, estimateTokens, estimateMessagesTokens, DEFAULT_COMPACTION_CONFIG, type CompactableMessage, } from '../../desktop/src/lib/context-compactor'; import { resetMemoryManager } from '../../desktop/src/lib/agent-memory'; import { resetAgentIdentityManager } from '../../desktop/src/lib/agent-identity'; import { resetMemoryExtractor } from '../../desktop/src/lib/memory-extractor'; // === Mock localStorage === const localStorageMock = (() => { let store: Record = {}; 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); // === Helpers === function makeMessages(count: number, contentLength: number = 100): CompactableMessage[] { const msgs: CompactableMessage[] = []; for (let i = 0; i < count; i++) { msgs.push({ role: i % 2 === 0 ? 'user' : 'assistant', content: '测试消息内容'.repeat(Math.ceil(contentLength / 6)).slice(0, contentLength), id: `msg_${i}`, timestamp: new Date(Date.now() - (count - i) * 60000), }); } return msgs; } function makeLargeConversation(targetTokens: number): CompactableMessage[] { const msgs: CompactableMessage[] = []; let totalTokens = 0; let i = 0; while (totalTokens < targetTokens) { const content = i % 2 === 0 ? `用户问题 ${i}: 请帮我分析一下这个技术方案的可行性,包括性能、安全性和可维护性方面` : `助手回答 ${i}: 好的,我来从三个维度分析这个方案。首先从性能角度来看,这个方案使用了异步处理机制,能够有效提升吞吐量。其次从安全性方面,建议增加输入验证和权限控制。最后从可维护性来看,模块化设计使得后续修改更加方便。`; msgs.push({ role: i % 2 === 0 ? 'user' : 'assistant', content, id: `msg_${i}`, timestamp: new Date(Date.now() - (1000 - i) * 60000), }); totalTokens = estimateMessagesTokens(msgs); i++; } return msgs; } // ============================================= // Token Estimation Tests // ============================================= describe('Token Estimation', () => { it('returns 0 for empty string', () => { expect(estimateTokens('')).toBe(0); }); it('estimates CJK text at ~1.5 tokens per char', () => { const text = '你好世界测试'; const tokens = estimateTokens(text); // 6 CJK chars × 1.5 = 9 expect(tokens).toBe(9); }); it('estimates English text at ~0.3 tokens per char', () => { const text = 'hello world test'; const tokens = estimateTokens(text); // Roughly: 13 ASCII chars × 0.3 + 2 spaces × 0.25 ≈ 4.4 expect(tokens).toBeGreaterThan(3); expect(tokens).toBeLessThan(10); }); it('estimates mixed CJK+English text', () => { const text = '用户的项目叫 ZCLAW Desktop'; const tokens = estimateTokens(text); expect(tokens).toBeGreaterThan(5); }); it('estimateMessagesTokens includes framing overhead', () => { const msgs: CompactableMessage[] = [ { role: 'user', content: '你好' }, { role: 'assistant', content: '你好!' }, ]; const tokens = estimateMessagesTokens(msgs); // Content tokens + framing (4 per message × 2) expect(tokens).toBeGreaterThan(estimateTokens('你好') + estimateTokens('你好!')); }); }); // ============================================= // ContextCompactor Tests // ============================================= describe('ContextCompactor', () => { let compactor: ContextCompactor; beforeEach(() => { localStorageMock.clear(); resetContextCompactor(); resetMemoryManager(); resetAgentIdentityManager(); resetMemoryExtractor(); compactor = new ContextCompactor(); }); describe('checkThreshold', () => { it('returns none urgency for small conversations', () => { const msgs = makeMessages(4); const check = compactor.checkThreshold(msgs); expect(check.shouldCompact).toBe(false); expect(check.urgency).toBe('none'); }); it('returns soft urgency when approaching threshold', () => { const msgs = makeLargeConversation(DEFAULT_COMPACTION_CONFIG.softThresholdTokens); const check = compactor.checkThreshold(msgs); expect(check.shouldCompact).toBe(true); expect(check.urgency).toBe('soft'); }); it('returns hard urgency when exceeding hard threshold', () => { const msgs = makeLargeConversation(DEFAULT_COMPACTION_CONFIG.hardThresholdTokens); const check = compactor.checkThreshold(msgs); expect(check.shouldCompact).toBe(true); expect(check.urgency).toBe('hard'); }); it('reports current token count', () => { const msgs = makeMessages(10); const check = compactor.checkThreshold(msgs); expect(check.currentTokens).toBeGreaterThan(0); }); }); describe('compact', () => { it('retains keepRecentMessages recent messages', async () => { const config = { keepRecentMessages: 4 }; const comp = new ContextCompactor(config); const msgs = makeMessages(20); const result = await comp.compact(msgs, 'agent-1'); // Should have: 1 summary + 4 recent = 5 expect(result.retainedCount).toBe(5); expect(result.compactedMessages).toHaveLength(5); expect(result.compactedMessages[0].role).toBe('system'); // summary }); it('generates a summary that mentions message count', async () => { const msgs = makeMessages(20); const result = await compactor.compact(msgs, 'agent-1'); expect(result.summary).toContain('压缩'); expect(result.summary).toContain('条消息'); }); it('reduces token count significantly', async () => { const msgs = makeLargeConversation(16000); const result = await compactor.compact(msgs, 'agent-1'); expect(result.tokensAfterCompaction).toBeLessThan(result.tokensBeforeCompaction); }); it('preserves most recent messages in order', async () => { const msgs: CompactableMessage[] = [ { role: 'user', content: 'old message 1', id: 'old1' }, { role: 'assistant', content: 'old reply 1', id: 'old2' }, { role: 'user', content: 'old message 2', id: 'old3' }, { role: 'assistant', content: 'old reply 2', id: 'old4' }, { role: 'user', content: 'recent message 1', id: 'recent1' }, { role: 'assistant', content: 'recent reply 1', id: 'recent2' }, { role: 'user', content: 'recent message 2', id: 'recent3' }, { role: 'assistant', content: 'recent reply 2', id: 'recent4' }, ]; const comp = new ContextCompactor({ keepRecentMessages: 4 }); const result = await comp.compact(msgs, 'agent-1'); // Last 4 messages should be preserved const retained = result.compactedMessages.slice(1); // skip summary expect(retained).toHaveLength(4); expect(retained[0].content).toBe('recent message 1'); expect(retained[3].content).toBe('recent reply 2'); }); it('handles empty message list', async () => { const result = await compactor.compact([], 'agent-1'); expect(result.retainedCount).toBe(1); // just the summary expect(result.summary).toContain('对话开始'); }); it('handles fewer messages than keepRecentMessages', async () => { const msgs = makeMessages(3); const result = await compactor.compact(msgs, 'agent-1'); // All messages kept + summary expect(result.compactedMessages.length).toBeLessThanOrEqual(msgs.length + 1); }); }); describe('memoryFlush', () => { it('returns 0 when disabled', async () => { const comp = new ContextCompactor({ memoryFlushEnabled: false }); const flushed = await comp.memoryFlush(makeMessages(10), 'agent-1'); expect(flushed).toBe(0); }); it('extracts memories from conversation messages', async () => { const msgs: CompactableMessage[] = [ { role: 'user', content: '我的公司叫字节跳动,我在做AI项目' }, { role: 'assistant', content: '好的,了解了。' }, { role: 'user', content: '我喜欢简洁的代码风格' }, { role: 'assistant', content: '明白。' }, { role: 'user', content: '帮我看看这个问题' }, { role: 'assistant', content: '好的。' }, ]; const flushed = await compactor.memoryFlush(msgs, 'agent-1'); // Should extract at least some memories expect(flushed).toBeGreaterThanOrEqual(0); // May or may not match patterns }); }); describe('generateSummary (via compact)', () => { it('includes topic extraction from user messages', async () => { const msgs: CompactableMessage[] = [ { role: 'user', content: '帮我分析一下React性能优化方案' }, { role: 'assistant', content: '好的,React性能优化主要从以下几个方面入手:1. 使用React.memo 2. 使用useMemo' }, { role: 'user', content: '那TypeScript的类型推导呢?' }, { role: 'assistant', content: 'TypeScript类型推导是一个重要特性...' }, ...makeMessages(4), // pad to exceed keepRecentMessages ]; const comp = new ContextCompactor({ keepRecentMessages: 2 }); const result = await comp.compact(msgs, 'agent-1'); // Summary should mention topics expect(result.summary).toContain('讨论主题'); }); it('includes technical context when code blocks present', async () => { const msgs: CompactableMessage[] = [ { role: 'user', content: '帮我写一个函数' }, { role: 'assistant', content: '好的,这是实现:\n```typescript\nfunction hello() { return "world"; }\n```' }, ...makeMessages(6), ]; const comp = new ContextCompactor({ keepRecentMessages: 2 }); const result = await comp.compact(msgs, 'agent-1'); expect(result.summary).toContain('技术上下文'); }); }); describe('buildCompactionPrompt', () => { it('generates a valid LLM prompt', () => { const msgs: CompactableMessage[] = [ { role: 'user', content: '帮我优化数据库查询' }, { role: 'assistant', content: '好的,我建议使用索引...' }, ]; const prompt = compactor.buildCompactionPrompt(msgs); expect(prompt).toContain('压缩为简洁摘要'); expect(prompt).toContain('优化数据库'); expect(prompt).toContain('用户'); expect(prompt).toContain('助手'); }); }); describe('config management', () => { it('uses default config', () => { const config = compactor.getConfig(); expect(config.softThresholdTokens).toBe(15000); expect(config.keepRecentMessages).toBe(6); }); it('allows config updates', () => { compactor.updateConfig({ softThresholdTokens: 10000 }); expect(compactor.getConfig().softThresholdTokens).toBe(10000); }); it('accepts partial config in constructor', () => { const comp = new ContextCompactor({ keepRecentMessages: 10 }); const config = comp.getConfig(); expect(config.keepRecentMessages).toBe(10); expect(config.softThresholdTokens).toBe(15000); // default preserved }); }); });