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

import { describe, it, expect } from 'vitest';
import { prune, filterCompacted } from '../../../src/context/prune.js';
import {
  COMPACTED_PLACEHOLDER,
  SUMMARY_MARKER,
  COMPACTED_MARKER,
  type CompressionConfig,
} from '../../../src/context/types.js';
import type { ModelMessage } from 'ai';

// 创建测试用的工具结果消息
function createToolResultMessage(
  toolCallId: string,
  result: unknown,
  compacted = false
): ModelMessage {
  const content: Record<string, unknown> = {
    type: 'tool-result',
    toolCallId,
    toolName: 'test_tool',
    result,
  };

  if (compacted) {
    content[COMPACTED_MARKER] = {
      compactedAt: Date.now(),
      originalSize: 100,
    };
  }

  return {
    role: 'tool',
    content: [content],
  } as ModelMessage;
}

// 创建测试用的工具调用消息
function createToolCallMessage(toolCallId: string): ModelMessage {
  return {
    role: 'assistant',
    content: [
      {
        type: 'tool-call',
        toolCallId,
        toolName: 'test_tool',
        args: { param: 'value' },
      },
    ],
  } as ModelMessage;
}

// 创建用户消息
function createUserMessage(content: string): ModelMessage {
  return { role: 'user', content };
}

// 创建助手消息
function createAssistantMessage(content: string): ModelMessage {
  return { role: 'assistant', content };
}

// 创建摘要消息
function createSummaryMessage(): ModelMessage {
  return {
    role: 'assistant',
    content: `${SUMMARY_MARKER}\n## 对话摘要\n\n这是一个摘要\n${SUMMARY_MARKER}`,
  };
}

// 创建大工具结果（指定 token 数）
function createLargeToolResult(toolCallId: string, sizeInChars: number): ModelMessage {
  // 约 4 字符/token（英文）
  const content = 'a'.repeat(sizeInChars);
  return createToolResultMessage(toolCallId, { output: content });
}

