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.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",
            },
            clear=False,
        ):
            trace, _, _ = await self._prepare_new("legacy")
        self.assertEqual("legacy", trace.context["agent_mode"])
        self.assertNotIn("resource_budget", trace.context)

    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(2, trace.context["agent_mode_revision"])
        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_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)


if __name__ == "__main__":
    unittest.main()