strix/tests/test_tui_backend_controller.py

from __future__ import annotations

import argparse
import asyncio
import os
import threading
from pathlib import Path

import pytest

from strix.config import apply_config_override, loader
from strix.config.settings import DEFAULT_MAX_TURNS
from strix.core.agents import COMPACTION_FAILED, COMPACTION_SUCCESS, COMPACTION_UNAVAILABLE
from strix.interface.tui.backend.controller import TuiController


COMPACT_COMMAND = "/compact"
COMPRESS_COMMAND = "/compress"
COMPACTION_AGENT_ID = "compact-agent"
COMPACTION_SUCCESS_MESSAGE = "Context compaction complete."
COMPACTION_UNAVAILABLE_MESSAGE = "Context compaction could not be completed for this agent."
COMPACTION_FAILED_MESSAGE = "Context compaction failed. Please try again."
THREAD_JOIN_TIMEOUT_SECONDS = 1.0


class _SendingCoordinator:
    def __init__(self, delivered: bool = True) -> None:
        self.delivered = delivered
        self.messages: list[tuple[str, dict[str, object]]] = []

    async def send(self, agent_id: str, message: dict[str, object]) -> bool:
        self.messages.append((agent_id, message))
        return self.delivered


def args() -> argparse.Namespace:
    return argparse.Namespace(
        needs_setup=True,
        targets_info=[],
        instruction=None,
        scan_mode="deep",
        max_budget_usd=None,
        max_turns=DEFAULT_MAX_TURNS,
        scope_mode="auto",
        diff_base=None,
        local_sources=[],
        diff_scope={"active": False},
        user_explicit_instruction=None,
        run_name=None,
    )


@pytest.fixture(autouse=True)
def isolated_config(tmp_path: Path) -> None:
    for key in (
        "STRIX_LLM",
        "OPENAI_API_KEY",
        "ANTHROPIC_API_KEY",
        "LLM_API_KEY",
        "LLM_API_BASE",
        "AZURE_API_KEY",
        "AZURE_API_BASE",
        "AZURE_API_VERSION",
    ):
        os.environ.pop(key, None)
    apply_config_override(tmp_path / "config.json")


@pytest.mark.asyncio
async def test_setup_state_is_serializable() -> None:
    controller = TuiController(args())
    await controller.handle("setup.add_target", {"target": "https://example.com"})
    await controller.handle("setup.set_instruction", {"instruction": "focus on auth"})
    snapshot = controller.snapshot()
    assert snapshot["targets"] == ["https://example.com"]
    assert snapshot["instruction"] == "focus on auth"
    assert snapshot["scan_state"] == "setup"
    assert snapshot["scan_mode"] == "deep"
    assert snapshot["max_budget_usd"] is None
    assert snapshot["max_turns"] == 500
    assert snapshot["scope_mode"] == "auto"
    assert snapshot["diff_base"] is None


@pytest.mark.asyncio
async def test_connections_snapshot_reflects_the_pushed_mcp_roster() -> None:
    controller = TuiController(args())
    # A run with no MCP connections carries an empty roster, so the sidebar
    # omits the panel entirely.
    assert controller.snapshot()["connections"] == []

    controller.set_mcp_connections(
        [
            {"name": "supabase", "tool_count": 3, "dead": False},
            {"name": "vercel", "tool_count": 1, "dead": True},
        ]
    )
    assert controller.snapshot()["connections"] == [
        {"name": "supabase", "tool_count": 3, "dead": False},
        {"name": "vercel", "tool_count": 1, "dead": True},
    ]


@pytest.mark.asyncio
async def test_setup_instruction_starts_from_cli_and_can_be_cleared() -> None:
    setup_args = args()
    setup_args.instruction = "  CLI instruction  "
    controller = TuiController(setup_args)

    assert controller.snapshot()["instruction"] == "CLI instruction"

    result = await controller.handle("setup.set_instruction", {"instruction": ""})

    assert result == {"instruction": ""}
    assert controller.snapshot()["instruction"] == ""


