vibe/tests/cli/test_ui_session_resume.py

from __future__ import annotations

from pathlib import Path
import time
from unittest.mock import patch

import pytest

from tests.cli.plan_offer.adapters.fake_whoami_gateway import FakeWhoAmIGateway
from tests.conftest import (
    build_test_agent_loop,
    build_test_vibe_app,
    build_test_vibe_config,
)
from tests.update_notifier.adapters.fake_update_cache_repository import (
    FakeUpdateCacheRepository,
)
from tests.update_notifier.adapters.fake_update_gateway import FakeUpdateGateway
from vibe.cli.plan_offer.ports.whoami_gateway import WhoAmIPlanType, WhoAmIResponse
from vibe.cli.textual_ui.widgets.messages import (
    AssistantMessage,
    UserMessage,
    WhatsNewMessage,
)
from vibe.cli.textual_ui.widgets.tools import ToolCallMessage, ToolResultMessage
from vibe.cli.update_notifier import UpdateCache
from vibe.core.config import VibeConfig
from vibe.core.types import FunctionCall, LLMMessage, Role, ToolCall


@pytest.mark.asyncio
async def test_ui_displays_messages_when_resuming_session(
    vibe_config: VibeConfig,
) -> None:
    """Test that messages are properly displayed when resuming a session."""
    agent_loop = build_test_agent_loop(config=vibe_config)

    # Simulate a previous session with messages
    user_msg = LLMMessage(role=Role.user, content="Hello, how are you?")
    assistant_msg = LLMMessage(
        role=Role.assistant,
        content="I'm doing well, thank you!",
        tool_calls=[
            ToolCall(
                id="tool_call_1",
                index=0,
                function=FunctionCall(
                    name="read", arguments='{"file_path": "test.txt"}'
                ),
            )
        ],
    )
    tool_result_msg = LLMMessage(
        role=Role.tool,
        content="File content here",
        name="read",
        tool_call_id="tool_call_1",
    )

    agent_loop.messages.extend([user_msg, assistant_msg, tool_result_msg])

    app = build_test_vibe_app(agent_loop=agent_loop)

    async with app.run_test() as pilot:
        # Wait for the app to initialize and rebuild history
        await pilot.pause(0.5)

        # Verify user message is displayed
        user_messages = app.query(UserMessage)
        assert len(user_messages) == 1
        assert user_messages[0]._content == "Hello, how are you?"

        # Verify assistant message is displayed
        assistant_messages = app.query(AssistantMessage)
        assert len(assistant_messages) == 1
        assert assistant_messages[0]._content == "I'm doing well, thank you!"

        # Verify tool call message is displayed
        tool_call_messages = app.query(ToolCallMessage)
        assert len(tool_call_messages) == 1
        assert tool_call_messages[0]._tool_name == "read"

        # Verify tool result message is displayed
        tool_result_messages = app.query(ToolResultMessage)
        assert len(tool_result_messages) == 1
        assert tool_result_messages[0].tool_name == "read"
        assert tool_result_messages[0]._content == "File content here"


@pytest.mark.asyncio
async def test_ui_does_not_display_messages_when_only_system_messages_exist(
    vibe_config: VibeConfig,
) -> None:
    """Test that no messages are displayed when only system messages exist."""
    agent_loop = build_test_agent_loop(config=vibe_config)

    # Only system messages
    system_msg = LLMMessage(role=Role.system, content="System prompt")
    agent_loop.messages.append(system_msg)

    app = build_test_vibe_app(agent_loop=agent_loop)

    async with app.run_test() as pilot:
        await pilot.pause(0.5)

        # Verify no user or assistant messages are displayed
        user_messages = app.query(UserMessage)
        assert len(user_messages) == 0

        assistant_messages = app.query(AssistantMessage)
        assert len(assistant_messages) == 0


