Agent/tests/test_agent_mode.py

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