tracer-cloud--opensre/tests/interactive_shell/references/test_docs_reference.py

"""Tests for the documentation-grounding helpers used by the interactive shell."""

from __future__ import annotations

from pathlib import Path

import pytest

from core.agent_harness.grounding import docs_reference
from core.agent_harness.grounding._cache import excerpt
from core.agent_harness.grounding.docs_reference import (
    DocPage,
    DocsReference,
    _query_tokens,
    build_docs_index,
    find_relevant_docs,
)


def _write_doc(root: Path, relpath: str, content: str) -> None:
    path = root / relpath
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")


def _seed_docs(root: Path) -> None:
    _write_doc(
        root,
        "datadog.mdx",
        '---\ntitle: "Datadog"\n---\n\n'
        "### Step 1: Create API Key\n\n"
        "In Datadog, create an API Key under organizational settings.\n\n"
        "### Step 2: Configure OpenSRE\n\n"
        "Set DD_API_KEY and DD_APP_KEY in your environment.\n",
    )
    _write_doc(
        root,
        "deployment.mdx",
        '---\ntitle: "Deployment"\n---\n\n'
        "OpenSRE can deploy to Railway or EC2.\n\n"
        "Use `opensre remote` to connect to a deployed agent.\n",
    )
    _write_doc(
        root,
        "quickstart.mdx",
        '---\ntitle: "Quickstart"\n---\n\nInstall OpenSRE and run your first investigation.\n',
    )
    _write_doc(
        root,
        "tutorials/investigating-task-failures.mdx",
        "# Investigating task failures\n\n"
        "Walk through how to investigate a failed task using OpenSRE.\n",
    )
    # Asset content under skip dirs MUST be excluded from the index.
    _write_doc(root, "images/datadog.mdx", "should be skipped")
    _write_doc(root, "assets/anything.mdx", "should be skipped")


class TestDiscoverDocs:
    def test_returns_empty_list_when_root_missing(self, tmp_path: Path) -> None:
        missing = tmp_path / "no-docs"
        assert DocsReference().discover(missing) == []

    def test_walks_root_and_skips_asset_dirs(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        slugs = {p.slug for p in pages}
        assert "datadog" in slugs
        assert "deployment" in slugs
        assert "quickstart" in slugs
        assert "investigating-task-failures" in slugs
        # Anything under images/ or assets/ must be skipped.
        relpaths = {p.relpath for p in pages}
        assert all(not r.startswith("images/") for r in relpaths)
        assert all(not r.startswith("assets/") for r in relpaths)

    def test_extracts_title_from_frontmatter(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        by_slug = {p.slug: p for p in pages}
        assert by_slug["datadog"].title == "Datadog"
        assert by_slug["deployment"].title == "Deployment"

    def test_falls_back_to_first_heading_when_no_frontmatter(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        by_slug = {p.slug: p for p in pages}
        assert by_slug["investigating-task-failures"].title == "Investigating task failures"

    def test_strips_frontmatter_from_body(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        by_slug = {p.slug: p for p in pages}
        # Frontmatter delimiters and the title line must NOT appear in the body
        # (they would otherwise leak into the LLM grounding context).
        assert "title:" not in by_slug["datadog"].body
        assert by_slug["datadog"].body.lstrip().startswith("###")


class TestQueryTokens:
    def test_strips_stopwords_and_short_tokens(self) -> None:
        tokens = _query_tokens("How do I configure Datadog?")
        # Stopwords are removed, 'datadog' / 'configure' remain.
        assert "datadog" in tokens
        assert "configure" in tokens
        assert "how" not in tokens
        assert "do" not in tokens
        assert "i" not in tokens

    def test_drops_opensre_brand_token(self) -> None:
        # Every doc mentions "opensre" so it would otherwise dominate ranking.
        tokens = _query_tokens("how do I install opensre")
        assert "opensre" not in tokens
        assert "install" in tokens

    def test_keeps_two_letter_tokens(self) -> None:
        tokens = _query_tokens("how do I tune ai vm sizing")
        assert "ai" in tokens
        assert "vm" in tokens


class TestFindRelevantDocs:
    def test_empty_query_returns_empty(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        # Query with only stopwords should not match anything.
        assert find_relevant_docs("how do I", pages) == []

    def test_ranks_datadog_page_first_for_datadog_query(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        results = find_relevant_docs("how do I configure Datadog?", pages)
        assert results, "expected at least one match"
        assert results[0].slug == "datadog"

    def test_ranks_deployment_page_first_for_deploy_query(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        results = find_relevant_docs("how do I deploy this?", pages)
        assert results
        assert results[0].slug == "deployment"

    def test_caps_results_at_top_n(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        results = find_relevant_docs("install configure deploy investigate", pages, top_n=2)
        assert len(results) <= 2

    def test_nested_page_with_weak_match_is_not_dropped_by_depth(self, tmp_path: Path) -> None:
        """A page whose only match is a single body token, nested deep enough
        that the depth penalty equals or exceeds its raw score, must still
        surface as a lower-ranked result instead of being excluded entirely.

        Regression: previously the depth penalty was applied unconditionally
        before the score>0 filter, so a page with raw_score=1 at depth=2
        scored -1 and was dropped from results.
        """
        # Page nested 2 levels deep whose only match for "masking" is a single
        # body-token mention. raw_score == 1, depth == 2, so without clamping
        # the final score would be -1 and the page would be filtered out.
        _write_doc(
            tmp_path,
            "tutorials/advanced/notes.mdx",
            "# Notes\n\nWe briefly mention masking in this tutorial.\n",
        )
        pages = DocsReference().discover(tmp_path)
        results = find_relevant_docs("masking", pages)
        slugs = [p.slug for p in results]
        assert "notes" in slugs, (
            "weak nested match must still surface, not be dropped by depth alone"
        )


class TestBuildDocsReferenceText:
    def test_returns_empty_when_no_docs_present(self, tmp_path: Path) -> None:
        # Point at a non-existent docs root.
        assert DocsReference().build_text("anything", root=tmp_path / "missing") == ""

    def test_includes_relevant_doc_excerpt_and_index(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        text = DocsReference().build_text("how do I configure Datadog?", root=tmp_path)
        assert "datadog.mdx" in text
        assert "API Key" in text
        # The compact index of all pages must always be appended so the LLM
        # can suggest other relevant pages even when one ranked highest.
        assert "docs index" in text
        assert "deployment.mdx" in text

    def test_truncates_to_max_chars(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        text = DocsReference().build_text("Datadog", max_chars=120, root=tmp_path)
        assert len(text) <= 200
        assert "truncated" in text


class TestBuildDocsIndex:
    def test_lists_all_pages_with_titles(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        pages = DocsReference().discover(tmp_path)
        index = build_docs_index(pages)
        assert "datadog.mdx: Datadog" in index
        assert "deployment.mdx: Deployment" in index

    def test_returns_empty_string_for_no_pages(self) -> None:
        assert build_docs_index([]) == ""


class TestExcerpt:
    def test_returns_full_body_when_short(self) -> None:
        body = "Short body."
        assert excerpt(body, max_chars=100) == "Short body."

    def test_truncates_long_body_with_marker(self) -> None:
        body = ("paragraph one. " * 10) + "\n\n" + ("paragraph two. " * 10)
        out = excerpt(body, max_chars=80)
        assert "truncated" in out


class TestDocPageDataclass:
    def test_is_hashable_and_immutable(self) -> None:
        page = DocPage(slug="x", relpath="x.mdx", title="X", body="hello")
        # frozen dataclasses are hashable, so they can be stored in sets.
        assert page in {page}


class TestDocsGroundingCache:
    def test_cache_maxsize_matches_implementation(self) -> None:
        stats = DocsReference().stats()
        assert stats.maxsize == 32

    def test_repeated_discover_hits_parse_cache(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        ref = DocsReference()
        ref.discover(tmp_path)
        info1 = ref.stats()
        ref.discover(tmp_path)
        info2 = ref.stats()
        assert info2.hits == info1.hits + 1
        assert info2.misses == info1.misses

    def test_invalidate_resets_stats(self, tmp_path: Path) -> None:
        _seed_docs(tmp_path)
        ref = DocsReference()
        ref.discover(tmp_path)
        ref.discover(tmp_path)
        assert ref.stats().hits >= 1
        ref.invalidate()
        cleared = ref.stats()
        assert cleared.hits == 0
        assert cleared.misses == 0
        assert cleared.currsize == 0

    def test_file_edit_invalidates_and_refreshes_content(self, tmp_path: Path) -> None:
        ref = DocsReference()
        _write_doc(
            tmp_path,
            "datadog.mdx",
            '---\ntitle: "Datadog"\n---\n\nOld content.\n',
        )
        pages1 = ref.discover(tmp_path)
        assert any("Old content" in p.body for p in pages1)

        datadog = tmp_path / "datadog.mdx"
        datadog.write_text(
            '---\ntitle: "Datadog"\n---\n\nNew refreshed content.\n',
            encoding="utf-8",
        )
        pages2 = ref.discover(tmp_path)
        assert any("New refreshed content" in p.body for p in pages2)
        assert not any("Old content" in p.body for p in pages2)

    def test_single_tree_walk_per_discover_call(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Regression: avoid a second full walk inside the parse path on a miss."""
        _seed_docs(tmp_path)
        calls = 0
        real_iter = docs_reference._iter_doc_files

        def _spy(root: Path) -> list[Path]:
            nonlocal calls
            calls += 1
            return real_iter(root)

        monkeypatch.setattr(docs_reference, "_iter_doc_files", _spy)
        ref = DocsReference()
        ref.discover(tmp_path)
        assert calls == 1
        ref.discover(tmp_path)
        assert calls == 2