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

"""Phase 37 E2E (Playwright, mock-only): agent document tools (list + read).

Story: ``.agent/user_stories/agent-document-tools.md``
Run in isolation (DB must be up: ``podman compose up -d db``):

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

MOCK-ONLY suite: ``E2E_REAL_LLM=1`` is not supported — the real ``turbo``
does whatever it does with the tools, while this story's gate is the
deterministic marker flow in ``tests/e2e/mock_llm.py`` (user message
contains ``use your tools`` **and** the system prompt carries the
``<tools>`` section of the HIGH prompt):

1. request 1 (``tools`` offered, no tool results yet) → streams ONLY
   ``tool_calls`` deltas calling ``list_documents`` (id ``call_0``, no
   arguments, ``finish_reason: "tool_calls"``);
2. request 2 (a ``tool``-role catalog result in the messages) → streams a
   ``tool_calls`` delta calling ``read_document`` on the FIRST catalog
   line (id ``call_1``);
3. request 3 (no ``tools`` parameter, the read result in the messages) →
   the content answer ``Read <source/path>. <first 80 chars of the read
   document's content>`` — so the suite can assert the read document
   reached the model and landed in the answer.

KB fixture — reproduces the TODO failure (``aws-route53.md`` references
``example-record-file.json`` "for the exact JSON shape of
reseelink.json" but does not include it):

* ``Homelab/aws-route53.md`` — seeded with one chunk whose embedding is
  the mock's own bag-of-words vector (genuine token overlap: the marker
  question cosines ≈0.69 against it, well past the E2E 0.30 threshold,
  and it FTS-matches too) → the only RETRIEVABLE document, i.e. the
  grounded context;
* ``Deployments/example-record-file.json`` — the JSON shape, indexed
  (a ``documents`` row: it is in the agent's catalog, readable, and a
  source-chip target) but seeded WITHOUT chunks. In a real
  hundreds-of-document KB the file would simply fail to rank into the
  top-2 context; with a two-document corpus every chunk would rank, so
  "not in context" is expressed as "no retrieval candidates". Its
  ``(source, path)`` also sorts FIRST in the catalog
  (``Deployments`` < ``Homelab``) — which is exactly the line the mock
  parses out of the listing and reads.

Test → story mapping (Playwright Mapping Rule):
1. ``test_marker_question_lists_reads_and_quotes`` — the SSE carries
   ``tool`` frames (list, then read, ahead of any delta), the UI shows
   the "calling tool" label while a tool runs, the bubble shows both
   tool lines, the final answer quotes the read document, and the
   source chips include the read document (viewer link).
2. ``test_tool_lines_re_render_after_reload`` — the persisted record
   (phase 14) re-renders the tool lines.
3. ``test_plain_grounded_question_has_no_tool_frames`` — no marker → no
   ``tool`` frames, the answer renders exactly as today (regression
   inside the story file).
4. ``test_deflected_question_has_no_tool_frames`` — the tools are
   grounded-only: a deflected turn runs none.
"""
from __future__ import annotations

import hashlib
import json
import re
import time
from collections.abc import Callable
from datetime import UTC, datetime

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

from app.db import SessionLocal
from app.models import Chunk, Document, QueryLog
from tests.e2e.mock_llm import embed_text

# --------------------------------------------------------------------------
# Fixture documents (deterministic, token-controlled)
# --------------------------------------------------------------------------

SEED_SOURCE = "Homelab"
SEED_PATH = "aws-route53.md"
SEED_SP = f"{SEED_SOURCE}/{SEED_PATH}"

READ_SOURCE = "Deployments"
READ_PATH = "example-record-file.json"
READ_SP = f"{READ_SOURCE}/{READ_PATH}"

#: The retrievable document: references the JSON file "for the exact JSON
#: shape of reeselink.json" but never includes it (the TODO failure).
#: The repeated record-file lines carry the marker question's key tokens
#: (aws, route53, hosted, zone, reeselink, json, exact, shape) — verified
#: ≈0.69 cosine against the mock's embeddings (E2E threshold 0.30) plus
#: FTS hits, so the turn is solidly grounded.
ROUTE53_CONTENT = (
    "# AWS Route 53 Notes\n\n"
    "## Record file\n\n"
    + (
        "The aws route53 hosted zone for reeselink keeps every record in "
        "reseelink.json — the exact JSON shape of reeselink.json is "
        "documented in example-record-file.json.\n"
    )
    * 10
    + "\n## Sync job\n\n"
    "A cron job pushes reeselink.json to the aws route53 hosted zone "
    "every fifteen minutes; the diff is applied through the route53 api.\n"
)

