mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
888 lines
33 KiB
Python
888 lines
33 KiB
Python
"""
|
||
主应用入口
|
||
采用模块化架构设计
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import time
|
||
from collections.abc import Awaitable, Callable
|
||
from contextlib import asynccontextmanager
|
||
from dataclasses import dataclass, field
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
import uvicorn
|
||
from fastapi import FastAPI, HTTPException, Request
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
|
||
from src.api.admin import python_admin_router
|
||
from src.api.announcements import router as announcement_router
|
||
|
||
# API路由
|
||
from src.api.auth import router as auth_router
|
||
from src.api.dashboard import router as dashboard_router
|
||
from src.api.monitoring import router as monitoring_router
|
||
from src.api.payment import router as payment_router
|
||
from src.api.user_me import router as me_router
|
||
from src.api.wallet import router as wallet_router
|
||
from src.clients.http_client import close_http_clients
|
||
|
||
# 核心模块
|
||
from src.config import config
|
||
from src.core.exceptions import ExceptionHandlers, ProxyException
|
||
from src.core.logger import logger
|
||
from src.core.modules import get_module_registry
|
||
from src.database import init_db
|
||
from src.middleware.plugin_middleware import PluginMiddleware
|
||
from src.plugins.manager import get_plugin_manager
|
||
|
||
if TYPE_CHECKING:
|
||
from redis.asyncio.client import Redis
|
||
|
||
from src.core.modules.base import ModuleDefinition
|
||
from src.plugins.manager import PluginManager
|
||
from src.services.model.fetch_scheduler import ModelFetchScheduler
|
||
from src.services.provider_keys.pool_quota_probe_scheduler import PoolQuotaProbeScheduler
|
||
from src.services.rate_limit.concurrency_manager import ConcurrencyManager
|
||
from src.services.rate_limit.user_rpm_limiter import UserRpmLimiter
|
||
from src.services.system.maintenance_scheduler import MaintenanceScheduler
|
||
from src.services.system.scheduler import TaskScheduler
|
||
from src.services.task.polling.task_poller import TaskPollerService
|
||
from src.services.usage.quota_scheduler import QuotaScheduler
|
||
from src.utils.task_coordinator import StartupTaskCoordinator
|
||
|
||
|
||
async def initialize_providers() -> None:
|
||
"""从数据库初始化提供商（仅用于日志记录，使用轻量查询）"""
|
||
from sqlalchemy import func
|
||
from sqlalchemy.orm import Session
|
||
|
||
from src.database.database import create_session
|
||
from src.models.database import Provider, ProviderEndpoint
|
||
|
||
try:
|
||
db: Session = create_session()
|
||
|
||
try:
|
||
# 使用聚合查询代替全量加载 ORM 对象，大幅减少内存占用
|
||
results = (
|
||
db.query(
|
||
Provider.name,
|
||
func.count(ProviderEndpoint.id).label("total"),
|
||
func.count(func.nullif(ProviderEndpoint.is_active, False)).label("active"),
|
||
)
|
||
.outerjoin(ProviderEndpoint, Provider.id == ProviderEndpoint.provider_id)
|
||
.filter(Provider.is_active.is_(True))
|
||
.group_by(Provider.id, Provider.name, Provider.provider_priority)
|
||
.order_by(Provider.provider_priority.asc())
|
||
.all()
|
||
)
|
||
|
||
if not results:
|
||
logger.warning("数据库中未找到活跃的提供商")
|
||
return
|
||
|
||
logger.info(f"从数据库加载了 {len(results)} 个活跃提供商")
|
||
for name, total, active in results:
|
||
logger.info(f"提供商: {name} (端点: {active}/{total})")
|
||
|
||
finally:
|
||
db.close()
|
||
|
||
except Exception:
|
||
logger.exception("从数据库初始化提供商失败")
|
||
|
||
|
||
@dataclass
|
||
class LifecycleState:
|
||
"""应用生命周期阶段共享的运行时状态。"""
|
||
|
||
redis_client: Redis | None = None
|
||
concurrency_manager: ConcurrencyManager | None = None
|
||
user_rpm_limiter: UserRpmLimiter | None = None
|
||
plugin_manager: PluginManager | None = None
|
||
available_modules: list[ModuleDefinition] = field(default_factory=list)
|
||
task_coordinator: StartupTaskCoordinator | None = None
|
||
quota_scheduler: QuotaScheduler | None = None
|
||
maintenance_scheduler: MaintenanceScheduler | None = None
|
||
model_fetch_scheduler: ModelFetchScheduler | None = None
|
||
pool_quota_probe_scheduler: PoolQuotaProbeScheduler | None = None
|
||
task_poller: TaskPollerService | None = None
|
||
task_scheduler: TaskScheduler | None = None
|
||
warmup_task: asyncio.Task[None] | None = None
|
||
|
||
|
||
async def _stop_service_on_lock_lost(
|
||
state: LifecycleState,
|
||
*,
|
||
lock_name: str,
|
||
service_name: str,
|
||
state_attr: str,
|
||
stop: Callable[[], Awaitable[Any]],
|
||
) -> None:
|
||
logger.warning("检测到 {} 的 leader 锁已丢失，停止本实例的 {}", lock_name, service_name)
|
||
try:
|
||
await stop()
|
||
except Exception:
|
||
logger.exception("丢失 {} leader 锁后停止 {} 失败", lock_name, service_name)
|
||
return
|
||
|
||
setattr(state, state_attr, None)
|
||
|
||
|
||
def _make_lock_lost_callback(
|
||
state: LifecycleState,
|
||
*,
|
||
lock_name: str,
|
||
service_name: str,
|
||
state_attr: str,
|
||
stop: Callable[[], Awaitable[Any]],
|
||
) -> Callable[[str], Awaitable[None]]:
|
||
async def _callback(_name: str) -> None:
|
||
await _stop_service_on_lock_lost(
|
||
state,
|
||
lock_name=lock_name,
|
||
service_name=service_name,
|
||
state_attr=state_attr,
|
||
stop=stop,
|
||
)
|
||
|
||
return _callback
|
||
|
||
|
||
def _configure_uvicorn_access_log() -> None:
|
||
"""禁用 uvicorn access 日志（在子进程中执行）。"""
|
||
import logging
|
||
|
||
logging.getLogger("uvicorn.access").setLevel(logging.CRITICAL)
|
||
logging.getLogger("uvicorn.access").disabled = True
|
||
|
||
|
||
def _log_startup_banner() -> None:
|
||
logger.info("=" * 60)
|
||
from src import __version__
|
||
|
||
logger.info(f"AI Proxy v{__version__} - GlobalModel Architecture")
|
||
logger.info("=" * 60)
|
||
|
||
|
||
def _validate_security_or_raise() -> None:
|
||
"""启动前安全配置校验。"""
|
||
security_errors = config.validate_security_config()
|
||
if security_errors:
|
||
for error in security_errors:
|
||
logger.error(f"[SECURITY] {error}")
|
||
if config.environment == "production":
|
||
raise RuntimeError(
|
||
"Security configuration errors detected. "
|
||
"Please fix the following issues before starting in production:\n"
|
||
+ "\n".join(f" - {e}" for e in security_errors)
|
||
)
|
||
|
||
|
||
async def _initialize_core_infrastructure(state: LifecycleState) -> None:
|
||
"""初始化数据库、缓存、并发与基础后台组件。"""
|
||
# 记录启动警告（密码、连接池、JWT 等）
|
||
config.log_startup_warnings()
|
||
|
||
# 初始化数据库
|
||
logger.info("初始化数据库...")
|
||
init_db()
|
||
|
||
# 从数据库初始化提供商
|
||
await initialize_providers()
|
||
|
||
# 全局HTTP客户端池按需初始化，避免启动阶段预分配连接池资源
|
||
logger.info("全局HTTP客户端池采用按需初始化")
|
||
|
||
# 初始化全局Redis客户端（可根据配置降级为内存模式）
|
||
logger.info("初始化全局Redis客户端...")
|
||
from src.clients.redis_client import get_redis_client
|
||
|
||
try:
|
||
state.redis_client = await get_redis_client(require_redis=config.require_redis)
|
||
if state.redis_client:
|
||
logger.info("[OK] Redis客户端初始化成功，缓存亲和性功能已启用")
|
||
else:
|
||
logger.warning(
|
||
"[WARN] Redis未启用或连接失败，将使用内存缓存亲和性（仅适用于单实例/开发环境）"
|
||
)
|
||
except RuntimeError as e:
|
||
if config.require_redis:
|
||
logger.exception("[ERROR] Redis连接失败，应用启动中止")
|
||
raise
|
||
logger.warning(f"Redis连接失败，但配置允许降级，将继续使用内存模式: {e}")
|
||
state.redis_client = None
|
||
|
||
# 初始化并发管理器（内部会使用Redis）
|
||
logger.info("初始化并发管理器...")
|
||
from src.services.rate_limit.concurrency_manager import get_concurrency_manager
|
||
|
||
state.concurrency_manager = await get_concurrency_manager()
|
||
|
||
logger.info("初始化用户/API Key RPM 限流器...")
|
||
from src.services.rate_limit.user_rpm_limiter import get_user_rpm_limiter
|
||
|
||
state.user_rpm_limiter = await get_user_rpm_limiter()
|
||
|
||
# 初始化批量提交器（提升数据库并发能力）
|
||
logger.info("初始化批量提交器...")
|
||
from src.core.batch_committer import init_batch_committer
|
||
|
||
await init_batch_committer()
|
||
logger.info("[OK] 批量提交器已启动，数据库写入性能优化已启用")
|
||
|
||
# 初始化 Codex 配额异步同步器（请求路径仅投递事件）
|
||
logger.info("初始化 Codex 配额异步同步器...")
|
||
from src.services.provider_keys.codex_quota_sync_dispatcher import (
|
||
init_codex_quota_sync_dispatcher,
|
||
)
|
||
|
||
await init_codex_quota_sync_dispatcher()
|
||
logger.info("[OK] Codex 配额异步同步器已启动")
|
||
|
||
# 初始化 Usage 队列消费者（可选）
|
||
if config.usage_queue_enabled and config.usage_queue_python_consumer_enabled:
|
||
logger.info("初始化 Usage 队列消费者...")
|
||
from src.services.usage.consumer_streams import start_usage_queue_consumer
|
||
|
||
await start_usage_queue_consumer()
|
||
elif config.usage_queue_enabled:
|
||
logger.info("Usage 队列消费者已切到 Rust gateway，Python 宿主跳过启动")
|
||
|
||
|
||
async def _initialize_plugins_and_modules(app: FastAPI, state: LifecycleState) -> None:
|
||
"""初始化插件系统、模块系统，并注册路由与钩子。"""
|
||
# 初始化插件系统
|
||
logger.info("初始化插件系统...")
|
||
state.plugin_manager = get_plugin_manager()
|
||
init_results = await state.plugin_manager.initialize_all()
|
||
successful = sum(1 for success in init_results.values() if success)
|
||
logger.info(f"插件初始化完成: {successful}/{len(init_results)} 个插件成功启动")
|
||
|
||
# 注册格式转换器
|
||
logger.info("注册格式转换器...")
|
||
from src.core.api_format.conversion.registry import register_default_normalizers
|
||
|
||
register_default_normalizers()
|
||
|
||
# 初始化功能模块系统
|
||
logger.info("初始化功能模块系统...")
|
||
from src.modules import ALL_MODULES
|
||
|
||
module_registry = get_module_registry()
|
||
|
||
# 注入配置后端，消除 core/modules→services 的运行时 lazy import
|
||
from src.services.system.config import SystemConfigService
|
||
|
||
module_registry.set_config_backend(SystemConfigService) # type: ignore[arg-type]
|
||
|
||
for module in ALL_MODULES:
|
||
module_registry.register(module)
|
||
|
||
# 注册模块钩子
|
||
from src.core.modules.hooks import get_hook_dispatcher
|
||
|
||
hook_dispatcher = get_hook_dispatcher()
|
||
for module in ALL_MODULES:
|
||
for hook_name, handler in module.hooks.items():
|
||
hook_dispatcher.register(hook_name, module.metadata.name, handler)
|
||
|
||
# 注册可用模块的路由
|
||
# 注意：模块的 router 自带 prefix，api_prefix 字段仅用于日志和文档
|
||
state.available_modules = module_registry.get_available_modules()
|
||
for module in state.available_modules:
|
||
if module.router_factory:
|
||
router = module.router_factory()
|
||
app.include_router(router)
|
||
prefix = module.metadata.api_prefix or "(default)"
|
||
logger.info(f"模块 [{module.metadata.name}] 路由已注册: {prefix}")
|
||
|
||
# 执行启动钩子
|
||
if module.on_startup:
|
||
await module.on_startup()
|
||
|
||
logger.info(f"功能模块初始化完成: {len(state.available_modules)}/{len(ALL_MODULES)} 个模块可用")
|
||
|
||
# 显式 bootstrap provider plugins（注册 envelope/enricher 等）
|
||
# 使 core/provider_oauth_utils 不需要在运行时 lazy import services 层
|
||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||
|
||
ensure_providers_bootstrapped()
|
||
|
||
# 显式触发 parsers 注册（使 core/stream_types 不需要 lazy import api 层）
|
||
from src.api.handlers.base.parsers import register_default_parsers
|
||
|
||
register_default_parsers()
|
||
|
||
|
||
def _warmup_lazy_request_dependencies(provider_types: list[str] | None = None) -> int:
|
||
"""预热懒加载链路，减少首个真实请求的导入与实例化抖动。
|
||
|
||
逐个 adapter 独立 try-except，确保单个失败不影响其余组件的预热。
|
||
"""
|
||
import importlib
|
||
|
||
warmed = 0
|
||
|
||
try:
|
||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||
|
||
ensure_providers_bootstrapped(provider_types=provider_types)
|
||
warmed += 1
|
||
except Exception as exc:
|
||
logger.warning("预热 provider bootstrap 失败: {}", exc)
|
||
|
||
try:
|
||
from src.services.health.monitor import get_health_monitor
|
||
|
||
get_health_monitor()
|
||
warmed += 1
|
||
except Exception as exc:
|
||
logger.warning("预热 health monitor 失败: {}", exc)
|
||
|
||
adapter_specs: list[tuple[str, str]] = [
|
||
("src.api.handlers.openai", "OpenAIChatAdapter"),
|
||
("src.api.handlers.openai_cli", "OpenAICliAdapter"),
|
||
("src.api.handlers.openai_cli", "OpenAICompactAdapter"),
|
||
("src.api.handlers.claude", "ClaudeChatAdapter"),
|
||
("src.api.handlers.claude", "ClaudeTokenCountAdapter"),
|
||
("src.api.handlers.claude_cli", "ClaudeCliAdapter"),
|
||
("src.api.handlers.gemini", "GeminiChatAdapter"),
|
||
("src.api.handlers.gemini_cli", "GeminiCliAdapter"),
|
||
]
|
||
for module_path, class_name in adapter_specs:
|
||
try:
|
||
mod = importlib.import_module(module_path)
|
||
cls = getattr(mod, class_name)
|
||
cls()
|
||
warmed += 1
|
||
except Exception as exc:
|
||
logger.warning("预热 {} 失败: {}", class_name, exc)
|
||
|
||
return warmed
|
||
|
||
|
||
async def _run_startup_warmup(app: FastAPI) -> None:
|
||
"""异步执行启动预热，结果写入 app.state 供 /readyz 读取。"""
|
||
app.state.startup_warmup_started_at = time.monotonic()
|
||
provider_types = config.startup_warmup_provider_types
|
||
logger.info(
|
||
"启动预热任务开始（provider_types={}）",
|
||
provider_types if provider_types else "auto",
|
||
)
|
||
|
||
try:
|
||
warmed_count = await asyncio.to_thread(
|
||
_warmup_lazy_request_dependencies,
|
||
provider_types,
|
||
)
|
||
app.state.startup_warmup_status = "ready"
|
||
elapsed_ms = int((time.monotonic() - app.state.startup_warmup_started_at) * 1000)
|
||
logger.info("启动预热完成（adapters={}, elapsed_ms={}）", warmed_count, elapsed_ms)
|
||
except Exception as exc:
|
||
app.state.startup_warmup_status = "failed"
|
||
app.state.startup_warmup_error = str(exc)
|
||
logger.exception("启动预热失败，/readyz 将返回 503")
|
||
finally:
|
||
app.state.startup_warmup_finished_at = time.monotonic()
|
||
app.state.startup_warmup_done.set()
|
||
|
||
|
||
def _schedule_startup_warmup(app: FastAPI, state: LifecycleState) -> None:
|
||
"""初始化预热状态并按配置启动后台预热任务。"""
|
||
app.state.startup_warmup_done = asyncio.Event()
|
||
app.state.startup_warmup_status = "pending"
|
||
app.state.startup_warmup_error = None
|
||
app.state.startup_warmup_started_at = None
|
||
app.state.startup_warmup_finished_at = None
|
||
|
||
if not config.startup_warmup_enabled:
|
||
app.state.startup_warmup_status = "disabled"
|
||
app.state.startup_warmup_done.set()
|
||
logger.info("启动预热已禁用（STARTUP_WARMUP_ENABLED=false）")
|
||
return
|
||
|
||
state.warmup_task = asyncio.create_task(_run_startup_warmup(app), name="startup_warmup")
|
||
|
||
|
||
async def _start_background_services(state: LifecycleState) -> None:
|
||
"""启动调度器与后台轮询服务。"""
|
||
# 启动月卡额度重置调度器（仅一个 worker 执行）
|
||
logger.info("启动月卡额度重置调度器...")
|
||
from src.services.model.fetch_scheduler import get_model_fetch_scheduler
|
||
from src.services.provider_keys.pool_quota_probe_scheduler import (
|
||
get_pool_quota_probe_scheduler,
|
||
)
|
||
from src.services.system.maintenance_scheduler import (
|
||
get_maintenance_scheduler,
|
||
should_start_python_maintenance_scheduler,
|
||
)
|
||
from src.services.task.polling.task_poller import get_task_poller
|
||
from src.services.usage.quota_scheduler import get_quota_scheduler
|
||
from src.utils.task_coordinator import StartupTaskCoordinator
|
||
|
||
state.task_coordinator = StartupTaskCoordinator(state.redis_client)
|
||
|
||
# 启动额度调度器
|
||
if not config.quota_scheduler_python_enabled:
|
||
logger.info("月卡额度重置调度器已切到 Rust gateway，Python 宿主跳过启动")
|
||
state.quota_scheduler = None
|
||
else:
|
||
quota_scheduler_active = await state.task_coordinator.acquire("quota_scheduler")
|
||
if quota_scheduler_active:
|
||
state.quota_scheduler = get_quota_scheduler()
|
||
await state.quota_scheduler.start()
|
||
state.task_coordinator.register_lock_lost_callback(
|
||
"quota_scheduler",
|
||
_make_lock_lost_callback(
|
||
state,
|
||
lock_name="quota_scheduler",
|
||
service_name="月卡额度重置调度器",
|
||
state_attr="quota_scheduler",
|
||
stop=state.quota_scheduler.stop,
|
||
),
|
||
)
|
||
else:
|
||
logger.info("检测到其他 worker 已运行额度调度器，本实例跳过")
|
||
state.quota_scheduler = None
|
||
|
||
# 启动维护调度器
|
||
if not should_start_python_maintenance_scheduler():
|
||
logger.info("维护调度器已无剩余 Python owner，Python 宿主跳过启动")
|
||
state.maintenance_scheduler = None
|
||
else:
|
||
maintenance_scheduler_active = await state.task_coordinator.acquire("maintenance_scheduler")
|
||
if maintenance_scheduler_active:
|
||
state.maintenance_scheduler = get_maintenance_scheduler()
|
||
logger.info("启动系统维护调度器...")
|
||
await state.maintenance_scheduler.start()
|
||
state.task_coordinator.register_lock_lost_callback(
|
||
"maintenance_scheduler",
|
||
_make_lock_lost_callback(
|
||
state,
|
||
lock_name="maintenance_scheduler",
|
||
service_name="系统维护调度器",
|
||
state_attr="maintenance_scheduler",
|
||
stop=state.maintenance_scheduler.stop,
|
||
),
|
||
)
|
||
else:
|
||
logger.info("检测到其他 worker 已运行维护调度器，本实例跳过")
|
||
state.maintenance_scheduler = None
|
||
|
||
# 启动模型自动获取调度器
|
||
if not config.model_fetch_scheduler_python_enabled:
|
||
logger.info("模型自动获取调度器已切到 Rust gateway，Python 宿主跳过启动")
|
||
state.model_fetch_scheduler = None
|
||
else:
|
||
model_fetch_scheduler_active = await state.task_coordinator.acquire("model_fetch_scheduler")
|
||
if model_fetch_scheduler_active:
|
||
state.model_fetch_scheduler = get_model_fetch_scheduler()
|
||
logger.info("启动模型自动获取调度器...")
|
||
await state.model_fetch_scheduler.start()
|
||
state.task_coordinator.register_lock_lost_callback(
|
||
"model_fetch_scheduler",
|
||
_make_lock_lost_callback(
|
||
state,
|
||
lock_name="model_fetch_scheduler",
|
||
service_name="模型自动获取调度器",
|
||
state_attr="model_fetch_scheduler",
|
||
stop=state.model_fetch_scheduler.stop,
|
||
),
|
||
)
|
||
else:
|
||
logger.info("检测到其他 worker 已运行模型获取调度器，本实例跳过")
|
||
state.model_fetch_scheduler = None
|
||
|
||
# 启动号池额度主动探测调度器
|
||
pool_quota_probe_scheduler_active = await state.task_coordinator.acquire(
|
||
"pool_quota_probe_scheduler"
|
||
)
|
||
if pool_quota_probe_scheduler_active:
|
||
state.pool_quota_probe_scheduler = get_pool_quota_probe_scheduler()
|
||
logger.info("启动号池额度主动探测调度器...")
|
||
await state.pool_quota_probe_scheduler.start()
|
||
state.task_coordinator.register_lock_lost_callback(
|
||
"pool_quota_probe_scheduler",
|
||
_make_lock_lost_callback(
|
||
state,
|
||
lock_name="pool_quota_probe_scheduler",
|
||
service_name="号池额度主动探测调度器",
|
||
state_attr="pool_quota_probe_scheduler",
|
||
stop=state.pool_quota_probe_scheduler.stop,
|
||
),
|
||
)
|
||
else:
|
||
logger.info("检测到其他 worker 已运行号池额度主动探测调度器，本实例跳过")
|
||
state.pool_quota_probe_scheduler = None
|
||
|
||
# 启动异步任务轮询服务（当前仅视频）
|
||
if not config.video_task_python_poller_enabled:
|
||
logger.info("TaskPoller（video）已切到 Rust gateway，Python 宿主跳过启动")
|
||
state.task_poller = None
|
||
else:
|
||
task_poller_active = await state.task_coordinator.acquire("task_poller:video")
|
||
if task_poller_active:
|
||
state.task_poller = get_task_poller()
|
||
logger.info("启动 TaskPoller（video）...")
|
||
await state.task_poller.start()
|
||
state.task_coordinator.register_lock_lost_callback(
|
||
"task_poller:video",
|
||
_make_lock_lost_callback(
|
||
state,
|
||
lock_name="task_poller:video",
|
||
service_name="TaskPoller（video）",
|
||
state_attr="task_poller",
|
||
stop=state.task_poller.stop,
|
||
),
|
||
)
|
||
else:
|
||
logger.info("检测到其他 worker 已运行 TaskPoller（video），本实例跳过")
|
||
state.task_poller = None
|
||
|
||
# 启动统一的定时任务调度器
|
||
from src.services.system.scheduler import get_scheduler
|
||
|
||
state.task_scheduler = get_scheduler()
|
||
state.task_scheduler.start()
|
||
|
||
|
||
async def _run_python_host_startup(app: FastAPI) -> LifecycleState:
|
||
"""执行 Python 宿主运行时启动流程。"""
|
||
_configure_uvicorn_access_log()
|
||
_log_startup_banner()
|
||
_validate_security_or_raise()
|
||
|
||
state = LifecycleState()
|
||
await _initialize_core_infrastructure(state)
|
||
await _initialize_plugins_and_modules(app, state)
|
||
_schedule_startup_warmup(app, state)
|
||
|
||
logger.info(f"服务启动成功: http://{config.host}:{config.port}")
|
||
logger.info("=" * 60)
|
||
|
||
await _start_background_services(state)
|
||
return state
|
||
|
||
|
||
async def _run_python_host_shutdown(state: LifecycleState) -> None:
|
||
"""执行 Python 宿主运行时关闭流程。"""
|
||
logger.info("正在关闭服务...")
|
||
|
||
# 停止启动预热任务
|
||
if state.warmup_task and not state.warmup_task.done():
|
||
logger.info("停止启动预热任务...")
|
||
state.warmup_task.cancel()
|
||
try:
|
||
await asyncio.wait_for(state.warmup_task, timeout=5.0)
|
||
except (asyncio.CancelledError, asyncio.TimeoutError):
|
||
pass
|
||
|
||
# 停止 Codex 配额异步同步器（停止前会 flush 待同步事件）
|
||
logger.info("停止 Codex 配额异步同步器...")
|
||
from src.services.provider_keys.codex_quota_sync_dispatcher import (
|
||
shutdown_codex_quota_sync_dispatcher,
|
||
)
|
||
|
||
await shutdown_codex_quota_sync_dispatcher()
|
||
logger.info("[OK] Codex 配额异步同步器已停止")
|
||
|
||
# 停止批量提交器（确保所有待提交的数据都被保存）
|
||
logger.info("停止批量提交器...")
|
||
from src.core.batch_committer import shutdown_batch_committer
|
||
|
||
await shutdown_batch_committer()
|
||
logger.info("[OK] 批量提交器已停止，所有待提交数据已保存")
|
||
|
||
# 停止 Usage 队列消费者
|
||
if config.usage_queue_enabled and config.usage_queue_python_consumer_enabled:
|
||
logger.info("停止 Usage 队列消费者...")
|
||
from src.services.usage.consumer_streams import stop_usage_queue_consumer
|
||
|
||
await stop_usage_queue_consumer()
|
||
|
||
# 停止维护调度器
|
||
if state.maintenance_scheduler:
|
||
logger.info("停止系统维护调度器...")
|
||
await state.maintenance_scheduler.stop()
|
||
if state.task_coordinator:
|
||
await state.task_coordinator.release("maintenance_scheduler")
|
||
|
||
# 停止月卡额度重置调度器，并释放分布式锁
|
||
if state.quota_scheduler:
|
||
logger.info("停止月卡额度重置调度器...")
|
||
await state.quota_scheduler.stop()
|
||
if state.task_coordinator:
|
||
await state.task_coordinator.release("quota_scheduler")
|
||
|
||
# 停止模型自动获取调度器
|
||
if state.model_fetch_scheduler:
|
||
logger.info("停止模型自动获取调度器...")
|
||
await state.model_fetch_scheduler.stop()
|
||
if state.task_coordinator:
|
||
await state.task_coordinator.release("model_fetch_scheduler")
|
||
|
||
if state.pool_quota_probe_scheduler:
|
||
logger.info("停止号池额度主动探测调度器...")
|
||
await state.pool_quota_probe_scheduler.stop()
|
||
if state.task_coordinator:
|
||
await state.task_coordinator.release("pool_quota_probe_scheduler")
|
||
|
||
if state.task_poller:
|
||
logger.info("停止 TaskPoller（video）...")
|
||
await state.task_poller.stop()
|
||
if state.task_coordinator:
|
||
await state.task_coordinator.release("task_poller:video")
|
||
|
||
# 停止统一的定时任务调度器
|
||
logger.info("停止定时任务调度器...")
|
||
if state.task_scheduler:
|
||
state.task_scheduler.stop()
|
||
|
||
# 关闭插件系统
|
||
logger.info("关闭插件系统...")
|
||
if state.plugin_manager:
|
||
await state.plugin_manager.shutdown_all()
|
||
|
||
# 关闭功能模块
|
||
logger.info("关闭功能模块...")
|
||
for module in state.available_modules:
|
||
if module.on_shutdown:
|
||
await module.on_shutdown()
|
||
|
||
# 关闭并发管理器
|
||
logger.info("关闭并发管理器...")
|
||
if state.concurrency_manager:
|
||
await state.concurrency_manager.close()
|
||
|
||
logger.info("关闭用户/API Key RPM 限流器...")
|
||
if state.user_rpm_limiter:
|
||
await state.user_rpm_limiter.close()
|
||
|
||
# 关闭全局Redis客户端
|
||
logger.info("关闭全局Redis客户端...")
|
||
from src.clients.redis_client import close_redis_client
|
||
|
||
await close_redis_client()
|
||
|
||
# 关闭HTTP客户端池
|
||
logger.info("关闭HTTP客户端池...")
|
||
await close_http_clients()
|
||
|
||
logger.info("服务已关闭")
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI) -> Any:
|
||
"""FastAPI 生命周期只管理 Python-only 宿主状态。"""
|
||
state = await _run_python_host_startup(app)
|
||
try:
|
||
yield # 应用运行期间
|
||
finally:
|
||
await _run_python_host_shutdown(state)
|
||
|
||
|
||
from src import __version__ as app_version
|
||
|
||
# OpenAPI Tags 元数据定义
|
||
openapi_tags = [
|
||
{
|
||
"name": "Authentication",
|
||
"description": "用户认证相关接口，包括登录、注册、令牌刷新等",
|
||
},
|
||
{
|
||
"name": "User Profile",
|
||
"description": "用户个人信息管理，包括 API 密钥、使用统计、偏好设置等",
|
||
},
|
||
{
|
||
"name": "Management Tokens",
|
||
"description": "管理令牌，用于 CLI 工具等外部应用的认证",
|
||
},
|
||
{
|
||
"name": "Dashboard",
|
||
"description": "仪表盘统计数据，包括请求量、Token 用量、成本等概览信息",
|
||
},
|
||
{
|
||
"name": "Announcements",
|
||
"description": "系统公告管理",
|
||
},
|
||
{
|
||
"name": "Monitoring",
|
||
"description": "用户监控与审计日志查询",
|
||
},
|
||
{
|
||
"name": "Admin - Users",
|
||
"description": "用户管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - Providers",
|
||
"description": "提供商管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - Endpoints",
|
||
"description": "端点管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - Models",
|
||
"description": "模型管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - API Keys",
|
||
"description": "API 密钥管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - Usage",
|
||
"description": "使用统计管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - Monitoring",
|
||
"description": "系统监控（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - Security",
|
||
"description": "安全配置管理（管理员）",
|
||
},
|
||
{
|
||
"name": "Admin - System",
|
||
"description": "系统配置管理（管理员）",
|
||
},
|
||
]
|
||
|
||
|
||
def _register_exception_handlers(app: FastAPI) -> None:
|
||
"""注册 Python 宿主全局异常处理器。"""
|
||
# 注意：异常处理器的注册顺序很重要，必须先注册更通用的异常类型，再注册具体的
|
||
# ProxyException 处理器的启用由配置控制：
|
||
# - propagate_provider_exceptions=True (默认): 不注册，让异常传播到路由层以记录 provider_request_headers
|
||
# - propagate_provider_exceptions=False: 注册全局处理器统一处理
|
||
if not config.propagate_provider_exceptions:
|
||
app.add_exception_handler(ProxyException, ExceptionHandlers.handle_proxy_exception) # type: ignore[arg-type]
|
||
app.add_exception_handler(Exception, ExceptionHandlers.handle_generic_exception) # type: ignore[arg-type]
|
||
app.add_exception_handler(HTTPException, ExceptionHandlers.handle_http_exception) # type: ignore[arg-type]
|
||
|
||
|
||
def _register_frontdoor_replaceable_middlewares(app: FastAPI) -> None:
|
||
"""注册未来可由 Rust frontdoor 接管的宿主壳层中间件。"""
|
||
# CORS配置 - 使用环境变量配置允许的域名
|
||
# 生产环境必须通过 CORS_ORIGINS 环境变量显式指定允许的域名
|
||
# 开发环境默认允许本地前端访问
|
||
if config.cors_origins:
|
||
# CORS_ORIGINS=* 时自动禁用 credentials（浏览器规范要求）
|
||
allow_credentials = config.cors_allow_credentials and "*" not in config.cors_origins
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=config.cors_origins, # 使用配置的白名单
|
||
allow_credentials=allow_credentials,
|
||
allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
|
||
allow_headers=["*"],
|
||
expose_headers=["*"],
|
||
)
|
||
logger.info(
|
||
f"CORS已启用,允许的源: {config.cors_origins}, credentials: {allow_credentials}"
|
||
)
|
||
else:
|
||
# 没有配置CORS源,不允许跨域
|
||
logger.warning(
|
||
f"CORS未配置,不允许跨域请求。如需启用CORS,请设置 CORS_ORIGINS 环境变量(当前环境: {config.environment})"
|
||
)
|
||
|
||
|
||
def _register_python_host_middlewares(app: FastAPI) -> None:
|
||
"""注册仍依附于 Python 宿主的全局中间件。"""
|
||
# 添加插件中间件（包含认证、审计、速率限制等功能）
|
||
app.add_middleware(PluginMiddleware)
|
||
|
||
|
||
def _register_python_only_host_routes(app: FastAPI) -> None:
|
||
"""注册在当前阶段仍需留在 Python 的路由面。"""
|
||
app.include_router(auth_router) # 认证相关
|
||
app.include_router(python_admin_router) # 管理员端点
|
||
app.include_router(me_router) # 用户个人端点
|
||
app.include_router(wallet_router) # 钱包端点
|
||
app.include_router(payment_router) # 支付回调端点
|
||
app.include_router(announcement_router) # 公告系统
|
||
app.include_router(dashboard_router) # 仪表盘端点
|
||
app.include_router(monitoring_router) # 监控端点
|
||
|
||
|
||
app = FastAPI(
|
||
title="Aether AI Gateway",
|
||
version=app_version,
|
||
lifespan=lifespan,
|
||
docs_url="/docs" if config.docs_enabled else None,
|
||
redoc_url="/redoc" if config.docs_enabled else None,
|
||
openapi_url="/openapi.json" if config.docs_enabled else None,
|
||
openapi_tags=openapi_tags,
|
||
)
|
||
|
||
_register_exception_handlers(app)
|
||
_register_frontdoor_replaceable_middlewares(app)
|
||
_register_python_host_middlewares(app)
|
||
_register_python_only_host_routes(app)
|
||
|
||
|
||
def main() -> Any:
|
||
# 初始化新日志系统
|
||
debug_mode = config.environment == "development"
|
||
# 日志系统已在导入时自动初始化
|
||
|
||
# Parse log level
|
||
log_level = config.log_level.split()[0].lower()
|
||
if log_level not in ["debug", "info", "warning", "error", "critical"]:
|
||
log_level = "info"
|
||
|
||
# 自定义uvicorn日志配置,完全禁用access日志
|
||
uvicorn_log_config = {
|
||
"version": 1,
|
||
"disable_existing_loggers": False,
|
||
"formatters": {
|
||
"default": {
|
||
"format": "%(levelprefix)s %(message)s",
|
||
"use_colors": True,
|
||
},
|
||
},
|
||
"handlers": {
|
||
"default": {
|
||
"formatter": "default",
|
||
"class": "logging.StreamHandler",
|
||
"stream": "ext://sys.stderr",
|
||
},
|
||
},
|
||
"loggers": {
|
||
"uvicorn": {"handlers": ["default"], "level": log_level.upper()},
|
||
"uvicorn.error": {"level": log_level.upper()},
|
||
"uvicorn.access": {"handlers": [], "level": "CRITICAL"}, # 禁用access日志
|
||
},
|
||
}
|
||
|
||
# Start server
|
||
# 根据环境设置热重载
|
||
is_dev = config.environment == "development"
|
||
uvicorn.run(
|
||
"src.main:app",
|
||
host=config.host,
|
||
port=config.port,
|
||
log_level=log_level,
|
||
reload=is_dev,
|
||
reload_dirs=["src"] if is_dev else None,
|
||
access_log=False, # 禁用 uvicorn 访问日志，使用自定义中间件
|
||
log_config=uvicorn_log_config, # 使用自定义日志配置
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 使用安全的方式清屏，避免命令注入风险
|
||
try:
|
||
import os
|
||
|
||
if os.name == "nt": # Windows
|
||
os.system("cls")
|
||
else: # Unix/Linux/MacOS
|
||
print("\033[2J\033[H", end="") # ANSI escape sequence
|
||
except:
|
||
pass # 清屏失败不影响程序运行
|
||
|
||
main()
|