import { describe, it, expect, beforeEach, vi } from 'vitest'; import { CompressionManager } from '../../../src/context/manager.js'; import { DEFAULT_COMPRESSION_CONFIG, SUMMARY_MARKER } from '../../../src/context/types.js'; import type { ModelMessage, LanguageModel } from 'ai'; // 创建测试用消息 function createUserMessage(content: string): ModelMessage { return { role: 'user', content }; } function createAssistantMessage(content: string): ModelMessage { return { role: 'assistant', content }; } function createSummaryMessage(summary: string): ModelMessage { return { role: 'assistant', content: `${SUMMARY_MARKER}\n## 对话摘要\n\n${summary}\n${SUMMARY_MARKER}`, }; } // 创建大量消息以超过阈值 function createLargeConversation(count: number): ModelMessage[] { const messages: ModelMessage[] = []; for (let i = 0; i < count; i++) { messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`)); messages.push(createAssistantMessage(`Response ${i}: ${'b'.repeat(100)}`)); } return messages; } describe('CompressionManager - 压缩管理器', () => { let manager: CompressionManager; beforeEach(() => { manager = new CompressionManager(); }); describe('配置管理', () => { it('使用默认配置', () => { const config = manager.getConfig(); expect(config.contextLimit).toBe(DEFAULT_COMPRESSION_CONFIG.contextLimit); expect(config.outputReserve).toBe(DEFAULT_COMPRESSION_CONFIG.outputReserve); expect(config.overflowThreshold).toBe(DEFAULT_COMPRESSION_CONFIG.overflowThreshold); }); it('自定义配置覆盖默认值', () => { const customManager = new CompressionManager({ contextLimit: 50000, outputReserve: 5000, }); const config = customManager.getConfig(); expect(config.contextLimit).toBe(50000); expect(config.outputReserve).toBe(5000); // 未指定的使用默认值 expect(config.overflowThreshold).toBe(DEFAULT_COMPRESSION_CONFIG.overflowThreshold); }); it('updateConfig 更新配置', () => { manager.updateConfig({ pruneProtect: 10000 }); const config = manager.getConfig(); expect(config.pruneProtect).toBe(10000); }); }); describe('calculateUsage - 计算 token 使用情况', () => { it('空消息返回零使用', () => { const usage = manager.calculateUsage([]); expect(usage.input).toBe(0); expect(usage.usagePercent).toBe(0); }); it('计算简单消息的使用量', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi there'), ]; const usage = manager.calculateUsage(messages); expect(usage.input).toBeGreaterThan(0); expect(usage.usagePercent).toBeGreaterThan(0); expect(usage.usagePercent).toBeLessThan(100); }); it('大量消息使用量更高', () => { const smallMessages: ModelMessage[] = [ createUserMessage('Hello'), ]; const largeMessages = createLargeConversation(50); const smallUsage = manager.calculateUsage(smallMessages); const largeUsage = manager.calculateUsage(largeMessages); expect(largeUsage.input).toBeGreaterThan(smallUsage.input); }); it('使用量百分比不超过 100', () => { // 创建超大对话 const messages = createLargeConversation(500); const usage = manager.calculateUsage(messages); expect(usage.usagePercent).toBeLessThanOrEqual(100); }); }); describe('shouldCompress - 判断是否需要压缩', () => { it('小对话不需要压缩', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; expect(manager.shouldCompress(messages)).toBe(false); }); it('大对话可能需要压缩', () => { // 创建足够大的对话以超过阈值 const messages = createLargeConversation(200); // 取决于配置的阈值 const usage = manager.calculateUsage(messages); const threshold = manager.getConfig().overflowThreshold * 100; const shouldCompress = usage.usagePercent >= threshold; expect(manager.shouldCompress(messages)).toBe(shouldCompress); }); }); describe('isOverflow - 判断是否溢出', () => { it('小对话不溢出', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; expect(manager.isOverflow(messages)).toBe(false); }); }); describe('prune - 执行裁剪', () => { it('空消息不裁剪', () => { const result = manager.prune([]); expect(result.messages).toHaveLength(0); expect(result.freedTokens).toBe(0); }); it('小对话不裁剪', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; const result = manager.prune(messages); expect(result.messages).toEqual(messages); expect(result.freedTokens).toBe(0); }); }); describe('compact - 执行压缩', () => { it('无模型时使用简单压缩', async () => { const messages = createLargeConversation(50); const result = await manager.compact(messages); // 简单压缩会根据配置决定是否压缩 expect(result.messages).toBeDefined(); expect(result.freedTokens).toBeGreaterThanOrEqual(0); }); it('设置模型后使用 AI 压缩', async () => { // 创建 mock 模型 const mockModel = { doGenerate: vi.fn().mockResolvedValue({ text: '这是一个摘要', }), } as unknown as LanguageModel; manager.setModel(mockModel); const messages = createLargeConversation(50); const result = await manager.compact(messages); expect(result.messages).toBeDefined(); }); }); describe('compress - 自动压缩', () => { it('小对话不压缩', async () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; const result = await manager.compress(messages); expect(result.messages).toEqual(messages); expect(result.freedTokens).toBe(0); }); it('返回正确的压缩类型', async () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; const result = await manager.compress(messages); // 小对话不压缩时返回 'none' expect(['prune', 'compaction', 'both', 'none']).toContain(result.type); }); }); describe('forceCompress - 强制压缩', () => { it('消息数量少于 4 条不压缩', async () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; const result = await manager.forceCompress(messages); expect(result.messages).toEqual(messages); expect(result.freedTokens).toBe(0); }); it('强制压缩大对话', async () => { const messages = createLargeConversation(20); const result = await manager.forceCompress(messages); // 强制压缩应该减少消息数量 expect(result.messages.length).toBeLessThanOrEqual(messages.length); }); }); describe('filterCompacted - 过滤已压缩内容', () => { it('不修改普通消息', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; const result = manager.filterCompacted(messages); expect(result).toEqual(messages); }); }); describe('isSummaryMessage - 检测摘要消息', () => { it('检测摘要消息', () => { const summary = createSummaryMessage('这是摘要'); expect(manager.isSummaryMessage(summary)).toBe(true); }); it('普通消息不是摘要', () => { const normal = createAssistantMessage('普通回复'); expect(manager.isSummaryMessage(normal)).toBe(false); }); }); describe('formatUsage - 格式化使用情况', () => { it('格式化空消息', () => { const formatted = manager.formatUsage([]); expect(formatted).toContain('/'); expect(formatted).toContain('('); expect(formatted).toContain('%)'); }); it('格式化包含消息的使用情况', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi there'), ]; const formatted = manager.formatUsage(messages); expect(typeof formatted).toBe('string'); expect(formatted).toContain('/'); }); }); }); describe('compressionManager 单例', () => { it('导出单例实例', async () => { const { compressionManager } = await import('../../../src/context/manager.js'); expect(compressionManager).toBeDefined(); expect(compressionManager).toBeInstanceOf(CompressionManager); }); });