47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
|
|
from app.core.contracts import RuntimeEvent
|
|
from app.events.event_bus import EventBus
|
|
|
|
|
|
class StreamingManager:
|
|
"""Simple in-process projection from event bus to websocket consumers."""
|
|
|
|
def __init__(self, event_bus: EventBus) -> None:
|
|
self._event_bus = event_bus
|
|
self._subscribers: dict[str, list[StreamSubscriber]] = defaultdict(list)
|
|
self._event_bus.subscribe(self._on_event)
|
|
|
|
def replay_events(self, task_id: str) -> list[RuntimeEvent]:
|
|
return self._event_bus.list_for_task(task_id)
|
|
|
|
def subscribe(self, task_id: str) -> asyncio.Queue[RuntimeEvent]:
|
|
queue: asyncio.Queue[RuntimeEvent] = asyncio.Queue()
|
|
self._subscribers[task_id].append(
|
|
StreamSubscriber(loop=asyncio.get_running_loop(), queue=queue)
|
|
)
|
|
return queue
|
|
|
|
def unsubscribe(self, task_id: str, queue: asyncio.Queue[RuntimeEvent]) -> None:
|
|
listeners = self._subscribers.get(task_id, [])
|
|
for listener in list(listeners):
|
|
if listener.queue is queue:
|
|
listeners.remove(listener)
|
|
break
|
|
if not listeners and task_id in self._subscribers:
|
|
del self._subscribers[task_id]
|
|
|
|
def _on_event(self, event: RuntimeEvent) -> None:
|
|
for listener in list(self._subscribers.get(event.task_id, [])):
|
|
listener.loop.call_soon_threadsafe(listener.queue.put_nowait, event)
|
|
|
|
|
|
@dataclass
|
|
class StreamSubscriber:
|
|
loop: asyncio.AbstractEventLoop
|
|
queue: asyncio.Queue[RuntimeEvent]
|