"""Unit: the grounded-turn agent loop (phase 37, ``app.rag.agent``). A scripted fake LLM (canned stream sequences) + monkeypatched ``list_catalog`` / ``find_document`` — no database, no network. Covers the loop mechanics: the list → read → answer happy path (event order, holder state, the tools staying offered on every request — phase 45 removed the per-tool budgets, the assistant/tool message history), the kill switch (``agent_max_rounds=0`` single-call path), the round cap forcing a final no-tools answer (an always-calling stream and an always-rejected stream), re-lists and multi-reads executing without budgets, dedupe, unknown tool / missing args / unknown path, the ```` prompt section (HIGH only), and the phase-67 per-round retries (a dead-then-recovered round restarts before its first piece with a ``RetryPiece``; a mid-stream drop stays terminal — locked A2; the forced final no-tools call retries too; ``llm_retries=0`` is one plain attempt; retries are invisible to the round cap; consumer abandon mid-retry-sleep leaks nothing). """ from __future__ import annotations import asyncio import json import logging import uuid from collections.abc import AsyncGenerator, AsyncIterator from copy import deepcopy from typing import Any, cast import pytest from sqlalchemy.orm import Session from app.config import Settings from app.models import Document from app.rag import agent from app.rag.agent import ( AGENT_TOOLS, AgentHolder, run_agent, ) from app.rag.llm import LLMClient, LLMError, RetryPiece, StreamPiece, ToolCallPiece from app.rag.prompts import TOOLS_SECTION, _base, build_deflect_prompt, build_high_prompt def _settings(**kwargs: Any) -> Settings: kwargs.setdefault("_env_file", None) return Settings(**kwargs) # pyright: ignore[reportCallIssue] def _doc(source: str, path: str, title: str = "Title", content: str = "CONTENT") -> Document: return Document( id=uuid.uuid4(), source=source, path=path, full_path=f"/tmp/{path}", title=title, content=content, content_hash="0" * 64, ) class ScriptedLLM: """Canned stream sequences; records every ``chat_stream`` request so the tests can assert on the messages and the ``tools`` passthrough.""" def __init__(self, *streams: list[StreamPiece | ToolCallPiece]) -> None: self.streams: list[list[StreamPiece | ToolCallPiece]] = list(streams) self.requests: list[tuple[list[dict[str, Any]], list[dict[str, Any]] | None]] = [] async def chat_stream( self, messages: list[dict[str, str]], tools: list[dict[str, Any]] | None = None, ) -> AsyncIterator[StreamPiece | ToolCallPiece]: self.requests.append((deepcopy(messages), tools)) if not self.streams: raise AssertionError("ScriptedLLM ran out of canned streams") for piece in self.streams.pop(0): yield piece async def _run( llm: ScriptedLLM | FailingLLM, holder: AgentHolder, settings: Settings, seed_docs: list[Document] | None = None, ) -> list[StreamPiece | ToolCallPiece | RetryPiece]: out: list[StreamPiece | ToolCallPiece | RetryPiece] = [] async for piece in run_agent( cast("LLMClient", llm), cast("Session", None), system_prompt="SYSTEM_PROMPT", user_message="QUESTION", seed_docs=seed_docs or [], settings=settings, holder=holder, ): out.append(piece) return out # ---------- AGENT_TOOLS shape ---------- def test_agent_tools_names_and_parameters() -> None: by_name = {t["function"]["name"]: t for t in AGENT_TOOLS} assert len(AGENT_TOOLS) == 3 # list / read / search (phase 68) assert set(by_name) == {"list_documents", "read_document", "search_documents"} assert all(t["type"] == "function" for t in AGENT_TOOLS) list_params = by_name["list_documents"]["function"]["parameters"] assert list_params["type"] == "object" assert list_params["properties"] == {} # no parameters read_params = by_name["read_document"]["function"]["parameters"] assert read_params["required"] == ["source", "path"] assert set(read_params["properties"]) == {"source", "path"} # The model repeatedly conflated the two fields — passing the # combined 'source/path' string as 'source' — so the read_document # description pins the split rule with a worked example. assert by_name["read_document"]["function"]["description"] == ( "Add the full content of one more indexed document to your " "context. A document is identified by the (source, path) pair " "exactly as shown in the list_documents output: 'source' is " "the top-level source name only (e.g. 'homelab'), 'path' is " "the file path inside that source (e.g. " "'active/container_caddy/caddy.md'). If you only have a " "combined 'source/path' string (as in search_documents " "results), split it at the FIRST '/': the part before is the " "source, the part after is the path. Example: " "read_document(source='homelab', " "path='active/container_caddy/caddy.md')." ) # The parameter descriptions define the split: source = before the # first '/', path = after it. assert read_params["properties"]["source"]["description"] == ( "Top-level source name only (e.g. 'homelab') — the part BEFORE " "the first '/' of a combined 'source/path' string, exactly as " "shown after 'source: ' in the list_documents output. Must not " "contain '/' itself — do not pass the full source/path here." ) assert read_params["properties"]["path"]["description"] == ( "File path relative to the source directory (e.g. " "'active/container_caddy/caddy.md') — the part AFTER the first " "'/' of a combined 'source/path' string, exactly as shown " "after 'path: ' in the list_documents output. Must not start " "with the source name." ) # Phase 68: search_documents — the third tool, a locator (locked # A5); its description maps result lines back onto the split. search = by_name["search_documents"]["function"] assert search["description"] == ( "Search every indexed document for an exact string " "(case-insensitive) and return up to 20 matching lines as " "'source/path:line: text' — use this to locate content, " "then read_document the winner (each result line's " "'source/path' splits at the first '/': the part before " "is the source, the part after is the path). Optionally " "pass 'source' and 'path' (as shown in list_documents) " "to search one document only." ) search_params = search["parameters"] assert search_params["type"] == "object" assert search_params["required"] == ["pattern"] assert set(search_params["properties"]) == {"pattern", "source", "path"} assert search_params["properties"]["pattern"]["description"] == ( "The exact text to search for (a plain substring, not a regex)" ) # Shared constants: search's source/path params ARE read_document's # (one definition, no drift between the two tools). assert search_params["properties"]["source"] is read_params["properties"]["source"] assert search_params["properties"]["path"] is read_params["properties"]["path"] # ---------- happy path: list → read → answer ---------- def test_list_then_read_then_answer( monkeypatch: pytest.MonkeyPatch, ) -> None: catalog = [ ("Deployments", "backups.md", "Backup Strategy"), ("Homelab", "aws-route53.md", "AWS Route53 Records"), ] monkeypatch.setattr(agent, "list_catalog", lambda db: catalog) target = _doc("Homelab", "aws-route53.md", "AWS Route53 Records", "R53-CONTENT") monkeypatch.setattr(agent, "find_document", lambda db, source, path: target) seed = [_doc("Homelab", "kubernetes.md", "Kubernetes", "K8S-CONTENT")] holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="list_documents", arguments={})], [ ToolCallPiece( id="call_2", name="read_document", arguments={"source": "Homelab", "path": "aws-route53.md"}, ) ], [StreamPiece("thinking", "hmm "), StreamPiece("content", "Done! ")], ) pieces = asyncio.run(_run(llm, holder, _settings(), seed_docs=seed)) # Event order: tool pieces before the answer content/thinking. assert [type(p) for p in pieces] == [ ToolCallPiece, ToolCallPiece, StreamPiece, StreamPiece, ] assert pieces[0] == ToolCallPiece(id="call_1", name="list_documents", arguments={}) assert isinstance(pieces[1], ToolCallPiece) assert pieces[1].name == "read_document" assert pieces[3] == StreamPiece("content", "Done! ") # The read document is recorded for done.sources / query_log (task 04). assert holder.read_docs == [target] assert holder.tool_calls == 2 # Phase 45: no per-tool budgets — the tools stay offered on every # request (the round cap, not spent budgets, bounds the loop), so # the answer request still carries them (2 rounds < default cap 10). assert llm.requests[0][1] == AGENT_TOOLS assert llm.requests[1][1] == AGENT_TOOLS assert llm.requests[2][1] == AGENT_TOOLS assert len(llm.requests) == 3 # The follow-up request carries the assistant tool-call + tool result. msgs = llm.requests[1][0] assert msgs[0] == {"role": "system", "content": "SYSTEM_PROMPT"} assert msgs[1] == {"role": "user", "content": "QUESTION"} assert msgs[2] == { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_1", "type": "function", "function": {"name": "list_documents", "arguments": "{}"}, } ], } assert msgs[3] == { "role": "tool", "tool_call_id": "call_1", "content": ( "2 documents:\n" "source: Deployments | path: backups.md | title: Backup Strategy\n" "source: Homelab | path: aws-route53.md | title: AWS Route53 Records" ), } # The second follow-up request carries the read call + the FULL text. msgs = llm.requests[2][0] assert msgs[4]["role"] == "assistant" assert msgs[4]["tool_calls"][0]["id"] == "call_2" assert json.loads(msgs[4]["tool_calls"][0]["function"]["arguments"]) == { "source": "Homelab", "path": "aws-route53.md", } assert msgs[5] == { "role": "tool", "tool_call_id": "call_2", "content": "Document Homelab/aws-route53.md:\nR53-CONTENT", # full text, no cap } def test_empty_catalog_listing_says_zero_documents(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(agent, "list_catalog", lambda db: []) holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="list_documents", arguments={})], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == "0 documents:\n" assert holder.tool_calls == 1 def test_content_and_tool_call_in_one_stream_keeps_both( monkeypatch: pytest.MonkeyPatch, ) -> None: """Rare stream with content AND a tool call: the content stays (it was already emitted) and the tool still runs.""" monkeypatch.setattr(agent, "list_catalog", lambda db: []) holder = AgentHolder() llm = ScriptedLLM( [ StreamPiece("content", "Let me check "), ToolCallPiece(id="call_1", name="list_documents", arguments={}), ], [StreamPiece("content", "the answer")], ) pieces = asyncio.run(_run(llm, holder, _settings())) assert [type(p) for p in pieces] == [StreamPiece, ToolCallPiece, StreamPiece] assert holder.tool_calls == 1 # the tool ran despite the content assert llm.requests[1][0][3]["content"] == "0 documents:\n" # ---------- round cap (phase 45: replaces the per-tool budgets) ---------- def test_always_list_bounded_by_round_cap(monkeypatch: pytest.MonkeyPatch) -> None: """A model that keeps calling ``list_documents`` gets exactly ``agent_max_rounds`` tool rounds, then one forced ``tools=None`` request streams the answer — the cap is the only forced exit.""" monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")]) listing = "1 documents:\nsource: S | path: a.md | title: A" holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="list_documents", arguments={})], [ToolCallPiece(id="call_2", name="list_documents", arguments={})], [ToolCallPiece(id="call_3", name="list_documents", arguments={})], [StreamPiece("content", "forced answer")], ) pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=3))) assert [type(p) for p in pieces] == [ ToolCallPiece, ToolCallPiece, ToolCallPiece, StreamPiece, ] assert len(llm.requests) == 4 # 3 tool rounds + the forced answer # The three tool rounds were offered the tools… assert llm.requests[0][1] == AGENT_TOOLS assert llm.requests[1][1] == AGENT_TOOLS assert llm.requests[2][1] == AGENT_TOOLS # …and the forced final request carries no tools, whatever is left. assert llm.requests[3][1] is None # Every re-list executed and counted. assert holder.tool_calls == 3 # The final request carries all three executed listings as history. final_msgs = llm.requests[3][0] assert len(final_msgs) == 8 # 2 + 3 rounds × (assistant + tool) assert final_msgs[3]["content"] == listing assert final_msgs[5]["content"] == listing assert final_msgs[7]["content"] == listing def test_zero_max_rounds_is_one_request_without_tools() -> None: """``agent_max_rounds=0`` — the kill switch: exactly one request, ``tools=None``, no tool lines, no history growth (byte-identical to the pre-phase-37 path).""" holder = AgentHolder() llm = ScriptedLLM([StreamPiece("thinking", "t "), StreamPiece("content", "direct answer")]) pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=0))) assert [type(p) for p in pieces] == [StreamPiece, StreamPiece] assert len(llm.requests) == 1 assert llm.requests[0][1] is None assert llm.requests[0][0] == [ {"role": "system", "content": "SYSTEM_PROMPT"}, {"role": "user", "content": "QUESTION"}, ] assert holder.read_docs == [] and holder.tool_calls == 0 def test_rejected_read_spam_runs_to_round_cap( monkeypatch: pytest.MonkeyPatch, ) -> None: """Every call rejected (unknown path — "No document at …"): rejections no longer end the loop early via budgets — the round cap bounds them and forces the final no-tools answer.""" monkeypatch.setattr(agent, "find_document", lambda db, source, path: None) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S", "path": "ghost.md"}, ) ], [ ToolCallPiece( id="call_2", name="read_document", arguments={"source": "S", "path": "ghost.md"}, ) ], [ ToolCallPiece( id="call_3", name="read_document", arguments={"source": "S", "path": "ghost.md"}, ) ], [StreamPiece("content", "forced answer")], ) asyncio.run(_run(llm, holder, _settings(agent_max_rounds=3))) assert len(llm.requests) == 4 # 3 rejected rounds + the forced answer assert llm.requests[0][1] == AGENT_TOOLS assert llm.requests[1][1] == AGENT_TOOLS assert llm.requests[2][1] == AGENT_TOOLS assert llm.requests[3][1] is None # the forced final request: no tools assert holder.read_docs == [] and holder.tool_calls == 0 # nothing executed refusal = "No document at S/ghost.md — check the list_documents output." assert llm.requests[1][0][3]["content"] == refusal assert llm.requests[2][0][5]["content"] == refusal assert llm.requests[3][0][7]["content"] == refusal # ---------- unlimited calls: re-lists and multi-reads (phase 45) ---------- def test_relist_executes_and_counts(monkeypatch: pytest.MonkeyPatch) -> None: """Re-lists execute — a second ``list_documents`` in one turn returns the catalog again and counts in ``tool_calls`` (no budget to exhaust).""" catalog = [ ("Deployments", "backups.md", "Backup Strategy"), ("Homelab", "aws-route53.md", "AWS Route53 Records"), ] monkeypatch.setattr(agent, "list_catalog", lambda db: catalog) holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="list_documents", arguments={})], [ToolCallPiece(id="call_2", name="list_documents", arguments={})], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.tool_calls == 2 # both re-lists executed and counted listing = ( "2 documents:\n" "source: Deployments | path: backups.md | title: Backup Strategy\n" "source: Homelab | path: aws-route53.md | title: AWS Route53 Records" ) # The answer request carries the catalog a second time as a tool result. assert llm.requests[2][0][3]["content"] == listing # first listing assert llm.requests[2][0][5]["content"] == listing # the re-list assert llm.requests[2][1] == AGENT_TOOLS # still offered (no budgets) def test_multi_read_executes_without_budgets( monkeypatch: pytest.MonkeyPatch, ) -> None: """Reads are no longer budgeted either — two different documents can be read in one turn (re-reading the same one is still deduped via ALREADY_IN_CONTEXT — see the rejection tests).""" a = _doc("S", "a.md", "A", "A-CONTENT") b = _doc("S", "b.md", "B", "B-CONTENT") monkeypatch.setattr( agent, "find_document", lambda db, source, path: {"a.md": a, "b.md": b}[path] ) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S", "path": "a.md"} ) ], [ ToolCallPiece( id="call_2", name="read_document", arguments={"source": "S", "path": "b.md"} ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.read_docs == [a, b] # both reads appended, in order assert holder.tool_calls == 2 assert llm.requests[1][0][3]["content"] == "Document S/a.md:\nA-CONTENT" assert llm.requests[2][0][5]["content"] == "Document S/b.md:\nB-CONTENT" assert llm.requests[2][1] == AGENT_TOOLS # the second read was still offered # ---------- rejections (non-budget; the round cap bounds their repetition) ---------- def test_reading_a_seed_doc_is_already_in_context( monkeypatch: pytest.MonkeyPatch, ) -> None: seed = [_doc("Homelab", "kubernetes.md", "Kubernetes", "K8S-CONTENT")] def _boom(*_a: Any, **_k: Any) -> None: raise AssertionError("find_document must not be called for a seeded doc") monkeypatch.setattr(agent, "list_catalog", lambda db: []) monkeypatch.setattr(agent, "find_document", _boom) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "Homelab", "path": "kubernetes.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings(), seed_docs=seed)) assert holder.read_docs == [] and holder.tool_calls == 0 assert llm.requests[1][0][3]["content"] == agent.ALREADY_IN_CONTEXT # Rejected → the tools are still offered on the next request (the # round cap is the only bound). assert llm.requests[1][1] == AGENT_TOOLS def test_reading_an_already_read_doc_is_deduped( monkeypatch: pytest.MonkeyPatch, ) -> None: doc = _doc("S", "a.md", "A", "A-CONTENT") monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S", "path": "a.md"} ) ], [ ToolCallPiece( id="call_2", name="read_document", arguments={"source": "S", "path": "a.md"} ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.read_docs == [doc] # appended exactly once assert holder.tool_calls == 1 assert llm.requests[2][0][5]["content"] == agent.ALREADY_IN_CONTEXT # Rejected → the tools are still offered on the next request… assert llm.requests[2][1] == AGENT_TOOLS def test_unknown_path_refused( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(agent, "find_document", lambda db, source, path: None) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S", "path": "ghost.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.read_docs == [] and holder.tool_calls == 0 assert ( llm.requests[1][0][3]["content"] == "No document at S/ghost.md — check the list_documents output." ) assert llm.requests[1][1] == AGENT_TOOLS # tools stay offered (cap bounds) def test_unknown_tool_name_refused( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(agent, "list_catalog", lambda db: []) holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="delete_universe", arguments={"x": 1})], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.read_docs == [] and holder.tool_calls == 0 assert llm.requests[1][0][3]["content"] == agent.UNKNOWN_TOOL assert llm.requests[1][1] == AGENT_TOOLS # rejected → tools stay offered @pytest.mark.parametrize( ("arguments", "label"), [ ({}, "no arguments"), ({"source": "S"}, "path missing"), ({"path": "p.md"}, "source missing"), ({"source": "", "path": "p.md"}, "empty source"), ({"source": "S", "path": " "}, "blank path"), ({"source": 7, "path": "p.md"}, "non-string source"), ], ) def test_read_document_missing_arguments_refused( monkeypatch: pytest.MonkeyPatch, arguments: dict[str, Any], label: str ) -> None: def _boom(*_a: Any, **_k: Any) -> None: raise AssertionError(f"find_document must not be called ({label})") monkeypatch.setattr(agent, "list_catalog", lambda db: []) monkeypatch.setattr(agent, "find_document", _boom) holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="read_document", arguments=arguments)], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.read_docs == [] and holder.tool_calls == 0 assert llm.requests[1][0][3]["content"] == agent.MISSING_READ_ARGS assert llm.requests[1][1] == AGENT_TOOLS # ---------- combined 'source/path' self-correction ---------- # The model treats the combined 'source/path' string (search result # lines, read-result headers, refusals) as the document's identity and # sometimes passes it as 'source'. _resolve_document splits it at the # first slash (source names are directory basenames — they can never # contain '/'); a refusal for a still-unknown split teaches the split. def test_read_combined_source_is_split_and_read( monkeypatch: pytest.MonkeyPatch, ) -> None: """source='S/a/b.md' (the combined form) + path='a/b.md': the exact lookup misses, the first-slash split hits — the read executes, the holder records the document, and the result header carries the TRUE source/path (not the model's raw arguments).""" doc = _doc("S", "a/b.md", "B", "B-CONTENT") calls: list[tuple[str, str]] = [] def _find(db: Any, source: str, path: str) -> Document | None: calls.append((source, path)) return doc if (source, path) == ("S", "a/b.md") else None monkeypatch.setattr(agent, "find_document", _find) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S/a/b.md", "path": "a/b.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) # Exact pair first, then the first-slash split (no third attempt). assert calls == [("S/a/b.md", "a/b.md"), ("S", "a/b.md")] assert holder.read_docs == [doc] assert holder.tool_calls == 1 assert llm.requests[1][0][3]["content"] == "Document S/a/b.md:\nB-CONTENT" def test_read_combined_source_split_at_later_slash( monkeypatch: pytest.MonkeyPatch, ) -> None: """source='S/a' + path='b.md' — a split at a LATER slash (source carried source + leading directory, path the remainder): the exact lookup and the first-slash split miss, the continuation candidate (source, split/path) hits.""" doc = _doc("S", "a/b.md", "B", "B-CONTENT") calls: list[tuple[str, str]] = [] def _find(db: Any, source: str, path: str) -> Document | None: calls.append((source, path)) return doc if (source, path) == ("S", "a/b.md") else None monkeypatch.setattr(agent, "find_document", _find) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S/a", "path": "b.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert calls == [("S/a", "b.md"), ("S", "a"), ("S", "a/b.md")] assert holder.read_docs == [doc] assert holder.tool_calls == 1 assert llm.requests[1][0][3]["content"] == "Document S/a/b.md:\nB-CONTENT" def test_read_combined_source_unknown_teaches_the_split( monkeypatch: pytest.MonkeyPatch, ) -> None: """A combined source that matches nothing — even split — gets the EDUCATIONAL refusal: it names the corrected arguments instead of repeating the combined form (the old generic line reinforced the mistake).""" monkeypatch.setattr(agent, "find_document", lambda db, source, path: None) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S/a/b.md", "path": "a/b.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.read_docs == [] and holder.tool_calls == 0 assert llm.requests[1][0][3]["content"] == ( "source must not contain '/': for 'S/a/b.md' call " "read_document(source='S', path='a/b.md')." ) assert llm.requests[1][1] == AGENT_TOOLS # rejected → tools stay offered def test_read_combined_source_for_seed_doc_is_already_in_context( monkeypatch: pytest.MonkeyPatch, ) -> None: """The combined form of a document ALREADY in context: the raw pair cannot match the dedupe set, so the split resolves it — and it is still rejected as already-in-context (no duplicate read_docs entry, no re-read into the context).""" seed = [_doc("S", "a.md", "A", "A-CONTENT")] def _find(db: Any, source: str, path: str) -> Document | None: return seed[0] if (source, path) == ("S", "a.md") else None monkeypatch.setattr(agent, "find_document", _find) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="read_document", arguments={"source": "S/a.md", "path": "a.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings(), seed_docs=seed)) assert holder.read_docs == [] and holder.tool_calls == 0 assert llm.requests[1][0][3]["content"] == agent.ALREADY_IN_CONTEXT def test_search_scoped_combined_source_is_split( monkeypatch: pytest.MonkeyPatch, ) -> None: """A scoped search whose 'source' carries the combined form resolves through the split — the search runs on the right document.""" doc = _doc("S", "a.md", "A", "needle here") def _find(db: Any, source: str, path: str) -> Document | None: return doc if (source, path) == ("S", "a.md") else None monkeypatch.setattr(agent, "find_document", _find) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "needle", "source": "S/a.md", "path": "a.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == "S/a.md:1: needle here" assert holder.tool_calls == 1 assert holder.read_docs == [] # searched doc did not enter the context def test_search_scoped_combined_source_unknown_teaches_the_split( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(agent, "find_document", lambda db, source, path: None) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "x", "source": "S/ghost.md", "path": "ghost.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == ( "source must not contain '/': for 'S/ghost.md' use " "source='S', path='ghost.md'." ) assert holder.tool_calls == 0 and holder.read_docs == [] # a refusal # ---------- search_documents (phase 68, locked A5/A6) ---------- def test_grep_document_case_insensitive_line_numbers() -> None: """Case-insensitive fixed substring, 1-based line numbers, file order, repeated matches within a line collapse to one match (grep semantics).""" content = "The NEEDLE is here\nno hit\nneedle again\nNEEDLE NEEDLE\n" assert agent.grep_document(content, "NEEDLE") == [ (1, "The NEEDLE is here"), (3, "needle again"), (4, "NEEDLE NEEDLE"), ] def test_grep_document_rstrips_lines_and_empty_content() -> None: assert agent.grep_document("hello \t\nworld ", "WORLD") == [(2, "world")] assert agent.grep_document("", "x") == [] assert agent.grep_document("no newlines", "NO") == [(1, "no newlines")] assert agent.grep_document("a\nb\n", "MISSING") == [] def test_search_whole_kb_grep_style_output( monkeypatch: pytest.MonkeyPatch, ) -> None: """Whole-KB search: catalog order, `source/path:line: text` lines, case-insensitive; the call counts in ``tool_calls`` and never touches ``read_docs``; the tools stay offered on the answer request.""" d1 = _doc("Alpha", "a/one.md", "One", "first\nNEEDLE in one\nlast") d2 = _doc("Beta", "b/two.md", "Two", "no hit\nneedle in two\n") monkeypatch.setattr(agent, "all_documents", lambda db: [d1, d2]) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "needle"} ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == ( "Alpha/a/one.md:2: NEEDLE in one\n" "Beta/b/two.md:2: needle in two" ) assert holder.tool_calls == 1 assert holder.read_docs == [] # locked A5: a search adds no context assert llm.requests[1][1] == AGENT_TOOLS # tools stay offered def test_search_capped_at_20_matches_in_catalog_order( monkeypatch: pytest.MonkeyPatch, ) -> None: """The 20-match cap is GLOBAL across documents in catalog order, and the scan stops once it is hit (a 35-match corpus yields exactly 20).""" d1 = _doc("S", "a.md", "A", "\n".join(f"hit-{i}" for i in range(15))) d2 = _doc("S", "b.md", "B", "\n".join(f"hit-{i}" for i in range(20))) monkeypatch.setattr(agent, "all_documents", lambda db: [d1, d2]) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "hit-"} ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) lines = llm.requests[1][0][3]["content"].split("\n") assert len(lines) == agent.SEARCH_MAX_MATCHES assert lines[0] == "S/a.md:1: hit-0" assert lines[14] == "S/a.md:15: hit-14" # all of a.md assert lines[15] == "S/b.md:1: hit-0" # then b.md, in order assert lines[19] == "S/b.md:5: hit-4" # cut at the global cap assert holder.tool_calls == 1 def test_search_truncates_match_lines_at_200_chars( monkeypatch: pytest.MonkeyPatch, ) -> None: """A 300-char match line yields exactly 200 chars of it (no crash).""" d1 = _doc("S", "a.md", "A", "top\n" + "x" * 300 + " NEEDLE tail") monkeypatch.setattr(agent, "all_documents", lambda db: [d1]) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "needle"} ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert ( llm.requests[1][0][3]["content"] == f"S/a.md:2: {'x' * agent.SEARCH_LINE_LIMIT}" ) assert holder.tool_calls == 1 def test_search_scoped_to_one_document(monkeypatch: pytest.MonkeyPatch) -> None: """Scoped search: only the named document is loaded (find_document), ``all_documents`` never runs, and the match line carries its path.""" d1 = _doc("S", "a.md", "A", "needle here") def _find(db: Any, source: str, path: str) -> Document | None: if (source, path) == ("S", "a.md"): return d1 raise AssertionError( f"find_document({source}, {path}) — the scoped " "search must not load any other document" ) def _boom(*_a: Any, **_k: Any) -> None: raise AssertionError("all_documents must not run for a scoped search") monkeypatch.setattr(agent, "find_document", _find) monkeypatch.setattr(agent, "all_documents", _boom) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "needle", "source": "S", "path": "a.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == "S/a.md:1: needle here" assert holder.tool_calls == 1 assert holder.read_docs == [] # searched doc did not enter the context def test_search_scoped_missing_document_refused( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(agent, "find_document", lambda db, source, path: None) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "x", "source": "S", "path": "ghost.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert ( llm.requests[1][0][3]["content"] == "No document at S/ghost.md — check the list_documents output." ) assert holder.tool_calls == 0 and holder.read_docs == [] # a refusal @pytest.mark.parametrize( ("arguments", "label"), [ ({}, "no arguments"), ({"pattern": ""}, "empty pattern"), ({"pattern": " "}, "whitespace pattern"), ({"pattern": 42}, "non-string pattern"), ({"pattern": None}, "null pattern"), ({"pattern": "x", "source": "S"}, "source without path"), ({"pattern": "x", "path": "a.md"}, "path without source"), ], ) def test_search_missing_arguments_refused( monkeypatch: pytest.MonkeyPatch, arguments: dict[str, Any], label: str ) -> None: """Unusable pattern OR a half-specified source/path pair → the missing-args refusal, with no DB access at all.""" def _boom(*_a: Any, **_k: Any) -> None: raise AssertionError(f"no DB access for a refused search ({label})") monkeypatch.setattr(agent, "all_documents", _boom) monkeypatch.setattr(agent, "find_document", _boom) holder = AgentHolder() llm = ScriptedLLM( [ToolCallPiece(id="call_1", name="search_documents", arguments=arguments)], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == agent.MISSING_SEARCH_ARGS assert holder.tool_calls == 0 and holder.read_docs == [] assert llm.requests[1][1] == AGENT_TOOLS # rejected → tools stay offered def test_search_no_matches_whole_kb(monkeypatch: pytest.MonkeyPatch) -> None: """Zero hits across the KB → the no-match line (pattern quoted); the search still executed, so it counts — and never adds context.""" monkeypatch.setattr( agent, "all_documents", lambda db: [_doc("S", "a.md", "A", "nothing here")] ) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "zebra"} ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == ( "No matches for 'zebra' in the knowledge base." ) assert holder.tool_calls == 1 assert holder.read_docs == [] def test_search_no_matches_scoped(monkeypatch: pytest.MonkeyPatch) -> None: doc = _doc("S", "a.md", "A", "nothing here") monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "zebra", "source": "S", "path": "a.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == "No matches for 'zebra' in S/a.md." assert holder.tool_calls == 1 assert holder.read_docs == [] def test_search_no_match_truncates_long_pattern( monkeypatch: pytest.MonkeyPatch, ) -> None: """A pattern longer than 100 chars is truncated in the no-match line (kept short); the search itself still runs on the full pattern.""" monkeypatch.setattr(agent, "all_documents", lambda db: []) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "p" * 150}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert llm.requests[1][0][3]["content"] == ( f"No matches for '{'p' * 100}' in the knowledge base." ) assert holder.tool_calls == 1 def test_search_counts_but_never_adds_context( monkeypatch: pytest.MonkeyPatch, ) -> None: """The locate-then-read workflow: a search finds the document but does NOT add it — the subsequent read_document does (and is not rejected as already-in-context, because the search touched nothing).""" doc = _doc("S", "a.md", "A", "needle here") monkeypatch.setattr(agent, "all_documents", lambda db: [doc]) monkeypatch.setattr(agent, "find_document", lambda db, source, path: doc) holder = AgentHolder() llm = ScriptedLLM( [ ToolCallPiece( id="call_1", name="search_documents", arguments={"pattern": "needle"} ) ], [ ToolCallPiece( id="call_2", name="read_document", arguments={"source": "S", "path": "a.md"}, ) ], [StreamPiece("content", "ans")], ) asyncio.run(_run(llm, holder, _settings())) assert holder.tool_calls == 2 # search + read, both executed assert holder.read_docs == [doc] # only the read added context (A5) assert llm.requests[2][0][5]["content"] == "Document S/a.md:\nneedle here" # ---------- retries inside the agent loop (phase 67, locked A2) ---------- class FailingLLM: """A scripted fake whose Nth ``chat_stream`` call yields pieces and then raises (phase 67): ``attempts`` is a list of ``(pieces, error)`` — an error after zero pieces = "the endpoint died before the first token"; after some pieces = a mid-stream drop. Records every request's messages/tools and the indices of the attempts whose stream teardown ran (``closed``).""" def __init__( self, attempts: list[tuple[list[StreamPiece | ToolCallPiece], Exception | None]], ) -> None: self.attempts = list(attempts) self.requests: list[tuple[list[dict[str, Any]], list[dict[str, Any]] | None]] = [] #: Indices of attempts whose stream teardown has run. self.closed: list[int] = [] def chat_stream( self, messages: list[dict[str, str]], tools: list[dict[str, Any]] | None = None, ) -> AsyncIterator[StreamPiece | ToolCallPiece]: index = len(self.requests) pieces, error = ( self.attempts[index] if index < len(self.attempts) else ([], LLMError("script exhausted")) ) self.requests.append( (deepcopy(messages), deepcopy(tools) if tools is not None else None) ) return self._attempt(index, pieces, error) async def _attempt( self, index: int, pieces: list[StreamPiece | ToolCallPiece], error: Exception | None, ) -> AsyncIterator[StreamPiece | ToolCallPiece]: try: for piece in pieces: yield piece if error is not None: raise error finally: self.closed.append(index) def _record_sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]: """Monkeypatch ``asyncio.sleep`` (what ``chat_stream_retried`` awaits for the flat pre-retry delay) and record every awaited delay.""" sleeps: list[float] = [] async def fake_sleep(seconds: float) -> None: sleeps.append(seconds) monkeypatch.setattr(asyncio, "sleep", fake_sleep) return sleeps def test_round_retried_before_first_piece( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: """A tool round that dies before its first piece is restarted with the same messages: the stream carries a RetryPiece BEFORE the tool call, the tool executes, the final answer streams, and the per-call log line is still emitted exactly once (retries are invisible to the loop).""" monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")]) holder = AgentHolder() llm = FailingLLM( [ ([], LLMError("connection refused")), ([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None), ([StreamPiece("content", "Done!")], None), ] ) sleeps = _record_sleeps(monkeypatch) with caplog.at_level(logging.INFO, logger="app.agent"): pieces = asyncio.run( _run(llm, holder, _settings(agent_max_rounds=2, llm_retry_delay=2.5)) ) assert pieces == [ RetryPiece(2, 4), # default llm_retries=3 → 4 attempts ToolCallPiece(id="call_1", name="list_documents", arguments={}), StreamPiece("content", "Done!"), ] assert holder.tool_calls == 1 assert holder.read_docs == [] # The restart is byte-identical: same messages, same tools offered. assert len(llm.requests) == 3 assert llm.requests[0] == llm.requests[1] assert llm.requests[0][1] == AGENT_TOOLS assert llm.requests[2][1] == AGENT_TOOLS # the answer round still offered # The flat delay was awaited exactly once, before the retry. assert sleeps == [2.5] tool_logs = [r for r in caplog.records if r.getMessage().startswith("agent tool=")] assert len(tool_logs) == 1 # the retry did not re-run the tool or log assert tool_logs[0].getMessage() == "agent tool=list_documents args={} round=1/2" def test_round_failure_after_first_piece_is_terminal( monkeypatch: pytest.MonkeyPatch, ) -> None: """Locked A2: a round that already streamed a piece fails the turn — the LLMError propagates out of ``run_agent``, no RetryPiece, no sleep, no second request, and the holder is untouched (the tool never ran).""" monkeypatch.setattr(agent, "list_catalog", lambda db: []) holder = AgentHolder() llm = FailingLLM( [([StreamPiece("content", "partial ")], LLMError("mid-stream drop"))] ) sleeps = _record_sleeps(monkeypatch) async def drain() -> list[StreamPiece | ToolCallPiece | RetryPiece]: out: list[StreamPiece | ToolCallPiece | RetryPiece] = [] with pytest.raises(LLMError, match="mid-stream drop"): async for piece in run_agent( cast("LLMClient", llm), cast("Session", None), system_prompt="SYSTEM_PROMPT", user_message="QUESTION", seed_docs=[], settings=_settings(), holder=holder, ): out.append(piece) return out out = asyncio.run(drain()) assert out == [StreamPiece("content", "partial ")] # no RetryPiece assert len(llm.requests) == 1 # no retry assert sleeps == [] assert holder.read_docs == [] and holder.tool_calls == 0 def test_forced_final_no_tools_call_is_retried( monkeypatch: pytest.MonkeyPatch, ) -> None: """The forced final request (round cap reached) goes through the same retry rule: a failure before its first piece yields a RetryPiece and restarts with ``tools=None``; the answer from the retry streams.""" monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")]) holder = AgentHolder() llm = FailingLLM( [ ([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None), ([ToolCallPiece(id="call_2", name="list_documents", arguments={})], None), ([], LLMError("down at the cap")), ([StreamPiece("content", "forced answer")], None), ] ) pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=2))) assert [type(p) for p in pieces] == [ ToolCallPiece, ToolCallPiece, RetryPiece, StreamPiece, ] assert pieces[2] == RetryPiece(2, 4) assert pieces[3] == StreamPiece("content", "forced answer") assert len(llm.requests) == 4 # 2 tool rounds + the final + its retry # The forced final (and its retry) carry no tools, whatever is left. assert llm.requests[2][1] is None assert llm.requests[3][1] is None # …and the restart is byte-identical. assert llm.requests[2][0] == llm.requests[3][0] assert holder.tool_calls == 2 def test_zero_retries_is_one_plain_attempt( monkeypatch: pytest.MonkeyPatch, ) -> None: """The kill-switch path (``llm_retries=0``): a dead round raises immediately — one request, no RetryPiece, no sleep (pre-phase-67 behavior).""" monkeypatch.setattr(agent, "list_catalog", lambda db: []) holder = AgentHolder() llm = FailingLLM([([], LLMError("connection refused"))]) sleeps = _record_sleeps(monkeypatch) async def drain() -> list[StreamPiece | ToolCallPiece | RetryPiece]: out: list[StreamPiece | ToolCallPiece | RetryPiece] = [] with pytest.raises(LLMError, match="connection refused"): async for piece in run_agent( cast("LLMClient", llm), cast("Session", None), system_prompt="SYSTEM_PROMPT", user_message="QUESTION", seed_docs=[], settings=_settings(llm_retries=0), holder=holder, ): out.append(piece) return out out = asyncio.run(drain()) assert out == [] # nothing streamed, no RetryPiece assert len(llm.requests) == 1 assert sleeps == [] assert holder.read_docs == [] and holder.tool_calls == 0 def test_abandon_mid_retry_sleep_leaks_nothing( monkeypatch: pytest.MonkeyPatch, ) -> None: """Consumer abandon while a retried round is parked in the pre-retry sleep (client disconnect): the driving task is cancelled cleanly, the production teardown ``aclose()`` on ``run_agent`` does not raise, the inner attempt's stream was torn down, and the retry never starts.""" entered = asyncio.Event() async def parking_sleep(seconds: float) -> None: entered.set() await asyncio.Event().wait() # park until the abandon arrives monkeypatch.setattr(asyncio, "sleep", parking_sleep) monkeypatch.setattr(agent, "list_catalog", lambda db: []) holder = AgentHolder() llm = FailingLLM( [([], LLMError("endpoint down")), ([StreamPiece("content", "never")], None)] ) async def run() -> None: gen = run_agent( cast("LLMClient", llm), cast("Session", None), system_prompt="SYSTEM_PROMPT", user_message="QUESTION", seed_docs=[], settings=_settings(), holder=holder, ) async def consumer() -> list[StreamPiece | ToolCallPiece | RetryPiece]: return [p async for p in gen] task = asyncio.ensure_future(consumer()) await entered.wait() # the round's retry is parked in the sleep assert not task.done() task.cancel() # client disconnect: the driving task is cancelled with pytest.raises(asyncio.CancelledError): await task # Production teardown (phase 48 pattern): must not raise. ``run_agent`` # is an async generator despite its AsyncIterator annotation. await cast( "AsyncGenerator[StreamPiece | ToolCallPiece | RetryPiece, None]", gen ).aclose() asyncio.run(run()) assert len(llm.requests) == 1 # the retry never started assert llm.closed == [0] # attempt 1's inner stream was torn down assert holder.read_docs == [] and holder.tool_calls == 0 def test_retries_are_invisible_to_the_round_cap( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: """A failing-then-succeeding round consumes ONE round: with a cap of 2, the retried first round and the second tool round fill the cap — the forced final follows the SECOND call, and the log lines read round=1/2 and round=2/2.""" monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")]) holder = AgentHolder() llm = FailingLLM( [ ([], LLMError("down")), ([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None), ([ToolCallPiece(id="call_2", name="list_documents", arguments={})], None), ([StreamPiece("content", "forced answer")], None), ] ) with caplog.at_level(logging.INFO, logger="app.agent"): pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=2))) assert [type(p) for p in pieces] == [ RetryPiece, ToolCallPiece, ToolCallPiece, StreamPiece, ] assert len(llm.requests) == 4 # 2 (round 1 + its retry) + 1 + the forced final assert llm.requests[3][1] is None # the forced final, after round 2 assert holder.tool_calls == 2 msgs = [r.getMessage() for r in caplog.records] assert "agent tool=list_documents args={} round=1/2" in msgs assert "agent tool=list_documents args={} round=2/2" in msgs assert any("round cap reached (rounds=2)" in m for m in msgs) # ---------- prompts: section (HIGH only) ---------- def test_high_prompt_carries_tools_section_after_documents() -> None: prompt = build_high_prompt([_doc("S", "a.md", "A", "A-CONTENT")]) assert TOOLS_SECTION in prompt assert "call `list_documents`" in prompt assert "then `read_document` to pull in exactly one more document" in prompt assert "do not read more than one extra document" in prompt # After the mode body: follows . assert prompt.index("") < prompt.index("") assert prompt.rstrip().endswith("") def test_high_prompt_tools_section_with_notes_and_kb() -> None: prompt = build_high_prompt( [_doc("S", "a.md", "A", "A-CONTENT")], notes=["be concise"], kb_overview="- KB" ) assert prompt.index("") < prompt.index("") assert prompt.index("") < prompt.index("") assert prompt.index("") < prompt.index("") def test_low_prompt_is_byte_identical_and_tool_free() -> None: expected = ( _base("LOW") + "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest " "your notes come to the question. They are titles only; do not pretend " "they answer it. Use them to propose 2-3 alternative questions.\n" + "- T1\n- T2" ) assert build_deflect_prompt(["T1", "T2"]) == expected for prompt in ( build_deflect_prompt(["T1"]), build_deflect_prompt(["T1"], notes=["be concise"]), build_deflect_prompt(["T1"], kb_overview="- KB"), build_deflect_prompt(["T1"], notes=["be concise"], kb_overview="- KB"), ): assert "" not in prompt assert TOOLS_SECTION not in prompt