ai-terminal-assistant/tests/unit/context/manager-extended.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import type { ModelMessage, LanguageModel } from 'ai';

// Mock prune module
const mockPrune = vi.fn();
const mockFilterCompacted = vi.fn();
vi.mock('../../../src/context/prune.js', () => ({
  prune: (...args: unknown[]) => mockPrune(...args),
  filterCompacted: (...args: unknown[]) => mockFilterCompacted(...args),
}));

// Mock compaction module
const mockCompact = vi.fn();
const mockSimpleCompact = vi.fn();
const mockIsSummaryMessage = vi.fn();
vi.mock('../../../src/context/compaction.js', () => ({
  compact: (...args: unknown[]) => mockCompact(...args),
  simpleCompact: (...args: unknown[]) => mockSimpleCompact(...args),
  isSummaryMessage: (...args: unknown[]) => mockIsSummaryMessage(...args),
}));

// Mock token counter
const mockEstimateMessages = vi.fn();
const mockFormat = vi.fn();
vi.mock('../../../src/context/token-counter.js', () => ({
  TokenCounter: {
    estimateMessages: (...args: unknown[]) => mockEstimateMessages(...args),
    format: (...args: unknown[]) => mockFormat(...args),
  },
}));

import { CompressionManager } from '../../../src/context/manager.js';

