fix: add integration tests for agentic loop WebSocket broadcasts

Verify that plan_ready, step_complete, step_adapted, and task_complete
events flow through the real WebSocketManager to connected WS clients.
Also tests error resilience when WS disconnects mid-loop.

Fixes #445

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

tests/integrations/test_agentic_ws_broadcast.py

@@ -1,285 +0,0 @@
-"""Integration tests for agentic loop WebSocket broadcasts.
-
-Verifies that ``run_agentic_loop`` pushes the correct sequence of events
-through the real ``ws_manager`` and that connected (mock) WebSocket clients
-receive every broadcast with the expected payloads.
-"""
-
-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
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _mock_run(content: str):
-    m = MagicMock()
-    m.content = content
-    return m
-
-
-def _ws_client() -> AsyncMock:
-    """Return a fake WebSocket that records sent messages."""
-    return AsyncMock()
-
-
-def _collected_events(ws: AsyncMock) -> list[dict]:
-    """Extract parsed JSON events from a mock WebSocket's send_text calls."""
-    return [json.loads(call.args[0]) for call in ws.send_text.call_args_list]
-
-
-# ---------------------------------------------------------------------------
-# Tests
-# ---------------------------------------------------------------------------
-
-
-class TestAgenticLoopBroadcastSequence:
-    """Events arrive at WS clients in the correct order with expected data."""
-
-    @pytest.mark.asyncio
-    async def test_successful_run_broadcasts_plan_steps_complete(self):
-        """A successful 2-step loop emits plan_ready → 2× step_complete → task_complete."""
-        mgr = WebSocketManager()
-        ws = _ws_client()
-        mgr._connections = [ws]
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Gather data\n2. Summarise"),
-                _mock_run("Gathered 10 records"),
-                _mock_run("Summary written"),
-            ]
-        )
-
-        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("Gather and summarise", max_steps=2)
-
-        assert result.status == "completed"
-
-        events = _collected_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_plan_ready_payload(self):
-        """plan_ready contains task_id, task, steps list, and total count."""
-        mgr = WebSocketManager()
-        ws = _ws_client()
-        mgr._connections = [ws]
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Alpha\n2. Beta"),
-                _mock_run("Alpha done"),
-                _mock_run("Beta 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("Two steps")
-
-        plan_event = _collected_events(ws)[0]
-        assert plan_event["event"] == "agentic.plan_ready"
-        data = plan_event["data"]
-        assert data["task_id"] == result.task_id
-        assert data["task"] == "Two steps"
-        assert data["steps"] == ["Alpha", "Beta"]
-        assert data["total"] == 2
-
-    @pytest.mark.asyncio
-    async def test_step_complete_payload(self):
-        """step_complete carries step number, total, description, and result."""
-        mgr = WebSocketManager()
-        ws = _ws_client()
-        mgr._connections = [ws]
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Only step"),
-                _mock_run("Step result text"),
-            ]
-        )
-
-        with (
-            patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
-            patch("infrastructure.ws_manager.handler.ws_manager", mgr),
-        ):
-            await run_agentic_loop("Single step", max_steps=1)
-
-        step_event = _collected_events(ws)[1]
-        assert step_event["event"] == "agentic.step_complete"
-        data = step_event["data"]
-        assert data["step"] == 1
-        assert data["total"] == 1
-        assert data["description"] == "Only step"
-        assert "Step result text" in data["result"]
-
-    @pytest.mark.asyncio
-    async def test_task_complete_payload(self):
-        """task_complete has status, steps_completed, summary, and duration_ms."""
-        mgr = WebSocketManager()
-        ws = _ws_client()
-        mgr._connections = [ws]
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Do it"),
-                _mock_run("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("Quick", max_steps=1)
-
-        complete_event = _collected_events(ws)[-1]
-        assert complete_event["event"] == "agentic.task_complete"
-        data = complete_event["data"]
-        assert data["status"] == "completed"
-        assert data["steps_completed"] == 1
-        assert isinstance(data["duration_ms"], int)
-        assert data["duration_ms"] >= 0
-        assert data["summary"]
-
-
-class TestAdaptationBroadcast:
-    """Adapted steps emit step_adapted events."""
-
-    @pytest.mark.asyncio
-    async def test_adapted_step_broadcasts_step_adapted(self):
-        """A failed-then-adapted step emits agentic.step_adapted."""
-        mgr = WebSocketManager()
-        ws = _ws_client()
-        mgr._connections = [ws]
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Risky step"),
-                Exception("disk full"),
-                _mock_run("Used /tmp instead"),
-            ]
-        )
-
-        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("Adapt test", max_steps=1)
-
-        events = _collected_events(ws)
-        event_names = [e["event"] for e in events]
-        assert "agentic.step_adapted" in event_names
-
-        adapted = next(e for e in events if e["event"] == "agentic.step_adapted")
-        assert adapted["data"]["error"] == "disk full"
-        assert adapted["data"]["adaptation"]
-        assert result.steps[0].status == "adapted"
-
-
-class TestMultipleClients:
-    """All connected clients receive every broadcast."""
-
-    @pytest.mark.asyncio
-    async def test_two_clients_receive_all_events(self):
-        mgr = WebSocketManager()
-        ws1 = _ws_client()
-        ws2 = _ws_client()
-        mgr._connections = [ws1, ws2]
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Step A"),
-                _mock_run("A 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("Multi-client", max_steps=1)
-
-        events1 = _collected_events(ws1)
-        events2 = _collected_events(ws2)
-        assert len(events1) == len(events2) == 3  # plan + step + complete
-        assert [e["event"] for e in events1] == [e["event"] for e in events2]
-
-
-class TestEventHistory:
-    """Broadcasts are recorded in ws_manager event history."""
-
-    @pytest.mark.asyncio
-    async def test_events_appear_in_history(self):
-        mgr = WebSocketManager()
-
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Only"),
-                _mock_run("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("History test", max_steps=1)
-
-        history_events = [e.event for e in mgr.event_history]
-        assert "agentic.plan_ready" in history_events
-        assert "agentic.step_complete" in history_events
-        assert "agentic.task_complete" in history_events
-
-
-class TestBroadcastGracefulDegradation:
-    """Loop completes even when ws_manager is unavailable."""
-
-    @pytest.mark.asyncio
-    async def test_loop_succeeds_when_broadcast_fails(self):
-        """ImportError from ws_manager doesn't crash the loop."""
-        mock_agent = MagicMock()
-        mock_agent.run = MagicMock(
-            side_effect=[
-                _mock_run("1. Do it"),
-                _mock_run("Done"),
-            ]
-        )
-
-        with (
-            patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
-            patch(
-                "infrastructure.ws_manager.handler.ws_manager",
-                new_callable=lambda: MagicMock,
-            ) as broken_mgr,
-        ):
-            broken_mgr.broadcast = AsyncMock(side_effect=RuntimeError("ws down"))
-            result = await run_agentic_loop("Resilient task", max_steps=1)
-
-        assert result.status == "completed"
-        assert len(result.steps) == 1

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

tests/unit/test_memory_embeddings.py

@@ -99,9 +99,16 @@ class TestKeywordOverlap:
 
 
 class TestEmbedText:
+    def setup_method(self):
+        self._saved_model = emb.EMBEDDING_MODEL
+        emb.EMBEDDING_MODEL = None
+
+    def teardown_method(self):
+        emb.EMBEDDING_MODEL = self._saved_model
+
     def test_uses_fallback_when_model_disabled(self):
-        with patch.object(emb, "_get_embedding_model", return_value=False):
-            vec = embed_text("test")
+        emb.EMBEDDING_MODEL = False
+        vec = embed_text("test")
         assert len(vec) == 128  # hash fallback dimension
 
     def test_uses_model_when_available(self):
@@ -109,9 +116,9 @@ class TestEmbedText:
         mock_encoding.tolist.return_value = [0.1, 0.2, 0.3]
         mock_model = MagicMock()
         mock_model.encode.return_value = mock_encoding
+        emb.EMBEDDING_MODEL = mock_model
 
-        with patch.object(emb, "_get_embedding_model", return_value=mock_model):
-            result = embed_text("test")
+        result = embed_text("test")
         assert result == pytest.approx([0.1, 0.2, 0.3])
         mock_model.encode.assert_called_once_with("test")