/tests/unit/test_chat_ask_invariants.py

"""Regression tests for ``ChatAPI.ask`` after the chat-transport refactor.

These assertions pin down the new contract:

- ``ask`` uses ``self._reqid.next_reqid()`` for the URL ``_reqid`` param (the
  ``_reqid_counter`` property + deprecation gesture were retired in the
  session-shrink arc; this test now guards against any new
  ``DeprecationWarning`` escaping ``_chat/api.py``).
- ``authuser=`` is present on the chat URL when ``account_email`` is set on
  the auth tokens, mirroring the batchexecute path in
  ``RpcExecutor.build_url``. Previously omitted entirely on the chat endpoint.
- Concurrent ``asyncio.gather(ask*3)`` produces three distinct reqid values.
- 401 mid-chat triggers a refresh, and the post-refresh attempt's body
  carries the refreshed CSRF token (snapshot-per-attempt invariant).
- ``NOTEBOOKLM_BL`` env override still works after the move to
  :mod:`notebooklm._env`.
"""

from __future__ import annotations

import asyncio
import json
import re
import warnings
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock
from urllib.parse import parse_qs, unquote, urlparse

import httpx
import pytest

from notebooklm import NotebookLMClient
from notebooklm._chat import ChatAPI
from notebooklm._request_types import AuthSnapshot
from notebooklm._runtime.config import DEFAULT_CHAT_RESPONSE_MAX_BYTES
from notebooklm.auth import AuthTokens
from notebooklm.exceptions import ChatError
from tests._helpers.client_factory import build_client_shell_for_tests
from tests.unit.conftest import install_post_as_stream

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_answer_response_body(
    answer: str = "Refactor answer is long enough.",
    *,
    server_conv_id: str = "server-refactor-conv",
) -> bytes:
    """Build a minimal valid streaming chat response.

    A ``server_conv_id`` is always present at ``first[2][0]`` because
    ``ChatAPI.ask`` requires the server to assign the id for new
    conversations (issue #659); responses lacking one raise ``ChatError``.
    """
    inner_json = json.dumps([[answer, None, [server_conv_id, 12345], None, [1]]])
    chunk_json = json.dumps([["wrb.fr", None, inner_json]])
    return f")]}}'\n{len(chunk_json)}\n{chunk_json}\n".encode()


def _extract_query_param(url: str, key: str) -> str | None:
    qs = parse_qs(urlparse(url).query, keep_blank_values=True)
    values = qs.get(key)
    return values[0] if values else None


# ---------------------------------------------------------------------------
# chat timeout routing
# ---------------------------------------------------------------------------


