feat(core): 实现 Token 消耗统计系统
- 扩展 SessionStats schema 添加 token 统计字段 - 添加 TokenUsageInfo 类型和 ChatResult.usage 字段 - AgentMessageHandler 从 AI SDK response 提取 usage - AgentExecutor 返回 usage 到执行结果 - 新增 TokenStatsManager 管理统计: - updateSessionStats: 更新会话 token 统计 - mergeChildSessionStats: 合并子会话统计到父会话 - getSessionStats/getProjectStats: 查询统计 - Agent.chat() 完成后自动更新统计 - Task 工具完成后合并子会话统计 - 新增 REST API: /api/stats/sessions/:id, /api/stats/projects/:id - 添加 TokenStatsManager 单元测试 (12 tests)
This commit is contained in:
@@ -0,0 +1,372 @@
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
|
||||
// Mock storage functions before imports
|
||||
vi.mock('../../../src/session/storage/session.js', () => ({
|
||||
get: vi.fn(),
|
||||
update: vi.fn(),
|
||||
updateStats: vi.fn(),
|
||||
list: vi.fn(),
|
||||
listByProject: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../src/session/storage/index.js', () => ({
|
||||
read: vi.fn(),
|
||||
update: vi.fn(),
|
||||
StorageNotFoundError: class StorageNotFoundError extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = 'StorageNotFoundError';
|
||||
}
|
||||
},
|
||||
}));
|
||||
|
||||
import * as SessionStorage from '../../../src/session/storage/session.js';
|
||||
import { TokenStatsManager } from '../../../src/session/token-stats-manager.js';
|
||||
import type { TokenUsageInfo } from '../../../src/types/index.js';
|
||||
|
||||
describe('TokenStatsManager', () => {
|
||||
let manager: TokenStatsManager;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
manager = new TokenStatsManager();
|
||||
});
|
||||
|
||||
describe('updateSessionStats', () => {
|
||||
it('会话不存在时不更新', async () => {
|
||||
const usage: TokenUsageInfo = {
|
||||
promptTokens: 1000,
|
||||
completionTokens: 500,
|
||||
totalTokens: 1500,
|
||||
};
|
||||
|
||||
vi.mocked(SessionStorage.get).mockResolvedValue(null);
|
||||
|
||||
await manager.updateSessionStats('project-1', 'session-1', usage);
|
||||
|
||||
expect(SessionStorage.get).toHaveBeenCalledWith('project-1', 'session-1');
|
||||
expect(SessionStorage.updateStats).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('更新会话统计 - 已存在的会话', async () => {
|
||||
const usage: TokenUsageInfo = {
|
||||
promptTokens: 1000,
|
||||
completionTokens: 500,
|
||||
totalTokens: 1500,
|
||||
};
|
||||
|
||||
vi.mocked(SessionStorage.get).mockResolvedValue({
|
||||
id: 'session-1',
|
||||
projectId: 'project-1',
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 5,
|
||||
inputTokens: 2000,
|
||||
outputTokens: 1000,
|
||||
totalTokens: 3000,
|
||||
updatedAt: Date.now() - 1000,
|
||||
},
|
||||
});
|
||||
|
||||
await manager.updateSessionStats('project-1', 'session-1', usage);
|
||||
|
||||
expect(SessionStorage.updateStats).toHaveBeenCalledWith(
|
||||
'project-1',
|
||||
'session-1',
|
||||
expect.objectContaining({
|
||||
messageCount: 6, // 5 + 1
|
||||
inputTokens: 3000, // 2000 + 1000
|
||||
outputTokens: 1500, // 1000 + 500
|
||||
totalTokens: 4500, // 3000 + 1500
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('处理缓存 token', async () => {
|
||||
const usage: TokenUsageInfo = {
|
||||
promptTokens: 1000,
|
||||
completionTokens: 500,
|
||||
totalTokens: 1500,
|
||||
cacheReadInputTokens: 200,
|
||||
cacheCreationInputTokens: 100,
|
||||
};
|
||||
|
||||
vi.mocked(SessionStorage.get).mockResolvedValue({
|
||||
id: 'session-1',
|
||||
projectId: 'project-1',
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 0,
|
||||
inputTokens: 0,
|
||||
outputTokens: 0,
|
||||
totalTokens: 0,
|
||||
updatedAt: Date.now(),
|
||||
},
|
||||
});
|
||||
|
||||
await manager.updateSessionStats('project-1', 'session-1', usage);
|
||||
|
||||
expect(SessionStorage.updateStats).toHaveBeenCalledWith(
|
||||
'project-1',
|
||||
'session-1',
|
||||
expect.objectContaining({
|
||||
cacheReadTokens: 200,
|
||||
cacheWriteTokens: 100,
|
||||
})
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('mergeChildSessionStats', () => {
|
||||
it('合并子会话统计到父会话', async () => {
|
||||
vi.mocked(SessionStorage.get).mockImplementation(async (projectId: string, sessionId: string) => {
|
||||
if (sessionId === 'child-session') {
|
||||
return {
|
||||
id: 'child-session',
|
||||
projectId,
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 3,
|
||||
inputTokens: 500,
|
||||
outputTokens: 300,
|
||||
totalTokens: 800,
|
||||
updatedAt: Date.now(),
|
||||
},
|
||||
};
|
||||
}
|
||||
if (sessionId === 'parent-session') {
|
||||
return {
|
||||
id: 'parent-session',
|
||||
projectId,
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 10,
|
||||
inputTokens: 2000,
|
||||
outputTokens: 1000,
|
||||
totalTokens: 3000,
|
||||
updatedAt: Date.now(),
|
||||
},
|
||||
};
|
||||
}
|
||||
return null;
|
||||
});
|
||||
|
||||
await manager.mergeChildSessionStats('project-1', 'parent-session', 'child-session');
|
||||
|
||||
expect(SessionStorage.updateStats).toHaveBeenCalledWith(
|
||||
'project-1',
|
||||
'parent-session',
|
||||
expect.objectContaining({
|
||||
childSessionsTokens: expect.objectContaining({
|
||||
inputTokens: 500,
|
||||
outputTokens: 300,
|
||||
totalTokens: 800,
|
||||
}),
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('累加子会话统计', async () => {
|
||||
vi.mocked(SessionStorage.get).mockImplementation(async (projectId: string, sessionId: string) => {
|
||||
if (sessionId === 'child-session-2') {
|
||||
return {
|
||||
id: 'child-session-2',
|
||||
projectId,
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 2,
|
||||
inputTokens: 300,
|
||||
outputTokens: 200,
|
||||
totalTokens: 500,
|
||||
updatedAt: Date.now(),
|
||||
},
|
||||
};
|
||||
}
|
||||
if (sessionId === 'parent-session') {
|
||||
return {
|
||||
id: 'parent-session',
|
||||
projectId,
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 10,
|
||||
inputTokens: 2000,
|
||||
outputTokens: 1000,
|
||||
totalTokens: 3000,
|
||||
childSessionsTokens: {
|
||||
inputTokens: 500,
|
||||
outputTokens: 300,
|
||||
totalTokens: 800,
|
||||
},
|
||||
updatedAt: Date.now(),
|
||||
},
|
||||
};
|
||||
}
|
||||
return null;
|
||||
});
|
||||
|
||||
await manager.mergeChildSessionStats('project-1', 'parent-session', 'child-session-2');
|
||||
|
||||
expect(SessionStorage.updateStats).toHaveBeenCalledWith(
|
||||
'project-1',
|
||||
'parent-session',
|
||||
expect.objectContaining({
|
||||
childSessionsTokens: expect.objectContaining({
|
||||
inputTokens: 800, // 500 + 300
|
||||
outputTokens: 500, // 300 + 200
|
||||
totalTokens: 1300, // 800 + 500
|
||||
}),
|
||||
})
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getSessionStats', () => {
|
||||
it('返回完整的会话统计', async () => {
|
||||
const now = Date.now();
|
||||
vi.mocked(SessionStorage.get).mockResolvedValue({
|
||||
id: 'session-1',
|
||||
projectId: 'project-1',
|
||||
createdAt: now,
|
||||
updatedAt: now,
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 10,
|
||||
inputTokens: 5000,
|
||||
outputTokens: 2500,
|
||||
totalTokens: 7500,
|
||||
cacheReadTokens: 1000,
|
||||
cacheWriteTokens: 500,
|
||||
childSessionsTokens: {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
},
|
||||
updatedAt: now,
|
||||
},
|
||||
});
|
||||
|
||||
const stats = await manager.getSessionStats('project-1', 'session-1');
|
||||
|
||||
expect(stats).toMatchObject({
|
||||
self: {
|
||||
inputTokens: 5000,
|
||||
outputTokens: 2500,
|
||||
totalTokens: 7500,
|
||||
cacheReadTokens: 1000,
|
||||
cacheWriteTokens: 500,
|
||||
},
|
||||
children: {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
},
|
||||
total: {
|
||||
inputTokens: 6000, // 5000 + 1000
|
||||
outputTokens: 3000, // 2500 + 500
|
||||
totalTokens: 9000, // 7500 + 1500
|
||||
},
|
||||
messageCount: 10,
|
||||
});
|
||||
});
|
||||
|
||||
it('会话不存在时返回 null', async () => {
|
||||
vi.mocked(SessionStorage.get).mockResolvedValue(null);
|
||||
|
||||
const stats = await manager.getSessionStats('project-1', 'nonexistent');
|
||||
|
||||
expect(stats).toBeNull();
|
||||
});
|
||||
|
||||
it('没有子会话统计时返回空的 children', async () => {
|
||||
vi.mocked(SessionStorage.get).mockResolvedValue({
|
||||
id: 'session-1',
|
||||
projectId: 'project-1',
|
||||
createdAt: Date.now(),
|
||||
updatedAt: Date.now(),
|
||||
workdir: '/test',
|
||||
stats: {
|
||||
messageCount: 5,
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
updatedAt: Date.now(),
|
||||
},
|
||||
});
|
||||
|
||||
const stats = await manager.getSessionStats('project-1', 'session-1');
|
||||
|
||||
expect(stats?.children).toEqual({
|
||||
inputTokens: 0,
|
||||
outputTokens: 0,
|
||||
totalTokens: 0,
|
||||
});
|
||||
expect(stats?.total).toEqual({
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatTokens', () => {
|
||||
it('格式化小数字', () => {
|
||||
expect(manager.formatTokens(100)).toBe('100');
|
||||
expect(manager.formatTokens(999)).toBe('999');
|
||||
});
|
||||
|
||||
it('格式化千级数字', () => {
|
||||
expect(manager.formatTokens(1000)).toBe('1.0k');
|
||||
expect(manager.formatTokens(1500)).toBe('1.5k');
|
||||
expect(manager.formatTokens(9999)).toBe('10.0k');
|
||||
});
|
||||
|
||||
it('格式化百万级数字', () => {
|
||||
expect(manager.formatTokens(1000000)).toBe('1.00M');
|
||||
expect(manager.formatTokens(1500000)).toBe('1.50M');
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatSessionStats', () => {
|
||||
it('格式化会话统计为人类可读字符串', () => {
|
||||
const stats = {
|
||||
self: {
|
||||
inputTokens: 5000,
|
||||
outputTokens: 2500,
|
||||
totalTokens: 7500,
|
||||
},
|
||||
children: {
|
||||
inputTokens: 1000,
|
||||
outputTokens: 500,
|
||||
totalTokens: 1500,
|
||||
},
|
||||
total: {
|
||||
inputTokens: 6000,
|
||||
outputTokens: 3000,
|
||||
totalTokens: 9000,
|
||||
},
|
||||
messageCount: 10,
|
||||
updatedAt: Date.now(),
|
||||
};
|
||||
|
||||
const formatted = manager.formatSessionStats(stats);
|
||||
|
||||
expect(formatted).toContain('6.0k'); // total input
|
||||
expect(formatted).toContain('3.0k'); // total output
|
||||
expect(formatted).toContain('9.0k'); // total
|
||||
expect(formatted).toContain('1.5k'); // children
|
||||
expect(formatted).toContain('10'); // message count
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user