This commit is contained in:
2026-05-17 23:09:56 +08:00
parent 1b4f4c836e
commit ddc285b8f4
36 changed files with 4552 additions and 872 deletions
+16 -6
View File
@@ -2,6 +2,7 @@ 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
@@ -12,7 +13,7 @@ class StreamingManager:
def __init__(self, event_bus: EventBus) -> None:
self._event_bus = event_bus
self._subscribers: dict[str, list[asyncio.Queue[RuntimeEvent]]] = defaultdict(list)
self._subscribers: dict[str, list[StreamSubscriber]] = defaultdict(list)
self._event_bus.subscribe(self._on_event)
def replay_events(self, task_id: str) -> list[RuntimeEvent]:
@@ -20,17 +21,26 @@ class StreamingManager:
def subscribe(self, task_id: str) -> asyncio.Queue[RuntimeEvent]:
queue: asyncio.Queue[RuntimeEvent] = asyncio.Queue()
self._subscribers[task_id].append(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, [])
if queue in listeners:
listeners.remove(queue)
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 queue in self._subscribers.get(event.task_id, []):
queue.put_nowait(event)
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]