diff --git a/.changeset/audio-stream-transformation-hook.md b/.changeset/audio-stream-transformation-hook.md new file mode 100644 index 000000000..9c630ca75 --- /dev/null +++ b/.changeset/audio-stream-transformation-hook.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Add audio stream transformation hook diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 5ffd89823..e9eecbc6d 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -1009,7 +1009,7 @@ }, { "scenarioDirName": "livekit-instrumentation", - "label": "LiveKit cascade instrumentation and PCM recording (streaming STT, tools, interruption recovery; automatic and manual)", + "label": "LiveKit cascade instrumentation, PCM recording, and attachment stream hooks (streaming STT, tools, interruption recovery; automatic and manual)", "metadataScenario": "livekit-instrumentation", "variants": [ { diff --git a/e2e/scenarios/livekit-instrumentation/scenario.mjs b/e2e/scenarios/livekit-instrumentation/scenario.mjs index 9681db5b0..339d1298d 100644 --- a/e2e/scenarios/livekit-instrumentation/scenario.mjs +++ b/e2e/scenarios/livekit-instrumentation/scenario.mjs @@ -42,7 +42,119 @@ if (process.env.LIVEKIT_SEGMENTS === "1") { return response; }); } -const logger = initLogger({ projectName: scopedName("e2e-livekit") }); +// Buffer only one file in encoder-owned state to produce a valid WAV header. +// File boundaries are chosen by this hook, independently of incoming chunks. +const hookStats = { chunks: 0, files: 0, closes: 0 }; +const attachmentStreamHook = + process.env.LIVEKIT_STREAM_HOOK === "1" + ? async function* (source) { + const { format, targetFileInputBytes, signal } = source; + assert.equal(format.contentType, "audio/pcm"); + assert.equal(format.audio.sampleFormat, "s16le"); + const target = Math.max(4, Math.floor(targetFileInputBytes / 4) * 4); + let parts = []; + let inputBytes = 0; + let offsetBytes = 0; + function finishFile() { + const data = parts; + const size = inputBytes / 2; + const header = Buffer.alloc(44); + header.write("RIFF", 0); + header.writeUInt32LE(36 + size, 4); + header.write("WAVEfmt ", 8); + header.writeUInt32LE(16, 16); + header.writeUInt16LE(1, 20); + header.writeUInt16LE(format.audio.channels, 22); + header.writeUInt32LE(format.audio.sampleRate, 24); + header.writeUInt32LE( + format.audio.sampleRate * format.audio.channels, + 28, + ); + header.writeUInt16LE(format.audio.channels, 32); + header.writeUInt16LE(8, 34); + header.write("data", 36); + header.writeUInt32LE(size, 40); + const bytesPerMs = + (format.audio.sampleRate * format.audio.channels * 2) / 1000; + const timeRange = { + startMs: offsetBytes / bytesPerMs, + endMs: (offsetBytes + inputBytes) / bytesPerMs, + }; + const filename = `encoded-${hookStats.files++}.wav`; + offsetBytes += inputBytes; + inputBytes = 0; + parts = []; + return { + filename, + timeRange, + chunks: (async function* () { + yield header; + yield* data; + })(), + format: { + contentType: "audio/wav", + audio: { + channels: format.audio.channels, + sampleRate: format.audio.sampleRate, + }, + }, + }; + } + try { + for await (const chunk of source.chunks) { + hookStats.chunks++; + for (let at = 0; at < chunk.byteLength; ) { + const length = Math.min( + target - inputBytes, + chunk.byteLength - at, + ); + const pcm8 = new Uint8Array(length / 2); + const view = new DataView( + chunk.buffer, + chunk.byteOffset + at, + length, + ); + for (let i = 0; i < pcm8.length; i++) + pcm8[i] = (view.getInt16(i * 2, true) >> 8) + 128; + parts.push(pcm8); + inputBytes += length; + at += length; + if (inputBytes === target) yield finishFile(); + } + } + if (inputBytes) yield finishFile(); + } finally { + hookStats.closes++; + assert.equal(signal.aborted, false); + assert.equal(parts.length, 0); + } + } + : undefined; +const uploadedWavs = []; +if (attachmentStreamHook) { + const originalFetch = globalThis.fetch; + globalThis.fetch = async (input, init) => { + if (init?.method === "PUT" && init.body instanceof Blob) { + const bytes = Buffer.from(await init.body.arrayBuffer()); + assert.equal(bytes.toString("ascii", 0, 4), "RIFF"); + assert.equal(bytes.toString("ascii", 8, 12), "WAVE"); + assert.equal(bytes.toString("ascii", 36, 40), "data"); + assert.equal(bytes.readUInt32LE(4), bytes.length - 8); + assert.equal(bytes.readUInt32LE(40), bytes.length - 44); + assert.equal(bytes.readUInt16LE(22), 2); + assert.equal(bytes.readUInt16LE(34), 8); + uploadedWavs.push({ + samples: (bytes.length - 44) / 2, + sampleRate: bytes.readUInt32LE(24), + }); + } + return originalFetch(input, init); + }; +} +const logger = initLogger({ + projectName: scopedName("e2e-livekit"), + attachmentStreamHook, +}); const providerMode = process.env.LIVEKIT_PROVIDER === "1"; class Output extends voice.AudioOutput { samples = 0; @@ -271,6 +383,13 @@ try { if (streamingStt) await session.input.audio?.close(); await session.close(); await logger.flush(); + if (attachmentStreamHook) + console.log( + JSON.stringify({ + hookStats, + uploadedWavs, + }), + ); if (provider) assert.equal(telemetry.tracer.getProvider(), provider); // LiveKit's job process exits right after its shutdown callbacks, like this; // anything not finished by then is lost. diff --git a/e2e/scenarios/livekit-instrumentation/scenario.test.ts b/e2e/scenarios/livekit-instrumentation/scenario.test.ts index 57e6afd6c..0d7143f40 100644 --- a/e2e/scenarios/livekit-instrumentation/scenario.test.ts +++ b/e2e/scenarios/livekit-instrumentation/scenario.test.ts @@ -184,6 +184,66 @@ for (const variant of variants) { 120000, ); + for (const mode of ["auto", "wrapped"] as const) + test.sequential( + `${variant} ${mode} attachment stream hook`, + async () => { + await withScenarioHarness(async (harness) => { + const result = await harness.runNodeScenarioDir( + scenarioOptions(variant, { + mode, + cassette: "order", + env: { LIVEKIT_STREAM_HOOK: "1", LIVEKIT_SEGMENTS: "1" }, + }), + ); + const events = assertAudioTrace(harness.events()); + const root = findLatestSpan(events, "session")!; + const recordings = root.metadata?.["audio.recordings"] as { + state: string; + mime_type: string; + duration_ms: number; + sample_format?: string; + }[]; + expect(recordings.length).toBeGreaterThan(1); + for (const recording of recordings) { + expect(recording).toMatchObject({ + state: "ready", + mime_type: "audio/wav", + }); + expect(recording).not.toHaveProperty("sample_format"); + } + const attachments = ( + root.row.input as { + audio: Record; + } + ).audio; + expect(Object.values(attachments)).toHaveLength(recordings.length); + for (const attachment of Object.values(attachments)) { + expect(attachment.content_type).toBe("audio/wav"); + expect(attachment.filename).toMatch(/\.wav$/); + } + const stats = JSON.parse( + result.stdout + .split("\n") + .find((line) => line.startsWith('{"hookStats":'))!, + ); + expect(stats.hookStats).toMatchObject({ + files: recordings.length, + closes: 1, + }); + expect(stats.hookStats.chunks).toBeGreaterThan(0); + expect(stats.uploadedWavs).toHaveLength(recordings.length); + for (const [index, wav] of stats.uploadedWavs.entries()) { + expect((wav.samples / wav.sampleRate) * 1000).toBeCloseTo( + recordings[index].duration_ms, + 5, + ); + } + }); + }, + 120000, + ); + test.sequential( `${variant} recording requires attachment capture`, async () => { diff --git a/js/src/attachment-stream.test.ts b/js/src/attachment-stream.test.ts new file mode 100644 index 000000000..31e109f4d --- /dev/null +++ b/js/src/attachment-stream.test.ts @@ -0,0 +1,338 @@ +import { + beforeAll, + beforeEach, + afterEach, + expect, + expectTypeOf, + test, + vi, +} from "vitest"; +import { + ATTACHMENT_STREAM_HOOK, + AttachmentStreamProcessor, + type AttachmentStreamHook, + type AttachmentStreamSource, + type AttachmentStreamOptions, +} from "./attachment-stream"; +import { + _exportsForTestingOnly as logging, + initLogger, + startSpan, + type InitLoggerOptions, + type Logger, + type Span, +} from "./logger"; +import { configureNode } from "./node/config"; + +configureNode(); +beforeAll(() => logging.simulateLoginForTests()); +beforeEach(() => logging.useTestBackgroundLogger()); +afterEach(() => { + logging.clearTestBackgroundLogger(); + vi.restoreAllMocks(); +}); + +const format = { contentType: "application/octet-stream" }; +async function* bytes(...values: number[][]) { + for (const value of values) yield new Uint8Array(value); +} + +test("calls the factory once per stream and drains each file before requesting another", async () => { + const events: string[] = []; + const hook = vi.fn(async function* (source: AttachmentStreamSource) { + expect(source.format).toEqual(format); + expect(source.targetFileInputBytes).toBe(4); + let file = 0; + try { + for await (const chunk of source.chunks) { + const index = file++; + events.push(`file ${index}`); + yield { + format, + chunks: (async function* () { + yield chunk.subarray(0, 1); + yield chunk.subarray(1); + events.push(`drained ${index}`); + })(), + }; + } + } finally { + events.push("closed"); + } + }); + const output: number[][] = []; + for (let i = 0; i < 2; i++) { + const processor = new AttachmentStreamProcessor(hook, format, 4, () => 100); + expect( + await processor.run(bytes([1, 2], [3]), async (_, data) => { + output.push([...new Uint8Array(await data.arrayBuffer())]); + events.push("uploaded"); + }), + ).toBe(true); + expect(processor.signal.aborted).toBe(false); + expect(processor.bufferedByteLength).toBe(0); + } + expect(hook).toHaveBeenCalledTimes(2); + expect(output).toEqual([[1, 2], [3], [1, 2], [3]]); + expect(events.slice(0, 7)).toEqual([ + "file 0", + "drained 0", + "uploaded", + "file 1", + "drained 1", + "uploaded", + "closed", + ]); +}); + +test("permits output before all input is available and a file larger than the target", async () => { + let consumed = 0; + const processor = new AttachmentStreamProcessor( + async function* ({ chunks }) { + yield { + format: { contentType: "image/png" }, + chunks: (async function* () { + yield new Uint8Array([137, 80]); + for await (const chunk of chunks) { + consumed++; + yield chunk; + } + })(), + }; + }, + format, + 1, + () => 100, + ); + const uploaded = vi.fn(async (_, data: Blob) => { + expect(consumed).toBe(2); + expect([...new Uint8Array(await data.arrayBuffer())]).toEqual([ + 137, 80, 1, 2, 3, + ]); + }); + await processor.run(bytes([1], [2, 3]), uploaded); + expect(uploaded).toHaveBeenCalledTimes(1); +}); + +test("can asynchronously decline without consuming input", async () => { + const input = bytes([1, 2]); + const processor = new AttachmentStreamProcessor( + async () => undefined, + format, + 1, + () => 100, + ); + expect(await processor.run(input, vi.fn())).toBe(false); + expect(await input.next()).toEqual({ + done: false, + value: new Uint8Array([1, 2]), + }); +}); + +test("declining after reading input fails instead of falling back to incomplete source bytes", async () => { + const processor = new AttachmentStreamProcessor( + async ({ chunks }) => { + await chunks[Symbol.asyncIterator]().next(); + return undefined; + }, + format, + 1, + () => 100, + ); + await expect(processor.run(bytes([1]), vi.fn())).rejects.toThrow( + "attachment_stream_declined_after_read", + ); + expect(processor.signal.aborted).toBe(true); +}); + +test("an empty output intentionally suppresses capture and aborts the unused source", async () => { + const processor = new AttachmentStreamProcessor( + async function* () {}, + format, + 1, + () => 100, + ); + const uploaded = vi.fn(); + expect(await processor.run(bytes([1]), uploaded)).toBe(true); + expect(uploaded).not.toHaveBeenCalled(); + expect(processor.signal.aborted).toBe(true); +}); + +test("rejects multiple consumers of the source", async () => { + const processor = new AttachmentStreamProcessor( + async function* ({ chunks }) { + chunks[Symbol.asyncIterator](); + chunks[Symbol.asyncIterator](); + }, + format, + 1, + () => 100, + ); + await expect(processor.run(bytes([1]), vi.fn())).rejects.toThrow( + "attachment_stream_already_consumed", + ); +}); + +test("bounds retained output and closes the inner, outer and source iterators on failure", async () => { + const closed: string[] = []; + const processor = new AttachmentStreamProcessor( + async function* ({ chunks, signal }) { + try { + yield { + format, + chunks: (async function* () { + try { + for await (const chunk of chunks) { + yield chunk; + yield new Uint8Array(10); + } + } finally { + expect(signal.aborted).toBe(true); + closed.push("file"); + } + })(), + }; + } finally { + closed.push("files"); + } + }, + format, + 1, + () => 8, + ); + const source = (async function* () { + try { + yield new Uint8Array([1]); + } finally { + closed.push("source"); + } + })(); + const uploaded = vi.fn(); + await expect(processor.run(source, uploaded)).rejects.toThrow( + "capture_byte_limit", + ); + expect(closed).toEqual(["source", "file", "files"]); + expect(uploaded).not.toHaveBeenCalled(); + expect(processor.bufferedByteLength).toBe(0); +}); + +test("does not close the source twice when closing output also closes input", async () => { + const sourceReturn = vi.fn(async () => ({ + done: true as const, + value: undefined, + })); + const input = { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false as const, value: new Uint8Array([1]) }), + return: sourceReturn, + }), + }; + const processor = new AttachmentStreamProcessor( + async function* ({ chunks }) { + yield { + format, + chunks: (async function* () { + for await (const _ of chunks) yield new Uint8Array(10); + })(), + }; + }, + format, + 1, + () => 8, + ); + await expect(processor.run(input, vi.fn())).rejects.toThrow( + "capture_byte_limit", + ); + expect(sourceReturn).toHaveBeenCalledTimes(1); +}); + +test("cleans up iterators whose next rejects and contains cleanup errors", async () => { + const closeFile = vi.fn(async () => { + throw new Error("private cleanup"); + }); + const closeFiles = vi.fn(async () => ({ + done: true as const, + value: undefined, + })); + const processor = new AttachmentStreamProcessor( + () => ({ + [Symbol.asyncIterator]: () => ({ + next: async () => ({ + done: false, + value: { + format, + chunks: { + [Symbol.asyncIterator]: () => ({ + next: async () => { + throw new Error("private failure"); + }, + return: closeFile, + }), + }, + }, + }), + return: closeFiles, + }), + }), + format, + 1, + () => 100, + ); + await expect(processor.run(bytes([1]), vi.fn())).rejects.toThrow( + "private failure", + ); + expect(closeFile).toHaveBeenCalledTimes(1); + expect(closeFiles).toHaveBeenCalledTimes(1); +}); + +test("logger hooks are typed and inherited without sharing configuration through SDK state", () => { + const hook: AttachmentStreamHook = async function* (source) { + expectTypeOf(source).toEqualTypeOf(); + yield { format: source.format, chunks: source.chunks }; + }; + const options: InitLoggerOptions = { attachmentStreamHook: hook }; + const first = initLogger({ + ...options, + projectId: "first", + setCurrent: false, + }); + const second = initLogger({ projectId: "second", setCurrent: false }); + expect(first.loggingState).toBe(second.loggingState); + expectTypeOf(first).toEqualTypeOf>(); + expectTypeOf( + initLogger({ + asyncFlush: false, + attachmentStreamHook: hook, + setCurrent: false, + }), + ).toEqualTypeOf>(); + const child: Span & AttachmentStreamOptions = first + .startSpan() + .startSpan() + .startSpan(); + expect(child[ATTACHMENT_STREAM_HOOK]).toBe(hook); + const other: Span & AttachmentStreamOptions = second.startSpan(); + expect(other[ATTACHMENT_STREAM_HOOK]).toBeUndefined(); +}); + +test("spans continuing an exported parent don't use a logger's hook", async () => { + const hook: AttachmentStreamHook = async () => undefined; + const logger = initLogger({ + projectId: "current", + attachmentStreamHook: hook, + }); + const exported = await initLogger({ + projectId: "exported", + setCurrent: false, + }) + .startSpan() + .export(); + const spans: (Span & AttachmentStreamOptions)[] = [ + startSpan({ parent: exported }), + logger.startSpan({ parent: exported }), + logger.startSpan().startSpan({ parent: exported }), + ]; + expect(logger[ATTACHMENT_STREAM_HOOK]).toBe(hook); + for (const span of spans) + expect(span[ATTACHMENT_STREAM_HOOK]).toBeUndefined(); +}); diff --git a/js/src/attachment-stream.ts b/js/src/attachment-stream.ts new file mode 100644 index 000000000..6007fe25e --- /dev/null +++ b/js/src/attachment-stream.ts @@ -0,0 +1,230 @@ +import { debugLogger } from "./debug-logger"; + +/** Describes incoming or transformed attachment bytes. */ +export interface AttachmentStreamFormat { + contentType: string; + audio?: { + sampleRate: number; + channels: number; + /** For example, `s16le` for signed 16-bit little-endian PCM. */ + sampleFormat?: string; + }; +} + +export interface AttachmentStreamSource { + format: AttachmentStreamFormat; + /** Ordered, instrumentation-owned bytes. Single consumer; bytes may be modified. */ + chunks: AsyncIterable; + /** Aborted when processing fails or output ends before the source does. */ + signal: AbortSignal; + /** Desired file size in input bytes. A hint, never a mandatory file boundary. */ + targetFileInputBytes: number; +} + +export interface AttachmentStreamFile { + /** Audio output must describe `audio` and keep the source's channel count. */ + format: AttachmentStreamFormat; + /** Defaults to a name derived from the content type. */ + filename?: string; + /** + * Bytes of one complete file. The SDK exhausts this iterable before requesting + * the next file. Do not modify bytes after yielding them. + */ + chunks: AsyncIterable; + /** + * Source time covered by this file, relative to the logical stream's start. + * Files must be contiguous; yield a file without bytes to omit a range. + * May be finalized while producing chunks; read after chunks are exhausted. + * Required when a stream produces more than one file. A single file + * consuming the entire source may omit it. + */ + timeRange?: { startMs: number; endMs: number }; +} + +/** + * Called once per captured logical stream. Keep encoder state in a closure or + * generator. Register with initLogger({ attachmentStreamHook }); capture remains + * opt-in, and ordinary Attachment instances are unaffected. + * + * Return undefined before consuming input to use default handling. An empty + * iterable intentionally suppresses capture. Selected hooks never fall back to + * uploading original bytes on failure; completed files remain valid. + * + * Output is currently buffered per file for upload. SDK-owned input and output + * count against capture limits; customer encoder and worker allocations do not. + * Hooks must cooperate with cancellation and settle pending work. Use finally + * for cleanup, and never await flush() from inside a hook. flush() sends + * pending logs first, then waits for uploads already started and for closed + * streams to drain; it does not finish an active stream or force the hook to + * end a file. Integrations that finalize on close, such as LiveKit sessions, + * also wait for the closed stream to drain. + * + * Spans continuing an explicit parent and experiment spans don't use hooks. + */ +export type AttachmentStreamHook = ( + source: AttachmentStreamSource, +) => + | AsyncIterable + | undefined + | Promise | undefined>; + +// Internal symbols also work when instrumentation and spans use different SDK bundles. +export const ATTACHMENT_STREAM_HOOK = Symbol.for( + "braintrust.attachment-stream-hook", +); + +export interface AttachmentStreamOptions { + readonly [ATTACHMENT_STREAM_HOOK]?: AttachmentStreamHook; +} + +/** An SDK-owned failure code, safe to report unlike customer hook errors. */ +export class AttachmentStreamError extends Error {} + +/** Consumes hook-owned file boundaries. The adapter owns source timing and capture. */ +export class AttachmentStreamProcessor { + private readonly controller = new AbortController(); + private inputBytes = 0; + private outputBytes = 0; + inputEnded = false; + + constructor( + private readonly hook: AttachmentStreamHook, + private readonly format: AttachmentStreamFormat, + private readonly targetFileInputBytes: number, + private readonly availableBytes: () => number, + ) {} + + get signal(): AbortSignal { + return this.controller.signal; + } + + get bufferedByteLength(): number { + return this.inputBytes + this.outputBytes; + } + + async run( + input: AsyncIterable, + onFile: (file: AttachmentStreamFile, data: Blob) => Promise, + ): Promise { + let inputIterator: AsyncIterator | undefined; + let files: AsyncIterator | undefined; + let chunks: AsyncIterator | undefined; + let inputClosed = false; + const source: AttachmentStreamSource = { + format: { + ...this.format, + ...(this.format.audio ? { audio: { ...this.format.audio } } : {}), + }, + signal: this.signal, + targetFileInputBytes: this.targetFileInputBytes, + chunks: { + [Symbol.asyncIterator]: () => { + if (inputIterator) + throw new AttachmentStreamError( + "attachment_stream_already_consumed", + ); + const iterator = input[Symbol.asyncIterator](); + inputIterator = iterator; + return { + next: async () => { + this.inputBytes = 0; + this.signal.throwIfAborted(); + const result = await iterator.next(); + this.signal.throwIfAborted(); + if (result.done) this.inputEnded = true; + else { + this.inputBytes = result.value.byteLength; + if (this.bufferedByteLength > this.availableBytes()) + throw new AttachmentStreamError("capture_byte_limit"); + } + return result; + }, + return: async () => { + this.inputBytes = 0; + // run() owns final cleanup; don't call return twice. + inputClosed = true; + return iterator.return + ? await iterator.return() + : { done: true, value: undefined }; + }, + }; + }, + }, + }; + let selected = false; + try { + this.signal.throwIfAborted(); + const output = await this.hook(source); + if (output === undefined) { + if (inputIterator) + throw new AttachmentStreamError( + "attachment_stream_declined_after_read", + ); + return false; + } + selected = true; + files = output[Symbol.asyncIterator](); + while (true) { + this.signal.throwIfAborted(); + const next = await files.next(); + this.signal.throwIfAborted(); + if (next.done) { + files = undefined; + break; + } + const file = next.value; + if ( + !file.format?.contentType || + (file.filename !== undefined && !file.filename) + ) + throw new AttachmentStreamError("attachment_stream_invalid_file"); + const parts: Uint8Array[] = []; + chunks = file.chunks[Symbol.asyncIterator](); + while (true) { + const part = await chunks.next(); + this.signal.throwIfAborted(); + if (part.done) { + chunks = undefined; + break; + } + if (!(part.value instanceof Uint8Array)) + throw new AttachmentStreamError("attachment_stream_invalid_chunk"); + if ( + this.bufferedByteLength + part.value.byteLength > + this.availableBytes() + ) + throw new AttachmentStreamError("capture_byte_limit"); + if (part.value.byteLength) { + parts.push(part.value); + this.outputBytes += part.value.byteLength; + } + } + const data = new Blob(parts, { type: file.format.contentType }); + parts.length = 0; + await onFile(file, data); + this.outputBytes = 0; + } + return true; + } catch (error) { + this.controller.abort(); + throw error; + } finally { + // Wake pending input reads before closing generators. Early output completion + // intentionally abandons the remainder of the source. + if (selected && !this.inputEnded) this.controller.abort(); + for (const iterator of [chunks, files, inputIterator]) { + // Closing the file generator can itself close the source. Check after + // each return(), rather than scheduling duplicate source cleanup. + if (iterator === inputIterator && inputClosed) continue; + try { + await iterator?.return?.(); + } catch { + // Customer errors may contain media or secrets. + debugLogger.warn("Attachment stream hook cleanup failed."); + } + } + this.inputBytes = 0; + this.outputBytes = 0; + } + } +} diff --git a/js/src/instrumentation/plugins/livekit-plugin.test.ts b/js/src/instrumentation/plugins/livekit-plugin.test.ts index d05361e97..c63de69c4 100644 --- a/js/src/instrumentation/plugins/livekit-plugin.test.ts +++ b/js/src/instrumentation/plugins/livekit-plugin.test.ts @@ -8,6 +8,10 @@ import { test, vi, } from "vitest"; +import type { + AttachmentStreamHook, + AttachmentStreamSource, +} from "../../attachment-stream"; import { EventEmitter } from "node:events"; import { z } from "zod/v3"; import { configureNode } from "../../node/config"; @@ -1290,11 +1294,24 @@ describe("recording files", () => { channel, }); /** A recording whose wall-clock hold no longer limits uploads. */ - function recording(chunkBytes = DEFAULT_STREAMING_ATTACHMENT_BUFFER_BYTES) { + function recording( + chunkBytes = DEFAULT_STREAMING_ATTACHMENT_BUFFER_BYTES, + hook?: AttachmentStreamHook, + captures?: { input: boolean; output: boolean }, + ) { const log = vi.fn(); const turn = vi.fn(); - const r = new Recording({ spanId: "root", log } as any, chunkBytes, (id) => - id === "turn" ? ({ span: { log: turn } } as any) : undefined, + const owner = initLogger({ + projectId: "recording-project", + attachmentStreamHook: hook, + }); + const span = owner.startSpan(); + vi.spyOn(span, "log").mockImplementation(log); + const r = new Recording( + span, + chunkBytes, + (id) => (id === "turn" ? ({ span: { log: turn } } as any) : undefined), + captures, ); now += 3_600_000; const last = (fn: typeof log, key: string) => @@ -1304,7 +1321,12 @@ describe("recording files", () => { .at(-1); return { r, + owner, log, + sources: () => + (last(log, "audio.recordings") as any[]).map((f) => + f.sources.map((source: any) => source.channel_index), + ), files: () => (last(log, "audio.recordings") as any[]).map((f) => [ f.id, @@ -1320,6 +1342,395 @@ describe("recording files", () => { ]), }; } + // A test encoder chooses its own frame-aligned boundaries, independently of + // incoming chunk sizes. Real encoders may also wait for codec-safe boundaries. + async function* segmentPcm( + source: AttachmentStreamSource, + size = source.targetFileInputBytes, + ) { + let parts: Uint8Array[] = []; + let buffered = 0; + let offset = 0; + for await (const chunk of source.chunks) { + for (let at = 0; at < chunk.length; ) { + const count = Math.min(size - buffered, chunk.length - at); + parts.push(chunk.subarray(at, at + count)); + at += count; + buffered += count; + if (buffered === size) { + const data = parts; + yield { + format: source.format, + timeRange: { + startMs: offset / 96, + endMs: (offset + buffered) / 96, + }, + chunks: (async function* () { + yield* data; + })(), + }; + offset += buffered; + buffered = 0; + parts = []; + } + } + } + if (buffered) + yield { + format: source.format, + timeRange: { startMs: offset / 96, endMs: (offset + buffered) / 96 }, + chunks: (async function* () { + yield* parts; + })(), + }; + } + + test("hooks preserve bytes, closure state, chosen file boundaries and selections", async () => { + const closed = vi.fn(); + const hook = vi.fn(async function* (source: AttachmentStreamSource) { + let count = 0; + try { + // Deliberately exceed the SDK's target; the hook owns file boundaries. + for await (const file of segmentPcm(source, 96000)) { + count++; + yield file; + } + } finally { + closed(count); + } + }); + const { r, files, selections } = recording(9600, hook); + r.select("turn", 0, 2000, 0); + r.record(packet(0, 48000)); + r.advance(); + await vi.waitFor(() => expect(uploads).toHaveLength(2)); + const expected = Buffer.alloc(48000 * 4); + for (let i = 0; i < 48000; i++) + expected.writeInt16LE(((i * 31) % 30000) - 15000, i * 4); + expect(Buffer.concat(uploads)).toEqual(expected); + await r.finish(); + expect(hook).toHaveBeenCalledTimes(1); + expect(closed).toHaveBeenCalledExactlyOnceWith(2); + expect(files()).toEqual([ + ["call-0000", "ready", 0, 1000], + ["call-0001", "ready", 1000, 1000], + ]); + expect(selections()).toEqual([ + ["call-0000", 0, 1000], + ["call-0001", 0, 1000], + ]); + }); + + test("consumes partial input incrementally but does not force a file boundary on flush", async () => { + const read = vi.fn(); + const { r, owner } = recording(9600, async function* (source) { + yield { + format: source.format, + chunks: (async function* () { + for await (const chunk of source.chunks) { + read(chunk.length); + yield chunk; + } + })(), + }; + }); + r.record(packet(0, 4800)); + r.advance(); + await vi.waitFor(() => expect(read).toHaveBeenCalled()); + await owner.flush(); + expect(uploads).toHaveLength(0); + await r.finish(); + expect(uploads[0]).toHaveLength(19200); + }); + + test("hooks are not started for recordings without audio", async () => { + const hook = vi.fn(async function* (source: AttachmentStreamSource) { + yield { format: source.format, chunks: source.chunks }; + }); + const { r, files } = recording(9600, hook); + await r.finish(); + expect(hook).not.toHaveBeenCalled(); + expect(uploads).toHaveLength(0); + expect(r.reason).toBeUndefined(); + expect(files()).toEqual([["call", "omitted", undefined, undefined]]); + }); + + test("logger flush sends pending rows, then waits for a closed stream's final processing and upload", async () => { + let release!: () => void; + const waiting = new Promise((resolve) => { + release = resolve; + }); + const { r, owner, files } = recording(96000, async function* (source) { + await waiting; + yield { format: source.format, chunks: source.chunks }; + }); + const sent = vi.spyOn(logger, "flush"); + r.record(packet(0)); + const finish = r.finish(); + const flushed = vi.fn(); + const done = owner.flush().then(flushed); + await Promise.resolve(); + expect(sent).toHaveBeenCalledTimes(1); + expect(flushed).not.toHaveBeenCalled(); + release(); + await done; + expect(sent).toHaveBeenCalledTimes(2); + await finish; + expect(files()).toEqual([["call-0000", "ready", 0, 1000]]); + }); + + test("logger flush waits for default uploads already started", async () => { + let release!: () => void; + const waiting = new Promise((resolve) => { + release = resolve; + }); + const upload = Attachment.prototype.upload; + vi.spyOn(Attachment.prototype, "upload").mockImplementationOnce( + async function (this: Attachment) { + await waiting; + return upload.call(this); + }, + ); + const owner = initLogger({ projectId: "default-project" }); + const r = new Recording(owner.startSpan(), 9600, () => undefined); + now += 3_600_000; + r.record(packet(0, 4800)); + r.advance(); + const flushed = vi.fn(); + const done = owner.flush().then(flushed); + await Promise.resolve(); + expect(flushed).not.toHaveBeenCalled(); + release(); + await done; + expect(uploads[0]).toHaveLength(9600); + await r.finish(); + }); + + test("sync-flush logging does not wait for attachment uploads", async () => { + let release!: () => void; + const waiting = new Promise((resolve) => { + release = resolve; + }); + const upload = Attachment.prototype.upload; + vi.spyOn(Attachment.prototype, "upload").mockImplementationOnce( + async function (this: Attachment) { + await waiting; + return upload.call(this); + }, + ); + const owner = initLogger({ + projectId: "sync-project", + asyncFlush: false, + }); + const r = new Recording(owner.startSpan(), 9600, () => undefined); + now += 3_600_000; + r.record(packet(0, 4800)); + r.advance(); + await owner.traced(() => undefined); + expect(uploads).toHaveLength(0); + release(); + await r.finish(); + expect(uploads[0]).toHaveLength(9600); + }); + + test.each([false, true])( + "declining a hook retains default PCM handling (async=%s)", + async (asyncDecline) => { + const hook = vi.fn(() => + asyncDecline ? Promise.resolve(undefined) : undefined, + ); + const { r, files } = recording(9600, hook); + r.record(packet(0, 4800)); + await r.finish(); + expect(hook).toHaveBeenCalledTimes(1); + expect(files()).toEqual([ + ["call-0000", "ready", 0, 100], + ["call-0001", "ready", 100, 100], + ]); + expect(uploads.map((b) => b.length)).toEqual([9600, 9600]); + }, + ); + + test.each([ + ["chunk", "attachment_stream_hook_failed"], + ["files", "attachment_stream_hook_failed"], + ["format", "attachment_stream_invalid_audio"], + ["range", "attachment_stream_invalid_audio"], + ["gap", "attachment_stream_invalid_audio"], + ])( + "hook %s failure stops capture without source fallback", + async (failure, reason) => { + const closed = vi.fn(); + const { r, files, selections } = recording( + 9600, + async function* (source) { + try { + if (failure === "files") throw new Error("private payload"); + yield { + format: + failure === "format" + ? { contentType: "audio/flac" } + : source.format, + timeRange: + failure === "range" + ? { startMs: 100, endMs: 0 } + : failure === "gap" + ? { startMs: 50, endMs: 100 } + : undefined, + chunks: (async function* () { + for await (const chunk of source.chunks) { + if (failure === "chunk") throw new Error("private payload"); + yield chunk; + } + })(), + }; + } finally { + closed(); + } + }, + ); + r.select("turn", 0, 100, 0); + r.record(packet(0, 4800)); + await r.finish(); + expect(r.active).toBe(false); + expect(uploads).toHaveLength(0); + expect(files().every(([, state]) => state === "omitted")).toBe(true); + expect(selections() ?? []).toEqual([]); + expect(r.reason).toBe(reason); + expect(closed).toHaveBeenCalledTimes(1); + }, + ); + + test("can omit one file and retain its source time for the next file", async () => { + const { r, files } = recording(9600, async function* (source) { + let index = 0; + for await (const file of segmentPcm(source)) { + yield index++ ? file : { ...file, chunks: (async function* () {})() }; + } + }); + r.record(packet(0, 4800)); + await r.finish(); + expect(files()).toEqual([ + ["call-0000", "omitted", 0, undefined], + ["call-0001", "ready", 100, 100], + ]); + expect(uploads).toHaveLength(1); + expect(r.reason).toBeUndefined(); + }); + + test("preserves leading silence so hook timing uses the session clock", async () => { + const { r, files, selections } = recording(9600, segmentPcm); + r.select("turn", 100, 200, 0); + r.record(packet(100, 2400)); + await r.finish(); + expect(uploads[0]).toEqual(new Uint8Array(9600)); + expect(files()).toEqual([ + ["call-0000", "ready", 0, 100], + ["call-0001", "ready", 100, 100], + ]); + expect(selections()).toEqual([["call-0001", 0, 100]]); + }); + + test("slow hooks stop capturing at the capture limit and finish with the audio captured so far", async () => { + let release!: () => void; + const released = new Promise((resolve) => { + release = resolve; + }); + const waiting = vi.fn(); + const closed = vi.fn(); + let consumed = 0; + const { r, files } = recording(8 * 1024 * 1024, async function* (source) { + try { + yield { + format: source.format, + chunks: (async function* () { + for await (const chunk of source.chunks) { + waiting(); + await released; + consumed += chunk.length; + } + yield new Uint8Array(1); + })(), + }; + } finally { + closed(source.signal.aborted); + } + }); + const samples = 2 * 1024 * 1024; + let recorded = 1; + r.record(packet(0, samples)); + r.advance(); + await vi.waitFor(() => expect(waiting).toHaveBeenCalled()); + while (r.record(packet((recorded * samples) / 24, samples))) recorded++; + expect(r.reason).toBe("capture_byte_limit"); + release(); + await r.finish(); + expect(consumed).toBe(recorded * samples * 4); + expect(uploads).toHaveLength(1); + expect(files()[0]).toEqual([ + "call-0000", + "ready", + 0, + (recorded * samples) / 24, + ]); + expect(closed).toHaveBeenCalledExactlyOnceWith(false); + }); + + test.each([ + [{ input: true, output: false }, [0]], + [{ input: false, output: true }, [1]], + [{ input: true, output: true }, [0, 1]], + ])( + "hook files report the captured directions as sources (%o)", + async (captures, channels) => { + const { r, sources } = recording(9600, segmentPcm, captures); + r.record(packet(0, 2400, channels[0])); + await r.finish(); + expect(sources()).toEqual([channels]); + }, + ); + + test("a soft stop drains the captured audio into the hook's unfinished file", async () => { + const { r, files } = recording(96000, async function* (source) { + yield { format: source.format, chunks: source.chunks }; + }); + r.record(packet(0, 4800)); + r.advance(); + r.stop("unfinished_output_segment"); + await r.finish(); + expect(uploads.map((b) => b.length)).toEqual([4800 * 4]); + expect(files()[0]).toEqual(["call-0000", "ready", 0, 200]); + expect(r.reason).toBe("unfinished_output_segment"); + }); + + test("failure in a later file preserves previously completed attachments", async () => { + const { r, files } = recording(9600, async function* (source) { + for await (const file of segmentPcm(source)) { + yield file; + throw new Error("private payload"); + } + }); + r.record(packet(0, 4800)); + await r.finish(); + expect(uploads).toHaveLength(1); + expect(files()[0]).toEqual(["call-0000", "ready", 0, 100]); + expect(r.reason).toBe("attachment_stream_hook_failed"); + }); + + test("ending output before the input ends truncates the recording", async () => { + const { r, files } = recording(9600, async function* (source) { + for await (const file of segmentPcm(source)) { + yield file; + return; + } + }); + r.record(packet(0, 4800)); + await r.finish(); + expect(uploads).toHaveLength(1); + expect(files()[0]).toEqual(["call-0000", "ready", 0, 100]); + expect(r.reason).toBe("attachment_stream_output_ended"); + }); + test("publishes attachment references and selections only after upload succeeds", async () => { let complete!: () => void; const blocked = new Promise((resolve) => { diff --git a/js/src/instrumentation/plugins/livekit-plugin.ts b/js/src/instrumentation/plugins/livekit-plugin.ts index 6218632f9..2c191d9f1 100644 --- a/js/src/instrumentation/plugins/livekit-plugin.ts +++ b/js/src/instrumentation/plugins/livekit-plugin.ts @@ -22,7 +22,14 @@ import { runInstrumentation } from "../core/observe-result"; import { livekitChannels } from "./livekit-interceptors"; import { v4 as uuidv4 } from "uuid"; import { runWithAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; -import { startSpan, withCurrent, type Span } from "../../logger"; +import { Attachment, startSpan, withCurrent, type Span } from "../../logger"; +import { + ATTACHMENT_STREAM_HOOK, + AttachmentStreamError, + AttachmentStreamProcessor, + type AttachmentStreamOptions, + type AttachmentStreamFile, +} from "../../attachment-stream"; import { withSpanInstrumentationName, INSTRUMENTATION_NAMES, @@ -30,7 +37,10 @@ import { import { isObject } from "../../../util"; import { zodToJsonSchema } from "../../zod/utils"; import { debugLogger } from "../../debug-logger"; -import { isEnvFlagEnabled } from "../../wrappers/attachment-utils"; +import { + getExtensionFromMediaType, + isEnvFlagEnabled, +} from "../../wrappers/attachment-utils"; import { DEFAULT_STREAMING_ATTACHMENT_BUFFER_BYTES } from "../config"; import { PcmAudio, @@ -1142,7 +1152,13 @@ type AudioFile = { start: number; end: number; channels: number[]; - metadata: NonNullable>["metadata"]; + metadata: { + mime_type: string; + sample_rate_hz: number; + channel_count: number; + sample_format?: string; + duration_ms: number; + }; state: "pending" | "ready" | "omitted"; reason?: string; }; @@ -1156,6 +1172,10 @@ class Recording { /** Session-clock playout that each pending output run may still add audio before. */ readonly holds = new Map(); private pcm: PcmAudio; + private processor?: AttachmentStreamProcessor; + private transformJob?: Promise; + private inputChanged?: () => void; + private consumedThroughMs = 0; private files: AudioFile[] = []; private formats = [ new Map(), @@ -1188,6 +1208,30 @@ class Recording { Math.max(32 * 1024 * 1024, 2 * chunkBytes), ), }); + const hook = (span as Span & AttachmentStreamOptions)[ + ATTACHMENT_STREAM_HOOK + ]; + if (hook) + this.processor = new AttachmentStreamProcessor( + hook, + { + contentType: "audio/pcm", + audio: { + sampleRate: SAMPLE_RATE, + channels: 2, + sampleFormat: "s16le", + }, + }, + chunkBytes, + () => this.availableBytes, + ); + } + + /** Capture budget left after buffered and staged PCM; the processor adds its own bytes. */ + private get availableBytes(): number { + return ( + this.pcm.maxBufferBytes - this.pcm.bufferedByteLength - this.stagedBytes + ); } get active(): boolean { @@ -1195,6 +1239,8 @@ class Recording { } stop(reason: string): void { this.reason ??= reason; + // Capture has ended: a hook drains what was already captured, like finish(). + this.inputChanged?.(); } /** Validates a LiveKit frame, returning its samples while capture continues. */ @@ -1230,8 +1276,8 @@ class Recording { ((packet.pcm.length / packet.channels) * SAMPLE_RATE) / packet.rate, ) * 2; if ( - this.pcm.bufferedByteLength + this.stagedBytes + bytes > - this.pcm.maxBufferBytes + bytes + (this.processor?.bufferedByteLength ?? 0) > + this.availableBytes ) { this.stop("capture_byte_limit"); return false; @@ -1259,8 +1305,8 @@ class Recording { const data = this.samples(frame); if (!data) return; if ( - this.stagedBytes + this.pcm.bufferedByteLength + data.byteLength > - this.pcm.maxBufferBytes + data.byteLength + (this.processor?.bufferedByteLength ?? 0) > + this.availableBytes ) { this.stop("capture_byte_limit"); return; @@ -1377,15 +1423,22 @@ class Recording { } finish(): Promise { - return (this.finishing ??= (async () => { - this.closed = true; - this.upload(); - while (this.job) await this.job; - this.publishFiles(); - })()); + return (this.finishing ??= this.track( + (async () => { + this.closed = true; + this.upload(); + await this.transformJob; + while (this.job) await this.job; + this.publishFiles(); + })(), + )); } private upload(): void { + if (this.processor) { + this.uploadTransformed(); + return; + } if ( this.job || !this.pcm.bufferedByteLength || @@ -1393,15 +1446,14 @@ class Recording { ) return; const id = `call-${String(this.files.length).padStart(4, "0")}`; - let result: ReturnType; + let result: ReturnType; try { - result = this.pcm.digest({ + result = this.pcm.read({ maxBytes: this.chunkBytes, endOffsetMs: this.closed ? undefined : Math.min(this.confirmedThroughMs, this.pcm.endOffsetMs), flush: this.closed, - filename: `${id}.pcm`, }); } catch (error) { this.stop( @@ -1410,7 +1462,6 @@ class Recording { return; } if (!result) return; - const { attachment } = result; const file: AudioFile = { id, start: result.startOffsetMs, @@ -1421,35 +1472,196 @@ class Recording { }; this.files.push(file); this.publishFiles(); - // One ordinary SDK attachment upload at a time; capture continues into the next digest. - this.job = (async () => { - try { - const status = await attachment.upload(); - file.state = status.upload_status === "done" ? "ready" : "omitted"; - } catch { - file.state = "omitted"; - } - if (file.state === "ready") { - this.span.log({ input: { audio: { [id]: attachment.reference } } }); - } else { - file.reason = "recording_export_failed"; - this.stop(file.reason); - } - this.publishFiles(); - if (file.state === "ready") - this.publish( - [...this.selections] - .filter(([, list]) => - list.some((s) => s.start < file.end && s.end > file.start), - ) - .map(([owner]) => owner), - ); - })().finally(() => { + // One ordinary SDK attachment upload at a time; capture continues into the next read. + const attachment = new Attachment({ + data: result.data, + filename: `${id}.pcm`, + contentType: "audio/pcm", + state: this.span.state(), + }); + this.job = this.track(this.uploadFile(file, attachment)).finally(() => { this.job = undefined; this.upload(); }); } + private track(job: Promise): Promise { + const jobs = this.span.state().attachmentJobs; + jobs.add(job); + void job.then( + () => jobs.delete(job), + () => jobs.delete(job), + ); + return job; + } + + /** A continuous PCM clock, including silence, independent of output file boundaries. */ + private async *sourceChunks(): AsyncGenerator { + const signal = this.processor!.signal; + const wake = () => this.inputChanged?.(); + signal.addEventListener("abort", wake); + try { + while (true) { + signal.throwIfAborted(); + const ended = !this.active; + const endOffsetMs = ended + ? this.pcm.endOffsetMs + : Math.min(this.confirmedThroughMs, this.pcm.endOffsetMs); + if (endOffsetMs > this.consumedThroughMs) { + const chunk = this.pcm.read({ + maxBytes: 64 * 1024, + endOffsetMs, + flush: true, + continuous: true, + }); + if (chunk) { + this.consumedThroughMs = chunk.endOffsetMs; + yield chunk.data; + continue; + } + } + if (ended) return; + await new Promise((resolve) => { + this.inputChanged = resolve; + }); + this.inputChanged = undefined; + } + } finally { + signal.removeEventListener("abort", wake); + this.inputChanged = undefined; + } + } + + private uploadTransformed(): void { + this.inputChanged?.(); + if (this.transformJob || !this.pcm.bufferedByteLength) return; + const processor = this.processor!; + this.transformJob = (async () => { + try { + const handled = await processor.run( + this.sourceChunks(), + (output, data) => this.uploadTransformedFile(processor, output, data), + ); + if (!handled) { + this.processor = undefined; + this.upload(); + return; + } + if (!processor.inputEnded) { + this.stop("attachment_stream_output_ended"); + this.pcm.discard(); + } + } catch (error) { + // Never log customer exception payloads or fall back to the original bytes. + this.stop( + error instanceof AttachmentStreamError + ? error.message + : "attachment_stream_hook_failed", + ); + this.pcm.discard(); + debugLogger.warn( + "Attachment stream processing failed; capture stopped.", + ); + } finally { + this.publishFiles(); + } + })(); + } + + private async uploadTransformedFile( + processor: AttachmentStreamProcessor, + output: AttachmentStreamFile, + data: Blob, + ): Promise { + const audio = output.format.audio; + // Explicit source ranges support encoder read-ahead. Files are contiguous so + // every gap is reported; omissions are files without bytes. Only a + // whole-source file can infer its range from input exhaustion. + let range = output.timeRange; + if (range == null && processor.inputEnded && !this.files.length) { + range = { startMs: 0, endMs: this.consumedThroughMs }; + } + if ( + !audio || + audio.channels !== 2 || + !Number.isSafeInteger(audio.sampleRate) || + audio.sampleRate <= 0 || + !range || + !Number.isFinite(range.startMs) || + !Number.isFinite(range.endMs) || + Math.abs(range.startMs - (this.files.at(-1)?.end ?? 0)) > 1e-6 || + range.endMs <= range.startMs || + range.endMs > this.consumedThroughMs + 1e-6 + ) + throw new AttachmentStreamError("attachment_stream_invalid_audio"); + const id = `call-${String(this.files.length).padStart(4, "0")}`; + const file: AudioFile = { + id, + start: range.startMs, + end: range.endMs, + // The normalized stereo source keeps one channel per captured direction. + channels: [0, 1].filter((channel) => + channel ? this.captures.output : this.captures.input, + ), + metadata: { + mime_type: output.format.contentType, + sample_rate_hz: audio.sampleRate, + channel_count: audio.channels, + ...(audio.sampleFormat ? { sample_format: audio.sampleFormat } : {}), + duration_ms: range.endMs - range.startMs, + }, + state: data.size ? "pending" : "omitted", + }; + if (!data.size) file.reason = "attachment_stream_hook_omitted"; + this.files.push(file); + this.publishFiles(); + if (!data.size) return; + // flush() waits for completed files being uploaded, not an open input + // stream. finish() additionally registers the stream's final drain. + await this.track( + this.uploadFile( + file, + new Attachment({ + data, + filename: + output.filename ?? + `${id}.${getExtensionFromMediaType(output.format.contentType)}`, + contentType: output.format.contentType, + state: this.span.state(), + }), + ), + ); + if (file.state === "omitted") + throw new AttachmentStreamError("recording_export_failed"); + } + + private async uploadFile( + file: AudioFile, + attachment: Attachment, + ): Promise { + try { + const status = await attachment.upload(); + file.state = status.upload_status === "done" ? "ready" : "omitted"; + } catch { + file.state = "omitted"; + } + if (file.state === "ready") { + this.span.log({ input: { audio: { [file.id]: attachment.reference } } }); + } else { + file.reason = "recording_export_failed"; + this.stop(file.reason); + } + this.publishFiles(); + if (file.state === "ready") + this.publish( + [...this.selections] + .filter(([, list]) => + list.some((s) => s.start < file.end && s.end > file.start), + ) + .map(([owner]) => owner), + ); + } + private publishFiles(): void { const recordings: Record[] = this.files.map((f) => ({ id: f.id, diff --git a/js/src/logger.ts b/js/src/logger.ts index b124ce5c9..c54ae6a82 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -1,4 +1,9 @@ /// +import { + ATTACHMENT_STREAM_HOOK, + type AttachmentStreamHook, + type AttachmentStreamOptions, +} from "./attachment-stream"; import { TypedAttachmentProcessor } from "./typed-attachments"; import { @@ -763,6 +768,8 @@ export class BraintrustState { private _apiConn: HTTPConnection | null = null; private _proxyConn: HTTPConnection | null = null; + /** Attachment work started by this state's spans, awaited by flush(). */ + public readonly attachmentJobs = new Set>(); public promptCache: PromptCache; public parametersCache: ParametersCache; public spanCache: SpanCache; @@ -1880,6 +1887,7 @@ export interface LinkArgs { export interface LogOptions { asyncFlush?: IsAsyncFlush; + attachmentStreamHook?: AttachmentStreamHook; computeMetadataArgs?: Record; linkArgs?: LinkArgs; } @@ -2899,6 +2907,7 @@ function startSpanParentArgs(args: { } export class Logger implements Exportable { + readonly [ATTACHMENT_STREAM_HOOK]?: AttachmentStreamHook; private state: BraintrustState; private lazyMetadata: LazyValue; private _asyncFlush: IsAsyncFlush | undefined; @@ -2917,6 +2926,7 @@ export class Logger implements Exportable { lazyMetadata: LazyValue, logOptions: LogOptions = {}, ) { + this[ATTACHMENT_STREAM_HOOK] = logOptions.attachmentStreamHook; this.lazyMetadata = lazyMetadata; this._asyncFlush = logOptions.asyncFlush; this.computeMetadataArgs = logOptions.computeMetadataArgs; @@ -2985,7 +2995,7 @@ export class Logger implements Exportable { return ret as Ret; } else { return (async () => { - await this.flush(); + await this.state.bgLogger().flush(); return ret; })() as Ret; } @@ -3031,7 +3041,10 @@ export class Logger implements Exportable { }); type Ret = PromiseUnless; if (this.asyncFlush === true) return span.id as Ret; - return this.flush().then(() => span.id) as Ret; + return this.state + .bgLogger() + .flush() + .then(() => span.id) as Ret; } /** @@ -3143,7 +3156,7 @@ export class Logger implements Exportable { } else { return (async () => { const awaitedRet = await ret; - await this.flush(); + await this.state.bgLogger().flush(); return awaitedRet; })() as Ret; } @@ -3163,6 +3176,9 @@ export class Logger implements Exportable { private startSpanImpl(args?: StartSpanArgs): Span { return new SpanImpl({ + [ATTACHMENT_STREAM_HOOK]: args?.parent + ? undefined + : this[ATTACHMENT_STREAM_HOOK], ...args, // Sometimes `args` gets passed directly into this function, and it contains an undefined value for `state`. // To ensure that we always use this logger's state, we override the `state` argument no matter what. @@ -3242,7 +3258,7 @@ export class Logger implements Exportable { * Flush any pending logs to the server. */ async flush(): Promise { - return await this.state.bgLogger().flush(); + return await flush({ state: this.state }); } get asyncFlush(): IsAsyncFlush | undefined { @@ -5486,6 +5502,8 @@ type AsyncFlushArg = { }; export type InitLoggerOptions = FullLoginOptions & { + /** Transform captured attachment streams for this logger and its child spans. */ + attachmentStreamHook?: AttachmentStreamHook; projectName?: string; projectId?: string; _createInProjectGroup?: string; @@ -5569,6 +5587,7 @@ export function initLogger( ); const ret = new Logger(state, lazyMetadata, { + attachmentStreamHook: options.attachmentStreamHook, asyncFlush, computeMetadataArgs, linkArgs, @@ -7366,6 +7385,11 @@ export function _internalStartSpanWithContext< */ export async function flush(options?: OptionalStateArg): Promise { const state = options?.state ?? _globalState; + // Send pending rows before waiting on attachment work, which may run customer + // hooks, then send the attachment references that work logged. + await state.bgLogger().flush(); + if (!state.attachmentJobs.size) return; + await Promise.allSettled(state.attachmentJobs); return await state.bgLogger().flush(); } @@ -8683,6 +8707,7 @@ function _resolveSpanIds( * We suggest using one of the various `traced` methods, instead of creating Spans directly. See {@link Span.startSpan} for full details. */ export class SpanImpl implements Span { + readonly [ATTACHMENT_STREAM_HOOK]?: AttachmentStreamHook; private _state: BraintrustState; private isMerge: boolean; @@ -8720,8 +8745,10 @@ export class SpanImpl implements Span { propagatedState?: PropagatedState | undefined; } & Omit & InitialSpanWriteAsMergeArg & - InternalSpanContextArg, + InternalSpanContextArg & + AttachmentStreamOptions, ) { + this[ATTACHMENT_STREAM_HOOK] = args[ATTACHMENT_STREAM_HOOK]; this._state = args.state; this._propagatedState = args.propagatedState; const instrumentationName = @@ -8951,6 +8978,9 @@ export class SpanImpl implements Span { ? undefined : { spanId: this._spanId, rootSpanId: this._rootSpanId }; return new SpanImpl({ + [ATTACHMENT_STREAM_HOOK]: args?.parent + ? undefined + : this[ATTACHMENT_STREAM_HOOK], state: this._state, ...args, ...startSpanParentArgs({ @@ -8976,6 +9006,9 @@ export class SpanImpl implements Span { rootSpanId: this._rootSpanId, }; return new SpanImpl({ + [ATTACHMENT_STREAM_HOOK]: args?.parent + ? undefined + : this[ATTACHMENT_STREAM_HOOK], state: this._state, ...args, ...startSpanParentArgs({ diff --git a/js/src/pcm-audio.test.ts b/js/src/pcm-audio.test.ts index 891a3bb0e..2ba3ed95d 100644 --- a/js/src/pcm-audio.test.ts +++ b/js/src/pcm-audio.test.ts @@ -4,9 +4,8 @@ import { PcmAudio } from "./pcm-audio"; const MAX = Number.MAX_SAFE_INTEGER; // Without a chunk-sized limit, flush consumes everything through the cut. const all = { maxBytes: MAX, flush: true }; -async function bytes(result: ReturnType) { - return Buffer.from(await (await result!.attachment.data()).arrayBuffer()); -} +const bytes = (result: ReturnType) => + Buffer.from(result!.data); const samples = (data: Buffer) => Array.from({ length: data.length / 2 }, (_, i) => data.readInt16LE(i * 2)); const mono = (sampleRate: number, offsetMs = 0, channel = 0) => ({ @@ -16,7 +15,29 @@ const mono = (sampleRate: number, offsetMs = 0, channel = 0) => ({ offsetMs, }); -test("positioned sources are downmixed and resampled onto separate channels with silent gaps", async () => { +test.each([1, 7, 96, 192])( + "continuous reads preserve the source clock and silence with a %i-byte budget", + (maxBytes) => { + const pcm = new PcmAudio({ sampleRate: 1000, channels: 2 }); + pcm.update(new Int16Array([1, 2, 3]), mono(1000, 203, 0)); + pcm.update(new Int16Array([4, 5, 6]), mono(1000, 242, 1)); + const parts: Uint8Array[] = []; + let end = 0; + while (pcm.bufferedByteLength) { + const result = pcm.read({ maxBytes, continuous: true, flush: true })!; + expect(result.startOffsetMs).toBe(end); + end = result.endOffsetMs; + parts.push(result.data); + } + const expected = Buffer.alloc(245 * 4); + [1, 2, 3].forEach((n, i) => expected.writeInt16LE(n, (203 + i) * 4)); + [4, 5, 6].forEach((n, i) => expected.writeInt16LE(n, (242 + i) * 4 + 2)); + expect(Buffer.concat(parts)).toEqual(expected); + expect(end).toBe(245); + }, +); + +test("positioned sources are downmixed and resampled onto separate channels with silent gaps", () => { const audio = new PcmAudio({ sampleRate: 24000, channels: 2, @@ -28,10 +49,8 @@ test("positioned sources are downmixed and resampled onto separate channels with offsetMs: 0, }); audio.update(new Int16Array([-1, -2]), mono(24000, 2 / 24, 1)); - const result = audio.digest({ ...all, filename: "speech.pcm" }); - expect(samples(await bytes(result))).toEqual([ - 200, 0, 300, 0, 400, -1, 400, -2, - ]); + const result = audio.read(all); + expect(samples(bytes(result))).toEqual([200, 0, 300, 0, 400, -1, 400, -2]); expect(result).toMatchObject({ channels: [0, 1], startOffsetMs: 0, @@ -43,31 +62,27 @@ test("positioned sources are downmixed and resampled onto separate channels with duration_ms: 4 / 24, }, }); - expect(result!.attachment.reference).toMatchObject({ - filename: "speech.pcm", - content_type: "audio/pcm", - }); expect(audio.bufferedByteLength).toBe(0); }); -test("updates copy input and digests cannot be changed by later writes", async () => { +test("updates copy input and reads cannot be changed by later writes", () => { const audio = new PcmAudio({ sampleRate: 1000, channels: 1 }); const source = new Int16Array([1, 2, 3, 4]); audio.update(source, mono(1000)); source.fill(0); // Later writes win at overlaps. audio.update(new Int16Array([9, 8]), mono(1000, 1)); - const first = audio.digest({ ...all, endOffsetMs: 2 }); - expect(samples(await bytes(first))).toEqual([1, 9]); + const first = audio.read({ ...all, endOffsetMs: 2 }); + expect(samples(bytes(first))).toEqual([1, 9]); expect(() => audio.update(new Int16Array([7]), mono(1000, 1))).toThrow( "capture_clock_discontinuity", ); - const rest = audio.digest(all); - expect(samples(await bytes(rest))).toEqual([8, 4]); + const rest = audio.read(all); + expect(samples(bytes(rest))).toEqual([8, 4]); expect(rest!.startOffsetMs).toBe(first!.endOffsetMs); }); -test("cuts preserve resampled audio sample-for-sample", async () => { +test("cuts preserve resampled audio sample-for-sample", () => { const whole = new PcmAudio({ sampleRate: 24000, channels: 2 }); const split = new PcmAudio({ sampleRate: 24000, channels: 2 }); const source = Int16Array.from( @@ -77,12 +92,12 @@ test("cuts preserve resampled audio sample-for-sample", async () => { for (const audio of [whole, split]) audio.update(source, mono(44100, 0, 1)); const pieces = []; for (const endOffsetMs of [7 / 24, 13.5, 33.125, 99, 100]) - pieces.push(await bytes(split.digest({ ...all, endOffsetMs }))); - expect(Buffer.concat(pieces)).toEqual(await bytes(whole.digest(all))); + pieces.push(bytes(split.read({ ...all, endOffsetMs }))); + expect(Buffer.concat(pieces)).toEqual(bytes(whole.read(all))); expect(split.bufferedByteLength).toBe(0); }); -test("bounds storage and output allocation without consuming data on failure", async () => { +test("bounds storage and output allocation without consuming data on failure", () => { const audio = new PcmAudio({ sampleRate: 1000, channels: 1, @@ -93,11 +108,11 @@ test("bounds storage and output allocation without consuming data on failure", a "capture_byte_limit", ); expect(audio.bufferedByteLength).toBe(4); - expect(samples(await bytes(audio.digest(all)))).toEqual([1, 2]); + expect(samples(bytes(audio.read(all)))).toEqual([1, 2]); audio.update(new Int16Array([7]), mono(1000, 1000000)); - expect(() => audio.digest(all)).toThrow("capture_byte_limit"); - expect(audio.digest({ ...all, endOffsetMs: 1000000 })).toBeUndefined(); - expect(samples(await bytes(audio.digest(all)))).toEqual([7]); + expect(() => audio.read(all)).toThrow("capture_byte_limit"); + expect(audio.read({ ...all, endOffsetMs: 1000000 })).toBeUndefined(); + expect(samples(bytes(audio.read(all)))).toEqual([7]); }); test("rejects invalid formats, incomplete frames, cuts and byte limits", () => { @@ -116,32 +131,28 @@ test("rejects invalid formats, incomplete frames, cuts and byte limits", () => { ).toThrow("unsupported_audio_format"); audio.update(new Int16Array([1, 2]), stereo); for (const endOffsetMs of [1000, -1]) - expect(() => audio.digest({ ...all, endOffsetMs })).toThrow( - "invalid_digest_offset", + expect(() => audio.read({ ...all, endOffsetMs })).toThrow( + "invalid_read_offset", ); for (const maxBytes of [0, -1, 1.5, NaN, Infinity, MAX + 1]) - expect(() => audio.digest({ maxBytes })).toThrow( - "invalid_digest_byte_limit", - ); + expect(() => audio.read({ maxBytes })).toThrow("invalid_read_byte_limit"); expect(audio.bufferedByteLength).toBe(2); }); -test("byte-limited digests wait for confirmed full chunks and flush a final remainder", async () => { +test("byte-limited reads wait for confirmed full chunks and flush a final remainder", () => { const audio = new PcmAudio({ sampleRate: 1000, channels: 2 }); const source = Int16Array.from({ length: 6 }, (_, i) => i + 1); audio.update(source, mono(1000)); - expect(audio.digest({ maxBytes: 8, endOffsetMs: 1 })).toBeUndefined(); - const first = audio.digest({ maxBytes: 8, endOffsetMs: 3 }); + expect(audio.read({ maxBytes: 8, endOffsetMs: 1 })).toBeUndefined(); + const first = audio.read({ maxBytes: 8, endOffsetMs: 3 }); expect(first).toMatchObject({ startOffsetMs: 0, endOffsetMs: 2 }); - expect(audio.digest({ maxBytes: 8, endOffsetMs: 3 })).toBeUndefined(); - const second = audio.digest({ maxBytes: 8, endOffsetMs: 3, flush: true }); + expect(audio.read({ maxBytes: 8, endOffsetMs: 3 })).toBeUndefined(); + const second = audio.read({ maxBytes: 8, endOffsetMs: 3, flush: true }); expect(second).toMatchObject({ startOffsetMs: 2, endOffsetMs: 3 }); - const third = audio.digest({ maxBytes: 8 }); - expect(audio.digest({ maxBytes: 8 })).toBeUndefined(); - const last = audio.digest({ maxBytes: 8, flush: true }); - const data = Buffer.concat( - await Promise.all([first, second, third, last].map(bytes)), - ); + const third = audio.read({ maxBytes: 8 }); + expect(audio.read({ maxBytes: 8 })).toBeUndefined(); + const last = audio.read({ maxBytes: 8, flush: true }); + const data = Buffer.concat([first, second, third, last].map(bytes)); expect(samples(data)).toEqual([1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0]); expect(audio.bufferedByteLength).toBe(0); }); @@ -152,25 +163,25 @@ test.each([ { channels: 3, maxBytes: 13, sizes: [12, 12, 6] }, ])( "byte limits align to complete frames for $channels channels", - async ({ channels, maxBytes, sizes }) => { + ({ channels, maxBytes, sizes }) => { const audio = new PcmAudio({ sampleRate: 1000, channels }); audio.update(new Int16Array([1, 2, 3, 4, 5]), mono(1000)); const files: Buffer[] = []; while (audio.bufferedByteLength) - files.push(await bytes(audio.digest({ maxBytes, flush: true }))); + files.push(bytes(audio.read({ maxBytes, flush: true }))); expect(files.map((f) => f.length)).toEqual(sizes); }, ); -test("byte-limited digests skip empty intervals and retain silence inside a file", async () => { +test("byte-limited reads skip empty intervals and retain silence inside a file", () => { const audio = new PcmAudio({ sampleRate: 1000, channels: 1 }); audio.update(new Int16Array([7]), mono(1000, 5)); audio.update(new Int16Array([9]), mono(1000, 8)); - const first = audio.digest({ maxBytes: 4 }); + const first = audio.read({ maxBytes: 4 }); expect(first).toMatchObject({ startOffsetMs: 4, endOffsetMs: 6 }); - expect(samples(await bytes(first))).toEqual([0, 7]); - expect(audio.digest({ maxBytes: 4 })).toBeUndefined(); - const last = audio.digest({ maxBytes: 4, flush: true }); + expect(samples(bytes(first))).toEqual([0, 7]); + expect(audio.read({ maxBytes: 4 })).toBeUndefined(); + const last = audio.read({ maxBytes: 4, flush: true }); expect(last).toMatchObject({ startOffsetMs: 8, endOffsetMs: 9 }); - expect(samples(await bytes(last))).toEqual([9]); + expect(samples(bytes(last))).toEqual([9]); }); diff --git a/js/src/pcm-audio.ts b/js/src/pcm-audio.ts index 1728d23cb..30fd2c955 100644 --- a/js/src/pcm-audio.ts +++ b/js/src/pcm-audio.ts @@ -1,5 +1,3 @@ -import { Attachment } from "./logger"; - const MAX_CHUNKS = 100_000; type Chunk = { samples: Int16Array; start: number; channel: number }; @@ -7,7 +5,7 @@ type Chunk = { samples: Int16Array; start: number; channel: number }; /** * Incremental signed 16-bit little-endian PCM on a shared sample clock. * update() copies, resamples, and downmixes a positioned source onto one - * output channel; digest() consumes a prefix into an ordinary attachment. + * output channel; read() consumes a prefix as owned bytes. */ export class PcmAudio { private chunks: Chunk[] = []; @@ -41,6 +39,13 @@ export class PcmAudio { get bufferedByteLength() { return this.bytes; } + + /** Release unprocessed input after capture has irrecoverably failed. */ + discard(): void { + this.chunks = []; + this.bytes = 0; + } + get endOffsetMs() { return samplesToMs(this.end, this.sampleRate); } @@ -108,36 +113,38 @@ export class PcmAudio { /** * Consume at most maxBytes through an optional confirmed offset; future * updates may continue. Waits for a full chunk unless flush is set. Cuts align - * to complete sample frames (at least one), skipping wholly empty intervals. + * to complete sample frames (at least one), skipping wholly empty intervals + * unless continuous reads preserve silence on the source clock. */ - digest(options: { + read(options: { maxBytes: number; endOffsetMs?: number; flush?: boolean; - filename?: string; + continuous?: boolean; }) { if (!Number.isSafeInteger(options.maxBytes) || options.maxBytes <= 0) - throw new Error("invalid_digest_byte_limit"); + throw new Error("invalid_read_byte_limit"); let end = options.endOffsetMs === undefined ? this.end : Math.round(msToSamples(options.endOffsetMs, this.sampleRate)); if (!Number.isSafeInteger(end) || end < this.start || end > this.end) - throw new Error("invalid_digest_offset"); + throw new Error("invalid_read_offset"); const frames = Math.max( 1, Math.floor(options.maxBytes / (this.channels * 2)), ); - const first = this.chunks.reduce((min, c) => Math.min(min, c.start), end); - // Skip empty chunks without allocating silence or emitting attachments. - this.start += Math.floor((first - this.start) / frames) * frames; + if (!options.continuous) { + const first = this.chunks.reduce((min, c) => Math.min(min, c.start), end); + this.start += Math.floor((first - this.start) / frames) * frames; + } const cut = this.start + frames; if (end < cut && !options.flush) return; end = Math.min(end, cut); const start = this.start; if (end === start) return; const included = this.chunks.filter((c) => c.start < end); - if (!included.length) { + if (!included.length && !options.continuous) { this.start = end; return; } @@ -161,18 +168,14 @@ export class PcmAudio { if (chunk.start >= end) return [chunk]; const consumed = Math.min(chunk.samples.length, end - chunk.start); this.bytes -= consumed * 2; - // Copy the tail so it doesn't retain the already digested prefix's allocation. + // Copy the tail so it doesn't retain the already read prefix's allocation. return consumed === chunk.samples.length ? [] : [{ ...chunk, start: end, samples: chunk.samples.slice(consumed) }]; }); this.start = end; return { - attachment: new Attachment({ - data: bytes, - filename: options.filename ?? "audio.pcm", - contentType: "audio/pcm", - }), + data: bytes, metadata: { mime_type: "audio/pcm", sample_rate_hz: this.sampleRate, diff --git a/js/src/wrappers/attachment-utils.ts b/js/src/wrappers/attachment-utils.ts index 4282736ab..8d13a8db6 100644 --- a/js/src/wrappers/attachment-utils.ts +++ b/js/src/wrappers/attachment-utils.ts @@ -29,6 +29,7 @@ export function getExtensionFromMediaType(mediaType: string): string { "audio/mpeg": "mp3", "audio/flac": "flac", "audio/basic": "mulaw", + "audio/pcm": "pcm", "audio/mp4": "m4a", "audio/wav": "wav", "audio/webm": "webm",