Переглянути джерело

增加五层 recursive 上下文传递集成测试

使用真实 Runner、TraceStore、任务协议、Validator 和资源预算模拟 root 到 depth 5 的完整调用链。

断言根锚点不漂移、引用逐层授权、depth 5 工具边界、逐级审核和最终资源账本均符合 Recursive 约束。
SamLee 20 годин тому
батько
коміт
b9ea50d27a
1 змінених файлів з 381 додано та 0 видалено
  1. 381 0
      tests/test_recursive_context_integration.py

+ 381 - 0
tests/test_recursive_context_integration.py

@@ -0,0 +1,381 @@
+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()