brain-of-reese/tests/integration/test_read_line_ranges_handbook.py

"""Integration: phase 131 (task 03) — the owner's handbook repro.

The exact incident from ``TODO.md`` L7: the agent was asked "What are
the components of the find traps spell?" over the real D&D player's
handbook (``test_sources/players_handbook.pdf`` — the untracked fixture
from ``pdf_tests.tar``), it ``grep``'d the text and FOUND it "including
line number", and then dead-ended: "However, read takes a path, not a
line number." This suite proves the fixed loop closes the dance on the
REAL document: the handbook is indexed through the REAL importer
(:func:`app.rag.importer.import_sources` — the deterministic in-process
:class:`~tests.fakes.FakeEmbedder` stands in for the embed/summary
calls, so the ~1200-chunk index build stays fast and offline), and a
scripted fake LLM (the ``tests/unit/test_agent.py`` round pattern)
drives the turn with the owner's ``BOR_READ_MAX_CHARS=24000`` in force:

* the **grep-first** repro — round 1 a ``grep`` (``pattern="find
  traps"``, lowercase — grep is case-insensitive), round 2 a ``read``
  with ``offset`` = the line the REAL grep result reported (parsed,
  never hardcoded) + ``limit`` = 40, round 3 the answer: the read
  result's content, verbatim (the phase-94 "the read document reached
  the context" echo convention — the loop's reply IS the slice).
  Exactly 3 rounds, ``tool_calls == 2`` (the owner's stuck loop made 2
  and dead-ended — the fixed loop closes in 3).
* the **plain-read-first** variant — round 1 a plain ``read`` (the
  handbook truncates at 24000 with the REWRITTEN phase-131 notice,
  which names ``offset``/``limit`` — the dead-end is gone), round 2
  the grep, round 3 the range read of the ALREADY-TRUNCATED document
  (the phase-141 COVERAGE dedupe matrix's truncated-then-range row,
  end-to-end on the real document — it PROCEEDS; the
  ``REREAD_PARTIAL`` teaching — the phase-131 ``REREAD_TRUNCATED``
  line, retired phase 141 — is never needed because the model went
  straight to the range), round 4 the echo answer.

The expected find-traps text is DERIVED from the fixture at test time
(extract the PDF's text with the phase-126
:func:`app.rag.pdf.extract_pdf_text`, grep it the same way the model
does — :func:`app.rag.agent.grep_document`'s case-insensitive fixed
substring — and take the passage's first 60 chars): no hardcoded
handbook wording, so the test stays true if the fixture is
re-extracted.

House KB pattern (``test_agent_tools.py`` / ``test_importer_e2e.py``):
the ``documents`` table is truncated around each test, so the whole-KB
``grep`` sees EXACTLY the handbook — the grep's line numbers are then
the extracted text's line numbers, one-to-one. The untracked fixture
is only ever COPIED into a temp source dir, never mutated.

Requires:  podman compose up -d db
"""
from __future__ import annotations

import asyncio
import re
import shutil
from collections.abc import AsyncIterator, Iterator
from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast

import pytest
from sqlalchemy import delete, select, text
from sqlalchemy.orm import Session

from app.config import Settings
from app.db import SessionLocal
from app.models import Document, GitSource
from app.rag import agent
from app.rag.agent import (
    AGENT_TOOLS,
    READ_TRUNCATION_NOTICE,
    REREAD_PARTIAL,
    AgentHolder,
    run_agent,
)
from app.rag.importer import import_sources
from app.rag.llm import (
    LLMClient,
    RetryPiece,
    StreamPiece,
    ToolCallPiece,
    ToolResultPiece,
)
from app.rag.pdf import extract_pdf_text
from app.rag.retriever import TRUNCATION_MARKER
from tests.fakes import FakeEmbedder

if TYPE_CHECKING:
    from app.rag.scaffolding import ScaffoldingFilter

#: The repo root (this file lives in ``tests/integration/``).
REPO_ROOT = Path(__file__).resolve().parents[2]

#: The owner's real fixture (untracked — unpacked from
#: ``test_sources/pdf_tests.tar``): the D&D player's handbook.
HANDBOOK = REPO_ROOT / "test_sources" / "players_handbook.pdf"

