import { describe, it, expect, vi } from 'vitest'; import { isSummaryMessage, simpleCompact } from '../../../src/context/compaction.js'; import { SUMMARY_MARKER, type CompressionConfig, } from '../../../src/context/types.js'; import type { ModelMessage } 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 createMessageWithTextParts(texts: string[]): ModelMessage { return { role: 'assistant', content: texts.map((text) => ({ type: 'text', text })), } as ModelMessage; } describe('isSummaryMessage - 检测摘要消息', () => { it('字符串内容包含摘要标记返回 true', () => { const message = createSummaryMessage('这是摘要内容'); expect(isSummaryMessage(message)).toBe(true); }); it('字符串内容不包含摘要标记返回 false', () => { const message = createAssistantMessage('普通助手消息'); expect(isSummaryMessage(message)).toBe(false); }); it('数组内容包含摘要标记返回 true', () => { const message = createMessageWithTextParts([ '一些文本', `${SUMMARY_MARKER}\n摘要内容\n${SUMMARY_MARKER}`, ]); expect(isSummaryMessage(message)).toBe(true); }); it('数组内容不包含摘要标记返回 false', () => { const message = createMessageWithTextParts(['文本1', '文本2']); expect(isSummaryMessage(message)).toBe(false); }); it('用户消息不是摘要消息', () => { const message = createUserMessage(SUMMARY_MARKER); // 虽然包含标记,但这种情况下 isSummaryMessage 还是会返回 true // 因为它只检查内容是否包含标记 expect(isSummaryMessage(message)).toBe(true); }); it('空数组内容返回 false', () => { const message: ModelMessage = { role: 'assistant', content: [], }; expect(isSummaryMessage(message)).toBe(false); }); it('非文本部分不匹配', () => { const message: ModelMessage = { role: 'assistant', content: [ { type: 'tool-call', toolCallId: 'call_1', toolName: 'test', args: { key: SUMMARY_MARKER }, }, ], } as ModelMessage; expect(isSummaryMessage(message)).toBe(false); }); }); describe('simpleCompact - 简单压缩', () => { const testConfig: CompressionConfig = { contextLimit: 1000, outputReserve: 100, pruneProtect: 200, // 保护最近 200 tokens pruneMinimum: 50, overflowThreshold: 0.85, }; describe('基本压缩行为', () => { it('空消息数组不压缩', () => { const result = simpleCompact([], testConfig); expect(result.messages).toHaveLength(0); expect(result.freedTokens).toBe(0); }); it('消息在保护范围内不压缩', () => { const messages: ModelMessage[] = [ createUserMessage('Hello'), createAssistantMessage('Hi'), ]; const result = simpleCompact(messages, testConfig); // 消息很短,应该在保护范围内 expect(result.freedTokens).toBe(0); expect(result.messages).toEqual(messages); }); it('压缩超出保护范围的消息', () => { // 创建大量消息超出保护范围 const messages: ModelMessage[] = []; for (let i = 0; i < 50; i++) { messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`)); messages.push(createAssistantMessage(`Response ${i}: ${'b'.repeat(100)}`)); } const result = simpleCompact(messages, testConfig); // 应该压缩一些旧消息 expect(result.freedTokens).toBeGreaterThan(0); expect(result.messages.length).toBeLessThan(messages.length); }); }); describe('保留消息数量', () => { it('至少保留 2 条消息(正常模式)', () => { const messages: ModelMessage[] = []; for (let i = 0; i < 10; i++) { messages.push(createUserMessage(`Long message ${'a'.repeat(500)}`)); } const result = simpleCompact(messages, testConfig); // 即使压缩,也至少保留 2 条 expect(result.messages.length).toBeGreaterThanOrEqual(2 + 1); // 2 保留 + 1 摘要 }); it('强制模式下至少保留 1 条消息', () => { const forceConfig: CompressionConfig = { ...testConfig, pruneProtect: 0, // 强制模式 }; const messages: ModelMessage[] = []; for (let i = 0; i < 10; i++) { messages.push(createUserMessage(`Message ${i}`)); } const result = simpleCompact(messages, forceConfig); // 强制模式下至少保留 1 条消息 expect(result.messages.length).toBeGreaterThanOrEqual(2); // 1 保留 + 1 摘要 }); }); describe('摘要消息生成', () => { it('压缩后生成摘要消息', () => { const messages: ModelMessage[] = []; for (let i = 0; i < 50; i++) { messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`)); } // 禁用首条保护以便测试摘要消息在第一位 const result = simpleCompact(messages, testConfig, { protectFirstPair: false }); if (result.freedTokens > 0) { // 第一条消息应该是摘要 expect(isSummaryMessage(result.messages[0])).toBe(true); } }); it('摘要消息包含移除数量信息', () => { const messages: ModelMessage[] = []; for (let i = 0; i < 50; i++) { messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`)); } // 禁用首条保护以便测试摘要消息 const result = simpleCompact(messages, testConfig, { protectFirstPair: false }); if (result.freedTokens > 0) { const summaryContent = result.messages[0].content as string; expect(summaryContent).toContain('对话历史已压缩'); expect(summaryContent).toContain('条消息'); } }); }); describe('保护范围计算', () => { it('短消息全部在保护范围内', () => { const messages: ModelMessage[] = [ createUserMessage('Hi'), createAssistantMessage('Hello'), createUserMessage('How?'), createAssistantMessage('Good!'), ]; const result = simpleCompact(messages, testConfig); // 短消息应该全部保留 expect(result.messages).toEqual(messages); }); }); describe('不修改原数组', () => { it('原消息数组不被修改', () => { const messages: ModelMessage[] = []; for (let i = 0; i < 20; i++) { messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`)); } const originalLength = messages.length; simpleCompact(messages, testConfig); expect(messages.length).toBe(originalLength); }); }); }); // 注意: compact 函数需要真实的 LanguageModel, // 这里只测试 simpleCompact 和辅助函数 // compact 的测试应该在集成测试中进行 describe('Compaction 边界情况', () => { it('单条消息不压缩', () => { const messages: ModelMessage[] = [ createUserMessage('Single message'), ]; const testConfig: CompressionConfig = { contextLimit: 100, outputReserve: 10, pruneProtect: 50, pruneMinimum: 10, overflowThreshold: 0.85, }; const result = simpleCompact(messages, testConfig); expect(result.messages).toEqual(messages); expect(result.freedTokens).toBe(0); }); it('两条消息不压缩(最小保留)', () => { const messages: ModelMessage[] = [ createUserMessage('First'), createAssistantMessage('Second'), ]; const testConfig: CompressionConfig = { contextLimit: 100, outputReserve: 10, pruneProtect: 10, // 很小的保护范围 pruneMinimum: 1, overflowThreshold: 0.85, }; const result = simpleCompact(messages, testConfig); // 即使配置很激进,也至少保留 2 条 expect(result.messages.length).toBeGreaterThanOrEqual(2); }); });