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

import { describe, it, expect, vi } from 'vitest';
import { isSummaryMessage, simpleCompact } from '../../../src/context/compaction.js';
import {
  SUMMARY_MARKER,
  type CompressionConfig,
} from '../../../src/context/types.js';
import type { ModelMessage } 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 createMessageWithTextParts(texts: string[]): ModelMessage {
  return {
    role: 'assistant',
    content: texts.map((text) => ({ type: 'text', text })),
  } as ModelMessage;
}

describe('isSummaryMessage - 检测摘要消息', () => {
  it('字符串内容包含摘要标记返回 true', () => {
    const message = createSummaryMessage('这是摘要内容');

    expect(isSummaryMessage(message)).toBe(true);
  });

  it('字符串内容不包含摘要标记返回 false', () => {
    const message = createAssistantMessage('普通助手消息');

    expect(isSummaryMessage(message)).toBe(false);
  });

  it('数组内容包含摘要标记返回 true', () => {
    const message = createMessageWithTextParts([
      '一些文本',
      `${SUMMARY_MARKER}\n摘要内容\n${SUMMARY_MARKER}`,
    ]);

    expect(isSummaryMessage(message)).toBe(true);
  });

  it('数组内容不包含摘要标记返回 false', () => {
    const message = createMessageWithTextParts(['文本1', '文本2']);

    expect(isSummaryMessage(message)).toBe(false);
  });

  it('用户消息不是摘要消息', () => {
    const message = createUserMessage(SUMMARY_MARKER);

    // 虽然包含标记，但这种情况下 isSummaryMessage 还是会返回 true
    // 因为它只检查内容是否包含标记
    expect(isSummaryMessage(message)).toBe(true);
  });

  it('空数组内容返回 false', () => {
    const message: ModelMessage = {
      role: 'assistant',
      content: [],
    };

    expect(isSummaryMessage(message)).toBe(false);
  });

  it('非文本部分不匹配', () => {
    const message: ModelMessage = {
      role: 'assistant',
      content: [
        {
          type: 'tool-call',
          toolCallId: 'call_1',
          toolName: 'test',
          args: { key: SUMMARY_MARKER },
        },
      ],
    } as ModelMessage;

    expect(isSummaryMessage(message)).toBe(false);
  });
});

describe('simpleCompact - 简单压缩', () => {
  const testConfig: CompressionConfig = {
    contextLimit: 1000,
    outputReserve: 100,
    pruneProtect: 200, // 保护最近 200 tokens
    pruneMinimum: 50,
    overflowThreshold: 0.85,
  };

  describe('基本压缩行为', () => {
    it('空消息数组不压缩', () => {
      const result = simpleCompact([], testConfig);

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

    it('消息在保护范围内不压缩', () => {
      const messages: ModelMessage[] = [
        createUserMessage('Hello'),
        createAssistantMessage('Hi'),
      ];

      const result = simpleCompact(messages, testConfig);

      // 消息很短，应该在保护范围内
      expect(result.freedTokens).toBe(0);
      expect(result.messages).toEqual(messages);
    });

    it('压缩超出保护范围的消息', () => {
      // 创建大量消息超出保护范围
      const messages: ModelMessage[] = [];
      for (let i = 0; i < 50; i++) {
        messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`));
        messages.push(createAssistantMessage(`Response ${i}: ${'b'.repeat(100)}`));
      }

      const result = simpleCompact(messages, testConfig);

      // 应该压缩一些旧消息
      expect(result.freedTokens).toBeGreaterThan(0);
      expect(result.messages.length).toBeLessThan(messages.length);
    });
  });

  describe('保留消息数量', () => {
    it('至少保留 2 条消息（正常模式）', () => {
      const messages: ModelMessage[] = [];
      for (let i = 0; i < 10; i++) {
        messages.push(createUserMessage(`Long message ${'a'.repeat(500)}`));
      }

      const result = simpleCompact(messages, testConfig);

      // 即使压缩，也至少保留 2 条
      expect(result.messages.length).toBeGreaterThanOrEqual(2 + 1); // 2 保留 + 1 摘要
    });

    it('强制模式下至少保留 1 条消息', () => {
      const forceConfig: CompressionConfig = {
        ...testConfig,
        pruneProtect: 0, // 强制模式
      };

      const messages: ModelMessage[] = [];
      for (let i = 0; i < 10; i++) {
        messages.push(createUserMessage(`Message ${i}`));
      }

      const result = simpleCompact(messages, forceConfig);

      // 强制模式下至少保留 1 条消息
      expect(result.messages.length).toBeGreaterThanOrEqual(2); // 1 保留 + 1 摘要
    });
  });

  describe('摘要消息生成', () => {
    it('压缩后生成摘要消息', () => {
      const messages: ModelMessage[] = [];
      for (let i = 0; i < 50; i++) {
        messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`));
      }

      // 禁用首条保护以便测试摘要消息在第一位
      const result = simpleCompact(messages, testConfig, { protectFirstPair: false });

      if (result.freedTokens > 0) {
        // 第一条消息应该是摘要
        expect(isSummaryMessage(result.messages[0])).toBe(true);
      }
    });

    it('摘要消息包含移除数量信息', () => {
      const messages: ModelMessage[] = [];
      for (let i = 0; i < 50; i++) {
        messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`));
      }

      // 禁用首条保护以便测试摘要消息
      const result = simpleCompact(messages, testConfig, { protectFirstPair: false });

      if (result.freedTokens > 0) {
        const summaryContent = result.messages[0].content as string;
        expect(summaryContent).toContain('对话历史已压缩');
        expect(summaryContent).toContain('条消息');
      }
    });
  });

  describe('保护范围计算', () => {
    it('短消息全部在保护范围内', () => {
      const messages: ModelMessage[] = [
        createUserMessage('Hi'),
        createAssistantMessage('Hello'),
        createUserMessage('How?'),
        createAssistantMessage('Good!'),
      ];

      const result = simpleCompact(messages, testConfig);

      // 短消息应该全部保留
      expect(result.messages).toEqual(messages);
    });
  });

  describe('不修改原数组', () => {
    it('原消息数组不被修改', () => {
      const messages: ModelMessage[] = [];
      for (let i = 0; i < 20; i++) {
        messages.push(createUserMessage(`Message ${i}: ${'a'.repeat(100)}`));
      }

      const originalLength = messages.length;
      simpleCompact(messages, testConfig);

      expect(messages.length).toBe(originalLength);
    });
  });
});

// 注意: compact 函数需要真实的 LanguageModel，
// 这里只测试 simpleCompact 和辅助函数
// compact 的测试应该在集成测试中进行

describe('Compaction 边界情况', () => {
  it('单条消息不压缩', () => {
    const messages: ModelMessage[] = [
      createUserMessage('Single message'),
    ];

    const testConfig: CompressionConfig = {
      contextLimit: 100,
      outputReserve: 10,
      pruneProtect: 50,
      pruneMinimum: 10,
      overflowThreshold: 0.85,
    };

    const result = simpleCompact(messages, testConfig);

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

  it('两条消息不压缩（最小保留）', () => {
    const messages: ModelMessage[] = [
      createUserMessage('First'),
      createAssistantMessage('Second'),
    ];

    const testConfig: CompressionConfig = {
      contextLimit: 100,
      outputReserve: 10,
      pruneProtect: 10, // 很小的保护范围
      pruneMinimum: 1,
      overflowThreshold: 0.85,
    };

    const result = simpleCompact(messages, testConfig);

    // 即使配置很激进，也至少保留 2 条
    expect(result.messages.length).toBeGreaterThanOrEqual(2);
  });
});