whisper_web/tests/server-live.test.ts

import { afterAll, describe, expect, test } from "bun:test";
import { mkdtemp, readdir, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { createLiveWebSocket, type LiveSessionData, type LiveSocket } from "../src/server/live";
import type { Sidecar, SidecarSegment, SidecarTranscribeRequest } from "../src/server/transcribe";
import type { SessionUser } from "../src/server/auth";

// D-019: the server-live WebSocket bridge, tested against a fake Sidecar
// (createLiveWebSocket takes it as an injected dependency, mirroring
// createTranscribeHandlers in tests/transcribe.test.ts) — no Python sidecar,
// no real model, but real LiveTranscriber/segmenter logic and real temp files.

const USER: SessionUser = { id: 1, username: "tester" };
const RATE = 16_000;

function speech(seconds: number): Float32Array {
  const out = new Float32Array(Math.round(seconds * RATE));
  for (let i = 0; i < out.length; i++) out[i] = i % 2 === 0 ? 0.1 : -0.1;
  return out;
}

function silence(seconds: number): Float32Array {
  return new Float32Array(Math.round(seconds * RATE));
}

function toBuffer(audio: Float32Array): Buffer {
  return Buffer.from(audio.buffer, audio.byteOffset, audio.byteLength);
}

async function waitFor(
  predicate: () => boolean | Promise<boolean>,
  timeoutMs = 3000,
): Promise<void> {
  const start = performance.now();
  while (!(await predicate())) {
    if (performance.now() - start > timeoutMs) {
      throw new Error("timed out waiting for condition");
    }
    await new Promise((resolve) => setTimeout(resolve, 5));
  }
}

const tempDirs: string[] = [];
async function makeTempDir(): Promise<string> {
  const dir = await mkdtemp(join(tmpdir(), "live-ws-test-"));
  tempDirs.push(dir);
  return dir;
}
afterAll(async () => {
  await Promise.all(tempDirs.map((dir) => rm(dir, { recursive: true, force: true })));
});

interface FakeSidecar extends Sidecar {
  requests: SidecarTranscribeRequest[];
  audioPayloads: Float32Array[];
}

function fakeSidecar(
  behavior: { segments?: SidecarSegment[]; error?: string } = {},
): FakeSidecar {
  const segments = behavior.segments ?? [{ start: 0, end: 1, text: " hello world" }];
  const sidecar: FakeSidecar = {
    requests: [],
    audioPayloads: [],
    async status() {
      return { available: true };
    },
    async transcribe(request, onSegment) {
      sidecar.requests.push(request);
      sidecar.audioPayloads.push(
        new Float32Array(await Bun.file(request.audioPath).arrayBuffer()),
      );
      if (behavior.error) throw new Error(behavior.error);
      for (const segment of segments) onSegment(segment);
      return { cancelled: false };
    },
  };
  return sidecar;
}

type FakeSocket = LiveSocket & {
  sent: { type: string; [key: string]: unknown }[];
  closed: boolean;
};

function fakeSocket(data: LiveSessionData): FakeSocket {
  const sent: { type: string; [key: string]: unknown }[] = [];
  const ws = {
    data,
    sent,
    closed: false,
    send(message: string) {
      sent.push(JSON.parse(message));
    },
    close() {
      ws.closed = true;
    },
  };
  return ws as unknown as FakeSocket;
}

function sessionData(overrides: Partial<LiveSessionData> = {}): LiveSessionData {
  return {
    user: USER,
    model: "tiny",
    task: "transcribe",
    tempDir: "/tmp",
    ...overrides,
  };
}

describe("createLiveWebSocket", () => {
  test("a committed segment writes the exact PCM bytes and passes model/language/task/mode through", async () => {
    const tempDir = await makeTempDir();
    const sidecar = fakeSidecar();
    const handler = createLiveWebSocket({ sidecar });
    const ws = fakeSocket(sessionData({ model: "base", language: "da", task: "translate", tempDir }));

    handler.open(ws);
    handler.message(ws, toBuffer(speech(1)));
    handler.message(ws, toBuffer(silence(1.5))); // trailing silence commits the segment

    // The fake records the request before it finishes reading the PCM file —
    // wait for the payload itself, or a loaded machine races past it.
    await waitFor(() => sidecar.audioPayloads.length > 0);
    expect(sidecar.requests[0]).toEqual(
      // Live passes are short: sequential decoding, no batching/VAD (D-WS5a).
      expect.objectContaining({ model: "base", language: "da", task: "translate", mode: "live" }),
    );
    // The bytes written to the temp file are exactly the committed segment's
    // samples — not, e.g., a whole shared buffer (see the byteOffset/byteLength fix).
    const payload = sidecar.audioPayloads[0]!;
    expect(payload.length).toBeGreaterThan(0);
    for (const sample of payload) expect(Math.abs(sample)).toBeLessThanOrEqual(0.1 + 1e-6);

    await waitFor(() =>
      ws.sent.some((m) => m.type === "update" && (m.update as any).committedText === "hello world"),
    );

    // The temp file is cleaned up once the pass finishes.
    await waitFor(async () => (await readdir(tempDir)).length === 0);
  });

  test("every pass of a session carries the session's vocabulary hint (D-WS2a)", async () => {
    const tempDir = await makeTempDir();
    const sidecar = fakeSidecar();
    const handler = createLiveWebSocket({ sidecar });
    const ws = fakeSocket(sessionData({ prompt: "Maxwell, Voices of the Void", tempDir }));
    handler.open(ws);
    handler.message(ws, toBuffer(speech(1)));
    handler.message(ws, toBuffer(silence(1.5)));
    await waitFor(() => sidecar.requests.length > 0);
    for (const request of sidecar.requests) {
      expect(request.prompt).toBe("Maxwell, Voices of the Void");
    }
    await waitFor(async () => (await readdir(tempDir)).length === 0);
  });

  test("a sidecar error surfaces on the update stream instead of crashing the session", async () => {
    const sidecar = fakeSidecar({ error: "model load failed" });
    const handler = createLiveWebSocket({ sidecar });
    const ws = fakeSocket(sessionData({ tempDir: await makeTempDir() }));

    handler.open(ws);
    handler.message(ws, toBuffer(speech(1)));
    handler.message(ws, toBuffer(silence(1.5)));

    await waitFor(() =>
      ws.sent.some((m) => m.type === "update" && (m.update as any).error === "model load failed"),
    );
  });

  test('a "stop" text frame drains, sends the result, and closes the socket', async () => {
    const sidecar = fakeSidecar({ segments: [{ start: 0, end: 1, text: " done" }] });
    const handler = createLiveWebSocket({ sidecar });
    const ws = fakeSocket(sessionData({ tempDir: await makeTempDir() }));

    handler.open(ws);
    handler.message(ws, toBuffer(speech(1)));
    handler.message(ws, JSON.stringify({ type: "stop" }));

    await waitFor(() => ws.closed);
    const result = ws.sent.find((m) => m.type === "result");
    expect(result?.result).toEqual({ text: "done", chunks: [{ text: " done", timestamp: [0, 1] }] });
  });

  test("a malformed text frame is ignored rather than throwing", () => {
    const handler = createLiveWebSocket({ sidecar: fakeSidecar() });
    const ws = fakeSocket(sessionData({ tempDir: "/tmp" }));
    handler.open(ws);
    expect(() => handler.message(ws, "not json")).not.toThrow();
    expect(ws.closed).toBe(false);
  });

  test("close() drains the session without throwing even if never stopped", async () => {
    const handler = createLiveWebSocket({ sidecar: fakeSidecar() });
    const ws = fakeSocket(sessionData({ tempDir: await makeTempDir() }));
    handler.open(ws);
    handler.message(ws, toBuffer(speech(0.5)));
    expect(() => handler.close(ws)).not.toThrow();
  });
});