src/engine/runtime-generate.test.ts

import { describe, expect, it, vi } from 'vitest';
import type {
  ChatCompletionChunk,
  ChatCompletionResponse,
} from '../../vendor/wllama-bonsai/esm/index.js';
import type { EngineEvent } from './protocol';
import { BrowserEngineRuntime } from './runtime';

interface CompletionOptions {
  onData(chunk: ChatCompletionChunk): void;
  stream?: boolean;
  logprobs?: boolean;
  top_logprobs?: number;
  tools?: unknown[];
  tool_choice?: unknown;
  max_tokens?: number;
  return_progress?: boolean;
  timings_per_token?: boolean;
}

type GenerationProgressEvent = Extract<EngineEvent, { event: 'generation' }>;

interface RuntimeInternals {
  wllamaFlavor: 'compat' | 'jspi' | null;
  wllama: {
    isModelLoaded(): boolean;
    createChatCompletion(options: CompletionOptions): Promise<void | ChatCompletionResponse>;
  } | null;
  loaded: {
    manifest: unknown;
    backend: 'wasm' | 'webgpu';
    model: {
      id: '1_7b';
      displayName: string;
      cpuFallback: true;
      runtimePolicy: {
        flashAttention: false;
        tokenEmbeddingOnWebGPU: boolean;
        requireSingleWebGPUGraph: boolean;
      };
    };
  } | null;
}

function chunk(
  content: string,
  metadata: Partial<Pick<ChatCompletionChunk, 'usage' | 'timings'>> = {},
): ChatCompletionChunk {
  return {
    id: 'chunk',
    object: 'chat.completion.chunk',
    created: 0,
    model: 'fixture',
    choices: [{
      index: 0,
      delta: { content },
      finish_reason: content ? null : 'stop',
      logprobs: null,
    }],
    ...metadata,
  };
}

function tracedChunk(
  content: string,
  id: number,
  metadata: Partial<Pick<ChatCompletionChunk, 'usage' | 'timings'>> = {},
): ChatCompletionChunk {
  const value = chunk(content, metadata);
  const choice = value.choices[0];
  if (!choice) throw new Error('Chunk fixture requires one choice.');
  choice.logprobs = {
    content: [{
      id,
      token: content,
      logprob: -0.1,
      bytes: null,
      top_logprobs: [
        { id: id + 1_004, token: 'fifth', logprob: -4, bytes: null },
        { id: id + 1_002, token: 'third', logprob: -2, bytes: null },
        { id, token: content, logprob: -0.1, bytes: null },
        { id: id + 1_003, token: 'fourth', logprob: -3, bytes: null },
        { id: id + 1_001, token: 'second', logprob: -1, bytes: null },
      ],
    }],
    refusal: null,
  };
  return value;
}

function reasoningTracedChunk(
  reasoningContent: string,
  id: number,
  metadata: Partial<Pick<ChatCompletionChunk, 'usage' | 'timings'>> = {},
): ChatCompletionChunk {
  const value = tracedChunk('', id, metadata);
  const choice = value.choices[0];
  if (!choice) throw new Error('Chunk fixture requires one choice.');
  choice.finish_reason = null;
  (choice.delta as unknown as Record<string, unknown>).reasoning_content = reasoningContent;
  return value;
}

describe('BrowserEngineRuntime generation telemetry', () => {
  it('retains usage and timings when a trailing stream chunk omits metadata', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    let completionOptions: CompletionOptions | null = null;
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        completionOptions = options;
        options.onData(chunk('Ready', {
          usage: { prompt_tokens: 9, completion_tokens: 1, total_tokens: 10 },
          timings: {
            cache_n: 0,
            prompt_n: 9,
            prompt_ms: 90,
            prompt_per_token_ms: 10,
            prompt_per_second: 100,
            predicted_n: 1,
            predicted_ms: 20,
            predicted_per_token_ms: 20,
            predicted_per_second: 50,
          },
        }));
        options.onData(chunk(''));
      },
    };

    const result = await runtime.generate(
      'request',
      { messages: [{ role: 'user', content: 'Ready?' }] },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.usage).toEqual({ promptTokens: 9, completionTokens: 1, totalTokens: 10 });
    expect(result.timings).toEqual({ promptTokensPerSecond: 100, predictedTokensPerSecond: 50 });
    expect(result.reasoningText).toBe('');
    expect(result.tokenIds).toBeNull();
    expect(result.tokenTrace).toBeNull();
    expect(completionOptions).not.toHaveProperty('logprobs');
    expect(completionOptions).not.toHaveProperty('top_logprobs');
    expect(completionOptions).toMatchObject({ max_tokens: 4_096 });
  });

  it('streams separate live rates without trusting unstable first-token timings', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    let completionOptions: CompletionOptions | null = null;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        completionOptions = options;
        const earlyPrefill = chunk('', {
          timings: {
            cache_n: 2,
            prompt_n: 7,
            prompt_ms: 0.001,
            prompt_per_token_ms: 0.001 / 5,
            prompt_per_second: 5_000_000,
            predicted_n: 0,
            predicted_ms: 0,
            predicted_per_token_ms: 0,
            predicted_per_second: 0,
          },
        }) as ChatCompletionChunk & {
          prompt_progress: { total: number; cache: number; processed: number; time_ms: number };
        };
        const earlyPrefillChoice = earlyPrefill.choices[0];
        if (!earlyPrefillChoice) throw new Error('Prefill fixture requires one choice.');
        earlyPrefillChoice.finish_reason = null;
        earlyPrefill.prompt_progress = { total: 12, cache: 2, processed: 7, time_ms: 1 };
        options.onData(earlyPrefill);
        const prefill = chunk('') as ChatCompletionChunk & {
          prompt_progress: { total: number; cache: number; processed: number; time_ms: number };
        };
        const prefillChoice = prefill.choices[0];
        if (!prefillChoice) throw new Error('Prefill fixture requires one choice.');
        prefillChoice.finish_reason = null;
        prefill.prompt_progress = { total: 12, cache: 2, processed: 7, time_ms: 100 };
        options.onData(prefill);
        options.onData(chunk('A', {
          timings: {
            cache_n: 2,
            prompt_n: 12,
            prompt_ms: 100,
            prompt_per_token_ms: 10,
            prompt_per_second: 100,
            predicted_n: 1,
            predicted_ms: 0.001,
            predicted_per_token_ms: 0.001,
            predicted_per_second: 1_000_000,
          },
        }));
        options.onData(chunk('B', {
          usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 },
          timings: {
            cache_n: 2,
            prompt_n: 12,
            prompt_ms: 100,
            prompt_per_token_ms: 10,
            prompt_per_second: 100,
            predicted_n: 3,
            predicted_ms: 200,
            predicted_per_token_ms: 200 / 3,
            predicted_per_second: 15,
          },
        }));
        options.onData(chunk(''));
      },
    };
    const progress: GenerationProgressEvent[] = [];

    const result = await runtime.generate(
      'request-live-progress',
      { messages: [{ role: 'user', content: 'Count.' }] },
      new AbortController().signal,
      (event) => {
        if (event.event === 'generation') progress.push(event);
      },
    );

    expect(completionOptions).toHaveProperty('return_progress', true);
    expect(completionOptions).toHaveProperty('timings_per_token', true);
    expect(progress.map(({ phase }) => phase)).toEqual([
      'prefill',
      'prefill',
      'prefill',
      'decode',
      'decode',
    ]);
    expect(progress[1]).toMatchObject({
      promptProcessed: 7,
      promptTotal: 12,
      promptCached: 2,
      elapsedMs: 1,
      promptTokensPerSecond: 0,
      decodeTokensPerSecond: 0,
    });
    expect(progress[2]).toMatchObject({
      promptProcessed: 7,
      promptTotal: 12,
      promptCached: 2,
      elapsedMs: 100,
      promptTokensPerSecond: 50,
      decodeTokensPerSecond: 0,
    });
    expect(progress[3]).toMatchObject({
      completionTokens: 1,
      elapsedMs: 0.001,
      decodeTokensPerSecond: 0,
    });
    expect(progress[4]).toMatchObject({ completionTokens: 3, elapsedMs: 200 });
    expect(progress[4]?.decodeTokensPerSecond).toBeCloseTo(10, 3);
    expect(result.text).toBe('AB');
    expect(result.usage).toEqual({ promptTokens: 12, completionTokens: 3, totalTokens: 15 });
  });

  it('uses token arrival time for live decode telemetry in the compat runtime', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    let completionOptions: CompletionOptions | null = null;
    internals.wllamaFlavor = 'compat';
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        completionOptions = options;
        options.onData(chunk('A'));
        options.onData(chunk('B', {
          usage: { prompt_tokens: 8, completion_tokens: 2, total_tokens: 10 },
          timings: {
            cache_n: 0,
            prompt_n: 8,
            prompt_ms: 200,
            prompt_per_token_ms: 25,
            prompt_per_second: 40,
            predicted_n: 2,
            predicted_ms: 500,
            predicted_per_token_ms: 250,
            predicted_per_second: 4,
          },
        }));
        options.onData(chunk(''));
      },
    };
    const now = vi.spyOn(performance, 'now')
      .mockReturnValueOnce(1_000)
      .mockReturnValueOnce(1_250);
    const progress: GenerationProgressEvent[] = [];

    try {
      const result = await runtime.generate(
        'request-compat-progress',
        { messages: [{ role: 'user', content: 'Count.' }] },
        new AbortController().signal,
        (event) => {
          if (event.event === 'generation') progress.push(event);
        },
      );

      expect(completionOptions).not.toHaveProperty('timings_per_token');
      expect(completionOptions).not.toHaveProperty('return_progress');
      expect(progress.map(({ phase }) => phase)).toEqual([
        'prefill',
        'decode',
        'decode',
      ]);
      expect(progress[1]).toMatchObject({ completionTokens: 1, decodeTokensPerSecond: 0 });
      expect(progress[2]).toMatchObject({ completionTokens: 2, decodeTokensPerSecond: 4 });
      expect(result.timings).toEqual({ promptTokensPerSecond: 40, predictedTokensPerSecond: 4 });
    } finally {
      now.mockRestore();
    }
  });

  it('reconciles incomplete streamed usage with llama.cpp timing counts', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        onData(chunk('one, two, three', {
          usage: { prompt_tokens: 3, completion_tokens: 1, total_tokens: 4 },
          timings: {
            cache_n: 0,
            prompt_n: 9,
            prompt_ms: 90,
            prompt_per_token_ms: 10,
            prompt_per_second: 100,
            predicted_n: 64,
            predicted_ms: 1_280,
            predicted_per_token_ms: 20,
            predicted_per_second: 50,
          },
        }));
        onData(chunk(''));
      },
    };

    const result = await runtime.generate(
      'request-incomplete-usage',
      { messages: [{ role: 'user', content: 'Count.' }] },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.usage).toEqual({ promptTokens: 9, completionTokens: 64, totalTokens: 73 });
    expect(result.tokenIds).toBeNull();
    expect(result.tokenTrace).toBeNull();
  });

  it('returns sampled token ids with selected logprobs and sorted top-five candidates', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    let completionOptions: CompletionOptions | null = null;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        completionOptions = options;
        options.onData(tracedChunk('one', 101));
        options.onData(tracedChunk(' two', 202, {
          usage: { prompt_tokens: 4, completion_tokens: 2, total_tokens: 6 },
          timings: {
            cache_n: 0,
            prompt_n: 4,
            prompt_ms: 40,
            prompt_per_token_ms: 10,
            prompt_per_second: 100,
            predicted_n: 2,
            predicted_ms: 40,
            predicted_per_token_ms: 20,
            predicted_per_second: 50,
          },
        }));
        options.onData(chunk(''));
      },
    };

    const result = await runtime.generate(
      'request-token-ids',
      { messages: [{ role: 'user', content: 'Count.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.tokenIds).toEqual([101, 202]);
    expect(result.tokenTrace).toEqual([
      {
        selected: { id: 101, logprob: -0.1 },
        topCandidates: [
          { id: 101, logprob: -0.1 },
          { id: 1_102, logprob: -1 },
          { id: 1_103, logprob: -2 },
          { id: 1_104, logprob: -3 },
          { id: 1_105, logprob: -4 },
        ],
      },
      {
        selected: { id: 202, logprob: -0.1 },
        topCandidates: [
          { id: 202, logprob: -0.1 },
          { id: 1_203, logprob: -1 },
          { id: 1_204, logprob: -2 },
          { id: 1_205, logprob: -3 },
          { id: 1_206, logprob: -4 },
        ],
      },
    ]);
    expect(result.tokenTraceAccounting).toEqual({
      usageCompletionTokens: 2,
      tracedTokens: 2,
      delta: 0,
    });
    expect(completionOptions).toHaveProperty('logprobs', true);
    expect(completionOptions).toHaveProperty('top_logprobs', 5);
  });

  it('uses a non-streaming completion for a tool-call token trace', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    let completionOptions: CompletionOptions | null = null;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        completionOptions = options;
        const trace = tracedChunk('', 303);
        return {
          id: 'completion',
          object: 'chat.completion',
          created: 0,
          model: 'fixture',
          choices: [{
            index: 0,
            message: {
              role: 'assistant',
              content: null,
              tool_calls: [{
                id: 'call_memory',
                type: 'function',
                function: {
                  name: 'memory',
                  arguments: '{"operation":"set","key":"color","value":"blue"}',
                },
              }],
            },
            finish_reason: 'tool_calls',
            logprobs: trace.choices[0]?.logprobs ?? null,
          }],
          usage: { prompt_tokens: 30, completion_tokens: 1, total_tokens: 31 },
        };
      },
    };

    const result = await runtime.generate(
      'request-tool-trace',
      {
        messages: [{ role: 'user', content: 'Remember blue.' }],
        tools: [{
          type: 'function',
          function: {
            name: 'memory',
            parameters: { type: 'object', properties: {} },
          },
        }],
        toolChoice: 'auto',
        returnTokenIds: true,
      },
      new AbortController().signal,
      () => undefined,
    );

    expect(completionOptions).toMatchObject({
      stream: false,
      logprobs: true,
      top_logprobs: 5,
      tool_choice: 'auto',
    });
    expect(completionOptions).not.toHaveProperty('onData');
    expect(result.finishReason).toBe('tool_calls');
    expect(result.tokenIds).toEqual([303]);
    expect(result.toolCalls).toEqual([{
      id: 'call_memory',
      type: 'function',
      function: {
        name: 'memory',
        arguments: '{"operation":"set","key":"color","value":"blue"}',
      },
    }]);
  });

  it('streams a required HTML artifact call with the full completion allowance', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    let completionOptions: CompletionOptions | null = null;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        completionOptions = options;
        options.onData({
          id: 'artifact-chunk',
          object: 'chat.completion.chunk',
          created: 0,
          model: 'fixture',
          choices: [{
            index: 0,
            delta: {
              content: null,
              tool_calls: [{
                index: 0,
                id: 'call_artifact',
                type: 'function',
                function: {
                  name: 'html_artifact',
                  arguments: '{"html":"<h1>Artifact proof</h1>"}',
                },
              }],
            },
            finish_reason: 'tool_calls',
            logprobs: null,
          }],
        } as ChatCompletionChunk);
      },
    };

    const events: EngineEvent[] = [];
    const result = await runtime.generate(
      'request-required-artifact',
      {
        messages: [{ role: 'user', content: 'Create an HTML app.' }],
        tools: [{
          type: 'function',
          function: {
            name: 'html_artifact',
            parameters: {
              type: 'object',
              properties: { html: { type: 'string' } },
              required: ['html'],
            },
          },
        }],
        toolChoice: 'required',
        maxTokens: 4_096,
      },
      new AbortController().signal,
      (event) => events.push(event),
    );

    expect(completionOptions).toMatchObject({
      stream: true,
      max_tokens: 4_096,
      tool_choice: 'required',
    });
    expect(completionOptions).not.toHaveProperty('logprobs');
    expect(result.finishReason).toBe('tool_calls');
    expect(result.toolCalls).toEqual([{
      id: 'call_artifact',
      type: 'function',
      function: {
        name: 'html_artifact',
        arguments: '{"html":"<h1>Artifact proof</h1>"}',
      },
    }]);
    expect(events).toContainEqual({
      type: 'event',
      requestId: 'request-required-artifact',
      event: 'tool-call',
      index: 0,
      id: 'call_artifact',
      name: 'html_artifact',
      argumentCharacters: 34,
    });
  });

  it('does not execute a partial tool call when generation reaches the token limit', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async (options) => {
        options.onData({
          id: 'partial-artifact',
          object: 'chat.completion.chunk',
          created: 0,
          model: 'fixture',
          choices: [{
            index: 0,
            delta: {
              content: null,
              tool_calls: [{
                index: 0,
                id: 'call_partial',
                type: 'function',
                function: { name: 'html_artifact', arguments: '{"html":"<main>' },
              }],
            },
            finish_reason: 'length',
            logprobs: null,
          }],
        } as ChatCompletionChunk);
      },
    };

    const result = await runtime.generate(
      'request-partial-artifact',
      {
        messages: [{ role: 'user', content: 'Create an HTML app.' }],
        tools: [{
          type: 'function',
          function: {
            name: 'html_artifact',
            parameters: { type: 'object', properties: {} },
          },
        }],
        toolChoice: 'required',
      },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.finishReason).toBe('length');
    expect(result.toolCalls).toEqual([]);
    expect(result.toolCallError).toContain('arguments are not valid JSON');
  });

  it('keeps reasoning-only output separate from visible text while preserving its token trace', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        onData(reasoningTracedChunk('private reasoning', 303, {
          usage: { prompt_tokens: 4, completion_tokens: 1, total_tokens: 5 },
        }));
        onData(chunk(''));
      },
    };
    const streamed: Array<{ text: string; reasoningDelta?: string }> = [];

    const result = await runtime.generate(
      'request-reasoning-trace',
      { messages: [{ role: 'user', content: 'Think.' }], returnTokenIds: true },
      new AbortController().signal,
      (event) => {
        if (event.event === 'token') streamed.push(event);
      },
    );

    expect(result.text).toBe('');
    expect(result.reasoningText).toBe('private reasoning');
    expect(result.tokenIds).toEqual([303]);
    expect(result.tokenTrace).toHaveLength(1);
    expect(streamed).toEqual([{ type: 'event', requestId: 'request-reasoning-trace', event: 'token', text: '', reasoningDelta: 'private reasoning' }]);
  });

  it('fails loudly when a sampled token id is non-integer', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        onData(tracedChunk('bad', 1.5));
      },
    };

    await expect(runtime.generate(
      'request-invalid-token-id',
      { messages: [{ role: 'user', content: 'Count.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    )).rejects.toMatchObject({ code: 'INVALID_TOKEN_ID_TRACE' });
  });

  it('fails loudly when a sampled logprob entry omits its token id', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        const malformed = tracedChunk('missing', 123);
        const entry = malformed.choices[0]?.logprobs?.content?.[0];
        if (!entry) throw new Error('Malformed trace fixture requires one logprob entry.');
        Reflect.deleteProperty(entry, 'id');
        onData(malformed);
      },
    };

    await expect(runtime.generate(
      'request-missing-token-id-field',
      { messages: [{ role: 'user', content: 'Count.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    )).rejects.toMatchObject({
      code: 'INVALID_TOKEN_ID_TRACE',
      details: { index: 0, id: undefined },
    });
  });

  it('records usage drift without treating telemetry as the token-identity source', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        onData(chunk('missing', {
          usage: { prompt_tokens: 4, completion_tokens: 1, total_tokens: 5 },
        }));
      },
    };

    const result = await runtime.generate(
      'request-missing-token-id',
      { messages: [{ role: 'user', content: 'Count.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.tokenIds).toEqual([]);
    expect(result.tokenTraceAccounting).toEqual({
      usageCompletionTokens: 1,
      tracedTokens: 0,
      delta: -1,
    });
  });

  it('accepts one untraced terminal stop token while preserving visible token ids', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        onData(tracedChunk('visible', 707));
        onData(chunk('', {
          usage: { prompt_tokens: 4, completion_tokens: 2, total_tokens: 6 },
        }));
      },
    };

    const result = await runtime.generate(
      'request-terminal-stop-token',
      { messages: [{ role: 'user', content: 'Stop.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.finishReason).toBe('stop');
    expect(result.usage?.completionTokens).toBe(2);
    expect(result.tokenIds).toEqual([707]);
    expect(result.tokenTrace).toHaveLength(1);
    expect(result.tokenTraceAccounting).toEqual({
      usageCompletionTokens: 2,
      tracedTokens: 1,
      delta: -1,
    });
  });

  it('accepts one final traced token emitted after terminal usage accounting', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        onData(tracedChunk('first', 808));
        const final = tracedChunk(' final', 909, {
          usage: { prompt_tokens: 4, completion_tokens: 1, total_tokens: 5 },
        });
        const choice = final.choices[0];
        if (!choice) throw new Error('Terminal fixture requires one choice.');
        choice.finish_reason = 'length';
        onData(final);
      },
    };

    const result = await runtime.generate(
      'request-terminal-late-trace',
      { messages: [{ role: 'user', content: 'Continue.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    );

    expect(result.finishReason).toBe('length');
    expect(result.usage?.completionTokens).toBe(1);
    expect(result.tokenIds).toEqual([808, 909]);
    expect(result.tokenTrace).toHaveLength(2);
    expect(result.tokenTraceAccounting).toEqual({
      usageCompletionTokens: 1,
      tracedTokens: 2,
      delta: 1,
    });
  });

  it('fails loudly when a sampled token does not have exactly five top candidates', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        const malformed = tracedChunk('short', 404);
        malformed.choices[0]?.logprobs?.content?.[0]?.top_logprobs.pop();
        onData(malformed);
      },
    };

    await expect(runtime.generate(
      'request-short-top-logprobs',
      { messages: [{ role: 'user', content: 'Count.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    )).rejects.toMatchObject({
      code: 'INVALID_TOKEN_LOGPROB_TRACE',
      details: { index: 0, expected: 5, observed: 4 },
    });
  });

  it('fails loudly when a top candidate has an invalid logprob', async () => {
    const runtime = new BrowserEngineRuntime();
    const internals = runtime as unknown as RuntimeInternals;
    internals.loaded = {
      manifest: {},
      backend: 'wasm',
      model: {
        id: '1_7b',
        displayName: 'Fixture Bonsai',
        cpuFallback: true,
        runtimePolicy: {
          flashAttention: false,
          tokenEmbeddingOnWebGPU: true,
          requireSingleWebGPUGraph: false,
        },
      },
    };
    internals.wllama = {
      isModelLoaded: () => true,
      createChatCompletion: async ({ onData }) => {
        const malformed = tracedChunk('invalid', 505);
        const candidate = malformed.choices[0]?.logprobs?.content?.[0]?.top_logprobs[1];
        if (!candidate) throw new Error('Malformed trace fixture requires a top candidate.');
        candidate.logprob = Number.NaN;
        onData(malformed);
      },
    };

    await expect(runtime.generate(
      'request-invalid-top-logprob',
      { messages: [{ role: 'user', content: 'Count.' }], returnTokenIds: true },
      new AbortController().signal,
      () => undefined,
    )).rejects.toMatchObject({
      code: 'INVALID_TOKEN_LOGPROB_TRACE',
      details: { index: 0, field: 'topCandidates[1].logprob' },
    });
  });
});