#: The referenced document: the exact JSON shape. Its FIRST line is longer
#: than 80 chars, so the mock's first-80-chars quote is newline-free (the
#: rendered-text assertions below match it verbatim). Pinned by the assert
#: below.
RECORD_FILE_CONTENT = (
    '{ "version": 3, "comment": "ReeseLink hosted zone records — the exact '
    'JSON shape of reeselink.json",\n'
    '  "hosted_zone_id": "Z0RESEELINK01",\n'
    '  "record_sets": [\n'
    '    { "name": "www.reeselink.example", "type": "A", "ttl": 300,\n'
    '      "resource_records": [ { "value": "10.0.0.20" } ] },\n'
    '    { "name": "api.reeselink.example", "type": "CNAME", "ttl": 300,\n'
    '      "resource_records": [ { "value": "www.reeselink.example" } ] }\n'
    "  ]\n"
    "}\n"
)
assert "\n" not in RECORD_FILE_CONTENT[:80]  # the quote must stay one line

MARKER_QUESTION = (
    "Use your tools: what is the exact JSON shape of reeselink.json "
    "for my aws route53 hosted zone?"
)
PLAIN_QUESTION = (
    "How does my aws route53 sync job push reeselink.json to the "
    "hosted zone?"
)
DEFLECT_QUESTION = "tell me about quantum wormhole cooling"

MOCK_ANSWER_MARKER = "Deterministic mock answer for E2E"
DEFLECT_PHRASE = r"haven't done anything like that"
ANSWER_PREFIX = f"Read {READ_SP}."
ANSWER_QUOTE = RECORD_FILE_CONTENT[:80]
READ_CHIP_HREF = f"/document.html?source={READ_SOURCE}&path={READ_PATH}&back=%2F"

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


def _seed(db: Session) -> None:
    """The two-document pair from the TODO (see the module docstring)."""
    md = Document(
        source=SEED_SOURCE,
        path=SEED_PATH,
        full_path=f"/tmp/{SEED_PATH}",
        title="AWS Route 53 Notes",
        content=ROUTE53_CONTENT,
        content_hash=hashlib.sha256(ROUTE53_CONTENT.encode()).hexdigest(),
        indexed_at=datetime.now(UTC),
    )
    db.add(md)
    db.flush()
    # One chunk carrying the mock's own embedding → genuine token overlap
    # between the marker question and this document.
    db.add(
        Chunk(
            document_id=md.id,
            position=0,
            content=ROUTE53_CONTENT,
            embedding=embed_text(ROUTE53_CONTENT),
        )
    )
    # The referenced JSON: indexed, catalogued, readable — but NO chunks,
    # so retrieval never puts it in context (the failure the tools fix).
    db.add(
        Document(
            source=READ_SOURCE,
            path=READ_PATH,
            full_path=f"/tmp/{READ_PATH}",
            title="Example Record File",
            content=RECORD_FILE_CONTENT,
            content_hash=hashlib.sha256(RECORD_FILE_CONTENT.encode()).hexdigest(),
            indexed_at=datetime.now(UTC),
        )
    )


def _reset_db(seed: Callable[[Session], None] | None = None) -> None:
    """Truncate the KB (plus the prompt-shaping tables), then re-seed.

    ``steering_notes`` / ``kb_overview`` are truncated too, so the HIGH
    prompt is exactly ``<relevance>`` + ``<documents>`` + ``<tools>``
    regardless of leftovers from other suites — byte-stable prompts,
    byte-stable answers.
    """
    with SessionLocal() as db:
        db.execute(
            text("TRUNCATE chunks, documents, query_log, steering_notes, kb_overview")
        )
        db.commit()
        if seed is not None:
            seed(db)
            db.commit()


def _last_query_log() -> QueryLog:
    with SessionLocal() as db:
        rows = db.scalars(select(QueryLog)).all()
    assert len(rows) == 1, f"expected exactly one query_log row, got {len(rows)}"
    return rows[0]


# --------------------------------------------------------------------------
# Page helpers
# --------------------------------------------------------------------------

