ForcePilot/backend/package/yuxi/channel/gateway/sse.py

"""Gateway HTTP SSE 端点 — 为不支持 WebSocket 的客户端提供备选流式通道。

复用 stream_agent_chat 的流式输出能力，通过 HTTP SSE (text/event-stream) 推送给浏览器客户端。

SSE 事件类型对齐 CowAgent 标准：
  reasoning / delta / tool_start / tool_end / message_end /
  phase / image / file / video / done / error

路由：
  POST /api/sse/chat     — 发送聊天消息，返回 request_id
  GET  /api/sse/stream   — 订阅 SSE 事件流
  GET  /api/poll         — Polling 降级轮询
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid

from fastapi import APIRouter, Header, Query, Request
from fastapi.responses import StreamingResponse

from yuxi.channel.protocols import SseEventType
from yuxi.storage.postgres.manager import pg_manager

logger = logging.getLogger(__name__)

SSE_KEEP_ALIVE_SEC = 15.0
SSE_QUEUE_TTL_SEC = 600
MAX_QUEUE_SIZE = 256


def _sse_frame(event_type: str, data: dict | str) -> str:
    content = json.dumps(data, ensure_ascii=False) if isinstance(data, dict) else data
    return f"event: {event_type}\ndata: {content}\n\n"


def _sse_comment(comment: str) -> str:
    return f": {comment}\n\n"


class GatewaySseEndpoint:
    def __init__(self):
        self._queues: dict[str, asyncio.Queue[dict]] = {}
        self._tasks: dict[str, asyncio.Task] = {}
        self._last_active: dict[str, float] = {}

    def create(self, rid: str, q: asyncio.Queue[dict], task: asyncio.Task) -> None:
        self._queues[rid] = q
        self._tasks[rid] = task
        self._last_active[rid] = time.monotonic()

    def get(self, rid: str) -> asyncio.Queue[dict] | None:
        return self._queues.get(rid)

    def touch(self, rid: str) -> None:
        self._last_active[rid] = time.monotonic()

    def remove(self, rid: str) -> None:
        self._queues.pop(rid, None)
        task = self._tasks.pop(rid, None)
        if task and not task.done():
            task.cancel()
        self._last_active.pop(rid, None)

    def cleanup_stale(self) -> int:
        now = time.monotonic()
        stale = [rid for rid, ts in self._last_active.items() if now - ts > SSE_QUEUE_TTL_SEC]
        for rid in stale:
            self.remove(rid)
        return len(stale)


gateway_sse_endpoint = GatewaySseEndpoint()


async def _run_chat_feed(
    query: str,
    agent_config_id: int,
    thread_id: str | None,
    image_content: str | None,
    current_user,
    db,
    q: asyncio.Queue[dict],
    rid: str,
) -> None:
    from yuxi.services.chat_service import stream_agent_chat

    meta = {
        "source": "gateway_sse",
        "channel_type": "sse",
        "account_id": "default",
        "request_id": rid,
    }

    try:
        async for chunk in stream_agent_chat(
            query=query,
            agent_config_id=agent_config_id,
            thread_id=thread_id,
            meta=meta,
            image_content=image_content,
            current_user=current_user,
            db=db,
        ):
            try:
                data = json.loads(chunk.decode("utf-8"))
            except (json.JSONDecodeError, UnicodeDecodeError):
                continue

            status = data.get("status", "")
            content = data.get("response", "")

            if status == "error":
                await q.put({"type": SseEventType.ERROR, "data": {"message": data.get("error_message", content)}})
                await q.put({"type": SseEventType.DONE, "data": {}})
                return

            if status == "init":
                continue

            if status in ("streaming", "reasoning"):
                event_type = SseEventType.REASONING if data.get("reasoning") else SseEventType.DELTA
                if content:
                    await q.put({"type": event_type, "data": {"content": content}})

            if status == "thinking":
                await q.put({"type": SseEventType.PHASE, "data": {"content": content or "thinking"}})

            if status == "finished":
                final_content = data.get("final_response", content)
                await q.put(
                    {
                        "type": SseEventType.MESSAGE_END,
                        "data": {"content": final_content, "thread_id": data.get("thread_id", thread_id)},
                    }
                )

        await q.put({"type": SseEventType.DONE, "data": {}})
    except asyncio.CancelledError:
        pass
    except Exception:
        logger.exception("SSE chat stream failed for %s", rid)
        await q.put({"type": SseEventType.ERROR, "data": {"message": "SSE stream error"}})
        await q.put({"type": SseEventType.DONE, "data": {}})


async def _resolve_sse_user(db, authorization: str | None):
    from sqlalchemy import select

    from server.utils.auth_utils import AuthUtils
    from yuxi.storage.postgres.models_business import User

    if authorization and authorization.startswith("Bearer "):
        token = authorization[7:]
        try:
            payload = AuthUtils.verify_access_token(token)
            user_id = payload.get("sub")
            if user_id:
                result = await db.execute(select(User).where(User.id == int(user_id)))
                user = result.scalar_one_or_none()
                if user:
                    return user
        except Exception:
            pass

    result = await db.execute(
        select(User).where(User.is_deleted == 0).order_by(User.id).limit(1)
    )
    user = result.scalar_one_or_none()
    if user:
        logger.warning("SSE endpoint using fallback user id=%s (no valid auth provided)", user.id)
    return user


router = APIRouter(prefix="/api/sse", tags=["sse"])


@router.post("/chat")
async def sse_chat(
    query: str = Query(...),
    agent_config_id: int = Query(...),
    thread_id: str | None = Query(None),
    image_content: str | None = Query(None),
    authorization: str | None = Header(None),
):
    rid = str(uuid.uuid4())

    async with pg_manager.get_async_session_context() as db:
        current_user = await _resolve_sse_user(db, authorization)

        q: asyncio.Queue[dict] = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)

        task = asyncio.create_task(
            _run_chat_feed(
                query=query,
                agent_config_id=agent_config_id,
                thread_id=thread_id,
                image_content=image_content,
                current_user=current_user,
                db=db,
                q=q,
                rid=rid,
            ),
            name=f"sse_chat:{rid}",
        )
        gateway_sse_endpoint.create(rid, q, task)

    return {"request_id": rid, "stream_url": f"/api/sse/stream?request_id={rid}"}


@router.get("/stream")
async def sse_stream(request: Request, request_id: str = Query(...)):
    q = gateway_sse_endpoint.get(request_id)
    if q is None:
        return StreamingResponse(
            iter([_sse_frame(SseEventType.ERROR, {"message": "request_id invalid or expired"})]),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
            status_code=404,
        )

    gateway_sse_endpoint.touch(request_id)

    async def generate():
        try:
            while True:
                if await request.is_disconnected():
                    logger.info("SSE client disconnected: %s", request_id)
                    break
                try:
                    event = await asyncio.wait_for(q.get(), timeout=SSE_KEEP_ALIVE_SEC)
                    yield _sse_frame(str(event["type"]), event["data"])
                    if event["type"] == SseEventType.DONE:
                        break
                except TimeoutError:
                    yield _sse_comment("keepalive")
        finally:
            gateway_sse_endpoint.remove(request_id)

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


# ── OPT-7: 实时日志 SSE ──────────────────────────────────


@router.get("/logs")
async def sse_logs(
    request: Request,
    levels: str | None = Query(None, description="comma-separated: DEBUG,INFO,WARNING,ERROR"),
):
    logger.info("SSE log stream requested (not yet implemented), levels=%s", levels)

    async def generate():
        yield _sse_frame(SseEventType.ERROR, {"message": "SSE log streaming is not yet implemented"})
        yield _sse_frame(SseEventType.DONE, {})

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )