from __future__ import annotations import json import shutil import tempfile import unittest from dataclasses import replace from pathlib import Path from unittest import mock from jsonschema import Draft202012Validator from docforge.errors import DocForgeError from docforge.index import ProjectIndex from docforge.models import ProjectState from docforge.policy import compose_effective_policy from docforge.project import Project from docforge.retrieval import ( BASE_RELATION_CATEGORIES, MAX_TASK_CANDIDATE_EDGES, MAX_TASK_EVIDENCE, RELATION_CATEGORIES, TASK_KINDS, TASK_REQUIREMENTS, build_retrieval_plan, relation_category, ) ROOT = Path(__file__).resolve().parents[1] FIXTURES = ROOT / "tests" / "fixtures" CAPSULE_SCHEMA = json.loads( (ROOT / "schemas" / "context-capsule.schema.json").read_text(encoding="utf-8") ) class TaskRetrievalTests(unittest.TestCase): def copy_fixture(self, name: str, destination: Path) -> Path: root = destination / name shutil.copytree(FIXTURES / name, root) return root @staticmethod def effective_policy() -> dict[str, object]: return compose_effective_policy( selected_mode="read", capability_source="explicit", no_ast=False, diagnostics=False, render_configured=True, application_enabled=False, ).as_dict() def plan( self, project: Project, *, task_kind: str = "change", task: str = "Change the editing workflow", focus_node_id: str | None = "guide.workflow", budget: int | None = None, limit: int | None = None, ): return build_retrieval_plan( project.descriptor, task_kind=task_kind, task=task, focus_node_id=focus_node_id, budget=budget, limit=limit, effective_policy=self.effective_policy(), ) def test_every_task_plan_is_closed_deterministic_and_hash_bound(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("alpha", Path(directory))) hashes = {} relation_hashes = set() for task_kind in TASK_KINDS: first = self.plan(project, task_kind=task_kind) second = self.plan(project, task_kind=task_kind) self.assertEqual(first, second) self.assertEqual(TASK_REQUIREMENTS[task_kind], first.category_order[:1]) self.assertEqual( { *BASE_RELATION_CATEGORIES, "unclassified", }, set(first.category_order), ) self.assertEqual(64, len(first.plan_hash)) relation_hashes.update( step.relation_set_hash for step in first.steps if step.relation_scope == "project_allowed" ) hashes[task_kind] = first.plan_hash self.assertEqual(len(TASK_KINDS), len(set(hashes.values()))) self.assertEqual(1, len(relation_hashes)) self.assertEqual( { "change": ("dependency",), "implementation": ("implementation",), "failure": ("execution",), "ownership": ("structure",), "test": ("evidence",), "operation": ("execution",), "release": ("evidence",), }, TASK_REQUIREMENTS, ) mapped = [ relation for relations in RELATION_CATEGORIES.values() for relation in relations ] self.assertEqual(len(mapped), len(set(mapped))) self.assertEqual( { "structure": ("contains", "defined_in", "defines", "owns"), "implementation": ( "implemented_by", "implements", "inherits", "inherits_from", ), "dependency": ("depends_on", "imports"), "execution": ( "activates", "calls", "dispatches_to", "launches", ), "data": ("reads", "writes"), "evidence": ( "documents", "governs", "proves", "tested_by", "verifies", ), "context": ("relates_to",), "unclassified": (), }, RELATION_CATEGORIES, ) self.assertEqual("dependency", relation_category("depends_on")) self.assertEqual("evidence", relation_category("proves")) self.assertEqual("unclassified", relation_category("owns_database")) self.assertEqual("unclassified", relation_category("when_true")) def test_exact_capsule_is_schema_valid_stable_and_explainable(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("alpha", Path(directory))) index = ProjectIndex(project) index.build() plan = self.plan(project) first = index.task_context(plan) second = index.task_context(plan) self.assertEqual(first, second) capsule = first["capsule"] Draft202012Validator(CAPSULE_SCHEMA).validate(capsule) self.assertEqual("complete", capsule["state"]) self.assertEqual("resolved", capsule["focus"]["state"]) self.assertEqual("guide.workflow", capsule["focus"]["node_id"]) self.assertEqual( ["guide.workflow", "guide.foundation", "proof.validation"], [item["node_id"] for item in capsule["evidence"]], ) dependency = capsule["evidence"][1]["relationship_path"][0] self.assertEqual("depends_on", dependency["relation"]) self.assertEqual("dependency", dependency["category"]) self.assertEqual("outgoing", dependency["direction"]) workflow_reasons = capsule["evidence"][0]["relationship_reasons"] self.assertIn( ("depends_on", "outgoing"), { (relationship["relation"], relationship["direction"]) for relationship in workflow_reasons }, ) self.assertIn( ("proves", "incoming"), { (relationship["relation"], relationship["direction"]) for relationship in workflow_reasons }, ) foundation_reason = capsule["evidence"][1]["relationship_reasons"][0] self.assertEqual("depends_on", foundation_reason["relation"]) self.assertEqual("incoming", foundation_reason["direction"]) self.assertNotIn( "no_selected_evidence", [gap["code"] for gap in capsule["gaps"]], ) self.assertEqual(64, len(capsule["collection_hash"])) self.assertEqual(64, len(capsule["capsule_hash"])) def test_gaps_distinguish_undeclared_complete_and_incomplete_checks(self) -> None: with tempfile.TemporaryDirectory() as directory: root = self.copy_fixture("alpha", Path(directory)) project = Project.open(root) index = ProjectIndex(project) index.build() undeclared = index.task_context(self.plan(project, task_kind="implementation"))[ "capsule" ] self.assertIn( "category_not_declared", [gap["code"] for gap in undeclared["gaps"]], ) descriptor = root / ".docforge" / "project.toml" descriptor.write_text( descriptor.read_text(encoding="utf-8").replace( '"returns_to"]', '"returns_to", "calls"]', ), encoding="utf-8", ) declared_project = Project.open(root) declared_index = ProjectIndex(declared_project) declared_index.build() missing = declared_index.task_context( self.plan( declared_project, task_kind="failure", focus_node_id="guide.foundation", ) )["capsule"] self.assertIn( "no_selected_evidence", [gap["code"] for gap in missing["gaps"]], ) incomplete = declared_index.task_context(self.plan(declared_project, budget=1))[ "capsule" ] self.assertEqual("incomplete", incomplete["state"]) self.assertTrue(incomplete["omissions"]) self.assertIn( "evidence_incomplete", [gap["code"] for gap in incomplete["gaps"]], ) self.assertNotIn( "no_selected_evidence", [gap["code"] for gap in incomplete["gaps"]], ) def test_unclassified_relation_is_preserved_without_guessed_semantics(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("awesome-ski-game", Path(directory))) index = ProjectIndex(project) index.build() plan = build_retrieval_plan( project.descriptor, task_kind="change", task="Change the first descent session", focus_node_id="session.first-descent", budget=None, limit=None, effective_policy=self.effective_policy(), ) capsule = index.task_context(plan)["capsule"] self.assertIn("informs", capsule["summary"]["unclassified_relations"]) relation = next( relationship for item in capsule["evidence"] for relationship in item["relationship_reasons"] if relationship["relation"] == "informs" ) self.assertEqual("unclassified", relation["category"]) self.assertIn( "unclassified_relation", [gap["code"] for gap in capsule["gaps"]], ) def test_plan_is_compact_at_large_valid_relation_scale_and_limits_are_fixed(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("alpha", Path(directory))) relation_names = tuple(f"relation-{position:05d}" for position in range(33_005)) descriptor = replace( project.descriptor, allowed_relations=relation_names, limits=replace( project.descriptor.limits, max_results=2**63, ), ) plan = build_retrieval_plan( descriptor, task_kind="change", task="Exercise a very large valid relation policy", focus_node_id="guide.workflow", budget=None, limit=None, effective_policy=self.effective_policy(), ) self.assertLessEqual(plan.max_evidence, MAX_TASK_EVIDENCE) self.assertLessEqual(plan.max_candidate_edges, MAX_TASK_CANDIDATE_EDGES) traversal = [step for step in plan.steps if step.operation in {"outgoing", "incoming"}] self.assertEqual(2, len(traversal)) self.assertTrue(all(step.relation_scope == "project_allowed" for step in traversal)) self.assertTrue(all(len(step.relation_set_hash or "") == 64 for step in traversal)) self.assertLess(len(json.dumps(plan.as_dict())), 5_000) def test_executor_rejects_tampered_public_plan_objects(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("alpha", Path(directory))) index = ProjectIndex(project) index.build() plan = self.plan(project) tampered = ( replace(plan, steps=plan.steps[:-1]), replace(plan, requirements=()), replace(plan, task_query="Different task"), replace(plan, category_order=tuple(reversed(plan.category_order))), replace(plan, max_evidence=plan.max_evidence + 1), replace(plan, effective_policy_hash="0" * 64), replace(plan, plan_hash="0" * 64), ) for candidate in tampered: with self.subTest(candidate=candidate): with self.assertRaises(DocForgeError) as invalid: index.task_context(candidate) self.assertEqual("invalid_retrieval_plan", invalid.exception.code) def test_work_and_unclassified_limits_are_explicit_and_bounded(self) -> None: with tempfile.TemporaryDirectory() as directory: root = self.copy_fixture("alpha", Path(directory)) relation_names = [f"relation-{position:03d}" for position in range(450)] descriptor = root / ".docforge" / "project.toml" raw_descriptor = descriptor.read_text(encoding="utf-8") raw_relations = ", ".join( json.dumps(relation) for relation in ( "depends_on", "proves", "supersedes", "relates_to", "returns_to", *relation_names, ) ) descriptor.write_text( raw_descriptor.replace( '"depends_on", "proves", "supersedes", "relates_to", "returns_to"', raw_relations, ), encoding="utf-8", ) workflow = root / "docs" / "content" / "workflow.md" raw_workflow = workflow.read_text(encoding="utf-8") relationships = "".join( f'{relation} = ["guide.foundation"]\n' for relation in relation_names ) workflow.write_text( raw_workflow.replace("+++\n\nEditors", f"{relationships}+++\n\nEditors", 1), encoding="utf-8", ) project = Project.open(root) index = ProjectIndex(project) index.build() capsule = index.task_context(self.plan(project))["capsule"] Draft202012Validator(CAPSULE_SCHEMA).validate(capsule) omission_codes = {omission["code"] for omission in capsule["omissions"]} self.assertIn("edge_examination_limit", omission_codes) self.assertIn("unclassified_relation_limit", omission_codes) self.assertEqual( project.descriptor.limits.max_results, len(capsule["summary"]["unclassified_relations"]), ) self.assertLessEqual( capsule["summary"]["examined_edge_count"], MAX_TASK_CANDIDATE_EDGES, ) def test_lexical_focus_blocks_missing_and_ambiguous_selection(self) -> None: with tempfile.TemporaryDirectory() as directory: root = self.copy_fixture("alpha", Path(directory)) project = Project.open(root) index = ProjectIndex(project) index.build() missing = index.task_context( self.plan( project, task="Words absent from every indexed node", focus_node_id=None, ) )["capsule"] self.assertEqual("blocked", missing["state"]) self.assertEqual("not_found", missing["focus"]["state"]) self.assertEqual("focus_not_found", missing["gaps"][0]["code"]) duplicate = root / "docs" / "content" / "workflow-copy.md" duplicate.write_text( (root / "docs" / "content" / "workflow.md") .read_text(encoding="utf-8") .replace('id = "guide.workflow"', 'id = "guide.workflow-copy"'), encoding="utf-8", ) duplicate_project = Project.open(root) duplicate_index = ProjectIndex(duplicate_project) duplicate_index.build() ambiguous = duplicate_index.task_context( self.plan( duplicate_project, task="Editing workflow", focus_node_id=None, ) )["capsule"] self.assertEqual("blocked", ambiguous["state"]) self.assertEqual("ambiguous", ambiguous["focus"]["state"]) self.assertEqual("focus_ambiguous", ambiguous["gaps"][0]["code"]) def test_plan_rejects_unbounded_or_untyped_inputs(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("alpha", Path(directory))) cases = ( ({"task_kind": "unknown"}, "invalid_task_kind"), ({"task": " "}, "invalid_task_focus"), ({"focus_node_id": ""}, "invalid_task_focus"), ({"budget": True}, "invalid_budget"), ({"budget": 0}, "invalid_budget"), ({"limit": True}, "invalid_limit"), ({"limit": project.descriptor.limits.max_results + 1}, "invalid_limit"), ) for overrides, code in cases: arguments = { "task_kind": "change", "task": "Change the editing workflow", "focus_node_id": "guide.workflow", "budget": None, "limit": None, **overrides, } with self.subTest(arguments=overrides): with self.assertRaises(DocForgeError) as invalid: self.plan(project, **arguments) self.assertEqual(code, invalid.exception.code) def test_final_generation_change_rejects_the_whole_capsule(self) -> None: with tempfile.TemporaryDirectory() as directory: project = Project.open(self.copy_fixture("alpha", Path(directory))) index = ProjectIndex(project) built = index.build() current = ProjectState( revision=built["revision"], source_hash=built["source_hash"], ) changed = ProjectState( revision=current.revision, source_hash="0" * 64, ) with ( mock.patch.object( project, "incremental_state", side_effect=(current, changed), ), self.assertRaises(DocForgeError) as stale, ): index.task_context(self.plan(project)) self.assertEqual("source_changed", stale.exception.code)