Agent/tests/test_recursive_validator_lifecycle_complex.py

import asyncio
import json
import tempfile
import unittest

from cyber_agent.core.artifacts import (
    ArtifactRef,
    ValidationMaterial,
    canonical_material_content,
    material_content_hash,
)
from cyber_agent.core.context_policy import persist_root_task_anchor
from cyber_agent.core.resource_budget import (
    RESOURCE_BUDGET_CONTEXT_KEY,
    ResourceBudget,
    ResourceBudgetController,
)
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.core.task_protocol import (
    TaskReport,
    ensure_task_protocol,
    new_task_protocol,
    pending_review_entry,
    replace_task_brief,
)
from cyber_agent.core.validation import (
    ValidationPolicy,
    ValidatorSettings,
    persist_validation_policy,
)
from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
from cyber_agent.tools.builtin.task_protocol import review_task_result
from cyber_agent.tools.models import ToolResult
from cyber_agent.tools.registry import ToolRegistry
from cyber_agent.trace.goal_models import Goal, GoalTree
from cyber_agent.trace.models import Message, Trace
from cyber_agent.trace.store import FileSystemTraceStore


ANCHOR = {
    "objective": "发布只使用官方数字的完整建议",
    "completion_criteria": ["必须使用官方12%", "必须说明30%无官方支持"],
    "constraints": ["不得把转载当成官方来源"],
}


def brief(depth, *, evidence=False, version=1):
    value = "12%" if version == 2 else "30%"
    return {
        "objective": f"depth-{depth} 核实局部数字 {value}",
        "reason": f"depth-{depth - 1} 需要直属子任务的证据",
        "completion_criteria": [f"确认 {value} 是否有官方支持"],
        "expected_outputs": ["一条可追溯结论"],
        "constraints": ["不得猜测"],
        "validation_scopes": ["evidence"] if evidence else [],
    }


def report(trace_id, summary):
    return TaskReport.model_validate({
        "child_trace_id": trace_id,
        "summary": summary,
        "outcome": "satisfied",
        "validation": {"hard_passed": True, "open_issues": []},
        "next_step_suggestion": {
            "direction": "ASCEND",
            "reason": "当前局部任务已完成",
        },
        "outputs": [{"conclusion": summary}],
        "evidence": [{"source": "validator-will-check"}],
        "remaining_issues": [],
    })


def validation_packet(messages):
    return next(
        json.loads(message["content"])
        for message in messages
        if message.get("role") == "user"
        and "recursive_validation_protocol" in str(message.get("content", ""))
    )


def check_response(packet, outcome="passed", issue=None, evidence_refs=None):
    scope = packet["validation_scope"]
    return {
        "content": json.dumps({
            "scope": scope,
            "outcome": outcome,
            "checks": [
                {
                    "check_id": item["check_id"],
                    "status": outcome,
                    "evidence_refs": list(evidence_refs or []),
                    "issue": None if outcome == "passed" else issue,
                }
                for item in packet["validation_plan"]["checks"]
                if item["scope"] == scope
            ],
            "reason": "验收通过" if outcome == "passed" else issue,
            "retry_from": None if outcome == "passed" else (
                "task_definition" if scope in {"task", "root"} else scope
            ),
        }, ensure_ascii=False),
        "tool_calls": [],
        "finish_reason": "stop",
        "prompt_tokens": 3,
        "completion_tokens": 2,
        "cost": 0.0001,
    }


def tool_call(name, arguments):
    return {
        "content": "",
        "tool_calls": [{
            "id": f"{name}-call",
            "type": "function",
            "function": {
                "name": name,
                "arguments": json.dumps(arguments, ensure_ascii=False),
            },
        }],
        "finish_reason": "tool_calls",
        "prompt_tokens": 3,
        "completion_tokens": 2,
        "cost": 0.0001,
    }


class SearchProvider:
    async def search(self, _query, max_results):
        return [
            {
                "title": "官方数字",
                "link": "https://official.example/12",
                "snippet": "12%",
            },
            {
                "title": "二次转载",
                "link": "https://repost.example/30",
                "snippet": "30%",
            },
        ][:max_results]


async def page_fetcher(url, resolver=None):
    del resolver
    return {
        "source_id": "src-official-12",
        "url": url,
        "final_url": url,
        "title": "官方数字",
        "content_type": "text/plain",
        "retrieved_at": "2026-07-17T00:00:00+00:00",
        "content_sha256": "d" * 64,
        "text": "官方页面只支持12%，没有30%。",
        "truncated": False,
        "untrusted_material": True,
    }


class FiveLevelValidatorLifecycleTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp.name)
        self.root_id = "complex-root"
        self.ids = [self.root_id, *[f"depth-{depth}" for depth in range(1, 6)]]
        budget = ResourceBudget()
        root_context = {
            "agent_mode": "recursive",
            "agent_mode_revision": 2,
            "agent_depth": 0,
            "root_trace_id": self.root_id,
            RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
            "task_protocol": new_task_protocol(),
        }
        persist_root_task_anchor(root_context, ANCHOR)
        persist_validation_policy(
            root_context,
            ValidationPolicy(),
            ValidatorSettings(search_provider="serper"),
        )
        await self.store.create_trace(Trace(
            trace_id=self.root_id,
            mode="agent",
            task="root",
            uid="user-1",
            model="fake",
            context=root_context,
        ))
        await self.store.add_message(Message.create(
            trace_id=self.root_id,
            role="user",
            sequence=1,
            content="生成可发布建议",
        ))
        await self.store.update_trace(self.root_id, head_sequence=1)
        for depth in range(1, 6):
            context = {
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "agent_depth": depth,
                "root_trace_id": self.root_id,
                "created_by_tool": "agent",
                "task_protocol": new_task_protocol(
                    brief(depth, evidence=(depth == 5))
                ),
            }
            persist_root_task_anchor(context, ANCHOR)
            await self.store.create_trace(Trace(
                trace_id=self.ids[depth],
                mode="agent",
                task=f"depth-{depth}",
                parent_trace_id=self.ids[depth - 1],
                parent_goal_id="1",
                uid="user-1",
                model="fake",
                context=context,
            ))
            await self.store.add_message(Message.create(
                trace_id=self.ids[depth],
                role="user",
                sequence=1,
                content=f"execute depth-{depth}",
            ))
            await self.store.update_trace(self.ids[depth], head_sequence=1)
        await ResourceBudgetController(self.store).initialize(
            self.root_id,
            budget,
            initial_agents=6,
        )

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

    async def test_deepest_evidence_failure_revises_only_direct_child(self):
        async def llm_call(**kwargs):
            messages = kwargs["messages"]
            packet = validation_packet(messages)
            scope = packet["validation_scope"]
            if scope == "evidence":
                tool_names = [
                    message.get("name")
                    for message in messages
                    if message.get("role") == "tool"
                ]
                if not tool_names:
                    return tool_call(
                        "validator_web_search",
                        {"query": "official 30% 12%", "max_results": 5},
                    )
                if tool_names[-1] == "validator_web_search":
                    return tool_call(
                        "validator_open_url",
                        {"url": "https://official.example/12"},
                    )
                summary = packet["task_report"]["summary"]
                if "30%" in summary:
                    return check_response(
                        packet,
                        outcome="failed",
                        issue="官方来源只支持12%，不支持30%",
                        evidence_refs=["src-official-12"],
                    )
                return check_response(
                    packet,
                    evidence_refs=["src-official-12"],
                )
            return check_response(packet)

        runner = AgentRunner(
            trace_store=self.store,
            llm_call=llm_call,
            validator_search_provider=SearchProvider(),
            validator_page_fetcher=page_fetcher,
        )
        child_id = self.ids[5]
        parent_id = self.ids[4]
        first_report = report(child_id, "官方支持30%")
        child = await self.store.get_trace(child_id)
        ensure_task_protocol(child.context)["task_report"] = first_report.model_dump()
        await self.store.update_trace(child_id, context=child.context)

        first = await runner.validate_recursive_trace(
            child_id,
            task_report=first_report.model_dump(mode="json"),
        )
        self.assertEqual("failed", first.result.outcome)
        self.assertEqual(
            ["failed", "passed"],
            [item.outcome for item in first.result.scope_results],
        )

        parent = await self.store.get_trace(parent_id)
        parent_state = ensure_task_protocol(parent.context)
        parent_state["pending_reviews"][child_id] = pending_review_entry(
            goal_id="1",
            report=first_report,
            validation_result=first.result.model_dump(mode="json"),
            received_at_sequence=10,
        )
        await self.store.update_trace(parent_id, context=parent.context)
        tree = GoalTree(
            mission="depth-4",
            goals=[Goal(id="1", description="verify depth-5", status="pending_review")],
            current_id=None,
        )
        await self.store.update_goal_tree(parent_id, tree)
        reviewed = await review_task_result(
            child_trace_id=child_id,
            decision="REVISE_CHILD",
            reason="改为官方可支持的12%",
            approved_next_task=brief(5, evidence=True, version=2),
            context={
                "store": self.store,
                "trace_id": parent_id,
                "goal_id": "1",
                "goal_tree": tree,
                "sequence": 11,
            },
        )
        self.assertEqual("completed", reviewed["status"])
        child = await self.store.get_trace(child_id)
        child_state = ensure_task_protocol(child.context)
        self.assertIsNone(child_state["task_report"])
        self.assertIsNone(child_state["task_report_validation"])
        self.assertEqual(1, len(child_state["report_history"]))
        parent = await self.store.get_trace(parent_id)
        parent_state = ensure_task_protocol(parent.context)
        self.assertEqual("REVISE_CHILD", parent_state["next_actions"][0]["decision"])

        revised_brief = parent_state["next_actions"][0]["task_brief"]
        replace_task_brief(child_state, revised_brief, effective_at_sequence=12)
        second_report = report(child_id, "官方只支持12%")
        child_state["task_report"] = second_report.model_dump()
        parent_state["next_actions"].clear()
        await self.store.update_trace(child_id, context=child.context)
        await self.store.update_trace(parent_id, context=parent.context)
        second = await runner.validate_recursive_trace(
            child_id,
            task_report=second_report.model_dump(mode="json"),
        )
        self.assertEqual("passed", second.result.outcome)
        self.assertNotEqual(first.result.plan_hash, second.result.plan_hash)
        self.assertEqual(
            4,
            len(await self.store.list_traces(
                parent_trace_id=child_id,
                created_by_tool="validator",
                limit=10,
            )),
        )
        usage = await ResourceBudgetController(self.store).get_usage(self.root_id)
        self.assertEqual(6, usage.total_agents)

    async def test_parent_stop_during_validator_page_fetch_discards_tool_result(self):
        page_started = asyncio.Event()
        release_page = asyncio.Event()
        page_text = "官方页面只支持12%，没有30%。"

        async def blocking_page_fetcher(url, resolver=None):
            del resolver
            page_started.set()
            await release_page.wait()
            return {
                "source_id": "src-official-12",
                "url": url,
                "final_url": url,
                "title": "官方数字",
                "content_type": "text/plain",
                "retrieved_at": "2026-07-17T00:00:00+00:00",
                "content_sha256": "d" * 64,
                "text": page_text,
                "truncated": False,
                "untrusted_material": True,
            }

        llm_calls = 0

        async def llm_call(**kwargs):
            nonlocal llm_calls
            llm_calls += 1
            packet = validation_packet(kwargs["messages"])
            self.assertEqual("evidence", packet["validation_scope"])
            return tool_call(
                "validator_open_url",
                {"url": "https://official.example/12"},
            )

        runner = AgentRunner(
            trace_store=self.store,
            llm_call=llm_call,
            validator_page_fetcher=blocking_page_fetcher,
        )
        root_event = asyncio.Event()
        runner._cancel_events[self.root_id] = root_event
        runner._recursive_active_traces[self.root_id] = root_event
        registered = []
        for parent_id, child_id in zip(self.ids, self.ids[1:]):
            registered.append((
                child_id,
                runner.register_recursive_child(parent_id, child_id),
            ))

        child_id = self.ids[5]
        payload = report(child_id, "官方只支持12%").model_dump(mode="json")
        payload["source_urls"] = ["https://official.example/12"]
        task = asyncio.create_task(runner.validate_recursive_trace(
            child_id,
            task_report=payload,
        ))
        await asyncio.wait_for(page_started.wait(), timeout=1)
        self.assertTrue(await runner.stop(self.root_id))
        release_page.set()
        run = await asyncio.wait_for(task, timeout=2)

        self.assertEqual("error", run.result.outcome)
        self.assertEqual(1, llm_calls)
        validators = await self.store.list_traces(
            parent_trace_id=child_id,
            created_by_tool="validator",
            limit=10,
        )
        evidence_trace = next(
            trace for trace in validators
            if trace.context["validation_scope"] == "evidence"
        )
        messages = await self.store.get_trace_messages(evidence_trace.trace_id)
        self.assertNotIn("tool", [message.role for message in messages])
        child = await self.store.get_trace(child_id)
        cache = ensure_task_protocol(child.context)["task_report_validation"]
        self.assertEqual("error", cache["aggregate_result"]["outcome"])
        usage = await ResourceBudgetController(self.store).get_usage(self.root_id)
        self.assertEqual(len(page_text), usage.validation_material_chars)

        for trace_id, event in reversed(registered):
            runner.release_recursive_trace(trace_id, event)
        runner.release_recursive_trace(self.root_id, root_event)

    async def test_all_children_pass_but_root_must_revise_its_own_output(self):
        root_validation_attempt = 0
        ordinary_calls = 0
        validated_versions = []
        scripts = {
            "v1": "建议使用官方12%。",
            "v2": "建议使用官方12%；30%只见于转载，无官方支持。",
        }
        refs = {
            version: ArtifactRef(
                artifact_id="script:final",
                version=version,
                content_hash=material_content_hash(content),
                kind="script",
                mime_type="text/plain",
            )
            for version, content in scripts.items()
        }

        class ScriptResolver:
            async def resolve(self, ref, root_trace_id, uid):
                return ValidationMaterial(
                    **ref.model_dump(mode="json"),
                    root_trace_id=root_trace_id,
                    uid=uid,
                    content=scripts[ref.version],
                )

        registry = ToolRegistry()

        async def publish_script(version: str):
            return ToolResult(
                title=f"script {version}",
                output=scripts[version],
                artifact_refs=[refs[version].model_dump(mode="json")],
            )

        registry.register(publish_script)

        async def llm_call(**kwargs):
            nonlocal root_validation_attempt, ordinary_calls
            messages = kwargs["messages"]
            if messages and "recursive_validation_protocol" in messages[0]["content"]:
                packet = validation_packet(messages)
                if packet["validation_scope"] == "root":
                    root_validation_attempt += 1
                    validated_versions.append([
                        item["version"] for item in packet["materials"]
                    ])
                    if root_validation_attempt == 1:
                        return check_response(
                            packet,
                            outcome="failed",
                            issue="根输出遗漏了30%无官方支持的风险说明",
                        )
                return check_response(packet)
            ordinary_calls += 1
            if ordinary_calls in {1, 3}:
                version = "v1" if ordinary_calls == 1 else "v2"
                return tool_call("publish_script", {"version": version})
            return {
                "content": scripts["v1" if ordinary_calls == 2 else "v2"],
                "tool_calls": [],
                "finish_reason": "stop",
                "prompt_tokens": 3,
                "completion_tokens": 2,
                "cost": 0.0001,
            }

        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=registry,
            llm_call=llm_call,
            artifact_resolver=ScriptResolver(),
        )
        for depth in range(1, 6):
            trace_id = self.ids[depth]
            child = await self.store.get_trace(trace_id)
            child_state = ensure_task_protocol(child.context)
            if depth == 5:
                replace_task_brief(
                    child_state,
                    brief(depth, evidence=False),
                    effective_at_sequence=2,
                )
            child_report = report(trace_id, f"depth-{depth} 局部交付准确")
            child_state["task_report"] = child_report.model_dump()
            await self.store.update_trace(trace_id, context=child.context)
            run = await runner.validate_recursive_trace(
                trace_id,
                task_report=child_report.model_dump(mode="json"),
            )
            self.assertEqual("passed", run.result.outcome)

        result = await runner.run_result(
            [],
            RunConfig(
                trace_id=self.root_id,
                tools=["publish_script"],
                tool_groups=[],
                enable_research_flow=False,
                knowledge=KnowledgeConfig(
                    enable_extraction=False,
                    enable_completion_extraction=False,
                    enable_injection=False,
                ),
            ),
        )
        self.assertEqual("completed", result["status"])
        root = await self.store.get_trace(self.root_id)
        state = ensure_task_protocol(root.context)
        self.assertEqual(2, state["root_validation_attempts"])
        self.assertEqual(
            ["failed", "passed"],
            [item["outcome"] for item in state["root_validation_history"]],
        )
        self.assertTrue(state["root_validation_passed"])
        self.assertEqual(4, ordinary_calls)
        self.assertEqual(2, root_validation_attempt)
        self.assertEqual([["v1"], ["v1", "v2"]], validated_versions)
        usage = await ResourceBudgetController(self.store).get_usage(self.root_id)
        self.assertEqual(6, usage.total_agents)
        expected_material_chars = (
            2 * len(canonical_material_content(scripts["v1"]))
            + len(canonical_material_content(scripts["v2"]))
        )
        self.assertEqual(expected_material_chars, usage.validation_material_chars)

        restarted_runner = AgentRunner(
            trace_store=self.store,
            tool_registry=registry,
            llm_call=llm_call,
            artifact_resolver=ScriptResolver(),
        )
        cached = await restarted_runner.validate_recursive_trace(
            self.root_id,
            scope="root",
            candidate_output=scripts["v2"],
            root_validator=True,
        )
        self.assertTrue(cached.cached)
        after_cache = await ResourceBudgetController(self.store).get_usage(self.root_id)
        self.assertEqual(
            expected_material_chars,
            after_cache.validation_material_chars,
        )


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