tests/unit/test_cli_rendering.py

"""Regression tests for interactive CLI rendering and research model routing."""

import asyncio
import sys
from io import StringIO
from types import SimpleNamespace

import pytest
from rich.console import Console

import agent.main as main_mod
from agent.tools.research_tool import _get_research_model
from agent.utils import terminal_display


def test_direct_anthropic_research_model_stays_off_bedrock():
    assert (
        _get_research_model("anthropic/claude-opus-4-6")
        == "anthropic/claude-sonnet-4-6"
    )


def test_bedrock_anthropic_research_model_stays_on_bedrock():
    assert (
        _get_research_model("bedrock/us.anthropic.claude-opus-4-6-v1")
        == "bedrock/us.anthropic.claude-sonnet-4-6"
    )


def test_non_anthropic_research_model_is_unchanged():
    assert _get_research_model("openai/gpt-5.4") == "openai/gpt-5.4"


def test_help_output_keeps_descriptions_aligned(monkeypatch):
    output = StringIO()
    console = Console(
        file=output,
        color_system=None,
        theme=terminal_display._THEME,
        width=120,
    )
    monkeypatch.setattr(terminal_display, "_console", console)

    terminal_display.print_help()

    lines = [line.rstrip() for line in output.getvalue().splitlines() if line.strip()]
    description_columns = []
    for command, args, description in terminal_display.HELP_ROWS:
        line = next(line for line in lines if command in line)
        if args:
            assert args in line
        description_columns.append(line.index(description))

    assert len(set(description_columns)) == 1


def test_help_output_recomputes_widths_from_rows():
    rows = terminal_display.HELP_ROWS + (
        ("/longer-command", "[longer-args]", "Synthetic help row"),
    )
    output = StringIO()
    Console(
        file=output,
        color_system=None,
        theme=terminal_display._THEME,
        width=140,
    ).print(terminal_display.format_help_text(rows))

    lines = [line.rstrip() for line in output.getvalue().splitlines() if line.strip()]
    description_columns = [
        next(line for line in lines if command in line).index(description)
        for command, _args, description in rows
    ]

    assert len(set(description_columns)) == 1


def test_subagent_display_does_not_spawn_background_redraw(monkeypatch):
    calls: list[object] = []

    def _unexpected_future(*args, **kwargs):
        calls.append((args, kwargs))
        raise AssertionError("background redraw task should not be created")

    monkeypatch.setattr("asyncio.ensure_future", _unexpected_future)
    monkeypatch.setattr(
        terminal_display,
        "_console",
        SimpleNamespace(file=StringIO(), width=100),
    )

    mgr = terminal_display.SubAgentDisplayManager()
    mgr.start("agent-1", "research")
    mgr.add_call("agent-1", '▸ hf_papers  {"operation": "search"}')
    mgr.clear("agent-1")

    assert calls == []


def test_cli_forwards_model_flag_to_interactive_main(monkeypatch):
    seen: dict[str, object] = {}

    async def fake_main(*, model=None, sandbox_tools=False):
        seen["model"] = model
        seen["sandbox_tools"] = sandbox_tools

    monkeypatch.setattr(sys, "argv", ["ml-intern", "--model", "openai/gpt-5.5"])
    monkeypatch.setattr(main_mod, "main", fake_main)

    main_mod.cli()

    assert seen["model"] == "openai/gpt-5.5"
    assert seen["sandbox_tools"] is False


def test_cli_forwards_sandbox_flag_to_interactive_main(monkeypatch):
    seen: dict[str, object] = {}

    async def fake_main(*, model=None, sandbox_tools=False):
        seen["model"] = model
        seen["sandbox_tools"] = sandbox_tools

    monkeypatch.setattr(sys, "argv", ["ml-intern", "--sandbox-tools"])
    monkeypatch.setattr(main_mod, "main", fake_main)

    main_mod.cli()

    assert seen == {"model": None, "sandbox_tools": True}


def test_cli_forwards_sandbox_flag_to_headless_main(monkeypatch):
    seen: dict[str, object] = {}

    async def fake_headless_main(
        prompt,
        *,
        model=None,
        max_iterations=None,
        stream=True,
        sandbox_tools=False,
    ):
        seen.update(
            {
                "prompt": prompt,
                "model": model,
                "max_iterations": max_iterations,
                "stream": stream,
                "sandbox_tools": sandbox_tools,
            }
        )

    monkeypatch.setattr(
        sys,
        "argv",
        ["ml-intern", "--sandbox-tools", "--no-stream", "train a model"],
    )
    monkeypatch.setattr(main_mod, "headless_main", fake_headless_main)

    main_mod.cli()

    assert seen == {
        "prompt": "train a model",
        "model": None,
        "max_iterations": None,
        "stream": False,
        "sandbox_tools": True,
    }


