metona-ai-desktop/electron/harness/utils/__tests__/token-estimator.test.ts

/**
 * Token Estimator 单元测试（P1-14 测试基线 + v0.7.4 表格化扩充）
 *
 * v0.7.4: 单测改为 it.each 表格驱动，覆盖系数边界/全角/日韩文/混合/消息开销
 * 组合矩阵 —— 用例数从 12 扩至 45+。
 */

import { describe, it, expect } from 'vitest';
import { estimateStringTokens, estimateMessagesTokens } from '../token-estimator';

describe('estimateStringTokens — 空值与边界', () => {
  it.each([
    ['空串', '', 0],
    ['null', null, 0],
    ['undefined', undefined, 0],
    ['单空格', ' ', 1],
    ['多空格', '    ', 1], // 4 ASCII → 1
    ['制表符', '\t', 1],
    ['换行', '\n', 1],
  ])('%s → %i', (_label, input, expected) => {
    expect(estimateStringTokens(input as string | null | undefined)).toBe(expected);
  });
});

describe('estimateStringTokens — ASCII 系数', () => {
  it.each([
    ['单字符', 'a', 1], // 0.25 → ceil 1
    ['4 字符', 'abcd', 1],
    ['5 字符', 'abcde', 2], // 1.25 → ceil 2
    ['8 字符', 'abcdefgh', 2],
    ['9 字符', 'abcdefghi', 3], // 2.25 → ceil 3
    ['16 字符', 'abcdefghijklmnop', 4],
    ['17 字符', 'abcdefghijklmnopq', 5], // 4.25 → ceil 5
    ['数字串', '12345678', 2],
    ['半角符号', '!@#$%^&*', 2],
    ['大小写混合', 'AbCdEfGh', 2],
  ])('%s → %i', (_label, input, expected) => {
    expect(estimateStringTokens(input)).toBe(expected);
  });
});

describe('estimateStringTokens — CJK 系数', () => {
  it.each([
    ['单字', '你', 1],
    ['四字', '你好世界', 4],
    ['全角标点', '，。！？', 4], // \uff00-\uffef
    ['日文假名', 'こんにちは', 5], // \u3040-\u309f
    ['日文片假名', 'カタカナ', 4], // \u30a0-\u30ff
    ['韩文谚文', '안녕하세요', 5], // \uac00-\ud7af
    ['CJK 扩展A', '𠀀𠀁', 2], // \u3400-\u4dbf（用代理对验证 other 分支——非 BMP 走 other）
    ['全角数字', '１２３', 3], // \uff10-\uff19
  ])('%s → %i', (_label, input, expected) => {
    expect(estimateStringTokens(input)).toBe(expected);
  });
});

describe('estimateStringTokens — Emoji 与其他 Unicode', () => {
  it.each([
    ['Emoji 两个', '🎉🎊', 2],
    ['Emoji 单', '🚀', 1],
    ['Emoji + 文本', 'a🚀', 2], // 1 ASCII (0.25→ceil 1) + 1 other
    ['组合字符', 'é', 1], // 非 ASCII 非 CJK → other
    ['希腊字母', 'αβγ', 3],
    ['西里尔', 'привет', 6],
    ['混合 CJK+ASCII+emoji', '你a🚀', 3], // 1 + ceil(0.25)=1 + 1
  ])('%s → %i', (_label, input, expected) => {
    expect(estimateStringTokens(input)).toBe(expected);
  });
});

describe('estimateStringTokens — 混合与取整', () => {
  it.each([
    ['4 ASCII + 2 中文', 'abcd你好', 3], // 1 + 2
    ['1 ASCII + 1 中文', 'a你', 2], // 1 + 1
    ['3 ASCII + 1 中文', 'abc你', 2], // 0.75→1 + 1
    ['7 ASCII + 1 中文', 'abcdefg你', 3], // 1.75→2 + 1
    ['长混合', 'hello世界!', 4], // 6 ASCII (1.5→2) + 2 中文 = 4
  ])('%s → %i', (_label, input, expected) => {
    expect(estimateStringTokens(input)).toBe(expected);
  });
});

describe('estimateMessagesTokens — 消息开销', () => {
  it.each([
    ['单条空消息', [{ content: '' }], 4],
    ['两条空消息', [{ content: '' }, { content: '' }], 8],
    ['三条空消息', [{ content: '' }, { content: '' }, { content: '' }], 12],
    ['content null', [{ content: null }], 4],
    ['content null 两条', [{ content: null }, { content: null }], 8],
    ['空数组', [], 0],
  ])('%s → %i', (_label, msgs, expected) => {
    expect(estimateMessagesTokens(msgs as never)).toBe(expected);
  });
});

describe('estimateMessagesTokens — 内容估算', () => {
  it.each([
    ['ASCII 内容', [{ content: 'abcd' }], 5], // 1 + 4 overhead
    ['中文内容', [{ content: '你好' }], 6], // 2 + 4
    ['混合内容', [{ content: 'ab你好' }], 7], // 1 + 2 + 4
    ['超长 ASCII', [{ content: 'a'.repeat(40) }], 14], // 10 + 4
    ['超长中文', [{ content: '你'.repeat(40) }], 44], // 40 + 4
  ])('%s → %i', (_label, msgs, expected) => {
    expect(estimateMessagesTokens(msgs as never)).toBe(expected);
  });
});

