AstroResearch/src/main.rs

// src/main.rs

use anyhow::Context;
use axum::{
    http::HeaderValue,
    routing::{get, post},
    Router,
};
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::Arc;
use tower_http::cors::CorsLayer;
use tower_http::services::ServeDir;
use tower_http::set_header::SetResponseHeaderLayer;
use tracing::{error, info, warn};

use astroresearch::agent::skills::SkillRegistry;
use astroresearch::api::handlers::{self, AppState};
use astroresearch::clients::ads::AdsClient;
use astroresearch::clients::arxiv::ArxivClient;
use astroresearch::clients::cds::vizier::VizierClient;
use astroresearch::clients::desi::DesiClient;
use astroresearch::clients::gaia::GaiaClient;
use astroresearch::clients::lamost::LamostClient;
use astroresearch::clients::llm::{EmbeddingClient, LlmClient};
use astroresearch::clients::qiniu::QiniuClient;
use astroresearch::clients::sdss::SdssClient;
use astroresearch::services::download::Downloader;
use astroresearch::services::translation::Dictionary;
use astroresearch::Config;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // 1. 初始化日志记录器并保留异步写保护 Guard
    let _logging_guards = astroresearch::services::logging::init_logging()?;

    // 1.2 设置全局 panic hook — 记录 panic 信息到日志而非 stderr
    // 防止 Agent 工具执行中的 panic 导致 tokio 任务静默失败
    std::panic::set_hook(Box::new(|panic| {
        let payload = panic.payload();
        let msg = if let Some(s) = payload.downcast_ref::<&str>() {
            s.to_string()
        } else if let Some(s) = payload.downcast_ref::<String>() {
            s.clone()
        } else {
            "Box<dyn Any>".to_string()
        };
        let location = panic
            .location()
            .map(|l| format!(" at {}:{}", l.file(), l.line()))
            .unwrap_or_default();
        tracing::error!("PANIC{}: {}", location, msg);
    }));

    info!("正在启动 AstroResearch 天文学文献辅助系统后端服务...");

    // 1.5 静态注册 sqlite-vec 自动扩展
    // SAFETY: sqlite3_vec_init 的函数签名严格符合 SQLite C API 自动扩展
    // 回调规范 (sqlite3*, char**, const sqlite3_api_routines*)。transmute
    // 将该函数指针转换为 sqlite3_auto_extension 所需的 Option<fn()> 类型。
    // 该注册必须在任何数据库连接开启之前执行，保证所有 Connection
    // 自动拥有 vec0 虚拟表能力。
    unsafe {
        libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::<
            *const (),
            unsafe extern "C" fn(
                *mut libsqlite3_sys::sqlite3,
                *mut *const i8,
                *const libsqlite3_sys::sqlite3_api_routines,
            ) -> i32,
        >(
            sqlite_vec::sqlite3_vec_init as *const ()
        )));
    }
    info!("sqlite-vec 自动扩展注册完成。");

    // 2. 加载环境变量配置
    let config = Config::from_env();
    info!(
        "系统配置成功载入。本地 SQLite 连接串: {}",
        config.database_url
    );
    // 密码安全检查
    let pwd = &config.admin_password;
    let default_pwds = [
        "admin",
        "fmq",
        "password",
        "123456",
        "12345678",
        "admin123",
        "astroresearch",
    ];
    if pwd.is_empty() {
        anyhow::bail!("启动失败: ADMIN_PASSWORD 未设置。请在 .env 文件中设置一个强密码后重试。");
    } else if pwd.len() < 6 {
        warn!(
            "⚠️  安全警告: ADMIN_PASSWORD 长度仅 {} 个字符，属弱密码。请设至少 8 位的强密码。",
            pwd.len()
        );
    } else if default_pwds.contains(&pwd.as_str()) {
        warn!("⚠️  安全警告: ADMIN_PASSWORD 使用了常见弱密码。请更换为自定义强密码。");
    }

    // 创建本地馆藏物理文件夹分类结构
    std::fs::create_dir_all(&config.library_dir).unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("PDF")).unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("HTML")).unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("Markdown")).unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("Translation")).unwrap_or_default();
    // 光谱文件目录（Telescope/{望远镜}/...，各 service 的 persist_bytes 会自动创建深层子目录）
    std::fs::create_dir_all(config.library_dir.join("Telescope").join("lamost"))
        .unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("Telescope").join("gaia")).unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("Telescope").join("sdss")).unwrap_or_default();
    std::fs::create_dir_all(config.library_dir.join("Telescope").join("desi")).unwrap_or_default();
    // Agent Skills 目录
    std::fs::create_dir_all(&config.skills_dir).unwrap_or_default();

    // 3. 初始化本地 SQLite 数据库连接池（开启外键约束并启用 WAL 模式与繁忙等待）
    let options = SqliteConnectOptions::from_str(&config.database_url)?
        .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
        .busy_timeout(std::time::Duration::from_secs(10))
        .foreign_keys(true)
        .create_if_missing(true);

    let pool = SqlitePoolOptions::new()
        .max_connections(5)
        .connect_with(options)
        .await?;

    info!("SQLite 数据库连接已建立（外键约束已启用）。");

    // 4. 自动执行数据库迁移脚本
    info!("开始执行 SQL 表结构迁移...");
    sqlx::migrate!("./migrations").run(&pool).await?;
    info!("数据库迁移执行完成，主表准备就绪。");

    // 4.5 动态创建 vec0 向量虚拟表（维度可由环境变量 EMBEDDING_DIM 控制）
    let embedding_dim: usize = std::env::var("EMBEDDING_DIM")
        .unwrap_or_else(|_| "1536".to_string())
        .parse()
        .unwrap_or(1536);

    // 检测并自愈：如果已存在的 vec_paper_chunks 维度与当前配置不一致，则重建该虚拟表并清空切片内容
    let existing_sql: Option<(String,)> = sqlx::query_as(
        "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'vec_paper_chunks'",
    )
    .fetch_optional(&pool)
    .await?;

    if let Some((sql,)) = existing_sql {
        let expected_pattern = format!("float[{}]", embedding_dim);
        if !sql.contains(&expected_pattern) {
            let force_rebuild = std::env::var("EMBEDDING_DIM_FORCE_REBUILD")
                .map(|v| v == "1" || v == "true")
                .unwrap_or(false);

            if !force_rebuild {
                anyhow::bail!(
                    "向量表维度不匹配 (当前: {}, 期望: {})。\n\
                     此操作将删除所有向量嵌入和文本切片数据且不可恢复。\n\
                     如确认需要重建，请设置环境变量 EMBEDDING_DIM_FORCE_REBUILD=1 后重新启动。",
                    sql,
                    expected_pattern
                );
            }

            warn!(
                "⚠️  EMBEDDING_DIM_FORCE_REBUILD=1 已设置，开始重建向量表。\
                 当前维度: {}, 期望维度: {}。所有现有数据将被清除！",
                sql, expected_pattern
            );
            sqlx::query("DROP TABLE IF EXISTS vec_paper_chunks")
                .execute(&pool)
                .await?;
            sqlx::query("DELETE FROM paper_chunks_content")
                .execute(&pool)
                .await?;
            warn!(
                "向量表和切片内容已清空，正在按新维度 {} 重建...",
                embedding_dim
            );
        }
    }

    let create_vec_table = format!(
        "CREATE VIRTUAL TABLE IF NOT EXISTS vec_paper_chunks USING vec0(embedding float[{}])",
        embedding_dim
    );
    sqlx::query(&create_vec_table).execute(&pool).await?;
    info!("vec_paper_chunks 虚拟表就绪（维度={}）。", embedding_dim);

    // 5. 异步加载天文学专业名词对照词表
    let mut dict = Dictionary::new();
    if let Err(e) = dict.load_from_file("dictionary.txt") {
        error!("天文学名词词表加载失败: {}", e);
    }

    // 6. 初始化并配置全部 API 与下载客户端
    let qiniu = QiniuClient::new(
        config.qiniu_ak.clone(),
        config.qiniu_sk.clone(),
        config.qiniu_bucket.clone(),
        config.qiniu_domain.clone(),
    );

    let ads = AdsClient::new(config.ads_api_key.clone()).context("构建 ADS 客户端失败")?;
    let arxiv = ArxivClient::new().context("构建 arXiv 客户端失败")?;
    let vizier = VizierClient::new(&config.vizier_tap_url, config.vizier_timeout_secs)
        .context("构建 VizieR 客户端失败")?;
    let lamost = LamostClient::new(&config.lamost_base_url, config.lamost_timeout_secs)
        .context("构建 LAMOST 客户端失败")?;
    let gaia = GaiaClient::new(
        &config.gaia_tap_url,
        &config.gaia_datalink_url,
        config.gaia_timeout_secs,
    )
    .context("构建 Gaia 客户端失败")?;
    let sdss = SdssClient::new(&config.sdss_tap_url, config.sdss_timeout_secs)
        .context("构建 SDSS 客户端失败")?;
    let desi = DesiClient::new(&config.desi_tap_url, config.desi_timeout_secs)
        .context("构建 DESI 客户端失败")?;
    let downloader = Downloader::new().context("构建 HTTP 下载客户端失败")?;
    let llm = LlmClient::new(
        config.llm_api_key.clone(),
        config.llm_api_base.clone(),
        config.llm_model.clone(),
    )
    .context("构建 LLM 客户端失败")?;
    let medium_llm = LlmClient::new(
        config.llm_medium_api_key.clone(),
        config.llm_medium_api_base.clone(),
        config.llm_medium_model.clone(),
    )
    .context("构建 Medium LLM 客户端失败")?;
    let fast_llm = LlmClient::new(
        config.llm_fast_api_key.clone(),
        config.llm_fast_api_base.clone(),
        config.llm_fast_model.clone(),
    )
    .context("构建 Fast LLM 客户端失败")?;
    let vision_llm = if !config.llm_vision_model.is_empty() {
        let key = if config.llm_vision_api_key.is_empty() {
            config.llm_api_key.clone()
        } else {
            config.llm_vision_api_key.clone()
        };
        let base = if config.llm_vision_api_base.is_empty() {
            config.llm_api_base.clone()
        } else {
            config.llm_vision_api_base.clone()
        };
        info!("视觉模型已启用: {}", config.llm_vision_model);
        Some(
            LlmClient::new(key, base, config.llm_vision_model.clone())
                .context("构建 Vision LLM 客户端失败")?,
        )
    } else {
        None
    };
    let embedding = EmbeddingClient::new(
        config.embedding_api_key.clone(),
        config.embedding_api_base.clone(),
        config.embedding_model.clone(),
    )
    .context("构建 Embedding 客户端失败")?;

    let skill_registry = Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(
        config.skills_dir.clone(),
    )));
    {
        let mut reg = skill_registry.write().await;
        reg.refresh();
    }
    info!(
        "SkillRegistry 初始化完成，加载 {} 个 skill。",
        skill_registry.read().await.len()
    );

    // 启动文件监听（热更新 skills）
    let _watcher_handle = SkillRegistry::start_watcher(skill_registry.clone());

    let app_state = Arc::new(AppState {
        config: config.clone(),
        db: pool,
        dict,
        qiniu,
        ads,
        arxiv,
        vizier,
        lamost,
        gaia,
        sdss,
        desi,
        llm,
        medium_llm,
        fast_llm,
        vision_llm,
        embedding,
        downloader,
        http_client: reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(30))
            .connect_timeout(std::time::Duration::from_secs(10))
            .build()
            .context("Failed to create shared HTTP client")?,
        harvest_status: Arc::new(tokio::sync::Mutex::new(
            astroresearch::services::batch::MetaSyncStatus::new(),
        )),
        batch_status: Arc::new(tokio::sync::Mutex::new(
            astroresearch::services::batch::AssetBatchStatus::new(),
        )),
        active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
        cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
        skill_registry,
        pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
        pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
        sse_broadcast: {
            let (tx, _) = tokio::sync::broadcast::channel(256);
            Some(tx)
        },
        memory_manager: Arc::new(tokio::sync::Mutex::new(
            astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()),
        )),
        session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
            std::collections::HashMap::new(),
        )),
        sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
        login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
    });

    // 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管
    let allowed_methods = vec![
        axum::http::Method::GET,
        axum::http::Method::POST,
        axum::http::Method::PUT,
        axum::http::Method::DELETE,
        axum::http::Method::OPTIONS,
    ];

    let cors_permissive = CorsLayer::new()
        .allow_origin(tower_http::cors::AllowOrigin::predicate(
            |_origin, _parts| {
                // [Code Review 豁免说明]:
                // 这里的跨域策略被设置为无条件允许 (true)，这是由于书签采集脚本（Bookmarklet）
                // 需要在任意外部学术站点（如 arxiv.org, nature.com 等）向本服务发起 /upload 等请求。
                // 虽然该策略允许凭据 (allow_credentials) 可能构成潜在的跨站请求伪造 (CSRF) 面，
                // 但暴露的 API 仅限文件上传和元数据写入，不涉及任何提权、敏感信息泄露或 RCE 风险。
                // 强制添加白名单会极大增加维护成本，因此该策略是有意为之的安全平衡。
                true
            },
        ))
        .allow_methods(allowed_methods.clone())
        .allow_headers(tower_http::cors::AllowHeaders::mirror_request())
        .allow_credentials(true);

    let cors_local = CorsLayer::new()
        .allow_origin(tower_http::cors::AllowOrigin::predicate(
            |origin, _parts| {
                if let Ok(origin_str) = origin.to_str() {
                    // 支持 http:// 或 https:// 协议下的本地回环及局域网私有网段 IP 来源（允许任意端口，如 Vite dev 端口 :5173）
                    let host_part = if let Some(stripped) = origin_str.strip_prefix("http://") {
                        Some(stripped)
                    } else {
                        origin_str.strip_prefix("https://")
                    };

                    if let Some(host_and_port) = host_part {
                        let host = host_and_port.split(':').next().unwrap_or("");
                        if host == "localhost"
                            || host == "127.0.0.1"
                            || host == "[::1]"
                            || host.starts_with("192.168.")
                            || host.starts_with("10.")
                            || host.starts_with("169.254.")
                        {
                            return true;
                        }
                        if host.starts_with("172.") {
                            if let Some(second_octet_str) =
                                host.strip_prefix("172.").and_then(|s| s.split('.').next())
                            {
                                if let Ok(second_octet) = second_octet_str.parse::<u8>() {
                                    if (16..=31).contains(&second_octet) {
                                        return true;
                                    }
                                }
                            }
                        }
                    }
                }
                false
            },
        ))
        .allow_methods(allowed_methods)
        .allow_headers(tower_http::cors::AllowHeaders::mirror_request())
        .allow_credentials(true);

    let public_routes = Router::new()
        .route("/auth/login", post(handlers::login))
        .layer(cors_local.clone());

    // 专门为浏览器书签采集脚本保留的、允许任何外站跨域的 API 路由
    let bookmarklet_routes = Router::new()
        .route(
            "/upload",
            post(handlers::upload_paper_file)
                .layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024)),
        )
        .route(
            "/active_bibcode",
            get(handlers::get_active_bibcode).post(handlers::set_active_bibcode),
        )
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
            astroresearch::api::auth::auth_middleware,
        ))
        .layer(cors_permissive);

    let protected_routes = Router::new()
        .route("/search", get(handlers::search_papers))
        .route("/search/history", get(handlers::search_history))
        .route("/download", post(handlers::download_paper))
        .route("/no_resource", post(handlers::mark_no_resource))
        .route("/parse", post(handlers::parse_paper))
        .route("/translate", post(handlers::translate_paper))
        .route("/embed", post(handlers::embed_paper))
        .route("/citations", get(handlers::get_citation_network))
        .route("/paper", get(handlers::get_paper_detail))
        .route("/library", get(handlers::get_library))
        .route("/export", post(handlers::export_citations))
        .route("/notes", post(handlers::create_note))
        .route("/notes", get(handlers::get_notes))
        .route("/notes", axum::routing::delete(handlers::delete_note))
        .route("/sync/meta/count", get(handlers::get_meta_sync_count))
        .route("/sync/meta/run", post(handlers::run_meta_sync))
        .route("/sync/meta/status", get(handlers::get_meta_sync_status))
        .route("/batch/asset/run", post(handlers::run_asset_batch))
        .route("/batch/asset/stop", post(handlers::stop_asset_batch))
        .route("/batch/asset/status", get(handlers::get_asset_batch_status))
        .route("/sync/queries", get(handlers::get_sync_queries))
        .route(
            "/sync/queries/:id",
            axum::routing::delete(handlers::delete_sync_query),
        )
        .route("/target/query", get(handlers::query_target))
        .route("/target/associate", post(handlers::associate_target))
        .route("/target/extract", post(handlers::extract_paper_targets))
        .route("/target/list", get(handlers::list_targets))
        // 天文星表查询路由（VizieR TAP + VO Cone Search）
        .route("/catalog/vizier", get(handlers::vizier_query))
        .route("/catalog/vizier/table", get(handlers::vizier_table))
        .route("/catalog/cone", get(handlers::cone_search))
        .route("/catalog/crossmatch", get(handlers::cross_match))
        // 统一光谱下载（跨 LAMOST/Gaia/SDSS，支持坐标模式与标识符模式）
        .route(
            "/catalog/spectrum/download",
            get(handlers::spectrum_download),
        )
        .route("/catalog/spectrum/list", get(handlers::spectrum_list))
        // 智能体路由
        .route("/chat/agent", post(handlers::chat_agent))
        .route("/chat/modes", get(handlers::get_agent_modes))
        .route("/chat/metrics", get(handlers::get_agent_metrics))
        .route("/chat/sessions", get(handlers::list_sessions))
        .route(
            "/chat/sessions/:id",
            get(handlers::get_session).delete(handlers::delete_session),
        )
        .route("/chat/sessions/:id/stop", post(handlers::stop_agent))
        .route("/chat/sessions/:id/audit", get(handlers::get_session_audit))
        .route("/chat/sessions/:id/branch", post(handlers::branch_session))
        .route("/chat/sessions/:id/retry", post(handlers::retry_session))
        .route("/chat/sessions/:id/rewind", post(handlers::rewind_session))
        .route(
            "/chat/sessions/:id/rewind/restore",
            post(handlers::restore_rewound_session),
        )
        .route("/chat/questions", get(handlers::get_pending_questions))
        .route("/chat/answer", post(handlers::answer_question))
        .route(
            "/chat/sessions/:id/permissions",
            get(handlers::get_pending_permissions),
        )
        .route(
            "/chat/sessions/:id/permissions/respond",
            post(handlers::respond_permission),
        )
        .route(
            "/chat/sessions/:id/permissions/rules",
            get(handlers::list_permission_rules).post(handlers::update_permission_rules),
        )
        .route(
            "/chat/sessions/:id/permissions/mode",
            axum::routing::put(handlers::update_permission_mode),
        )
        .route("/auth/logout", post(handlers::logout))
        .route("/auth/check", get(handlers::check_auth))
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
            astroresearch::api::auth::auth_middleware,
        ))
        .layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024))
        .layer(cors_local.clone());

    let api_routes = Router::new()
        .merge(public_routes)
        .merge(bookmarklet_routes)
        .merge(protected_routes);

    // 静态文件资源代理托管（当前端打包至 dashboard/dist 后，直接挂载到主域名根路由）
    let serve_dir = ServeDir::new("dashboard/dist").fallback(tower_http::services::ServeFile::new(
        "dashboard/dist/index.html",
    ));

    let protected_files = Router::new()
        .fallback_service(ServeDir::new(&config.library_dir))
        .layer(axum::middleware::from_fn(pdf_inline_middleware))
        .layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
            astroresearch::api::auth::auth_middleware,
        ))
        .layer(cors_local);

    let app = Router::new()
        .nest("/api", api_routes)
        .nest_service("/api/files", protected_files)
        .fallback_service(serve_dir)
        .layer(tower_http::trace::TraceLayer::new_for_http())
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::X_CONTENT_TYPE_OPTIONS,
            HeaderValue::from_static("nosniff"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::X_FRAME_OPTIONS,
            HeaderValue::from_static("SAMEORIGIN"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            axum::http::header::STRICT_TRANSPORT_SECURITY,
            HeaderValue::from_static("max-age=31536000; includeSubDomains"),
        ))
        .with_state(app_state);

    let addr = SocketAddr::from(([0, 0, 0, 0], config.port));
    info!("天文学科研服务已成功监听 http://localhost:{}", config.port);

    let listener = tokio::net::TcpListener::bind(addr).await?;
    let server = axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    );

    // 优雅关停：监听 SIGTERM / Ctrl+C
    let shutdown_signal = async {
        tokio::signal::ctrl_c()
            .await
            .expect("Failed to install Ctrl+C handler");
        info!("收到关停信号，正在关闭...");
    };

    server
        .with_graceful_shutdown(shutdown_signal)
        .await
        .expect("Server error");

    info!("服务已安全关闭。");
    Ok(())
}

// 针对 PDF 文件请求的中间件：强行设置 Content-Type 和 Content-Disposition 头，防止在移动端/手机浏览器上被强制下载。
async fn pdf_inline_middleware(
    req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Result<axum::response::Response<axum::body::Body>, axum::http::StatusCode> {
    let path = req.uri().path().to_lowercase();
    let is_pdf = path.ends_with(".pdf");

    let mut response = next.run(req).await;

    if is_pdf && response.status().is_success() {
        // 获取文件名
        let filename = std::path::Path::new(&path)
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("document.pdf");

        // 强行设置 Content-Type 为 application/pdf
        response.headers_mut().insert(
            axum::http::header::CONTENT_TYPE,
            axum::http::HeaderValue::from_static("application/pdf"),
        );

        // 强行设置 Content-Disposition 为 inline
        let disposition_val = format!("inline; filename=\"{}\"", filename);
        if let Ok(hv) = axum::http::HeaderValue::from_str(&disposition_val) {
            response
                .headers_mut()
                .insert(axum::http::header::CONTENT_DISPOSITION, hv);
        }
    }

    Ok(response)
}