roboco/tests/unit/api/test_websocket_bridge.py

"""websocket_bridge coverage — event handlers + bridge starter.

The handlers fan events from the Redis-stream bus to per-recipient WebSocket
connections. We don't need real Redis or sockets; we patch `manager` and the
`broadcast_*` helpers so each handler exercises its branches against
in-memory state.
"""

from __future__ import annotations

from unittest.mock import AsyncMock, patch
from uuid import uuid4

import pytest
from roboco.api.websocket_bridge import (
    _handle_agent_event,
    _handle_notification_sent,
    _handle_rate_limit_event,
    _handle_session_event,
    _handle_usage_event,
    register_websocket_bridge_handlers,
    start_websocket_bridge,
)
from roboco.models.events import Event, EventType


def _evt(event_type: EventType, data: dict, source_agent: str | None = None) -> Event:
    return Event(type=event_type, data=data, source_agent=source_agent)


# ---------------------------------------------------------------------------
# _handle_notification_sent
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_handle_notification_sent_skips_when_missing_ids() -> None:
    """Incomplete event (missing recipient/notification IDs) → log + return."""
    event = _evt(EventType.NOTIFICATION_SENT, {})  # No notification_id/recipient_id
    with patch("roboco.api.websocket_bridge.broadcast_notification") as bcast:
        await _handle_notification_sent(event)
    bcast.assert_not_called()


@pytest.mark.asyncio
async def test_handle_notification_sent_skips_invalid_uuid() -> None:
    """Invalid UUID strings → log error + return without broadcasting."""
    event = _evt(
        EventType.NOTIFICATION_SENT,
        {"notification_id": "not-a-uuid", "recipient_id": str(uuid4())},
    )
    with patch("roboco.api.websocket_bridge.broadcast_notification") as bcast:
        await _handle_notification_sent(event)
    bcast.assert_not_called()


@pytest.mark.asyncio
async def test_handle_notification_sent_skips_when_no_connections() -> None:
    """Recipient has no WS connections → no broadcast."""
    nid = uuid4()
    rid = uuid4()
    event = _evt(
        EventType.NOTIFICATION_SENT,
        {
            "notification_id": str(nid),
            "recipient_id": str(rid),
            "type": "blocker",
            "subject": "x",
            "priority": "high",
        },
    )
    with (
        patch("roboco.api.websocket_bridge.broadcast_notification") as bcast,
        patch("roboco.api.websocket_bridge.manager") as mgr,
    ):
        mgr.notification_connections = {}  # No connections for any agent.
        await _handle_notification_sent(event)
    bcast.assert_not_called()


@pytest.mark.asyncio
async def test_handle_notification_sent_broadcasts_when_connected() -> None:
    """Recipient has WS connection → broadcast_notification called."""
    nid = uuid4()
    rid = uuid4()
    event = _evt(
        EventType.NOTIFICATION_SENT,
        {
            "notification_id": str(nid),
            "recipient_id": str(rid),
            "type": "qa_ready",
            "subject": "Task ready",
            "priority": "normal",
        },
    )
    bcast = AsyncMock()
    with (
        patch("roboco.api.websocket_bridge.broadcast_notification", bcast),
        patch("roboco.api.websocket_bridge.manager") as mgr,
    ):
        mgr.notification_connections = {rid: {"socket-1"}}  # Has a connection.
        await _handle_notification_sent(event)
    bcast.assert_awaited_once()
    call_kwargs = bcast.await_args.kwargs
    assert call_kwargs["notification_id"] == nid
    assert call_kwargs["agent_ids"] == [rid]


@pytest.mark.asyncio
async def test_handle_notification_acked_broadcasts_using_agent_id() -> None:
    """ACKED events carry `agent_id`, not `recipient_id`; the shared handler
    must still forward (to the acking agent) rather than log 'Incomplete
    notification event' on every acknowledgement."""
    nid = uuid4()
    aid = uuid4()
    event = _evt(
        EventType.NOTIFICATION_ACKED,
        {"notification_id": str(nid), "agent_id": str(aid), "ack_type": "read"},
    )
    bcast = AsyncMock()
    with (
        patch("roboco.api.websocket_bridge.broadcast_notification", bcast),
        patch("roboco.api.websocket_bridge.manager") as mgr,
    ):
        mgr.notification_connections = {aid: {"socket-1"}}
        await _handle_notification_sent(event)
    bcast.assert_awaited_once()
    call_kwargs = bcast.await_args.kwargs
    assert call_kwargs["notification_id"] == nid
    assert call_kwargs["agent_ids"] == [aid]


# ---------------------------------------------------------------------------
# _handle_session_event
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_handle_session_event_skips_missing_session_id() -> None:
    event = _evt(EventType.SESSION_CREATED, {})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_to_session = AsyncMock()
        await _handle_session_event(event)
        mgr.broadcast_to_session.assert_not_called()