@pytest.mark.asyncio
async def test_ui_displays_multiple_user_assistant_turns(
    vibe_config: VibeConfig,
) -> None:
    """Test that multiple conversation turns are properly displayed."""
    agent_loop = build_test_agent_loop(config=vibe_config)

    # Multiple conversation turns
    messages = [
        LLMMessage(role=Role.user, content="First question"),
        LLMMessage(role=Role.assistant, content="First answer"),
        LLMMessage(role=Role.user, content="Second question"),
        LLMMessage(role=Role.assistant, content="Second answer"),
    ]

    agent_loop.messages.extend(messages)

    app = build_test_vibe_app(agent_loop=agent_loop)

    async with app.run_test() as pilot:
        await pilot.pause(0.5)

        # Verify all messages are displayed
        user_messages = app.query(UserMessage)
        assert len(user_messages) == 2
        assert user_messages[0]._content == "First question"
        assert user_messages[1]._content == "Second question"

        assistant_messages = app.query(AssistantMessage)
        assert len(assistant_messages) == 2
        assert assistant_messages[0]._content == "First answer"
        assert assistant_messages[1]._content == "Second answer"


@pytest.mark.asyncio
async def test_ui_displays_messages_when_resuming_in_dangerous_directory(
    monkeypatch: pytest.MonkeyPatch, vibe_config: VibeConfig
) -> None:
    monkeypatch.setattr(
        "vibe.cli.textual_ui.app.is_dangerous_directory",
        lambda: (True, "You are in the home directory"),
    )

    agent_loop = build_test_agent_loop(config=vibe_config)
    agent_loop.messages.extend([
        LLMMessage(role=Role.user, content="Hello from a previous run"),
        LLMMessage(role=Role.assistant, content="Welcome back!"),
    ])

    app = build_test_vibe_app(agent_loop=agent_loop)

    async with app.run_test() as pilot:
        await pilot.pause(0.5)

        user_messages = app.query(UserMessage)
        assistant_messages = app.query(AssistantMessage)

        assert len(user_messages) == 1
        assert user_messages[0]._content == "Hello from a previous run"
        assert len(assistant_messages) == 1
        assert assistant_messages[0]._content == "Welcome back!"


@pytest.mark.asyncio
async def test_ui_rebuilds_history_when_whats_new_is_shown(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    # we have to define an api key to make sure we display the Plan Offer message
    monkeypatch.setenv("MISTRAL_API_KEY", "api-key")
    config = build_test_vibe_config(enable_update_checks=True)
    agent_loop = build_test_agent_loop(config=config)
    agent_loop.messages.extend([
        LLMMessage(role=Role.user, content="Hello from the previous session."),
        LLMMessage(role=Role.assistant, content="Welcome back!"),
    ])
    update_cache = UpdateCache(
        latest_version="1.0.0",
        stored_at_timestamp=int(time.time()),
        seen_whats_new_version=None,
    )
    update_cache_repository = FakeUpdateCacheRepository(update_cache=update_cache)
    plan_offer_gateway = FakeWhoAmIGateway(
        WhoAmIResponse(
            plan_type=WhoAmIPlanType.API,
            plan_name="FREE",
            prompt_switching_to_pro_plan=False,
        )
    )
    app = build_test_vibe_app(
        agent_loop=agent_loop,
        update_notifier=FakeUpdateGateway(update=None),
        update_cache_repository=update_cache_repository,
        plan_offer_gateway=plan_offer_gateway,
        current_version="1.0.0",
        config=config,
    )

    with patch("vibe.cli.update_notifier.whats_new.VIBE_ROOT", tmp_path):
        whats_new_file = tmp_path / "whats_new.md"
        whats_new_file.write_text("# What's New\n\n- Feature 1")

        async with app.run_test() as pilot:
            await pilot.pause(0.5)
            whats_new_message = app.query_one(WhatsNewMessage)
            message = app.query_one(UserMessage)
            assistant_message = app.query_one(AssistantMessage)
            messages_area = app.query_one("#messages")
            children = list(messages_area.children)

    assert message._content == "Hello from the previous session."
    assert whats_new_message is not None
    assert whats_new_message.parent is not messages_area
    assert children == [message, assistant_message]