refactor(livekit): Record audio through shared audio pipeline

.changeset/livekit-audio-pipeline.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+refactor(livekit): Record audio through shared audio pipeline

js/src/audio-middleware.test.ts

@@ -0,0 +1,194 @@
+import { expect, test, vi } from "vitest";
+import {
+  runMiddleware,
+  type AudioMiddleware,
+  type AudioSegmentMeta,
+} from "./audio-middleware";
+
+const meta: AudioSegmentMeta = {
+  index: 3,
+  startOffsetMs: 1000,
+  durationMs: 2,
+  sampleRate: 1000,
+  channels: [{ sources: ["a"] }, { sources: ["b"] }],
+  mimeType: "audio/pcm",
+  encoding: "pcm",
+  attributes: {
+    "pcm.sample_format": "s16le",
+    "pcm.channel_layout": "interleaved",
+  },
+};
+/** Two stereo frames: a = [2, 4], b = [6, 8]. */
+const pcm = () => new Uint8Array(Int16Array.from([2, 6, 4, 8]).buffer);
+const samples = (data: Uint8Array) =>
+  Array.from(new Int16Array(data.slice().buffer));
+
+test("without middleware, the segment and metadata come back unchanged", async () => {
+  const result = await runMiddleware([], pcm(), meta);
+  expect(samples(result.segment)).toEqual([2, 6, 4, 8]);
+  expect(result.meta).toEqual(meta);
+});
+
+test("middleware runs outermost first, transforms input and observes results on the way back", async () => {
+  const calls: string[] = [];
+  const outer: AudioMiddleware = async (segment, meta, next) => {
+    calls.push(`outer in ${meta.encoding}`);
+    const result = await next(segment, meta);
+    calls.push(`outer out ${result.meta.encoding}`);
+    return result;
+  };
+  const reverse: AudioMiddleware = async (segment, meta, next) => {
+    calls.push(`reverse in ${meta.encoding}`);
+    const result = await next(segment.slice().reverse(), {
+      ...meta,
+      encoding: "reversed",
+      mimeType: "application/octet-stream",
+    });
+    calls.push(`reverse out ${result.meta.encoding}`);
+    return result;
+  };
+  const result = await runMiddleware([outer, reverse], pcm(), meta);
+  expect(calls).toEqual([
+    "outer in pcm",
+    "reverse in pcm",
+    "reverse out reversed",
+    "outer out reversed",
+  ]);
+  expect([...result.segment]).toEqual([...pcm()].reverse());
+  expect(result.meta).toMatchObject({
+    index: 3,
+    startOffsetMs: 1000,
+    encoding: "reversed",
+    mimeType: "application/octet-stream",
+  });
+});
+
+test("middleware can short-circuit the rest of the stack", async () => {
+  const skipped = vi.fn<AudioMiddleware>();
+  const replace: AudioMiddleware = (_segment, meta) => ({
+    segment: Uint8Array.of(7),
+    meta: { ...meta, encoding: "replaced" },
+  });
+  const result = await runMiddleware([replace, skipped], pcm(), meta);
+  expect(skipped).not.toHaveBeenCalled();
+  expect([...result.segment]).toEqual([7]);
+  expect(result.meta.encoding).toBe("replaced");
+});
+
+test("attributes and state pass custom values to later middleware and the result", async () => {
+  let seen: unknown[] = [];
+  const tag: AudioMiddleware = (segment, meta, next) =>
+    next(segment, {
+      ...meta,
+      attributes: { ...meta.attributes, "acme.key_id": "k-1" },
+      state: { ...meta.state, "acme.key": Uint8Array.of(1, 2) },
+    });
+  const read: AudioMiddleware = (segment, meta, next) => {
+    seen = [meta.attributes?.["acme.key_id"], meta.state?.["acme.key"]];
+    return next(segment, meta);
+  };
+  const result = await runMiddleware([tag, read], pcm(), meta);
+  expect(seen).toEqual(["k-1", Uint8Array.of(1, 2)]);
+  expect(result.meta).toMatchObject({
+    attributes: { "acme.key_id": "k-1" },
+    state: { "acme.key": Uint8Array.of(1, 2) },
+  });
+});
+
+test("segments passed down may take any form; the stack must end with bytes", async () => {
+  type Words = { words: string[] };
+  const split: AudioMiddleware<Uint8Array, Words> = (segment, meta, next) =>
+    next({ words: ["two", "frames", String(segment.length)] }, meta);
+  const join: AudioMiddleware<Words, Uint8Array> = (segment, meta, next) =>
+    next(new TextEncoder().encode(segment.words.join(" ")), {
+      ...meta,
+      mimeType: "text/plain",
+    });
+  const result = await runMiddleware([split, join], pcm(), meta);
+  expect(new TextDecoder().decode(result.segment)).toBe("two frames 8");
+  await expect(runMiddleware([split], pcm(), meta)).rejects.toThrow(
+    "invalid_audio_middleware_result",
+  );
+});
+
+test("middleware that changes the layout describes it in the metadata", async () => {
+  // Averages the two channels into one, keeping every source on channel 0.
+  const mixdown: AudioMiddleware = (segment, meta, next) => {
+    const stereo = new Int16Array(segment.slice().buffer);
+    const mono = Int16Array.from(
+      { length: stereo.length / 2 },
+      (_, i) => (stereo[i * 2] + stereo[i * 2 + 1]) / 2,
+    );
+    return next(new Uint8Array(mono.buffer), {
+      ...meta,
+      channels: [{ sources: meta.channels.flatMap((c) => c.sources) }],
+    });
+  };
+  const result = await runMiddleware([mixdown], pcm(), meta);
+  expect(samples(result.segment)).toEqual([4, 6]);
+  expect(result.meta.channels).toEqual([{ sources: ["a", "b"] }]);
+});
+
+test.each<[string, AudioMiddleware]>([
+  [
+    "throws",
+    () => {
+      throw new Error("boom");
+    },
+  ],
+  [
+    "rejects",
+    async () => {
+      throw new Error("boom");
+    },
+  ],
+  [
+    "mutates metadata",
+    (segment, meta, next) => {
+      (meta as { encoding: string }).encoding = "changed";
+      return next(segment, meta);
+    },
+  ],
+  [
+    "calls next twice",
+    async (segment, meta, next) => {
+      await next(segment, meta);
+      return next(segment, meta);
+    },
+  ],
+  [
+    "changes timeline placement",
+    (segment, meta, next) =>
+      next(segment, { ...meta, startOffsetMs: meta.startOffsetMs + 1 }),
+  ],
+  [
+    "changes segment identity",
+    async (segment, meta) => ({
+      segment,
+      meta: { ...meta, index: meta.index + 1 },
+    }),
+  ],
+  [
+    "returns a non-byte segment",
+    async (_segment, meta) => ({
+      segment: "pcm" as unknown as Uint8Array,
+      meta,
+    }),
+  ],
+  ["returns nothing", (async () => undefined) as unknown as AudioMiddleware],
+])("middleware that %s rejects the run", async (_name, failing) => {
+  await expect(runMiddleware([failing], pcm(), meta)).rejects.toThrow();
+});
+
+test("downstream failures reach middleware that ignores next() without unhandled rejections", async () => {
+  const detached: AudioMiddleware = (segment, meta, next) => {
+    void next(segment, meta);
+    return { segment, meta };
+  };
+  const failing: AudioMiddleware = () => {
+    throw new Error("boom");
+  };
+  const result = await runMiddleware([detached, failing], pcm(), meta);
+  await new Promise((resolve) => setTimeout(resolve, 0));
+  expect(samples(result.segment)).toEqual([2, 6, 4, 8]);
+});

