Add --save-audio flag to record streamed PCM to WAV (#186)

Adds a `--save-audio` flag to `assembly stream` that tees the streamed
PCM audio to a WAV file while transcribing, enabling callers to keep a
verbatim recording of the audio sent to the streaming API without owning
capture themselves.

## Key changes

- **New `aai_cli/streaming/record.py` module**: Implements `tee_wav()`
generator that yields PCM chunks unchanged while writing them to disk as
a 16-bit mono WAV, and `validate_target()` to check the output path's
parent directory exists before streaming starts.

- **`--save-audio` CLI option**: Added to
`aai_cli/commands/stream/__init__.py` with validation that the path is
writable and the parent directory exists.

- **Validation & mutual exclusivity**: 
- Rejects `--save-audio` with `--from-stdin` (can't tee multiple sources
to one file)
- Rejects `--save-audio` with `--system-audio` (mic and system streams
can't share one file)
- Rejects `--save-audio` with `--show-code` (generated SDK code doesn't
tee audio)
- Validates the output path's parent directory exists before
credentials/audio are opened, so path errors surface cleanly upfront

- **Integration into streaming session**: `StreamSession.stream_one()`
wraps the audio iterable with `tee_wav()` when `save_audio` is set,
writing at the source's true sample rate.

- **Comprehensive test coverage**: 
- Unit tests for `tee_wav()` behavior (unchanged passthrough, valid WAV
output, early close handling, empty streams, unopenable paths)
- Integration tests for the flag (tees streamed PCM correctly, not
written when unset, rejects incompatible flag combinations, validates
parent directory)

## Implementation details

The tee is implemented as a generator that writes each chunk to disk and
yields it onward unchanged, ensuring the API receives exactly what was
captured. The WAV header's length fields are patched on generator close
(including via `GeneratorExit` on Ctrl-C), so even interrupted runs
leave valid, playable WAV files. The recording is always mono 16-bit PCM
at the source's sample rate, matching the shape the streaming API
receives.

https://claude.ai/code/session_01MiPAW6mr1pYQuAE123HGxD

Co-authored-by: Claude <noreply@anthropic.com>

Author: alexkroman

aai_cli/commands/stream/__init__.py

@@ -33,6 +33,7 @@
             ("Stream a list of files in turn", "ls *.wav | assembly stream --from-stdin"),
             ("Stream the hosted sample", "assembly stream --sample"),
             ("Label speakers in the live transcript", "assembly stream --speaker-labels"),
+            ("Save a WAV of the audio while streaming", "assembly stream --save-audio out.wav"),
             (
                 "Boost domain terms with keyterm prompts",
                 'assembly stream --keyterms-prompt "AssemblyAI" --keyterms-prompt "Claude"',
@@ -82,6 +83,16 @@ def stream(
         help="macOS only: stream system/app audio without the microphone",
         rich_help_panel=help_panels.OPT_CAPTURE,
     ),
+    save_audio: Path | None = typer.Option(
+        None,
+        "--save-audio",
+        help="Tee the streamed PCM to PATH as a 16-bit mono WAV while transcribing",
+        rich_help_panel=help_panels.OPT_CAPTURE,
+        dir_okay=False,
+        # Click guardrail; flipping it changes no behavior a unit test can observe
+        # (and the writable check is a no-op under the test runner's root uid).
+        writable=True,  # pragma: no mutate
+    ),
     # model & input
     speech_model: SpeechModel = typer.Option(
         DEFAULT_SPEECH_MODEL,
@@ -355,5 +366,6 @@ def stream(
         config_file=config_file,
         output_field=output_field,
         show_code=show_code,
+        save_audio=save_audio,
     )
     run_with_options(ctx, stream_exec.run_stream, opts, json=json_out)

aai_cli/commands/stream/_exec.py

@@ -22,7 +22,7 @@
 from aai_cli.core import choices, client, config_builder, stdio, youtube
 from aai_cli.core.errors import UsageError, mutually_exclusive
 from aai_cli.core.microphone import MicrophoneSource
-from aai_cli.streaming import turn_presets
+from aai_cli.streaming import record, turn_presets
 from aai_cli.streaming.macos import MacSystemAudioSource
 from aai_cli.streaming.render import StreamRenderer
 from aai_cli.streaming.session import (
@@ -85,6 +85,7 @@ class StreamOptions:
     config_file: Path | None
     output_field: choices.TextOrJson | None
     show_code: bool
+    save_audio: Path | None
 
     def source_options(self) -> SourceOptions:
         """The audio-input subset, in the shape the validation/dispatch helpers read."""
@@ -245,6 +246,11 @@ def _collect_batch_sources(opts: StreamOptions, *, text_mode: bool) -> list[str]
         ("--show-code", opts.show_code),
         suggestion="--show-code renders one source; pass a single file or URL.",
     )
+    mutually_exclusive(
+        ("--from-stdin", True),
+        ("--save-audio", opts.save_audio is not None),
+        suggestion="--save-audio tees one stream; run a single source to record it.",
+    )
     mutually_exclusive(
         ("--llm", bool(opts.llm_prompt)),
         ("-o text", text_mode),
@@ -305,12 +311,25 @@ def run_stream(opts: StreamOptions, state: AppState, *, json_mode: bool) -> None
     base_flags = opts.base_flags()
 
     if opts.show_code:
+        if opts.save_audio is not None:
+            raise UsageError(
+                "--save-audio cannot be combined with --show-code; the generated SDK "
+                "code does not tee audio to disk."
+            )
         _print_show_code(opts, sources, base_flags, text_mode=text_mode)
         return
 
     # Validate the requested sources (including that a local file exists) before
     # credentials, so a typo'd path reads as "file not found" — not as a login.
     validate_sources(sources, has_llm=bool(opts.llm_prompt), text_mode=text_mode)
+    if opts.save_audio is not None:
+        if sources.from_system_audio:
+            raise UsageError(
+                "--save-audio cannot be combined with --system-audio; the mic and system "
+                "streams can't share one file.",
+                suggestion="Record a single source (mic, file, URL, or - on stdin).",
+            )
+        record.validate_target(opts.save_audio)
     if sources.from_file and not sources.from_stdin:
         client.resolve_audio_source(sources.source, sample=sources.sample)
     api_key = state.resolve_api_key()
@@ -326,6 +345,7 @@ def run_stream(opts: StreamOptions, state: AppState, *, json_mode: bool) -> None
         llm_prompts=llm_prompts,
         model=opts.model,
         max_tokens=opts.max_tokens,
+        save_audio=opts.save_audio,
         llm_interval=opts.llm_interval,
     )
     _dispatch(session, sources)

aai_cli/streaming/record.py

@@ -0,0 +1,67 @@
+"""Tee streamed PCM to a WAV file — backs `assembly stream --save-audio PATH`.
+
+The whole point is a verbatim recording of exactly the bytes sent to the streaming
+API, so a caller (e.g. an ensemble that compares the live turns against an async
+re-transcribe) can keep the audio without owning capture itself. The tee never alters
+what's transcribed: it writes each chunk to disk and yields it onward unchanged.
+"""
+
+from __future__ import annotations
+
+import wave
+from collections.abc import Generator, Iterable
+from pathlib import Path
+
+from aai_cli.core.errors import CLIError
+from aai_cli.streaming.sources import PCM16_SAMPLE_WIDTH_BYTES
+
+
+def validate_target(path: Path) -> None:
+    """Reject a ``--save-audio`` path whose parent directory is missing, before streaming.
+
+    Run before credentials/audio are opened so a bad path reads as a path error up
+    front, not after a session has already started recording into the void.
+    """
+    parent = path.parent
+    if not parent.is_dir():
+        raise CLIError(
+            f"Cannot save audio to {path}: {parent} is not a directory.",
+            error_type="save_audio_path",
+            exit_code=2,
+            suggestion="Create the directory first, or pass a path under an existing one.",
+        )
+
+
+def tee_wav(audio: Iterable[bytes], path: Path, *, rate: int) -> Generator[bytes, None, None]:
+    """Yield every PCM16 chunk from ``audio`` unchanged while writing it to ``path`` as WAV.
+
+    The recording is mono 16-bit PCM at ``rate`` — the same shape the streaming API
+    receives. The header's length fields are patched when the iterable is exhausted or
+    closed early (Ctrl-C raises ``GeneratorExit`` at the ``yield``), so even an
+    interrupted run leaves a valid, playable WAV of the audio captured so far.
+    """
+    try:
+        # Open the handle ourselves (rather than letting wave.open(str) do it): a bad
+        # path then fails here cleanly, with no half-built Wave_write whose __del__ would
+        # later raise an "ignored in __del__" warning during GC.
+        handle = path.open("wb")
+    except OSError as exc:
+        raise CLIError(
+            f"Cannot open {path} for writing: {exc}",
+            error_type="save_audio_path",
+            exit_code=2,
+        ) from exc
+    try:
+        # The Wave_write context manager closes (flushes + patches the length fields from
+        # what was actually written) on exit, so the file is a valid WAV even when the
+        # generator is closed mid-stream (Ctrl-C). The outer finally then closes the
+        # handle we opened — after the patch — since wave only closes handles it opened.
+        with wave.open(handle, "wb") as wav:
+            wav.setnchannels(1)
+            wav.setsampwidth(PCM16_SAMPLE_WIDTH_BYTES)
+            wav.setframerate(rate)
+            for chunk in audio:
+                wav.writeframesraw(chunk)
+                yield chunk
+    finally:
+        handle.close()

aai_cli/streaming/session.py

@@ -17,6 +17,7 @@
     UsageError,
     mutually_exclusive,
 )
+from aai_cli.streaming import record
 from aai_cli.streaming.render import StreamRenderer, speaker_prefix
 from aai_cli.ui import output
 from aai_cli.ui.follow import FollowRenderer
@@ -137,6 +138,9 @@ class StreamSession:
     llm_prompts: list[str]
     model: str
     max_tokens: int
+    # When set, tee the streamed PCM to this path as a WAV (see record.tee_wav). Only
+    # the single-source path sets it — the parallel/batch callers reject --save-audio.
+    save_audio: Path | None = None
     # Seconds between --llm summary refreshes; <=0 re-runs the chain on every turn.
     llm_interval: float = 0.0
     # Monotonic clock, injectable so the interval throttle is deterministic in tests.
@@ -242,6 +246,9 @@ def _maybe_summarize(self, *, final: bool = False) -> None:
     def stream_one(
         self, audio: Iterable[bytes], rate: int, *, source_label: str | None = None
     ) -> None:
+        if self.save_audio is not None:
+            # Tee verbatim to disk at the source's true rate before it hits the wire.
+            audio = record.tee_wav(audio, self.save_audio, rate=rate)
         flags = self.base_flags | {"sample_rate": rate}
         if source_label == "you":
             # The microphone captures you alone, so never diarize it into separate

tests/__snapshots__/test_snapshots_help_run.ambr

@@ -731,6 +731,9 @@
   │ --system-audio-only                              macOS only: stream          │
   │                                                  system/app audio without    │
   │                                                  the microphone              │
+  │ --save-audio               FILE                  Tee the streamed PCM to     │
+  │                                                  PATH as a 16-bit mono WAV   │
+  │                                                  while transcribing          │
   ╰──────────────────────────────────────────────────────────────────────────────╯
   ╭─ Model & Language ───────────────────────────────────────────────────────────╮
   │ --speech-model              [universal-streaming-m  Streaming speech model   │
@@ -832,6 +835,8 @@
    $ assembly stream --sample
    Label speakers in the live transcript
    $ assembly stream --speaker-labels
+   Save a WAV of the audio while streaming
+   $ assembly stream --save-audio out.wav
    Boost domain terms with keyterm prompts
    $ assembly stream --keyterms-prompt "AssemblyAI" --keyterms-prompt "Claude"
    Summarize action items live as you talk

tests/test_stream_exec.py

@@ -9,14 +9,16 @@
 from __future__ import annotations
 
 import dataclasses
+import wave
+from pathlib import Path
 
 import pytest
 
 from aai_cli.app.context import AppState
 from aai_cli.commands.stream import DEFAULT_SPEECH_MODEL
 from aai_cli.commands.stream import _exec as stream_exec
 from aai_cli.core import config, llm
-from aai_cli.core.errors import UsageError
+from aai_cli.core.errors import CLIError, UsageError
 from aai_cli.streaming.turn_presets import TurnDetectionPreset
 
 # The CLI's flag defaults, as data. Tests override per-case with dataclasses.replace.
@@ -60,6 +62,7 @@
     config_file=None,
     output_field=None,
     show_code=False,
+    save_audio=None,
 )
 
 
@@ -170,6 +173,7 @@ def test_stream_options_are_immutable():
         {"from_stdin": True, "device": 2},  # mic-only capture flags
         {"from_stdin": True, "sample_rate": 44100},
         {"from_stdin": True, "show_code": True},  # renders one source
+        {"from_stdin": True, "save_audio": Path("out.wav")},  # tees one stream
     ],
 )
 def test_from_stdin_rejects_incompatible_flags(overrides):
@@ -222,3 +226,84 @@ def fake_stream_batch(sources, *, make_session, open_source, renderer, json_mode
         dataclasses.replace(DEFAULTS, from_stdin=True), AppState(), json_mode=True
     )
     assert seen["sources"] == ["a.wav", "b.wav"]
+
+
+# --- --save-audio (tee the streamed PCM to a WAV) --------------------------
+class RecordingMic(FakeMic):
+    """A mic that yields known PCM so the tee'd WAV's contents can be asserted."""
+
+    PCM = b"\x01\x02\x03\x04\x05\x06\x07\x08"
+
+    def __iter__(self):
+        return iter([self.PCM])
+
+
+def test_save_audio_tees_streamed_pcm_to_a_wav(monkeypatch, tmp_path):
+    # The bytes the streaming API receives are also written to --save-audio, verbatim,
+    # as a 16-bit mono WAV at the source's sample rate.
+    config.set_api_key("default", "sk_live")
+    out = tmp_path / "rec.wav"
+
+    def fake_stream_audio(api_key, source, *, params, **_kwargs):
+        # Draining the iterable is what drives the tee — mirror the real SDK consuming it.
+        sent = b"".join(source)
+        assert sent == RecordingMic.PCM  # the API still sees the unaltered audio
+
+    monkeypatch.setattr(stream_exec.client, "stream_audio", fake_stream_audio)
+    monkeypatch.setattr(stream_exec, "MicrophoneSource", RecordingMic)
+
+    stream_exec.run_stream(
+        dataclasses.replace(DEFAULTS, save_audio=out), AppState(), json_mode=True
+    )
+
+    assert out.is_file()
+    with wave.open(str(out), "rb") as w:
+        assert w.getnchannels() == 1
+        assert w.getsampwidth() == 2
+        assert w.getframerate() == 16000  # FakeMic's reported rate
+        assert w.readframes(w.getnframes()) == RecordingMic.PCM
+
+
+def test_save_audio_not_written_when_flag_unset(monkeypatch, tmp_path):
+    # Without --save-audio, the default run leaves no stray WAV behind (kills a mutant
+    # that tees unconditionally).
+    config.set_api_key("default", "sk_live")
+    monkeypatch.setattr(stream_exec.client, "stream_audio", lambda *a, **k: b"".join(a[1]))
+    monkeypatch.setattr(stream_exec, "MicrophoneSource", RecordingMic)
+
+    stream_exec.run_stream(DEFAULTS, AppState(), json_mode=True)
+
+    assert list(tmp_path.glob("*.wav")) == []
+
+
+def test_save_audio_rejects_system_audio():
+    # The mic + system streams can't share one file, so the combo is a usage error
+    # (raised before credentials).
+    with pytest.raises(UsageError):
+        stream_exec.run_stream(
+            dataclasses.replace(DEFAULTS, save_audio=Path("rec.wav"), system_audio=True),
+            AppState(),
+            json_mode=False,
+        )
+
+
+def test_save_audio_rejects_show_code():
+    # --show-code emits SDK code that doesn't tee audio, so the combo is rejected.
+    with pytest.raises(UsageError):
+        stream_exec.run_stream(
+            dataclasses.replace(DEFAULTS, save_audio=Path("rec.wav"), show_code=True),
+            AppState(),
+            json_mode=False,
+        )
+
+
+def test_save_audio_rejects_missing_parent_dir(tmp_path):
+    # A path under a directory that doesn't exist is a clean path error, before auth.
+    config.set_api_key("default", "sk_live")
+    with pytest.raises(CLIError) as excinfo:
+        stream_exec.run_stream(
+            dataclasses.replace(DEFAULTS, save_audio=tmp_path / "nope" / "rec.wav"),
+            AppState(),
+            json_mode=False,
+        )
+    assert excinfo.value.error_type == "save_audio_path"

tests/test_streaming_record.py

@@ -0,0 +1,67 @@
+"""Unit tests for aai_cli.streaming.record — the --save-audio WAV tee."""
+
+from __future__ import annotations
+
+import wave
+
+import pytest
+
+from aai_cli.core.errors import CLIError
+from aai_cli.streaming import record
+
+
+def _read_wav(path):
+    with wave.open(str(path), "rb") as w:
+        return w.getnchannels(), w.getsampwidth(), w.getframerate(), w.readframes(w.getnframes())
+
+
+def test_tee_wav_yields_chunks_unchanged(tmp_path):
+    chunks = [b"\x01\x02", b"\x03\x04\x05\x06"]
+    out = list(record.tee_wav(iter(chunks), tmp_path / "a.wav", rate=16000))
+    assert out == chunks  # the tee must not alter what's streamed onward
+
+
+def test_tee_wav_writes_a_valid_wav_with_the_source_rate(tmp_path):
+    path = tmp_path / "a.wav"
+    list(record.tee_wav(iter([b"\x01\x02", b"\x03\x04"]), path, rate=44100))
+    channels, width, rate, frames = _read_wav(path)
+    assert channels == 1
+    assert width == 2
+    assert rate == 44100  # the declared source rate, not a hardcoded default
+    assert frames == b"\x01\x02\x03\x04"
+
+
+def test_tee_wav_finalizes_a_valid_wav_on_early_close(tmp_path):
+    # Ctrl-C closes the generator mid-stream; the partial file must still be valid WAV.
+    path = tmp_path / "a.wav"
+    gen = record.tee_wav(iter([b"\x01\x02", b"\x03\x04"]), path, rate=16000)
+    assert next(gen) == b"\x01\x02"  # consume only the first chunk
+    gen.close()  # raises GeneratorExit at the yield -> finally closes the WAV
+    _channels, _width, _rate, frames = _read_wav(path)
+    assert frames == b"\x01\x02"  # only the consumed chunk landed
+
+
+def test_tee_wav_empty_stream_writes_a_zero_length_wav(tmp_path):
+    path = tmp_path / "a.wav"
+    assert list(record.tee_wav(iter([]), path, rate=16000)) == []
+    _channels, _width, _rate, frames = _read_wav(path)
+    assert frames == b""
+
+
+def test_tee_wav_unopenable_path_is_a_clean_error(tmp_path):
+    # Pointing at a directory can't be opened for writing -> a CLIError, not a raw OSError.
+    with pytest.raises(CLIError) as excinfo:
+        # tee_wav opens lazily on first iteration, so the generator must be started.
+        next(record.tee_wav(iter([b"\x01\x02"]), tmp_path, rate=16000))
+    assert excinfo.value.error_type == "save_audio_path"
+
+
+def test_validate_target_accepts_an_existing_directory(tmp_path):
+    record.validate_target(tmp_path / "rec.wav")  # parent exists -> no raise
+
+
+def test_validate_target_rejects_a_missing_parent_directory(tmp_path):
+    with pytest.raises(CLIError) as excinfo:
+        record.validate_target(tmp_path / "nope" / "rec.wav")
+    assert excinfo.value.error_type == "save_audio_path"
+    assert excinfo.value.exit_code == 2