"""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