Agent/tests/test_recursive_validation_core.py

import json
import tempfile
import unittest

from pydantic import ValidationError

from cyber_agent.core.validation import (
    LLMValidator,
    ScopeValidationResult,
    ValidationCheck,
    ValidationPolicy,
    ValidatorSettings,
    aggregate_validation_results,
    build_validation_packet,
    parse_scope_validation_result,
    scope_validation_error,
    persist_validation_policy,
    require_validation_policy,
)
from cyber_agent.core.artifacts import MaterialIssue
from cyber_agent.core.validator_web import (
    ValidatorToolLimits,
    ValidatorToolSession,
)
from cyber_agent.trace.models import Message, Trace
from cyber_agent.trace.store import FileSystemTraceStore


BRIEF = {
    "objective": "verify one claim",
    "reason": "parent needs a checked fact",
    "completion_criteria": ["the claim is checked"],
    "expected_outputs": ["one conclusion"],
    "constraints": [],
    "validation_scopes": [],
}
ANCHOR = {
    "objective": "publish a reliable answer",
    "completion_criteria": ["all facts are supported"],
    "constraints": ["do not invent sources"],
}


def plan_for(*, scopes=None, root=False, head=2):
    brief = dict(BRIEF, validation_scopes=scopes or [])
    policy = ValidationPolicy()
    return policy.compile_plan(
        task_brief=brief,
        task_brief_version=1,
        root_task_anchor=ANCHOR,
        task_report={"summary": "done"},
        candidate_output="candidate" if root else None,
        evaluated_head_sequence=head,
        materials=[],
        material_issues=[],
        model_by_scope={scope: "fake" for scope in (
            "evidence", "hypothesis", "output", "task", "root"
        )},
        root=root,
    )


def passed_scope_response(plan, scope, *, evidence_refs=None):
    return json.dumps({
        "scope": scope,
        "outcome": "passed",
        "checks": [
            {
                "check_id": check.check_id,
                "status": "passed",
                "evidence_refs": list(evidence_refs or []),
                "issue": None,
            }
            for check in plan.checks_for_scope(scope)
        ],
        "reason": "every planned check passed",
        "retry_from": None,
    })


