agentos/tests/test_cli.py

from __future__ import annotations

import io
import json
import tempfile
import threading
import time
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from typing import Any

from agentos.background import BackgroundExecutionManager
from agentos.cli import _handle_event
from agentos.config import RuntimeConfig
from agentos.model import DeterministicModelClient
from agentos.runtime import AgentRuntime
from agentos.types import ActionPlan, ExecutionOption, InputEvent, SandboxResult, utc_now


class FakeSandbox:
    def run(self, **kwargs: Any) -> SandboxResult:
        workspace = Path(kwargs["workspace"])
        (workspace / "result.txt").write_text("done\n", encoding="utf-8")
        now = utc_now()
        return SandboxResult(
            instance_id=kwargs["instance_id"],
            linux_uid=kwargs["linux_uid"],
            container_name="fake-container",
            exit_code=0,
            stdout="",
            stderr="",
            started_at=now,
            finished_at=now,
        )


class ConcurrentFakeSandbox(FakeSandbox):
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self.active = 0
        self.max_active = 0

    def run(self, **kwargs: Any) -> SandboxResult:
        with self._lock:
            self.active += 1
            self.max_active = max(self.max_active, self.active)
        try:
            time.sleep(0.05)
            return super().run(**kwargs)
        finally:
            with self._lock:
                self.active -= 1


class FourOptionModel(DeterministicModelClient):
    def propose_options(self, user_input, context):
        del user_input, context
        return [
            ExecutionOption(
                option_id=str(index),
                title=f"方案 {index}",
                approach=f"执行方案 {index}",
                risk="high" if index == 1 else "low",
                expected_changes=[f"option-{index}.txt"],
            )
            for index in range(1, 5)
        ]

    def build_plan(self, user_input, option, context):
        del user_input, context
        filename = f"option-{option.option_id}.txt"
        return ActionPlan(
            summary=f"生成 {filename}",
            script=f"touch {filename}",
            success_criteria=[f"{filename} 存在"],
            memory_notes=[],
        )


class CliExecutionRoutingTests(unittest.TestCase):
    def _runtime(self, directory: str) -> AgentRuntime:
        config = RuntimeConfig(
            state_dir=Path(directory),
            model_command=("model",),
            uid_range_start=282_000,
            uid_range_size=16,
        )
        runtime = AgentRuntime(config, DeterministicModelClient())
        runtime.sandbox = FakeSandbox()  # type: ignore[assignment]
        runtime.initialize()
        return runtime

    def test_major_operation_executes_all_options(self) -> None:
        with tempfile.TemporaryDirectory() as directory, redirect_stdout(io.StringIO()):
            runtime = self._runtime(directory)
            result = _handle_event(
                runtime,
                InputEvent(event_id="major-input", text="执行重大架构迁移"),
                choice=1,
                decision="discard",
            )

            states = {row["state"] for row in runtime.memory.list_worldlines()}
            self.assertEqual(result, 0)
            self.assertEqual(states, {"discarded", "kept"})
            self.assertEqual(len(runtime.memory.list_worldlines()), 2)
            self.assertEqual(runtime.memory.list_memories(), [])

    def test_routine_operation_executes_only_selected_option(self) -> None:
        with tempfile.TemporaryDirectory() as directory, redirect_stdout(io.StringIO()):
            runtime = self._runtime(directory)
            result = _handle_event(
                runtime,
                InputEvent(event_id="routine-input", text="记录一段文本"),
                choice=1,
                decision="discard",
            )

            rows = runtime.memory.list_worldlines()
            self.assertEqual(result, 0)
            self.assertEqual(len(rows), 1)
            self.assertEqual(rows[0]["state"], "discarded")

    def test_major_operation_forks_every_option_and_runs_alternatives_in_parallel(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            config = RuntimeConfig(
                state_dir=Path(directory),
                model_command=("model",),
                uid_range_start=282_100,
                uid_range_size=16,
                background_workers=4,
            )
            runtime = AgentRuntime(config, FourOptionModel())
            sandbox = ConcurrentFakeSandbox()
            runtime.sandbox = sandbox  # type: ignore[assignment]
            runtime.initialize()
            notifications: list[str] = []
            background = BackgroundExecutionManager(
                runtime,
                max_workers=4,
                notifier=notifications.append,
            )
            output = io.StringIO()
            try:
                with redirect_stdout(output):
                    result = _handle_event(
                        runtime,
                        InputEvent(event_id="four-options", text="执行重大架构迁移"),
                        choice=3,
                        decision="discard",
                        background=background,
                    )
                background.wait()
            finally:
                background.shutdown()

            rows = runtime.memory.list_worldlines()
            states_by_option = {
                json.loads(row["option_json"])["option_id"]: row["state"] for row in rows
            }
            self.assertEqual(result, 0)
            self.assertEqual(len(rows), 4)
            self.assertEqual(states_by_option["3"], "discarded")
            self.assertEqual(
                {states_by_option[key] for key in ("1", "2", "4")},
                {"kept"},
            )
            self.assertGreaterEqual(sandbox.max_active, 2)
            self.assertEqual(len(notifications), 3)
            self.assertEqual(output.getvalue().count("\n世界线:"), 1)


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