|
|
@@ -0,0 +1,907 @@
|
|
|
+import asyncio
|
|
|
+import json
|
|
|
+import os
|
|
|
+import tempfile
|
|
|
+import unittest
|
|
|
+from unittest.mock import patch
|
|
|
+
|
|
|
+from cyber_agent.core.agent_mode import RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY
|
|
|
+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 ensure_task_protocol, new_task_protocol
|
|
|
+from cyber_agent.tools import get_tool_registry
|
|
|
+from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
|
|
|
+from cyber_agent.tools.builtin.subagent import agent
|
|
|
+from cyber_agent.tools.registry import ToolRegistry
|
|
|
+from cyber_agent.trace.models import Trace
|
|
|
+from cyber_agent.trace.store import FileSystemTraceStore
|
|
|
+
|
|
|
+
|
|
|
+ROOT_CRITERIA = ["return a checked final answer"]
|
|
|
+CHILD_BRIEF = {
|
|
|
+ "objective": "check one concrete fact",
|
|
|
+ "reason": "the root needs independently checked evidence",
|
|
|
+ "completion_criteria": ["return one supported conclusion"],
|
|
|
+ "expected_outputs": ["one checked conclusion"],
|
|
|
+ "constraints": ["do not guess"],
|
|
|
+}
|
|
|
+
|
|
|
+
|
|
|
+def knowledge_disabled():
|
|
|
+ return KnowledgeConfig(
|
|
|
+ enable_extraction=False,
|
|
|
+ enable_completion_extraction=False,
|
|
|
+ enable_injection=False,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def recursive_env(**overrides):
|
|
|
+ values = {
|
|
|
+ "AGENT_MODE": "recursive",
|
|
|
+ "AGENT_RESOURCE_BUDGET_ENABLED": "true",
|
|
|
+ "AGENT_MAX_TOTAL_AGENTS": "50",
|
|
|
+ "AGENT_MAX_LLM_CALLS": "150",
|
|
|
+ "AGENT_MAX_TOTAL_TOKENS": "1500000",
|
|
|
+ "AGENT_MAX_TOTAL_COST_USD": "15",
|
|
|
+ "AGENT_MAX_DURATION_SECONDS": "3600",
|
|
|
+ "AGENT_RESERVED_FINAL_CALLS": "1",
|
|
|
+ }
|
|
|
+ values.update(overrides)
|
|
|
+ return patch.dict(os.environ, values, clear=False)
|
|
|
+
|
|
|
+
|
|
|
+def tool_names(schemas):
|
|
|
+ return {
|
|
|
+ schema["function"]["name"]
|
|
|
+ for schema in schemas or []
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
+def is_validator_call(messages):
|
|
|
+ return any(
|
|
|
+ message.get("role") == "system"
|
|
|
+ and "independent validator" in str(message.get("content", ""))
|
|
|
+ for message in messages
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def validator_response(messages, outcome="passed"):
|
|
|
+ packet = json.loads(messages[-1]["content"])
|
|
|
+ scope = packet["validation_scope"]
|
|
|
+ passed = outcome == "passed"
|
|
|
+ return {
|
|
|
+ "content": json.dumps({
|
|
|
+ "outcome": outcome,
|
|
|
+ "scope": scope,
|
|
|
+ "reason": "the persisted record was checked",
|
|
|
+ "issues": [] if passed else ["the answer lacks required support"],
|
|
|
+ "retry_from": None if passed else "evidence",
|
|
|
+ }),
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ "prompt_tokens": 3,
|
|
|
+ "completion_tokens": 2,
|
|
|
+ "cost": 0.001,
|
|
|
+ }
|
|
|
+
|
|
|
+
|
|
|
+class RootValidationIntegrationTest(unittest.IsolatedAsyncioTestCase):
|
|
|
+ async def asyncSetUp(self):
|
|
|
+ self.temp_dir = tempfile.TemporaryDirectory()
|
|
|
+ self.store = FileSystemTraceStore(self.temp_dir.name)
|
|
|
+
|
|
|
+ async def asyncTearDown(self):
|
|
|
+ self.temp_dir.cleanup()
|
|
|
+
|
|
|
+ def config(self, **overrides):
|
|
|
+ values = {
|
|
|
+ "tools": [],
|
|
|
+ "tool_groups": [],
|
|
|
+ "enable_research_flow": False,
|
|
|
+ "root_completion_criteria": ROOT_CRITERIA,
|
|
|
+ "knowledge": knowledge_disabled(),
|
|
|
+ }
|
|
|
+ values.update(overrides)
|
|
|
+ return RunConfig(**values)
|
|
|
+
|
|
|
+ async def test_root_completes_only_after_independent_validator_passes(self):
|
|
|
+ calls = []
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ calls.append(kwargs)
|
|
|
+ if is_validator_call(kwargs["messages"]):
|
|
|
+ return validator_response(kwargs["messages"])
|
|
|
+ return {
|
|
|
+ "content": "checked final answer",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ "prompt_tokens": 5,
|
|
|
+ "completion_tokens": 4,
|
|
|
+ "cost": 0.002,
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
|
|
|
+ with recursive_env():
|
|
|
+ result = await runner.run_result(
|
|
|
+ messages=[{"role": "user", "content": "answer carefully"}],
|
|
|
+ config=self.config(),
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("completed", result["status"])
|
|
|
+ root = await self.store.get_trace(result["trace_id"])
|
|
|
+ state = ensure_task_protocol(root.context)
|
|
|
+ self.assertTrue(state["root_validation_passed"])
|
|
|
+ self.assertEqual(1, state["root_validation_attempts"])
|
|
|
+ validators = await self.store.list_traces(
|
|
|
+ parent_trace_id=root.trace_id,
|
|
|
+ created_by_tool="validator",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(validators))
|
|
|
+ self.assertEqual(root.trace_id, validators[0].context["root_trace_id"])
|
|
|
+ self.assertEqual([], validators[0].tools)
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
|
|
|
+ self.assertEqual(1, usage.total_agents)
|
|
|
+ self.assertEqual(2, usage.llm_calls)
|
|
|
+ self.assertEqual(14, usage.total_tokens)
|
|
|
+ self.assertEqual(2, len(calls))
|
|
|
+
|
|
|
+ async def test_root_gets_one_revision_then_fails_on_second_rejection(self):
|
|
|
+ ordinary_calls = 0
|
|
|
+ validation_calls = 0
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ nonlocal ordinary_calls, validation_calls
|
|
|
+ if is_validator_call(kwargs["messages"]):
|
|
|
+ validation_calls += 1
|
|
|
+ return validator_response(kwargs["messages"], outcome="failed")
|
|
|
+ ordinary_calls += 1
|
|
|
+ return {
|
|
|
+ "content": f"candidate-{ordinary_calls}",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
|
|
|
+ with recursive_env():
|
|
|
+ result = await runner.run_result(
|
|
|
+ messages=[{"role": "user", "content": "answer carefully"}],
|
|
|
+ config=self.config(),
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("failed", result["status"])
|
|
|
+ self.assertEqual(2, ordinary_calls)
|
|
|
+ self.assertEqual(2, validation_calls)
|
|
|
+ root = await self.store.get_trace(result["trace_id"])
|
|
|
+ state = ensure_task_protocol(root.context)
|
|
|
+ self.assertEqual(2, state["root_validation_attempts"])
|
|
|
+ self.assertEqual(
|
|
|
+ ["failed", "failed"],
|
|
|
+ [item["outcome"] for item in state["root_validation_history"]],
|
|
|
+ )
|
|
|
+ self.assertFalse(state["root_validation_passed"])
|
|
|
+
|
|
|
+ async def test_root_can_complete_after_one_validator_requested_revision(self):
|
|
|
+ ordinary_calls = 0
|
|
|
+ validation_outcomes = iter(["failed", "passed"])
|
|
|
+ second_candidate_messages = None
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ nonlocal ordinary_calls, second_candidate_messages
|
|
|
+ if is_validator_call(kwargs["messages"]):
|
|
|
+ return validator_response(
|
|
|
+ kwargs["messages"],
|
|
|
+ outcome=next(validation_outcomes),
|
|
|
+ )
|
|
|
+ ordinary_calls += 1
|
|
|
+ if ordinary_calls == 2:
|
|
|
+ second_candidate_messages = kwargs["messages"]
|
|
|
+ return {
|
|
|
+ "content": f"candidate-{ordinary_calls}",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
|
|
|
+ with recursive_env():
|
|
|
+ result = await runner.run_result(
|
|
|
+ messages=[{"role": "user", "content": "revise if needed"}],
|
|
|
+ config=self.config(),
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("completed", result["status"])
|
|
|
+ root = await self.store.get_trace(result["trace_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(any(
|
|
|
+ message.get("role") == "user"
|
|
|
+ and "framework-owned ValidationResult" in str(message.get("content"))
|
|
|
+ for message in second_candidate_messages
|
|
|
+ ))
|
|
|
+ persisted_messages = await self.store.get_trace_messages(root.trace_id)
|
|
|
+ self.assertTrue(any(
|
|
|
+ message.role == "user"
|
|
|
+ and "framework-owned ValidationResult" in str(message.content)
|
|
|
+ for message in persisted_messages
|
|
|
+ ))
|
|
|
+
|
|
|
+ calls_after_completion = 0
|
|
|
+
|
|
|
+ async def must_not_call(**_kwargs):
|
|
|
+ nonlocal calls_after_completion
|
|
|
+ calls_after_completion += 1
|
|
|
+ raise AssertionError("a third validation attempt must not start")
|
|
|
+
|
|
|
+ resumed = AgentRunner(trace_store=self.store, llm_call=must_not_call)
|
|
|
+ with recursive_env():
|
|
|
+ with self.assertRaisesRegex(ValueError, "two independent"):
|
|
|
+ await resumed.run_result(
|
|
|
+ [],
|
|
|
+ self.config(trace_id=root.trace_id),
|
|
|
+ )
|
|
|
+ self.assertEqual(0, calls_after_completion)
|
|
|
+
|
|
|
+ async def test_root_validator_reads_current_tool_result_and_candidate(self):
|
|
|
+ registry = ToolRegistry()
|
|
|
+
|
|
|
+ async def evidence_tool():
|
|
|
+ return {"verified": "persisted-tool-result"}
|
|
|
+
|
|
|
+ registry.register(evidence_tool, groups=["test"])
|
|
|
+ ordinary_calls = 0
|
|
|
+ validator_packet = None
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ nonlocal ordinary_calls, validator_packet
|
|
|
+ if is_validator_call(kwargs["messages"]):
|
|
|
+ validator_packet = json.loads(kwargs["messages"][-1]["content"])
|
|
|
+ return validator_response(kwargs["messages"])
|
|
|
+ ordinary_calls += 1
|
|
|
+ if ordinary_calls == 1:
|
|
|
+ return {
|
|
|
+ "content": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "evidence-call",
|
|
|
+ "type": "function",
|
|
|
+ "function": {
|
|
|
+ "name": "evidence_tool",
|
|
|
+ "arguments": "{}",
|
|
|
+ },
|
|
|
+ }],
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
+ }
|
|
|
+ return {
|
|
|
+ "content": "candidate-from-real-tool-result",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=registry,
|
|
|
+ llm_call=llm_call,
|
|
|
+ )
|
|
|
+ with recursive_env():
|
|
|
+ result = await runner.run_result(
|
|
|
+ [{"role": "user", "content": "use the evidence tool"}],
|
|
|
+ self.config(tools=["evidence_tool"]),
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("completed", result["status"])
|
|
|
+ trajectory = validator_packet["trajectory"]
|
|
|
+ self.assertTrue(any(
|
|
|
+ item.get("role") == "tool"
|
|
|
+ and "persisted-tool-result" in str(item.get("content"))
|
|
|
+ for item in trajectory
|
|
|
+ ))
|
|
|
+ self.assertTrue(any(
|
|
|
+ item.get("role") == "assistant"
|
|
|
+ and "candidate-from-real-tool-result" in str(item.get("content"))
|
|
|
+ for item in trajectory
|
|
|
+ ))
|
|
|
+
|
|
|
+ async def test_new_recursive_root_requires_explicit_completion_criteria(self):
|
|
|
+ async def llm_call(**_kwargs):
|
|
|
+ raise AssertionError("LLM must not be called")
|
|
|
+
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
|
|
|
+ with recursive_env():
|
|
|
+ with self.assertRaisesRegex(ValueError, "root_completion_criteria"):
|
|
|
+ await runner.run_result(
|
|
|
+ messages=[{"role": "user", "content": "missing criteria"}],
|
|
|
+ config=self.config(root_completion_criteria=[]),
|
|
|
+ )
|
|
|
+ self.assertEqual([], await self.store.list_traces(limit=10))
|
|
|
+
|
|
|
+ async def test_old_recursive_revision_two_trace_is_not_silently_migrated(self):
|
|
|
+ trace_id = "old-recursive-experiment"
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=trace_id,
|
|
|
+ mode="agent",
|
|
|
+ context={
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "agent_depth": 0,
|
|
|
+ "root_trace_id": trace_id,
|
|
|
+ "task_protocol": {},
|
|
|
+ },
|
|
|
+ ))
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=lambda **_: None)
|
|
|
+ with self.assertRaisesRegex(ValueError, "create a new trace"):
|
|
|
+ await runner.run_result([], RunConfig(trace_id=trace_id))
|
|
|
+
|
|
|
+ async def test_validator_trace_cannot_be_continued_as_an_agent(self):
|
|
|
+ validator_id = "root@validator-fixture"
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=validator_id,
|
|
|
+ mode="agent",
|
|
|
+ agent_type="validator",
|
|
|
+ tools=[],
|
|
|
+ status="completed",
|
|
|
+ context={
|
|
|
+ "created_by_tool": "validator",
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "root_trace_id": "root",
|
|
|
+ },
|
|
|
+ ))
|
|
|
+ calls = 0
|
|
|
+
|
|
|
+ async def llm_call(**_kwargs):
|
|
|
+ nonlocal calls
|
|
|
+ calls += 1
|
|
|
+ return {"content": "must not run", "tool_calls": []}
|
|
|
+
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
|
|
|
+ with self.assertRaisesRegex(ValueError, "cannot be continued"):
|
|
|
+ await runner.run_result(
|
|
|
+ [{"role": "user", "content": "gain tools"}],
|
|
|
+ RunConfig(trace_id=validator_id),
|
|
|
+ )
|
|
|
+ self.assertEqual(0, calls)
|
|
|
+ validator = await self.store.get_trace(validator_id)
|
|
|
+ self.assertEqual("completed", validator.status)
|
|
|
+ self.assertEqual([], validator.tools)
|
|
|
+
|
|
|
+
|
|
|
+class RecursiveBudgetIntegrationTest(unittest.IsolatedAsyncioTestCase):
|
|
|
+ async def asyncSetUp(self):
|
|
|
+ self.temp_dir = tempfile.TemporaryDirectory()
|
|
|
+ self.store = FileSystemTraceStore(self.temp_dir.name)
|
|
|
+
|
|
|
+ async def asyncTearDown(self):
|
|
|
+ self.temp_dir.cleanup()
|
|
|
+
|
|
|
+ async def test_response_over_token_budget_never_executes_its_tool_call(self):
|
|
|
+ dangerous_calls = 0
|
|
|
+ registry = ToolRegistry()
|
|
|
+
|
|
|
+ async def dangerous():
|
|
|
+ nonlocal dangerous_calls
|
|
|
+ dangerous_calls += 1
|
|
|
+ return "must not run"
|
|
|
+
|
|
|
+ registry.register(dangerous, groups=["test"])
|
|
|
+
|
|
|
+ async def llm_call(**_kwargs):
|
|
|
+ return {
|
|
|
+ "content": "over-budget response",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "dangerous-call",
|
|
|
+ "type": "function",
|
|
|
+ "function": {"name": "dangerous", "arguments": "{}"},
|
|
|
+ }],
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
+ "prompt_tokens": 8,
|
|
|
+ "completion_tokens": 8,
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=registry,
|
|
|
+ llm_call=llm_call,
|
|
|
+ )
|
|
|
+ config = RunConfig(
|
|
|
+ tools=["dangerous"],
|
|
|
+ tool_groups=[],
|
|
|
+ enable_research_flow=False,
|
|
|
+ root_completion_criteria=ROOT_CRITERIA,
|
|
|
+ knowledge=knowledge_disabled(),
|
|
|
+ )
|
|
|
+ with recursive_env(AGENT_MAX_TOTAL_TOKENS="10"):
|
|
|
+ result = await runner.run_result(
|
|
|
+ [{"role": "user", "content": "do not overspend"}],
|
|
|
+ config,
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("failed", result["status"])
|
|
|
+ self.assertEqual(0, dangerous_calls)
|
|
|
+ root = await self.store.get_trace(result["trace_id"])
|
|
|
+ self.assertEqual("budget_exhausted:tokens", root.context["termination_reason"])
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
|
|
|
+ self.assertEqual(16, usage.total_tokens)
|
|
|
+ self.assertEqual("budget_exhausted:tokens", usage.exhausted_reason)
|
|
|
+ assistant = [
|
|
|
+ message
|
|
|
+ for message in await self.store.get_trace_messages(root.trace_id)
|
|
|
+ if message.role == "assistant"
|
|
|
+ ]
|
|
|
+ self.assertIsNone(assistant[-1].content["tool_calls"])
|
|
|
+
|
|
|
+ async def test_agent_budget_denial_blocks_child_but_preserves_root_validator(self):
|
|
|
+ ordinary_calls = 0
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ nonlocal ordinary_calls
|
|
|
+ if is_validator_call(kwargs["messages"]):
|
|
|
+ return validator_response(kwargs["messages"])
|
|
|
+ ordinary_calls += 1
|
|
|
+ if ordinary_calls == 1:
|
|
|
+ return {
|
|
|
+ "content": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "delegate",
|
|
|
+ "type": "function",
|
|
|
+ "function": {
|
|
|
+ "name": "agent",
|
|
|
+ "arguments": json.dumps({"task_brief": CHILD_BRIEF}),
|
|
|
+ },
|
|
|
+ }],
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
+ }
|
|
|
+ return {
|
|
|
+ "content": "root can still finish without a child",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=get_tool_registry(),
|
|
|
+ llm_call=llm_call,
|
|
|
+ )
|
|
|
+ config = RunConfig(
|
|
|
+ tools=["agent"],
|
|
|
+ tool_groups=[],
|
|
|
+ enable_research_flow=False,
|
|
|
+ root_completion_criteria=ROOT_CRITERIA,
|
|
|
+ knowledge=knowledge_disabled(),
|
|
|
+ )
|
|
|
+ with recursive_env(AGENT_MAX_TOTAL_AGENTS="1"):
|
|
|
+ result = await runner.run_result(
|
|
|
+ [{"role": "user", "content": "try one delegated check"}],
|
|
|
+ config,
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("completed", result["status"])
|
|
|
+ root = await self.store.get_trace(result["trace_id"])
|
|
|
+ children = await self.store.list_traces(
|
|
|
+ parent_trace_id=root.trace_id,
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual([], children)
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
|
|
|
+ self.assertEqual(1, usage.total_agents)
|
|
|
+ self.assertEqual("budget_exhausted:agents", usage.exhausted_reason)
|
|
|
+ self.assertTrue(ensure_task_protocol(root.context)["root_validation_passed"])
|
|
|
+
|
|
|
+ async def test_persisted_child_is_not_refunded_when_goal_write_fails(self):
|
|
|
+ root_id = "partial-child-create-root"
|
|
|
+ budget = ResourceBudget()
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=root_id,
|
|
|
+ mode="agent",
|
|
|
+ uid="user-1",
|
|
|
+ model="fake",
|
|
|
+ context={
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "agent_depth": 0,
|
|
|
+ "root_trace_id": root_id,
|
|
|
+ "root_completion_criteria": ROOT_CRITERIA,
|
|
|
+ RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
|
|
|
+ "task_protocol": new_task_protocol(),
|
|
|
+ },
|
|
|
+ ))
|
|
|
+ await ResourceBudgetController(self.store).initialize(root_id, budget)
|
|
|
+ original_update_goal_tree = self.store.update_goal_tree
|
|
|
+
|
|
|
+ async def fail_child_goal_write(trace_id, goal_tree):
|
|
|
+ if trace_id != root_id:
|
|
|
+ raise RuntimeError("simulated child GoalTree write failure")
|
|
|
+ return await original_update_goal_tree(trace_id, goal_tree)
|
|
|
+
|
|
|
+ self.store.update_goal_tree = fail_child_goal_write
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=get_tool_registry(),
|
|
|
+ llm_call=lambda **_: None,
|
|
|
+ )
|
|
|
+
|
|
|
+ with self.assertRaisesRegex(RuntimeError, "GoalTree write failure"):
|
|
|
+ await agent(
|
|
|
+ task_brief=CHILD_BRIEF,
|
|
|
+ context={
|
|
|
+ "store": self.store,
|
|
|
+ "trace_id": root_id,
|
|
|
+ "runner": runner,
|
|
|
+ RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent"],
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+ children = await self.store.list_traces(
|
|
|
+ parent_trace_id=root_id,
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(children))
|
|
|
+ self.assertEqual("failed", children[0].status)
|
|
|
+ self.assertIn("Child initialization failed", children[0].error_message)
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root_id)
|
|
|
+ self.assertEqual(2, usage.total_agents)
|
|
|
+
|
|
|
+ async def test_parent_goal_projection_failure_closes_precreated_child(self):
|
|
|
+ root_id = "parent-goal-projection-root"
|
|
|
+ budget = ResourceBudget()
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=root_id,
|
|
|
+ mode="agent",
|
|
|
+ uid="user-1",
|
|
|
+ model="fake",
|
|
|
+ context={
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "agent_depth": 0,
|
|
|
+ "root_trace_id": root_id,
|
|
|
+ "root_completion_criteria": ROOT_CRITERIA,
|
|
|
+ RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
|
|
|
+ "task_protocol": new_task_protocol(),
|
|
|
+ },
|
|
|
+ ))
|
|
|
+ await ResourceBudgetController(self.store).initialize(root_id, budget)
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=get_tool_registry(),
|
|
|
+ llm_call=lambda **_: None,
|
|
|
+ )
|
|
|
+
|
|
|
+ async def fail_parent_projection(*_args, **_kwargs):
|
|
|
+ raise RuntimeError("simulated parent Goal projection failure")
|
|
|
+
|
|
|
+ with patch(
|
|
|
+ "cyber_agent.tools.builtin.subagent._update_goal_start",
|
|
|
+ new=fail_parent_projection,
|
|
|
+ ):
|
|
|
+ with self.assertRaisesRegex(RuntimeError, "projection failure"):
|
|
|
+ await agent(
|
|
|
+ task_brief=CHILD_BRIEF,
|
|
|
+ context={
|
|
|
+ "store": self.store,
|
|
|
+ "trace_id": root_id,
|
|
|
+ "runner": runner,
|
|
|
+ RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent"],
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+ children = await self.store.list_traces(
|
|
|
+ parent_trace_id=root_id,
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(children))
|
|
|
+ self.assertEqual("failed", children[0].status)
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root_id)
|
|
|
+ self.assertEqual(2, usage.total_agents)
|
|
|
+
|
|
|
+ async def test_approved_action_write_failure_closes_precreated_child(self):
|
|
|
+ root_id = "approved-action-write-root"
|
|
|
+ budget = ResourceBudget()
|
|
|
+ context = {
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "agent_depth": 0,
|
|
|
+ "root_trace_id": root_id,
|
|
|
+ "root_completion_criteria": ROOT_CRITERIA,
|
|
|
+ RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
|
|
|
+ "task_protocol": new_task_protocol(),
|
|
|
+ }
|
|
|
+ state = ensure_task_protocol(context)
|
|
|
+ state["next_actions"] = [{
|
|
|
+ "decision": "ACCEPT_REFINE",
|
|
|
+ "child_trace_id": "reviewed-child",
|
|
|
+ "goal_id": None,
|
|
|
+ "task_brief": CHILD_BRIEF,
|
|
|
+ "created_at_sequence": 1,
|
|
|
+ }]
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=root_id,
|
|
|
+ mode="agent",
|
|
|
+ uid="user-1",
|
|
|
+ model="fake",
|
|
|
+ context=context,
|
|
|
+ ))
|
|
|
+ await ResourceBudgetController(self.store).initialize(root_id, budget)
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=get_tool_registry(),
|
|
|
+ llm_call=lambda **_: None,
|
|
|
+ )
|
|
|
+ original_update_trace = self.store.update_trace
|
|
|
+ rejected_parent_write = False
|
|
|
+
|
|
|
+ async def fail_approved_action_write(trace_id, **updates):
|
|
|
+ nonlocal rejected_parent_write
|
|
|
+ if trace_id == root_id and "context" in updates and not rejected_parent_write:
|
|
|
+ rejected_parent_write = True
|
|
|
+ raise RuntimeError("simulated approved action write failure")
|
|
|
+ return await original_update_trace(trace_id, **updates)
|
|
|
+
|
|
|
+ self.store.update_trace = fail_approved_action_write
|
|
|
+ with self.assertRaisesRegex(RuntimeError, "approved action write failure"):
|
|
|
+ await agent(
|
|
|
+ task_brief=CHILD_BRIEF,
|
|
|
+ context={
|
|
|
+ "store": self.store,
|
|
|
+ "trace_id": root_id,
|
|
|
+ "runner": runner,
|
|
|
+ RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: ["agent"],
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+ children = await self.store.list_traces(
|
|
|
+ parent_trace_id=root_id,
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(children))
|
|
|
+ self.assertEqual("failed", children[0].status)
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root_id)
|
|
|
+ self.assertEqual(2, usage.total_agents)
|
|
|
+
|
|
|
+
|
|
|
+class RecursiveLifecycleIntegrationTest(unittest.IsolatedAsyncioTestCase):
|
|
|
+ async def asyncSetUp(self):
|
|
|
+ self.temp_dir = tempfile.TemporaryDirectory()
|
|
|
+ self.store = FileSystemTraceStore(self.temp_dir.name)
|
|
|
+
|
|
|
+ async def asyncTearDown(self):
|
|
|
+ self.temp_dir.cleanup()
|
|
|
+
|
|
|
+ async def test_child_report_is_validated_before_parent_review(self):
|
|
|
+ child_submitted = False
|
|
|
+ root_delegated = False
|
|
|
+ root_reviewed = False
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ nonlocal child_submitted, root_delegated, root_reviewed
|
|
|
+ messages = kwargs["messages"]
|
|
|
+ if is_validator_call(messages):
|
|
|
+ return validator_response(messages)
|
|
|
+ names = tool_names(kwargs.get("tools"))
|
|
|
+ user_text = "\n".join(
|
|
|
+ str(message.get("content", ""))
|
|
|
+ for message in messages
|
|
|
+ if message.get("role") == "user"
|
|
|
+ )
|
|
|
+ if "submit_task_report" in names and not child_submitted:
|
|
|
+ child_submitted = True
|
|
|
+ return {
|
|
|
+ "content": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "submit-report",
|
|
|
+ "type": "function",
|
|
|
+ "function": {
|
|
|
+ "name": "submit_task_report",
|
|
|
+ "arguments": json.dumps({
|
|
|
+ "task_report": {
|
|
|
+ "summary": "one fact was checked",
|
|
|
+ "outcome": "satisfied",
|
|
|
+ "validation": {
|
|
|
+ "hard_passed": True,
|
|
|
+ "open_issues": [],
|
|
|
+ },
|
|
|
+ "next_step_suggestion": {
|
|
|
+ "direction": "NONE",
|
|
|
+ "reason": "the delegated check is complete",
|
|
|
+ },
|
|
|
+ "outputs": [{"kind": "answer", "value": "checked"}],
|
|
|
+ "evidence": [{"kind": "tool", "value": "persisted"}],
|
|
|
+ "remaining_issues": [],
|
|
|
+ }
|
|
|
+ }),
|
|
|
+ },
|
|
|
+ }],
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
+ }
|
|
|
+ if "# Task Brief" in user_text:
|
|
|
+ return {
|
|
|
+ "content": "structured child report submitted",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ }
|
|
|
+ if names == {"review_task_result"}:
|
|
|
+ roots = [
|
|
|
+ trace
|
|
|
+ for trace in await self.store.list_traces(limit=20)
|
|
|
+ if trace.parent_trace_id is None
|
|
|
+ ]
|
|
|
+ root = roots[0]
|
|
|
+ pending = ensure_task_protocol(root.context)["pending_reviews"]
|
|
|
+ child_id = next(iter(pending))
|
|
|
+ root_reviewed = True
|
|
|
+ return {
|
|
|
+ "content": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "review-report",
|
|
|
+ "type": "function",
|
|
|
+ "function": {
|
|
|
+ "name": "review_task_result",
|
|
|
+ "arguments": json.dumps({
|
|
|
+ "child_trace_id": child_id,
|
|
|
+ "decision": "ASCEND",
|
|
|
+ "reason": "the validated child satisfies this goal",
|
|
|
+ }),
|
|
|
+ },
|
|
|
+ }],
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
+ }
|
|
|
+ if "agent" in names and not root_delegated:
|
|
|
+ root_delegated = True
|
|
|
+ return {
|
|
|
+ "content": "",
|
|
|
+ "tool_calls": [{
|
|
|
+ "id": "delegate-child",
|
|
|
+ "type": "function",
|
|
|
+ "function": {
|
|
|
+ "name": "agent",
|
|
|
+ "arguments": json.dumps({"task_brief": CHILD_BRIEF}),
|
|
|
+ },
|
|
|
+ }],
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
+ }
|
|
|
+ return {
|
|
|
+ "content": "final answer based on the validated child",
|
|
|
+ "tool_calls": [],
|
|
|
+ "finish_reason": "stop",
|
|
|
+ }
|
|
|
+
|
|
|
+ runner = AgentRunner(
|
|
|
+ trace_store=self.store,
|
|
|
+ tool_registry=get_tool_registry(),
|
|
|
+ llm_call=llm_call,
|
|
|
+ )
|
|
|
+ config = RunConfig(
|
|
|
+ tools=["agent", "submit_task_report", "review_task_result"],
|
|
|
+ tool_groups=[],
|
|
|
+ enable_research_flow=False,
|
|
|
+ root_completion_criteria=ROOT_CRITERIA,
|
|
|
+ knowledge=knowledge_disabled(),
|
|
|
+ )
|
|
|
+ with recursive_env():
|
|
|
+ result = await runner.run_result(
|
|
|
+ [{"role": "user", "content": "perform one delegated check"}],
|
|
|
+ config,
|
|
|
+ )
|
|
|
+
|
|
|
+ self.assertEqual("completed", result["status"])
|
|
|
+ self.assertTrue(root_reviewed)
|
|
|
+ root = await self.store.get_trace(result["trace_id"])
|
|
|
+ business_children = await self.store.list_traces(
|
|
|
+ parent_trace_id=root.trace_id,
|
|
|
+ created_by_tool="agent",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(business_children))
|
|
|
+ child = business_children[0]
|
|
|
+ child_validators = await self.store.list_traces(
|
|
|
+ parent_trace_id=child.trace_id,
|
|
|
+ created_by_tool="validator",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ root_validators = await self.store.list_traces(
|
|
|
+ parent_trace_id=root.trace_id,
|
|
|
+ created_by_tool="validator",
|
|
|
+ limit=10,
|
|
|
+ )
|
|
|
+ self.assertEqual(1, len(child_validators))
|
|
|
+ self.assertEqual(1, len(root_validators))
|
|
|
+ self.assertEqual(
|
|
|
+ child.trace_id,
|
|
|
+ child_validators[0].context["evaluated_trace_id"],
|
|
|
+ )
|
|
|
+ child_state = ensure_task_protocol(child.context)
|
|
|
+ self.assertEqual(
|
|
|
+ child_validators[0].trace_id,
|
|
|
+ child_state["task_report_validation"]["validation_result"][
|
|
|
+ "validator_trace_id"
|
|
|
+ ],
|
|
|
+ )
|
|
|
+ root_state = ensure_task_protocol(root.context)
|
|
|
+ self.assertEqual({}, root_state["pending_reviews"])
|
|
|
+ self.assertEqual(child.trace_id, root_state["reviews"][0]["child_trace_id"])
|
|
|
+ usage = await ResourceBudgetController(self.store).get_usage(root.trace_id)
|
|
|
+ self.assertEqual(2, usage.total_agents)
|
|
|
+ self.assertEqual(7, usage.llm_calls)
|
|
|
+
|
|
|
+ async def test_parent_stop_reaches_validator_after_child_runner_finished(self):
|
|
|
+ root_id = "validator-cancel-root"
|
|
|
+ child_id = "validator-cancel-child"
|
|
|
+ budget = ResourceBudget()
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=root_id,
|
|
|
+ mode="agent",
|
|
|
+ model="fake-model",
|
|
|
+ context={
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "agent_depth": 0,
|
|
|
+ "root_trace_id": root_id,
|
|
|
+ "root_completion_criteria": ROOT_CRITERIA,
|
|
|
+ RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
|
|
|
+ "task_protocol": new_task_protocol(),
|
|
|
+ },
|
|
|
+ ))
|
|
|
+ await self.store.create_trace(Trace(
|
|
|
+ trace_id=child_id,
|
|
|
+ mode="agent",
|
|
|
+ model="fake-model",
|
|
|
+ parent_trace_id=root_id,
|
|
|
+ context={
|
|
|
+ "agent_mode": "recursive",
|
|
|
+ "agent_mode_revision": 2,
|
|
|
+ "agent_depth": 1,
|
|
|
+ "root_trace_id": root_id,
|
|
|
+ "created_by_tool": "agent",
|
|
|
+ "task_protocol": new_task_protocol(CHILD_BRIEF),
|
|
|
+ },
|
|
|
+ ))
|
|
|
+ await ResourceBudgetController(self.store).initialize(
|
|
|
+ root_id,
|
|
|
+ budget,
|
|
|
+ initial_agents=2,
|
|
|
+ )
|
|
|
+ started = asyncio.Event()
|
|
|
+ release = asyncio.Event()
|
|
|
+
|
|
|
+ async def llm_call(**kwargs):
|
|
|
+ started.set()
|
|
|
+ await release.wait()
|
|
|
+ return validator_response(kwargs["messages"])
|
|
|
+
|
|
|
+ runner = AgentRunner(trace_store=self.store, llm_call=llm_call)
|
|
|
+ root_event = asyncio.Event()
|
|
|
+ runner._cancel_events[root_id] = root_event
|
|
|
+ runner._recursive_active_traces[root_id] = root_event
|
|
|
+ validation_task = asyncio.create_task(runner.validate_recursive_trace(
|
|
|
+ child_id,
|
|
|
+ scope="task",
|
|
|
+ task_brief=CHILD_BRIEF,
|
|
|
+ task_report={"summary": "candidate"},
|
|
|
+ ))
|
|
|
+ await asyncio.wait_for(started.wait(), timeout=1)
|
|
|
+
|
|
|
+ self.assertIn(child_id, runner._active_children[root_id])
|
|
|
+ self.assertTrue(await runner.stop(root_id))
|
|
|
+ release.set()
|
|
|
+ validation_run = await asyncio.wait_for(validation_task, timeout=1)
|
|
|
+
|
|
|
+ self.assertEqual("error", validation_run.result.outcome)
|
|
|
+ validator = await self.store.get_trace(validation_run.trace_id)
|
|
|
+ self.assertEqual("failed", validator.status)
|
|
|
+ self.assertEqual(child_id, validator.parent_trace_id)
|
|
|
+ self.assertEqual(root_id, validator.context["root_trace_id"])
|
|
|
+ self.assertNotIn(child_id, runner._recursive_active_traces)
|
|
|
+ self.assertNotIn(validation_run.trace_id, runner._recursive_active_traces)
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ unittest.main()
|
SamLee