@pytest.mark.asyncio
async def test_handle_session_event_skips_invalid_uuid() -> None:
    event = _evt(EventType.SESSION_CREATED, {"session_id": "bad-uuid"})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_to_session = AsyncMock()
        await _handle_session_event(event)
        mgr.broadcast_to_session.assert_not_called()


@pytest.mark.asyncio
async def test_handle_session_event_skips_when_no_connections() -> None:
    sid = uuid4()
    event = _evt(EventType.SESSION_CREATED, {"session_id": str(sid)})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.session_connections = {}
        mgr.broadcast_to_session = AsyncMock()
        await _handle_session_event(event)
        mgr.broadcast_to_session.assert_not_called()


@pytest.mark.asyncio
async def test_handle_session_event_broadcasts() -> None:
    sid = uuid4()
    event = _evt(EventType.SESSION_CLOSED, {"session_id": str(sid)})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.session_connections = {sid: {"sock-1"}}
        mgr.broadcast_to_session = AsyncMock()
        await _handle_session_event(event)
        mgr.broadcast_to_session.assert_awaited_once()
        # Payload includes the trailing piece of the event-type ('closed').
        call_args = mgr.broadcast_to_session.await_args
        assert call_args.args[0] == sid
        assert call_args.args[1]["type"] == "session.closed"


# ---------------------------------------------------------------------------
# _handle_agent_event
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_handle_agent_event_skips_when_no_agent_id() -> None:
    event = _evt(EventType.AGENT_SPAWNED, {})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_to_agent_watchers = AsyncMock()
        await _handle_agent_event(event)
        mgr.broadcast_to_agent_watchers.assert_not_called()


@pytest.mark.asyncio
async def test_handle_agent_event_skips_invalid_uuid() -> None:
    event = _evt(EventType.AGENT_SPAWNED, {"agent_id": "bad"})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_to_agent_watchers = AsyncMock()
        await _handle_agent_event(event)
        mgr.broadcast_to_agent_watchers.assert_not_called()


@pytest.mark.asyncio
async def test_handle_agent_event_uses_source_agent_fallback() -> None:
    """When data has no agent_id, falls back to event.source_agent."""
    aid = uuid4()
    event = _evt(EventType.AGENT_STOPPED, {}, source_agent=str(aid))
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.agent_connections = {aid: {"sock"}}
        mgr.broadcast_to_agent_watchers = AsyncMock()
        await _handle_agent_event(event)
        mgr.broadcast_to_agent_watchers.assert_awaited_once()


@pytest.mark.asyncio
async def test_handle_agent_event_skips_when_no_connections() -> None:
    aid = uuid4()
    event = _evt(EventType.AGENT_SPAWNED, {"agent_id": str(aid)})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.agent_connections = {}
        mgr.broadcast_to_agent_watchers = AsyncMock()
        await _handle_agent_event(event)
        mgr.broadcast_to_agent_watchers.assert_not_called()


@pytest.mark.asyncio
async def test_handle_agent_event_broadcasts() -> None:
    aid = uuid4()
    event = _evt(EventType.AGENT_RESUMED, {"agent_id": str(aid)})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.agent_connections = {aid: {"sock"}}
        mgr.broadcast_to_agent_watchers = AsyncMock()
        await _handle_agent_event(event)
        mgr.broadcast_to_agent_watchers.assert_awaited_once()
        call_args = mgr.broadcast_to_agent_watchers.await_args
        assert call_args.args[0] == aid
        assert call_args.args[1]["type"] == "agent.resumed"


# ---------------------------------------------------------------------------
# _handle_rate_limit_event
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_handle_rate_limit_hit_broadcasts_to_system() -> None:
    """RATE_LIMIT_HIT → broadcast_system tagged with the type, payload intact."""
    retry_after = 60.0
    event = _evt(
        EventType.RATE_LIMIT_HIT,
        {
            "provider": "anthropic",
            "affectedAgents": ["be-dev-1"],
            "retryAfterSeconds": retry_after,
            "timestamp": "2026-06-11T00:00:00+00:00",
        },
    )
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_system = AsyncMock()
        await _handle_rate_limit_event(event)
    mgr.broadcast_system.assert_awaited_once()
    msg = mgr.broadcast_system.await_args.args[0]
    assert msg["type"] == "RATE_LIMIT_HIT"
    assert msg["provider"] == "anthropic"
    assert msg["affectedAgents"] == ["be-dev-1"]
    assert msg["retryAfterSeconds"] == retry_after


