ForcePilot/backend/package/yuxi/channel/interfaces/sse/endpoint.py

from __future__ import annotations

import asyncio
import json
import logging
import time
from dataclasses import dataclass

from fastapi import HTTPException, Request
from fastapi.responses import StreamingResponse

from yuxi.channel.domain.port.metrics_port import MetricsPort
from yuxi.channel.domain.port.sse_push_port import SsePushPort

logger = logging.getLogger(__name__)


@dataclass
class _SseConnection:
    session_id: str
    queue: asyncio.Queue
    last_active_at: float = 0.0


class SseEndpoint(SsePushPort):
    STALE_THRESHOLD_SECONDS = 300
    CLEANUP_INTERVAL_SECONDS = 60
    MAX_CONNECTIONS_PER_SESSION = 5

    def __init__(self, *, max_connections: int = 1000, metrics: MetricsPort | None = None) -> None:
        self._connections: dict[str, list[_SseConnection]] = {}
        self._max_connections = max_connections
        self._cleanup_task: asyncio.Task | None = None
        self._metrics = metrics

    @property
    def connection_count(self) -> int:
        return sum(len(conns) for conns in self._connections.values())

    async def start(self) -> None:
        logger.info("SSE endpoint started, max_connections=%d", self._max_connections)
        self._cleanup_task = asyncio.create_task(self._periodic_cleanup())

    async def stop(self) -> None:
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
            self._cleanup_task = None
        for conns in self._connections.values():
            for conn in conns:
                try:
                    conn.queue.put_nowait({"type": "shutdown", "reason": "server_stopping"})
                except asyncio.QueueFull:
                    pass
        await asyncio.sleep(1.0)
        self._connections.clear()

    async def subscribe(self, session_id: str, request: Request) -> StreamingResponse:
        if self.connection_count >= self._max_connections:
            await self._cleanup_stale()
            if self.connection_count >= self._max_connections:
                raise HTTPException(status_code=503, detail="SSE connection limit reached")

        session_conns = self._connections.get(session_id, [])
        if len(session_conns) >= self.MAX_CONNECTIONS_PER_SESSION:
            raise HTTPException(status_code=429, detail="too many SSE connections for this session")

        now = time.monotonic()
        queue: asyncio.Queue = asyncio.Queue(maxsize=256)
        conn = _SseConnection(session_id=session_id, queue=queue, last_active_at=now)
        self._connections.setdefault(session_id, []).append(conn)
        await self._update_connection_count()

        async def event_stream():
            try:
                yield f"event: connected\ndata: {json.dumps({'session_id': session_id}, ensure_ascii=False)}\n\n"
                while True:
                    if await request.is_disconnected():
                        break
                    try:
                        data = await asyncio.wait_for(queue.get(), timeout=30.0)
                        conn.last_active_at = time.monotonic()
                        yield f"event: message\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
                    except TimeoutError:
                        conn.last_active_at = time.monotonic()
                        yield "event: ping\ndata: {}\n\n"
            finally:
                conns = self._connections.get(session_id)
                if conns:
                    try:
                        conns.remove(conn)
                    except ValueError:
                        pass
                    if not conns:
                        self._connections.pop(session_id, None)
                await self._update_connection_count()

        return StreamingResponse(event_stream(), media_type="text/event-stream")

    async def push_event(self, session_id: str, data: dict) -> bool:
        conns = self._connections.get(session_id)
        if not conns:
            return False
        delivered = False
        for conn in conns:
            try:
                conn.queue.put_nowait(data)
                conn.last_active_at = time.monotonic()
                delivered = True
            except asyncio.QueueFull:
                pass
        return delivered

    async def broadcast_shutdown(self) -> None:
        for conns in self._connections.values():
            for conn in conns:
                try:
                    conn.queue.put_nowait({"type": "shutdown", "reason": "gateway_shutting_down"})
                except asyncio.QueueFull:
                    pass

    async def _periodic_cleanup(self) -> None:
        try:
            while True:
                await asyncio.sleep(self.CLEANUP_INTERVAL_SECONDS)
                await self._cleanup_stale()
        except asyncio.CancelledError:
            pass

    async def _update_connection_count(self) -> None:
        if self._metrics:
            try:
                await self._metrics.set_sse_connections(self.connection_count)
            except Exception:
                pass

    async def _cleanup_stale(self) -> None:
        now = time.monotonic()
        for sid in list(self._connections):
            conns = self._connections[sid]
            stale = [c for c in conns if now - c.last_active_at > self.STALE_THRESHOLD_SECONDS]
            for c in stale:
                conns.remove(c)
            if not conns:
                del self._connections[sid]

        if self.connection_count >= self._max_connections:
            all_conns: list[tuple[str, _SseConnection]] = []
            for sid, conns in self._connections.items():
                for c in conns:
                    all_conns.append((sid, c))
            all_conns.sort(key=lambda x: x[1].last_active_at)
            evict_count = len(all_conns) // 4
            for sid, conn in all_conns[:evict_count]:
                conns = self._connections.get(sid)
                if conns:
                    try:
                        conns.remove(conn)
                    except ValueError:
                        pass
                    if not conns:
                        self._connections.pop(sid, None)