swarm-assistent/test/context.test.js

import { describe, it } from 'node:test';
import assert from 'node:assert/strict';
import {
  emptyContextMemory,
  normalizeContextMemory,
  charsToTokens,
  estimateBudget,
  shouldCompress,
  assembleModelMessages,
  messagesToFold,
  mergeSummary,
  conversationMemoryBlock,
} from '../src/context.js';

describe('context.js', () => {
  it('charsToTokens ceil divide', () => {
    assert.equal(charsToTokens(32, 3.2), 10);
    assert.equal(charsToTokens(0, 3.2), 0);
  });

  it('estimateBudget prefers prompt_eval_count', () => {
    const est = estimateBudget({
      systemChars: 3200,
      historyChars: 3200,
      memoryChars: 0,
      numCtx: 16384,
      numPredict: 3072,
      charsPerToken: 3.2,
      compressAt: 0.7,
    });
    assert.equal(est.fromEval, false);
    assert.ok(est.estimated > 0);
    assert.equal(est.used, est.estimated);

    const fact = estimateBudget({
      ...est,
      systemChars: 3200,
      historyChars: 3200,
      promptEvalCount: 14000,
      numCtx: 16384,
      numPredict: 3072,
      compressAt: 0.7,
    });
    assert.equal(fact.fromEval, true);
    assert.equal(fact.used, 14000);
    assert.equal(fact.level, 'hot');
  });

  it('shouldCompress only when over threshold and uncovered > keep', () => {
    const budget = estimateBudget({
      systemChars: 20000,
      historyChars: 20000,
      numCtx: 16384,
      numPredict: 3072,
      compressAt: 0.7,
      charsPerToken: 3.2,
    });
    assert.equal(shouldCompress(budget, emptyContextMemory(), 20, { keepMessages: 8 }), true);
    assert.equal(shouldCompress(budget, emptyContextMemory(), 6, { keepMessages: 8 }), false);
    assert.equal(
      shouldCompress(budget, { untilCount: 12, summary: 'x' }, 20, { keepMessages: 8 }),
      false,
    );
  });

  it('assembleModelMessages skips covered prefix then keeps last window', () => {
    const history = [];
    for (let i = 0; i < 12; i++) {
      history.push({ role: i % 2 === 0 ? 'user' : 'assistant', content: `m${i}` });
    }
    const msgs = assembleModelMessages(history, { untilCount: 4 }, 2);
    assert.equal(msgs.length, 4);
    assert.equal(msgs[0].content, 'm8');
    assert.equal(msgs[3].content, 'm11');
  });

  it('messagesToFold leaves keepTurns raw', () => {
    const history = Array.from({ length: 10 }, (_, i) => ({
      role: i % 2 === 0 ? 'user' : 'assistant',
      content: `m${i}`,
    }));
    const fold = messagesToFold(history, emptyContextMemory(), 2);
    assert.equal(fold.length, 6);
    assert.equal(fold[0].content, 'm0');
    assert.equal(fold[5].content, 'm5');
  });

  it('normalize + merge + conversation block', () => {
    const m = normalizeContextMemory({ summary: '  hello  ', untilCount: 4, uiCollapsed: true });
    assert.equal(m.summary, 'hello');
    assert.equal(m.uiCollapsed, true);
    assert.equal(mergeSummary('old', 'new'), 'new');
    assert.equal(mergeSummary('old', ''), 'old');
    const block = conversationMemoryBlock(m, 100);
    assert.equal(block.summary, 'hello');
    assert.equal(block.until_count, 4);
    assert.equal(conversationMemoryBlock(emptyContextMemory()), null);
  });
});