ForcePilot/backend/server/routers/ws_chat_router.py

from __future__ import annotations

import asyncio
import json
import time
import uuid

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from jose import JWTError

from yuxi.channels.manager import channel_manager
from yuxi.channels.models import (
    ChannelIdentity,
    ChannelMessage,
    ChannelType,
    ChatType,
    EventType,
    MessageType,
)
from yuxi.utils.logging_config import logger

from server.utils.auth_utils import AuthUtils

MAX_CONNECTIONS_PER_USER = 3
HEARTBEAT_INTERVAL = 30
IDLE_TIMEOUT = 60
MAX_MSG_PER_SEC = 5

ws_chat = APIRouter(tags=["ws_chat"])


class ConnectionRegistry:
    def __init__(self):
        self._connections: dict[str, list[WebSocket]] = {}

    def add(self, user_id: str, ws: WebSocket) -> bool:
        if user_id not in self._connections:
            self._connections[user_id] = []
        if len(self._connections[user_id]) >= MAX_CONNECTIONS_PER_USER:
            return False
        self._connections[user_id].append(ws)
        return True

    def remove(self, user_id: str, ws: WebSocket) -> None:
        if user_id in self._connections:
            try:
                self._connections[user_id].remove(ws)
            except ValueError:
                pass
            if not self._connections[user_id]:
                del self._connections[user_id]

    def count(self, user_id: str) -> int:
        return len(self._connections.get(user_id, []))


connections = ConnectionRegistry()


@ws_chat.websocket("/ws/chat/{client_id}")
async def websocket_chat(ws: WebSocket, client_id: str, token: str = Query(None)):
    user_id = None
    try:
        if not token:
            await ws.close(code=4001, reason="Missing token")
            return

        try:
            payload = AuthUtils.verify_access_token(token)
            user_id = str(payload.get("sub"))
            if not user_id:
                await ws.close(code=4001, reason="Invalid token")
                return
        except (JWTError, ValueError):
            await ws.close(code=4001, reason="Invalid token")
            return

        if not connections.add(user_id, ws):
            await ws.close(code=4001, reason="Too many connections")
            return

        await ws.accept()
        await ws.send_json({
            "type": "connected",
            "payload": {"client_id": client_id, "user_id": user_id},
            "timestamp": _now_iso(),
        })

        last_msg_time = time.monotonic()
        rate_bucket: list[float] = []

        async def heartbeat_sender():
            while True:
                await asyncio.sleep(HEARTBEAT_INTERVAL)
                try:
                    await ws.send_json({"type": "pong", "payload": {}, "timestamp": _now_iso()})
                except Exception:
                    break

        async def idle_monitor():
            nonlocal last_msg_time
            while True:
                await asyncio.sleep(IDLE_TIMEOUT)
                if time.monotonic() - last_msg_time > IDLE_TIMEOUT:
                    logger.info(f"WS idle timeout for client {client_id}")
                    try:
                        await ws.close(code=4002, reason="Idle timeout")
                    except Exception:
                        pass
                    break

        heartbeat_task = asyncio.create_task(heartbeat_sender())
        idle_task = asyncio.create_task(idle_monitor())

        try:
            while True:
                raw = await ws.receive_text()
                last_msg_time = time.monotonic()

                data = json.loads(raw)
                msg_type = data.get("type")

                if msg_type == "ping":
                    await ws.send_json({"type": "pong", "payload": {}, "timestamp": _now_iso()})
                    continue

                if msg_type != "message":
                    continue

                now = time.monotonic()
                rate_bucket = [t for t in rate_bucket if now - t < 1.0]
                if len(rate_bucket) >= MAX_MSG_PER_SEC:
                    await ws.send_json({
                        "type": "error",
                        "payload": {"code": 429, "message": "Rate limit exceeded"},
                        "timestamp": _now_iso(),
                    })
                    continue
                rate_bucket.append(now)

                payload = data.get("payload", {})
                content = payload.get("content", "")
                thread_id = payload.get("thread_id", str(uuid.uuid4()))
                agent_config_id = payload.get("agent_config_id", 1)

                identity = ChannelIdentity(
                    channel_id="webchat",
                    channel_type=ChannelType.WEBCHAT,
                    channel_user_id=user_id,
                    channel_chat_id=thread_id,
                    channel_message_id=f"ws_{uuid.uuid4().hex[:12]}",
                )
                message = ChannelMessage(
                    identity=identity,
                    event_type=EventType.MESSAGE_RECEIVED,
                    message_type=MessageType.TEXT,
                    chat_type=ChatType.DIRECT,
                    content=content,
                    metadata={
                        "agent_config_id": agent_config_id,
                        "thread_id": thread_id,
                        "user_id": user_id,
                    },
                )

                request_id = identity.channel_message_id
                await ws.send_json({
                    "type": "typing",
                    "payload": {"status": "agent", "request_id": request_id},
                    "timestamp": _now_iso(),
                })

                try:
                    await channel_manager.dispatch_inbound(message)
                except Exception as e:
                    logger.error(f"WS message processing error: {e}")
                    await ws.send_json({
                        "type": "error",
                        "payload": {
                            "request_id": request_id,
                            "code": 500,
                            "message": f"处理出错: {str(e)}",
                        },
                        "timestamp": _now_iso(),
                    })

                await ws.send_json({
                    "type": "done",
                    "payload": {"request_id": request_id},
                    "timestamp": _now_iso(),
                })

        except WebSocketDisconnect:
            logger.info(f"WS client {client_id} disconnected")
        except Exception as e:
            logger.error(f"WS error for client {client_id}: {e}")
        finally:
            heartbeat_task.cancel()
            idle_task.cancel()
            try:
                await heartbeat_task
            except asyncio.CancelledError:
                pass
            try:
                await idle_task
            except asyncio.CancelledError:
                pass

    except Exception as e:
        logger.error(f"WS handshake error: {e}")
    finally:
        if user_id:
            connections.remove(user_id, ws)


def _now_iso() -> str:
    from yuxi.utils.datetime_utils import format_utc_datetime, utc_now_naive

    return format_utc_datetime(utc_now_naive())