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

"""Integration: phase 141 (task 03) — the owner's stuck two-range
scenario, proven end to end.

The exact incident from ``TODO.md`` L3: the model was reading
different RANGES of one large document (the D&D "Hail of Thorns"
spell entry) and every second range dead-ended with
``ALREADY_IN_CONTEXT`` — "The system is telling me the document is
already in my context, but I haven't actually seen the full text …
The read calls keep returning 'Already in your context'". Pre-phase,
the pair-set dedupe recorded ANY successful read (a range read
included) as if the FULL document were in context, so every attempt
to read a DIFFERENT range was refused with a false claim. Phase 141
(the coverage rework, tasks 01/02) makes the dedupe decide from the
lines each read actually put in the model's context: a range read
PROCEEDS unless its lines are fully covered, a plain re-read of a
partially-covered doc gets the pinned
:data:`~app.rag.agent.REREAD_PARTIAL` teaching (the covered span(s)
named) instead of the false full-text claim.

This suite locks the fix as two scripted reproductions over ONE
generated text document (200 lines, two distinct marker lines far
apart — the content is GENERATED in the test, no binary fixture: the
untracked owner PDF belongs to the separate phase-131 handbook
repro, untouched). The document is inserted directly through the
house KB pattern (``test_agent_tools.py`` — the ``documents`` table
truncated around the test, so the read resolves against exactly this
doc; the agent's read path only reads ``documents``), and a scripted
fake LLM (the ``tests/unit/test_agent.py`` round-shape pattern)
drives ``run_agent``'s loop against the real DB:

* the **pure dedupe variant** (the doc UNDER the read cap — no
  truncation anywhere): round 1 ``read(path, offset=1, limit=40)``
  (range A), round 2 ``read(path, offset=<marker B's line − 10>,
  limit=30)`` (range B — the PRE-PHASE code refused THIS round with
  ``ALREADY_IN_CONTEXT``), round 3 the EXACT re-read of range A
  (refused — now with a claim that is true: those lines ARE in the
  prompt), round 4 the answer echoing the marker lines the delivered
  tool results carry (the read-echo convention — the model can only
  echo what reached it; pre-phase, the echo could contain only
  marker A).
* the **truncated variant** (the doc OVER the cap — the phase-131
  truncated-then-range row, now via coverage): round 1 a plain
  ``read(path)`` (truncated at a cap DERIVED from the generated
  content — just past marker A's line + a mid-line cut of the next
  line — and it covers only its cut span), round 2 the SAME range-B
  read (PROCEEDS — the truncated first read covered only its first
  lines), round 3 a PLAIN re-read (the generalized teaching —
  :data:`REREAD_PARTIAL` naming the covered span; the retired
  phase-131 char-based ``REREAD_TRUNCATED`` line is gone), round 4
  the echo answer (both markers — the truncated first read carried
  marker A, the range read carried marker B).

The marker A line sits INSIDE range A (lines 1–40) — the task's
"near line 50" hint is unsatisfiable with the pinned round shape
(``offset=1, limit=40``) and the pinned echo assertion (the answer
must echo marker A, and the model can only echo what reached it):
line 30 keeps it ten lines clear of range A's end while the markers
stay far apart (120 lines).

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

import asyncio
import uuid
from collections.abc import AsyncIterator, Iterator
from copy import deepcopy
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, cast

import pytest
from sqlalchemy import text
from sqlalchemy.orm import Session

from app.config import Settings
from app.db import SessionLocal
from app.models import Document
from app.rag import agent
from app.rag.agent import (
    READ_TRUNCATION_NOTICE,
    AgentHolder,
    run_agent,
)
from app.rag.llm import (
    LLMClient,
    RetryPiece,
    StreamPiece,
    ToolCallPiece,
    ToolResultPiece,
)
from app.rag.retriever import TRUNCATION_MARKER

if TYPE_CHECKING:
    from app.rag.scaffolding import ScaffoldingFilter

#: The fixture document's fixed creation date (phase 106, D5) — the
#: ``read`` result's second line formats its UTC date part; a fixed
#: value keeps the pins deterministic.
_FIXTURE_CREATED_AT = datetime(2024, 6, 15, 12, 0, 0, tzinfo=UTC)

#: The document's identity (the house KB pattern: one source, one
#: doc — the owner's spell entry).
SOURCE = "Spellbook"
DOC_PATH = "hail-of-thorns.txt"
IDENTITY = f"{SOURCE}/{DOC_PATH}"

#: The generated document's shape (see :func:`build_two_range_content`).
TOTAL_LINES = 200
#: Range A (round 1/3 of the pure-dedupe variant): lines 1–40.
RANGE_A_OFFSET = 1
RANGE_A_LIMIT = 40
#: Range B (round 2 of both variants): 30 lines starting 10 lines
#: before marker B (``offset = marker B's line − 10``).
RANGE_B_LIMIT = 30
#: The marker lines' positions (marker A inside range A — see the
#: module docstring for why the task's "near line 50" hint cannot
#: stand with the pinned round shape + echo assertion).
MARKER_A_LINE = 30
MARKER_B_LINE = 150

#: The two distinct marker lines (each a complete line of the
#: generated document — the echo collects whole lines).
MARKER_A = "MARKER_A_hail-of-thorns-spell-entry-start"
MARKER_B = "MARKER_B_hail-of-thorns-spell-entry-end"

#: How far past marker A's line the truncated variant's cut runs
#: (see :func:`_truncated_first_cap`): the newline + eight chars of
#: the following line — a MID-LINE cut (the counted-from-truth rule
#: the code uses).
TRUNCATED_FIRST_CUT_TAIL = 8

#: The turn's user message (the owner's question in spirit).
QUESTION = "What does the Hail of Thorns spell entry say?"


def build_two_range_content() -> str:
    """The generated document (no binary fixture): exactly
    :data:`TOTAL_LINES` filler lines with the two marker lines in
    place — marker A on line :data:`MARKER_A_LINE` (inside range A —
    the round-1 slice must carry it for the echo convention to prove
    range A reached the model) and marker B on line
    :data:`MARKER_B_LINE` (the round-2 slice starts 10 lines before
    it)."""
    lines = []
    for i in range(1, TOTAL_LINES + 1):
        if i == MARKER_A_LINE:
            lines.append(MARKER_A)
        elif i == MARKER_B_LINE:
            lines.append(MARKER_B)
        else:
            lines.append(f"line {i:03d} filler filler filler")
    return "\n".join(lines)


def _settings(**kwargs: Any) -> Settings:
    """The turn's settings: ``_env_file=None`` keeps an operator's
    local ``.env`` out of the pins; ``read_max_chars`` at the default
    (the doc fits far under it) unless the test overrides it."""
    kwargs["_env_file"] = None
    return Settings(**kwargs)  # pyright: ignore[reportCallIssue]


def _truncated_first_cap(content: str) -> int:
    """The truncated variant's read cap, DERIVED from the generated
    *content* (the no-hardcoded-offsets rule): just past marker A's
    line — the marker must sit fully INSIDE the cut (the echo of both
    markers needs it in the truncated first read; the marker lines
    are longer than the filler lines, so a literal char cap would
    drift) — plus :data:`TRUNCATED_FIRST_CUT_TAIL` chars of the
    following line (a mid-line cut — the counted-from-truth rule)."""
    return content.index(MARKER_A) + len(MARKER_A) + 1 + TRUNCATED_FIRST_CUT_TAIL


class TwoRangeAgentLLM:
    """The scripted two-range loop (the ``test_agent.py`` round-shape
    pattern) — deterministic; both variants share one class:

    ``truncated_first=False`` (the pure dedupe variant):

    * round 1 → ``read(path, offset=1, limit=40)`` (range A);
    * round 2 → ``read(path, offset=<marker B's line − 10>,
      limit=30)`` (range B — the pre-phase false refusal);
    * round 3 → ``read(path, offset=1, limit=40)`` again (the exact
      re-covered range A);
    * round 4 → the answer: the marker lines the delivered tool
      results carry, in delivery order (the read-echo convention —
      the model can only echo what reached it).

    ``truncated_first=True`` (the truncated variant):

    * round 1 → plain ``read(path)`` (truncates at the cap);
    * round 2 → the same range-B read (proceeds via coverage);
    * round 3 → a plain ``read(path)`` (the ``REREAD_PARTIAL``
      teaching);
    * round 4 → the echo answer.

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

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

    def _next_call(self, index: int) -> ToolCallPiece | None:
        """The scripted call for request *index* — ``None`` on the
        final (echo) round."""
        range_b = {
            "path": self.identity,
            "offset": self.range_b_offset,
            "limit": RANGE_B_LIMIT,
        }
        if not self.truncated_first:
            if index == 0:
                return ToolCallPiece(
                    id="call_1",
                    name="read",
                    arguments={
                        "path": self.identity,
                        "offset": RANGE_A_OFFSET,
                        "limit": RANGE_A_LIMIT,
                    },
                )
            if index == 1:
                return ToolCallPiece(id="call_2", name="read", arguments=range_b)
            if index == 2:
                return ToolCallPiece(
                    id="call_3",
                    name="read",
                    arguments={
                        "path": self.identity,
                        "offset": RANGE_A_OFFSET,
                        "limit": RANGE_A_LIMIT,
                    },
                )
            return None
        if index == 0:
            return ToolCallPiece(id="call_1", name="read", arguments={"path": self.identity})
        if index == 1:
            return ToolCallPiece(id="call_2", name="read", arguments=range_b)
        if index == 2:
            return ToolCallPiece(id="call_3", name="read", arguments={"path": self.identity})
        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)
        if call is not None:
            self.calls.append(call)
            yield call
            return
        # The echo round: the marker lines the DELIVERED tool results
        # carry, in delivery order — the model can only echo what
        # reached it (pre-phase, range B's slice never did, so the
        # echo could contain only marker A).
        marker_lines = [
            line
            for m in messages
            if m.get("role") == "tool" and isinstance(m.get("content"), str)
            for line in m["content"].splitlines()
            if MARKER_A in line or MARKER_B in line
        ]
        content = "\n".join(marker_lines)
        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: TwoRangeAgentLLM,
    holder: AgentHolder,
    settings: Settings,
) -> list[StreamPiece | ToolCallPiece | RetryPiece | ToolResultPiece]:
    """Consume one scripted ``run_agent`` turn. SEC-14-04: a
    short-lived session per tool call — the same pattern as
    production and the phase-70/131 integration suites."""
    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 two_range_doc(db: Session) -> Iterator[Document]:
    """The generated two-range document, indexed as a single row
    (the house KB pattern — ``test_agent_tools.py``: the ``documents``
    table truncated around the test, so the read resolves against
    exactly this doc). Inserted directly (no importer — the content
    is generated in the test, there is no file to import, and the
    agent's read path only reads ``documents``)."""
    db.execute(text("TRUNCATE chunks, documents, folder_summaries"))
    db.commit()
    doc = Document(
        id=uuid.uuid4(),
        source=SOURCE,
        path=DOC_PATH,
        full_path=f"/tmp/{SOURCE}/{DOC_PATH}",
        title="Hail of Thorns",
        content=build_two_range_content(),
        content_hash="0" * 64,
        created_at=_FIXTURE_CREATED_AT,
    )
    db.add(doc)
    db.commit()
    yield doc
    db.execute(text("TRUNCATE chunks, documents, folder_summaries"))
    db.commit()
    db.expire_all()


