brain-of-reese/tests/e2e/test_tool_budget_telemetry.py

"""Phase 160 E2E (Playwright, mock-only): the agent knows its
tool-call budget (and never dumps tool calls into the answer).

Source: ``TODO.md`` L4 — "The LLM need to know how many tool calls
it has left. During long research sessions it will sometimes end the
conversation with tool calls that just get dumped into chat because
it ran out of turns. …"
Run in isolation (DB must be up: ``podman compose up -d db``):

    uv run pytest tests/e2e/test_tool_budget_telemetry.py -v --no-cov

MOCK-ONLY suite: ``E2E_REAL_LLM=1`` is not supported — the contract
is the deterministic ROUND-CAP marker flow in ``tests/e2e/mock_llm.py``
(``ROUND_CAP_TRIGGER`` = ``run out of tool calls``, detected on the
LAST user message only — the mock's existing invariant: the last
``user`` message is the current question, the earlier non-system
messages are prior turns):

* while ``tools`` are offered (the app issues such a request only
  while the round budget is unspent — i.e. no tool result answers
  the last emitted call yet) the mock NEVER answers — it streams
  ONLY ``tool_calls`` deltas calling ``ls`` with no arguments,
  incrementing ids ``call_cap_1``, ``call_cap_2``, … (the next id
  is derived from the ids already in the messages' ``assistant``
  tool calls), ``finish_reason: "tool_calls"``, no content — the
  turn is guaranteed to hit ``BOR_AGENT_MAX_ROUNDS``;
* on the request WITHOUT ``tools`` (the forced final) the mock
  streams a ``tool_calls`` delta ANYWAY — ``read`` on the seed
  path (id ``call_cap_dump``, NO content) — the exact dumped-call
  shape the phase kills (the app REJECTS it: never executed, never
  shown, counted in ``AgentHolder.rejected_calls``);
* the ``continue`` question (exactly ``continue`` — the UI's
  ``CONTINUE_QUESTION`` — with the trigger in an EARLIER user
  message; the client history rides in the stateless A10 request)
  gets the stable ``ROUND_CAP_CONTINUE_ANSWER`` (a fixed sentence
  naming the seed document — the ``SCOPED_SEARCH_PATH``
  suite-constant convention: the literal "continue" shares no tokens
  with the KB, so the continuation turn deflects and its prompt
  carries no ``<documents>`` block; the sentence names the document
  directly).

The e2e app server boots with ``BOR_AGENT_MAX_ROUNDS=3`` (the
module-scoped ``app_server`` override — the conftest env shape with
one swap): the cap flow burns exactly three executed rounds, and
the tool results carry ``[tool calls remaining: 2]``, ``[1 tool
call remaining — make it count, then answer]`` and the spent line
(``[No tool calls remain — answer now from what you have, or tell
the user to say "continue" to keep researching.]`` — the
``app.rag.agent`` constants, imported so the pin cannot drift).

KB fixture (the ``test_failed_turn_retry.py`` TRUNCATE-then-seed
pattern): ONE fixture document (``homelab/kubernetes.md``) — the
cap question's tokens clear the mock-calibrated gate (≈0.44 against
the mock's token-overlap embeddings, past the 0.30 threshold —
HIGH without any FTS help; the E2E app pins the mock-calibrated
values in conftest).

Test → contract mapping (``00_phase.md`` §Completion criteria):

1. ``test_tool_results_carry_the_remaining_count`` — server-side
   pin at the wire level: the mock's recorded request log (the
   ``/v1/e2e/captured`` pair — reset per turn) shows each
   successive ``tool``-role message ending with the decreasing
   trailer (``[tool calls remaining: 2]``, the last-call line, the
   spent line), the incrementing ``call_cap_N`` ids in the
   ``assistant`` messages, and the final request carrying the spent
   tool result + the coached system message (no ``tools`` key);
2. ``test_capped_turn_never_shows_tool_lines_after_the_cap`` —
   with ``BOR_AGENT_MAX_ROUNDS=3``: the chat shows EXACTLY three
   tool lines (one per executed round), the dump call produces NO
   tool line (no ``Reading <seed path>`` line anywhere), the answer
   bubble holds the non-empty ``ROUND_CAP_FALLBACK_ANSWER`` (the
   deterministic backstop — the mock streams no content on the
   forced final), the turn settles with ``done`` (never the error
   banner) and the SSE capture has no ``error`` frame;
3. ``test_done_flags_budget_exhausted`` — the captured SSE
   stream's ``done`` frame carries ``budget_exhausted: true`` on
   the capped turn; a normal grounded turn's ``done`` carries NO
   ``budget_exhausted`` key at all (absent, never false — the
   byte-identical control case in the same file); the capped
   answer's Continue button is gone once the newer turn lands;
4. ``test_continue_button_runs_a_fresh_turn`` — the capped answer
   bubble shows **Continue researching**; clicking it appends
   exactly ONE user bubble ``continue``, the mock answers with
   ``ROUND_CAP_CONTINUE_ANSWER`` (a real continuation answer — a
   fixed sentence naming the seed document, not a tool call), the
   earlier tool lines do not duplicate (still exactly three — the
   continuation turn streams no tool frames), and ``bor.chat.v1``
   holds ``[user, brain(flagged), user("continue"), brain]`` (the
   flagged record carries ``budget_exhausted: true``);
5. ``test_continue_is_inert_while_streaming_and_single_shot`` —
   double-clicking Continue (``page.dblclick`` — the second click
   lands after the first handler removed every ``.continue-btn``
   and started the turn; the in-flight no-op guard + the
   remove-first clear make a double-send impossible) starts ONE
   turn (counted via ``page.on("request")`` — exactly one
   ``/api/chat`` POST; the single ``continue`` user bubble + the
   single continuation answer pin it);
6. ``test_flag_survives_a_refresh`` — ``page.reload()`` after a
   capped turn: the flagged answer restores WITH the Continue
   button (the restore path keys on ``m.budget_exhausted``) and
   clicking it still works (the mock's continuation answer
   streams);
7. negative: ``test_deflected_turn_gets_no_continue_button`` — a
   deflected question (A8 LOW mode — token-disjoint from the KB:
   ``Who composed Rhapsody in Blue?`` ≈0.05, zero FTS hits)
   settles deflected with NO ``budget_exhausted`` key, NO ``tool``
   frame and NO Continue button.
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import subprocess
import sys
import time
from collections.abc import Iterator
from datetime import UTC, datetime
from pathlib import Path

import httpx
import pytest
from playwright.sync_api import Page, expect
from sqlalchemy import text
from sqlalchemy.orm import Session

from app.config import Settings as _Settings
from app.db import SessionLocal
from app.models import Chunk, Document
from app.rag.agent import (
    ROUND_CAP_FALLBACK_ANSWER,
    ROUND_CAP_WRAPUP_INSTRUCTION,
    TOOL_BUDGET_LAST,
    TOOL_BUDGET_REMAINING,
    TOOL_BUDGET_SPENT,
)
from e2e.auth_helpers import login
from e2e.conftest import (
    ADMIN_PASSWORD,
    SESSION_SECRET,
    USE_REAL_LLM,
    _wait_http,
)
from tests.e2e.mock_llm import ROUND_CAP_CONTINUE_ANSWER, embed_text

REPO = Path(__file__).resolve().parents[2]

#: This module app binds its OWN port (the conftest session app
#: owns its port in a combined run — a same-port second uvicorn
#: dies on bind and would drive the wrong server; the
#: test_big_read_progress.py convention, env-overridable).
APP_PORT = int(os.environ.get("E2E_APP_PORT_TOOLBUDGET", "8129"))
APP_URL = f"http://127.0.0.1:{APP_PORT}"

#: The cap the whole suite runs under (the module app's env
#: override): the cap flow burns exactly three executed rounds —
#: the trailer math runs over a 3-round budget (2 → 1 → spent).
MAX_ROUNDS = 3

# --------------------------------------------------------------------------
# Seed + questions (see the module docstring for the gate notes)
# --------------------------------------------------------------------------

SEED_SOURCE = "docs"
SEED_PATH = "homelab/kubernetes.md"
SEED_SP = f"{SEED_SOURCE}/{SEED_PATH}"
KUB_CONTENT = (REPO / "tests" / "fixtures" / "docs" / SEED_PATH).read_text()

#: The cap-flow question: the tokens clear the mock-calibrated
#: gate (≈0.44 against the mock's token-overlap embeddings, past
#: the 0.30 threshold — HIGH without any FTS help), and it carries
#: ``ROUND_CAP_TRIGGER``.
CAP_QUESTION = (
    "Explain my kubernetes cluster: talos linux nodes, cilium "
    "networking, local-path-provisioner and longhorn storage. "
    "run out of tool calls"
)
assert "run out of tool calls" in CAP_QUESTION.lower()

#: The normal grounded control (test 3 — the byte-identical
#: ``done`` case): token overlap with the seed (≈0.21, past the
#: 0.15 lexical floor) + FTS hits — grounded, no trigger, no tool
#: calls (the mock answers plain).
CONTROL_QUESTION = "How is my Kubernetes cluster set up?"

#: The negative (test 7 — A8 LOW): token-disjoint from the KB
#: (≈0.05, zero FTS hits — deflected; the phase-112 calibrated
#: question).
DEFLECT_QUESTION = "Who composed Rhapsody in Blue?"

STORAGE_KEY = "bor.chat.v1"

# --------------------------------------------------------------------------
# The module app (conftest env shape + the BOR_AGENT_MAX_ROUNDS=3
# swap)
# --------------------------------------------------------------------------


@pytest.fixture(scope="module")
def app_server(mock_llm: int) -> Iterator[str]:
    """The app under test — the conftest's session-app env shape
    (including the phase-61/62 leak guards) with one swap:
    ``BOR_AGENT_MAX_ROUNDS=3`` (the cap flow's denominator — the
    trailer math runs over a 3-round budget)."""
    env = dict(os.environ)
    env.pop("DEBUGPY", None)
    env["BOR_ENVIRONMENT"] = "e2e"
    env["BOR_STATIC_DIR"] = str(REPO / "frontend")
    env["BOR_LLM_BASE_URL"] = (
        "https://aipi.reeseapps.com/v1"
        if USE_REAL_LLM
        else f"http://127.0.0.1:{mock_llm}/v1"
    )
    # The E2E mock's token-overlap embeddings have their own score
    # distribution (phase 09) — the mock-calibrated values (the
    # conftest pattern).
    env["BOR_RELEVANCE_THRESHOLD"] = "0.30"
    env["BOR_LEXICAL_SUPPORT_FLOOR"] = "0.15"
    env["BOR_SOURCE_USEFULNESS_FLOOR"] = "0.15"
    # Phase 67: instant retry waits; the real retry budget (the
    # conftest pattern).
    env["BOR_LLM_RETRY_DELAY"] = "0"
    env["BOR_LLM_RETRIES"] = str(_Settings.model_fields["llm_retries"].default)
    env.setdefault(
        "BOR_DATABASE_URL",
        "postgresql+psycopg://reese:reese@localhost:5432/brain_of_reese",
    )
    # Phase 16: admin auth must be set or create_app() refuses to boot.
    env["BOR_ADMIN_PASSWORD"] = ADMIN_PASSWORD
    env["BOR_SESSION_SECRET"] = SESSION_SECRET
    # Phase 61/62 leak guards (the conftest pattern): force the code
    # defaults so an operator's local (gitignored) ``.env`` cannot
    # leak into the app under test.
    env["BOR_DOCS_REPO"] = ""
    env["BOR_SUGGESTIONS"] = json.dumps(
        _Settings.model_fields["suggestions"].default
    )
    env["BOR_INPUT_PLACEHOLDER"] = _Settings.model_fields["input_placeholder"].default
    env["BOR_FOOTER_TEXT"] = _Settings.model_fields["footer_text"].default
    # The suite's swap: the cap flow burns exactly three executed
    # rounds (the trailer math: 2 → 1 → spent).
    env["BOR_AGENT_MAX_ROUNDS"] = str(MAX_ROUNDS)
    proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "app.main:app",
         "--host", "127.0.0.1", "--port", str(APP_PORT), "--log-level", "warning"],
        cwd=REPO,
        env=env,
    )
    try:
        _wait_http(f"{APP_URL}/api/health")
        yield APP_URL
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()


@pytest.fixture(scope="module")
def app_url(app_server: str) -> str:
    return app_server


# --------------------------------------------------------------------------
# DB seeding (TRUNCATE-then-seed, cf. test_failed_turn_retry.py)
# --------------------------------------------------------------------------


def _seed(db: Session) -> None:
    """The single fixture document (see the module docstring)."""
    md = Document(
        source=SEED_SOURCE,
        path=SEED_PATH,
        full_path=f"/tmp/{SEED_PATH}",
        title="Kubernetes",
        content=KUB_CONTENT,
        content_hash=hashlib.sha256(KUB_CONTENT.encode()).hexdigest(),
        indexed_at=datetime.now(UTC),
    )
    db.add(md)
    db.flush()
    # One chunk carrying the mock's own embedding → genuine token
    # overlap for the grounded questions (the FTS path carries them
    # to HIGH).
    db.add(
        Chunk(
            document_id=md.id,
            position=0,
            content=KUB_CONTENT,
            embedding=embed_text(KUB_CONTENT),
        )
    )


def _reset_db() -> None:
    """Truncate the KB (plus the prompt-shaping tables), then
    re-seed.

    ``steering_notes`` / ``kb_overview`` are truncated too, so the
    prompts are byte-stable regardless of leftovers from other
    suites.
    """
    with SessionLocal() as db:
        db.execute(
            text(
                "TRUNCATE chunks, documents, query_log, steering_notes, "
                "kb_overview"
            )
        )
        db.commit()
        _seed(db)
        db.commit()


@pytest.fixture(autouse=True)
def clean_chats() -> Iterator[None]:
    """The send auto-saves a ``saved_chats`` row per turn (phase 55) —
    truncate around every test so the suite starts from (and leaves)
    an empty deployment (the phase-80/103 isolation pattern)."""
    with SessionLocal() as db:
        db.execute(text("TRUNCATE saved_chats"))
        db.commit()
    yield
    # The auto-save is fire-and-forget (the browser never awaits the
    # PUT): a short margin lets the last turn's PUT land server-side
    # before the TRUNCATE, so the teardown never 500s an in-flight
    # upsert (the "Could not refresh instance" race). This runs AFTER
    # the page closes (LIFO finalization), so the PUT is either done
    # or aborted — never in flight.
    time.sleep(0.75)
    with SessionLocal() as db:
        db.execute(text("TRUNCATE saved_chats"))
        db.commit()


# --------------------------------------------------------------------------
# Page helpers (the SSE capture — the phase-37/67 pattern) + the
# mock's request-log capture
# --------------------------------------------------------------------------

SSE_HOOK = """
() => {
  if (window.__sseInstalled) return;
  window.__sseInstalled = true;
  window.__sseFrames = [];
  const origFetch = window.fetch;
  window.fetch = async function (...args) {
    const res = await origFetch.apply(this, args);
    try {
      const url = typeof args[0] === 'string' ? args[0] : args[0].url;
      if (url.includes('/api/chat')) {
        res.clone().text().then((bodyText) => {
          for (const block of bodyText.split('\\n\\n')) {
            const line = block.trim();
            if (line.startsWith('data: ')) {
              window.__sseFrames.push(line.slice(6));
            }
          }
        });
      }
    } catch (e) { /* non-clonable responses: ignored */ }
    return res;
  };
}
"""


def _install_sse_hook(page: Page) -> None:
    page.evaluate(SSE_HOOK)


def _frames(page: Page, terminal: str = "done") -> list[dict]:
    """The captured SSE frames of the CURRENT turn, once the
    *terminal* frame lands."""
    deadline = time.monotonic() + 30.0
    while True:
        raw = page.evaluate("() => window.__sseFrames || []")
        parsed = [json.loads(line) for line in raw if line]
        if any(f.get("type") == terminal for f in parsed):
            return parsed
        if time.monotonic() > deadline:
            raise AssertionError(
                f"SSE hook captured no `{terminal}` frame (frames so far: "
                f"{len(parsed)}) — hook install failed?"
            )
        time.sleep(0.05)


def _stored(page: Page) -> dict:
    """The persisted ``bor.chat.v1`` conversation (localStorage — the
    phase-14 store the restore reads on refresh)."""
    raw = page.evaluate(f"() => localStorage.getItem('{STORAGE_KEY}')")
    assert raw is not None, "the conversation key must exist in localStorage"
    return json.loads(raw)


def _submit(page: Page, question: str) -> None:
    page.fill("#message-input", question)
    page.click("#send-btn")
    # The user bubble lands synchronously with the submit handler.
    expect(page.locator(".msg.user .bubble").last).to_contain_text(question)


def _wait_settled(page: Page) -> None:
    """The turn is complete: answer text in the bubble, button
    recovered (the phase-48 settle convention)."""
    expect(page.locator(".msg.brain .bubble").last).not_to_have_text(
        "", timeout=30_000
    )
    expect(page.locator("#send-btn")).to_be_enabled(timeout=30_000)
    expect(page.locator("#send-label")).to_have_text("Send", timeout=30_000)


def _mock_capture_reset(mock_port: int) -> None:
    """Clear the mock's request capture so a suite asserts on
    EXACTLY the requests of the turn it is about to drive (the
    phase-123 capture pair)."""
    r = httpx.post(
        f"http://127.0.0.1:{mock_port}/v1/e2e/captured/reset", timeout=5
    )
    assert r.status_code == 200


def _mock_captured(mock_port: int) -> list[dict]:
    """The mock's recorded chat-completions request bodies (oldest
    first) — the wire-level pin surface for test 1."""
    r = httpx.get(
        f"http://127.0.0.1:{mock_port}/v1/e2e/captured", timeout=5
    )
    assert r.status_code == 200
    return r.json()


def _tool_msgs(messages: list[dict]) -> list[str]:
    """The ``tool``-role result contents of one request, in order."""
    return [
        str(m.get("content") or "")
        for m in messages
        if m.get("role") == "tool"
    ]


# --------------------------------------------------------------------------
# 1. The budget trailer on every executed tool result (wire level)
# --------------------------------------------------------------------------


def test_tool_results_carry_the_remaining_count(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _mock_capture_reset(mock_llm)

    _submit(page, CAP_QUESTION)
    _wait_settled(page)

    # Wire level (the mock's recorded request log — the app's
    # EXACT requests): exactly four chat requests — three executed
    # tool rounds + the forced final.
    bodies = _mock_captured(mock_llm)
    assert len(bodies) == 4, len(bodies)

    # Request 1: no tool results yet; the tools are offered.
    assert bodies[0].get("tools"), "request 1 must offer the tools"
    assert _tool_msgs(bodies[0]["messages"]) == []

    # Round 1's result (in request 2): the decreasing trailer —
    # ``[tool calls remaining: 2]`` (two of three left after the
    # call it is reading).
    msgs2 = _tool_msgs(bodies[1]["messages"])
    assert len(msgs2) == 1, msgs2
    assert msgs2[0].endswith(TOOL_BUDGET_REMAINING.format(n=2)), msgs2[0]

    # Round 2's result (in request 3): the last-call line (one left
    # — "make it count, then answer").
    msgs3 = _tool_msgs(bodies[2]["messages"])
    assert len(msgs3) == 2, msgs3
    assert msgs3[0].endswith(TOOL_BUDGET_REMAINING.format(n=2)), msgs3[0]
    assert msgs3[1].endswith(TOOL_BUDGET_LAST), msgs3[1]

    # The boundary call's result (in request 4 — the forced final):
    # the SPENT line (the budget is gone; say "continue").
    msgs4 = _tool_msgs(bodies[3]["messages"])
    assert len(msgs4) == 3, msgs4
    assert msgs4[2].endswith(TOOL_BUDGET_SPENT), msgs4[2]

    # The mock's incrementing ids (the flow's stateless
    # derivation): call_cap_1 → call_cap_2 → call_cap_3 in the
    # assistant messages.
    asst_ids = [
        bodies[i + 1]["messages"][2 * i + 2]["tool_calls"][0]["id"]
        for i in range(3)
    ]
    assert asst_ids == ["call_cap_1", "call_cap_2", "call_cap_3"], asst_ids

    # The final request is the no-tools one, coached: no ``tools``
    # key (or null), the system message carries the wrap-up
    # instruction, and the user question + the tool history stay in
    # place.
    assert bodies[3].get("tools") is None
    systems = [
        str(m.get("content") or "")
        for m in bodies[3]["messages"]
        if m.get("role") == "system"
    ]
    assert any(ROUND_CAP_WRAPUP_INSTRUCTION in s for s in systems), systems
    users = [
        str(m.get("content") or "")
        for m in bodies[3]["messages"]
        if m.get("role") == "user"
    ]
    assert CAP_QUESTION in users, users


# --------------------------------------------------------------------------
# 2. Past the cap: zero tool lines after the cap, the fallback
#    answer, no error
# --------------------------------------------------------------------------


def test_capped_turn_never_shows_tool_lines_after_the_cap(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _install_sse_hook(page)

    _submit(page, CAP_QUESTION)
    _wait_settled(page)

    frames = _frames(page)

    # The turn settles with ``done`` — never the terminal ``error``
    # frame (the pre-phase-160 cap path's MalformedReplyError class
    # is gone; the budget fallback is the terminal).
    assert not [f for f in frames if f.get("type") == "error"], frames
    assert frames[-1]["type"] == "done"

    # Wire: EXACTLY three ``tool`` frames — one per executed round
    # (the cap). The dump call (the forced final's structured
    # ``read`` on the seed path) produces NO frame — it was
    # rejected (never executed, never shown).
    tool_frames = [f for f in frames if f.get("type") == "tool"]
    assert len(tool_frames) == MAX_ROUNDS, tool_frames
    assert all(f["name"] == "ls" for f in tool_frames), tool_frames
    assert not [f for f in tool_frames if f.get("argument") == SEED_SP]

    # UI: exactly three tool lines (all the no-arg top-level
    # listings); the dump call's ``Reading <seed path>`` line is
    # NOWHERE — the rejected call never renders.
    lines = page.locator(".msg.brain .tool-call")
    expect(lines).to_have_count(3)
    for i in range(3):
        expect(lines.nth(i)).to_contain_text("Listing documents")
    msg_text = page.locator(".msg.brain").last.text_content() or ""
    assert "Reading " not in msg_text, msg_text

    # The answer: the deterministic backstop line (the mock streams
    # no content on the forced final — only the rejected dump
    # call) — non-empty and readable.
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(ROUND_CAP_FALLBACK_ANSWER)

    # No error banner (the turn settled, never the failure path).
    expect(page.locator("#kb-banner")).to_be_hidden()


# --------------------------------------------------------------------------
# 3. done.budget_exhausted: true exactly on the capped turn, ABSENT
#    (never false) on the normal grounded control
# --------------------------------------------------------------------------


def test_done_flags_budget_exhausted(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _install_sse_hook(page)

    # The capped turn: ``done`` carries ``budget_exhausted: true``.
    _submit(page, CAP_QUESTION)
    _wait_settled(page)
    frames = _frames(page)
    done = next(f for f in frames if f.get("type") == "done")
    assert done["budget_exhausted"] is True, done
    assert done["deflected"] is False
    # The flagged answer offers the Continue affordance (task 04).
    expect(page.locator(".continue-btn")).to_have_count(1)

    # The normal grounded control: the ``done`` frame carries NO
    # ``budget_exhausted`` key at all (absent, never false — the
    # byte-identical off rule, PLAN §13).
    page.evaluate("() => { window.__sseFrames = []; }")
    _submit(page, CONTROL_QUESTION)
    _wait_settled(page)
    frames = _frames(page)
    done = next(f for f in frames if f.get("type") == "done")
    assert "budget_exhausted" not in done, done
    assert done["deflected"] is False

    # The newer turn's landing cleared the capped answer's Continue
    # button (the remove-first pattern — a stale "continue" would
    # append to a history the conversation moved past).
    expect(page.locator(".continue-btn")).to_have_count(0)


# --------------------------------------------------------------------------
# 4. Continue researching: one fresh turn, the real continuation
#    answer, the record shape
# --------------------------------------------------------------------------


def test_continue_button_runs_a_fresh_turn(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _install_sse_hook(page)

    _submit(page, CAP_QUESTION)
    _wait_settled(page)
    expect(page.locator(".msg.brain .bubble").last).to_contain_text(
        ROUND_CAP_FALLBACK_ANSWER
    )

    # The flagged answer offers **Continue researching**.
    btn = page.locator(".continue-btn")
    expect(btn).to_have_count(1)
    expect(btn).to_contain_text("Continue researching")

    # Fresh capture window: the continuation turn's frames only
    # (the cap turn's frames are still in the hook's buffer).
    page.evaluate("() => { window.__sseFrames = []; }")
    btn.click()

    # Exactly ONE ``continue`` user bubble (the plain send).
    expect(page.locator(".msg.user .bubble")).to_have_count(2)
    expect(page.locator(".msg.user .bubble").last).to_have_text("continue")

    # … and the mock's REAL continuation answer (a fixed sentence
    # naming the seed document — not a tool call).
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(ROUND_CAP_CONTINUE_ANSWER, timeout=30_000)
    _wait_settled(page)

    # The continuation turn streams NO tool frames — the earlier
    # tool lines do not duplicate (still exactly three).
    frames = _frames(page)
    assert [f for f in frames if f.get("type") == "tool"] == []
    expect(page.locator(".msg.brain .tool-call")).to_have_count(3)
    done = next(f for f in frames if f.get("type") == "done")
    # A FRESH turn: the budget is whole again — the flag is absent
    # (the continuation turn deflects — the literal "continue"
    # shares no tokens with the KB, so the A8 gate is LOW).
    assert "budget_exhausted" not in done, done
    assert done["deflected"] is True

    # The record: [user, brain(flagged), user("continue"), brain].
    stored = _stored(page)
    assert [m["who"] for m in stored["messages"]] == [
        "user",
        "brain",
        "user",
        "brain",
    ]
    assert stored["messages"][1]["budget_exhausted"] is True
    assert stored["messages"][1]["text"] == ROUND_CAP_FALLBACK_ANSWER
    assert stored["messages"][2]["text"] == "continue"
    assert ROUND_CAP_CONTINUE_ANSWER in stored["messages"][3]["text"]
    assert not stored["messages"][3].get("budget_exhausted")


# --------------------------------------------------------------------------
# 5. Single-shot: double-clicking Continue starts ONE turn
# --------------------------------------------------------------------------


def test_continue_is_inert_while_streaming_and_single_shot(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _install_sse_hook(page)

    _submit(page, CAP_QUESTION)
    _wait_settled(page)
    expect(page.locator(".continue-btn")).to_have_count(1)

    # Count the ``/api/chat`` POSTs (the single-shot pin) — the
    # exact path (the auto-save's ``PUT /api/chats/…`` must not
    # match: ``/api/chats`` is not ``/api/chat``).
    posts: list[int] = []
    page.on(
        "request",
        lambda req: (
            posts.append(1)
            if req.method == "POST"
            and req.url.split("?", 1)[0].endswith("/api/chat")
            else None
        ),
    )

    # The double-click: the FIRST click runs the handler (removes
    # every ``.continue-btn`` and starts the turn); the SECOND
    # click lands after the removal — no second target (and the
    # in-flight no-op guard would swallow it anyway) — ONE turn.
    page.locator(".continue-btn").dblclick()
    _wait_settled(page)
    assert len(posts) == 1, posts

    # ONE ``continue`` user bubble + ONE continuation answer (no
    # double-send: a second turn would add a second user bubble +
    # a second brain answer).
    expect(page.locator(".msg.user .bubble")).to_have_count(2)
    expect(page.locator(".msg.brain")).to_have_count(2)
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(ROUND_CAP_CONTINUE_ANSWER, timeout=30_000)


# --------------------------------------------------------------------------
# 6. The flag survives a refresh (the restore path)
# --------------------------------------------------------------------------


def test_flag_survives_a_refresh(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _install_sse_hook(page)

    _submit(page, CAP_QUESTION)
    _wait_settled(page)
    expect(page.locator(".msg.brain .bubble").last).to_contain_text(
        ROUND_CAP_FALLBACK_ANSWER
    )
    expect(page.locator(".continue-btn")).to_have_count(1)

    # The refresh: the flagged answer restores WITH the Continue
    # button (the restore path keys on ``m.budget_exhausted``).
    page.reload()
    expect(page.locator("#messages > .msg")).to_have_count(2)
    expect(page.locator(".msg.user .bubble")).to_have_text(CAP_QUESTION)
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(ROUND_CAP_FALLBACK_ANSWER, timeout=30_000)
    btn = page.locator(".continue-btn")
    expect(btn).to_have_count(1)
    expect(btn).to_contain_text("Continue researching")

    # The restored button WORKS: clicking it runs the continuation
    # turn (the mock's continuation answer streams).
    btn.click()
    expect(page.locator(".msg.user .bubble")).to_have_count(2)
    expect(page.locator(".msg.brain .bubble").last).to_contain_text(
        ROUND_CAP_CONTINUE_ANSWER, timeout=30_000
    )
    _wait_settled(page)

    # The record: [user, brain(flagged), user("continue"), brain].
    stored = _stored(page)
    assert [m["who"] for m in stored["messages"]] == [
        "user",
        "brain",
        "user",
        "brain",
    ]
    assert stored["messages"][1]["budget_exhausted"] is True
    assert stored["messages"][2]["text"] == "continue"
    assert ROUND_CAP_CONTINUE_ANSWER in stored["messages"][3]["text"]


# --------------------------------------------------------------------------
# 7. Negative: a deflected turn — no flag, no Continue button
# --------------------------------------------------------------------------


def test_deflected_turn_gets_no_continue_button(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    _reset_db()
    page.set_default_timeout(30_000)
    login(page, app_url, next="/")
    _install_sse_hook(page)

    _submit(page, DEFLECT_QUESTION)
    _wait_settled(page)

    frames = _frames(page)
    done = next(f for f in frames if f.get("type") == "done")
    # Settled DEFLECTED (A8 LOW), with NO ``budget_exhausted`` key
    # (the flag is cap-reached turns only) and NO ``tool`` frame
    # (a tools=None request never emits one).
    assert done["deflected"] is True, done
    assert "budget_exhausted" not in done, done
    assert [f for f in frames if f.get("type") == "tool"] == []

    # The deflected bubble — the honest copy, and NO Continue
    # button (the affordance is cap-reached answers only).
    last = page.locator(".msg.brain").last
    expect(last).to_have_class(re.compile(r"is-deflected"))
    expect(page.locator(".msg.brain .tool-call")).to_have_count(0)
    expect(page.locator(".continue-btn")).to_have_count(0)
    expect(page.locator(".msg.brain .bubble").last).to_contain_text(
        "I haven't done anything like that"
    )