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 }); }); });