@pytest.mark.asyncio
async def test_setup_controls_reject_changes_after_start() -> None:
    controller = TuiController(args())
    controller.setup_mode = False
    controller.scan_started = True

    with pytest.raises(RuntimeError, match="can no longer be changed"):
        await controller.handle("setup.add_target", {"target": "https://example.com"})


@pytest.mark.asyncio
async def test_large_target_list_reports_truncated_snapshot_count() -> None:
    controller = TuiController(args())

    for index in range(20):
        await controller.handle("setup.add_target", {"target": f"https://target-{index}.example"})
    added = await controller.handle("setup.add_target", {"target": "https://last.example"})
    snapshot = controller.snapshot()

    assert added == {"target": "https://last.example", "total": 21}
    assert snapshot["target_count"] == 21
    # The snapshot only carries a bounded prefix of the list.
    assert len(snapshot["targets"]) == 16


def test_state_populates_model_warning_for_non_frontier_model() -> None:
    os.environ["STRIX_LLM"] = "openai/gpt-3.5-turbo"
    loader._cached = None

    warning = TuiController(args()).snapshot()["model_warning"]

    assert "openai/gpt-3.5-turbo" in warning
    assert "not a recommended frontier model" in warning


def test_setup_restores_prepared_cli_targets() -> None:
    setup_args = args()
    setup_args.targets_info = [
        {"type": "web", "details": {}, "original": "https://example.com"},
        {"type": "local_code", "details": {}, "original": "/workspace/source"},
    ]

    controller = TuiController(setup_args)

    assert controller.snapshot()["targets"] == ["https://example.com", "/workspace/source"]


@pytest.mark.asyncio
async def test_start_validates_model_before_callback() -> None:
    started = False

    async def start() -> None:
        nonlocal started
        started = True

    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.add_target", {"target": "https://example.com"})
    with pytest.raises(ValueError, match="No model configured"):
        await controller.handle("setup.start", {})
    assert started is False


@pytest.mark.asyncio
async def test_start_launches_with_a_configured_model() -> None:
    started = False

    async def start() -> None:
        nonlocal started
        started = True

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    loader._cached = None
    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.add_target", {"target": "https://example.com"})

    result = await controller.handle("setup.start", {})

    assert result == {"started": True}
    assert started is True


@pytest.mark.asyncio
async def test_start_without_target_requires_mount_consent() -> None:
    started = False

    async def start() -> None:
        nonlocal started
        started = True

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)

    # Mounting the working directory is never silent.
    with pytest.raises(ValueError, match="No target set"):
        await controller.handle("setup.start", {})
    assert started is False
    assert controller.targets == []
    assert controller.workspace_mount is None


@pytest.mark.asyncio
async def test_target_less_start_enters_live_view_and_waits_for_the_mount() -> None:
    """Nothing is prepared until the live-view confirmation is answered."""
    started = False

    async def start() -> None:
        nonlocal started
        started = True

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)

    result = await controller.handle("setup.start", {"mount_working_dir": True})

    assert result == {"started": True}
    # The live view is up so the prompt can be shown there, but the scan has not
    # been prepared and nothing is mounted yet.
    assert started is False
    assert controller.setup_mode is False
    assert controller.scan_state == "preparing"
    assert controller.pending_workspace_mount == str(Path.cwd())
    assert controller.workspace_mount is None
    assert controller.snapshot()["pending_mount"] == str(Path.cwd())


@pytest.mark.asyncio
async def test_confirming_the_mount_starts_the_scan_without_a_target() -> None:
    started = False

    async def start() -> None:
        nonlocal started
        started = True

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.start", {"mount_working_dir": True})

    result = await controller.handle("setup.confirm_mount", {"approved": True})

    assert result == {"approved": True}
    assert started is True
    # Mounted as a workspace: the scan genuinely has no target, so the
    # instruction is the only source of truth.
    assert controller.workspace_mount == str(Path.cwd())
    assert controller.targets == []
    assert controller.scan_state == "running"
    assert controller.snapshot()["pending_mount"] == ""