@pytest.mark.asyncio
async def test_handle_rate_limit_lifted_broadcasts_to_system() -> None:
    event = _evt(
        EventType.RATE_LIMIT_LIFTED,
        {"provider": "anthropic", "timestamp": "2026-06-11T00:01:00+00:00"},
    )
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_system = AsyncMock()
        await _handle_rate_limit_event(event)
    msg = mgr.broadcast_system.await_args.args[0]
    assert msg["type"] == "RATE_LIMIT_LIFTED"
    assert msg["provider"] == "anthropic"


@pytest.mark.asyncio
async def test_handle_rate_limit_ignores_unrelated_event() -> None:
    """A non-rate-limit event type is a no-op (defensive guard)."""
    event = _evt(EventType.AGENT_SPAWNED, {"provider": "anthropic"})
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_system = AsyncMock()
        await _handle_rate_limit_event(event)
    mgr.broadcast_system.assert_not_called()


# ---------------------------------------------------------------------------
# _handle_usage_event
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_handle_usage_update_broadcasts_to_system() -> None:
    """USAGE_UPDATE event → broadcast_system tagged USAGE_UPDATE + data fields."""
    expected_input = 100
    expected_output = 50
    event = _evt(
        EventType.USAGE_UPDATE,
        {
            "agent_id": "be-dev-1",
            "task_id": "task-abc",
            "input_tokens": expected_input,
            "output_tokens": expected_output,
            "model": "claude-sonnet-4-6",
            "timestamp": "2026-06-11T00:00:00+00:00",
        },
    )
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_system = AsyncMock()
        await _handle_usage_event(event)
    mgr.broadcast_system.assert_awaited_once()
    msg = mgr.broadcast_system.await_args.args[0]
    assert msg["type"] == "USAGE_UPDATE"
    assert msg["agent_id"] == "be-dev-1"
    assert msg["input_tokens"] == expected_input
    assert msg["output_tokens"] == expected_output


@pytest.mark.asyncio
async def test_handle_usage_snapshot_broadcasts_to_system() -> None:
    """USAGE_SNAPSHOT event → broadcast_system tagged USAGE_SNAPSHOT + aggregate."""
    expected_input = 500
    event = _evt(
        EventType.USAGE_SNAPSHOT,
        {
            "period": "live",
            "totals": {"input_tokens": expected_input, "output_tokens": 200},
            "cost_estimate": 0.0025,
            "by_agent": [
                {
                    "agent_id": "be-dev-1",
                    "input_tokens": 500,
                    "output_tokens": 200,
                    "model": "sonnet",
                    "cost_estimate": 0.0025,
                }
            ],
            "timestamp": "2026-06-11T00:01:00+00:00",
        },
    )
    with patch("roboco.api.websocket_bridge.manager") as mgr:
        mgr.broadcast_system = AsyncMock()
        await _handle_usage_event(event)
    mgr.broadcast_system.assert_awaited_once()
    msg = mgr.broadcast_system.await_args.args[0]
    assert msg["type"] == "USAGE_SNAPSHOT"
    assert msg["period"] == "live"
    assert msg["totals"]["input_tokens"] == expected_input
    assert len(msg["by_agent"]) == 1


# ---------------------------------------------------------------------------
# Registration + start
# ---------------------------------------------------------------------------


def test_register_websocket_bridge_handlers_subscribes_all_event_types() -> None:
    """Registration wires up all handler categories, including usage events."""

    class _FakeBus:
        def __init__(self) -> None:
            self.subscribed: list[tuple[EventType, object]] = []

        def subscribe(self, event_type: EventType, handler: object) -> None:
            self.subscribed.append((event_type, handler))

    fake = _FakeBus()
    with patch("roboco.api.websocket_bridge.get_event_bus", return_value=fake):
        register_websocket_bridge_handlers()
    types = [t for t, _ in fake.subscribed]
    # All 14 expected event types appear at least once.
    assert EventType.NOTIFICATION_SENT in types
    assert EventType.NOTIFICATION_ACKED in types
    assert EventType.SESSION_CREATED in types
    assert EventType.SESSION_CLOSED in types
    assert EventType.SESSION_TIMEOUT in types
    assert EventType.AGENT_SPAWNED in types
    assert EventType.AGENT_STOPPED in types
    assert EventType.AGENT_WAITING in types
    assert EventType.AGENT_RESUMED in types
    assert EventType.AGENT_ERROR in types
    assert EventType.RATE_LIMIT_HIT in types
    assert EventType.RATE_LIMIT_LIFTED in types
    # Usage events forwarded to /ws/system.
    assert EventType.USAGE_UPDATE in types
    assert EventType.USAGE_SNAPSHOT in types


@pytest.mark.asyncio
async def test_start_websocket_bridge_registers_handlers() -> None:
    """start_websocket_bridge() calls register_websocket_bridge_handlers."""
    with patch("roboco.api.websocket_bridge.register_websocket_bridge_handlers") as reg:
        await start_websocket_bridge()
    reg.assert_called_once()