app/routers/stream.py

from __future__ import annotations

import asyncio
import json
import logging
from typing import AsyncIterator

from fastapi import APIRouter, HTTPException
from sse_starlette.sse import EventSourceResponse

from app.services import jobs
from app.services.jobs import SENTINEL, StreamEvent

logger = logging.getLogger(__name__)
router = APIRouter(tags=["stream"])

HEARTBEAT_SECONDS = 15.0


@router.get("/stream/{job_id}")
async def stream(job_id: str) -> EventSourceResponse:
    job = jobs.store().get(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="job not found")

    async def event_source() -> AsyncIterator[dict[str, str]]:
        while True:
            try:
                item = await asyncio.wait_for(job.queue.get(), timeout=HEARTBEAT_SECONDS)
            except asyncio.TimeoutError:
                yield {"event": "heartbeat", "data": "{}"}
                if job.finished.is_set():
                    return
                continue

            if item is SENTINEL:
                return

            assert isinstance(item, StreamEvent)
            payload = (
                item.data.model_dump()
                if hasattr(item.data, "model_dump")
                else item.data
            )
            yield {"event": item.type, "data": json.dumps(payload, default=str)}

            if item.type in ("report_complete", "error"):
                return

    return EventSourceResponse(event_source())