| 1 | import asyncio |
| 2 | import json |
| 3 | from collections import defaultdict |
| 4 | from typing import Any, AsyncIterator, Dict, Set |
| 5 | |
| 6 | _subscribers: Dict[str, Set[asyncio.Queue[dict[str, Any]]]] = defaultdict(set) |
| 7 | |
| 8 | |
| 9 | def publish_task_event(task_id: str, event: dict[str, Any]) -> None: |
| 10 | queues = list(_subscribers.get(task_id, set())) |
| 11 | if not queues: |
| 12 | return |
| 13 | payload = {"task_id": task_id, **event} |
| 14 | for queue in queues: |
| 15 | if queue.full(): |
| 16 | try: |
| 17 | queue.get_nowait() |
| 18 | except asyncio.QueueEmpty: |
| 19 | pass |
| 20 | queue.put_nowait(payload) |
| 21 | |
| 22 | |
| 23 | async def task_event_stream(task_id: str, initial_event: dict[str, Any] | None = None) -> AsyncIterator[str]: |
| 24 | queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=100) |
| 25 | _subscribers[task_id].add(queue) |
| 26 | try: |
| 27 | if initial_event: |
| 28 | yield _format_sse(initial_event) |
| 29 | while True: |
| 30 | try: |
| 31 | event = await asyncio.wait_for(queue.get(), timeout=20) |
| 32 | yield _format_sse(event) |
| 33 | if event.get("type") in {"completed", "failed"}: |
| 34 | return |
| 35 | except asyncio.TimeoutError: |
| 36 | yield ": heartbeat\n\n" |
| 37 | finally: |
| 38 | _subscribers[task_id].discard(queue) |
| 39 | if not _subscribers[task_id]: |
| 40 | _subscribers.pop(task_id, None) |
| 41 | |
| 42 | |
| 43 | def _format_sse(event: dict[str, Any]) -> str: |
| 44 | return f"data: {json.dumps(event, ensure_ascii=False)}\n\n" |
| 45 |