External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
da563487b89c5035ad7863a4acfe7e880510f331
roboco/tests/unit/api/test_websocket_bridge.py
T
Renzo FGitHubRenn F
 da563487b8 Wave 2 features: A2A live view (CEO chime-in + reply budget) and prompter memory (#297) 
* feat(a2a): live view — watch fleet conversations, CEO chime-in, reply budget

A2A_MESSAGE_SENT published from A2AService.send (excerpt-capped) and
fanned through the existing /ws/system bridge; CEO-only admin REST for
conversations/messages + a reply route on the publish-bearing send path;
panel /a2a page with live transcript and a composer gated on task-linked
conversations. The matrix gains its one asymmetric rule: CEO may message
anyone, nobody may target the CEO — and agent replies inside a
CEO-opened conversation are hard-budgeted to one per CEO message
(per conversation, per agent), rejected with wait-don't-retry guidance.
Built subagent-driven (Sonnet 5), reviewed; v1 seams documented in the
map delta.

* feat(prompter): intake remembers the task history

Intake spawns now carry a per-project chronological digest of recent
tasks (capped: 15 lines/project, 4000 chars total — ~300-1000 tokens)
merged into the ambient layer, and the interviewer gets a bounded
search_past_tasks tool (one shared implementation behind the grok MCP
tool and the Claude SDK in-process tool) to check precedent
mid-conversation. Informational memory only — the sequencing analyzer
keeps ownership of ordering. Built subagent-driven (Sonnet 5), reviewed;
pre-existing conventions-ambient MegaTask-scope gap flagged, untouched.

---------

Co-authored-by: Renn F <rennf93@users.noreply.github.com>
2026-07-03 00:07:55 +02:00

517 lines
19 KiB
Python
Raw Blame History

"""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_a2a_message_event,
_handle_agent_event,
_handle_message_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()
assert bcast.await_args is not None
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()
assert bcast.await_args is not None
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_message_event
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_handle_message_event_skips_missing_ids() -> None:
"""A MESSAGE_SENT event with no session_id/message_id → no broadcast."""
event = _evt(EventType.MESSAGE_SENT, {"content": "x"})
with patch("roboco.api.websocket_bridge.manager") as mgr:
mgr.broadcast_to_session = AsyncMock()
mgr.broadcast_to_channel = AsyncMock()
await _handle_message_event(event)
mgr.broadcast_to_session.assert_not_called()
mgr.broadcast_to_channel.assert_not_called()
@pytest.mark.asyncio
async def test_handle_message_event_skips_invalid_uuid() -> None:
event = _evt(
EventType.MESSAGE_SENT,
{"session_id": "bad", "channel_id": "bad", "message_id": "bad"},
)
with patch("roboco.api.websocket_bridge.manager") as mgr:
mgr.broadcast_to_session = AsyncMock()
mgr.broadcast_to_channel = AsyncMock()
await _handle_message_event(event)
mgr.broadcast_to_session.assert_not_called()
mgr.broadcast_to_channel.assert_not_called()
@pytest.mark.asyncio
async def test_handle_message_event_broadcasts_to_session_and_channel() -> None:
"""A MESSAGE_SENT event fans out a `message.new` payload to both the
session stream and the channel stream — the payload the panel's
useChannelStream/useSessionStream filters on."""
sid = uuid4()
cid = uuid4()
mid = uuid4()
aid = uuid4()
event = _evt(
EventType.MESSAGE_SENT,
{
"session_id": str(sid),
"channel_id": str(cid),
"message_id": str(mid),
"agent_id": str(aid),
"content": "hello",
"message_type": "dialogue",
"timestamp": "2026-06-30T00:00:00+00:00",
},
)
with patch("roboco.api.websocket_bridge.manager") as mgr:
mgr.broadcast_to_session = AsyncMock()
mgr.broadcast_to_channel = AsyncMock()
await _handle_message_event(event)
mgr.broadcast_to_session.assert_awaited_once()
mgr.broadcast_to_channel.assert_awaited_once()
s_payload = mgr.broadcast_to_session.await_args.args[1]
c_payload = mgr.broadcast_to_channel.await_args.args[1]
assert s_payload["type"] == "message.new"
assert c_payload["type"] == "message.new"
assert s_payload["message_id"] == str(mid)
assert s_payload["session_id"] == str(sid)
assert s_payload["channel_id"] == str(cid)
assert s_payload["content"] == "hello"
assert mgr.broadcast_to_session.await_args.args[0] == sid
assert mgr.broadcast_to_channel.await_args.args[0] == cid
@pytest.mark.asyncio
async def test_handle_message_event_skips_when_no_connections() -> None:
"""No subscribers on either stream → broadcast helpers still called (they
no-op internally), but the UUIDs must resolve without error."""
sid = uuid4()
cid = uuid4()
event = _evt(
EventType.MESSAGE_SENT,
{
"session_id": str(sid),
"channel_id": str(cid),
"message_id": str(uuid4()),
"agent_id": str(uuid4()),
"content": "x",
"message_type": "dialogue",
},
)
with patch("roboco.api.websocket_bridge.manager") as mgr:
mgr.broadcast_to_session = AsyncMock()
mgr.broadcast_to_channel = AsyncMock()
await _handle_message_event(event)
# Forwarder always dispatches; the manager no-ops on empty sets.
mgr.broadcast_to_session.assert_awaited_once()
mgr.broadcast_to_channel.assert_awaited_once()
# ---------------------------------------------------------------------------
# _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_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
# ---------------------------------------------------------------------------
# _handle_a2a_message_event
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_handle_a2a_message_event_broadcasts_to_system() -> None:
"""An A2A_MESSAGE_SENT event is forwarded to /ws/system as an
`a2a.message` frame — the CEO's live view of every agent-to-agent chat."""
event = _evt(
EventType.A2A_MESSAGE_SENT,
{
"conversation_id": "conv-1",
"message_id": "msg-1",
"task_id": "task-1",
"from_agent": "be-dev-1",
"to_agent": "be-qa",
"skill": "code_review",
"body_excerpt": "please review",
"timestamp": "2026-07-02T00:00:00+00:00",
},
)
with patch("roboco.api.websocket_bridge.manager") as mgr:
mgr.broadcast_system = AsyncMock()
await _handle_a2a_message_event(event)
mgr.broadcast_system.assert_awaited_once()
msg = mgr.broadcast_system.await_args.args[0]
assert msg["type"] == "a2a.message"
assert msg["conversation_id"] == "conv-1"
assert msg["message_id"] == "msg-1"
assert msg["task_id"] == "task-1"
assert msg["from_agent"] == "be-dev-1"
assert msg["to_agent"] == "be-qa"
assert msg["skill"] == "code_review"
assert msg["body_excerpt"] == "please review"
assert msg["timestamp"] == "2026-07-02T00:00:00+00:00"
# ---------------------------------------------------------------------------
# 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 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_SNAPSHOT in types
# Message delivery forwarded to /ws/channels + /ws/sessions.
assert EventType.MESSAGE_SENT in types
# A2A live chat forwarded to /ws/system (CEO live view).
assert EventType.A2A_MESSAGE_SENT 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()
Reference in New Issue View Git Blame Copy Permalink