FabledScribe/tests/test_note_pipeline.py

"""The notes corpus on the one retrieval pipeline (milestone 456 step 4).

`run_note_arm` is driven directly here, with the search and both recorders
stubbed through `NoteIO`: the stage order, the invariants the rule arms
already pin (#3497's call row before any return, #3752's fresh-only cut,
#4101's ledger-renders-not-removes), the write path's withheld menu (#3739),
and the two reserved slots in their load-bearing order. The builders that
compose these arms keep their own tests; these pin the stages they share.
"""
from __future__ import annotations

import ast
from pathlib import Path
from unittest.mock import MagicMock

import pytest

from scribe.services import retrieval_pipeline as rp
from scribe.services.lessons import LESSON_NOTE_TYPE
from scribe.services.retrieval_registry import FAN_OUT_SITES, POINTS
from scribe.services.retrieval_surfaces import SURFACES
from tests.helpers import fake_note


def _io(route):
    """A NoteIO whose search answers by the kinds asked for, and records calls."""
    calls: list[dict] = []

    async def search(_uid, _q, **kw):
        calls.append(kw)
        return list(route(kw))

    return rp.NoteIO(search=search, record_retrieval=MagicMock(),
                     record_surfaced=MagicMock()), calls


def _rows(mock, source):
    return [c.kwargs for c in mock.call_args_list if c.kwargs["source"] == source]


def _note(nid, score=None, **kw):
    note = fake_note(id=nid, title=f"record {nid}", user_id=1, **kw)
    return (score, note) if score is not None else note


async def _run(arm, route, *, budget=3, **moment):
    io, calls = _io(route)
    result = await rp.run_note_arm(
        arm, rp.NoteMoment(user_id=1, query="q", project_id=moment.pop("project_id", 2),
                           **moment),
        floor=0.5, budget=budget, io=io,
    )
    return result, io, calls


# ── the call row (#3497, #3752) ───────────────────────────────────────────


@pytest.mark.asyncio
async def test_a_call_that_found_nothing_is_still_logged():
    result, io, _ = await _run(rp.AUTO_INJECT, lambda kw: [])
    assert result.menu == []
    (row,) = _rows(io.record_retrieval, "auto_inject")
    assert row["results"] == [] and row["suppressed"] == 0
    io.record_surfaced.assert_not_called()


@pytest.mark.asyncio
async def test_repeats_and_named_records_are_counted_not_reported():
    hits = [_note(11, 0.80), _note(22, 0.79), _note(33, 0.78)]
    result, io, _ = await _run(
        rp.AUTO_INJECT, lambda kw: [] if kw.get("note_type") else hits,
        seen=frozenset({11}), named=frozenset({33}),
    )
    (row,) = _rows(io.record_retrieval, "auto_inject")
    assert [int(n.id) for _s, n in row["results"]] == [22]
    assert row["suppressed"] == 2
    # The repeat stays on the menu (#4101); the named record does not — it is
    # shown by its own block.
    assert [int(n.id) for _s, n in result.menu] == [11, 22]
    (surfaced,) = _rows(io.record_surfaced, "auto_inject")
    assert surfaced["note_ids"] == [22]


def test_every_notes_stage_logs_before_it_can_return():
    """Structural, like the rule arms' guard: no `return` in a notes stage
    precedes its call row. The via-lesson arm's one exception is the return
    before any search runs — nothing was asked, so there is nothing to log."""
    tree = ast.parse(Path("src/scribe/services/retrieval_pipeline.py").read_text())
    fns = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.AsyncFunctionDef)}
    for name in ("run_note_arm", "_reserve_note_slot", "run_via_lesson_arm"):
        fn = fns[name]
        logged = [n.lineno for n in ast.walk(fn) if isinstance(n, ast.Call)
                  and getattr(n.func, "attr", None) == "record_retrieval"]
        searched = [n.lineno for n in ast.walk(fn) if isinstance(n, ast.Call)
                    and getattr(n.func, "attr", None) == "search"]
        assert logged and searched, f"{name} no longer searches and logs"
        early = [n.lineno for n in ast.walk(fn) if isinstance(n, ast.Return)
                 and min(searched) < n.lineno < min(logged)]
        assert not early, f"{name} can return at line {early[0]} before logging"