class TestChatTimeoutRouting:
    def test_client_uses_chat_specific_timeout_by_default(self):
        auth = AuthTokens(cookies={"SID": "x"}, csrf_token="csrf", session_id="sid")
        client = NotebookLMClient(auth, timeout=75.0)

        assert client.chat._chat_timeout == 180.0
        assert client.chat._chat_response_max_bytes == DEFAULT_CHAT_RESPONSE_MAX_BYTES

    def test_client_chat_timeout_none_inherits_transport_timeout(self):
        auth = AuthTokens(cookies={"SID": "x"}, csrf_token="csrf", session_id="sid")
        client = NotebookLMClient(auth, timeout=75.0, chat_timeout=None)

        assert client.chat._chat_timeout is None

    def test_client_chat_response_max_bytes_none_inherits_shared_rpc_cap(self):
        auth = AuthTokens(cookies={"SID": "x"}, csrf_token="csrf", session_id="sid")
        client = NotebookLMClient(auth, timeout=75.0, chat_response_max_bytes=None)

        assert client.chat._chat_response_max_bytes is None

    def test_client_chat_timeout_override_wins(self):
        auth = AuthTokens(cookies={"SID": "x"}, csrf_token="csrf", session_id="sid")
        client = NotebookLMClient(
            auth,
            timeout=75.0,
            chat_timeout=180.0,
            chat_response_max_bytes=123456,
        )

        assert client.chat._chat_timeout == 180.0
        assert client.chat._chat_response_max_bytes == 123456

    @pytest.mark.parametrize("value", [0, -1])
    def test_client_rejects_invalid_chat_response_max_bytes(self, value: int):
        auth = AuthTokens(cookies={"SID": "x"}, csrf_token="csrf", session_id="sid")

        with pytest.raises(ValueError, match="chat_response_max_bytes must be >= 1"):
            NotebookLMClient(auth, chat_response_max_bytes=value)

    @pytest.mark.asyncio
    async def test_ask_passes_chat_read_timeout_response_cap_and_disables_timeout_retry(self):
        """``ask`` uses the chat-specific read window without retrying timed-out streams."""
        transport = SimpleNamespace(
            perform_authed_post=AsyncMock(
                return_value=httpx.Response(
                    200,
                    request=httpx.Request("POST", "https://example.test/chat"),
                    content=_make_answer_response_body(),
                )
            )
        )
        chat = ChatAPI(
            rpc=SimpleNamespace(),
            transport=transport,
            reqid=SimpleNamespace(next_reqid=AsyncMock(return_value=100000)),
            loop_guard=SimpleNamespace(assert_bound_loop=lambda: None),
            chat_timeout=45.0,
            chat_response_max_bytes=987654,
        )

        result = await chat.ask(
            "nb-1",
            "Q?",
            source_ids=["s1"],
            conversation_id="conv-1",
        )

        assert result.answer == "Refactor answer is long enough."
        assert transport.perform_authed_post.await_args.kwargs.get("read_timeout") == 45.0
        assert transport.perform_authed_post.await_args.kwargs.get("max_response_bytes") == 987654
        assert (
            transport.perform_authed_post.await_args.kwargs.get("disable_read_timeout_retries")
            is True
        )


# ---------------------------------------------------------------------------
# authuser= URL parameter
# ---------------------------------------------------------------------------


class TestChatAuthuserParam:
    """``authuser=`` was previously omitted entirely on the chat endpoint."""

    @pytest.mark.asyncio
    async def test_authuser_set_when_account_email_provided(
        self, httpx_mock, mock_get_conversation_id
    ):
        """When ``account_email`` is set on auth, chat URL carries authuser=email."""
        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
            authuser=2,
            account_email="user@example.com",
        )

        httpx_mock.add_response(
            url=re.compile(r".*GenerateFreeFormStreamed.*"),
            content=_make_answer_response_body(),
            method="POST",
        )
        mock_get_conversation_id()  # issue #659 post-ask round-trip

        async with NotebookLMClient(auth) as client:
            await client.chat.ask("nb_x", "Q?", source_ids=["s1"])

        # Filter for the chat-ask request — the post-ask hPTbtc request also
        # lands on the same authuser query, but we want to assert the chat
        # leg specifically.
        request = next(
            r for r in httpx_mock.get_requests() if "GenerateFreeFormStreamed" in str(r.url)
        )
        # Email is preferred over the integer index because it survives
        # browser-account reordering.
        assert _extract_query_param(str(request.url), "authuser") == "user@example.com"

    @pytest.mark.asyncio
    async def test_authuser_set_when_only_authuser_index(
        self, httpx_mock, mock_get_conversation_id
    ):
        """When only ``authuser`` is non-zero (no email), still emit the index."""
        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
            authuser=3,
        )

        httpx_mock.add_response(
            url=re.compile(r".*GenerateFreeFormStreamed.*"),
            content=_make_answer_response_body(),
            method="POST",
        )
        mock_get_conversation_id()

        async with NotebookLMClient(auth) as client:
            await client.chat.ask("nb_x", "Q?", source_ids=["s1"])

        request = next(
            r for r in httpx_mock.get_requests() if "GenerateFreeFormStreamed" in str(r.url)
        )
        assert _extract_query_param(str(request.url), "authuser") == "3"

    @pytest.mark.asyncio
    async def test_authuser_absent_for_default_profile(self, httpx_mock, mock_get_conversation_id):
        """No ``authuser=`` on the URL when authuser=0 and no email — matches the
        previous-contract default-profile behavior (don't churn the existing single-account
        contract)."""
        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
            # authuser defaults to 0, account_email defaults to None
        )

        httpx_mock.add_response(
            url=re.compile(r".*GenerateFreeFormStreamed.*"),
            content=_make_answer_response_body(),
            method="POST",
        )
        mock_get_conversation_id()

        async with NotebookLMClient(auth) as client:
            await client.chat.ask("nb_x", "Q?", source_ids=["s1"])

        request = next(
            r for r in httpx_mock.get_requests() if "GenerateFreeFormStreamed" in str(r.url)
        )
        assert _extract_query_param(str(request.url), "authuser") is None


