tests/integration/engine.test.ts

/**
 * Integration tests: run the real generation loop against a deterministic
 * mock model (no WebGPU in CI) and verify that each watermark's detector
 * statistically separates watermarked generations from baseline ones.
 * Multiple seeds are used - a single lucky example is not accepted as proof.
 */

import { describe, it, expect } from 'vitest';
import { runGenerationLoop, type LoopModel } from '../../src/lib/model/generation-engine';
import { applyGreenBias, KIRCHENBAUER_DEFAULTS } from '../../src/lib/watermark/kirchenbauer';
import { deriveKeys, gumbelMaxChoose, TEXTSEAL_DEFAULTS } from '../../src/lib/watermark/textseal';
import { detectKirchenbauer, type TokenizedText } from '../../src/lib/detectors/kirchenbauer-detector';
import { detectTextseal } from '../../src/lib/detectors/textseal-detector';
import { detectKsemstamp } from '../../src/lib/detectors/ksemstamp-detector';
import {
  acceptSentence,
  KSEMSTAMP_DEFAULTS,
  type Centroids,
} from '../../src/lib/watermark/ksemstamp';
import { RandomStream, mix64 } from '../../src/lib/utils/rng';
import { keyFromString } from '../../src/lib/utils/hashing';
import { splitSentences } from '../../src/lib/utils/sentences';
import { mean } from '../../src/lib/utils/stats';

const VOCAB = 256;
const PERIOD = 0; // token 0 decodes to "."

/**
 * Deterministic mock LM: logits depend only on (last context token, position),
 * so identical prefixes yield identical distributions - like a real LM.
 * The period token gets a strong boost every ~7th position to create
 * sentence structure.
 */
function makeMockModel(): LoopModel {
  return {
    vocabSize: VOCAB,
    eosTokenIds: [VOCAB - 1],
    decode(ids: number[]): string {
      let out = '';
      for (const id of ids) {
        if (id === PERIOD) out = out.trimEnd() + '. ';
        else out += `W${id} `;
      }
      return out.trim();
    },
    async forward(inputIds, fullSeqLen, past) {
      const lastToken = inputIds[inputIds.length - 1];
      const stream = new RandomStream(mix64((BigInt(lastToken) << 20n) ^ BigInt(fullSeqLen)));
      const logits = new Float32Array(VOCAB);
      for (let v = 0; v < VOCAB; v++) logits[v] = stream.next() * 4;
      if (fullSeqLen % 7 === 0) logits[PERIOD] += 6; // sentence boundary pressure
      else logits[PERIOD] -= 4;
      logits[VOCAB - 1] = -20; // no early EOS
      return { logits, past: (past as number ?? 0) + 1 };
    },
    disposePast() {
      /* no GPU tensors in the mock */
    },
  };
}

function asTokenized(_model: LoopModel, ids: number[]): TokenizedText {
  // Offsets don't matter for statistics; use index-based spans.
  return {
    tokenIds: ids,
    offsets: ids.map((_, i) => [i, i + 1] as [number, number]),
    labels: ids.map((t) => (t === PERIOD ? '.' : `W${t}`)),
  };
}

const PROMPT = [3, 14, 15, 92, 65];
const CFG = { temperature: 1.0, topP: 0.95, maxNewTokens: 80, topKTrace: 4 };
const SEEDS = [1, 2, 3, 4, 5];

describe('generation loop + Kirchenbauer end to end', () => {
  it('same seed reproduces the baseline run exactly', async () => {
    const model = makeMockModel();
    const a = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: 11 });
    const b = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: 11 });
    expect(a.tokenIds).toEqual(b.tokenIds);
  });

  it('watermarked mean z clearly exceeds baseline mean z across seeds', async () => {
    const model = makeMockModel();
    const params = KIRCHENBAUER_DEFAULTS;
    const zBase: number[] = [];
    const zWm: number[] = [];
    for (const seed of SEEDS) {
      const base = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: seed });
      zBase.push(detectKirchenbauer(asTokenized(model, base.tokenIds), params).statistic);

      const wm = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: seed }, {
        transformLogits: (_i, ctx, logits) => applyGreenBias(logits, ctx[ctx.length - 1], params),
      });
      zWm.push(detectKirchenbauer(asTokenized(model, wm.tokenIds), params).statistic);
    }
    expect(mean(zWm)).toBeGreaterThan(4);
    expect(mean(zBase)).toBeLessThan(2);
    expect(mean(zWm)).toBeGreaterThan(mean(zBase) + 3);
  });
});