# ── the band, and the arm failing open ────────────────────────────────────


@pytest.mark.asyncio
async def test_the_band_narrows_the_menu_after_the_row_is_written():
    """Notes log BEFORE the band (#2085) — the row is the candidate set a
    floor is tuned against, the surfacing rows are what the reader saw."""
    hits = [_note(1, 0.90), _note(2, 0.70)]
    result, io, _ = await _run(rp.WRITE_PATH, lambda kw: hits)
    (row,) = _rows(io.record_retrieval, "write_path")
    assert [int(n.id) for _s, n in row["results"]] == [1, 2]
    assert [int(n.id) for _s, n in result.menu] == [1]
    (surfaced,) = _rows(io.record_surfaced, "write_path_semantic")
    assert surfaced["note_ids"] == [1]


@pytest.mark.asyncio
async def test_a_search_that_raises_costs_the_arm_not_the_response():
    async def boom(*_a, **_kw):
        raise RuntimeError("index unavailable")

    io = rp.NoteIO(search=boom, record_retrieval=MagicMock(), record_surfaced=MagicMock())
    result = await rp.run_note_arm(
        rp.AUTO_INJECT, rp.NoteMoment(user_id=1, query="q", project_id=None),
        floor=0.5, budget=3, io=io,
    )
    assert result.menu == [] and result.answered == []


@pytest.mark.asyncio
async def test_a_recorder_that_raises_costs_only_its_row():
    hits = [_note(1, 0.90)]
    io, _ = _io(lambda kw: [] if kw.get("note_type") else hits)
    io.record_retrieval.side_effect = RuntimeError("telemetry down")
    result = await rp.run_note_arm(
        rp.AUTO_INJECT, rp.NoteMoment(user_id=1, query="q", project_id=None),
        floor=0.5, budget=3, io=io,
    )
    assert [int(n.id) for _s, n in result.menu] == [1]


# ── the write path's withheld menu (#3739) ────────────────────────────────


@pytest.mark.asyncio
async def test_this_responses_menu_is_withheld_but_a_pulled_one_is_still_scored():
    pulled, listed = 5, 6
    answered = [_note(pulled, 0.91), _note(7, 0.88)]
    result, io, calls = await _run(
        rp.WRITE_PATH, lambda kw: answered, budget=2,
        in_menu=frozenset({pulled, listed}), still_scored=frozenset({pulled}),
    )
    (call,) = calls
    # The listed one never reaches the search; the pulled one does, and the
    # limit covers it so the menu still gets its full budget.
    assert call["exclude_ids"] == {listed}
    assert call["limit"] == 3
    assert [int(n.id) for _s, n in result.answered] == [pulled, 7]
    assert [int(n.id) for _s, n in result.menu] == [7]
    # Something was withheld after the search answered, so the score it
    # reported may be one of ours: not measured on this call.
    (row,) = _rows(io.record_retrieval, "write_path")
    assert row["best_available"] is None and row["best_available_id"] is None


@pytest.mark.asyncio
async def test_the_prompt_menu_never_sends_the_ledger_into_its_search():
    _result, _io_, calls = await _run(
        rp.AUTO_INJECT, lambda kw: [], seen=frozenset({11}),
    )
    assert "exclude_ids" not in calls[0]


# ── the reserved slots, in their order ────────────────────────────────────


def _routed(main, reuse=(), lesson=()):
    def route(kw):
        kinds = kw.get("note_type") or ()
        if LESSON_NOTE_TYPE in kinds:
            return lesson
        if kinds:
            return reuse
        return main
    return route