#: Records every value #send-label takes during the turn (a
#: MutationObserver on the element), so the transient "Calling tool…"
#: state is captured deterministically — no polling race.
LABEL_RECORDER = """
() => {
  if (window.__labelsInstalled) return;
  window.__labelsInstalled = true;
  window.__labels = [];
  const el = document.querySelector('#send-label');
  if (!el) return;
  const rec = (v) => {
    const l = window.__labels;
    if (!l.length || l[l.length - 1] !== v) l.push(v);
  };
  rec(el.textContent);
  new MutationObserver(() => rec(el.textContent)).observe(el, {
    childList: true,
    subtree: true,
  });
}
"""

#: Captures the raw SSE ``data:`` payloads of the /api/chat stream
#: (a response clone read in the background) — wire-level assertions for
#: the ``tool`` frames, independent of the UI rendering.
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_page_hooks(page: Page) -> None:
    """Install both hooks on the loaded page (post-goto, pre-submit).

    The fetch wrapper only needs to be in place before the turn's
    ``fetch("/api/chat")`` call; the label observer needs the rendered
    ``#send-label``. (``add_init_script`` would not do — it binds to the
    NEXT navigation, and the story page is navigated exactly once.)
    """
    page.evaluate(SSE_HOOK)
    page.evaluate(LABEL_RECORDER)


def _frames(page: Page) -> list[dict]:
    """The captured SSE frames, once the hook's background read settles.

    The hook reads ``res.clone().text()`` in a background promise that
    resolves right after the stream closes — poll briefly until the
    final ``done`` frame lands (fail loud if the hook captured nothing).
    """
    deadline = time.monotonic() + 10.0
    while True:
        raw = page.evaluate("() => window.__sseFrames || []")
        parsed = [json.loads(line) for line in raw if line]
        if any(f.get("type") == "done" for f in parsed):
            return parsed
        if time.monotonic() > deadline:
            raise AssertionError(
                f"SSE hook captured no `done` frame (frames so far: "
                f"{len(parsed)}) — hook install failed?"
            )
        time.sleep(0.05)


