ai-terminal-assistant/packages/core/tests/unit/session/token-stats-manager.test.ts

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