""" SSEManager — pub/sub hub for Server-Sent Events. Task runner calls publish() → all connected SSE clients watching that task_id receive the event. """ import asyncio import json from dataclasses import dataclass, field from datetime import datetime, timezone @dataclass class SSEEvent: type: str # "task_event" | "task_status" | "task_result" task_id: str data: dict def encode(self) -> str: payload = json.dumps({"task_id": self.task_id, **self.data}) return f"event: {self.type}\ndata: {payload}\n\n" class SSEManager: def __init__(self) -> None: # task_id -> list[Queue] self._subscribers: dict[str, list[asyncio.Queue]] = {} self._lock = asyncio.Lock() async def subscribe(self, task_id: str, queue: asyncio.Queue) -> None: async with self._lock: self._subscribers.setdefault(task_id, []).append(queue) async def unsubscribe(self, task_id: str, queue: asyncio.Queue) -> None: async with self._lock: subs = self._subscribers.get(task_id, []) try: subs.remove(queue) except ValueError: pass if not subs: self._subscribers.pop(task_id, None) async def publish(self, event: SSEEvent) -> None: async with self._lock: queues = list(self._subscribers.get(event.task_id, [])) for queue in queues: try: queue.put_nowait(event) except asyncio.QueueFull: pass # slow client, drop event def subscriber_count(self, task_id: str) -> int: return len(self._subscribers.get(task_id, [])) # Module-level singleton shared across all requests sse_manager = SSEManager()