def test_two_ranges_of_one_doc_both_reach_the_model(two_range_doc: Document) -> None:
    """The owner's stuck loop, fixed — the pure dedupe case (the doc
    UNDER the read cap: no truncation anywhere): range A (lines 1–40)
    → range B (the PRE-PHASE pair-set dedupe refused THIS round with
    ``ALREADY_IN_CONTEXT`` — "the document is already in my context"
    — false: only range A's slice had reached the model) → the EXACT
    re-read of range A (refused — the claim in the byte-identical
    line is now TRUE: those lines ARE in the prompt) → the answer
    echoing BOTH marker lines (range A's slice carried marker A,
    range B's slice carried marker B — the read-echo convention: the
    model can only echo what reached it; pre-phase, the echo could
    contain only marker A)."""
    doc = two_range_doc
    marker_b_line = doc.content.splitlines().index(MARKER_B) + 1
    llm = TwoRangeAgentLLM(
        IDENTITY, truncated_first=False, range_b_offset=marker_b_line - 10
    )
    holder = AgentHolder()
    pieces = asyncio.run(_consume(llm, holder, _settings()))

    # Four rounds: range A → range B → re-read range A → echo answer.
    assert len(llm.requests) == 4
    assert [call.name for call in llm.calls] == ["read", "read", "read"]
    assert llm.calls[0].arguments == {
        "path": IDENTITY,
        "offset": RANGE_A_OFFSET,
        "limit": RANGE_A_LIMIT,
    }
    assert llm.calls[1].arguments == {
        "path": IDENTITY,
        "offset": marker_b_line - 10,
        "limit": RANGE_B_LIMIT,
    }
    assert llm.calls[2].arguments == {
        "path": IDENTITY,
        "offset": RANGE_A_OFFSET,
        "limit": RANGE_A_LIMIT,
    }

    # Round 1: range A's slice — the true extent on the third line +
    # marker A inside.
    range_a_result = llm.requests[1][0][3]
    assert range_a_result["role"] == "tool"
    assert isinstance(range_a_result["content"], str)
    total_lines = len(doc.content.splitlines())
    assert range_a_result["content"].startswith(
        f"Document {IDENTITY}:\ndate: 2024-06-15\nlines: 1-40 of {total_lines}\n"
    )
    assert MARKER_A in range_a_result["content"]

    # Round 2: range B's slice — NOT a refusal (the owner's bug: the
    # pre-phase code refused it with ALREADY_IN_CONTEXT), the true
    # extent + marker B inside.
    range_b_result = llm.requests[2][0][5]
    assert range_b_result["role"] == "tool"
    assert isinstance(range_b_result["content"], str)
    assert range_b_result["content"] != agent.ALREADY_IN_CONTEXT
    assert (
        f"lines: {marker_b_line - 10}-{marker_b_line - 10 + RANGE_B_LIMIT - 1} "
        f"of {total_lines}\n" in range_b_result["content"]
    )
    assert MARKER_B in range_b_result["content"]

    # Round 3: the EXACT re-covered range A — refused (the claim in
    # the byte-identical line is now true).
    reread_result = llm.requests[3][0][7]
    assert reread_result["role"] == "tool"
    assert reread_result["content"].startswith(agent.ALREADY_IN_CONTEXT)

    # Round 4: the answer echoes BOTH markers (the bug: pre-phase,
    # the echo could contain only marker A — range B's slice never
    # reached the model).
    answer = _answer(pieces)
    assert answer == f"{MARKER_A}\n{MARKER_B}"

    # Bookkeeping: the doc committed to the context exactly once; the
    # two executed range reads counted (the refusal counted nothing);
    # no truncation (the doc is under the cap); the coverage spans
    # are the two ranges, in read order.
    assert [d.id for d in holder.read_docs] == [doc.id]
    assert holder.tool_calls == 2
    assert holder.read_truncations == []
    assert holder.read_coverage == [
        (SOURCE, DOC_PATH, RANGE_A_OFFSET, RANGE_A_LIMIT),
        (SOURCE, DOC_PATH, marker_b_line - 10, marker_b_line - 10 + RANGE_B_LIMIT - 1),
    ]