#: The imported identity: the source dir is named ``test_sources``
#: (the importer's ``source = root.name``) and the PDF sits directly
#: inside it.
SOURCE = "test_sources"
PDF_NAME = "players_handbook.pdf"
IDENTITY = f"{SOURCE}/{PDF_NAME}"

#: The owner's exact question (``TODO.md`` L7) — the turn's user
#: message; the scripted loop answers it through the tools.
QUESTION = "What are the components of the find traps spell?"

#: The owner's env: ``BOR_READ_MAX_CHARS=24000`` (the handbook is far
#: longer, so a plain read of it truncates).
READ_CAP = 24_000

#: The owner's pattern — lowercase (grep is case-insensitive).
PATTERN = "find traps"

#: The scripted range read's line count (both variants).
RANGE_LIMIT = 40

pytestmark = pytest.mark.skipif(
    not HANDBOOK.is_file(),
    reason=(
        "test_sources/players_handbook.pdf is missing — unpack "
        "test_sources/pdf_tests.tar (the untracked owner fixture)"
    ),
)


def _settings() -> Settings:
    """The turn's settings: the owner's read cap in force and the code
    default import extensions plus ``pdf`` (the owner's
    ``BOR_IMPORT_EXTENSIONS`` names it — the handbook is an indexed
    source). ``_env_file=None`` keeps an operator's local ``.env`` out
    of the pins."""
    defaults = Settings.model_fields["import_extensions"].default
    kwargs: dict[str, Any] = {
        "_env_file": None,
        "import_extensions": f"{defaults},pdf",
        "read_max_chars": READ_CAP,
    }
    return Settings(**kwargs)  # pyright: ignore[reportCallIssue]


def _expected_passage() -> str:
    """The first 60 chars of the handbook's find-traps passage, DERIVED
    from the fixture at test time (the task's no-hardcoded-wording
    rule): the PDF's extracted text (the same text the importer stored)
    is grepped the same way the model does —
    :func:`app.rag.agent.grep_document`, the tool's case-insensitive
    fixed-substring semantics — and the passage is the
    :data:`RANGE_LIMIT` lines from the FIRST grep match (the exact
    window the scripted loop reads). If the fixture is re-extracted the
    line numbers and wording are re-derived, so the assertion stays
    true."""
    content = extract_pdf_text(HANDBOOK)
    first_line, _ = agent.grep_document(content, PATTERN)[0]
    lines = content.splitlines()
    return "\n".join(lines[first_line - 1 : first_line - 1 + RANGE_LIMIT])[:60]


def _first_grep_line(messages: list[dict[str, Any]], identity: str) -> int:
    """The first line number the loop's grep reported for *identity* —
    parsed from the REAL grep result the loop produced (the task's
    no-hardcoded-line rule). The last ``tool`` message carries the
    grep result, one ``source/path:line: text`` line per match; the
    first match line for the handbook names the line the scripted
    range read takes."""
    last_tool = next(m for m in reversed(messages) if m.get("role") == "tool")
    content = last_tool["content"]
    assert isinstance(content, str)
    for line in content.splitlines():
        if line.startswith(f"{identity}:"):
            match = re.match(rf"^{re.escape(identity)}:(\d+):", line)
            assert match is not None, f"unparseable grep line: {line!r}"
            return int(match.group(1))
    raise AssertionError(f"no grep result line for {identity!r}")


