AstrBot/astrbot/dashboard/routes/log.py

import asyncio
import json
import time
from collections.abc import AsyncGenerator

from quart import Response as QuartResponse
from quart import make_response, request

from astrbot.core import LogBroker, logger

from .route import Response, Route, RouteContext


def _format_log_sse(log: dict, ts: float) -> str:
    """辅助函数：格式化 SSE 消息"""
    payload = {"type": "log", **log}
    return f"id: {ts}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n"


def _coerce_log_timestamp(value: object) -> float | None:
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        try:
            return float(value)
        except ValueError:
            return None
    return None


class LogRoute(Route):
    def __init__(self, context: RouteContext, log_broker: LogBroker) -> None:
        super().__init__(context)
        self.log_broker = log_broker
        self.app.add_url_rule("/api/live-log", view_func=self.log, methods=["GET"])
        self.app.add_url_rule(
            "/api/log-history",
            view_func=self.log_history,
            methods=["GET"],
        )
        self.app.add_url_rule(
            "/api/trace/settings",
            view_func=self.get_trace_settings,
            methods=["GET"],
        )
        self.app.add_url_rule(
            "/api/trace/settings",
            view_func=self.update_trace_settings,
            methods=["POST"],
        )

    async def _replay_cached_logs(
        self,
        last_event_id: str,
    ) -> AsyncGenerator[str, None]:
        """辅助生成器：重放缓存的日志"""
        try:
            last_ts = float(last_event_id)
            cached_logs = list(self.log_broker.log_cache)
            for log_item in cached_logs:
                log_ts = _coerce_log_timestamp(log_item.get("time"))
                if log_ts is None:
                    continue
                if log_ts > last_ts:
                    yield _format_log_sse(log_item, log_ts)
        except ValueError:
            pass
        except Exception as e:
            logger.error(f"Log SSE 补发历史错误: {e}")

    async def log(self) -> QuartResponse:
        last_event_id = request.headers.get("Last-Event-ID")

        async def stream():
            queue = None
            try:
                if last_event_id:
                    async for event in self._replay_cached_logs(last_event_id):
                        yield event
                queue = self.log_broker.register()
                while True:
                    message = await queue.get()
                    current_ts = _coerce_log_timestamp(message.get("time"))
                    if current_ts is None:
                        current_ts = time.time()
                    yield _format_log_sse(message, current_ts)
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error(f"Log SSE 连接错误: {e}")
            finally:
                if queue:
                    self.log_broker.unregister(queue)

        response = await make_response(
            stream(),
            {
                "Content-Type": "text/event-stream",
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "Transfer-Encoding": "chunked",
            },
        )
        response.timeout = None
        return response

    async def log_history(self):
        """获取日志历史"""
        try:
            logs = list(self.log_broker.log_cache)
            return Response().ok(data={"logs": logs}).to_json()
        except Exception as e:
            logger.error(f"获取日志历史失败: {e}")
            return Response().error(f"获取日志历史失败: {e}").to_json()

    async def get_trace_settings(self):
        """获取 Trace 设置"""
        try:
            trace_enable = self.config.get("trace_enable", True)
            return Response().ok(data={"trace_enable": trace_enable}).to_json()
        except Exception as e:
            logger.error(f"获取 Trace 设置失败: {e}")
            return Response().error(f"获取 Trace 设置失败: {e}").to_json()

    async def update_trace_settings(self):
        """更新 Trace 设置"""
        try:
            data = await request.json
            if data is None:
                return Response().error("请求数据为空").to_json()
            trace_enable = data.get("trace_enable")
            if trace_enable is not None:
                self.config["trace_enable"] = bool(trace_enable)
                self.config.save_config()
            return Response().ok(message="Trace 设置已更新").to_json()
        except Exception as e:
            logger.error(f"更新 Trace 设置失败: {e}")
            return Response().error(f"更新 Trace 设置失败: {e}").to_json()