@pytest.mark.asyncio
async def test_declining_the_mount_runs_without_one() -> None:
    started = 0

    async def start() -> None:
        nonlocal started
        started += 1

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.start", {"mount_working_dir": True})

    result = await controller.handle("setup.confirm_mount", {"approved": False})

    assert result == {"approved": False}
    # Declining skips the directory; it does not abandon the scan.
    assert started == 1
    assert controller.workspace_mount is None
    assert controller.pending_workspace_mount is None
    assert controller.setup_mode is False
    assert controller.scan_started is True
    assert controller.scan_state == "running"


@pytest.mark.asyncio
async def test_approving_the_mount_runs_with_it() -> None:
    started = 0

    async def start() -> None:
        nonlocal started
        started += 1

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.start", {"mount_working_dir": True})

    result = await controller.handle("setup.confirm_mount", {"approved": True})

    assert result == {"approved": True}
    assert started == 1
    assert controller.workspace_mount == str(Path.cwd())
    assert controller.scan_state == "running"


@pytest.mark.asyncio
async def test_confirm_mount_requires_a_pending_request() -> None:
    controller = TuiController(args())

    with pytest.raises(RuntimeError, match="No mount confirmation is pending"):
        await controller.handle("setup.confirm_mount", {"approved": True})


def test_snapshot_exposes_working_directory() -> None:
    controller = TuiController(args())

    assert controller.snapshot()["working_dir"] == str(Path.cwd())
    assert controller.snapshot()["pending_mount"] == ""


@pytest.mark.asyncio
async def test_user_message_updates_live_agent_projection_immediately() -> None:
    coordinator = _SendingCoordinator()
    controller = TuiController(args(), coordinator=coordinator)
    controller.setup_mode = False
    controller.scan_started = True
    controller.scan_loop = asyncio.get_running_loop()
    controller.live_view.upsert_agent(
        "root",
        name="Strix",
        status="failed",
        error_message="provider rejected request",
    )

    result = await controller.handle(
        "agent.send_message",
        {"agent_id": "root", "message": "try again"},
    )

    assert result == {"sent": True}
    assert coordinator.messages == [
        ("root", {"from": "user", "content": "try again", "type": "instruction"})
    ]
    agent = controller.live_view.agents["root"]
    assert agent["status"] == "waiting"
    assert "error_message" not in agent


@pytest.mark.asyncio
async def test_start_verifies_the_model_before_a_targeted_launch() -> None:
    order: list[str] = []

    async def verify() -> None:
        order.append("verify")

    async def start() -> None:
        order.append("start")

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start, on_verify=verify)
    await controller.handle("setup.add_target", {"target": "https://example.com"})

    await controller.handle("setup.start", {})

    assert order == ["verify", "start"]


@pytest.mark.asyncio
async def test_start_verifies_the_model_before_a_bare_prompt_leaves_setup() -> None:
    """A bare prompt gets the same model check as a named target, while the
    setup log is still on screen to show the outcome."""
    verified = 0

    async def verify() -> None:
        nonlocal verified
        verified += 1

    async def start() -> None:
        return None

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start, on_verify=verify)

    await controller.handle("setup.start", {"mount_working_dir": True})

    assert verified == 1
    assert controller.setup_mode is False
    assert controller.pending_workspace_mount == str(Path.cwd())


@pytest.mark.asyncio
async def test_failed_model_check_keeps_the_start_screen() -> None:
    async def verify() -> None:
        raise RuntimeError("Model connection failed: timed out")

    async def start() -> None:
        pytest.fail("the scan must not start when the model check fails")

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start, on_verify=verify)

    with pytest.raises(RuntimeError, match="Model connection failed"):
        await controller.handle("setup.start", {"mount_working_dir": True})

    # Still on the start screen, so the error lands in the setup log and the
    # user can retry; no run was prepared behind a stuck live view.
    assert controller.setup_mode is True
    assert controller.scan_started is False
    assert controller.scan_state == "setup"
    assert controller.pending_workspace_mount is None


@pytest.mark.asyncio
async def test_confirmed_mount_launch_failure_is_reported_in_the_live_view() -> None:
    async def start() -> None:
        raise ValueError("Scan preparation failed")

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.start", {"mount_working_dir": True})

    with pytest.raises(ValueError, match="Scan preparation failed"):
        await controller.handle("setup.confirm_mount", {"approved": True})

    assert controller.scan_state == "failed"
    assert controller.error == "Scan preparation failed"


