from __future__ import annotations import asyncio from collections.abc import Callable import logging import pytest from tests.mock.utils import mock_llm_chunk from tests.stubs.fake_backend import FakeBackend from vibe.cli.turn_summary import ( NARRATOR_MODEL, NoopTurnSummary, TurnSummaryResult, TurnSummaryTracker, create_narrator_backend, ) from vibe.core.config import ModelConfig, ProviderConfig, VibeConfig from vibe.core.llm.backend.mistral import MistralBackend from vibe.core.types import AssistantEvent, Backend, ToolStreamEvent, UserMessageEvent _TEST_MODEL = ModelConfig(name="test-model", provider="test", alias="test-model") def _noop_callback(result: TurnSummaryResult) -> None: pass class TestCreateNarratorBackend: def test_uses_mistral_provider(self, monkeypatch): monkeypatch.setenv("MISTRAL_API_KEY", "test-key") config = VibeConfig() result = create_narrator_backend(config) assert result is not None backend, model = result assert isinstance(backend, MistralBackend) assert model is NARRATOR_MODEL def test_uses_custom_provider_base_url(self, monkeypatch): monkeypatch.setenv("MISTRAL_API_KEY", "test-key") custom_provider = ProviderConfig( name="mistral", api_base="https://on-prem.example.com/v1", api_key_env_var="MISTRAL_API_KEY", backend=Backend.MISTRAL, ) config = VibeConfig(providers=[custom_provider]) result = create_narrator_backend(config) assert result is not None backend, model = result assert isinstance(backend, MistralBackend) assert backend._provider.api_base == custom_provider.api_base def test_returns_none_when_api_key_missing(self, monkeypatch): monkeypatch.setenv("MISTRAL_API_KEY", "test-key") config = VibeConfig() monkeypatch.delenv("MISTRAL_API_KEY") assert create_narrator_backend(config) is None def test_returns_none_when_provider_missing(self): config = VibeConfig(providers=[]) assert create_narrator_backend(config) is None class TestTrack: def _make_tracker(self, backend: FakeBackend | None = None) -> TurnSummaryTracker: return TurnSummaryTracker( backend=backend or FakeBackend(), model=_TEST_MODEL, on_summary=_noop_callback, ) def test_assistant_event(self): tracker = self._make_tracker() tracker.start_turn("test") tracker.track(AssistantEvent(content="chunk1")) tracker.track(AssistantEvent(content="chunk2")) assert tracker._data is not None assert tracker._data.assistant_fragments == ["chunk1", "chunk2"] def test_assistant_event_empty_content_ignored(self): tracker = self._make_tracker() tracker.start_turn("test") tracker.track(AssistantEvent(content="")) assert tracker._data is not None assert tracker._data.assistant_fragments == [] def test_start_turn_preserves_full_message(self): tracker = self._make_tracker() long_msg = "a" * 1500 tracker.start_turn(long_msg) assert tracker._data is not None assert len(tracker._data.user_message) == 1500 def test_start_turn_increments_generation(self): tracker = self._make_tracker() assert tracker.generation == 0 tracker.start_turn("turn 1") assert tracker.generation == 1 tracker.start_turn("turn 2") assert tracker.generation == 2 def test_cancel_turn_clears_data(self): tracker = self._make_tracker() tracker.start_turn("test") assert tracker._data is not None tracker.cancel_turn() assert tracker._data is None def test_set_error_stores_message(self): tracker = self._make_tracker() tracker.start_turn("test") tracker.set_error("rate limit exceeded") assert tracker._data is not None assert tracker._data.error == "rate limit exceeded" def test_set_error_without_start_is_noop(self): tracker = self._make_tracker() tracker.set_error("should be ignored") assert tracker._data is None def test_cancel_turn_without_start_is_noop(self): tracker = self._make_tracker() tracker.cancel_turn() assert tracker._data is None def test_unrelated_events_ignored(self): tracker = self._make_tracker() tracker.start_turn("test") tracker.track(UserMessageEvent(content="hi", message_id="m1")) tracker.track( ToolStreamEvent(tool_name="bash", message="output", tool_call_id="tc1") ) assert tracker._data is not None assert tracker._data.assistant_fragments == [] class TestTurnSummaryTracker: def _make_tracker( self, backend: FakeBackend, on_summary: Callable[[TurnSummaryResult], None] = _noop_callback, ) -> TurnSummaryTracker: return TurnSummaryTracker( backend=backend, model=_TEST_MODEL, on_summary=on_summary ) @pytest.mark.asyncio async def test_track_accumulates_events(self): backend = FakeBackend(mock_llm_chunk(content="summary")) tracker = self._make_tracker(backend) tracker.start_turn("hello") tracker.track(AssistantEvent(content="chunk1")) tracker.track(AssistantEvent(content="chunk2")) assert tracker._data is not None assert tracker._data.assistant_fragments == ["chunk1", "chunk2"] @pytest.mark.asyncio async def test_end_turn_fires_summary(self): backend = FakeBackend(mock_llm_chunk(content="the summary")) tracker = self._make_tracker(backend) tracker.start_turn("do something") tracker.track(AssistantEvent(content="response")) tracker.end_turn() await asyncio.sleep(0.1) assert len(backend.requests_messages) == 1 summary_msgs = backend.requests_messages[0] assert len(summary_msgs) == 2 assert summary_msgs[0].role.value == "system" assert summary_msgs[1].role.value == "user" assert summary_msgs[1].content is not None assert "do something" in summary_msgs[1].content @pytest.mark.asyncio async def test_end_turn_clears_state(self): backend = FakeBackend(mock_llm_chunk(content="summary")) tracker = self._make_tracker(backend) tracker.start_turn("hello") tracker.end_turn() assert tracker._data is None @pytest.mark.asyncio async def test_track_without_start_is_noop(self): backend = FakeBackend(mock_llm_chunk(content="summary")) tracker = self._make_tracker(backend) tracker.track(AssistantEvent(content="ignored")) assert tracker._data is None @pytest.mark.asyncio async def test_end_turn_without_start_is_noop(self): backend = FakeBackend(mock_llm_chunk(content="summary")) tracker = self._make_tracker(backend) tracker.end_turn() assert len(backend.requests_messages) == 0 @pytest.mark.asyncio async def test_end_turn_after_cancel_is_noop(self): backend = FakeBackend(mock_llm_chunk(content="summary")) tracker = self._make_tracker(backend) tracker.start_turn("hello") tracker.cancel_turn() tracker.end_turn() await asyncio.sleep(0.1) assert len(backend.requests_messages) == 0 @pytest.mark.asyncio async def test_on_summary_callback_called(self): backend = FakeBackend(mock_llm_chunk(content="the summary text")) received: list[TurnSummaryResult] = [] def capture(result: TurnSummaryResult) -> None: received.append(result) tracker = self._make_tracker(backend, on_summary=capture) tracker.start_turn("hello") tracker.track(AssistantEvent(content="response")) tracker.end_turn() await asyncio.sleep(0.1) assert len(received) == 1 assert received[0].summary == "the summary text" assert received[0].generation == tracker.generation @pytest.mark.asyncio async def test_backend_error_calls_callback_with_none(self): backend = FakeBackend(exception_to_raise=RuntimeError("backend down")) received: list[TurnSummaryResult] = [] def capture(result: TurnSummaryResult) -> None: received.append(result) tracker = self._make_tracker(backend, on_summary=capture) tracker.start_turn("hello") tracker.end_turn() await asyncio.sleep(0.2) assert len(received) == 1 assert received[0].summary is None @pytest.mark.asyncio async def test_backend_error_logged_no_crash(self, caplog): backend = FakeBackend(exception_to_raise=RuntimeError("backend down")) tracker = self._make_tracker(backend) with caplog.at_level(logging.WARNING, logger="vibe"): tracker.start_turn("hello") tracker.end_turn() await asyncio.sleep(0.2) assert "Turn summary generation failed" in caplog.text @pytest.mark.asyncio async def test_close_cancels_pending_tasks(self): backend = FakeBackend(mock_llm_chunk(content="summary")) tracker = self._make_tracker(backend) tracker.start_turn("hello") tracker.end_turn() assert len(tracker._tasks) == 1 await tracker.close() assert len(tracker._tasks) == 0 @pytest.mark.asyncio async def test_error_only_turn_includes_error_in_summary(self): backend = FakeBackend(mock_llm_chunk(content="error summary")) received: list[TurnSummaryResult] = [] def capture(result: TurnSummaryResult) -> None: received.append(result) tracker = self._make_tracker(backend, on_summary=capture) tracker.start_turn("do something") tracker.set_error("Rate limit exceeded") cancel = tracker.end_turn() await asyncio.sleep(0.1) assert cancel is not None assert len(backend.requests_messages) == 1 prompt_content = backend.requests_messages[0][1].content assert prompt_content is not None assert "do something" in prompt_content assert "## Error" in prompt_content assert "Rate limit exceeded" in prompt_content assert "## Assistant Response" not in prompt_content assert len(received) == 1 assert received[0].summary == "error summary" @pytest.mark.asyncio async def test_error_with_partial_response_includes_both(self): backend = FakeBackend(mock_llm_chunk(content="partial error summary")) tracker = self._make_tracker(backend) tracker.start_turn("do something") tracker.track(AssistantEvent(content="partial response")) tracker.set_error("connection lost") tracker.end_turn() await asyncio.sleep(0.1) assert len(backend.requests_messages) == 1 prompt_content = backend.requests_messages[0][1].content assert prompt_content is not None assert "## Assistant Response" in prompt_content assert "partial response" in prompt_content assert "## Error" in prompt_content assert "connection lost" in prompt_content @pytest.mark.asyncio async def test_stale_summary_has_old_generation(self): backend = FakeBackend(mock_llm_chunk(content="stale summary")) received: list[TurnSummaryResult] = [] def capture(result: TurnSummaryResult) -> None: received.append(result) tracker = self._make_tracker(backend, on_summary=capture) tracker.start_turn("turn 1") tracker.end_turn() tracker.start_turn("turn 2") await asyncio.sleep(0.1) assert len(received) == 1 assert received[0].generation == 1 assert tracker.generation == 2 assert received[0].generation != tracker.generation class TestNoopTurnSummary: def test_all_methods_callable(self): noop = NoopTurnSummary() noop.start_turn("hello") noop.track(AssistantEvent(content="chunk")) noop.set_error("some error") noop.cancel_turn() noop.end_turn() def test_generation_is_zero(self): noop = NoopTurnSummary() assert noop.generation == 0 @pytest.mark.asyncio async def test_close_is_noop(self): noop = NoopTurnSummary() await noop.close()