@pytest.mark.asyncio
async def test_interactive_main_applies_model_override_before_banner(monkeypatch):
    class StopAfterBanner(Exception):
        pass

    def fake_banner(*, model=None, hf_user=None, tool_runtime=None):
        assert model == "openai/gpt-5.5"
        assert hf_user == "tester"
        assert tool_runtime == "local filesystem"
        raise StopAfterBanner

    monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0)
    monkeypatch.setattr(main_mod, "PromptSession", lambda: object())
    monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: "hf-token")
    monkeypatch.setattr(main_mod, "_get_hf_user", lambda _token: "tester")
    monkeypatch.setattr(
        main_mod,
        "load_config",
        lambda _path, **_kwargs: SimpleNamespace(
            model_name="moonshotai/Kimi-K2.6",
            mcpServers={},
            tool_runtime="local",
        ),
    )
    monkeypatch.setattr(main_mod, "print_banner", fake_banner)

    with pytest.raises(StopAfterBanner):
        await main_mod.main(model="openai/gpt-5.5")


@pytest.mark.asyncio
async def test_local_model_local_runtime_skips_hf_token_prompt(monkeypatch):
    class StopAfterBanner(Exception):
        pass

    async def fail_prompt(_prompt_session):
        raise AssertionError("local model with local tools should not prompt")

    def fake_banner(*, model=None, hf_user=None, tool_runtime=None):
        assert model == "llamacpp/model"
        assert hf_user is None
        assert tool_runtime == "local filesystem"
        raise StopAfterBanner

    monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0)
    monkeypatch.setattr(main_mod, "PromptSession", lambda: object())
    monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: None)
    monkeypatch.setattr(main_mod, "_prompt_and_save_hf_token", fail_prompt)
    monkeypatch.setattr(main_mod, "_get_hf_user", lambda _token: None)
    monkeypatch.setattr(
        main_mod,
        "load_config",
        lambda _path, **_kwargs: SimpleNamespace(
            model_name="llamacpp/model",
            mcpServers={},
            tool_runtime="local",
        ),
    )
    monkeypatch.setattr(main_mod, "print_banner", fake_banner)

    with pytest.raises(StopAfterBanner):
        await main_mod.main()


@pytest.mark.asyncio
async def test_local_model_sandbox_runtime_prompts_for_hf_token(monkeypatch):
    class StopAfterBanner(Exception):
        pass

    prompted = False

    async def fake_prompt(_prompt_session):
        nonlocal prompted
        prompted = True
        return "hf-token"

    def fake_banner(*, model=None, hf_user=None, tool_runtime=None):
        assert model == "llamacpp/model"
        assert hf_user == "tester"
        assert tool_runtime == "HF sandbox"
        raise StopAfterBanner

    monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0)
    monkeypatch.setattr(main_mod, "PromptSession", lambda: object())
    monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: None)
    monkeypatch.setattr(main_mod, "_prompt_and_save_hf_token", fake_prompt)
    monkeypatch.setattr(main_mod, "_get_hf_user", lambda _token: "tester")
    monkeypatch.setattr(
        main_mod,
        "load_config",
        lambda _path, **_kwargs: SimpleNamespace(
            model_name="llamacpp/model",
            mcpServers={},
            tool_runtime="local",
        ),
    )
    monkeypatch.setattr(main_mod, "print_banner", fake_banner)

    with pytest.raises(StopAfterBanner):
        await main_mod.main(sandbox_tools=True)

    assert prompted is True


@pytest.mark.asyncio
async def test_interactive_main_passes_sandbox_runtime_to_tool_router(monkeypatch):
    class StopAfterToolRouter(Exception):
        pass

    seen: dict[str, object] = {}

    class FakeGateway:
        def __init__(self, _config):
            pass

        async def start(self):
            pass

    class FakeToolRouter:
        def __init__(self, mcp_servers, *, hf_token=None, local_mode=True):
            seen["mcp_servers"] = mcp_servers
            seen["hf_token"] = hf_token
            seen["local_mode"] = local_mode
            raise StopAfterToolRouter

    from agent.core import hf_router_catalog

    monkeypatch.setattr(main_mod.os, "system", lambda *_args, **_kwargs: 0)
    monkeypatch.setattr(main_mod, "PromptSession", lambda: object())
    monkeypatch.setattr(main_mod, "resolve_hf_token", lambda: "hf-token")
    monkeypatch.setattr(main_mod, "_get_hf_user", lambda _token: "tester")
    monkeypatch.setattr(main_mod, "print_banner", lambda **_kwargs: None)
    monkeypatch.setattr(hf_router_catalog, "prewarm", lambda: None)
    monkeypatch.setattr(
        main_mod,
        "load_config",
        lambda _path, **_kwargs: SimpleNamespace(
            model_name="llamacpp/model",
            mcpServers={"server": object()},
            messaging=SimpleNamespace(default_auto_destinations=lambda: []),
            tool_runtime="local",
        ),
    )
    monkeypatch.setattr(main_mod, "NotificationGateway", FakeGateway)
    monkeypatch.setattr(main_mod, "ToolRouter", FakeToolRouter)

    with pytest.raises(StopAfterToolRouter):
        await main_mod.main(sandbox_tools=True)

    assert seen["hf_token"] == "hf-token"
    assert seen["local_mode"] is False


@pytest.mark.asyncio
async def test_initial_sandbox_preload_waits_before_prompt():
    waited = False

    async def preload():
        nonlocal waited
        await asyncio.sleep(0)
        waited = True

    task = asyncio.create_task(preload())
    await main_mod._wait_for_initial_sandbox_preload(
        [SimpleNamespace(sandbox_preload_task=task)]
    )

    assert waited is True