AstroResearch/src/api/mod.rs

// src/api/mod.rs
use crate::agent::memory::MemoryManager;
use crate::agent::skills::SkillRegistry;
use crate::clients::ads::AdsClient;
use crate::clients::arxiv::ArxivClient;
use crate::clients::cds::vizier::VizierClient;
use crate::clients::desi::DesiClient;
use crate::clients::gaia::GaiaClient;
use crate::clients::lamost::LamostClient;
use crate::clients::llm::{EmbeddingClient, LlmClient};
use crate::clients::qiniu::QiniuClient;
use crate::clients::sdss::SdssClient;
use crate::services::download::Downloader;
use crate::services::translation::Dictionary;
use crate::Config;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, oneshot, Mutex, RwLock};

/// 提供给前端的 SSE 事件
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum AppEvent {
    #[serde(rename = "user_question")]
    UserQuestion { data: String },
}

/// 待回答的用户问题（供 ask_user 工具和 API 端点共享）
#[derive(Debug)]
pub struct PendingQuestion {
    pub question_json: String,
    pub answer_tx: oneshot::Sender<crate::agent::tools::ask_user::UserAnswer>,
}

/// 用户对权限请求的响应
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PermissionResponse {
    pub tool_call_id: String,
    pub allowed: bool,
    #[serde(default)]
    pub allow_always: bool,
}

/// 待处理的权限请求（供权限系统与 API 端点共享）
#[derive(Debug)]
pub struct PendingPermission {
    pub tool_call_id: String,
    pub tool_name: String,
    pub message: String,
    pub arguments: serde_json::Value,
    pub response_tx: oneshot::Sender<PermissionResponse>,
}

// 全局共享的 Axum 应用上下文状态
pub struct AppState {
    pub config: Config,
    pub db: SqlitePool,
    pub dict: Dictionary,
    pub qiniu: QiniuClient,
    pub ads: AdsClient,
    pub arxiv: ArxivClient,
    /// VizieR TAP 星表查询客户端（CDS VizieR TAP 服务）
    pub vizier: VizierClient,
    /// LAMOST 观测数据客户端（ConeSearch + FITS.gz 下载，服务光谱/光变等产品）
    pub lamost: LamostClient,
    /// Gaia 观测数据客户端（TAP + DataLink，服务光谱/光变等产品）
    pub gaia: GaiaClient,
    /// SDSS 观测数据客户端（Data Lab TAP + SAS，服务光谱/APOGEE 等产品）
    pub sdss: SdssClient,
    /// DESI 观测数据客户端（Data Lab TAP + HEALPix coadd SAS，服务光谱等产品）
    pub desi: DesiClient,
    /// 观测数据 fetcher 注册表（跨源 × 跨产品类型的统一获取入口）
    pub observation_registry: Arc<crate::services::observation::ObservationRegistry>,
    pub llm: LlmClient,
    pub medium_llm: LlmClient,
    pub fast_llm: LlmClient,
    /// 专用视觉模型（LLM_VISION_MODEL 配置时启用图片分析工具）
    pub vision_llm: Option<LlmClient>,
    pub embedding: EmbeddingClient,
    pub downloader: Downloader,
    /// 共享 HTTP 客户端（带超时），避免每次请求新建连接池
    pub http_client: reqwest::Client,
    pub harvest_status: Arc<tokio::sync::Mutex<crate::services::batch::MetaSyncStatus>>,
    pub batch_status: Arc<tokio::sync::Mutex<crate::services::batch::AssetBatchStatus>>,
    pub active_bibcode: Arc<tokio::sync::Mutex<Option<String>>>,
    pub cancelled_runs: Arc<dashmap::DashMap<String, ()>>,
    pub skill_registry: Arc<RwLock<SkillRegistry>>,
    /// ask_user 工具 — 待回答的问题
    pub pending_questions: Arc<Mutex<HashMap<String, PendingQuestion>>>,
    /// 权限检查 — 待处理的权限确认请求
    pub pending_permissions: Arc<Mutex<HashMap<String, PendingPermission>>>,
    /// 会话级权限检查器注册表（按 session_id 隔离，支持 API 动态添加/移除规则）
    pub session_permission_checkers:
        Arc<dashmap::DashMap<String, crate::agent::runtime::permission::PermissionChecker>>,
    /// SSE 广播通道（agent 运行时向所有连接的客户端推送事件）
    pub sse_broadcast: Option<broadcast::Sender<AppEvent>>,
    /// 项目记忆管理器（跨会话持久化）
    pub memory_manager: Arc<tokio::sync::Mutex<MemoryManager>>,
    /// 活跃的登录会话 Token 及其最后活跃时间（单用户内存管理）
    pub sessions: Arc<RwLock<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
    /// 登录速率限制器（IP -> 最近失败次数 + 首次失败时间）
    pub login_rate_limiter:
        Arc<Mutex<std::collections::HashMap<String, (u32, std::time::Instant)>>>,
    /// 上传速率限制器（token -> 最近上传次数 + 首次上传时间）
    pub upload_rate_limiter:
        Arc<Mutex<std::collections::HashMap<String, (u32, std::time::Instant)>>>,
}

