zclaw_openfang/crates/zclaw-saas/src/main.rs

//! ZCLAW SaaS 服务入口

use axum::extract::State;
use tower_http::timeout::TimeoutLayer;
use tracing::info;
use zclaw_saas::{config::SaaSConfig, db::init_db, state::AppState};
use zclaw_saas::workers::WorkerDispatcher;
use zclaw_saas::workers::log_operation::LogOperationWorker;
use zclaw_saas::workers::cleanup_refresh_tokens::CleanupRefreshTokensWorker;
use zclaw_saas::workers::cleanup_rate_limit::CleanupRateLimitWorker;
use zclaw_saas::workers::record_usage::RecordUsageWorker;
use zclaw_saas::workers::update_last_used::UpdateLastUsedWorker;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "zclaw_saas=debug,tower_http=debug".into()),
        )
        .init();

    let config = SaaSConfig::load()?;
    info!("SaaS config loaded: {}:{}", config.server.host, config.server.port);

    let db = init_db(&config.database.url).await?;
    info!("Database initialized");

    // 初始化 Worker 调度器 + 注册所有 Worker
    let mut dispatcher = WorkerDispatcher::new(db.clone());
    dispatcher.register(LogOperationWorker);
    dispatcher.register(CleanupRefreshTokensWorker);
    dispatcher.register(CleanupRateLimitWorker);
    dispatcher.register(RecordUsageWorker);
    dispatcher.register(UpdateLastUsedWorker);
    info!("Worker dispatcher initialized (5 workers registered)");

    let state = AppState::new(db.clone(), config.clone(), dispatcher)?;

    // 启动声明式 Scheduler（从 TOML 配置读取定时任务）
    let scheduler_config = &config.scheduler;
    zclaw_saas::scheduler::start_scheduler(scheduler_config, db.clone(), state.worker_dispatcher.clone_ref());
    info!("Scheduler started with {} jobs", scheduler_config.jobs.len());

    // 启动内置 DB 清理任务（设备清理等不通过 Worker 的任务）
    zclaw_saas::scheduler::start_db_cleanup_tasks(db.clone());

    // 启动内存中的 rate limit 条目清理
    let rate_limit_state = state.clone();
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(std::time::Duration::from_secs(300));
        loop {
            interval.tick().await;
            rate_limit_state.cleanup_rate_limit_entries();
        }
    });

    let app = build_router(state).await;

    let listener = tokio::net::TcpListener::bind(format!("{}:{}", config.server.host, config.server.port))
        .await?;
    info!("SaaS server listening on {}:{}", config.server.host, config.server.port);

    axum::serve(listener, app.into_make_service_with_connect_info::<std::net::SocketAddr>()).await?;
    Ok(())
}

async fn health_handler(State(state): State<AppState>) -> axum::Json<serde_json::Value> {
    let db_healthy = sqlx::query_scalar::<_, i32>("SELECT 1")
        .fetch_one(&state.db)
        .await
        .is_ok();

    let status = if db_healthy { "healthy" } else { "degraded" };
    let _code = if db_healthy { 200 } else { 503 };

    axum::Json(serde_json::json!({
        "status": status,
        "database": db_healthy,
        "timestamp": chrono::Utc::now().to_rfc3339(),
        "version": env!("CARGO_PKG_VERSION"),
    }))
}

async fn build_router(state: AppState) -> axum::Router {
    use axum::middleware;
    use tower_http::cors::{Any, CorsLayer};
    use tower_http::trace::TraceLayer;

    use axum::http::HeaderValue;
    let cors = {
        let config = state.config.read().await;
        let is_dev = std::env::var("ZCLAW_SAAS_DEV")
            .map(|v| v == "true" || v == "1")
            .unwrap_or(false);
        if config.server.cors_origins.is_empty() {
            if is_dev {
                CorsLayer::new()
                    .allow_origin(Any)
                    .allow_methods(Any)
                    .allow_headers(Any)
            } else {
                tracing::error!("生产环境必须配置 server.cors_origins，不能使用 allow_origin(Any)");
                panic!("生产环境必须配置 server.cors_origins 白名单。开发环境可设置 ZCLAW_SAAS_DEV=true 绕过。");
            }
        } else {
            let origins: Vec<HeaderValue> = config.server.cors_origins.iter()
                .filter_map(|o: &String| o.parse::<HeaderValue>().ok())
                .collect();
            CorsLayer::new()
                .allow_origin(origins)
                .allow_methods([
                    axum::http::Method::GET,
                    axum::http::Method::POST,
                    axum::http::Method::PUT,
                    axum::http::Method::PATCH,
                    axum::http::Method::DELETE,
                    axum::http::Method::OPTIONS,
                ])
                .allow_headers([
                    axum::http::header::AUTHORIZATION,
                    axum::http::header::CONTENT_TYPE,
                    axum::http::HeaderName::from_static("x-request-id"),
                ])
        }
    };

    let public_routes = zclaw_saas::auth::routes()
        .route("/api/health", axum::routing::get(health_handler));

    let protected_routes = zclaw_saas::auth::protected_routes()
        .merge(zclaw_saas::account::routes())
        .merge(zclaw_saas::model_config::routes())
        .merge(zclaw_saas::relay::routes())
        .merge(zclaw_saas::migration::routes())
        .merge(zclaw_saas::role::routes())
        .merge(zclaw_saas::prompt::routes())
        .merge(zclaw_saas::agent_template::routes())
        .merge(zclaw_saas::telemetry::routes())
        .layer(middleware::from_fn_with_state(
            state.clone(),
            zclaw_saas::middleware::api_version_middleware,
        ))
        .layer(middleware::from_fn_with_state(
            state.clone(),
            zclaw_saas::middleware::request_id_middleware,
        ))
        .layer(middleware::from_fn_with_state(
            state.clone(),
            zclaw_saas::middleware::rate_limit_middleware,
        ))
        .layer(middleware::from_fn_with_state(
            state.clone(),
            zclaw_saas::auth::auth_middleware,
        ));

    axum::Router::new()
        .merge(public_routes)
        .merge(protected_routes)
        .layer(TimeoutLayer::new(std::time::Duration::from_secs(30)))
        .layer(TraceLayer::new_for_http())
        .layer(cors)
        .with_state(state)
}