From 6c033ff19a770a942c92e31066a2eb05f81b8e72 Mon Sep 17 00:00:00 2001 From: David Elner Date: Tue, 3 Nov 2026 20:07:50 -0500 Subject: [PATCH 01/12] feat(audio): add shared session recording infrastructure Add positioned PCM recording, middleware, progressive v3 export, bounded native workers and direct attachment uploads. Include PCM/WAV support, cancellation and upload lifecycle coverage, and audio-only attachment request timeouts. --- py/noxfile.py | 6 + py/src/braintrust/_audio/__init__.py | 1 + py/src/braintrust/_audio/middleware.py | 111 +++++++++ py/src/braintrust/_audio/recorder.py | 277 +++++++++++++++++++++++ py/src/braintrust/_audio/session.py | 203 +++++++++++++++++ py/src/braintrust/_audio/test_audio.py | 277 +++++++++++++++++++++++ py/src/braintrust/_audio/test_session.py | 148 ++++++++++++ py/src/braintrust/_audio/types.py | 69 ++++++ py/src/braintrust/_audio/worker.py | 55 +++++ py/src/braintrust/audio.py | 52 +++++ py/src/braintrust/logger.py | 12 +- py/src/braintrust/test_logger.py | 28 +++ 12 files changed, 1236 insertions(+), 3 deletions(-) create mode 100644 py/src/braintrust/_audio/__init__.py create mode 100644 py/src/braintrust/_audio/middleware.py create mode 100644 py/src/braintrust/_audio/recorder.py create mode 100644 py/src/braintrust/_audio/session.py create mode 100644 py/src/braintrust/_audio/test_audio.py create mode 100644 py/src/braintrust/_audio/test_session.py create mode 100644 py/src/braintrust/_audio/types.py create mode 100644 py/src/braintrust/_audio/worker.py create mode 100644 py/src/braintrust/audio.py diff --git a/py/noxfile.py b/py/noxfile.py index 5ecb41e7..47cc5bca 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -878,6 +878,12 @@ def is_importable(name): _run_core_tests(session) +@nox.session() +def test_audio(session): + _install_test_deps(session) + _run_tests(session, "braintrust/_audio") + + # Not a default session: it needs a wheel from `make build`, so a bare `nox` # (`make test`) must not pick it up. @nox.session(default=False) diff --git a/py/src/braintrust/_audio/__init__.py b/py/src/braintrust/_audio/__init__.py new file mode 100644 index 00000000..d4d03155 --- /dev/null +++ b/py/src/braintrust/_audio/__init__.py @@ -0,0 +1 @@ +"""Internal shared audio infrastructure. No framework or codec dependencies.""" diff --git a/py/src/braintrust/_audio/middleware.py b/py/src/braintrust/_audio/middleware.py new file mode 100644 index 00000000..aec8e6f1 --- /dev/null +++ b/py/src/braintrust/_audio/middleware.py @@ -0,0 +1,111 @@ +import importlib +import io +import logging +import wave +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import replace +from fractions import Fraction +from typing import Any, Protocol + +from braintrust.audio import AudioEncoding + +from .types import AudioResult, SegmentMeta +from .worker import run_native + + +AudioNext = Callable[[Any, SegmentMeta], Awaitable[AudioResult]] + + +class AudioMiddleware(Protocol): + async def __call__(self, payload: Any, meta: SegmentMeta, next: AudioNext) -> AudioResult: ... + + +def _validate(meta: SegmentMeta, original: SegmentMeta): + if not isinstance(meta, SegmentMeta) or ( + meta.index != original.index + or Fraction(meta.start_sample, meta.sample_rate_hz) != Fraction(original.start_sample, original.sample_rate_hz) + or Fraction(meta.sample_count, meta.sample_rate_hz) != Fraction(original.sample_count, original.sample_rate_hz) + or meta.channels != original.channels + ): + raise ValueError("Audio middleware changed segment identity, placement, duration, or sources") + if ( + meta.encoding == "pcm" + and meta.mime_type == "audio/pcm" + and ( + meta.attributes.get("pcm.sample_format") != "s16le" + or meta.attributes.get("pcm.channel_layout") != "interleaved" + ) + ): + raise ValueError("Raw PCM requires sample format and channel layout") + + +async def run_middleware(stack: Sequence[AudioMiddleware], payload: bytes, meta: SegmentMeta) -> AudioResult: + async def run(index, data, current): + _validate(current, meta) + if index == len(stack): + result = AudioResult(data, current) + else: + called = False + + def next(data, new_meta): + nonlocal called + if called: + raise ValueError("Audio middleware called next more than once") + called = True + return run(index + 1, data, new_meta) + + result = await stack[index](data, current, next) + if not isinstance(result, AudioResult) or not isinstance(result.data, bytes): + raise ValueError("Audio middleware must return final bytes") + _validate(result.meta, meta) + if result.meta.encoding == "pcm" and result.meta.mime_type == "audio/pcm": + if len(result.data) != result.meta.sample_count * len(result.meta.channels) * 2: + raise ValueError("PCM byte length does not match playable duration") + return result + + return await run(0, payload, meta) + + +def _wav(data, meta): + buffer = io.BytesIO() + with wave.Wave_write(buffer) as output: + output.setnchannels(len(meta.channels)) + output.setsampwidth(2) + output.setframerate(meta.sample_rate_hz) + output.writeframes(data) + return buffer.getvalue() + + +async def wav_middleware(data, meta, next): + data = await run_native(_wav, data, meta) + return await next(data, replace(meta, mime_type="audio/wav")) + + +_warned_missing_extension = False + + +def encoding_stack(config: AudioEncoding | None, *, channels: int | None = None) -> tuple[AudioMiddleware, ...]: + global _warned_missing_extension + selected = config or AudioEncoding() + if selected.codec == "pcm": + return (wav_middleware,) if selected.container == "wav" else () + if channels is not None: + if channels not in (1, 2): + raise ValueError("Opus/Ogg supports mono or stereo recordings") + if selected.bitrate_bps > 256000 * channels: + raise ValueError("Opus bitrate exceeds 256000 bits/s per output channel") + try: + extension = importlib.import_module("braintrust_audio") + stack = tuple(extension.create_stack(selected, contract_version=1)) + if not all(callable(layer) for layer in stack): + raise ValueError("Invalid audio extension middleware contract") + return stack + except Exception as error: + if config is not None: + raise ValueError("Opus/Ogg requires a compatible braintrust-audio installation") from error + if not _warned_missing_extension: + _warned_missing_extension = True + logging.getLogger("braintrust").warning( + "braintrust-audio is unavailable; recording raw PCM. Install braintrust-audio for Opus/Ogg." + ) + return () diff --git a/py/src/braintrust/_audio/recorder.py b/py/src/braintrust/_audio/recorder.py new file mode 100644 index 00000000..de3e120f --- /dev/null +++ b/py/src/braintrust/_audio/recorder.py @@ -0,0 +1,277 @@ +import asyncio +import logging +import math +import sys +from array import array +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from fractions import Fraction + +from braintrust.audio import AudioRecordingOptions + +from .middleware import AudioMiddleware, run_middleware +from .types import AudioChannel, AudioResult, RecordingResult, SegmentMeta +from .worker import run_native + + +SAMPLE_RATE = 48000 +NS = 1_000_000_000 + + +@dataclass(frozen=True) +class Packet: + source: str + data: bytes + rate: int + channels: int + start: Fraction # Recording samples; retaining the fraction preserves run phase. + + @property + def count(self): + return len(self.data) // (2 * self.channels) + + @property + def end(self): + return self.start + Fraction(self.count * SAMPLE_RATE, self.rate) + + +def _samples(data): + result = array("h") + result.frombytes(data) + if sys.byteorder != "little": + result.byteswap() + return result + + +def render(packets: tuple[Packet, ...], sources: tuple[str, ...], start: int, end: int): + """Sparse placement and phase-continuous linear resampling, in a worker. + + Phase is derived from absolute rational positions, never rounded per packet. + A following packet supplies the interpolation sample at a packet boundary. + Silence is materialized for only the segment being processed. + """ + output = array("h", [0]) * ((end - start) * len(sources)) + represented = [] + for channel, source in enumerate(sources): + stream = [packet for packet in packets if packet.source == source] + present = False + for index, packet in enumerate(stream): + first = max(start, math.ceil(packet.start)) + last = min(end, math.ceil(packet.end)) + if first >= last: + continue + present = True + samples = _samples(packet.data) + # Downmix before interpolation. The next packet belongs to the same + # run only when its position, rate and channel count agree exactly. + mono = [ + sum(samples[i : i + packet.channels]) // packet.channels + for i in range(0, len(samples), packet.channels) + ] + following = mono[-1] + if index + 1 < len(stream): + other = stream[index + 1] + if other.start == packet.end and other.rate == packet.rate and other.channels == packet.channels: + frame = _samples(other.data[: other.channels * 2]) + following = sum(frame) // other.channels + phase = (Fraction(first) - packet.start) * packet.rate / SAMPLE_RATE + step = Fraction(packet.rate, SAMPLE_RATE) + denominator = math.lcm(phase.denominator, step.denominator) + numerator = phase.numerator * (denominator // phase.denominator) + increment = step.numerator * (denominator // step.denominator) + for position in range(first, last): + sample, remainder = divmod(numerator, denominator) + value = mono[sample] + next_value = mono[sample + 1] if sample + 1 < len(mono) else following + output[(position - start) * len(sources) + channel] = ( + value + (next_value - value) * remainder // denominator + ) + numerator += increment + represented.append(AudioChannel((source,) if present else ())) + if sys.byteorder != "little": + output.byteswap() + return output.tobytes(), tuple(represented) + + +class AudioRecorder: + """Accept positioned PCM without waiting for normalization, encoding or IO. + + ``advance`` is a promise by the adapter: no source will add audio before its + watermark. Thresholds request flushing; they never reject admitted samples. + """ + + def __init__( + self, + *, + sources: tuple[str, ...], + options: AudioRecordingOptions, + middleware: tuple[AudioMiddleware, ...], + on_segment: Callable[[AudioResult], Awaitable[None]], + on_omitted: Callable[[SegmentMeta, str], None], + ): + if not sources or len(set(sources)) != len(sources): + raise ValueError("Recording sources must be nonempty and unique") + self.sources = tuple(sources) + self.options = options + self.middleware = middleware + self.on_segment = on_segment + self.on_omitted = on_omitted + self._packets: list[Packet] = [] + self._watermark = 0 + self._sealed = 0 + self._end = 0 + self._index = 0 + self._task: asyncio.Task | None = None + self._finishing = False + self._discarded = False + self.reason: str | None = None + self._warned_backlog = False + self._processing: SegmentMeta | None = None + self._gaps: list[tuple[int, int, str]] = [] + + @property + def buffered_bytes(self): + return sum(len(packet.data) for packet in self._packets) + + def record( + self, source_id: str, pcm: bytes, *, sample_rate_hz: int, channels: int, offset_ns: int | Fraction + ) -> bool: + if self._finishing or self.reason: + return False + if source_id not in self.sources or sample_rate_hz <= 0 or channels <= 0 or len(pcm) % (channels * 2): + raise ValueError("Invalid positioned PCM") + if not pcm: + return True + start = Fraction(offset_ns) * SAMPLE_RATE / NS + if start < self._sealed or start < 0: + raise ValueError("Audio arrived before the sealed recording watermark") + packet = Packet(source_id, bytes(pcm), sample_rate_hz, channels, start) + self._packets.append(packet) + self._end = max(self._end, math.ceil(packet.end)) + if self.options.max_duration_seconds is not None: + limit = round(self.options.max_duration_seconds * SAMPLE_RATE) + if self._end >= limit: + self._end = limit + self.stop("duration_limit") + buffered = self.buffered_bytes + if buffered > self.options.flush_threshold_bytes * 4 and not self._warned_backlog: + self._warned_backlog = True + logging.getLogger("braintrust").warning( + "Audio recording backlog exceeds its flush target; capture continues" + ) + self._pump() + return True + + def advance(self, through_ns: int | Fraction): + self._watermark = max(self._watermark, math.floor(Fraction(through_ns) * SAMPLE_RATE / NS)) + self._pump() + + def unavailable(self, start_ns, end_ns, reason): + """Known bounds of unavailable audio; never encode it as real silence.""" + start = max(0, math.floor(Fraction(start_ns) * SAMPLE_RATE / NS)) + end = math.ceil(Fraction(end_ns) * SAMPLE_RATE / NS) + if end <= start or self._discarded: + return + if start < self._sealed: + raise ValueError("Unavailable audio arrived before the sealed watermark") + self._gaps.append((start, end, reason)) + self._end = max(self._end, end) + self._pump() + + def stop(self, reason: str): + self.reason = self.reason or reason + + def discard(self, reason: str): + """Stop and erase queued PCM; in-flight workers retain their own copies.""" + if self._discarded: + return + self.stop(reason) + self._discarded = True + if self._processing is not None: + self.on_omitted(self._processing, reason) + if self._end > self._sealed: + self.on_omitted( + SegmentMeta( + self._index, + self._sealed, + self._end - self._sealed, + SAMPLE_RATE, + tuple(AudioChannel(()) for _ in self.sources), + ), + reason, + ) + self._index += 1 + self._packets.clear() + + def _cut(self): + eligible = min(self._end, self._watermark) + target = self._sealed + max(1, round(self.options.segment_duration_seconds * SAMPLE_RATE)) + if eligible <= self._sealed: + return None + if eligible >= target: + return target + if self._finishing or self.buffered_bytes >= self.options.flush_threshold_bytes: + return eligible + return None + + def _pump(self): + if self._task is None or self._task.done(): + if self._cut() is not None: + self._task = asyncio.create_task(self._drain()) + + async def _drain(self): + while not self._discarded and (end := self._cut()) is not None: + start = self._sealed + meta = SegmentMeta( + self._index, start, end - start, SAMPLE_RATE, tuple(AudioChannel(()) for _ in self.sources) + ) + self._index += 1 + self._processing = meta + # Snapshot references before detaching accepted buffers. Native work + # owns them even if the awaiting coroutine is cancelled. + packets = tuple(self._packets) + self._sealed = end + gap = next((gap for gap in self._gaps if gap[0] < end and gap[1] > start), None) + if gap is not None: + self.on_omitted(meta, gap[2]) + self._packets = [packet for packet in self._packets if packet.end > end] + self._gaps = [gap for gap in self._gaps if gap[1] > end] + self._processing = None + continue + try: + data, channels = await run_native(render, packets, self.sources, start, end) + meta = SegmentMeta(meta.index, start, end - start, SAMPLE_RATE, channels) + result = await run_middleware(self.middleware, data, meta) + if not self._discarded: + await self.on_segment(result) + else: + self.on_omitted(meta, self.reason or "recording_stopped") + except asyncio.CancelledError: + self.on_omitted(meta, "shutdown_timeout") + raise + except Exception: + self.stop("encoding_or_upload_failed") + self.on_omitted(meta, self.reason) + # Already accepted audio is unavailable; retain honest placement + # rather than claim the remainder of the call was recorded. + if self._end > end: + self.on_omitted( + SegmentMeta(self._index, end, self._end - end, SAMPLE_RATE, meta.channels), self.reason + ) + self._index += 1 + self._packets.clear() + self._sealed = self._end + self._processing = None + self._discarded = True + return + self._packets = [packet for packet in self._packets if packet.end > end] + self._processing = None + + async def finish(self) -> RecordingResult: + self._finishing = True + # The adapter settles pending playback before calling finish. + self._watermark = max(self._watermark, self._end) + self._pump() + if self._task is not None: + await asyncio.shield(self._task) + return RecordingResult(self._index, self.reason) diff --git a/py/src/braintrust/_audio/session.py b/py/src/braintrust/_audio/session.py new file mode 100644 index 00000000..3b93d38f --- /dev/null +++ b/py/src/braintrust/_audio/session.py @@ -0,0 +1,203 @@ +import asyncio +import logging +import time +from collections.abc import Callable + +from braintrust.audio import AudioRecordingOptions +from braintrust.env import BraintrustEnv +from braintrust.logger import Attachment + +from .middleware import encoding_stack +from .recorder import AudioRecorder +from .types import AudioResult, RecordingResult, SegmentMeta, thaw +from .worker import run_native + + +class AudioSessionRecorder: + """Schema registries and recorder-local direct attachment export. + + ``allowed`` must include capture opt-ins and effective framework privacy. It + is checked again before upload. Trace customizers cannot veto audio uploads. + """ + + def __init__( + self, + span, + *, + sources: tuple[str, ...] = (), + options: AudioRecordingOptions | None = None, + allowed: Callable[[], bool] = lambda: True, + ): + self.span = span + self.wall_ns, self.monotonic_ns = time.time_ns(), time.monotonic_ns() + self.allowed = allowed + self.options = options or AudioRecordingOptions() + self._participants: dict[object, str] = {} + self._speakers: dict[object, str] = {} + self._terminal: set[int] = set() + self._inflight: dict[int, SegmentMeta] = {} + self._closed = False + self._sources = {source: {} for source in sources} + self.recorder = ( + AudioRecorder( + sources=sources, + options=self.options, + middleware=encoding_stack(self.options.encoding, channels=len(sources)), + on_segment=self._export, + on_omitted=self.omitted, + ) + if sources + else None + ) + self.span.log(metadata=self.session_metadata()) + + def session_metadata(self): + return { + "audio.schema.version": "3.0.0", + "audio.timeline": {"origin_unix_ms": self.wall_ns // 1_000_000}, + "audio.sources": dict(self._sources), + } + + def source(self, source_id, member=None): + if source_id not in self._sources: + self._sources[source_id] = dict(member or {}) + self.span.log(metadata={"audio.sources": {source_id: self._sources[source_id]}}) + return source_id + + def elapsed_ns(self): + return time.monotonic_ns() - self.monotonic_ns + + def from_wall(self, seconds: float): + return max(0, round(seconds * 1_000_000_000) - self.wall_ns) + + def participant(self, native_key, member): + return self._identity(self._participants, native_key, member, "audio.participants", "participant") + + def speaker(self, source_id, native_key, member): + return self._identity(self._speakers, (source_id, native_key), member, "audio.speakers", "speaker") + + def _identity(self, registry, key, member, field, prefix): + if not self.allowed(): + return None + if key not in registry: + identifier = f"{prefix}-{len(registry)}" + registry[key] = identifier + self.span.log(metadata={field: {identifier: member}}) + return registry[key] + + def record(self, source_id, pcm, **format_and_time): + if self.recorder is None or not self.allowed(): + if self.recorder is not None: + self.recorder.discard("privacy_restricted") + return False + return self.recorder.record(source_id, pcm, **format_and_time) + + def advance(self, through_ns): + if self.recorder is not None: + self.recorder.advance(through_ns) + + def unavailable(self, start_ns, end_ns, reason): + if self.recorder is not None: + self.recorder.unavailable(start_ns, end_ns, reason) + + def omitted(self, meta: SegmentMeta, reason: str): + if meta.index in self._terminal: + return + self._terminal.add(meta.index) + self._inflight.pop(meta.index, None) + self.span.log( + metadata={ + "audio.segments": { + f"part-{meta.index}": { + "state": "omitted", + **meta.placement(), + "reason": reason, + } + } + } + ) + + async def _export(self, result: AudioResult): + meta = result.meta + self._inflight[meta.index] = meta + identifier = f"part-{meta.index}" + if self._closed or not self.allowed(): + self.omitted(meta, "privacy_restricted" if not self._closed else "shutdown_timeout") + return + self.span.log(metadata={"audio.segments": {identifier: {"state": "pending", **meta.placement()}}}) + extension = {"audio/ogg": "ogg", "audio/wav": "wav", "audio/pcm": "pcm"}[meta.mime_type] + attachment = Attachment(data=result.data, filename=f"{identifier}.{extension}", content_type=meta.mime_type) + attachment._upload_timeout_seconds = BraintrustEnv.HTTP_TIMEOUT.get(60.0) + + def upload(): + # Recheck after waiting for a native worker slot. + if self._closed or not self.allowed(): + return None + return attachment.upload() + + try: + status = await run_native(upload) + except Exception: + self.omitted(meta, "upload_failed") + raise + if meta.index in self._terminal: + return + if status is None or not self.allowed(): + self.omitted(meta, "privacy_restricted") + elif status.get("upload_status") != "done": + self.omitted(meta, "upload_failed") + raise RuntimeError("Audio upload failed") + else: + self._terminal.add(meta.index) + self._inflight.pop(meta.index, None) + self.span.log( + input={"audio": {identifier: attachment.reference}}, + metadata={ + "audio.segments": { + identifier: { + "state": "ready", + **meta.placement(), + "attachment": {"ref": f"/input/audio/{identifier}"}, + "mime_type": meta.mime_type, + "encoding": meta.encoding, + "sample_rate_hz": meta.sample_rate_hz, + "channels": { + str(i): {"sources": list(channel.sources)} for i, channel in enumerate(meta.channels) + }, + "attributes": thaw(meta.attributes), + } + } + }, + ) + + async def finish(self) -> RecordingResult: + result = RecordingResult(0) + loop = asyncio.get_running_loop() + deadline = loop.time() + self.options.shutdown_timeout_seconds + try: + if self.recorder is not None: + result = await asyncio.wait_for(self.recorder.finish(), self.options.shutdown_timeout_seconds) + except (asyncio.TimeoutError, asyncio.CancelledError): + self._closed = True + for meta in tuple(self._inflight.values()): + self.omitted(meta, "shutdown_timeout") + if self.recorder is not None: + self.recorder.discard("shutdown_timeout") + if self.recorder._task is not None: + self.recorder._task.cancel() + await asyncio.gather(self.recorder._task, return_exceptions=True) + result = RecordingResult(len(self._terminal), "shutdown_timeout") + finally: + self._closed = True + if result.reason: + self.span.log(metadata={"contrib.braintrust.recording_stop_reason": result.reason}) + # Terminal updates are queued before flushing. No exporter changes. + flushing = loop.run_in_executor(None, self.span.flush) + flushing.add_done_callback(lambda completed: completed.exception() if not completed.cancelled() else None) + try: + await asyncio.wait_for(asyncio.shield(flushing), max(0, deadline - loop.time())) + except asyncio.TimeoutError: + result = RecordingResult(result.segments, result.reason or "logger_flush_timeout") + except Exception: + logging.getLogger("braintrust").warning("Audio terminal metadata flush failed") + return result diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py new file mode 100644 index 00000000..ef6d8ef9 --- /dev/null +++ b/py/src/braintrust/_audio/test_audio.py @@ -0,0 +1,277 @@ +import asyncio +import io +import subprocess +import sys +import threading +import wave +from array import array +from dataclasses import replace +from fractions import Fraction + +import pytest +from braintrust._audio.middleware import encoding_stack, run_middleware +from braintrust._audio.recorder import AudioRecorder +from braintrust._audio.types import AudioChannel, AudioResult, SegmentMeta +from braintrust._audio.worker import run_native +from braintrust.audio import AudioEncoding, AudioRecordingOptions + + +def pcm(values): + return array("h", values).tobytes() + + +def recorder(*, duration=0.01, threshold=1024 * 1024, middleware=(), on_segment=None): + segments, omitted = [], [] + + async def collect(result): + segments.append(result) + if on_segment is not None: + await on_segment(result) + + instance = AudioRecorder( + sources=("input", "output"), + options=AudioRecordingOptions( + segment_duration_seconds=duration, + flush_threshold_bytes=threshold, + ), + middleware=middleware, + on_segment=collect, + on_omitted=lambda *args: omitted.append(args), + ) + return instance, segments, omitted + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rate", [16000, 24000, 44100, 48000]) +async def test_packet_splits_keep_exact_phase_and_length(rate): + async def record(split): + instance, segments, omitted = recorder() + count = rate // 5 + values = [i % 1000 for i in range(count)] + for start in range(0, count, split): + instance.record( + "input", + pcm(values[start : start + split]), + sample_rate_hz=rate, + channels=1, + offset_ns=Fraction(start * 1_000_000_000, rate), + ) + instance.advance(200_000_000) + await instance.finish() + assert not omitted + assert sum(s.meta.sample_count for s in segments) == 9600 + return b"".join(s.data for s in segments) + + assert await record(137) == await record(rate) + + +@pytest.mark.asyncio +async def test_sparse_timeline_downmix_and_channel_alignment(): + instance, segments, omitted = recorder(duration=0.1) + instance.record("input", pcm([200, 400] * 480), sample_rate_hz=48000, channels=2, offset_ns=5_000_000) + instance.record("output", pcm([800] * 240), sample_rate_hz=24000, channels=1, offset_ns=20_000_000) + await instance.finish() + samples = array("h") + samples.frombytes(segments[0].data) + assert len(samples) == 1440 * 2 + assert samples[240 * 2 : 720 * 2 : 2] == array("h", [300] * 480) + assert samples[960 * 2 + 1 :: 2] == array("h", [800] * 480) + assert samples[: 240 * 2] == array("h", [0] * 480) + assert not omitted + + +@pytest.mark.asyncio +async def test_size_threshold_waits_for_watermark_and_congestion_preserves_samples(): + gate = asyncio.Event() + + async def slow_upload(_): + await gate.wait() + + instance, segments, omitted = recorder(duration=60, threshold=32, on_segment=slow_upload) + for index in range(10): + instance.record("input", pcm([index] * 480), sample_rate_hz=48000, channels=1, offset_ns=index * 10_000_000) + await asyncio.sleep(0) + assert not segments + instance.advance(50_000_000) + for _ in range(50): + if segments: + break + await asyncio.sleep(0.005) + assert len(segments) == 1 + for index in range(10, 20): + instance.record("input", pcm([index] * 480), sample_rate_hz=48000, channels=1, offset_ns=index * 10_000_000) + instance.advance(200_000_000) + gate.set() + await instance.finish() + samples = array("h") + samples.frombytes(b"".join(s.data for s in segments)) + assert samples[::2] == array("h", (value for index in range(20) for value in [index] * 480)) + assert not omitted + + +@pytest.mark.asyncio +async def test_middleware_validates_duration_sources_bytes_and_deep_freeze(): + meta = SegmentMeta( + 0, + 0, + 48, + 48000, + (AudioChannel(("input",)),), + attributes={ + "nested": {"list": [1]}, + "pcm.sample_format": "s16le", + "pcm.channel_layout": "interleaved", + }, + ) + with pytest.raises(TypeError): + meta.attributes["nested"]["list"][0] = 2 + + async def invalid(data, current, next): + return AudioResult(data, replace(current, sample_count=49)) + + with pytest.raises(ValueError, match="identity, placement"): + await run_middleware((invalid,), bytes(96), meta) + + async def twice(data, current, next): + await next(data, current) + return await next(data, current) + + meta = replace(meta, attributes={"pcm.sample_format": "s16le", "pcm.channel_layout": "interleaved"}) + with pytest.raises(ValueError, match="more than once"): + await run_middleware((twice,), bytes(96), meta) + + +@pytest.mark.asyncio +async def test_wav_requires_no_optional_backend_and_preserves_tiny_tail(): + meta = SegmentMeta(0, 0, 137, 48000, (AudioChannel(("input",)),)) + result = await run_middleware(encoding_stack(AudioEncoding(codec="pcm", container="wav")), pcm(range(137)), meta) + with wave.open(io.BytesIO(result.data)) as decoded: + assert decoded.getnframes() == 137 + assert decoded.readframes(137) == pcm(range(137)) + + +@pytest.mark.asyncio +async def test_native_worker_cancellation_keeps_buffer_owned_until_completion(): + started, release, completed = threading.Event(), threading.Event(), threading.Event() + payload = bytearray(b"owned") + + def native(data): + started.set() + release.wait(5) + assert data == b"owned" + completed.set() + + task = asyncio.create_task(run_native(native, payload)) + await asyncio.to_thread(started.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + assert await asyncio.to_thread(completed.wait, 5) + + +@pytest.mark.asyncio +async def test_native_worker_admits_waiting_sessions_in_order_and_removes_cancelled_waiters(): + from braintrust._audio import worker + + started = [threading.Event() for _ in range(4)] + release = [threading.Event() for _ in range(3)] + + def blocking(index): + started[index].set() + assert release[index].wait(5) + + tasks = [asyncio.create_task(run_native(blocking, index)) for index in range(2)] + try: + for event in started[:2]: + assert await asyncio.to_thread(event.wait, 5) + first = asyncio.create_task(run_native(blocking, 2)) + cancelled = asyncio.create_task(run_native(lambda: pytest.fail("Cancelled job was submitted"))) + last = asyncio.create_task(run_native(started[3].set)) + tasks.extend([first, cancelled, last]) + + async def wait_for_queue(): + while worker._waiters is None or len(worker._waiters) != 3: + await asyncio.sleep(0.005) + + await asyncio.wait_for(wait_for_queue(), 5) + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + release[0].set() + assert await asyncio.to_thread(started[2].wait, 5) + assert not started[3].is_set() + release[2].set() + assert await asyncio.to_thread(started[3].wait, 5) + finally: + for event in release: + event.set() + await asyncio.gather(*tasks, return_exceptions=True) + + +def test_options_reject_unsupported_configuration(): + with pytest.raises(ValueError): + AudioEncoding(codec="opus/ogg") + with pytest.raises(ValueError): + AudioRecordingOptions(segment_duration_seconds=float("nan")) + + +@pytest.mark.asyncio +async def test_unavailable_audio_is_omitted_instead_of_encoded_as_silence(): + instance, segments, omitted = recorder(duration=0.01) + instance.record("input", pcm([100] * 1440), sample_rate_hz=48000, channels=1, offset_ns=0) + instance.unavailable(10_000_000, 20_000_000, "playback_timing_unavailable") + instance.advance(30_000_000) + await instance.finish() + assert [segment.meta.index for segment in segments] == [0, 2] + assert len(omitted) == 1 + assert omitted[0][0].placement() == {"start_offset_ms": 10, "duration_ms": 10} + assert omitted[0][1] == "playback_timing_unavailable" + + +def test_missing_extension_warns_once_and_explicit_opus_fails(caplog, monkeypatch): + from braintrust._audio import middleware + + monkeypatch.setattr(middleware, "_warned_missing_extension", False) + assert encoding_stack(None) == () + assert encoding_stack(None) == () + assert len([record for record in caplog.records if "recording raw PCM" in record.message]) == 1 + with pytest.raises(ValueError, match="requires a compatible"): + encoding_stack(AudioEncoding(codec="opus", container="ogg")) + + +def test_tracing_only_session_has_no_media_imports_or_audio_workers(): + subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +from braintrust.logger import NOOP_SPAN +from braintrust._audio.session import AudioSessionRecorder +from braintrust._audio import worker +session = AudioSessionRecorder(NOOP_SPAN) +assert session.recorder is None +assert worker._executor is None +assert not {'av', 'braintrust_audio', 'livekit', 'numpy'} & sys.modules.keys() +""", + ], + check=True, + ) + + +@pytest.mark.asyncio +async def test_terminal_transform_failure_never_reencodes_omitted_tail_as_silence(): + async def fail(data, meta, next): + raise RuntimeError("terminal codec failure") + + instance, segments, omitted = recorder(duration=0.01, middleware=(fail,)) + instance.record("input", pcm([100] * 1440), sample_rate_hz=48000, channels=1, offset_ns=0) + instance.advance(30_000_000) + await instance.finish() + await instance.finish() + assert not segments + assert [meta.placement() for meta, _ in omitted] == [ + {"start_offset_ms": 0, "duration_ms": 10}, + {"start_offset_ms": 10, "duration_ms": 20}, + ] diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py new file mode 100644 index 00000000..6eddb4f6 --- /dev/null +++ b/py/src/braintrust/_audio/test_session.py @@ -0,0 +1,148 @@ +import asyncio +import threading + +import pytest +from braintrust import SpanCustomizer, logger, set_span_customizers +from braintrust._audio.session import AudioSessionRecorder +from braintrust.audio import AudioEncoding, AudioRecordingOptions +from braintrust.test_helpers import init_test_logger + + +@pytest.fixture +def memory(monkeypatch): + monkeypatch.setattr(logger, "_compute_logger_metadata", logger._compute_logger_metadata) + init_test_logger("audio-tests") + with logger._internal_with_memory_background_logger() as bgl: + monkeypatch.setattr(logger._state, "_global_bg_logger", logger.LazyValue(lambda: bgl, use_mutex=False)) + yield bgl + set_span_customizers(None) + + +def recording(**kwargs): + span = logger.start_span(name="session") + return AudioSessionRecorder( + span, + sources=("session-input",), + options=AudioRecordingOptions( + encoding=AudioEncoding(codec="pcm", container=None), + **kwargs, + ), + ) + + +@pytest.mark.asyncio +async def test_session_source_registry_is_framework_independent(memory): + span = logger.start_span(name="standalone") + session = AudioSessionRecorder(span) + session.source("microphone-A", {"label": "Input"}) + session.source("speaker-B") + assert session.session_metadata()["audio.sources"] == { + "microphone-A": {"label": "Input"}, + "speaker-B": {}, + } + await session.finish() + metadata = memory.pop()[0]["metadata"] + assert set(metadata["audio.sources"]) == {"microphone-A", "speaker-B"} + + +@pytest.mark.asyncio +async def test_ready_only_after_success_and_reference_is_not_reuploaded(memory, monkeypatch): + entered, release = threading.Event(), threading.Event() + uploads = [] + + def upload(attachment): + uploads.append(attachment) + entered.set() + release.wait(5) + return {"upload_status": "done"} + + monkeypatch.setattr(logger.Attachment, "upload", upload) + session = recording() + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + finishing = asyncio.create_task(session.finish()) + assert await asyncio.to_thread(entered.wait, 5) + records = memory.pop() + assert records[0]["metadata"]["audio.segments"]["part-0"]["state"] == "pending" + assert "audio" not in records[0].get("input", {}) + release.set() + await finishing + records = memory.pop() + segment = records[0]["metadata"]["audio.segments"]["part-0"] + assert segment["state"] == "ready" + assert records[0]["input"]["audio"]["part-0"] == uploads[0].reference + assert not memory.upload_attempts + assert len(uploads) == 1 + + +@pytest.mark.asyncio +async def test_customizer_removal_does_not_veto_upload(memory, monkeypatch): + uploads = [] + + class RemoveAudio(SpanCustomizer): + def on_span_export(self, data): + data.pop("input", None) + return data + + def upload(attachment): + uploads.append(attachment) + return {"upload_status": "done"} + + set_span_customizers([RemoveAudio()]) + monkeypatch.setattr(logger.Attachment, "upload", upload) + session = recording() + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + await session.finish() + assert len(uploads) == 1 + assert "input" not in memory.pop()[0] + + +@pytest.mark.asyncio +async def test_privacy_tightening_before_upload_discards_audio(memory, monkeypatch): + def forbidden(_): + pytest.fail("Restricted audio must never upload") + + monkeypatch.setattr(logger.Attachment, "upload", forbidden) + allowed = True + session = recording() + session.allowed = lambda: allowed + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + allowed = False + await session.finish() + segment = memory.pop()[0]["metadata"]["audio.segments"]["part-0"] + assert segment == {"state": "omitted", "reason": "privacy_restricted", "start_offset_ms": 0, "duration_ms": 1} + + +@pytest.mark.asyncio +async def test_terminal_upload_failure_preserves_known_placement(memory, monkeypatch): + monkeypatch.setattr(logger.Attachment, "upload", lambda _: {"upload_status": "error"}) + session = recording() + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + result = await session.finish() + assert result.reason == "encoding_or_upload_failed" + segment = memory.pop()[0]["metadata"]["audio.segments"]["part-0"] + assert segment["state"] == "omitted" + assert segment["reason"] == "upload_failed" + assert segment["duration_ms"] == 1 + assert "attachment" not in segment + + +@pytest.mark.asyncio +async def test_shutdown_deadline_fences_late_success(memory, monkeypatch): + entered, release = threading.Event(), threading.Event() + + def upload(_): + entered.set() + release.wait(5) + return {"upload_status": "done"} + + monkeypatch.setattr(logger.Attachment, "upload", upload) + session = recording(shutdown_timeout_seconds=0.02) + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + result = await session.finish() + assert result.reason == "shutdown_timeout" + assert entered.is_set() + records = memory.pop() + assert records[0]["metadata"]["audio.segments"]["part-0"]["state"] == "omitted" + release.set() + await asyncio.sleep(0.02) + assert not memory.pop() diff --git a/py/src/braintrust/_audio/types.py b/py/src/braintrust/_audio/types.py new file mode 100644 index 00000000..6bf9d91e --- /dev/null +++ b/py/src/braintrust/_audio/types.py @@ -0,0 +1,69 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any + + +def freeze(value): + if isinstance(value, Mapping): + return MappingProxyType({key: freeze(item) for key, item in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(freeze(item) for item in value) + if value is None or isinstance(value, (str, bool, int, float)): + return value + raise TypeError("Audio attributes must contain JSON values") + + +def thaw(value): + if isinstance(value, Mapping): + return {key: thaw(item) for key, item in value.items()} + if isinstance(value, tuple): + return [thaw(item) for item in value] + return value + + +@dataclass(frozen=True) +class AudioChannel: + sources: tuple[str, ...] + + def __post_init__(self): + object.__setattr__(self, "sources", tuple(self.sources)) + + +@dataclass(frozen=True) +class SegmentMeta: + index: int + start_sample: int + sample_count: int + sample_rate_hz: int + channels: tuple[AudioChannel, ...] + mime_type: str = "audio/pcm" + encoding: str = "pcm" + attributes: Mapping[str, Any] = field( + default_factory=lambda: {"pcm.sample_format": "s16le", "pcm.channel_layout": "interleaved"} + ) + + def __post_init__(self): + object.__setattr__(self, "channels", tuple(self.channels)) + object.__setattr__(self, "attributes", freeze(self.attributes)) + if self.index < 0 or self.start_sample < 0 or self.sample_count <= 0 or self.sample_rate_hz <= 0: + raise ValueError("Invalid audio segment placement") + if not self.channels or not self.mime_type or not self.encoding: + raise ValueError("Invalid audio segment format") + + def placement(self): + start = round(self.start_sample * 1000 / self.sample_rate_hz) + end = round((self.start_sample + self.sample_count) * 1000 / self.sample_rate_hz) + return {"start_offset_ms": start, "duration_ms": end - start} + + +@dataclass(frozen=True) +class AudioResult: + data: bytes + meta: SegmentMeta + + +@dataclass(frozen=True) +class RecordingResult: + segments: int + reason: str | None = None diff --git a/py/src/braintrust/_audio/worker.py b/py/src/braintrust/_audio/worker.py new file mode 100644 index 00000000..68e73999 --- /dev/null +++ b/py/src/braintrust/_audio/worker.py @@ -0,0 +1,55 @@ +"""Bound native submission; cancellation keeps worker-owned buffers alive.""" + +import asyncio +import os +import threading +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from functools import partial + + +_lock = threading.Lock() +_executor: ThreadPoolExecutor | None = None +_pid: int | None = None +_slots: threading.Semaphore | None = None +_waiters: deque | None = None + + +async def run_native(function, *args): + global _executor, _pid, _slots, _waiters + with _lock: + if _executor is None or _pid != os.getpid(): + _executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="braintrust-audio") + _slots = threading.Semaphore(2) + _waiters = deque() + _pid = os.getpid() + executor, slots, waiters = _executor, _slots, _waiters + assert slots is not None + assert waiters is not None + # Waiting sessions retain their buffers. Never queue unbounded native jobs + # or block the voice loop on a threading semaphore. + request = object() + with _lock: + waiters.append(request) + try: + while True: + with _lock: + if waiters[0] is request and slots.acquire(blocking=False): + waiters.popleft() + break + await asyncio.sleep(0.005) + except BaseException: + with _lock: + waiters.remove(request) + raise + try: + future = executor.submit(partial(function, *args)) + except BaseException: + slots.release() + raise + future.add_done_callback(lambda _: slots.release()) + wrapped = asyncio.wrap_future(future) + # A cancelled waiter cannot consume a late native failure. Retrieve it + # without logging potentially private exception details. + wrapped.add_done_callback(lambda completed: completed.exception() if not completed.cancelled() else None) + return await asyncio.shield(wrapped) diff --git a/py/src/braintrust/audio.py b/py/src/braintrust/audio.py new file mode 100644 index 00000000..10fc587a --- /dev/null +++ b/py/src/braintrust/audio.py @@ -0,0 +1,52 @@ +"""Lightweight, framework-independent audio recording configuration. + +Configuration never enables capture. Integrations require their dedicated audio +opt-ins and enforce their framework's privacy policy independently. +""" + +import math +from dataclasses import dataclass + + +@dataclass(frozen=True, kw_only=True) +class AudioEncoding: + codec: str = "opus" + container: str | None = "ogg" + bitrate_bps: int = 64000 + complexity: int = 10 + application: str = "audio" + + def __post_init__(self): + if (self.codec, self.container) not in (("opus", "ogg"), ("pcm", None), ("pcm", "wav")): + raise ValueError("Supported audio formats are opus/ogg, pcm, and pcm/wav") + if ( + type(self.bitrate_bps) is not int + or type(self.complexity) is not int + or not 500 <= self.bitrate_bps <= 512000 + or not 0 <= self.complexity <= 10 + ): + raise ValueError("Invalid Opus bitrate or complexity") + if self.application not in ("audio", "voip", "lowdelay"): + raise ValueError("Invalid Opus application") + if self.codec == "pcm" and (self.bitrate_bps, self.complexity, self.application) != (64000, 10, "audio"): + raise ValueError("Opus options cannot be used with PCM") + + +@dataclass(frozen=True, kw_only=True) +class AudioRecordingOptions: + segment_duration_seconds: float = 60 + flush_threshold_bytes: int = 32 * 1024 * 1024 + encoding: AudioEncoding | None = None + max_duration_seconds: float | None = None + shutdown_timeout_seconds: float = 30 + + def __post_init__(self): + for value in (self.segment_duration_seconds, self.shutdown_timeout_seconds, self.max_duration_seconds): + if value is not None and (not math.isfinite(value) or value <= 0): + raise ValueError("Audio durations must be positive and finite") + if self.segment_duration_seconds < 1 / 48000: + raise ValueError("An audio segment must contain at least one sample") + if type(self.flush_threshold_bytes) is not int or self.flush_threshold_bytes <= 0: + raise ValueError("Audio flush threshold must be positive") + if self.encoding is not None and not isinstance(self.encoding, AudioEncoding): + raise TypeError("Audio encoding must be an AudioEncoding instance") diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index f89c4852..a0b9ced8 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -3830,6 +3830,9 @@ def __init__( "key": str(uuid.uuid4()), } self._data_debug_string = data if isinstance(data, str) else "" + # Internal audio uploads opt into finite request timeouts. Ordinary + # attachments retain their existing transport behavior. + self._upload_timeout_seconds: float | None = None self._data = self._init_data(data) self._uploader = self._init_uploader() @@ -3861,6 +3864,9 @@ def debug_info(self) -> Mapping[str, Any]: return {"input_data": self._data_debug_string, "reference": self._reference} def _init_uploader(self) -> LazyValue[AttachmentStatus]: + def request_options(): + return {"timeout": self._upload_timeout_seconds} if self._upload_timeout_seconds is not None else {} + def do_upload(api_conn: HTTPConnection, org_id: str) -> Mapping[str, Any]: assert self._reference["type"] == "braintrust_attachment" @@ -3872,7 +3878,7 @@ def do_upload(api_conn: HTTPConnection, org_id: str) -> Mapping[str, Any]: } try: - metadata_response = api_conn.post("/attachment", json=request_params) + metadata_response = api_conn.post("/attachment", json=request_params, **request_options()) metadata_response.raise_for_status() metadata = metadata_response.json() except Exception as e: @@ -3893,7 +3899,7 @@ def do_upload(api_conn: HTTPConnection, org_id: str) -> Mapping[str, Any]: # TODO multipart upload. try: obj_conn = HTTPConnection(base_url="", adapter=_http_adapter) - obj_response = obj_conn.put(signed_url, headers=headers, data=data) + obj_response = obj_conn.put(signed_url, headers=headers, data=data, **request_options()) obj_response.raise_for_status() except Exception as e: raise RuntimeError(f"Failed to upload attachment to object store: {e}") from e @@ -3925,7 +3931,7 @@ def error_wrapper() -> AttachmentStatus: "status": status, } try: - status_response = api_conn.post("/attachment/status", json=request_params) + status_response = api_conn.post("/attachment/status", json=request_params, **request_options()) status_response.raise_for_status() except Exception as e: raise RuntimeError(f"Couldn't log attachment status: {e}") from e diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 514e4918..0142b977 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -2543,6 +2543,34 @@ def broken_masking_function(data): assert "private" not in json.dumps(record) +@pytest.mark.parametrize("timeout", [None, 12.5]) +def test_attachment_internal_timeout_preserves_ordinary_behavior(monkeypatch, with_simulate_login, timeout): + requests_seen = [] + + def post(connection, path, *args, **kwargs): + requests_seen.append(kwargs.get("timeout")) + response = requests.Response() + response.status_code = 200 + response._content = json.dumps( + {"signedUrl": "https://objectstore.example/upload", "headers": {}} if path == "/attachment" else {} + ).encode() + return response + + def put(connection, path, *args, **kwargs): + requests_seen.append(kwargs.get("timeout")) + response = requests.Response() + response.status_code = 200 + return response + + monkeypatch.setattr(logger.HTTPConnection, "post", post) + monkeypatch.setattr(logger.HTTPConnection, "put", put) + attachment = Attachment(data=b"audio", filename="part.pcm", content_type="audio/pcm") + attachment._upload_timeout_seconds = timeout + assert attachment.upload()["upload_status"] == "done" + assert attachment.upload()["upload_status"] == "done" + assert requests_seen == [timeout, timeout, timeout] + + def test_attachment_unreadable_path_logs_warning(caplog): with caplog.at_level(logging.WARNING, logger="braintrust"): Attachment( From 5d5cb23e0b28693ff0abd7d1258a60d5c5010abf Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:09:19 -0500 Subject: [PATCH 02/12] ref: remove braintrust.audio public api --- py/src/braintrust/{audio.py => _audio/config.py} | 0 py/src/braintrust/_audio/middleware.py | 2 +- py/src/braintrust/_audio/recorder.py | 2 +- py/src/braintrust/_audio/session.py | 2 +- py/src/braintrust/_audio/test_audio.py | 2 +- py/src/braintrust/_audio/test_session.py | 2 +- 6 files changed, 5 insertions(+), 5 deletions(-) rename py/src/braintrust/{audio.py => _audio/config.py} (100%) diff --git a/py/src/braintrust/audio.py b/py/src/braintrust/_audio/config.py similarity index 100% rename from py/src/braintrust/audio.py rename to py/src/braintrust/_audio/config.py diff --git a/py/src/braintrust/_audio/middleware.py b/py/src/braintrust/_audio/middleware.py index aec8e6f1..439632c8 100644 --- a/py/src/braintrust/_audio/middleware.py +++ b/py/src/braintrust/_audio/middleware.py @@ -7,7 +7,7 @@ from fractions import Fraction from typing import Any, Protocol -from braintrust.audio import AudioEncoding +from braintrust._audio.config import AudioEncoding from .types import AudioResult, SegmentMeta from .worker import run_native diff --git a/py/src/braintrust/_audio/recorder.py b/py/src/braintrust/_audio/recorder.py index de3e120f..7eef47ba 100644 --- a/py/src/braintrust/_audio/recorder.py +++ b/py/src/braintrust/_audio/recorder.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from fractions import Fraction -from braintrust.audio import AudioRecordingOptions +from braintrust._audio.config import AudioRecordingOptions from .middleware import AudioMiddleware, run_middleware from .types import AudioChannel, AudioResult, RecordingResult, SegmentMeta diff --git a/py/src/braintrust/_audio/session.py b/py/src/braintrust/_audio/session.py index 3b93d38f..8fa65410 100644 --- a/py/src/braintrust/_audio/session.py +++ b/py/src/braintrust/_audio/session.py @@ -3,7 +3,7 @@ import time from collections.abc import Callable -from braintrust.audio import AudioRecordingOptions +from braintrust._audio.config import AudioRecordingOptions from braintrust.env import BraintrustEnv from braintrust.logger import Attachment diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index ef6d8ef9..ec4b75d0 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -9,11 +9,11 @@ from fractions import Fraction import pytest +from braintrust._audio.config import AudioEncoding, AudioRecordingOptions from braintrust._audio.middleware import encoding_stack, run_middleware from braintrust._audio.recorder import AudioRecorder from braintrust._audio.types import AudioChannel, AudioResult, SegmentMeta from braintrust._audio.worker import run_native -from braintrust.audio import AudioEncoding, AudioRecordingOptions def pcm(values): diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py index 6eddb4f6..082f4db5 100644 --- a/py/src/braintrust/_audio/test_session.py +++ b/py/src/braintrust/_audio/test_session.py @@ -3,8 +3,8 @@ import pytest from braintrust import SpanCustomizer, logger, set_span_customizers +from braintrust._audio.config import AudioEncoding, AudioRecordingOptions from braintrust._audio.session import AudioSessionRecorder -from braintrust.audio import AudioEncoding, AudioRecordingOptions from braintrust.test_helpers import init_test_logger From 346466bafb146541ff5e10107eb5160e9b7ae8c0 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:10:48 -0500 Subject: [PATCH 03/12] ref: re-raise cancellations we should not swallow --- py/src/braintrust/_audio/session.py | 4 ++- py/src/braintrust/_audio/test_session.py | 32 ++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/_audio/session.py b/py/src/braintrust/_audio/session.py index 8fa65410..3bd0fc96 100644 --- a/py/src/braintrust/_audio/session.py +++ b/py/src/braintrust/_audio/session.py @@ -177,7 +177,7 @@ async def finish(self) -> RecordingResult: try: if self.recorder is not None: result = await asyncio.wait_for(self.recorder.finish(), self.options.shutdown_timeout_seconds) - except (asyncio.TimeoutError, asyncio.CancelledError): + except (asyncio.TimeoutError, asyncio.CancelledError) as error: self._closed = True for meta in tuple(self._inflight.values()): self.omitted(meta, "shutdown_timeout") @@ -187,6 +187,8 @@ async def finish(self) -> RecordingResult: self.recorder._task.cancel() await asyncio.gather(self.recorder._task, return_exceptions=True) result = RecordingResult(len(self._terminal), "shutdown_timeout") + if isinstance(error, asyncio.CancelledError): + raise finally: self._closed = True if result.reason: diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py index 082f4db5..3df6787d 100644 --- a/py/src/braintrust/_audio/test_session.py +++ b/py/src/braintrust/_audio/test_session.py @@ -126,6 +126,38 @@ async def test_terminal_upload_failure_preserves_known_placement(memory, monkeyp assert "attachment" not in segment +@pytest.mark.asyncio +async def test_finish_propagates_cancellation_after_cleanup(memory): + entered, cleaned_up = asyncio.Event(), asyncio.Event() + + async def blocking(data, meta, next): + entered.set() + try: + await asyncio.Event().wait() + finally: + cleaned_up.set() + + session = recording() + session.recorder.middleware = (blocking,) + session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + finishing = asyncio.create_task(session.finish()) + try: + await asyncio.wait_for(entered.wait(), 5) + finishing.cancel() + with pytest.raises(asyncio.CancelledError): + await finishing + assert cleaned_up.is_set() + assert finishing.cancelled() + assert session.recorder.buffered_bytes == 0 + assert not session.record("session-input", bytes(96), sample_rate_hz=48000, channels=1, offset_ns=0) + segment = memory.pop()[0]["metadata"]["audio.segments"]["part-0"] + assert segment["state"] == "omitted" + assert "attachment" not in segment + finally: + finishing.cancel() + await asyncio.gather(finishing, return_exceptions=True) + + @pytest.mark.asyncio async def test_shutdown_deadline_fences_late_success(memory, monkeypatch): entered, release = threading.Event(), threading.Event() From 2c94fce9d133faf43df9478eeba409bb9edd8b01 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:13:15 -0500 Subject: [PATCH 04/12] ref: do not access private recorder api --- py/src/braintrust/_audio/recorder.py | 7 +++++++ py/src/braintrust/_audio/session.py | 5 +---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/py/src/braintrust/_audio/recorder.py b/py/src/braintrust/_audio/recorder.py index 7eef47ba..98817102 100644 --- a/py/src/braintrust/_audio/recorder.py +++ b/py/src/braintrust/_audio/recorder.py @@ -267,6 +267,13 @@ async def _drain(self): self._packets = [packet for packet in self._packets if packet.end > end] self._processing = None + async def abort(self, reason: str) -> None: + """Discard queued audio and wait for the background task to cancel.""" + self.discard(reason) + if self._task is not None: + self._task.cancel() + await asyncio.gather(self._task, return_exceptions=True) + async def finish(self) -> RecordingResult: self._finishing = True # The adapter settles pending playback before calling finish. diff --git a/py/src/braintrust/_audio/session.py b/py/src/braintrust/_audio/session.py index 3bd0fc96..0253bb4d 100644 --- a/py/src/braintrust/_audio/session.py +++ b/py/src/braintrust/_audio/session.py @@ -182,10 +182,7 @@ async def finish(self) -> RecordingResult: for meta in tuple(self._inflight.values()): self.omitted(meta, "shutdown_timeout") if self.recorder is not None: - self.recorder.discard("shutdown_timeout") - if self.recorder._task is not None: - self.recorder._task.cancel() - await asyncio.gather(self.recorder._task, return_exceptions=True) + await self.recorder.abort("shutdown_timeout") result = RecordingResult(len(self._terminal), "shutdown_timeout") if isinstance(error, asyncio.CancelledError): raise From 764d3729142334ec6d7581c037c2298ce2a62eda Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:14:46 -0500 Subject: [PATCH 05/12] fix: stronger typing on config --- py/src/braintrust/_audio/config.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/py/src/braintrust/_audio/config.py b/py/src/braintrust/_audio/config.py index 10fc587a..4b9d615b 100644 --- a/py/src/braintrust/_audio/config.py +++ b/py/src/braintrust/_audio/config.py @@ -6,15 +6,16 @@ import math from dataclasses import dataclass +from typing import Literal @dataclass(frozen=True, kw_only=True) class AudioEncoding: - codec: str = "opus" - container: str | None = "ogg" + codec: Literal["opus", "pcm"] = "opus" + container: Literal["ogg", "wav"] | None = "ogg" bitrate_bps: int = 64000 complexity: int = 10 - application: str = "audio" + application: Literal["audio", "voip", "lowdelay"] = "audio" def __post_init__(self): if (self.codec, self.container) not in (("opus", "ogg"), ("pcm", None), ("pcm", "wav")): From 471cc97354a8737b9a5cc8139dc6d7346b10fe7a Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:49:21 -0500 Subject: [PATCH 06/12] ref: simplify serialization - clean up some export logic - reduce repeated middleware processing --- py/src/braintrust/_audio/middleware.py | 41 ++++++++++++------------ py/src/braintrust/_audio/session.py | 6 ++-- py/src/braintrust/_audio/test_audio.py | 32 ++++++++++++++++++ py/src/braintrust/_audio/test_session.py | 2 ++ py/src/braintrust/_audio/types.py | 8 ----- 5 files changed, 57 insertions(+), 32 deletions(-) diff --git a/py/src/braintrust/_audio/middleware.py b/py/src/braintrust/_audio/middleware.py index 439632c8..e964be2a 100644 --- a/py/src/braintrust/_audio/middleware.py +++ b/py/src/braintrust/_audio/middleware.py @@ -43,27 +43,26 @@ async def run_middleware(stack: Sequence[AudioMiddleware], payload: bytes, meta: async def run(index, data, current): _validate(current, meta) if index == len(stack): - result = AudioResult(data, current) - else: - called = False - - def next(data, new_meta): - nonlocal called - if called: - raise ValueError("Audio middleware called next more than once") - called = True - return run(index + 1, data, new_meta) - - result = await stack[index](data, current, next) - if not isinstance(result, AudioResult) or not isinstance(result.data, bytes): - raise ValueError("Audio middleware must return final bytes") - _validate(result.meta, meta) - if result.meta.encoding == "pcm" and result.meta.mime_type == "audio/pcm": - if len(result.data) != result.meta.sample_count * len(result.meta.channels) * 2: - raise ValueError("PCM byte length does not match playable duration") - return result - - return await run(0, payload, meta) + return AudioResult(data, current) + called = False + + def next(data, new_meta): + nonlocal called + if called: + raise ValueError("Audio middleware called next more than once") + called = True + return run(index + 1, data, new_meta) + + return await stack[index](data, current, next) + + result = await run(0, payload, meta) + if not isinstance(result, AudioResult) or not isinstance(result.data, bytes): + raise ValueError("Audio middleware must return final bytes") + _validate(result.meta, meta) + if result.meta.encoding == "pcm" and result.meta.mime_type == "audio/pcm": + if len(result.data) != result.meta.sample_count * len(result.meta.channels) * 2: + raise ValueError("PCM byte length does not match playable duration") + return result def _wav(data, meta): diff --git a/py/src/braintrust/_audio/session.py b/py/src/braintrust/_audio/session.py index 0253bb4d..970b1045 100644 --- a/py/src/braintrust/_audio/session.py +++ b/py/src/braintrust/_audio/session.py @@ -9,7 +9,7 @@ from .middleware import encoding_stack from .recorder import AudioRecorder -from .types import AudioResult, RecordingResult, SegmentMeta, thaw +from .types import AudioResult, RecordingResult, SegmentMeta from .worker import run_native @@ -162,9 +162,9 @@ def upload(): "encoding": meta.encoding, "sample_rate_hz": meta.sample_rate_hz, "channels": { - str(i): {"sources": list(channel.sources)} for i, channel in enumerate(meta.channels) + str(i): {"sources": channel.sources} for i, channel in enumerate(meta.channels) }, - "attributes": thaw(meta.attributes), + "attributes": meta.attributes, } } }, diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index ec4b75d0..a2442333 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -141,6 +141,38 @@ async def twice(data, current, next): await run_middleware((twice,), bytes(96), meta) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("data", "message"), + [(None, "final bytes"), (bytearray(96), "final bytes"), (bytes(94), "PCM byte length")], +) +async def test_middleware_validates_final_result_after_nested_layers(data, message): + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) + + async def passthrough(data, current, next): + return await next(data, current) + + async def invalid(payload, current, next): + return None if data is None else AudioResult(data, current) + + with pytest.raises(ValueError, match=message): + await run_middleware((passthrough, passthrough, invalid), bytes(96), meta) + + +@pytest.mark.asyncio +async def test_middleware_rejects_invalid_metadata_before_calling_downstream(): + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) + + async def invalid(data, current, next): + return await next(data, replace(current, sample_count=49)) + + async def downstream(data, current, next): + pytest.fail("Invalid segment metadata reached downstream middleware") + + with pytest.raises(ValueError, match="identity, placement"): + await run_middleware((invalid, downstream), bytes(96), meta) + + @pytest.mark.asyncio async def test_wav_requires_no_optional_backend_and_preserves_tiny_tail(): meta = SegmentMeta(0, 0, 137, 48000, (AudioChannel(("input",)),)) diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py index 3df6787d..78808375 100644 --- a/py/src/braintrust/_audio/test_session.py +++ b/py/src/braintrust/_audio/test_session.py @@ -69,6 +69,8 @@ def upload(attachment): records = memory.pop() segment = records[0]["metadata"]["audio.segments"]["part-0"] assert segment["state"] == "ready" + assert segment["channels"] == {"0": {"sources": ["session-input"]}} + assert segment["attributes"] == {"pcm.sample_format": "s16le", "pcm.channel_layout": "interleaved"} assert records[0]["input"]["audio"]["part-0"] == uploads[0].reference assert not memory.upload_attempts assert len(uploads) == 1 diff --git a/py/src/braintrust/_audio/types.py b/py/src/braintrust/_audio/types.py index 6bf9d91e..4c675518 100644 --- a/py/src/braintrust/_audio/types.py +++ b/py/src/braintrust/_audio/types.py @@ -14,14 +14,6 @@ def freeze(value): raise TypeError("Audio attributes must contain JSON values") -def thaw(value): - if isinstance(value, Mapping): - return {key: thaw(item) for key, item in value.items()} - if isinstance(value, tuple): - return [thaw(item) for item in value] - return value - - @dataclass(frozen=True) class AudioChannel: sources: tuple[str, ...] From eca20b73eaa44deb577215d0b87a89b80f36807c Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:53:34 -0500 Subject: [PATCH 07/12] use a frozen mapping for less expensive operations --- py/src/braintrust/_audio/test_audio.py | 23 +++++++++++++++++++++++ py/src/braintrust/_audio/types.py | 25 +++++++++++++++++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index a2442333..55f85e4c 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -7,6 +7,7 @@ from array import array from dataclasses import replace from fractions import Fraction +from types import MappingProxyType import pytest from braintrust._audio.config import AudioEncoding, AudioRecordingOptions @@ -109,6 +110,28 @@ async def slow_upload(_): assert not omitted +@pytest.mark.parametrize("mapping", [dict, MappingProxyType]) +def test_segment_attributes_are_reused_and_isolated_from_caller_mutations(mapping): + attributes = {"nested": {"list": [1]}} + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),), attributes=mapping(attributes)) + changed = replace(meta, mime_type="audio/wav") + assert changed.attributes is meta.attributes + + attributes["nested"]["list"].append(2) + attributes["extra"] = True + assert meta.attributes == {"nested": {"list": (1,)}} + with pytest.raises(TypeError): + meta.attributes["extra"] = True + with pytest.raises(TypeError): + meta.attributes["nested"]["extra"] = True + + replacement = {"nested": {"list": [3]}} + changed = replace(meta, attributes=mapping(replacement)) + replacement["nested"]["list"].append(4) + assert changed.attributes is not meta.attributes + assert changed.attributes == {"nested": {"list": (3,)}} + + @pytest.mark.asyncio async def test_middleware_validates_duration_sources_bytes_and_deep_freeze(): meta = SegmentMeta( diff --git a/py/src/braintrust/_audio/types.py b/py/src/braintrust/_audio/types.py index 4c675518..f8dcd8b3 100644 --- a/py/src/braintrust/_audio/types.py +++ b/py/src/braintrust/_audio/types.py @@ -1,12 +1,33 @@ -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import Any +@dataclass(frozen=True, eq=False) +class _FrozenMapping(Mapping[str, Any]): + """A snapshot whose values have been recursively frozen.""" + + _data: Mapping[str, Any] + + def __post_init__(self): + object.__setattr__(self, "_data", MappingProxyType({key: freeze(value) for key, value in self._data.items()})) + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._data) + + def __len__(self) -> int: + return len(self._data) + + def freeze(value): + if isinstance(value, _FrozenMapping): + return value if isinstance(value, Mapping): - return MappingProxyType({key: freeze(item) for key, item in value.items()}) + return _FrozenMapping(value) if isinstance(value, (list, tuple)): return tuple(freeze(item) for item in value) if value is None or isinstance(value, (str, bool, int, float)): From 90ed867683c6cf27453053b3c614920d0d08e8d4 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 12:59:42 -0500 Subject: [PATCH 08/12] test: combine redundant tests --- py/noxfile.py | 6 ---- py/src/braintrust/_audio/test_audio.py | 41 ++++++++---------------- py/src/braintrust/_audio/test_session.py | 1 - 3 files changed, 14 insertions(+), 34 deletions(-) diff --git a/py/noxfile.py b/py/noxfile.py index 47cc5bca..5ecb41e7 100644 --- a/py/noxfile.py +++ b/py/noxfile.py @@ -878,12 +878,6 @@ def is_importable(name): _run_core_tests(session) -@nox.session() -def test_audio(session): - _install_test_deps(session) - _run_tests(session, "braintrust/_audio") - - # Not a default session: it needs a wheel from `make build`, so a bare `nox` # (`make test`) must not pick it up. @nox.session(default=False) diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index 55f85e4c..d26afed5 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -124,6 +124,8 @@ def test_segment_attributes_are_reused_and_isolated_from_caller_mutations(mappin meta.attributes["extra"] = True with pytest.raises(TypeError): meta.attributes["nested"]["extra"] = True + with pytest.raises(TypeError): + meta.attributes["nested"]["list"][0] = 2 replacement = {"nested": {"list": [3]}} changed = replace(meta, attributes=mapping(replacement)) @@ -133,50 +135,35 @@ def test_segment_attributes_are_reused_and_isolated_from_caller_mutations(mappin @pytest.mark.asyncio -async def test_middleware_validates_duration_sources_bytes_and_deep_freeze(): - meta = SegmentMeta( - 0, - 0, - 48, - 48000, - (AudioChannel(("input",)),), - attributes={ - "nested": {"list": [1]}, - "pcm.sample_format": "s16le", - "pcm.channel_layout": "interleaved", - }, - ) - with pytest.raises(TypeError): - meta.attributes["nested"]["list"][0] = 2 - - async def invalid(data, current, next): - return AudioResult(data, replace(current, sample_count=49)) - - with pytest.raises(ValueError, match="identity, placement"): - await run_middleware((invalid,), bytes(96), meta) +async def test_middleware_rejects_calling_next_twice(): + meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) async def twice(data, current, next): await next(data, current) return await next(data, current) - meta = replace(meta, attributes={"pcm.sample_format": "s16le", "pcm.channel_layout": "interleaved"}) with pytest.raises(ValueError, match="more than once"): await run_middleware((twice,), bytes(96), meta) @pytest.mark.asyncio @pytest.mark.parametrize( - ("data", "message"), - [(None, "final bytes"), (bytearray(96), "final bytes"), (bytes(94), "PCM byte length")], + ("data", "sample_count", "message"), + [ + (None, 48, "final bytes"), + (bytearray(96), 48, "final bytes"), + (bytes(94), 48, "PCM byte length"), + (bytes(96), 49, "identity, placement"), + ], ) -async def test_middleware_validates_final_result_after_nested_layers(data, message): +async def test_middleware_validates_final_result_after_nested_layers(data, sample_count, message): meta = SegmentMeta(0, 0, 48, 48000, (AudioChannel(("input",)),)) async def passthrough(data, current, next): return await next(data, current) async def invalid(payload, current, next): - return None if data is None else AudioResult(data, current) + return None if data is None else AudioResult(data, replace(current, sample_count=sample_count)) with pytest.raises(ValueError, match=message): await run_middleware((passthrough, passthrough, invalid), bytes(96), meta) @@ -206,7 +193,7 @@ async def test_wav_requires_no_optional_backend_and_preserves_tiny_tail(): @pytest.mark.asyncio -async def test_native_worker_cancellation_keeps_buffer_owned_until_completion(): +async def test_native_work_completes_after_awaiter_cancellation(): started, release, completed = threading.Event(), threading.Event(), threading.Event() payload = bytearray(b"owned") diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py index 78808375..0eedb1af 100644 --- a/py/src/braintrust/_audio/test_session.py +++ b/py/src/braintrust/_audio/test_session.py @@ -10,7 +10,6 @@ @pytest.fixture def memory(monkeypatch): - monkeypatch.setattr(logger, "_compute_logger_metadata", logger._compute_logger_metadata) init_test_logger("audio-tests") with logger._internal_with_memory_background_logger() as bgl: monkeypatch.setattr(logger._state, "_global_bg_logger", logger.LazyValue(lambda: bgl, use_mutex=False)) From a67658b923d3eb9c2507e24e8a0640a7a560c2f2 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 13:07:56 -0500 Subject: [PATCH 09/12] ref: use a FIFO queue instead of sleeping --- py/src/braintrust/_audio/test_audio.py | 29 +++++++----- py/src/braintrust/_audio/worker.py | 61 +++++++++++++++++++------- 2 files changed, 63 insertions(+), 27 deletions(-) diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index d26afed5..121fdfab 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -213,11 +213,18 @@ def native(data): @pytest.mark.asyncio -async def test_native_worker_admits_waiting_sessions_in_order_and_removes_cancelled_waiters(): - from braintrust._audio import worker - +@pytest.mark.parametrize("separate_loop", [False, True]) +async def test_native_worker_admits_waiting_sessions_in_order_and_removes_cancelled_waiters(separate_loop): started = [threading.Event() for _ in range(4)] release = [threading.Event() for _ in range(3)] + loop = asyncio.get_running_loop() + queued = asyncio.Event() + + async def first_waiter(): + task = asyncio.create_task(run_native(blocking, 2)) + # Notify after the task has had a turn to request admission on its loop. + asyncio.get_running_loop().call_soon(loop.call_soon_threadsafe, queued.set) + return await task def blocking(index): started[index].set() @@ -227,16 +234,18 @@ def blocking(index): try: for event in started[:2]: assert await asyncio.to_thread(event.wait, 5) - first = asyncio.create_task(run_native(blocking, 2)) + first = asyncio.create_task( + asyncio.to_thread(asyncio.run, first_waiter()) if separate_loop else first_waiter() + ) + tasks.append(first) + await asyncio.wait_for(queued.wait(), 5) cancelled = asyncio.create_task(run_native(lambda: pytest.fail("Cancelled job was submitted"))) last = asyncio.create_task(run_native(started[3].set)) - tasks.extend([first, cancelled, last]) - - async def wait_for_queue(): - while worker._waiters is None or len(worker._waiters) != 3: - await asyncio.sleep(0.005) + tasks.extend([cancelled, last]) - await asyncio.wait_for(wait_for_queue(), 5) + queue_ready = asyncio.Event() + loop.call_soon(queue_ready.set) + await asyncio.wait_for(queue_ready.wait(), 5) cancelled.cancel() with pytest.raises(asyncio.CancelledError): await cancelled diff --git a/py/src/braintrust/_audio/worker.py b/py/src/braintrust/_audio/worker.py index 68e73999..df9d7ca5 100644 --- a/py/src/braintrust/_audio/worker.py +++ b/py/src/braintrust/_audio/worker.py @@ -12,11 +12,36 @@ _executor: ThreadPoolExecutor | None = None _pid: int | None = None _slots: threading.Semaphore | None = None -_waiters: deque | None = None +_waiters: deque[asyncio.Future[None]] | None = None + + +def _admit(request: asyncio.Future[None], slots: threading.Semaphore, waiters: deque[asyncio.Future[None]]) -> None: + # This callback runs on the waiter's event loop, including cancellation checks. + if request.cancelled(): + _release(slots, waiters) + else: + request.set_result(None) + + +def _release(slots: threading.Semaphore, waiters: deque[asyncio.Future[None]]) -> None: + # Transfer the occupied slot directly to the oldest waiter. + while True: + with _lock: + if not waiters: + slots.release() + return + request = waiters.popleft() + try: + request.get_loop().call_soon_threadsafe(_admit, request, slots, waiters) + return + except RuntimeError: + # A closed event loop cannot accept the slot. Try the next waiter. + continue async def run_native(function, *args): global _executor, _pid, _slots, _waiters + request = asyncio.get_running_loop().create_future() with _lock: if _executor is None or _pid != os.getpid(): _executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="braintrust-audio") @@ -24,30 +49,32 @@ async def run_native(function, *args): _waiters = deque() _pid = os.getpid() executor, slots, waiters = _executor, _slots, _waiters - assert slots is not None - assert waiters is not None - # Waiting sessions retain their buffers. Never queue unbounded native jobs - # or block the voice loop on a threading semaphore. - request = object() - with _lock: - waiters.append(request) + assert slots is not None + assert waiters is not None + # Waiting sessions retain their buffers; only admitted jobs reach the executor. + if slots.acquire(blocking=False): + request.set_result(None) + else: + waiters.append(request) try: - while True: - with _lock: - if waiters[0] is request and slots.acquire(blocking=False): - waiters.popleft() - break - await asyncio.sleep(0.005) + await request except BaseException: with _lock: - waiters.remove(request) + try: + waiters.remove(request) + except ValueError: + pass + # If admission won the race with cancellation, return its reserved slot. + # Otherwise a pending _admit callback owns the handoff, if any. + if not request.cancelled(): + _release(slots, waiters) raise try: future = executor.submit(partial(function, *args)) except BaseException: - slots.release() + _release(slots, waiters) raise - future.add_done_callback(lambda _: slots.release()) + future.add_done_callback(lambda _: _release(slots, waiters)) wrapped = asyncio.wrap_future(future) # A cancelled waiter cannot consume a late native failure. Retrieve it # without logging potentially private exception details. From 14aa263c5fb527bb3e90dc5b4542691c3e3d687d Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 13:57:27 -0500 Subject: [PATCH 10/12] fix: gh pr review --- py/src/braintrust/_audio/recorder.py | 40 +++++++++++-- py/src/braintrust/_audio/test_audio.py | 81 +++++++++++++++++++++++++- 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/py/src/braintrust/_audio/recorder.py b/py/src/braintrust/_audio/recorder.py index 98817102..17f108d6 100644 --- a/py/src/braintrust/_audio/recorder.py +++ b/py/src/braintrust/_audio/recorder.py @@ -43,6 +43,12 @@ def _samples(data): return result +def _continues(packet: Packet, following: Packet | None): + return following is not None and ( + following.start == packet.end and following.rate == packet.rate and following.channels == packet.channels + ) + + def render(packets: tuple[Packet, ...], sources: tuple[str, ...], start: int, end: int): """Sparse placement and phase-continuous linear resampling, in a worker. @@ -71,7 +77,7 @@ def render(packets: tuple[Packet, ...], sources: tuple[str, ...], start: int, en following = mono[-1] if index + 1 < len(stream): other = stream[index + 1] - if other.start == packet.end and other.rate == packet.rate and other.channels == packet.channels: + if _continues(packet, other): frame = _samples(other.data[: other.channels * 2]) following = sum(frame) // other.channels phase = (Fraction(first) - packet.start) * packet.rate / SAMPLE_RATE @@ -114,6 +120,9 @@ def __init__( self.sources = tuple(sources) self.options = options self.middleware = middleware + self._limit = ( + round(options.max_duration_seconds * SAMPLE_RATE) if options.max_duration_seconds is not None else None + ) self.on_segment = on_segment self.on_omitted = on_omitted self._packets: list[Packet] = [] @@ -145,14 +154,13 @@ def record( start = Fraction(offset_ns) * SAMPLE_RATE / NS if start < self._sealed or start < 0: raise ValueError("Audio arrived before the sealed recording watermark") + if self._limit is not None and start >= self._limit: + return False packet = Packet(source_id, bytes(pcm), sample_rate_hz, channels, start) self._packets.append(packet) self._end = max(self._end, math.ceil(packet.end)) - if self.options.max_duration_seconds is not None: - limit = round(self.options.max_duration_seconds * SAMPLE_RATE) - if self._end >= limit: - self._end = limit - self.stop("duration_limit") + if self._limit is not None: + self._end = min(self._end, self._limit) buffered = self.buffered_bytes if buffered > self.options.flush_threshold_bytes * 4 and not self._warned_backlog: self._warned_backlog = True @@ -164,12 +172,18 @@ def record( def advance(self, through_ns: int | Fraction): self._watermark = max(self._watermark, math.floor(Fraction(through_ns) * SAMPLE_RATE / NS)) + # Only the adapter's watermark establishes that every source has had a + # chance to deliver its in-window audio. + if self._limit is not None and self._watermark >= self._limit: + self.stop("duration_limit") self._pump() def unavailable(self, start_ns, end_ns, reason): """Known bounds of unavailable audio; never encode it as real silence.""" start = max(0, math.floor(Fraction(start_ns) * SAMPLE_RATE / NS)) end = math.ceil(Fraction(end_ns) * SAMPLE_RATE / NS) + if self._limit is not None: + end = min(end, self._limit) if end <= start or self._discarded: return if start < self._sealed: @@ -205,6 +219,18 @@ def discard(self, reason: str): def _cut(self): eligible = min(self._end, self._watermark) + if not self._finishing and self.reason != "duration_limit": + following: dict[str, Packet] = {} + for packet in reversed(self._packets): + if self._watermark <= packet.end and not _continues(packet, following.get(packet.source)): + # A packet starting exactly at the watermark may still + # arrive. Keep its predecessor's interpolation tail until + # that lookahead arrives, the watermark passes the boundary, + # or capture is finalized. Partial flushes retain the packet + # because its end remains beyond the sealed position. + last_sample = packet.start + Fraction((packet.count - 1) * SAMPLE_RATE, packet.rate) + eligible = min(eligible, math.floor(last_sample) + 1) + following[packet.source] = packet target = self._sealed + max(1, round(self.options.segment_duration_seconds * SAMPLE_RATE)) if eligible <= self._sealed: return None @@ -278,6 +304,8 @@ async def finish(self) -> RecordingResult: self._finishing = True # The adapter settles pending playback before calling finish. self._watermark = max(self._watermark, self._end) + if self._limit is not None and self._end >= self._limit: + self.stop("duration_limit") self._pump() if self._task is not None: await asyncio.shield(self._task) diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index 121fdfab..88fc5917 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -21,7 +21,7 @@ def pcm(values): return array("h", values).tobytes() -def recorder(*, duration=0.01, threshold=1024 * 1024, middleware=(), on_segment=None): +def recorder(*, duration=0.01, threshold=1024 * 1024, middleware=(), on_segment=None, max_duration=None): segments, omitted = [], [] async def collect(result): @@ -34,6 +34,7 @@ async def collect(result): options=AudioRecordingOptions( segment_duration_seconds=duration, flush_threshold_bytes=threshold, + max_duration_seconds=max_duration, ), middleware=middleware, on_segment=collect, @@ -66,6 +67,84 @@ async def record(split): assert await record(137) == await record(rate) +@pytest.mark.asyncio +@pytest.mark.parametrize("advance", [False, True]) +@pytest.mark.parametrize("count", [480, 960]) +async def test_duration_limit_waits_for_all_sources(advance, count): + instance, segments, omitted = recorder(max_duration=0.01) + assert instance.record("input", pcm([100] * count), sample_rate_hz=48000, channels=1, offset_ns=0) + assert instance.record("output", pcm([200] * 480), sample_rate_hz=48000, channels=1, offset_ns=0) + assert not instance.record("output", pcm([300] * 480), sample_rate_hz=48000, channels=1, offset_ns=10_000_000) + if advance: + instance.advance(10_000_000) + assert instance.reason == "duration_limit" + assert not instance.record("input", pcm([400]), sample_rate_hz=48000, channels=1, offset_ns=10_000_000) + result = await instance.finish() + samples = array("h") + samples.frombytes(b"".join(s.data for s in segments)) + assert samples == array("h", [100, 200] * 480) + assert result.reason == "duration_limit" + assert not omitted + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rate", [16000, 24000, 44100, 48000]) +@pytest.mark.parametrize("threshold", [1, 1024 * 1024]) +async def test_progressive_resampling_waits_for_packet_lookahead(rate, threshold): + async def record(progressive): + instance, segments, omitted = recorder(threshold=threshold) + for index, value in enumerate((0, 10000, -10000)): + assert instance.record( + "input", + pcm([value] * (rate // 100)), + sample_rate_hz=rate, + channels=1, + offset_ns=index * 10_000_000, + ) + if progressive: + instance.advance((index + 1) * 10_000_000) + if instance._task is not None: + await instance._task + if index == 1: + assert segments # Lookahead releases audio before finish. + await instance.finish() + assert not omitted + return b"".join(s.data for s in segments) + + assert await record(True) == await record(False) + + +@pytest.mark.asyncio +async def test_resampling_tail_released_when_watermark_passes_packet_end(): + instance, segments, omitted = recorder() + instance.record("input", pcm([100] * 240), sample_rate_hz=24000, channels=1, offset_ns=0) + instance.advance(10_000_000) + if instance._task is not None: + await instance._task + assert not segments + # No contiguous packet can arrive once the watermark passes its start. + instance.advance(10_100_000) + await instance._task + assert len(segments) == 1 + samples = array("h") + samples.frombytes(segments[0].data) + assert samples == array("h", [100, 0] * 480) + await instance.finish() + assert not omitted + + +@pytest.mark.asyncio +async def test_unavailable_audio_is_clamped_to_duration_limit(): + instance, segments, omitted = recorder(max_duration=0.01) + instance.unavailable(0, 20_000_000, "unavailable") + instance.unavailable(20_000_000, 30_000_000, "unavailable") + result = await instance.finish() + assert not segments + assert len(omitted) == 1 + assert omitted[0][0].sample_count == 480 + assert result.reason == "duration_limit" + + @pytest.mark.asyncio async def test_sparse_timeline_downmix_and_channel_alignment(): instance, segments, omitted = recorder(duration=0.1) From a8fab87e536f3872bbca2b107aa13f200347b0f2 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 14:09:44 -0500 Subject: [PATCH 11/12] make sure ordering works --- py/src/braintrust/_audio/recorder.py | 5 ++++- py/src/braintrust/_audio/test_audio.py | 30 ++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/_audio/recorder.py b/py/src/braintrust/_audio/recorder.py index 17f108d6..44d80e39 100644 --- a/py/src/braintrust/_audio/recorder.py +++ b/py/src/braintrust/_audio/recorder.py @@ -3,6 +3,7 @@ import math import sys from array import array +from bisect import insort_right from collections.abc import Awaitable, Callable from dataclasses import dataclass from fractions import Fraction @@ -157,7 +158,9 @@ def record( if self._limit is not None and start >= self._limit: return False packet = Packet(source_id, bytes(pcm), sample_rate_hz, channels, start) - self._packets.append(packet) + # Rendering and flush eligibility both need each source's lookahead in + # timeline order, even when packets arrive out of order. + insort_right(self._packets, packet, key=lambda packet: packet.start) self._end = max(self._end, math.ceil(packet.end)) if self._limit is not None: self._end = min(self._end, self._limit) diff --git a/py/src/braintrust/_audio/test_audio.py b/py/src/braintrust/_audio/test_audio.py index 88fc5917..76b3c030 100644 --- a/py/src/braintrust/_audio/test_audio.py +++ b/py/src/braintrust/_audio/test_audio.py @@ -7,6 +7,7 @@ from array import array from dataclasses import replace from fractions import Fraction +from itertools import permutations from types import MappingProxyType import pytest @@ -114,6 +115,35 @@ async def record(progressive): assert await record(True) == await record(False) +@pytest.mark.asyncio +@pytest.mark.parametrize("order", list(permutations(range(3)))) +@pytest.mark.parametrize("progressive", [False, True]) +async def test_resampling_is_independent_of_packet_arrival_order(order, progressive): + async def record(indices): + instance, segments, omitted = recorder() + for index in indices: + assert instance.record( + "input", + pcm([index * 10000] * 240), + sample_rate_hz=24000, + channels=1, + offset_ns=index * 10_000_000, + ) + if progressive: + instance.advance(20_000_000) + await instance._task + assert len(segments) == 2 # Both boundaries have contiguous lookahead. + await instance.finish() + assert not omitted + return b"".join(s.data for s in segments) + + expected = await record(range(3)) + samples = array("h") + samples.frombytes(expected) + assert samples[959 * 2] == 15000 + assert await record(order) == expected + + @pytest.mark.asyncio async def test_resampling_tail_released_when_watermark_passes_packet_end(): instance, segments, omitted = recorder() From b6ec925afa6c9c0f8cfc4681343877d020b38b05 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 6 Nov 2026 15:05:04 -0500 Subject: [PATCH 12/12] fix test --- py/src/braintrust/_audio/test_session.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/py/src/braintrust/_audio/test_session.py b/py/src/braintrust/_audio/test_session.py index 0eedb1af..2a1f1130 100644 --- a/py/src/braintrust/_audio/test_session.py +++ b/py/src/braintrust/_audio/test_session.py @@ -10,6 +10,8 @@ @pytest.fixture def memory(monkeypatch): + # init_test_logger replaces this global; restore it after each test. + monkeypatch.setattr(logger, "_compute_logger_metadata", logger._compute_logger_metadata) init_test_logger("audio-tests") with logger._internal_with_memory_background_logger() as bgl: monkeypatch.setattr(logger._state, "_global_bg_logger", logger.LazyValue(lambda: bgl, use_mutex=False))