95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
"""SSE routes for live dashboard updates."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING
|
|
|
|
from fastapi import APIRouter
|
|
from sse_starlette.sse import EventSourceResponse
|
|
|
|
if TYPE_CHECKING:
|
|
from ..db import Database
|
|
from ..notifications import NotificationBus
|
|
from ..scheduler import Scheduler
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/sse")
|
|
|
|
_notification_bus: NotificationBus | None = None
|
|
_scheduler: Scheduler | None = None
|
|
_db: Database | None = None
|
|
|
|
|
|
def setup(notification_bus, scheduler, db):
|
|
global _notification_bus, _scheduler, _db
|
|
_notification_bus = notification_bus
|
|
_scheduler = scheduler
|
|
_db = db
|
|
|
|
|
|
@router.get("/notifications")
|
|
async def sse_notifications():
|
|
"""Stream new notifications as they arrive."""
|
|
listener_id = f"sse-notif-{id(asyncio.current_task())}"
|
|
|
|
# Subscribe to notification bus
|
|
queue: asyncio.Queue = asyncio.Queue()
|
|
loop = asyncio.get_event_loop()
|
|
|
|
if _notification_bus:
|
|
def on_notify(msg, cat):
|
|
loop.call_soon_threadsafe(
|
|
queue.put_nowait, {"message": msg, "category": cat}
|
|
)
|
|
_notification_bus.subscribe(listener_id, on_notify)
|
|
|
|
async def generate():
|
|
try:
|
|
while True:
|
|
try:
|
|
notif = await asyncio.wait_for(queue.get(), timeout=30)
|
|
yield {
|
|
"event": "notification",
|
|
"data": json.dumps(notif),
|
|
}
|
|
except TimeoutError:
|
|
yield {"event": "heartbeat", "data": ""}
|
|
finally:
|
|
if _notification_bus:
|
|
_notification_bus.unsubscribe(listener_id)
|
|
|
|
return EventSourceResponse(generate())
|
|
|
|
|
|
@router.get("/loops")
|
|
async def sse_loops():
|
|
"""Push loop timestamps + active executions every 15s."""
|
|
async def generate():
|
|
while True:
|
|
data = {"loops": {}, "executions": {}}
|
|
if _scheduler:
|
|
ts = _scheduler.get_loop_timestamps()
|
|
data["loops"] = ts
|
|
# Serialize active executions (datetime -> str)
|
|
raw_exec = _scheduler.get_active_executions()
|
|
execs = {}
|
|
for tid, info in raw_exec.items():
|
|
execs[tid] = {
|
|
"name": info.get("name", ""),
|
|
"tool": info.get("tool", ""),
|
|
"started_at": info["started_at"].isoformat()
|
|
if isinstance(info.get("started_at"), datetime)
|
|
else str(info.get("started_at", "")),
|
|
"thread": info.get("thread", ""),
|
|
}
|
|
data["executions"] = execs
|
|
yield {"event": "loops", "data": json.dumps(data)}
|
|
await asyncio.sleep(15)
|
|
|
|
return EventSourceResponse(generate())
|