# ---------------------------------------------------------------------------
# next_reqid + DeprecationWarning silence
# ---------------------------------------------------------------------------


class TestChatReqid:
    """``ChatAPI.ask`` must call ``self._reqid.next_reqid()`` — not poke
    ``_reqid_counter`` directly, which would emit ``DeprecationWarning``."""

    @pytest.mark.asyncio
    async def test_ask_uses_next_reqid_no_deprecation_warning(
        self, httpx_mock, mock_get_conversation_id
    ):
        """No ``DeprecationWarning`` is emitted by ``_chat/api.py`` during ask()."""
        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
        )

        httpx_mock.add_response(
            url=re.compile(r".*GenerateFreeFormStreamed.*"),
            content=_make_answer_response_body(),
            method="POST",
        )
        mock_get_conversation_id()

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            async with NotebookLMClient(auth) as client:
                await client.chat.ask("nb_x", "Q?", source_ids=["s1"])

        chat_dep_warnings = [
            w
            for w in caught
            if issubclass(w.category, DeprecationWarning)
            and "_reqid_counter" in str(w.message)
            and "_chat" in str(w.filename)
            and "api.py" in str(w.filename)
        ]
        assert chat_dep_warnings == [], (
            f"_chat/api.py must not emit _reqid_counter DeprecationWarning; "
            f"got: {[(str(w.filename), str(w.message)) for w in chat_dep_warnings]}"
        )

    @pytest.mark.asyncio
    async def test_concurrent_asks_produce_distinct_reqids(
        self, httpx_mock, mock_get_conversation_id
    ):
        """``asyncio.gather(ask*3)`` → three distinct ``_reqid`` URL values.

        Previously, the body did ``self._core._collaborators.reqid_counter += 100000`` under
        a read-modify-write race (``self._core`` was the pre-Phase-2 attribute
        name, now ``self._runtime``); under concurrent gather() this collapsed
        to a single reqid value. ``runtime.next_reqid()`` serializes the
        increment under an asyncio.Lock, restoring monotonic distinct ids.
        """
        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
        )

        # One response per gathered chat-ask. pytest_httpx replays in order.
        for _ in range(3):
            httpx_mock.add_response(
                url=re.compile(r".*GenerateFreeFormStreamed.*"),
                content=_make_answer_response_body(),
                method="POST",
            )
        # Plus the three post-ask hPTbtc round-trips (issue #659).
        mock_get_conversation_id(reusable=True)

        async with NotebookLMClient(auth) as client:
            await asyncio.gather(
                client.chat.ask("nb_x", "Q1", source_ids=["s1"]),
                client.chat.ask("nb_x", "Q2", source_ids=["s1"]),
                client.chat.ask("nb_x", "Q3", source_ids=["s1"]),
            )

        reqids = [
            _extract_query_param(str(req.url), "_reqid")
            for req in httpx_mock.get_requests()
            if "GenerateFreeFormStreamed" in str(req.url)
        ]
        assert len(reqids) == 3
        assert all(r is not None for r in reqids)
        assert len(set(reqids)) == 3, f"reqids must be distinct, got {reqids}"


# ---------------------------------------------------------------------------
# 401 mid-chat → snapshot regenerated for retry
# ---------------------------------------------------------------------------