describe('prune - 消息裁剪策略', () => {
  // 测试配置：小值便于测试
  const testConfig: CompressionConfig = {
    contextLimit: 1000,
    outputReserve: 100,
    pruneProtect: 200, // 保护最近 200 tokens
    pruneMinimum: 50, // 至少释放 50 tokens 才执行
    overflowThreshold: 0.85,
  };

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

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

    it('只有用户消息不裁剪', () => {
      const messages: ModelMessage[] = [
        createUserMessage('Hello'),
        createUserMessage('How are you?'),
      ];

      const result = prune(messages, testConfig);

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

    it('保护范围内的工具结果不裁剪', () => {
      // 创建小的工具结果（在保护范围内）
      const messages: ModelMessage[] = [
        createUserMessage('Use tool'),
        createToolCallMessage('call_1'),
        createToolResultMessage('call_1', { output: 'small result' }),
      ];

      const result = prune(messages, testConfig);

      // 不应该裁剪
      expect(result.freedTokens).toBe(0);
    });

    it('裁剪超出保护范围的工具结果', () => {
      // 创建多个工具结果，超出保护范围
      const messages: ModelMessage[] = [
        createUserMessage('Task 1'),
        createToolCallMessage('call_1'),
        createLargeToolResult('call_1', 1000), // ~250 tokens, 超出 pruneProtect
        createUserMessage('Task 2'),
        createToolCallMessage('call_2'),
        createLargeToolResult('call_2', 400), // ~100 tokens
        createUserMessage('Task 3'),
        createToolCallMessage('call_3'),
        createToolResultMessage('call_3', { output: 'recent' }), // 最近的，在保护范围内
      ];

      const result = prune(messages, testConfig);

      // 应该裁剪旧的大工具结果
      expect(result.freedTokens).toBeGreaterThan(0);
    });
  });

  describe('摘要消息边界', () => {
    it('遇到摘要消息停止裁剪', () => {
      const messages: ModelMessage[] = [
        createSummaryMessage(), // 摘要消息
        createUserMessage('After summary'),
        createToolCallMessage('call_1'),
        createLargeToolResult('call_1', 2000), // 大结果
      ];

      const result = prune(messages, testConfig);

      // 因为遇到摘要消息，不会继续向前裁剪
      // 但是摘要后面的大工具结果如果超出保护范围仍会被裁剪
      expect(result.messages[0]).toEqual(messages[0]); // 摘要保留
    });

    it('摘要消息前的内容不裁剪', () => {
      const messages: ModelMessage[] = [
        createUserMessage('Before summary'),
        createToolCallMessage('call_old'),
        createLargeToolResult('call_old', 2000), // 摘要前的大结果
        createSummaryMessage(),
        createUserMessage('After summary'),
      ];

      const result = prune(messages, testConfig);

      // 摘要前的内容因为遇到摘要边界停止
      expect(result.messages.length).toBe(messages.length);
    });
  });

  describe('已压缩内容处理', () => {
    it('遇到已压缩的工具结果停止', () => {
      const messages: ModelMessage[] = [
        createUserMessage('Task 1'),
        createToolCallMessage('call_1'),
        createToolResultMessage('call_1', COMPACTED_PLACEHOLDER, true), // 已压缩
        createUserMessage('Task 2'),
        createToolCallMessage('call_2'),
        createLargeToolResult('call_2', 1000), // 新的大结果
      ];

      const result = prune(messages, testConfig);

      // 遇到已压缩的结果应该停止继续向前
      expect(result.messages).toBeDefined();
    });
  });

  describe('最小裁剪量检查', () => {
    it('释放量不足最小值时不执行裁剪', () => {
      // 使用较大的 pruneMinimum
      const strictConfig: CompressionConfig = {
        ...testConfig,
        pruneMinimum: 10000, // 要求至少释放 10000 tokens
      };

      const messages: ModelMessage[] = [
        createUserMessage('Task'),
        createToolCallMessage('call_1'),
        createLargeToolResult('call_1', 400), // 只有 ~100 tokens
      ];

      const result = prune(messages, strictConfig);

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

  describe('深拷贝验证', () => {
    it('不修改原消息数组', () => {
      const messages: ModelMessage[] = [
        createUserMessage('Task'),
        createToolCallMessage('call_1'),
        createLargeToolResult('call_1', 2000),
        createUserMessage('Recent'),
      ];

      const originalLength = messages.length;
      const originalFirstMessage = { ...messages[0] };

      prune(messages, testConfig);

      expect(messages.length).toBe(originalLength);
      expect(messages[0]).toEqual(originalFirstMessage);
    });
  });
});

describe('filterCompacted - 过滤已压缩内容', () => {
  it('不修改普通消息', () => {
    const messages: ModelMessage[] = [
      createUserMessage('Hello'),
      createAssistantMessage('Hi there'),
    ];

    const result = filterCompacted(messages);

    expect(result).toEqual(messages);
  });

  it('不修改未压缩的工具结果', () => {
    const messages: ModelMessage[] = [
      createToolResultMessage('call_1', { output: 'normal result' }),
    ];

    const result = filterCompacted(messages);

    expect(result).toEqual(messages);
  });

  it('将已压缩工具结果替换为占位符', () => {
    const messages: ModelMessage[] = [
      createToolResultMessage('call_1', 'original', true), // 已压缩
    ];

    const result = filterCompacted(messages);

    const content = result[0].content as { result: unknown }[];
    expect(content[0].result).toBe(COMPACTED_PLACEHOLDER);
  });

  it('混合内容正确处理', () => {
    const messages: ModelMessage[] = [
      createUserMessage('Task'),
      createToolResultMessage('call_1', 'original', true), // 已压缩
      createToolResultMessage('call_2', { output: 'normal' }), // 未压缩
      createAssistantMessage('Done'),
    ];

    const result = filterCompacted(messages);

    expect(result).toHaveLength(4);
    // 第一个工具结果应该被替换
    const compactedContent = result[1].content as { result: unknown }[];
    expect(compactedContent[0].result).toBe(COMPACTED_PLACEHOLDER);
    // 第二个工具结果应该保持不变
    const normalContent = result[2].content as { result: unknown }[];
    expect(normalContent[0].result).toEqual({ output: 'normal' });
  });

  it('保留消息的其他属性', () => {
    const messages: ModelMessage[] = [
      {
        role: 'tool',
        content: [
          {
            type: 'tool-result',
            toolCallId: 'call_1',
            toolName: 'my_tool',
            result: 'data',
            [COMPACTED_MARKER]: { compactedAt: 123, originalSize: 100 },
          },
        ],
      } as ModelMessage,
    ];

    const result = filterCompacted(messages);

    const content = result[0].content as { toolCallId: string; toolName: string }[];
    expect(content[0].toolCallId).toBe('call_1');
    expect(content[0].toolName).toBe('my_tool');
  });
});