describe('generation loop + TextSeal end to end', () => {
  it('watermarked p-values are tiny; baseline is not detected (across seeds)', async () => {
    const model = makeMockModel();
    const params = TEXTSEAL_DEFAULTS;
    const keys = deriveKeys('integration-secret');
    let detectedWm = 0;
    let detectedBase = 0;
    for (const seed of SEEDS) {
      const router = new RandomStream(BigInt(seed) + 7777n);
      const wm = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: seed }, {
        sampleOverride: (_i, ctx, probs) => gumbelMaxChoose(probs, ctx, keys, params, router),
      });
      const detW = detectTextseal(asTokenized(model, wm.tokenIds), keys, params, { windowL0: 20 });
      if (detW.watermarked) detectedWm++;

      const base = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: seed });
      const detB = detectTextseal(asTokenized(model, base.tokenIds), keys, params, { windowL0: 20 });
      if (detB.watermarked) detectedBase++;
    }
    expect(detectedWm).toBe(SEEDS.length);
    expect(detectedBase).toBe(0);
  });

  it('detection fails with the wrong key', async () => {
    const model = makeMockModel();
    const params = TEXTSEAL_DEFAULTS;
    const keys = deriveKeys('right-key');
    const wrong = deriveKeys('wrong-key');
    const router = new RandomStream(1n);
    const wm = await runGenerationLoop(model, PROMPT, { ...CFG, baseSeed: 9 }, {
      sampleOverride: (_i, ctx, probs) => gumbelMaxChoose(probs, ctx, keys, params, router),
    });
    const det = detectTextseal(asTokenized(model, wm.tokenIds), wrong, params, { windowL0: 20 });
    expect(det.watermarked).toBe(false);
  });
});