class TestChatRefreshRetry:
    """Snapshot-per-attempt invariant: ``build_request`` is invoked once
    per attempt with a *fresh* ``AuthSnapshot``, so the retry body carries
    the post-refresh CSRF token rather than replaying the stale pre-refresh
    body."""

    @pytest.mark.asyncio
    async def test_post_refresh_retry_uses_fresh_csrf_in_body(self, monkeypatch):
        """401 → refresh callback rotates CSRF → retry body contains new token."""
        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="OLD_CSRF",
            session_id="OLD_SID",
        )

        async def refresh() -> AuthTokens:
            # Mutate the live auth tokens — the next snapshot picks this up.
            auth.csrf_token = "NEW_CSRF"
            auth.session_id = "NEW_SID"
            return auth

        core = build_client_shell_for_tests(
            auth=auth, refresh_callback=refresh, refresh_retry_delay=0.0
        )
        await core.__aenter__()
        try:
            observed_bodies: list[str] = []
            call_count = {"n": 0}

            async def fake_post(url, *args, **kwargs):  # type: ignore[no-untyped-def]
                # The post-ask ``hPTbtc`` request from issue #659 also goes
                # through this fake_post. Identify it by URL and return a
                # minimal RPC response that decodes to a valid conv_id.
                if "batchexecute" in str(url):
                    rpc_body = (
                        ")]}'\n"
                        '63\n[["wrb.fr","hPTbtc","[[[\\"real-conv-id-from-hptbtc\\"]]]",null,null]]'
                    )
                    return httpx.Response(
                        200,
                        request=httpx.Request("POST", url),
                        content=rpc_body.encode(),
                    )

                # Chat-ask path: capture the body and exercise the retry contract.
                body = kwargs.get("content")
                if isinstance(body, bytes):
                    body = body.decode()
                observed_bodies.append(body or "")

                call_count["n"] += 1
                if call_count["n"] == 1:
                    # First attempt: 401 → triggers refresh path.
                    response = httpx.Response(401, request=httpx.Request("POST", url), content=b"")
                    raise httpx.HTTPStatusError("401", request=response.request, response=response)
                # Second attempt (after refresh): return a valid answer.
                return httpx.Response(
                    200,
                    request=httpx.Request("POST", url),
                    content=_make_answer_response_body(),
                )

            assert core._collaborators.kernel.http_client is not None
            install_post_as_stream(
                monkeypatch, core._collaborators.kernel.get_http_client(), fake_post
            )

            # Wave 8 of session-decoupling (ADR-0014 Rule 2 Corollary):
            # ``ChatAPI`` takes its four direct collaborators by keyword
            # arg. Wired here from the real ``Session`` under test so the
            # refresh path exercises the production transport/rpc/reqid
            # collaborators end-to-end.
            # Stage B1 PR 2 deleted the Stage A accessors
            # (``Session.session_transport`` / ``Session.collaborators``);
            # read the private slots directly instead.
            api = ChatAPI(
                rpc=core._rpc_executor,
                transport=core._composed.transport,
                reqid=core._collaborators.reqid,
                loop_guard=core._collaborators.lifecycle,
            )
            result = await api.ask("nb_x", "Q?", source_ids=["s1"])

            assert call_count["n"] == 2
            assert "Refactor answer is long enough." in result.answer

            # First attempt body carries OLD_CSRF (pre-refresh snapshot).
            assert "at=OLD_CSRF" in observed_bodies[0]
            assert "at=NEW_CSRF" not in observed_bodies[0]
            # Second attempt body carries NEW_CSRF (post-refresh snapshot)
            # — this is the snapshot-per-attempt contract surfacing
            # through chat_aware_authed_post.
            assert "at=NEW_CSRF" in observed_bodies[1]
            assert "at=OLD_CSRF" not in observed_bodies[1]
        finally:
            await core.close()


# ---------------------------------------------------------------------------
# NOTEBOOKLM_BL override still works after the move to _env.py
# ---------------------------------------------------------------------------


