feat: 添加完整的单元测试套件

- 新增 vitest 测试框架配置
- 添加 54 个测试文件,共 951 个测试用例
- 覆盖核心模块:
  - Agent: executor, registry, config-loader, permission-merger
  - Context: manager, compaction, prune, token-counter
  - Permission: manager, bash/file/git/web checkers, wildcard
  - Session: manager, storage
  - Tools: filesystem (12个), git (10个), web, shell, todo, task
  - LSP: client, server, language
  - Utils: config, diff
  - UI: terminal
This commit is contained in:
2025-12-11 14:45:24 +08:00
parent f4df6483a6
commit 729fb2d42a
58 changed files with 14320 additions and 3 deletions
+279
View File
@@ -0,0 +1,279 @@
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);
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);
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);
});
});
+300
View File
@@ -0,0 +1,300 @@
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);
expect(['prune', 'compaction', 'both']).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);
});
});
+307
View File
@@ -0,0 +1,307 @@
import { describe, it, expect } from 'vitest';
import { prune, filterCompacted } from '../../../src/context/prune.js';
import {
COMPACTED_PLACEHOLDER,
SUMMARY_MARKER,
COMPACTED_MARKER,
type CompressionConfig,
} from '../../../src/context/types.js';
import type { ModelMessage } from 'ai';
// 创建测试用的工具结果消息
function createToolResultMessage(
toolCallId: string,
result: unknown,
compacted = false
): ModelMessage {
const content: Record<string, unknown> = {
type: 'tool-result',
toolCallId,
toolName: 'test_tool',
result,
};
if (compacted) {
content[COMPACTED_MARKER] = {
compactedAt: Date.now(),
originalSize: 100,
};
}
return {
role: 'tool',
content: [content],
} as ModelMessage;
}
// 创建测试用的工具调用消息
function createToolCallMessage(toolCallId: string): ModelMessage {
return {
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId,
toolName: 'test_tool',
args: { param: 'value' },
},
],
} as ModelMessage;
}
// 创建用户消息
function createUserMessage(content: string): ModelMessage {
return { role: 'user', content };
}
// 创建助手消息
function createAssistantMessage(content: string): ModelMessage {
return { role: 'assistant', content };
}
// 创建摘要消息
function createSummaryMessage(): ModelMessage {
return {
role: 'assistant',
content: `${SUMMARY_MARKER}\n## 对话摘要\n\n这是一个摘要\n${SUMMARY_MARKER}`,
};
}
// 创建大工具结果(指定 token 数)
function createLargeToolResult(toolCallId: string, sizeInChars: number): ModelMessage {
// 约 4 字符/token(英文)
const content = 'a'.repeat(sizeInChars);
return createToolResultMessage(toolCallId, { output: content });
}
describe('prune - 消息裁剪策略', () => {
// 测试配置:小值便于测试
const testConfig: CompressionConfig = {
contextLimit: 1000,
outputReserve: 100,
pruneProtect: 200, // 保护最近 200 tokens
pruneMinimum: 50, // 至少释放 50 tokens 才执行
overflowThreshold: 0.85,
};
describe('基本裁剪行为', () => {
it('空消息数组不裁剪', () => {
const result = prune([], testConfig);
expect(result.messages).toHaveLength(0);
expect(result.freedTokens).toBe(0);
});
it('只有用户消息不裁剪', () => {
const messages: ModelMessage[] = [
createUserMessage('Hello'),
createUserMessage('How are you?'),
];
const result = prune(messages, testConfig);
expect(result.messages).toEqual(messages);
expect(result.freedTokens).toBe(0);
});
it('保护范围内的工具结果不裁剪', () => {
// 创建小的工具结果(在保护范围内)
const messages: ModelMessage[] = [
createUserMessage('Use tool'),
createToolCallMessage('call_1'),
createToolResultMessage('call_1', { output: 'small result' }),
];
const result = prune(messages, testConfig);
// 不应该裁剪
expect(result.freedTokens).toBe(0);
});
it('裁剪超出保护范围的工具结果', () => {
// 创建多个工具结果,超出保护范围
const messages: ModelMessage[] = [
createUserMessage('Task 1'),
createToolCallMessage('call_1'),
createLargeToolResult('call_1', 1000), // ~250 tokens, 超出 pruneProtect
createUserMessage('Task 2'),
createToolCallMessage('call_2'),
createLargeToolResult('call_2', 400), // ~100 tokens
createUserMessage('Task 3'),
createToolCallMessage('call_3'),
createToolResultMessage('call_3', { output: 'recent' }), // 最近的,在保护范围内
];
const result = prune(messages, testConfig);
// 应该裁剪旧的大工具结果
expect(result.freedTokens).toBeGreaterThan(0);
});
});
describe('摘要消息边界', () => {
it('遇到摘要消息停止裁剪', () => {
const messages: ModelMessage[] = [
createSummaryMessage(), // 摘要消息
createUserMessage('After summary'),
createToolCallMessage('call_1'),
createLargeToolResult('call_1', 2000), // 大结果
];
const result = prune(messages, testConfig);
// 因为遇到摘要消息,不会继续向前裁剪
// 但是摘要后面的大工具结果如果超出保护范围仍会被裁剪
expect(result.messages[0]).toEqual(messages[0]); // 摘要保留
});
it('摘要消息前的内容不裁剪', () => {
const messages: ModelMessage[] = [
createUserMessage('Before summary'),
createToolCallMessage('call_old'),
createLargeToolResult('call_old', 2000), // 摘要前的大结果
createSummaryMessage(),
createUserMessage('After summary'),
];
const result = prune(messages, testConfig);
// 摘要前的内容因为遇到摘要边界停止
expect(result.messages.length).toBe(messages.length);
});
});
describe('已压缩内容处理', () => {
it('遇到已压缩的工具结果停止', () => {
const messages: ModelMessage[] = [
createUserMessage('Task 1'),
createToolCallMessage('call_1'),
createToolResultMessage('call_1', COMPACTED_PLACEHOLDER, true), // 已压缩
createUserMessage('Task 2'),
createToolCallMessage('call_2'),
createLargeToolResult('call_2', 1000), // 新的大结果
];
const result = prune(messages, testConfig);
// 遇到已压缩的结果应该停止继续向前
expect(result.messages).toBeDefined();
});
});
describe('最小裁剪量检查', () => {
it('释放量不足最小值时不执行裁剪', () => {
// 使用较大的 pruneMinimum
const strictConfig: CompressionConfig = {
...testConfig,
pruneMinimum: 10000, // 要求至少释放 10000 tokens
};
const messages: ModelMessage[] = [
createUserMessage('Task'),
createToolCallMessage('call_1'),
createLargeToolResult('call_1', 400), // 只有 ~100 tokens
];
const result = prune(messages, strictConfig);
expect(result.freedTokens).toBe(0);
expect(result.messages).toEqual(messages);
});
});
describe('深拷贝验证', () => {
it('不修改原消息数组', () => {
const messages: ModelMessage[] = [
createUserMessage('Task'),
createToolCallMessage('call_1'),
createLargeToolResult('call_1', 2000),
createUserMessage('Recent'),
];
const originalLength = messages.length;
const originalFirstMessage = { ...messages[0] };
prune(messages, testConfig);
expect(messages.length).toBe(originalLength);
expect(messages[0]).toEqual(originalFirstMessage);
});
});
});
describe('filterCompacted - 过滤已压缩内容', () => {
it('不修改普通消息', () => {
const messages: ModelMessage[] = [
createUserMessage('Hello'),
createAssistantMessage('Hi there'),
];
const result = filterCompacted(messages);
expect(result).toEqual(messages);
});
it('不修改未压缩的工具结果', () => {
const messages: ModelMessage[] = [
createToolResultMessage('call_1', { output: 'normal result' }),
];
const result = filterCompacted(messages);
expect(result).toEqual(messages);
});
it('将已压缩工具结果替换为占位符', () => {
const messages: ModelMessage[] = [
createToolResultMessage('call_1', 'original', true), // 已压缩
];
const result = filterCompacted(messages);
const content = result[0].content as { result: unknown }[];
expect(content[0].result).toBe(COMPACTED_PLACEHOLDER);
});
it('混合内容正确处理', () => {
const messages: ModelMessage[] = [
createUserMessage('Task'),
createToolResultMessage('call_1', 'original', true), // 已压缩
createToolResultMessage('call_2', { output: 'normal' }), // 未压缩
createAssistantMessage('Done'),
];
const result = filterCompacted(messages);
expect(result).toHaveLength(4);
// 第一个工具结果应该被替换
const compactedContent = result[1].content as { result: unknown }[];
expect(compactedContent[0].result).toBe(COMPACTED_PLACEHOLDER);
// 第二个工具结果应该保持不变
const normalContent = result[2].content as { result: unknown }[];
expect(normalContent[0].result).toEqual({ output: 'normal' });
});
it('保留消息的其他属性', () => {
const messages: ModelMessage[] = [
{
role: 'tool',
content: [
{
type: 'tool-result',
toolCallId: 'call_1',
toolName: 'my_tool',
result: 'data',
[COMPACTED_MARKER]: { compactedAt: 123, originalSize: 100 },
},
],
} as ModelMessage,
];
const result = filterCompacted(messages);
const content = result[0].content as { toolCallId: string; toolName: string }[];
expect(content[0].toolCallId).toBe('call_1');
expect(content[0].toolName).toBe('my_tool');
});
});
+349
View File
@@ -0,0 +1,349 @@
import { describe, it, expect } from 'vitest';
import { TokenCounter } from '../../../src/context/token-counter.js';
import type { ModelMessage } from 'ai';
describe('TokenCounter - Token 计数器', () => {
describe('estimateText - 文本估算', () => {
it('空文本返回 0', () => {
expect(TokenCounter.estimateText('')).toBe(0);
expect(TokenCounter.estimateText(null as unknown as string)).toBe(0);
expect(TokenCounter.estimateText(undefined as unknown as string)).toBe(0);
});
it('纯英文估算(约 4 字符/token', () => {
// 40 个字符 / 4 = 10 tokens
const text = 'This is a test message with some words.';
const tokens = TokenCounter.estimateText(text);
expect(tokens).toBe(Math.ceil(text.length / 4));
});
it('纯中文估算(约 1.5 字符/token', () => {
// 6 个中文字符 / 1.5 = 4 tokens
const text = '这是测试文本';
const tokens = TokenCounter.estimateText(text);
expect(tokens).toBe(Math.ceil(6 / 1.5));
});
it('中英混合估算', () => {
// 中文 4 个 + 英文 10 个
// 4/1.5 + 10/4 = 2.67 + 2.5 = 5.17 -> 6
const text = '测试test文本text';
const tokens = TokenCounter.estimateText(text);
// 4 个中文: 4/1.5 = 2.67
// 8 个其他: 8/4 = 2
// 总计: ceil(4.67) = 5
expect(tokens).toBeGreaterThan(0);
expect(tokens).toBeLessThan(text.length); // 应该小于字符数
});
it('代码片段估算', () => {
const code = `function hello() {
console.log("Hello World");
return true;
}`;
const tokens = TokenCounter.estimateText(code);
expect(tokens).toBeGreaterThan(0);
// 代码主要是英文,约 4 字符/token
expect(tokens).toBeLessThan(code.length);
});
it('长文本估算', () => {
const longText = 'a'.repeat(10000);
const tokens = TokenCounter.estimateText(longText);
// 10000 / 4 = 2500
expect(tokens).toBe(2500);
});
});
describe('estimateContent - 内容估算', () => {
it('字符串内容', () => {
const content = 'Hello World';
const tokens = TokenCounter.estimateContent(content);
expect(tokens).toBe(TokenCounter.estimateText(content));
});
it('数组内容 - 纯文本部分', () => {
const content = ['Hello', 'World'];
const tokens = TokenCounter.estimateContent(content);
const expected = TokenCounter.estimateText('Hello') + TokenCounter.estimateText('World');
expect(tokens).toBe(expected);
});
it('数组内容 - text 对象', () => {
const content = [{ type: 'text', text: 'Hello World' }];
const tokens = TokenCounter.estimateContent(content);
expect(tokens).toBe(TokenCounter.estimateText('Hello World'));
});
it('数组内容 - tool-result', () => {
const content = [
{
type: 'tool-result',
toolCallId: 'call_123',
toolName: 'read_file',
result: { success: true, output: 'file content' },
},
];
const tokens = TokenCounter.estimateContent(content);
const expectedText = JSON.stringify({ success: true, output: 'file content' });
expect(tokens).toBe(TokenCounter.estimateText(expectedText));
});
it('数组内容 - tool-call', () => {
const content = [
{
type: 'tool-call',
toolCallId: 'call_123',
toolName: 'read_file',
args: { path: '/test.txt' },
},
];
const tokens = TokenCounter.estimateContent(content);
const argsText = JSON.stringify({ path: '/test.txt' });
// 工具调用增加 20 token 开销
expect(tokens).toBe(TokenCounter.estimateText(argsText) + 20);
});
it('混合内容', () => {
const content = [
{ type: 'text', text: 'Processing file' },
{
type: 'tool-call',
toolCallId: 'call_1',
toolName: 'read_file',
args: { path: '/a.txt' },
},
];
const tokens = TokenCounter.estimateContent(content);
expect(tokens).toBeGreaterThan(0);
});
it('空数组返回 0', () => {
expect(TokenCounter.estimateContent([])).toBe(0);
});
it('非字符串非数组返回 0', () => {
expect(TokenCounter.estimateContent(null as unknown as string)).toBe(0);
expect(TokenCounter.estimateContent(123 as unknown as string)).toBe(0);
});
});
describe('estimateMessage - 单条消息估算', () => {
it('用户消息', () => {
const message: ModelMessage = {
role: 'user',
content: 'Hello',
};
const tokens = TokenCounter.estimateMessage(message);
// 4 (角色开销) + 内容 tokens
expect(tokens).toBe(4 + TokenCounter.estimateText('Hello'));
});
it('助手消息', () => {
const message: ModelMessage = {
role: 'assistant',
content: 'I can help you with that.',
};
const tokens = TokenCounter.estimateMessage(message);
expect(tokens).toBe(4 + TokenCounter.estimateText('I can help you with that.'));
});
it('系统消息', () => {
const message: ModelMessage = {
role: 'system',
content: 'You are a helpful assistant.',
};
const tokens = TokenCounter.estimateMessage(message);
expect(tokens).toBe(4 + TokenCounter.estimateText('You are a helpful assistant.'));
});
it('工具消息', () => {
const message: ModelMessage = {
role: 'tool',
content: [
{
type: 'tool-result',
toolCallId: 'call_123',
toolName: 'bash',
result: { success: true, output: 'done' },
},
],
};
const tokens = TokenCounter.estimateMessage(message);
expect(tokens).toBeGreaterThan(4); // 至少有角色开销
});
});
describe('estimateMessages - 消息数组估算', () => {
it('空数组返回 0', () => {
expect(TokenCounter.estimateMessages([])).toBe(0);
});
it('单条消息', () => {
const messages: ModelMessage[] = [{ role: 'user', content: 'Hello' }];
const tokens = TokenCounter.estimateMessages(messages);
// 消息 tokens + 3 (分隔开销)
expect(tokens).toBe(TokenCounter.estimateMessage(messages[0]) + 3);
});
it('多条消息', () => {
const messages: ModelMessage[] = [
{ role: 'user', content: 'Hello' },
{ role: 'assistant', content: 'Hi there!' },
{ role: 'user', content: 'How are you?' },
];
const tokens = TokenCounter.estimateMessages(messages);
const msgTokens = messages.reduce((sum, m) => sum + TokenCounter.estimateMessage(m), 0);
const separatorTokens = messages.length * 3;
expect(tokens).toBe(msgTokens + separatorTokens);
});
it('包含工具调用的对话', () => {
const messages: ModelMessage[] = [
{ role: 'user', content: '读取 /tmp/test.txt' },
{
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId: 'call_1',
toolName: 'read_file',
args: { path: '/tmp/test.txt' },
},
],
},
{
role: 'tool',
content: [
{
type: 'tool-result',
toolCallId: 'call_1',
toolName: 'read_file',
result: { success: true, output: 'file content here' },
},
],
},
{ role: 'assistant', content: '文件内容是: file content here' },
];
const tokens = TokenCounter.estimateMessages(messages);
expect(tokens).toBeGreaterThan(0);
// 应该能处理复杂的消息结构
});
});
describe('format - 格式化显示', () => {
it('小于 1000 显示原数', () => {
expect(TokenCounter.format(0)).toBe('0');
expect(TokenCounter.format(100)).toBe('100');
expect(TokenCounter.format(999)).toBe('999');
});
it('大于等于 1000 显示 k 单位', () => {
expect(TokenCounter.format(1000)).toBe('1.0k');
expect(TokenCounter.format(1500)).toBe('1.5k');
expect(TokenCounter.format(10000)).toBe('10.0k');
expect(TokenCounter.format(100000)).toBe('100.0k');
});
it('小数精度', () => {
expect(TokenCounter.format(1234)).toBe('1.2k');
expect(TokenCounter.format(1250)).toBe('1.3k'); // 四舍五入
expect(TokenCounter.format(12345)).toBe('12.3k');
});
});
});
describe('TokenCounter 实际场景测试', () => {
it('典型对话 token 估算', () => {
const messages: ModelMessage[] = [
{
role: 'system',
content:
'You are a helpful coding assistant. Help users with programming tasks.',
},
{
role: 'user',
content: '请帮我写一个 Python 函数来计算斐波那契数列',
},
{
role: 'assistant',
content: `好的,这是一个计算斐波那契数列的 Python 函数:
\`\`\`python
def fibonacci(n):
if n <= 0:
return []
elif n == 1:
return [0]
elif n == 2:
return [0, 1]
fib = [0, 1]
for i in range(2, n):
fib.append(fib[i-1] + fib[i-2])
return fib
\`\`\``,
},
];
const tokens = TokenCounter.estimateMessages(messages);
expect(tokens).toBeGreaterThan(100); // 应该有一定数量的 tokens
expect(tokens).toBeLessThan(1000); // 但不会太多
});
it('大量工具调用的 token 估算', () => {
const messages: ModelMessage[] = [];
// 模拟 10 轮工具调用
for (let i = 0; i < 10; i++) {
messages.push({
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId: `call_${i}`,
toolName: 'bash',
args: { command: `echo "iteration ${i}"` },
},
],
});
messages.push({
role: 'tool',
content: [
{
type: 'tool-result',
toolCallId: `call_${i}`,
toolName: 'bash',
result: { success: true, output: `iteration ${i}` },
},
],
});
}
const tokens = TokenCounter.estimateMessages(messages);
expect(tokens).toBeGreaterThan(0);
// 20 条消息应该有合理的 token 数
expect(TokenCounter.format(tokens)).toBeDefined();
});
it('上下文窗口占用估算', () => {
// 模拟 200k token 上下文窗口
const maxContextTokens = 200000;
// 创建一个大消息
const largeContent = 'a'.repeat(40000); // 约 10k tokens
const messages: ModelMessage[] = [
{ role: 'user', content: largeContent },
];
const tokens = TokenCounter.estimateMessages(messages);
const usagePercent = (tokens / maxContextTokens) * 100;
expect(usagePercent).toBeLessThan(10); // 应该占用不到 10%
expect(usagePercent).toBeGreaterThan(0);
});
});