78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
from collections.abc import AsyncGenerator
|
|
from typing import Any
|
|
|
|
from app.logging_utils import structured_log
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RunEventPublisher:
|
|
def __init__(self) -> None:
|
|
# Maps run_id to a set of subscriber queues
|
|
self._subscribers: dict[int, set[asyncio.Queue]] = {}
|
|
|
|
def subscribe(self, run_id: int) -> asyncio.Queue:
|
|
if run_id not in self._subscribers:
|
|
self._subscribers[run_id] = set()
|
|
queue: asyncio.Queue[Any] = asyncio.Queue()
|
|
self._subscribers[run_id].add(queue)
|
|
structured_log(
|
|
logger,
|
|
"debug",
|
|
"runs.event_subscriber_added",
|
|
run_id=run_id,
|
|
subscriber_count=len(self._subscribers[run_id]),
|
|
)
|
|
return queue
|
|
|
|
def unsubscribe(self, run_id: int, queue: asyncio.Queue) -> None:
|
|
if run_id in self._subscribers:
|
|
self._subscribers[run_id].discard(queue)
|
|
if not self._subscribers[run_id]:
|
|
self._subscribers.pop(run_id, None)
|
|
|
|
async def publish(self, run_id: int, event_type: str, data: dict[str, Any]) -> None:
|
|
if run_id not in self._subscribers:
|
|
return
|
|
|
|
message = {"type": event_type, "data": data}
|
|
|
|
# Fan-out to all active subscribers for this run
|
|
for queue in list(self._subscribers[run_id]):
|
|
try:
|
|
queue.put_nowait(message)
|
|
except asyncio.QueueFull:
|
|
structured_log(
|
|
logger,
|
|
"warning",
|
|
"runs.event_subscriber_queue_full",
|
|
run_id=run_id,
|
|
)
|
|
|
|
|
|
run_events = RunEventPublisher()
|
|
|
|
|
|
async def event_generator(run_id: int) -> AsyncGenerator[str, None]:
|
|
queue = run_events.subscribe(run_id)
|
|
try:
|
|
while True:
|
|
# Wait for a new event
|
|
message = await queue.get()
|
|
# Server-Sent Events format: "event: <type>\ndata: <json>\n\n"
|
|
event_type = message["type"]
|
|
data_str = json.dumps(message["data"])
|
|
yield f"event: {event_type}\ndata: {data_str}\n\n"
|
|
except asyncio.CancelledError:
|
|
structured_log(
|
|
logger,
|
|
"debug",
|
|
"runs.event_stream_disconnected",
|
|
run_id=run_id,
|
|
)
|
|
raise
|
|
finally:
|
|
run_events.unsubscribe(run_id, queue)
|