describe('generation loop + k-SemStamp end to end', () => {
  // Synthetic semantic space: embedding of a sentence = centroid of
  // hash(sentence) % K, plus deterministic jitter. Different candidate
  // sentences land in different clusters, exactly what rejection needs.
  const K = 8;
  const DIM = 16;
  const centroids: Centroids = {
    vectors: Array.from({ length: K }, (_, c) => {
      const v = new Array(DIM).fill(0);
      v[c] = 1;
      v[(c + K) % DIM] = c % 2 === 0 ? 0.1 : -0.1;
      const norm = Math.hypot(...v);
      return v.map((x) => x / norm);
    }),
    dim: DIM,
    encoder: 'mock',
    corpus: 'mock',
  };
  const params = { ...KSEMSTAMP_DEFAULTS, k: K, gamma: 0.25, margin: 0.0, maxTrials: 25 };

  function mockEmbed(text: string): number[] {
    let h = keyFromString(text);
    const c = Number(h % BigInt(K));
    const base = centroids.vectors[c];
    const jitterStream = new RandomStream(h);
    const v = base.map((x) => x + (jitterStream.next() - 0.5) * 0.05);
    const norm = Math.hypot(...v);
    return v.map((x) => x / norm);
  }

  it('rejection sampling produces a chain the detector accepts; baseline does not', async () => {
    const model = makeMockModel();
    let totalRetries = 0;
    const zWm: number[] = [];
    const zBase: number[] = [];

    for (const seed of SEEDS) {
      let prevCluster = 0;
      const wm = await runGenerationLoop(
        model,
        PROMPT,
        { ...CFG, maxNewTokens: 120, baseSeed: seed },
        {
          maxSentenceTrials: params.maxTrials,
          onSentenceEnd: async (text, _idx, attempt) => {
            const emb = mockEmbed(text);
            const res = acceptSentence(emb, prevCluster, centroids, params);
            if (res.accepted || attempt >= params.maxTrials) {
              prevCluster = res.assignment.clusterId;
              return { accept: true };
            }
            return { accept: false };
          },
        },
      );
      totalRetries += wm.retries;
      const sentsW = splitSentences(wm.text);
      const embsW = sentsW.map((s) => mockEmbed(s.text));
      if (sentsW.length >= 3) {
        zWm.push(detectKsemstamp(sentsW, embsW, centroids, params).statistic);
      }

      const base = await runGenerationLoop(model, PROMPT, {
        ...CFG,
        maxNewTokens: 120,
        baseSeed: seed,
      });
      const sentsB = splitSentences(base.text);
      const embsB = sentsB.map((s) => mockEmbed(s.text));
      if (sentsB.length >= 3) {
        zBase.push(detectKsemstamp(sentsB, embsB, centroids, params).statistic);
      }
    }

    expect(totalRetries).toBeGreaterThan(0); // rejection actually happened
    expect(zWm.length).toBeGreaterThan(2);
    expect(mean(zWm)).toBeGreaterThan(2);
    expect(mean(zWm)).toBeGreaterThan(mean(zBase) + 1.5);
  });

  it('rollback restores the exact accepted prefix (no stray tokens)', async () => {
    const model = makeMockModel();
    let calls = 0;
    const result = await runGenerationLoop(
      model,
      PROMPT,
      { ...CFG, maxNewTokens: 40, baseSeed: 3 },
      {
        maxSentenceTrials: 3,
        onSentenceEnd: async () => {
          calls++;
          // Reject the first attempt of every sentence, accept the second.
          return { accept: calls % 2 === 0 };
        },
      },
    );
    expect(result.retries).toBeGreaterThan(0);
    // decoded text must equal decode of tokenIds (consistency after rollbacks)
    expect(result.text.trim()).toBe(model.decode(result.tokenIds));
    // steps must be aligned with tokenIds
    expect(result.steps.length).toBe(result.tokenIds.length);
    expect(result.steps.map((s) => s.chosenTokenId)).toEqual(result.tokenIds);
  });
});

describe('streaming callback', () => {
  it('reports every accepted token once, in order', async () => {
    const model = makeMockModel();
    const seen: Array<{ index: number; id: number }> = [];
    const result = await runGenerationLoop(model, PROMPT, {
      ...CFG,
      maxNewTokens: 30,
      baseSeed: 5,
      onToken: (step) => seen.push({ index: step.index, id: step.chosenTokenId }),
    });
    expect(seen.map((s) => s.id)).toEqual(result.tokenIds);
    expect(seen.map((s) => s.index)).toEqual(result.tokenIds.map((_, i) => i));
  });

  it('rewinds the index when a sentence is rejected, so a listener can undo it', async () => {
    const model = makeMockModel();
    const indices: number[] = [];
    let calls = 0;
    const result = await runGenerationLoop(
      model,
      PROMPT,
      { ...CFG, maxNewTokens: 40, baseSeed: 3, onToken: (step) => indices.push(step.index) },
      {
        maxSentenceTrials: 3,
        onSentenceEnd: async () => {
          calls++;
          return { accept: calls % 2 === 0 };
        },
      },
    );

    // A rollback happened, so the stream is not monotonically increasing...
    expect(result.retries).toBeGreaterThan(0);
    const wentBackwards = indices.some((v, i) => i > 0 && v <= indices[i - 1]);
    expect(wentBackwards).toBe(true);

    // ...yet replaying it with "drop everything from this index" reconstructs
    // exactly the tokens the loop kept.
    const rebuilt: number[] = [];
    let cursor = 0;
    for (const step of indices) {
      rebuilt.length = step;
      rebuilt.push(cursor++);
    }
    expect(rebuilt.length).toBe(result.tokenIds.length);
  });
});

