ad0d5c9520
Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Peter Evers <peter.evers@mistral.ai> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
557 lines
19 KiB
Python
557 lines
19 KiB
Python
"""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)
|