"""Multi-source streaming-chat VCR cassette. This module records and replays a single ``chat.ask`` call against a fresh scratch notebook that carries **five text sources**, with all five source IDs passed explicitly to ``ask``. The point is to lock in the on-wire shape of the **nested-list source-id encoder** for the chat endpoint — the audit flagged this path as fragile because the encoding has two distinct nesting depths (``[[id]]`` per source, then collected into an outer list) and the per-method test coverage prior to this cassette only exercised the zero-source and one-source cases via the default-sources branch (``ask`` calls ``get_source_ids`` itself). What this guards ---------------- ``_chat._build_chat_request`` builds the params with:: sources_array = nest_source_ids(source_ids, 2) params = [sources_array, question, conversation_history, ...] ``nest_source_ids(ids, 2)`` wraps each id with two inner lists, giving ``[[[id1]], [[id2]], [[id3]], [[id4]], [[id5]]]``. A regression that silently flattens to depth 1 (``[[id1], [id2], ...]``) or depth 3 (``[[[[id1]]], ...]``) would still type-check and still POST, but the server would return an error or — worse — silently drop sources. The cassette captures the recorded ``f.req`` body so the unit-level ``test_chat_request_carries_five_source_ids_nested`` assertion has a real on-wire reference to compare against. Recording --------- :: NOTEBOOKLM_VCR_RECORD=1 uv run pytest \\ tests/integration/test_chat_multi_source_vcr.py -v -s The recording flow uses **its own scratch notebook**, NOT the shared ``NOTEBOOKLM_GENERATION_NOTEBOOK_ID``, so parallel cassette-recording agents don't trip over one another. The scratch notebook and its sources are deleted at the end of the recording run; only the chat ``ask`` interaction is captured inside the cassette context. Replay ------ Replay opts into the ``freq`` body matcher on top of the default ``method/scheme/host/port/path`` matchers. The streaming-chat endpoint at ``/GenerateFreeFormStreamed`` does not carry ``rpcids``, so without ``freq`` two chat POSTs would be indistinguishable. ``freq`` decodes ``f.req`` and disambiguates by param count + ``notebook_id`` at slot 7, which is exactly what we need: the cassette's recorded ``notebook_id`` (a now-deleted scratch UUID) is read out of the cassette and passed back into ``ask`` so the replay matches the recording byte-for-byte at the matcher's chosen slots. """ from __future__ import annotations import json import re import sys import uuid from pathlib import Path from typing import Any from urllib.parse import parse_qs import pytest import yaml from notebooklm import NotebookLMClient from tests.integration.conftest import _vcr_record_mode, get_vcr_auth, skip_no_cassettes from tests.vcr_config import notebooklm_vcr pytestmark = [pytest.mark.vcr, skip_no_cassettes] CASSETTE_NAME = "chat_ask_multi_source.yaml" CASSETTE_PATH = Path(__file__).parent.parent / "cassettes" / CASSETTE_NAME # Number of sources recorded into the cassette. Five is the smallest value # that comfortably exceeds the "two or three sources" gut-check most chat # call sites use during dev and still keeps the cassette small (text # sources, one short paragraph each). The ``acceptance`` clause in the # plan requires 5+; we record exactly 5. SOURCE_COUNT = 5 # Question text. Kept short, generic, and PII-free — the question text # is NOT scrubbed by ``cassette_patterns.py``, so we deliberately avoid # anything that could resemble an account-linked phrase. QUESTION = "Summarize all sources in one short paragraph." # ============================================================================= # Cassette parsing helpers (used by both recording-time and replay tests) # ============================================================================= def _find_chat_interaction(cassette: dict[str, Any]) -> dict[str, Any]: """Locate the single ``GenerateFreeFormStreamed`` interaction inside the cassette. The cassette context wraps exactly one chat POST; the helper raises a clear AssertionError if a future recording accidentally captures more or fewer interactions. """ chat_interactions = [ interaction for interaction in cassette.get("interactions", []) if "GenerateFreeFormStreamed" in interaction.get("request", {}).get("uri", "") ] assert len(chat_interactions) == 1, ( f"expected exactly one GenerateFreeFormStreamed interaction in {CASSETTE_NAME}, " f"found {len(chat_interactions)}" ) return chat_interactions[0] def _decode_freq_params(body: str | bytes) -> list[Any]: """Decode the form-encoded ``f.req`` body into its param list. Mirrors the logic in ``vcr_config._freq_body_matcher._extract_freq``; duplicated here so the replay test does not depend on vcr_config's private helper. """ if isinstance(body, bytes): body = body.decode("utf-8") qs = parse_qs(body) f_req_values = qs.get("f.req", []) assert f_req_values, f"f.req not found in body: {body[:200]!r}" outer = json.loads(f_req_values[0]) assert isinstance(outer, list) and len(outer) >= 2, "f.req envelope malformed" inner = outer[1] assert isinstance(inner, str), "f.req inner JSON missing" params = json.loads(inner) assert isinstance(params, list), "f.req params not a list" return params def _load_cassette_params() -> tuple[list[Any], str]: """Return ``(decoded_params, notebook_id)`` from the recorded chat POST. Reads the cassette from disk so the replay test stays decoupled from any state the recording test left behind. """ assert CASSETTE_PATH.exists(), ( f"cassette missing: {CASSETTE_PATH}. " "Re-record with NOTEBOOKLM_VCR_RECORD=1 — see module docstring." ) with CASSETTE_PATH.open(encoding="utf-8") as fh: cassette = yaml.safe_load(fh) interaction = _find_chat_interaction(cassette) body = interaction["request"]["body"] params = _decode_freq_params(body) # Bounds-guard ``params[7]`` so a malformed/short cassette fails with a # message that identifies the helper rather than a bare ``IndexError``. # The caller's ``len(params) == 9`` assertion would otherwise fire too # late (after this helper has already crashed). assert len(params) > 7, ( f"cassette params too short: expected at least 8 slots, got {len(params)}. " f"Re-record with NOTEBOOKLM_VCR_RECORD=1." ) notebook_id = params[7] assert isinstance(notebook_id, str), f"slot 7 (notebook_id) not a string: {notebook_id!r}" return params, notebook_id def _extract_source_ids_from_nested(sources_array: list[Any]) -> list[str]: """Flatten the depth-2 nested ``[[[id]], ...]`` structure back to ids. ``nest_source_ids(ids, 2)`` wraps each id into ``[[id]]`` and collects them into an outer list. This helper undoes that exact shape and raises a clear AssertionError if any slot deviates — making the replay test fail loudly on an encoder regression rather than silently returning an empty list. """ extracted: list[str] = [] for index, slot in enumerate(sources_array): assert isinstance(slot, list) and len(slot) == 1, ( f"sources_array slot {index} is not a single-element list: {slot!r}" ) inner = slot[0] assert isinstance(inner, list) and len(inner) == 1, ( f"sources_array slot {index} inner is not a single-element list: {inner!r}" ) source_id = inner[0] assert isinstance(source_id, str) and source_id, ( f"sources_array slot {index} id is not a non-empty string: {source_id!r}" ) extracted.append(source_id) return extracted # ============================================================================= # Recording-time scratch notebook setup # ============================================================================= # # Setup and teardown live OUTSIDE the cassette context. The cassette # context wraps only the ``ask`` call, so the cassette captures exactly # one streaming-chat POST and nothing else — no notebook-create, no # source-add, no notebook-delete clutter to mismatch on replay. async def _seed_scratch_notebook( client: NotebookLMClient, ) -> tuple[str, list[str]]: """Create a fresh notebook with ``SOURCE_COUNT`` text sources. Returns ``(notebook_id, source_ids)``. Each source carries a short, PII-free paragraph; the title embeds a UUID so a re-run never collides with a leftover source from a prior aborted recording (the UUID also rules out title-based dedupe heuristics). """ title = f"T8.E6 scratch multi-source chat ({uuid.uuid4()})" notebook = await client.notebooks.create(title) notebook_id = notebook.id # Five short, generic paragraphs. Distinct enough that the answer # has something to summarize, but boring enough that no scrubbing # corner case fires on the recorded request bodies. contents = [ ( "Bicycles are human-powered, pedal-driven vehicles with two wheels " "attached to a frame. They are widely used for transport, recreation, " "and exercise across many countries." ), ( "Lighthouses are towers built near coastlines to emit light and warn " "ships of dangerous shoals. Their construction varies with terrain " "but the navigational role is consistent worldwide." ), ( "Tea is a beverage prepared by pouring hot water over cured leaves " "of the Camellia sinensis plant. Different processing methods yield " "green, black, oolong, and white tea varieties." ), ( "Sourdough bread is leavened by naturally occurring yeasts and " "lactobacilli in a flour-and-water starter. The long fermentation " "produces the characteristic tangy flavor and chewy crumb." ), ( "Origami is the Japanese art of paper folding. Starting from a " "single sheet, a sequence of folds produces shapes ranging from " "simple cranes to complex modular geometric forms." ), ] assert len(contents) == SOURCE_COUNT, "contents and SOURCE_COUNT must agree" source_ids: list[str] = [] for index, content in enumerate(contents, start=1): source = await client.sources.add_text( notebook_id, title=f"Source {index} ({uuid.uuid4()})", content=content, ) source_ids.append(source.id) # Wait until all sources are ready so the chat call actually has # something to retrieve. The text-source ingestion is normally fast # but not instant; ``wait_for_sources`` polls server-side state. await client.sources.wait_for_sources(notebook_id, source_ids, timeout=120.0) return notebook_id, source_ids async def _teardown_scratch_notebook(client: NotebookLMClient, notebook_id: str) -> None: """Delete the scratch notebook so the recording leaves no residue. Failures here are logged but not raised — a leaked scratch notebook is a maintainer hygiene issue, not a correctness failure for the cassette we just recorded. The maintainer can prune leftovers manually via ``notebooklm list``. """ try: await client.notebooks.delete(notebook_id) except Exception as exc: # noqa: BLE001 — best-effort cleanup print( f"WARNING: failed to delete scratch notebook {notebook_id}: {exc}", file=sys.stderr, ) # ============================================================================= # Replay test — runs from the recorded cassette under normal CI # ============================================================================= class TestChatMultiSource: """Multi-source streaming-chat ask: recording + replay.""" @pytest.mark.vcr @pytest.mark.asyncio async def test_ask_with_five_explicit_sources(self) -> None: """A chat ask against five explicit sources round-trips successfully. Recording-mode behavior: create scratch notebook + five sources outside the cassette context, enter the cassette context for exactly the ``ask`` call, then exit and delete the notebook. Replay-mode behavior: read the recorded ``notebook_id`` and ``source_ids`` out of the cassette, pass them back into ``ask`` so the request body matches the cassette's recorded payload at every slot the ``freq`` matcher inspects (param count + slot 7 notebook_id), then assert the recorded answer is non-empty. """ auth = await get_vcr_auth() async with NotebookLMClient(auth) as client: if _vcr_record_mode: notebook_id, source_ids = await _seed_scratch_notebook(client) try: with notebooklm_vcr.use_cassette( CASSETTE_NAME, # Opt-in to the ``freq`` body matcher. The # streaming-chat POST has no ``rpcids`` query # param so the body-aware matcher is the only way # to disambiguate it from any future chat # cassette that hits the same path. match_on=[ "method", "scheme", "host", "port", "path", "freq", ], ): result = await client.chat.ask( notebook_id, QUESTION, source_ids=source_ids, ) finally: await _teardown_scratch_notebook(client, notebook_id) else: # Replay path: discover the recorded inputs from the # cassette so the request the client builds matches the # request the cassette holds. Single helper call — # ``params`` carries the sources_array at slot 0 and # ``recorded_notebook_id`` is what slot 7 already # carries (validated by the helper). params, recorded_notebook_id = _load_cassette_params() recorded_source_ids = _extract_source_ids_from_nested(params[0]) with notebooklm_vcr.use_cassette( CASSETTE_NAME, match_on=[ "method", "scheme", "host", "port", "path", "freq", ], ): result = await client.chat.ask( recorded_notebook_id, QUESTION, source_ids=recorded_source_ids, ) # The streaming-chat endpoint always returns a non-empty answer # for a well-formed multi-source ask. Empty would indicate a # server-side error path leaked through ``AskResult``. assert result is not None assert result.answer is not None assert isinstance(result.answer, str) assert result.answer.strip(), "chat answer should be non-empty" assert result.conversation_id is not None def test_cassette_carries_all_five_source_ids_with_correct_nesting( self, ) -> None: """The recorded ``f.req`` body has all five source-ids at depth 2. This is the regression guard for the fragile nested-list encoder (audit C-class finding on ``_build_chat_request``). Failure modes the test catches: 1. ``sources_array`` collapsed to depth 1 — slot[0] would NOT be a single-element list-of-list; the ``_extract_source_ids_from_nested`` parser raises. 2. Source list truncated — the recorded length deviates from ``SOURCE_COUNT``. 3. Source IDs lost their string shape — the parser raises because a slot's terminal value is not a non-empty ``str``. The cassette is the source of truth; we parse it directly rather than relying on the live client's behavior so the assertion is independent of the implementation under test. """ params, _notebook_id = _load_cassette_params() # Param count must match the 9-slot shape ``_build_chat_request`` # documents. A drift here is itself a regression worth shouting # about (stale-cassette regression class). assert len(params) == 9, ( f"chat f.req param count drift: expected 9, got {len(params)}. params={params!r}" ) sources_array = params[0] assert isinstance(sources_array, list), ( f"slot 0 (sources_array) must be a list, got {type(sources_array).__name__}" ) assert len(sources_array) == SOURCE_COUNT, ( f"recorded sources_array has {len(sources_array)} entries, expected {SOURCE_COUNT}" ) # Round-trip through the strict parser — raises with a clear # message if any slot deviates from ``[[id]]``. extracted_ids = _extract_source_ids_from_nested(sources_array) assert len(extracted_ids) == SOURCE_COUNT # Source-ids should look like UUIDs (Google-issued source IDs # are UUIDs). The exact value of each id is not asserted — # rerecording yields fresh ids and the test must stay green. uuid_pat = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$") for sid in extracted_ids: assert uuid_pat.match(sid), f"recorded source id is not a UUID: {sid!r}" # Slot positions that ``_build_chat_request`` documents — assert # them so any reshuffle of the 9-param shape fails this test # immediately rather than silently changing wire format. assert params[1] == QUESTION, ( f"slot 1 (question) drift: expected {QUESTION!r}, got {params[1]!r}" ) assert params[3] == [2, None, [1], [1]], f"slot 3 (mode tuple) drift: got {params[3]!r}" assert params[8] == 1, f"slot 8 (always-1) drift: got {params[8]!r}"