describe('CompressionManager - 压缩管理器扩展测试', () => {
  let manager: CompressionManager;

  const createMessages = (count: number): ModelMessage[] => {
    return Array.from({ length: count }, (_, i) => ({
      role: 'user' as const,
      content: `Message ${i}`,
    }));
  };

  beforeEach(() => {
    vi.clearAllMocks();
    manager = new CompressionManager();

    // 默认 mock 返回值
    mockEstimateMessages.mockReturnValue(1000);
    mockFormat.mockReturnValue('1K');
    mockPrune.mockReturnValue({ messages: [], freedTokens: 0 });
    mockFilterCompacted.mockImplementation((msgs) => msgs);
    mockCompact.mockResolvedValue({ messages: [], freedTokens: 0 });
    mockSimpleCompact.mockReturnValue({ messages: [], freedTokens: 0 });
    mockIsSummaryMessage.mockReturnValue(false);
  });

  describe('constructor - 构造函数', () => {
    it('使用默认配置', () => {
      const config = manager.getConfig();

      expect(config.contextLimit).toBeDefined();
      expect(config.outputReserve).toBeDefined();
      expect(config.overflowThreshold).toBeDefined();
    });

    it('合并自定义配置', () => {
      const customManager = new CompressionManager({
        contextLimit: 100000,
        outputReserve: 5000,
      });

      const config = customManager.getConfig();
      expect(config.contextLimit).toBe(100000);
      expect(config.outputReserve).toBe(5000);
    });
  });

  describe('setModel - 设置模型', () => {
    it('设置模型用于摘要生成', () => {
      const mockModel = {} as LanguageModel;
      manager.setModel(mockModel);

      // 模型被设置后 compact 应该使用它
      expect(manager['model']).toBe(mockModel);
    });
  });

  describe('updateConfig - 更新配置', () => {
    it('更新部分配置', () => {
      const originalConfig = manager.getConfig();

      manager.updateConfig({ contextLimit: 200000 });

      const updatedConfig = manager.getConfig();
      expect(updatedConfig.contextLimit).toBe(200000);
      expect(updatedConfig.outputReserve).toBe(originalConfig.outputReserve);
    });
  });

  describe('calculateUsage - 计算使用量', () => {
    it('计算 token 使用情况', () => {
      mockEstimateMessages.mockReturnValue(50000);

      const messages = createMessages(10);
      const usage = manager.calculateUsage(messages);

      expect(usage.input).toBe(50000);
      expect(usage.contextLimit).toBeDefined();
      expect(usage.available).toBeDefined();
      expect(usage.usagePercent).toBeGreaterThanOrEqual(0);
    });

    it('使用百分比不超过 100', () => {
      // 模拟超出限制
      mockEstimateMessages.mockReturnValue(300000);

      const usage = manager.calculateUsage([]);

      expect(usage.usagePercent).toBeLessThanOrEqual(100);
    });
  });

  describe('shouldCompress - 是否需要压缩', () => {
    it('未超过阈值返回 false', () => {
      mockEstimateMessages.mockReturnValue(10000);

      const result = manager.shouldCompress([]);

      expect(result).toBe(false);
    });

    it('超过阈值返回 true', () => {
      // 设置接近阈值的使用量
      mockEstimateMessages.mockReturnValue(150000);

      const result = manager.shouldCompress([]);

      expect(result).toBe(true);
    });
  });

  describe('isOverflow - 是否溢出', () => {
    it('未溢出返回 false', () => {
      mockEstimateMessages.mockReturnValue(10000);

      const result = manager.isOverflow([]);

      expect(result).toBe(false);
    });

    it('溢出返回 true', () => {
      mockEstimateMessages.mockReturnValue(500000);

      const result = manager.isOverflow([]);

      expect(result).toBe(true);
    });
  });

  describe('prune - 执行裁剪', () => {
    it('调用 prune 函数', () => {
      const messages = createMessages(5);
      mockPrune.mockReturnValue({ messages: messages.slice(2), freedTokens: 1000 });

      const result = manager.prune(messages);

      expect(mockPrune).toHaveBeenCalledWith(messages, expect.any(Object));
      expect(result.freedTokens).toBe(1000);
    });
  });

  describe('compact - 执行压缩', () => {
    it('有模型时使用 AI 压缩', async () => {
      const mockModel = {} as LanguageModel;
      manager.setModel(mockModel);
      mockCompact.mockResolvedValue({ messages: [], freedTokens: 2000 });

      const messages = createMessages(5);
      const result = await manager.compact(messages);

      expect(mockCompact).toHaveBeenCalled();
      expect(result.freedTokens).toBe(2000);
    });

    it('无模型时使用简单压缩', async () => {
      mockSimpleCompact.mockReturnValue({ messages: [], freedTokens: 500 });

      const messages = createMessages(5);
      const result = await manager.compact(messages);

      expect(mockSimpleCompact).toHaveBeenCalled();
      expect(result.freedTokens).toBe(500);
    });
  });

  describe('compress - 自动压缩', () => {
    it('先 prune 后不需要 compact', async () => {
      mockEstimateMessages.mockReturnValue(10000); // 低于阈值
      mockPrune.mockReturnValue({ messages: [], freedTokens: 500 });

      const messages = createMessages(5);
      const result = await manager.compress(messages);

      expect(result.type).toBe('prune');
      expect(mockPrune).toHaveBeenCalled();
      expect(mockCompact).not.toHaveBeenCalled();
    });

    it('prune 后仍需 compact', async () => {
      mockEstimateMessages.mockReturnValue(150000); // 高于阈值
      mockPrune.mockReturnValue({ messages: createMessages(3), freedTokens: 500 });
      mockSimpleCompact.mockReturnValue({ messages: [], freedTokens: 1000 });

      const messages = createMessages(5);
      const result = await manager.compress(messages);

      expect(result.type).toBe('both');
      expect(mockPrune).toHaveBeenCalled();
    });

    it('只 compact 时类型为 compaction', async () => {
      mockEstimateMessages.mockReturnValue(150000);
      mockPrune.mockReturnValue({ messages: createMessages(5), freedTokens: 0 });
      mockSimpleCompact.mockReturnValue({ messages: [], freedTokens: 1000 });

      const messages = createMessages(5);
      const result = await manager.compress(messages);

      expect(result.type).toBe('compaction');
    });
  });

  describe('forceCompress - 强制压缩', () => {
    it('消息太少不压缩', async () => {
      const messages = createMessages(3);
      const result = await manager.forceCompress(messages);

      expect(result.messages).toEqual(messages);
      expect(result.freedTokens).toBe(0);
    });

    it('足够消息时执行压缩', async () => {
      mockEstimateMessages.mockReturnValue(10000);
      mockPrune.mockReturnValue({ messages: createMessages(3), freedTokens: 500 });
      mockSimpleCompact.mockReturnValue({ messages: createMessages(2), freedTokens: 300 });

      const messages = createMessages(10);
      const result = await manager.forceCompress(messages);

      expect(mockPrune).toHaveBeenCalled();
    });

    it('有模型时尝试 AI 压缩', async () => {
      const mockModel = {} as LanguageModel;
      manager.setModel(mockModel);
      mockEstimateMessages.mockReturnValue(10000);
      mockPrune.mockReturnValue({ messages: createMessages(5), freedTokens: 500 });
      mockCompact.mockResolvedValue({ messages: createMessages(2), freedTokens: 1000 });

      const messages = createMessages(10);
      await manager.forceCompress(messages);

      expect(mockCompact).toHaveBeenCalled();
    });

    it('AI 压缩失败时回退到简单压缩', async () => {
      const mockModel = {} as LanguageModel;
      manager.setModel(mockModel);
      mockEstimateMessages.mockReturnValue(10000);
      mockPrune.mockReturnValue({ messages: createMessages(5), freedTokens: 500 });
      mockCompact.mockRejectedValue(new Error('AI error'));
      mockSimpleCompact.mockReturnValue({ messages: createMessages(2), freedTokens: 800 });

      const messages = createMessages(10);
      const result = await manager.forceCompress(messages);

      expect(mockSimpleCompact).toHaveBeenCalled();
    });
  });

  describe('filterCompacted - 过滤压缩内容', () => {
    it('调用 filterCompacted 函数', () => {
      const messages = createMessages(5);
      mockFilterCompacted.mockReturnValue(messages.slice(1));

      const result = manager.filterCompacted(messages);

      expect(mockFilterCompacted).toHaveBeenCalledWith(messages);
    });
  });

  describe('isSummaryMessage - 检查摘要消息', () => {
    it('调用 isSummaryMessage 函数', () => {
      mockIsSummaryMessage.mockReturnValue(true);

      const message: ModelMessage = { role: 'assistant', content: 'summary' };
      const result = manager.isSummaryMessage(message);

      expect(mockIsSummaryMessage).toHaveBeenCalledWith(message);
      expect(result).toBe(true);
    });
  });

  describe('formatUsage - 格式化使用情况', () => {
    it('返回格式化字符串', () => {
      mockEstimateMessages.mockReturnValue(50000);
      mockFormat.mockImplementation((n) => `${Math.round(n / 1000)}K`);

      const messages = createMessages(5);
      const result = manager.formatUsage(messages);

      expect(result).toMatch(/\d+K\/\d+K/);
      expect(result).toContain('%');
    });
  });
});