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

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

      // 小对话不压缩时返回 'none'
      expect(['prune', 'compaction', 'both', 'none']).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);
  });
});