"""Tests for deferred initialization: _complete_init, _wait_for_init, integrate_mcp idempotency.""" from __future__ import annotations import asyncio import threading from unittest.mock import AsyncMock, MagicMock, patch import pytest from tests.conftest import build_test_agent_loop, build_test_vibe_config from tests.mock.utils import mock_llm_chunk from tests.stubs.fake_backend import FakeBackend from tests.stubs.fake_connector_registry import FakeConnectorRegistry from tests.stubs.fake_mcp_registry import FakeMCPRegistry from vibe.core import agent_loop as agent_loop_module from vibe.core.agent_loop import AgentLoop from vibe.core.config import MCPStdio from vibe.core.tools.manager import ToolManager from vibe.core.tools.mcp.tools import RemoteTool def _build_uninitiated_loop(**kwargs): """Build a test loop with defer_heavy_init=True but without auto-starting the init thread.""" with patch.object(AgentLoop, "_start_deferred_init"): return build_test_agent_loop(defer_heavy_init=True, **kwargs) # --------------------------------------------------------------------------- # _complete_init # --------------------------------------------------------------------------- def _run_init(loop: AgentLoop) -> None: """Run _complete_init in a thread (matching production behavior) and wait.""" thread = threading.Thread(target=loop._complete_init, daemon=True) loop._deferred_init_thread = thread thread.start() thread.join() class TestCompleteInit: def test_success_sets_init_complete(self) -> None: loop = _build_uninitiated_loop() assert not loop.is_initialized _run_init(loop) assert loop.is_initialized assert loop._init_error is None def test_failure_sets_init_complete_and_stores_error(self) -> None: loop = _build_uninitiated_loop() error = RuntimeError("mcp boom") with patch.object(loop.tool_manager, "integrate_all", side_effect=error): _run_init(loop) assert loop.is_initialized assert loop._init_error is error def test_mcp_failure_sets_init_error(self) -> None: mcp_server = MCPStdio(name="test-server", transport="stdio", command="echo") config = build_test_vibe_config(mcp_servers=[mcp_server]) loop = _build_uninitiated_loop(config=config) with patch.object( loop.tool_manager._mcp_registry, "get_tools_async", side_effect=RuntimeError("mcp discovery boom"), ): _run_init(loop) assert loop.is_initialized assert isinstance(loop._init_error, RuntimeError) assert str(loop._init_error) == "mcp discovery boom" # --------------------------------------------------------------------------- # wait_until_ready # --------------------------------------------------------------------------- class TestWaitForInit: @pytest.mark.asyncio async def test_returns_immediately_when_already_complete(self) -> None: loop = build_test_agent_loop(defer_heavy_init=True) await loop.wait_until_ready() # should not block assert loop.is_initialized @pytest.mark.asyncio async def test_waits_for_background_thread(self) -> None: loop = build_test_agent_loop(defer_heavy_init=True) await loop.wait_until_ready() assert loop.is_initialized @pytest.mark.asyncio async def test_raises_stored_error(self) -> None: loop = _build_uninitiated_loop() error = RuntimeError("init failed") with patch.object(loop.tool_manager, "integrate_all", side_effect=error): loop._complete_init() with pytest.raises(RuntimeError, match="init failed"): await loop.wait_until_ready() @pytest.mark.asyncio async def test_raises_error_for_every_caller(self) -> None: loop = _build_uninitiated_loop() error = RuntimeError("once only") with patch.object(loop.tool_manager, "integrate_all", side_effect=error): loop._complete_init() with pytest.raises(RuntimeError): await loop.wait_until_ready() with pytest.raises(RuntimeError): await loop.wait_until_ready() # --------------------------------------------------------------------------- # integrate_mcp idempotency # --------------------------------------------------------------------------- class TestIntegrateMcpIdempotency: def test_second_call_is_noop(self) -> None: mcp_server = MCPStdio(name="test-server", transport="stdio", command="echo") config = build_test_vibe_config(mcp_servers=[mcp_server]) registry = FakeMCPRegistry() manager = ToolManager(lambda: config, mcp_registry=registry, defer_mcp=True) manager.integrate_mcp() tools_after_first = dict(manager.registered_tools) # Spy on the registry to ensure get_tools is not called again. registry.get_tools = MagicMock(wraps=registry.get_tools) manager.integrate_mcp() registry.get_tools.assert_not_called() assert manager.registered_tools == tools_after_first def test_flag_not_set_when_no_servers(self) -> None: config = build_test_vibe_config(mcp_servers=[]) manager = ToolManager(lambda: config, defer_mcp=True) manager.integrate_mcp() # No servers means the method returns early without setting the flag, # so a future call with servers would still run discovery. assert not manager._mcp_integrated class TestRefreshRemoteTools: @pytest.mark.asyncio async def test_refresh_rediscovers_mcp_and_connector_tools(self) -> None: mcp_server = MCPStdio(name="srv", transport="stdio", command="echo") config = build_test_vibe_config(mcp_servers=[mcp_server]) registry = FakeMCPRegistry() registry.get_tools_async = AsyncMock(wraps=registry.get_tools_async) connector_registry = FakeConnectorRegistry({ "alpha": [RemoteTool(name="search", description="Search alpha")] }) manager = ToolManager( lambda: config, mcp_registry=registry, connector_registry=connector_registry, defer_mcp=True, ) await manager.refresh_remote_tools_async() assert "srv_fake_tool" in manager.registered_tools assert "connector_alpha_search" in manager.registered_tools connector_registry._fake_connectors = { "beta": [RemoteTool(name="list", description="List beta")] } await manager.refresh_remote_tools_async() assert registry.get_tools_async.await_count == 2 assert "srv_fake_tool" in manager.registered_tools assert "connector_alpha_search" not in manager.registered_tools assert "connector_beta_list" in manager.registered_tools class TestDeferredInitPublicMethods: @pytest.mark.asyncio async def test_act_waits_for_deferred_init(self) -> None: loop = build_test_agent_loop( defer_heavy_init=True, backend=FakeBackend(mock_llm_chunk(content="hello")) ) events = [event async for event in loop.act("Hello")] assert loop.is_initialized assert [event.content for event in events if hasattr(event, "content")][ -1 ] == "hello" @pytest.mark.asyncio async def test_reload_with_initial_messages_waits_for_deferred_init(self) -> None: loop = build_test_agent_loop(defer_heavy_init=True) await loop.reload_with_initial_messages() assert loop.is_initialized @pytest.mark.asyncio async def test_switch_agent_waits_for_deferred_init(self) -> None: loop = build_test_agent_loop(defer_heavy_init=True) await loop.switch_agent("plan") assert loop.is_initialized assert loop.agent_profile.name == "plan" @pytest.mark.asyncio async def test_clear_history_waits_for_deferred_init(self) -> None: loop = build_test_agent_loop( defer_heavy_init=True, backend=FakeBackend(mock_llm_chunk(content="hello")) ) [_ async for _ in loop.act("Hello")] await loop.clear_history() assert loop.is_initialized assert len(loop.messages) == 1 @pytest.mark.asyncio async def test_compact_waits_for_deferred_init(self) -> None: loop = build_test_agent_loop( defer_heavy_init=True, backend=FakeBackend([ [mock_llm_chunk(content="hello")], [mock_llm_chunk(content="summary")], ]), ) [_ async for _ in loop.act("Hello")] summary = await loop.compact() assert loop.is_initialized assert summary == "summary" @pytest.mark.asyncio async def test_inject_user_context_waits_for_deferred_init(self) -> None: loop = build_test_agent_loop(defer_heavy_init=True) await loop.inject_user_context("context") assert loop.is_initialized assert loop.messages[-1].content == "context" # --------------------------------------------------------------------------- # start_initialize_experiments / wait_until_ready experiment gating # --------------------------------------------------------------------------- class TestStartInitializeExperiments: @pytest.mark.asyncio async def test_does_not_block_caller(self) -> None: loop = build_test_agent_loop() gate = asyncio.Event() async def slow_init() -> None: await gate.wait() with patch.object(loop, "initialize_experiments", side_effect=slow_init): loop.start_initialize_experiments() task = loop._experiments_task assert task is not None assert not task.done() gate.set() await task @pytest.mark.asyncio async def test_is_idempotent(self) -> None: loop = build_test_agent_loop() init_mock = AsyncMock() with patch.object(loop, "initialize_experiments", new=init_mock): loop.start_initialize_experiments() first_task = loop._experiments_task loop.start_initialize_experiments() second_task = loop._experiments_task assert first_task is second_task assert first_task is not None await first_task assert init_mock.await_count == 1 @pytest.mark.asyncio async def test_sets_pending_telemetry_flags(self) -> None: loop = build_test_agent_loop() with patch.object(loop, "initialize_experiments", new=AsyncMock()): loop.start_initialize_experiments() assert loop._pending_new_session_telemetry is True assert loop._ready_telemetry_pending is True task = loop._experiments_task assert task is not None await task @pytest.mark.asyncio async def test_refreshes_system_prompt_when_experiments_update(self) -> None: loop = build_test_agent_loop() refresh_mock = AsyncMock() with ( patch.object( agent_loop_module, "session_initialize_experiments", new=AsyncMock(return_value=True), ), patch.object(loop, "refresh_system_prompt", new=refresh_mock), ): await loop.initialize_experiments() refresh_mock.assert_awaited_once() @pytest.mark.asyncio async def test_does_not_refresh_system_prompt_when_experiments_unchanged( self, ) -> None: loop = build_test_agent_loop() refresh_mock = AsyncMock() with ( patch.object( agent_loop_module, "session_initialize_experiments", new=AsyncMock(return_value=False), ), patch.object(loop, "refresh_system_prompt", new=refresh_mock), ): await loop.initialize_experiments() refresh_mock.assert_not_awaited() class TestWaitUntilReadyJoinsExperiments: @pytest.mark.asyncio async def test_joins_in_flight_task(self) -> None: loop = build_test_agent_loop() gate = asyncio.Event() completed = False async def slow_init() -> None: nonlocal completed await gate.wait() completed = True with patch.object(loop, "initialize_experiments", side_effect=slow_init): loop.start_initialize_experiments() async def release() -> None: await asyncio.sleep(0.01) gate.set() asyncio.create_task(release()) await loop.wait_until_ready() assert completed is True task = loop._experiments_task assert task is not None assert task.done() @pytest.mark.asyncio async def test_emits_new_session_telemetry_once(self) -> None: loop = build_test_agent_loop() emit_new_session = MagicMock() with ( patch.object(loop, "initialize_experiments", new=AsyncMock()), patch.object(loop, "emit_new_session_telemetry", new=emit_new_session), ): loop.start_initialize_experiments() await loop.wait_until_ready() await loop.wait_until_ready() emit_new_session.assert_called_once() assert loop._pending_new_session_telemetry is False @pytest.mark.asyncio async def test_emits_ready_telemetry_when_only_experiments_deferred(self) -> None: loop = build_test_agent_loop() emit_ready = MagicMock() with ( patch.object(loop, "initialize_experiments", new=AsyncMock()), patch.object(loop, "emit_ready_telemetry", new=emit_ready), ): loop.start_initialize_experiments() await loop.wait_until_ready() await loop.wait_until_ready() emit_ready.assert_called_once() ((duration,), _) = emit_ready.call_args assert isinstance(duration, int) assert duration >= 0 assert loop._ready_telemetry_pending is False @pytest.mark.asyncio async def test_does_not_emit_new_session_when_only_hydrating(self) -> None: loop = build_test_agent_loop() emit_new_session = MagicMock() with ( patch.object(loop, "hydrate_experiments_from_session", new=AsyncMock()), patch.object(loop, "emit_new_session_telemetry", new=emit_new_session), ): await loop.hydrate_experiments_from_session() await loop.wait_until_ready() emit_new_session.assert_not_called() assert loop._pending_new_session_telemetry is False @pytest.mark.asyncio async def test_no_op_when_nothing_deferred(self) -> None: loop = build_test_agent_loop() emit_ready = MagicMock() emit_new_session = MagicMock() with ( patch.object(loop, "emit_ready_telemetry", new=emit_ready), patch.object(loop, "emit_new_session_telemetry", new=emit_new_session), ): await loop.wait_until_ready() emit_ready.assert_not_called() emit_new_session.assert_not_called() class TestACloseCancelsExperimentsTask: @pytest.mark.asyncio async def test_cancels_in_flight_task(self) -> None: loop = build_test_agent_loop() gate = asyncio.Event() async def never_completing() -> None: await gate.wait() with patch.object(loop, "initialize_experiments", side_effect=never_completing): loop.start_initialize_experiments() task = loop._experiments_task assert task is not None assert not task.done() await loop.aclose() assert task.done() assert task.cancelled() @pytest.mark.asyncio async def test_does_not_cancel_completed_task(self) -> None: loop = build_test_agent_loop() with patch.object(loop, "initialize_experiments", new=AsyncMock()): loop.start_initialize_experiments() task = loop._experiments_task assert task is not None await task await loop.aclose() assert task.done() assert not task.cancelled() class TestCycleAgentDuringInit: @pytest.mark.asyncio async def test_shift_tab_during_experiments_init_does_not_crash(self) -> None: """Regression: shift+tab during init crashed with RuntimeError("await wasn't used with future"). _cycle_agent ran switch_agent via asyncio.run() in a thread worker, creating a second event loop. wait_until_ready then tried to await _experiments_task (owned by the main Textual loop) from that new loop. The fix uses asyncio.run_coroutine_threadsafe() to schedule on the main loop instead. This test presses shift+tab while experiments are still initializing and asserts the worker completes without error. """ from tests.conftest import build_test_vibe_app gate = asyncio.Event() async def slow_init() -> None: await gate.wait() agent_loop = build_test_agent_loop() app = build_test_vibe_app(agent_loop=agent_loop) async with app.run_test() as pilot: with patch.object( agent_loop, "initialize_experiments", side_effect=slow_init ): agent_loop._experiments_task = None # reset so start_ re-fires agent_loop.start_initialize_experiments() assert agent_loop._experiments_task is not None assert not agent_loop._experiments_task.done() # Press shift+tab while experiments are still running. await pilot.press("shift+tab") await pilot.pause(0.05) # Unblock experiments so switch_agent can complete. gate.set() # wait_for_complete raises WorkerFailed if the thread worker # crashed — which is exactly what happened before the fix. await pilot.app.workers.wait_for_complete() assert agent_loop.agent_profile.name == "plan" class TestActGatesOnExperiments: @pytest.mark.asyncio async def test_act_awaits_experiments_before_llm_call(self) -> None: loop = build_test_agent_loop( backend=FakeBackend(mock_llm_chunk(content="hello")) ) gate = asyncio.Event() finished_init = False async def slow_init() -> None: nonlocal finished_init await gate.wait() finished_init = True with patch.object(loop, "initialize_experiments", side_effect=slow_init): loop.start_initialize_experiments() async def release() -> None: await asyncio.sleep(0.01) gate.set() asyncio.create_task(release()) events = [event async for event in loop.act("Hello")] assert finished_init is True assert any(getattr(event, "content", None) == "hello" for event in events)