import { expect, test } from "bun:test";import { mkdirSync, mkdtempSync, realpathSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { Effect, Fiber, Schema } from "effect";import {  CHATTERBOX_BACKEND,  CHATTERBOX_MODEL_REVISION,  CHATTERBOX_PACKAGE_VERSION,} from "@local-studio/contracts/speech";import {  ChatterboxWorkerClient,  SpeechWorkerError,  SpeechWorkerRequestSchema,  spawnNodeSpeechWorker,  type SpeechWorkerRequest,  type SpeechWorkerSpawnOptions,  type SpeechWorkerTransport,} from "./worker-client";import { chatterboxRuntimePaths } from "./runtime"; const GPU_UUID = "GPU-01234567-89ab-cdef-0123-456789abcdef"; class FakeTransport implements SpeechWorkerTransport {  readonly writes: string[] = [];  killed = false;  inputClosed = false;  exitOnClose = true;  exitOnKill = true;  closeFailure: Error | null = null;  onWrite: ((request: SpeechWorkerRequest) => void) | null = null;  private exited = false;  private readonly lineListeners = new Set<(line: string) => void>();  private readonly stderrListeners = new Set<(line: string) => void>();  private readonly errorListeners = new Set<(error: Error) => void>();  private readonly exitListeners = new Set<    (code: number | null, signal: NodeJS.Signals | null) => void  >();   write(line: string): void {    this.writes.push(line);    const request = Schema.decodeUnknownSync(SpeechWorkerRequestSchema)(JSON.parse(line));    this.onWrite?.(request);  }   closeInput(): void {    this.inputClosed = true;    if (this.closeFailure) throw this.closeFailure;    if (this.exitOnClose) queueMicrotask(() => this.emitExit(0, null));  }   kill(): void {    this.killed = true;    if (this.exitOnKill) queueMicrotask(() => this.emitExit(null, "SIGKILL"));  }   onLine(listener: (line: string) => void): () => void {    this.lineListeners.add(listener);    return () => this.lineListeners.delete(listener);  }   onStderr(listener: (line: string) => void): () => void {    this.stderrListeners.add(listener);    return () => this.stderrListeners.delete(listener);  }   onError(listener: (error: Error) => void): () => void {    this.errorListeners.add(listener);    return () => this.errorListeners.delete(listener);  }   onExit(listener: (code: number | null, signal: NodeJS.Signals | null) => void): () => void {    this.exitListeners.add(listener);    return () => this.exitListeners.delete(listener);  }   emitLine(value: unknown): void {    const line = typeof value === "string" ? value : JSON.stringify(value);    this.lineListeners.forEach((listener) => listener(line));  }   emitExit(code: number | null, signal: NodeJS.Signals | null): void {    if (this.exited) return;    this.exited = true;    this.exitListeners.forEach((listener) => listener(code, signal));  }} const readyFrame = {  type: "ready",  backend: CHATTERBOX_BACKEND,  package_version: CHATTERBOX_PACKAGE_VERSION,  model_revision: CHATTERBOX_MODEL_REVISION,  cuda_devices: 1,  sample_rate: 24000,} as const; const fixture = (): { directory: string; voicePath: string; workerPath: string } => {  const directory = mkdtempSync(join(tmpdir(), "local-studio-speech-worker-"));  const workerPath = join(directory, "worker.py");  const paths = chatterboxRuntimePaths(directory, workerPath);  const voicePath = join(paths.voiceDirectory, "voice.wav");  mkdirSync(paths.voiceDirectory, { recursive: true });  writeFileSync(workerPath, "");  writeFileSync(voicePath, "voice");  return { directory, voicePath, workerPath };}; const speechError = async (promise: Promise<unknown>): Promise<SpeechWorkerError> => {  try {    await promise;  } catch (error) {    if (error instanceof SpeechWorkerError) return error;    throw error;  }  throw new Error("Expected speech operation to fail");}; test("pins the worker environment and frames synthesize and shutdown requests", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  const spawns: SpeechWorkerSpawnOptions[] = [];  const ids = ["speech-one", "shutdown-one"];  try {    transport.onWrite = (request): void => {      if (request.type === "synthesize") {        transport.emitLine({          type: "synthesize",          id: request.id,          output_path: request.output_path,          sample_rate: 24000,        });      } else {        transport.emitLine({ type: "shutdown", id: request.id });      }    };    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => ids.shift() ?? "unused",      environment: {        PATH: "/usr/bin",        LOCAL_STUDIO_API_KEY: "controller-secret",        HF_TOKEN: "hub-secret",      },      spawnWorker: (options): SpeechWorkerTransport => {        spawns.push(options);        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });     const result = await client.synthesize({ text: "Hello from Local Studio", voicePath });    expect(result.sampleRate).toBe(24000);    expect(spawns).toHaveLength(1);    expect(spawns[0]?.command).toBe(client.paths.pythonPath);    expect(spawns[0]?.args).toEqual(["-u", workerPath]);    expect(spawns[0]?.env["CUDA_VISIBLE_DEVICES"]).toBe(GPU_UUID);    expect(spawns[0]?.env["CUDA_DEVICE_ORDER"]).toBe("PCI_BUS_ID");    expect(spawns[0]?.env["HF_HOME"]).toBe(client.paths.cacheDirectory);    expect(spawns[0]?.env["LOCAL_STUDIO_API_KEY"]).toBeUndefined();    expect(spawns[0]?.env["HF_TOKEN"]).toBeUndefined();    expect(spawns[0]?.env["LOCAL_STUDIO_SPEECH_VOICES_DIR"]).toBeUndefined();    expect(spawns[0]?.env["LOCAL_STUDIO_SPEECH_OUTPUTS_DIR"]).toBeUndefined();    expect(transport.writes[0]?.endsWith("\n")).toBe(true);    expect(JSON.parse(transport.writes[0] ?? "")).toEqual({      type: "synthesize",      id: "speech-one",      text: "Hello from Local Studio",      voice_path: realpathSync(voicePath),      output_path: result.path,    });     await client.shutdown();    expect(JSON.parse(transport.writes[1] ?? "")).toEqual({      type: "shutdown",      id: "shutdown-one",    });    expect(transport.inputClosed).toBe(true);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("fails an exited request once and restarts only on the next explicit call", async () => {  const { directory, voicePath, workerPath } = fixture();  const firstTransport = new FakeTransport();  const secondTransport = new FakeTransport();  const transports = [firstTransport, secondTransport];  let spawnCount = 0;  try {    firstTransport.onWrite = (request): void => {      if (request.type === "synthesize") firstTransport.emitExit(1, null);    };    secondTransport.onWrite = (request): void => {      if (request.type === "synthesize") {        secondTransport.emitLine({          type: "synthesize",          id: request.id,          output_path: request.output_path,          sample_rate: 24000,        });      }    };    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => `speech-${spawnCount}`,      spawnWorker: (): SpeechWorkerTransport => {        const transport = transports[spawnCount];        if (!transport) throw new Error("Unexpected worker spawn");        spawnCount += 1;        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });     const firstError = await speechError(client.synthesize({ text: "First", voicePath }));    expect(firstError.message).toContain("exited with code 1");    expect(spawnCount).toBe(1);    expect(firstTransport.writes).toHaveLength(1);     expect((await client.synthesize({ text: "Second", voicePath })).sampleRate).toBe(24000);    expect(spawnCount).toBe(2);    expect(firstTransport.writes).toHaveLength(1);    expect(secondTransport.writes).toHaveLength(1);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("does not finish shutdown until the worker has released its process", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  transport.exitOnClose = false;  transport.exitOnKill = false;  try {    transport.onWrite = (request): void => {      if (request.type === "synthesize") {        transport.emitLine({          type: "synthesize",          id: request.id,          output_path: request.output_path,          sample_rate: 24000,        });      } else {        transport.emitLine({ type: "shutdown", id: request.id });      }    };    const ids = ["lease-speech", "lease-shutdown"];    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => ids.shift() ?? "unused",      shutdownGraceMs: 1_000,      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });    await client.synthesize({ text: "Lease safe", voicePath });     let settled = false;    const shutdown = client.shutdown().then((): void => {      settled = true;    });    await Promise.resolve();    await Promise.resolve();    expect(transport.inputClosed).toBe(true);    expect(settled).toBe(false);     transport.emitExit(0, null);    await shutdown;    expect(settled).toBe(true);    expect(transport.killed).toBe(false);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("kills a timed out worker without replaying the request", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  try {    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => "speech-timeout",      synthesisTimeoutMs: 10,      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });     const error = await speechError(client.synthesize({ text: "Wait forever", voicePath }));    expect(error.code).toBe("timeout");    expect(transport.killed).toBe(true);    expect(transport.writes).toHaveLength(1);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("shutdown settles an earlier killed worker before the lease can be released", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  transport.exitOnKill = false;  try {    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => "speech-termination",      synthesisTimeoutMs: 10,      shutdownKillTimeoutMs: 1_000,      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });    expect((await speechError(client.synthesize({ text: "Timeout", voicePath }))).code).toBe(      "timeout",    );     let settled = false;    const shutdown = client.shutdown().then((): void => {      settled = true;    });    await Promise.resolve();    await Promise.resolve();    expect(settled).toBe(false);     transport.emitExit(null, "SIGKILL");    await shutdown;    expect(settled).toBe(true);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("waits for killed process exit before surfacing a close failure", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  transport.closeFailure = new Error("stdin closed");  transport.exitOnKill = false;  try {    transport.onWrite = (request): void => {      if (request.type === "synthesize") {        transport.emitLine({          type: "synthesize",          id: request.id,          output_path: request.output_path,          sample_rate: 24000,        });      } else {        transport.emitLine({ type: "shutdown", id: request.id });      }    };    const ids = ["close-speech", "close-shutdown"];    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => ids.shift() ?? "unused",      shutdownKillTimeoutMs: 1_000,      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });    await client.synthesize({ text: "Close safely", voicePath });     let settled = false;    const shutdown = client.shutdown().then(      (): SpeechWorkerError => {        settled = true;        throw new Error("Expected shutdown to fail");      },      (error: unknown): SpeechWorkerError => {        settled = true;        if (error instanceof SpeechWorkerError) return error;        throw error;      },    );    await Promise.resolve();    await Promise.resolve();    expect(settled).toBe(false);    expect(transport.killed).toBe(true);     transport.emitExit(null, "SIGKILL");    expect((await shutdown).message).toContain("stdin closed");  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("kills a worker that emits a malformed JSONL frame", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  try {    transport.onWrite = (request): void => {      if (request.type === "synthesize") transport.emitLine("not-json");    };    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => "speech-invalid",      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });     const error = await speechError(client.synthesize({ text: "Invalid frame", voicePath }));    expect(error.code).toBe("protocol");    expect(transport.killed).toBe(true);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("holds the permit until an interrupted worker confirms exit", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  transport.exitOnKill = false;  try {    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => "speech-interrupted",      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });    const fiber = Effect.runFork(client.synthesizeEffect({ text: "Interrupt me", voicePath }));    while (transport.writes.length === 0) await Promise.resolve();     let settled = false;    const interrupted = Effect.runPromise(Fiber.interrupt(fiber)).then((): void => {      settled = true;    });    await Promise.resolve();    await Promise.resolve();    expect(transport.killed).toBe(true);    expect(settled).toBe(false);     transport.emitExit(null, "SIGKILL");    await interrupted;    expect(settled).toBe(true);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("terminates an active synthesis without waiting for its permit", async () => {  const { directory, voicePath, workerPath } = fixture();  const transport = new FakeTransport();  transport.exitOnKill = false;  try {    const client = new ChatterboxWorkerClient({      dataDirectory: directory,      gpuUuid: GPU_UUID,      workerPath,      randomId: (): string => "terminate-active",      spawnWorker: (): SpeechWorkerTransport => {        queueMicrotask(() => transport.emitLine(readyFrame));        return transport;      },    });    const synthesis = client.synthesize({ text: "Stop now", voicePath });    await Promise.resolve();    await Promise.resolve();     let terminated = false;    const termination = client.terminate().then((): void => {      terminated = true;    });    await expect(synthesis).rejects.toThrow("interrupted");    expect(transport.killed).toBe(true);    expect(terminated).toBe(false);     transport.emitExit(null, "SIGKILL");    await termination;    expect(terminated).toBe(true);  } finally {    rmSync(directory, { recursive: true, force: true });  }}); test("kills a real worker transport before an oversized frame can accumulate", async () => {  const transport = spawnNodeSpeechWorker({    command: process.execPath,

Lines 1–500 of 521. Load more