Agent/tests/test_recursive_context_integration.py

import json
import os
import re
import tempfile
import unittest
from unittest.mock import patch

from cyber_agent.core.context_policy import (
    canonical_json,
    context_ref_descriptors,
    require_root_task_anchor,
)
from cyber_agent.core.resource_budget import ResourceBudgetController
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.core.task_protocol import ensure_task_protocol
from cyber_agent.tools import get_tool_registry
from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
from cyber_agent.trace.store import FileSystemTraceStore


ROOT_ANCHOR = {
    "objective": "发布只使用官方数字的建议",
    "completion_criteria": [
        "最终建议采用官方12%",
        "明确30%仅见于转载且无官方支持",
    ],
    "constraints": ["不得编造证据"],
}


def recursive_env():
    return patch.dict(os.environ, {
        "AGENT_MODE": "recursive",
        "AGENT_RESOURCE_BUDGET_ENABLED": "true",
        "AGENT_MAX_TOTAL_AGENTS": "10",
        "AGENT_MAX_LLM_CALLS": "80",
        "AGENT_MAX_TOTAL_TOKENS": "100000",
        "AGENT_MAX_TOTAL_COST_USD": "10",
        "AGENT_MAX_DURATION_SECONDS": "600",
        "AGENT_RESERVED_FINAL_CALLS": "8",
        "AGENT_VALIDATOR_SEARCH_PROVIDER": "serper",
    }, clear=False)


def knowledge_disabled():
    return KnowledgeConfig(
        enable_extraction=False,
        enable_completion_extraction=False,
        enable_injection=False,
    )


def schema_names(schemas):
    return {
        schema["function"]["name"]
        for schema in schemas or []
    }


def message_text(messages):
    parts = []
    for message in messages:
        content = message.get("content", "")
        if isinstance(content, str):
            parts.append(content)
        elif isinstance(content, list):
            parts.extend(
                item.get("text", "")
                for item in content
                if isinstance(item, dict) and item.get("type") == "text"
            )
    return "\n".join(parts)


def trace_depth(messages):
    match = re.search(r"## Objective\s*\n\s*depth-(\d+)", message_text(messages))
    return int(match.group(1)) if match else 0


def last_tool_name(messages):
    for message in reversed(messages):
        if message.get("role") == "tool":
            return message.get("name")
    return None


def has_tool_result(messages, name):
    return any(
        message.get("role") == "tool" and message.get("name") == name
        for message in messages
    )


def ready_progress_call(call_index):
    return tool_call(
        "update_task_progress",
        {
            "expected_revision": 1,
            "progress": {
                "phase": "ready_to_submit",
                "questions": [],
                "blockers": [],
                "findings": [],
                "hypotheses": [],
                "work_items": [],
                "decision_rationale": "The delegated evidence chain is complete.",
            },
        },
        call_index,
    )


def pending_child_id(messages):
    for message in reversed(messages):
        if message.get("role") != "tool":
            continue
        content = message.get("content", "")
        text = content if isinstance(content, str) else json.dumps(content)
        match = re.search(r'"child_trace_id"\s*:\s*"([^"]+)"', text)
        if match:
            return match.group(1)
    raise AssertionError("agent tool result did not contain child_trace_id")


def available_ref(messages):
    text = message_text(messages)
    match = re.search(
        r'"ref_id":"([^"]+)","source_trace_id":"[^"]+",'
        r'"summary":"[^"]+","version":"([0-9a-f]{64})"',
        text,
    )
    if not match:
        # 字段按 sort_keys 排序，但摘要可能含有转义字符；用宽松回退只解析 ID/版本。
        ref_id = re.search(r'"ref_id":"([^"]+)"', text)
        version = re.search(r'"version":"([0-9a-f]{64})"', text)
        if not ref_id or not version:
            raise AssertionError("no authorized ContextRef in child task prompt")
        return ref_id.group(1), version.group(1)
    return match.group(1), match.group(2)