def test_truncated_plain_read_then_range_b_proceeds_and_plain_reread_teaches(
    two_range_doc: Document,
) -> None:
    """The truncated variant (the doc OVER the cap — the phase-131
    truncated-then-range row, now via coverage): the plain read
    truncates at the cap (it covers only its cut span — the count
    from the ACTUAL cut string, the same truth rule the code uses)
    and carries the pinned notice; the range-B read PROCEEDS (the
    truncated first read covered only its first lines — pre-phase the
    row ran through the ``truncated_shown`` map, now through
    ``read_coverage``); the plain re-read gets the generalized
    ``REREAD_PARTIAL`` teaching (BOTH covered spans named in line
    form — the cut span and the already-read range B; the retired
    phase-131 char-based ``REREAD_TRUNCATED`` line is gone); the
    answer echoes both markers (the truncated first read carried
    marker A, the range read carried marker B)."""
    doc = two_range_doc
    marker_b_line = doc.content.splitlines().index(MARKER_B) + 1
    cap = _truncated_first_cap(doc.content)  # derived — see the helper
    covered = len(doc.content[:cap].splitlines())  # the cut's truth
    llm = TwoRangeAgentLLM(
        IDENTITY, truncated_first=True, range_b_offset=marker_b_line - 10
    )
    holder = AgentHolder()
    pieces = asyncio.run(_consume(llm, holder, _settings(read_max_chars=cap)))

    # The cut's shape the assertions below lean on: marker A sits
    # INSIDE the cut (the echo of both markers needs it in the
    # truncated first read), marker B does not (the range-B read is
    # what carries it).
    assert covered >= MARKER_A_LINE and covered < MARKER_B_LINE

    # Four rounds: plain read → range B → plain re-read → echo.
    assert len(llm.requests) == 4
    assert [call.name for call in llm.calls] == ["read", "read", "read"]
    # Round 1: the PLAIN read (no range arguments at all).
    assert llm.calls[0].arguments == {"path": IDENTITY}
    assert llm.calls[1].arguments == {
        "path": IDENTITY,
        "offset": marker_b_line - 10,
        "limit": RANGE_B_LIMIT,
    }
    assert llm.calls[2].arguments == {"path": IDENTITY}

    # Round 1's result: cut at the cap, the shared marker + the
    # pinned notice, marker A inside the cut; the truncation surfaced
    # to the API layer as a ToolResultPiece (the SSE
    # ``tool_result`` frame's wire value).
    plain_result = llm.requests[1][0][3]
    assert plain_result["role"] == "tool"
    assert isinstance(plain_result["content"], str)
    notice = READ_TRUNCATION_NOTICE.format(shown=cap, total=len(doc.content))
    assert plain_result["content"].startswith(f"Document {IDENTITY}:\n")
    # Phase 160 (task 01): the budget-telemetry trailer rides the
    # truncated result — the marker + notice end the RAW part.
    raw = plain_result["content"].removesuffix(agent._budget_note(9))
    assert raw.endswith(f"\n{TRUNCATION_MARKER}\n{notice}")
    assert MARKER_A in plain_result["content"]
    truncations = [p for p in pieces if isinstance(p, ToolResultPiece)]
    assert [(p.argument, p.chars_shown, p.chars_total) for p in truncations] == [
        (IDENTITY, cap, len(doc.content))
    ]

    # Round 2: the range-B read of the ALREADY-TRUNCATED doc
    # PROCEEDS — the slice, not a refusal.
    range_b_result = llm.requests[2][0][5]
    assert range_b_result["role"] == "tool"
    assert isinstance(range_b_result["content"], str)
    assert range_b_result["content"] != agent.ALREADY_IN_CONTEXT
    total_lines = len(doc.content.splitlines())
    assert (
        f"lines: {marker_b_line - 10}-{marker_b_line - 10 + RANGE_B_LIMIT - 1} "
        f"of {total_lines}\n" in range_b_result["content"]
    )
    assert MARKER_B in range_b_result["content"]

    # Round 3: the plain re-read of the PARTIALLY-covered doc — the
    # generalized teaching naming BOTH covered spans in line form
    # (the merged cut span AND the already-read range B — the
    # multi-span rendering, in span order), the true total from the
    # resolved doc.
    reread_result = llm.requests[3][0][7]
    assert reread_result["role"] == "tool"
    assert reread_result["content"].startswith(agent.REREAD_PARTIAL.format(
        covered=f"1–{covered}, {marker_b_line - 10}–{marker_b_line - 10 + RANGE_B_LIMIT - 1}",
        source=SOURCE,
        path=DOC_PATH,
        total=total_lines,
        ))

    # Round 4: the answer echoes both markers.
    answer = _answer(pieces)
    assert answer == f"{MARKER_A}\n{MARKER_B}"

    # Bookkeeping: the doc once; the two successful reads counted (the
    # teaching refusal counted nothing); the truncation record is the
    # plain read's tuple; the coverage spans — the cut span, then the
    # range — in read order.
    assert [d.id for d in holder.read_docs] == [doc.id]
    assert holder.tool_calls == 2
    assert holder.read_truncations == [(IDENTITY, cap, len(doc.content))]
    assert holder.read_coverage == [
        (SOURCE, DOC_PATH, 1, covered),
        (SOURCE, DOC_PATH, marker_b_line - 10, marker_b_line - 10 + RANGE_B_LIMIT - 1),
    ]