class HandbookAgentLLM:
    """The scripted owner-scenario loop (the ``test_agent.py`` round
    pattern) — deterministic, and it parses the REAL line number from
    the grep result the loop produced (never a hardcoded line):

    ``plain_first=False`` (the grep-first repro — test 1):

    * round 1 → ``grep`` (``pattern="find traps"``);
    * round 2 → ``read`` (``path`` = the handbook identity, ``offset``
      = the first grep-reported line, ``limit`` = 40);
    * round 3 → the answer: the read result's content, verbatim (the
      phase-94 echo convention — the loop's reply IS the slice).

    ``plain_first=True`` (the plain-read-first variant — test 2):

    * round 1 → plain ``read`` (``path``);
    * round 2 → ``grep`` (``pattern="find traps"``);
    * round 3 → ``read`` (``path``, ``offset`` = the first grep-
      reported line, ``limit`` = 40) — the coverage dedupe matrix's
      truncated-then-range row on the real document;
    * round 4 → the answer: the read result's content, verbatim.

    Every ``chat_stream`` request's messages and tools are recorded
    (``requests``); the scripted :class:`ToolCallPiece` s, in call
    order, in ``calls``. The canned answer pieces are fed through the
    round's :class:`ScaffoldingFilter` exactly like
    ``LLMClient.chat_stream`` (the ``test_agent.py`` ScriptedLLM
    pattern)."""

    def __init__(self, identity: str, *, plain_first: bool) -> None:
        self.identity = identity
        self.plain_first = plain_first
        self.requests: list[
            tuple[list[dict[str, Any]], list[dict[str, Any]] | None]
        ] = []
        self.calls: list[ToolCallPiece] = []

    def _next_call(self, index: int, messages: list[dict[str, Any]]) -> ToolCallPiece | None:
        """The scripted call for request *index* — ``None`` on the
        final (answer) round. The range read's ``offset`` is parsed
        from the grep result just delivered."""
        if self.plain_first:
            if index == 0:
                return ToolCallPiece(id="call_1", name="read", arguments={"path": self.identity})
            if index == 1:
                return ToolCallPiece(id="call_2", name="grep", arguments={"pattern": PATTERN})
            if index == 2:
                return ToolCallPiece(
                    id="call_3",
                    name="read",
                    arguments={
                        "path": self.identity,
                        "offset": _first_grep_line(messages, self.identity),
                        "limit": RANGE_LIMIT,
                    },
                )
        else:
            if index == 0:
                return ToolCallPiece(id="call_1", name="grep", arguments={"pattern": PATTERN})
            if index == 1:
                return ToolCallPiece(
                    id="call_2",
                    name="read",
                    arguments={
                        "path": self.identity,
                        "offset": _first_grep_line(messages, self.identity),
                        "limit": RANGE_LIMIT,
                    },
                )
        return None

    async def chat_stream(
        self,
        messages: list[dict[str, str]],
        tools: list[dict[str, Any]] | None = None,
        scaffolding: ScaffoldingFilter | None = None,
    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
        self.requests.append((deepcopy(messages), deepcopy(tools)))
        call = self._next_call(len(self.requests) - 1, messages)
        if call is not None:
            self.calls.append(call)
            yield call
            return
        # The answer round: the last tool result (the range read's
        # result) — the loop's reply IS the slice, verbatim.
        last_tool = next(m for m in reversed(messages) if m.get("role") == "tool")
        content = last_tool["content"]
        assert isinstance(content, str)
        if scaffolding is None:
            yield StreamPiece("content", content)
            return
        cleaned = scaffolding.feed(content)
        if cleaned:
            yield StreamPiece("content", cleaned)
        tail = scaffolding.flush()
        if tail:
            yield StreamPiece("content", tail)


async def _consume(
    llm: HandbookAgentLLM,
    holder: AgentHolder,
    settings: Settings,
) -> list[StreamPiece | ToolCallPiece | RetryPiece | ToolResultPiece]:
    """Consume one scripted ``run_agent`` turn (the owner's question,
    no seed docs). SEC-14-04: a short-lived session per tool call —
    the same pattern as production and the phase-70 integration suite."""
    out: list[StreamPiece | ToolCallPiece | RetryPiece | ToolResultPiece] = []
    async for piece in run_agent(
        cast("LLMClient", llm),
        lambda: SessionLocal(),
        system_prompt="SYSTEM_PROMPT",
        user_message=QUESTION,
        seed_docs=[],
        settings=settings,
        holder=holder,
    ):
        out.append(piece)
    return out


def _answer(pieces: list[StreamPiece | ToolCallPiece | RetryPiece | ToolResultPiece]) -> str:
    """The turn's final answer text (the content pieces)."""
    return "".join(
        piece.text for piece in pieces if isinstance(piece, StreamPiece) and piece.kind == "content"
    )


@pytest.fixture()
def handbook_kb(db: Session, tmp_path: Path) -> Iterator[Document]:
    """The real handbook, indexed once per test through the REAL
    importer: a temp source dir containing ``test_sources/
    players_handbook.pdf`` (a COPY — the untracked fixture is never
    mutated), registered as a local source, imported with the
    deterministic in-process :class:`~tests.fakes.FakeEmbedder` (the
    ~1200-chunk embed is instant — the whole test stays well under the
    task's ~60 s bar). Yields the indexed
    :class:`~app.models.Document` row.

    The house KB pattern: ``chunks``/``documents``/``folder_summaries``
    are truncated around the test, so the whole-KB ``grep`` sees
    EXACTLY the handbook (its line numbers are the extracted text's
    line numbers, one-to-one)."""
    db.execute(text("TRUNCATE chunks, documents, folder_summaries"))
    db.commit()
    root = tmp_path / SOURCE
    root.mkdir()
    shutil.copy2(HANDBOOK, root / PDF_NAME)
    row = GitSource(url=str(root), kind="local", path=str(root))
    db.add(row)
    db.commit()
    fake = FakeEmbedder(dim=768)
    fake.settings = _settings()  # the owner's cap + the pdf token
    asyncio.run(import_sources([root], fake, session=db))
    doc = db.scalar(
        select(Document).where(Document.source == SOURCE, Document.path == PDF_NAME)
    )
    assert doc is not None, "the handbook did not import"
    assert len(doc.content) > READ_CAP, "the handbook must outgrow the owner's read cap"
    yield doc
    db.execute(text("TRUNCATE chunks, documents, folder_summaries"))
    db.execute(delete(GitSource).where(GitSource.id == row.id))
    db.commit()
    db.expire_all()


def test_grep_then_range_read_answers_the_find_traps_question(
    handbook_kb: Document,
) -> None:
    """The owner's stuck loop, fixed — grep → range read → answer:
    exactly 3 rounds with ``tool_calls == 2`` (the owner's stuck loop
    made 2 and dead-ended), the read call carries the grep-reported
    ``offset`` + ``limit`` (parsed, not literal), and the final answer
    — the read result, verbatim — contains the handbook's REAL
    find-traps text (derived from the fixture at test time)."""
    doc = handbook_kb
    expected = _expected_passage()

    llm = HandbookAgentLLM(IDENTITY, plain_first=False)
    holder = AgentHolder()
    pieces = asyncio.run(_consume(llm, holder, _settings()))

    # Exactly 3 rounds: grep → range read → answer (the fixed loop
    # closes where the owner's stuck loop dead-ended).
    assert len(llm.requests) == 3
    assert [call.name for call in llm.calls] == ["grep", "read"]
    # The tools stay offered on every request (phase 45 — no budgets).
    assert llm.requests[0][1] == AGENT_TOOLS
    assert llm.requests[2][1] == AGENT_TOOLS

    # Round 1: the lowercase pattern (grep is case-insensitive).
    assert llm.calls[0].arguments == {"pattern": PATTERN}
    # The grep result the loop produced: the handbook's REAL match
    # lines (``source/path:line: text``), the first naming the line
    # the range read must take.
    grep_result = llm.requests[1][0][-1]
    assert grep_result["role"] == "tool"
    assert isinstance(grep_result["content"], str)
    assert grep_result["content"].startswith(f"{IDENTITY}:")
    grep_line = _first_grep_line(llm.requests[1][0], IDENTITY)

    # Round 2: the read call carried offset/limit — offset is the
    # grep-reported line (parsed from the real grep result, never a
    # literal), limit 40.
    read_call = llm.calls[1]
    assert read_call.arguments == {
        "path": IDENTITY,
        "offset": grep_line,
        "limit": RANGE_LIMIT,
    }

    # The range read's result: the byte-identical header lines, the
    # true extent on the third line, the slice below.
    read_result = llm.requests[2][0][-1]
    assert read_result["role"] == "tool"
    assert isinstance(read_result["content"], str)
    total_lines = len(doc.content.splitlines())
    assert read_result["content"].startswith(f"Document {IDENTITY}:\n")
    assert (
        f"lines: {grep_line}-{grep_line + RANGE_LIMIT - 1} of {total_lines}\n"
        in read_result["content"]
    )

    # Round 3: the answer IS the read result (the phase-94 echo
    # convention) and it carries the REAL find-traps text.
    answer = _answer(pieces)
    assert answer == read_result["content"]
    assert expected in answer

    # Executed-call bookkeeping: both calls counted; the doc read into
    # the context exactly once.
    assert holder.tool_calls == 2
    assert [d.id for d in holder.read_docs] == [doc.id]
    assert holder.read_truncations == []  # the range slice fits the cap


def test_plain_read_truncates_then_range_read_answers_the_find_traps_question(
    handbook_kb: Document,
) -> None:
    """The plain-read-first variant on the REAL document: round 1's
    plain read truncates at the owner's 24000 cap with the REWRITTEN
    phase-131 notice (it names ``offset``/``limit`` — the dead-end is
    gone) and surfaces a truncation signal; round 2 greps; round 3's
    range read of the ALREADY-TRUNCATED document PROCEEDS (the dedupe
    matrix's truncated-then-range row, end-to-end — the
    ``REREAD_TRUNCATED`` teaching is never needed because the model
    went straight to the range); round 4 answers with the slice."""
    doc = handbook_kb
    expected = _expected_passage()

    llm = HandbookAgentLLM(IDENTITY, plain_first=True)
    holder = AgentHolder()
    pieces = asyncio.run(_consume(llm, holder, _settings()))

    # Four rounds: plain read → grep → range read → answer.
    assert len(llm.requests) == 4
    assert [call.name for call in llm.calls] == ["read", "grep", "read"]
    # Round 1: the PLAIN read (no range arguments at all).
    assert llm.calls[0].arguments == {"path": IDENTITY}

    # Round 1's result: cut at the owner's cap, the shared marker +
    # the REWRITTEN notice (``{total}`` = the handbook's true length,
    # ``{shown}`` = the cap) — the notice names the offset/limit
    # workflow the loop then follows.
    plain_result = llm.requests[1][0][-1]
    assert plain_result["role"] == "tool"
    assert isinstance(plain_result["content"], str)
    notice = READ_TRUNCATION_NOTICE.format(shown=READ_CAP, total=len(doc.content))
    assert plain_result["content"].startswith(f"Document {IDENTITY}:\n")
    assert plain_result["content"].endswith(f"\n{TRUNCATION_MARKER}\n{notice}")
    # The truncation surfaced to the API layer as a ToolResultPiece
    # (the SSE ``tool_result`` frame's wire value).
    truncations = [p for p in pieces if isinstance(p, ToolResultPiece)]
    assert [(p.argument, p.chars_shown, p.chars_total) for p in truncations] == [
        (IDENTITY, READ_CAP, len(doc.content))
    ]

    # Round 3: the range read of the truncated doc — offset is the
    # grep-reported line (parsed), limit 40 — and the coverage dedupe
    # matrix lets it PROCEED (no ``REREAD_PARTIAL`` teaching anywhere
    # in the loop's tool results — the partial-teaching line, re-pinned
    # from the retired phase-131 ``REREAD_TRUNCATED``, is derived from
    # the REAL document: the truncated plain read covers line 1
    # through the last line inside the cut, the total is the true line
    # count).
    assert llm.calls[1].arguments == {"pattern": PATTERN}
    grep_line = _first_grep_line(llm.requests[2][0], IDENTITY)
    read_call = llm.calls[2]
    assert read_call.arguments == {
        "path": IDENTITY,
        "offset": grep_line,
        "limit": RANGE_LIMIT,
    }
    covered_lines = len(doc.content[:READ_CAP].splitlines())
    reread = REREAD_PARTIAL.format(
        covered=f"1–{covered_lines}",
        source=SOURCE,
        path=PDF_NAME,
        total=len(doc.content.splitlines()),
    )
    tool_results = [
        m["content"]
        for messages, _tools in llm.requests
        for m in messages
        if m.get("role") == "tool" and isinstance(m["content"], str)
    ]
    assert reread not in tool_results
    # The range read's result: the true extent on the third line.
    read_result = llm.requests[3][0][-1]
    assert read_result["role"] == "tool"
    assert isinstance(read_result["content"], str)
    total_lines = len(doc.content.splitlines())
    assert read_result["content"].startswith(f"Document {IDENTITY}:\n")
    assert (
        f"lines: {grep_line}-{grep_line + RANGE_LIMIT - 1} of {total_lines}\n"
        in read_result["content"]
    )

    # Round 4: the answer IS the range read's result and it carries
    # the REAL find-traps text.
    answer = _answer(pieces)
    assert answer == read_result["content"]
    assert expected in answer

    # All three calls executed (the truncated plain read, the grep,
    # the range read); the doc is committed to the context exactly
    # once; the holder's truncation record is the plain read's tuple.
    assert holder.tool_calls == 3
    assert [d.id for d in holder.read_docs] == [doc.id]
    assert holder.read_truncations == [(IDENTITY, READ_CAP, len(doc.content))]