@pytest.mark.asyncio
async def test_reuse_evicts_the_last_line_and_the_lesson_extends():
    main = [_note(1, 0.70), _note(2, 0.69), _note(3, 0.68)]
    reuse = [_note(9, 0.60, note_type="snippet")]
    lesson = [_note(42, 0.58, note_type=LESSON_NOTE_TYPE)]
    result, io, calls = await _run(
        rp.AUTO_INJECT, _routed(main, reuse, lesson), budget=3,
    )
    # Reuse took the last of three; the lesson was added as a fourth.
    assert [int(n.id) for _s, n in result.menu] == [1, 2, 9, 42]
    assert result.slot_ids == {"reuse_slot": 9, "lesson_slot": 42}
    # Reuse ran first: the lesson query was told about the snippet.
    assert [c.get("note_type") for c in calls] == [
        None, ("snippet", "process"), (LESSON_NOTE_TYPE,),
    ]
    assert 9 in calls[2]["exclude_ids"]
    # Each slot logs under its own name.
    assert _rows(io.record_retrieval, "reuse_slot") and _rows(io.record_retrieval, "lesson_slot")
    # The lesson books its own surfacing and never the arm's; the reuse line
    # is counted under the arm, as it always was.
    assert _rows(io.record_surfaced, "lesson_slot")[0]["note_ids"] == [42]
    assert _rows(io.record_surfaced, "auto_inject")[0]["note_ids"] == [1, 2, 9]


@pytest.mark.asyncio
async def test_a_slot_stands_down_when_its_kind_won_on_score():
    main = [_note(9, 0.80, note_type="snippet"), _note(42, 0.79, note_type=LESSON_NOTE_TYPE)]
    _result, io, calls = await _run(rp.AUTO_INJECT, _routed(main))
    assert len(calls) == 1
    assert not _rows(io.record_retrieval, "reuse_slot")
    assert not _rows(io.record_retrieval, "lesson_slot")


@pytest.mark.asyncio
async def test_a_slot_spends_itself_on_a_repeat_without_booking_it_again():
    """A record shown in an EARLIER call is exactly what a slot may spend
    itself on (#4101) — it is rendered, but it is not a new surfacing."""
    main = [_note(1, 0.70)]
    lesson = [_note(42, 0.58, note_type=LESSON_NOTE_TYPE)]
    result, io, _ = await _run(
        rp.AUTO_INJECT, _routed(main, lesson=lesson), seen=frozenset({42}),
    )
    assert 42 in [int(n.id) for _s, n in result.menu]
    assert not _rows(io.record_surfaced, "lesson_slot")
    (row,) = _rows(io.record_retrieval, "lesson_slot")
    assert row["results"] == [] and row["suppressed"] == 1


@pytest.mark.asyncio
async def test_the_write_path_reserves_no_slot():
    main = [_note(1, 0.80)]
    _result, _io_, calls = await _run(rp.WRITE_PATH, _routed(main, [_note(9, 0.7)]))
    assert len(calls) == 1


# ── the specs are the declarations ────────────────────────────────────────


def test_every_notes_source_is_registered_and_every_arm_is_tunable():
    for source in (*rp.NOTE_SOURCES, *rp.NOTE_SURFACED_SOURCES):
        assert source in POINTS, f"{source} records telemetry but is not registered"
    for arm in rp.NOTE_ARMS:
        assert arm.source in SURFACES, f"{arm.source} has no floor or budget to tune"
    # And the slots are measurable, never tunable: a budget of 1 is their feature.
    for slot in rp.NOTE_SLOTS:
        assert slot.source not in SURFACES


def test_the_pipelines_fan_out_sites_declare_every_notes_source():
    logged = FAN_OUT_SITES[
        "scribe/services/retrieval_pipeline.py::record_retrieval(source=source)"]
    surfaced = FAN_OUT_SITES[
        "scribe/services/retrieval_pipeline.py::record_surfaced(source=source)"]
    assert set(rp.NOTE_SOURCES) <= set(logged)
    assert set(rp.NOTE_SURFACED_SOURCES) == set(surfaced)
    # Not the arm's: the slot whose line it is, and the lookups' own site.
    assert "write_path_semantic" not in FAN_OUT_SITES[
        "scribe/services/plugin_context.py::record_surfaced(source=arm)"]


def test_the_filters_say_which_kinds_and_whose_records():
    assert rp.note_search_filters(rp.AUTO_INJECT) == {
        "include_global_kinds": True, "scope": "browse",
    }
    assert rp.note_search_filters(rp.WRITE_PATH) == {
        "note_type": ("snippet", "note", LESSON_NOTE_TYPE), "task_kind": "issue",
        "include_global_kinds": True, "scope": "browse",
    }