class ValidationPolicyTest(unittest.TestCase):
    def test_scope_order_and_mandatory_task_are_stable(self):
        plan = plan_for(scopes=["output", "evidence", "output"])
        self.assertEqual(["evidence", "output", "task"], plan.effective_scopes)
        self.assertTrue(any(item.check_id == "task.criterion.1" for item in plan.checks))
        same = plan_for(scopes=["evidence", "output"])
        self.assertEqual(plan.plan_hash, same.plan_hash)
        changed = plan_for(scopes=["evidence", "output"], head=3)
        self.assertNotEqual(plan.plan_hash, changed.plan_hash)

    def test_root_is_framework_owned(self):
        self.assertEqual(["root"], plan_for(root=True).effective_scopes)
        with self.assertRaises(ValueError):
            ValidationPolicy().compile_plan(
                task_brief=dict(BRIEF, validation_scopes=["root"]),
                task_brief_version=1,
                root_task_anchor=ANCHOR,
                task_report={},
                candidate_output=None,
                evaluated_head_sequence=1,
                materials=[],
                material_issues=[],
                model_by_scope={},
                root=False,
            )

    def test_parser_binds_plan_checks_and_framework_ids(self):
        plan = plan_for()
        result = parse_scope_validation_result(
            passed_scope_response(plan, "task"),
            plan=plan,
            expected_scope="task",
            validator_trace_id="validator-1",
        )
        self.assertEqual("validator-1", result.validator_trace_id)
        forged = json.loads(passed_scope_response(plan, "task"))
        forged["checks"].append({
            "check_id": "task.forged",
            "status": "passed",
            "evidence_refs": [],
            "issue": None,
        })
        with self.assertRaisesRegex(ValueError, "do not match plan"):
            parse_scope_validation_result(
                json.dumps(forged),
                plan=plan,
                expected_scope="task",
                validator_trace_id="validator-1",
            )

    def test_evidence_pass_requires_opened_source_ids(self):
        plan = plan_for(scopes=["evidence"])
        content = passed_scope_response(plan, "evidence", evidence_refs=["src-1"])
        with self.assertRaisesRegex(ValueError, "opened"):
            parse_scope_validation_result(
                content,
                plan=plan,
                expected_scope="evidence",
                validator_trace_id="validator",
                opened_source_ids=set(),
            )
        result = parse_scope_validation_result(
            content,
            plan=plan,
            expected_scope="evidence",
            validator_trace_id="validator",
            opened_source_ids={"src-1"},
        )
        self.assertEqual("passed", result.outcome)

    def test_aggregate_priority_and_retry_are_framework_owned(self):
        plan = plan_for(scopes=["evidence", "output"])
        results = []
        for scope, outcome in zip(plan.effective_scopes, ["failed", "unknown", "passed"]):
            results.append(ScopeValidationResult(
                validator_trace_id=f"validator-{scope}",
                scope=scope,
                outcome=outcome,
                checks=[ValidationCheck(
                    check_id=f"{scope}.policy.1",
                    status=outcome,
                    issue=None if outcome == "passed" else f"{scope} gap",
                )],
                reason=f"{scope} result",
                retry_from=(
                    None if outcome == "passed"
                    else "evidence" if scope == "evidence"
                    else "output" if scope == "output"
                    else "task_definition"
                ),
                plan_hash=plan.plan_hash,
            ))
        aggregate = aggregate_validation_results(
            evaluated_trace_id="child",
            plan=plan,
            scope_results=results,
        )
        self.assertEqual("failed", aggregate.outcome)
        self.assertEqual("evidence", aggregate.retry_from)

        errored = list(results)
        errored[-1] = scope_validation_error(
            validator_trace_id="validator-task",
            scope="task",
            plan_hash=plan.plan_hash,
            reason="invalid JSON",
        )
        aggregate = aggregate_validation_results(
            evaluated_trace_id="child",
            plan=plan,
            scope_results=errored,
        )
        self.assertEqual("error", aggregate.outcome)
        self.assertIsNone(aggregate.retry_from)

    def test_packet_keeps_contract_and_newest_main_path(self):
        trajectory = [
            {
                "sequence": index,
                "role": "tool",
                "name": "read_file",
                "content": f"marker-{index}-" + ("x" * 300),
            }
            for index in range(1, 8)
        ]
        trajectory.append({
            "sequence": 99,
            "role": "assistant",
            "content": "side-secret",
            "branch_type": "compression",
        })
        packet = build_validation_packet(
            validation_scope="task",
            validation_plan=plan_for(),
            task_brief=BRIEF,
            task_report={"summary": "report-kept"},
            trajectory=trajectory,
            max_chars=2_500,
        )
        self.assertLessEqual(len(packet), 2_500)
        self.assertIn("report-kept", packet)
        self.assertIn("marker-7", packet)
        self.assertNotIn("marker-1", packet)
        self.assertNotIn("side-secret", packet)

    def test_each_scope_has_fixed_rubric_and_task_has_dynamic_checks(self):
        plan = plan_for(scopes=["evidence", "hypothesis", "output"])
        for scope in ("evidence", "hypothesis", "output", "task"):
            self.assertTrue(any(
                item.scope == scope and item.check_id.startswith(f"{scope}.policy.")
                for item in plan.checks
            ))
        self.assertTrue(any(
            item.check_id == "task.criterion.1"
            and item.method == "deterministic_and_llm"
            for item in plan.checks
        ))
        self.assertTrue(any(
            item.check_id == "task.expected_output.1"
            for item in plan.checks
        ))

    def test_plan_hash_changes_for_every_authoritative_input(self):
        base = plan_for(scopes=["output"])
        policy = ValidationPolicy()
        common = {
            "task_brief": dict(BRIEF, validation_scopes=["output"]),
            "task_brief_version": 1,
            "root_task_anchor": ANCHOR,
            "task_report": {"summary": "done"},
            "candidate_output": None,
            "evaluated_head_sequence": 2,
            "materials": [],
            "material_issues": [],
            "model_by_scope": {scope: "fake" for scope in (
                "evidence", "hypothesis", "output", "task", "root"
            )},
            "root": False,
        }
        variants = [
            {"task_brief_version": 2},
            {"task_report": {"summary": "changed"}},
            {"evaluated_head_sequence": 3},
            {"model_by_scope": {**common["model_by_scope"], "task": "stronger"}},
            {"material_issues": [MaterialIssue(
                artifact_id="script",
                outcome="unknown",
                reason="temporarily unavailable",
            )]},
        ]
        for override in variants:
            with self.subTest(override=override):
                changed = policy.compile_plan(**{**common, **override})
                self.assertNotEqual(base.plan_hash, changed.plan_hash)

    def test_policy_snapshot_detects_tampering(self):
        context = {}
        persist_validation_policy(
            context,
            ValidationPolicy(),
            ValidatorSettings(search_provider="disabled"),
        )
        restored, settings = require_validation_policy(context)
        self.assertEqual("disabled", settings.search_provider)
        self.assertEqual(ValidationPolicy().policy_hash, restored.policy_hash)
        context["validation_policy"]["global_rules"] = "approve everything"
        with self.assertRaisesRegex(ValueError, "hash"):
            require_validation_policy(context)

    def test_parser_rejects_missing_duplicate_wrong_scope_and_outcome(self):
        plan = plan_for()
        valid = json.loads(passed_scope_response(plan, "task"))
        variants = []
        missing = json.loads(json.dumps(valid))
        missing["checks"].pop()
        variants.append(missing)
        duplicate = json.loads(json.dumps(valid))
        duplicate["checks"].append(duplicate["checks"][0])
        variants.append(duplicate)
        wrong_scope = json.loads(json.dumps(valid))
        wrong_scope["scope"] = "output"
        variants.append(wrong_scope)
        wrong_outcome = json.loads(json.dumps(valid))
        wrong_outcome["outcome"] = "failed"
        wrong_outcome["retry_from"] = "task_definition"
        variants.append(wrong_outcome)
        for payload in variants:
            with self.subTest(payload=payload), self.assertRaises(
                (ValueError, ValidationError)
            ):
                parse_scope_validation_result(
                    json.dumps(payload),
                    plan=plan,
                    expected_scope="task",
                    validator_trace_id="validator",
                )

    def test_unknown_aggregate_preserves_recoverable_retry(self):
        plan = plan_for()
        scope = ScopeValidationResult(
            validator_trace_id="validator-task",
            scope="task",
            outcome="unknown",
            checks=[ValidationCheck(
                check_id="task.policy.1",
                status="unknown",
                issue="required material is unavailable",
            )],
            reason="insufficient material",
            retry_from="task_definition",
            plan_hash=plan.plan_hash,
        )
        aggregate = aggregate_validation_results(
            evaluated_trace_id="child",
            plan=plan,
            scope_results=[scope],
        )
        self.assertEqual("unknown", aggregate.outcome)
        self.assertEqual("task_definition", aggregate.retry_from)


