|
|
@@ -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()
|