def tool_call(name, arguments, call_index):
    return {
        "content": "",
        "tool_calls": [{
            "id": f"call-{call_index}-{name}",
            "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 FiveLevelRecursiveContextTest(unittest.IsolatedAsyncioTestCase):
    async def test_real_runner_keeps_anchor_refs_permissions_and_reviews_through_depth_five(self):
        with tempfile.TemporaryDirectory() as directory:
            store = FileSystemTraceStore(directory)
            call_count = 0
            token_total = 0
            read_depths = []
            observed_tools = {}
            validator_packets = []
            validator_packet_keys = set()
            search_queries = []

            class SearchProvider:
                async def search(self, query, max_results):
                    search_queries.append((query, max_results))
                    return [
                        {
                            "title": "转载页面",
                            "link": "https://repost.example/30-percent",
                            "snippet": "未经核实的30%",
                        },
                        {
                            "title": "官方页面",
                            "link": "https://official.example/12-percent",
                            "snippet": "官方数字12%",
                        },
                    ][: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": "c" * 64,
                    "text": "官方数字是12%，不支持30%。",
                    "truncated": False,
                    "untrusted_material": True,
                }

            async def fake_llm(**kwargs):
                nonlocal call_count, token_total
                call_count += 1
                messages = kwargs["messages"]
                tools = schema_names(kwargs.get("tools"))
                depth = trace_depth(messages)
                observed_tools.setdefault(depth, []).append(tools)

                if any(
                    message.get("role") == "system"
                    and "recursive_validation_protocol" in str(message.get("content", ""))
                    for message in messages
                ):
                    packet = next(
                        json.loads(message["content"])
                        for message in messages
                        if message.get("role") == "user"
                        and "recursive_validation_protocol" in str(message.get("content", ""))
                    )
                    scope = packet["validation_scope"]
                    packet_key = (packet["validation_plan"]["plan_hash"], scope)
                    if packet_key not in validator_packet_keys:
                        validator_packet_keys.add(packet_key)
                        validator_packets.append(packet)
                    if scope == "root":
                        self.assertIn("官方12%", packet["candidate_output"])
                        self.assertIn("30%仅见于转载", packet["candidate_output"])
                    else:
                        brief = packet["task_brief"]
                        local_depth = brief["context"]["local_depth"]
                        report_text = json.dumps(
                            packet["task_report"], ensure_ascii=False,
                        )
                        self.assertIn(
                            self._semantic_result(local_depth),
                            report_text,
                        )
                        self.assertIn(
                            scope,
                            [*brief["validation_scopes"], "task"],
                        )
                        task_criteria = [
                            item["criterion"]
                            for item in packet["validation_plan"]["checks"]
                            if item["check_id"].startswith("task.criterion.")
                        ]
                        self.assertEqual(
                            brief["completion_criteria"],
                            task_criteria,
                        )
                        self.assertTrue(
                            set(ROOT_ANCHOR["completion_criteria"]).isdisjoint(
                                task_criteria
                            )
                        )
                    validator_last_tool = last_tool_name(messages)
                    if scope == "evidence" and validator_last_tool is None:
                        response = tool_call(
                            "validator_web_search",
                            {"query": "官方数字 12% 30%", "max_results": 5},
                            call_count,
                        )
                        token_total += 5
                        return response
                    if (
                        scope == "evidence"
                        and validator_last_tool == "validator_web_search"
                    ):
                        response = tool_call(
                            "validator_open_url",
                            {"url": "https://official.example/12-percent"},
                            call_count,
                        )
                        token_total += 5
                        return response
                    evidence_refs = (
                        ["src-official-12"] if scope == "evidence" else []
                    )
                    response = {
                        "content": json.dumps({
                            "outcome": "passed",
                            "scope": scope,
                            "checks": [
                                {
                                    "check_id": item["check_id"],
                                    "status": "passed",
                                    "evidence_refs": evidence_refs,
                                    "issue": None,
                                }
                                for item in packet["validation_plan"]["checks"]
                                if item["scope"] == scope
                            ],
                            "reason": "已核对持久化轨迹、标准和输出",
                            "retry_from": None,
                        }, ensure_ascii=False),
                        "tool_calls": [],
                        "finish_reason": "stop",
                        "prompt_tokens": 3,
                        "completion_tokens": 2,
                        "cost": 0.0001,
                    }
                    token_total += 5
                    return response

                last_tool = last_tool_name(messages)
                if tools == {"review_task_result", "read_context_ref"}:
                    response = tool_call(
                        "review_task_result",
                        {
                            "child_trace_id": pending_child_id(messages),
                            "decision": "ASCEND",
                            "reason": f"depth-{depth + 1} 已通过独立验收",
                        },
                        call_count,
                    )
                elif depth > 0 and "submit_task_report" not in tools:
                    response = {
                        "content": f"depth-{depth} TaskReport 已提交",
                        "tool_calls": [],
                        "finish_reason": "stop",
                        "prompt_tokens": 3,
                        "completion_tokens": 2,
                        "cost": 0.0001,
                    }
                elif last_tool == "review_task_result":
                    response = ready_progress_call(call_count)
                elif last_tool == "update_task_progress":
                    if depth == 0:
                        response = {
                            "content": (
                                "最终建议采用官方12%；30%仅见于转载，"
                                "没有官方支持，不能作为发布依据。"
                            ),
                            "tool_calls": [],
                            "finish_reason": "stop",
                            "prompt_tokens": 3,
                            "completion_tokens": 2,
                            "cost": 0.0001,
                        }
                    else:
                        response = tool_call(
                            "submit_task_report",
                            {"task_report": self._report(depth)},
                            call_count,
                        )
                elif depth in {4, 5} and last_tool != "read_context_ref":
                    ref_id, version = available_ref(messages)
                    read_depths.append(depth)
                    response = tool_call(
                        "read_context_ref",
                        {"ref_id": ref_id, "version": version},
                        call_count,
                    )
                elif depth == 5:
                    if has_tool_result(messages, "update_task_progress"):
                        response = tool_call(
                            "submit_task_report",
                            {"task_report": self._report(depth)},
                            call_count,
                        )
                    else:
                        response = ready_progress_call(call_count)
                else:
                    next_depth = depth + 1
                    response = tool_call(
                        "agent",
                        {"task_brief": self._brief(next_depth)},
                        call_count,
                    )
                token_total += 5
                return response

            runner = AgentRunner(
                trace_store=store,
                tool_registry=get_tool_registry(),
                llm_call=fake_llm,
                validator_search_provider=SearchProvider(),
                validator_page_fetcher=page_fetcher,
            )
            config = RunConfig(
                tools=[
                    "agent",
                    "submit_task_report",
                    "review_task_result",
                    "read_context_ref",
                    "update_task_progress",
                ],
                tool_groups=[],
                enable_research_flow=False,
                root_task_anchor=ROOT_ANCHOR,
                knowledge=knowledge_disabled(),
                child_execution_mode="sequential",
            )
            read_stats_before = get_tool_registry().get_stats("read_context_ref")[
                "read_context_ref"
            ]["call_count"]
            with recursive_env():
                result = await runner.run_result(
                    [{"role": "user", "content": "请逐层拆解并验收根任务"}],
                    config,
                )

            self.assertEqual("completed", result["status"])
            traces = await store.list_traces(limit=100)
            business = [
                trace for trace in traces
                if trace.context.get("created_by_tool") == "agent"
                or trace.trace_id == result["trace_id"]
            ]
            business.sort(key=lambda trace: trace.context["agent_depth"])
            self.assertEqual(list(range(6)), [trace.context["agent_depth"] for trace in business])
            self.assertEqual(6, len(business))
            for index, trace in enumerate(business):
                self.assertEqual(result["trace_id"], trace.context["root_trace_id"])
                self.assertEqual(
                    canonical_json(ROOT_ANCHOR),
                    canonical_json(require_root_task_anchor(trace.context).model_dump(mode="json")),
                )
                messages = await store.get_trace_messages(trace.trace_id)
                first_user = next(message for message in messages if message.role == "user")
                first_user_text = (
                    first_user.content
                    if isinstance(first_user.content, str)
                    else json.dumps(first_user.content, ensure_ascii=False)
                )
                self.assertEqual(1, first_user_text.count("# Root Task Anchor"))
                metrics = trace.context["context_access"]["metrics"]
                self.assertGreater(metrics["root_anchor_chars"], 0)
                self.assertEqual(
                    len(context_ref_descriptors(trace.context)),
                    metrics["authorized_ref_count"],
                )
                if index:
                    self.assertEqual(business[index - 1].trace_id, trace.parent_trace_id)
                    state = ensure_task_protocol(trace.context)
                    self.assertEqual(1, state["task_brief_version"])
                    self.assertEqual(
                        ["不得编造证据", *[f"约束-{level}" for level in range(1, index + 1)]],
                        state["task_brief"]["constraints"],
                    )
                    self.assertEqual(
                        [f"直接父级结论 depth-{index - 1}"],
                        state["task_brief"]["parent_findings"],
                    )
                    self.assertEqual(
                        {"local_depth": index},
                        state["task_brief"]["context"],
                    )
                    self.assertEqual([], state["task_brief"]["context_refs"])
                    descriptors = context_ref_descriptors(trace.context)
                    expected_kinds = (
                        ["reviewed_task_result"]
                        if index == 1
                        else ["task_brief"]
                        if index == 5
                        else ["task_brief", "reviewed_task_result"]
                    )
                    self.assertEqual(
                        expected_kinds,
                        [item["kind"] for item in descriptors],
                    )

            self.assertEqual([4, 5], read_depths)
            read_stats_after = get_tool_registry().get_stats("read_context_ref")[
                "read_context_ref"
            ]["call_count"]
            self.assertEqual(2, read_stats_after - read_stats_before)
            self.assertTrue(any("agent" in tools for tools in observed_tools[4]))
            self.assertTrue(all("evaluate" not in tools and "bash_command" not in tools for tools in observed_tools[5]))
            self.assertTrue(all("agent" not in tools for tools in observed_tools[5]))
            self.assertTrue(any("read_context_ref" in tools for tools in observed_tools[5]))

            validators = [
                trace for trace in traces
                if trace.context.get("created_by_tool") == "validator"
            ]
            self.assertEqual(11, len(validators))
            self.assertEqual(11, len(validator_packets))
            self.assertTrue(all(packet["root_task_anchor"] == ROOT_ANCHOR for packet in validator_packets))
            self.assertEqual(2, len(search_queries))
            self.assertEqual(
                ["evidence", "hypothesis", "output", "root", "task"],
                sorted({packet["validation_scope"] for packet in validator_packets}),
            )
            packets_with_real_ref_reads = [
                packet
                for packet in validator_packets
                if any(
                    item.get("role") == "tool"
                    and item.get("name") == "read_context_ref"
                    for item in packet["trajectory"]
                )
            ]
            self.assertEqual(4, len(packets_with_real_ref_reads))
            self.assertTrue(ensure_task_protocol(business[0].context)["root_validation_passed"])
            usage = await ResourceBudgetController(store).get_usage(result["trace_id"])
            self.assertEqual(6, usage.total_agents)
            self.assertEqual(call_count, usage.llm_calls)
            self.assertEqual(token_total, usage.total_tokens)
            self.assertEqual(4, usage.validation_tool_calls)
            self.assertGreater(usage.validation_material_chars, 0)

    @staticmethod
    def _brief(depth):
        definitions = {
            1: (
                "汇总最终发布建议，只采用通过验证的官方数字",
                "建议采用官方12%并说明30%没有官方支持",
                "可发布的最终建议",
                ["output"],
            ),
            2: (
                "判断提升30%的假设是否获得证据支持",
                "区分30%转载说法与12%官方数字",
                "假设取舍结论",
                ["hypothesis"],
            ),
            3: (
                "生成可用和禁用的数字表达",
                "可用表达只采用官方12%，禁用官方30%表达",
                "可用与禁用表达清单",
                ["output"],
            ),
            4: (
                "建立数字与来源的对应关系",
                "明确官方来源对应12%、转载对应30%",
                "数字来源对应表",
                ["evidence"],
            ),
            5: (
                "核实官方页面真实支持的数字",
                "确认官方只支持12%且不支持30%",
                "官方数字核查结论",
                ["evidence"],
            ),
        }
        objective, criterion, expected_output, validation_scopes = definitions[depth]
        return {
            "objective": f"depth-{depth} {objective}",
            "reason": f"depth-{depth - 1} 需要直属子任务结果",
            "completion_criteria": [criterion],
            "expected_outputs": [expected_output],
            "parent_findings": [f"直接父级结论 depth-{depth - 1}"],
            "context": {"local_depth": depth},
            "constraints": [f"约束-{depth}"],
            "validation_scopes": validation_scopes,
        }

    @staticmethod
    def _report(depth):
        result = FiveLevelRecursiveContextTest._semantic_result(depth)
        return {
            "summary": result,
            "outcome": "satisfied",
            "validation": {"hard_passed": True, "open_issues": []},
            "next_step_suggestion": {
                "direction": "ASCEND",
                "reason": "当前层已满足完成标准",
            },
            "outputs": [{"depth": depth, "result": result}],
            "evidence": [{"depth": depth, "source": "官方页面与转载页面"}],
            "remaining_issues": [],
        }

    @staticmethod
    def _semantic_result(depth):
        return {
            1: "最终建议采用官方12%，并明确30%仅见于转载、无官方支持",
            2: "30%假设无官方证据，采用12%有官方来源",
            3: "可用表达：官方12%；禁用表达：官方30%",
            4: "官方来源→12%；二次转载→30%，二者不可混用",
            5: "官方页面只支持12%，不支持30%",
        }[depth]


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