class TestChatBlOverride:
    """Single-source-of-truth for the ``bl`` parameter lives in ``_env.py``.
    The ``NOTEBOOKLM_BL`` override must still flow through to the chat URL.
    """

    @pytest.mark.asyncio
    async def test_custom_bl_env_appears_in_url(
        self, httpx_mock, monkeypatch, mock_get_conversation_id
    ):
        monkeypatch.setenv("NOTEBOOKLM_BL", "boq_labs-custom_99999999.00_p0")

        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
        )

        httpx_mock.add_response(
            url=re.compile(r".*GenerateFreeFormStreamed.*"),
            content=_make_answer_response_body(),
            method="POST",
        )
        mock_get_conversation_id()

        async with NotebookLMClient(auth) as client:
            await client.chat.ask("nb_x", "Q?", source_ids=["s1"])

        request = next(
            r for r in httpx_mock.get_requests() if "GenerateFreeFormStreamed" in str(r.url)
        )
        assert _extract_query_param(str(request.url), "bl") == "boq_labs-custom_99999999.00_p0"

    @pytest.mark.asyncio
    async def test_default_bl_is_pinned_constant(
        self, httpx_mock, monkeypatch, mock_get_conversation_id
    ):
        """With ``NOTEBOOKLM_BL`` unset, the URL falls back to the pinned default.

        The expected literal is duplicated here on purpose: importing
        ``DEFAULT_BL`` from the SUT and asserting equality would be a
        tautology — any wrong-value edit to ``_env.DEFAULT_BL`` would still
        pass. The literal pin catches that.
        """
        monkeypatch.delenv("NOTEBOOKLM_BL", raising=False)

        auth = AuthTokens(
            cookies={"SID": "x"},
            csrf_token="csrf",
            session_id="sid",
        )

        httpx_mock.add_response(
            url=re.compile(r".*GenerateFreeFormStreamed.*"),
            content=_make_answer_response_body(),
            method="POST",
        )
        mock_get_conversation_id()

        async with NotebookLMClient(auth) as client:
            await client.chat.ask("nb_x", "Q?", source_ids=["s1"])

        request = next(
            r for r in httpx_mock.get_requests() if "GenerateFreeFormStreamed" in str(r.url)
        )
        assert (
            _extract_query_param(str(request.url), "bl")
            == "boq_labs-tailwind-frontend_20260301.03_p0"
        )


# ---------------------------------------------------------------------------
# _build_chat_request direct unit-level coverage
# ---------------------------------------------------------------------------


class TestChatNewConversationLocks:
    """Notebook-scoped locks use the same per-key cache style as conversation locks."""

    def _factory(self) -> ChatAPI:
        from unittest.mock import MagicMock

        loop_guard = MagicMock()
        loop_guard.assert_bound_loop = MagicMock()
        return ChatAPI(
            rpc=MagicMock(),
            transport=MagicMock(),
            reqid=MagicMock(),
            loop_guard=loop_guard,
        )

    def test_same_notebook_reuses_new_conversation_lock(self):
        chat = self._factory()

        lock_a = chat._get_new_conversation_lock("nb-1")
        lock_b = chat._get_new_conversation_lock("nb-1")

        assert lock_a is lock_b

    def test_different_notebooks_get_distinct_new_conversation_locks(self):
        chat = self._factory()

        lock_a = chat._get_new_conversation_lock("nb-1")
        lock_b = chat._get_new_conversation_lock("nb-2")

        assert lock_a is not lock_b

    @pytest.mark.asyncio
    async def test_failed_post_ask_hptbtc_lookup_releases_new_conversation_lock(self):
        class HptbtcFailureChatAPI(ChatAPI):
            def __init__(self, *, lookup_results: list[str | ChatError], **kwargs: Any) -> None:
                super().__init__(**kwargs)
                self._lookup_results = iter(lookup_results)
                self.lookup_count = 0

            async def get_conversation_id(self, notebook_id: str) -> str | None:
                self.lookup_count += 1
                result = next(self._lookup_results)
                if isinstance(result, ChatError):
                    raise result
                return result

        async def fake_perform_authed_post(*args: Any, **kwargs: Any) -> httpx.Response:
            return httpx.Response(
                200,
                request=httpx.Request("POST", "https://notebooklm.google.com/_/LabsTailwindUi"),
                content=_make_answer_response_body(),
            )

        chat = HptbtcFailureChatAPI(
            rpc=SimpleNamespace(),
            transport=SimpleNamespace(
                perform_authed_post=AsyncMock(side_effect=fake_perform_authed_post)
            ),
            reqid=SimpleNamespace(next_reqid=AsyncMock(side_effect=[100000, 200000])),
            loop_guard=SimpleNamespace(assert_bound_loop=lambda: None),
            lookup_results=[ChatError("hPTbtc lookup failed"), "conv-after-failure"],
        )
        new_conversation_lock = chat._get_new_conversation_lock("nb-1")

        with pytest.raises(ChatError, match="hPTbtc lookup failed"):
            await chat.ask("nb-1", "first ask", source_ids=["s1"])

        assert not new_conversation_lock.locked()

        result = await asyncio.wait_for(
            chat.ask("nb-1", "second ask", source_ids=["s1"]),
            timeout=1.0,
        )

        assert result.conversation_id == "conv-after-failure"
        assert result.answer == "Refactor answer is long enough."
        assert chat.lookup_count == 2


