KiloStar/kilostar/api/__init__.py

#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import os
from typing import Dict

from fastapi import FastAPI, WebSocket, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

_STANDALONE = os.environ.get("KILOSTAR_MODE", "distributed") == "standalone"
if not _STANDALONE:
    from ray import serve

from .agent import agent_router
from .auth import auth_router
from .system import system_router
from .platform.frontend import client_router
from .platform.onebot import onebot_router
from .provider import provider_router
from .resource import resource_router
from .workflow import workflow_router
from .chat import chat_router
from kilostar.utils.error import (
    KiloStarError,
    BusinessError,
    InfraError,
)
from kilostar.utils.logger import get_logger
from kilostar.utils.request_context import (
    bind_request_id,
    new_request_id,
    reset_request_id,
)
from kilostar.utils.i18n import t

_api_logger = get_logger("api")


def _get_locale(request: Request) -> str | None:
    """从请求头解析首选语言，供异常 handler 使用。"""
    return request.headers.get("accept-language") or None

app = FastAPI()

_cors_origins_env = os.environ.get("KILOSTAR_CORS_ORIGINS", "")
_is_dev = os.environ.get("KILOSTAR_ENV", "production").lower() in ("dev", "development")
if not _cors_origins_env and _is_dev:
    _cors_origins_env = "*"
elif not _cors_origins_env:
    _cors_origins_env = "http://localhost:8000"

_cors_origins = [o.strip() for o in _cors_origins_env.split(",") if o.strip()]
_allow_credentials = "*" not in _cors_origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_credentials=_allow_credentials,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def request_id_middleware(request: Request, call_next):
    """请求级 ``request_id`` 注入。

    入口策略：``X-Request-Id`` 头存在则继承（便于网关/前端串联调用链），
    否则生成新的 UUID。退出时把它写到响应头，方便客户端日志对账。
    contextvars 让同一请求生命周期内所有协程的日志都自动带上这个 ID。
    """
    incoming = request.headers.get("X-Request-Id", "").strip()
    request_id = incoming or new_request_id()
    token = bind_request_id(request_id)
    try:
        response = await call_next(request)
    finally:
        reset_request_id(token)
    response.headers["X-Request-Id"] = request_id
    return response

app.include_router(system_router)  # 健康探针 + 系统信息
app.include_router(client_router)  # 客户端路径
app.include_router(onebot_router)  # OneBot v11 路径
app.include_router(auth_router)  # 用户路径
app.include_router(provider_router)  # 供应商路径
app.include_router(resource_router)  # 资源路径
app.include_router(agent_router)  # agent路径
app.include_router(workflow_router)  # workflow路径
app.include_router(chat_router)  # chat路径


@app.exception_handler(BusinessError)
async def business_error_handler(request: Request, exc: BusinessError):
    """业务可预期错误：按 ``http_status`` 返回 4xx，附 ``code`` + 异常消息。"""
    return JSONResponse(
        status_code=exc.http_status,
        content={"code": exc.code, "message": str(exc) or exc.code},
    )


@app.exception_handler(InfraError)
async def infra_error_handler(request: Request, exc: InfraError):
    """系统失败错误：落日志后返回脱敏的 5xx。"""
    _api_logger.exception(
        f"InfraError on {request.method} {request.url.path}: {exc}"
    )
    loc = _get_locale(request)
    return JSONResponse(
        status_code=exc.http_status,
        content={"code": exc.code, "message": t("internal_error", accept_language=loc)},
    )


@app.exception_handler(Exception)
async def unhandled_exception_handler(request: Request, exc: Exception):
    """全局兜底：未预期的异常落日志后返回脱敏的 500，避免泄露 traceback。"""
    _api_logger.exception(
        f"Unhandled exception on {request.method} {request.url.path}: {exc}"
    )
    loc = _get_locale(request)
    return JSONResponse(
        status_code=500,
        content={"code": "internal_error", "message": t("internal_error", accept_language=loc)},
    )


base_dir = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
frontend_dir = os.path.join(base_dir, "frontend", "dist")

if os.path.exists(frontend_dir):
    app.mount(
        "/assets",
        StaticFiles(directory=os.path.join(frontend_dir, "assets")),
        name="assets",
    )

    @app.get("/favicon.svg", include_in_schema=False)
    async def serve_favicon():
        return FileResponse(os.path.join(frontend_dir, "favicon.svg"))

    @app.get("/icons.svg", include_in_schema=False)
    async def serve_icons():
        return FileResponse(os.path.join(frontend_dir, "icons.svg"))

    @app.get("/{full_path:path}", include_in_schema=False)
    async def serve_frontend(full_path: str):
        # 【重要安全修复】避免拦截不存在的 API 路由。如果是调用了不存在的 /api/ 接口，直接返回 404，不返回前端页面
        if full_path.startswith("api/"):
            return JSONResponse(
                status_code=404, content={"detail": "API endpoint not found"}
            )

        index_path = os.path.join(frontend_dir, "index.html")
        if os.path.exists(index_path):
            return FileResponse(index_path)
        return JSONResponse(
            status_code=404, content={"detail": t("frontend_not_found")}
        )
else:
    import logging

    logging.getLogger("kilostar").warning(
        f"Frontend dist folder not found at {frontend_dir}. Skipping frontend mount."
    )


if not _STANDALONE:
    @serve.deployment
    @serve.ingress(app)
    class KiloStarGateway:
        gateway: Dict[str, WebSocket]

        def __init__(self):
            self.gateway = {}