@pytest.mark.asyncio
async def test_start_rejects_concurrent_and_repeated_submissions() -> None:
    entered = asyncio.Event()
    release = asyncio.Event()

    async def start() -> None:
        entered.set()
        await release.wait()

    os.environ["STRIX_LLM"] = "anthropic/claude-sonnet-4"
    os.environ["ANTHROPIC_API_KEY"] = "test-key"
    loader._cached = None
    controller = TuiController(args(), on_start=start)
    await controller.handle("setup.add_target", {"target": "https://example.com"})

    first_start = asyncio.create_task(controller.handle("setup.start", {}))
    await entered.wait()
    with pytest.raises(RuntimeError, match="already starting or running"):
        await controller.handle("setup.start", {})
    release.set()
    await first_start
    with pytest.raises(RuntimeError, match="already starting or running"):
        await controller.handle("setup.start", {})


@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["completed", "failed", "crashed", "stopped"])
async def test_stop_rejects_terminal_agents(status: str) -> None:
    class Coordinator:
        def __init__(self) -> None:
            self.calls: list[str] = []

        async def cancel_descendants_graceful(self, agent_id: str) -> bool:
            self.calls.append(agent_id)
            return True

    coordinator = Coordinator()
    controller = TuiController(args(), coordinator=coordinator)
    controller.set_runtime(scan_loop=asyncio.get_running_loop())
    controller.live_view.upsert_agent("agent-1", name="Agent", status=status)

    with pytest.raises(RuntimeError, match=f"cannot be stopped while {status}"):
        await controller.handle("agent.stop", {"agent_id": "agent-1"})

    assert coordinator.calls == []


@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["running", "waiting", "budget_paused"])
async def test_stop_allows_active_agents(status: str) -> None:
    class Coordinator:
        def __init__(self) -> None:
            self.calls: list[str] = []

        async def cancel_descendants_graceful(self, agent_id: str) -> bool:
            self.calls.append(agent_id)
            return True

    coordinator = Coordinator()
    controller = TuiController(args(), coordinator=coordinator)
    controller.set_runtime(scan_loop=asyncio.get_running_loop())
    controller.live_view.upsert_agent("agent-1", name="Agent", status=status)

    result = await controller.handle("agent.stop", {"agent_id": "agent-1"})

    assert result == {"stopped": True}
    assert coordinator.calls == ["agent-1"]


@pytest.mark.asyncio
async def test_stop_handles_coordinator_rejection_after_stale_active_projection() -> None:
    class Coordinator:
        async def cancel_descendants_graceful(self, _agent_id: str) -> bool:
            return False

    controller = TuiController(args(), coordinator=Coordinator())
    controller.set_runtime(scan_loop=asyncio.get_running_loop())
    controller.live_view.upsert_agent("agent-1", name="Agent", status="running")

    with pytest.raises(RuntimeError, match="no longer active"):
        await controller.handle("agent.stop", {"agent_id": "agent-1"})


@pytest.mark.asyncio
@pytest.mark.parametrize("command", [COMPACT_COMMAND, COMPRESS_COMMAND])
async def test_compact_command_compacts_without_sending_message(command: str) -> None:
    class Coordinator:
        def __init__(self) -> None:
            self.compacted: list[str] = []
            self.sent: list[tuple[str, dict[str, str]]] = []

        async def compact_agent_session(self, agent_id: str) -> str:
            self.compacted.append(agent_id)
            return COMPACTION_SUCCESS

        async def send(self, agent_id: str, message: dict[str, str]) -> bool:
            self.sent.append((agent_id, message))
            return True

    coordinator = Coordinator()
    controller = TuiController(args(), coordinator=coordinator)
    controller.set_runtime(scan_loop=asyncio.get_running_loop())

    result = await controller.handle(
        "agent.send_message",
        {"agent_id": COMPACTION_AGENT_ID, "message": command},
    )

    assert result == {"compacted": True}
    assert coordinator.compacted == [COMPACTION_AGENT_ID]
    assert coordinator.sent == []
    events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID)
    assert [event["data"]["content"] for event in events] == [
        command,
        COMPACTION_SUCCESS_MESSAGE,
    ]