def _tool_frames(frames: list[dict]) -> list[dict]:
    return [f for f in frames if f.get("type") == "tool"]


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."""
    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")


# --------------------------------------------------------------------------
# 1. The marker question: list → read → quoted answer, "calling tool" UI
# --------------------------------------------------------------------------


def test_marker_question_lists_reads_and_quotes(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    page.set_default_timeout(30_000)
    _reset_db(_seed)
    page.goto(app_url)
    _install_page_hooks(page)

    _submit(page, MARKER_QUESTION)
    # While a tool runs the button carries the "calling tool" label: the
    # first `tool` frame sets it and it holds until the FIRST answer
    # delta (the agent loop completes before the answer stream) — so the
    # poll issued right after the click must catch it inside that window.
    expect(page.locator("#send-label")).to_have_text("Calling tool…", timeout=20_000)
    _wait_settled(page)

    # The label transition is also recorded deterministically (no race):
    # Thinking… → Calling tool… → … → Send.
    labels = page.evaluate("() => window.__labels")
    assert "Calling tool…" in labels, labels
    assert labels.index("Calling tool…") > labels.index("Thinking…")

    # Wire level: exactly two `tool` frames — list then read — and both
    # ahead of the first `delta` frame.
    frames = _frames(page)
    assert _tool_frames(frames) == [
        {"type": "tool", "name": "list_documents", "argument": None},
        {"type": "tool", "name": "read_document", "argument": READ_SP},
    ]
    first_delta = next(i for i, f in enumerate(frames) if f.get("type") == "delta")
    assert all(
        i < first_delta for i, f in enumerate(frames) if f.get("type") == "tool"
    )
    done = next(f for f in frames if f.get("type") == "done")
    assert done["deflected"] is False
    assert [(s["source"], s["path"]) for s in done["sources"]] == [
        (SEED_SOURCE, SEED_PATH),
        (READ_SOURCE, READ_PATH),
    ]

    # Both tool lines, in order, above the answer.
    lines = page.locator(".msg.brain .tool-call")
    expect(lines).to_have_count(2)
    expect(lines.nth(0)).to_contain_text("Listing documents")
    expect(lines.nth(1)).to_contain_text("Reading ")
    expect(lines.nth(1)).to_contain_text(READ_SP)

    # The final answer quotes the read document (the mock's deterministic
    # quote: "Read <source/path>. <first 80 chars of its content>").
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(ANSWER_PREFIX)
    expect(bubble).to_contain_text(ANSWER_QUOTE)

    # Source chips: the retrieval doc AND the read doc (deduped, in
    # order) — the read chip links to the viewer.
    chips = page.locator(".msg.brain .source-chip")
    expect(chips).to_have_count(2)
    expect(chips.nth(0)).to_contain_text(SEED_SP)
    chip_read = page.locator(".msg.brain .source-chip", has_text=READ_PATH)
    expect(chip_read).to_have_count(1)
    expect(chip_read.first).to_have_attribute("href", READ_CHIP_HREF)

    # Durable record: grounded, both sources logged (retrieval + read).
    row = _last_query_log()
    assert row.question == MARKER_QUESTION
    assert row.deflected is False
    assert row.sources == f"{SEED_SP}, {READ_SP}"


# --------------------------------------------------------------------------
# 2. Persistence: the tool lines re-render after a reload (phase 14)
# --------------------------------------------------------------------------


def test_tool_lines_re_render_after_reload(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    page.set_default_timeout(30_000)
    _reset_db(_seed)
    page.goto(app_url)

    _submit(page, MARKER_QUESTION)
    _wait_settled(page)
    expect(page.locator(".msg.brain .tool-call")).to_have_count(2)

    page.reload()
    expect(page.locator("#empty-state")).to_be_hidden()

    # The persisted record re-renders BOTH tool lines, in saved order,
    # through the same append helper as the live frames.
    restored = page.locator(".msg.brain .tool-call")
    expect(restored).to_have_count(2)
    expect(restored.nth(0)).to_contain_text("Listing documents")
    expect(restored.nth(1)).to_contain_text("Reading ")
    expect(restored.nth(1)).to_contain_text(READ_SP)

    # Answer + the read-document chip are intact (phase-14 restore path).
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(ANSWER_PREFIX)
    expect(bubble).to_contain_text(ANSWER_QUOTE)
    chip_read = page.locator(".msg.brain .source-chip", has_text=READ_PATH)
    expect(chip_read).to_have_count(1)
    expect(chip_read.first).to_have_attribute("href", READ_CHIP_HREF)


# --------------------------------------------------------------------------
# 3. Regression: a plain grounded question (no marker) takes the
#    no-tool path — the answer renders exactly as today
# --------------------------------------------------------------------------


def test_plain_grounded_question_has_no_tool_frames(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    page.set_default_timeout(30_000)
    _reset_db(_seed)
    page.goto(app_url)
    _install_page_hooks(page)

    _submit(page, PLAIN_QUESTION)
    _wait_settled(page)

    # No tool frames on the wire, no tool lines in the UI.
    assert _tool_frames(_frames(page)) == []
    expect(page.locator(".tool-call")).to_have_count(0)

    # The standard grounded answer, citing the retrieval doc only — the
    # referenced JSON stays OUT of the sources (it was never read).
    bubble = page.locator(".msg.brain .bubble").last
    expect(bubble).to_contain_text(PLAIN_QUESTION)
    expect(bubble).to_contain_text(MOCK_ANSWER_MARKER)
    chips = page.locator(".msg.brain .source-chip")
    expect(chips).to_have_count(1)
    expect(chips.first).to_contain_text(SEED_SP)

    row = _last_query_log()
    assert row.question == PLAIN_QUESTION
    assert row.deflected is False
    assert row.sources == SEED_SP


# --------------------------------------------------------------------------
# 4. Grounded-only scope: a deflected turn runs no tools at all
# --------------------------------------------------------------------------


def test_deflected_question_has_no_tool_frames(
    page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
    page.set_default_timeout(30_000)
    _reset_db(_seed)
    page.goto(app_url)
    _install_page_hooks(page)

    _submit(page, DEFLECT_QUESTION)
    _wait_settled(page)

    # The honesty gate fired — and no tool frames / tool lines came with
    # it (the LOW prompt never carries the tools).
    last = page.locator(".msg.brain").last
    expect(last).to_have_class(re.compile(r"is-deflected"))
    expect(last.locator(".bubble")).to_contain_text(
        re.compile(DEFLECT_PHRASE, re.IGNORECASE)
    )
    assert _tool_frames(_frames(page)) == []
    expect(page.locator(".tool-call")).to_have_count(0)

    row = _last_query_log()
    assert row.question == DEFLECT_QUESTION
    assert row.deflected is True