Agent/agent/trace/logs_websocket.py

"""
Logs WebSocket - 实时推送后端日志到前端
"""

import asyncio
import logging
from typing import Set
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from datetime import datetime


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


# 存储所有连接的WebSocket客户端
_clients: Set[WebSocket] = set()


class WebSocketLogHandler(logging.Handler):
    """自定义日志处理器，将日志推送到WebSocket客户端"""

    def emit(self, record: logging.LogRecord):
        """发送日志记录到所有连接的客户端"""
        try:
            log_entry = self.format(record)
            # 构造日志消息
            message = {
                "timestamp": datetime.now().isoformat(),
                "level": record.levelname,
                "name": record.name,
                "message": log_entry,
            }
            # 异步发送到所有客户端
            asyncio.create_task(broadcast_log(message))
        except Exception:
            self.handleError(record)


async def broadcast_log(message: dict):
    """广播日志消息到所有连接的客户端"""
    disconnected = set()
    for client in _clients:
        try:
            await client.send_json(message)
        except Exception:
            disconnected.add(client)

    # 移除断开连接的客户端
    for client in disconnected:
        _clients.discard(client)


@router.websocket("/watch")
async def logs_websocket(websocket: WebSocket):
    """
    日志WebSocket端点

    客户端连接后，实时接收后端日志
    """
    await websocket.accept()
    _clients.add(websocket)

    try:
        # 发送欢迎消息
        await websocket.send_json({
            "timestamp": datetime.now().isoformat(),
            "level": "INFO",
            "name": "logs_websocket",
            "message": "Connected to logs stream",
        })

        # 保持连接，等待客户端断开
        while True:
            # 接收客户端消息（用于保持连接）
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        _clients.discard(websocket)


def setup_websocket_logging(level=logging.INFO):
    """
    设置WebSocket日志处理器

    将根日志器的日志推送到WebSocket客户端
    """
    handler = WebSocketLogHandler()
    handler.setLevel(level)

    # 设置日志格式
    formatter = logging.Formatter(
        "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S"
    )
    handler.setFormatter(formatter)

    # 添加到根日志器
    root_logger = logging.getLogger()
    root_logger.addHandler(handler)

    return handler