describe('token text is decoded in context', () => {
  /**
   * A byte-level tokenizer, like the real ones: '하' (3 UTF-8 bytes) is split
   * across two tokens, so neither half decodes to anything on its own.
   */
  const TOKEN_BYTES: number[][] = [
    [0xed, 0x95], // '하' bytes 1-2
    [0x98], //       '하' byte 3
    [0xeb, 0x82], // '나' bytes 1-2
    [0x98], //       '나' byte 3
    [0x2e, 0x20], // '. '
    [0x61], //       'a'
  ];
  const EOS = TOKEN_BYTES.length;
  const decoder = new TextDecoder();

  function byteModel(): LoopModel {
    return {
      vocabSize: EOS + 1,
      eosTokenIds: [EOS],
      decode(ids: number[]): string {
        return decoder.decode(Uint8Array.from(ids.flatMap((id) => TOKEN_BYTES[id] ?? [])));
      },
      async forward(_inputIds, fullSeqLen, past) {
        // Emit the halves of a character in order, so the stream is valid text.
        const logits = new Float32Array(EOS + 1).fill(-20);
        const cycle = [0, 1, 2, 3, 5, 4];
        logits[cycle[(fullSeqLen - 1) % cycle.length]] = 10;
        return { logits, past: ((past as number) ?? 0) + 1 };
      },
      disposePast() {},
    };
  }

  it('emits whole characters, never half of one', async () => {
    const model = byteModel();
    const pieces: string[] = [];
    const result = await runGenerationLoop(model, [5], {
      ...CFG,
      maxNewTokens: 24,
      baseSeed: 2,
      onToken: (s) => {
        pieces.length = s.index;
        pieces.push(s.chosenTokenText);
      },
    });

    expect(result.text).toContain('하');
    expect(result.text).toContain('나');
    // Half a character is never handed out: a token either completes one or
    // contributes nothing yet.
    for (const piece of pieces) expect(piece).not.toContain('\ufffd');
    expect(pieces.some((p) => p === '')).toBe(true);
  });

  it('per-token texts concatenate to exactly the final text', async () => {
    const model = byteModel();
    const pieces: string[] = [];
    const result = await runGenerationLoop(model, [5], {
      ...CFG,
      maxNewTokens: 24,
      baseSeed: 7,
      onToken: (s) => {
        pieces.length = s.index;
        pieces.push(s.chosenTokenText);
      },
    });

    // Trailing partial character (generation stopping mid-character) is the
    // only thing the pieces legitimately lack.
    const settled = result.text.replace(/\ufffd+$/, '');
    expect(pieces.join('')).toBe(settled);
    expect(result.steps.map((s) => s.chosenTokenText).join('')).toBe(settled);
  });
});

describe('retry temperature is bounded', () => {
  /** Max entropy the sampler saw - a direct read on how hot it ran. */
  async function hottest(step: number, max: number): Promise<number> {
    const model = makeMockModel();
    let peak = 0;
    await runGenerationLoop(
      model,
      PROMPT,
      {
        ...CFG,
        maxNewTokens: 60,
        baseSeed: 4,
        retryTemperatureStep: step,
        retryTemperatureMax: max,
        onToken: (s) => {
          peak = Math.max(peak, s.entropy ?? 0);
        },
      },
      {
        maxSentenceTrials: 20,
        // Reject deep into the retries so the escalation has room to run away.
        onSentenceEnd: async (_t, _i, attempt) => ({ accept: attempt >= 12 }),
      },
    );
    return peak;
  }

  it('a runaway step is clamped to the ceiling', async () => {
    const noEscalation = await hottest(0, 1);
    const clamped = await hottest(1.0, 1.5); // would reach 12x unbounded
    const uncapped = await hottest(1.0, 100);

    // Hotter sampling means higher entropy, but the mock's 256-token vocab
    // caps entropy at ln(256), so compare positions rather than ratios.
    expect(clamped).toBeGreaterThan(noEscalation);
    expect(uncapped).toBeGreaterThan(clamped);
    // The ceiling keeps the run nearer the un-escalated baseline than the
    // runaway one - which is the whole point of having it.
    expect(clamped - noEscalation).toBeLessThan(uncapped - noEscalation);
  });
});