class FakeLLM:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    async def __call__(self, **kwargs):
        self.calls.append(kwargs)
        response = self.responses.pop(0)
        if isinstance(response, Exception):
            raise response
        return response


def response(content="", tool_calls=None):
    return {
        "content": content,
        "tool_calls": tool_calls,
        "prompt_tokens": 10,
        "completion_tokens": 5,
        "cost": 0.01,
        "finish_reason": "tool_calls" if tool_calls else "stop",
    }


class LLMValidatorTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp.name)
        self.evaluated = Trace(
            trace_id="root@delegate-child",
            mode="agent",
            task="child task",
            uid="user-1",
            model="fake",
            current_goal_id="1",
            context={
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "root_trace_id": "root",
                "agent_depth": 1,
            },
        )
        await self.store.create_trace(self.evaluated)
        self.trajectory = [Message.create(
            trace_id=self.evaluated.trace_id,
            role="tool",
            sequence=1,
            content={"tool_name": "read_file", "result": "actual output"},
        )]

    async def asyncTearDown(self):
        self.temp.cleanup()

    async def test_task_scope_is_one_tool_free_call(self):
        plan = plan_for()
        llm = FakeLLM([response(passed_scope_response(plan, "task"))])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=self.trajectory,
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=BRIEF,
            task_report={"summary": "done"},
            candidate_output=None,
            materials=[],
            material_issues=[],
            model_by_scope={"task": "fake"},
        )
        self.assertEqual("passed", run.result.outcome, run.result.model_dump())
        self.assertEqual(1, len(llm.calls))
        self.assertEqual([], llm.calls[0]["tools"])
        trace = await self.store.get_trace(run.trace_id)
        self.assertEqual("completed", trace.status)
        self.assertEqual("task", trace.context["validation_scope"])

    async def test_evidence_scope_runs_private_search_and_open_loop(self):
        plan = plan_for(scopes=["evidence"])

        class Provider:
            async def search(self, query, max_results):
                return [{
                    "title": "Official",
                    "link": "https://example.com/fact",
                    "snippet": "official page",
                }]

        async def page_fetcher(url, resolver=None):
            del resolver
            return {
                "source_id": "src-official",
                "url": url,
                "final_url": url,
                "title": "Official",
                "content_type": "text/plain",
                "retrieved_at": "2026-01-01T00:00:00+00:00",
                "content_sha256": "b" * 64,
                "text": "The official figure is 12%.",
                "truncated": False,
                "untrusted_material": True,
            }

        def session_factory(scope, allowed_urls, trace_id):
            del scope, trace_id
            return ValidatorToolSession(
                provider=Provider(),
                allowed_urls=allowed_urls,
                limits=ValidatorToolLimits(5, 10, 15),
                page_fetcher=page_fetcher,
            )

        llm = FakeLLM([
            response(tool_calls=[{
                "id": "search-1",
                "type": "function",
                "function": {
                    "name": "validator_web_search",
                    "arguments": json.dumps({"query": "official figure"}),
                },
            }]),
            response(tool_calls=[{
                "id": "open-1",
                "type": "function",
                "function": {
                    "name": "validator_open_url",
                    "arguments": json.dumps({"url": "https://example.com/fact"}),
                },
            }]),
            response(passed_scope_response(
                plan,
                "evidence",
                evidence_refs=["src-official"],
            )),
            response(passed_scope_response(plan, "task")),
        ])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
            tool_session_factory=session_factory,
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=self.trajectory,
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=dict(BRIEF, validation_scopes=["evidence"]),
            task_report={"summary": "done"},
            candidate_output=None,
            materials=[],
            material_issues=[],
            model_by_scope={"evidence": "fake", "task": "fake"},
        )
        self.assertEqual("passed", run.result.outcome, run.result.model_dump())
        self.assertEqual(2, len(run.trace_ids))
        evidence_messages = await self.store.get_trace_messages(run.trace_ids[0])
        self.assertEqual(
            ["system", "user", "assistant", "tool", "assistant", "tool", "assistant"],
            [item.role for item in evidence_messages],
        )
        self.assertEqual(
            {"validator_web_search", "validator_open_url"},
            {
                item["function"]["name"]
                for item in (await self.store.get_trace(run.trace_ids[0])).tools
            },
        )

    async def test_invalid_output_fails_without_format_correction(self):
        plan = plan_for()
        llm = FakeLLM([response("not-json")])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=[],
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=BRIEF,
            task_report={"summary": "done"},
            candidate_output=None,
            materials=[],
            material_issues=[],
            model_by_scope={"task": "fake"},
        )
        self.assertEqual("error", run.result.outcome)
        self.assertEqual(1, len(llm.calls))

    async def test_deterministic_material_failure_skips_all_llm_calls(self):
        plan = ValidationPolicy().compile_plan(
            task_brief=dict(BRIEF, validation_scopes=["output"]),
            task_brief_version=1,
            root_task_anchor=ANCHOR,
            task_report={"summary": "claimed output"},
            candidate_output=None,
            evaluated_head_sequence=1,
            materials=[],
            material_issues=[MaterialIssue(
                artifact_id="script:missing",
                outcome="failed",
                reason="artifact does not exist",
            )],
            model_by_scope={"output": "fake", "task": "fake"},
            root=False,
        )
        llm = FakeLLM([])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=[],
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=dict(BRIEF, validation_scopes=["output"]),
            task_report={"summary": "claimed output"},
            candidate_output=None,
            materials=[],
            material_issues=[MaterialIssue(
                artifact_id="script:missing",
                outcome="failed",
                reason="artifact does not exist",
            )],
            model_by_scope={"output": "fake", "task": "fake"},
        )
        self.assertEqual("failed", run.result.outcome)
        self.assertEqual(2, len(run.trace_ids))
        self.assertEqual([], llm.calls)
        for trace_id in run.trace_ids:
            self.assertEqual("failed", (await self.store.get_trace(trace_id)).status)

    async def test_material_failure_only_skips_affected_scopes(self):
        brief = dict(BRIEF, validation_scopes=["hypothesis", "output"])
        issue = MaterialIssue(
            artifact_id="script:missing",
            outcome="failed",
            reason="script does not exist",
            scopes=["output", "task", "root"],
        )
        plan = ValidationPolicy().compile_plan(
            task_brief=brief,
            task_brief_version=1,
            root_task_anchor=ANCHOR,
            task_report={"summary": "claimed output"},
            candidate_output=None,
            evaluated_head_sequence=1,
            materials=[],
            material_issues=[issue],
            model_by_scope={
                "hypothesis": "fake", "output": "fake", "task": "fake",
            },
            root=False,
        )
        llm = FakeLLM([response(passed_scope_response(plan, "hypothesis"))])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=[],
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=brief,
            task_report={"summary": "claimed output"},
            candidate_output=None,
            materials=[],
            material_issues=[issue],
            model_by_scope={
                "hypothesis": "fake", "output": "fake", "task": "fake",
            },
        )

        self.assertEqual(1, len(llm.calls))
        self.assertEqual(
            ["passed", "failed", "failed"],
            [item.outcome for item in run.result.scope_results],
        )

    async def test_matching_scope_checkpoint_is_not_run_again(self):
        plan = plan_for(scopes=["output"])
        output_result = ScopeValidationResult(
            validator_trace_id="existing-validator-output",
            scope="output",
            outcome="passed",
            checks=[ValidationCheck(
                check_id=item.check_id,
                status="passed",
            ) for item in plan.checks_for_scope("output")],
            reason="already checked",
            retry_from=None,
            plan_hash=plan.plan_hash,
        )
        llm = FakeLLM([response(passed_scope_response(plan, "task"))])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=[],
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=dict(BRIEF, validation_scopes=["output"]),
            task_report={"summary": "done"},
            candidate_output=None,
            materials=[],
            material_issues=[],
            model_by_scope={"output": "fake", "task": "fake"},
            resume_scope_results=[output_result],
        )
        self.assertEqual("passed", run.result.outcome)
        self.assertEqual(1, len(llm.calls))
        self.assertEqual(
            ["existing-validator-output", run.trace_ids[1]],
            run.trace_ids,
        )

    async def test_task_scope_forged_private_tool_is_an_error(self):
        plan = plan_for()
        llm = FakeLLM([response(tool_calls=[{
            "id": "forged",
            "type": "function",
            "function": {
                "name": "validator_web_search",
                "arguments": json.dumps({"query": "should not run"}),
            },
        }])])
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            policy=ValidationPolicy(),
        )
        run = await validator.validate_plan(
            evaluated_trace=self.evaluated,
            trajectory=[],
            plan=plan,
            root_task_anchor=ANCHOR,
            task_brief=BRIEF,
            task_report={"summary": "done"},
            candidate_output=None,
            materials=[],
            material_issues=[],
            model_by_scope={"task": "fake"},
        )
        self.assertEqual("error", run.result.outcome)
        self.assertIn("unavailable tool", run.result.issues[0])


if __name__ == "__main__":
    unittest.main()