js/src/audio-middleware.ts

@@ -0,0 +1,147 @@
+import { isObject } from "../util";
+
+/**
+ * Describes what a segment is and where it sits on the recording timeline, in
+ * terms any audio format shares. How it is encoded belongs in the segment.
+ */
+export interface AudioSegmentMeta {
+  /** Zero-based position of this segment within its recording. */
+  readonly index: number;
+  /** Segment start on the recording timeline. Transforms must preserve it. */
+  readonly startOffsetMs: number;
+  /** Playable duration of the segment. */
+  readonly durationMs: number;
+  readonly sampleRate: number;
+  /**
+   * Every output channel in order, each with the sources heard in it at any
+   * point in this segment (empty when none were).
+   */
+  readonly channels: readonly AudioChannel[];
+  /** Media or container type of the segment bytes, such as `audio/pcm`. */
+  readonly mimeType: string;
+  /** Audio codec, such as `pcm` or `opus`. */
+  readonly encoding: string;
+  /**
+   * Format details of the resulting file, published with it, such as
+   * `pcm.sample_format`. Use namespaced keys (`"acme.key_id"`).
+   */
+  readonly attributes?: Readonly<Record<string, AudioAttributeValue>>;
+  /** Private values for later middleware, including binary data. Never published. */
+  readonly state?: Readonly<Record<string, unknown>>;
+}
+
+type AudioAttributeValue =
+  | string
+  | number
+  | boolean
+  | null
+  | readonly AudioAttributeValue[]
+  | { readonly [key: string]: AudioAttributeValue };
+
+/** One output channel: the recorder source keys, as given by the integration, it carries. */
+type AudioChannel = { readonly sources: readonly string[] };
+
+export type AudioResult = { segment: Uint8Array; meta: AudioSegmentMeta };
+
+/**
+ * Transforms recorded audio segments, Rack/Koa style. Call `next` at most once
+ * to run the rest of the stack; it resolves to the final segment and metadata.
+ * Metadata is frozen: pass or return a new object, such as `{ ...meta }`.
+ *
+ * The stack receives PCM bytes. Each layer may pass a different segment type
+ * down (`Down`), such as codec packets for a container, but the stack must end
+ * with bytes, and results returned up are always bytes.
+ *
+ * Middleware runs once per segment on the main thread, one segment at a time per
+ * recording, while recordings run concurrently. Keep per-segment state inside
+ * the call. A throw or rejection omits that segment and stops the recording.
+ *
+ * @example
+ * ```ts
+ * const logEncoding: AudioMiddleware = async (segment, meta, next) => {
+ *   const result = await next(segment, meta);
+ *   console.log(meta.encoding, "->", result.meta.encoding);
+ *   return result;
+ * };
+ * ```
+ */
+export type AudioMiddleware<In = Uint8Array, Down = In> = (
+  segment: In,
+  meta: AudioSegmentMeta,
+  next: (segment: Down, meta: AudioSegmentMeta) => Promise<AudioResult>,
+) => AudioResult | Promise<AudioResult>;
+
+/** A stack of middleware whose segment types chain from layer to layer. */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any -- Each layer has its own segment types.
+export type AudioMiddlewareStack = readonly AudioMiddleware<any, any>[];
+
+/**
+ * Runs one segment through the stack and resolves to the bytes and metadata
+ * that come back out. Rejects if a layer throws or breaks the contract.
+ */
+export function runMiddleware(
+  stack: AudioMiddlewareStack,
+  segment: Uint8Array,
+  meta: AudioSegmentMeta,
+): Promise<AudioResult> {
+  const original = freeze(meta);
+  const run = async (
+    layer: number,
+    segment: unknown,
+    meta: unknown,
+  ): Promise<AudioResult> => {
+    const middleware = stack[layer];
+    // The end of the stack returns what reached it, which must be bytes.
+    if (!middleware) return validate({ segment, meta }, original);
+    if (!preservesIdentity(meta, original))
+      throw new Error("invalid_audio_middleware_result");
+    let called = false;
+    const next = (segment: unknown, meta: AudioSegmentMeta) => {
+      if (called) throw new Error("next() called multiple times");
+      called = true;
+      const result = run(layer + 1, segment, meta);
+      // Middleware may leave `next` unawaited; it still sees the rejection if awaited.
+      result.catch(() => {});
+      return result;
+    };
+    return validate(await middleware(segment, freeze(meta), next), original);
+  };
+  return run(0, segment, original);
+}
+
+function freeze(meta: AudioSegmentMeta): AudioSegmentMeta {
+  if (Object.isFrozen(meta)) return meta;
+  return Object.freeze({
+    ...meta,
+    channels: Object.freeze(
+      meta.channels.map((c) =>
+        Object.freeze({ sources: Object.freeze([...c.sources]) }),
+      ),
+    ),
+    ...(meta.attributes && {
+      attributes: Object.freeze({ ...meta.attributes }),
+    }),
+    ...(meta.state && { state: Object.freeze({ ...meta.state }) }),
+  });
+}
+
+/** Segment identity and timeline placement survive every transform. */
+function validate(result: unknown, original: AudioSegmentMeta): AudioResult {
+  const segment = isObject(result) ? result.segment : undefined;
+  const meta = isObject(result) ? result.meta : undefined;
+  if (!(segment instanceof Uint8Array) || !preservesIdentity(meta, original))
+    throw new Error("invalid_audio_middleware_result");
+  return { segment, meta: freeze(meta) };
+}
+
+function preservesIdentity(
+  meta: unknown,
+  original: AudioSegmentMeta,
+): meta is AudioSegmentMeta {
+  return (
+    isObject(meta) &&
+    meta.index === original.index &&
+    meta.startOffsetMs === original.startOffsetMs &&
+    Array.isArray(meta.channels)
+  );
+}