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": ["五层任务都经过直接父级审核", "根结果通过独立验收"],
    "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": "1",
    }, 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 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 = []

            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 "independent validator" in str(message.get("content", ""))
                    for message in messages
                ):
                    packet = json.loads(messages[-1]["content"])
                    validator_packets.append(packet)
                    response = {
                        "content": json.dumps({
                            "outcome": "passed",
                            "scope": packet["validation_scope"],
                            "reason": "已核对持久化轨迹、标准和输出",
                            "issues": [],
                            "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":
                    if depth == 0:
                        response = {
                            "content": "根任务的五层结果已逐层审核完成",
                            "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:
                    response = tool_call(
                        "submit_task_report",
                        {"task_report": self._report(depth)},
                        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,
            )
            config = RunConfig(
                tools=[
                    "agent",
                    "submit_task_report",
                    "review_task_result",
                    "read_context_ref",
                ],
                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(6, len(validators))
            self.assertEqual(6, len(validator_packets))
            self.assertTrue(all(packet["root_task_anchor"] == ROOT_ANCHOR 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(2, 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)

    @staticmethod
    def _brief(depth):
        return {
            "objective": f"depth-{depth}",
            "reason": f"depth-{depth - 1} 需要直属子任务结果",
            "completion_criteria": [f"depth-{depth} 结果通过验收"],
            "expected_outputs": [f"depth-{depth} 可追溯结论"],
            "parent_findings": [f"直接父级结论 depth-{depth - 1}"],
            "context": {"local_depth": depth},
            "constraints": [f"约束-{depth}"],
        }

    @staticmethod
    def _report(depth):
        return {
            "summary": f"depth-{depth} 局部任务完成",
            "outcome": "satisfied",
            "validation": {"hard_passed": True, "open_issues": []},
            "next_step_suggestion": {
                "direction": "ASCEND",
                "reason": "当前层已满足完成标准",
            },
            "outputs": [{"depth": depth, "result": f"结论-{depth}"}],
            "evidence": [{"depth": depth, "source": "fake-persisted-trace"}],
            "remaining_issues": [],
        }


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