// 统一标准化的文献格式，用于向前端传输
pub use crate::services::paper::StandardPaper;

pub mod agent;
pub mod auth;
pub mod catalog;
pub mod error;
pub mod notes;
pub mod observation;
pub mod papers;
pub mod permissions;
pub mod search;
pub mod sync;
pub mod targets;

// 提供兼容的 handlers 命名空间，避免修改 main.rs 里的导入
pub mod handlers {
    pub use super::agent::{
        answer_question, branch_session, chat_agent, delete_session, get_agent_metrics,
        get_agent_modes, get_pending_permissions, get_pending_questions, get_session,
        get_session_audit, list_sessions, respond_permission, restore_rewound_session,
        retry_session, rewind_session, stop_agent, AgentChatRequest, BranchResponse,
        RestoreResponse, RetryResponse, RewindRequest, RewindResponse, SessionListParams,
    };
    pub use super::auth::{check_auth, login, logout};
    pub use super::catalog::{
        cone_search, vizier_query, vizier_table, ConeSearchParams, VizierQueryParams,
        VizierTableParams,
    };
    pub use super::notes::{
        create_note, delete_note, get_notes, CreateNoteRequest, DeleteNoteParams, GetNotesParams,
        NoteRecord,
    };
    pub use super::observation::{
        observation_capabilities, observation_download, observation_list, observation_search,
        DownloadMode, ObservationDownloadRequest, ObservationListParams, ObservationListResponse,
        ObservationSearchParams,
    };
    pub use super::papers::{
        download_paper, embed_paper, export_citations, get_active_bibcode, get_citation_network,
        get_library, get_paper_detail, mark_no_resource, parse_paper, search_papers,
        set_active_bibcode, translate_paper, upload_paper_file, DownloadRequest, EmbedRequest,
        EmbedResponse, ExportRequest, ExportResponse, MarkNoResourceRequest, PaperDetailResponse,
        ParseRequest, ParseResponse, SearchParams, TranslateRequest, TranslateResponse,
    };
    pub use super::permissions::{
        list_permission_rules, update_permission_mode, update_permission_rules,
    };
    pub use super::search::{search as search_history, SearchParams as SearchHistoryParams};
    pub use super::sync::{
        delete_sync_query, get_asset_batch_status, get_meta_sync_count, get_meta_sync_status,
        get_sync_queries, run_asset_batch, run_meta_sync, stop_asset_batch, AssetBatchRunRequest,
        MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest,
    };
    pub use super::targets::{
        associate_target, extract_paper_targets, list_targets, query_target, AssociateResponse,
        AssociateTargetRequest, ExtractTargetsRequest, ExtractTargetsResponse, TargetListParams,
        TargetQueryParams,
    };
    pub use super::{AppState, StandardPaper};
    pub use crate::services::batch::SavedSyncQuery;
    pub use crate::services::paper::{
        check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard,
        get_paper_from_db, save_paper_to_db, save_paper_to_db_tx,
    };
    pub use crate::services::session::{
        AgentMetrics as AgentMetricsResponse, AuditLogEntry, MessageRecord, SessionDetail,
        SessionSummary,
    };
}