class TestBuildChatRequestFactory:
    """Direct unit tests for the new ``ChatAPI._build_chat_request`` factory.

    Bypassing the full ``ask`` plumbing keeps these checks focused on the
    URL/body assembly contract that ``chat_aware_authed_post`` relies on.
    """

    def _factory(self) -> ChatAPI:
        # Wave 8 of session-decoupling (ADR-0014 Rule 2 Corollary):
        # ``ChatAPI`` takes direct collaborators by keyword arg. Pure
        # ``_build_chat_request`` exercise — none of these collaborators
        # are touched, so they are bare ``MagicMock()`` placeholders.
        from unittest.mock import MagicMock

        return ChatAPI(
            rpc=MagicMock(),
            transport=MagicMock(),
            reqid=MagicMock(),
            loop_guard=MagicMock(),
        )

    def test_build_request_omits_authuser_for_default_profile(self):
        chat = self._factory()
        snapshot = AuthSnapshot(
            csrf_token="csrf",
            session_id="sid",
            authuser=0,
            account_email=None,
        )
        url, body, headers = chat._build_chat_request(
            snapshot=snapshot,
            notebook_id="nb_x",
            question="Q?",
            source_ids=["s1"],
            conversation_history=None,
            conversation_id="conv-1",
            reqid=200000,
        )
        assert _extract_query_param(url, "authuser") is None
        assert _extract_query_param(url, "_reqid") == "200000"
        assert "at=csrf" in body
        assert headers == {}

    def test_build_request_authuser_email_wins_over_index(self):
        chat = self._factory()
        snapshot = AuthSnapshot(
            csrf_token="csrf",
            session_id="sid",
            authuser=5,
            account_email="me@example.com",
        )
        url, _, _ = chat._build_chat_request(
            snapshot=snapshot,
            notebook_id="nb_x",
            question="Q?",
            source_ids=["s1"],
            conversation_history=None,
            conversation_id="conv-1",
            reqid=300000,
        )
        # Email is preferred when present — matches ``format_authuser_value``.
        assert _extract_query_param(url, "authuser") == "me@example.com"

    def test_build_request_omits_at_when_csrf_blank(self):
        chat = self._factory()
        snapshot = AuthSnapshot(
            csrf_token="",
            session_id="sid",
            authuser=0,
            account_email=None,
        )
        _, body, _ = chat._build_chat_request(
            snapshot=snapshot,
            notebook_id="nb_x",
            question="Q?",
            source_ids=["s1"],
            conversation_history=None,
            conversation_id="conv-1",
            reqid=400000,
        )
        assert "at=" not in body

    def test_build_request_source_encoding_is_triple_nested(self):
        chat = self._factory()
        snapshot = AuthSnapshot(
            csrf_token="csrf",
            session_id="sid",
            authuser=0,
            account_email=None,
        )
        _, body, _ = chat._build_chat_request(
            snapshot=snapshot,
            notebook_id="nb_x",
            question="Q?",
            source_ids=["s1", "s2"],
            conversation_history=None,
            conversation_id="conv-1",
            reqid=500000,
        )
        match = re.search(r"f\.req=([^&]+)", body)
        assert match is not None
        f_req_data: list[Any] = json.loads(unquote(match.group(1)))
        params: list[Any] = json.loads(f_req_data[1])
        assert params[0] == [[["s1"]], [["s2"]]]