AstrBot/astrbot/_internal/geteway/server.py

"""
AstrBot Gateway - FastAPI server for the dashboard backend.

Provides REST API endpoints and WebSocket connections for the frontend dashboard.
The gateway acts as the communication bridge between the dashboard and the orchestrator.
"""

from __future__ import annotations

import json
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any

from astrbot import logger
from astrbot._internal.abc.base_astrbot_gateway import BaseAstrbotGateway
from astrbot._internal.abc.base_astrbot_orchestrator import BaseAstrbotOrchestrator
from astrbot._internal.geteway.ws_manager import WebSocketManager

if TYPE_CHECKING:
    from fastapi import FastAPI, WebSocket, WebSocketDisconnect
else:
    try:
        from fastapi import FastAPI, WebSocket, WebSocketDisconnect
    except ImportError:
        logger.warning("FastAPI not installed, gateway unavailable.")
        FastAPI = None
        WebSocket = None
        WebSocketDisconnect = None
from fastapi.middleware.cors import CORSMiddleware

log = logger


class AstrbotGateway(BaseAstrbotGateway):
    """
    FastAPI-based gateway server for AstrBot.

    Handles:
    - REST API endpoints for configuration and stats
    - WebSocket connections for real-time communication
    - CORS middleware for dashboard access
    """

    def __init__(self, orchestrator: BaseAstrbotOrchestrator) -> None:
        self.orchestrator = orchestrator
        self.ws_manager = WebSocketManager()
        self._app: FastAPI | None = None
        self._host = "0.0.0.0"
        self._port = 8765

    async def serve(self) -> None:
        """
        Start the gateway server.

        Creates and runs a FastAPI application with WebSocket support.
        """
        if FastAPI is None:
            raise RuntimeError("FastAPI is not installed")
        log.info(f"Starting AstrBot Gateway on {self._host}:{self._port}")

        @asynccontextmanager
        async def lifespan(app: FastAPI):
            log.info("Gateway server started.")
            yield
            await self.ws_manager.broadcast({"type": "server_shutdown"})
            log.info("Gateway server stopped.")

        self._app = FastAPI(
            title="AstrBot Gateway",
            description="Backend API for AstrBot dashboard",
            version="1.0.0",
            lifespan=lifespan,
        )
        self._app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        self._setup_routes()
        import uvicorn

        config = uvicorn.Config(
            self._app, host=self._host, port=self._port, log_level="info"
        )
        server = uvicorn.Server(config)
        await server.serve()

    def _setup_routes(self) -> None:
        """Set up API routes."""
        if self._app is None:
            return
        from fastapi import APIRouter

        @self._app.get("/health")
        async def health():
            return {"status": "ok"}

        @self._app.websocket("/ws")
        async def websocket_endpoint(ws: WebSocket):
            await self.ws_manager.connect(ws)
            try:
                while True:
                    data = await ws.receive_text()
                    try:
                        message = json.loads(data)
                        response = await self._handle_ws_message(message)
                        if response:
                            await ws.send_json(response)
                    except json.JSONDecodeError:
                        await ws.send_json({"error": "Invalid JSON"})
            except WebSocketDisconnect:
                self.ws_manager.disconnect(ws)

        stats_router = APIRouter(prefix="/api/stats", tags=["stats"])

        @stats_router.get("/overview")
        async def get_overview():
            return await self._get_stats_overview()

        self._app.include_router(stats_router)
        inspector_router = APIRouter(prefix="/api/inspector", tags=["inspector"])

        @inspector_router.get("/stars")
        async def list_stars():
            return await self._list_stars()

        @inspector_router.get("/stars/{star_name}")
        async def get_star(star_name: str):
            return await self._get_star_detail(star_name)

        self._app.include_router(inspector_router)
        memory_router = APIRouter(prefix="/api/memory", tags=["memory"])

        @memory_router.get("/")
        async def get_memory():
            return await self._get_memory_info()

        self._app.include_router(memory_router)

    async def _handle_ws_message(
        self, message: dict[str, Any]
    ) -> dict[str, Any] | None:
        """
        Handle an incoming WebSocket message.

        Args:
            message: Parsed JSON message from the client

        Returns:
            Response message to send back, or None for no response
        """
        msg_type = message.get("type")
        data = message.get("data", {})
        if msg_type == "ping":
            return {"type": "pong", "data": {}}
        if msg_type == "call_tool":
            return await self._handle_call_tool(data)
        if msg_type == "get_stars":
            return {"type": "stars_list", "data": await self._list_stars()}
        return {
            "type": "error",
            "data": {"message": f"Unknown message type: {msg_type}"},
        }

    async def _handle_call_tool(self, data: dict[str, Any]) -> dict[str, Any]:
        """Handle a tool call request via WebSocket."""
        star_name = data.get("star")
        tool_name = data.get("tool")
        arguments = data.get("arguments", {})
        if not star_name or not tool_name:
            return {
                "type": "tool_result",
                "data": {"error": "Missing star or tool name"},
            }
        try:
            result = await self.orchestrator.abp.call_star_tool(
                star_name, tool_name, arguments
            )
            return {"type": "tool_result", "data": {"result": result}}
        except Exception as e:
            return {"type": "tool_result", "data": {"error": str(e)}}

    async def _get_stats_overview(self) -> dict[str, Any]:
        """Get overview statistics."""
        return {
            "stars_count": len(self.orchestrator.abp._stars),
            "lsp_connected": self.orchestrator.lsp._connected,
            "mcp_sessions": getattr(self.orchestrator.mcp, "session", None) is not None,
            "acp_clients": len(getattr(self.orchestrator.acp, "_clients", [])),
        }

    async def _list_stars(self) -> list[dict[str, Any]]:
        """List all registered stars."""
        stars = []
        for name in self.orchestrator.abp._stars:
            stars.append({"name": name, "status": "active"})
        return stars

    async def _get_star_detail(self, star_name: str) -> dict[str, Any]:
        """Get details of a specific star."""
        star = self.orchestrator.abp._stars.get(star_name)
        if not star:
            return {"error": f"Star '{star_name}' not found"}
        return {"name": star_name, "status": "active"}

    async def _get_memory_info(self) -> dict[str, Any]:
        """Get memory usage information."""
        import gc

        gc.collect()
        return {"gc_objects": len(gc.get_objects()), "python_memory": "N/A"}

    def set_listen_address(self, host: str, port: int) -> None:
        """Set the listen address for the gateway server."""
        self._host = host
        self._port = port