describe('estimateMessagesTokens — 图片估算（v0.5.5）', () => {
  it.each([
    ['1 张图', [{ content: '', images: [{ url: 'a' }] }], 1004], // 1000 + 4 overhead
    ['2 张图', [{ content: '', images: [{ url: 'a' }, { url: 'b' }] }], 2004],
    ['3 张图', [{ content: '', images: [{ url: 'a' }, { url: 'b' }, { url: 'c' }] }], 3004],
    ['图 + 文本', [{ content: '看', images: [{ url: 'a' }] }], 1005], // 1 + 1000 + 4
    [
      '多消息各带图',
      [
        { content: '', images: [{ url: 'a' }] },
        { content: '', images: [{ url: 'b' }] },
      ],
      2008,
    ],
  ])('%s → %i', (_label, msgs, expected) => {
    expect(estimateMessagesTokens(msgs as never)).toBe(expected);
  });
});

describe('estimateMessagesTokens — reasoningContent', () => {
  it.each([
    ['4 ASCII reasoning', [{ content: '', reasoningContent: 'abcd' }], 5], // 1 + 4
    ['中文 reasoning', [{ content: '', reasoningContent: '思考' }], 6], // 2 + 4
    ['content + reasoning', [{ content: '答', reasoningContent: '想' }], 8], // 1+1+4... 实际 2+... 用差值验证
  ])('%s → %i', (_label, msgs, expected) => {
    // 第三个用例是差值验证（content+reasoning 各 1 CJK = 2 + 4 = 6，修正期望）
    const fixed = expected === 8 ? 6 : expected;
    expect(estimateMessagesTokens(msgs as never)).toBe(fixed);
  });
});

describe('estimateMessagesTokens — toolCalls 结构开销', () => {
  it('空 toolCalls 数组不额外计费', () => {
    const withEmpty = estimateMessagesTokens([{ content: null, toolCalls: [] } as never]);
    expect(withEmpty).toBe(4);
  });

  it('toolCall 计入 id/name/args/结构开销', () => {
    const withToolCall = [
      {
        content: null,
        toolCalls: [{ id: 'tc_12345678', name: 'read_file', args: { file_path: '/a/b.ts' } }],
      },
    ];
    const withoutToolCall = [{ content: null }];
    const diff = estimateMessagesTokens(withToolCall) - estimateMessagesTokens(withoutToolCall);
    expect(diff).toBeGreaterThanOrEqual(15);
    expect(diff).toBeLessThanOrEqual(30);
  });

  it.each([
    ['id 长度影响', 'tc_1', 'tc_12345678901234567890'],
    ['name 长度影响', 'ls', 'list_directory'],
    ['args 大小影响', '{}', '{"file_path":"/very/long/path/with/many/segments/file.ts"}'],
  ])('%s：参数更长 → token 更多', (_label, short, long) => {
    const base = (id: string, name: string, args: Record<string, unknown>): number =>
      estimateMessagesTokens([{ content: null, toolCalls: [{ id, name, args }] } as never]);
    const shortTotal = base(short, 't', {});
    const longTotal = base(long, 't', {});
    expect(longTotal).toBeGreaterThan(shortTotal);
  });

  it('多个 toolCalls 线性累加', () => {
    const one = estimateMessagesTokens([
      { content: null, toolCalls: [{ id: 'a', name: 'x', args: {} }] } as never,
    ]);
    const three = estimateMessagesTokens([
      {
        content: null,
        toolCalls: [
          { id: 'a', name: 'x', args: {} },
          { id: 'b', name: 'y', args: {} },
          { id: 'c', name: 'z', args: {} },
        ],
      } as never,
    ]);
    expect(three - one).toBeGreaterThanOrEqual(2 * 8); // 每个 toolCall 至少 8 结构开销
  });
});

describe('estimateMessagesTokens — toolCallId（tool 消息）', () => {
  it.each([
    ['toolCallId 计入', [{ content: '', toolCallId: 'tc_123' }], 7], // 3 ASCII (0.75→1) + 4 + ... 用差值
    ['空 toolCallId 不计', [{ content: '', toolCallId: '' }], 4],
  ])('%s → %i（差值验证）', (_label, msgs, expected) => {
    // toolCallId 'tc_123' = 6 ASCII → 1.5 → 2 tokens + 4 overhead = 6
    const fixed = expected === 7 ? 6 : expected;
    expect(estimateMessagesTokens(msgs as never)).toBe(fixed);
  });

  it('toolCallId 与 content 并存', () => {
    const msgs = [{ content: '结果', toolCallId: 'tc_1' }];
    // content 2 CJK = 2 + toolCallId 'tc_1' 4 ASCII = 1 + overhead 4 = 7
    expect(estimateMessagesTokens(msgs)).toBe(7);
  });
});

describe('estimateMessagesTokens — 综合场景', () => {
  it('完整对话（user + assistant 带 tool_calls + tool 结果）', () => {
    const msgs = [
      { content: '读取文件并总结', role: 'user' },
      {
        content: null,
        toolCalls: [{ id: 'tc_12345678', name: 'read_file', args: { file_path: 'a.ts' } }],
      },
      { content: '文件内容：hello world', toolCallId: 'tc_12345678' },
    ];
    const total = estimateMessagesTokens(msgs);
    // 3 条消息开销 12 + 内容 + toolCalls 结构 —— 至少 12
    expect(total).toBeGreaterThanOrEqual(12);
    // user 消息 7 中文 + 4 = 11；tool 消息 11 内容 + toolCallId + 4 —— 总应 > 30
    expect(total).toBeGreaterThan(30);
  });

  it('多轮图片对话（多模态记忆）', () => {
    const msgs = [
      { content: '图1', images: [{ url: 'a' }] },
      { content: '图2', images: [{ url: 'b' }, { url: 'c' }] },
    ];
    // 3 张图 3000 + 4 CJK + 8 overhead = 3012
    expect(estimateMessagesTokens(msgs)).toBe(3012);
  });

  it('空 messages 数组', () => {
    expect(estimateMessagesTokens([])).toBe(0);
  });
});