diff --git a/tests/timmy/test_agentic_loop_integration.py b/tests/timmy/test_agentic_loop_integration.py new file mode 100644 index 00000000..eaa985c6 --- /dev/null +++ b/tests/timmy/test_agentic_loop_integration.py @@ -0,0 +1,281 @@ +"""Integration tests for agentic loop WebSocket broadcasts. + +Verifies that agentic loop events (plan_ready, step_complete, task_complete) +flow through the real WebSocketManager to connected clients, and that the +loop survives a WS disconnect mid-execution. + +Fixes #445 +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from infrastructure.ws_manager.handler import WebSocketManager +from timmy.agentic_loop import run_agentic_loop + + +def _mock_run(content: str): + """Create a mock return value for agent.run().""" + m = MagicMock() + m.content = content + return m + + +def _fake_ws() -> AsyncMock: + """Return an AsyncMock that behaves like a WebSocket connection.""" + ws = AsyncMock() + ws.send_text = AsyncMock() + return ws + + +def _collect_events(ws_mock: AsyncMock) -> list[dict]: + """Extract parsed JSON events from a mock WebSocket's send_text calls.""" + events = [] + for call in ws_mock.send_text.call_args_list: + raw = call.args[0] if call.args else call.kwargs.get("data", "") + events.append(json.loads(raw)) + return events + + +# --------------------------------------------------------------------------- +# Integration: events reach a real WebSocketManager +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_plan_ready_event_received_by_ws_client(): + """A connected WS client receives an agentic.plan_ready event.""" + mgr = WebSocketManager() + ws = _fake_ws() + mgr._connections = [ws] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Gather data\n2. Summarise"), + _mock_run("Data gathered"), + _mock_run("Summary written"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + await run_agentic_loop("Gather and summarise", max_steps=2) + + events = _collect_events(ws) + plan_events = [e for e in events if e["event"] == "agentic.plan_ready"] + assert len(plan_events) == 1 + assert plan_events[0]["data"]["steps"] == ["Gather data", "Summarise"] + assert plan_events[0]["data"]["total"] == 2 + + +@pytest.mark.asyncio +async def test_step_complete_events_received_by_ws_client(): + """A connected WS client receives agentic.step_complete for each step.""" + mgr = WebSocketManager() + ws = _fake_ws() + mgr._connections = [ws] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Step A\n2. Step B"), + _mock_run("A done"), + _mock_run("B done"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + await run_agentic_loop("Do A and B", max_steps=2) + + events = _collect_events(ws) + step_events = [e for e in events if e["event"] == "agentic.step_complete"] + assert len(step_events) == 2 + assert step_events[0]["data"]["step"] == 1 + assert step_events[1]["data"]["step"] == 2 + + +@pytest.mark.asyncio +async def test_task_complete_event_received_by_ws_client(): + """A connected WS client receives agentic.task_complete at the end.""" + mgr = WebSocketManager() + ws = _fake_ws() + mgr._connections = [ws] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Only step"), + _mock_run("Done"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + result = await run_agentic_loop("One step task", max_steps=1) + + events = _collect_events(ws) + task_events = [e for e in events if e["event"] == "agentic.task_complete"] + assert len(task_events) == 1 + assert task_events[0]["data"]["status"] == "completed" + assert task_events[0]["data"]["task_id"] == result.task_id + + +@pytest.mark.asyncio +async def test_all_event_types_received_in_order(): + """Client receives plan_ready → step_complete(s) → task_complete in order.""" + mgr = WebSocketManager() + ws = _fake_ws() + mgr._connections = [ws] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. First\n2. Second"), + _mock_run("First done"), + _mock_run("Second done"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + await run_agentic_loop("Two steps", max_steps=2) + + events = _collect_events(ws) + event_names = [e["event"] for e in events] + assert event_names == [ + "agentic.plan_ready", + "agentic.step_complete", + "agentic.step_complete", + "agentic.task_complete", + ] + + +@pytest.mark.asyncio +async def test_adapted_step_broadcasts_step_adapted(): + """When a step fails and adapts, client receives agentic.step_adapted.""" + mgr = WebSocketManager() + ws = _fake_ws() + mgr._connections = [ws] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Risky step"), + Exception("network error"), + _mock_run("Adapted approach worked"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + await run_agentic_loop("Risky task", max_steps=1) + + events = _collect_events(ws) + event_names = [e["event"] for e in events] + assert "agentic.step_adapted" in event_names + adapted = [e for e in events if e["event"] == "agentic.step_adapted"][0] + assert adapted["data"]["error"] == "network error" + + +# --------------------------------------------------------------------------- +# Resilience: WS disconnect mid-loop +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_ws_disconnect_mid_loop_does_not_crash(): + """If a WS client disconnects during the loop, the loop completes.""" + mgr = WebSocketManager() + ws = _fake_ws() + # First send succeeds (plan_ready), then WS dies + ws.send_text = AsyncMock(side_effect=[None, ConnectionError("gone")]) + mgr._connections = [ws] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Step A\n2. Step B"), + _mock_run("A done"), + _mock_run("B done"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + result = await run_agentic_loop("Survive disconnect", max_steps=2) + + # Loop completed despite WS failure + assert result.status == "completed" + assert len(result.steps) == 2 + + +@pytest.mark.asyncio +async def test_no_ws_connections_does_not_crash(): + """Loop completes normally when no WS clients are connected.""" + mgr = WebSocketManager() + # No connections at all + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Solo step"), + _mock_run("Done alone"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + result = await run_agentic_loop("No audience", max_steps=1) + + assert result.status == "completed" + assert len(result.steps) == 1 + + +@pytest.mark.asyncio +async def test_multiple_ws_clients_all_receive_events(): + """All connected WS clients receive the same broadcast events.""" + mgr = WebSocketManager() + ws1 = _fake_ws() + ws2 = _fake_ws() + mgr._connections = [ws1, ws2] + + mock_agent = MagicMock() + mock_agent.run = MagicMock( + side_effect=[ + _mock_run("1. Shared step"), + _mock_run("Shared result"), + ] + ) + + with ( + patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent), + patch("infrastructure.ws_manager.handler.ws_manager", mgr), + ): + await run_agentic_loop("Broadcast to all", max_steps=1) + + events1 = _collect_events(ws1) + events2 = _collect_events(ws2) + + names1 = [e["event"] for e in events1] + names2 = [e["event"] for e in events2] + assert names1 == names2 + assert "agentic.plan_ready" in names1 + assert "agentic.task_complete" in names1