@pytest.mark.asyncio
async def test_compact_command_reports_unavailable_compaction() -> None:
    class Coordinator:
        async def compact_agent_session(self, _agent_id: str) -> str:
            return COMPACTION_UNAVAILABLE

    controller = TuiController(args(), coordinator=Coordinator())
    controller.set_runtime(scan_loop=asyncio.get_running_loop())

    result = await controller.handle(
        "agent.send_message",
        {"agent_id": COMPACTION_AGENT_ID, "message": COMPACT_COMMAND},
    )

    assert result == {"compacted": False}
    events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID)
    assert events[-1]["data"]["content"] == COMPACTION_UNAVAILABLE_MESSAGE


@pytest.mark.asyncio
async def test_compact_command_reports_failed_compaction() -> None:
    class Coordinator:
        async def compact_agent_session(self, _agent_id: str) -> str:
            return COMPACTION_FAILED

    controller = TuiController(args(), coordinator=Coordinator())
    controller.set_runtime(scan_loop=asyncio.get_running_loop())

    result = await controller.handle(
        "agent.send_message",
        {"agent_id": COMPACTION_AGENT_ID, "message": COMPACT_COMMAND},
    )

    assert result == {"compacted": False}
    events = controller.live_view.events_for_agent(COMPACTION_AGENT_ID)
    assert events[-1]["data"]["content"] == COMPACTION_FAILED_MESSAGE


@pytest.mark.asyncio
async def test_scan_loop_runner_rejects_missing_loop() -> None:
    controller = TuiController(args())
    coroutine = asyncio.sleep(0)
    try:
        with pytest.raises(RuntimeError, match="Scan loop is not ready"):
            await controller._run_on_scan_loop(coroutine)
    finally:
        coroutine.close()


@pytest.mark.asyncio
async def test_scan_loop_runner_dispatches_to_background_loop() -> None:
    controller = TuiController(args())
    scan_loop = asyncio.new_event_loop()
    thread = threading.Thread(target=scan_loop.run_forever)
    thread.start()
    try:
        controller.set_runtime(scan_loop=scan_loop)
        result = await controller._run_on_scan_loop(asyncio.sleep(0, result=COMPACTION_SUCCESS))
    finally:
        scan_loop.call_soon_threadsafe(scan_loop.stop)
        thread.join(timeout=THREAD_JOIN_TIMEOUT_SECONDS)
        scan_loop.close()

    assert result == COMPACTION_SUCCESS


@pytest.mark.asyncio
async def test_unknown_command_is_rejected() -> None:
    controller = TuiController(args())
    with pytest.raises(ValueError, match="Unknown command"):
        await controller.handle("nope", {})


def test_messages_are_sanitized_and_agents_are_collection_only() -> None:
    controller = TuiController(args())
    controller.add_message("replace\x1b]52;c;Y2xpcA==\x07 key\x85")
    for index in range(40):
        controller.live_view.upsert_agent(f"agent-{index}", name=f"Agent {index}")

    snapshot = controller.snapshot()

    assert "agents" not in snapshot
    assert [message["text"] for message in snapshot["messages"]] == ["replace key"]
    assert len(controller.collection("agents")) == 40


@pytest.mark.asyncio
async def test_existing_viewer_is_reopened_and_closed(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    opened: list[str] = []

    class ViewerServer:
        shutdown_called = False
        close_called = False

        def shutdown(self) -> None:
            self.shutdown_called = True

        def server_close(self) -> None:
            self.close_called = True

    controller = TuiController(args())
    controller.viewer_status = "running"
    controller.viewer_url = "http://127.0.0.1:1234/?token=test"
    server = ViewerServer()
    controller._viewer_httpd = server
    monkeypatch.setattr("strix.interface.tui.backend.controller.webbrowser.open", opened.append)

    result = await controller.handle("viewer.open", {})
    controller.close_viewer()

    assert result == {"status": "running", "url": controller.viewer_url}
    assert opened == [controller.viewer_url]
    assert server.shutdown_called is True
    assert server.close_called is True