fmq/AstroResearch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: LAMOST DR12-14 与子版本体系接入、观测层安全加固与并发异步化

Browse Source 
- LAMOST 新增 DR12/13/14 及子版本（v0/v1.0/v1.1/v2.0）维度，Internal 发布标记需登录认证并前端灰显，release×subtype 交叉约束下沉至 capabilities 统一声明
- ObservationFetcher trait 扩展版本/认证/交叉约束能力声明，version 参数贯穿 client→service→API→Agent tool→前端全链路
- 安全：observation cache SQL 全参数绑定 + LIKE 转义、cone_cache_hash 加长度前缀防碰撞、DESI survey/program 白名单防穿越
- 异步化：persist/cached_files_total_size/maybe_persist_tool_result迁移到 tokio::fs；cancelled_runs 与 session_permission_checkers改用 DashMap；auth 读锁优先 + 60s 节流
- Gaia 去 native-tls 改禁用连接池规避 UnexpectedEof，reqwest 移除 native-tls feature
- 重构：Source/ProductType from_str 集中解析、download 模块拆分为 try_download_pdf/html、AgentRuntime::init 抽取共享逻辑
- 部署：新增 deploy.sh 一键打包推送脚本、catch-panic 启用
This commit is contained in:
fmq
2026-07-07 21:16:46 +08:00
parent 2f1fd19d74
commit 2c8d0b8f8b
75 changed files with 2482 additions and 1370 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/runtime/executor/helpers.rs
+1
View File
@@ -138,6 +138,7 @@ pub(super) async fn process_single_result(
max_output_chars,
&tool_results_dir,
)
.await
};
// PostToolUse hook
src/agent/runtime/executor/mod.rs
+1 -2
View File
@@ -538,8 +538,7 @@ pub async fn execute_parallel(
let cancel_handle = tokio::spawn(async move {
loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
let locked = app_state_ref.cancelled_runs.lock().await;
if locked.contains(&sid_ref) {
if app_state_ref.cancelled_runs.contains_key(&sid_ref) {
cancel_flag.store(true, Ordering::SeqCst);
return;
}
src/agent/runtime/mod.rs
+44 -109
View File
@@ -96,10 +96,12 @@ pub struct AgentRuntime {
}
impl AgentRuntime {
/// 创建新的运行时实例
pub fn new(app_state: Arc<AppState>) -> Self {
let mut config = AgentConfig::default();
let mode_registry = ModeRegistry::builtins();
/// 共享初始化逻辑：根据 config + mode 构建完整运行时。
fn init(
app_state: Arc<AppState>,
mut config: AgentConfig,
mode_registry: ModeRegistry,
) -> Self {
let mode = mode_registry.get(&config.mode).copied().unwrap_or_else(|| {
tracing::warn!("[AgentRuntime] 未知模式 '{}'，回退到默认模式", config.mode);
mode_registry
@@ -123,14 +125,13 @@ impl AgentRuntime {
tool_registry.add_tool(Box::new(crate::agent::tools::memory::SaveMemoryTool::new(
app_state.memory_manager.clone(),
)));
// 替换 DelegateResearchTool 为带有 permission_checker 的版本（SSE 通道通过 ToolContext 注入）
// 替换 DelegateResearchTool 为带有 permission_checker 的版本
tool_registry.replace_tool(Box::new(
crate::agent::tools::subagent::SubAgentTool::new_with_hooks(
None,
permission_checker.clone(),
),
));
// 会话级权限检查器将在 run_react_loop 中按 session_id 注册
// 视觉模型可用时注册 analyze_image 工具
if app_state.vision_llm.is_some() {
@@ -140,9 +141,9 @@ impl AgentRuntime {
// ── 应用模式的工具集过滤 ──
apply_mode_tool_filter(&mut tool_registry, mode);
// 初始化 checkpoint 管理器
// 初始化 checkpoint 管理器（存储在 library_dir/.checkpoints 下）
let checkpoint_enabled = true;
let checkpoint_store = app_state.config.library_dir.join("..").join(".checkpoints");
let checkpoint_store = app_state.config.library_dir.join(".checkpoints");
let checkpoint_manager = Arc::new(checkpoint::CheckpointManager::new(
std::fs::canonicalize(&checkpoint_store).unwrap_or(checkpoint_store),
checkpoint_enabled,
@@ -168,81 +169,28 @@ impl AgentRuntime {
}
}
/// 创建带自定义配置的运行时实例
pub fn with_config(app_state: Arc<AppState>, mut config: AgentConfig) -> Self {
/// 创建新的运行时实例
pub fn new(app_state: Arc<AppState>) -> Self {
let config = AgentConfig::default();
let mode_registry = ModeRegistry::builtins();
let mode = mode_registry.get(&config.mode).copied().unwrap_or_else(|| {
tracing::warn!("[AgentRuntime] 未知模式 '{}'，回退到默认模式", config.mode);
mode_registry
.get(ModeRegistry::default_id())
.copied()
.unwrap_or(&modes::default::DEFAULT_MODE)
});
// 合并模式配置预设
apply_mode_config(&mut config, mode);
Self::init(app_state, config, mode_registry)
}
let queue = Arc::new(BgNotificationQueue::new());
let metrics_data = Arc::new(tokio::sync::Mutex::new(super::hooks::MetricsData::default()));
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new(
config.denial_max_consecutive,
config.denial_max_total,
)));
let skill_registry = app_state.skill_registry.clone();
let mut tool_registry = ToolRegistry::new_with_queue(Some(queue.clone()), skill_registry);
tool_registry.add_tool(Box::new(crate::agent::tools::memory::SaveMemoryTool::new(
app_state.memory_manager.clone(),
)));
tool_registry.replace_tool(Box::new(
crate::agent::tools::subagent::SubAgentTool::new_with_hooks(
None,
permission_checker.clone(),
),
));
// 会话级权限检查器将在 run_react_loop 中按 session_id 注册
// 视觉模型可用时注册 analyze_image 工具
if app_state.vision_llm.is_some() {
tool_registry.add_tool(Box::new(crate::agent::tools::astro::AnalyzeImageTool));
}
// ── 应用模式的工具集过滤 ──
apply_mode_tool_filter(&mut tool_registry, mode);
// 初始化 checkpoint 管理器
let checkpoint_enabled = true;
let checkpoint_store = app_state.config.library_dir.join("..").join(".checkpoints");
let checkpoint_manager = Arc::new(checkpoint::CheckpointManager::new(
std::fs::canonicalize(&checkpoint_store).unwrap_or(checkpoint_store),
checkpoint_enabled,
));
AgentRuntime {
app_state,
config,
tool_registry,
bg_notification_queue: queue,
metrics_data,
compaction_breaker: Arc::new(std::sync::Mutex::new(
circuit_breaker::CompactionCircuitBreaker::new(),
)),
permission_checker,
denial_tracker,
read_file_state: Arc::new(std::sync::Mutex::new(file_cache::FileStateCache::new())),
prompt_cache: std::sync::Mutex::new(SystemPromptCache::new()),
collapse_log: Arc::new(std::sync::Mutex::new(compact::collapse::CollapseLog::new())),
checkpoint_manager,
mode,
mode_registry,
}
/// 创建带自定义配置的运行时实例
pub fn with_config(app_state: Arc<AppState>, config: AgentConfig) -> Self {
let mode_registry = ModeRegistry::builtins();
Self::init(app_state, config, mode_registry)
}
/// 返回当前运行指标快照（锁异常时返回默认值并记录警告）
pub fn get_metrics(&self) -> super::hooks::MetricsData {
self.metrics_data.try_lock().map(|m| m.clone()).unwrap_or_else(|_| {
warn!("[AgentRuntime] 获取指标锁失败，返回默认值");
super::hooks::MetricsData::default()
})
self.metrics_data
.try_lock()
.map(|m| m.clone())
.unwrap_or_else(|_| {
warn!("[AgentRuntime] 获取指标锁失败，返回默认值");
super::hooks::MetricsData::default()
})
}
/// 设置是否启用 LLM 思考模式（向后兼容，优先使用 mode 设置）
@@ -456,7 +404,7 @@ impl AgentRuntime {
// Phase 4: 会话收尾（传入实际的终止原因 + trajectory 导出参数）
let system_prompt = self.system_prompt();
let model_name = self.app_state.llm.model();
let model_name = self.app_state.llm.model().await;
finalize::finalize_turn(
db,
&session_info.session_id,
@@ -491,12 +439,10 @@ impl AgentRuntime {
let turn_index = session_info.turn_index;
// 注册当前会话的权限检查器（如不存在则从全局配置初始化）
{
let mut checkers = self.app_state.session_permission_checkers.write().await;
checkers
.entry(sid.clone())
.or_insert_with(|| (*self.permission_checker).clone());
}
self.app_state
.session_permission_checkers
.entry(sid.clone())
.or_insert_with(|| (*self.permission_checker).clone());
let tool_defs = self.tool_registry.definitions();
let mut duplicate_detector = DuplicateDetector::default();
@@ -522,10 +468,7 @@ impl AgentRuntime {
self.checkpoint_manager.new_turn();
// 检查用户取消
let is_cancelled = {
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(sid)
};
let is_cancelled = self.app_state.cancelled_runs.remove(sid).is_some();
if is_cancelled {
warn!("[AgentRuntime] 用户手动中止了会话 {} 的智能体执行", sid);
@@ -868,14 +811,11 @@ impl AgentRuntime {
}
// 并行执行工具（带权限检查、checkpoint 和分区器）
let session_checker_snapshot = {
self.app_state
.session_permission_checkers
.read()
.await
.get(sid)
.cloned()
};
let session_checker_snapshot = self
.app_state
.session_permission_checkers
.get(sid)
.map(|r| r.value().clone());
let exec_result = executor::execute_parallel(
&prepared_calls,
&self.tool_registry,
@@ -947,8 +887,7 @@ impl AgentRuntime {
}
if exec_result.was_cancelled {
let mut locked = self.app_state.cancelled_runs.lock().await;
locked.remove(sid);
self.app_state.cancelled_runs.remove(sid);
warn!(
"[AgentRuntime] 工具执行期间被用户手动中止，会话 ID: {}",
sid
@@ -1016,8 +955,7 @@ impl AgentRuntime {
match output.status {
StreamStatus::Success => return Some(output),
StreamStatus::Cancelled => {
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id);
self.app_state.cancelled_runs.remove(session_id);
warn!(
"[AgentRuntime] 流式调用期间被用户手动中止，会话 ID: {}",
session_id
@@ -1077,7 +1015,7 @@ impl AgentRuntime {
"[AgentRuntime] 连续 {} 次过载，切换到备用模型: {}",
consecutive_overloads, fallback
);
llm.set_model(fallback.clone());
llm.set_model(fallback.clone()).await;
consecutive_overloads = 0;
} else if !self.app_state.config.llm_fallback_chain.is_empty() {
let idx = ((consecutive_overloads as usize - 3)
@@ -1088,15 +1026,14 @@ impl AgentRuntime {
"[AgentRuntime] 连续 {} 次过载，从链中切换: {}",
consecutive_overloads, alt
);
llm.set_model(alt.clone());
llm.set_model(alt.clone()).await;
consecutive_overloads = 0;
}
}
}
// 检查用户取消
let cancelled = self.app_state.cancelled_runs.lock().await;
if cancelled.contains(session_id) {
if self.app_state.cancelled_runs.contains_key(session_id) {
warn!("[AgentRuntime] 退避重试期间被用户取消");
let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(),
@@ -1123,8 +1060,7 @@ impl AgentRuntime {
return Some(retry_output);
}
StreamStatus::Cancelled => {
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id);
self.app_state.cancelled_runs.remove(session_id);
return None;
}
StreamStatus::Error(_) => {
@@ -1232,8 +1168,7 @@ impl AgentRuntime {
return Some(retry_output);
}
StreamStatus::Cancelled => {
let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id);
self.app_state.cancelled_runs.remove(session_id);
warn!("[AgentRuntime] 恢复期间被用户中止");
let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(),
@@ -1400,7 +1335,7 @@ impl AgentRuntime {
};
let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
let model_name = self.app_state.llm.model().to_string();
let model_name = self.app_state.llm.model_name_snapshot();
let mut lines = vec![
"# 环境信息".to_string(),
src/agent/runtime/session.rs
+9 -5
View File
@@ -83,7 +83,7 @@ pub async fn create_or_resume_session(
)
.bind(&new_id)
.bind("")
.bind(llm.model())
.bind(llm.model().await)
.bind(mode)
.execute(db)
.await?;
@@ -544,7 +544,9 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result
let forked_at = last_active_id.unwrap_or(0);
// 4. 创建新会话
// 4. 创建新会话与复制消息（同一个事务内）
let mut tx = db.begin().await?;
let branch_id = uuid::Uuid::new_v4().to_string();
let branch_title = if title.is_empty() {
format!("分支 (来自 {})", &session_id[..8.min(session_id.len())])
@@ -565,7 +567,7 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result
.bind(&branch_title)
.bind(session_id)
.bind(serde_json::to_string(&branch_meta).unwrap_or_default())
.execute(db)
.execute(&mut *tx)
.await?;
// 5. 复制所有 active=1 的消息到新会话
@@ -573,7 +575,7 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result
"SELECT COUNT(*) FROM agent_messages WHERE session_id = ? AND active = 1",
)
.bind(session_id)
.fetch_one(db)
.fetch_one(&mut *tx)
.await?;
sqlx::query(
@@ -588,9 +590,11 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result
)
.bind(&branch_id)
.bind(session_id)
.execute(db)
.execute(&mut *tx)
.await?;
tx.commit().await?;
info!(
"[Session] 分叉完成: parent={}, branch={}, copied={} messages, forked_at={}",
session_id, branch_id, copied, forked_at
src/agent/runtime/streaming.rs
+2 -3
View File
@@ -43,7 +43,7 @@ pub async fn process_llm_stream(
tx: &mpsc::UnboundedSender<AgentStreamEvent>,
step: usize,
session_id: &str,
cancelled_runs: Arc<tokio::sync::Mutex<std::collections::HashSet<String>>>,
cancelled_runs: Arc<dashmap::DashMap<String, ()>>,
enable_thinking: bool,
) -> StreamOutput {
// 1. 发起 LLM 流式调用
@@ -76,8 +76,7 @@ pub async fn process_llm_stream(
let cancel_fut = async {
loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
let cancelled = cancelled_runs.lock().await;
if cancelled.contains(&sid) {
if cancelled_runs.contains_key(&sid) {
return;
}
}
src/agent/runtime/streaming_executor.rs
+7 -10
View File
@@ -590,11 +590,8 @@ mod tests {
ads,
arxiv,
vizier,
lamost: crate::clients::lamost::LamostClient::new(
"https://www.lamost.org",
60,
)
.unwrap(),
lamost: crate::clients::lamost::LamostClient::new("https://www.lamost.org", 60)
.unwrap(),
gaia: crate::clients::gaia::GaiaClient::new(
"https://gea.esac.esa.int/tap-server/tap",
"https://gea.esac.esa.int/data-server",
@@ -621,7 +618,7 @@ mod tests {
harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())),
batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
cancelled_runs: Arc::new(dashmap::DashMap::new()),
skill_registry: Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(PathBuf::from(
"/tmp/sk",
)))),
@@ -629,16 +626,16 @@ mod tests {
pending_permissions: Arc::new(
tokio::sync::Mutex::new(std::collections::HashMap::new()),
),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
session_permission_checkers: Arc::new(dashmap::DashMap::new()),
sse_broadcast: None,
memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from(
"/tmp/test_mem",
)))),
sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
upload_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
upload_rate_limiter: Arc::new(
tokio::sync::Mutex::new(std::collections::HashMap::new()),
),
});
ToolContext::new(app_state)