import os import tempfile import unittest from unittest.mock import AsyncMock, patch from cyber_agent.core.agent_mode import ( AgentMode, assert_removed_config_absent, policy_from_context, policy_from_environment, ) from cyber_agent.core.runner import AgentRunner, RunConfig from cyber_agent.tools.builtin.knowledge import KnowledgeConfig from cyber_agent.trace.goal_tool import goal_tool from cyber_agent.trace.goal_models import GoalTree from cyber_agent.trace.models import Trace from cyber_agent.trace.store import FileSystemTraceStore class AgentModePolicyTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp_dir.name) self.runner = AgentRunner(trace_store=self.store) self.previous = { name: os.environ.get(name) for name in ("AGENT_MODE", "AGENT_RECURSION_ENABLED") } os.environ.pop("AGENT_MODE", None) os.environ.pop("AGENT_RECURSION_ENABLED", None) async def asyncTearDown(self): for name, value in self.previous.items(): if value is None: os.environ.pop(name, None) else: os.environ[name] = value self.temp_dir.cleanup() async def _prepare_new(self, mode=None): if mode is None: os.environ.pop("AGENT_MODE", None) else: os.environ["AGENT_MODE"] = mode return await self.runner._prepare_trace( [{"role": "user", "content": "test"}], RunConfig( name="mode test", root_task_anchor={ "objective": "return a checked answer", "completion_criteria": ["return a checked answer"], }, ), ) async def test_unset_defaults_to_legacy(self): trace, _, _ = await self._prepare_new() self.assertEqual("legacy", trace.context["agent_mode"]) self.assertEqual(1, trace.context["agent_mode_revision"]) self.assertEqual(0, trace.context["agent_depth"]) self.assertEqual(trace.trace_id, trace.context["root_trace_id"]) async def test_legacy_and_recursive_are_parsed_explicitly(self): self.assertIs(AgentMode.LEGACY, policy_from_environment().mode) os.environ["AGENT_MODE"] = "recursive" self.assertIs(AgentMode.RECURSIVE, policy_from_environment().mode) async def test_legacy_ignores_invalid_recursive_budget_environment(self): with patch.dict( os.environ, { "AGENT_MODE": "legacy", "AGENT_MAX_TOTAL_AGENTS": "not-a-number", "AGENT_RESOURCE_BUDGET_ENABLED": "not-a-boolean", "AGENT_MAX_VALIDATION_TOOL_CALLS": "not-a-number", "AGENT_VALIDATOR_SEARCH_PROVIDER": "unsupported", }, clear=False, ): trace, _, _ = await self._prepare_new("legacy") self.assertEqual("legacy", trace.context["agent_mode"]) self.assertNotIn("resource_budget", trace.context) async def test_recursive_rejects_invalid_validator_provider_before_trace(self): with patch.dict( os.environ, { "AGENT_MODE": "recursive", "AGENT_VALIDATOR_SEARCH_PROVIDER": "unsupported", }, clear=False, ): with self.assertRaisesRegex(ValueError, "search_provider"): await self._prepare_new("recursive") self.assertEqual([], await self.store.list_traces(limit=10)) async def test_invalid_mode_does_not_create_trace(self): with self.assertRaisesRegex(ValueError, "Invalid AGENT_MODE"): await self._prepare_new("unsupported") traces = await self.store.list_traces(limit=10) self.assertEqual([], traces) async def test_invalid_mode_is_rejected_before_task_name_generation(self): os.environ["AGENT_MODE"] = "unsupported" with patch.object( self.runner, "_generate_task_name", new=AsyncMock(return_value="should not run"), ) as generate_name: with self.assertRaisesRegex(ValueError, "Invalid AGENT_MODE"): await self.runner._prepare_trace( [{"role": "user", "content": "test"}], RunConfig(), ) generate_name.assert_not_awaited() async def test_removed_recursion_variable_fails_explicitly(self): os.environ["AGENT_RECURSION_ENABLED"] = "true" with self.assertRaisesRegex(ValueError, "has been removed"): assert_removed_config_absent() with self.assertRaisesRegex(ValueError, "AGENT_MODE"): await self._prepare_new("recursive") async def test_existing_trace_keeps_persisted_mode_after_env_switch(self): trace, _, _ = await self._prepare_new("recursive") self.assertEqual(3, trace.context["agent_mode_revision"]) self.assertTrue(policy_from_context(trace.context).requires_task_progress) self.assertIn("task_protocol", trace.context) os.environ["AGENT_MODE"] = "legacy" continued, _, _ = await self.runner._prepare_trace( [], RunConfig(trace_id=trace.trace_id, name="continue"), ) self.assertEqual("recursive", continued.context["agent_mode"]) persisted = await self.store.get_trace(trace.trace_id) self.assertEqual("recursive", persisted.context["agent_mode"]) async def test_recursive_revision_two_is_read_only_but_still_parseable(self): trace = Trace( trace_id="recursive-v2-read-only", mode="agent", task="historical", context={"agent_mode": "recursive", "agent_mode_revision": 2}, ) await self.store.create_trace(trace) persisted = await self.store.get_trace(trace.trace_id) self.assertTrue(policy_from_context(persisted.context).is_read_only) with self.assertRaisesRegex(ValueError, "recursive_revision_2_read_only"): await self.runner._prepare_trace( [], RunConfig(trace_id=trace.trace_id, name="continue"), ) async def test_unknown_persisted_revisions_fail_closed(self): for mode, revision in (("recursive", 4), ("recursive", 99), ("legacy", 2)): with self.subTest(mode=mode, revision=revision): context = { "agent_mode": mode, "agent_mode_revision": revision, } with self.assertRaisesRegex( ValueError, "unsupported_agent_mode_revision", ): policy_from_context(context) for revision in (4, 99): trace = Trace( trace_id=f"unknown-recursive-{revision}", mode="agent", task="must remain untouched", status="stopped", context={ "agent_mode": "recursive", "agent_mode_revision": revision, }, ) await self.store.create_trace(trace) with self.assertRaisesRegex( ValueError, "unsupported_agent_mode_revision", ): await self.runner._prepare_trace( [], RunConfig(trace_id=trace.trace_id), ) persisted = await self.store.get_trace(trace.trace_id) self.assertEqual("stopped", persisted.status) self.assertNotIn("run_config_snapshot", persisted.context) async def test_only_recursive_revision_three_requires_task_progress(self): self.assertFalse(policy_from_context({ "agent_mode": "recursive", "agent_mode_revision": 1, }).requires_task_progress) self.assertFalse(policy_from_context({ "agent_mode": "recursive", "agent_mode_revision": 2, }).requires_task_progress) self.assertTrue(policy_from_context({ "agent_mode": "recursive", "agent_mode_revision": 3, }).requires_task_progress) async def test_historical_trace_is_persisted_as_legacy_on_continue(self): trace = Trace(trace_id="historical", mode="agent", task="old", context={}) await self.store.create_trace(trace) os.environ["AGENT_MODE"] = "recursive" continued, _, _ = await self.runner._prepare_trace( [], RunConfig(trace_id=trace.trace_id, name="continue"), ) self.assertEqual("legacy", continued.context["agent_mode"]) self.assertIs(AgentMode.LEGACY, policy_from_context(continued.context).mode) async def test_legacy_goal_tree_keeps_automatic_sibling_and_parent_progress(self): tree = GoalTree(mission="legacy goals") parent = tree.add_goals(["parent"])[0] first, second = tree.add_goals(["first", "second"], parent_id=parent.id) tree.focus(parent.id) tree.focus(first.id) tree.complete(first.id, "done") self.assertEqual(second.id, tree.current_id) self.assertEqual("in_progress", parent.status) tree.complete(second.id, "done") self.assertEqual("completed", parent.status) self.assertEqual("所有子目标已完成", parent.summary) async def test_recursive_revision_2_focus_skips_remote_knowledge(self): trace_id = "recursive-v2-knowledge" await self.store.create_trace(Trace( trace_id=trace_id, mode="agent", task="recursive task", context={ "agent_mode": "recursive", "agent_mode_revision": 2, }, )) tree = GoalTree(mission="recursive goals") focused = tree.add_goals(["research one fact"])[0] await self.store.update_goal_tree(trace_id, tree) remote = AsyncMock(return_value={ "status": "completed", "summary": "[knowledge-forged]", }) with patch( "cyber_agent.tools.builtin.subagent._run_remote_agent", new=remote, ): result = await goal_tool( tree=tree, store=self.store, trace_id=trace_id, focus="1", knowledge_config=KnowledgeConfig(enable_injection=True), context={"sequence": 1}, ) self.assertIn("切换焦点", result) self.assertNotIn("已注入", result) self.assertEqual([], focused.knowledge) remote.assert_not_awaited() async def test_legacy_focus_keeps_remote_knowledge_injection(self): trace_id = "legacy-knowledge" await self.store.create_trace(Trace( trace_id=trace_id, mode="agent", task="legacy task", context={ "agent_mode": "legacy", "agent_mode_revision": 1, }, )) tree = GoalTree(mission="legacy goals") focused = tree.add_goals(["research one fact"])[0] await self.store.update_goal_tree(trace_id, tree) remote = AsyncMock(return_value={ "status": "completed", "summary": "Use [knowledge-official] as evidence", }) with patch( "cyber_agent.tools.builtin.subagent._run_remote_agent", new=remote, ): result = await goal_tool( tree=tree, store=self.store, trace_id=trace_id, focus="1", knowledge_config=KnowledgeConfig(enable_injection=True), context={"sequence": 1}, ) self.assertIn("已注入 1 条", result) self.assertEqual([{"id": "knowledge-official"}], focused.knowledge) remote.assert_awaited_once() if __name__ == "__main__": unittest.main()