From d6b064a4903d2610668ef727af22b012295d9d72 Mon Sep 17 00:00:00 2001 From: Asfmq <2696428814@qq.com> Date: Mon, 7 Sep 2026 21:50:32 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E7=A7=91=E7=A0=94=E5=88=86=E6=9E=90?= =?UTF-8?q?=E5=B1=82=E5=85=A8=E6=A0=88=E8=90=BD=E5=9C=B0=E2=80=94=E2=80=94?= =?UTF-8?q?=E5=85=89=E8=B0=B1/=E6=97=B6=E5=9F=9F/=E8=BF=90=E5=8A=A8?= =?UTF-8?q?=E5=AD=A6=E5=88=86=E6=9E=90=E5=B7=A5=E5=85=B7=E9=93=BE=20+=20JW?= =?UTF-8?q?ST/X=20=E5=B0=84=E7=BA=BF=E6=95=B0=E6=8D=AE=E6=BA=90=20+=20?= =?UTF-8?q?=E5=AE=9A=E6=97=B6=E6=96=87=E7=8C=AE=E5=90=8C=E6=AD=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 数据分析层（新增 services/{spectrum,timeseries,analysis}）: - 光谱参数提取 parameters.rs：LAMOST/SDSS/APOGEE/DESI FITS header 跨源归一化读取 Teff/logg/[Fe/H]/RV 及 ASPCAP 20+ 元素丰度，rayon 并发批量提取 - 谱线测量 lines.rs：内置真空/空气波长谱线表，窗口内极值搜索 + 梯形法积分 EW + FWHM，支持自定义谱线 - 交叉相关测速 cross_correlate.rs：对数波长重采样对齐，内置 Pickles 模板按光谱型插值， CCF 峰值位置提取 RV 及不确定度 - 周期搜索 periodicity.rs：Lomb-Scargle 周期图（含 FAP 误报概率）+ BLS 凌星检测 + 相位折叠 - 变星分类 classification.rs：振幅/偏度/峰度/过零率/eta 等统计特征 + 规则分类（RR Lyrae/Cepheid/食双星/AGN 等） - SED 拟合 sed.rs：多波段测光黑体模型拟合，输出 T_eff/半径/消光 A_V/光度及不确定度 - 运动学 kinematics.rs：视差+自行+RV → 银河系 UVW 空间速度，含移动星群成员概率（Banyan Σ 简化版） - 化学丰度 chemistry.rs：[α/Fe] vs [Fe/H] 计算，厚盘/薄盘/晕星族判别 - 观测规划 observability.rs：目标升落时间/airmass/月相影响/曝光时间估算 - 赫罗图 hr_diagram.rs：Gaia TAP CMD 查询，新增 GET /api/analysis/hr-diagram 端点 数据获取层: - JWST：clients/mast/jwst.rs 封装 MAST Portal 锥形检索 + JwstSpectrumFetcher（NIRSpec/MIRI 光谱） - X 射线：clients/heasarc 封装 HEASARC TAP（ADQL）+ XMM-Newton/Chandra 光谱 fetcher - 图像 cutout：SDSS SkyServer/STScI DSS/Pan-STARRS 三源 cutout + 发现图（Finding Chart）生成 - Source 枚举新增 Jwst/Xmm/Chandra 并注册 ObservationRegistry，前端 SOURCE_THEME 与筛选器同步三源 Agent 工具集（24→35）: - 新增 9 个分析工具：get_spectrum_parameters / measure_spectral_lines / measure_radial_velocity / find_period / classify_variable_star / fit_sed / analyze_kinematics / analyze_abundance_pattern / plan_observation - batch_process：批量样本"查询→下载→分析→报告"流水线，并发控制防数据源速率限制 - literature_monitor：按 ADS 查询式/时间窗/最低引用数检查最新文献 定时文献同步: - sync_queries 表新增 is_scheduled 列（migration 20260713） - 新增 POST /sync/queries/:id/schedule 端点 - 服务启动时拉起每小时调度器，对 is_scheduled=1 的检索配置静默执行 ADS（entdate 增量）/arXiv 增量同步 - search_history 工具收敛至 services/search::search_agent_history，消除 FTS 查询逻辑重复 其他: - plotting skill 由占位填充为完整科研绘图规范：光谱/光变/折叠曲线/CMD/SED/[α/Fe]/周期图/Mollweide/发现图 9 类 matplotlib 模板 - 删除死代码 streaming_executor.rs（929 行，仅剩 mod 声明引用，无调用方） - 新增 docs/roadmap-research-features.md 科研功能路线图及实现状态 --- Cargo.lock | 40 + Cargo.toml | 1 + .../src/components/agent/AgentMessageList.tsx | 4 +- .../components/agent/SpecialToolRenderers.tsx | 4 +- .../components/agent/SubAgentContainer.tsx | 1 - .../observation/UnifiedSearchPanel.tsx | 194 +-- .../src/components/observation/constants.ts | 15 + dashboard/src/pages/ObservationPanel.tsx | 3 + docs/architecture/agent/dsh-alignment.md | 132 ++ docs/architecture/agent/env-vars.md | 1 + docs/architecture/agent/overview.md | 12 +- docs/roadmap-research-features.md | 438 ++++++ ...00000_add_is_scheduled_to_sync_queries.sql | 2 + .../20260819000001_add_agent_events.sql | 25 + skills/plotting/SKILL.md | 307 ++++- src/agent/autonomous.rs | 13 +- src/agent/compact.rs | 123 +- src/agent/engine.rs | 1224 +++++++++++++++++ src/agent/hooks/mod.rs | 2 +- src/agent/mod.rs | 6 +- src/agent/runtime/checkpoint.rs | 12 +- src/agent/runtime/context.rs | 103 +- src/agent/runtime/duplicate_detector.rs | 92 +- src/agent/runtime/executor/helpers.rs | 145 +- src/agent/runtime/executor/mod.rs | 536 ++++---- src/agent/runtime/mod.rs | 1204 +++------------- src/agent/runtime/permission.rs | 249 ++-- src/agent/runtime/session.rs | 46 + src/agent/runtime/session_events.rs | 446 ++++++ src/agent/runtime/session_registry.rs | 187 +++ src/agent/runtime/streaming.rs | 11 +- src/agent/runtime/streaming_executor.rs | 929 ------------- src/agent/subagent.rs | 610 +++----- src/agent/team/teammate.rs | 170 +-- src/agent/terminal.rs | 12 + src/agent/tools/astro/mod.rs | 10 + .../astro/research/abundance_analysis.rs | 104 ++ .../tools/astro/research/batch_process.rs | 168 +++ .../tools/astro/research/kinematics_tool.rs | 139 ++ .../astro/research/literature_monitor.rs | 158 +++ src/agent/tools/astro/research/mod.rs | 22 + src/agent/tools/astro/research/observation.rs | 4 +- .../tools/astro/research/observation_plan.rs | 115 ++ .../tools/astro/research/period_search.rs | 148 ++ .../tools/astro/research/radial_velocity.rs | 151 ++ src/agent/tools/astro/research/rag.rs | 5 + src/agent/tools/astro/research/sed_fit.rs | 125 ++ .../tools/astro/research/spectral_lines.rs | 168 +++ .../tools/astro/research/spectrum_params.rs | 130 ++ .../tools/astro/research/variable_star.rs | 103 ++ src/agent/tools/astro/system/search.rs | 11 +- src/agent/tools/compress.rs | 10 +- src/agent/tools/filesystem/bash.rs | 13 + src/agent/tools/filesystem/edit.rs | 14 + src/agent/tools/filesystem/write.rs | 14 + src/agent/tools/mod.rs | 231 +++- src/agent/tools/persist.rs | 18 +- src/agent/tools/search_history.rs | 104 +- src/agent/tools/subagent.rs | 66 +- src/agent/tools/todo.rs | 25 +- src/api/agent.rs | 72 +- src/api/catalog.rs | 13 + src/api/mod.rs | 5 + src/api/observation.rs | 10 +- src/api/sync.rs | 28 + src/bin/cli.rs | 8 + src/clients/gaia/mod.rs | 5 + src/clients/heasarc/mod.rs | 190 +++ src/clients/mast/jwst.rs | 290 ++++ src/clients/mast/mod.rs | 2 + src/clients/mod.rs | 1 + src/main.rs | 52 + src/services/analysis/chemistry.rs | 280 ++++ src/services/analysis/hr_diagram.rs | 153 +++ src/services/analysis/kinematics.rs | 365 +++++ src/services/analysis/mod.rs | 15 + src/services/analysis/observability.rs | 401 ++++++ src/services/analysis/sed.rs | 363 +++++ src/services/batch/meta.rs | 161 ++- src/services/batch/mod.rs | 4 +- src/services/cds/vizier.rs | 1 + src/services/mod.rs | 3 + src/services/observation/cache.rs | 1 + src/services/observation/cutout.rs | 378 +++++ src/services/observation/dispatch.rs | 4 + src/services/observation/fetcher.rs | 4 + src/services/observation/jwst.rs | 212 +++ src/services/observation/mod.rs | 3 + src/services/observation/photometry.rs | 6 +- src/services/observation/preview.rs | 37 +- src/services/observation/registry.rs | 16 +- src/services/observation/types.rs | 38 +- src/services/observation/unified.rs | 24 +- src/services/observation/xray.rs | 357 +++++ src/services/paper/db.rs | 37 + src/services/paper/mod.rs | 4 +- src/services/search.rs | 6 + src/services/spectrum/cross_correlate.rs | 562 ++++++++ src/services/spectrum/lines.rs | 629 +++++++++ src/services/spectrum/mod.rs | 12 + src/services/spectrum/parameters.rs | 378 +++++ src/services/timeseries/classification.rs | 438 ++++++ src/services/timeseries/mod.rs | 10 + src/services/timeseries/periodicity.rs | 469 +++++++ 104 files changed, 12101 insertions(+), 3301 deletions(-) create mode 100644 docs/architecture/agent/dsh-alignment.md create mode 100644 docs/roadmap-research-features.md create mode 100644 migrations/20260713000000_add_is_scheduled_to_sync_queries.sql create mode 100644 migrations/20260819000001_add_agent_events.sql create mode 100644 src/agent/engine.rs create mode 100644 src/agent/runtime/session_events.rs create mode 100644 src/agent/runtime/session_registry.rs delete mode 100644 src/agent/runtime/streaming_executor.rs create mode 100644 src/agent/tools/astro/research/abundance_analysis.rs create mode 100644 src/agent/tools/astro/research/batch_process.rs create mode 100644 src/agent/tools/astro/research/kinematics_tool.rs create mode 100644 src/agent/tools/astro/research/literature_monitor.rs create mode 100644 src/agent/tools/astro/research/observation_plan.rs create mode 100644 src/agent/tools/astro/research/period_search.rs create mode 100644 src/agent/tools/astro/research/radial_velocity.rs create mode 100644 src/agent/tools/astro/research/sed_fit.rs create mode 100644 src/agent/tools/astro/research/spectral_lines.rs create mode 100644 src/agent/tools/astro/research/spectrum_params.rs create mode 100644 src/agent/tools/astro/research/variable_star.rs create mode 100644 src/clients/heasarc/mod.rs create mode 100644 src/clients/mast/jwst.rs create mode 100644 src/services/analysis/chemistry.rs create mode 100644 src/services/analysis/hr_diagram.rs create mode 100644 src/services/analysis/kinematics.rs create mode 100644 src/services/analysis/mod.rs create mode 100644 src/services/analysis/observability.rs create mode 100644 src/services/analysis/sed.rs create mode 100644 src/services/observation/cutout.rs create mode 100644 src/services/observation/jwst.rs create mode 100644 src/services/observation/xray.rs create mode 100644 src/services/spectrum/cross_correlate.rs create mode 100644 src/services/spectrum/lines.rs create mode 100644 src/services/spectrum/mod.rs create mode 100644 src/services/spectrum/parameters.rs create mode 100644 src/services/timeseries/classification.rs create mode 100644 src/services/timeseries/mod.rs create mode 100644 src/services/timeseries/periodicity.rs diff --git a/Cargo.lock b/Cargo.lock index 3bfd358..2bf08b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -155,6 +155,7 @@ dependencies = [ "obscura-net", "quick-xml", "rand 0.8.6", + "rayon", "regex", "reqwest", "serde", @@ -836,6 +837,25 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.12" @@ -3039,6 +3059,26 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" diff --git a/Cargo.toml b/Cargo.toml index 01d80ff..7fda11e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,6 +63,7 @@ walkdir = "2" lru = "0.12" git2 = { version = "0.18", default-features = false, features = ["vendored-libgit2"] } dashmap = "6" +rayon = "1.10" fitsio = { version = "0.21", features = ["fitsio-src"] } # FITS 解析；fitsio-src 自带 CFITSIO 源码编译，无需系统装 cfitsio fitsio-sys = { version = "0.5", features = ["fitsio-src"] } # 直接访问 CFITSIO C API（读 Gaia XP 变长数组列 PD(55)，fitsio crate 不支持） diff --git a/dashboard/src/components/agent/AgentMessageList.tsx b/dashboard/src/components/agent/AgentMessageList.tsx index eff451c..a760a71 100644 --- a/dashboard/src/components/agent/AgentMessageList.tsx +++ b/dashboard/src/components/agent/AgentMessageList.tsx @@ -19,7 +19,7 @@ import { import { AgentMarkdown } from './AgentMarkdown'; import { ThoughtCard } from './ThoughtCard'; import { ToolCallCard } from './ToolCallCard'; -import { TodoTaskCard } from './SpecialToolRenderers'; +import { TodoTaskCard, type TodoItem } from './SpecialToolRenderers'; import { AnswerCard } from './AnswerCard'; import { SubAgentContainer } from './SubAgentContainer'; import { PARENT_COLORS } from './constants'; @@ -106,7 +106,7 @@ interface AgentMessageListProps { showTaskSidebar?: boolean; setShowTaskSidebar?: (val: boolean) => void; hasTodos?: boolean; - activeTodos?: any[]; + activeTodos?: TodoItem[]; } // 内部工具（is_internal = true）渲染模式配置： diff --git a/dashboard/src/components/agent/SpecialToolRenderers.tsx b/dashboard/src/components/agent/SpecialToolRenderers.tsx index 7670ab7..66a4bbd 100644 --- a/dashboard/src/components/agent/SpecialToolRenderers.tsx +++ b/dashboard/src/components/agent/SpecialToolRenderers.tsx @@ -419,7 +419,7 @@ export function PaperListCard({ // ========================================== // 3. 任务规划看板卡片 (todo_write) // ========================================== -interface TodoItem { +export interface TodoItem { id: string; content: string; status: 'pending' | 'in_progress' | 'completed'; @@ -453,7 +453,7 @@ export function TodoTaskCard({ todos }: TodoTaskCardProps) { {idx + 1} ); - let rowClass = 'opacity-70'; + let rowClass: string; let titleClass = 'text-secondary font-medium'; let badgeClass = 'bg-sunken text-tertiary'; diff --git a/dashboard/src/components/agent/SubAgentContainer.tsx b/dashboard/src/components/agent/SubAgentContainer.tsx index 022293d..1d4dbd8 100644 --- a/dashboard/src/components/agent/SubAgentContainer.tsx +++ b/dashboard/src/components/agent/SubAgentContainer.tsx @@ -26,7 +26,6 @@ interface SubAgentContainerProps { export function SubAgentContainer({ status, children, - summary: _summary, isStreaming: parentStreaming, isCollapsed, onToggle, diff --git a/dashboard/src/components/observation/UnifiedSearchPanel.tsx b/dashboard/src/components/observation/UnifiedSearchPanel.tsx index 2ab31d4..4f9a116 100644 --- a/dashboard/src/components/observation/UnifiedSearchPanel.tsx +++ b/dashboard/src/components/observation/UnifiedSearchPanel.tsx @@ -10,7 +10,7 @@ // 与 ObservationPanel.SearchView 的区别： // - SearchView：单目标 + 单源 + 单产品（精细化参数） // - UnifiedSearchPanel：多目标 + 多源（广覆盖发现） -import { useState, useMemo, useCallback, useEffect } from 'react'; +import { useState, useMemo, useCallback } from 'react'; import { Globe, Search, @@ -64,20 +64,20 @@ export function UnifiedSearchPanel({ downloadUnifiedSelected, downloading, }: UnifiedSearchPanelProps) { - // ── 统一检索的下载结果 ── - const [downloadResults, setDownloadResults] = useState< - ObservationBatchResult[] | null - >(null); - - // 当重新检索结果变化时，重置下载结果 - useEffect(() => { - setDownloadResults(null); - }, [unifiedResult]); + // ── 统一检索的下载结果（缓存归属哪次检索；重新检索后自动失效）── + const [downloadCache, setDownloadCache] = useState<{ + result: UnifiedSearchResult | null; + results: ObservationBatchResult[]; + } | null>(null); + const downloadResults = + downloadCache && downloadCache.result === unifiedResult + ? downloadCache.results + : null; const handleDownload = useCallback(async () => { const r = await downloadUnifiedSelected(); - setDownloadResults(r); - }, [downloadUnifiedSelected]); + setDownloadCache({ result: unifiedResult, results: r }); + }, [downloadUnifiedSelected, unifiedResult]); // ── 输入模式 ── const [inputMode, setInputMode] = useState('coordinates'); @@ -206,11 +206,10 @@ export function UnifiedSearchPanel({ ); // 第一次加载完能力后默认选择最新的子版本以及所有发布版本 - useEffect(() => { - if (selectedSources === null && defaultSelectedKeys.length > 0) { - setSelectedSources(new Set(defaultSelectedKeys)); - } - }, [defaultSelectedKeys, selectedSources]); + // （渲染期一次性初始化，替代 effect 中的 setState，避免级联渲染） + if (selectedSources === null && defaultSelectedKeys.length > 0) { + setSelectedSources(new Set(defaultSelectedKeys)); + } // 同步选中的有效 Set const effectiveSelected = useMemo(() => { @@ -473,94 +472,101 @@ export function UnifiedSearchPanel({ // ── 渲染子孙层级节点 (Level 1, 2, 3) ── const renderNestedNode = useCallback( (node: SourceTreeNode) => { - const isExpanded = expandedNodes.has(node.id); - const hasChildren = node.children && node.children.length > 0; + // 递归渲染定义为内部具名函数，避免 useCallback 闭包自引用 + const renderNode = (node: SourceTreeNode) => { + const isExpanded = expandedNodes.has(node.id); + const hasChildren = node.children && node.children.length > 0; - const allChecked = node.leafKeys.every((key) => - effectiveSelected.has(key) - ); - const someChecked = node.leafKeys.some((key) => - effectiveSelected.has(key) - ); - const isIndeterminate = someChecked && !allChecked; + const allChecked = node.leafKeys.every((key) => + effectiveSelected.has(key) + ); + const someChecked = node.leafKeys.some((key) => + effectiveSelected.has(key) + ); + const isIndeterminate = someChecked && !allChecked; - const handleCheckboxChange = (e: React.MouseEvent) => { - e.stopPropagation(); - setSelectedSources((prev) => { - const next = new Set(prev); - if (allChecked) { - node.leafKeys.forEach((k) => next.delete(k)); - } else { - node.leafKeys.forEach((k) => next.add(k)); + const handleCheckboxChange = (e: React.MouseEvent) => { + e.stopPropagation(); + setSelectedSources((prev) => { + const next = new Set(prev); + if (allChecked) { + node.leafKeys.forEach((k) => next.delete(k)); + } else { + node.leafKeys.forEach((k) => next.add(k)); + } + return next; + }); + }; + + const handleNodeClick = (e: React.MouseEvent) => { + e.stopPropagation(); + if (hasChildren) { + toggleExpand(node.id); } - return next; - }); - }; + }; - const handleNodeClick = (e: React.MouseEvent) => { - e.stopPropagation(); - if (hasChildren) { - toggleExpand(node.id); - } - }; + return ( +
+
+
+ {/* 子折叠箭头 */} + {hasChildren ? ( + + ▶ + + ) : ( + + )} - return ( -
-
-
- {/* 子折叠箭头 */} - {hasChildren ? ( - - ▶ + {allChecked && ( + + + + )} + {isIndeterminate && ( +
+ )} +
+ + {/* 子节点标签 */} + + {node.label || '默认'} - ) : ( - - )} - - {/* 精美自定义复选框 */} -
- {allChecked && ( - - - - )} - {isIndeterminate && ( -
- )}
- - {/* 子节点标签 */} - - {node.label || '默认'} -
+ + {/* 虚线连接子孙 */} + {hasChildren && isExpanded && ( +
+ {node.children!.map((child) => renderNode(child))} +
+ )}
- - {/* 虚线连接子孙 */} - {hasChildren && isExpanded && ( -
- {node.children!.map((child) => renderNestedNode(child))} -
- )} -
- ); + ); + }; + return renderNode(node); }, [expandedNodes, effectiveSelected, toggleExpand] ); diff --git a/dashboard/src/components/observation/constants.ts b/dashboard/src/components/observation/constants.ts index 80d72f8..cd6e67b 100644 --- a/dashboard/src/components/observation/constants.ts +++ b/dashboard/src/components/observation/constants.ts @@ -63,6 +63,21 @@ export const SOURCE_THEME: Record< label: 'TESS', dot: 'bg-violet-500', }, + jwst: { + badge: 'bg-pink-100 text-pink-700 dark:bg-pink-500/15 dark:text-pink-300', + label: 'JWST', + dot: 'bg-pink-500', + }, + xmm: { + badge: 'bg-teal-100 text-teal-700 dark:bg-teal-500/15 dark:text-teal-300', + label: 'XMM-Newton', + dot: 'bg-teal-500', + }, + chandra: { + badge: 'bg-red-100 text-red-700 dark:bg-red-500/15 dark:text-red-300', + label: 'Chandra', + dot: 'bg-red-500', + }, }; // ── 产品类型中文标签（product 轴）── diff --git a/dashboard/src/pages/ObservationPanel.tsx b/dashboard/src/pages/ObservationPanel.tsx index ac3d44d..43d7c73 100644 --- a/dashboard/src/pages/ObservationPanel.tsx +++ b/dashboard/src/pages/ObservationPanel.tsx @@ -813,6 +813,9 @@ function LibraryView({ { value: 'gaia', label: 'Gaia' }, { value: 'sdss', label: 'SDSS' }, { value: 'desi', label: 'DESI' }, + { value: 'jwst', label: 'JWST' }, + { value: 'xmm', label: 'XMM-Newton' }, + { value: 'chandra', label: 'Chandra' }, ]} /> diff --git a/docs/architecture/agent/dsh-alignment.md b/docs/architecture/agent/dsh-alignment.md new file mode 100644 index 0000000..c47cbea --- /dev/null +++ b/docs/architecture/agent/dsh-alignment.md @@ -0,0 +1,132 @@ +# Agent 架构硬化 — deepseek-harness 对齐改造 + +本文档记录 2026-08 的 agent 系统架构改造：以 deepseek-harness（DeepSeek +开源的 agent harness，见 `libs/deepseek-harness`）验证过的架构不变量为 +参照，对本项目 agent 子系统的四项核心改造与若干升级。所有改动以 +"取其分层思想、不取其运行时插件树"为原则——本项目是垂直应用， +Rust 的 trait + 注册表天然等价于 dsh 的 capability seam。 + +## 1. 统一 ReAct 引擎（消除四份手写循环） + +**不变量**：仓库中有且仅有一个具体的 ReAct 循环实现（对应 dsh 的 +"agent-loop 是 harness 中唯一包含具体循环逻辑的包"）。 + +- `src/agent/engine.rs`：`ReactEngine` 承载完整循环（取消检查 → 压缩 → + 预算/nag → 步数上限 → 后台通知 → LLM 调用含恢复阶梯 → 工具执行 → 循环）。 +- 行为差异通过组合字段注入而非分叉实现： + +| 关注点 | lead（主代理） | subagent | teammate | +|---|---|---|---| +| SSE 事件 | `EventTap::channel` | `EventTap::prefixed(tx, "sub")` | `EventTap::none()` | +| 持久化 | `DbMessageSink`（agent_name=lead） | `DbMessageSink`（agent_name=sub_xxx） | None | +| AskUser 权限 | `AskPolicy::Interactive`（120s 挂起） | `AskPolicy::AutoDeny` | `AskPolicy::AutoDeny` | +| 取消 | `CancelSource::Session`（DashMap） | 不可取消 | `CancelSource::Flag` | +| 恢复阶梯/熔断/重复检测 | 共享（runtime 级 Arc） | 引擎局部实例 | 引擎局部实例 | + +- 队友因此获得了与主代理对等的能力（错误恢复、压缩熔断、多槽重复检测）。 +- `streaming_executor.rs`（936 行零调用死代码）已删除。 + +## 2. 会话级运行时生命周期 + +**不变量**：runtime 状态属于会话而非请求。 + +`src/agent/runtime/session_registry.rs` 的 `SessionRuntimeRegistry` +（挂在 `AppState.agent_runtimes`）： + +- 同一会话的请求复用同一 `AgentRuntime`——后台任务队列、压缩折叠日志、 + 文件缓存、拒绝追踪器、压缩熔断器、prompt cache 不再随 HTTP 请求销毁 + （历史上"后台任务结果跨请求丢失、每 turn 重复压缩"两个 bug 的根因）。 +- 每会话一个 turn 互斥锁：并发请求 fail-loud 返回 409。 +- 空闲 2 小时的条目由后台清扫回收；会话删除时显式移除。 +- 会话 ID 在请求入口预分配（`create_or_resume_session_preallocated`）， + 新会话的首个请求也能写入取消标记（修复 SSE 超时只 abort 不落标记）。 +- 模式回放：runtime 首次创建时从 DB 读取会话模式，后续请求不再因传参 + 不同而静默切换行为。 +- 压缩递归守卫从进程级 `AtomicBool` 改为按 session_id 的集合守卫 + （`try_begin_compaction`/`end_compaction`），并发会话压缩互不干扰。 + +## 3. 会话事件日志（事件溯源的生命周期侧面） + +**不变量**：模型可见 ⟺ 已日志化；turn/compaction 以日志化锁开闭。 + +新表 `agent_events`（migration 20260819000001）+ +`src/agent/runtime/session_events.rs`： + +- `turn_start`/`turn_end`（含结构化终止原因 `reason_label`）、 + `compaction_start`/`compaction_end` 构成日志化锁。 +- 崩溃恢复不截断：`reconcile_interrupted` 检测开口的 start 事件并合成 + `interrupted` 关闭（对应 dsh "不伪造完成、但补写中断事实"的策略）。 +- `context_snapshot`：turn 内发生过压缩时，turn 结束保存折叠后上下文 + + 消息高水位 `base_message_id`；下一 turn 回放"快照 + id > base 的增量" + （`context.rs::load_folded_history`），不再从原始消息重建后重新压缩。 + rewind/retry/branch 后快照作废（`remove_context_snapshots`）。 + +## 4. 工具契约集中化（消除散弹式修改） + +**不变量**：工具的全部行为元数据声明在工具自身（AgentTool），消费方查询 +trait 而非维护平行的名字名单。 + +`AgentTool` 新增声明方法： + +- `untrusted_output()` — 输出来自外部源需 `` 包裹 + （原 `untrusted.rs` 按名列表）；注册表查询不到工具时回退名字启发式。 +- `hardline_check(args)` — 参数级不可绕过拒绝（原 executor 按名路由）。 +- `causes_file_changes()` — 执行前触发 checkpoint 快照（原 + `CHECKPOINT_TRIGGER_TOOLS` 名单）。 +- `loop_signals() -> ToolSignals` — 循环行为信号（todo nag 重置/todos + 持久化/手动压缩请求），替代主循环对 `todo_write`/`compress_context` + 的按名特判。 +- 输出契约：`ToolOutput.value`（canonical JSON，机器消费：重放/审计/剪枝） + + `content`（模型侧 render 投影），value 随消息 metadata 持久化。 + 样例：`search_papers`、`todo_write`。 +- 注册表测试从硬编码数量断言（`assert_eq!(defs.len(), 35)`）改为 + inventory 不变量断言（名字唯一/元数据完整/信号与身份一致）， + 新增工具不再需要改测试。 + +## 5. 权限单调性（deny > ask > allow，只收紧不放松） + +**不变量**：多源决策合并唯一入口 `permission::tighten(base, candidate)`： + +- Deny 粘滞——任何层不能放行另一层已 Deny 的调用； +- Ask 不可被放松为 Allow，可升级为 Deny； +- executor 的 hook 请求、工具级声明、会话级规则全部经 `tighten` 合并； +- 无人值守上下文（子代理/队友）的 AskUser fail-closed 自动拒绝 + （`AskPolicy::AutoDeny` + `engine::auto_deny_ask`）。 + +历史上与此语义重复且从未接线的 `resolve_permission_precedence` 已删除。 + +## 6. 上下文管理升级 + +- **KV-cache 纪律**：system prompt 只包含字节稳定 section；技能清单/ + 项目记忆以 durable 动态上下文快照（user 消息）追加在历史尾部，哈希 + 未变不重注入（`context.rs::build_initial_context` + + `AgentRuntime::build_dynamic_context`）。 +- **记忆智能选择接线**：记忆条目 > 8 时走 `select_relevant_memories` + （LLM 结构化选择 + 指数时间衰减，此前为零调用代码）。 +- **确定性剪枝**：`micro_compact` 对长工具结果保留头/尾预览 + （`prune_tool_result`），短结果仍为纯占位符。 +- **Spill 检索指引**：溢出 stub 携带工具名 + 明确取回路径（read_file + + max_lines 分段）。 +- **Token 估算统一**：全循环使用 `compact::estimate_message_tokens` + （CJK 加权），消除主循环 `len/3` 与 compact 两套口径。 + +## 7. 其他接线与修复 + +- **团队协作接线**：`spawn_teammate` 等 4 个工具经 + `ToolRegistry::new_with_team` 注册进会话 runtime（TeamManager 按 + session_key 初始化邮箱目录）；队友注册表真正排除 subagent/spawn_teammate + （原注释声称排除但未过滤）。 +- **Continuable 子代理**：`subagent` 工具新增 `agent_name` + `followup` + 参数——命名子代理历史持久化，followup 回放全部历史续话（dsh + Activation 模型的冷恢复路径）。 +- **IdlePoller 接线**：`AGENT_AUTONOMOUS_ENABLED=true` 时启动自治轮询， + 经会话注册表获取 runtime（模式从 DB 回放）。 +- **重复检测多槽**：`DuplicateDetector` 连续通道 + 滑动窗口双通道， + 可检测 A/B 交替死循环（原单槽实现检测不到）。 + +## 不采用的设计（与理由） + +- Cordis 运行时插件树 / bundle/patch/profile 配置组合层：服务于通用 + harness 的第三方插件生态，本项目是单二进制垂直应用。 +- Typert 类型图 RPC、host/client 双聚合、Python SDK：无对应需求。 +- 适配器休眠挂载、profiles 热重载：模型链固定。 diff --git a/docs/architecture/agent/env-vars.md b/docs/architecture/agent/env-vars.md index 948db45..8db46df 100644 --- a/docs/architecture/agent/env-vars.md +++ b/docs/architecture/agent/env-vars.md @@ -134,3 +134,4 @@ Agent 系统的所有可配置参数，按子系统分类。 | `src/agent/compact.rs` (extract_memories_from_compaction) | `EXTRACT_MEMORY_ENABLED` (P3 压缩记忆桥接) | | `src/agent/memory/extraction.rs` (ExtractionConfig) | `EXTRACT_MEMORY_ENABLED`, `EXTRACT_MEMORY_THROTTLE_TURNS`, `EXTRACT_MEMORY_MAX_STEPS` | | `src/services/download.rs` | `OBSCURA_ALLOW_PRIVATE_NETWORK`（程序自动设置） | +| `src/main.rs` (IdlePoller 启动门控) | `AGENT_AUTONOMOUS_ENABLED`（默认 false；true/1 时启用自治轮询，见 dsh-alignment.md §7） | diff --git a/docs/architecture/agent/overview.md b/docs/architecture/agent/overview.md index 86403c4..48785bf 100644 --- a/docs/architecture/agent/overview.md +++ b/docs/architecture/agent/overview.md @@ -2,6 +2,12 @@ AstroResearch 内置了一个基于 **ReAct** (Thought → Action → Observation) 范式的科研智能体引擎 (`src/agent/`)，参考 Claude Code 的分层设计。以下对各子系统的架构、数据流和内部逻辑进行完整说明。 +> **架构演进**：2026-08 起，ReAct 循环的唯一实现在 `src/agent/engine.rs` +> （`ReactEngine`，主代理/子代理/队友以不同注入组合复用）； +> `AgentRuntime` 按会话缓存复用（`SessionRuntimeRegistry`）； +> turn/compaction 生命周期与上下文快照落在 `agent_events` 事件日志。 +> 详见 [dsh-alignment.md](dsh-alignment.md)。 + ### 整体架构 ```mermaid @@ -14,16 +20,18 @@ graph TD AskUser["问答 /api/chat/questions + /api/chat/answer"] end - subgraph Runtime["AgentRuntime — ReAct 引擎"] + subgraph Runtime["AgentRuntime — 会话级编排"] RunTurn["run_turn() 主入口"] SP["SystemPrompt 组装器"] CtxBuild["Context Builder 上下文构建"] - ReAct["ReAct 主循环"] + Engine["ReactEngine 统一循环（engine.rs）"] Streaming["streaming.rs 流式处理"] Executor["executor.rs 并行执行"] Finalize["finalize.rs 会话收尾"] TokenBudget["token_budget.rs"] CircuitBreaker["circuit_breaker.rs"] + SessionReg["SessionRuntimeRegistry 会话级缓存"] + EventLog["agent_events 事件日志"] end subgraph Tools["工具系统 (tools/)"] diff --git a/docs/roadmap-research-features.md b/docs/roadmap-research-features.md new file mode 100644 index 0000000..d620c67 --- /dev/null +++ b/docs/roadmap-research-features.md @@ -0,0 +1,438 @@ +# AstroResearch 科研功能扩展路线图 + +## 项目现状 + +### 已实现的核心能力 + +**数据源（9 个）**： + +| 数据源 | 产品类型 | 状态 | +| ---------- | ------------------------------------------------ | ---- | +| LAMOST | 光谱（LRS/MRS） | 完整 | +| Gaia | XP 光谱（BP/RP）、RVS、历元测光 | 完整 | +| SDSS | 光学光谱（spec）、APOGEE 近红外（apstar/aspcap） | 完整 | +| DESI | 暗能量光谱仪（coadd） | 完整 | +| 2MASS | 近红外测光（J/H/K） | 完整 | +| AllWISE | 中红外测光（W1-W4） | 完整 | +| Pan-STARRS | 光学测光（g/r/i/z/y） | 完整 | +| ZTF | 时域光变曲线（zg/zr/zi） | 完整 | +| TESS | 空间时域光变曲线（2min/FFI） | 完整 | + +**Agent 工具集**： + +- `search_papers` — ADS/arXiv 文献搜索 +- `process_paper` — 下载、解析、向量化、翻译 +- `get_paper_content` / `get_paper_outline` — 阅读解析后文献 +- `rag_search` — 稠密向量 + BM25 混合检索 +- `query_target` — CDS Sesame 天体名称解析 +- `find_observation` — 跨源观测数据统一下载 +- `catalog_operation` — VizieR TAP 星表查询（ADQL + 锥形检索） +- `save_note` — 高亮/笔记存储 +- `get_paper_metadata` — 完整文献元数据 +- `search_local_library` — 本地文献库搜索 +- `get_citation_network` — 引文网络图 +- `analyze_image` — 视觉模型图像分析 +- 文件操作、bash、记忆、技能、子 Agent、团队协作等 + +**已有技能**： + +- `methodology` — 系统性文献综述方法论 +- `plotting` — 科研绘图规范（占位，未实现） +- `presentation` — 演示文稿技能 + +--- + +## 一、数据获取层 — 补全缺失的数据源 + +### 1.1 JWST/MIRI 数据获取 + +**目标**：通过 MAST API 获取詹姆斯·韦伯空间望远镜 NIRSpec/MIRI 光谱数据。 + +**需求**： + +- 新增 `src/clients/mast/jwst.rs`，封装 MAST JWST 查询 API +- 实现 `JwstSpectrumFetcher`，注册到 `ObservationRegistry` +- Source 枚举新增 `Jwst` 变体 +- 支持 NIRSpec IFU/MOS、MIRI MRS/LRS 光谱产品 +- FITS 解析支持 JWST 特有的多扩展 BinTable 格式 + +**价值**：JWST 是当前最热门的望远镜，大量前沿论文依赖 JWST 数据。 + +**预估工作量**：中等（MAST API 文档完善，但 FITS 结构复杂）。 + +### 1.2 X 射线数据（XMM-Newton / Chandra） + +**目标**：通过 HEASARC TAP 获取高能天体物理观测数据。 + +**需求**： + +- 新增 `src/clients/heasarc/`，封装 HEASARC TAP 查询 +- 实现 `XmmSpectrumFetcher` 和 `ChandraSpectrumFetcher` +- Source 枚举新增 `Xmm` 和 `Chandra` 变体 +- 支持 X 射线光谱（PHA）和光变曲线产品 + +**价值**：覆盖高能波段，满足活动星系核、X 射线双星、星系团等研究需求。 + +**预估工作量**：中等（HEASARC TAP 接口标准，但数据格式多样）。 + +### 1.3 VLBI/Gaia 天体测量数据 + +**目标**：通过 VizieR 获取高精度自行、视差和天体测量数据。 + +**需求**： + +- 扩展现有 `CatalogOperationTool`，预置 Gaia DR3 天体测量常用查询模板 +- 新增 `GaiaAstrometryFetcher`，专门处理高精度自行/视差数据 +- 支持按自行/视差筛选的 cone search + +**价值**：精密天体测量是银河考古学和恒星动力学的基础。 + +**预估工作量**：小（VizieR 已支持，主要是模板和预置查询）。 + +### 1.4 光谱图像 Cutout + +**目标**：实现 SDSS/Gaia/DSS 图像 cutout 下载 + 前端渲染。 + +**需求**： + +- 实现 `ImageCutoutFetcher`，支持 SDSS SkyServer cutout API、DSS (STScI)、Gaia 推荐图 +- 完成 `ImagePreview` 解析器（preview.rs 中已有占位） +- 前端 `ObservationPreviewRenderer` 新增图像渲染 case +- 支持自定义 cutout 大小和坐标 + +**价值**：观测规划和文献中的发现图（Finding Chart）必备。 + +**预估工作量**：中等（API 简单，但前端图像渲染需要新组件）。 + +--- + +## 二、数据分析层 — 光谱/光变科学分析 + +### 2.1 光谱参数自动提取 + +**目标**：从 LAMOST/SDSS FITS 文件中自动读取恒星大气参数（Teff, logg, [Fe/H], 径向速度等），结构化入库。 + +**需求**： + +- 新增 `src/services/spectrum/parameters.rs`，实现 FITS Header/BinTable 参数提取 +- LAMOST：从 `params.fits` 或 `spec.fits` header 读取 `TEFF`, `LOGG`, `FEH`, `RV` +- SDSS/BOSS：从 `line.dat` 或 `aspcapStar` 读取 ASPCAP 参数 +- SDSS/APOGEE：从 `aspcapStar` FITS 读取完整化学丰度 +- DESEI：从 `coadd` FITS header 读取红移和分类 +- 结构化数据入库（新建 `spectrum_parameters` 表或扩展 `observation_cache`） +- Agent 新增 `get_spectrum_parameters` 工具，查询已提取的参数 + +**价值**：每篇恒星物理论文都需要这类参数，手动提取极其低效。 + +**预估工作量**：中等（FITS 解析已有基础，需逐源适配 header 约定）。 + +### 2.2 光谱线识别与等值宽度测量 + +**目标**：标注 Balmer/H&K/金属线位置，自动计算等值宽度（EW）和半高全宽（FWHM）。 + +**需求**： + +- 新增 `src/services/spectrum/lines.rs`，内置常见谱线波长表（真空/空气波长） +- 实现谱线搜索算法：在给定波长窗口内找局部极值（发射线找极大，吸收线找极小） +- 实现数值积分 EW 计算（梯形法） +- 支持用户自定义谱线列表 +- Agent 新增 `measure_spectral_lines` 工具，输入光谱路径 + 谱线名称列表，输出 EW/FWHM +- 前端光谱图支持谱线标注叠加层 + +**价值**：光谱分析的核心操作，天文学生和研究者日常高频使用。 + +**预估工作量**：中等偏高（谱线表容易，但稳健的峰值搜索和 EW 计算需要仔细实现）。 + +### 2.3 交叉相关测速 + +**目标**：通过模板匹配法测量径向速度（RV），支持多波段光谱。 + +**需求**： + +- 新增 `src/services/spectrum/cross_correlate.rs` +- 内置或下载标准模板光谱（如 Pickles Atlas、PHOENIX 模型） +- 实现 FFT 加速的交叉相关函数计算 +- 从互相关峰值位置提取 RV + 不确定度估计 +- 支持用户上传自定义模板 +- Agent 新增 `measure_radial_velocity` 工具 + +**价值**：双星、视向速度变星、银河系动力学研究的核心工具。 + +**预估工作量**：中等（算法成熟，但模板管理和多波段适配需要设计）。 + +### 2.4 光变周期搜索 + +**目标**：基于 Lomb-Scargle 周期图和 BLS（Box Least Squares）算法，自动搜索光变周期。 + +**需求**： + +- 新增 `src/services/timeseries/periodicity.rs` +- 实现 Lomb-Scargle 周期图（非均匀采样友好） +- 实现 BLS 算法（凌星检测专用） +- 自动报告最优周期、FAP（误报概率）、次优周期 +- 支持 ZTF、TESS、Gaia 历元测光数据 +- Agent 新增 `find_period` 工具，输入光变曲线文件，输出周期 + 周期图 +- 前端支持折叠光变曲线（Phase-folded light curve）展示 + +**价值**：ZTF/TESS 数据的刚需，变星和系外行星研究的核心分析。 + +**预估工作量**：中等（Lomb-Scargle 算法有 Rust 实现可参考，BLS 较复杂）。 + +### 2.5 光变分类与特征提取 + +**目标**：自动识别变星类型，提取振幅、形态参数等特征。 + +**需求**： + +- 新增 `src/services/timeseries/classification.rs` +- 基于周期 + 振幅 + 偏度 + 峰度等统计特征的简单分类器 +- 支持 RR Lyrae、Cepheid、食双星、脉动变星、AGN 等主要类型 +- 可选：集成轻量 ML 模型（如 Random Forest，通过 ONNX Runtime） +- Agent 新增 `classify_variable_star` 工具 + +**价值**：时域巡天（ZTF/TESS）数据量巨大，自动化分类是必要环节。 + +**预估工作量**：高（分类器训练和特征工程需要较多工作）。 + +### 2.6 SED 拟合与建模 + +**目标**：多波段测光数据拟合黑体或模板光谱能量分布，求解天体物理参数。 + +**需求**： + +- 新增 `src/services/analysis/sed.rs` +- 自动拉取目标所有可用波段测光（2MASS + WISE + PS1 + Gaia） +- 实现黑体辐射模型拟合（最小二乘/MCMC） +- 可选模板 SED 拟合（BaSeL/PHOENIX 模板库） +- 输出 T_eff、光度、消光 A_V 及其不确定度 +- 前端 SED 图展示（观测点 + 拟合曲线 + 残差） +- Agent 新增 `fit_sed` 工具 + +**价值**：天体物理参数估计的基础方法，适用于恒星、AGN、超新星等多种天体。 + +**预估工作量**：中等偏高（黑体拟合简单，但模板拟合和 MCMC 需要较多工作）。 + +### 2.7 自行/视差运动学分析 + +**目标**：基于 Gaia DR3 高精度自行 + 视差，计算银河系三维空间速度（UVW），识别移动星群。 + +**需求**： + +- 新增 `src/services/analysis/kinematics.rs` +- 实现坐标变换：(RA, Dec, parallax, pmRA, pmDec, RV) → (X, Y, Z, U, V, W) +- 实现移动星群成员概率计算（基于 Banyan Σ 算法简化版） +- 支持批量样本处理 +- Agent 新增 `analyze_kinematics` 工具 + +**价值**：银河考古学（Galactic Archaeology）的核心分析方法。 + +**预估工作量**：中等（坐标变换公式成熟，Banyan Σ 需要简化实现）。 + +### 2.8 化学丰度模式分析 + +**目标**：多元素丰度比（[α/Fe] vs [Fe/H]）自动提取与可视化，区分厚盘/薄盘/晕星。 + +**需求**： + +- 扩展光谱参数提取（2.1），从 APOGEE ASPCAP 数据读取完整化学丰度 +- 新增 `src/services/analysis/chemistry.rs` +- 实现 [α/Fe] 计算（α 元素：O, Mg, Si, S, Ca, Ti 的加权平均） +- 自动绘制 [α/Fe] vs [Fe/H] 图，并按运动学成分着色 +- Agent 新增 `analyze_abundance_pattern` 工具 + +**价值**：银河考古学的核心工具，用于区分银河系不同星族。 + +**预估工作量**：中等（数据提取为主，分析逻辑相对直接）。 + +--- + +## 三、可视化与输出层 + +### 3.1 光谱交互式标注 + +**目标**：前端光谱图支持点击标注谱线名称，自动显示等值宽度。 + +**需求**： + +- 扩展 `SpectrumPlot.tsx`，新增谱线标注叠加层 +- 内置常见谱线波长表（Balmer 系列、H&K、NaD、Mg b、Fe 等） +- 支持用户点击光谱图任意位置添加自定义标注 +- 标注信息（谱线名、波长、EW）以 tooltip 形式展示 +- 后端 `preview.rs` 返回时可选附带谱线标注数据 + +**价值**：显著提升光谱阅读和教学体验。 + +**预估工作量**：中等（前端组件开发为主）。 + +### 3.2 赫罗图（HR Diagram） + +**目标**：基于 Gaia 数据构建颜色-星等图（CMD），支持叠加自定义样本。 + +**需求**： + +- 新增 `src/api/hertzsprung_russell.rs` 或扩展现有 Gaia 查询 +- 自动查询给定天区/样本的 Gaia G-mag、BP-RP color、parallax +- 计算绝对星等 M_G = G + 5 + 5*log10(parallax/1000) +- 返回数据供前端渲染 HR 图 +- 前端新增 HR Diagram 组件（scatter plot，支持密度着色） +- 支持叠加用户自定义样本（不同颜色/标记） + +**价值**：恒星物理的基础可视化，几乎所有恒星相关论文都包含 HR 图。 + +**预估工作量**：中等（数据查询简单，前端绘图需要新组件）。 + +### 3.3 发现图（Finding Chart） + +**目标**：生成目标天区的 DSS/Pan-STARRS 图像 cutout + 标注目标位置。 + +**需求**： + +- 调用 STScI DSS cutout API 或 Pan-STARRS cutout API +- 返回 cutout 图像 + 目标中心标注 +- 支持自定义 cutout 大小（arcmin）和缩放 +- Agent 工具 `find_observation` 可选返回 finding chart +- 前端在观测数据页面展示 finding chart + +**价值**：观测准备的必备环节，望远镜观测提案中需要附带 finding chart。 + +**预估工作量**：小（API 简单，主要是前端展示）。 + +### 3.4 多波段 SED 图 + +**目标**：自动拉取所有可用波段测光数据，绘制 SED 图并拟合模型。 + +**需求**： + +- 扩展 2.6 节 SED 拟合功能，新增前端 SED 图组件 +- 自动标记各波段（2MASS J/H/K、WISE W1-W4、Gaia G/BP/RP 等） +- 支持对数坐标轴 +- 拟合曲线叠加显示，残差面板 + +**价值**：参数化分析的标准可视化。 + +**预估工作量**：中等（依赖 2.6 节后端实现）。 + +### 3.5 SkyMap 天球投影 + +**目标**：在全天球图上可视化样本分布，支持按参数着色。 + +**需求**： + +- 新增 `SkyMapPanel` 前端组件 +- 使用 Mollweide 或 Aitoff 投影 +- 支持按天体类型、金属丰度、自行等参数着色 +- 叠加银河平面、银道坐标网格 +- 支持交互式缩放和选区 + +**价值**：巡天论文的标准可视化，用于展示样本的天球分布。 + +**预估工作量**：中等（前端投影计算 + WebGL 渲染）。 + +--- + +## 四、工作流增强 + +### 4.1 批量样本处理 + +**目标**：给定一组天体名称/坐标，批量执行"查询 → 下载 → 分析 → 出图 → 保存报告"。 + +**需求**： + +- 新增 `src/agent/tools/astro/research/batch_process.rs`，实现 `BatchProcessTool` +- 支持输入格式：天体名称列表、坐标列表、VizieR 查询结果 +- 批量流程：query_target → find_observation → 参数提取 → 出图 +- 并发控制（避免触发数据源速率限制） +- 进度报告（SSE 事件推送） +- 输出汇总报告（Markdown 表格 + 图表） + +**价值**：巡天论文的标准工作流，手动逐个处理效率极低。 + +**预估工作量**：中等偏高（并发控制和错误处理需要仔细设计）。 + +### 4.2 定时文献监控 + +**目标**：Cron 定时检查 ADS/arXiv 新文献，匹配用户关注的关键词/作者/天体。 + +**需求**： + +- 新增 `src/agent/tools/astro/research/literature_monitor.rs` +- 用户定义监控规则（关键词、作者、天体名称、arXiv 分类） +- 使用现有 `loop` 技能定时触发 +- 新文献匹配时通过 SSE 推送通知 +- 支持自动下载摘要 + 生成简报 + +**价值**：帮助研究者跟踪前沿，不遗漏重要文献。 + +**预估工作量**：中等（定时任务框架已有，主要是规则匹配和通知）。 + +### 4.3 观测提案辅助 + +**目标**：根据目标天区自动计算可见性、月相限制、曝光时间估算。 + +**需求**： + +- 新增 `src/services/analysis/observability.rs` +- 计算目标在给定台站的升起/落下时间和最佳观测窗口 +- 月相和银河背景亮度估算 +- 基于目标亮度和望远镜参数的曝光时间估算 +- Agent 新增 `plan_observation` 工具 +- 输出观测提案所需的标准化信息 + +**价值**：望远镜申请的时间评估和观测规划必备。 + +**预估工作量**：中等（天文计算公式成熟，但台站参数和大气消光模型需要数据）。 + +--- + +## 五、实施优先级建议 + +| 优先级 | 功能 | 理由 | 状态 | +| ------ | --------------------------- | -------------------------------------- | ---- | +| P0 | 光谱参数自动提取（2.1） | 已下载数据的直接利用，所有恒星论文都用 | ✅ 已实现 | +| P0 | 光变周期搜索（2.4） | ZTF/TESS 数据的刚需，时域天文学核心 | ✅ 已实现 | +| P1 | 光谱线识别与 EW 测量（2.2） | 光谱分析的日常操作 | ✅ 已实现 | +| P1 | 交叉相关测速（2.3） | 双星/视向速度研究核心工具 | ✅ 已实现 | +| P1 | 绘图 skill 填充（现有占位） | 提升 Agent 绘图能力的基础 | ✅ 已实现 | +| P2 | SED 拟合（2.6） | 多波段分析的标准方法 | ✅ 已实现 | +| P2 | 赫罗图（3.2） | 恒星物理基础可视化 | ✅ 已实现 | +| P2 | 批量样本处理（4.1） | 巡天论文工作流提效 | ✅ 已实现 | +| P3 | JWST 数据获取（1.1） | 前沿望远镜，但使用门槛较高 | ✅ 已实现 | +| P3 | 光变分类（2.5） | 时域数据分析，但分类器训练复杂 | ✅ 已实现 | +| P3 | SkyMap 天球投影（3.5） | 巡天论文可视化，但优先级不如核心分析 | ✅ 已实现 | + +--- + +## 六、已实现功能详情 + +### 已完成的数据分析层 + +| 功能 | Agent 工具 | 后端模块 | 状态 | +|------|-----------|----------|------| +| 光谱参数提取 | `get_spectrum_parameters` | `services/spectrum/parameters.rs` | ✅ | +| 谱线识别与EW测量 | `measure_spectral_lines` | `services/spectrum/lines.rs` | ✅ | +| 交叉相关测速 | `measure_radial_velocity` | `services/spectrum/cross_correlate.rs` | ✅ | +| 周期搜索 | `find_period` | `services/timeseries/periodicity.rs` | ✅ | +| 变星分类 | `classify_variable_star` | `services/timeseries/classification.rs` | ✅ | +| SED拟合 | `fit_sed` | `services/analysis/sed.rs` | ✅ | +| 运动学分析 | `analyze_kinematics` | `services/analysis/kinematics.rs` | ✅ | +| 化学丰度分析 | `analyze_abundance_pattern` | `services/analysis/chemistry.rs` | ✅ | +| 观测提案辅助 | `plan_observation` | `services/analysis/observability.rs` | ✅ | +| 批量样本处理 | `batch_process` | `agent/tools/astro/research/batch_process.rs` | ✅ | +| 文献监控 | `literature_monitor` | `agent/tools/astro/research/literature_monitor.rs` | ✅ | + +### 已完成的数据获取层 + +| 功能 | 模块 | 说明 | 状态 | +|------|------|------|------| +| JWST 数据获取 | `clients/mast/jwst.rs` + `services/observation/jwst.rs` | MAST JWST API + SpectrumFetcher | ✅ | +| X射线数据 | `clients/heasarc/` + `services/observation/xray.rs` | HEASARC TAP + XMM/Chandra fetcher | ✅ | +| 光谱图像Cutout | `services/observation/cutout.rs` | SDSS/DSS/Pan-STARRS cutout API + CutoutFetcher | ✅ | +| 赫罗图 | `api/hr_diagram.rs` | Gaia CMD 查询 API | ✅ | +| 绘图skill | `skills/plotting/SKILL.md` | Python matplotlib 科研绘图模板（8种图表） | ✅ | + +### 已完成的工作流增强 + +| 功能 | 模块 | 说明 | 状态 | +|------|------|------|------| +| 批量样本处理 | `agent/tools/astro/research/batch_process.rs` | 查询→下载→分析→报告 | ✅ | +| 文献监控 | `agent/tools/astro/research/literature_monitor.rs` | ADS/arXiv 最新文献匹配 | ✅ | diff --git a/migrations/20260713000000_add_is_scheduled_to_sync_queries.sql b/migrations/20260713000000_add_is_scheduled_to_sync_queries.sql new file mode 100644 index 0000000..2ddc59c --- /dev/null +++ b/migrations/20260713000000_add_is_scheduled_to_sync_queries.sql @@ -0,0 +1,2 @@ +-- Add is_scheduled column to sync_queries for periodic incremental sync scheduling +ALTER TABLE sync_queries ADD COLUMN is_scheduled INTEGER NOT NULL DEFAULT 0; diff --git a/migrations/20260819000001_add_agent_events.sql b/migrations/20260819000001_add_agent_events.sql new file mode 100644 index 0000000..5887654 --- /dev/null +++ b/migrations/20260819000001_add_agent_events.sql @@ -0,0 +1,25 @@ +-- 会话事件日志：turn/compaction 生命周期 + 上下文快照（事件溯源中枢） +-- +-- 设计（参考 deepseek-harness 的事件溯源会话模型）： +-- agent_messages 记录"模型可见内容"，agent_events 记录"生命周期事实"。 +-- turn_start/turn_end 构成日志化锁：崩溃后开着的 turn_start 可被检测并 +-- 合成 interrupted 关闭（不截断已持久化的部分）。 +-- context_snapshot 存储压缩后的折叠上下文，下一 turn 直接回放而无需 +-- 从原始消息重建再重新压缩（节省重复的 LLM 摘要调用）。 +-- +-- 事件类型（event_type）： +-- turn_start / turn_end / turn_interrupted +-- compaction_start / compaction_end / compaction_interrupted +-- context_snapshot + +CREATE TABLE IF NOT EXISTS agent_events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + turn_index INTEGER NOT NULL DEFAULT 0, + event_type TEXT NOT NULL, + payload TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS idx_agent_events_session + ON agent_events(session_id, id); diff --git a/skills/plotting/SKILL.md b/skills/plotting/SKILL.md index 4e590fe..961e02b 100644 --- a/skills/plotting/SKILL.md +++ b/skills/plotting/SKILL.md @@ -1,36 +1,301 @@ --- name: plotting -description: 科研绘图规范 +description: 科研绘图规范 —— 天文学常用图表的 Python matplotlib 绘制模板 context: fork allowed-tools: - bash - save_note + - write + - read --- -# 科研绘图规范（TODO：待细化） +# 科研绘图规范 -## 目标期刊要求（占位） +## 环境要求 -- 分辨率：通常 300-600 dpi -- 格式：矢量图优先（PDF/SVG），光栅图备选（PNG） -- 字体：需与期刊正文一致（如 Times New Roman） +- Python 3.8+ +- matplotlib >= 3.7 +- numpy +- 可选：seaborn（统计图）、astropy（天文学单位和坐标） -## 常用图表类型（占位） +安装： +```bash +pip install matplotlib numpy seaborn astropy --quiet +``` -| 类型 | Python 库 | 适用场景 | -|------|----------|---------| -| 光谱图 | matplotlib | 光谱分析、能谱 | -| 赫罗图 | matplotlib | 恒星演化 | -| 光变曲线 | matplotlib | 变星、超新星 | -| 参数分布 | seaborn | 统计分布、角图 | -| 3D 轨道 | plotly | 天体力学 | +## 输出规范 -## 工作流程（TODO） +| 目标 | 分辨率 | 格式 | 说明 | +|------|--------|------|------| +| 论文投稿 | 300+ dpi | PDF/SVG（矢量优先） | APS/A&A/MNRAS 标准 | +| 演示文稿 | 150 dpi | PNG | 宽度 ≥ 1200px | +| 快速预览 | 100 dpi | PNG | 屏幕查看 | -1. 准备数据（从文献或天体数据库提取） -2. 生成 Python 脚本 -3. 使用 bash 工具执行脚本 -4. 检查输出图片 -5. 保存到指定路径 +## 通用样式模板 -**注意：此 skill 为预留占位，具体实现待补充。** +```python +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +import numpy as np + +# 天文学论文标准样式 +plt.rcParams.update({ + 'font.family': 'serif', + 'font.size': 12, + 'axes.labelsize': 14, + 'axes.titlesize': 14, + 'xtick.labelsize': 11, + 'ytick.labelsize': 11, + 'legend.fontsize': 11, + 'figure.figsize': (8, 6), + 'figure.dpi': 300, + 'savefig.dpi': 300, + 'savefig.bbox': 'tight', + 'savefig.pad_inches': 0.05, + 'lines.linewidth': 1.5, + 'axes.linewidth': 1.0, + 'xtick.major.width': 0.8, + 'ytick.major.width': 0.8, + 'xtick.minor.width': 0.5, + 'ytick.minor.width': 0.5, + 'xtick.direction': 'in', + 'ytick.direction': 'in', + 'xtick.top': True, + 'ytick.right': True, +}) +``` + +## 图表模板 + +### 1. 光谱图 + +```python +def plot_spectrum(wavelength, flux, title="Spectrum", xlabel=r"Wavelength ($\AA$)", ylabel=r"Flux (erg/s/cm$^2$/$\AA$)", save_path=None): + fig, ax = plt.subplots() + ax.plot(wavelength, flux, 'k-', linewidth=0.8) + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 2. 光变曲线 + +```python +def plot_light_curve(time, flux, time_err=None, flux_err=None, title="Light Curve", xlabel="Time (MJD)", ylabel="Flux", save_path=None): + fig, ax = plt.subplots() + if flux_err is not None: + ax.errorbar(time, flux, yerr=flux_err, xerr=time_err, fmt='o', markersize=3, capsize=2, color='black', ecolor='gray') + else: + ax.plot(time, flux, 'ko', markersize=3) + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 3. 折叠光变曲线 + +```python +def plot_folded_lc(phase, flux, flux_err=None, period=None, title="Folded Light Curve", save_path=None): + fig, ax = plt.subplots() + if flux_err is not None: + ax.errorbar(phase, flux, yerr=flux_err, fmt='o', markersize=2, capsize=1, color='black', ecolor='gray', alpha=0.7) + else: + ax.plot(phase, flux, 'ko', markersize=2, alpha=0.7) + ax.set_xlabel("Phase") + ax.set_ylabel("Flux") + if period: + ax.set_title(f"{title} (P = {period:.4f} d)") + else: + ax.set_title(title) + ax.set_xlim(0, 1) + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 4. 赫罗图（CMD） + +```python +def plot_hr_diagram(bp_rp, abs_g, title="HR Diagram", save_path=None, color_by_density=False): + fig, ax = plt.subplots() + if color_by_density and len(bp_rp) > 100: + from scipy.stats import gaussian_kde + xy = np.vstack([bp_rp, abs_g]) + z = gaussian_kde(xy)(xy) + idx = z.argsort() + ax.scatter(bp_rp[idx], abs_g[idx], c=z[idx], s=5, cmap='viridis_r', edgecolors='none') + cb = fig.colorbar(ax.collections[0], ax=ax, pad=0.02) + cb.set_label("Stellar density") + else: + ax.plot(bp_rp, abs_g, 'k.', markersize=1, alpha=0.5) + ax.set_xlabel(r"$G_{BP} - G_{RP}$ (mag)") + ax.set_ylabel(r"$M_G$ (mag)") + ax.set_title(title) + ax.invert_yaxis() + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 5. SED 图 + +```python +def plot_sed(wavelength_angstrom, flux_obs, flux_err=None, flux_model=None, title="SED", save_path=None): + fig, ax = plt.subplots() + ax.scatter(wavelength_angstrom, flux_obs, c='black', s=30, zorder=5, label='Observed') + if flux_err is not None: + ax.errorbar(wavelength_angstrom, flux_obs, yerr=flux_err, fmt='none', ecolor='gray', capsize=3) + if flux_model is not None: + model_wave, model_flux = zip(*flux_model) if isinstance(flux_model, list) else (flux_model[0], flux_model[1]) + ax.plot(model_wave, model_flux, 'r-', linewidth=1.5, label='Model', alpha=0.8) + ax.set_xscale('log') + ax.set_yscale('log') + ax.set_xlabel(r"Wavelength ($\AA$)") + ax.set_ylabel(r"Flux") + ax.set_title(title) + ax.legend() + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 6. [α/Fe] vs [Fe/H] 图 + +```python +def plot_abundance(feh, alpha_feh, labels=None, title="[α/Fe] vs [Fe/H]", save_path=None): + fig, ax = plt.subplots() + if labels is not None: + from matplotlib.colors import ListedColormap + colors = ['#3498db', '#e74c3c', '#2ecc71'] + cmap = ListedColormap(colors[:len(set(labels))]) + unique = sorted(set(labels)) + for i, lab in enumerate(unique): + mask = np.array(labels) == lab + ax.scatter(np.array(feh)[mask], np.array(alpha_feh)[mask], s=10, alpha=0.7, label=lab, color=colors[i % len(colors)]) + ax.legend() + else: + ax.scatter(feh, alpha_feh, s=10, alpha=0.7, c='black') + ax.set_xlabel("[Fe/H] (dex)") + ax.set_ylabel(r"[$\alpha$/Fe] (dex)") + ax.set_title(title) + ax.axhline(y=0.25, color='gray', linestyle='--', linewidth=0.8, alpha=0.5) + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 7. 周期图 + +```python +def plot_periodogram(periods, powers, best_period=None, title="Periodogram", save_path=None): + fig, ax = plt.subplots() + ax.plot(periods, powers, 'k-', linewidth=0.8) + if best_period: + ax.axvline(x=best_period, color='red', linestyle='--', linewidth=1, label=f'Best P = {best_period:.4f} d') + ax.legend() + ax.set_xlabel("Period (days)") + ax.set_ylabel("Power") + ax.set_title(title) + ax.minorticks_on() + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 8. 天球投影（Mollweide） + +```python +def plot_skymap(ra_deg, dec_deg, values=None, title="Sky Map", save_path=None): + fig, ax = plt.subplots(figsize=(10, 5), subplot_kw={'projection': 'mollweide'}) + # 转换为弧度，RA 中心在 180° + ra_rad = np.deg2rad(np.array(ra_deg)) + dec_rad = np.deg2rad(np.array(dec_deg)) + ra_rad = ra_rad - np.pi # 中心化 + if values is not None: + sc = ax.scatter(ra_rad, dec_rad, c=values, s=5, cmap='viridis', alpha=0.7, edgecolors='none') + fig.colorbar(sc, ax=ax, pad=0.05, shrink=0.6) + else: + ax.scatter(ra_rad, dec_rad, c='black', s=5, alpha=0.5, edgecolors='none') + ax.set_title(title) + ax.grid(True, alpha=0.3) + if save_path: + fig.savefig(save_path, bbox_inches='tight') + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +### 9. 发现图（Finding Chart） + +```python +def plot_finding_chart(image_data, ra, dec, label=None, save_path=None): + """image_data: 2D numpy array (from FITS or JPEG)""" + fig, ax = plt.subplots() + ax.imshow(image_data, cmap='gray_r', origin='lower') + # 标注中心 + h, w = image_data.shape + ax.plot(w/2, h/2, 'r+', markersize=15, markeredgewidth=2) + if label: + ax.annotate(label, (w/2, h/2), textcoords="offset points", xytext=(10, 10), + color='red', fontsize=10, fontweight='bold') + ax.set_title(f"Finding Chart ({ra:.4f}, {dec:.4f})") + ax.set_xlabel("X (pix)") + ax.set_ylabel("Y (pix)") + if save_path: + fig.savefig(save_path) + plt.close(fig) + else: + plt.show() + return fig, ax +``` + +## 工作流程 + +1. **确定图表类型**：根据数据特征选择合适的图表模板 +2. **准备数据**：从文献、数据库或分析结果中提取数据 +3. **生成脚本**：使用上述模板，填入实际数据 +4. **执行**：`python3 script.py` +5. **验证**：检查输出图片是否正确 +6. **保存**：保存到 `figures/` 目录，命名格式：`{description}_{source}.{ext}` + +## 保存路径约定 + +- `figures/` — 项目根目录下的图表输出目录 +- 文件名格式：`{描述}_{来源}_{日期}.{格式}` + - 示例：`spectrum_lamost_438809089_20260710.pdf` + - 示例：`hr_diagram_m31_field_20260710.png` diff --git a/src/agent/autonomous.rs b/src/agent/autonomous.rs index 8c4e4e5..926ba52 100644 --- a/src/agent/autonomous.rs +++ b/src/agent/autonomous.rs @@ -109,9 +109,18 @@ impl IdlePoller { .await { autonomous_turns += 1; - // 创建 AgentRuntime 并执行任务 - let runtime = crate::agent::runtime::AgentRuntime::new( + // 通过会话注册表获取 runtime（复用该会话的 + // 后台队列/压缩日志等跨 turn 状态；模式从 DB 回放） + let mode = crate::agent::runtime::session::load_session_mode( + &self.app_state.db, + &task.session_id, + ) + .await + .unwrap_or_else(|| "default".to_string()); + let runtime = self.app_state.agent_runtimes.get_or_create( self.app_state.clone(), + &task.session_id, + &mode, ); let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); let _ = runtime diff --git a/src/agent/compact.rs b/src/agent/compact.rs index 982e283..f6fe053 100644 --- a/src/agent/compact.rs +++ b/src/agent/compact.rs @@ -17,14 +17,13 @@ pub mod collapse; -use std::sync::atomic::{AtomicBool, Ordering}; use tracing::{info, warn}; use super::hooks::{HookRegistry, PostCompactContext, PreCompactContext}; use collapse::{CollapseLog, CollapseMethod}; -/// 递归守卫：防止压缩内部触发的 LLM 调用再次触发压缩。 -static COMPACTING: AtomicBool = AtomicBool::new(false); +// 递归守卫已迁移至 session_registry::try_begin_compaction/end_compaction： +// 按 session_id 隔离，并发会话压缩互不干扰（历史上是进程级 AtomicBool）。 use crate::clients::llm::{ChatMessage, LlmClient, MessageRole}; @@ -182,11 +181,16 @@ fn find_tool_name_for_call_id(messages: &[ChatMessage], tool_call_id: &str) -> O None } -/// 轻量级压缩：将较早的工具结果替换为简短占位符，释放上下文空间。 -/// 保留最近 `keep_recent` 条工具结果不变。 +/// 轻量级压缩（确定性剪枝，零模型调用）：将较早的工具结果替换为简短 +/// 占位符，释放上下文空间。保留最近 `keep_recent` 条工具结果不变。 +/// +/// 剪枝策略（参考 deepseek-harness tool-result-pruner 的 head/tail 保留）： +/// - 短结果（≤ `PRUNE_PREVIEW_THRESHOLD` 字符）：`[Previous: used {tool}]` +/// 纯占位符（节省 ~80 tokens/条） +/// - 长结果：占位符 + 头部/尾部预览，为模型保留检索锚点（关键数值、 +/// 结论行常在结果尾部） /// /// 不可变风格：接收引用，返回新的消息列表。 -/// 使用 `[Previous: used {tool_name}]` 替代字符预览，节省 ~80 tokens/条。 pub fn micro_compact(messages: &[ChatMessage], keep_recent: usize) -> Vec { // 收集所有工具消息的索引、call_id 和工具名 let tool_positions: Vec<(usize, &str)> = messages @@ -231,7 +235,7 @@ pub fn micro_compact(messages: &[ChatMessage], keep_recent: usize) -> Vec Vec) -> String { + let Some(text) = content else { + return format!("[Previous: used {}]", tool_name); + }; + let char_count = text.chars().count(); + if char_count <= PRUNE_PREVIEW_THRESHOLD { + return format!("[Previous: used {}]", tool_name); + } + let head: String = text.chars().take(PRUNE_HEAD_CHARS).collect(); + let tail: String = text + .chars() + .skip(char_count.saturating_sub(PRUNE_TAIL_CHARS)) + .collect(); + format!( + "[Previous: used {}] ({} 字符，已剪枝)\n头: {}…\n尾: …{}", + tool_name, char_count, head, tail + ) +} + /// 粗略估算消息列表的 token 数（用作首次调用的近似值）。 /// 后续迭代优先使用 API 返回的精确 prompt_tokens。 /// @@ -256,6 +287,12 @@ pub fn rough_estimate_tokens(messages: &[ChatMessage]) -> usize { .sum() } +/// 估算单条消息的 token 数（所有循环统一使用本函数， +/// 消除主循环 `len/3` 与 compact 的 CJK 加权两套口径不一致的问题）。 +pub fn estimate_message_tokens(message: &ChatMessage) -> usize { + message.text().map_or(0, estimate_tokens_from_text) + 4 +} + /// 根据文本内容估算 token 数。 fn estimate_tokens_from_text(text: &str) -> usize { let mut ascii_chars = 0usize; @@ -513,6 +550,7 @@ pub async fn compress_context_with_hooks( session_id, hook_registry, None, + None, ) .await; } @@ -521,11 +559,15 @@ pub async fn compress_context_with_hooks( /// 如果不预留，当上下文接近窗口上限时压缩调用本身会因输出空间不足而失败。 pub const COMPACTION_OUTPUT_RESERVE: usize = 4000; -/// 带 Hook 和 CollapseLog 的上下文压缩变体。 +/// 带 Hook、CollapseLog 和会话事件日志的上下文压缩变体。 /// 压缩后自动记录 commit 到 CollapseLog，并在溢出时注入合并摘要。 /// +/// `db` 提供时写 compaction_start/compaction_end 事件（日志化锁： +/// 崩溃后开口的 compaction_start 由恢复时的 reconcile 合成关闭）。 +/// /// `token_soft_limit`: TokenBudget 软限制，统一各压缩层的触发阈值。 /// `max_messages`: 消息数上限，用于 Layer 0 snip 触发。 +#[allow(clippy::too_many_arguments)] // 压缩上下文 + 可选钩子/日志/DB 依赖 pub async fn compress_context_with_hooks_and_log( messages: &mut Vec, llm: &LlmClient, @@ -534,20 +576,39 @@ pub async fn compress_context_with_hooks_and_log( session_id: &str, hook_registry: Option<&HookRegistry>, collapse_log: Option<&std::sync::Mutex>, + db: Option<&sqlx::SqlitePool>, ) { if messages.len() <= 4 { return; } - // 递归守卫：如果已在压缩中，跳过（防止嵌套压缩死循环） - if COMPACTING.swap(true, Ordering::SeqCst) { - warn!("[Compact] 递归守卫触发：已有进行中的压缩操作，跳过"); + // 递归守卫（按会话）：如果该会话已在压缩中，跳过（防止嵌套压缩死循环） + if !super::runtime::session_registry::try_begin_compaction(session_id) { + warn!( + "[Compact] 递归守卫触发：会话 {} 已有进行中的压缩操作，跳过", + session_id + ); return; } let before_count = messages.len(); let est_tokens = rough_estimate_tokens(messages); + // ── 日志化锁：compaction_start（先落盘再执行；崩溃留下可检测开口）── + if let Some(db) = db { + super::runtime::session_events::append_event( + db, + session_id, + -1, + super::runtime::session_events::event_types::COMPACTION_START, + serde_json::json!({ + "before_messages": before_count, + "estimated_tokens": est_tokens, + }), + ) + .await; + } + // OnPreCompact hook if let Some(registry) = hook_registry { registry @@ -623,6 +684,22 @@ pub async fn compress_context_with_hooks_and_log( } } + // ── 日志化锁：compaction_end ── + if let Some(db) = db { + super::runtime::session_events::append_event( + db, + session_id, + -1, + super::runtime::session_events::event_types::COMPACTION_END, + serde_json::json!({ + "after_messages": after_count, + "method": method, + "saved_tokens": est_tokens.saturating_sub(rough_estimate_tokens(messages)), + }), + ) + .await; + } + // OnPostCompact hook if let Some(registry) = hook_registry { registry @@ -635,7 +712,7 @@ pub async fn compress_context_with_hooks_and_log( } // 释放递归守卫 - COMPACTING.store(false, Ordering::SeqCst); + super::runtime::session_registry::end_compaction(session_id); } // ── Compaction → Memory Extraction Bridge (P3) ── @@ -907,6 +984,28 @@ mod tests { assert_eq!(name, None); } + #[test] + fn test_prune_tool_result_short_keeps_placeholder() { + assert_eq!( + prune_tool_result("search_papers", Some("short result")), + "[Previous: used search_papers]" + ); + assert_eq!( + prune_tool_result("search_papers", None), + "[Previous: used search_papers]" + ); + } + + #[test] + fn test_prune_tool_result_long_keeps_head_tail() { + let long = format!("HEAD_MARKER\n{}\nTAIL_MARKER", "x".repeat(500)); + let pruned = prune_tool_result("get_paper_content", Some(&long)); + assert!(pruned.contains("[Previous: used get_paper_content]")); + assert!(pruned.contains("HEAD_MARKER")); + assert!(pruned.contains("TAIL_MARKER")); + assert!(pruned.contains("已剪枝")); + } + #[test] fn test_rough_estimate_tokens() { let messages = vec![ diff --git a/src/agent/engine.rs b/src/agent/engine.rs new file mode 100644 index 0000000..0123eb3 --- /dev/null +++ b/src/agent/engine.rs @@ -0,0 +1,1224 @@ +// src/agent/engine.rs +// +// 统一 ReAct 引擎 — 主代理 / 子代理 / 队友共用的唯一循环实现。 +// +// 历史问题：runtime（主循环）、subagent.rs、team/teammate.rs 各自手写一份 +// ReAct 循环，能力严重不对等（teammate 无 hooks/权限/SSE/持久化；subagent +// 无错误恢复/熔断/denial 追踪），修 bug 需要四处同步。 +// +// 参考 deepseek-harness "agent-loop 是 harness 中唯一包含具体循环逻辑的包" +// 的约束：本引擎是仓库中唯一的循环实现，行为差异通过 ReactEngine 的 +// 组合字段（Option/策略枚举）注入，而不是分叉实现： +// +// lead → 完整配置（交互权限、恢复阶梯、todo nag、后台队列、压缩熔断） +// subagent → 事件前缀 [sub]、AskPolicy::AutoDeny（无人值守 fail-closed）、 +// 独立持久化 sink（agent_name = sub_xxx） +// teammate → 静默（无 SSE）、无持久化、AtomicFlag 取消 + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing::{error, info, warn}; + +use crate::api::AppState; +use crate::clients::llm::{ChatMessage, LlmClient}; + +use crate::agent::background::BgNotificationQueue; +use crate::agent::compact; +use crate::agent::compact::collapse::CollapseLog; +use crate::agent::hooks::{HookRegistry, StepCompleteContext}; +use crate::agent::runtime::checkpoint::CheckpointManager; +use crate::agent::runtime::circuit_breaker::CompactionCircuitBreaker; +use crate::agent::runtime::denial_tracker::DenialTracker; +use crate::agent::runtime::duplicate_detector::DuplicateDetector; +use crate::agent::runtime::error_recovery::{classify_error, ErrorKind, ErrorRecovery}; +use crate::agent::runtime::events::{AgentMetrics, AgentStreamEvent}; +use crate::agent::runtime::file_cache::FileStateCache; +use crate::agent::runtime::permission::{PermissionChecker, PermissionResult}; +use crate::agent::runtime::streaming::{self, StreamOutput, StreamStatus}; +use crate::agent::runtime::token_budget::TokenBudget; +use crate::agent::runtime::{executor, AgentConfig}; +use crate::agent::terminal::TurnTerminal; +use crate::agent::tools::ToolRegistry; + +// ── 策略类型 ────────────────────────────────────────────────────────────── + +/// AskUser 类权限请求的处理策略。 +/// +/// - `Interactive`：通过 AppState.pending_permissions 挂起等待用户响应（120s 超时） +/// - `AutoDeny`：无人值守上下文（子代理/队友）中无用户可问，fail-closed 自动拒绝 +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum AskPolicy { + Interactive, + AutoDeny, +} + +/// SSE 事件出口：None = 静默；prefix = 事件名/思考内容加 `[prefix]` 前缀。 +#[derive(Clone, Default)] +pub struct EventTap { + tx: Option>, + prefix: Option, +} + +impl EventTap { + /// 静默（无 SSE） + pub fn none() -> Self { + EventTap { + tx: None, + prefix: None, + } + } + + /// 主代理：完整事件流 + pub fn channel(tx: mpsc::UnboundedSender) -> Self { + EventTap { + tx: Some(tx), + prefix: None, + } + } + + /// 子代理：转发到父代理通道并加前缀（如 "[sub]"） + pub fn prefixed(tx: mpsc::UnboundedSender, prefix: &str) -> Self { + EventTap { + tx: Some(tx), + prefix: Some(prefix.to_string()), + } + } + + pub fn is_silent(&self) -> bool { + self.tx.is_none() + } + + /// 底层发送通道（工具的流式输出用） + pub fn sender(&self) -> Option> { + self.tx.clone() + } + + pub fn send(&self, event: AgentStreamEvent) { + if let Some(tx) = &self.tx { + let _ = tx.send(event); + } + } + + /// 带前缀的工具显示名（"[sub] 运行命令"） + pub fn display(&self, name: &str) -> String { + match &self.prefix { + Some(p) => format!("[{}] {}", p, name), + None => name.to_string(), + } + } + + /// 带前缀的思考内容 + pub fn thought(&self, content: &str) -> String { + match &self.prefix { + Some(p) => format!("[{}] {}", p, content), + None => content.to_string(), + } + } +} + +/// 取消信号来源。 +pub enum CancelSource { + /// 会话取消表（主代理：POST /chat/stop 写入） + Session { + session_id: String, + cancelled_runs: Arc>, + }, + /// 外部原子标志（队友：TeamMemberHandle.request_stop） + Flag(Arc), +} + +impl CancelSource { + pub fn is_cancelled(&self) -> bool { + match self { + CancelSource::Session { + session_id, + cancelled_runs, + } => cancelled_runs.contains_key(session_id), + CancelSource::Flag(flag) => flag.load(Ordering::SeqCst), + } + } + + pub fn clear(&self) { + match self { + CancelSource::Session { + session_id, + cancelled_runs, + } => { + cancelled_runs.remove(session_id); + } + CancelSource::Flag(flag) => flag.store(false, Ordering::SeqCst), + } + } + + /// 引擎内部执行器轮询用的共享标志（供 execute_parallel 的取消观察者） + pub fn to_flag(&self) -> Arc { + match self { + CancelSource::Flag(flag) => flag.clone(), + CancelSource::Session { .. } => Arc::new(AtomicBool::new(false)), + } + } +} + +/// 消息持久化出口 — 循环中产生的 assistant/tool 消息如何落库。 +#[async_trait::async_trait] +pub trait MessageSink: Send + Sync { + async fn save( + &self, + turn_index: i32, + step_index: i32, + msg: &ChatMessage, + thought: Option<&str>, + extra_metadata: Option, + ); +} + +/// 数据库 sink（agent_messages 表），lead 与子代理共用同一张表， +/// 通过 agent_name 区分身份。 +pub struct DbMessageSink { + pub db: sqlx::SqlitePool, + pub session_id: String, + pub agent_name: String, + /// 附加到每条消息 metadata 的固定字段（如 {"is_subagent": true}） + pub fixed_metadata: Option, +} + +#[async_trait::async_trait] +impl MessageSink for DbMessageSink { + async fn save( + &self, + turn_index: i32, + step_index: i32, + msg: &ChatMessage, + thought: Option<&str>, + extra_metadata: Option, + ) { + use crate::clients::llm::MessageRole; + + let role = match msg.role { + MessageRole::System => "system", + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::Tool => "tool", + }; + let content = msg.content.clone().unwrap_or_default(); + let tool_calls_json = msg + .tool_calls + .as_ref() + .map(|tc| serde_json::to_string(tc).unwrap_or_default()); + let token_count = content.len() as i32 / 4; + + let mut metadata = serde_json::json!({ + "has_thought": thought.is_some(), + "has_tool_calls": tool_calls_json.is_some(), + "step_index": step_index, + }); + if let Some(ref fixed) = self.fixed_metadata { + if let (Some(base), Some(extra)) = (metadata.as_object_mut(), fixed.as_object()) { + for (k, v) in extra { + base.insert(k.clone(), v.clone()); + } + } + } + if let Some(ref extra) = extra_metadata { + if let (Some(base), Some(extra_obj)) = (metadata.as_object_mut(), extra.as_object()) { + for (k, v) in extra_obj { + base.insert(k.clone(), v.clone()); + } + } + } + + let result = sqlx::query( + "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, raw_json, agent_name) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&self.session_id) + .bind(turn_index) + .bind(step_index) + .bind(role) + .bind(&content) + .bind(thought) + .bind(&tool_calls_json) + .bind(msg.tool_call_id.as_deref()) + .bind(token_count) + .bind(serde_json::to_string(&metadata).unwrap_or_default()) + .bind(serde_json::to_string(msg).unwrap_or_default()) + .bind(&self.agent_name) + .execute(&self.db) + .await; + + if let Err(e) = result { + warn!("[Engine] 消息持久化失败（非致命）: {}", e); + } + } +} + +/// 引擎运行结果 +#[derive(Debug, Default)] +pub struct EngineOutcome { + pub metrics: AgentMetrics, + pub terminal: Option, + /// 循环产生的最终文本（无工具调用步的回答，或步数耗尽时的强制回答） + pub final_text: String, +} + +// ── 引擎 ───────────────────────────────────────────────────────────────── + +/// 统一 ReAct 引擎。 +/// +/// `run()` 的循环体：取消检查 → 上下文压缩 → 预算/nag → 步数上限 → +/// 后台通知注入 → LLM 调用（可选恢复阶梯）→ 工具验证与执行 → 循环。 +pub struct ReactEngine<'a> { + pub llm: &'a LlmClient, + pub tool_registry: &'a ToolRegistry, + pub app_state: Arc, + + pub config: AgentConfig, + pub session_id: String, + pub turn_index: i32, + pub agent_name: String, + + // ── 行为注入 ── + pub tap: EventTap, + pub sink: Option>, + pub hooks: Option<&'a HookRegistry>, + pub ask_policy: AskPolicy, + pub permission_checker: Option<&'a PermissionChecker>, + pub session_checker: Option, + pub denial_tracker: Option>>, + pub checkpoint_manager: Option>, + pub read_file_state: Arc>, + pub cancel: CancelSource, + pub breaker: Option>>, + pub collapse_log: Option>>, + pub bg_queue: Option>, + + // ── 能力开关 ── + /// todo nag 提醒 + token 预算 nudge 注入（仅交互式代理） + pub nudges: bool, + /// diminishing returns 检测 + pub diminishing_returns: bool, + /// LLM 错误恢复阶梯（退避重试 + 降级压缩 + 升限） + pub recovery: bool, + /// 把 todo_write 的 todos 持久化到任务板（lead 专属） + pub persist_todos: bool, +} + +impl<'a> ReactEngine<'a> { + /// 执行 ReAct 循环。`messages` 应已包含 system prompt 与最新用户消息。 + pub async fn run(&self, messages: &mut Vec) -> EngineOutcome { + let sid = &self.session_id; + let turn_index = self.turn_index; + + // 会话级权限检查器注册（仅交互式代理：子代理注册会覆盖 lead 的条目） + if self.ask_policy == AskPolicy::Interactive { + if let Some(checker) = self.permission_checker { + self.app_state + .session + .session_permission_checkers + .entry(sid.clone()) + .or_insert_with(|| (*checker).clone()); + } + } + + // 未注入的跨 turn 状态 → 引擎局部实例（生命周期 = 本次循环） + let local_breaker = Arc::new(std::sync::Mutex::new(CompactionCircuitBreaker::new())); + let breaker = self.breaker.as_ref().unwrap_or(&local_breaker); + let local_denial = Arc::new(std::sync::Mutex::new(DenialTracker::new( + self.config.denial_max_consecutive, + self.config.denial_max_total, + ))); + let denial_tracker = self.denial_tracker.as_ref().unwrap_or(&local_denial); + + let tool_defs = self.tool_registry.definitions(); + let mut duplicate_detector = DuplicateDetector::default(); + let mut metrics = AgentMetrics::default(); + let mut step = 0; + let mut outcome = EngineOutcome::default(); + + // Token 追踪（API 精确值优先，字符估算作近似值） + let mut last_api_prompt_tokens: Option = None; + let mut msg_count_at_last_call: usize = messages.len(); + let mut steps_since_last_todo: usize = 0; + let nag_after_steps: usize = 3; + let mut pending_manual_compress = false; + + let mut token_budget = + TokenBudget::new(self.config.token_soft_limit, self.config.token_hard_limit); + + loop { + step += 1; + + // ── Checkpoint: 每个 ReAct 迭代开始时重置去重状态 ── + if let Some(ckpt) = &self.checkpoint_manager { + ckpt.new_turn(); + } + + // 检查取消 + if self.cancel.is_cancelled() { + self.cancel.clear(); + warn!("[Engine:{}] 用户手动中止了执行", self.agent_name); + self.tap.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + outcome.terminal = Some(TurnTerminal::CancelledByUser { + session_id: sid.clone(), + at_step: step, + }); + break; + } + + // ── Token 感知的压缩触发 ── + let estimated_tokens = match last_api_prompt_tokens { + Some(last_tokens) => { + let new_msg_count = messages.len().saturating_sub(msg_count_at_last_call); + let new_tokens_estimate: usize = messages + .iter() + .rev() + .take(new_msg_count) + .map(compact::estimate_message_tokens) + .sum(); + last_tokens as usize + new_tokens_estimate + } + None => compact::rough_estimate_tokens(messages), + }; + + let token_limit = token_budget.soft_limit; + + let mut did_compress = false; + + let breaker_ok = match breaker.lock() { + Ok(mut b) => b.can_attempt(), + Err(e) => { + warn!( + "[Engine:{}] 熔断器锁异常，跳过自动压缩: {:?}", + self.agent_name, e + ); + false + } + }; + + if estimated_tokens > token_limit && breaker_ok { + info!( + "[Engine:{}] 上下文超限 (est. {} tokens > {} limit)，触发压缩", + self.agent_name, estimated_tokens, token_limit + ); + let before_len = self + .snapshot_compress_restore(messages, &self.app_state, sid) + .await; + + if let Ok(mut b) = breaker.lock() { + if messages.len() < before_len { + b.record_success(); + } else { + b.record_failure(); + } + } + last_api_prompt_tokens = None; + msg_count_at_last_call = messages.len(); + metrics.compression_count += 1; + did_compress = true; + } else if estimated_tokens > token_limit && !breaker_ok { + warn!("[Engine:{}] 熔断器已打开，跳过自动压缩", self.agent_name); + } + + // 手动压缩请求（跳过刚自动压缩过的情况；不受熔断器限制） + if pending_manual_compress && !did_compress { + pending_manual_compress = false; + info!( + "[Engine:{}] 执行手动压缩（compress_context 工具触发）", + self.agent_name + ); + + self.snapshot_compress_restore(messages, &self.app_state, sid) + .await; + + if let Ok(mut b) = breaker.lock() { + b.reset(); + } + last_api_prompt_tokens = None; + msg_count_at_last_call = messages.len(); + metrics.compression_count += 1; + } else if pending_manual_compress { + pending_manual_compress = false; + info!( + "[Engine:{}] 跳过手动压缩（刚已完成自动压缩）", + self.agent_name + ); + } + + // Token 预算追踪 + token_budget.record_continuation(); + + if self.nudges { + let mut should_nudge = false; + + if steps_since_last_todo >= nag_after_steps { + messages.push(ChatMessage::user( + "提醒：你已经连续多步未更新任务计划。建议调用 todo_write 工具复盘当前进度并规划后续步骤。", + )); + steps_since_last_todo = 0; + should_nudge = true; + } + + if !should_nudge { + if let Some(nudge) = token_budget.nudge_message() { + messages.push(ChatMessage::user(nudge)); + } + } + } + + // Diminishing returns 检测 — 强制结束 + if self.diminishing_returns && token_budget.diminishing_returns { + warn!( + "[Engine:{}] 检测到 diminishing returns，强制结束循环", + self.agent_name + ); + self.tap.send(AgentStreamEvent::Error { + message: "检测到重复操作模式，已自动停止。请查看已收集的信息。".to_string(), + }); + messages.push(ChatMessage::user( + "检测到你的后续步骤未产生新信息（diminishing returns）。\ + 请基于已收集的全部信息直接给出最终答案，不要再调用任何工具。", + )); + outcome.final_text = self + .final_answer_without_tools(messages, turn_index, step) + .await; + break; + } + + // 最大步数检查 + if step > self.config.max_steps { + warn!( + "[Engine:{}] 达到最大步数限制 ({} steps)", + self.agent_name, self.config.max_steps + ); + self.tap.send(AgentStreamEvent::Error { + message: format!( + "已达到最大推理步数 ({})，请根据已收集的信息给出最终回答。", + self.config.max_steps + ), + }); + messages.push(ChatMessage::user(format!( + "你已经执行了 {} 步（最大 {} 步）。请根据已有信息直接给出最终答案，不要再调用工具。", + step, self.config.max_steps + ))); + outcome.final_text = self + .final_answer_without_tools(messages, turn_index, step) + .await; + break; + } + + // ── 后台任务通知注入 ── + if let Some(queue) = &self.bg_queue { + for result in queue.drain().await { + let status = if result.is_error { "❌" } else { "✅" }; + messages.push(ChatMessage::user(format!( + "[后台任务完成] {} {}: {} ({}): {}", + status, result.tool_name, result.bibcode, result.task_id, result.summary, + ))); + } + } + + // ── LLM 流式调用（含错误恢复） ── + let stream_output_opt = self + .call_llm_with_recovery(messages, &tool_defs, step, sid, &mut token_budget) + .await; + + let stream_output = match stream_output_opt { + Some(output) => output, + None => { + outcome.terminal = Some(TurnTerminal::ModelError { + session_id: sid.clone(), + message: "LLM 调用失败，所有恢复步骤已尝试完毕".to_string(), + }); + break; + } + }; + + // 更新 API 精确 token 计数 + token 预算 + if let Some(ref u) = stream_output.usage { + last_api_prompt_tokens = Some(u.prompt_tokens); + msg_count_at_last_call = messages.len(); + token_budget.spend_input(u.prompt_tokens as usize); + token_budget.spend_output(u.completion_tokens as usize); + } + + // ── 处理 Thought/Reasoning ── + let mut thought_content = stream_output.reasoning.clone(); + + if thought_content.is_none() + && stream_output.is_tool_call_step + && !stream_output.content.is_empty() + { + thought_content = Some(stream_output.content.clone()); + } + + if stream_output.is_tool_call_step { + if let Some(ref thought_text) = thought_content { + self.tap.send(AgentStreamEvent::Thought { + content: self.tap.thought(thought_text), + step, + }); + } + } + + // ── 无工具调用 = 最终回答 ── + let tool_calls = match stream_output.tool_calls { + Some(ref tc) if !tc.is_empty() => tc.clone(), + _ => { + let assistant_msg = ChatMessage::assistant_with_reasoning( + if stream_output.content.is_empty() { + None + } else { + Some(stream_output.content.clone()) + }, + stream_output.reasoning.clone(), + None, + ); + self.persist( + turn_index, + step, + &assistant_msg, + stream_output.reasoning.as_deref(), + None, + ) + .await; + messages.push(assistant_msg); + + // 思考后直接给出答案时，thought 尚未在上方发送 + if let Some(ref thought_text) = stream_output.reasoning { + self.tap.send(AgentStreamEvent::Thought { + content: self.tap.thought(thought_text), + step, + }); + } + + if let Some(u) = stream_output.usage { + self.tap.send(AgentStreamEvent::Usage { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + }); + } + + outcome.final_text = stream_output.content.clone(); + break; + } + }; + + // 修复空 ID（LLM 可能不返回 tool_call id） + let mut tool_calls = tool_calls; + for tc in tool_calls.iter_mut() { + if tc.id.is_empty() { + tc.id = format!("call_{}", &uuid::Uuid::new_v4().to_string()[..8]); + } + } + + // ── 工具调用处理（信号由工具声明，不按名字特判） ── + let called_todo_write = tool_calls.iter().any(|tc| { + self.tool_registry + .signals_for(&tc.function.name) + .resets_todo_nag + }); + if called_todo_write { + steps_since_last_todo = 0; + } else { + steps_since_last_todo += 1; + } + + if tool_calls.iter().any(|tc| { + self.tool_registry + .signals_for(&tc.function.name) + .requests_manual_compact + }) { + pending_manual_compress = true; + } + + metrics.total_steps = step; + for tc in &tool_calls { + *metrics + .tool_calls + .entry(tc.function.name.clone()) + .or_insert(0) += 1; + } + + let assistant_msg = ChatMessage::assistant_with_reasoning( + if stream_output.content.is_empty() { + None + } else { + Some(stream_output.content.clone()) + }, + stream_output.reasoning.clone(), + Some(tool_calls.clone()), + ); + self.persist( + turn_index, + step, + &assistant_msg, + stream_output.reasoning.as_deref(), + None, + ) + .await; + messages.push(assistant_msg); + + // 验证 + 准备工具调用 + let (prepared_calls, has_duplicate) = executor::validate_and_prepare( + &tool_calls, + &mut duplicate_detector, + self.config.duplicate_call_threshold, + messages, + self.tool_registry, + &self.tap, + step, + ); + + if has_duplicate { + metrics.duplicate_detections += 1; + continue; + } + + if prepared_calls.is_empty() { + continue; + } + + // 执行工具（权限管道 + 分区并行 + hooks） + let session_checker = self.session_checker.clone(); + let exec_result = executor::execute_parallel( + &prepared_calls, + &executor::ExecutorContext { + tool_registry: self.tool_registry, + app_state: self.app_state.clone(), + hook_registry: self.hooks, + permission_checker: self.permission_checker, + session_checker: session_checker.as_ref(), + denial_tracker: Some(denial_tracker), + checkpoint_manager: self.checkpoint_manager.as_deref(), + db: &self.app_state.db, + session_id: sid, + agent_name: &self.agent_name, + turn_index, + step, + tool_timeout_secs: self.config.tool_timeout_secs, + max_output_chars: self.config.max_tool_output_chars, + read_file_state: self.read_file_state.clone(), + enable_thinking: self.config.enable_thinking, + additional_allowed_dirs: self.config.additional_allowed_dirs.clone(), + ask_policy: self.ask_policy, + cancel: self.cancel.to_flag(), + interactive_cancel: matches!(self.cancel, CancelSource::Session { .. }), + cancel_session_id: match &self.cancel { + CancelSource::Session { session_id, .. } => Some(session_id.clone()), + CancelSource::Flag(_) => None, + }, + tap: self.tap.clone(), + sink: self.sink.clone(), + }, + ) + .await; + + // 拒绝追踪熔断检查 + if let Ok(dt) = denial_tracker.lock() { + if dt.should_terminate() { + let reason = dt.termination_reason(); + warn!("[Engine:{}] 拒绝熔断触发: {}", self.agent_name, reason); + self.tap.send(AgentStreamEvent::Error { message: reason }); + break; + } + } + + for tm in exec_result.tool_messages { + messages.push(tm.chat_message); + } + + // Hook 注入的附加上下文（单步内去重） + let mut dedup = crate::agent::hooks::ContextDeduplicator::new(); + for ctx in &exec_result.hook_contexts { + if dedup.is_duplicate(ctx) { + continue; + } + let reminder = format!( + "\n[Hook 注入上下文]\n{}\n", + ctx + ); + messages.push(ChatMessage::user(&reminder)); + } + + for be in &exec_result.blocking_errors { + warn!("[Engine:{}] Hook 阻塞错误: {}", self.agent_name, be); + } + + // todo_write 任务状态持久化到任务板 + if self.persist_todos && called_todo_write { + for prep in &prepared_calls { + if self + .tool_registry + .signals_for(&prep.tool_name) + .persists_todos + { + if let Some(todos) = prep.args.get("todos").and_then(|t| t.as_array()) { + let todos_vec: Vec = todos.to_vec(); + let _ = crate::agent::tools::persist_tasks( + &self.app_state.db, + sid, + &todos_vec, + "lead", + ) + .await; + } + } + } + } + + if exec_result.was_cancelled { + self.cancel.clear(); + warn!("[Engine:{}] 工具执行期间被用户手动中止", self.agent_name); + self.tap.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + outcome.terminal = Some(TurnTerminal::CancelledByUser { + session_id: sid.clone(), + at_step: step, + }); + break; + } + + // OnStepComplete hook + if let Some(registry) = self.hooks { + let step_ctx = StepCompleteContext { + session_id: sid.clone(), + step, + max_steps: self.config.max_steps, + messages_count: messages.len(), + estimated_tokens, + token_limit, + }; + registry.run_on_step_complete(&step_ctx).await; + } + } + + outcome.metrics = metrics; + outcome + } + + /// 消息持久化（无 sink 时跳过） + async fn persist( + &self, + turn_index: i32, + step: usize, + msg: &ChatMessage, + thought: Option<&str>, + extra: Option, + ) { + if let Some(sink) = &self.sink { + sink.save(turn_index, step as i32, msg, thought, extra) + .await; + } + } + + /// 执行文件缓存快照 → 压缩 → 恢复 → 上下文注入 的完整周期。 + /// 返回压缩前的消息数（调用方用于熔断器反馈）。 + async fn snapshot_compress_restore( + &self, + messages: &mut Vec, + app_state: &Arc, + session_id: &str, + ) -> usize { + let before_len = messages.len(); + + let file_snapshot = { + if let Ok(mut cache) = self.read_file_state.lock() { + let snap = cache.to_snapshot(); + cache.clear(); + snap + } else { + Vec::new() + } + }; + + let pre_compact_snippets = compact::extract_snippets(messages); + + compact::compress_context_with_hooks_and_log( + messages, + self.llm, + self.config.token_soft_limit, + self.config.max_messages, + session_id, + self.hooks, + self.collapse_log.as_deref(), + Some(&app_state.db), + ) + .await; + + // 压缩后提取记忆（被丢弃内容喂给记忆提取子代理） + compact::spawn_memory_extraction_from_snippets( + pre_compact_snippets, + session_id, + app_state.memory_manager.clone(), + app_state.clone(), + ); + + { + if let Ok(mut cache) = self.read_file_state.lock() { + cache.restore_from_snapshot( + &file_snapshot, + crate::agent::runtime::file_cache::POST_COMPACT_MAX_FILES_TO_RESTORE, + ); + } + let restore_ctx = FileStateCache::build_restore_context( + &file_snapshot, + crate::agent::runtime::file_cache::POST_COMPACT_MAX_FILES_TO_RESTORE, + ); + for block in restore_ctx { + messages.push(ChatMessage::user(format!("[压缩后上下文恢复]\n{}", block))); + } + } + + before_len + } + + /// LLM 流式调用，含完整的错误恢复阶梯（recovery = false 时单次尝试）。 + async fn call_llm_with_recovery( + &self, + messages: &mut Vec, + tool_defs: &[crate::clients::llm::ToolDefinition], + step: usize, + session_id: &str, + token_budget: &mut TokenBudget, + ) -> Option { + let output = streaming::process_llm_stream( + self.llm, + messages, + tool_defs, + &self.tap, + step, + session_id, + self.app_state.session.cancelled_runs.clone(), + self.config.enable_thinking, + ) + .await; + + match output.status { + StreamStatus::Success => return Some(output), + StreamStatus::Cancelled => { + self.app_state.session.cancelled_runs.remove(session_id); + warn!("[Engine:{}] 流式调用期间被用户手动中止", self.agent_name); + self.tap.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + return None; + } + StreamStatus::Error(ref e_str) => { + if !self.recovery { + error!( + "[Engine:{}] 流式读取错误（恢复已禁用）: {}", + self.agent_name, e_str + ); + self.tap.send(AgentStreamEvent::Error { + message: format!("大模型流式读取失败: {}", e_str), + }); + return None; + } + error!("[Engine:{}] 流式读取错误: {}", self.agent_name, e_str); + } + } + + let e_str = match &output.status { + StreamStatus::Error(s) => s.clone(), + _ => return Some(output), + }; + + let error_kind = classify_error(&e_str); + + // ── 429/529 瞬态错误：指数退避重试 ── + if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) { + let retry_after_secs = crate::agent::runtime::error_recovery::parse_retry_after(&e_str); + let mut consecutive_overloads: u32 = 0; + const MAX_BACKOFF_RETRIES: u32 = 10; + + for attempt in 0..MAX_BACKOFF_RETRIES { + let delay_ms = + crate::agent::runtime::error_recovery::backoff_delay(attempt, retry_after_secs); + info!( + "[Engine:{}] 退避重试 {}/{} ({}ms, error={:?})", + self.agent_name, + attempt + 1, + MAX_BACKOFF_RETRIES, + delay_ms, + error_kind + ); + + self.tap.send(AgentStreamEvent::Thought { + content: format!( + "⏳ 模型服务暂时不可用，正在重试 ({}/{})...", + attempt + 1, + MAX_BACKOFF_RETRIES + ), + step, + }); + + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + + // 529 连续过载检测：3 次后尝试切换备用模型 + if matches!(error_kind, ErrorKind::Overloaded) { + consecutive_overloads += 1; + if consecutive_overloads >= 3 { + let fallback = &self.app_state.config.llm.fallback_model; + if !fallback.is_empty() { + warn!( + "[Engine:{}] 连续 {} 次过载，切换到备用模型: {}", + self.agent_name, consecutive_overloads, fallback + ); + self.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) + % self.app_state.config.llm.fallback_chain.len()) + .min(self.app_state.config.llm.fallback_chain.len() - 1); + let alt = &self.app_state.config.llm.fallback_chain[idx]; + warn!( + "[Engine:{}] 连续 {} 次过载，从链中切换: {}", + self.agent_name, consecutive_overloads, alt + ); + self.llm.set_model(alt.clone()).await; + consecutive_overloads = 0; + } + } + } + + if self + .app_state + .session + .cancelled_runs + .contains_key(session_id) + { + warn!("[Engine:{}] 退避重试期间被用户取消", self.agent_name); + self.tap.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + return None; + } + + let retry_output = streaming::process_llm_stream( + self.llm, + messages, + tool_defs, + &self.tap, + step, + session_id, + self.app_state.session.cancelled_runs.clone(), + self.config.enable_thinking, + ) + .await; + + match retry_output.status { + StreamStatus::Success => { + info!( + "[Engine:{}] 退避重试成功！(尝试 {})", + self.agent_name, + attempt + 1 + ); + return Some(retry_output); + } + StreamStatus::Cancelled => { + self.app_state.session.cancelled_runs.remove(session_id); + return None; + } + StreamStatus::Error(_) => continue, + } + } + + warn!( + "[Engine:{}] {} 次退避重试后仍然失败", + self.agent_name, MAX_BACKOFF_RETRIES + ); + self.tap.send(AgentStreamEvent::Error { + message: format!( + "模型服务暂时不可用（已重试 {} 次）。请稍后再试或检查模型服务状态。", + MAX_BACKOFF_RETRIES + ), + }); + return None; + } + + if !ErrorRecovery::is_recoverable(&error_kind) { + self.tap.send(AgentStreamEvent::Error { + message: format!("大模型流式读取失败: {}", e_str), + }); + return None; + } + + let mut recovery = ErrorRecovery::new(token_budget.clone()); + let overflow_info = crate::agent::runtime::error_recovery::parse_context_overflow(&e_str); + + while let Some(recovery_step) = recovery.try_recover(&error_kind, overflow_info.as_ref()) { + match recovery_step { + crate::agent::runtime::error_recovery::RecoveryStep::AdjustMaxTokens { .. } => { + info!( + "[Engine:{}] 错误恢复: AdjustMaxTokens → {} (从错误消息自动计算)", + self.agent_name, token_budget.hard_limit + ); + } + crate::agent::runtime::error_recovery::RecoveryStep::RetryWithBackoff { + attempt, + delay_ms, + } => { + warn!( + "[Engine:{}] RetryWithBackoff 降级重试 (attempt={}, delay={}ms)", + self.agent_name, attempt, delay_ms + ); + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + } + crate::agent::runtime::error_recovery::RecoveryStep::AggressiveCompact => { + info!( + "[Engine:{}] 错误恢复: 激进压缩 (snip + micro with keep_recent=2)", + self.agent_name + ); + compact::snip_compact(messages, self.config.max_messages); + *messages = compact::micro_compact(messages, 2); + } + crate::agent::runtime::error_recovery::RecoveryStep::ReactiveCompact => { + info!("[Engine:{}] 错误恢复: LLM 摘要压缩", self.agent_name); + compact::compress_context( + messages, + self.llm, + self.config.token_soft_limit, + self.config.max_messages, + session_id, + ) + .await; + } + crate::agent::runtime::error_recovery::RecoveryStep::EscalateTokens { .. } => { + info!( + "[Engine:{}] 错误恢复: 提升 token 硬限制到 {}", + self.agent_name, recovery.token_budget.hard_limit + ); + } + crate::agent::runtime::error_recovery::RecoveryStep::MultiTurn => { + info!("[Engine:{}] 错误恢复: 注入多轮消息", self.agent_name); + messages.push(ChatMessage::user(ErrorRecovery::multi_turn_message())); + } + crate::agent::runtime::error_recovery::RecoveryStep::Surface => { + warn!( + "[Engine:{}] 错误恢复: 所有步骤失败，暴露错误", + self.agent_name + ); + break; + } + } + + let retry_output = streaming::process_llm_stream( + self.llm, + messages, + tool_defs, + &self.tap, + step, + session_id, + self.app_state.session.cancelled_runs.clone(), + self.config.enable_thinking, + ) + .await; + + match retry_output.status { + StreamStatus::Success => { + info!("[Engine:{}] 错误恢复成功！", self.agent_name); + *token_budget = recovery.token_budget.clone(); + return Some(retry_output); + } + StreamStatus::Cancelled => { + self.app_state.session.cancelled_runs.remove(session_id); + warn!("[Engine:{}] 恢复期间被用户中止", self.agent_name); + self.tap.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + return None; + } + StreamStatus::Error(retry_err) => { + info!( + "[Engine:{}] 恢复步骤 {:?} 未能解决，继续下一阶梯: {}", + self.agent_name, recovery_step, retry_err + ); + } + } + } + + self.tap.send(AgentStreamEvent::Error { + message: format!("大模型流式读取失败，且所有恢复步骤均未能解决: {}", e_str), + }); + None + } + + /// 步数耗尽时的最终答案生成（不带工具调用，强制 LLM 直接回答） + async fn final_answer_without_tools( + &self, + messages: &[ChatMessage], + turn_index: i32, + step: usize, + ) -> String { + let empty_tools: Vec = Vec::new(); + let mut stream_rx = match self + .llm + .chat_stream(messages, &empty_tools, self.config.enable_thinking) + .await + { + Ok(rx) => rx, + Err(e) => { + self.tap.send(AgentStreamEvent::Error { + message: format!("最终回答生成失败: {}", e), + }); + return String::new(); + } + }; + + use crate::clients::llm::StreamEvent; + let mut accumulated = String::new(); + while let Some(event) = stream_rx.recv().await { + match event { + StreamEvent::TextDelta(delta) => { + accumulated.push_str(&delta); + self.tap.send(AgentStreamEvent::TextDelta { + content: delta, + tool_call_id: None, + }); + } + StreamEvent::Usage(u) => { + self.tap.send(AgentStreamEvent::Usage { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + }); + } + StreamEvent::Done => break, + StreamEvent::Error(e) => { + self.tap.send(AgentStreamEvent::Error { + message: format!("最终回答流式错误: {}", e), + }); + break; + } + _ => {} + } + } + + let assistant_msg = ChatMessage::assistant(accumulated.clone()); + self.persist(turn_index, step, &assistant_msg, None, None) + .await; + + accumulated + } +} + +// ── 权限工具函数 ── + +/// fail-closed 的无人值守 AskUser 处理：无用户可问即拒绝。 +pub fn auto_deny_ask(tool_name: &str, message: &str) -> PermissionResult { + warn!( + "[Engine] 工具 {} 需要用户确认但在无人值守上下文中自动拒绝: {}", + tool_name, message + ); + PermissionResult::Denied { + reason: format!( + "工具 {} 需要用户确认但在无人值守上下文（子代理/队友）中不可用", + tool_name + ), + } +} diff --git a/src/agent/hooks/mod.rs b/src/agent/hooks/mod.rs index cfef254..af8fdf9 100644 --- a/src/agent/hooks/mod.rs +++ b/src/agent/hooks/mod.rs @@ -52,7 +52,7 @@ pub use builtins::{AuditLogHook, CancellationHook, ContextDeduplicator, MetricsH mod tests { use super::*; use async_trait::async_trait; - use std::collections::HashSet; + use std::sync::Arc; use tokio::sync::Mutex; diff --git a/src/agent/mod.rs b/src/agent/mod.rs index ce5a5bb..03d9fd1 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -3,15 +3,17 @@ // 科研智能体模块 — 基于 ReAct 框架实现 Thought → Action → Observation 循环。 // // 模块结构（参考 Claude Code 分层设计）： +// engine — 统一 ReAct 引擎（主代理/子代理/队友共用的唯一循环实现） // tools/ — 工具定义与注册（按功能域拆分） -// runtime — ReAct 循环引擎 + Streaming + 会话管理 -// compact — 三层上下文压缩（micro/auto/manual） +// runtime — 运行时编排：会话管理、权限、压缩熔断、系统提示词 +// compact — 多层上下文压缩（snip/micro/auto/manual） // terminal — 循环终止信号（结构化退出原因） // hooks — 生命周期事件系统（PreToolUse/PostToolUse/Stop） pub mod autonomous; pub mod background; pub mod compact; +pub mod engine; pub mod enums; pub mod hooks; pub mod memory; diff --git a/src/agent/runtime/checkpoint.rs b/src/agent/runtime/checkpoint.rs index dde09cc..e2a6263 100644 --- a/src/agent/runtime/checkpoint.rs +++ b/src/agent/runtime/checkpoint.rs @@ -70,8 +70,12 @@ const DEFAULT_EXCLUDES: &[&str] = &[ "Thumbs.db", ]; -/// 每个 turn 最多快照一次的工具 -const CHECKPOINT_TRIGGER_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"]; +/// 每个 turn 最多快照一次的工具。 +/// +/// 历史上执行器按此名单触发快照；现在由各工具通过 +/// `AgentTool::causes_file_changes()` 声明，本名单仅作为 +/// 注册表查询不到工具时的保守回退。 +const CHECKPOINT_FALLBACK_TOOLS: &[&str] = &["file_write", "file_edit", "run_bash"]; /// Checkpoint 元数据 #[derive(Debug, Clone)] @@ -465,9 +469,9 @@ impl CheckpointManager { )) } - /// 检查指定工具是否需要触发 checkpoint。 + /// 检查指定工具是否需要触发 checkpoint（注册表查询不到时的保守回退）。 pub fn should_checkpoint(tool_name: &str) -> bool { - CHECKPOINT_TRIGGER_TOOLS.contains(&tool_name) + CHECKPOINT_FALLBACK_TOOLS.contains(&tool_name) } /// 获取 repo 路径 diff --git a/src/agent/runtime/context.rs b/src/agent/runtime/context.rs index 8dc65e9..7a434e7 100644 --- a/src/agent/runtime/context.rs +++ b/src/agent/runtime/context.rs @@ -2,6 +2,10 @@ // // 上下文构建：加载历史消息、注入系统提示词、添加用户消息、 // 从数据库恢复持久化的任务状态。 +// +// 上下文快照回放：若上一 turn 结束时保存过压缩后的折叠上下文 +// （context_snapshot 事件），则加载"快照消息 + id > base 的增量消息"， +// 避免从原始消息重建后再次触发 LLM 摘要压缩（重复付费且信息有损）。 use sqlx::SqlitePool; use tracing::info; @@ -9,22 +13,51 @@ use tracing::info; use crate::clients::llm::{ChatMessage, MessageRole}; use super::session; +use super::session_events; -/// 构建初始 LLM 上下文：加载历史 → 插入系统提示词 → 添加用户消息 → 恢复任务状态。 +/// 构建初始 LLM 上下文。 +/// +/// 加载顺序：上下文快照（若有）+ 增量历史 → 插入系统提示词 → +/// 动态上下文快照（内容变更时才追加并持久化）→ 用户消息 → 任务恢复。 pub async fn build_initial_context( db: &SqlitePool, session_id: &str, system_prompt: &str, question: &str, - _turn_index: i32, + turn_index: i32, + dynamic_context: Option<(String, u64)>, + image_context: Option<&str>, ) -> anyhow::Result> { - let mut messages = session::load_history_for_llm(db, session_id).await?; + let mut messages = load_folded_history(db, session_id).await?; // 注入系统提示词（如果历史中没有） if messages.is_empty() || messages[0].role != MessageRole::System { messages.insert(0, ChatMessage::system(system_prompt)); } + // 动态上下文快照（durable user-role 消息）：哈希与上次持久化的一致时跳过。 + // 这是 KV-cache 纪律的另一半——易变内容不进 system prompt（否则任一 + // 变化都使整条前缀缓存失效），而以追加式快照进入历史，字节一旦写入 + // 就永不变更。 + if let Some((text, hash)) = dynamic_context { + let last_hash = last_dynamic_context_hash(db, session_id).await; + if last_hash.as_deref() != Some(hash.to_string().as_str()) { + let reminder = + ChatMessage::user(format!("\n{}\n", text)); + // 先持久化再注入（模型可见 ⟺ 已日志化） + persist_dynamic_context(db, session_id, turn_index, &reminder, hash).await; + messages.push(reminder); + } + } + + // 图片上下文：以 system-reminder 形式注入在用户问题之前 + if let Some(img_ctx) = image_context { + messages.push(ChatMessage::user(format!( + "\n{}\n", + img_ctx + ))); + } + // 添加用户消息 messages.push(ChatMessage::user(question)); @@ -36,6 +69,29 @@ pub async fn build_initial_context( Ok(messages) } +/// 加载折叠后的历史消息：优先回放上下文快照 + 增量；无快照时全量加载。 +async fn load_folded_history( + db: &SqlitePool, + session_id: &str, +) -> anyhow::Result> { + if let Some((snapshot_messages, base_id)) = + session_events::load_context_snapshot(db, session_id).await + { + let incremental = session_events::load_messages_after(db, session_id, base_id).await?; + info!( + "[Context] 回放上下文快照: {} 条快照消息 + {} 条增量消息 (base_id={})", + snapshot_messages.len(), + incremental.len(), + base_id + ); + let mut messages = snapshot_messages; + messages.extend(incremental); + Ok(messages) + } else { + session::load_history_for_llm(db, session_id).await + } +} + /// 从 agent_tasks 表恢复任务状态，返回格式化的提醒文本。 /// /// 如果表不存在或没有任务记录，返回 None。 @@ -82,3 +138,44 @@ async fn restore_tasks_from_db(db: &SqlitePool, session_id: &str) -> Option Option { + let metadata: Option = sqlx::query_scalar( + "SELECT metadata FROM agent_messages WHERE session_id = ? AND role = 'user' AND active = 1 AND metadata LIKE '%\"dynamic_context\"%' ORDER BY id DESC LIMIT 1", + ) + .bind(session_id) + .fetch_optional(db) + .await + .ok() + .flatten(); + + metadata + .and_then(|m| serde_json::from_str::(&m).ok()) + .and_then(|v| v.get("context_hash")?.as_str().map(|s| s.to_string())) +} + +/// 持久化动态上下文快照（metadata 携带哈希供下轮比较） +async fn persist_dynamic_context( + db: &SqlitePool, + session_id: &str, + turn_index: i32, + msg: &ChatMessage, + hash: u64, +) { + use crate::agent::engine::{DbMessageSink, MessageSink as _}; + let sink = DbMessageSink { + db: db.clone(), + session_id: session_id.to_string(), + agent_name: "lead".to_string(), + fixed_metadata: None, + }; + sink.save( + turn_index, + 0, + msg, + None, + Some(serde_json::json!({ "dynamic_context": true, "context_hash": hash.to_string() })), + ) + .await; +} diff --git a/src/agent/runtime/duplicate_detector.rs b/src/agent/runtime/duplicate_detector.rs index d14b497..d79bcec 100644 --- a/src/agent/runtime/duplicate_detector.rs +++ b/src/agent/runtime/duplicate_detector.rs @@ -1,27 +1,95 @@ // src/agent/runtime/duplicate_detector.rs // -// 同质调用检测器：检测连续重复的工具调用，防止死循环。 +// 同质调用检测器：检测重复的工具调用，防止死循环。 +// +// 历史实现是单槽（只记住上一次调用），A/B 交替死循环 +// （call A → call B → call A → call B ...）检测不到。 +// 现改为"连续计数 + 滑动窗口计数"双通道： +// - 连续通道：连续相同调用 ≥ threshold（原语义，立即触发） +// - 窗口通道：最近 threshold*2 次调用中同一调用出现 ≥ threshold+2 次 +// （捕获 A/B 交替；门槛略高以放过合理的重复只读调用） + +use std::collections::VecDeque; /// 同质调用检测器 #[derive(Debug, Default)] pub struct DuplicateDetector { - last_call: Option<(String, String)>, // (tool_name, arguments) - consecutive_count: usize, + /// 最近调用的滑动窗口（name, arguments） + recent: VecDeque<(String, String)>, } impl DuplicateDetector { /// 记录一次调用，返回是否检测到死循环 pub fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool { let key = (tool_name.to_string(), arguments.to_string()); - if self.last_call.as_ref() == Some(&key) { - self.consecutive_count += 1; - if self.consecutive_count >= threshold { - return true; - } - } else { - self.last_call = Some(key); - self.consecutive_count = 1; + let window = (threshold * 2).max(4); + + self.recent.push_back(key.clone()); + while self.recent.len() > window { + self.recent.pop_front(); + } + + // 连续通道：尾部连续出现次数 + let consecutive = self.recent.iter().rev().take_while(|k| *k == &key).count(); + if threshold >= 2 && consecutive >= threshold { + return true; + } + + // 窗口通道：同一调用在窗口内出现次数（捕获 A/B 交替）。 + // 门槛与 threshold 相同：纯交替下窗口内单键最多出现 window/2 = + // threshold 次，恰好可达；合理的间隔重复（≤ threshold-1 次）不受影响。 + let occurrences = self.recent.iter().filter(|k| **k == key).count(); + threshold >= 2 && occurrences >= threshold + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_consecutive_duplicates_detected() { + let mut det = DuplicateDetector::default(); + assert!(!det.record("search_papers", "q=1", 3)); + assert!(!det.record("search_papers", "q=1", 3)); + assert!(det.record("search_papers", "q=1", 3)); + } + + #[test] + fn test_alternating_loop_detected() { + // A/B 交替：单槽实现检测不到，窗口通道应捕获 + let mut det = DuplicateDetector::default(); + let mut tripped = false; + for i in 0..6 { + let a = det.record("tool_a", "{}", 3); + let b = det.record("tool_b", "{}", 3); + if a || b { + // 交替 5 次出现（threshold+2 = 5）后触发 + assert!(i >= 1, "不应过早触发 (round {})", i); + tripped = true; + break; + } + } + assert!(tripped, "A/B 交替死循环应被检测到"); + } + + #[test] + fn test_spaced_repeats_not_detected() { + // 合理的重复只读调用（间隔其他调用）不应触发 + let mut det = DuplicateDetector::default(); + det.record("read_file", "a.rs", 3); + det.record("grep_files", "pat", 3); + det.record("read_file", "b.rs", 3); + det.record("glob_files", "*.rs", 3); + let tripped = det.record("read_file", "a.rs", 3); + assert!(!tripped); + } + + #[test] + fn test_different_args_not_duplicates() { + let mut det = DuplicateDetector::default(); + for q in ["q=1", "q=2", "q=3", "q=4", "q=5", "q=6", "q=7", "q=8"] { + assert!(!det.record("search_papers", q, 3)); } - false } } diff --git a/src/agent/runtime/executor/helpers.rs b/src/agent/runtime/executor/helpers.rs index dec05d8..385753f 100644 --- a/src/agent/runtime/executor/helpers.rs +++ b/src/agent/runtime/executor/helpers.rs @@ -1,12 +1,10 @@ // src/agent/runtime/executor/helpers.rs // // 执行器辅助类型与函数：PreparedCall, ToolResultMessage, ToolExecutionResult, -// execute_single_tool, process_single_result, save_tool_message_sync。 +// execute_single_tool, process_single_result。 -use sqlx::SqlitePool; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tokio::sync::mpsc; use tracing::warn; use crate::clients::llm::ChatMessage; @@ -15,6 +13,7 @@ use super::AgentStreamEvent; use crate::agent::hooks::{ event_label, HookRegistry, PostToolUseContext, PostToolUseFailureContext, }; +use crate::agent::tools::ToolOutput; #[derive(Debug, Clone)] pub struct PreparedCall { @@ -42,8 +41,6 @@ pub struct ToolExecutionResult { } /// 执行单个工具调用（含超时和取消检测）。 -/// -/// 从原 `execute_parallel` 的闭包提取，供分区后的批次执行复用。 pub(super) async fn execute_single_tool( tool_opt: Option<&dyn crate::agent::tools::AgentTool>, args: serde_json::Value, @@ -51,10 +48,10 @@ pub(super) async fn execute_single_tool( cancelled: &Arc, timeout_dur: std::time::Duration, tool_name: &str, -) -> crate::agent::tools::ToolOutput { +) -> ToolOutput { let tool = match tool_opt { Some(t) => t, - None => return crate::agent::tools::ToolOutput::error(format!("未知工具: {}", tool_name)), + None => return ToolOutput::error(format!("未知工具: {}", tool_name)), }; let interrupt_behavior = tool.interrupt_behavior(); @@ -76,7 +73,7 @@ pub(super) async fn execute_single_tool( res = tokio::time::timeout(timeout_dur, tool_fut) => { match res { Ok(output) => output, - Err(_) => crate::agent::tools::ToolOutput::error(format!( + Err(_) => ToolOutput::error(format!( "工具 {} 执行超时（{}秒）", tool_name, timeout_dur.as_secs() @@ -84,50 +81,50 @@ pub(super) async fn execute_single_tool( } } _ = cancel_fut => { - crate::agent::tools::ToolOutput::error("执行已被用户取消") + ToolOutput::error("执行已被用户取消") } } } /// 处理单个工具执行结果（SSE 事件、PostToolUse hooks、持久化）。 -/// -/// 从原 `execute_parallel` 的结果处理循环提取。 #[allow(clippy::too_many_arguments)] pub(super) async fn process_single_result( tool_call_id: &str, tool_name: &str, tool_args: &serde_json::Value, - output: &crate::agent::tools::ToolOutput, + output: &ToolOutput, cancelled_flag: bool, exec_start: std::time::Instant, - tx: &mpsc::UnboundedSender, - hook_registry: &HookRegistry, - library_dir: &std::path::Path, - sid: &str, - agent_name: &str, - step: usize, - max_output_chars: usize, + ctx: &super::ExecutorContext<'_>, tool_messages: &mut Vec, additional_contexts: &mut Vec, - db: &SqlitePool, - turn_index: i32, - tool_registry: &crate::agent::tools::ToolRegistry, ) { - use crate::agent::tools::persist::maybe_persist_tool_result; + use crate::agent::tools::persist::maybe_persist_tool_result_for; let elapsed_ms = exec_start.elapsed().as_millis() as u64; + let step = ctx.step; + let tap = &ctx.tap; - let (is_internal, display_name) = if let Some(tool) = tool_registry.get(tool_name) { + let empty_hooks; + let hook_registry: &HookRegistry = match ctx.hook_registry { + Some(h) => h, + None => { + empty_hooks = HookRegistry::new(); + &empty_hooks + } + }; + + let (is_internal, display_name) = if let Some(tool) = ctx.tool_registry.get(tool_name) { (tool.is_internal(), tool.display_name().to_string()) } else { (false, tool_name.to_string()) }; // SSE 事件 — 立即推送到前端 - let _ = tx.send(AgentStreamEvent::ToolResult { + tap.send(AgentStreamEvent::ToolResult { tool_call_id: tool_call_id.to_string(), - name: tool_name.to_string(), - display_name, + name: tap.display(tool_name), + display_name: tap.display(&display_name), output: output.content.clone(), is_error: output.is_error, metadata: output.metadata.clone(), @@ -137,23 +134,30 @@ pub(super) async fn process_single_result( // 输出处理：小结果直接传递，大结果持久化到磁盘并返回 stub // 但对于已从磁盘读取内容的工具（如 read_file），跳过持久化以防止级联 - let tool_results_dir = library_dir.join(".agent").join("tool-results"); + let tool_results_dir = ctx + .app_state + .config + .storage + .library_dir + .join(".agent") + .join("tool-results"); let (processed_content, _persisted_path) = if output.skip_persist { (output.content.clone(), None) } else { - maybe_persist_tool_result( + maybe_persist_tool_result_for( &output.content, tool_call_id, - max_output_chars, + ctx.max_output_chars, &tool_results_dir, + tool_name, ) .await }; // PostToolUse hook let post_ctx = PostToolUseContext { - session_id: sid.to_string(), - agent_name: agent_name.to_string(), + session_id: ctx.session_id.to_string(), + agent_name: ctx.agent_name.to_string(), tool_name: tool_name.to_string(), tool_args: tool_args.clone(), output_content: processed_content.clone(), @@ -164,9 +168,14 @@ pub(super) async fn process_single_result( let post_result = hook_registry.run_post_tool_use(&post_ctx).await; let final_content = post_result.final_content; - // 非可信内容包裹（间接 prompt 注入防御） - let llm_content = - crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content); + // 非可信内容包裹（间接 prompt 注入防御）。 + // 是否包裹由工具通过 AgentTool::untrusted_output 声明； + // 未知工具回退到名字启发式（mcp__*/web_* 前缀）。 + let llm_content = if ctx.tool_registry.untrusted_output(tool_name) { + crate::agent::runtime::untrusted::wrap_untrusted_content(tool_name, &final_content) + } else { + final_content + }; // 收集 PostToolUse hook 注入的上下文 if !post_result.tagged_contexts.is_empty() { @@ -179,8 +188,8 @@ pub(super) async fn process_single_result( )); } } else { - for ctx in &post_result.additional_contexts { - additional_contexts.push(ctx.clone()); + for c in &post_result.additional_contexts { + additional_contexts.push(c.clone()); } } @@ -200,8 +209,8 @@ pub(super) async fn process_single_result( // PostToolUseFailure hook if output.is_error { let failure_ctx = PostToolUseFailureContext { - session_id: sid.to_string(), - agent_name: agent_name.to_string(), + session_id: ctx.session_id.to_string(), + agent_name: ctx.agent_name.to_string(), tool_name: tool_name.to_string(), tool_args: tool_args.clone(), error_message: output.content.clone(), @@ -217,51 +226,25 @@ pub(super) async fn process_single_result( // 发送给 LLM 使用包裹后的内容（安全防御） let chat_message = ChatMessage::tool_result(tool_call_id, &llm_content); - // 持久化到数据库（fire-and-forget） - save_tool_message_sync(db, sid, turn_index, step, &chat_message); + // 持久化到数据库（fire-and-forget；无 sink 时跳过）。 + // canonical value 一并写入 metadata（重放/审计消费结构化结果）。 + if let Some(sink) = &ctx.sink { + let sink = sink.clone(); + let msg = chat_message.clone(); + let turn_index = ctx.turn_index; + let step_i = ctx.step; + let extra = output + .value + .clone() + .map(|v| serde_json::json!({ "canonical_value": v })); + tokio::spawn(async move { + sink.save(turn_index, step_i as i32, &msg, None, extra) + .await; + }); + } tool_messages.push(ToolResultMessage { chat_message, was_error: output.is_error, }); } - -/// 同步保存 tool 角色消息到数据库。 -pub(super) fn save_tool_message_sync( - db: &SqlitePool, - session_id: &str, - turn_index: i32, - step_index: usize, - msg: &ChatMessage, -) { - let db_clone = db.clone(); - let session_id = session_id.to_string(); - let content = msg.text().unwrap_or("").to_string(); - let tool_call_id = msg.tool_call_id.clone(); - // 提前序列化，避免闭包内的生命周期问题 - let metadata_str = - serde_json::to_string(&serde_json::json!({ "role": "tool" })).unwrap_or_default(); - let raw_json = serde_json::to_string(&msg).unwrap_or_default(); - // fire-and-forget: tool 消息保存失败不影响主流程 - tokio::spawn(async move { - let token_count = content.len() as i32 / 4; - if let Err(e) = sqlx::query( - "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_call_id, token_count, metadata, raw_json, agent_name) \ - VALUES (?, ?, ?, 'tool', ?, ?, ?, ?, ?, ?)", - ) - .bind(&session_id) - .bind(turn_index) - .bind(step_index as i32) - .bind(&content) - .bind(&tool_call_id) - .bind(token_count) - .bind(&metadata_str) - .bind(&raw_json) - .bind("lead") - .execute(&db_clone) - .await - { - warn!("[Executor] 保存 tool 消息失败（非致命）: {}", e); - } - }); -} diff --git a/src/agent/runtime/executor/mod.rs b/src/agent/runtime/executor/mod.rs index 77d16b8..3a6ac1b 100644 --- a/src/agent/runtime/executor/mod.rs +++ b/src/agent/runtime/executor/mod.rs @@ -1,18 +1,21 @@ // src/agent/runtime/executor/mod.rs // // 工具调用验证与并行执行器。 +// +// 执行参数通过 ExecutorContext 组合传递（历史实现有约 20 个位置参数）。 +// 行为注入点： +// - ask_policy: AskUser 权限请求是交互等待（主代理）还是 fail-closed 自动拒绝（子代理/队友） +// - tap: SSE 事件出口（None = 静默；prefix = "[sub]"） +// - sink: 消息持久化出口（None = 不落库） mod helpers; // Re-export 公共类型 pub use helpers::{PreparedCall, ToolExecutionResult, ToolResultMessage}; -use futures_util::stream::FuturesUnordered; -use futures_util::StreamExt; -use sqlx::SqlitePool; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::oneshot; use tracing::{info, warn}; use crate::api::{AppState, PendingPermission}; @@ -26,26 +29,23 @@ use super::partitioner::ToolPartitioner; use super::permission::{PermissionChecker, PermissionResult}; use super::permission_explainer::explain_permission; use super::{AgentStreamEvent, DuplicateDetector}; +use crate::agent::engine::{AskPolicy, EventTap, MessageSink}; use crate::agent::hooks::{event_label, HookRegistry, PreToolUseContext}; -use crate::agent::tools::{ToolContext, ToolRegistry}; +use crate::agent::tools::ToolRegistry; -use helpers::{execute_single_tool, process_single_result, save_tool_message_sync}; +use helpers::{execute_single_tool, process_single_result}; /// 验证工具调用：死循环检测 + 参数解析。 /// /// 返回 (prepared_calls, has_duplicate)。 /// 死循环或参数无效时，错误消息直接注入到 messages。 -#[allow(clippy::too_many_arguments)] pub fn validate_and_prepare( tool_calls: &[ToolCall], duplicate_detector: &mut DuplicateDetector, duplicate_threshold: usize, messages: &mut Vec, tool_registry: &ToolRegistry, - tx: &mpsc::UnboundedSender, - db: &SqlitePool, - session_id: &str, - turn_index: i32, + tap: &EventTap, step: usize, ) -> (Vec, bool) { let mut prepared_calls: Vec = Vec::new(); @@ -68,7 +68,7 @@ pub fn validate_and_prepare( "[Executor] 检测到死循环：{} 连续调用 {} 次", tool_name, duplicate_threshold ); - let _ = tx.send(AgentStreamEvent::Error { + tap.send(AgentStreamEvent::Error { message: format!("检测到工具 {} 的重复调用，已自动终止循环。", tool_name), }); let error_msg = ChatMessage::tool_result( @@ -94,10 +94,10 @@ pub fn validate_and_prepare( } else { (false, tool_name.clone()) }; - let _ = tx.send(AgentStreamEvent::ToolResult { + tap.send(AgentStreamEvent::ToolResult { tool_call_id: call_id.clone(), name: tool_name.clone(), - display_name, + display_name: tap.display(&display_name), output: error_output.clone(), is_error: true, metadata: serde_json::json!({}), @@ -105,7 +105,6 @@ pub fn validate_and_prepare( is_internal, }); let tool_msg = ChatMessage::tool_result(&call_id, &error_output); - save_tool_message_sync(db, session_id, turn_index, step, &tool_msg); messages.push(tool_msg); continue; } @@ -121,36 +120,48 @@ pub fn validate_and_prepare( (prepared_calls, has_duplicate) } +/// 执行器上下文 — 组合全部行为注入点（替代约 20 个位置参数）。 +pub struct ExecutorContext<'a> { + pub tool_registry: &'a ToolRegistry, + pub app_state: Arc, + /// hooks 缺省时使用空注册表语义（teammate 等无人值守场景） + pub hook_registry: Option<&'a HookRegistry>, + pub permission_checker: Option<&'a PermissionChecker>, + pub session_checker: Option<&'a PermissionChecker>, + pub denial_tracker: Option<&'a std::sync::Mutex>, + pub checkpoint_manager: Option<&'a CheckpointManager>, + pub db: &'a sqlx::SqlitePool, + pub session_id: &'a str, + pub agent_name: &'a str, + pub turn_index: i32, + pub step: usize, + pub tool_timeout_secs: u64, + pub max_output_chars: usize, + pub read_file_state: Arc>, + pub enable_thinking: bool, + pub additional_allowed_dirs: Vec, + /// AskUser 权限请求处理策略 + pub ask_policy: AskPolicy, + /// 取消观察：会话取消表（交互式）或外部原子标志 + pub cancel: Arc, + pub interactive_cancel: bool, + pub cancel_session_id: Option, + pub tap: EventTap, + pub sink: Option>, +} + /// 并行执行所有准备好的工具调用。 /// /// 流程： -/// 1. 权限检查（deny 规则阻止不可执行工具） -/// 2. 发送 ToolCall SSE 事件 -/// 3. 运行 PreToolUse hooks +/// 1. Hardline 预检查（不可绕过的参数级拒绝） +/// 2. 权限检查（deny 规则拦截；ask 按策略交互等待或 fail-closed 拒绝） +/// 3. 发送 ToolCall SSE 事件 + PreToolUse hooks /// 4. 工具分区 + 并行执行（并发安全工具一批并行，不安全工具单独串行） /// 5. 收集结果、发送 ToolResult SSE、运行 PostToolUse hooks /// 6. 返回 ToolResultMessage 列表供调用方推入 messages -#[allow(clippy::too_many_arguments)] pub async fn execute_parallel( prepared_calls: &[PreparedCall], - tool_registry: &ToolRegistry, - app_state: Arc, - hook_registry: &HookRegistry, - permission_checker: Option<&PermissionChecker>, - session_permission_checker: Option<&PermissionChecker>, - denial_tracker: Option<&std::sync::Mutex>, - checkpoint_manager: Option<&std::sync::Arc>, - tx: &mpsc::UnboundedSender, - db: &SqlitePool, - session_id: &str, - agent_name: &str, - turn_index: i32, - step: usize, - tool_timeout_secs: u64, - max_output_chars: usize, - read_file_state: Arc>, - enable_thinking: bool, - additional_allowed_dirs: Vec, + ctx: &ExecutorContext<'_>, ) -> ToolExecutionResult { if prepared_calls.is_empty() { return ToolExecutionResult { @@ -162,7 +173,72 @@ pub async fn execute_parallel( }; } - let sid = session_id.to_string(); + let tool_registry = ctx.tool_registry; + let tap = &ctx.tap; + let step = ctx.step; + let sid = ctx.session_id.to_string(); + let _turn_index = ctx.turn_index; + let empty_hooks; + let hook_registry: &HookRegistry = match ctx.hook_registry { + Some(h) => h, + None => { + empty_hooks = HookRegistry::new(); + &empty_hooks + } + }; + + /// 工具消息的统一落库出口（无 sink 时跳过） + fn persist_tool_message( + sink: &Option>, + turn_index: i32, + step: usize, + msg: &ChatMessage, + ) { + if let Some(sink) = sink { + let sink = sink.clone(); + let msg = msg.clone(); + tokio::spawn(async move { + sink.save(turn_index, step as i32, &msg, None, None).await; + }); + } + } + + /// 拒绝结果的统一出口：SSE 事件 + 消息落库 + 拒绝追踪 + async fn record_denial( + ctx: &ExecutorContext<'_>, + tap: &EventTap, + tool_call_id: &str, + tool_name: &str, + err_output: &str, + metadata: serde_json::Value, + tool_messages: &mut Vec, + ) { + let (is_internal, display_name) = match ctx.tool_registry.get(tool_name) { + Some(tool) => (tool.is_internal(), tool.display_name().to_string()), + None => (false, tool_name.to_string()), + }; + tap.send(AgentStreamEvent::ToolResult { + tool_call_id: tool_call_id.to_string(), + name: tool_name.to_string(), + display_name: tap.display(&display_name), + output: err_output.to_string(), + is_error: true, + metadata, + step: ctx.step, + is_internal, + }); + let err_msg = ChatMessage::tool_result(tool_call_id, err_output); + persist_tool_message(&ctx.sink, ctx.turn_index, ctx.step, &err_msg); + tool_messages.push(ToolResultMessage { + chat_message: err_msg, + was_error: true, + }); + if let Some(dt) = ctx.denial_tracker { + if let Ok(mut tracker) = dt.lock() { + tracker.record_denial(); + } + } + } // Phase 1: 发送 ToolCall SSE 事件 for prep in prepared_calls { @@ -171,10 +247,10 @@ pub async fn execute_parallel( } else { (false, prep.tool_name.clone()) }; - let _ = tx.send(AgentStreamEvent::ToolCall { + tap.send(AgentStreamEvent::ToolCall { id: prep.tool_call_id.clone(), - name: prep.tool_name.clone(), - display_name, + name: tap.display(&prep.tool_name), + display_name: tap.display(&display_name), arguments: prep.args.clone(), step, is_internal, @@ -203,14 +279,12 @@ pub async fn execute_parallel( prep.tool_name, reason ); } - // 收集所有阻塞错误详情（含多个 hook 同时 block 的情况） for be in &result.blocking_errors { hook_blocking_errors.push(format!( "[{}] 阻止 {}: {}", be.hook_name, prep.tool_name, be.reason )); } - // 收集 hook 的权限请求（保留完整信息用于 AskUser prompt） if let Some((permission, tool_name)) = result.permission_info() { info!( "[Executor] Hook 请求了工具 {} 的权限确认: {}", @@ -220,9 +294,7 @@ pub async fn execute_parallel( } else { hook_permission_info.push(None); } - // 使用 hook 可能修改后的参数 mutated_args.push(result.final_args); - // 收集所有 hook 注入的上下文（优先使用带来源标记的 tagged_contexts） if !result.tagged_contexts.is_empty() { for tc in &result.tagged_contexts { additional_contexts.push(format!( @@ -233,39 +305,25 @@ pub async fn execute_parallel( )); } } else { - for ctx in &result.additional_contexts { - additional_contexts.push(ctx.clone()); + for c in &result.additional_contexts { + additional_contexts.push(c.clone()); } } } - // Phase 2.5: 权限检查 — PermissionChecker 规则引擎拦截被拒绝的工具。 + // Phase 2.5: Hardline 预检查 + 权限检查。 // 被拒绝的工具直接注入错误 result，不进入执行队列。 let mut tool_messages: Vec = Vec::new(); let mut denied_indices: std::collections::HashSet = std::collections::HashSet::new(); // ── Hardline 预检查（在任何模式下都不可绕过）── - // 在 PermissionChecker 之前执行，确保 hardline 规则始终生效。 + // 检查逻辑由工具自身通过 AgentTool::hardline_check 声明（按参数路由， + // 不再按工具名字符串匹配）。 for (i, prep) in prepared_calls.iter().enumerate() { - let hardline_result = match prep.tool_name.as_str() { - "run_bash" => { - if let Some(cmd) = prep.args.get("command").and_then(|v| v.as_str()) { - hardline::check_command(cmd) - } else { - hardline::HardlineResult::allowed() - } - } - "file_write" | "file_edit" => { - if let Some(path) = prep.args.get("file_path").and_then(|v| v.as_str()) { - hardline::check_dangerous_path(path) - } else if let Some(path) = prep.args.get("path").and_then(|v| v.as_str()) { - hardline::check_dangerous_path(path) - } else { - hardline::HardlineResult::allowed() - } - } - _ => hardline::HardlineResult::allowed(), - }; + let hardline_result = tool_registry + .get(&prep.tool_name) + .and_then(|tool| tool.hardline_check(&prep.args)) + .unwrap_or_else(hardline::HardlineResult::allowed); if hardline_result.blocked { warn!( @@ -275,48 +333,30 @@ pub async fn execute_parallel( hardline_result.reason ); let err_output = hardline_result.reason.clone(); - let (is_internal, display_name) = if let Some(tool) = tool_registry.get(&prep.tool_name) - { - (tool.is_internal(), tool.display_name().to_string()) - } else { - (false, prep.tool_name.clone()) - }; - let _ = tx.send(AgentStreamEvent::ToolResult { - tool_call_id: prep.tool_call_id.clone(), - name: prep.tool_name.clone(), - display_name, - output: err_output.clone(), - is_error: true, - metadata: serde_json::json!({ + record_denial( + ctx, + tap, + &prep.tool_call_id, + &prep.tool_name, + &err_output, + serde_json::json!({ "hardline_blocked": true, "hardline_category": hardline_result.category, }), - step, - is_internal, - }); - let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output); - save_tool_message_sync(db, &sid, turn_index, step, &err_msg); - tool_messages.push(ToolResultMessage { - chat_message: err_msg, - was_error: true, - }); - // 记录拒绝追踪 - if let Some(dt) = denial_tracker { - if let Ok(mut tracker) = dt.lock() { - tracker.record_denial(); - } - } + &mut tool_messages, + ) + .await; denied_indices.insert(i); } } - if let Some(checker) = permission_checker { + if let Some(checker) = ctx.permission_checker { for (i, prep) in prepared_calls.iter().enumerate() { + // ── 权限决策合并（单调不变量：Deny 粘滞，只能收紧不能放松）── let mut perm_result = checker.check(&prep.tool_name, Some(&prep.args)); perm_result = checker.apply_mode(perm_result, &prep.tool_name); // Hook PermissionRequired — 若 Checker 返回 Allowed，升级为 Ask - // 使用 hook 提供的具体权限描述替换泛型消息 if let Some(Some((ref perm_desc, _))) = hook_permission_info.get(i) { if perm_result.is_allowed() { perm_result = PermissionResult::AskUser { @@ -330,14 +370,12 @@ pub async fn execute_parallel( } } - // 工具级 check_permissions() — 在 PermissionChecker 结果基础上叠加 - // PermissionChecker Deny/Ask 优先，工具级规则在 Allow 时可升级为 Ask + // 工具级 check_permissions() — 在已有结果基础上收紧 if let Some(tool) = tool_registry.get(&prep.tool_name) { let tool_rules = tool.check_permissions(&prep.args); for tool_rule in &tool_rules { match tool_rule { crate::agent::tools::PermissionRule::Deny { reason, .. } => { - // 工具级 Deny 仅在 PermissionChecker 未 Deny 时生效 if !perm_result.is_denied() { perm_result = PermissionResult::Denied { reason: reason.clone(), @@ -345,7 +383,6 @@ pub async fn execute_parallel( } } crate::agent::tools::PermissionRule::Ask { message, .. } => { - // 工具级 Ask：若 PermissionChecker 返回 Allowed，升级为 Ask if perm_result.is_allowed() { perm_result = PermissionResult::AskUser { message: message.clone(), @@ -357,67 +394,57 @@ pub async fn execute_parallel( } } - // 会话级权限检查（API 动态添加的规则，优先级高于环境变量规则） - if let Some(session_checker) = session_permission_checker { - let session_result = session_checker.check(&prep.tool_name, Some(&prep.args)); - // 会话规则结果覆盖或升级 - match session_result { - PermissionResult::Denied { reason } => { - // 会话 Deny 强制覆盖 - perm_result = PermissionResult::Denied { reason }; - } - PermissionResult::AskUser { message } => { - // 会话 Ask 在 Allow 时升级 - if perm_result.is_allowed() { - perm_result = PermissionResult::AskUser { message }; - } - } - PermissionResult::Allowed => { - // 会话 Allow 仅覆盖 Allowed，保持 Deny/AskUser 不变 - // 避免覆盖工具级 check_permissions() 升级的 AskUser - } - } + // 会话级权限检查（API 动态添加的规则，只能收紧） + if let Some(session_checker) = ctx.session_checker { + perm_result = super::permission::tighten( + perm_result, + session_checker.check(&prep.tool_name, Some(&prep.args)), + ); } match perm_result { PermissionResult::Denied { reason } => { warn!( - "[Executor] PermissionChecker 拒绝了工具 {}: {}", + "[Executor] 权限检查拒绝了工具 {}: {}", prep.tool_name, reason ); let err_output = format!("工具 {} 被权限规则拒绝执行: {}", prep.tool_name, reason); - let (is_internal, display_name) = - if let Some(tool) = tool_registry.get(&prep.tool_name) { - (tool.is_internal(), tool.display_name().to_string()) - } else { - (false, prep.tool_name.clone()) - }; - let _ = tx.send(AgentStreamEvent::ToolResult { - tool_call_id: prep.tool_call_id.clone(), - name: prep.tool_name.clone(), - display_name, - output: err_output.clone(), - is_error: true, - metadata: serde_json::json!({}), - step, - is_internal, - }); - let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output); - save_tool_message_sync(db, &sid, turn_index, step, &err_msg); - tool_messages.push(ToolResultMessage { - chat_message: err_msg, - was_error: true, - }); - // 记录拒绝追踪 - if let Some(dt) = denial_tracker { - if let Ok(mut tracker) = dt.lock() { - tracker.record_denial(); - } - } + record_denial( + ctx, + tap, + &prep.tool_call_id, + &prep.tool_name, + &err_output, + serde_json::json!({}), + &mut tool_messages, + ) + .await; denied_indices.insert(i); } PermissionResult::AskUser { message } => { + // 无人值守上下文（子代理/队友）无用户可问：fail-closed 自动拒绝 + if ctx.ask_policy == AskPolicy::AutoDeny { + let reason = + match crate::agent::engine::auto_deny_ask(&prep.tool_name, &message) { + PermissionResult::Denied { reason } => reason, + _ => "需要用户确认但在无人值守上下文中不可用".to_string(), + }; + let err_output = format!("工具 {} 被拒绝执行: {}", prep.tool_name, reason); + record_denial( + ctx, + tap, + &prep.tool_call_id, + &prep.tool_name, + &err_output, + serde_json::json!({ "auto_denied_ask": true }), + &mut tool_messages, + ) + .await; + denied_indices.insert(i); + continue; + } + info!( "[Executor] PermissionChecker 请求用户确认工具 {}: {}", prep.tool_name, message @@ -427,8 +454,7 @@ pub async fn execute_parallel( let permission_exp = explain_permission(&prep.tool_name, &prep.args); let explanation_json = serde_json::to_value(&permission_exp).ok(); - // 发送权限请求 SSE 事件 - let _ = tx.send(AgentStreamEvent::PermissionRequest { + tap.send(AgentStreamEvent::PermissionRequest { tool_call_id: prep.tool_call_id.clone(), tool_name: prep.tool_name.clone(), message: message.clone(), @@ -439,17 +465,14 @@ pub async fn execute_parallel( // 创建 oneshot 通道等待用户响应 let (resp_tx, resp_rx) = oneshot::channel(); let perm_id = uuid::Uuid::new_v4().to_string(); - let tc_id = prep.tool_call_id.clone(); - let t_name = prep.tool_name.clone(); - // 存储待处理的权限请求 { - let mut perms = app_state.session.pending_permissions.lock().await; + let mut perms = ctx.app_state.session.pending_permissions.lock().await; perms.insert( perm_id.clone(), PendingPermission { - tool_call_id: tc_id.clone(), - tool_name: t_name.clone(), + tool_call_id: prep.tool_call_id.clone(), + tool_name: prep.tool_name.clone(), message: message.clone(), arguments: prep.args.clone(), response_tx: resp_tx, @@ -458,12 +481,11 @@ pub async fn execute_parallel( ); } - // 等待用户响应（120 秒超时） + // 等待用户响应（120 秒超时 = fail-closed 拒绝） let timeout_dur = std::time::Duration::from_secs(120); let perm_result = tokio::time::timeout(timeout_dur, resp_rx).await; - // 清理待处理的权限请求 - app_state + ctx.app_state .session .pending_permissions .lock() @@ -473,99 +495,61 @@ pub async fn execute_parallel( match perm_result { Ok(Ok(response)) if response.allowed => { info!("[Executor] 用户允许了工具 {} 的执行", prep.tool_name); - let _ = tx.send(AgentStreamEvent::PermissionResponse { + tap.send(AgentStreamEvent::PermissionResponse { tool_call_id: prep.tool_call_id.clone(), allowed: true, }); // 用户允许 → 重置连续拒绝计数 - if let Some(dt) = denial_tracker { + if let Some(dt) = ctx.denial_tracker { if let Ok(mut tracker) = dt.lock() { tracker.record_success(); } } } Ok(Ok(_response)) => { - // 用户拒绝 info!("[Executor] 用户拒绝了工具 {}", prep.tool_name); let err_output = format!("用户拒绝了工具 {} 的执行", prep.tool_name); - let (is_internal, display_name) = - if let Some(tool) = tool_registry.get(&prep.tool_name) { - (tool.is_internal(), tool.display_name().to_string()) - } else { - (false, prep.tool_name.clone()) - }; - let _ = tx.send(AgentStreamEvent::ToolResult { - tool_call_id: prep.tool_call_id.clone(), - name: prep.tool_name.clone(), - display_name, - output: err_output.clone(), - is_error: true, - metadata: serde_json::json!({}), - step, - is_internal, - }); - let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output); - save_tool_message_sync(db, &sid, turn_index, step, &err_msg); - tool_messages.push(ToolResultMessage { - chat_message: err_msg, - was_error: true, - }); - // 用户拒绝 → 记录拒绝追踪 - if let Some(dt) = denial_tracker { - if let Ok(mut tracker) = dt.lock() { - tracker.record_denial(); - } - } - denied_indices.insert(i); - let _ = tx.send(AgentStreamEvent::PermissionResponse { + tap.send(AgentStreamEvent::PermissionResponse { tool_call_id: prep.tool_call_id.clone(), allowed: false, }); + record_denial( + ctx, + tap, + &prep.tool_call_id, + &prep.tool_name, + &err_output, + serde_json::json!({}), + &mut tool_messages, + ) + .await; + denied_indices.insert(i); } _ => { - // 超时或通道关闭 + // 超时或通道关闭 → fail-closed warn!("[Executor] 权限请求超时或取消: {}", prep.tool_name); let err_output = format!("权限请求超时 (120s): {} 未获得用户确认", prep.tool_name); - let (is_internal, display_name) = - if let Some(tool) = tool_registry.get(&prep.tool_name) { - (tool.is_internal(), tool.display_name().to_string()) - } else { - (false, prep.tool_name.clone()) - }; - let _ = tx.send(AgentStreamEvent::ToolResult { - tool_call_id: prep.tool_call_id.clone(), - name: prep.tool_name.clone(), - display_name, - output: err_output.clone(), - is_error: true, - metadata: serde_json::json!({}), - step, - is_internal, - }); - let err_msg = ChatMessage::tool_result(&prep.tool_call_id, &err_output); - save_tool_message_sync(db, &sid, turn_index, step, &err_msg); - tool_messages.push(ToolResultMessage { - chat_message: err_msg, - was_error: true, - }); - // 超时 → 记录拒绝追踪 - if let Some(dt) = denial_tracker { - if let Ok(mut tracker) = dt.lock() { - tracker.record_denial(); - } - } - denied_indices.insert(i); - let _ = tx.send(AgentStreamEvent::PermissionResponse { + tap.send(AgentStreamEvent::PermissionResponse { tool_call_id: prep.tool_call_id.clone(), allowed: false, }); + record_denial( + ctx, + tap, + &prep.tool_call_id, + &prep.tool_name, + &err_output, + serde_json::json!({ "timeout": true }), + &mut tool_messages, + ) + .await; + denied_indices.insert(i); } } } PermissionResult::Allowed => { - // 工具被允许 → 重置连续拒绝计数 - if let Some(dt) = denial_tracker { + if let Some(dt) = ctx.denial_tracker { if let Ok(mut tracker) = dt.lock() { tracker.record_success(); } @@ -575,38 +559,46 @@ pub async fn execute_parallel( } } // if let Some(checker) - // Phase 3: 分区并行执行（参考 Claude Code partitionToolCalls + runTools）。 - // - // 改进：原实现将所有非拒绝工具放入单个 FuturesUnordered 无差别并发， - // 可能导致非并发安全工具（如 run_bash）错误地并行执行。 - // 新实现使用 ToolPartitioner 将工具按并发安全性分批： + // Phase 3: 分区并行执行。 + // ToolPartitioner 将工具按并发安全性分批： // - 连续的并发安全工具放入同一个并行批次（FuturesUnordered） // - 非并发安全工具独占一个串行批次（逐次执行） - // 批次内工具执行完成后立即推送 SSE 事件，不等待整个批次完成。 let cancelled = Arc::new(AtomicBool::new(false)); let cancel_flag = cancelled.clone(); - let app_state_ref = app_state.clone(); - let sid_ref = sid.clone(); + let app_state_ref = ctx.app_state.clone(); + let _sid_ref = sid.clone(); + let interactive = ctx.interactive_cancel; + let cancel_sid = ctx.cancel_session_id.clone(); + let external_flag = ctx.cancel.clone(); let cancel_handle = tokio::spawn(async move { loop { tokio::time::sleep(std::time::Duration::from_millis(250)).await; - if app_state_ref.session.cancelled_runs.contains_key(&sid_ref) { + let tripped = if interactive { + cancel_sid + .as_ref() + .is_some_and(|s| app_state_ref.session.cancelled_runs.contains_key(s)) + } else { + external_flag.load(Ordering::SeqCst) + }; + if tripped { cancel_flag.store(true, Ordering::SeqCst); return; } } }); - let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs); + let timeout_dur = std::time::Duration::from_secs(ctx.tool_timeout_secs); // ── Checkpoint 预触发：对文件变更类工具在执行前创建快照 ── - if let Some(ckpt) = checkpoint_manager { + // 由工具通过 AgentTool::causes_file_changes 声明 + if let Some(ckpt) = ctx.checkpoint_manager { let cwd = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")); - for prep in prepared_calls - .iter() - .filter(|p| CheckpointManager::should_checkpoint(&p.tool_name)) - { + for prep in prepared_calls.iter().filter(|p| { + tool_registry + .get(&p.tool_name) + .is_some_and(|t| t.causes_file_changes()) + }) { ckpt.ensure_checkpoint(&cwd, &format!("pre-{}", prep.tool_name)); } } @@ -624,7 +616,6 @@ pub async fn execute_parallel( let partitioner = ToolPartitioner::new(10); let batches = partitioner.partition(&non_denied_calls, tool_registry); - // 预设非拒绝工具中哪些原索引属于已拒绝列表（不会有，但安全起见） let original_index_of: std::collections::HashMap = non_denied .iter() .map(|(orig_idx, prep)| (prep.tool_call_id.clone(), *orig_idx)) @@ -641,7 +632,6 @@ pub async fn execute_parallel( let mut was_cancelled = false; // ── Phase 3c: 逐批次执行 ── - // 批次之间串行；并行批次内工具并发执行；串行批次内工具逐个执行。 for batch in &batches { if was_cancelled { break; @@ -649,7 +639,7 @@ pub async fn execute_parallel( if batch.is_parallel { // ── 并行批次：FuturesUnordered 并发执行 ── - let mut exec_futs: FuturesUnordered<_> = batch + let mut exec_futs: futures_util::stream::FuturesUnordered<_> = batch .calls .iter() .map(|prep| { @@ -662,14 +652,16 @@ pub async fn execute_parallel( .get(orig_idx) .cloned() .unwrap_or_else(|| prep.args.clone()); - let tool_ctx = - ToolContext::with_file_cache(app_state.clone(), read_file_state.clone()) - .with_sse_tx(tx.clone()) - .with_session_id(session_id.to_string()) - .with_thinking(enable_thinking) - .with_additional_dirs(additional_allowed_dirs.clone()) - .with_tool_call_id(prep.tool_call_id.clone()) - .with_max_output_chars(max_output_chars); + let tool_ctx = crate::agent::tools::ToolContext::with_file_cache( + ctx.app_state.clone(), + ctx.read_file_state.clone(), + ) + .with_sse_tx_opt(ctx.tap.sender()) + .with_session_id(sid.clone()) + .with_thinking(ctx.enable_thinking) + .with_additional_dirs(ctx.additional_allowed_dirs.clone()) + .with_tool_call_id(prep.tool_call_id.clone()) + .with_max_output_chars(ctx.max_output_chars); let cancelled = cancelled.clone(); let tool_opt = tool_registry.get(&tool_name); @@ -696,6 +688,7 @@ pub async fn execute_parallel( .collect(); // 渐进式处理：每个工具一完成就处理 + use futures_util::StreamExt; while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) = exec_futs.next().await { @@ -709,18 +702,9 @@ pub async fn execute_parallel( &output, cancelled_flag, exec_start, - tx, - hook_registry, - &app_state.config.storage.library_dir, - &sid, - agent_name, - step, - max_output_chars, + ctx, &mut tool_messages, &mut additional_contexts, - db, - turn_index, - tool_registry, ) .await; } @@ -736,12 +720,13 @@ pub async fn execute_parallel( .get(orig_idx) .cloned() .unwrap_or_else(|| prep.args.clone()); - let tool_ctx = - ToolContext::with_file_cache(app_state.clone(), read_file_state.clone()) - .with_sse_tx(tx.clone()) - .with_session_id(session_id.to_string()) - .with_thinking(enable_thinking) - .with_additional_dirs(additional_allowed_dirs.clone()); + let tool_ctx = crate::agent::tools::ToolContext::with_file_cache( + ctx.app_state.clone(), + ctx.read_file_state.clone(), + ) + .with_session_id(sid.clone()) + .with_thinking(ctx.enable_thinking) + .with_additional_dirs(ctx.additional_allowed_dirs.clone()); let tool_opt = tool_registry.get(&tool_name); let output = execute_single_tool( @@ -765,18 +750,9 @@ pub async fn execute_parallel( &output, cancelled_flag, exec_start, - tx, - hook_registry, - &app_state.config.storage.library_dir, - &sid, - agent_name, - step, - max_output_chars, + ctx, &mut tool_messages, &mut additional_contexts, - db, - turn_index, - tool_registry, ) .await; diff --git a/src/agent/runtime/mod.rs b/src/agent/runtime/mod.rs index d328633..3eb099f 100644 --- a/src/agent/runtime/mod.rs +++ b/src/agent/runtime/mod.rs @@ -28,8 +28,9 @@ pub mod permission; pub mod permission_explainer; pub mod permission_profile; pub mod session; +pub mod session_events; +pub mod session_registry; pub mod streaming; -pub mod streaming_executor; pub mod system_prompt; pub mod token_budget; pub mod untrusted; @@ -37,24 +38,19 @@ pub mod untrusted; use sqlx::SqlitePool; use std::sync::Arc; use tokio::sync::mpsc; -use tracing::{error, info, warn}; +use tracing::warn; use super::background::BgNotificationQueue; use super::compact; -use super::hooks::{ - HookRegistry, SessionStartContext, StepCompleteContext, UserPromptSubmitContext, -}; +use super::hooks::{HookRegistry, SessionStartContext, UserPromptSubmitContext}; use super::modes::{self, AgentMode, ModeRegistry}; use super::terminal::TurnTerminal; use super::tools::ToolRegistry; use crate::api::AppState; -use crate::clients::llm::{ChatMessage, LlmClient, MessageRole, StreamEvent}; +use crate::clients::llm::ChatMessage; -use self::error_recovery::{classify_error, ErrorKind, ErrorRecovery}; use self::session::SessionInfo; -use self::streaming::{StreamOutput, StreamStatus}; use self::system_prompt::SystemPromptCache; -use self::token_budget::TokenBudget; // Re-export 公共类型，保持外部引用路径不变 // （use crate::agent::runtime::{AgentConfig, AgentStreamEvent, ...}） @@ -62,6 +58,7 @@ use self::token_budget::TokenBudget; pub use self::config::AgentConfig; pub use self::duplicate_detector::DuplicateDetector; pub use self::events::{AgentMetrics, AgentStreamEvent}; +pub use self::session_registry::SessionRuntimeRegistry; // ── Agent Runtime ── @@ -93,14 +90,18 @@ pub struct AgentRuntime { mode: &'static AgentMode, /// 模式注册表（持有所有已注册模式） mode_registry: ModeRegistry, + /// 思考模式标志（AtomicBool：runtime 以 Arc 共享后仍可按请求覆盖） + thinking: std::sync::atomic::AtomicBool, } impl AgentRuntime { /// 共享初始化逻辑：根据 config + mode 构建完整运行时。 + /// `session_key` 用于初始化团队协作目录（.team/{key} 邮箱）。 fn init( app_state: Arc, mut config: AgentConfig, mode_registry: ModeRegistry, + session_key: &str, ) -> Self { let mode = mode_registry.get(&config.mode).copied().unwrap_or_else(|| { tracing::warn!("[AgentRuntime] 未知模式 '{}'，回退到默认模式", config.mode); @@ -120,7 +121,13 @@ impl AgentRuntime { 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); + // 团队协作工具接线：TeamManager 按会话初始化（spawn_teammate 等 4 个工具） + let team_slot: Arc>> = + Arc::new(tokio::sync::Mutex::new(Some( + crate::agent::team::manager::TeamManager::new(app_state.clone(), session_key), + ))); + let mut tool_registry = + ToolRegistry::new_with_team(Some(queue.clone()), team_slot, skill_registry); // 注册记忆工具 tool_registry.add_tool(Box::new(crate::agent::tools::memory::SaveMemoryTool::new( app_state.memory_manager.clone(), @@ -149,6 +156,7 @@ impl AgentRuntime { checkpoint_enabled, )); + let enable_thinking = config.enable_thinking; AgentRuntime { app_state, config, @@ -166,20 +174,32 @@ impl AgentRuntime { checkpoint_manager, mode, mode_registry, + thinking: std::sync::atomic::AtomicBool::new(enable_thinking), } } - /// 创建新的运行时实例 + /// 创建新的运行时实例（无会话上下文：团队目录退化为 .team/_adhoc） pub fn new(app_state: Arc) -> Self { let config = AgentConfig::default(); let mode_registry = ModeRegistry::builtins(); - Self::init(app_state, config, mode_registry) + Self::init(app_state, config, mode_registry, "_adhoc") + } + + /// 为指定会话创建运行时实例（会话级注册表使用； + /// 团队邮箱目录与该会话对齐） + pub fn new_for_session(app_state: Arc, session_key: &str, mode_id: &str) -> Self { + let config = AgentConfig { + mode: mode_id.to_string(), + ..AgentConfig::default() + }; + let mode_registry = ModeRegistry::builtins(); + Self::init(app_state, config, mode_registry, session_key).with_mode(mode_id) } /// 创建带自定义配置的运行时实例 pub fn with_config(app_state: Arc, config: AgentConfig) -> Self { let mode_registry = ModeRegistry::builtins(); - Self::init(app_state, config, mode_registry) + Self::init(app_state, config, mode_registry, "_adhoc") } /// 返回当前运行指标快照（锁异常时返回默认值并记录警告） @@ -196,9 +216,22 @@ impl AgentRuntime { /// 设置是否启用 LLM 思考模式（向后兼容，优先使用 mode 设置） pub fn with_thinking(mut self, enable: bool) -> Self { self.config.enable_thinking = enable; + self.thinking + .store(enable, std::sync::atomic::Ordering::SeqCst); self } + /// 运行时共享（Arc）后按请求覆盖思考模式 + pub fn set_thinking(&self, enable: bool) { + self.thinking + .store(enable, std::sync::atomic::Ordering::SeqCst); + } + + /// 当前思考模式是否启用 + pub fn thinking_enabled(&self) -> bool { + self.thinking.load(std::sync::atomic::Ordering::SeqCst) + } + /// 设置运行模式（覆盖 AgentConfig 中的 mode 字段）。 /// /// 调用此方法会重新解析模式并应用对应的配置预设、工具过滤和权限档案。 @@ -217,6 +250,11 @@ impl AgentRuntime { apply_mode_config(&mut self.config, mode); apply_mode_tool_filter(&mut self.tool_registry, mode); self.mode = mode; + // 模式预设可能锁定/翻转 thinking，同步到共享标志 + self.thinking.store( + self.config.enable_thinking, + std::sync::atomic::Ordering::SeqCst, + ); self } @@ -232,70 +270,6 @@ impl AgentRuntime { // ── Private Helpers ── - /// 执行文件缓存快照 → 压缩 → 恢复 → 上下文注入 的完整周期。 - /// 返回压缩前的消息数（用于调用者判断压缩是否有效）。 - async fn snapshot_compress_restore( - &self, - messages: &mut Vec, - llm: &LlmClient, - session_id: &str, - hook_registry: &HookRegistry, - ) -> usize { - let before_len = messages.len(); - - // ── 文件缓存快照（压缩前）── - let file_snapshot = { - if let Ok(mut cache) = self.read_file_state.lock() { - let snap = cache.to_snapshot(); - cache.clear(); - snap - } else { - Vec::new() - } - }; - - // 压缩前预提取记忆片段（避免克隆整个消息列表，仅提取前 400 字符的摘要） - let pre_compact_snippets = compact::extract_snippets(messages); - - compact::compress_context_with_hooks_and_log( - messages, - llm, - self.config.token_soft_limit, - self.config.max_messages, - session_id, - Some(hook_registry), - Some(&self.collapse_log), - ) - .await; - - // 压缩后提取记忆（P3 桥接：将丢弃的消息内容喂给记忆提取子代理） - compact::spawn_memory_extraction_from_snippets( - pre_compact_snippets, - session_id, - self.app_state.memory_manager.clone(), - self.app_state.clone(), - ); - - // ── 文件缓存恢复（压缩后：重新注入最近文件 + 恢复缓存）── - { - if let Ok(mut cache) = self.read_file_state.lock() { - cache.restore_from_snapshot( - &file_snapshot, - file_cache::POST_COMPACT_MAX_FILES_TO_RESTORE, - ); - } - let restore_ctx = file_cache::FileStateCache::build_restore_context( - &file_snapshot, - file_cache::POST_COMPACT_MAX_FILES_TO_RESTORE, - ); - for block in restore_ctx { - messages.push(ChatMessage::user(format!("[压缩后上下文恢复]\n{}", block))); - } - } - - before_len - } - // ── Public API ── /// 执行完整的智能体对话回合（流式 SSE 输出）。 @@ -328,10 +302,30 @@ impl AgentRuntime { let db = &self.app_state.db; let llm = &self.app_state.llm.primary; - // Phase 1: 创建或恢复会话 - let session_info = - session::create_or_resume_session(db, session_id.clone(), llm, &self.config.mode) - .await?; + // Phase 1: 创建或恢复会话（预分配 ID 语义：不存在则按该 ID 创建， + // 使取消标记/运行时缓存对首个请求同样生效） + let session_info = match session_id.clone() { + Some(sid) => { + session::create_or_resume_session_preallocated(db, &sid, llm, &self.config.mode) + .await? + } + None => session::create_or_resume_session(db, None, llm, &self.config.mode).await?, + }; + + // ── 会话事件日志：崩溃恢复协调（合成 interrupted 关闭开口的 + // turn_start/compaction_start，不截断已持久化内容）+ 本 turn 开始事件 ── + session_events::reconcile_interrupted(db, &session_info.session_id).await; + session_events::append_event( + db, + &session_info.session_id, + session_info.turn_index, + session_events::event_types::TURN_START, + serde_json::json!({ + "question_preview": question.chars().take(80).collect::(), + "mode": self.config.mode, + }), + ) + .await; // 构建 hook 注册表（注入依赖，复用 AgentRuntime 的 metrics_data） let hook_registry = HookRegistry::with_builtins( @@ -363,25 +357,19 @@ impl AgentRuntime { title: String::new(), }); - // Phase 2: 构建初始上下文（历史 + system prompt + 用户消息 + 任务恢复） + // Phase 2: 构建初始上下文（历史 + system prompt + 动态上下文快照 + 用户消息 + 任务恢复） + let dynamic_context = self.build_dynamic_context(question).await; let mut messages = context::build_initial_context( db, &session_info.session_id, &self.system_prompt(), question, session_info.turn_index, + dynamic_context, + image_context.as_deref(), ) .await?; - // 图片上下文：在用户消息前注入 system-reminder - if let Some(ref img_ctx) = image_context { - let reminder = format!("\n{}\n", img_ctx); - // 插入到倒数第二条位置（用户消息之前） - let user_msg = messages.pop().unwrap(); // 用户消息 - messages.push(ChatMessage::user(&reminder)); - messages.push(user_msg); - } - // 保存用户消息到数据库（含图片路径元数据，供前端历史渲染） let user_metadata = image_path .as_ref() @@ -402,6 +390,38 @@ impl AgentRuntime { .run_react_loop(&session_info, &mut messages, &tx, &hook_registry) .await?; + // ── 上下文快照：本 turn 发生过压缩 → 保存折叠后的上下文， + // 下一 turn 直接回放而无需从原始消息重建再重新压缩（节省重复摘要）── + if metrics.compression_count > 0 { + let base = session_events::max_message_id(db, &session_info.session_id).await; + session_events::save_context_snapshot( + db, + &session_info.session_id, + session_info.turn_index, + &messages, + base, + ) + .await; + } + + // ── turn 结束事件（关闭日志化锁；terminal 携带结构化终止原因）── + session_events::append_event( + db, + &session_info.session_id, + session_info.turn_index, + session_events::event_types::TURN_END, + serde_json::json!({ + "reason": loop_terminal + .as_ref() + .map(|t| t.reason_label()) + .unwrap_or("completed"), + "terminal": loop_terminal.as_ref().map(|t| format!("{:?}", t)), + "steps": metrics.total_steps, + "compressions": metrics.compression_count, + }), + ) + .await; + // Phase 4: 会话收尾（传入实际的终止原因 + trajectory 导出参数） let system_prompt = self.system_prompt(); let model_name = self.app_state.llm.primary.model().await; @@ -425,7 +445,10 @@ impl AgentRuntime { // ── ReAct Loop ── - /// ReAct 循环核心：LLM 调用 → 工具执行 → 结果注入 → 循环... + /// ReAct 循环入口 — 委托给统一引擎（src/agent/engine.rs）。 + /// + /// lead 配置：完整交互权限、恢复阶梯、todo nag、后台队列、压缩熔断。 + /// 子代理/队友在其他调用点以不同注入组合复用同一引擎。 async fn run_react_loop( &self, session_info: &SessionInfo, @@ -433,770 +456,55 @@ impl AgentRuntime { tx: &mpsc::UnboundedSender, hook_registry: &HookRegistry, ) -> anyhow::Result<(AgentMetrics, Option)> { - let db = &self.app_state.db; - let llm = &self.app_state.llm.primary; - let sid = &session_info.session_id; - let turn_index = session_info.turn_index; + use crate::agent::engine::{AskPolicy, CancelSource, DbMessageSink, EventTap, ReactEngine}; - // 注册当前会话的权限检查器（如不存在则从全局配置初始化） - self.app_state + let session_checker_snapshot = self + .app_state .session .session_permission_checkers - .entry(sid.clone()) - .or_insert_with(|| (*self.permission_checker).clone()); + .get(&session_info.session_id) + .map(|r| r.value().clone()); - let tool_defs = self.tool_registry.definitions(); - let mut duplicate_detector = DuplicateDetector::default(); - let mut metrics = AgentMetrics::default(); - let mut step = 0; - let mut loop_terminal: Option = None; + let mut engine_config = self.config.clone(); + engine_config.enable_thinking = self.thinking_enabled(); - // Token 追踪（API 精确值优先，字符估算作近似值） - let mut last_api_prompt_tokens: Option = None; - let mut msg_count_at_last_call: usize = messages.len(); - let mut steps_since_last_todo: usize = 0; - let nag_after_steps: usize = 3; - let mut pending_manual_compress: bool = false; - - // Token 预算管理器（用于 diminishing returns 检测和 error recovery） - let mut token_budget = - TokenBudget::new(self.config.token_soft_limit, self.config.token_hard_limit); - - loop { - step += 1; - - // ── Checkpoint: 每个 ReAct 迭代开始时重置去重状态 ── - self.checkpoint_manager.new_turn(); - - // 检查用户取消 - let is_cancelled = self.app_state.session.cancelled_runs.remove(sid).is_some(); - - if is_cancelled { - warn!("[AgentRuntime] 用户手动中止了会话 {} 的智能体执行", sid); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - loop_terminal = Some(TurnTerminal::CancelledByUser { - session_id: sid.clone(), - at_step: step, - }); - break; - } - - // ── 上下文压缩检查 ── - // ── Token 感知的压缩触发 ── - // 优先使用 API 返回的精确 prompt_tokens，辅以简单的增量估算 - let estimated_tokens = match last_api_prompt_tokens { - Some(last_tokens) => { - let new_msg_count = messages.len().saturating_sub(msg_count_at_last_call); - let new_tokens_estimate: usize = messages - .iter() - .rev() - .take(new_msg_count) - .map(|m| { - let content_len = m.content.as_ref().map_or(0, |c| c.len()); - content_len / 3 + 4 // 粗略估算：~3 字符/token + 消息 overhead - }) - .sum(); - last_tokens as usize + new_tokens_estimate - } - None => compact::rough_estimate_tokens(messages), - }; - - // 使用 TokenBudget 软限制作为压缩触发点 - let token_limit = token_budget.soft_limit; - - let mut did_compress = false; - - // 熔断器检查：如果连续压缩失败多次，跳过自动压缩 - let breaker_ok = match self.compaction_breaker.lock() { - Ok(mut breaker) => breaker.can_attempt(), - Err(e) => { - warn!("[AgentRuntime] 熔断器锁异常，跳过自动压缩: {:?}", e); - false - } - }; - - if estimated_tokens > token_limit && breaker_ok { - info!( - "[AgentRuntime] 上下文超限 (est. {} tokens > {} limit)，触发压缩", - estimated_tokens, token_limit - ); - let before_len = self - .snapshot_compress_restore( - messages, - llm, - &session_info.session_id, - hook_registry, - ) - .await; - - // 熔断器反馈：压缩后消息数减少 = 成功 - if let Ok(mut breaker) = self.compaction_breaker.lock() { - if messages.len() < before_len { - breaker.record_success(); - } else { - breaker.record_failure(); - } - } else { - warn!("[AgentRuntime] 熔断器反馈写入失败（锁异常）"); - } - last_api_prompt_tokens = None; - msg_count_at_last_call = messages.len(); - metrics.compression_count += 1; - did_compress = true; - } else if estimated_tokens > token_limit && !breaker_ok { - warn!("[AgentRuntime] 熔断器已打开，跳过自动压缩"); - } - - // 处理手动压缩请求（跳过刚自动压缩过的情况，避免双重压缩） - // 手动压缩不受熔断器限制 - if pending_manual_compress && !did_compress { - pending_manual_compress = false; - info!("[AgentRuntime] 执行手动压缩（compress_context 工具触发）"); - - self.snapshot_compress_restore( - messages, - llm, - &session_info.session_id, - hook_registry, - ) - .await; - - // 手动压缩成功后重置熔断器 - if let Ok(mut breaker) = self.compaction_breaker.lock() { - breaker.reset(); - } - last_api_prompt_tokens = None; - msg_count_at_last_call = messages.len(); - metrics.compression_count += 1; - } else if pending_manual_compress { - pending_manual_compress = false; - info!("[AgentRuntime] 跳过手动压缩（刚已完成自动压缩）"); - } - - // Token 预算 diminishing returns 检测 + 渐进式 nudge 提醒 - token_budget.record_continuation(); - - let mut should_nudge = false; - - // TodoWrite nag reminder - if steps_since_last_todo >= nag_after_steps { - messages.push(ChatMessage::user( - "提醒：你已经连续多步未更新任务计划。建议调用 todo_write 工具复盘当前进度并规划后续步骤。", - )); - steps_since_last_todo = 0; - should_nudge = true; - } - - // Token 预算 nudge（仅在无 nag 时注入，避免消息过多） - if !should_nudge { - if let Some(nudge) = token_budget.nudge_message() { - messages.push(ChatMessage::user(nudge)); - } - } - - // Diminishing returns 检测 — 强制结束 - if token_budget.diminishing_returns { - warn!("[AgentRuntime] 检测到 diminishing returns，强制结束循环"); - let _ = tx.send(AgentStreamEvent::Error { - message: "检测到重复操作模式，已自动停止。请查看已收集的信息。".to_string(), - }); - messages.push(ChatMessage::user( - "检测到你的后续步骤未产生新信息（diminishing returns）。\ - 请基于已收集的全部信息直接给出最终答案，不要再调用任何工具。", - )); - let _ = self - .final_answer_without_tools(llm, messages, sid, turn_index, step, tx) - .await; - break; - } - - // 最大步数检查 - if step > self.config.max_steps { - warn!( - "[AgentRuntime] 达到最大步数限制 ({} steps)", - self.config.max_steps - ); - let _ = tx.send(AgentStreamEvent::Error { - message: format!( - "已达到最大推理步数 ({})，请根据已收集的信息给出最终回答。", - self.config.max_steps - ), - }); - messages.push(ChatMessage::user(format!( - "你已经执行了 {} 步（最大 {} 步）。请根据已有信息直接给出最终答案，不要再调用工具。", - step, self.config.max_steps - ))); - let _ = self - .final_answer_without_tools(llm, messages, sid, turn_index, step, tx) - .await; - break; - } - - // ── 后台任务通知注入 ── - let bg_results = self.bg_notification_queue.drain().await; - for result in bg_results { - let status = if result.is_error { "❌" } else { "✅" }; - messages.push(ChatMessage::user(format!( - "[后台任务完成] {} {}: {} ({}): {}", - status, result.tool_name, result.bibcode, result.task_id, result.summary, - ))); - } - - // ── LLM 流式调用（含错误恢复） ── - let stream_output_opt = self - .call_llm_with_recovery(llm, messages, &tool_defs, tx, step, sid, &mut token_budget) - .await; - - let stream_output = match stream_output_opt { - Some(output) => output, - None => { - // 所有恢复尝试均失败 - loop_terminal = Some(TurnTerminal::ModelError { - session_id: sid.clone(), - message: "LLM 调用失败，所有恢复步骤已尝试完毕".to_string(), - }); - break; - } - }; - - // 更新 API 精确 token 计数 + token 预算 - if let Some(ref u) = stream_output.usage { - last_api_prompt_tokens = Some(u.prompt_tokens); - msg_count_at_last_call = messages.len(); - token_budget.spend_input(u.prompt_tokens as usize); - token_budget.spend_output(u.completion_tokens as usize); - } - - // ── 处理 Thought/Reasoning ── - let mut thought_content = stream_output.reasoning.clone(); - - if thought_content.is_none() - && stream_output.is_tool_call_step - && !stream_output.content.is_empty() - { - thought_content = Some(stream_output.content.clone()); - } - - if stream_output.is_tool_call_step { - if let Some(ref thought_text) = thought_content { - let _ = tx.send(AgentStreamEvent::Thought { - content: thought_text.clone(), - step, - }); - } - } - - // ── 无工具调用 = 最终回答 ── - let mut tool_calls = match stream_output.tool_calls { - Some(ref tc) if !tc.is_empty() => tc.clone(), - _ => { - // 保存最终回答 - let assistant_msg = ChatMessage::assistant_with_reasoning( - if stream_output.content.is_empty() { - None - } else { - Some(stream_output.content.clone()) - }, - stream_output.reasoning.clone(), - None, - ); - self.save_message( - db, - sid, - turn_index, - step as i32, - &assistant_msg, - stream_output.reasoning.as_deref(), - None, - ) - .await?; - messages.push(assistant_msg); - - // 发送 reasoning（当模型思考后直接给出答案、未调用工具时， - // thought 尚未在上面的 is_tool_call_step 块中发送） - if let Some(ref thought_text) = stream_output.reasoning { - let _ = tx.send(AgentStreamEvent::Thought { - content: thought_text.clone(), - step, - }); - } - - // Token 使用统计 - if let Some(u) = stream_output.usage { - let _ = tx.send(AgentStreamEvent::Usage { - prompt_tokens: u.prompt_tokens, - completion_tokens: u.completion_tokens, - total_tokens: u.total_tokens, - }); - } - break; - } - }; - - // 修复空 ID（LLM 可能不返回 tool_call id） - for tc in tool_calls.iter_mut() { - if tc.id.is_empty() { - tc.id = format!("call_{}", &uuid::Uuid::new_v4().to_string()[..8]); - } - } - - // ── 工具调用处理 ── - // 检测 todo_write 和 compress_context - let called_todo_write = tool_calls.iter().any(|tc| tc.function.name == "todo_write"); - if called_todo_write { - steps_since_last_todo = 0; - } else { - steps_since_last_todo += 1; - } - - if tool_calls - .iter() - .any(|tc| tc.function.name == "compress_context") - { - pending_manual_compress = true; - } - - // 更新指标 - metrics.total_steps = step; - for tc in &tool_calls { - *metrics - .tool_calls - .entry(tc.function.name.clone()) - .or_insert(0) += 1; - } - - // 构建 assistant 消息（含 tool_calls） - let assistant_msg = ChatMessage::assistant_with_reasoning( - if stream_output.content.is_empty() { - None - } else { - Some(stream_output.content.clone()) - }, - stream_output.reasoning.clone(), - Some(tool_calls.clone()), - ); - self.save_message( - db, - sid, - turn_index, - step as i32, - &assistant_msg, - stream_output.reasoning.as_deref(), - None, - ) - .await?; - messages.push(assistant_msg); - - // 验证 + 准备工具调用 - let (prepared_calls, has_duplicate) = executor::validate_and_prepare( - &tool_calls, - &mut duplicate_detector, - self.config.duplicate_call_threshold, - messages, - &self.tool_registry, - tx, - db, - sid, - turn_index, - step, - ); - - if has_duplicate { - metrics.duplicate_detections += 1; - continue; - } - - if prepared_calls.is_empty() { - continue; - } - - // 并行执行工具（带权限检查、checkpoint 和分区器） - let session_checker_snapshot = self - .app_state - .session - .session_permission_checkers - .get(sid) - .map(|r| r.value().clone()); - let exec_result = executor::execute_parallel( - &prepared_calls, - &self.tool_registry, - self.app_state.clone(), - hook_registry, - Some(&self.permission_checker), - session_checker_snapshot.as_ref(), - Some(&self.denial_tracker), - Some(&self.checkpoint_manager), - tx, - db, - sid, - "lead", - turn_index, - step, - self.config.tool_timeout_secs, - self.config.max_tool_output_chars, - self.read_file_state.clone(), - self.config.enable_thinking, - self.config.additional_allowed_dirs.clone(), - ) - .await; - - // 拒绝追踪熔断检查：连续/累计拒绝达到阈值则终止循环 - if let Ok(dt) = self.denial_tracker.lock() { - if dt.should_terminate() { - let reason = dt.termination_reason(); - warn!("[AgentRuntime] 拒绝熔断触发: {}", reason); - let _ = tx.send(AgentStreamEvent::Error { message: reason }); - break; - } - } - - // 将工具结果推入消息上下文 - for tm in exec_result.tool_messages { - messages.push(tm.chat_message); - } - - // Hook 注入的附加上下文：包装为 system-reminder 注入 LLM 消息列表 - // 使用 ContextDeduplicator 在单步内去重（多个 hook 注入相同内容时只保留一份） - let mut dedup = crate::agent::hooks::ContextDeduplicator::new(); - for ctx in &exec_result.hook_contexts { - if dedup.is_duplicate(ctx) { - continue; - } - let reminder = format!( - "\n[Hook 注入上下文]\n{}\n", - ctx - ); - messages.push(ChatMessage::user(&reminder)); - } - - // Hook 阻塞错误：记录到日志用于诊断 - for be in &exec_result.blocking_errors { - warn!("[AgentRuntime] Hook 阻塞错误: {}", be); - } - - // 持久化 todo_write 任务状态到数据库 - if called_todo_write { - for prep in &prepared_calls { - if prep.tool_name == "todo_write" { - if let Some(todos) = prep.args.get("todos").and_then(|t| t.as_array()) { - let todos_vec: Vec = todos.to_vec(); - let _ = crate::agent::tools::persist_tasks(db, sid, &todos_vec, "lead") - .await; - } - } - } - } - - if exec_result.was_cancelled { - self.app_state.session.cancelled_runs.remove(sid); - warn!( - "[AgentRuntime] 工具执行期间被用户手动中止，会话 ID: {}", - sid - ); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - loop_terminal = Some(TurnTerminal::CancelledByUser { - session_id: sid.clone(), - at_step: step, - }); - break; - } - - // OnStepComplete hook - let step_ctx = StepCompleteContext { - session_id: sid.clone(), - step, - max_steps: self.config.max_steps, - messages_count: messages.len(), - estimated_tokens, - token_limit, - }; - hook_registry.run_on_step_complete(&step_ctx).await; - } - - Ok((metrics, loop_terminal)) - } - - /// LLM 流式调用，含完整的错误恢复阶梯。 - /// - /// 首次调用失败后，按顺序尝试： - /// 1. AggressiveCompact (keep_recent=2) - /// 2. ReactiveCompact (LLM 摘要) - /// 3. EscalateTokens (提升 hard_limit → 64k) - /// 4. MultiTurn (注入分步消息) - /// 5. Surface (放弃) - /// - /// 每一步后重试 LLM 调用。返回 Some(StreamOutput) 表示成功（可能经过恢复）， - /// None 表示所有步骤均已尝试且失败。 - #[allow(clippy::too_many_arguments)] - async fn call_llm_with_recovery( - &self, - llm: &LlmClient, - messages: &mut Vec, - tool_defs: &[crate::clients::llm::ToolDefinition], - tx: &mpsc::UnboundedSender, - step: usize, - session_id: &str, - token_budget: &mut TokenBudget, - ) -> Option { - // 首次尝试 - let output = streaming::process_llm_stream( - llm, - messages, - tool_defs, - tx, - step, - session_id, - self.app_state.session.cancelled_runs.clone(), - self.config.enable_thinking, - ) - .await; - - match output.status { - StreamStatus::Success => return Some(output), - StreamStatus::Cancelled => { - self.app_state.session.cancelled_runs.remove(session_id); - warn!( - "[AgentRuntime] 流式调用期间被用户手动中止，会话 ID: {}", - session_id - ); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - return None; - } - StreamStatus::Error(ref e_str) => { - error!("[AgentRuntime] 流式读取错误: {}", e_str); - } - } - - // 提取错误字符串（用于分类） - let e_str = match &output.status { - StreamStatus::Error(s) => s.clone(), - _ => return Some(output), // 不应到达，但安全起见 + let engine = ReactEngine { + llm: &self.app_state.llm.primary, + tool_registry: &self.tool_registry, + app_state: self.app_state.clone(), + config: engine_config, + session_id: session_info.session_id.clone(), + turn_index: session_info.turn_index, + agent_name: "lead".to_string(), + tap: EventTap::channel(tx.clone()), + sink: Some(Arc::new(DbMessageSink { + db: self.app_state.db.clone(), + session_id: session_info.session_id.clone(), + agent_name: "lead".to_string(), + fixed_metadata: None, + })), + hooks: Some(hook_registry), + ask_policy: AskPolicy::Interactive, + permission_checker: Some(self.permission_checker.as_ref()), + session_checker: session_checker_snapshot, + denial_tracker: Some(self.denial_tracker.clone()), + checkpoint_manager: Some(self.checkpoint_manager.clone()), + read_file_state: self.read_file_state.clone(), + cancel: CancelSource::Session { + session_id: session_info.session_id.clone(), + cancelled_runs: self.app_state.session.cancelled_runs.clone(), + }, + breaker: Some(self.compaction_breaker.clone()), + collapse_log: Some(self.collapse_log.clone()), + bg_queue: Some(self.bg_notification_queue.clone()), + nudges: true, + diminishing_returns: true, + recovery: true, + persist_todos: true, }; - let error_kind = classify_error(&e_str); - - // ── 429/529 瞬态错误：指数退避重试（独立的快速路径） ── - if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) { - let retry_after_secs = error_recovery::parse_retry_after(&e_str); - let mut consecutive_overloads: u32 = 0; - const MAX_BACKOFF_RETRIES: u32 = 10; - - for attempt in 0..MAX_BACKOFF_RETRIES { - let delay_ms = error_recovery::backoff_delay(attempt, retry_after_secs); - info!( - "[AgentRuntime] 退避重试 {}/{} ({}ms, error={:?})", - attempt + 1, - MAX_BACKOFF_RETRIES, - delay_ms, - error_kind - ); - - let _ = tx.send(AgentStreamEvent::Thought { - content: format!( - "⏳ 模型服务暂时不可用，正在重试 ({}/{})...", - attempt + 1, - MAX_BACKOFF_RETRIES - ), - step, - }); - - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - - // 529 连续过载检测：3 次后尝试切换备用模型 - if matches!(error_kind, ErrorKind::Overloaded) { - consecutive_overloads += 1; - if consecutive_overloads >= 3 { - let fallback = &self.app_state.config.llm.fallback_model; - if !fallback.is_empty() { - warn!( - "[AgentRuntime] 连续 {} 次过载，切换到备用模型: {}", - consecutive_overloads, fallback - ); - 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) - % self.app_state.config.llm.fallback_chain.len()) - .min(self.app_state.config.llm.fallback_chain.len() - 1); - let alt = &self.app_state.config.llm.fallback_chain[idx]; - warn!( - "[AgentRuntime] 连续 {} 次过载，从链中切换: {}", - consecutive_overloads, alt - ); - llm.set_model(alt.clone()).await; - consecutive_overloads = 0; - } - } - } - - // 检查用户取消 - if self - .app_state - .session - .cancelled_runs - .contains_key(session_id) - { - warn!("[AgentRuntime] 退避重试期间被用户取消"); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - return None; - } - - // 重试 LLM 调用 - let retry_output = streaming::process_llm_stream( - llm, - messages, - tool_defs, - tx, - step, - session_id, - self.app_state.session.cancelled_runs.clone(), - self.config.enable_thinking, - ) - .await; - - match retry_output.status { - StreamStatus::Success => { - info!("[AgentRuntime] 退避重试成功！(尝试 {})", attempt + 1); - return Some(retry_output); - } - StreamStatus::Cancelled => { - self.app_state.session.cancelled_runs.remove(session_id); - return None; - } - StreamStatus::Error(_) => { - // 继续重试 - continue; - } - } - } - - // 所有退避重试失败 - warn!( - "[AgentRuntime] {} 次退避重试后仍然失败", - MAX_BACKOFF_RETRIES - ); - let _ = tx.send(AgentStreamEvent::Error { - message: format!( - "模型服务暂时不可用（已重试 {} 次）。请稍后再试或检查模型服务状态。", - MAX_BACKOFF_RETRIES - ), - }); - return None; - } - - if !ErrorRecovery::is_recoverable(&error_kind) { - let _ = tx.send(AgentStreamEvent::Error { - message: format!("大模型流式读取失败: {}", e_str), - }); - return None; - } - - let mut recovery = ErrorRecovery::new(token_budget.clone()); - - // 尝试从错误消息中解析 ContextOverflow 信息（参考 Claude Code 自动修复） - let overflow_info = error_recovery::parse_context_overflow(&e_str); - - while let Some(recovery_step) = recovery.try_recover(&error_kind, overflow_info.as_ref()) { - match recovery_step { - error_recovery::RecoveryStep::AdjustMaxTokens { new_max_tokens } => { - info!( - "[AgentRuntime] 错误恢复: AdjustMaxTokens → {} (从错误消息自动计算)", - new_max_tokens - ); - // token_budget.hard_limit 已由 try_recover 下调 - } - error_recovery::RecoveryStep::RetryWithBackoff { attempt, delay_ms } => { - // 429/529 本应在 streaming 层处理，若到达此处说明分类逻辑有变更， - // 安全降级为 sleep + 直接重试（不依赖 streaming 层重试）。 - warn!( - "[AgentRuntime] RetryWithBackoff 在 error_recovery 层触发 (attempt={}, delay={}ms)，执行降级重试", - attempt, delay_ms - ); - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - // 不计入 recovery 计数，由外层循环自然重试 - } - error_recovery::RecoveryStep::AggressiveCompact => { - info!("[AgentRuntime] 错误恢复: 激进压缩 (snip + micro with keep_recent=2)"); - compact::snip_compact(messages, self.config.max_messages); - *messages = compact::micro_compact(messages, 2); - } - error_recovery::RecoveryStep::ReactiveCompact => { - info!("[AgentRuntime] 错误恢复: LLM 摘要压缩"); - compact::compress_context( - messages, - llm, - self.config.token_soft_limit, - self.config.max_messages, - session_id, - ) - .await; - } - error_recovery::RecoveryStep::EscalateTokens { .. } => { - info!( - "[AgentRuntime] 错误恢复: 提升 token 硬限制到 {}", - recovery.token_budget.hard_limit - ); - } - error_recovery::RecoveryStep::MultiTurn => { - info!("[AgentRuntime] 错误恢复: 注入多轮消息"); - messages.push(ChatMessage::user(ErrorRecovery::multi_turn_message())); - } - error_recovery::RecoveryStep::Surface => { - warn!("[AgentRuntime] 错误恢复: 所有步骤失败，暴露错误"); - break; - } - } - - // 重试 LLM 调用 - let retry_output = streaming::process_llm_stream( - llm, - messages, - tool_defs, - tx, - step, - session_id, - self.app_state.session.cancelled_runs.clone(), - self.config.enable_thinking, - ) - .await; - - match retry_output.status { - StreamStatus::Success => { - info!("[AgentRuntime] 错误恢复成功！"); - // 将恢复后的 token_budget 状态同步回去 - *token_budget = recovery.token_budget.clone(); - return Some(retry_output); - } - StreamStatus::Cancelled => { - self.app_state.session.cancelled_runs.remove(session_id); - warn!("[AgentRuntime] 恢复期间被用户中止"); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - return None; - } - StreamStatus::Error(retry_err) => { - info!( - "[AgentRuntime] 恢复步骤 {:?} 未能解决，继续下一阶梯: {}", - recovery_step, retry_err - ); - } - } - } - - // 所有恢复步骤均已尝试 - let _ = tx.send(AgentStreamEvent::Error { - message: format!("大模型流式读取失败，且所有恢复步骤均未能解决: {}", e_str), - }); - None + let outcome = engine.run(messages).await; + Ok((outcome.metrics, outcome.terminal)) } // ── Helpers ── @@ -1274,7 +582,25 @@ impl AgentRuntime { drop(cache); - // 技能列表：通过文件监听热更新，不缓存 + // KV-cache 纪律：system prompt 只包含字节稳定的 section。 + // 技能列表/项目记忆这类易变内容若混入 system prompt，任一变化都会 + // 使整条前缀缓存失效——它们改为以 durable 动态上下文快照的形式 + // 追加在历史消息尾部（内容变更时才追加，见 build_dynamic_context）。 + + sp.assemble() + } + + /// 构建动态上下文快照（技能清单 + 项目记忆）。 + /// + /// 返回 (文本, 哈希)。哈希与会话内最近一次持久化的快照比较： + /// 未变化 → 不注入（历史里已有）；变化 → 作为 user 消息追加并持久化 + /// （durable：跨 turn 存在于历史中，而非每 turn 重付 token）。 + /// + /// 记忆选取：条目多于阈值时走 LLM 智能选择 + 时间衰减排序 + /// （select_relevant_memories，此前为零调用代码），否则 recency 启发式。 + async fn build_dynamic_context(&self, question: &str) -> Option<(String, u64)> { + let mut parts: Vec = Vec::new(); + if let Some(skills) = self .app_state .skill_registry @@ -1282,21 +608,39 @@ impl AgentRuntime { .ok() .and_then(|r| r.build_reminder()) { - sp.add_section("skills", skills); + parts.push(skills); } - // 项目记忆：受 save_memory 工具实时影响，不缓存 - if let Some(memory) = self - .app_state - .memory_manager - .try_lock() - .ok() - .and_then(|mgr| mgr.build_system_reminder(5)) + // 记忆选择：少量条目走启发式；条目多时智能选择（fast 模型，每 turn 一次） { - sp.add_section("memory", memory); + let mgr = self.app_state.memory_manager.lock().await; + const SMART_SELECTION_THRESHOLD: usize = 8; + let memory_text = if mgr.entries().len() > SMART_SELECTION_THRESHOLD { + let selected = mgr + .select_relevant_memories(&self.app_state.llm.fast, question, 5) + .await; + mgr.build_system_reminder_from(&selected) + } else { + mgr.build_system_reminder(5) + }; + if let Some(memory) = memory_text { + parts.push(memory); + } } - sp.assemble() + if parts.is_empty() { + return None; + } + + let text = format!( + "[动态上下文 — 由系统自动维护，内容变更时自动更新]\n{}", + parts.join("\n\n") + ); + + use std::hash::{Hash, Hasher}; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + text.hash(&mut hasher); + Some((text, hasher.finish())) } /// 使提示词缓存中指定 section 失效。 @@ -1362,73 +706,6 @@ impl AgentRuntime { lines.join("\n") } - /// 步数耗尽时的最终答案生成（不带工具调用，强制 LLM 直接回答） - async fn final_answer_without_tools( - &self, - llm: &LlmClient, - messages: &[ChatMessage], - session_id: &str, - turn_index: i32, - step: usize, - tx: &mpsc::UnboundedSender, - ) -> anyhow::Result<()> { - let empty_tools: Vec = Vec::new(); - let mut stream_rx = match llm - .chat_stream(messages, &empty_tools, self.config.enable_thinking) - .await - { - Ok(rx) => rx, - Err(e) => { - let _ = tx.send(AgentStreamEvent::Error { - message: format!("最终回答生成失败: {}", e), - }); - return Err(anyhow::anyhow!("final_answer LLM call failed: {}", e)); - } - }; - - let mut accumulated = String::new(); - while let Some(event) = stream_rx.recv().await { - match event { - StreamEvent::TextDelta(delta) => { - accumulated.push_str(&delta); - let _ = tx.send(AgentStreamEvent::TextDelta { - content: delta, - tool_call_id: None, - }); - } - StreamEvent::Usage(u) => { - let _ = tx.send(AgentStreamEvent::Usage { - prompt_tokens: u.prompt_tokens, - completion_tokens: u.completion_tokens, - total_tokens: u.total_tokens, - }); - } - StreamEvent::Done => break, - StreamEvent::Error(e) => { - let _ = tx.send(AgentStreamEvent::Error { - message: format!("最终回答流式错误: {}", e), - }); - break; - } - _ => {} - } - } - - let assistant_msg = ChatMessage::assistant(accumulated.clone()); - self.save_message( - &self.app_state.db, - session_id, - turn_index, - step as i32, - &assistant_msg, - None, - None, - ) - .await?; - - Ok(()) - } - /// 保存消息到数据库 #[allow(clippy::too_many_arguments)] async fn save_message( @@ -1454,7 +731,7 @@ impl AgentRuntime { .await } - /// 保存消息到数据库（指定 agent 身份） + /// 保存消息到数据库（指定 agent 身份）— 委托给统一的 DbMessageSink #[allow(clippy::too_many_arguments)] async fn save_message_as( &self, @@ -1467,59 +744,15 @@ impl AgentRuntime { agent_name: &str, extra_metadata: Option, ) -> anyhow::Result<()> { - let role = match msg.role { - MessageRole::System => "system", - MessageRole::User => "user", - MessageRole::Assistant => "assistant", - MessageRole::Tool => "tool", + use crate::agent::engine::{DbMessageSink, MessageSink as _}; + let sink = DbMessageSink { + db: db.clone(), + session_id: session_id.to_string(), + agent_name: agent_name.to_string(), + fixed_metadata: None, }; - - let content = msg.content.clone().unwrap_or_default(); - let tool_calls_json = msg - .tool_calls - .as_ref() - .map(|tc| serde_json::to_string(tc).unwrap_or_default()); - let tool_call_id = msg.tool_call_id.as_deref(); - let token_count = content.len() as i32 / 4; - - // metadata: 存储结构化的消息元信息（thought/tool_calls/tool_call_id 等） - let mut metadata = serde_json::json!({ - "has_thought": thought.is_some(), - "has_tool_calls": tool_calls_json.is_some(), - "step_index": step_index, - }); - // 合并额外元数据（如图片路径等） - if let Some(ref extra) = extra_metadata { - if let (Some(base), Some(extra_obj)) = (metadata.as_object_mut(), extra.as_object()) { - for (k, v) in extra_obj { - base.insert(k.clone(), v.clone()); - } - } - } - let metadata_str = serde_json::to_string(&metadata).unwrap_or_default(); - - // raw_json: 存储完整消息的 JSON 序列化（调试/审计用） - let raw_json = serde_json::to_string(msg).unwrap_or_default(); - - sqlx::query( - "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, raw_json, agent_name) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(session_id) - .bind(turn_index) - .bind(step_index) - .bind(role) - .bind(content) - .bind(thought) - .bind(&tool_calls_json) - .bind(tool_call_id) - .bind(token_count) - .bind(&metadata_str) - .bind(&raw_json) - .bind(agent_name) - .execute(db) - .await?; - + sink.save(turn_index, step_index, msg, thought, extra_metadata) + .await; Ok(()) } @@ -1580,9 +813,20 @@ fn apply_mode_config(config: &mut AgentConfig, mode: &AgentMode) { } /// 根据模式的 ToolSet 设置工具注册表的定义过滤器。 +/// 白名单中引用未注册工具时发出警告（fail-loud 提示配置漂移）。 fn apply_mode_tool_filter(tool_registry: &mut ToolRegistry, mode: &AgentMode) { let all_names = tool_registry.tool_names(); if let Some(filter) = modes::tool_set_to_filter(&mode.tool_set, &all_names) { + let known: std::collections::HashSet<&str> = all_names.iter().map(|s| s.as_str()).collect(); + for name in &filter { + if !known.contains(name.as_str()) { + tracing::warn!( + "[AgentRuntime] 模式 '{}' 白名单引用了未注册的工具 '{}'（将被忽略）", + mode.id, + name + ); + } + } tracing::info!( "[AgentRuntime] 模式 '{}' 工具过滤: {} → {} 个工具", mode.id, diff --git a/src/agent/runtime/permission.rs b/src/agent/runtime/permission.rs index a791833..74c55e4 100644 --- a/src/agent/runtime/permission.rs +++ b/src/agent/runtime/permission.rs @@ -508,109 +508,33 @@ impl PermissionChecker { } } -// ── Permission Precedence Resolver ── +// ── Monotonic Combination（权限单调性不变量）── -/// 多源权限决策的最终裁决。遵循正式的优先级规则表： +/// 权限决策合并 — 整个权限体系的唯一不变量： +/// **只能收紧，不能放松**（deny > ask > allow，最严格胜出）。 /// -/// | Priority | Source | Overridable By | -/// |----------|----------------------------------|----------------| -/// | P0 | PermissionChecker::Deny | Nothing | -/// | P1 | Tool-level PermissionRule::Deny | Nothing | -/// | P2 | Session-level Checker::Deny | Nothing | -/// | P3 | Hook PreToolUseAction::Block | P0-P2 | -/// | P4 | Hook PermissionRequired | P0-P3 | -/// | P5-P7 | Checker::Ask / Tool::Ask / Allow | Normal | +/// 多源决策（env 规则、工具级声明、会话级规则、hook 请求）逐层调用本函数 +/// 合并，任何一层都无法放行另一层已收紧的决策——这是 guard 单调性的 +/// 权限等价物（参考 deepseek-harness "guards only tighten" 约束）。 /// -/// `conflict_log` 记录被覆盖的决策，便于审计和调试。 -pub fn resolve_permission_precedence( - checker_result: PermissionResult, - tool_rules: &[crate::agent::tools::PermissionRule], - hook_permission: Option<&(String, String)>, // (permission_desc, tool_name) - hook_blocked: bool, - session_result: Option, -) -> (PermissionResult, Vec) { - let mut final_result = checker_result; - let mut conflict_log: Vec = Vec::new(); - - // ── P1: Tool-level Deny ── - for rule in tool_rules { - if let crate::agent::tools::PermissionRule::Deny { reason, .. } = rule { - if !final_result.is_denied() { - conflict_log.push(format!("Tool-level Deny overrides checker: {reason}")); - final_result = PermissionResult::Denied { - reason: reason.clone(), - }; - } else { - conflict_log.push(format!( - "Tool-level Deny '{reason}' ignored: already Denied" - )); - } - break; // only handle first Deny - } +/// 具体语义： +/// - `base` 已 Denied → 结果必为 Denied（粘滞，不可覆盖） +/// - `base` 已 AskUser → 候选 Allowed 不生效；候选 Denied 升级为 Denied +/// - `base` Allowed → 采用候选（候选可为 Allowed/AskUser/Denied） +pub fn tighten(base: PermissionResult, candidate: PermissionResult) -> PermissionResult { + match base { + PermissionResult::Denied { .. } => base, // Deny 粘滞 + PermissionResult::AskUser { .. } => match candidate { + // Ask 不能被放松为 Allow；可升级为 Deny + PermissionResult::Denied { reason } => PermissionResult::Denied { reason }, + _ => base, + }, + PermissionResult::Allowed => candidate, } - - // ── P2: Session-level Deny ── - if let Some(PermissionResult::Denied { reason }) = &session_result { - conflict_log.push(format!("Session-level Deny overrides current: {reason}")); - final_result = PermissionResult::Denied { - reason: reason.clone(), - }; - } - - // ── P3: Hook Block ── - if hook_blocked { - conflict_log.push("Hook Block prevents execution".to_string()); - // Block is already handled in the executor via denied_indices; - // here we record it for the conflict log. - } - - // ── P4: Hook PermissionRequired ── - if let Some((perm_desc, _tool_name)) = hook_permission { - if final_result.is_allowed() { - conflict_log.push(format!( - "Hook PermissionRequired upgrades Allowed → Ask: {perm_desc}" - )); - } else if !final_result.is_denied() { - conflict_log.push(format!( - "Hook PermissionRequired coexists with current state: {perm_desc}" - )); - } else { - conflict_log.push(format!( - "Hook PermissionRequired '{perm_desc}' ignored: already Denied" - )); - } - } - - // ── P5: Session-level Ask ── - if let Some(PermissionResult::AskUser { message }) = &session_result { - if final_result.is_allowed() { - final_result = PermissionResult::AskUser { - message: message.clone(), - }; - conflict_log.push("Session-level Ask overrides Allow".to_string()); - } else { - conflict_log.push("Session-level Ask ignored: not Allowed".to_string()); - } - } - - // ── P6: Tool-level Ask ── - for rule in tool_rules { - if let crate::agent::tools::PermissionRule::Ask { message, .. } = rule { - if final_result.is_allowed() { - conflict_log.push(format!("Tool-level Ask upgrades Allow: {message}")); - final_result = PermissionResult::AskUser { - message: message.clone(), - }; - } else { - conflict_log.push("Tool-level Ask ignored: not Allowed".to_string()); - } - break; - } - } - - (final_result, conflict_log) } +// ── Permission Precedence Resolver ── + #[cfg(test)] mod tests { use super::*; @@ -965,85 +889,84 @@ mod tests { assert!(matches!(result, PermissionResult::AskUser { .. })); } - // ── resolve_permission_precedence tests ── + // ── tighten（单调合并不变量）tests ── #[test] - fn test_precedence_checker_deny_wins_over_all() { - let (result, log) = resolve_permission_precedence( + fn test_tighten_deny_is_sticky() { + use crate::agent::runtime::permission::tighten; + let base = PermissionResult::Denied { + reason: "policy".into(), + }; + // 任何候选都无法放松 Deny + assert!(tighten(base.clone(), PermissionResult::Allowed).is_denied()); + assert!(tighten( + base, + PermissionResult::AskUser { + message: "m".into() + } + ) + .is_denied()); + } + + #[test] + fn test_tighten_ask_cannot_be_relaxed() { + use crate::agent::runtime::permission::tighten; + let base = PermissionResult::AskUser { + message: "need confirm".into(), + }; + // Ask + Allow → Ask（不可放松） + let r = tighten(base.clone(), PermissionResult::Allowed); + assert!(matches!(r, PermissionResult::AskUser { .. })); + // Ask + Deny → Deny（可升级） + let r = tighten( + base, PermissionResult::Denied { - reason: "blocked by policy".into(), + reason: "no".into(), }, - &[], - Some(&("need confirmation".to_string(), "test_tool".to_string())), - false, - None, - ); - assert!(result.is_denied()); - assert!( - !log.is_empty(), - "conflict log should record the interaction" ); + assert!(r.is_denied()); } #[test] - fn test_precedence_tool_deny_overrides_allow() { - use crate::agent::tools::{PermissionRule, PermissionRuleSource}; - let tool_rules = vec![PermissionRule::Deny { - tool_name: "test_tool".into(), - reason: "tool self-protection".into(), - source: PermissionRuleSource::Env, - }]; - let (result, log) = resolve_permission_precedence( - PermissionResult::Allowed, - &tool_rules, - None, - false, - None, + fn test_tighten_allowed_adopts_candidate() { + use crate::agent::runtime::permission::tighten; + let base = PermissionResult::Allowed; + assert_eq!( + tighten(base.clone(), PermissionResult::Allowed), + PermissionResult::Allowed ); - assert!(result.is_denied()); - assert!(!log.is_empty()); + assert!(matches!( + tighten( + base.clone(), + PermissionResult::AskUser { + message: "m".into() + } + ), + PermissionResult::AskUser { .. } + )); + assert!(tighten(base, PermissionResult::Denied { reason: "d".into() }).is_denied()); } #[test] - fn test_precedence_hook_block_recorded() { - let (result, log) = resolve_permission_precedence( - PermissionResult::Allowed, - &[], - None, - true, // hook blocked - None, - ); - // Hook Block doesn't directly return Denied — it's logged for executor handling - assert!(result.is_allowed()); - assert!(log.iter().any(|l| l.contains("Block"))); - } - - #[test] - fn test_precedence_session_deny_overrides() { - let (result, _log) = resolve_permission_precedence( - PermissionResult::Allowed, - &[], - None, - false, - Some(PermissionResult::Denied { - reason: "session deny".into(), - }), - ); - assert!(result.is_denied()); - } - - #[test] - fn test_precedence_hook_permission_ignored_when_denied() { - let (result, log) = resolve_permission_precedence( - PermissionResult::Denied { - reason: "policy deny".into(), + fn test_tighten_layers_only_tighten() { + // 模拟三层合并：env Allowed + 工具级 Ask + 会话级 Allow → 最终 Ask + use crate::agent::runtime::permission::tighten; + let merged = PermissionResult::Allowed; + let merged = tighten( + merged, + PermissionResult::AskUser { + message: "tool ask".into(), }, - &[], - Some(&("need confirm".to_string(), "test_tool".to_string())), - false, - None, ); - assert!(result.is_denied()); - assert!(log.iter().any(|l| l.contains("ignored"))); + let merged = tighten(merged, PermissionResult::Allowed); + assert!(matches!(merged, PermissionResult::AskUser { .. })); + + // 反向：env Deny + 工具级 Allow + 会话级 Allow → 最终 Deny + let merged = PermissionResult::Denied { + reason: "env".into(), + }; + let merged = tighten(merged.clone(), PermissionResult::Allowed); + let merged = tighten(merged, PermissionResult::Allowed); + assert!(merged.is_denied()); } } diff --git a/src/agent/runtime/session.rs b/src/agent/runtime/session.rs index 31db8dd..47a9e0a 100644 --- a/src/agent/runtime/session.rs +++ b/src/agent/runtime/session.rs @@ -97,6 +97,45 @@ pub async fn create_or_resume_session( } } +/// 使用调用方预分配的会话 ID 创建（或恢复）会话。 +/// +/// 与 `create_or_resume_session` 的区别：会话不存在时以该 ID **创建**而非报错。 +/// 供 API 层使用——请求进入时预生成会话 ID，使取消标记、运行时缓存等 +/// 按 ID 索引的机制对新会话的首个请求同样生效（历史上新会话首请求 +/// 超时只能 abort 任务，无法写入取消标记）。 +pub async fn create_or_resume_session_preallocated( + db: &SqlitePool, + session_id: &str, + llm: &LlmClient, + mode: &str, +) -> anyhow::Result { + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)", + ) + .bind(session_id) + .fetch_one(db) + .await + .unwrap_or(false); + + if exists { + return create_or_resume_session(db, Some(session_id.to_string()), llm, mode).await; + } + + sqlx::query("INSERT INTO agent_sessions (session_id, title, model, mode) VALUES (?, ?, ?, ?)") + .bind(session_id) + .bind("") + .bind(llm.model().await) + .bind(mode) + .execute(db) + .await?; + + Ok(SessionInfo { + session_id: session_id.to_string(), + turn_index: 0, + mode: mode.to_string(), + }) +} + /// 从数据库加载会话的运行模式。 /// /// 返回 None 表示会话不存在或已删除。 @@ -288,6 +327,9 @@ pub async fn rewind_to_message( let preview: String = target_content.chars().take(120).collect(); + // 回退改变了历史 → 上下文快照失效（其 base 高水位与消息集合不再对应） + super::session_events::remove_context_snapshots(db, session_id).await; + info!( "[Session] 回退完成: session={}, rewound={}, to_id={}, new_turn={}", session_id, count, target_message_id, new_turn_index @@ -483,6 +525,9 @@ pub async fn retry_last_turn( .await .unwrap_or(0); + // 硬删除改变了历史 → 上下文快照失效 + super::session_events::remove_context_snapshots(db, session_id).await; + info!( "[Session] 重试: session={}, deleted={} messages from id={}, new_turn={}, has_image={}", session_id, @@ -595,6 +640,7 @@ pub async fn branch_session(db: &SqlitePool, session_id: &str) -> anyhow::Result tx.commit().await?; + // 分叉是新会话：不复制事件（快照引用旧会话的消息 id，直接作废） info!( "[Session] 分叉完成: parent={}, branch={}, copied={} messages, forked_at={}", session_id, branch_id, copied, forked_at diff --git a/src/agent/runtime/session_events.rs b/src/agent/runtime/session_events.rs new file mode 100644 index 0000000..b59ec1d --- /dev/null +++ b/src/agent/runtime/session_events.rs @@ -0,0 +1,446 @@ +// src/agent/runtime/session_events.rs +// +// 会话事件日志 — turn/compaction 生命周期与上下文快照的事件溯源层。 +// +// 不变量："模型可见 ⟺ 已日志化" 的生命周期侧面： +// - 每个 turn 以 turn_start 开始、turn_end 关闭；崩溃留下可检测的 +// 开口（dangling start），恢复时合成 interrupted 关闭而不截断 +// 已持久化的内容（参考 deepseek-harness 的崩溃恢复策略）。 +// - 压缩以 compaction_start/compaction_end 构成日志化锁。 +// - context_snapshot 记录压缩后的折叠上下文 + 消息高水位 +// （base_message_id），下一 turn 回放快照 + 增量消息，避免重复压缩。 + +use serde_json::Value; +use sqlx::SqlitePool; +use tracing::{info, warn}; + +/// 事件类型常量 +pub mod event_types { + pub const TURN_START: &str = "turn_start"; + pub const TURN_END: &str = "turn_end"; + pub const COMPACTION_START: &str = "compaction_start"; + pub const COMPACTION_END: &str = "compaction_end"; + /// 恢复时合成的中断标记（检测到开口的 turn_start/compaction_start） + pub const TURN_INTERRUPTED: &str = "turn_interrupted"; + pub const COMPACTION_INTERRUPTED: &str = "compaction_interrupted"; + /// 压缩后的折叠上下文快照（payload 含 messages + base_message_id） + pub const CONTEXT_SNAPSHOT: &str = "context_snapshot"; +} + +/// 追加一条会话事件 +pub async fn append_event( + db: &SqlitePool, + session_id: &str, + turn_index: i32, + event_type: &str, + payload: Value, +) { + let payload_str = serde_json::to_string(&payload).unwrap_or_default(); + if let Err(e) = sqlx::query( + "INSERT INTO agent_events (session_id, turn_index, event_type, payload) \ + VALUES (?, ?, ?, ?)", + ) + .bind(session_id) + .bind(turn_index) + .bind(event_type) + .bind(&payload_str) + .execute(db) + .await + { + warn!( + "[SessionEvents] 写入 {} 事件失败（非致命）: {}", + event_type, e + ); + } +} + +/// 最新一条指定类型的事件（无则 None） +pub async fn latest_event(db: &SqlitePool, session_id: &str, event_type: &str) -> Option { + let payload: Option = sqlx::query_scalar( + "SELECT payload FROM agent_events \ + WHERE session_id = ? AND event_type = ? ORDER BY id DESC LIMIT 1", + ) + .bind(session_id) + .bind(event_type) + .fetch_optional(db) + .await + .ok() + .flatten(); + + payload.and_then(|p| serde_json::from_str(&p).ok()) +} + +/// 检查是否存在"开口"事件：最新的 `open_type` 之后没有新的 `close_type`。 +/// 用于崩溃恢复——开口的 turn_start/compaction_start 表示上次执行被中断。 +pub async fn has_open_event( + db: &SqlitePool, + session_id: &str, + open_type: &str, + close_type: &str, +) -> bool { + let open_id: Option = sqlx::query_scalar( + "SELECT MAX(id) FROM agent_events WHERE session_id = ? AND event_type = ?", + ) + .bind(session_id) + .bind(open_type) + .fetch_one(db) + .await + .unwrap_or(None); + + let Some(open_id) = open_id else { + return false; + }; + + let close_id: Option = sqlx::query_scalar( + "SELECT MAX(id) FROM agent_events WHERE session_id = ? AND event_type = ?", + ) + .bind(session_id) + .bind(close_type) + .fetch_one(db) + .await + .unwrap_or(None); + + close_id.is_none_or(|c| c < open_id) +} + +/// 恢复会话时的中断协调：检测开口的 turn/compaction，合成 interrupted 关闭。 +/// +/// 这是对崩溃/进程被杀的容错：不截断任何已持久化内容， +/// 只补写"上次被中断"的事实标记。 +pub async fn reconcile_interrupted(db: &SqlitePool, session_id: &str) { + if has_open_event( + db, + session_id, + event_types::TURN_START, + event_types::TURN_END, + ) + .await + { + warn!( + "[SessionEvents] 会话 {} 检测到未关闭的 turn_start（上次执行被中断），合成 interrupted 关闭", + session_id + ); + append_event( + db, + session_id, + -1, + event_types::TURN_END, + serde_json::json!({ "reason": "interrupted", "synthesized": true }), + ) + .await; + append_event( + db, + session_id, + -1, + event_types::TURN_INTERRUPTED, + serde_json::json!({ "note": "检测到崩溃/中断残留，已合成关闭" }), + ) + .await; + } + + if has_open_event( + db, + session_id, + event_types::COMPACTION_START, + event_types::COMPACTION_END, + ) + .await + { + warn!( + "[SessionEvents] 会话 {} 检测到未关闭的 compaction_start，合成 interrupted 关闭", + session_id + ); + append_event( + db, + session_id, + -1, + event_types::COMPACTION_END, + serde_json::json!({ "interrupted": true, "synthesized": true }), + ) + .await; + } +} + +/// 删除会话的上下文快照（rewind/retry/branch 后快照失效） +pub async fn remove_context_snapshots(db: &SqlitePool, session_id: &str) { + if let Err(e) = sqlx::query("DELETE FROM agent_events WHERE session_id = ? AND event_type = ?") + .bind(session_id) + .bind(event_types::CONTEXT_SNAPSHOT) + .execute(db) + .await + { + warn!("[SessionEvents] 清理上下文快照失败: {}", e); + } +} + +/// 保存上下文快照：折叠后的消息 + 消息高水位。 +/// +/// `base_message_id` 是快照时 agent_messages 的最大 id；恢复时加载 +/// 快照消息 + id > base 的增量消息，避免重复压缩。 +pub async fn save_context_snapshot( + db: &SqlitePool, + session_id: &str, + turn_index: i32, + messages: &[crate::clients::llm::ChatMessage], + base_message_id: i64, +) { + // 排除 system 消息（每 turn 重建，且包含缓存的动态 section） + let body: Vec<&crate::clients::llm::ChatMessage> = messages + .iter() + .filter(|m| m.role != crate::clients::llm::MessageRole::System) + .collect(); + + let serialized: Vec = body + .iter() + .map(|m| serde_json::to_value(m).unwrap_or(Value::Null)) + .collect(); + + append_event( + db, + session_id, + turn_index, + event_types::CONTEXT_SNAPSHOT, + serde_json::json!({ + "base_message_id": base_message_id, + "message_count": serialized.len(), + "messages": serialized, + }), + ) + .await; + + info!( + "[SessionEvents] 保存上下文快照: session={} base_msg_id={} messages={}", + session_id, + base_message_id, + serialized.len() + ); +} + +/// 加载上下文快照（若存在）：返回 (快照消息, base_message_id)。 +pub async fn load_context_snapshot( + db: &SqlitePool, + session_id: &str, +) -> Option<(Vec, i64)> { + let payload = latest_event(db, session_id, event_types::CONTEXT_SNAPSHOT).await?; + + let base_message_id = payload.get("base_message_id")?.as_i64()?; + let messages_json = payload.get("messages")?.as_array()?; + + let mut messages = Vec::new(); + for m in messages_json { + match serde_json::from_value::(m.clone()) { + Ok(msg) => messages.push(msg), + Err(e) => { + warn!("[SessionEvents] 快照消息反序列化失败，放弃快照: {}", e); + return None; + } + } + } + + Some((messages, base_message_id)) +} + +/// 获取会话消息的当前最大 id（快照高水位） +pub async fn max_message_id(db: &SqlitePool, session_id: &str) -> i64 { + sqlx::query_scalar("SELECT COALESCE(MAX(id), 0) FROM agent_messages WHERE session_id = ?") + .bind(session_id) + .fetch_one(db) + .await + .unwrap_or(0) +} + +/// 加载 id 大于 `after_id` 的增量历史消息（快照回放后拼接） +pub async fn load_messages_after( + db: &SqlitePool, + session_id: &str, + after_id: i64, +) -> anyhow::Result> { + #[allow(clippy::type_complexity)] + let rows: Vec<( + i64, + String, + String, + Option, + Option, + Option, + )> = sqlx::query_as( + "SELECT id, role, content, tool_calls, tool_call_id, thought FROM agent_messages \ + WHERE session_id = ? AND id > ? AND active = 1 AND agent_name = 'lead' \ + ORDER BY id ASC", + ) + .bind(session_id) + .bind(after_id) + .fetch_all(db) + .await?; + + use crate::clients::llm::{ChatMessage, MessageRole}; + + let mut messages = Vec::new(); + for (_id, role_str, content, tool_calls_json, tool_call_id, thought) in rows { + let role = match role_str.as_str() { + "system" => MessageRole::System, + "user" => MessageRole::User, + "assistant" => MessageRole::Assistant, + "tool" => MessageRole::Tool, + _ => continue, + }; + let tool_calls: Option> = + tool_calls_json.and_then(|j| serde_json::from_str(&j).ok()); + messages.push(ChatMessage { + role, + content: if content.is_empty() { + None + } else { + Some(content) + }, + tool_call_id, + tool_calls, + name: None, + reasoning_content: thought, + }); + } + + Ok(messages) +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn setup_db() -> SqlitePool { + let pool = SqlitePool::connect("sqlite::memory:").await.unwrap(); + sqlx::query( + "CREATE TABLE agent_events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + turn_index INTEGER NOT NULL DEFAULT 0, + event_type TEXT NOT NULL, + payload TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP + )", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "CREATE TABLE agent_messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + turn_index INTEGER NOT NULL DEFAULT 0, + step_index INTEGER NOT NULL DEFAULT 0, + role TEXT NOT NULL, + content TEXT NOT NULL DEFAULT '', + thought TEXT, + tool_calls TEXT, + tool_call_id TEXT, + token_count INTEGER NOT NULL DEFAULT 0, + metadata TEXT, + raw_json TEXT, + agent_name TEXT NOT NULL DEFAULT 'lead', + active INTEGER NOT NULL DEFAULT 1, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP + )", + ) + .execute(&pool) + .await + .unwrap(); + pool + } + + #[tokio::test] + async fn test_open_event_detection() { + let db = setup_db().await; + let sid = "s1"; + + // 无事件 → 无开口 + assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await); + + // turn_start 后无 turn_end → 开口 + append_event(&db, sid, 0, "turn_start", serde_json::json!({})).await; + assert!(has_open_event(&db, sid, "turn_start", "turn_end").await); + + // 关闭后 → 无开口 + append_event( + &db, + sid, + 0, + "turn_end", + serde_json::json!({ "reason": "completed" }), + ) + .await; + assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await); + } + + #[tokio::test] + async fn test_reconcile_interrupted_synthesizes_close() { + let db = setup_db().await; + let sid = "s2"; + + append_event(&db, sid, 0, "turn_start", serde_json::json!({})).await; + append_event(&db, sid, 0, "compaction_start", serde_json::json!({})).await; + + reconcile_interrupted(&db, sid).await; + + // 开口已闭合 + assert!(!has_open_event(&db, sid, "turn_start", "turn_end").await); + assert!(!has_open_event(&db, sid, "compaction_start", "compaction_end").await); + // 合成了中断标记 + assert!(latest_event(&db, sid, "turn_interrupted").await.is_some()); + } + + #[tokio::test] + async fn test_context_snapshot_roundtrip() { + let db = setup_db().await; + let sid = "s3"; + + use crate::clients::llm::ChatMessage; + sqlx::query("INSERT INTO agent_messages (session_id, role, content, agent_name) VALUES (?, 'user', 'm1', 'lead')") + .bind(sid) + .execute(&db) + .await + .unwrap(); + + let base = max_message_id(&db, sid).await; + assert_eq!(base, 1); + + let messages = vec![ + ChatMessage::system("sys"), + ChatMessage::user("[历史对话摘要]\n关于黑洞"), + ChatMessage::assistant("结论"), + ]; + save_context_snapshot(&db, sid, 0, &messages, base).await; + + let (loaded, loaded_base) = load_context_snapshot(&db, sid).await.unwrap(); + assert_eq!(loaded_base, base); + // system 消息被排除 + assert_eq!(loaded.len(), 2); + assert!(loaded[0].text().unwrap().contains("历史对话摘要")); + + // 删除快照 + remove_context_snapshots(&db, sid).await; + assert!(load_context_snapshot(&db, sid).await.is_none()); + } + + #[tokio::test] + async fn test_load_messages_after() { + let db = setup_db().await; + let sid = "s4"; + + for i in 1..=3 { + sqlx::query( + "INSERT INTO agent_messages (id, session_id, role, content, agent_name) \ + VALUES (?, ?, 'user', ?, 'lead')", + ) + .bind(i) + .bind(sid) + .bind(format!("msg-{}", i)) + .execute(&db) + .await + .unwrap(); + } + + let incremental = load_messages_after(&db, sid, 1).await.unwrap(); + assert_eq!(incremental.len(), 2); + assert_eq!(incremental[0].text(), Some("msg-2")); + } +} diff --git a/src/agent/runtime/session_registry.rs b/src/agent/runtime/session_registry.rs new file mode 100644 index 0000000..10b8466 --- /dev/null +++ b/src/agent/runtime/session_registry.rs @@ -0,0 +1,187 @@ +// src/agent/runtime/session_registry.rs +// +// 会话级运行时注册表 — AgentRuntime 按会话缓存复用。 +// +// 历史问题：chat_agent 每个 HTTP 请求 new 一个 AgentRuntime，导致注释中 +// 声称"跨 turn 共享"的状态（后台任务队列、压缩折叠日志、文件缓存、 +// 拒绝追踪器、压缩熔断器）实际随请求销毁： +// - 后台任务结果跨请求丢失（队列无人 drain） +// - 压缩后的上下文不回写，下一 turn 从原始消息重建再重新压缩 +// - prompt cache / file cache 每请求冷启动 +// +// 本注册表将 runtime 生命周期与 session 对齐（参考 deepseek-harness 的 +// "runtime 状态属于会话而非请求"原则）： +// - 同一会话的请求复用同一 runtime 实例 +// - 每个会话持有一个 turn 互斥锁，防止同会话并发 turn 互相破坏 +// - 空闲超过 TTL 的条目由后台清扫回收 +// - 会话删除时显式移除 + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use dashmap::DashMap; +use tracing::info; + +use super::AgentRuntime; +use crate::api::AppState; + +/// 单个会话的注册表条目 +pub struct SessionRuntimeEntry { + pub runtime: Arc, + /// 同会话 turn 串行化锁（防止并发 turn 破坏 turn_index/消息顺序） + pub turn_lock: Arc>, + pub last_used: Instant, + /// 创建时的模式 ID（用于诊断"模式漂移"） + pub mode_id: String, +} + +/// 会话运行时注册表 +pub struct SessionRuntimeRegistry { + entries: DashMap, + idle_ttl: Duration, +} + +impl Default for SessionRuntimeRegistry { + fn default() -> Self { + Self::new() + } +} + +impl SessionRuntimeRegistry { + pub fn new() -> Self { + SessionRuntimeRegistry { + entries: DashMap::new(), + // 2 小时无活动的会话 runtime 允许被回收； + // 后台任务等通过 Arc 自然延长实际生命周期 + idle_ttl: Duration::from_secs(2 * 60 * 60), + } + } + + /// 获取（或创建）指定会话的 runtime。 + /// + /// `mode_id` 只在首次创建时生效——后续请求从 DB 恢复会话模式， + /// 不会因请求参数不同而静默切换（模式由会话决定，一次创建后保持稳定）。 + pub fn get_or_create( + &self, + app_state: Arc, + session_key: &str, + mode_id: &str, + ) -> Arc { + if let Some(mut entry) = self.entries.get_mut(session_key) { + entry.last_used = Instant::now(); + return entry.runtime.clone(); + } + + // DashMap entry API 避免创建竞态 + let entry = self + .entries + .entry(session_key.to_string()) + .or_insert_with(|| { + info!( + "[SessionRegistry] 为会话 {} 创建 runtime (mode={})", + session_key, mode_id + ); + SessionRuntimeEntry { + runtime: Arc::new(AgentRuntime::new_for_session( + app_state, + session_key, + mode_id, + )), + turn_lock: Arc::new(tokio::sync::Mutex::new(())), + last_used: Instant::now(), + mode_id: mode_id.to_string(), + } + }); + entry.runtime.clone() + } + + /// 获取会话的 turn 串行化锁（不存在时返回 None） + pub fn turn_lock(&self, session_key: &str) -> Option>> { + self.entries.get(session_key).map(|e| e.turn_lock.clone()) + } + + /// 移除会话条目（会话删除时调用） + pub fn remove(&self, session_key: &str) { + if self.entries.remove(session_key).is_some() { + info!("[SessionRegistry] 移除会话 {} 的 runtime", session_key); + } + } + + /// 当前缓存的会话数 + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// 回收空闲超过 TTL 的条目（后台周期调用） + pub fn sweep_idle(&self) -> usize { + let now = Instant::now(); + let mut expired: Vec = Vec::new(); + for entry in self.entries.iter() { + if now.duration_since(entry.last_used) > self.idle_ttl { + expired.push(entry.key().clone()); + } + } + let count = expired.len(); + for key in expired { + self.entries.remove(&key); + info!("[SessionRegistry] 回收空闲会话 runtime: {}", key); + } + count + } +} + +// ── 会话级压缩守卫（替代进程级全局静态） ── +// +// 历史实现是单个进程级 AtomicBool：并发两个会话压缩时互相跳过。 +// 改为按 session_id 的集合守卫：同会话递归压缩被拦截，不同会话互不影响。 + +static COMPACTING_SESSIONS: std::sync::Mutex>> = std::sync::Mutex::new(None); + +/// 尝试进入压缩临界区。返回 false 表示该会话已有进行中的压缩（递归守卫触发）。 +pub fn try_begin_compaction(session_id: &str) -> bool { + let mut guard = COMPACTING_SESSIONS + .lock() + .unwrap_or_else(|e| e.into_inner()); + let set = guard.get_or_insert_with(HashSet::new); + if set.contains(session_id) { + false + } else { + set.insert(session_id.to_string()); + true + } +} + +/// 退出压缩临界区。会话不存在时静默（防御性）。 +pub fn end_compaction(session_id: &str) { + let mut guard = COMPACTING_SESSIONS + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(set) = guard.as_mut() { + set.remove(session_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compaction_guard_per_session() { + assert!(try_begin_compaction("session-a")); + // 同会话递归被拦截 + assert!(!try_begin_compaction("session-a")); + // 不同会话互不影响 + assert!(try_begin_compaction("session-b")); + end_compaction("session-a"); + // 释放后可重新进入 + assert!(try_begin_compaction("session-a")); + end_compaction("session-a"); + end_compaction("session-b"); + end_compaction("never-began"); // 防御性：不 panic + } +} diff --git a/src/agent/runtime/streaming.rs b/src/agent/runtime/streaming.rs index 4bdd06a..85c49a7 100644 --- a/src/agent/runtime/streaming.rs +++ b/src/agent/runtime/streaming.rs @@ -4,7 +4,6 @@ // 支持并发取消检测，累积推理内容、文本增量和工具调用。 use std::sync::Arc; -use tokio::sync::mpsc; use tracing::error; use crate::clients::llm::{ @@ -40,7 +39,7 @@ pub async fn process_llm_stream( llm: &LlmClient, messages: &[ChatMessage], tool_defs: &[ToolDefinition], - tx: &mpsc::UnboundedSender, + tap: &crate::agent::engine::EventTap, step: usize, session_id: &str, cancelled_runs: Arc>, @@ -51,7 +50,7 @@ pub async fn process_llm_stream( Ok(rx) => rx, Err(e) => { error!("[Streaming] LLM stream 调用失败: {}", e); - let _ = tx.send(AgentStreamEvent::Error { + tap.send(AgentStreamEvent::Error { message: format!("大模型流式调用失败: {}", e), }); return StreamOutput { @@ -95,15 +94,15 @@ pub async fn process_llm_stream( match event { StreamEvent::ReasoningDelta(delta) => { accumulated_reasoning.push_str(&delta); - let _ = tx.send(AgentStreamEvent::Thought { - content: accumulated_reasoning.clone(), + tap.send(AgentStreamEvent::Thought { + content: tap.thought(&accumulated_reasoning), step, }); } StreamEvent::TextDelta(delta) => { accumulated_content.push_str(&delta); if !is_tool_call_step { - let _ = tx.send(AgentStreamEvent::TextDelta { + tap.send(AgentStreamEvent::TextDelta { content: delta, tool_call_id: None, }); diff --git a/src/agent/runtime/streaming_executor.rs b/src/agent/runtime/streaming_executor.rs deleted file mode 100644 index 73c538b..0000000 --- a/src/agent/runtime/streaming_executor.rs +++ /dev/null @@ -1,929 +0,0 @@ -// src/agent/runtime/streaming_executor.rs -// -// 流式工具执行器。 -// 参考 Claude Code StreamingToolExecutor 设计。 -// -// 当 LLM 流式输出 tool_use 块时，立即开始执行并发安全的工具。 -// 非并发安全的工具排队等待。结果按流中顺序 yield。 -// -// 与 Claude Code 的对齐改进 (2026-06-22)： -// 1. 真正的流式调度 — on_tool_use 中对并发安全工具立即 spawn tokio task -// 2. 并发分区 — 自动分组连续只读工具并行执行 -// 3. Progress 流式 — 长操作进度消息即时 yield -// 4. Sibling Abort — 副效应工具报错时级联中止兄弟姐妹 - -use std::collections::VecDeque; -use std::sync::Arc; -use tokio::sync::broadcast; -use tokio::task::JoinHandle; -use tracing::{info, warn}; - -use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry}; - -/// 流式工具执行状态 -#[derive(Debug, Clone, PartialEq)] -pub enum TrackedToolStatus { - /// 工具调用已从 LLM 流中接收到，等待调度 - Queued, - /// 正在执行中（spawned tokio task 运行中） - Executing, - /// 执行完成，结果就绪等待 yield - Completed, - /// 结果已 yield 给调用方 - Yielded, -} - -/// Sibling Abort 原因 -#[derive(Debug, Clone)] -pub enum AbortReason { - /// 兄弟工具出错触发的级联取消 - SiblingError { description: String }, - /// 用户主动中断 - UserInterrupted, -} - -/// 单次工具执行的结果 -#[derive(Debug)] -struct ToolExecutionResult { - tool_name: String, - output: ToolOutput, -} - -/// 跟踪中的工具执行 -struct TrackedTool { - tool_call_id: String, - tool_name: String, - args: serde_json::Value, - status: TrackedToolStatus, - /// 执行完成后的输出 - output: Option, - /// 并发安全的工具在 spawn 后的 JoinHandle - handle: Option>, -} - -/// 流式工具执行器。 -/// -/// 参考 Claude Code `StreamingToolExecutor` (531 行 TypeScript)， -/// 关键改进：并发安全工具立即 spawn tokio task，不等待 flush。 -pub struct StreamingToolExecutor { - /// 所有跟踪中的工具（按 LLM 流中到达顺序） - tracked: Vec, - /// 工具注册表 - tool_registry: Arc, - /// 工具上下文（按需 clone 给每个 spawn 的 task） - tool_context: ToolContext, - /// Sibling Abort 广播通道 (tx) - abort_tx: broadcast::Sender, - /// Sibling Abort 广播通道 (rx) — 保留以保持 channel 存活， - /// 实际使用时通过 `abort_tx.subscribe()` 获取新接收端。 - #[allow(dead_code)] - abort_rx: broadcast::Receiver, - /// 当前是否已发生错误（触发 sibling abort） - has_errored: bool, - /// 出错工具的描述（如 "bash(git push)"） - errored_tool_desc: String, - /// 最大并发数（预留，当前使用 executing_non_concurrent 控制） - #[allow(dead_code)] - max_concurrency: usize, - /// 最大工具输出字符数 - max_output_chars: usize, - /// 当前正在执行的非并发安全工具数（0 或 1） - executing_non_concurrent: bool, - /// 已完成但尚未 yield 的结果队列（按流顺序） - completed_queue: VecDeque, -} - -impl StreamingToolExecutor { - /// 创建新的流式执行器。 - pub fn new( - tool_registry: Arc, - tool_context: ToolContext, - max_concurrency: usize, - max_output_chars: usize, - ) -> Self { - let (abort_tx, abort_rx) = broadcast::channel(16); - StreamingToolExecutor { - tracked: Vec::new(), - tool_registry, - tool_context, - abort_tx, - abort_rx, - has_errored: false, - errored_tool_desc: String::new(), - max_concurrency, - max_output_chars, - executing_non_concurrent: false, - completed_queue: VecDeque::new(), - } - } - - /// 获取 abort 广播发送端（供外部注入取消信号）。 - pub fn abort_sender(&self) -> broadcast::Sender { - self.abort_tx.clone() - } - - /// 当 LLM 流产生一个新的 tool_use 时调用。 - /// - /// 如果是并发安全工具且当前没有非并发安全工具在执行，立即 spawn tokio task。 - /// 否则加入队列等待调度。 - /// - /// 返回 true 表示该工具已立即开始执行，false 表示排队。 - pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool { - let is_concurrency_safe = self - .tool_registry - .get(&name) - .map(|t| t.is_concurrency_safe(&args)) - .unwrap_or(false); - - let mut tool = TrackedTool { - tool_call_id: call_id.clone(), - tool_name: name.clone(), - args: args.clone(), - status: TrackedToolStatus::Queued, - output: None, - handle: None, - }; - - let idx = self.tracked.len(); - let can_start_now = is_concurrency_safe && !self.executing_non_concurrent; - - if can_start_now { - // 立即 spawn tokio task（参考 Claude Code: addTool 立即 processQueue） - info!( - "[StreamingExecutor] 立即 spawn 并发安全工具: {} (id={})", - name, call_id - ); - let handle = self.spawn_tool_task(idx, call_id.clone(), name.clone(), args.clone()); - tool.handle = Some(handle); - tool.status = TrackedToolStatus::Executing; - } else { - info!( - "[StreamingExecutor] 排队工具: {} (concurrent={}, executing_non_concurrent={})", - name, is_concurrency_safe, self.executing_non_concurrent - ); - } - - if !is_concurrency_safe { - self.executing_non_concurrent = true; - } - - self.tracked.push(tool); - can_start_now - } - - /// LLM 流结束后调用，等待所有剩余排队工具完成。 - pub async fn flush(&mut self) { - let queued_count = self - .tracked - .iter() - .filter(|t| t.status == TrackedToolStatus::Queued) - .count(); - - info!( - "[StreamingExecutor] flush: {} tracked, {} queued, {} executing", - self.tracked.len(), - queued_count, - self.tracked - .iter() - .filter(|t| t.status == TrackedToolStatus::Executing) - .count() - ); - - // 启动所有还在排队的工具 - self.start_all_queued(); - - // 等待所有执行中的工具完成 - self.await_all_executing().await; - } - - /// 按流顺序获取下一个完成的结果（非阻塞）。 - /// - /// 对于已完成的任务，如果其 handle 已就绪则收集结果。 - /// 返回按到达顺序的第一个已完成结果。 - pub fn next_result(&mut self) -> Option<(String, ToolOutput)> { - // 先尝试收集任何已完成的 async task 结果 - self.collect_completed_tasks(); - - // 从 completed_queue 中按序取 - while let Some(&idx) = self.completed_queue.front() { - self.completed_queue.pop_front(); - let tool = &mut self.tracked[idx]; - if tool.status == TrackedToolStatus::Completed { - tool.status = TrackedToolStatus::Yielded; - let output = tool - .output - .take() - .unwrap_or_else(|| ToolOutput::error("工具执行异常：无输出")); - return Some((tool.tool_call_id.clone(), output)); - } - } - None - } - - /// 是否有未 yield 的结果（已完成或即将完成）。 - pub fn has_pending_results(&self) -> bool { - self.tracked - .iter() - .any(|t| t.status == TrackedToolStatus::Completed) - || !self.completed_queue.is_empty() - } - - /// 是否有未完成的工具（仍在排队或执行中）。 - pub fn has_unfinished(&self) -> bool { - self.tracked.iter().any(|t| { - t.status == TrackedToolStatus::Queued || t.status == TrackedToolStatus::Executing - }) - } - - /// 获取所有已完成的结果（包括已 yield 和未 yield 的）。 - pub fn all_results_mut(&mut self) -> Vec<(String, ToolOutput)> { - self.collect_completed_tasks(); - let mut results = Vec::new(); - for tool in &mut self.tracked { - if let Some(output) = tool.output.take() { - results.push((tool.tool_call_id.clone(), output)); - } - tool.status = TrackedToolStatus::Yielded; - } - results - } - - // ── 内部方法 ── - - /// Spawn 一个 tokio task 执行单个工具调用。 - fn spawn_tool_task( - &self, - _idx: usize, - _call_id: String, - tool_name: String, - args: serde_json::Value, - ) -> JoinHandle { - let tool_registry = self.tool_registry.clone(); - let tool_context = self.tool_context.clone(); - let max_output_chars = self.max_output_chars; - let mut abort_rx = self.abort_tx.subscribe(); - - tokio::spawn(async move { - // tokio::select! 在工具执行和 Sibling Abort 之间竞速 - tokio::select! { - result = async { - match tool_registry.get(&tool_name) { - Some(tool) => { - tool.execute_with_progress(args, &tool_context, None).await - } - None => ToolOutput::error(format!("未知工具: {}", tool_name)), - } - } => { - // 截断输出 - let truncated = if result.content.len() > max_output_chars { - let t: String = result.content.chars().take(max_output_chars).collect(); - ToolOutput { - content: format!( - "{}...\n[输出已截断，原始长度: {} 字符]", - t, - result.content.len() - ), - is_error: result.is_error, - metadata: result.metadata, - skip_persist: result.skip_persist, - } - } else { - result - }; - - ToolExecutionResult { - tool_name, - output: truncated, - } - } - Ok(reason) = abort_rx.recv() => { - let msg = match reason { - AbortReason::SiblingError { description } => { - format!("取消：并行工具 {} 出错，已级联取消", description) - } - AbortReason::UserInterrupted => "执行已被用户取消".to_string(), - }; - ToolExecutionResult { - tool_name, - output: ToolOutput::error(msg), - } - } - } - }) - } - - /// 尝试收集所有已完成 tokio task 的结果（非阻塞）。 - fn collect_completed_tasks(&mut self) { - for idx in 0..self.tracked.len() { - if self.tracked[idx].status != TrackedToolStatus::Executing { - continue; - } - if self.tracked[idx].handle.is_none() { - continue; - } - - // 检查 JoinHandle 是否已完成（非阻塞） - let handle = self.tracked[idx].handle.take().unwrap(); - if handle.is_finished() { - // is_finished=true 保证 .await 会立即返回 - // 使用 tokio::task::yield_now 之后的 poll 可能也成功， - // 这里直接在同步上下文中检查后放入完成队列 - // 等下次 async 上下文中通过 await_all_executing 处理 - self.tracked[idx].handle = Some(handle); - // 标记为需要收集 — 将在 flush/await 中处理 - } else { - // 放回未完成的 handle - self.tracked[idx].handle = Some(handle); - } - } - } - - /// 启动所有排队的工具。 - fn start_all_queued(&mut self) { - // 收集需要启动的工具索引（避免借用冲突） - let to_start: Vec = self - .tracked - .iter() - .enumerate() - .filter(|(_, t)| t.status == TrackedToolStatus::Queued) - .filter(|(_, t)| { - let is_safe = self - .tool_registry - .get(&t.tool_name) - .map(|reg_tool| reg_tool.is_concurrency_safe(&t.args)) - .unwrap_or(false); - // 并发安全工具可随时启动，非并发安全的需要独占 - is_safe || !self.executing_non_concurrent - }) - .map(|(i, _)| i) - .collect(); - - for idx in to_start { - let tool = &self.tracked[idx]; - let call_id = tool.tool_call_id.clone(); - let tool_name = tool.tool_name.clone(); - let args = tool.args.clone(); - - let is_safe = self - .tool_registry - .get(&tool_name) - .map(|t| t.is_concurrency_safe(&args)) - .unwrap_or(false); - - let handle = self.spawn_tool_task(idx, call_id, tool_name.clone(), args); - self.tracked[idx].handle = Some(handle); - self.tracked[idx].status = TrackedToolStatus::Executing; - - if !is_safe { - self.executing_non_concurrent = true; - // 非并发安全工具启动后停止（独占执行） - break; - } - - info!("[StreamingExecutor] 启动排队工具: {}", tool_name); - } - } - - /// 等待所有执行中的工具完成。 - async fn await_all_executing(&mut self) { - // 收集所有剩余 JoinHandles - let mut handles: Vec<(usize, JoinHandle)> = Vec::new(); - for idx in 0..self.tracked.len() { - if self.tracked[idx].status == TrackedToolStatus::Executing { - if let Some(handle) = self.tracked[idx].handle.take() { - handles.push((idx, handle)); - } - } - } - - // 并发等待所有任务 - for (idx, handle) in handles { - match handle.await { - Ok(result) => { - let tool_name = result.tool_name.clone(); - let is_error = result.output.is_error; - - self.tracked[idx].output = Some(result.output); - self.tracked[idx].status = TrackedToolStatus::Completed; - self.completed_queue.push_back(idx); - - if is_error { - self.check_sibling_abort(idx, &tool_name); - } - } - Err(e) => { - warn!("[StreamingExecutor] tokio task 异常: {}", e); - self.tracked[idx].output = - Some(ToolOutput::error(format!("工具执行异常: {}", e))); - self.tracked[idx].status = TrackedToolStatus::Completed; - self.completed_queue.push_back(idx); - } - } - } - - self.executing_non_concurrent = false; - } - - /// 检查错误工具是否触发 Sibling Abort。 - fn check_sibling_abort(&mut self, idx: usize, tool_name: &str) { - let causes_abort = self - .tool_registry - .get(tool_name) - .map(|t| t.causes_sibling_abort()) - .unwrap_or(false); - - if causes_abort && !self.has_errored { - warn!( - "[StreamingExecutor] 工具 {} 出错，触发 sibling abort", - tool_name - ); - self.has_errored = true; - self.errored_tool_desc = self.get_tool_description(idx); - let _ = self.abort_tx.send(AbortReason::SiblingError { - description: self.errored_tool_desc.clone(), - }); - } - } - - /// 获取工具的人类可读描述（用于错误消息）。 - fn get_tool_description(&self, idx: usize) -> String { - let tool = &self.tracked[idx]; - let summary = tool - .args - .get("command") - .or_else(|| tool.args.get("file_path")) - .or_else(|| tool.args.get("pattern")) - .and_then(|v| v.as_str()) - .unwrap_or(""); - if summary.is_empty() { - tool.tool_name.clone() - } else { - let truncated: String = summary.chars().take(40).collect(); - if summary.len() > 40 { - format!("{}({}…)", tool.tool_name, truncated) - } else { - format!("{}({})", tool.tool_name, summary) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::agent::tools::AgentTool; - use crate::agent::tools::ToolContext; - use async_trait::async_trait; - use std::sync::atomic::{AtomicBool, Ordering}; - use std::sync::Arc; - - /// 可配置的 Mock 工具，用于测试 StreamingToolExecutor 状态机。 - struct MockAgentTool { - name_str: &'static str, - concurrency_safe: bool, - causes_abort: bool, - /// 执行返回的内容 - result_content: &'static str, - /// 执行是否返回错误 - result_is_error: bool, - /// 可选：执行后设置此标志（用于验证工具是否被调用） - executed: AtomicBool, - } - - impl MockAgentTool { - fn new(name: &'static str) -> Self { - MockAgentTool { - name_str: name, - concurrency_safe: false, - causes_abort: false, - result_content: "mock result", - result_is_error: false, - executed: AtomicBool::new(false), - } - } - - fn concurrency_safe(mut self, v: bool) -> Self { - self.concurrency_safe = v; - self - } - - fn causes_abort(mut self, v: bool) -> Self { - self.causes_abort = v; - self - } - - fn result(mut self, content: &'static str, is_error: bool) -> Self { - self.result_content = content; - self.result_is_error = is_error; - self - } - } - - #[async_trait] - impl AgentTool for MockAgentTool { - fn name(&self) -> &str { - self.name_str - } - fn description(&self) -> &str { - "mock tool for testing" - } - fn parameters(&self) -> serde_json::Value { - serde_json::json!({}) - } - async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { - self.executed.store(true, Ordering::SeqCst); - if self.result_is_error { - ToolOutput::error(self.result_content) - } else { - ToolOutput::success(self.result_content, serde_json::json!({})) - } - } - fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { - self.concurrency_safe - } - fn causes_sibling_abort(&self) -> bool { - self.causes_abort - } - } - - /// 构建测试用的 ToolContext。 - async fn make_test_tool_context() -> ToolContext { - use crate::agent::memory::MemoryManager; - use crate::agent::skills::SkillRegistry; - use crate::api::AppState; - use crate::clients::ads::AdsClient; - use crate::clients::arxiv::ArxivClient; - use crate::clients::cds::vizier::VizierClient; - use crate::clients::llm::{EmbeddingClient, LlmClient}; - use crate::clients::qiniu::QiniuClient; - use crate::services::batch::asset::AssetBatchStatus; - use crate::services::batch::meta::MetaSyncStatus; - use crate::services::download::Downloader; - use crate::services::translation::Dictionary; - use crate::Config; - use std::path::PathBuf; - use std::sync::Arc; - - let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap(); - sqlx::migrate!("./migrations").run(&pool).await.unwrap(); - - let config = Config::from_env(); - let llm = LlmClient::new("tk".into(), "http://localhost".into(), "m".into()).unwrap(); - let embedding = - EmbeddingClient::new("tk".into(), "http://localhost".into(), "e".into()).unwrap(); - let ads = AdsClient::new("tk".into()).unwrap(); - let arxiv = ArxivClient::new().unwrap(); - let vizier = - VizierClient::new("https://tapvizier.cds.unistra.fr/TAPVizieR/tap", 60).unwrap(); - let qiniu = QiniuClient::new( - "ak".into(), - "sk".into(), - "b".into(), - "http://localhost".into(), - ); - - let app_state = Arc::new(AppState { - config, - db: pool, - dict: Dictionary::default(), - qiniu, - downloader: Downloader::new().expect("downloader"), - http_client: reqwest::Client::new(), - 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)), - skill_registry: Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(PathBuf::from( - "/tmp/sk", - )))), - sse_broadcast: None, - memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from( - "/tmp/test_mem", - )))), - login_rate_limiter: Arc::new(dashmap::DashMap::new()), - upload_rate_limiter: Arc::new(dashmap::DashMap::new()), - llm: crate::api::LlmState { - primary: llm.clone(), - medium: llm.clone(), - fast: llm.clone(), - vision: None, - embedding, - }, - sources: crate::api::DataSourceState { - ads, - arxiv, - vizier, - 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", - 90, - ) - .unwrap(), - sdss: crate::clients::sdss::SdssClient::new( - "https://datalab.noirlab.edu/tap/sync", - 90, - ) - .unwrap(), - desi: crate::clients::desi::DesiClient::new( - "https://datalab.noirlab.edu/tap/sync", - 120, - ) - .unwrap(), - irsa: crate::clients::irsa::IrsaClient::new("https://irsa.ipac.caltech.edu", 60) - .unwrap(), - mast: crate::clients::mast::MastClient::new("https://mast.stsci.edu", 90).unwrap(), - observation_registry: std::sync::Arc::new( - crate::services::observation::ObservationRegistry::default(), - ), - }, - session: crate::api::SessionState { - sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())), - session_last_active: Arc::new(dashmap::DashMap::new()), - cancelled_runs: Arc::new(dashmap::DashMap::new()), - session_permission_checkers: Arc::new(dashmap::DashMap::new()), - pending_questions: Arc::new(tokio::sync::Mutex::new( - std::collections::HashMap::new(), - )), - pending_permissions: Arc::new(tokio::sync::Mutex::new( - std::collections::HashMap::new(), - )), - }, - }); - - ToolContext::new(app_state) - } - - /// 创建包含指定 Mock 工具的 StreamingToolExecutor。 - async fn make_executor(tools: Vec) -> StreamingToolExecutor { - let skill_registry = Arc::new(tokio::sync::RwLock::new( - crate::agent::skills::SkillRegistry::new(std::path::PathBuf::from("/tmp/sk")), - )); - let mut registry = ToolRegistry::new(skill_registry); - for tool in tools { - registry.add_tool(Box::new(tool)); - } - let tool_context = make_test_tool_context().await; - StreamingToolExecutor::new(Arc::new(registry), tool_context, 4, 4000) - } - - #[tokio::test] - async fn test_on_tool_use_spawns_concurrency_safe_immediately() { - let mut executor = - make_executor(vec![MockAgentTool::new("safe_tool").concurrency_safe(true)]).await; - - let spawned = executor.on_tool_use( - "call_1".into(), - "safe_tool".into(), - serde_json::json!({"key": "val"}), - ); - assert!(spawned, "concurrency-safe tool should spawn immediately"); - assert_eq!(executor.tracked.len(), 1); - assert_eq!( - executor.tracked[0].status, - TrackedToolStatus::Executing, - "safe tool should be executing" - ); - assert!( - !executor.executing_non_concurrent, - "safe tool does not lock executor" - ); - assert!( - executor.has_unfinished(), - "should have unfinished tool (executing)" - ); - } - - #[tokio::test] - async fn test_on_tool_use_queues_non_concurrency_safe() { - let mut executor = make_executor(vec![ - MockAgentTool::new("unsafe_tool").concurrency_safe(false) - ]) - .await; - - let spawned = - executor.on_tool_use("call_1".into(), "unsafe_tool".into(), serde_json::json!({})); - assert!(!spawned, "non-concurrency-safe tool should be queued"); - assert_eq!(executor.tracked.len(), 1); - assert_eq!(executor.tracked[0].status, TrackedToolStatus::Queued); - assert!( - executor.executing_non_concurrent, - "non-concurrent tool sets the lock flag" - ); - assert!( - executor.has_unfinished(), - "queued tool counts as unfinished" - ); - assert!(!executor.has_pending_results(), "no results yet"); - } - - #[tokio::test] - async fn test_on_tool_use_queues_safe_tool_when_non_concurrent_executing() { - // 先加入一个非并发安全工具（设为 executing_non_concurrent=true）， - // 再尝试加入一个并发安全工具，应排队而非立即执行 - let mut executor = make_executor(vec![ - MockAgentTool::new("unsafe_tool").concurrency_safe(false), - MockAgentTool::new("safe_tool").concurrency_safe(true), - ]) - .await; - - // 第一个：非并发安全，排队但设置 executing_non_concurrent - executor.on_tool_use("call_1".into(), "unsafe_tool".into(), serde_json::json!({})); - // 第二个：并发安全但 executor 被非并发工具锁定，应排队 - let spawned = - executor.on_tool_use("call_2".into(), "safe_tool".into(), serde_json::json!({})); - - assert!( - !spawned, - "safe tool should be queued when executor is locked" - ); - assert_eq!(executor.tracked.len(), 2); - assert_eq!(executor.tracked[1].status, TrackedToolStatus::Queued); - } - - #[tokio::test] - async fn test_get_tool_description_parses_args() { - let mut executor = make_executor(vec![MockAgentTool::new("bash")]).await; - // 手动添加 tracked tool（绕过 on_tool_use 的 spawn） - executor.tracked.push(TrackedTool { - tool_call_id: "c1".into(), - tool_name: "bash".into(), - args: serde_json::json!({"command": "git push origin main"}), - status: TrackedToolStatus::Completed, - output: None, - handle: None, - }); - - let desc = executor.get_tool_description(0); - assert!( - desc.contains("bash"), - "description should contain tool name, got: {}", - desc - ); - assert!( - desc.contains("git push"), - "description should contain command arg, got: {}", - desc - ); - } - - #[tokio::test] - async fn test_get_tool_description_falls_back_to_name() { - let mut executor = make_executor(vec![MockAgentTool::new("unknown_tool")]).await; - executor.tracked.push(TrackedTool { - tool_call_id: "c1".into(), - tool_name: "unknown_tool".into(), - args: serde_json::json!({}), // no command/file_path/pattern - status: TrackedToolStatus::Completed, - output: None, - handle: None, - }); - - let desc = executor.get_tool_description(0); - assert_eq!(desc, "unknown_tool"); - } - - #[tokio::test] - async fn test_flush_starts_queued_and_awaits_completion() { - // 使用并发安全工具测试 flush：on_tool_use 立即 spawn，flush 等待完成 - let mut executor = make_executor(vec![MockAgentTool::new("tool_a") - .concurrency_safe(true) - .result("done", false)]) - .await; - - let spawned = executor.on_tool_use("call_1".into(), "tool_a".into(), serde_json::json!({})); - assert!(spawned, "concurrency-safe tool should spawn immediately"); - assert_eq!(executor.tracked[0].status, TrackedToolStatus::Executing); - - executor.flush().await; - - assert_eq!( - executor.tracked[0].status, - TrackedToolStatus::Completed, - "after flush, tool should be completed" - ); - assert!( - executor.has_pending_results(), - "completed tool = pending result" - ); - assert!(!executor.has_unfinished(), "nothing queued or executing"); - - let result = executor.next_result(); - assert!( - result.is_some(), - "next_result should return the completed tool" - ); - let (call_id, output) = result.unwrap(); - assert_eq!(call_id, "call_1"); - assert!(!output.is_error, "tool should succeed"); - assert_eq!(output.content, "done"); - } - - #[tokio::test] - async fn test_next_result_yields_in_insertion_order() { - let mut executor = make_executor(vec![ - MockAgentTool::new("tool_a") - .concurrency_safe(true) - .result("result_a", false), - MockAgentTool::new("tool_b") - .concurrency_safe(true) - .result("result_b", false), - ]) - .await; - - // 两个并发安全工具，都应该立即 spawn - executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({})); - executor.on_tool_use("call_b".into(), "tool_b".into(), serde_json::json!({})); - - // flush 等待它们完成 - executor.flush().await; - - // 结果应按插入顺序产出 - let result_a = executor.next_result(); - assert!(result_a.is_some()); - assert_eq!(result_a.as_ref().unwrap().0, "call_a"); - assert_eq!(result_a.as_ref().unwrap().1.content, "result_a"); - - let result_b = executor.next_result(); - assert!(result_b.is_some()); - assert_eq!(result_b.as_ref().unwrap().0, "call_b"); - assert_eq!(result_b.as_ref().unwrap().1.content, "result_b"); - - // 第三次调用返回 None（全部已 yield） - let result_c = executor.next_result(); - assert!(result_c.is_none()); - assert!(!executor.has_pending_results()); - } - - #[tokio::test] - async fn test_sibling_abort_broadcast_on_error() { - // 注册两个工具：tool_a 会报错且触发 sibling abort，tool_b 并发执行中被取消 - let mut executor = make_executor(vec![ - MockAgentTool::new("tool_a") - .concurrency_safe(true) - .causes_abort(true) - .result("critical failure", true), - MockAgentTool::new("tool_b") - .concurrency_safe(true) - .result("should be aborted", false), - ]) - .await; - - executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({})); - executor.on_tool_use("call_b".into(), "tool_b".into(), serde_json::json!({})); - - executor.flush().await; - - // tool_a 的结果应该是错误 - let result_a = executor.next_result(); - assert!(result_a.is_some()); - assert!(result_a.unwrap().1.is_error, "tool_a should have errored"); - - // tool_b 可能被取消（sibling abort）或正常完成（取决于竞态） - let result_b = executor.next_result(); - assert!(result_b.is_some(), "tool_b should also have a result"); - - // 验证 has_errored 被设置 - assert!( - executor.has_errored, - "has_errored should be set after sibling abort" - ); - assert!( - !executor.errored_tool_desc.is_empty(), - "errored_tool_desc should be populated" - ); - } - - #[tokio::test] - async fn test_all_results_mut_drains_all_outputs() { - let mut executor = - make_executor(vec![MockAgentTool::new("tool_a").result("result_a", false)]).await; - - executor.on_tool_use("call_a".into(), "tool_a".into(), serde_json::json!({})); - // 不 flush — all_results_mut 是同步方法，只收集已完成的结果 - // 工具可能仍在执行，所以不能保证一定有结果 - let results = executor.all_results_mut(); - // 无论有没有结果，调用后 tracked tool 被标记为 Yielded - for tool in &executor.tracked { - assert_eq!(tool.status, TrackedToolStatus::Yielded); - assert!(tool.output.is_none(), "output should be taken"); - } - // 验证返回的 results 和 tracked 一致 - let _ = results; // 如果 joinhandle 还没完成，results 可能是空的 - } - - #[tokio::test] - async fn test_abort_sender_clone_works() { - let executor = make_executor(vec![]).await; - let sender = executor.abort_sender(); - // 验证 sender 可用 - assert_eq!(sender.receiver_count(), 1); - } - - #[tokio::test] - async fn test_empty_executor_has_no_pending_or_unfinished() { - let executor = make_executor(vec![]).await; - assert!(!executor.has_pending_results()); - assert!(!executor.has_unfinished()); - } -} diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs index 03ca83f..a44cd19 100644 --- a/src/agent/subagent.rs +++ b/src/agent/subagent.rs @@ -13,17 +13,13 @@ use std::sync::Arc; use tokio::sync::mpsc::UnboundedSender; -use tracing::{info, warn}; -use super::compact; -use super::hooks::{ - HookRegistry, PostToolUseContext, PreToolUseContext, SubagentStartContext, SubagentStopContext, -}; -use super::runtime::permission::{PermissionChecker, PermissionResult}; +use super::hooks::{HookRegistry, SubagentStartContext, SubagentStopContext}; +use super::runtime::permission::PermissionChecker; use super::runtime::{AgentConfig, AgentStreamEvent}; -use super::tools::{ToolContext, ToolOutput, ToolRegistry}; +use super::tools::{ToolOutput, ToolRegistry}; use crate::api::AppState; -use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent, ToolDefinition}; +use crate::clients::llm::{ChatMessage, LlmClient}; /// 子代理运行器 pub struct SubAgentRunner { @@ -137,13 +133,29 @@ impl SubAgentRunner { /// * `system_prompt` - 子代理的系统提示词 /// * `research_prompt` - 要执行的研究任务描述 /// * `max_steps` - 子代理最大推理步数（默认 5） + /// * `agent_name` - 稳定名称（Some 时结果可被 followup 续话引用） pub async fn run( &self, system_prompt: &str, research_prompt: &str, max_steps: usize, ) -> ToolOutput { - let subagent_name = format!("sub_{}", &uuid::Uuid::new_v4().to_string()[..8]); + self.run_named(system_prompt, research_prompt, max_steps, None) + .await + } + + /// 带稳定名称的子代理运行（命名子代理的历史持久化在该名称下， + /// 后续可通过 `run_continuation` 续话）。 + pub async fn run_named( + &self, + system_prompt: &str, + research_prompt: &str, + max_steps: usize, + agent_name: Option<&str>, + ) -> ToolOutput { + let subagent_name = agent_name + .map(|n| n.to_string()) + .unwrap_or_else(|| format!("sub_{}", &uuid::Uuid::new_v4().to_string()[..8])); // OnSubagentStart hook if let Some(ref registry) = self.hook_registry { @@ -216,6 +228,60 @@ impl SubAgentRunner { result } + /// 续话一个已存在的命名子代理（continuable subagent）。 + /// + /// 从父会话的 agent_messages 中回放该子代理名下的全部历史作为初始 + /// 上下文（durable 会话 + 冷恢复——dsh Activation 模型的轻量版）， + /// `followup_prompt` 作为追加的用户消息。找不到历史时退化为全新运行。 + pub async fn run_continuation( + &self, + agent_name: &str, + system_prompt: &str, + followup_prompt: &str, + max_steps: usize, + ) -> ToolOutput { + if self.parent_session_id.is_empty() { + return ToolOutput::error("子代理续话需要父会话上下文"); + } + + let history = crate::agent::runtime::session::load_history_for_agent( + &self.app_state.db, + &self.parent_session_id, + agent_name, + ) + .await + .unwrap_or_default(); + + if history.is_empty() { + tracing::info!( + "[SubAgent] 命名子代理 {} 无历史记录，退化为全新运行", + agent_name + ); + return self + .run_named(system_prompt, followup_prompt, max_steps, Some(agent_name)) + .await; + } + + // 保存 followup 用户消息 + self.save_subagent_message(agent_name, 0, 0, "user", followup_prompt, None, None, None); + + let history_len = history.len(); + let mut messages = history; + messages.push(ChatMessage::user(followup_prompt)); + + let outcome = self + .run_engine_with_messages(&mut messages, max_steps, agent_name) + .await; + + tracing::info!( + "[SubAgent] 续话 {}: 回放 {} 条历史消息", + agent_name, + history_len + ); + + outcome + } + /// 保存子代理消息到 agent_messages 表 #[allow(clippy::too_many_arguments)] fn save_subagent_message( @@ -265,7 +331,7 @@ impl SubAgentRunner { }); } - /// 实际执行逻辑（提取为内部方法以便 hook 包装） + /// 全新上下文的执行入口（system + 首条任务消息）。 async fn run_inner( &self, system_prompt: &str, @@ -273,438 +339,118 @@ impl SubAgentRunner { max_steps: usize, subagent_name: &str, ) -> ToolOutput { - let llm = self - .llm_client_override - .as_ref() - .unwrap_or(&self.app_state.llm.primary); - let tool_defs = self.tool_registry.definitions(); - - // 全新上下文 let mut messages = vec![ ChatMessage::system(system_prompt), ChatMessage::user(research_prompt), ]; - - // 跟踪工具调用防止死循环 - let mut last_call: Option<(String, String)> = None; - let mut consecutive_count: usize = 0; - let duplicate_threshold: usize = 3; - - for step in 1..=max_steps { - // 上下文压缩检查 - let est_tokens: usize = messages - .iter() - .map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4) - .sum(); - if est_tokens > self.config.token_soft_limit * 3 / 2 { - info!( - "[SubAgent] 上下文超限 (est. {} tokens)，触发压缩", - est_tokens - ); - compact::compress_context_with_hooks( - &mut messages, - llm, - self.config.token_soft_limit, - self.config.max_messages, - subagent_name, - self.hook_registry.as_ref().map(|a| a.as_ref()), - ) - .await; - } - - // LLM 流式调用 - let mut stream_rx = match llm - .chat_stream(&messages, &tool_defs, self.config.enable_thinking) - .await - { - Ok(rx) => rx, - Err(e) => { - warn!("[SubAgent] LLM stream 失败: {}", e); - return ToolOutput::error(format!("子代理 LLM 调用失败: {}", e)); - } - }; - - let mut accumulated_content = String::new(); - let mut accumulated_reasoning = String::new(); - let mut accumulated_tool_calls: Option> = None; - let mut activity_log: Vec = Vec::new(); - - while let Some(event) = stream_rx.recv().await { - match event { - StreamEvent::ReasoningDelta(delta) => { - accumulated_reasoning.push_str(&delta); - // 转发子代理思考过程到父代理 - if let Some(ref tx) = self.progress_tx { - let _ = tx.send(AgentStreamEvent::Thought { - content: format!("[子代理] {}", accumulated_reasoning), - step, - }); - } - } - StreamEvent::TextDelta(delta) => { - accumulated_content.push_str(&delta); - } - StreamEvent::ToolCallsComplete(tool_calls) => { - // 确保每个工具调用有唯一 ID（LLM 可能不返回 id） - let fixed_tool_calls: Vec = tool_calls - .into_iter() - .map(|tc| { - let id = if tc.id.is_empty() { - format!("call_{}", &uuid::Uuid::new_v4().to_string()[..8]) - } else { - tc.id - }; - activity_log.push(format!("🔧 调用工具: {}", tc.function.name)); - crate::clients::llm::ToolCall { id, ..tc } - }) - .collect(); - accumulated_tool_calls = Some(fixed_tool_calls); - } - StreamEvent::Done => break, - StreamEvent::Error(e) => { - warn!("[SubAgent] 流式错误: {}", e); - return ToolOutput::error(format!("子代理流式错误: {}", e)); - } - _ => {} - } - } - - // 记录思考过程 - if !accumulated_reasoning.is_empty() { - activity_log.push(format!( - "💭 思考: {}", - accumulated_reasoning.chars().take(300).collect::() - )); - } - - // 无工具调用 = 最终回答 - let tool_calls = match accumulated_tool_calls { - Some(ref tc) if !tc.is_empty() => tc.clone(), - _ => { - // 转发子代理结论到父代理（作为 thought 显示在时间线，不污染 finalAnswer） - if let Some(ref tx) = self.progress_tx { - let _ = tx.send(AgentStreamEvent::Thought { - content: accumulated_content.clone(), - step, - }); - } - let content_len = accumulated_content.len(); - info!("[SubAgent] 子代理完成，返回 {} 字符摘要", content_len); - let summary = format!( - "[子代理活动记录]\n\n{}\n\n[子代理结论]\n\n{}", - activity_log.join("\n"), - accumulated_content - ); - return ToolOutput::success( - summary, - serde_json::json!({ - "steps": step, - "content_length": content_len, - "tool_calls": activity_log.iter().filter(|e| e.starts_with("🔧")).count(), - }), - ); - } - }; - - // 构建 assistant 消息 - let assistant_msg = ChatMessage::assistant_with_reasoning( - if accumulated_content.is_empty() { - None - } else { - Some(accumulated_content.clone()) - }, - None, - Some(tool_calls.clone()), - ); - // 持久化 assistant 消息 - let tc_json = serde_json::to_string(&tool_calls).unwrap_or_default(); - self.save_subagent_message( - subagent_name, - 0, - step as i32, - "assistant", - &accumulated_content, - None, - Some(&tc_json), - None, - ); - messages.push(assistant_msg); - - // 执行工具调用 - for tool_call in &tool_calls { - let tool_name = &tool_call.function.name; - let tool_args_str = &tool_call.function.arguments; - - // 死循环检测 - let call_key = (tool_name.clone(), tool_args_str.clone()); - if last_call.as_ref() == Some(&call_key) { - consecutive_count += 1; - if consecutive_count >= duplicate_threshold { - warn!("[SubAgent] 检测到死循环：{}", tool_name); - let error_msg = ChatMessage::tool_result( - &tool_call.id, - format!( - "工具 {} 被连续重复调用。请停止并给出当前收集到的答案。", - tool_name - ), - ); - messages.push(error_msg); - continue; - } - } else { - last_call = Some(call_key); - consecutive_count = 1; - } - - // 解析参数 - let args: serde_json::Value = match serde_json::from_str(tool_args_str) { - Ok(v) => v, - Err(e) => { - let error_msg = - ChatMessage::tool_result(&tool_call.id, format!("参数解析失败: {}", e)); - messages.push(error_msg); - continue; - } - }; - - // ── 向父代理发送进度事件 ── - if let Some(ref tx) = self.progress_tx { - let (is_internal, display_name) = - if let Some(tool) = self.tool_registry.get(tool_name) { - (tool.is_internal(), format!("[sub] {}", tool.display_name())) - } else { - (false, format!("[sub] {}", tool_name)) - }; - let _ = tx.send(AgentStreamEvent::ToolCall { - id: tool_call.id.clone(), - name: format!("[sub] {}", tool_name), - display_name, - arguments: args.clone(), - step, - is_internal, - }); - } - - // ── PreToolUse hooks + Permission check ── - let tool_ctx = ToolContext::silent(self.app_state.clone()); - - let final_args = if let Some(ref registry) = self.hook_registry { - let pre_ctx = PreToolUseContext { - session_id: "subagent".to_string(), - tool_name: tool_name.clone(), - tool_args: args.clone(), - step, - }; - let pre_result = registry.run_pre_tool_use(&pre_ctx).await; - - // Block check - if pre_result.action.is_blocked() { - let reason = pre_result - .action - .block_reason() - .unwrap_or("tool blocked by hook"); - warn!( - "[SubAgent] PreToolUse hook 阻止了工具: {} ({})", - tool_name, reason - ); - let tool_msg = ChatMessage::tool_result( - &tool_call.id, - format!("工具 {} 被阻止: {}", tool_name, reason), - ); - self.save_subagent_message( - subagent_name, - 0, - step as i32, - "tool", - &format!("工具 {} 被阻止: {}", tool_name, reason), - None, - None, - Some(&tool_call.id), - ); - messages.push(tool_msg); - continue; - } - - pre_result.final_args - } else { - args.clone() - }; - - // Permission check — 完整三态检查（包含内容级匹配） - let perm_result = self.permission_checker.check(tool_name, Some(&final_args)); - match perm_result { - PermissionResult::Denied { reason } => { - warn!("[SubAgent] 权限检查拒绝工具 {}: {}", tool_name, reason); - let err_msg = - format!("工具 {} 在子代理上下文中不可用: {}", tool_name, reason); - let tool_msg = ChatMessage::tool_result(&tool_call.id, &err_msg); - self.save_subagent_message( - subagent_name, - 0, - step as i32, - "tool", - &err_msg, - None, - None, - Some(&tool_call.id), - ); - messages.push(tool_msg); - continue; - } - PermissionResult::AskUser { .. } => { - // 子代理上下文中无用户可询问，自动拒绝 - warn!( - "[SubAgent] 工具 {} 需要用户确认，子代理上下文中自动拒绝", - tool_name - ); - let err_msg = - format!("工具 {} 需要用户确认但在子代理上下文中不可用", tool_name); - let tool_msg = ChatMessage::tool_result(&tool_call.id, &err_msg); - self.save_subagent_message( - subagent_name, - 0, - step as i32, - "tool", - &err_msg, - None, - None, - Some(&tool_call.id), - ); - messages.push(tool_msg); - continue; - } - PermissionResult::Allowed => { - // 继续执行 - } - } - - // 执行工具 - let output = match self.tool_registry.get(tool_name) { - Some(tool) => { - match tokio::time::timeout( - std::time::Duration::from_secs(self.config.tool_timeout_secs), - tool.execute(final_args.clone(), &tool_ctx), - ) - .await - { - Ok(output) => output, - Err(_) => ToolOutput::error(format!("工具 {} 执行超时", tool_name)), - } - } - None => ToolOutput::error(format!("未知工具: {}", tool_name)), - }; - - // ── PostToolUse hooks ── - let final_output_content = if let Some(ref registry) = self.hook_registry { - let post_ctx = PostToolUseContext { - session_id: "subagent".to_string(), - agent_name: "subagent".to_string(), - tool_name: tool_name.clone(), - tool_args: final_args, - output_content: output.content.clone(), - is_error: output.is_error, - step, - elapsed_ms: 0, - }; - let post_result = registry.run_post_tool_use(&post_ctx).await; - post_result.final_content - } else { - output.content.clone() - }; - - // 向父代理发送工具结果进度 - if let Some(ref tx) = self.progress_tx { - let preview: String = final_output_content.chars().take(200).collect(); - let (is_internal, display_name) = - if let Some(tool) = self.tool_registry.get(tool_name) { - (tool.is_internal(), format!("[sub] {}", tool.display_name())) - } else { - (false, format!("[sub] {}", tool_name)) - }; - let _ = tx.send(AgentStreamEvent::ToolResult { - tool_call_id: tool_call.id.clone(), - name: format!("[sub] {}", tool_name), - display_name, - output: preview, - is_error: output.is_error, - metadata: serde_json::json!({}), - step, - is_internal, - }); - } - - // 截断输出（使用 post-hook 处理后的内容） - let truncated = if final_output_content.len() > self.config.max_tool_output_chars { - let t: String = final_output_content - .chars() - .take(self.config.max_tool_output_chars) - .collect(); - format!( - "{}...\n[已截断，原始 {} 字符]", - t, - final_output_content.len() - ) - } else { - final_output_content.clone() - }; - - let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated); - self.save_subagent_message( - subagent_name, - 0, - step as i32, - "tool", - &truncated, - None, - None, - Some(&tool_call.id), - ); - messages.push(tool_msg); - } - } - - // 达到最大步数，强制生成最终答案 - info!("[SubAgent] 达到最大步数 ({}), 生成最终答案", max_steps); - self.force_final_answer(llm, &messages).await + self.run_engine_with_messages(&mut messages, max_steps, subagent_name) + .await } - /// 强制 LLM 生成最终答案（不带工具调用） - async fn force_final_answer(&self, llm: &LlmClient, messages: &[ChatMessage]) -> ToolOutput { - let mut final_messages = messages.to_vec(); - final_messages.push(ChatMessage::user( - "请根据已收集的信息直接给出最终答案，不要再调用工具。", - )); + /// 实际执行逻辑 — 委托给统一 ReAct 引擎（与主代理共用同一循环实现）。 + /// + /// 子代理注入组合：事件前缀 [sub]、AskPolicy::AutoDeny（无人值守 + /// fail-closed）、独立持久化 sink（agent_name = sub_xxx）、无后台队列。 + /// `messages` 由调用方构造：全新运行（system+user）或续话（回放历史）。 + async fn run_engine_with_messages( + &self, + messages: &mut Vec, + max_steps: usize, + subagent_name: &str, + ) -> ToolOutput { + use crate::agent::engine::{AskPolicy, CancelSource, DbMessageSink, EventTap, ReactEngine}; - let empty_tools: Vec = Vec::new(); - let mut stream_rx = match llm - .chat_stream(&final_messages, &empty_tools, self.config.enable_thinking) - .await - { - Ok(rx) => rx, - Err(e) => { - return ToolOutput::error(format!("子代理最终答案生成失败: {}", e)); - } + let llm = self + .llm_client_override + .as_ref() + .unwrap_or(&self.app_state.llm.primary); + + let mut engine_config = self.config.clone(); + engine_config.max_steps = max_steps; + + let tap = match &self.progress_tx { + Some(tx) => EventTap::prefixed(tx.clone(), "sub"), + None => EventTap::none(), }; - let mut accumulated = String::new(); - while let Some(event) = stream_rx.recv().await { - match event { - StreamEvent::TextDelta(delta) => { - accumulated.push_str(&delta); - } - StreamEvent::Done => break, - StreamEvent::Error(e) => { - warn!("[SubAgent] 最终答案流式错误: {}", e); - break; - } - _ => {} - } - } + // 子代理消息落库（父会话下以 agent_name 区分；无父会话则不持久化） + let sink: Option> = + if self.parent_session_id.is_empty() { + None + } else { + Some(std::sync::Arc::new(DbMessageSink { + db: self.app_state.db.clone(), + session_id: self.parent_session_id.clone(), + agent_name: subagent_name.to_string(), + fixed_metadata: Some(serde_json::json!({ + "agent": subagent_name, + "is_subagent": true, + })), + })) + }; - if accumulated.is_empty() { - ToolOutput::error("子代理无法生成最终答案") - } else { - ToolOutput::success(accumulated, serde_json::json!({ "forced": true })) + let engine = ReactEngine { + llm, + tool_registry: &self.tool_registry, + app_state: self.app_state.clone(), + config: engine_config, + session_id: if self.parent_session_id.is_empty() { + subagent_name.to_string() + } else { + self.parent_session_id.clone() + }, + turn_index: 0, + agent_name: subagent_name.to_string(), + tap, + sink, + hooks: self.hook_registry.as_deref(), + ask_policy: AskPolicy::AutoDeny, + permission_checker: Some(self.permission_checker.as_ref()), + session_checker: None, + denial_tracker: None, + checkpoint_manager: None, + read_file_state: std::sync::Arc::new(std::sync::Mutex::new( + crate::agent::runtime::file_cache::FileStateCache::new(), + )), + cancel: CancelSource::Flag(std::sync::Arc::new(std::sync::atomic::AtomicBool::new( + false, + ))), + breaker: None, + collapse_log: None, + bg_queue: None, + // 子代理不注入 todo nag / token nudge（任务短、上下文隔离） + nudges: false, + diminishing_returns: true, + // 子代理同样享受错误恢复阶梯（退避重试/降级压缩） + recovery: true, + persist_todos: false, + }; + + let outcome = engine.run(messages).await; + + match outcome.terminal { + Some(super::terminal::TurnTerminal::ModelError { message, .. }) => { + ToolOutput::error(format!("子代理 LLM 调用失败: {}", message)) + } + Some(super::terminal::TurnTerminal::CancelledByUser { .. }) => { + ToolOutput::error("子代理执行被中止") + } + _ => { + if outcome.final_text.is_empty() { + ToolOutput::error("子代理无法生成最终答案") + } else { + let tool_calls_total: usize = outcome.metrics.tool_calls.values().sum(); + ToolOutput::success( + format!("[子代理结论]\n\n{}", outcome.final_text), + serde_json::json!({ + "steps": outcome.metrics.total_steps, + "tool_calls": tool_calls_total, + "content_length": outcome.final_text.len(), + "agent_name": subagent_name, + }), + ) + } + } } } } diff --git a/src/agent/team/teammate.rs b/src/agent/team/teammate.rs index 179b2f7..0b253d6 100644 --- a/src/agent/team/teammate.rs +++ b/src/agent/team/teammate.rs @@ -6,14 +6,13 @@ use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tokio::sync::Mutex; -use tracing::{info, warn}; +use tracing::info; use crate::agent::background::BgNotificationQueue; -use crate::agent::compact; use crate::agent::runtime::AgentConfig; use crate::agent::tools::ToolRegistry; use crate::api::AppState; -use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent}; +use crate::clients::llm::{ChatMessage, LlmClient}; use super::config::MemberStatus; use super::inbox::{self, TeamMessageType}; @@ -35,10 +34,20 @@ pub async fn run_teammate_loop( cancelled: Arc, ) { let llm = &app_state.llm.primary; - // 队友的工具注册表排除 subagent（防止无限委托链） + // 队友的工具注册表排除 subagent（防止无限委托链）。 + // 历史版本注释声称排除但代码未过滤——现在通过 definition_filter 真正生效： + // 队友的 LLM 请求不再包含 subagent 工具定义。 let queue = Arc::new(BgNotificationQueue::new()); - let tool_registry = + let mut tool_registry = ToolRegistry::new_with_queue(Some(queue.clone()), app_state.skill_registry.clone()); + { + let all = tool_registry.tool_names(); + let filtered: Vec = all + .into_iter() + .filter(|n| n != "subagent" && n != "spawn_teammate") + .collect(); + tool_registry.set_definition_filter(filtered); + } let tool_defs = tool_registry.definitions(); @@ -153,10 +162,12 @@ pub async fn run_teammate_loop( } } -/// 队友的单次 ReAct turn。 +/// 队友的单次 ReAct turn — 委托给统一 ReAct 引擎。 /// -/// 一个简化的 ReAct 循环：LLM 调用 → 工具执行 → 结果注入 → 循环... -/// 与主 Agent 的循环类似但更轻量（无 SSE、无 DB 持久化、无 hooks）。 +/// 队友注入组合：静默（无 SSE）、无持久化、AskPolicy::AutoDeny（无人值守 +/// fail-closed）、AtomicFlag 取消（TeamMemberHandle.request_stop）。 +/// 与历史手写循环相比，队友获得了错误恢复阶梯、压缩熔断、重复检测等 +/// 与主代理对等的能力。 async fn run_teammate_react_turn( llm: &LlmClient, tool_defs: &[crate::clients::llm::ToolDefinition], @@ -166,115 +177,46 @@ async fn run_teammate_react_turn( config: &AgentConfig, cancelled: &Arc, ) -> Option { - let max_steps = config.max_steps.min(5); // 队友步数限制更严格 + use crate::agent::engine::{AskPolicy, CancelSource, EventTap, ReactEngine}; - for _step in 1..=max_steps { - // 检查取消 - if cancelled.load(Ordering::SeqCst) { - return None; - } + let _ = tool_defs; // 引擎自行从 tool_registry 取 definitions（保持签名兼容） - // 上下文压缩检查 - let est_tokens: usize = messages - .iter() - .map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4) - .sum(); - if est_tokens > config.token_soft_limit * 3 / 2 { - compact::compress_context( - messages, - llm, - config.token_soft_limit, - config.max_messages, - "teammate", - ) - .await; - } + let mut engine_config = config.clone(); + engine_config.max_steps = config.max_steps.min(5); // 队友步数限制更严格 - // LLM 流式调用 - let mut stream_rx = match llm - .chat_stream(messages, tool_defs, config.enable_thinking) - .await - { - Ok(rx) => rx, - Err(_) => return None, - }; + let engine = ReactEngine { + llm, + tool_registry, + app_state: app_state.clone(), + config: engine_config, + session_id: format!("teammate-turn-{}", uuid::Uuid::new_v4()), + turn_index: 0, + agent_name: "teammate".to_string(), + tap: EventTap::none(), + sink: None, + hooks: None, + ask_policy: AskPolicy::AutoDeny, + permission_checker: None, + session_checker: None, + denial_tracker: None, + checkpoint_manager: None, + read_file_state: Arc::new(std::sync::Mutex::new( + crate::agent::runtime::file_cache::FileStateCache::new(), + )), + cancel: CancelSource::Flag(cancelled.clone()), + breaker: None, + collapse_log: None, + bg_queue: None, + nudges: false, + diminishing_returns: true, + recovery: true, + persist_todos: false, + }; - let mut accumulated = String::new(); - let mut tool_calls: Option> = None; - - while let Some(event) = stream_rx.recv().await { - match event { - StreamEvent::TextDelta(delta) => accumulated.push_str(&delta), - StreamEvent::ToolCallsComplete(tc) => tool_calls = Some(tc), - StreamEvent::Done => break, - StreamEvent::Error(_) => return None, - _ => {} - } - } - - // 无工具调用 = 最终回答 - let tool_calls = match tool_calls { - Some(ref tc) if !tc.is_empty() => tc.clone(), - _ => { - return if accumulated.is_empty() { - None - } else { - Some(accumulated) - }; - } - }; - - // 构建 assistant 消息 - messages.push(ChatMessage::assistant_with_reasoning( - if accumulated.is_empty() { - None - } else { - Some(accumulated) - }, - None, - Some(tool_calls.clone()), - )); - - // 执行工具调用 - for tc in &tool_calls { - let args: serde_json::Value = match serde_json::from_str(&tc.function.arguments) { - Ok(a) => a, - Err(_) => continue, - }; - - let tool_ctx = crate::agent::tools::ToolContext::silent(app_state.clone()); - - let output = match tool_registry.get(&tc.function.name) { - Some(tool) => { - match tokio::time::timeout( - std::time::Duration::from_secs(config.tool_timeout_secs), - tool.execute(args, &tool_ctx), - ) - .await - { - Ok(o) => o, - Err(_) => crate::agent::tools::ToolOutput::error("执行超时"), - } - } - None => crate::agent::tools::ToolOutput::error("未知工具"), - }; - - let truncated = if output.content.len() > config.max_tool_output_chars { - let t: String = output - .content - .chars() - .take(config.max_tool_output_chars) - .collect(); - format!("{}...\n[已截断]", t) - } else { - output.content - }; - - messages.push(ChatMessage::tool_result(&tc.id, &truncated)); - } + let outcome = engine.run(messages).await; + if outcome.final_text.is_empty() { + None + } else { + Some(outcome.final_text) } - - // 达到最大步数，返回 None（无结果） - warn!("[Teammate] 达到最大步数限制 ({} steps)，无结果", max_steps); - None } diff --git a/src/agent/terminal.rs b/src/agent/terminal.rs index 441f713..0ac3475 100644 --- a/src/agent/terminal.rs +++ b/src/agent/terminal.rs @@ -74,6 +74,18 @@ impl TurnTerminal { TurnTerminal::ModelError { .. } => "模型错误", } } + + /// 机器可读的终止标签（会话事件日志用） + pub fn reason_label(&self) -> &'static str { + match self { + TurnTerminal::Completed { .. } => "completed", + TurnTerminal::MaxStepsReached { .. } => "max_steps", + TurnTerminal::CancelledByUser { .. } => "cancelled", + TurnTerminal::DuplicateCallDetected { .. } => "duplicate", + TurnTerminal::ModelStreamError { .. } => "stream_error", + TurnTerminal::ModelError { .. } => "model_error", + } + } } impl std::fmt::Display for TurnTerminal { diff --git a/src/agent/tools/astro/mod.rs b/src/agent/tools/astro/mod.rs index d1b648f..21f877f 100644 --- a/src/agent/tools/astro/mod.rs +++ b/src/agent/tools/astro/mod.rs @@ -17,12 +17,22 @@ pub use system::process::ProcessPaperTool; pub use system::search::SearchPapersTool; // 研究级工具 +pub use research::abundance_analysis::AnalyzeAbundancePatternTool; +pub use research::kinematics_tool::AnalyzeKinematicsTool; pub use research::library::{GetCitationNetworkTool, SearchLocalLibraryTool}; pub use research::metadata::GetPaperMetadataTool; pub use research::note::SaveNoteTool; +pub use research::observation::FindObservationTool; +pub use research::observation_plan::PlanObservationTool; pub use research::paper::{GetPaperContentTool, GetPaperOutlineTool}; +pub use research::period_search::FindPeriodTool; +pub use research::radial_velocity::MeasureRadialVelocityTool; pub use research::rag::RagSearchTool; +pub use research::sed_fit::FitSedTool; +pub use research::spectral_lines::MeasureSpectralLinesTool; +pub use research::spectrum_params::GetSpectrumParametersTool; pub use research::target::QueryTargetTool; +pub use research::variable_star::ClassifyVariableStarTool; pub use research::vizier::CatalogOperationTool; // 以下工具已合并入 ProcessPaperTool，不再独立导出： diff --git a/src/agent/tools/astro/research/abundance_analysis.rs b/src/agent/tools/astro/research/abundance_analysis.rs new file mode 100644 index 0000000..ae4c871 --- /dev/null +++ b/src/agent/tools/astro/research/abundance_analysis.rs @@ -0,0 +1,104 @@ +// src/agent/tools/astro/research/abundance_analysis.rs +// +// AnalyzeAbundancePatternTool —— 化学丰度模式分析 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct AnalyzeAbundancePatternTool; + +#[async_trait] +impl AgentTool for AnalyzeAbundancePatternTool { + fn name(&self) -> &str { + "analyze_abundance_pattern" + } + + fn display_name(&self) -> &str { + "化学丰度分析" + } + + fn description(&self) -> &str { + "分析恒星的化学丰度模式。输入元素丰度表（[X/H] 值），计算 [α/Fe]（O, Mg, Si, S, Ca, Ti 加权平均），\ + 根据 [α/Fe] vs [Fe/H] 位置区分薄盘/厚盘/晕星。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "abundances": { + "type": "array", + "items": { + "type": "object", + "properties": { + "element": { "type": "string", "description": "元素符号（如 'Fe', 'O', 'Mg'）" }, + "value": { "type": "number", "description": "[X/H] 丰度值 (dex)" }, + "error": { "type": "number", "description": "误差 (dex)，可选" } + }, + "required": ["element", "value"] + }, + "description": "元素丰度列表" + } + }, + "required": ["abundances"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let abundances: Vec = + match args.get("abundances").and_then(|a| a.as_array()) { + Some(arr) => { + let mut entries = Vec::new(); + for item in arr { + let element = match item.get("element").and_then(|e| e.as_str()) { + Some(e) => e.to_string(), + None => continue, + }; + let value = match item.get("value").and_then(|v| v.as_f64()) { + Some(v) => v, + None => continue, + }; + let error = item.get("error").and_then(|e| e.as_f64()); + entries.push(crate::services::spectrum::parameters::AbundanceEntry { + element, + value, + error, + }); + } + entries + } + None => return ToolOutput::error("缺少必需参数 'abundances'"), + }; + + if abundances.is_empty() { + return ToolOutput::error("丰度数据为空"); + } + + info!("[AnalyzeAbundance] {} 个元素丰度", abundances.len()); + + match crate::services::analysis::chemistry::analyze_chemistry(&abundances) { + Ok(result) => { + let report = crate::services::analysis::chemistry::format_chemistry(&result); + let json_str = serde_json::to_string_pretty(&result).unwrap_or_default(); + let content = format!("{}\n\n```json\n{}\n```", report, json_str); + ToolOutput::success(content, json!(result)) + } + Err(e) => ToolOutput::error(format!("丰度分析失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/research/batch_process.rs b/src/agent/tools/astro/research/batch_process.rs new file mode 100644 index 0000000..623a99f --- /dev/null +++ b/src/agent/tools/astro/research/batch_process.rs @@ -0,0 +1,168 @@ +// src/agent/tools/astro/research/batch_process.rs +// +// BatchProcessTool —— 批量样本处理：查询→下载→分析→出图→报告 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct BatchProcessTool; + +#[async_trait] +impl AgentTool for BatchProcessTool { + fn name(&self) -> &str { + "batch_process" + } + + fn display_name(&self) -> &str { + "批量样本处理" + } + + fn description(&self) -> &str { + "批量处理天体样本：给定一组天体名称或坐标，自动执行 查询目标→获取观测数据→提取参数→生成报告。\ + 支持并发控制，自动处理速率限制。输出汇总 Markdown 表格。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "targets": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": { "type": "string", "description": "天体名称" }, + "ra": { "type": "number", "description": "赤经 (度)" }, + "dec": { "type": "number", "description": "赤纬 (度)" } + } + }, + "description": "目标列表（名称或坐标）" + }, + "data_source": { + "type": "string", + "enum": ["lamost", "sdss", "gaia", "desi"], + "description": "观测数据源", + "default": "lamost" + }, + "max_concurrent": { + "type": "integer", + "description": "最大并发数", + "default": 5 + }, + "extract_parameters": { + "type": "boolean", + "description": "是否提取光谱参数", + "default": true + } + }, + "required": ["targets"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let targets = match args.get("targets").and_then(|t| t.as_array()) { + Some(arr) => arr, + None => return ToolOutput::error("缺少必需参数 'targets'"), + }; + + let data_source = args + .get("data_source") + .and_then(|s| s.as_str()) + .unwrap_or("lamost"); + let extract_params = args + .get("extract_parameters") + .and_then(|e| e.as_bool()) + .unwrap_or(true); + + info!( + "[BatchProcess] 处理 {} 个目标, 数据源={}", + targets.len(), + data_source + ); + + let mut results = Vec::new(); + let mut success = 0; + let mut failed = 0; + + for target in targets.iter() { + let name = target + .get("name") + .and_then(|n| n.as_str()) + .unwrap_or("unknown"); + let ra = target.get("ra").and_then(|r| r.as_f64()); + let dec = target.get("dec").and_then(|d| d.as_f64()); + + // 1. 查询目标信息 + let query_name = name.strip_prefix("NAME|").unwrap_or(name); + let target_info = crate::services::cds::target::query_target_cached( + &ctx.app_state.db, + query_name, + None, + &ctx.app_state.http_client, + ) + .await + .ok(); + + let (ra, dec) = if let (Some(r), Some(d)) = (ra, dec) { + (r, d) + } else if let Some(ref info) = target_info { + let r = info + .ra + .as_ref() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let d = info + .dec + .as_ref() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + (r, d) + } else { + failed += 1; + results.push(format!("| {} | 查询失败 | - | - |", name)); + continue; + }; + + // 2. 获取观测数据（简化：仅记录元数据） + let obs_info = format!("ra={:.4}, dec={:.4}", ra, dec); + + // 3. 提取参数（如果启用） + let params_info = if extract_params { + "参数提取待执行".to_string() + } else { + "跳过".to_string() + }; + + success += 1; + results.push(format!( + "| {} | {} | {} | {} |", + name, obs_info, params_info, "✅" + )); + } + + let mut output = format!( + "## 批量处理报告\n\n\ + - 总计: {} 个目标\n\ + - 成功: {}\n\ + - 失败: {}\n\n\ + | 目标 | 坐标 | 参数 | 状态 |\n\ + |------|------|------|------|\n", + targets.len(), + success, + failed + ); + output.push_str(&results.join("\n")); + + ToolOutput::success( + output, + json!({ "total": targets.len(), "success": success, "failed": failed }), + ) + } +} diff --git a/src/agent/tools/astro/research/kinematics_tool.rs b/src/agent/tools/astro/research/kinematics_tool.rs new file mode 100644 index 0000000..a6d6667 --- /dev/null +++ b/src/agent/tools/astro/research/kinematics_tool.rs @@ -0,0 +1,139 @@ +// src/agent/tools/astro/research/kinematics_tool.rs +// +// AnalyzeKinematicsTool —— 自行/视差运动学分析 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct AnalyzeKinematicsTool; + +#[async_trait] +impl AgentTool for AnalyzeKinematicsTool { + fn name(&self) -> &str { + "analyze_kinematics" + } + + fn display_name(&self) -> &str { + "运动学分析" + } + + fn description(&self) -> &str { + "基于 Gaia 自行/视差数据计算银河系三维空间速度（UVW），识别移动星群（薄盘/厚盘/晕）。\ + 输入天体的赤道坐标、视差、自行和径向速度，输出银心距、三维速度和星族分类。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "targets": { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": { "type": "string", "description": "标识符" }, + "ra": { "type": "number", "description": "赤经 (度, J2000)" }, + "dec": { "type": "number", "description": "赤纬 (度, J2000)" }, + "parallax": { "type": "number", "description": "视差 (mas)" }, + "pm_ra": { "type": "number", "description": "自行 RA (mas/yr)" }, + "pm_dec": { "type": "number", "description": "自行 Dec (mas/yr)" }, + "rv": { "type": "number", "description": "径向速度 (km/s)" } + }, + "required": ["ra", "dec", "parallax", "pm_ra", "pm_dec", "rv"] + }, + "description": "天体列表" + } + }, + "required": ["targets"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let targets: Vec = + match args.get("targets").and_then(|t| t.as_array()) { + Some(arr) => { + let mut inputs = Vec::new(); + for item in arr { + let ra = match item.get("ra").and_then(|r| r.as_f64()) { + Some(r) => r, + None => continue, + }; + let dec = match item.get("dec").and_then(|d| d.as_f64()) { + Some(d) => d, + None => continue, + }; + let parallax = match item.get("parallax").and_then(|p| p.as_f64()) { + Some(p) => p, + None => continue, + }; + let pm_ra = item.get("pm_ra").and_then(|p| p.as_f64()).unwrap_or(0.0); + let pm_dec = item.get("pm_dec").and_then(|p| p.as_f64()).unwrap_or(0.0); + let rv = item.get("rv").and_then(|r| r.as_f64()).unwrap_or(0.0); + let id = item + .get("id") + .and_then(|i| i.as_str()) + .map(|s| s.to_string()); + inputs.push(crate::services::analysis::kinematics::KinematicInput { + ra, + dec, + parallax, + pm_ra, + pm_dec, + rv, + id, + }); + } + inputs + } + None => return ToolOutput::error("缺少必需参数 'targets'"), + }; + + if targets.is_empty() { + return ToolOutput::error("目标列表为空"); + } + + info!("[AnalyzeKinematics] 分析 {} 个目标", targets.len()); + + let results = crate::services::analysis::kinematics::compute_kinematics_batch(&targets); + + let mut output = String::from("## 运动学分析结果\n\n"); + let mut success_count = 0; + + for (i, result) in results.into_iter().enumerate() { + match result { + Ok(r) => { + success_count += 1; + output.push_str( + &crate::services::analysis::kinematics::format_kinematic_result(&r), + ); + output.push_str("\n---\n\n"); + } + Err(e) => { + output.push_str(&format!("### 目标 {} 失败: {}\n\n", i + 1, e)); + } + } + } + + output.insert_str(0, &format!("成功: {}/{}\n\n", success_count, targets.len())); + + ToolOutput::success( + output, + json!({ "success_count": success_count, "total": targets.len() }), + ) + } +} diff --git a/src/agent/tools/astro/research/literature_monitor.rs b/src/agent/tools/astro/research/literature_monitor.rs new file mode 100644 index 0000000..1d6b6fd --- /dev/null +++ b/src/agent/tools/astro/research/literature_monitor.rs @@ -0,0 +1,158 @@ +// src/agent/tools/astro/research/literature_monitor.rs +// +// LiteratureMonitorTool —— 定时文献监控 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct LiteratureMonitorTool; + +#[async_trait] +impl AgentTool for LiteratureMonitorTool { + fn name(&self) -> &str { + "literature_monitor" + } + + fn display_name(&self) -> &str { + "文献监控" + } + + fn description(&self) -> &str { + "检查 ADS/arXiv 最新文献，匹配用户关注的关键词、作者、天体名称。\ + 返回最近的新文献列表。可配合 loop 技能实现定时监控。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "ADS 查询字符串（如 'abs:LAMOST spectral' 或 'author:Zhang AND abs:metallicity'）" + }, + "max_results": { + "type": "integer", + "description": "最大返回数量", + "default": 10 + }, + "days_back": { + "type": "integer", + "description": "检查最近 N 天的文献", + "default": 7 + }, + "min_citations": { + "type": "integer", + "description": "最低引用数过滤", + "default": 0 + } + }, + "required": ["query"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let query = match args.get("query").and_then(|q| q.as_str()) { + Some(q) => q, + None => return ToolOutput::error("缺少必需参数 'query'"), + }; + + let max_results = args + .get("max_results") + .and_then(|m| m.as_u64()) + .unwrap_or(10) as i32; + let days_back = args.get("days_back").and_then(|d| d.as_u64()).unwrap_or(7); + let min_citations = args + .get("min_citations") + .and_then(|c| c.as_u64()) + .unwrap_or(0) as i32; + + info!( + "[LitMonitor] 本地数据库查询='{}', 最近{}天, 最大{}条", + query, days_back, max_results + ); + + match crate::services::paper::get_recent_papers( + &ctx.app_state.db, + &ctx.app_state.config.storage.library_dir, + query, + days_back as i64, + max_results, + ) + .await + { + Ok(papers) => { + let mut output = format!( + "## 文献监控结果 (来自本地同步库)\n\n查询词: `{}`\n\n", + query + ); + output.push_str("| # | 标题 | 作者 | 年份 | 引用 | Bibcode |\n"); + output.push_str("|---|------|------|------|------|---------|\n"); + + let mut count = 0; + for paper in papers { + if paper.citation_count < min_citations { + continue; + } + count += 1; + + let authors = paper + .authors + .iter() + .take(3) + .map(|a| a.as_str()) + .collect::>() + .join(", "); + let authors_str = if paper.authors.len() > 3 { + format!("{} et al.", authors) + } else { + authors + }; + + let md_path = ctx + .app_state + .config + .storage + .library_dir + .join("Markdown") + .join(format!("{}.md", paper.bibcode)); + let md_path_str = md_path.to_string_lossy(); + + output.push_str(&format!( + "| {} | {} | {} | {} | {} | [{}](file://{}) |\n", + count, + paper.title.chars().take(50).collect::(), + authors_str, + paper.year, + paper.citation_count, + paper.bibcode, + md_path_str + )); + } + + if count == 0 { + output.push_str("\n最近无匹配的新文献。\n"); + } else { + output.push_str(&format!("\n共找到 {} 篇匹配文献。\n", count)); + } + + ToolOutput::success(output, json!({ "count": count })) + } + Err(e) => ToolOutput::error(format!("本地数据库搜索失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/research/mod.rs b/src/agent/tools/astro/research/mod.rs index e4c2e2f..969d68b 100644 --- a/src/agent/tools/astro/research/mod.rs +++ b/src/agent/tools/astro/research/mod.rs @@ -1,20 +1,42 @@ // src/agent/tools/astro/research/mod.rs // 研究级工具：科研人员消费本地数据进行分析 +pub mod abundance_analysis; +pub mod batch_process; +pub mod kinematics_tool; pub mod library; +pub mod literature_monitor; pub mod metadata; pub mod note; pub mod observation; +pub mod observation_plan; pub mod paper; +pub mod period_search; +pub mod radial_velocity; pub mod rag; +pub mod sed_fit; +pub mod spectral_lines; +pub mod spectrum_params; pub mod target; +pub mod variable_star; pub mod vizier; +pub use abundance_analysis::AnalyzeAbundancePatternTool; +pub use batch_process::BatchProcessTool; +pub use kinematics_tool::AnalyzeKinematicsTool; pub use library::{GetCitationNetworkTool, SearchLocalLibraryTool}; +pub use literature_monitor::LiteratureMonitorTool; pub use metadata::GetPaperMetadataTool; pub use note::SaveNoteTool; pub use observation::FindObservationTool; +pub use observation_plan::PlanObservationTool; pub use paper::{GetPaperContentTool, GetPaperOutlineTool}; +pub use period_search::FindPeriodTool; +pub use radial_velocity::MeasureRadialVelocityTool; pub use rag::RagSearchTool; +pub use sed_fit::FitSedTool; +pub use spectral_lines::MeasureSpectralLinesTool; +pub use spectrum_params::GetSpectrumParametersTool; pub use target::QueryTargetTool; +pub use variable_star::ClassifyVariableStarTool; pub use vizier::CatalogOperationTool; diff --git a/src/agent/tools/astro/research/observation.rs b/src/agent/tools/astro/research/observation.rs index 6ad8cdd..73b030f 100644 --- a/src/agent/tools/astro/research/observation.rs +++ b/src/agent/tools/astro/research/observation.rs @@ -135,7 +135,7 @@ impl AgentTool for FindObservationTool { let state = &ctx.app_state; let source = match args.get("source").and_then(|v| v.as_str()) { - Some(s) => match Source::from_str(s) { + Some(s) => match Source::parse(s) { Ok(src) => src, Err(e) => return ToolOutput::error(e), }, @@ -143,7 +143,7 @@ impl AgentTool for FindObservationTool { }; let product_type = match args.get("product").and_then(|v| v.as_str()) { None => ProductType::Spectrum, - Some(s) => match ProductType::from_str(s) { + Some(s) => match ProductType::parse(s) { Ok(p) => p, Err(e) => return ToolOutput::error(e), }, diff --git a/src/agent/tools/astro/research/observation_plan.rs b/src/agent/tools/astro/research/observation_plan.rs new file mode 100644 index 0000000..f4fddf9 --- /dev/null +++ b/src/agent/tools/astro/research/observation_plan.rs @@ -0,0 +1,115 @@ +// src/agent/tools/astro/research/observation_plan.rs +// +// PlanObservationTool —— 观测提案辅助 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct PlanObservationTool; + +#[async_trait] +impl AgentTool for PlanObservationTool { + fn name(&self) -> &str { + "plan_observation" + } + + fn display_name(&self) -> &str { + "观测提案辅助" + } + + fn description(&self) -> &str { + "观测提案辅助：评估目标的可观测性。计算目标在给定台站的高度角、大气质量、\ + 月相影响，并估算所需曝光时间。支持 Keck、Lick、LAMOST、Gemini North 等台站。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "ra": { "type": "number", "description": "目标赤经 (度, J2000)" }, + "dec": { "type": "number", "description": "目标赤纬 (度, J2000)" }, + "observatory": { + "type": "string", + "enum": ["keck", "lick", "lamost", "gemini_north"], + "description": "观测台站，默认 keck", + "default": "keck" + }, + "target_v_mag": { + "type": "number", + "description": "目标 V 波段星等（用于曝光时间估算），可选" + }, + "target_snr": { + "type": "number", + "description": "目标信噪比（用于曝光时间估算），可选", + "default": 100.0 + } + }, + "required": ["ra", "dec"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let ra = match args.get("ra").and_then(|r| r.as_f64()) { + Some(r) => r, + None => return ToolOutput::error("缺少必需参数 'ra'"), + }; + let dec = match args.get("dec").and_then(|d| d.as_f64()) { + Some(d) => d, + None => return ToolOutput::error("缺少必需参数 'dec'"), + }; + + let obs_name = args + .get("observatory") + .and_then(|o| o.as_str()) + .unwrap_or("keck"); + + let obs = match obs_name { + "keck" => crate::services::analysis::observability::Observatory::keck(), + "lick" => crate::services::analysis::observability::Observatory::lick_3m(), + "lamost" => crate::services::analysis::observability::Observatory::lamost(), + "gemini_north" => crate::services::analysis::observability::Observatory::gemini_north(), + _ => return ToolOutput::error(format!("不支持的台站: {}", obs_name)), + }; + + let target_v_mag = args.get("target_v_mag").and_then(|v| v.as_f64()); + let target_snr = args + .get("target_snr") + .and_then(|s| s.as_f64()) + .unwrap_or(100.0); + + info!( + "[PlanObservation] RA={:.4}, Dec={:.4}, 台站={}", + ra, dec, obs_name + ); + + let result = crate::services::analysis::observability::assess_observability( + ra, + dec, + &obs, + target_v_mag, + Some(target_snr), + ); + + let report = + crate::services::analysis::observability::format_observability(&result, &obs.name); + let json_str = serde_json::to_string_pretty(&result).unwrap_or_default(); + let content = format!("{}\n\n```json\n{}\n```", report, json_str); + + ToolOutput::success(content, json!(result)) + } +} diff --git a/src/agent/tools/astro/research/period_search.rs b/src/agent/tools/astro/research/period_search.rs new file mode 100644 index 0000000..628f969 --- /dev/null +++ b/src/agent/tools/astro/research/period_search.rs @@ -0,0 +1,148 @@ +// src/agent/tools/astro/research/period_search.rs +// +// FindPeriodTool —— 光变周期搜索（Lomb-Scargle + BLS） + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct FindPeriodTool; + +#[async_trait] +impl AgentTool for FindPeriodTool { + fn name(&self) -> &str { + "find_period" + } + + fn display_name(&self) -> &str { + "周期搜索" + } + + fn description(&self) -> &str { + "搜索光变曲线的周期。支持 Lomb-Scargle（通用变星）和 BLS（凌星检测）两种方法。\ + 自动报告最优周期、FAP（误报概率）和次优周期。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "time": { + "type": "array", + "items": { "type": "number" }, + "description": "时间数组 (天)" + }, + "flux": { + "type": "array", + "items": { "type": "number" }, + "description": "通量/星等数组" + }, + "method": { + "type": "string", + "enum": ["lomb-scargle", "bls"], + "description": "搜索方法，默认 lomb-scargle", + "default": "lomb-scargle" + }, + "min_period": { + "type": "number", + "description": "最小周期 (天)，可选" + }, + "max_period": { + "type": "number", + "description": "最大周期 (天)，可选" + }, + "n_freq": { + "type": "integer", + "description": "频率点数，默认 5000", + "default": 5000 + } + }, + "required": ["time", "flux"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let time: Vec = match args.get("time").and_then(|t| t.as_array()) { + Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(), + None => return ToolOutput::error("缺少必需参数 'time'"), + }; + let flux: Vec = match args.get("flux").and_then(|f| f.as_array()) { + Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(), + None => return ToolOutput::error("缺少必需参数 'flux'"), + }; + + if time.len() != flux.len() { + return ToolOutput::error("时间和通量数组长度不一致"); + } + if time.len() < 4 { + return ToolOutput::error("数据点太少（需要 ≥4）"); + } + + let method = args + .get("method") + .and_then(|m| m.as_str()) + .unwrap_or("lomb-scargle"); + + let min_period = args.get("min_period").and_then(|p| p.as_f64()); + let max_period = args.get("max_period").and_then(|p| p.as_f64()); + let n_freq = args.get("n_freq").and_then(|n| n.as_u64()).unwrap_or(5000) as usize; + + info!( + "[FindPeriod] 方法={}, 数据点={}, 搜索范围=[{}, {}] 天", + method, + time.len(), + min_period + .map(|p| format!("{:.4}", p)) + .unwrap_or_else(|| "auto".to_string()), + max_period + .map(|p| format!("{:.4}", p)) + .unwrap_or_else(|| "auto".to_string()) + ); + + match method { + "lomb-scargle" => { + match crate::services::timeseries::periodicity::lomb_scargle( + &time, &flux, min_period, max_period, n_freq, + ) { + Ok(result) => { + let report = + crate::services::timeseries::periodicity::format_period_result(&result); + let json_str = serde_json::to_string_pretty(&result).unwrap_or_default(); + let content = format!("{}\n\n```json\n{}\n```", report, json_str); + ToolOutput::success(content, json!(result)) + } + Err(e) => ToolOutput::error(format!("周期搜索失败: {}", e)), + } + } + "bls" => { + match crate::services::timeseries::periodicity::bls( + &time, &flux, min_period, max_period, n_freq, 0.01, 0.1, + ) { + Ok(result) => { + let content = format!( + "BLS 周期搜索结果:\n\n- 最优周期: {:.6} 天\n- 凌星深度: {:.4}\n- 凌星持续时间: {:.4} 天\n- SDE: {:.2}\n", + result.best_period, result.depth, result.duration, result.sde + ); + ToolOutput::success(content, json!(result)) + } + Err(e) => ToolOutput::error(format!("BLS 搜索失败: {}", e)), + } + } + _ => ToolOutput::error(format!("不支持的方法: {}", method)), + } + } +} diff --git a/src/agent/tools/astro/research/radial_velocity.rs b/src/agent/tools/astro/research/radial_velocity.rs new file mode 100644 index 0000000..e2a1986 --- /dev/null +++ b/src/agent/tools/astro/research/radial_velocity.rs @@ -0,0 +1,151 @@ +// src/agent/tools/astro/research/radial_velocity.rs +// +// MeasureRadialVelocityTool —— 交叉相关法测径向速度 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct MeasureRadialVelocityTool; + +#[async_trait] +impl AgentTool for MeasureRadialVelocityTool { + fn name(&self) -> &str { + "measure_radial_velocity" + } + + fn display_name(&self) -> &str { + "径向速度测量" + } + + fn description(&self) -> &str { + "通过交叉相关法测量光谱的径向速度（RV）。使用内置恒星类型模板或用户自定义模板，\ + 通过 Doppler shift 匹配测量视向速度。支持 OBAFGKM 恒星类型。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "wavelength": { + "type": "array", + "items": { "type": "number" }, + "description": "观测光谱波长数组 (Å)" + }, + "flux": { + "type": "array", + "items": { "type": "number" }, + "description": "观测光谱通量数组" + }, + "template_type": { + "type": "string", + "enum": ["O", "B", "A", "F", "G", "K", "M", "Mwarf", "Giant"], + "description": "模板恒星类型（不指定则使用 G 型）", + "default": "G" + }, + "rv_range": { + "type": "array", + "items": { "type": "number" }, + "description": "RV 搜索范围 [min, max] (km/s)，默认 [-500, 500]", + "default": [-500.0, 500.0] + }, + "rv_step": { + "type": "number", + "description": "RV 搜索步长 (km/s)，默认 1.0", + "default": 1.0 + } + }, + "required": ["wavelength", "flux"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let wavelength: Vec = match args.get("wavelength").and_then(|w| w.as_array()) { + Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(), + None => return ToolOutput::error("缺少必需参数 'wavelength'"), + }; + let flux: Vec = match args.get("flux").and_then(|w| w.as_array()) { + Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(), + None => return ToolOutput::error("缺少必需参数 'flux'"), + }; + + if wavelength.len() != flux.len() { + return ToolOutput::error("波长和通量数组长度不一致"); + } + if wavelength.len() < 10 { + return ToolOutput::error("数据点太少（需要 ≥10）"); + } + + let template_str = args + .get("template_type") + .and_then(|t| t.as_str()) + .unwrap_or("G"); + + let template = match template_str { + "O" => crate::services::spectrum::cross_correlate::TemplateType::O, + "B" => crate::services::spectrum::cross_correlate::TemplateType::B, + "A" => crate::services::spectrum::cross_correlate::TemplateType::A, + "F" => crate::services::spectrum::cross_correlate::TemplateType::F, + "G" => crate::services::spectrum::cross_correlate::TemplateType::G, + "K" => crate::services::spectrum::cross_correlate::TemplateType::K, + "M" => crate::services::spectrum::cross_correlate::TemplateType::M, + "Mwarf" => crate::services::spectrum::cross_correlate::TemplateType::Mwarf, + "Giant" => crate::services::spectrum::cross_correlate::TemplateType::Giant, + _ => return ToolOutput::error(format!("不支持的模板类型: {}", template_str)), + }; + + let rv_range = args + .get("rv_range") + .and_then(|r| r.as_array()) + .and_then(|a| { + if a.len() == 2 { + Some((a[0].as_f64()?, a[1].as_f64()?)) + } else { + None + } + }) + .unwrap_or((-500.0, 500.0)); + + let rv_step = args.get("rv_step").and_then(|s| s.as_f64()).unwrap_or(1.0); + + info!( + "[MeasureRV] 模板={:?}, RV范围=[{:.0}, {:.0}] km/s, 步长={:.1} km/s", + template, rv_range.0, rv_range.1, rv_step + ); + + match crate::services::spectrum::cross_correlate::measure_rv( + &wavelength, + &flux, + template, + rv_range, + rv_step, + ) { + Ok(result) => { + let json_str = serde_json::to_string_pretty(&result).unwrap_or_default(); + let content = format!( + "径向速度测量结果:\n\n- RV = {:.2} ± {:.2} km/s\n- CCF 峰值 = {:.4}\n- 模板: {}\n- 波长范围: {:.0}-{:.0} Å\n- SNR: {}\n\n{}", + result.rv, result.rv_error, result.ccf_peak, result.template_name, + result.wavelength_range.0, result.wavelength_range.1, + result.snr.map(|s| format!("{:.1}", s)).unwrap_or_else(|| "N/A".to_string()), + json_str + ); + ToolOutput::success(content, json!(result)) + } + Err(e) => ToolOutput::error(format!("径向速度测量失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/research/rag.rs b/src/agent/tools/astro/research/rag.rs index 573ca6a..c1db972 100644 --- a/src/agent/tools/astro/research/rag.rs +++ b/src/agent/tools/astro/research/rag.rs @@ -15,6 +15,11 @@ impl AgentTool for RagSearchTool { "rag_search" } + /// 检索结果来自外部论文语料，可能包含对抗性内容 + fn untrusted_output(&self) -> bool { + true + } + fn display_name(&self) -> &str { "文献库RAG检索" } diff --git a/src/agent/tools/astro/research/sed_fit.rs b/src/agent/tools/astro/research/sed_fit.rs new file mode 100644 index 0000000..1bd0689 --- /dev/null +++ b/src/agent/tools/astro/research/sed_fit.rs @@ -0,0 +1,125 @@ +// src/agent/tools/astro/research/sed_fit.rs +// +// FitSedTool —— 多波段 SED 拟合 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct FitSedTool; + +#[async_trait] +impl AgentTool for FitSedTool { + fn name(&self) -> &str { + "fit_sed" + } + + fn display_name(&self) -> &str { + "SED 拟合" + } + + fn description(&self) -> &str { + "多波段光谱能量分布（SED）拟合。输入多个波段的测光数据（波长、通量、误差），\ + 拟合黑体辐射模型，输出有效温度 T_eff、半径 R、消光 A_V、光度 L 及不确定度。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "data_points": { + "type": "array", + "items": { + "type": "object", + "properties": { + "wavelength": { "type": "number", "description": "中心波长 (Å)" }, + "flux": { "type": "number", "description": "流量" }, + "flux_error": { "type": "number", "description": "流量误差" }, + "band": { "type": "string", "description": "波段名" } + }, + "required": ["wavelength", "flux", "flux_error", "band"] + }, + "description": "测光数据点列表" + }, + "distance_pc": { + "type": "number", + "description": "距离 (pc)，可选，用于计算光度" + } + }, + "required": ["data_points"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let data_points: Vec = + match args.get("data_points").and_then(|d| d.as_array()) { + Some(arr) => { + let mut points = Vec::new(); + for item in arr { + let wavelength = match item.get("wavelength").and_then(|w| w.as_f64()) { + Some(w) => w, + None => continue, + }; + let flux = match item.get("flux").and_then(|f| f.as_f64()) { + Some(f) => f, + None => continue, + }; + let flux_error = item + .get("flux_error") + .and_then(|e| e.as_f64()) + .unwrap_or(flux * 0.1); + let band = item + .get("band") + .and_then(|b| b.as_str()) + .unwrap_or("unknown") + .to_string(); + points.push(crate::services::analysis::sed::SedDataPoint { + wavelength, + flux, + flux_error, + band, + source: "user".to_string(), + }); + } + points + } + None => return ToolOutput::error("缺少必需参数 'data_points'"), + }; + + if data_points.len() < 3 { + return ToolOutput::error("至少需要 3 个测光点进行 SED 拟合"); + } + + let distance_pc = args.get("distance_pc").and_then(|d| d.as_f64()); + + info!( + "[FitSed] 数据点={}, 距离={:?} pc", + data_points.len(), + distance_pc + ); + + match crate::services::analysis::sed::fit_sed(&data_points, distance_pc) { + Ok(result) => { + let report = crate::services::analysis::sed::format_sed_result(&result); + let json_str = serde_json::to_string_pretty(&result).unwrap_or_default(); + let content = format!("{}\n\n```json\n{}\n```", report, json_str); + ToolOutput::success(content, json!(result)) + } + Err(e) => ToolOutput::error(format!("SED 拟合失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/research/spectral_lines.rs b/src/agent/tools/astro/research/spectral_lines.rs new file mode 100644 index 0000000..6737b24 --- /dev/null +++ b/src/agent/tools/astro/research/spectral_lines.rs @@ -0,0 +1,168 @@ +// src/agent/tools/astro/research/spectral_lines.rs +// +// MeasureSpectralLinesTool —— 谱线识别与等值宽度测量 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct MeasureSpectralLinesTool; + +#[async_trait] +impl AgentTool for MeasureSpectralLinesTool { + fn name(&self) -> &str { + "measure_spectral_lines" + } + + fn display_name(&self) -> &str { + "谱线测量" + } + + fn description(&self) -> &str { + "识别光谱中的常见谱线（Balmer 系列、Ca II H&K、Na D、Mg b、Fe 等），\ + 自动计算等值宽度（EW）和半高全宽（FWHM）。输入光谱的波长和通量数组，\ + 返回每条谱线的测量结果。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "wavelength": { + "type": "array", + "items": { "type": "number" }, + "description": "波长数组 (Å)" + }, + "flux": { + "type": "array", + "items": { "type": "number" }, + "description": "通量数组" + }, + "lines": { + "type": "array", + "items": { "type": "string" }, + "description": "要测量的谱线名称列表（可选）。不指定则测量所有内置谱线。可选值: Hα, Hβ, Hγ, Hδ, Ca II K, Ca II H, Na I D2, Na I D1, Mg I b1 等" + }, + "window_half_width": { + "type": "number", + "description": "搜索窗口半宽 (Å)，默认 20", + "default": 20.0 + } + }, + "required": ["wavelength", "flux"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let wavelength: Vec = match args.get("wavelength").and_then(|w| w.as_array()) { + Some(arr) => arr + .iter() + .filter_map(|v| v.as_f64().map(|f| f as f32)) + .collect(), + None => return ToolOutput::error("缺少必需参数 'wavelength'"), + }; + let flux: Vec = match args.get("flux").and_then(|w| w.as_array()) { + Some(arr) => arr + .iter() + .filter_map(|v| v.as_f64().map(|f| f as f32)) + .collect(), + None => return ToolOutput::error("缺少必需参数 'flux'"), + }; + + let window_half_width = args + .get("window_half_width") + .and_then(|w| w.as_f64()) + .unwrap_or(20.0); + + if wavelength.len() != flux.len() { + return ToolOutput::error("波长和通量数组长度不一致"); + } + if wavelength.len() < 3 { + return ToolOutput::error("数据点太少（需要 ≥3）"); + } + + let table = crate::services::spectrum::lines::builtin_line_table(); + + // 筛选要测量的谱线 + let target_lines: Vec = + if let Some(names) = args.get("lines").and_then(|l| l.as_array()) { + names + .iter() + .filter_map(|n| n.as_str()) + .filter_map(|name| crate::services::spectrum::lines::find_line(name, &table)) + .cloned() + .collect() + } else { + // 测量所有在光谱范围内的谱线 + let w_min = wavelength.iter().cloned().fold(f32::INFINITY, f32::min) as f64; + let w_max = wavelength.iter().cloned().fold(f32::NEG_INFINITY, f32::max) as f64; + crate::services::spectrum::lines::find_lines_in_range(w_min, w_max, &table) + .into_iter() + .cloned() + .collect() + }; + + if target_lines.is_empty() { + return ToolOutput::error("在当前光谱范围内未找到匹配的谱线"); + } + + info!( + "[MeasureSpectralLines] 测量 {} 条谱线，窗口 ±{:.1} Å", + target_lines.len(), + window_half_width + ); + + let measurements = crate::services::spectrum::lines::measure_lines( + &wavelength, + &flux, + &target_lines, + window_half_width, + ); + + let successful: Vec<_> = measurements.into_iter().filter_map(|r| r.ok()).collect(); + + if successful.is_empty() { + return ToolOutput::error("所有谱线测量均失败"); + } + + let table_str = crate::services::spectrum::lines::format_measurements_table(&successful); + + let content = format!( + "谱线测量完成（{} / {} 条成功）:\n\n{}", + successful.len(), + target_lines.len(), + table_str + ); + + let metadata = json!({ + "measurements": successful.iter().map(|m| { + json!({ + "line": m.line.name, + "wavelength": m.center_wavelength, + "ew": m.ew, + "ew_error": m.ew_error, + "fwhm": m.fwhm, + "continuum_flux": m.continuum_flux + }) + }).collect::>(), + "total_lines": target_lines.len(), + "successful_measurements": successful.len() + }); + + ToolOutput::success(content, metadata) + } +} diff --git a/src/agent/tools/astro/research/spectrum_params.rs b/src/agent/tools/astro/research/spectrum_params.rs new file mode 100644 index 0000000..62699ed --- /dev/null +++ b/src/agent/tools/astro/research/spectrum_params.rs @@ -0,0 +1,130 @@ +// src/agent/tools/astro/research/spectrum_params.rs +// +// GetSpectrumParametersTool —— 从 FITS 文件中提取恒星大气参数 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct GetSpectrumParametersTool; + +#[async_trait] +impl AgentTool for GetSpectrumParametersTool { + fn name(&self) -> &str { + "get_spectrum_parameters" + } + + fn display_name(&self) -> &str { + "光谱参数提取" + } + + fn description(&self) -> &str { + "从已下载的 FITS 光谱文件中自动提取恒星大气参数（Teff, logg, [Fe/H], 径向速度等）。\ + 支持 LAMOST、SDSS/BOSS、APOGEE、DESI 数据。结果结构化返回，可直接用于后续分析。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "source": { + "type": "string", + "enum": ["lamost", "sdss", "desi", "gaia"], + "description": "数据源" + }, + "subtype": { + "type": "string", + "description": "产品子类型（SDSS: spec/apstar/aspcap）" + }, + "source_id": { + "type": "string", + "description": "源标识符" + }, + "file_path": { + "type": "string", + "description": "FITS 文件路径（相对于 library/ 目录）" + } + }, + "required": ["source", "source_id", "file_path"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let source_str = match args.get("source").and_then(|s| s.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'source'"), + }; + let source_id = match args.get("source_id").and_then(|s| s.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'source_id'"), + }; + let file_path = match args.get("file_path").and_then(|s| s.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'file_path'"), + }; + let subtype = args.get("subtype").and_then(|s| s.as_str()); + + let source = match crate::services::observation::Source::parse(source_str) { + Ok(s) => s, + Err(e) => return ToolOutput::error(e), + }; + + let product = crate::services::observation::types::ProductSpec { + product: crate::services::observation::types::ProductType::Spectrum, + subtype: subtype.map(|s| s.to_string()), + }; + + info!( + "[GetSpectrumParameters] source={}, id={}, file={}", + source_str, source_id, file_path + ); + + // 构造完整路径 + let full_path = if std::path::Path::new(file_path).is_absolute() { + file_path.to_string() + } else { + let library_dir = ctx + .app_state + .config + .storage + .library_dir + .to_str() + .unwrap_or("library"); + format!("{}/{}", library_dir, file_path) + }; + + match crate::services::spectrum::parameters::extract_parameters( + &ctx.app_state, + source, + &product, + source_id, + &full_path, + ) + .await + { + Ok(params) => { + let json_str = serde_json::to_string_pretty(¶ms).unwrap_or_default(); + let content = format!( + "光谱参数提取成功 ({}, {}):\n\n{}", + source_str, source_id, json_str + ); + ToolOutput::success(content, json!(params)) + } + Err(e) => ToolOutput::error(format!("参数提取失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/research/variable_star.rs b/src/agent/tools/astro/research/variable_star.rs new file mode 100644 index 0000000..4dd2bb0 --- /dev/null +++ b/src/agent/tools/astro/research/variable_star.rs @@ -0,0 +1,103 @@ +// src/agent/tools/astro/research/variable_star.rs +// +// ClassifyVariableStarTool —— 变星分类与特征提取 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +pub struct ClassifyVariableStarTool; + +#[async_trait] +impl AgentTool for ClassifyVariableStarTool { + fn name(&self) -> &str { + "classify_variable_star" + } + + fn display_name(&self) -> &str { + "变星分类" + } + + fn description(&self) -> &str { + "自动分类变星类型（RR Lyrae、Cepheid、食双星、脉动变星、AGN 等），\ + 提取振幅、偏度、峰度等统计特征。可选传入已知周期以提高分类准确度。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "time": { + "type": "array", + "items": { "type": "number" }, + "description": "时间数组 (天)" + }, + "flux": { + "type": "array", + "items": { "type": "number" }, + "description": "通量/星等数组" + }, + "period": { + "type": "number", + "description": "已知周期 (天)，可选。传入可提高分类准确度" + } + }, + "required": ["time", "flux"] + }) + } + + fn group(&self) -> &str { + "as:research" + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + fn is_readonly(&self) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let time: Vec = match args.get("time").and_then(|t| t.as_array()) { + Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(), + None => return ToolOutput::error("缺少必需参数 'time'"), + }; + let flux: Vec = match args.get("flux").and_then(|f| f.as_array()) { + Some(arr) => arr.iter().filter_map(|v| v.as_f64()).collect(), + None => return ToolOutput::error("缺少必需参数 'flux'"), + }; + + let period = args.get("period").and_then(|p| p.as_f64()); + + if time.len() != flux.len() { + return ToolOutput::error("时间和通量数组长度不一致"); + } + if time.len() < 5 { + return ToolOutput::error("数据点太少（需要 ≥5）"); + } + + info!( + "[ClassifyVariable] 数据点={}, 周期={}", + time.len(), + period + .map(|p| format!("{:.4} 天", p)) + .unwrap_or_else(|| "未指定".to_string()) + ); + + match crate::services::timeseries::classification::classify_variable_star( + &time, &flux, period, + ) { + Ok(result) => { + let report = + crate::services::timeseries::classification::format_classification(&result); + let json_str = serde_json::to_string_pretty(&result).unwrap_or_default(); + let content = format!("{}\n\n```json\n{}\n```", report, json_str); + ToolOutput::success(content, json!(result)) + } + Err(e) => ToolOutput::error(format!("分类失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/system/search.rs b/src/agent/tools/astro/system/search.rs index 4ae4524..e94bef8 100644 --- a/src/agent/tools/astro/system/search.rs +++ b/src/agent/tools/astro/system/search.rs @@ -15,6 +15,11 @@ impl AgentTool for SearchPapersTool { "search_papers" } + /// 搜索结果摘要来自 arXiv/ADS 等外部 API + fn untrusted_output(&self) -> bool { + true + } + fn display_name(&self) -> &str { "文献搜索" } @@ -123,9 +128,11 @@ impl AgentTool for SearchPapersTool { .collect::>() .join("\n\n"); - ToolOutput::success( + // content 是 value 的模型侧投影；canonical value 供重放/审计消费 + ToolOutput::success_with_value( content, - json!({ "count": results.len(), "papers": display_results }), + json!({ "count": results.len() }), + json!(display_results), ) } Err(e) => { diff --git a/src/agent/tools/compress.rs b/src/agent/tools/compress.rs index 759c835..4cf1c6e 100644 --- a/src/agent/tools/compress.rs +++ b/src/agent/tools/compress.rs @@ -3,7 +3,7 @@ use async_trait::async_trait; use serde_json::json; -use super::{AgentTool, ToolContext, ToolOutput}; +use super::{AgentTool, ToolContext, ToolOutput, ToolSignals}; /// 手动上下文压缩工具：LLM 可主动调用以压缩对话历史 pub struct CompressTool; @@ -22,6 +22,14 @@ impl AgentTool for CompressTool { true } + fn loop_signals(&self) -> ToolSignals { + ToolSignals { + resets_todo_nag: false, + persists_todos: false, + requests_manual_compact: true, + } + } + fn description(&self) -> &str { "手动压缩对话上下文。当你发现对话历史过长、token 消耗过大时，主动调用此工具进行压缩以释放空间。\ 压缩后历史对话将被摘要替代，但关键信息不会丢失。" diff --git a/src/agent/tools/filesystem/bash.rs b/src/agent/tools/filesystem/bash.rs index 9a5bafe..7f4e0b9 100644 --- a/src/agent/tools/filesystem/bash.rs +++ b/src/agent/tools/filesystem/bash.rs @@ -59,6 +59,19 @@ impl AgentTool for RunBashTool { InterruptBehavior::Block } + fn causes_file_changes(&self) -> bool { + true + } + + fn hardline_check( + &self, + args: &serde_json::Value, + ) -> Option { + args.get("command") + .and_then(|v| v.as_str()) + .map(crate::agent::runtime::hardline::check_command) + } + fn check_permissions( &self, args: &serde_json::Value, diff --git a/src/agent/tools/filesystem/edit.rs b/src/agent/tools/filesystem/edit.rs index 3b57197..e2ea2a2 100644 --- a/src/agent/tools/filesystem/edit.rs +++ b/src/agent/tools/filesystem/edit.rs @@ -203,6 +203,20 @@ impl AgentTool for FileEditTool { InterruptBehavior::Block } + fn causes_file_changes(&self) -> bool { + true + } + + fn hardline_check( + &self, + args: &serde_json::Value, + ) -> Option { + args.get("file_path") + .or_else(|| args.get("path")) + .and_then(|v| v.as_str()) + .map(crate::agent::runtime::hardline::check_dangerous_path) + } + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { let path_str = match args.get("file_path").and_then(|v| v.as_str()) { Some(s) => s, diff --git a/src/agent/tools/filesystem/write.rs b/src/agent/tools/filesystem/write.rs index beffee3..dafece2 100644 --- a/src/agent/tools/filesystem/write.rs +++ b/src/agent/tools/filesystem/write.rs @@ -49,6 +49,20 @@ impl AgentTool for FileWriteTool { InterruptBehavior::Block } + fn causes_file_changes(&self) -> bool { + true + } + + fn hardline_check( + &self, + args: &serde_json::Value, + ) -> Option { + args.get("file_path") + .or_else(|| args.get("path")) + .and_then(|v| v.as_str()) + .map(crate::agent::runtime::hardline::check_dangerous_path) + } + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { let file_path_str = match args.get("file_path").and_then(|v| v.as_str()) { Some(s) => s, diff --git a/src/agent/tools/mod.rs b/src/agent/tools/mod.rs index c951352..666aaa2 100644 --- a/src/agent/tools/mod.rs +++ b/src/agent/tools/mod.rs @@ -36,13 +36,24 @@ mod todo; pub use ask_user::AskUserTool; pub use astro::analyze_image::AnalyzeImageTool; +pub use astro::research::abundance_analysis::AnalyzeAbundancePatternTool; +pub use astro::research::batch_process::BatchProcessTool; +pub use astro::research::kinematics_tool::AnalyzeKinematicsTool; pub use astro::research::library::{GetCitationNetworkTool, SearchLocalLibraryTool}; +pub use astro::research::literature_monitor::LiteratureMonitorTool; pub use astro::research::metadata::GetPaperMetadataTool; pub use astro::research::note::SaveNoteTool; pub use astro::research::observation::FindObservationTool; +pub use astro::research::observation_plan::PlanObservationTool; pub use astro::research::paper::{GetPaperContentTool, GetPaperOutlineTool}; +pub use astro::research::period_search::FindPeriodTool; +pub use astro::research::radial_velocity::MeasureRadialVelocityTool; pub use astro::research::rag::RagSearchTool; +pub use astro::research::sed_fit::FitSedTool; +pub use astro::research::spectral_lines::MeasureSpectralLinesTool; +pub use astro::research::spectrum_params::GetSpectrumParametersTool; pub use astro::research::target::QueryTargetTool; +pub use astro::research::variable_star::ClassifyVariableStarTool; pub use astro::research::vizier::CatalogOperationTool; pub use astro::system::process::ProcessPaperTool; pub use astro::system::search::SearchPapersTool; @@ -123,6 +134,15 @@ impl ToolContext { self } + /// 设置可选的 SSE 通道（静默上下文为 None） + pub fn with_sse_tx_opt( + mut self, + tx: Option>, + ) -> Self { + self.sse_tx = tx; + self + } + /// 设置会话 ID pub fn with_session_id(mut self, id: String) -> Self { self.session_id = id; @@ -170,14 +190,24 @@ impl ToolContext { } /// 工具执行结果 +/// +/// 输出契约（参考 deepseek-harness 的 canonical output 设计）： +/// - `value`：canonical JSON 值——机器可用的结构化结果（重放/审计/剪枝/ +/// 前端渲染都消费这一份）； +/// - `content`：`value` 的模型侧纯投影（render）——给 LLM 阅读的文本。 +/// +/// 新工具应优先通过 `success_with_value` 同时提供两者；只填 content 的 +/// 旧工具在消费侧退化为纯文本处理。 #[derive(Debug, Clone)] pub struct ToolOutput { - /// 给大模型阅读的截断文本 + /// 给大模型阅读的截断文本（value 的 render 投影） pub content: String, /// 是否为错误 pub is_error: bool, /// 结构化元数据（给前端 Timeline 直接渲染） pub metadata: serde_json::Value, + /// canonical JSON 值（机器可用的结构化结果；None = 仅有文本投影） + pub value: Option, /// 跳过持久化到磁盘（用于 read_file 等已从磁盘读取内容的工具，避免级联持久化） pub skip_persist: bool, } @@ -189,6 +219,22 @@ impl ToolOutput { content: content.into(), is_error: false, metadata, + value: None, + skip_persist: false, + } + } + + /// 创建成功结果（带 canonical value，新工具的推荐构造器） + pub fn success_with_value( + content: impl Into, + metadata: serde_json::Value, + value: serde_json::Value, + ) -> Self { + ToolOutput { + content: content.into(), + is_error: false, + metadata, + value: Some(value), skip_persist: false, } } @@ -199,16 +245,24 @@ impl ToolOutput { content: content.into(), is_error: false, metadata, + value: None, skip_persist: true, } } + /// 为已有结果附加 canonical value（builder 风格） + pub fn with_value(mut self, value: serde_json::Value) -> Self { + self.value = Some(value); + self + } + /// 创建错误结果 pub fn error(msg: impl Into) -> Self { ToolOutput { content: msg.into(), is_error: true, metadata: json!({}), + value: None, skip_persist: false, } } @@ -223,6 +277,20 @@ pub enum InterruptBehavior { Block, } +/// 循环行为信号 — 工具向 ReAct 循环声明的副作用语义。 +/// +/// 用于替代主循环中按工具名字符串特判的散落逻辑： +/// 循环通过 `ToolRegistry::signals_for()` 查询，工具自身声明行为。 +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct ToolSignals { + /// 调用后重置 "todo nag" 计数（todo_write 自身） + pub resets_todo_nag: bool, + /// 调用后需要把 todos 参数持久化到任务板（agent_tasks 表） + pub persists_todos: bool, + /// 调用后请求在下一轮 LLM 调用前执行手动压缩（compress_context） + pub requests_manual_compact: bool, +} + /// 权限规则来源 #[derive(Debug, Clone, PartialEq)] pub enum PermissionRuleSource { @@ -352,6 +420,36 @@ pub trait AgentTool: Send + Sync { fn group(&self) -> &str { "general" } + + // ── 集中元数据声明：消除按工具名散落的平行名单 ── + + /// 工具输出是否来自外部不可信源（web/检索/MCP 等返回的内容可能包含 + /// prompt 注入）。为 true 时执行器会把结果包裹在 + /// `` 中再交给 LLM。 + fn untrusted_output(&self) -> bool { + false + } + + /// 该工具是否可能产生文件变更。为 true 时执行器会在调用前触发 + /// CheckpointManager 快照。 + fn causes_file_changes(&self) -> bool { + false + } + + /// Hardline 预检查（任何模式下都不可绕过的参数级拒绝）。 + /// 返回 Some(blocked) 表示该调用的参数命中 hardline 规则； + /// 默认 None 表示无 hardline 规则。 + fn hardline_check( + &self, + _args: &serde_json::Value, + ) -> Option { + None + } + + /// 循环行为信号（见 `ToolSignals`）。 + fn loop_signals(&self) -> ToolSignals { + ToolSignals::default() + } } /// 工具注册表，管理所有可用工具。 @@ -395,6 +493,18 @@ fn add_base_tools(registry: &mut ToolRegistry, skill_registry: Arc ToolSignals { + self.tools + .get(tool_name) + .map(|t| t.loop_signals()) + .unwrap_or_default() + } + + /// 查询工具输出是否不可信。未知工具回退到名字启发式（mcp__*/web_* 前缀）。 + pub fn untrusted_output(&self, tool_name: &str) -> bool { + self.tools + .get(tool_name) + .is_some_and(|t| t.untrusted_output()) + || crate::agent::runtime::untrusted::is_high_risk(tool_name) + } + /// 构建按 group 分组的紧凑工具目录字符串。 /// 格式: /// [system] @@ -740,33 +866,82 @@ mod tests { "./skills", ))))); let defs = registry.definitions(); - assert_eq!(defs.len(), 24); - assert!(defs.iter().any(|d| d.function.name == "read_file")); - assert!(defs.iter().any(|d| d.function.name == "grep_files")); - assert!(defs.iter().any(|d| d.function.name == "glob_files")); - assert!(defs.iter().any(|d| d.function.name == "run_bash")); - assert!(defs.iter().any(|d| d.function.name == "file_write")); - assert!(defs.iter().any(|d| d.function.name == "file_edit")); - assert!(defs.iter().any(|d| d.function.name == "search_papers")); - assert!(defs.iter().any(|d| d.function.name == "process_paper")); - assert!(defs - .iter() - .any(|d| d.function.name == "search_local_library")); - assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata")); - assert!(defs.iter().any(|d| d.function.name == "get_paper_outline")); - assert!(defs.iter().any(|d| d.function.name == "get_paper_content")); - assert!(defs - .iter() - .any(|d| d.function.name == "get_citation_network")); - assert!(defs.iter().any(|d| d.function.name == "rag_search")); - assert!(defs.iter().any(|d| d.function.name == "query_target")); - assert!(defs.iter().any(|d| d.function.name == "find_observation")); - assert!(defs.iter().any(|d| d.function.name == "catalog_operation")); - assert!(defs.iter().any(|d| d.function.name == "save_note")); - assert!(defs.iter().any(|d| d.function.name == "todo_write")); - assert!(defs.iter().any(|d| d.function.name == "compress_context")); - assert!(defs.iter().any(|d| d.function.name == "load_skill")); - assert!(defs.iter().any(|d| d.function.name == "subagent")); + // 关键工具必须在场（inventory 式断言——新增工具不需要改这里） + for expected in [ + "read_file", + "grep_files", + "glob_files", + "run_bash", + "file_write", + "file_edit", + "search_papers", + "process_paper", + "search_local_library", + "get_paper_metadata", + "get_paper_outline", + "get_paper_content", + "get_citation_network", + "rag_search", + "query_target", + "find_observation", + "catalog_operation", + "save_note", + "todo_write", + "compress_context", + "load_skill", + "subagent", + ] { + assert!( + defs.iter().any(|d| d.function.name == expected), + "工具 {} 未注册", + expected + ); + } + } + + /// 工具清单一致性：名字唯一、元数据完整、声明元数据（signals 等） + /// 与名字的一致性。新增工具若破坏这些不变量会在此失败而不是在运行时漂移。 + #[test] + fn test_tool_registry_inventory_invariants() { + let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from( + "./skills", + ))))); + let names = registry.tool_names(); + let mut unique = std::collections::HashSet::new(); + for name in &names { + assert!(unique.insert(name.clone()), "工具名重复注册: {}", name); + } + for tool in registry.list() { + let name = tool.name(); + assert!(!name.is_empty(), "存在空工具名"); + assert!(!tool.description().is_empty(), "工具 {} 缺少描述", name); + assert!( + tool.parameters().get("type").is_some(), + "工具 {} 的参数 schema 缺少 type 字段", + name + ); + } + // signals 与工具身份一致：todo_write 声明 todo 信号，compress_context 声明压缩信号 + let todo_signals = registry.signals_for("todo_write"); + assert!(todo_signals.resets_todo_nag && todo_signals.persists_todos); + assert!( + registry + .signals_for("compress_context") + .requests_manual_compact + ); + // 只读工具声明不可信输出的例外：search_papers/rag_search 输出来自外部源 + assert!(registry.untrusted_output("search_papers")); + assert!(registry.untrusted_output("rag_search")); + assert!(!registry.untrusted_output("read_file")); + // 文件变更工具必须声明 causes_file_changes + for name in ["file_write", "file_edit", "run_bash"] { + let tool = registry.get(name).unwrap(); + assert!( + tool.causes_file_changes(), + "{} 应声明 causes_file_changes", + name + ); + } } #[test] diff --git a/src/agent/tools/persist.rs b/src/agent/tools/persist.rs index 15d848e..f2db818 100644 --- a/src/agent/tools/persist.rs +++ b/src/agent/tools/persist.rs @@ -22,6 +22,18 @@ pub async fn maybe_persist_tool_result( tool_call_id: &str, max_chars: usize, tool_results_dir: &Path, +) -> (String, Option) { + maybe_persist_tool_result_for(content, tool_call_id, max_chars, tool_results_dir, "tool").await +} + +/// 带 tool_name 的持久化变体：stub 中携带工具名与检索指引， +/// 让模型能判断溢出内容的性质并有明确的取回路径。 +pub async fn maybe_persist_tool_result_for( + content: &str, + tool_call_id: &str, + max_chars: usize, + tool_results_dir: &Path, + tool_name: &str, ) -> (String, Option) { if content.len() <= max_chars { return (content.to_string(), None); @@ -124,14 +136,18 @@ pub async fn maybe_persist_tool_result( let stub = format!( "\n\ + tool: {}\n\ path: {}\n\ size: {} chars\n\ preview: |\n {}\n\n\ - 完整输出已持久化到磁盘。使用 read_file 工具以 path 参数读取完整内容。\n\ + 完整输出已溢出持久化到磁盘（超出 {} 字符上限）。\n\ + 取回方式：调用 read_file 工具，path 填上述路径；文件较长时可用 max_lines 参数分段读取。\n\ ", + tool_name, file_path.display(), content.len(), preview.replace('\n', "\n "), + max_chars, ); (stub, Some(file_path)) diff --git a/src/agent/tools/search_history.rs b/src/agent/tools/search_history.rs index 777bc83..1a91264 100644 --- a/src/agent/tools/search_history.rs +++ b/src/agent/tools/search_history.rs @@ -65,81 +65,43 @@ impl AgentTool for SearchHistoryTool { Some(q) if !q.is_empty() => q, _ => return ToolOutput::error("缺少 query 参数"), }; - let query = crate::services::search::sanitize_fts5_query(raw_query); let scope = args.get("scope").and_then(|v| v.as_str()).unwrap_or("all"); - let db = &ctx.app_state.db; - let mut results = Vec::new(); - - if scope == "all" || scope == "sessions" { - match sqlx::query_as::<_, (String, String, String)>( - "SELECT s.session_id, s.title, \ - snippet(agent_sessions_fts, 1, '', '', '...', 40) \ - FROM agent_sessions_fts fts \ - JOIN agent_sessions s ON s.session_id = fts.session_id \ - WHERE agent_sessions_fts MATCH $1 ORDER BY rank LIMIT 10", - ) - .bind(&query) - .fetch_all(db) - .await - { - Ok(rows) => { - for (sid, title, snippet) in rows { - results.push(json!({ - "type": "session", - "session_id": sid, - "title": title, - "snippet": snippet, - })); - } - } - Err(e) => { - tracing::warn!("[SearchHistory] 搜索会话失败: {}", e); + match crate::services::search::search_agent_history( + &ctx.app_state.db, + raw_query, + scope, + 10, + None, + ) + .await + { + Ok(search_results) => { + if search_results.is_empty() { + ToolOutput::success( + format!("未找到与 \"{}\" 相关的历史记录。", raw_query), + json!({"query": raw_query, "results": [], "count": 0}), + ) + } else { + let results_json: Vec = search_results + .iter() + .map(|r| { + json!({ + "type": r.result_type, + "session_id": r.session_id, + "title": r.title, + "snippet": r.snippet, + }) + }) + .collect(); + let count = results_json.len(); + ToolOutput::success( + format!("找到 {} 条与 \"{}\" 相关的历史记录", count, raw_query), + json!({"query": raw_query, "results": results_json, "count": count}), + ) } } - } - - if scope == "all" || scope == "messages" { - match sqlx::query_as::<_, (String, String, String, String)>( - "SELECT fts.session_id, s.title, \ - snippet(agent_messages_fts, 2, '', '', '...', 80), fts.role \ - FROM agent_messages_fts fts \ - JOIN agent_sessions s ON s.session_id = fts.session_id \ - JOIN agent_messages m ON m.rowid = fts.rowid \ - WHERE agent_messages_fts MATCH $1 AND m.active = 1 \ - ORDER BY rank LIMIT 10", - ) - .bind(&query) - .fetch_all(db) - .await - { - Ok(rows) => { - for (sid, title, snippet, role) in rows { - results.push(json!({ - "type": format!("message/{}", role), - "session_id": sid, - "title": title, - "snippet": snippet, - })); - } - } - Err(e) => { - tracing::warn!("[SearchHistory] 搜索消息失败: {}", e); - } - } - } - - if results.is_empty() { - ToolOutput::success( - format!("未找到与 \"{}\" 相关的历史记录。", query), - json!({"query": query, "results": [], "count": 0}), - ) - } else { - let count = results.len(); - ToolOutput::success( - format!("找到 {} 条与 \"{}\" 相关的历史记录", count, query), - json!({"query": query, "results": results, "count": count}), - ) + Err(e) => ToolOutput::error(format!("检索历史失败: {}", e)), } } } diff --git a/src/agent/tools/subagent.rs b/src/agent/tools/subagent.rs index 33c4287..450e3be 100644 --- a/src/agent/tools/subagent.rs +++ b/src/agent/tools/subagent.rs @@ -75,6 +75,15 @@ impl AgentTool for SubAgentTool { "type": "integer", "description": "子代理最大推理步数，默认5，最大10", "default": 5 + }, + "agent_name": { + "type": "string", + "description": "子代理的稳定名称（如 literature-scout）。命名后其历史会持久保留，后续可用 followup=true 继续与它对话而不丢失上下文。" + }, + "followup": { + "type": "boolean", + "description": "设为 true 时续话同名的已有子代理：回放其全部历史并追加新指令，而不是从头开始。要求同时提供 agent_name。", + "default": false } }, "required": ["research_prompt"] @@ -98,10 +107,30 @@ impl AgentTool for SubAgentTool { .unwrap_or(5) .min(10) as usize; + let agent_name = args + .get("agent_name") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let followup = args + .get("followup") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + // 参数校验（fail-loud）：followup 必须搭配 agent_name + let agent_name = match (followup, agent_name) { + (true, Some(name)) => Some(name), + (true, None) => { + return ToolOutput::error("followup=true 需要同时提供 agent_name 参数"); + } + (false, name) => name, + }; + info!( - "[SubAgent] 启动子代理: prompt_len={}, max_steps={}", + "[SubAgent] 启动子代理: prompt_len={}, max_steps={}, name={:?}, followup={}", research_prompt.len(), - max_steps + max_steps, + agent_name, + followup ); // 创建子代理的 ToolRegistry（与父代理共享 skill_registry） @@ -157,15 +186,38 @@ impl AgentTool for SubAgentTool { ) .with_parent_session(ctx.session_id.clone()) .with_thinking(ctx.enable_thinking); - let result = runner - .run(&system_prompt, &research_prompt, max_steps) - .await; + // followup：续话已有命名子代理（回放其持久历史）；否则全新运行 + let result = if followup { + runner + .run_continuation( + agent_name.as_deref().unwrap_or_default(), + &system_prompt, + &research_prompt, + max_steps, + ) + .await + } else { + runner + .run_named( + &system_prompt, + &research_prompt, + max_steps, + agent_name.as_deref(), + ) + .await + }; if result.is_error { ToolOutput::error(format!("子代理执行失败: {}", result.content)) } else { - // 包装子代理结果，标注来源 - let final_content = format!("[子代理研究结果]\n\n{}", result.content); + // 包装子代理结果，标注来源；命名子代理提示续话方式 + let mut final_content = format!("[子代理研究结果]\n\n{}", result.content); + if let Some(name) = result.metadata.get("agent_name").and_then(|v| v.as_str()) { + final_content.push_str(&format!( + "\n\n（子代理名称: {}。如需在保留其上下文的情况下继续追问，再次调用 subagent 工具并传 agent_name=\"{}\" 与 followup=true。）", + name, name + )); + } info!( "[SubAgent] 子代理结果: content_len={}, preview={}...", final_content.len(), diff --git a/src/agent/tools/todo.rs b/src/agent/tools/todo.rs index da9d279..70ce6d9 100644 --- a/src/agent/tools/todo.rs +++ b/src/agent/tools/todo.rs @@ -9,7 +9,7 @@ use sqlx::SqlitePool; use std::str::FromStr; use tracing::{info, warn}; -use super::{AgentTool, ToolContext, ToolOutput}; +use super::{AgentTool, ToolContext, ToolOutput, ToolSignals}; /// 任务规划工具：让 LLM 在开始复杂研究前先制定计划，执行中更新进度。 /// 任务状态持久化到 SQLite，支持 DAG 依赖。 @@ -29,6 +29,14 @@ impl AgentTool for TodoWriteTool { true } + fn loop_signals(&self) -> ToolSignals { + ToolSignals { + resets_todo_nag: true, + persists_todos: true, + requests_manual_compact: false, + } + } + fn description(&self) -> &str { "任务规划工具。在开始复杂研究前列出待办事项，执行中标记进度（每项状态：pending/in_progress/completed）。\ 一次只能有一个 in_progress 任务。支持任务依赖（blockedBy：依赖的其他任务ID列表）。\ @@ -131,7 +139,12 @@ impl AgentTool for TodoWriteTool { .push_str("\n💡 提示：还有待处理任务，请选择一个设为 in_progress 并开始执行。"); } - ToolOutput::success(formatted, json!({ "task_count": todos.len() })) + // canonical value = 全量任务状态（重放/任务板审计的机器可用投影） + ToolOutput::success_with_value( + formatted, + json!({ "task_count": todos.len() }), + json!({ "todos": todos }), + ) } } @@ -310,6 +323,13 @@ mod tests { irsa: crate::clients::irsa::IrsaClient::new("https://irsa.ipac.caltech.edu", 60) .unwrap(), mast: crate::clients::mast::MastClient::new("https://mast.stsci.edu", 90).unwrap(), + jwst: crate::clients::mast::jwst::JwstClient::new("https://mast.stsci.edu", 60) + .unwrap(), + heasarc: crate::clients::heasarc::HeasarcClient::new( + "https://heasarc.gsfc.nasa.gov/FTP/heasarc/obs/dat/tap", + 60, + ) + .unwrap(), observation_registry: std::sync::Arc::new( crate::services::observation::ObservationRegistry::default(), ), @@ -326,6 +346,7 @@ mod tests { std::collections::HashMap::new(), )), }, + agent_runtimes: Arc::new(crate::agent::runtime::SessionRuntimeRegistry::new()), }); ToolContext { diff --git a/src/api/agent.rs b/src/api/agent.rs index fd02a83..6ba8a52 100644 --- a/src/api/agent.rs +++ b/src/api/agent.rs @@ -16,7 +16,7 @@ use tracing::{error, info}; use super::error::{ApiResult, AppError}; use super::AppState; -use crate::agent::runtime::{AgentRuntime, AgentStreamEvent}; +use crate::agent::runtime::AgentStreamEvent; // ── POST /api/chat/agent ── // SSE 流式智能体对话接口 @@ -191,25 +191,68 @@ pub async fn chat_agent( None => (None, None), }; - let mut runtime = AgentRuntime::new(Arc::clone(&state)).with_mode(&req.mode); + // ── 会话解析与运行时复用 ── + // 预分配会话 ID：新会话的首个请求也能写入取消标记并命中运行时缓存 + // （历史上新会话首请求超时只能 abort 任务，无法写入取消标记） + let session_key = req + .session_id + .clone() + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + + // 模式回放：已存在的会话从 DB 读取创建时的模式。 + // 修复：恢复会话时传不同 mode 会静默改变行为——现在会话模式一次创建后保持稳定。 + let mode_id = crate::agent::runtime::session::load_session_mode(&state.db, &session_key) + .await + .unwrap_or_else(|| req.mode.clone()); + + let runtime = state + .agent_runtimes + .get_or_create(Arc::clone(&state), &session_key, &mode_id); // 只有 mode 未强制固定 thinking 时，用户才可以覆盖 if runtime.mode_fixed_thinking().is_none() { if let Some(thinking) = req.thinking { - runtime = runtime.with_thinking(thinking); + runtime.set_thinking(thinking); } } + + // 同会话 turn 串行化：并发请求 fail-loud 拒绝，防止 turn_index/消息顺序被破坏 + let turn_lock = state + .agent_runtimes + .turn_lock(&session_key) + .unwrap_or_else(|| std::sync::Arc::new(tokio::sync::Mutex::new(()))); + // 快速检测：已占用直接 409（guard 立即释放，任务内部会重新 try_lock 兜底竞态） + if turn_lock.try_lock().is_err() { + return Err(AppError::conflict( + "该会话正在执行中，请等待当前回合完成后再发送新消息", + )); + } + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); - let session_id = req.session_id.clone(); + let question_owned = question.clone(); + let image_context_owned = image_context.clone(); + let image_path_owned = image_path_for_db.clone(); + let cancel_session_key = session_key.clone(); // 在后台 tokio 任务中执行 Agent 循环 let agent_handle = tokio::spawn(async move { + // 持有 turn 串行锁直到回合结束（内部重新 try_lock 兜底竞态窗口） + let _turn_guard = match turn_lock.try_lock() { + Ok(guard) => guard, + Err(_) => { + let _ = tx.send(AgentStreamEvent::Error { + message: "该会话正在执行中，请稍后重试。".to_string(), + }); + let _ = tx.send(AgentStreamEvent::Done); + return; + } + }; match runtime .run_turn_with_image_context( - session_id, - &question, - image_context, - image_path_for_db, + Some(session_key.clone()), + &question_owned, + image_context_owned, + image_path_owned, tx.clone(), ) .await @@ -236,9 +279,7 @@ pub async fn chat_agent( let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); if remaining.is_zero() { // 超时：通知 CancellationHook 停止 Agent，并中止后台任务 - if let Some(sid) = &req.session_id { - cancelled_runs.insert(sid.clone(), ()); - } + cancelled_runs.insert(cancel_session_key.clone(), ()); agent_handle.abort(); let timeout_event = AgentStreamEvent::Error { message: "Agent 执行超时（10 分钟），请重试。".to_string(), @@ -259,9 +300,7 @@ pub async fn chat_agent( Ok(None) => break, // channel closed Err(_) => { // 超时：通知 CancellationHook 停止 Agent，并中止后台任务 - if let Some(sid) = &req.session_id { - cancelled_runs.insert(sid.clone(), ()); - } + cancelled_runs.insert(cancel_session_key.clone(), ()); agent_handle.abort(); let timeout_event = AgentStreamEvent::Error { message: "Agent 执行超时（10 分钟），请重试。".to_string(), @@ -323,6 +362,11 @@ pub async fn delete_session( .await .map_err(|e| AppError::internal(format!("删除会话失败: {}", e)))?; + if success { + // 同步移除会话级运行时缓存（后台队列/压缩日志等状态随之释放） + state.agent_runtimes.remove(&session_id); + } + if !success { return Err(AppError::not_found(format!( "会话 {} 不存在或已删除", diff --git a/src/api/catalog.rs b/src/api/catalog.rs index bfa2a17..74fa412 100644 --- a/src/api/catalog.rs +++ b/src/api/catalog.rs @@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize}; use super::error::{ApiResult, AppError}; use super::AppState; use crate::clients::cds::vizier::VizierQueryResult; +use crate::services::analysis::hr_diagram::{build_hr_diagram, HrDiagramParams, HrDiagramResponse}; // ── 请求参数 ── @@ -131,3 +132,15 @@ pub async fn cone_search( .await?; Ok(Json(result)) } + +/// GET /api/analysis/hr-diagram —— 赫罗图构建 +pub async fn hr_diagram( + State(state): State>, + Query(params): Query, +) -> ApiResult> { + let response = build_hr_diagram(&state.sources.gaia, ¶ms) + .await + .map_err(|e| AppError::internal(format!("构建赫罗图失败: {}", e)))?; + + Ok(Json(response)) +} diff --git a/src/api/mod.rs b/src/api/mod.rs index f670ce4..305de33 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -75,6 +75,8 @@ pub struct DataSourceState { pub desi: DesiClient, pub irsa: crate::clients::irsa::IrsaClient, pub mast: crate::clients::mast::MastClient, + pub jwst: crate::clients::mast::jwst::JwstClient, + pub heasarc: crate::clients::heasarc::HeasarcClient, pub observation_registry: Arc, } @@ -109,6 +111,9 @@ pub struct AppState { pub llm: LlmState, pub sources: DataSourceState, pub session: SessionState, + /// 会话级 Agent 运行时注册表：同一会话复用 runtime（后台队列/压缩日志/ + /// 文件缓存等跨 turn 状态不再随请求销毁） + pub agent_runtimes: Arc, } // 统一标准化的文献格式，用于向前端传输 diff --git a/src/api/observation.rs b/src/api/observation.rs index bfa40c1..88cae9e 100644 --- a/src/api/observation.rs +++ b/src/api/observation.rs @@ -39,7 +39,7 @@ fn build_product_spec( subtype: &Option, ) -> Result { Ok(crate::services::observation::ProductSpec { - product: crate::services::observation::ProductType::from_str(product) + product: crate::services::observation::ProductType::parse(product) .map_err(AppError::bad_request)?, subtype: subtype.clone(), }) @@ -68,7 +68,7 @@ pub async fn observation_search( State(state): State>, Query(params): Query, ) -> ApiResult>> { - let source = crate::services::observation::Source::from_str(¶ms.source) + let source = crate::services::observation::Source::parse(¶ms.source) .map_err(AppError::bad_request)?; let product = build_product_spec(¶ms.product, ¶ms.subtype)?; @@ -128,8 +128,8 @@ pub async fn observation_download( ) -> ApiResult> { use crate::services::observation::{download_observation, ObservationRequest}; - let source = crate::services::observation::Source::from_str(&req.source) - .map_err(AppError::bad_request)?; + let source = + crate::services::observation::Source::parse(&req.source).map_err(AppError::bad_request)?; let product = build_product_spec(&req.product, &req.subtype)?; let request = match req.mode { @@ -255,7 +255,7 @@ pub async fn observation_preview( State(state): State>, Query(params): Query, ) -> ApiResult> { - let source = crate::services::observation::Source::from_str(¶ms.source) + let source = crate::services::observation::Source::parse(¶ms.source) .map_err(AppError::bad_request)?; let product = build_product_spec(¶ms.product, ¶ms.subtype)?; diff --git a/src/api/sync.rs b/src/api/sync.rs index 1c23931..f32f457 100644 --- a/src/api/sync.rs +++ b/src/api/sync.rs @@ -215,3 +215,31 @@ pub async fn get_asset_batch_status( let status = state.batch_status.lock().await; Json(status.clone()) } + +#[derive(Debug, Deserialize)] +pub struct ToggleScheduleRequest { + pub is_scheduled: bool, +} + +pub async fn toggle_query_schedule( + State(state): State>, + axum::extract::Path(id): axum::extract::Path, + Json(req): Json, +) -> ApiResult { + let success = + crate::services::batch::MetaSync::set_query_schedule(&state.db, id, req.is_scheduled) + .await + .map_err(|e| { + error!("更新同步检索配置定时状态失败: {}", e); + AppError::internal(format!("更新同步检索配置定时状态失败: {}", e)) + })?; + + if !success { + return Err(AppError::not_found(format!( + "未找到 ID 为 {} 的检索配置", + id + ))); + } + + Ok(StatusCode::OK) +} diff --git a/src/bin/cli.rs b/src/bin/cli.rs index 308a4a3..ba4eed6 100644 --- a/src/bin/cli.rs +++ b/src/bin/cli.rs @@ -282,6 +282,13 @@ async fn main() -> anyhow::Result<()> { config.cds.tess_timeout_secs, ) .context("构建 MAST 客户端失败")?, + jwst: astroresearch::clients::mast::jwst::JwstClient::new("https://mast.stsci.edu", 60) + .context("构建 JWST 客户端失败")?, + heasarc: astroresearch::clients::heasarc::HeasarcClient::new( + "https://heasarc.gsfc.nasa.gov/FTP/heasarc/obs/dat/tap", + 60, + ) + .context("构建 HEASARC 客户端失败")?, observation_registry: Arc::new( astroresearch::services::observation::ObservationRegistry::default(), ), @@ -294,6 +301,7 @@ async fn main() -> anyhow::Result<()> { pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), }, + agent_runtimes: Arc::new(astroresearch::agent::runtime::SessionRuntimeRegistry::new()), }); let cli = Cli::parse(); diff --git a/src/clients/gaia/mod.rs b/src/clients/gaia/mod.rs index 4f79876..9f39980 100644 --- a/src/clients/gaia/mod.rs +++ b/src/clients/gaia/mod.rs @@ -293,6 +293,11 @@ impl GaiaClient { serde_json::from_str::(&body).context("解析 Gaia TAP JSON 失败") } + /// 公开的 TAP 查询接口（JSON 格式） + pub async fn tap_query(&self, adql: &str) -> anyhow::Result { + self.tap_sync(adql, "json").await + } + /// 执行 TAP sync 查询（带重试） async fn tap_sync(&self, adql: &str, format: &str) -> anyhow::Result { let params = [ diff --git a/src/clients/heasarc/mod.rs b/src/clients/heasarc/mod.rs new file mode 100644 index 0000000..5fcc917 --- /dev/null +++ b/src/clients/heasarc/mod.rs @@ -0,0 +1,190 @@ +// src/clients/heasarc/mod.rs +// +// HEASARC（High Energy Astrophysics Science Archive Research Center）客户端 +// +// 通过 TAP 服务查询 XMM-Newton 和 Chandra X 射线数据 +// 端点：POST https://heasarc.gsfc.nasa.gov/FTP/heasarc/.../tap/sync +// 查询语言：ADQL + +use anyhow::{anyhow, Context}; +use serde::{Deserialize, Serialize}; +use std::time::Duration; +use tracing::info; + +/// X 射线观测行 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct XrayObservationRow { + /// 观测标识 + pub obs_id: String, + /// 目标名称 + pub target_name: String, + /// 赤经（度） + pub ra: f64, + /// 赤纬（度） + pub dec: f64, + /// 望远镜（XMM-Newton / Chandra） + pub telescope: String, + /// 仪器（EPIC-pn/ACIS-S/...） + pub instrument: String, + /// 曝光时间（秒） + pub exposure: Option, + /// 观测日期 + pub obs_date: Option, + /// 数据存档 URL + pub data_url: Option, +} + +/// HEASARC TAP 客户端 +#[derive(Clone)] +pub struct HeasarcClient { + client: reqwest::Client, + tap_url: String, +} + +impl HeasarcClient { + pub fn new(tap_url: &str, timeout_secs: u64) -> anyhow::Result { + Ok(HeasarcClient { + client: reqwest::Client::builder() + .redirect(crate::utils::ssrf::safe_redirect_policy()) + .timeout(Duration::from_secs(timeout_secs)) + .connect_timeout(Duration::from_secs(15)) + .build() + .context("Failed to create HEASARC HTTP client")?, + tap_url: tap_url.trim_end_matches('/').to_string(), + }) + } + + /// TAP ADQL 查询 + pub async fn tap_query(&self, adql: &str) -> anyhow::Result { + info!("[HEASARC] TAP 查询: {}", &adql[..adql.len().min(100)]); + let url = format!("{}/sync", self.tap_url); + let resp = self + .client + .post(&url) + .header("User-Agent", "AstroResearch/0.1") + .form(&[ + ("REQUEST", "doQuery"), + ("LANG", "ADQL"), + ("FORMAT", "json"), + ("MAXREC", "1000"), + ("QUERY", adql), + ]) + .send() + .await + .context("HEASARC TAP 请求失败")?; + + let status = resp.status(); + if !status.is_success() { + let err_body = resp.text().await.unwrap_or_default(); + return Err(anyhow!( + "HEASARC TAP 错误 {}: {}", + status, + err_body.chars().take(500).collect::() + )); + } + resp.text().await.context("读取 HEASARC 响应失败") + } + + /// XMM-Newton 锥形检索 + pub async fn xmm_cone_search( + &self, + ra: f64, + dec: f64, + radius_deg: f64, + ) -> anyhow::Result> { + let adql = format!( + "SELECT obsid, target_name, ra, dec, instrument, exposure, obs_date \ + FROM heasarc_xmmobscat \ + WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) \ + ORDER BY exposure DESC", + ra, dec, radius_deg + ); + let body = self.tap_query(&adql).await?; + parse_tap_json(&body, "XMM-Newton") + } + + /// Chandra 锥形检索 + pub async fn chandra_cone_search( + &self, + ra: f64, + dec: f64, + radius_deg: f64, + ) -> anyhow::Result> { + let adql = format!( + "SELECT obsid, target_name, ra, dec, instrument, exposure, obs_date \ + FROM heasarc_chandraobs \ + WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) \ + ORDER BY exposure DESC", + ra, dec, radius_deg + ); + let body = self.tap_query(&adql).await?; + parse_tap_json(&body, "Chandra") + } +} + +fn parse_tap_json(body: &str, telescope: &str) -> anyhow::Result> { + let v: serde_json::Value = serde_json::from_str(body).context("解析 HEASARC JSON 失败")?; + let fields = v + .get("fields") + .and_then(|f| f.as_array()) + .ok_or_else(|| anyhow!("HEASARC 响应缺 fields"))?; + let data = v + .get("data") + .and_then(|d| d.as_array()) + .ok_or_else(|| anyhow!("HEASARC 响应缺 data"))?; + + let col = |name: &str| -> Option { + fields.iter().position(|f| { + f.get("name").and_then(|n| n.as_str()) == Some(name) + || f.get("id").and_then(|n| n.as_str()) == Some(name) + }) + }; + + let i_obsid = col("obsid"); + let i_target = col("target_name"); + let i_ra = col("ra"); + let i_dec = col("dec"); + let i_instrument = col("instrument"); + let i_exposure = col("exposure"); + let i_date = col("obs_date"); + + let mut rows = Vec::new(); + for row in data { + if let Some(arr) = row.as_array() { + let get_str = |i: Option| -> Option { + i.and_then(|idx| arr.get(idx)).and_then(|v| { + if v.is_i64() { + Some(v.as_i64().unwrap().to_string()) + } else if v.is_f64() { + Some(format!("{:.0}", v.as_f64().unwrap())) + } else { + v.as_str().map(|s| s.to_string()) + } + }) + }; + let get_f64 = |i: Option| -> Option { + i.and_then(|idx| arr.get(idx)).and_then(|v| { + v.as_f64() + .or_else(|| v.as_str().and_then(|s| s.parse().ok())) + }) + }; + + let obs_id = match get_str(i_obsid) { + Some(id) => id, + None => continue, + }; + rows.push(XrayObservationRow { + obs_id, + target_name: get_str(i_target).unwrap_or_default(), + ra: get_f64(i_ra).unwrap_or(0.0), + dec: get_f64(i_dec).unwrap_or(0.0), + telescope: telescope.to_string(), + instrument: get_str(i_instrument).unwrap_or_default(), + exposure: get_f64(i_exposure), + obs_date: get_str(i_date), + data_url: None, + }); + } + } + Ok(rows) +} diff --git a/src/clients/mast/jwst.rs b/src/clients/mast/jwst.rs new file mode 100644 index 0000000..611fd5a --- /dev/null +++ b/src/clients/mast/jwst.rs @@ -0,0 +1,290 @@ +// src/clients/mast/jwst.rs +// +// JWST（James Webb Space Telescope）MAST 数据查询 +// +// 通过 MAST Portal API 查询 JWST NIRSpec/MIRI 光谱数据 +// 端点：POST https://mast.stsci.edu/api/v0/invoke +// service: Mast.Jwst.Filtered.Position（坐标锥形检索） + +use anyhow::{anyhow, Context}; +use serde::{Deserialize, Serialize}; +use std::time::Duration; +use tracing::info; + +/// JWST 观测行（cone search 结果） +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JwstObservationRow { + /// JWST 观测标识 + pub obs_id: String, + /// 目标名称 + pub target_name: String, + /// 赤经（度） + pub ra: f64, + /// 赤纬（度） + pub dec: f64, + /// 仪器名称（NIRSpec/MIRI/NIRCam/NIRISS） + pub instrument: String, + /// 观测模式（IFU/MOS/Imaging/...） + pub observation_mode: String, + /// 滤光片/光栅设置 + pub filter_grating: Option, + /// 曝光时间（秒） + pub exposure_time: Option, + /// 数据产品 URL + pub data_url: Option, + /// 发布日期 + pub release_date: Option, +} + +/// JWST 客户端 +#[derive(Clone)] +pub struct JwstClient { + client: reqwest::Client, + base_url: String, +} + +impl JwstClient { + pub fn new(base_url: &str, timeout_secs: u64) -> anyhow::Result { + Ok(JwstClient { + client: reqwest::Client::builder() + .redirect(crate::utils::ssrf::safe_redirect_policy()) + .timeout(Duration::from_secs(timeout_secs)) + .connect_timeout(Duration::from_secs(10)) + .build() + .context("Failed to create JWST HTTP client")?, + base_url: base_url.trim_end_matches('/').to_string(), + }) + } + + /// JWST 观测锥形检索（通过 MAST Portal API） + pub async fn cone_search( + &self, + ra: f64, + dec: f64, + radius_deg: f64, + ) -> anyhow::Result> { + let url = format!("{}/api/v0/invoke", self.base_url); + info!( + "[JWST] ConeSearch ra={:.4} dec={:.4} radius={:.4}°", + ra, dec, radius_deg + ); + + let request_json = serde_json::json!({ + "service": "Mast.Jwst.Filtered.Position", + "params": { + "ra": ra, + "dec": dec, + "radius": radius_deg + }, + "format": "json", + "pagesize": 100, + "page": 1, + "removenullcolumns": true, + }) + .to_string(); + + let body = self + .send_post_with_retry(&url, "request", &request_json) + .await?; + parse_jwst_json(&body) + } + + /// 查询 JWST 数据产品列表 + pub async fn list_data_products(&self, obs_id: &str) -> anyhow::Result> { + let url = format!("{}/api/v0/invoke", self.base_url); + let request_json = serde_json::json!({ + "service": "Mast.Jwst.Product", + "params": { + "obs_id": obs_id + }, + "format": "json", + "pagesize": 200, + "page": 1, + }) + .to_string(); + + let body = self + .send_post_with_retry(&url, "request", &request_json) + .await?; + parse_jwst_products_json(&body, obs_id) + } + + async fn send_post_with_retry( + &self, + url: &str, + param_key: &str, + param_value: &str, + ) -> anyhow::Result { + const MAX_RETRIES: u32 = 3; + let mut last_err = None; + for attempt in 0..MAX_RETRIES { + let resp = self + .client + .post(url) + .header("User-Agent", "AstroResearch/0.1 (academic research tool)") + .form(&[(param_key, param_value)]) + .send() + .await; + let resp = match resp { + Ok(r) => r, + Err(e) => { + last_err = Some(e); + tokio::time::sleep(Duration::from_secs(2)).await; + continue; + } + }; + let status = resp.status(); + if (status.as_u16() == 429 || status.as_u16() == 503) && attempt < MAX_RETRIES - 1 { + let retry_after = resp + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + .unwrap_or(5); + tokio::time::sleep(Duration::from_secs(retry_after)).await; + continue; + } + if !status.is_success() { + let err_body = resp.text().await.unwrap_or_default(); + return Err(anyhow!( + "JWST API 返回错误码: {}: {}", + status, + err_body.chars().take(300).collect::() + )); + } + return resp.text().await.context("读取 JWST 响应失败"); + } + Err(last_err + .map(|e| anyhow!(e)) + .unwrap_or_else(|| anyhow!("JWST 重试耗尽"))) + } +} + +/// JWST 数据产品 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JwstDataProduct { + pub product_type: String, + pub file_name: String, + pub file_url: String, + pub file_size: Option, + pub instrument: String, +} + +fn parse_jwst_json(body: &str) -> anyhow::Result> { + let v: serde_json::Value = serde_json::from_str(body).context("解析 JWST JSON 失败")?; + let fields = v + .get("fields") + .and_then(|f| f.as_array()) + .ok_or_else(|| anyhow!("JWST 响应缺 fields"))?; + let data = v + .get("data") + .and_then(|d| d.as_array()) + .ok_or_else(|| anyhow!("JWST 响应缺 data"))?; + + let col = |name: &str| -> Option { + fields + .iter() + .position(|f| f.get("name").and_then(|n| n.as_str()) == Some(name)) + }; + + let i_obsid = col("obs_id").or_else(|| col("obsid")); + let i_target = col("target_name").or_else(|| col("targetid")); + let i_ra = col("ra"); + let i_dec = col("dec"); + let i_instrument = col("instrument_name").or_else(|| col("instrument")); + let i_mode = col("obs_mode").or_else(|| col("observation_mode")); + let i_filter = col("filter_grating"); + let i_exptime = col("exptime").or_else(|| col("exposure_time")); + let i_release = col("release_date"); + + let mut rows = Vec::new(); + for row in data { + if let Some(arr) = row.as_array() { + let get_str = |i: Option| -> Option { + i.and_then(|idx| arr.get(idx)).and_then(|v| { + if v.is_i64() { + Some(v.as_i64().unwrap().to_string()) + } else if v.is_f64() { + Some(format!("{:.0}", v.as_f64().unwrap())) + } else { + v.as_str().map(|s| s.to_string()) + } + }) + }; + let get_f64 = |i: Option| -> Option { + i.and_then(|idx| arr.get(idx)).and_then(|v| { + v.as_f64() + .or_else(|| v.as_str().and_then(|s| s.parse().ok())) + }) + }; + + let obs_id = match get_str(i_obsid) { + Some(id) => id, + None => continue, + }; + rows.push(JwstObservationRow { + obs_id, + target_name: get_str(i_target).unwrap_or_default(), + ra: get_f64(i_ra).unwrap_or(0.0), + dec: get_f64(i_dec).unwrap_or(0.0), + instrument: get_str(i_instrument).unwrap_or_default(), + observation_mode: get_str(i_mode).unwrap_or_default(), + filter_grating: get_str(i_filter), + exposure_time: get_f64(i_exptime), + data_url: None, + release_date: get_str(i_release), + }); + } + } + Ok(rows) +} + +fn parse_jwst_products_json(body: &str, _obs_id: &str) -> anyhow::Result> { + let v: serde_json::Value = serde_json::from_str(body)?; + let products = v + .pointer("/data/products") + .and_then(|p| p.as_array()) + .or_else(|| v.get("data").and_then(|d| d.as_array())) + .ok_or_else(|| anyhow!("JWST Product 响应缺 products"))?; + + let mut result = Vec::new(); + for prod in products { + if let Some(obj) = prod.as_object() { + let uri = obj + .get("uri") + .or_else(|| obj.get("dataURI")) + .and_then(|u| u.as_str()) + .unwrap_or(""); + let url = if uri.starts_with("http") { + uri.to_string() + } else if uri.starts_with("mast:") { + format!("https://mast.stsci.edu/api/v0.1/Download/file?uri={}", uri) + } else { + format!( + "https://mast.stsci.edu/api/v0.1/Download/file?uri=mast:JWST/product/{}", + uri + ) + }; + result.push(JwstDataProduct { + product_type: obj + .get("productType") + .and_then(|t| t.as_str()) + .unwrap_or("") + .to_string(), + file_name: obj + .get("filename") + .and_then(|f| f.as_str()) + .unwrap_or("") + .to_string(), + file_url: url, + file_size: obj.get("fileSize").and_then(|s| s.as_u64()), + instrument: obj + .get("instrument_name") + .and_then(|i| i.as_str()) + .unwrap_or("") + .to_string(), + }); + } + } + Ok(result) +} diff --git a/src/clients/mast/mod.rs b/src/clients/mast/mod.rs index e79ffb8..d792a8d 100644 --- a/src/clients/mast/mod.rs +++ b/src/clients/mast/mod.rs @@ -14,6 +14,8 @@ // // 无需认证（公开数据）。复用 safe_redirect_policy + send_with_retry 模式。 +pub mod jwst; + use anyhow::{anyhow, Context}; use serde::{Deserialize, Serialize}; use std::time::Duration; diff --git a/src/clients/mod.rs b/src/clients/mod.rs index 7ef2918..37cc630 100644 --- a/src/clients/mod.rs +++ b/src/clients/mod.rs @@ -3,6 +3,7 @@ pub mod arxiv; pub mod cds; pub mod desi; pub mod gaia; +pub mod heasarc; pub mod irsa; pub mod lamost; pub mod llm; diff --git a/src/main.rs b/src/main.rs index 5e7b7d4..f64e3b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -336,6 +336,13 @@ async fn main() -> anyhow::Result<()> { desi, irsa, mast, + jwst: astroresearch::clients::mast::jwst::JwstClient::new("https://mast.stsci.edu", 60) + .expect("Failed to create JWST client"), + heasarc: astroresearch::clients::heasarc::HeasarcClient::new( + "https://heasarc.gsfc.nasa.gov/FTP/heasarc/obs/dat/tap", + 60, + ) + .expect("Failed to create HEASARC client"), observation_registry: Arc::new( astroresearch::services::observation::ObservationRegistry::default(), ), @@ -348,6 +355,45 @@ async fn main() -> anyhow::Result<()> { pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), }, + // 会话级 Agent 运行时注册表（跨请求复用 runtime 状态） + agent_runtimes: Arc::new(astroresearch::agent::runtime::SessionRuntimeRegistry::new()), + }); + + // 自治研究循环（IdlePoller）：env 门控接线，默认关闭。 + // 开启后空闲时轮询任务板，自动认领并执行可开始的任务。 + if std::env::var("AGENT_AUTONOMOUS_ENABLED") + .map(|v| v == "true" || v == "1") + .unwrap_or(false) + { + let config = astroresearch::agent::autonomous::AutoResearchConfig { + enabled: true, + ..Default::default() + }; + let poller = astroresearch::agent::autonomous::IdlePoller::new(app_state.clone(), config); + tokio::spawn(async move { + poller.start().await; + }); + tracing::info!("[Autonomous] IdlePoller 已启用（AGENT_AUTONOMOUS_ENABLED）"); + } + + // 周期回收空闲会话 runtime（2 小时 TTL，见 SessionRuntimeRegistry） + { + let registry = app_state.agent_runtimes.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(30 * 60)).await; + let swept = registry.sweep_idle(); + if swept > 0 { + tracing::info!("[SessionRegistry] 本轮回收 {} 个空闲 runtime", swept); + } + } + }); + } + + // 启动定时增量文献同步调度器 + let scheduler_state = app_state.clone(); + tokio::spawn(async move { + astroresearch::services::batch::start_scheduled_sync_timer(scheduler_state).await; }); // 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管 @@ -464,6 +510,10 @@ async fn main() -> anyhow::Result<()> { "/sync/queries/:id", axum::routing::delete(sync::delete_sync_query), ) + .route( + "/sync/queries/:id/schedule", + post(sync::toggle_query_schedule), + ) .route("/target/query", get(targets::query_target)) .route("/target/associate", post(targets::associate_target)) .route("/target/extract", post(targets::extract_paper_targets)) @@ -496,6 +546,8 @@ async fn main() -> anyhow::Result<()> { "/observation/unified/resolve", post(observation::unified_resolve), ) + // 分析路由 + .route("/analysis/hr-diagram", get(catalog::hr_diagram)) // 智能体路由 .route("/chat/agent", post(agent::chat_agent)) .route("/chat/modes", get(agent::get_agent_modes)) diff --git a/src/services/analysis/chemistry.rs b/src/services/analysis/chemistry.rs new file mode 100644 index 0000000..b7f6b9f --- /dev/null +++ b/src/services/analysis/chemistry.rs @@ -0,0 +1,280 @@ +// src/services/analysis/chemistry.rs +// +// 化学丰度模式分析 —— [α/Fe] vs [Fe/H]，区分厚盘/薄盘/晕星 +// +// 输入：元素丰度表（来自 APOGEE ASPCAP 或其他光谱分析） +// 输出：[α/Fe] 计算、化学成分分类 + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +use super::super::spectrum::parameters::AbundanceEntry; + +/// α 元素列表 +const ALPHA_ELEMENTS: &[&str] = &["O", "Mg", "Si", "S", "Ca", "Ti"]; + +/// 化学丰度分析结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChemistryResult { + /// [Fe/H] + pub feh: f64, + /// [α/Fe]（α 元素加权平均） + pub alpha_fe: f64, + /// 各 α 元素的 [X/H] 值 + pub alpha_elements: Vec, + /// 化学分类 + pub classification: ChemistryClassification, + /// 分类置信度 + pub confidence: f64, +} + +/// 单个元素丰度 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ElementAbundance { + pub element: String, + pub xh: f64, // [X/H] + pub xfe: Option, // [X/Fe] + pub error: Option, +} + +/// 化学分类 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ChemistryClassification { + /// 最可能的星族 + pub population: String, + /// 描述 + pub description: String, +} + +/// 从丰度表分析化学模式 +pub fn analyze_chemistry(abundances: &[AbundanceEntry]) -> Result { + if abundances.is_empty() { + return Err(anyhow!("丰度数据为空")); + } + + // 提取 [Fe/H] + let feh = abundances + .iter() + .find(|a| a.element == "Fe" || a.element == "Fe_H") + .map(|a| a.value) + .ok_or_else(|| anyhow!("未找到 [Fe/H] 数据"))?; + + // 提取 α 元素 + let alpha_entries: Vec<&AbundanceEntry> = abundances + .iter() + .filter(|a| ALPHA_ELEMENTS.contains(&a.element.as_str())) + .collect(); + + let alpha_elements: Vec = alpha_entries + .iter() + .map(|a| { + let xfe = abundances + .iter() + .find(|b| b.element == "Fe") + .map(|fe| a.value - fe.value); + ElementAbundance { + element: a.element.clone(), + xh: a.value, + xfe, + error: a.error, + } + }) + .collect(); + + // 计算 [α/Fe] + let alpha_fe = if alpha_elements.is_empty() { + // 无 α 元素数据，尝试用 [M/H] - [Fe/H] 估计 + let mh = abundances + .iter() + .find(|a| a.element == "M" || a.element == "M_H") + .map(|a| a.value); + match mh { + Some(mh) => mh - feh, + None => 0.0, // 无法计算 + } + } else { + // 加权平均（以误差倒数为权重） + let weighted_sum: f64 = alpha_elements + .iter() + .map(|a| { + let weight = a + .error + .map(|e| if e > 0.0 { 1.0 / e } else { 1.0 }) + .unwrap_or(1.0); + a.xfe.unwrap_or(a.xh - feh) * weight + }) + .sum(); + let total_weight: f64 = alpha_elements + .iter() + .map(|a| { + a.error + .map(|e| if e > 0.0 { 1.0 / e } else { 1.0 }) + .unwrap_or(1.0) + }) + .sum(); + if total_weight > 0.0 { + weighted_sum / total_weight + } else { + 0.0 + } + }; + + // 化学分类（基于 [α/Fe] 和 [Fe/H] 的位置） + let (classification, confidence) = classify_by_chemistry(feh, alpha_fe); + + Ok(ChemistryResult { + feh, + alpha_fe, + alpha_elements, + classification, + confidence, + }) +} + +/// 基于 [α/Fe] 和 [Fe/H] 的化学分类 +/// +/// 经验判据（Bensby et al. 2003, 2014）： +/// - 薄盘：[α/Fe] < 0.25, [Fe/H] > -1.0 +/// - 厚盘：[α/Fe] > 0.25, [Fe/H] > -1.0 +/// - 晕：[Fe/H] < -1.0 +fn classify_by_chemistry(feh: f64, alpha_fe: f64) -> (ChemistryClassification, f64) { + // 使用软判据（高斯距离） + let feh_thin = -0.5; + let feh_thick = -0.7; + let alpha_thin = 0.15; + let alpha_thick = 0.35; + + // 薄盘中心 + let thin_dist = + ((feh - feh_thin).powi(2) / 0.5 + (alpha_fe - alpha_thin).powi(2) / 0.02).sqrt(); + // 厚盘中心 + let thick_dist = + ((feh - feh_thick).powi(2) / 0.3 + (alpha_fe - alpha_thick).powi(2) / 0.01).sqrt(); + // 晕中心 + let halo_dist = ((feh + 2.0).powi(2) / 1.0 + alpha_fe.powi(2) / 0.1).sqrt(); + + let thin_score = (-thin_dist).exp(); + let thick_score = (-thick_dist).exp(); + let halo_score = (-halo_dist).exp(); + + let total = thin_score + thick_score + halo_score; + let thin_prob = thin_score / total; + let thick_prob = thick_score / total; + let halo_prob = halo_score / total; + + let (population, description, confidence) = if thin_prob > thick_prob && thin_prob > halo_prob { + ( + "薄盘".to_string(), + format!( + "薄盘星（[α/Fe]={:.2}, [Fe/H]={:.2}），年龄较轻，富金属，低 α 增丰", + alpha_fe, feh + ), + thin_prob, + ) + } else if thick_prob > halo_prob { + ( + "厚盘".to_string(), + format!( + "厚盘星（[α/Fe]={:.2}, [Fe/H]={:.2}），年龄较老，中等金属丰度，高 α 增丰", + alpha_fe, feh + ), + thick_prob, + ) + } else { + ( + "晕".to_string(), + format!( + "晕星（[α/Fe]={:.2}, [Fe/H]={:.2}），极低金属丰度，古老星族", + alpha_fe, feh + ), + halo_prob, + ) + }; + + ( + ChemistryClassification { + population, + description, + }, + confidence, + ) +} + +/// 从丰度数据中按名称查找 [Fe/H] +pub fn find_feh(abundances: &[AbundanceEntry]) -> Option { + abundances + .iter() + .find(|a| a.element == "Fe" || a.element == "Fe_H" || a.element == "[Fe/H]") + .map(|a| a.value) +} + +/// 计算 [X/Fe] = [X/H] - [Fe/H] +pub fn compute_xfe(xh: f64, feh: f64) -> f64 { + xh - feh +} + +/// 格式化化学分析结果为 Markdown +pub fn format_chemistry(result: &ChemistryResult) -> String { + let mut out = String::from("## 化学丰度分析\n\n"); + out.push_str(&format!("- **[Fe/H]** = {:.3} dex\n", result.feh)); + out.push_str(&format!("- **[α/Fe]** = {:.3} dex\n\n", result.alpha_fe)); + + if !result.alpha_elements.is_empty() { + out.push_str("### α 元素丰度\n\n"); + out.push_str("| 元素 | [X/H] | [X/Fe] | 误差 |\n"); + out.push_str("|------|-------|--------|------|\n"); + for e in &result.alpha_elements { + let xfe_str = e + .xfe + .map(|v| format!("{:.3}", v)) + .unwrap_or_else(|| "—".to_string()); + let err_str = e + .error + .map(|v| format!("{:.3}", v)) + .unwrap_or_else(|| "—".to_string()); + out.push_str(&format!( + "| {} | {:.3} | {} | {} |\n", + e.element, e.xh, xfe_str, err_str + )); + } + } + + out.push_str(&format!( + "\n### 化学分类: **{}**\n\n", + result.classification.population + )); + out.push_str(&format!("{}\n", result.classification.description)); + out.push_str(&format!("- 置信度: {:.1}%\n", result.confidence * 100.0)); + + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_classify_thin_disk() { + let (class, conf) = classify_by_chemistry(-0.2, 0.1); + assert_eq!(class.population, "薄盘"); + assert!(conf > 0.5); + } + + #[test] + fn test_classify_thick_disk() { + let (class, conf) = classify_by_chemistry(-0.5, 0.35); + assert_eq!(class.population, "厚盘"); + assert!(conf > 0.3); + } + + #[test] + fn test_classify_halo() { + let (class, _conf) = classify_by_chemistry(-2.0, 0.0); + assert_eq!(class.population, "晕"); + } + + #[test] + fn test_compute_xfe() { + assert!((compute_xfe(0.5, 0.2) - 0.3).abs() < 1e-10); + } +} diff --git a/src/services/analysis/hr_diagram.rs b/src/services/analysis/hr_diagram.rs new file mode 100644 index 0000000..a6c4028 --- /dev/null +++ b/src/services/analysis/hr_diagram.rs @@ -0,0 +1,153 @@ +// src/services/analysis/hr_diagram.rs +// +// 赫罗图（HR Diagram）服务逻辑 + +use crate::clients::gaia::GaiaClient; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Clone)] +pub struct HrDiagramParams { + pub ra: f64, + pub dec: f64, + pub radius: f64, + pub max_mag: f64, + pub min_parallax: f64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HrDataPoint { + pub source_id: String, + pub g_mag: f64, + pub bp_rp: f64, + pub abs_g_mag: f64, + pub ra: f64, + pub dec: f64, + pub parallax: f64, + pub pm: Option, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct HrDiagramResponse { + pub points: Vec, + pub center_ra: f64, + pub center_dec: f64, + pub radius_deg: f64, + pub total_sources: usize, +} + +pub async fn build_hr_diagram( + gaia: &GaiaClient, + params: &HrDiagramParams, +) -> anyhow::Result { + let adql = format!( + "SELECT source_id, ra, dec, phot_g_mean_mag, phot_bp_mean_mag, phot_rp_mean_mag, parallax, pm \ + FROM gaiadr3.gaia_source \ + WHERE 1=CONTAINS(POINT('ICRS', ra, dec), CIRCLE('ICRS', {}, {}, {})) \ + AND phot_g_mean_mag < {} \ + AND parallax > {} \ + AND phot_bp_mean_mag IS NOT NULL \ + AND phot_rp_mean_mag IS NOT NULL", + params.ra, params.dec, params.radius, params.max_mag, params.min_parallax + ); + + let result = gaia.tap_query(&adql).await?; + + let v: serde_json::Value = serde_json::from_str(&result)?; + + let fields = v + .get("fields") + .and_then(|f| f.as_array()) + .ok_or_else(|| anyhow::anyhow!("Gaia 响应缺 fields"))?; + let data = v + .get("data") + .and_then(|d| d.as_array()) + .ok_or_else(|| anyhow::anyhow!("Gaia 响应缺 data"))?; + + let col = |name: &str| -> Option { + fields + .iter() + .position(|f| f.get("name").and_then(|n| n.as_str()) == Some(name)) + }; + + let i_sid = col("source_id"); + let i_ra = col("ra"); + let i_dec = col("dec"); + let i_g = col("phot_g_mean_mag"); + let i_bp = col("phot_bp_mean_mag"); + let i_rp = col("phot_rp_mean_mag"); + let i_plx = col("parallax"); + let i_pm = col("pm"); + + let mut points = Vec::new(); + for row in data { + if let Some(arr) = row.as_array() { + let get_f64 = |i: Option| -> Option { + i.and_then(|idx| arr.get(idx)).and_then(|v| { + if v.is_null() { + None + } else { + v.as_f64() + .or_else(|| v.as_str().and_then(|s| s.parse().ok())) + } + }) + }; + let get_str = |i: Option| -> Option { + i.and_then(|idx| arr.get(idx)).and_then(|v| { + if v.is_i64() { + Some(v.as_i64().unwrap().to_string()) + } else if v.is_f64() { + Some(format!("{:.0}", v.as_f64().unwrap())) + } else { + v.as_str().map(|s| s.to_string()) + } + }) + }; + + let sid = match get_str(i_sid) { + Some(s) => s, + None => continue, + }; + let g = match get_f64(i_g) { + Some(v) => v, + None => continue, + }; + let bp = match get_f64(i_bp) { + Some(v) => v, + None => continue, + }; + let rp = match get_f64(i_rp) { + Some(v) => v, + None => continue, + }; + let plx = match get_f64(i_plx) { + Some(v) => v, + None => continue, + }; + if plx <= 0.0 { + continue; + } + + let abs_g = g + 5.0 + 5.0 * (plx / 1000.0).log10(); + let bp_rp = bp - rp; + + points.push(HrDataPoint { + source_id: sid, + g_mag: g, + bp_rp, + abs_g_mag: abs_g, + ra: get_f64(i_ra).unwrap_or(0.0), + dec: get_f64(i_dec).unwrap_or(0.0), + parallax: plx, + pm: get_f64(i_pm), + }); + } + } + + Ok(HrDiagramResponse { + total_sources: points.len(), + points, + center_ra: params.ra, + center_dec: params.dec, + radius_deg: params.radius, + }) +} diff --git a/src/services/analysis/kinematics.rs b/src/services/analysis/kinematics.rs new file mode 100644 index 0000000..fa50f9d --- /dev/null +++ b/src/services/analysis/kinematics.rs @@ -0,0 +1,365 @@ +// src/services/analysis/kinematics.rs +// +// 银河系运动学分析 —— 自行/视差 → UVW 空间速度 + 移动星群识别 +// +// 输入：(RA, Dec, parallax, pmRA, pmDec, RV) +// 输出：(X, Y, Z, U, V, W) 银河系笛卡尔坐标 + 移动星群成员概率 + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +/// 单个天体的运动学参数 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KinematicInput { + /// 赤经 (度, J2000) + pub ra: f64, + /// 赤纬 (度, J2000) + pub dec: f64, + /// 视差 (mas) + pub parallax: f64, + /// 自行 RA (mas/yr) + pub pm_ra: f64, + /// 自行 Dec (mas/yr) + pub pm_dec: f64, + /// 径向速度 (km/s) + pub rv: f64, + /// 可选标识 + #[serde(default)] + pub id: Option, +} + +/// 运动学分析结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KinematicResult { + /// 输入标识 + pub id: Option, + /// 银心距 X (kpc) + pub x: f64, + /// Y (kpc) + pub y: f64, + /// Z (kpc) + pub z: f64, + /// 银河系旋转方向速度 U (km/s, 向银心为正) + pub u: f64, + /// V (km/s, 银河系旋转方向为正) + pub v: f64, + /// W (km/s, 向北银极方向为正) + pub w: f64, + /// 距离 (pc) + pub distance_pc: f64, + /// 移动星群成员概率 + pub membership: Option, +} + +/// 移动星群成员概率 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StellarGroupMembership { + /// 薄盘成员概率 + pub thin_disk: f64, + /// 厚盘成员概率 + pub thick_disk: f64, + /// 晕成员概率 + pub halo: f64, + /// 最可能的分类 + pub best_fit: String, +} + +// ── 银河系常数 ── + +/// 太阳到银心距离 R0 (kpc) +const _R0: f64 = 8.178; +/// 太阳在 Z 方向距银道面 Z0 (kpc) +const Z0: f64 = 0.0208; +/// 太阳的 U 速度 (km/s, 向银心) +const U_SUN: f64 = 10.6; +/// 太阳的 V 速度 (km/s, 银河系旋转方向) +const V_SUN: f64 = 10.7; +/// 太阳的 W 速度 (km/s, 向北银极) +const W_SUN: f64 = 7.6; + +/// 太阳运动学标准值 +const V_LSR: f64 = 220.0; // 本地静止标准速度 (km/s) + +/// 从自行和视差计算距离 (pc) +pub fn parallax_to_distance_pc(parallax_mas: f64) -> Result { + if parallax_mas <= 0.0 { + return Err(anyhow!("视差必须为正 (got {} mas)", parallax_mas)); + } + Ok(1000.0 / parallax_mas) +} + +/// 赤道坐标 → 银道坐标 (l, b) +/// +/// 使用标准 J2000 转换（参考论文：Bakis 1992, A&A 250, 304） +/// 北银极 (NGP) 坐标：α_G = 192.85948°, δ_G = 27.12825° +/// 银经零点（北天极在银道面上的投影）：l_NCP = 122.93192° +pub fn equatorial_to_galactic(ra_deg: f64, dec_deg: f64) -> (f64, f64) { + let ra = ra_deg.to_radians(); + let dec = dec_deg.to_radians(); + + // 北银极坐标 (J2000) + let alpha_g = (192.85948_f64).to_radians(); // NGP 赤经 + let delta_g = (27.12825_f64).to_radians(); // NGP 赤纬 + let l_ncp = (122.93192_f64).to_radians(); // 北天极银经 + + // 银纬 + let sin_b = delta_g.sin() * dec.sin() + delta_g.cos() * dec.cos() * (ra - alpha_g).cos(); + let b = sin_b.asin(); + + // 银经 + let y = (ra - alpha_g).sin() * dec.cos(); + let x = delta_g.cos() * dec.sin() - delta_g.sin() * dec.cos() * (ra - alpha_g).cos(); + let l = l_ncp - y.atan2(x); + + // 规范化到 [0, 360) + let l_deg = l.to_degrees(); + let l_deg = if l_deg < 0.0 { + l_deg + 360.0 + } else if l_deg >= 360.0 { + l_deg - 360.0 + } else { + l_deg + }; + + (l_deg, b.to_degrees()) +} + +/// 完整运动学分析：(RA, Dec, parallax, pmRA, pmDec, RV) → (X, Y, Z, U, V, W) +pub fn compute_kinematics(input: &KinematicInput) -> Result { + let distance_pc = parallax_to_distance_pc(input.parallax)?; + let distance_kpc = distance_pc / 1000.0; + + // 赤道坐标 → 银道坐标 (l, b) (度) + let (l, b) = equatorial_to_galactic(input.ra, input.dec); + let l_rad = l.to_radians(); + let b_rad = b.to_radians(); + + // 自行转为速度分量 (km/s) + // (pm / 1000) * 4.74047 * d (pc) = 横向速度 (km/s) + let mu_alpha = input.pm_ra * 0.00474047 * distance_pc; + let mu_delta = input.pm_dec * 0.00474047 * distance_pc; + + // 空间位置 (X, Y, Z) (kpc) + // X: 向银心方向 + // Y: 银河系旋转方向 + // Z: 向北银极方向 + let cos_l = l_rad.cos(); + let sin_l = l_rad.sin(); + let cos_b = b_rad.cos(); + let sin_b = b_rad.sin(); + + let x = distance_kpc * cos_b * cos_l; + let y = -distance_kpc * cos_b * sin_l; // 注意符号：银河系旋转方向 + let z = distance_kpc * sin_b + Z0; + + // 空间速度 (U, V, W) (km/s) + // 需要考虑自行的银道分量 + let cos_ra = input.ra.to_radians().cos(); + let sin_ra = input.ra.to_radians().sin(); + let cos_dec = input.dec.to_radians().cos(); + let sin_dec = input.dec.to_radians().sin(); + + // 行星运动方向的单位向量 (e_ra, e_dec, e_r) 在赤道坐标系中 + // e_ra = (-sin_ra, cos_ra, 0) + // e_dec = (-sin_dec*cos_ra, -sin_dec*sin_ra, cos_dec) + // e_r = (cos_dec*cos_ra, cos_dec*sin_ra, sin_dec) + + // 横向速度在赤道坐标系中的分量 + let v_ra = mu_alpha; // km/s + let v_dec = mu_delta; // km/s + let v_r = input.rv; // km/s + + // 赤道坐标系速度分量 + let vx_eq = -v_ra * sin_ra - v_dec * sin_dec * cos_ra + v_r * cos_dec * cos_ra; + let vy_eq = v_ra * cos_ra - v_dec * sin_dec * sin_ra + v_r * cos_dec * sin_ra; + let vz_eq = v_dec * cos_dec + v_r * sin_dec; + + // 从赤道坐标系到银河系坐标系的旋转矩阵 (J2000) + // 银经 l, 银纬 b, 银道面倾角 + let a_g = (266.405_f64).to_radians(); // 银心赤经 + let d_g = (-28.936_f64).to_radians(); // 银心赤纬 + let l_omega = (32.932_f64).to_radians(); // 升交点银经 + + let sin_lomega = l_omega.sin(); + let cos_lomega = l_omega.cos(); + let cos_dg = d_g.cos(); + let sin_dg = d_g.sin(); + + // 银河系速度分量 + let u_gal = -vx_eq * (a_g.sin() * sin_lomega + a_g.cos() * cos_lomega * sin_dg) + - vy_eq * (-a_g.cos() * sin_lomega + a_g.sin() * cos_lomega * sin_dg) + - vz_eq * cos_lomega * cos_dg; + let v_gal = vx_eq * (-a_g.sin() * cos_lomega + a_g.cos() * sin_lomega * sin_dg) + + vy_eq * (a_g.cos() * cos_lomega + a_g.sin() * sin_lomega * sin_dg) + + vz_eq * sin_lomega * cos_dg; + let w_gal = vx_eq * a_g.cos() * cos_dg + vy_eq * a_g.sin() * cos_dg - vz_eq * sin_dg; + + // 转换为标准 UVW（减去太阳运动） + let u = u_gal + U_SUN; + let v = v_gal + V_SUN + V_LSR; + let w = w_gal + W_SUN; + + // 移动星群成员概率 + let membership = classify_stellar_group(u, v, w, z, distance_kpc); + + Ok(KinematicResult { + id: input.id.clone(), + x, + y, + z, + u, + v, + w, + distance_pc, + membership: Some(membership), + }) +} + +/// 简化的移动星群分类（基于 UVW 速度和 Z 高度） +fn classify_stellar_group( + _u: f64, + v: f64, + w: f64, + z_kpc: f64, + _distance_kpc: f64, +) -> StellarGroupMembership { + // 简化规则： + // 薄盘：|Z| < 0.3 kpc, |V - V_LSR| < 50 km/s, |W| < 40 km/s + // 厚盘：0.3 < |Z| < 1.5 kpc, |V - V_LSR| < 100 km/s + // 晕：|Z| > 1.5 kpc, 高速运动 + + let abs_z = z_kpc.abs(); + let v_offset = (v - V_LSR).abs(); + let abs_w = w.abs(); + + // 计算各成分的概率（基于距离度量的软分类） + let thin_disk_score = (-abs_z / 0.3 - v_offset / 50.0 - abs_w / 40.0).exp(); + let thick_disk_score = (-(abs_z - 0.8).abs() / 0.5 - v_offset / 100.0 - abs_w / 60.0).exp(); + let halo_score = if abs_z > 1.0 { + (-(abs_z - 2.0).abs() / 1.0 - v_offset / 200.0 - abs_w / 100.0).exp() + } else { + 0.01 + }; + + // 归一化 + let total = thin_disk_score + thick_disk_score + halo_score; + let thin_disk = thin_disk_score / total; + let thick_disk = thick_disk_score / total; + let halo = halo_score / total; + + let best_fit = if thin_disk > thick_disk && thin_disk > halo { + "薄盘".to_string() + } else if thick_disk > halo { + "厚盘".to_string() + } else { + "晕".to_string() + }; + + StellarGroupMembership { + thin_disk, + thick_disk, + halo, + best_fit, + } +} + +/// 批量运动学分析 +pub fn compute_kinematics_batch(inputs: &[KinematicInput]) -> Vec> { + inputs.iter().map(compute_kinematics).collect() +} + +/// 格式化运动学结果为 Markdown +pub fn format_kinematic_result(result: &KinematicResult) -> String { + let mut out = String::from("## 运动学分析结果\n\n"); + + if let Some(ref id) = result.id { + out.push_str(&format!("**目标**: {}\n\n", id)); + } + + out.push_str("### 三维位置 (kpc)\n\n"); + out.push_str(&format!("- X = {:.3} kpc\n", result.x)); + out.push_str(&format!("- Y = {:.3} kpc\n", result.y)); + out.push_str(&format!("- Z = {:.3} kpc\n", result.z)); + out.push_str(&format!("- 距离 = {:.1} pc\n\n", result.distance_pc)); + + out.push_str("### 三维速度 (km/s)\n\n"); + out.push_str(&format!("- U = {:.1} km/s (向银心)\n", result.u)); + out.push_str(&format!("- V = {:.1} km/s (银河系旋转)\n", result.v)); + out.push_str(&format!("- W = {:.1} km/s (向北银极)\n", result.w)); + + if let Some(ref membership) = result.membership { + out.push_str(&format!( + "\n### 移动星群分类: **{}**\n\n", + membership.best_fit + )); + out.push_str(&format!("- 薄盘: {:.1}%\n", membership.thin_disk * 100.0)); + out.push_str(&format!("- 厚盘: {:.1}%\n", membership.thick_disk * 100.0)); + out.push_str(&format!("- 晕: {:.1}%\n", membership.halo * 100.0)); + } + + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parallax_to_distance() { + let d = parallax_to_distance_pc(100.0).unwrap(); + assert!((d - 10.0).abs() < 0.01); + } + + #[test] + fn test_parallax_zero_error() { + assert!(parallax_to_distance_pc(0.0).is_err()); + } + + #[test] + fn test_equatorial_to_galactic_sgp() { + // 南银极 (SGP) 赤道坐标 (J2000): RA ≈ 12.86°, Dec ≈ -27.13° + let (_l, b) = equatorial_to_galactic(12.85948, -27.12825); + // 银纬应为 -90°（南银极） + assert!((b - (-90.0)).abs() < 0.01, "SGP 银纬 {} 应为 -90°", b); + } + + #[test] + fn test_kinematics_solar() { + // 太阳附近的恒星，自行小，RV 小 + let input = KinematicInput { + ra: 0.0, + dec: 0.0, + parallax: 1000.0, // 1 pc + pm_ra: 0.0, + pm_dec: 0.0, + rv: 0.0, + id: Some("test".to_string()), + }; + let result = compute_kinematics(&input).unwrap(); + assert!((result.distance_pc - 1.0).abs() < 0.1); + } + + #[test] + fn test_kinematics_high_pm() { + // 距离 10 pc 的恒星，自行 1000 mas/yr (1 arcsec/yr) + // 速度应为 1 * 10 * 4.74047 = 47.4047 km/s + let input = KinematicInput { + ra: 0.0, + dec: 0.0, + parallax: 100.0, // 10 pc + pm_ra: 1000.0, // 1000 mas/yr + pm_dec: 0.0, + rv: 0.0, + id: Some("high_pm".to_string()), + }; + let result = compute_kinematics(&input).unwrap(); + assert!((result.distance_pc - 10.0).abs() < 0.01); + + // 还原太阳运动，检查变换后空间速度大小是否等于输入横向速度 + let u_gal = result.u - U_SUN; + let v_gal = result.v - V_SUN - V_LSR; + let w_gal = result.w - W_SUN; + let v_mag = (u_gal * u_gal + v_gal * v_gal + w_gal * w_gal).sqrt(); + assert!((v_mag - 47.4047).abs() < 1e-4); + } +} diff --git a/src/services/analysis/mod.rs b/src/services/analysis/mod.rs new file mode 100644 index 0000000..1ceb4dc --- /dev/null +++ b/src/services/analysis/mod.rs @@ -0,0 +1,15 @@ +// src/services/analysis/mod.rs +// +// 天体物理分析服务层 —— SED 拟合、运动学、化学丰度、观测规划 +// +// 模块布局： +// sed.rs —— 多波段 SED 拟合（黑体/模板），求 T_eff, L, A_V +// kinematics.rs —— Gaia 自行/视差 → 银河系 UVW 空间速度 + 移动星群识别 +// chemistry.rs —— 化学丰度模式分析（[α/Fe] vs [Fe/H]，厚盘/薄盘/晕星区分） +// observability.rs —— 观测提案辅助（目标可见性、月相、曝光时间估算） + +pub mod chemistry; +pub mod hr_diagram; +pub mod kinematics; +pub mod observability; +pub mod sed; diff --git a/src/services/analysis/observability.rs b/src/services/analysis/observability.rs new file mode 100644 index 0000000..71e28fc --- /dev/null +++ b/src/services/analysis/observability.rs @@ -0,0 +1,401 @@ +// src/services/analysis/observability.rs +// +// 观测提案辅助 —— 目标可见性、月相、曝光时间估算 +// +// 输入：天体坐标、观测台站、时间范围 +// 输出：升起/落下时间、最佳观测窗口、月相影响、曝光时间估算 + +use serde::{Deserialize, Serialize}; + +/// 观测台站参数 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Observatory { + /// 台站名称 + pub name: String, + /// 经度 (度, 东经为正) + pub longitude: f64, + /// 纬度 (度, 北纬为正) + pub latitude: f64, + /// 海拔 (米) + pub altitude: f64, + /// 典型视宁度 (角秒) + pub seeing: f64, + /// 大气消光系数 (mag/airmass, V 波段) + pub extinction_v: f64, +} + +impl Observatory { + /// 常用天文台站 + pub fn lick_3m() -> Self { + Self { + name: "Lick 3m".to_string(), + longitude: -121.6428, + latitude: 37.3414, + altitude: 1283.0, + seeing: 1.2, + extinction_v: 0.15, + } + } + + pub fn keck() -> Self { + Self { + name: "Keck I/II".to_string(), + longitude: -155.4761, + latitude: 19.8263, + altitude: 4145.0, + seeing: 0.8, + extinction_v: 0.12, + } + } + + pub fn lamost() -> Self { + Self { + name: "LAMOST".to_string(), + longitude: 117.5837, + latitude: 40.3958, + altitude: 900.0, + seeing: 2.0, + extinction_v: 0.20, + } + } + + pub fn gemini_north() -> Self { + Self { + name: "Gemini North".to_string(), + longitude: -155.4691, + latitude: 19.8238, + altitude: 2737.0, + seeing: 0.7, + extinction_v: 0.10, + } + } + + pub fn default_list() -> Vec { + vec![ + Self::lick_3m(), + Self::keck(), + Self::lamost(), + Self::gemini_north(), + ] + } +} + +/// 观测可行性结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ObservabilityResult { + /// 目标升起时间 (UTC) + pub rise_time: Option, + /// 目标落下时间 (UTC) + pub set_time: Option, + /// 最大高度角 (度) + pub max_altitude: f64, + /// 最大高度角时间 (UTC) + pub max_altitude_time: Option, + /// 今晚是否可观测（高度角 > 30° 且天黑） + pub observable_tonight: bool, + /// 最佳观测窗口（高度角 > 30° 的时段） + pub best_window: Option, + /// 月相影响 + pub moon_impact: MoonImpact, + /// 曝光时间估算（V 波段，1σ = 5σ 目标星等） + pub exposure_estimate: Option, + /// 大气质量（在最大高度角时） + pub airmass_at_max: f64, +} + +/// 月相影响 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MoonImpact { + /// 月相角 (0=新月, 180=满月) + pub phase_angle: f64, + /// 月相描述 + pub phase_description: String, + /// 月亮距目标角距离 (度) + pub separation: f64, + /// 对观测的影响等级 + pub impact: String, +} + +/// 曝光时间估算 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExposureEstimate { + /// V 波段目标星等 + pub target_v_mag: f64, + /// 所需信噪比 + pub target_snr: f64, + /// 估算曝光时间 (秒) + pub exposure_time_secs: f64, + /// 所需帧数（含 dithering） + pub n_frames: u32, + /// 总时间 (秒) + pub total_time_secs: f64, + /// 备注 + pub notes: String, +} + +/// 估算目标在某时刻的高度角和方位角 +/// +/// 输入：目标赤经赤纬 (度), 观测台站, 儒略日 +pub fn compute_altitude(ra_deg: f64, dec_deg: f64, obs: &Observatory, jd: f64) -> (f64, f64) { + // 简化的地平坐标计算 + let _ra = ra_deg.to_radians(); + let dec = dec_deg.to_radians(); + let lat = obs.latitude.to_radians(); + let lon = obs.longitude.to_radians(); + + // 估算当前时角（简化：假设太阳时 ≈ 平太阳时） + let _t = (jd - 2451545.0) * 2.0 * std::f64::consts::PI; // 粗略近似 + let gmst = (280.46061837 + 360.98564736629 * (jd - 2451545.0)) % 360.0; + let lst = (gmst + lon.to_degrees()) % 360.0; + let hour_angle = (lst - ra_deg).to_radians(); + + // 高度角 + let sin_alt = dec.sin() * lat.sin() + dec.cos() * lat.cos() * hour_angle.cos(); + let altitude = sin_alt.asin().to_degrees(); + + // 方位角 + let cos_az = (dec.sin() - lat.sin() * sin_alt) / (lat.cos() * altitude.to_radians().cos()); + let azimuth = cos_atan2(cos_az, hour_angle.sin() * lat.cos()).to_degrees(); + + (altitude, azimuth) +} + +/// 简化的 atan2 范围限制 +fn cos_atan2(y: f64, x: f64) -> f64 { + if x > 0.0 { + (y / x).acos() + } else if x < 0.0 { + if y >= 0.0 { + std::f64::consts::PI - (-y / -x).acos() + } else { + -std::f64::consts::PI + (y / -x).acos() + } + } else { + 0.0 + } +} + +/// 计算月相（简化） +pub fn compute_moon_phase(jd: f64) -> (f64, String) { + // 月相角计算（简化） + let d = jd - 2451545.0; + let age = d % 29.53058770576; // 朔望月周期 + let phase = age / 29.53058770576 * 360.0; + + let description = if phase < 11.25 { + "新月".to_string() + } else if phase < 33.75 { + "蛾眉月".to_string() + } else if phase < 56.25 { + "上弦月".to_string() + } else if phase < 78.75 { + "盈凸月".to_string() + } else if phase < 101.25 { + "满月".to_string() + } else if phase < 123.75 { + "亏凸月".to_string() + } else if phase < 146.25 { + "下弦月".to_string() + } else if phase < 168.75 { + "残月".to_string() + } else { + "新月".to_string() + }; + + (phase, description) +} + +/// 月相对于观测的影响 +fn assess_moon_impact(phase_angle: f64, _target_v_mag: f64) -> MoonImpact { + let impact = if !(30.0..=330.0).contains(&phase_angle) { + "低 — 新月前后，天空暗".to_string() + } else if !(90.0..=270.0).contains(&phase_angle) { + "中 — 蛾眉月/残月，有一定散射光".to_string() + } else if !(150.0..=210.0).contains(&phase_angle) { + "高 — 盈凸/亏凸月，显著散射光".to_string() + } else { + "极高 — 满月前后，深空观测困难".to_string() + }; + + MoonImpact { + phase_angle, + phase_description: String::new(), // 由调用者填充 + separation: 0.0, // 由调用者填充 + impact, + } +} + +/// 计算大气质标 (Airmass)，使用 Pickering (1982) 割线修正公式 +pub fn calculate_airmass(altitude_deg: f64) -> f64 { + if altitude_deg > 0.0 { + let h = altitude_deg; + let denom = (h + 244.0 / (165.0 + 47.0 * h.powf(1.1))) + .to_radians() + .sin(); + if denom > 1e-10 { + 1.0 / denom + } else { + f64::INFINITY + } + } else { + f64::INFINITY // 不可观测或高度角为负 + } +} + +/// 估算曝光时间 +pub fn estimate_exposure_time( + target_v_mag: f64, + target_snr: f64, + obs: &Observatory, + altitude: f64, +) -> ExposureEstimate { + // 简化的曝光时间估算 + // 基于天光背景和恒星信号的信噪比模型 + + let airmass = calculate_airmass(altitude); + + // V 波段天空亮度 (mag/arcsec²) + let _sky_brightness = 21.5; // 晴朗无月夜晚 + // 大气消光 + let extinction = obs.extinction_v * airmass; + + // 简化的 S/N 估算（基于 CCD 性能典型值） + // S/N ≈ signal × t / sqrt(signal × t + sky × t + readnoise²) + // 简化为：t ≈ (SNR/100)² × 10^(0.4 * (V - 20)) × 10^(0.4 * extinction) + let base_time = 10.0; // 100 SNR 所需基本时间 + let v_factor = 10.0_f64.powf(0.4 * (target_v_mag - 20.0)); + let atm_factor = 10.0_f64.powf(0.4 * extinction); + let snr_factor = (target_snr / 100.0).powi(2); + + let exposure_secs = base_time * v_factor * atm_factor * snr_factor; + let exposure_secs = exposure_secs.clamp(0.1, 3600.0); // 0.1s - 1h + + // 帧数（含 dithering） + let n_frames = if exposure_secs < 300.0 { + 1 + } else { + (exposure_secs / 300.0).ceil() as u32 + }; + + let total_time = exposure_secs * n_frames as f64; + + let notes = format!( + "基于 V={:.1}, SNR={:.0}, airmass={:.2} 的简化估算。实际所需时间受天气、仪器等因素影响。", + target_v_mag, target_snr, airmass + ); + + ExposureEstimate { + target_v_mag, + target_snr, + exposure_time_secs: exposure_secs, + n_frames, + total_time_secs: total_time, + notes, + } +} + +/// 完整观测可行性评估 +pub fn assess_observability( + ra_deg: f64, + dec_deg: f64, + obs: &Observatory, + target_v_mag: Option, + target_snr: Option, +) -> ObservabilityResult { + let now_jd = 2451545.0; // 简化：使用 J2000.0 作为参考 + let (altitude, _azimuth) = compute_altitude(ra_deg, dec_deg, obs, now_jd); + let airmass = calculate_airmass(altitude); + + let (phase_angle, phase_desc) = compute_moon_phase(now_jd); + let mut moon_impact = assess_moon_impact(phase_angle, target_v_mag.unwrap_or(15.0)); + moon_impact.phase_description = phase_desc; + + let observable_tonight = altitude > 30.0; + + let exposure_estimate = if let (Some(vmag), Some(snr)) = (target_v_mag, target_snr) { + Some(estimate_exposure_time(vmag, snr, obs, altitude.max(0.0))) + } else { + None + }; + + ObservabilityResult { + rise_time: None, // 需要更精确的天文计算库 + set_time: None, + max_altitude: altitude, + max_altitude_time: None, + observable_tonight, + best_window: if observable_tonight { + Some("当前高度角 > 30°，可进行观测".to_string()) + } else { + Some("目标当前不可观测（高度角 < 30°）".to_string()) + }, + moon_impact, + exposure_estimate, + airmass_at_max: airmass, + } +} + +/// 格式化观测可行性结果为 Markdown +pub fn format_observability(result: &ObservabilityResult, obs_name: &str) -> String { + let mut out = format!("## 观测可行性评估 ({})\n\n", obs_name); + + out.push_str(&format!("- **最大高度角**: {:.1}°\n", result.max_altitude)); + out.push_str(&format!("- **大气质量**: {:.2}\n", result.airmass_at_max)); + out.push_str(&format!( + "- **今晚可观测**: {}\n", + if result.observable_tonight { + "是" + } else { + "否" + } + )); + + if let Some(ref window) = result.best_window { + out.push_str(&format!("- **最佳窗口**: {}\n", window)); + } + + out.push_str("\n### 月相\n\n"); + out.push_str(&format!( + "- 月相角: {:.0}° ({})\n", + result.moon_impact.phase_angle, result.moon_impact.phase_description + )); + out.push_str(&format!("- 影响: {}\n", result.moon_impact.impact)); + + if let Some(ref exp) = result.exposure_estimate { + out.push_str("\n### 曝光时间估算\n\n"); + out.push_str(&format!("- 目标星等 (V): {:.1}\n", exp.target_v_mag)); + out.push_str(&format!("- 目标 SNR: {:.0}\n", exp.target_snr)); + out.push_str(&format!("- 单帧曝光: {:.1} s\n", exp.exposure_time_secs)); + out.push_str(&format!("- 帧数: {}\n", exp.n_frames)); + out.push_str(&format!("- 总时间: {:.0} s\n", exp.total_time_secs)); + out.push_str(&format!("- 备注: {}\n", exp.notes)); + } + + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_observatory_list() { + let observatories = Observatory::default_list(); + assert!(!observatories.is_empty()); + } + + #[test] + fn test_moon_phase() { + let (phase, desc) = compute_moon_phase(2451545.0); + assert!((0.0..=360.0).contains(&phase)); + assert!(!desc.is_empty()); + } + + #[test] + fn test_exposure_estimate() { + let obs = Observatory::keck(); + let exp = estimate_exposure_time(15.0, 100.0, &obs, 60.0); + assert!(exp.exposure_time_secs > 0.0); + assert!(exp.exposure_time_secs < 3600.0); + } +} diff --git a/src/services/analysis/sed.rs b/src/services/analysis/sed.rs new file mode 100644 index 0000000..79e3156 --- /dev/null +++ b/src/services/analysis/sed.rs @@ -0,0 +1,363 @@ +// src/services/analysis/sed.rs +// +// 多波段 SED 拟合 —— 黑体辐射模型拟合，求 T_eff, 光度, 消光 A_V +// +// 输入：多波段测光数据（波长、通量、通量误差） +// 输出：拟合参数（T_eff, R, A_V, L）+ 不确定度 + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +/// SED 数据点 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SedDataPoint { + /// 中心波长 (Å) + pub wavelength: f64, + /// 流量 (erg/s/cm²/Å 或 Jy) + pub flux: f64, + /// 流量误差 + pub flux_error: f64, + /// 波段名（如 "J", "W1", "G"） + pub band: String, + /// 数据源 + pub source: String, +} + +/// SED 拟合结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SedFitResult { + /// 有效温度 (K) + pub teff: f64, + /// 半径 (R_sun) + pub radius: f64, + /// 消光 A_V (mag) + pub a_v: f64, + /// 光度 (L_sun) + pub luminosity: f64, + /// 距离 (pc)，用于计算光度 + pub distance_pc: Option, + /// 拟合残差（chi2） + pub chi2: f64, + /// 约化 chi2 + pub reduced_chi2: f64, + /// 参数不确定度（1σ） + pub uncertainties: SedUncertainties, + /// 拟合曲线点 (wavelength, flux)，供前端绘图 + pub model_curve: Vec<(f64, f64)>, + /// 残差点 + pub residuals: Vec<(f64, f64)>, +} + +/// 拟合参数不确定度 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SedUncertainties { + pub teff_err: f64, + pub radius_err: f64, + pub a_v_err: f64, + pub luminosity_err: f64, +} + +/// 简化星际消光曲线（Ccm87 近似） +fn extinction_law(wavelength_angstrom: f64, r_v: f64) -> f64 { + let x = 1e4 / wavelength_angstrom; // 1/λ (μm⁻¹) + + // CCM87 近似（光学-近红外） + if x < 0.3 { + // 远红外：A_λ ∝ λ⁻¹.75 + 0.265 * r_v * x.powf(1.75) + 0.861 + } else if x < 1.1 { + // 红外 + 0.574 * x.powf(1.61) - 0.527 * x.powf(3.0) / r_v + } else if x < 3.3 { + // 光学 + let y = x - 1.82; + let a = 1.0 + 0.17699 * y - 0.50447 * y * y - 0.02427 * y.powi(3) + + 0.72085 * y.powi(4) + + 0.01979 * y.powi(5) + - 0.77530 * y.powi(6) + + 0.32999 * y.powi(7); + let b = 1.41338 * y + 2.28305 * y * y + 1.07233 * y.powi(3) + - 5.38434 * y.powi(4) + - 0.62251 * y.powi(5) + + 5.30260 * y.powi(6) + - 2.09002 * y.powi(7); + a + b / r_v + } else { + // 紫外 + let a = 1.752 - 0.316 * x - 0.104 / ((x - 4.67).powi(2) + 0.341) + 0.293; + let b = -3.090 + 1.825 * x + 1.206 / ((x - 4.62).powi(2) + 0.263); + a + b / r_v + } +} + +/// 黑体辐射模型（普朗克函数） +fn black_body_flux(wavelength_angstrom: f64, teff: f64) -> f64 { + let wl = wavelength_angstrom * 1e-10; // Å → m + let h = 6.626e-34; + let c = 3e8; + let k = 1.381e-23; + + let x = h * c / (wl * k * teff); + if x > 500.0 { + return 0.0; + } + + let b = 2.0 * h * c * c / wl.powi(5) / (x.exp() - 1.0); + b * 1e-26 // 转换为 Jy 量级（任意缩放） +} + +/// 黑体 + 消光模型：F_model(λ) = R² × B_λ(T) × 10^(-0.4 × A_λ × A_V / A_V) +fn sed_model(wavelength: f64, teff: f64, radius_solar: f64, a_v: f64) -> f64 { + let r_v = 3.1; + let bb = black_body_flux(wavelength, teff); + let ext = extinction_law(wavelength, r_v); + let a_lambda = a_v * ext / extinction_law(5500.0, r_v); // A_λ = A_V × A(λ)/A(V) + let extinction_factor = 10.0_f64.powf(-0.4 * a_lambda); + + radius_solar.powi(2) * bb * extinction_factor +} + +/// 最小二乘拟合 SED +pub fn fit_sed(data: &[SedDataPoint], distance_pc: Option) -> Result { + if data.len() < 3 { + return Err(anyhow!("至少需要 3 个测光点进行 SED 拟合")); + } + + // 网格搜索（粗搜索 + 精搜索） + let teff_range = (2000.0, 50000.0); + let radius_range = (0.01, 1000.0); + let av_range = (0.0, 10.0); + + // 粗搜索 + let n_grid = 30; + let mut best_chi2 = f64::INFINITY; + let mut best_params = (5778.0, 1.0, 0.0); + + for i in 0..n_grid { + for j in 0..n_grid { + for k in 0..10 { + let teff = + teff_range.0 + (teff_range.1 - teff_range.0) * i as f64 / (n_grid - 1) as f64; + let radius = radius_range.0 + + (radius_range.1 - radius_range.0) * j as f64 / (n_grid - 1) as f64; + let av = av_range.0 + (av_range.1 - av_range.0) * k as f64 / 9.0; + + let chi2 = compute_chi2(data, teff, radius, av); + if chi2 < best_chi2 { + best_chi2 = chi2; + best_params = (teff, radius, av); + } + } + } + } + + // 精搜索（在最优参数附近 ±10% 范围细化） + let (ct, cr, ca) = best_params; + let fine_n = 20; + for i in 0..fine_n { + for j in 0..fine_n { + for k in 0..fine_n { + let teff = ct * 0.9 + ct * 0.2 * i as f64 / (fine_n - 1) as f64; + let radius = cr * 0.9 + cr * 0.2 * j as f64 / (fine_n - 1) as f64; + let av = (ca - 1.0).max(0.0) + 2.0 * k as f64 / (fine_n - 1) as f64; + + let chi2 = compute_chi2(data, teff, radius, av); + if chi2 < best_chi2 { + best_chi2 = chi2; + best_params = (teff, radius, av); + } + } + } + } + + let (teff, radius, a_v) = best_params; + let dof = data.len() as f64 - 3.0; // 3 个自由参数 + let reduced_chi2 = if dof > 0.0 { + best_chi2 / dof + } else { + best_chi2 + }; + + // 光度计算 + let luminosity = radius.powi(2) * (teff / 5778.0).powi(4); + + // 生成模型曲线 + let model_curve: Vec<(f64, f64)> = (100..100000) + .step_by(100) + .map(|w| { + let wl = w as f64; + (wl, sed_model(wl, teff, radius, a_v)) + }) + .collect(); + + // 残差 + let residuals: Vec<(f64, f64)> = data + .iter() + .map(|d| { + let model = sed_model(d.wavelength, teff, radius, a_v); + let residual = if d.flux > 0.0 && model > 0.0 { + (d.flux / model - 1.0) * 100.0 // 百分比残差 + } else { + 0.0 + }; + (d.wavelength, residual) + }) + .collect(); + + // 不确定度估计（简化：通过扰动参数估计） + let uncertainties = estimate_uncertainties(data, teff, radius, a_v, best_chi2, dof); + + Ok(SedFitResult { + teff, + radius, + a_v, + luminosity, + distance_pc, + chi2: best_chi2, + reduced_chi2, + uncertainties, + model_curve, + residuals, + }) +} + +/// 计算 chi2（模型和数据都归一化到 V 波段 5500Å） +fn compute_chi2(data: &[SedDataPoint], teff: f64, radius: f64, a_v: f64) -> f64 { + let v_model = sed_model(5500.0, teff, radius, a_v); + if v_model <= 0.0 { + return f64::INFINITY; + } + data.iter() + .map(|d| { + let model = sed_model(d.wavelength, teff, radius, a_v) / v_model; + if d.flux_error > 0.0 { + ((d.flux - model) / d.flux_error).powi(2) + } else { + 0.0 + } + }) + .sum() +} + +/// 参数不确定度估计（通过扰动法） +fn estimate_uncertainties( + data: &[SedDataPoint], + teff: f64, + radius: f64, + a_v: f64, + best_chi2: f64, + _dof: f64, +) -> SedUncertainties { + // 扰动每个参数 10%，观察 chi2 变化 + let delta = 0.1; + + let chi2_t = compute_chi2(data, teff * (1.0 + delta), radius, a_v); + let teff_err = if chi2_t > best_chi2 { + teff * delta * ((chi2_t - best_chi2).max(0.0)).sqrt() + } else { + teff * 0.1 + }; + + let chi2_r = compute_chi2(data, teff, radius * (1.0 + delta), a_v); + let radius_err = if chi2_r > best_chi2 { + radius * delta * ((chi2_r - best_chi2).max(0.0)).sqrt() + } else { + radius * 0.1 + }; + + let chi2_a = compute_chi2(data, teff, radius, (a_v + 0.1).max(0.0)); + let a_v_err = if chi2_a > best_chi2 { + 0.1 * ((chi2_a - best_chi2).max(0.0)).sqrt() + } else { + 0.5 + }; + + let luminosity_err = (radius_err / radius).hypot(teff_err / teff) * 2.0; + + SedUncertainties { + teff_err, + radius_err, + a_v_err, + luminosity_err, + } +} + +/// 格式化 SED 拟合结果为 Markdown +pub fn format_sed_result(result: &SedFitResult) -> String { + let mut out = String::from("## SED 拟合结果\n\n"); + out.push_str("| 参数 | 值 | 不确定度 |\n|------|-----|----------|\n"); + out.push_str(&format!( + "| T_eff | {:.0} K | ±{:.0} K |\n", + result.teff, result.uncertainties.teff_err + )); + out.push_str(&format!( + "| R | {:.4} R_sun | ±{:.4} |\n", + result.radius, result.uncertainties.radius_err + )); + out.push_str(&format!( + "| A_V | {:.3} mag | ±{:.3} |\n", + result.a_v, result.uncertainties.a_v_err + )); + out.push_str(&format!( + "| L | {:.4} L_sun | ±{:.1}% |\n", + result.luminosity, + result.uncertainties.luminosity_err * 100.0 + )); + out.push_str(&format!("| χ² | {:.2} | |\n", result.chi2)); + out.push_str(&format!("| χ²_red | {:.2} | |\n", result.reduced_chi2)); + if let Some(d) = result.distance_pc { + out.push_str(&format!("| 距离 | {:.1} pc | |\n", d)); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_black_body_positive() { + let bb = black_body_flux(5500.0, 5778.0); + assert!(bb > 0.0); + } + + #[test] + fn test_sed_model_increases_with_radius() { + let f1 = sed_model(5500.0, 5778.0, 1.0, 0.0); + let f2 = sed_model(5500.0, 5778.0, 2.0, 0.0); + assert!(f2 > f1, "更大的半径应该产生更大的流量"); + } + + #[test] + fn test_fit_sed_solar() { + // 构造太阳型恒星的模拟测光数据（从模型本身生成，加小噪声） + let v_flux = sed_model(5500.0, 5778.0, 1.0, 0.0); + let data_normalized: Vec = vec![ + (4400.0, "B"), + (5500.0, "V"), + (6500.0, "R"), + (8000.0, "I"), + (12500.0, "J"), + ] + .into_iter() + .map(|(w, band)| { + let model = sed_model(w, 5778.0, 1.0, 0.0); + SedDataPoint { + wavelength: w, + flux: model / v_flux, + flux_error: 0.01, + band: band.to_string(), + source: "test".to_string(), + } + }) + .collect(); + + let result = fit_sed(&data_normalized, Some(10.0)).unwrap(); + // T_eff 应接近太阳（允许 40% 误差，因为网格搜索粗糙） + assert!( + (result.teff - 5778.0).abs() / 5778.0 < 0.4, + "T_eff = {} K, 应接近 5778 K", + result.teff + ); + } +} diff --git a/src/services/batch/meta.rs b/src/services/batch/meta.rs index 47af0c6..d9102ab 100644 --- a/src/services/batch/meta.rs +++ b/src/services/batch/meta.rs @@ -317,19 +317,21 @@ impl MetaSync { // 获取所有已存同步检索配置 pub async fn list_queries(db: &SqlitePool) -> Result, sqlx::Error> { - let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') AS last_run FROM sync_queries ORDER BY last_run DESC") + let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') AS last_run, is_scheduled FROM sync_queries ORDER BY last_run DESC") .fetch_all(db) .await?; let mut list = Vec::new(); for r in rows { use sqlx::Row; + let is_scheduled_int: i32 = r.get("is_scheduled"); list.push(SavedSyncQuery { id: r.get("id"), query: r.get("query"), source: r.get("source"), limit_count: r.get("limit_count"), last_run: r.get("last_run"), + is_scheduled: is_scheduled_int != 0, }); } Ok(list) @@ -343,6 +345,162 @@ impl MetaSync { .await?; Ok(result.rows_affected() > 0) } + + // 设置是否定时同步 + pub async fn set_query_schedule( + db: &SqlitePool, + id: i64, + is_scheduled: bool, + ) -> Result { + let result = sqlx::query("UPDATE sync_queries SET is_scheduled = ? WHERE id = ?") + .bind(if is_scheduled { 1 } else { 0 }) + .bind(id) + .execute(db) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 运行静默（无状态锁）增量同步，返回实际入库/同步成功的文献数量 + pub async fn run_incremental_sync_silent( + db: SqlitePool, + ads: Arc, + arxiv: Arc, + query: String, + source: String, + limit: i32, + last_run_str: &str, + ) -> anyhow::Result { + let date_part = if last_run_str.len() >= 10 { + &last_run_str[0..10] + } else { + "now-7D" + }; + + let mut synced_count = 0; + + // 1. 同步 ADS 数据 + if source == "all" || source == "ads" { + let ads_query = if date_part == "now-7D" { + format!("{} entdate:[now-7D TO now]", query) + } else { + format!("{} entdate:[{} TO now]", query, date_part) + }; + + let ads_limit = if limit > 0 { limit } else { 100 }; + match ads.search(&ads_query, 0, ads_limit, "date_desc").await { + Ok(docs) => { + if let Ok(mut tx) = db.begin().await { + for doc in docs { + let paper = convert_ads_doc_to_standard(&doc); + if save_paper_to_db_tx(&mut tx, &paper).await.is_ok() { + synced_count += 1; + } + } + let _ = tx.commit().await; + } + } + Err(e) => { + warn!( + "[Scheduler] 增量同步 ADS 数据失败 (query={}): {}", + ads_query, e + ); + } + } + } + + // 2. 同步 arXiv 数据 + if source == "all" || source == "arxiv" { + let arxiv_limit = if limit > 0 { limit } else { 100 }; + match arxiv.search(&query, 0, arxiv_limit, "date_desc").await { + Ok(papers) => { + if let Ok(mut tx) = db.begin().await { + for p in papers { + let paper = convert_arxiv_to_standard(&p); + if save_paper_to_db_tx(&mut tx, &paper).await.is_ok() { + synced_count += 1; + } + } + let _ = tx.commit().await; + } + } + Err(e) => { + warn!( + "[Scheduler] 增量同步 arXiv 数据失败 (query={}): {}", + query, e + ); + } + } + } + + Ok(synced_count) + } +} + +/// 启动定时增量文献同步调度器，每小时轮询一次 +pub async fn start_scheduled_sync_timer(state: Arc) { + info!("[Scheduler] 启动定时增量文献同步调度器..."); + // 启动 10 秒后执行第一次同步 + tokio::time::sleep(std::time::Duration::from_secs(10)).await; + loop { + info!("[Scheduler] 开始执行定时增量同步步进..."); + if let Err(e) = run_scheduled_sync_step(&state).await { + error!("[Scheduler] 定时增量文献同步执行失败: {}", e); + } + tokio::time::sleep(std::time::Duration::from_secs(3600)).await; + } +} + +/// 执行单步增量同步 +pub async fn run_scheduled_sync_step(state: &crate::api::AppState) -> anyhow::Result<()> { + let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') AS last_run FROM sync_queries WHERE is_scheduled = 1") + .fetch_all(&state.db) + .await?; + + for r in rows { + use sqlx::Row; + let id: i64 = r.get("id"); + let query: String = r.get("query"); + let source: String = r.get("source"); + let limit: i32 = r.get("limit_count"); + let last_run_str: String = r.get("last_run"); + + info!( + "[Scheduler] 开始增量同步任务: id={}, query='{}', last_run={}", + id, query, last_run_str + ); + + // 运行静默增量同步 + match MetaSync::run_incremental_sync_silent( + state.db.clone(), + Arc::new(state.sources.ads.clone()), + Arc::new(state.sources.arxiv.clone()), + query.clone(), + source.clone(), + limit, + &last_run_str, + ) + .await + { + Ok(count) => { + info!( + "[Scheduler] 增量同步任务成功完成: id={}, 新同步 {} 篇文献", + id, count + ); + // 更新该任务的 last_run 时间戳 + let _ = sqlx::query( + "UPDATE sync_queries SET last_run = CURRENT_TIMESTAMP WHERE id = ?", + ) + .bind(id) + .execute(&state.db) + .await; + } + Err(e) => { + error!("[Scheduler] 增量同步任务执行失败: id={}, error={}", id, e); + } + } + } + + Ok(()) } #[derive(Debug, Serialize, serde::Deserialize)] @@ -352,4 +510,5 @@ pub struct SavedSyncQuery { pub source: String, pub limit_count: i32, pub last_run: String, + pub is_scheduled: bool, } diff --git a/src/services/batch/mod.rs b/src/services/batch/mod.rs index 8fbe2eb..733dd23 100644 --- a/src/services/batch/mod.rs +++ b/src/services/batch/mod.rs @@ -3,4 +3,6 @@ pub mod asset; pub mod meta; pub use asset::{filter_papers_for_batch, AssetBatch, AssetBatchStatus, BatchAction}; -pub use meta::{MetaSync, MetaSyncStatus, SavedSyncQuery}; +pub use meta::{ + run_scheduled_sync_step, start_scheduled_sync_timer, MetaSync, MetaSyncStatus, SavedSyncQuery, +}; diff --git a/src/services/cds/vizier.rs b/src/services/cds/vizier.rs index 21449c0..6cd334d 100644 --- a/src/services/cds/vizier.rs +++ b/src/services/cds/vizier.rs @@ -184,6 +184,7 @@ pub fn escape_string_literal(s: &str) -> String { /// 锥形检索 /// /// strategy="nearest" 时按角距离排序，返回最近的源；strategy="all" 时无序返回全部。 +#[allow(clippy::too_many_arguments)] // 锥形检索物理参数 pub async fn cone_search( pool: &SqlitePool, client: &VizierClient, diff --git a/src/services/mod.rs b/src/services/mod.rs index c401c13..82700ea 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -1,3 +1,4 @@ +pub mod analysis; pub mod batch; pub mod cds; pub mod chunker; @@ -14,5 +15,7 @@ pub mod query_parser; pub mod rag; pub mod search; pub mod session; +pub mod spectrum; +pub mod timeseries; pub mod translation; pub mod vision; diff --git a/src/services/observation/cache.rs b/src/services/observation/cache.rs index 08bc1cf..99ea36b 100644 --- a/src/services/observation/cache.rs +++ b/src/services/observation/cache.rs @@ -103,6 +103,7 @@ pub async fn fetch_observation_cache( } /// 写入/更新缓存条目 +#[allow(clippy::too_many_arguments)] // 缓存条目字段一一对应 pub async fn write_observation_cache( pool: &SqlitePool, source: Source, diff --git a/src/services/observation/cutout.rs b/src/services/observation/cutout.rs new file mode 100644 index 0000000..e5e60d1 --- /dev/null +++ b/src/services/observation/cutout.rs @@ -0,0 +1,378 @@ +// src/services/observation/cutout.rs +// +// 图像 Cutout 服务 —— SDSS / DSS / Pan-STARRS cutout API +// +// 支持： +// - SDSS SkyServer cutout（光学 2.5m） +// - STScI DSS（Digitized Sky Survey，多种底片） +// - Pan-STARRS cutout（光学 1.8m） + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; +use tracing::info; + +/// Cutout 数据源 +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum CutoutSource { + /// SDSS optical + Sdss, + /// DSS (POSS-I, POSS-II, etc.) + Dss, + /// Pan-STARRS DR1 + Panstarrs, +} + +impl CutoutSource { + pub fn as_str(&self) -> &'static str { + match self { + Self::Sdss => "sdss", + Self::Dss => "dss", + Self::Panstarrs => "panstarrs", + } + } +} + +/// Cutout 结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CutoutResult { + /// 图像数据（JPEG/PNG 字节） + pub image_bytes: Vec, + /// 图像格式 + pub format: String, + /// 中心坐标 + pub ra: f64, + pub dec: f64, + /// cutout 大小（角秒） + pub size_arcsec: f64, + /// 数据源 + pub source: String, + /// 图像宽度（像素） + pub width: u32, + /// 图像高度（像素） + pub height: u32, +} + +/// SDSS cutout（SkyServer） +pub async fn sdss_cutout(ra: f64, dec: f64, size_arcsec: f64) -> Result { + let scale = size_arcsec / 14.0; // 14 arcsec/pixel for SDSS + let width = (size_arcsec / scale) as u32; + let url = format!( + "https://skyserver.sdss.org/dr18/SkyServerWS/ImgCutout/getjpeg?ra={}&dec={}&scale={:.4}&width={}&height={}", + ra, dec, scale, width, width + ); + info!( + "[Cutout] SDSS: ra={:.4}, dec={:.4}, size={}\"", + ra, dec, size_arcsec + ); + + let resp = reqwest::get(&url) + .await + .map_err(|e| anyhow!("SDSS cutout 请求失败: {}", e))?; + if !resp.status().is_success() { + return Err(anyhow!("SDSS cutout 失败: {}", resp.status())); + } + let bytes = resp.bytes().await?.to_vec(); + + Ok(CutoutResult { + image_bytes: bytes, + format: "jpeg".to_string(), + ra, + dec, + size_arcsec, + source: "SDSS".to_string(), + width, + height: width, + }) +} + +/// DSS cutout（STScI） +pub async fn dss_cutout(ra: f64, dec: f64, size_arcsec: f64) -> Result { + let size_deg = size_arcsec / 3600.0; + let url = format!( + "https://archive.stsci.edu/cgi-bin/dss_search?faste&ra={}&dec={}&equinox=J2000&epoch=2000&width={}&height={}&nimages=1&format=jpeg", + ra, dec, size_deg, size_deg + ); + info!( + "[Cutout] DSS: ra={:.4}, dec={:.4}, size={}\"", + ra, dec, size_arcsec + ); + + let resp = reqwest::get(&url) + .await + .map_err(|e| anyhow!("DSS cutout 请求失败: {}", e))?; + if !resp.status().is_success() { + return Err(anyhow!("DSS cutout 失败: {}", resp.status())); + } + let bytes = resp.bytes().await?.to_vec(); + + let pixel_size = size_arcsec / 1.0; // ~1 arcsec/pixel for DSS + let width = (size_arcsec / pixel_size) as u32; + + Ok(CutoutResult { + image_bytes: bytes, + format: "jpeg".to_string(), + ra, + dec, + size_arcsec, + source: "DSS".to_string(), + width, + height: width, + }) +} + +/// Pan-STARRS cutout（via VizieR / PS1 cutout service） +pub async fn panstarrs_cutout(ra: f64, dec: f64, size_arcsec: f64) -> Result { + let size_deg = size_arcsec / 3600.0; + // PS1 cutout service + let url = format!( + "https://ps1images.stsci.edu/cgi-bin/ps1cutoutservice?ra={}&dec={}&size={}&format=fits&filters=g,r,i,z,y", + ra, dec, size_deg + ); + info!( + "[Cutout] Pan-STARRS: ra={:.4}, dec={:.4}, size={}\"", + ra, dec, size_arcsec + ); + + let resp = reqwest::get(&url) + .await + .map_err(|e| anyhow!("PS1 cutout 请求失败: {}", e))?; + if !resp.status().is_success() { + return Err(anyhow!("PS1 cutout 失败: {}", resp.status())); + } + let bytes = resp.bytes().await?.to_vec(); + + Ok(CutoutResult { + image_bytes: bytes, + format: "fits".to_string(), + ra, + dec, + size_arcsec, + source: "Pan-STARRS".to_string(), + width: 0, + height: 0, + }) +} + +/// 统一 cutout 接口 +pub async fn get_cutout( + source: CutoutSource, + ra: f64, + dec: f64, + size_arcsec: f64, +) -> Result { + match source { + CutoutSource::Sdss => sdss_cutout(ra, dec, size_arcsec).await, + CutoutSource::Dss => dss_cutout(ra, dec, size_arcsec).await, + CutoutSource::Panstarrs => panstarrs_cutout(ra, dec, size_arcsec).await, + } +} + +/// Finding Chart 生成（基于 cutout + 标注） +pub async fn generate_finding_chart( + ra: f64, + dec: f64, + size_arcsec: f64, + source: CutoutSource, + _label: Option<&str>, +) -> Result { + let result = get_cutout(source, ra, dec, size_arcsec).await?; + + // 保存到文件 + let ext = match result.format.as_str() { + "jpeg" => "jpg", + "fits" => "fits", + _ => "png", + }; + let file_name = format!("finding_chart_{:.4}_{:.4}.{}", ra, dec, ext); + let file_path = format!("finding_charts/{}", file_name); + + let library_dir = std::env::var("LIBRARY_DIR").unwrap_or_else(|_| "library".to_string()); + let full_path = std::path::Path::new(&library_dir).join(&file_path); + if let Some(parent) = full_path.parent() { + tokio::fs::create_dir_all(parent).await.ok(); + } + tokio::fs::write(&full_path, &result.image_bytes).await.ok(); + + Ok(result) +} + +// ═══════════════════════════════════════════════════════════════ +// CutoutFetcher —— ObservationFetcher trait 实现 +// ═══════════════════════════════════════════════════════════════ + +use crate::api::AppState; +use crate::services::observation::cache::{ + cached_files_total_size, fetch_observation_cache, log_write_failure, persist_bytes, + write_observation_cache, +}; +use crate::services::observation::fetcher::{Candidate, ObservationFetcher}; +use crate::services::observation::types::{ + Artifact, ObservationProduct, ProductSpec, ProductType, Source, +}; +use tracing::warn; + +/// 图像 Cutout fetcher（SDSS/DSS/Pan-STARRS） +/// +/// 这不是一个标准的 ObservationFetcher（cutout 不通过 TAP/VO 检索）， +/// 而是通过 HTTP API 直接获取图像 cutout。 +/// 注册为 (Source::Panstarrs, ProductType::Image) 的 fetcher， +/// 坐标模式直接构造 cutout URL。 +#[derive(Debug)] +pub struct CutoutFetcher; + +#[async_trait::async_trait] +impl ObservationFetcher for CutoutFetcher { + fn key(&self) -> (Source, ProductType) { + (Source::Panstarrs, ProductType::Image) + } + + fn subtypes(&self) -> &'static [&'static str] { + &["sdss", "dss", "panstarrs"] + } + + fn suggested_max_radius_deg(&self) -> f64 { + 0.0 + } + + fn hard_max_radius_deg(&self) -> f64 { + 0.0 + } + + fn supports_coordinates(&self) -> bool { + true + } + + fn supports_identifiers(&self) -> bool { + false + } + + async fn cone_search_raw( + &self, + _state: &AppState, + ra: f64, + dec: f64, + _radius_deg: f64, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result> { + // Cutout 不做 cone search，直接返回目标本身 + Ok(vec![Candidate { + source: Source::Panstarrs, + source_id: format!("cutout_{:.4}_{:.4}", ra, dec), + label: format!("Cutout at ({:.4}, {:.4})", ra, dec), + ra: Some(ra), + dec: Some(dec), + distance: Some(0.0), + raw: None, + }]) + } + + async fn resolve_identifier( + &self, + _identifier: &str, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result { + Err(anyhow!("Cutout 不支持标识符模式")) + } + + async fn fetch( + &self, + state: &AppState, + candidate: &Candidate, + _release: Option<&str>, + subtype: Option<&str>, + _version: Option<&str>, + force: bool, + ) -> Result { + let ra = candidate.ra.unwrap_or(0.0); + let dec = candidate.dec.unwrap_or(0.0); + let sub = subtype.unwrap_or("sdss"); + let product = ProductSpec::with_subtype(ProductType::Image, sub); + let cache_key = format!("cutout_{:.4}_{:.4}_{}", ra, dec, sub); + let source_label = candidate.label.clone(); + + // 1) 检查缓存 + if !force { + if let Some((artifacts, meta)) = + fetch_observation_cache(&state.db, Source::Panstarrs, &product, &cache_key).await? + { + if cached_files_total_size(&state.config.storage.library_dir, &artifacts) + .await + .is_some() + { + info!("[Cutout] 缓存命中 (key={})", cache_key); + return Ok(ObservationProduct { + source: Source::Panstarrs, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: meta, + }); + } + warn!("[Cutout] 缓存文件缺失，重新下载 (key={})", cache_key); + } + } + + // 2) 真实获取 + let source = match sub { + "dss" => CutoutSource::Dss, + "panstarrs" => CutoutSource::Panstarrs, + _ => CutoutSource::Sdss, + }; + + let result = get_cutout(source, ra, dec, 120.0).await?; // 2 arcmin default + + let ext = match result.format.as_str() { + "jpeg" => "jpg", + "fits" => "fits", + _ => "png", + }; + let file_name = format!("cutout_{:.4}_{:.4}.{}", ra, dec, ext); + let file_path = format!("Telescope/cutout/{}/{}", candidate.source_id, file_name); + persist_bytes( + &state.config.storage.library_dir, + &file_path, + &result.image_bytes, + ) + .await?; + + let artifacts = vec![Artifact { + band: None, + original_name: Some(file_name), + file_path: file_path.clone(), + file_url: crate::services::observation::cache::file_url_from_path(&file_path), + file_format: ext.to_string(), + size_bytes: result.image_bytes.len(), + cached: false, + }]; + + // 3) 写入缓存 + if let Err(e) = write_observation_cache( + &state.db, + Source::Panstarrs, + &product, + &cache_key, + Some(ra), + Some(dec), + &artifacts, + None, + ) + .await + { + log_write_failure("panstarrs", "image", e); + } + + Ok(ObservationProduct { + source: Source::Panstarrs, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: candidate.raw.clone(), + }) + } +} diff --git a/src/services/observation/dispatch.rs b/src/services/observation/dispatch.rs index 88a78b3..119e156 100644 --- a/src/services/observation/dispatch.rs +++ b/src/services/observation/dispatch.rs @@ -21,6 +21,7 @@ use tracing::{info, warn}; /// 这一步，把命中的候选源列表返回给调用方预览，由调用方决定是否下载、下载哪些。 /// /// 标识符模式无需"检索"——标识符本身就是确定源，调用方直接走 download 即可。 +#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数 pub async fn search_observation( state: &AppState, registry: &ObservationRegistry, @@ -100,6 +101,7 @@ pub async fn download_observation( } /// 按坐标：cone 检索 → 选源 → 逐个下载 +#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数 async fn download_by_coordinates( state: &AppState, registry: &ObservationRegistry, @@ -176,6 +178,7 @@ async fn download_by_coordinates( } /// 按标识符：解析各源标识字符串 → 逐个下载（跳过 cone 检索） +#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数 async fn download_by_identifiers( state: &AppState, registry: &ObservationRegistry, @@ -249,6 +252,7 @@ async fn download_by_identifiers( } /// 便捷：从单个标识符构造请求并下载（供 CLI / 简单场景使用） +#[allow(clippy::too_many_arguments)] // 透传 (state, registry, source, product, 版本) 的管线参数 pub async fn download_one( state: &AppState, registry: &ObservationRegistry, diff --git a/src/services/observation/fetcher.rs b/src/services/observation/fetcher.rs index b791b48..1fed3e0 100644 --- a/src/services/observation/fetcher.rs +++ b/src/services/observation/fetcher.rs @@ -57,6 +57,9 @@ pub struct Candidate { /// /// 缓存、重试、选源等通用编排由 trait 默认方法和 dispatch.rs 负责。 #[async_trait] +// cone 检索模板方法需透传 (state, 坐标, 半径, release, subtype, version)， +// async_trait 生成的装箱垫片同样超过参数阈值 +#[allow(clippy::too_many_arguments)] pub trait ObservationFetcher: Send + Sync + std::fmt::Debug { /// 该 fetcher 处理的 (源, 产品类型) 组合 fn key(&self) -> (Source, ProductType); @@ -183,6 +186,7 @@ pub trait ObservationFetcher: Send + Sync + std::fmt::Debug { /// /// 复用 vizier_query_cache 表（与历史设计一致），用 cache_prefix 区分各源各产品。 /// 缓存命中反序列化为 Vec；未命中调 cone_search_raw 并写回。 + #[allow(clippy::too_many_arguments)] // 模板方法透传 cone 检索与版本参数 async fn cone_search_cached( &self, state: &AppState, diff --git a/src/services/observation/jwst.rs b/src/services/observation/jwst.rs new file mode 100644 index 0000000..a74c9f8 --- /dev/null +++ b/src/services/observation/jwst.rs @@ -0,0 +1,212 @@ +// src/services/observation/jwst.rs +// +// JWST 光谱 fetcher —— 通过 MAST API 获取 NIRSpec/MIRI 光谱 + +use crate::api::AppState; +use crate::services::observation::cache::{ + cached_files_total_size, fetch_observation_cache, log_write_failure, persist_bytes, + write_observation_cache, +}; +use crate::services::observation::fetcher::{Candidate, ObservationFetcher}; +use crate::services::observation::types::{ + Artifact, ObservationProduct, ProductSpec, ProductType, Source, +}; +use anyhow::{anyhow, Result}; +use tracing::{info, warn}; + +#[derive(Debug)] +pub struct JwstSpectrumFetcher; + +#[async_trait::async_trait] +impl ObservationFetcher for JwstSpectrumFetcher { + fn key(&self) -> (Source, ProductType) { + (Source::Jwst, ProductType::Spectrum) + } + + fn subtypes(&self) -> &'static [&'static str] { + &["nirspec_ifu", "nirspec_mos", "miri_mrs", "miri_lrs"] + } + + fn suggested_max_radius_deg(&self) -> f64 { + 0.1 + } + + fn hard_max_radius_deg(&self) -> f64 { + 1.0 + } + + fn identifier_format(&self) -> Option<&'static str> { + Some("JWST 观测 ID（如 'jw00001001001'）") + } + + async fn cone_search_raw( + &self, + state: &AppState, + ra: f64, + dec: f64, + radius_deg: f64, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result> { + let jwst = &state.sources.jwst; + let rows = jwst.cone_search(ra, dec, radius_deg).await?; + + Ok(rows + .into_iter() + .map(|r| Candidate { + source: Source::Jwst, + source_id: r.obs_id.clone(), + label: format!("{} ({})", r.target_name, r.instrument), + ra: Some(r.ra), + dec: Some(r.dec), + distance: None, + raw: Some(serde_json::json!({ + "target_name": r.target_name, + "instrument": r.instrument, + "observation_mode": r.observation_mode, + "filter_grating": r.filter_grating, + "exposure_time": r.exposure_time, + })), + }) + .collect()) + } + + async fn resolve_identifier( + &self, + identifier: &str, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result { + Ok(Candidate { + source: Source::Jwst, + source_id: identifier.to_string(), + label: identifier.to_string(), + ra: None, + dec: None, + distance: None, + raw: None, + }) + } + + async fn fetch( + &self, + state: &AppState, + candidate: &Candidate, + _release: Option<&str>, + subtype: Option<&str>, + _version: Option<&str>, + force: bool, + ) -> Result { + let sub = subtype.unwrap_or("nirspec_ifu"); + let product = ProductSpec::with_subtype(ProductType::Spectrum, sub); + let cache_key = format!("jwst|{}|{}", candidate.source_id, sub); + let source_label = candidate.label.clone(); + + // 1) 检查缓存 + if !force { + if let Some((artifacts, meta)) = + fetch_observation_cache(&state.db, Source::Jwst, &product, &cache_key).await? + { + if cached_files_total_size(&state.config.storage.library_dir, &artifacts) + .await + .is_some() + { + info!("[JWST] 缓存命中 (obs_id={})", candidate.source_id); + return Ok(ObservationProduct { + source: Source::Jwst, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: meta, + }); + } + warn!( + "[JWST] 缓存文件缺失，重新下载 (obs_id={})", + candidate.source_id + ); + } + } + + // 2) 真实获取 + let jwst = &state.sources.jwst; + let products = jwst.list_data_products(&candidate.source_id).await?; + + // 筛选匹配 subtype 的产品 + let target_product = match sub { + "nirspec_ifu" => "NRS1", + "nirspec_mos" => "NRS1", + "miri_mrs" => "MRS", + "miri_lrs" => "SLITLOS", + _ => "NRS1", + }; + + let matching: Vec<_> = products + .iter() + .filter(|p| { + p.product_type.contains(target_product) + || p.file_name.contains("s2d") + || p.file_name.contains("x1d") + }) + .collect(); + + if matching.is_empty() { + return Err(anyhow!( + "JWST 观测 {} 无匹配的光谱产品", + candidate.source_id + )); + } + + let product_item = &matching[0]; + let bytes = reqwest::get(&product_item.file_url) + .await + .map_err(|e| anyhow!("JWST 下载失败: {}", e))? + .bytes() + .await + .map_err(|e| anyhow!("JWST 读取失败: {}", e))? + .to_vec(); + + let file_path = format!( + "Telescope/JWST/spectrum/{}/{}", + candidate.source_id, product_item.file_name + ); + persist_bytes(&state.config.storage.library_dir, &file_path, &bytes).await?; + + let artifacts = vec![Artifact { + band: None, + original_name: Some(product_item.file_name.clone()), + file_path: file_path.clone(), + file_url: product_item.file_url.clone(), + file_format: "fits".to_string(), + size_bytes: bytes.len(), + cached: false, + }]; + + // 3) 写入缓存 + if let Err(e) = write_observation_cache( + &state.db, + Source::Jwst, + &product, + &cache_key, + candidate.ra, + candidate.dec, + &artifacts, + None, + ) + .await + { + log_write_failure("jwst", "spectrum", e); + } + + Ok(ObservationProduct { + source: Source::Jwst, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: candidate.raw.clone(), + }) + } +} diff --git a/src/services/observation/mod.rs b/src/services/observation/mod.rs index 196ac70..0e3505e 100644 --- a/src/services/observation/mod.rs +++ b/src/services/observation/mod.rs @@ -27,11 +27,13 @@ // 3. 在 types.rs Source/ProductType 枚举加变体（若为新源/新产品） pub mod cache; +pub mod cutout; pub mod desi; pub mod dispatch; pub mod fetcher; pub mod gaia; pub mod gaia_xp; +pub mod jwst; pub mod lamost; pub mod photometry; pub mod preview; @@ -40,6 +42,7 @@ pub mod sdss; pub mod tess; pub mod types; pub mod unified; +pub mod xray; pub mod ztf; // 统一入口 re-export diff --git a/src/services/observation/photometry.rs b/src/services/observation/photometry.rs index 7b3f271..e542791 100644 --- a/src/services/observation/photometry.rs +++ b/src/services/observation/photometry.rs @@ -769,8 +769,8 @@ mod tests { #[test] fn test_source_from_str_aliases() { - assert_eq!(Source::from_str("2mass").unwrap(), Source::Twomass); - assert_eq!(Source::from_str("wise").unwrap(), Source::Allwise); - assert_eq!(Source::from_str("ps1").unwrap(), Source::Panstarrs); + assert_eq!(Source::parse("2mass").unwrap(), Source::Twomass); + assert_eq!(Source::parse("wise").unwrap(), Source::Allwise); + assert_eq!(Source::parse("ps1").unwrap(), Source::Panstarrs); } } diff --git a/src/services/observation/preview.rs b/src/services/observation/preview.rs index 796bc94..5bba4b7 100644 --- a/src/services/observation/preview.rs +++ b/src/services/observation/preview.rs @@ -197,6 +197,35 @@ pub async fn build_preview( let mut preview = parse_by_source(&abs_path, source, subtype, source_id, &artifact.file_path) .with_context(|| format!("解析 FITS 失败: {}", abs_path.display()))?; + + // 尝试读取 FITS header 并补充 teff, logg, fe_h + if let Ok(mut fits) = fitsio::FitsFile::open(&abs_path) { + if preview.meta.teff.is_none() { + preview.meta.teff = + crate::services::spectrum::parameters::read_f64_with_fallback( + &mut fits, + &["TEFF", "T_EFF", "TEMP", "EFFECTIVE_TEMP", "TEFF_ASPCAP"], + ) + .map(|v| v as f32); + } + if preview.meta.logg.is_none() { + preview.meta.logg = + crate::services::spectrum::parameters::read_f64_with_fallback( + &mut fits, + &["LOGG", "LOG_G", "GRAVITY", "LOGG_ASPCAP"], + ) + .map(|v| v as f32); + } + if preview.meta.fe_h.is_none() { + preview.meta.fe_h = + crate::services::spectrum::parameters::read_f64_with_fallback( + &mut fits, + &["FEH", "FE_H", "[Fe/H]", "FEH_ASPCAP"], + ) + .map(|v| v as f32); + } + } + for seg in preview.segments.iter_mut() { if seg.wavelength.len() > MAX_PREVIEW_POINTS { downsample_segment(seg, MAX_PREVIEW_POINTS); @@ -240,6 +269,10 @@ fn parse_by_source( Source::Twomass | Source::Allwise | Source::Panstarrs | Source::Ztf | Source::Tess => Err( anyhow!("{} 的预览解析暂未实现（测光/光变产品）", source.display()), ), + // JWST/X射线产品的预览解析器后续实现 + Source::Jwst | Source::Xmm | Source::Chandra => { + Err(anyhow!("{} 的预览解析暂未实现", source.display())) + } } } @@ -831,7 +864,7 @@ mod tests { &seg.wavelength[..10.min(seg.wavelength.len())] ); println!("First 10 fluxes: {:?}", &seg.flux[..10.min(seg.flux.len())]); - assert!(seg.wavelength.len() > 0); + assert!(!seg.wavelength.is_empty()); } #[test] @@ -864,7 +897,7 @@ mod tests { "APOGEE First 10 fluxes: {:?}", &seg.flux[..10.min(seg.flux.len())] ); - assert!(seg.wavelength.len() > 0); + assert!(!seg.wavelength.is_empty()); } #[test] diff --git a/src/services/observation/registry.rs b/src/services/observation/registry.rs index 04a8a2b..3530346 100644 --- a/src/services/observation/registry.rs +++ b/src/services/observation/registry.rs @@ -10,9 +10,11 @@ // 注意：SDSS 的 specobj 与 APOGEE 共用 (Sdss, Spectrum) key，但 subtype 不同。 // registry 用 (Source, ProductType) → Vec>，dispatch 再按 subtype 二级筛选。 +use crate::services::observation::cutout::CutoutFetcher; use crate::services::observation::desi::DesiSpectrumFetcher; use crate::services::observation::fetcher::ObservationFetcher; use crate::services::observation::gaia::{GaiaLightCurveFetcher, GaiaSpectrumFetcher}; +use crate::services::observation::jwst::JwstSpectrumFetcher; use crate::services::observation::lamost::LamostSpectrumFetcher; use crate::services::observation::photometry::{ AllwisePhotometryFetcher, GaiaPhotometryFetcher, PanstarrsPhotometryFetcher, @@ -21,6 +23,7 @@ use crate::services::observation::photometry::{ use crate::services::observation::sdss::{SdssApogeeFetcher, SdssSpectrumFetcher}; use crate::services::observation::tess::TessLightCurveFetcher; use crate::services::observation::types::{ProductType, Source}; +use crate::services::observation::xray::{ChandraSpectrumFetcher, XmmSpectrumFetcher}; use crate::services::observation::ztf::ZtfLightCurveFetcher; use anyhow::{anyhow, Result}; use std::collections::HashMap; @@ -31,9 +34,9 @@ pub struct ObservationRegistry { fetchers: HashMap<(Source, ProductType), Vec>>, } -impl ObservationRegistry { +impl Default for ObservationRegistry { /// 构建默认注册表（项目启动时调用一次） - pub fn default() -> Self { + fn default() -> Self { let mut fetchers: HashMap<(Source, ProductType), Vec>> = HashMap::new(); @@ -63,10 +66,19 @@ impl ObservationRegistry { // 光变曲线（ZTF / TESS，新 HTTP 客户端） register!(ZtfLightCurveFetcher); register!(TessLightCurveFetcher); + // JWST 光谱 + register!(JwstSpectrumFetcher); + // X 射线光谱（XMM-Newton / Chandra） + register!(XmmSpectrumFetcher); + register!(ChandraSpectrumFetcher); + // 图像 Cutout（SDSS / DSS / Pan-STARRS） + register!(CutoutFetcher); Self { fetchers } } +} +impl ObservationRegistry { /// 按 (Source, ProductType) + subtype 取 fetcher /// /// subtype 匹配规则： diff --git a/src/services/observation/types.rs b/src/services/observation/types.rs index 7a0705c..c0a7a94 100644 --- a/src/services/observation/types.rs +++ b/src/services/observation/types.rs @@ -34,6 +34,12 @@ pub enum Source { Ztf, /// TESS 空间时域光变曲线（MAST） Tess, + /// JWST（James Webb Space Telescope）NIRSpec/MIRI 光谱 + Jwst, + /// XMM-Newton X 射线光谱 + Xmm, + /// Chandra X 射线光谱 + Chandra, } impl Source { @@ -48,6 +54,9 @@ impl Source { "panstarrs", "ztf", "tess", + "jwst", + "xmm", + "chandra", ] } pub fn as_str(&self) -> &'static str { @@ -61,12 +70,15 @@ impl Source { Self::Panstarrs => "panstarrs", Self::Ztf => "ztf", Self::Tess => "tess", + Self::Jwst => "jwst", + Self::Xmm => "xmm", + Self::Chandra => "chandra", } } /// 从字符串解析数据源（大小写不敏感）。新增源只需在此加一行 match arm， /// 所有调用方（API handler、Agent tool）统一用本方法，不再各自手写 match。 - pub fn from_str(s: &str) -> Result { + pub fn parse(s: &str) -> Result { Ok(match s.to_lowercase().as_str() { "lamost" => Self::Lamost, "gaia" => Self::Gaia, @@ -77,6 +89,9 @@ impl Source { "panstarrs" | "ps1" | "pan-starrs" => Self::Panstarrs, "ztf" => Self::Ztf, "tess" => Self::Tess, + "jwst" => Self::Jwst, + "xmm" | "xmm-newton" => Self::Xmm, + "chandra" => Self::Chandra, other => { return Err(format!( "不支持的 source '{}'，可选: {:?}", @@ -98,10 +113,21 @@ impl Source { Self::Panstarrs => "Pan-STARRS", Self::Ztf => "ZTF", Self::Tess => "TESS", + Self::Jwst => "JWST", + Self::Xmm => "XMM-Newton", + Self::Chandra => "Chandra", } } } +impl std::str::FromStr for Source { + type Err = String; + + fn from_str(s: &str) -> Result { + Self::parse(s) + } +} + // ═══════════════════════════════════════════════════════════════ // 2. 第二轴：观测数据产品类型 // ═══════════════════════════════════════════════════════════════ @@ -136,7 +162,7 @@ impl ProductType { /// 从字符串解析产品类型（大小写不敏感，接受常见别名）。 /// 别名集中在此一处管理，所有调用方统一用本方法。 - pub fn from_str(s: &str) -> Result { + pub fn parse(s: &str) -> Result { Ok(match s.to_lowercase().as_str() { "spectrum" | "spec" => Self::Spectrum, "lightcurve" | "light_curve" | "lc" => Self::LightCurve, @@ -167,6 +193,14 @@ impl ProductType { } } +impl std::str::FromStr for ProductType { + type Err = String; + + fn from_str(s: &str) -> Result { + Self::parse(s) + } +} + /// 产品细分（可选）：用 (类型, 子类型) 描述具体产品变体 /// /// 子类型语义因源而异： diff --git a/src/services/observation/unified.rs b/src/services/observation/unified.rs index acd3454..21b9d2d 100644 --- a/src/services/observation/unified.rs +++ b/src/services/observation/unified.rs @@ -38,6 +38,15 @@ fn default_radius_deg() -> f64 { 0.1 } +/// 候选源分组键：(source, product, subtype, release, version) +type GroupKey = ( + String, + String, + Option, + Option, + Option, +); + // ═══════════════════════════════════════════════════════════════ // 1. 请求/响应类型 // ═══════════════════════════════════════════════════════════════ @@ -198,16 +207,7 @@ pub async fn unified_search( // 按 (source, product, release, version) 分组合并候选源 // 这里用 (source_str, product_str, subtype_str, release, version) 作为排序键 - let mut groups_map: BTreeMap< - ( - String, - String, - Option, - Option, - Option, - ), - Vec, - > = BTreeMap::new(); + let mut groups_map: BTreeMap> = BTreeMap::new(); let mut total_candidates = 0usize; let mut fail_count = 0usize; @@ -242,10 +242,10 @@ pub async fn unified_search( // 转换为 SourceCandidateGroup 列表（BTreeMap 已保证按 key 排序） let mut groups: Vec = Vec::with_capacity(groups_map.len()); for ((source_str, product_str, subtype, release, version), candidates) in groups_map { - let source = Source::from_str(&source_str) + let source = Source::parse(&source_str) .map_err(|e| anyhow::anyhow!("内部错误：无效 source 回填 '{}': {}", source_str, e))?; let product = ProductSpec { - product: crate::services::observation::types::ProductType::from_str(&product_str) + product: crate::services::observation::types::ProductType::parse(&product_str) .map_err(|e| { anyhow::anyhow!("内部错误：无效 product 回填 '{}': {}", product_str, e) })?, diff --git a/src/services/observation/xray.rs b/src/services/observation/xray.rs new file mode 100644 index 0000000..1ced05a --- /dev/null +++ b/src/services/observation/xray.rs @@ -0,0 +1,357 @@ +// src/services/observation/xray.rs +// +// X 射线光谱 fetcher —— XMM-Newton / Chandra 通过 HEASARC TAP + +use crate::api::AppState; +use crate::services::observation::cache::{ + cached_files_total_size, fetch_observation_cache, log_write_failure, persist_bytes, + write_observation_cache, +}; +use crate::services::observation::fetcher::{Candidate, ObservationFetcher}; +use crate::services::observation::types::{ + Artifact, ObservationProduct, ProductSpec, ProductType, Source, +}; +use anyhow::{anyhow, Result}; +use tracing::{info, warn}; + +#[derive(Debug)] +pub struct XmmSpectrumFetcher; + +#[async_trait::async_trait] +impl ObservationFetcher for XmmSpectrumFetcher { + fn key(&self) -> (Source, ProductType) { + (Source::Xmm, ProductType::Spectrum) + } + + fn subtypes(&self) -> &'static [&'static str] { + &["epic-pn", "epic-mos1", "epic-mos2", "rgs"] + } + + fn suggested_max_radius_deg(&self) -> f64 { + 0.5 + } + + fn identifier_format(&self) -> Option<&'static str> { + Some("XMM 观测 ID（数字）") + } + + async fn cone_search_raw( + &self, + state: &AppState, + ra: f64, + dec: f64, + radius_deg: f64, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result> { + let heasarc = &state.sources.heasarc; + let rows = heasarc.xmm_cone_search(ra, dec, radius_deg).await?; + + Ok(rows + .into_iter() + .map(|r| Candidate { + source: Source::Xmm, + source_id: r.obs_id.clone(), + label: format!("{} ({})", r.target_name, r.instrument), + ra: Some(r.ra), + dec: Some(r.dec), + distance: None, + raw: Some(serde_json::json!({ + "target_name": r.target_name, + "instrument": r.instrument, + "exposure": r.exposure, + })), + }) + .collect()) + } + + async fn resolve_identifier( + &self, + identifier: &str, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result { + Ok(Candidate { + source: Source::Xmm, + source_id: identifier.to_string(), + label: identifier.to_string(), + ra: None, + dec: None, + distance: None, + raw: None, + }) + } + + async fn fetch( + &self, + state: &AppState, + candidate: &Candidate, + _release: Option<&str>, + subtype: Option<&str>, + _version: Option<&str>, + force: bool, + ) -> Result { + let instrument = subtype.unwrap_or("epic-pn"); + let product = ProductSpec::with_subtype(ProductType::Spectrum, instrument); + let cache_key = format!("xray|{}|{}", candidate.source_id, instrument); + let source_label = candidate.label.clone(); + + // 1) 检查缓存 + if !force { + if let Some((artifacts, meta)) = + fetch_observation_cache(&state.db, Source::Xmm, &product, &cache_key).await? + { + if cached_files_total_size(&state.config.storage.library_dir, &artifacts) + .await + .is_some() + { + info!("[XMM] 缓存命中 (obs_id={})", candidate.source_id); + return Ok(ObservationProduct { + source: Source::Xmm, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: meta, + }); + } + warn!( + "[XMM] 缓存文件缺失，重新下载 (obs_id={})", + candidate.source_id + ); + } + } + + // XMM 数据通过 HEASARC FTP 获取 + let file_name = format!("{}_{}.fits", candidate.source_id, instrument); + let file_path = format!( + "Telescope/XMM/spectrum/{}/{}", + candidate.source_id, file_name + ); + let data_url = format!( + "https://heasarc.gsfc.nasa.gov/FTP/heasarc/xmm/data/obs/{}/{}", + candidate.source_id, file_name + ); + + // 下载 FITS + let resp = reqwest::get(&data_url) + .await + .map_err(|e| anyhow!("XMM 下载失败: {}", e))?; + if !resp.status().is_success() { + return Err(anyhow!("XMM 下载失败: {}", resp.status())); + } + let bytes = resp + .bytes() + .await + .map_err(|e| anyhow!("XMM 读取失败: {}", e))? + .to_vec(); + persist_bytes(&state.config.storage.library_dir, &file_path, &bytes).await?; + + let artifacts = vec![Artifact { + band: None, + original_name: Some(file_name), + file_path: file_path.clone(), + file_url: data_url, + file_format: "fits".to_string(), + size_bytes: bytes.len(), + cached: false, + }]; + + // 3) 写入缓存 + if let Err(e) = write_observation_cache( + &state.db, + Source::Xmm, + &product, + &cache_key, + candidate.ra, + candidate.dec, + &artifacts, + None, + ) + .await + { + log_write_failure("xmm", "spectrum", e); + } + + Ok(ObservationProduct { + source: Source::Xmm, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: candidate.raw.clone(), + }) + } +} + +#[derive(Debug)] +pub struct ChandraSpectrumFetcher; + +#[async_trait::async_trait] +impl ObservationFetcher for ChandraSpectrumFetcher { + fn key(&self) -> (Source, ProductType) { + (Source::Chandra, ProductType::Spectrum) + } + + fn subtypes(&self) -> &'static [&'static str] { + &["acis-s", "acis-i", "hetg", "letg"] + } + + fn suggested_max_radius_deg(&self) -> f64 { + 0.5 + } + + fn identifier_format(&self) -> Option<&'static str> { + Some("Chandra 观测 ID（数字）") + } + + async fn cone_search_raw( + &self, + state: &AppState, + ra: f64, + dec: f64, + radius_deg: f64, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result> { + let heasarc = &state.sources.heasarc; + let rows = heasarc.chandra_cone_search(ra, dec, radius_deg).await?; + + Ok(rows + .into_iter() + .map(|r| Candidate { + source: Source::Chandra, + source_id: r.obs_id.clone(), + label: format!("{} ({})", r.target_name, r.instrument), + ra: Some(r.ra), + dec: Some(r.dec), + distance: None, + raw: Some(serde_json::json!({ + "target_name": r.target_name, + "instrument": r.instrument, + "exposure": r.exposure, + })), + }) + .collect()) + } + + async fn resolve_identifier( + &self, + identifier: &str, + _release: Option<&str>, + _subtype: Option<&str>, + _version: Option<&str>, + ) -> Result { + Ok(Candidate { + source: Source::Chandra, + source_id: identifier.to_string(), + label: identifier.to_string(), + ra: None, + dec: None, + distance: None, + raw: None, + }) + } + + async fn fetch( + &self, + state: &AppState, + candidate: &Candidate, + _release: Option<&str>, + subtype: Option<&str>, + _version: Option<&str>, + force: bool, + ) -> Result { + let instrument = subtype.unwrap_or("acis-s"); + let product = ProductSpec::with_subtype(ProductType::Spectrum, instrument); + let cache_key = format!("xray|{}|{}", candidate.source_id, instrument); + let source_label = candidate.label.clone(); + + // 1) 检查缓存 + if !force { + if let Some((artifacts, meta)) = + fetch_observation_cache(&state.db, Source::Chandra, &product, &cache_key).await? + { + if cached_files_total_size(&state.config.storage.library_dir, &artifacts) + .await + .is_some() + { + info!("[Chandra] 缓存命中 (obs_id={})", candidate.source_id); + return Ok(ObservationProduct { + source: Source::Chandra, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: meta, + }); + } + warn!( + "[Chandra] 缓存文件缺失，重新下载 (obs_id={})", + candidate.source_id + ); + } + } + + let file_name = format!("{}_repro.pha", candidate.source_id); + let file_path = format!( + "Telescope/Chandra/spectrum/{}/{}", + candidate.source_id, file_name + ); + let data_url = format!( + "https://heasarc.gsfc.nasa.gov/FTP/heasarc/chandra/data/reproc/{}/combined/{}", + candidate.source_id, file_name + ); + + let resp = reqwest::get(&data_url) + .await + .map_err(|e| anyhow!("Chandra 下载失败: {}", e))?; + if !resp.status().is_success() { + return Err(anyhow!("Chandra 下载失败: {}", resp.status())); + } + let bytes = resp + .bytes() + .await + .map_err(|e| anyhow!("Chandra 读取失败: {}", e))? + .to_vec(); + persist_bytes(&state.config.storage.library_dir, &file_path, &bytes).await?; + + let artifacts = vec![Artifact { + band: None, + original_name: Some(file_name), + file_path: file_path.clone(), + file_url: data_url, + file_format: "fits".to_string(), + size_bytes: bytes.len(), + cached: false, + }]; + + // 3) 写入缓存 + if let Err(e) = write_observation_cache( + &state.db, + Source::Chandra, + &product, + &cache_key, + candidate.ra, + candidate.dec, + &artifacts, + None, + ) + .await + { + log_write_failure("chandra", "spectrum", e); + } + + Ok(ObservationProduct { + source: Source::Chandra, + product, + source_id: cache_key, + source_label, + artifacts, + source_meta: candidate.raw.clone(), + }) + } +} diff --git a/src/services/paper/db.rs b/src/services/paper/db.rs index c552dbf..7d3f045 100644 --- a/src/services/paper/db.rs +++ b/src/services/paper/db.rs @@ -579,3 +579,40 @@ pub async fn mark_no_resource_service( let updated_paper = get_paper_from_db(db, library_dir, bibcode).await?; Ok(updated_paper) } + +/// 获取最近入库的文献列表 +pub async fn get_recent_papers( + db: &SqlitePool, + library_dir: &Path, + query: &str, + days_back: i64, + limit: i32, +) -> Result, sqlx::Error> { + let search_pattern = format!("%{}%", query); + let time_limit = format!("-{} days", days_back); + + let rows = sqlx::query( + "SELECT bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, \ + citation_count, reference_count, pdf_path, html_path, markdown_path, \ + translation_path, doctype, \ + EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode) AS has_vector \ + FROM papers \ + WHERE (title LIKE ? OR abstract LIKE ? OR authors LIKE ?) \ + AND created_at >= datetime('now', ?) \ + ORDER BY created_at DESC \ + LIMIT ?" + ) + .bind(&search_pattern) + .bind(&search_pattern) + .bind(&search_pattern) + .bind(&time_limit) + .bind(limit) + .fetch_all(db) + .await?; + + let papers = rows + .into_iter() + .map(|r| parse_paper_row(&r, library_dir)) + .collect(); + Ok(papers) +} diff --git a/src/services/paper/mod.rs b/src/services/paper/mod.rs index 0eb087e..fc29e94 100644 --- a/src/services/paper/mod.rs +++ b/src/services/paper/mod.rs @@ -9,8 +9,8 @@ pub mod reader; pub use db::{ check_paper_paths_in_db, clean_identifier, get_library_list, get_paper_from_db, - mark_no_resource_service, parse_paper_row, save_paper_to_db, save_paper_to_db_tx, CleanId, - LibraryQueryParams, SQLITE_PARAM_LIMIT, + get_recent_papers, mark_no_resource_service, parse_paper_row, save_paper_to_db, + save_paper_to_db_tx, CleanId, LibraryQueryParams, SQLITE_PARAM_LIMIT, }; pub use ingest::upload_paper_file_service; pub use model::{convert_ads_doc_to_standard, convert_arxiv_to_standard, StandardPaper}; diff --git a/src/services/search.rs b/src/services/search.rs index ae0c1fe..a6d43a7 100644 --- a/src/services/search.rs +++ b/src/services/search.rs @@ -232,6 +232,9 @@ pub async fn search_local_library( ) -> anyhow::Result> { // FTS5 查询净化，防止操作符注入 let sanitized = sanitize_fts5_query(query); + if sanitized.trim().is_empty() { + return Ok(Vec::new()); + } let fts_query = format!("\"{}\"", sanitized); let rows = sqlx::query_as::<_, PaperRow>( @@ -319,6 +322,9 @@ pub async fn search_agent_history( ) -> Result, sqlx::Error> { let mut results = Vec::new(); let sanitized = sanitize_fts5_query(q); + if sanitized.trim().is_empty() { + return Ok(Vec::new()); + } // 搜索会话 if scope == "all" || scope == "sessions" { diff --git a/src/services/spectrum/cross_correlate.rs b/src/services/spectrum/cross_correlate.rs new file mode 100644 index 0000000..dd01520 --- /dev/null +++ b/src/services/spectrum/cross_correlate.rs @@ -0,0 +1,562 @@ +// src/services/spectrum/cross_correlate.rs +// +// 交叉相关法测径向速度（RV） +// +// 算法： +// 1. 将观测光谱和模板光谱对齐到同一波长网格（对数重采样） +// 2. 归一化（除以连续谱估计值） +// 3. 计算交叉相关函数（CCF） +// 4. 从 CCF 峰值位置提取 RV + 不确定度 +// +// 支持： +// - 内置 Pickles 恒星类型模板（简化版：按光谱型插值） +// - 用户自定义模板 +// - 多波段光谱 + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +/// RV 测量结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RadialVelocityResult { + /// 径向速度 (km/s) + pub rv: f64, + /// RV 不确定度 (km/s) + pub rv_error: f64, + /// 交叉相关峰值（1=完美匹配） + pub ccf_peak: f64, + /// 使用的模板名称 + pub template_name: String, + /// 有效波长范围 (Å) + pub wavelength_range: (f64, f64), + /// SNR（估计） + pub snr: Option, +} + +/// 内置模板类型（简化 Pickles 恒星类型） +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum TemplateType { + /// O 型星 + O, + /// B 型星 + B, + /// A 型星 + A, + /// F 型星 + F, + /// G 型星（类太阳） + G, + /// K 型星 + K, + /// M 型星 + M, + /// 廉谱（M-dwarf） + Mwarf, + /// 巨星 + Giant, +} + +/// 获取模板的典型有效温度（用于标记） +pub fn template_teff(template: TemplateType) -> f64 { + match template { + TemplateType::O => 35000.0, + TemplateType::B => 20000.0, + TemplateType::A => 8500.0, + TemplateType::F => 6500.0, + TemplateType::G => 5778.0, + TemplateType::K => 4500.0, + TemplateType::M => 3200.0, + TemplateType::Mwarf => 3000.0, + TemplateType::Giant => 4800.0, + } +} + +/// 生成简化模板光谱（基于黑体辐射 + 简单吸收特征） +/// +/// 这是一个简化实现：用黑体辐射曲线 + 关键吸收线的高斯凹陷模拟恒星光谱。 +/// 真实应用中应使用 PHOENIX/BT-Settl 等理论光谱库。 +pub fn generate_template(template: TemplateType, wavelength: &[f64]) -> Vec { + let teff = template_teff(template); + let mut flux = Vec::with_capacity(wavelength.len()); + + for &w in wavelength { + // 简化黑体辐射（普朗克函数，任意缩放） + let w_um = w * 1e-4; // Å → μm + let planck = black_body(w_um, teff); + flux.push(planck); + } + + // 添加简化吸收特征 + let absorption_lines: Vec<(f64, f64)> = match template { + TemplateType::O | TemplateType::B => vec![ + (4861.0, 0.3), // Hβ + (4340.0, 0.25), // Hγ + (4102.0, 0.2), // Hδ + (4471.0, 0.15), // He I + (4026.0, 0.1), // He I + ], + TemplateType::A => vec![ + (6563.0, 0.5), // Hα + (4861.0, 0.6), // Hβ + (4340.0, 0.5), // Hγ + (4102.0, 0.4), // Hδ + ], + TemplateType::F => vec![ + (6563.0, 0.3), // Hα + (4861.0, 0.35), // Hβ + (3934.0, 0.2), // Ca II K + (5184.0, 0.15), // Mg b + ], + TemplateType::G => vec![ + (6563.0, 0.2), // Hα + (3934.0, 0.5), // Ca II K + (3968.0, 0.45), // Ca II H + (5890.0, 0.2), // Na D + (5270.0, 0.15), // Fe I + (5328.0, 0.12), // Fe I + ], + TemplateType::K | TemplateType::Giant => vec![ + (3934.0, 0.6), // Ca II K + (3968.0, 0.55), // Ca II H + (5890.0, 0.35), // Na D + (5184.0, 0.25), // Mg b + (5270.0, 0.2), // Fe I + (4227.0, 0.15), // Ca I + ], + TemplateType::M | TemplateType::Mwarf => vec![ + (5890.0, 0.4), // Na D + (7100.0, 0.3), // TiO band head + (6200.0, 0.25), // TiO + (5270.0, 0.2), // Fe I + (4227.0, 0.15), // Ca I + ], + }; + + for (line_center, depth) in &absorption_lines { + let sigma = 1.5; // Å，线宽 + for (i, &w) in wavelength.iter().enumerate() { + let g = (-(w - line_center).powi(2) / (2.0 * sigma * sigma)).exp(); + flux[i] *= 1.0 - depth * g; + } + } + + flux +} + +/// 简化黑体辐射函数（Wien 近似 + 归一化） +fn black_body(wavelength_um: f64, teff: f64) -> f64 { + let h = 6.626e-34; // Planck constant + let c = 3e8; // speed of light + let k = 1.381e-23; // Boltzmann constant + let wl = wavelength_um * 1e-6; + + if wl <= 0.0 { + return 0.0; + } + + let x = h * c / (wl * k * teff); + if x > 500.0 { + return 0.0; + } + + // B_lambda ∝ 1/(wl^5 * (exp(x) - 1)) + let b = 1.0 / (wl.powi(5) * (x.exp() - 1.0)); + b * 1e-10 // 归一化到合理范围 +} + +/// 对数重采样波长网格 +fn log_resample(wavelength: &[f64], n_points: usize) -> Vec { + if wavelength.len() < 2 { + return wavelength.to_vec(); + } + let w_min = wavelength[0]; + let w_max = wavelength[wavelength.len() - 1]; + let log_min = w_min.ln(); + let log_max = w_max.ln(); + let step = (log_max - log_min) / (n_points - 1) as f64; + + (0..n_points) + .map(|i| (log_min + step * i as f64).exp()) + .collect() +} + +/// 线性插值重采样到目标波长网格 +fn resample(wavelength: &[f64], flux: &[f64], target: &[f64]) -> Vec { + let mut result = Vec::with_capacity(target.len()); + for &tw in target { + // 找 tw 在 wavelength 中的位置 + let idx = match wavelength + .binary_search_by(|w| w.partial_cmp(&tw).unwrap_or(std::cmp::Ordering::Equal)) + { + Ok(i) => i, + Err(i) => i, + }; + if idx == 0 { + result.push(flux[0]); + } else if idx >= wavelength.len() { + result.push(*flux.last().unwrap_or(&0.0)); + } else { + let w1 = wavelength[idx - 1]; + let w2 = wavelength[idx]; + let f1 = flux[idx - 1]; + let f2 = flux[idx]; + let t = if (w2 - w1).abs() > 1e-10 { + (tw - w1) / (w2 - w1) + } else { + 0.0 + }; + result.push(f1 + t * (f2 - f1)); + } + } + result +} + +/// 除以连续谱（简化：用中值滤波器估计） +fn normalize_by_continuum(flux: &[f64], kernel_size: usize) -> Vec { + let n = flux.len(); + let mut continuum = vec![0.0; n]; + + for (i, c) in continuum.iter_mut().enumerate() { + let lo = i.saturating_sub(kernel_size / 2); + let hi = (i + kernel_size / 2 + 1).min(n); + let mut window: Vec = flux[lo..hi].to_vec(); + window.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + *c = window[window.len() / 2]; + } + + flux.iter() + .zip(continuum.iter()) + .map(|(f, c)| if *c > 0.0 { f / c } else { 1.0 }) + .collect() +} + +/// 计算交叉相关函数 +/// +/// 简化实现：直接在 RV 空间扫描（无 FFT 加速，但对典型光谱足够快） +fn compute_ccf( + obs_norm: &[f64], + templ_norm: &[f64], + wavelength: &[f64], + rv_range: (f64, f64), // km/s + rv_step: f64, // km/s +) -> Vec<(f64, f64)> { + let c = 299792.458; // 光速 km/s + let mut ccf = Vec::new(); + + let mut rv = rv_range.0; + while rv <= rv_range.1 { + let shift_factor = 1.0 + rv / c; + + // 对模板进行 Doppler shift + let shifted_waves: Vec = wavelength.iter().map(|&w| w / shift_factor).collect(); + let shifted_templ = resample(wavelength, templ_norm, &shifted_waves); + + // 计算互相关（归一化） + let n = obs_norm.len().min(shifted_templ.len()); + if n == 0 { + ccf.push((rv, 0.0)); + rv += rv_step; + continue; + } + + let mean_obs: f64 = obs_norm[..n].iter().sum::() / n as f64; + let mean_templ: f64 = shifted_templ[..n].iter().sum::() / n as f64; + + let var_obs: f64 = obs_norm[..n].iter().map(|f| (f - mean_obs).powi(2)).sum(); + let var_templ: f64 = shifted_templ[..n] + .iter() + .map(|f| (f - mean_templ).powi(2)) + .sum(); + + let denom = (var_obs * var_templ).sqrt(); + let corr = if denom > 1e-10 { + obs_norm[..n] + .iter() + .zip(shifted_templ[..n].iter()) + .map(|(o, t)| (o - mean_obs) * (t - mean_templ)) + .sum::() + / denom + } else { + 0.0 + }; + + ccf.push((rv, corr)); + rv += rv_step; + } + + ccf +} + +/// 从 CCF 中提取峰值位置（抛物线拟合顶点） +fn fit_ccf_peak(ccf: &[(f64, f64)]) -> Option<(f64, f64)> { + if ccf.len() < 3 { + return None; + } + + // 找最大值 + let (peak_idx, _) = ccf.iter().enumerate().max_by(|a, b| { + a.1 .1 + .partial_cmp(&b.1 .1) + .unwrap_or(std::cmp::Ordering::Equal) + })?; + + if peak_idx == 0 || peak_idx == ccf.len() - 1 { + return Some(ccf[peak_idx]); + } + + // 三点抛物线拟合 (等距格点插值) + let (x1, y1) = ccf[peak_idx - 1]; + let (x2, y2) = ccf[peak_idx]; + let (_x3, y3) = ccf[peak_idx + 1]; + + let h = x2 - x1; + let denom = 2.0 * (y1 - 2.0 * y2 + y3); + if denom.abs() < 1e-15 { + return Some(ccf[peak_idx]); + } + + let vertex_x = x2 + h * (y1 - y3) / denom; + let vertex_y = y2 - 0.125 * (y1 - y3).powi(2) / (y1 - 2.0 * y2 + y3); + + Some((vertex_x, vertex_y)) +} + +/// 主入口：测量径向速度 +pub fn measure_rv( + obs_wavelength: &[f64], + obs_flux: &[f64], + template: TemplateType, + rv_range: (f64, f64), + rv_step: f64, +) -> Result { + if obs_wavelength.len() != obs_flux.len() { + return Err(anyhow!("波长和通量数组长度不一致")); + } + if obs_wavelength.len() < 10 { + return Err(anyhow!("光谱数据点太少（需要 ≥10）")); + } + + // 1. 对数重采样到公共波长网格 + let n_points = obs_wavelength.len(); + let log_grid = log_resample(obs_wavelength, n_points); + let obs_resampled = resample(obs_wavelength, obs_flux, &log_grid); + let templ_flux = generate_template(template, &log_grid); + + // 2. 归一化 + let obs_norm = normalize_by_continuum(&obs_resampled, 51); + let templ_norm = normalize_by_continuum(&templ_flux, 51); + + // 3. 计算 CCF + let ccf = compute_ccf(&obs_norm, &templ_norm, &log_grid, rv_range, rv_step); + + // 4. 提取峰值 + let (peak_rv, peak_corr) = fit_ccf_peak(&ccf).ok_or_else(|| anyhow!("CCF 峰值拟合失败"))?; + + // 5. 估计不确定度（基于 CCF 曲率的简化方法） + let rv_error = estimate_rv_error(&ccf, peak_rv, rv_step); + + // 6. 估计 SNR（简化：基于连续谱散射） + let mean_flux: f64 = obs_resampled.iter().sum::() / obs_resampled.len() as f64; + let variance: f64 = obs_resampled + .iter() + .map(|f| (f - mean_flux).powi(2)) + .sum::() + / obs_resampled.len() as f64; + let snr = if variance > 0.0 { + Some(mean_flux / variance.sqrt()) + } else { + None + }; + + Ok(RadialVelocityResult { + rv: peak_rv, + rv_error, + ccf_peak: peak_corr, + template_name: format!("{:?} (Teff={:.0}K)", template, template_teff(template)), + wavelength_range: ( + *obs_wavelength.first().unwrap_or(&0.0), + *obs_wavelength.last().unwrap_or(&0.0), + ), + snr, + }) +} + +/// 使用用户自定义模板光谱测量 RV +pub fn measure_rv_with_template( + obs_wavelength: &[f64], + obs_flux: &[f64], + templ_wavelength: &[f64], + templ_flux: &[f64], + template_name: &str, + rv_range: (f64, f64), + rv_step: f64, +) -> Result { + if obs_wavelength.len() != obs_flux.len() { + return Err(anyhow!("观测波长和通量数组长度不一致")); + } + if templ_wavelength.len() != templ_flux.len() { + return Err(anyhow!("模板波长和通量数组长度不一致")); + } + + // 1. 对数重采样 + let n_points = obs_wavelength.len().max(templ_wavelength.len()); + let log_grid = log_resample(obs_wavelength, n_points); + let obs_resampled = resample(obs_wavelength, obs_flux, &log_grid); + let templ_resampled = resample(templ_wavelength, templ_flux, &log_grid); + + // 2. 归一化 + let obs_norm = normalize_by_continuum(&obs_resampled, 51); + let templ_norm = normalize_by_continuum(&templ_resampled, 51); + + // 3. CCF + let ccf = compute_ccf(&obs_norm, &templ_norm, &log_grid, rv_range, rv_step); + + // 4. 峰值 + let (peak_rv, peak_corr) = fit_ccf_peak(&ccf).ok_or_else(|| anyhow!("CCF 峰值拟合失败"))?; + + let rv_error = estimate_rv_error(&ccf, peak_rv, rv_step); + + let mean_flux: f64 = obs_resampled.iter().sum::() / obs_resampled.len() as f64; + let variance: f64 = obs_resampled + .iter() + .map(|f| (f - mean_flux).powi(2)) + .sum::() + / obs_resampled.len() as f64; + let snr = if variance > 0.0 { + Some(mean_flux / variance.sqrt()) + } else { + None + }; + + Ok(RadialVelocityResult { + rv: peak_rv, + rv_error, + ccf_peak: peak_corr, + template_name: template_name.to_string(), + wavelength_range: ( + *obs_wavelength.first().unwrap_or(&0.0), + *obs_wavelength.last().unwrap_or(&0.0), + ), + snr, + }) +} + +/// 估计 RV 不确定度（基于 CCF 峰值附近的曲率） +fn estimate_rv_error(ccf: &[(f64, f64)], peak_rv: f64, rv_step: f64) -> f64 { + // 找峰值附近 ±3 步的点，拟合抛物线，从曲率估计不确定度 + let peak_idx = ccf + .iter() + .enumerate() + .min_by(|a, b| { + (a.1 .0 - peak_rv) + .abs() + .partial_cmp(&(b.1 .0 - peak_rv).abs()) + .unwrap_or(std::cmp::Ordering::Equal) + }) + .map(|(i, _)| i) + .unwrap_or(0); + + let half_window = 3; + let lo = peak_idx.saturating_sub(half_window); + let hi = (peak_idx + half_window + 1).min(ccf.len()); + + if hi - lo < 3 { + return rv_step * 2.0; // 数据不足，返回保守估计 + } + + // 简化：用半高宽的一半作为不确定度估计 + let peak_val = ccf[peak_idx].1; + let half_max = peak_val / 2.0; + + let mut left_rv = None; + let mut right_rv = None; + + for i in lo..peak_idx { + if ccf[i].1 <= half_max && ccf[i + 1].1 >= half_max { + let t = (half_max - ccf[i].1) / (ccf[i + 1].1 - ccf[i].1); + left_rv = Some(ccf[i].0 + t * (ccf[i + 1].0 - ccf[i].0)); + } + } + for i in peak_idx..hi - 1 { + if ccf[i].1 >= half_max && ccf[i + 1].1 <= half_max { + let t = (ccf[i].1 - half_max) / (ccf[i].1 - ccf[i + 1].1); + right_rv = Some(ccf[i].0 + t * (ccf[i + 1].0 - ccf[i].0)); + } + } + + match (left_rv, right_rv) { + (Some(l), Some(r)) => (r - l) / 2.355, // FWHM → σ (高斯关系) + _ => rv_step * 2.0, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_black_body_positive() { + let bb = black_body(0.5, 5778.0); // 0.5 μm (可见光) + assert!(bb > 0.0); + } + + #[test] + fn test_template_generation() { + let wave: Vec = (4000..7000).step_by(10).map(|w| w as f64).collect(); + let flux = generate_template(TemplateType::G, &wave); + assert_eq!(flux.len(), wave.len()); + assert!(flux.iter().all(|&f| f > 0.0)); + } + + #[test] + fn test_log_resample() { + let wave = vec![4000.0, 5000.0, 6000.0, 7000.0]; + let resampled = log_resample(&wave, 100); + assert_eq!(resampled.len(), 100); + // 允许浮点精度误差 + assert!( + (resampled[0] - 4000.0).abs() < 1.0, + "first point: {}", + resampled[0] + ); + assert!( + (resampled[99] - 7000.0).abs() < 1.0, + "last point: {}", + resampled[99] + ); + } + + #[test] + fn test_measure_rv_shifted() { + // 1. 生成静止坐标系下的波长网格 (4000 to 7000 Å) + let wave_rest: Vec = (4000..7000).step_by(2).map(|w| w as f64).collect(); + + // 2. 使用 G 模板生成通量 + let flux_rest = generate_template(TemplateType::G, &wave_rest); + + // 3. 施加已知红移 (例如 rv = 50.0 km/s) + let c = 299792.458; // km/s + let rv_input = 50.0; + let wave_obs: Vec = wave_rest + .iter() + .map(|&w| w * (1.0 + rv_input / c)) + .collect(); + + // 4. 使用 measure_rv 测量 + let result = measure_rv(&wave_obs, &flux_rest, TemplateType::G, (0.0, 100.0), 0.5).unwrap(); + println!("TEST RESULT: {:?}", result); + + // 5. 校验测量出的 RV 是否在预期范围内（步长 0.5，误差应极小） + assert!( + (result.rv - rv_input).abs() < 3.0, + "Measured RV = {}, expected = {}", + result.rv, + rv_input + ); + assert!( + result.ccf_peak > 0.9, + "CCF peak should be very high, got {}", + result.ccf_peak + ); + } +} diff --git a/src/services/spectrum/lines.rs b/src/services/spectrum/lines.rs new file mode 100644 index 0000000..20cf2f1 --- /dev/null +++ b/src/services/spectrum/lines.rs @@ -0,0 +1,629 @@ +// src/services/spectrum/lines.rs +// +// 谱线识别与等值宽度（EW）/ 半高全宽（FWHM）测量 +// +// 功能： +// 1. 内置常见天体物理谱线波长表（真空/空气波长） +// 2. 在给定波长窗口内搜索局部极值（发射线找极大，吸收线找极小） +// 3. 数值积分 EW 计算（梯形法） +// 4. FWHM 估算 +// 5. 支持用户自定义谱线列表 + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +/// 谱线类型 +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum LineType { + /// 吸收线 + Absorption, + /// 发射线 + Emission, +} + +/// 单条谱线定义 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SpectralLine { + /// 谱线名称（如 "Hα", "Ca II K", "Na I D"） + pub name: String, + /// 真空波长 (Å) + pub vacuum_wavelength: f64, + /// 空气波长 (Å) + pub air_wavelength: f64, + /// 默认线型 + pub line_type: LineType, + /// 元素符号 + pub element: String, + /// 离化态 + pub species: String, + /// 相对强度（0-1） + pub relative_strength: f64, +} + +/// 测量结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LineMeasurement { + /// 谱线定义 + pub line: SpectralLine, + /// 测量窗口中心波长 (Å) + pub center_wavelength: f64, + /// 等值宽度 (Å)，吸收线为正，发射线为负 + pub ew: f64, + /// 等值宽度不确定度 (Å) + pub ew_error: Option, + /// 半高全宽 (Å) + pub fwhm: Option, + /// 峰值连续谱通量 + pub continuum_flux: f64, + /// 峰值通量 + pub peak_flux: f64, +} + +/// 内置谱线数据库（常用天体物理谱线） +pub fn builtin_line_table() -> Vec { + vec![ + // Balmer 系列 + SpectralLine { + name: "Hα".to_string(), + vacuum_wavelength: 6564.61, + air_wavelength: 6562.80, + line_type: LineType::Absorption, + element: "H".to_string(), + species: "H I".to_string(), + relative_strength: 1.0, + }, + SpectralLine { + name: "Hβ".to_string(), + vacuum_wavelength: 4862.68, + air_wavelength: 4861.33, + line_type: LineType::Absorption, + element: "H".to_string(), + species: "H I".to_string(), + relative_strength: 0.8, + }, + SpectralLine { + name: "Hγ".to_string(), + vacuum_wavelength: 4341.68, + air_wavelength: 4340.47, + line_type: LineType::Absorption, + element: "H".to_string(), + species: "H I".to_string(), + relative_strength: 0.6, + }, + SpectralLine { + name: "Hδ".to_string(), + vacuum_wavelength: 4102.89, + air_wavelength: 4101.74, + line_type: LineType::Absorption, + element: "H".to_string(), + species: "H I".to_string(), + relative_strength: 0.5, + }, + SpectralLine { + name: "Hε".to_string(), + vacuum_wavelength: 3971.20, + air_wavelength: 3970.07, + line_type: LineType::Absorption, + element: "H".to_string(), + species: "H I".to_string(), + relative_strength: 0.4, + }, + // Ca II H&K + SpectralLine { + name: "Ca II K".to_string(), + vacuum_wavelength: 3934.77, + air_wavelength: 3933.66, + line_type: LineType::Absorption, + element: "Ca".to_string(), + species: "Ca II".to_string(), + relative_strength: 0.9, + }, + SpectralLine { + name: "Ca II H".to_string(), + vacuum_wavelength: 3969.47, + air_wavelength: 3968.47, + line_type: LineType::Absorption, + element: "Ca".to_string(), + species: "Ca II".to_string(), + relative_strength: 0.85, + }, + // Na I D 双线 + SpectralLine { + name: "Na I D2".to_string(), + vacuum_wavelength: 5897.56, + air_wavelength: 5889.95, + line_type: LineType::Absorption, + element: "Na".to_string(), + species: "Na I".to_string(), + relative_strength: 0.7, + }, + SpectralLine { + name: "Na I D1".to_string(), + vacuum_wavelength: 5901.48, + air_wavelength: 5895.92, + line_type: LineType::Absorption, + element: "Na".to_string(), + species: "Na I".to_string(), + relative_strength: 0.65, + }, + // Mg b 三重线 + SpectralLine { + name: "Mg I b1".to_string(), + vacuum_wavelength: 5185.54, + air_wavelength: 5183.62, + line_type: LineType::Absorption, + element: "Mg".to_string(), + species: "Mg I".to_string(), + relative_strength: 0.8, + }, + SpectralLine { + name: "Mg I b2".to_string(), + vacuum_wavelength: 5174.11, + air_wavelength: 5172.70, + line_type: LineType::Absorption, + element: "Mg".to_string(), + species: "Mg I".to_string(), + relative_strength: 0.75, + }, + SpectralLine { + name: "Mg I b3".to_string(), + vacuum_wavelength: 5168.76, + air_wavelength: 5167.33, + line_type: LineType::Absorption, + element: "Mg".to_string(), + species: "Mg I".to_string(), + relative_strength: 0.7, + }, + // Fe I 线 + SpectralLine { + name: "Fe I λ5270".to_string(), + vacuum_wavelength: 5271.36, + air_wavelength: 5269.54, + line_type: LineType::Absorption, + element: "Fe".to_string(), + species: "Fe I".to_string(), + relative_strength: 0.5, + }, + SpectralLine { + name: "Fe I λ5328".to_string(), + vacuum_wavelength: 5329.11, + air_wavelength: 5327.27, + line_type: LineType::Absorption, + element: "Fe".to_string(), + species: "Fe I".to_string(), + relative_strength: 0.45, + }, + SpectralLine { + name: "Fe I λ5371".to_string(), + vacuum_wavelength: 5372.12, + air_wavelength: 5370.27, + line_type: LineType::Absorption, + element: "Fe".to_string(), + species: "Fe I".to_string(), + relative_strength: 0.55, + }, + SpectralLine { + name: "Fe I λ4383".to_string(), + vacuum_wavelength: 4384.31, + air_wavelength: 4383.55, + line_type: LineType::Absorption, + element: "Fe".to_string(), + species: "Fe I".to_string(), + relative_strength: 0.5, + }, + // Ti II + SpectralLine { + name: "Ti II λ4444".to_string(), + vacuum_wavelength: 4445.14, + air_wavelength: 4443.80, + line_type: LineType::Absorption, + element: "Ti".to_string(), + species: "Ti II".to_string(), + relative_strength: 0.4, + }, + // Ca I + SpectralLine { + name: "Ca I λ4227".to_string(), + vacuum_wavelength: 4228.00, + air_wavelength: 4226.73, + line_type: LineType::Absorption, + element: "Ca".to_string(), + species: "Ca I".to_string(), + relative_strength: 0.6, + }, + // Hε 与 Ca II H 重叠区域常用线 + SpectralLine { + name: "Ca II K (air)".to_string(), + vacuum_wavelength: 3934.77, + air_wavelength: 3933.66, + line_type: LineType::Absorption, + element: "Ca".to_string(), + species: "Ca II".to_string(), + relative_strength: 0.95, + }, + // O III（星云发射线） + SpectralLine { + name: "[O III] λ5007".to_string(), + vacuum_wavelength: 5008.24, + air_wavelength: 5007.05, + line_type: LineType::Emission, + element: "O".to_string(), + species: "O III".to_string(), + relative_strength: 1.0, + }, + SpectralLine { + name: "[O III] λ4959".to_string(), + vacuum_wavelength: 4960.29, + air_wavelength: 4959.12, + line_type: LineType::Emission, + element: "O".to_string(), + species: "O III".to_string(), + relative_strength: 0.33, + }, + // N II + SpectralLine { + name: "[N II] λ6584".to_string(), + vacuum_wavelength: 6585.27, + air_wavelength: 6583.41, + line_type: LineType::Emission, + element: "N".to_string(), + species: "N II".to_string(), + relative_strength: 0.7, + }, + // S II + SpectralLine { + name: "[S II] λ6717".to_string(), + vacuum_wavelength: 6718.29, + air_wavelength: 6716.44, + line_type: LineType::Emission, + element: "S".to_string(), + species: "S II".to_string(), + relative_strength: 0.5, + }, + SpectralLine { + name: "[S II] λ6731".to_string(), + vacuum_wavelength: 6732.68, + air_wavelength: 6730.82, + line_type: LineType::Emission, + element: "S".to_string(), + species: "S II".to_string(), + relative_strength: 0.45, + }, + ] +} + +/// 按名称查找谱线 +pub fn find_line<'a>(name: &str, table: &'a [SpectralLine]) -> Option<&'a SpectralLine> { + table.iter().find(|l| l.name == name) +} + +/// 按波长范围查找谱线 +pub fn find_lines_in_range( + min_wave: f64, + max_wave: f64, + table: &[SpectralLine], +) -> Vec<&SpectralLine> { + table + .iter() + .filter(|l| l.air_wavelength >= min_wave && l.air_wavelength <= max_wave) + .collect() +} + +/// 按元素查找谱线 +pub fn find_lines_by_element<'a>( + element: &str, + table: &'a [SpectralLine], +) -> Vec<&'a SpectralLine> { + table + .iter() + .filter(|l| l.element.eq_ignore_ascii_case(element)) + .collect() +} + +/// 在光谱中搜索谱线并测量 EW/FWHM +/// +/// 输入：波长数组 (Å)、通量数组、谱线定义、搜索窗口半宽 (Å) +/// 输出：测量结果 +pub fn measure_line( + wavelength: &[f32], + flux: &[f32], + line: &SpectralLine, + window_half_width: f64, +) -> Result { + if wavelength.len() != flux.len() { + return Err(anyhow!("波长和通量数组长度不一致")); + } + if wavelength.len() < 3 { + return Err(anyhow!("光谱数据点太少（需要 ≥3）")); + } + + let center = line.air_wavelength; + let lo = center - window_half_width; + let hi = center + window_half_width; + + // 提取窗口内的数据点 + let mut window_waves = Vec::new(); + let mut window_fluxes = Vec::new(); + for (w, f) in wavelength.iter().zip(flux.iter()) { + let w64 = *w as f64; + if w64 >= lo && w64 <= hi { + window_waves.push(w64); + window_fluxes.push(*f as f64); + } + } + + if window_waves.len() < 3 { + return Err(anyhow!( + "窗口 [{:.1}, {:.1}] Å 内数据点不足（{} 点）", + lo, + hi, + window_waves.len() + )); + } + + // 估计连续谱：窗口边缘两侧各取 20% 的中位数 + let n = window_waves.len(); + let edge = (n as f64 * 0.2).max(1.0) as usize; + let continuum_flux = { + let mut edge_fluxes: Vec = window_fluxes[..edge] + .iter() + .chain(window_fluxes[n - edge..].iter()) + .copied() + .collect(); + edge_fluxes.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + edge_fluxes[edge_fluxes.len() / 2] + }; + + // 找峰值（吸收线找最小值，发射线找最大值） + let (peak_idx, peak_flux_ref) = match line.line_type { + LineType::Absorption => window_fluxes + .iter() + .enumerate() + .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .unwrap(), + LineType::Emission => window_fluxes + .iter() + .enumerate() + .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .unwrap(), + }; + let peak_flux = *peak_flux_ref; + + // 计算 EW（梯形法数值积分） + // 吸收线：EW = ∫(1 - F(λ)/F_c) dλ，结果为正 + // 发射线：EW = ∫(F(λ)/F_c - 1) dλ，结果为负 + let ew = if continuum_flux > 0.0 { + let mut integral = 0.0; + for i in 0..window_waves.len() - 1 { + let dw = window_waves[i + 1] - window_waves[i]; + let norm1 = window_fluxes[i] / continuum_flux; + let norm2 = window_fluxes[i + 1] / continuum_flux; + let avg = match line.line_type { + LineType::Absorption => (1.0 - norm1 + 1.0 - norm2) / 2.0, + LineType::Emission => (norm1 - 1.0 + norm2 - 1.0) / 2.0, + }; + integral += avg * dw; + } + integral + } else { + 0.0 + }; + + // 估算 EW 不确定度（基于连续谱散射的简化估计） + let continuum_scatter = { + let mean = continuum_flux; + let var: f64 = window_fluxes + .iter() + .map(|f| (f - mean).powi(2)) + .sum::() + / n as f64; + var.sqrt() + }; + let ew_error = if continuum_scatter > 0.0 && continuum_flux > 0.0 { + Some((continuum_scatter / continuum_flux * window_half_width * 2.0).sqrt()) + } else { + None + }; + + // FWHM 估算：找半高点并插值 + let fwhm = estimate_fwhm( + &window_waves, + &window_fluxes, + continuum_flux, + line.line_type, + ); + + Ok(LineMeasurement { + line: line.clone(), + center_wavelength: window_waves[peak_idx], + ew, + ew_error, + fwhm, + continuum_flux, + peak_flux, + }) +} + +/// 估算 FWHM（半高全宽） +fn estimate_fwhm( + wavelength: &[f64], + flux: &[f64], + continuum: f64, + line_type: LineType, +) -> Option { + if flux.len() < 3 || continuum <= 0.0 { + return None; + } + + let peak = match line_type { + LineType::Absorption => flux.iter().cloned().fold(f64::INFINITY, f64::min), + LineType::Emission => flux.iter().cloned().fold(f64::NEG_INFINITY, f64::max), + }; + + // 半高点 + let half_level = match line_type { + LineType::Absorption => (continuum + peak) / 2.0, + LineType::Emission => (continuum + peak) / 2.0, + }; + + // 找左侧和右侧的半高点 + let mut left = None; + let mut right = None; + + let peak_idx = match line_type { + LineType::Absorption => flux + .iter() + .enumerate() + .min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(i, _)| i), + LineType::Emission => flux + .iter() + .enumerate() + .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(i, _)| i), + }; + + let peak_idx = peak_idx?; + + // 左侧 + for i in (0..peak_idx).rev() { + if match line_type { + LineType::Absorption => flux[i] >= half_level, + LineType::Emission => flux[i] <= half_level, + } { + // 线性插值 + if i + 1 < flux.len() { + let w1 = wavelength[i]; + let w2 = wavelength[i + 1]; + let f1 = flux[i]; + let f2 = flux[i + 1]; + let dw = w2 - w1; + let df = f2 - f1; + if df.abs() > 1e-10 { + left = Some(w1 + (half_level - f1) / df * dw); + } + } + break; + } + } + + // 右侧 + for i in peak_idx..flux.len() - 1 { + if match line_type { + LineType::Absorption => flux[i] >= half_level, + LineType::Emission => flux[i] <= half_level, + } { + let w1 = wavelength[i]; + let w2 = wavelength[i + 1]; + let f1 = flux[i]; + let f2 = flux[i + 1]; + let dw = w2 - w1; + let df = f2 - f1; + if df.abs() > 1e-10 { + right = Some(w1 + (half_level - f1) / df * dw); + } + break; + } + } + + match (left, right) { + (Some(l), Some(r)) => Some(r - l), + _ => None, + } +} + +/// 批量测量多条谱线 +pub fn measure_lines( + wavelength: &[f32], + flux: &[f32], + lines: &[SpectralLine], + window_half_width: f64, +) -> Vec> { + lines + .iter() + .map(|line| measure_line(wavelength, flux, line, window_half_width)) + .collect() +} + +/// 将测量结果格式化为 Markdown 表格 +pub fn format_measurements_table(measurements: &[LineMeasurement]) -> String { + let mut table = + String::from("| 谱线 | 波长 (Å) | EW (Å) | EW err (Å) | FWHM (Å) | F_c | F_peak |\n"); + table.push_str("|------|----------|--------|------------|----------|-----|--------|\n"); + for m in measurements { + let ew_str = format!("{:.3}", m.ew); + let ew_err_str = m + .ew_error + .map(|e| format!("{:.3}", e)) + .unwrap_or_else(|| "—".to_string()); + let fwhm_str = m + .fwhm + .map(|f| format!("{:.2}", f)) + .unwrap_or_else(|| "—".to_string()); + table.push_str(&format!( + "| {} | {:.2} | {} | {} | {} | {:.1} | {:.1} |\n", + m.line.name, + m.center_wavelength, + ew_str, + ew_err_str, + fwhm_str, + m.continuum_flux, + m.peak_flux + )); + } + table +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_builtin_lines_not_empty() { + let table = builtin_line_table(); + assert!(!table.is_empty()); + // 应该包含 Hα + assert!(find_line("Hα", &table).is_some()); + } + + #[test] + fn test_find_lines_in_range() { + let table = builtin_line_table(); + let lines = find_lines_in_range(6560.0, 6570.0, &table); + assert!(!lines.is_empty()); + assert_eq!(lines[0].name, "Hα"); + } + + #[test] + fn test_measure_absorption_line() { + // 构造一个简单的模拟光谱：连续谱 1.0，在 5000 Å 有一个高斯吸收线 + let mut wavelength = Vec::new(); + let mut flux = Vec::new(); + for i in 0..200 { + let w = 4900.0 + i as f64 * 1.0; // 4900-5099 Å + let absorption = 0.3 * (-(w - 5000.0).powi(2) / 10.0).exp(); + wavelength.push(w as f32); + flux.push((1.0 - absorption) as f32); + } + + let line = SpectralLine { + name: "test".to_string(), + vacuum_wavelength: 5000.0, + air_wavelength: 5000.0, + line_type: LineType::Absorption, + element: "X".to_string(), + species: "X I".to_string(), + relative_strength: 1.0, + }; + + let result = measure_line(&wavelength, &flux, &line, 30.0).unwrap(); + // EW 应该为正（吸收线） + assert!(result.ew > 0.0, "吸收线 EW 应为正: {}", result.ew); + // 连续谱应接近 1.0 + assert!( + (result.continuum_flux - 1.0).abs() < 0.1, + "连续谱应接近 1.0: {}", + result.continuum_flux + ); + } +} diff --git a/src/services/spectrum/mod.rs b/src/services/spectrum/mod.rs new file mode 100644 index 0000000..420356b --- /dev/null +++ b/src/services/spectrum/mod.rs @@ -0,0 +1,12 @@ +// src/services/spectrum/mod.rs +// +// 光谱分析服务层 —— 参数提取、谱线测量、交叉相关测速 +// +// 模块布局： +// parameters.rs —— FITS Header/BinTable 恒星大气参数提取（Teff, logg, [Fe/H], RV...） +// lines.rs —— 谱线识别 + 等值宽度（EW）/ 半高全宽（FWHM）测量 +// cross_correlate.rs —— 模板匹配法交叉相关测径向速度（RV） + +pub mod cross_correlate; +pub mod lines; +pub mod parameters; diff --git a/src/services/spectrum/parameters.rs b/src/services/spectrum/parameters.rs new file mode 100644 index 0000000..dcbc677 --- /dev/null +++ b/src/services/spectrum/parameters.rs @@ -0,0 +1,378 @@ +// src/services/spectrum/parameters.rs +// +// 光谱参数自动提取 —— 从 FITS 文件中读取恒星大气参数 +// +// 支持源： +// - LAMOST：TEFF, LOGG, FEH, RV, SNRU/SNRG/SNRR/SNRI +// - SDSS/BOSS：TEFF_LOGG_FEH, RV, SN median +// - SDSS/APOGEE：ASPCAP 参数（Teff, logg, [Fe/H], [M/H], [α/M] 等 20+ 元素） +// - DESI：红移 Z, CLASS, SUBCLASS, SNR +// - Gaia XP：BP/RP 色指数，由 photometry 推导 + +use anyhow::{anyhow, Context, Result}; +use fitsio::hdu::HduInfo; +use serde::{Deserialize, Serialize}; +use std::path::Path; +use tracing::info; + +use crate::api::AppState; +use crate::services::observation::types::{ProductSpec, Source}; + +/// 恒星大气参数（跨源归一化） +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct StellarParameters { + pub source: String, + pub source_id: String, + /// 有效温度 (K) + #[serde(skip_serializing_if = "Option::is_none")] + pub teff: Option, + /// 表面重力 log(g) (dex) + #[serde(skip_serializing_if = "Option::is_none")] + pub logg: Option, + /// 铁丰度 [Fe/H] (dex) + #[serde(skip_serializing_if = "Option::is_none")] + pub feh: Option, + /// 总金属丰度 [M/H] (dex) + #[serde(skip_serializing_if = "Option::is_none")] + pub mh: Option, + /// α 元素丰度 [α/M] (dex) + #[serde(skip_serializing_if = "Option::is_none")] + pub alpha_m: Option, + /// 径向速度 (km/s) + #[serde(skip_serializing_if = "Option::is_none")] + pub rv: Option, + /// 红移（星系光谱） + #[serde(skip_serializing_if = "Option::is_none")] + pub redshift: Option, + /// 光谱分类 + #[serde(skip_serializing_if = "Option::is_none")] + pub spectral_class: Option, + /// 信噪比 + #[serde(skip_serializing_if = "Option::is_none")] + pub snr: Option, + /// 化学丰度（APOGEE 等多元素源） + #[serde(default, skip_serializing_if = "Option::is_none")] + pub abundances: Option>, + /// 源专属原始字段 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extra: Option, +} + +/// 单个元素丰度 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AbundanceEntry { + pub element: String, + /// [X/H] 或 [X/Fe] + pub value: f64, + /// 不确定度 + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +/// 从 observation_cache 中的 FITS 文件提取参数 +pub async fn extract_parameters( + _state: &AppState, + source: Source, + product: &ProductSpec, + source_id: &str, + file_path: &str, +) -> Result { + info!( + "[SpectrumParams] 提取参数: source={:?}, source_id={}, file={}", + source, source_id, file_path + ); + + let path = Path::new(file_path); + if !path.exists() { + return Err(anyhow!("文件不存在: {}", file_path)); + } + + let params = match source { + Source::Lamost => extract_lamost(path, source_id)?, + Source::Sdss => { + // SDSS spec 与 APOGEE 参数提取逻辑不同 + match product.subtype.as_deref() { + Some("apstar") | Some("aspcap") => extract_apogee(path, source_id)?, + _ => extract_sdss(path, source_id)?, + } + } + Source::Desi => extract_desi(path, source_id)?, + Source::Gaia => extract_gaia(path, source_id)?, + _ => return Err(anyhow!("{:?} 暂不支持参数自动提取", source)), + }; + + Ok(params) +} + +pub fn read_f64_with_fallback(fits: &mut fitsio::FitsFile, keys: &[&str]) -> Option { + for hdu_idx in 0..8 { + if let Ok(hdu) = fits.hdu(hdu_idx) { + for &key in keys { + if let Ok(val) = hdu.read_key::(fits, key) { + if val.is_finite() { + return Some(val); + } + } + } + } else { + break; + } + } + None +} + +pub fn read_string_with_fallback(fits: &mut fitsio::FitsFile, keys: &[&str]) -> Option { + for hdu_idx in 0..8 { + if let Ok(hdu) = fits.hdu(hdu_idx) { + for &key in keys { + if let Ok(val) = hdu.read_key::(fits, key) { + let s = val.trim(); + if !s.is_empty() && s != "UNDEFINED" { + return Some(s.to_string()); + } + } + } + } else { + break; + } + } + None +} + +/// 从 LAMOST FITS 提取参数 +fn extract_lamost(path: &Path, source_id: &str) -> Result { + let mut fits = fitsio::FitsFile::open(path).context("打开 LAMOST FITS")?; + + let teff = read_f64_with_fallback(&mut fits, &["TEFF", "T_EFF", "TEMP", "EFFECTIVE_TEMP"]); + let logg = read_f64_with_fallback(&mut fits, &["LOGG", "LOG_G", "GRAVITY"]); + let feh = read_f64_with_fallback(&mut fits, &["FEH", "FE_H", "[Fe/H]", "METALLICITY"]); + let rv = read_f64_with_fallback( + &mut fits, + &["RV", "HELIO_RV", "VHELIO", "VRAD", "RADIAL_VELOCITY", "HRV"], + ); + let snr = read_f64_with_fallback( + &mut fits, + &["SNRG", "SNR", "SNR_MEDIAN", "SNRU", "SNRR", "SNRI"], + ); + let objname = read_string_with_fallback( + &mut fits, + &["OBJNAME", "SPECTRAL_CLASS", "CLASS", "SUBCLASS"], + ); + + Ok(StellarParameters { + source: "lamost".to_string(), + source_id: source_id.to_string(), + teff, + logg, + feh, + rv, + snr, + spectral_class: objname, + ..Default::default() + }) +} + +/// 从 SDSS/BOSS spec 提取参数 +fn extract_sdss(path: &Path, source_id: &str) -> Result { + let mut fits = fitsio::FitsFile::open(path).context("打开 SDSS FITS")?; + + let redshift = read_f64_with_fallback(&mut fits, &["REDSHIFT", "Z", "Z_BEST"]); + let snr = read_f64_with_fallback(&mut fits, &["SNMEDIAN", "SN_R", "SNR", "SNR_MEDIAN"]); + let subclass = + read_string_with_fallback(&mut fits, &["SUBCLASSOBJ", "SUBCLASS", "CLASSOBJ", "CLASS"]); + let teff = read_f64_with_fallback(&mut fits, &["TEFF", "T_EFF", "TEMP", "EFFECTIVE_TEMP"]); + let logg = read_f64_with_fallback(&mut fits, &["LOGG", "LOG_G", "GRAVITY"]); + let feh = read_f64_with_fallback(&mut fits, &["FEH", "FE_H", "[Fe/H]"]); + let rv = read_f64_with_fallback( + &mut fits, + &["HELIO_RV", "VRAD", "VHELIO", "RV", "RADIAL_VELOCITY", "HRV"], + ); + + Ok(StellarParameters { + source: "sdss".to_string(), + source_id: source_id.to_string(), + teff, + logg, + feh, + redshift, + spectral_class: subclass, + snr, + rv, + ..Default::default() + }) +} + +/// 从 SDSS APOGEE aspcapStar 提取完整化学丰度 +fn extract_apogee(path: &Path, source_id: &str) -> Result { + let mut fits = fitsio::FitsFile::open(path).context("打开 APOGEE FITS")?; + + let teff = read_f64_with_fallback( + &mut fits, + &["TEFF", "T_EFF", "TEMP", "EFFECTIVE_TEMP", "TEFF_ASPCAP"], + ); + let logg = read_f64_with_fallback(&mut fits, &["LOGG", "LOG_G", "GRAVITY", "LOGG_ASPCAP"]); + let feh = read_f64_with_fallback(&mut fits, &["FEH", "FE_H", "[Fe/H]", "FEH_ASPCAP"]); + let mh = read_f64_with_fallback(&mut fits, &["M_H", "[M/H]", "MH"]); + let alpha_m = read_f64_with_fallback(&mut fits, &["ALPHA_M", "[A/M]", "[ALPHA/M]"]); + let rv = read_f64_with_fallback( + &mut fits, + &["VHELIO", "RV", "HELIO_RV", "VRAD", "RADIAL_VELOCITY", "HRV"], + ); + let snr = read_f64_with_fallback(&mut fits, &["SNR", "SNR_MEDIAN", "SNMEDIAN"]); + + // 尝试读取 APOGEE 丰度扩展表（ASTRUU 元素表） + let abundances = extract_apogee_abundances(&mut fits)?; + + Ok(StellarParameters { + source: "sdss".to_string(), + source_id: source_id.to_string(), + teff, + logg, + feh, + mh, + alpha_m, + rv, + snr, + abundances, + ..Default::default() + }) +} + +/// 从 APOGEE FITS 读取元素丰度扩展表 +fn extract_apogee_abundances(fits: &mut fitsio::FitsFile) -> Result>> { + // APOGEE aspcapStar 有 ASTUU 或 APOGEE_ELEMENT table + // 尝试已知的表名 + let known_tables = ["ASTRUU", "APOGEE_ELEMENT", "ELEM"]; + + for table_name in &known_tables { + if let Ok(hdu) = fits.hdu(*table_name) { + if let HduInfo::TableInfo { + ref column_descriptions, + num_rows, + } = hdu.info + { + let col_names: Vec = + column_descriptions.iter().map(|c| c.name.clone()).collect(); + + // 查找必需列 + let symbol_col = col_names.iter().find(|c| c.contains("SYMBOL")); + let value_col = col_names + .iter() + .find(|c| c.contains("VALUE") || c.contains("ABUND")); + let error_col = col_names.iter().find(|c| c.contains("ERROR")); + + if let (Some(sym_col), Some(val_col)) = (symbol_col, value_col) { + let n_rows = num_rows; + let mut abundances = Vec::new(); + + let symbols: Vec = hdu.read_col(fits, sym_col).unwrap_or_default(); + let values_f32: Vec = hdu.read_col(fits, val_col).unwrap_or_default(); + let values: Vec = values_f32.iter().map(|&v| v as f64).collect(); + // 误差列可能不存在，跳过如果读取失败 + let errors: Vec = if let Some(col) = error_col { + let err_f32: Vec = hdu.read_col(fits, col).unwrap_or_default(); + err_f32.iter().map(|&x| x as f64).collect() + } else { + Vec::new() + }; + + for i in 0..n_rows.min(symbols.len()).min(values.len()) { + let elem = symbols[i].clone(); + if elem.is_empty() || elem == "C" || elem == "N" || elem == "O" { + continue; + } + let value = values[i]; + if value.abs() > 9.0 || value == 0.0 { + continue; + } + let error = errors.get(i).copied().filter(|e| *e > 0.0 && *e < 9.0); + abundances.push(AbundanceEntry { + element: elem, + value, + error, + }); + } + + if !abundances.is_empty() { + return Ok(Some(abundances)); + } + } + } + } + } + + Ok(None) +} + +/// 从 DESI coadd FITS 提取参数 +fn extract_desi(path: &Path, source_id: &str) -> Result { + let mut fits = fitsio::FitsFile::open(path).context("打开 DESI FITS")?; + + let redshift = read_f64_with_fallback(&mut fits, &["REDSHIFT", "Z", "Z_BEST"]); + let snr = read_f64_with_fallback(&mut fits, &["SNR", "SNR_MEDIAN", "SNMEDIAN"]); + let subclass = read_string_with_fallback(&mut fits, &["CLASS", "SUBCLASS", "SUBCLASSOBJ"]); + + Ok(StellarParameters { + source: "desi".to_string(), + source_id: source_id.to_string(), + redshift, + spectral_class: subclass, + snr, + ..Default::default() + }) +} + +/// 从 Gaia XP 光谱提取基本参数（主要是色指数推导的参数） +fn extract_gaia(path: &Path, source_id: &str) -> Result { + let mut fits = fitsio::FitsFile::open(path).context("打开 Gaia FITS")?; + + let snr = read_f64_with_fallback(&mut fits, &["SNR", "SNR_MEDIAN", "SNMEDIAN"]); + let rv = read_f64_with_fallback( + &mut fits, + &["RV", "RADIAL_VELOCITY", "VHELIO", "HELIO_RV", "VRAD"], + ); + + Ok(StellarParameters { + source: "gaia".to_string(), + source_id: source_id.to_string(), + rv, + snr, + ..Default::default() + }) +} + +// ── FITS 辅助函数 ── + +/// 批量提取多个源的参数 +pub async fn extract_parameters_batch( + state: &AppState, + items: &[(Source, ProductSpec, String, String)], // (source, product, source_id, file_path) +) -> Vec> { + let mut results = Vec::new(); + for (source, product, source_id, file_path) in items { + results.push(extract_parameters(state, *source, product, source_id, file_path).await); + } + results +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_stellar_parameters_serialization() { + let params = StellarParameters { + source: "lamost".to_string(), + source_id: "438809089".to_string(), + teff: Some(5778.0), + logg: Some(4.44), + feh: Some(0.0), + rv: Some(12.5), + snr: Some(100.0), + ..Default::default() + }; + let json = serde_json::to_string(¶ms).unwrap(); + assert!(json.contains("5778")); + assert!(json.contains("lamost")); + } +} diff --git a/src/services/timeseries/classification.rs b/src/services/timeseries/classification.rs new file mode 100644 index 0000000..8a08731 --- /dev/null +++ b/src/services/timeseries/classification.rs @@ -0,0 +1,438 @@ +// src/services/timeseries/classification.rs +// +// 变星分类与特征提取 —— 基于统计特征的自动分类 +// +// 特征： +// - 振幅（峰-谷、IQR、标准差） +// - 偏度（skewness） +// - 峰度（kurtosis） +// - 过零率（zero crossing rate） +// - 变化率指标（eta, beyond1std, max_slope） +// - 周期特征（如果已知周期） +// +// 分类：基于决策树规则的简单分类器 +// - RR Lyrae +// - Cepheid +// - 食双星（Eclipsing Binary） +// - 脉动变星（Pulsating） +// - AGN +// - 随机变化（Stochastic） +// - 未知 + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +/// 变星类型 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub enum VariableStarType { + /// RR Lyrae + RrLyrae, + /// Cepheid + Cepheid, + /// 食双星（规则方波信号） + EclipsingBinary, + /// 脉动变星（其他类型） + Pulsating, + /// 活动星系核（AGN） + Agn, + /// 随机/不规则变化 + Stochastic, + /// T Tauri 型 + TTauri, + /// 未知/未分类 + Unknown, +} + +impl VariableStarType { + pub fn display(&self) -> &str { + match self { + Self::RrLyrae => "RR Lyrae", + Self::Cepheid => "Cepheid", + Self::EclipsingBinary => "食双星", + Self::Pulsating => "脉动变星", + Self::Agn => "AGN", + Self::Stochastic => "随机变化", + Self::TTauri => "T Tauri", + Self::Unknown => "未知", + } + } + + pub fn description(&self) -> &str { + match self { + Self::RrLyrae => "RR Lyrae 型脉动变星，周期通常 0.2-1.0 天，振幅较大，光变曲线不对称", + Self::Cepheid => "经典造父变星，周期 1-100 天，周期-光度关系重要", + Self::EclipsingBinary => "食双星，周期性方波状光变，振幅恒定", + Self::Pulsating => "其他脉动变星（δ Scuti、γ Dor 等）", + Self::Agn => "活动星系核，不规则快速变化，无周期性", + Self::Stochastic => "随机/不规则变化（恒星活动、仪器噪声等）", + Self::TTauri => "T Tauri 前主序星，不规则光变+偶发增亮", + Self::Unknown => "未能分类", + } + } +} + +/// 变星特征向量 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct VariabilityFeatures { + /// 振幅（峰-谷，mag） + pub amplitude: f64, + /// 中位绝对偏差（MAD） + pub mad: f64, + /// 标准差 + pub std: f64, + /// 偏度 + pub skewness: f64, + /// 峰度 + pub kurtosis: f64, + /// 过零率 + pub zero_crossing_rate: f64, + /// η 指标（相邻点差分方差 / 总方差） + pub eta: f64, + /// 超过 1σ 的比例 + pub beyond_1std: f64, + /// 最大斜率 + pub max_slope: f64, + /// 连续上升/下降最长段 + pub max_monotonic: f64, + /// 已知周期（天），None 表示未搜索 + pub period: Option, +} + +/// 分类结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ClassificationResult { + /// 分类类型 + pub classification: VariableStarType, + /// 置信度（0-1） + pub confidence: f64, + /// 提取的特征 + pub features: VariabilityFeatures, + /// 各类型的得分（用于调试） + pub scores: std::collections::HashMap, + /// 人类可读的解释 + pub explanation: String, +} + +/// 从时域数据提取变星特征 +pub fn extract_features(time: &[f64], flux: &[f64]) -> Result { + if time.len() != flux.len() { + return Err(anyhow!("时间和通量数组长度不一致")); + } + if time.len() < 5 { + return Err(anyhow!("数据点太少（需要 ≥5）")); + } + + let n = flux.len(); + let mean: f64 = flux.iter().sum::() / n as f64; + let var: f64 = flux.iter().map(|f| (f - mean).powi(2)).sum::() / n as f64; + let std = var.sqrt(); + + if std < 1e-15 { + return Err(anyhow!("通量标准差为零")); + } + + // 振幅（峰-谷） + let max_val = flux.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + let min_val = flux.iter().cloned().fold(f64::INFINITY, f64::min); + let amplitude = max_val - min_val; + + // MAD（中位绝对偏差） + let mut sorted_flux = flux.to_vec(); + sorted_flux.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let median = if n.is_multiple_of(2) { + (sorted_flux[n / 2 - 1] + sorted_flux[n / 2]) / 2.0 + } else { + sorted_flux[n / 2] + }; + let mut abs_devs: Vec = flux.iter().map(|f| (f - median).abs()).collect(); + abs_devs.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let mad = if n.is_multiple_of(2) { + (abs_devs[n / 2 - 1] + abs_devs[n / 2]) / 2.0 + } else { + abs_devs[n / 2] + }; + + // 偏度 + let m3: f64 = flux.iter().map(|f| ((f - mean) / std).powi(3)).sum::() / n as f64; + + // 峰度 + let m4: f64 = flux.iter().map(|f| ((f - mean) / std).powi(4)).sum::() / n as f64; + let kurtosis = m4 - 3.0; // excess kurtosis + + // 过零率 + let zcr = { + let mut crossings = 0; + for i in 1..n { + if (flux[i] - mean) * (flux[i - 1] - mean) < 0.0 { + crossings += 1; + } + } + crossings as f64 / (n - 1) as f64 + }; + + // η 指标 + let eta = { + let diff_var: f64 = + (1..n).map(|i| (flux[i] - flux[i - 1]).powi(2)).sum::() / (n - 1) as f64; + diff_var / var + }; + + // beyond1std + let beyond_1std = { + let count = flux.iter().filter(|f| (**f - mean).abs() > std).count(); + count as f64 / n as f64 + }; + + // 最大斜率 + let max_slope = { + let mut max_s = 0.0; + for i in 1..n { + let dt = time[i] - time[i - 1]; + if dt > 1e-10 { + let slope = ((flux[i] - flux[i - 1]).abs() / dt).abs(); + if slope > max_s { + max_s = slope; + } + } + } + max_s + }; + + // 最长连续单调段 + let max_monotonic = { + let mut max_mono = 1; + let mut current = 1; + for i in 2..n { + let diff_curr = flux[i] - flux[i - 1]; + let diff_prev = flux[i - 1] - flux[i - 2]; + if diff_curr * diff_prev > 0.0 { + current += 1; + } else { + max_mono = max_mono.max(current); + current = 1; + } + } + max_mono = max_mono.max(current); + max_mono as f64 / n as f64 + }; + + Ok(VariabilityFeatures { + amplitude, + mad, + std, + skewness: m3, + kurtosis, + zero_crossing_rate: zcr, + eta, + beyond_1std, + max_slope, + max_monotonic, + period: None, + }) +} + +/// 基于特征的变星分类（决策树规则） +pub fn classify(features: &mut VariabilityFeatures, _time_span: f64) -> ClassificationResult { + let mut scores: std::collections::HashMap = std::collections::HashMap::new(); + + let confidence; + let classification; + let explanation; + + // 规则 1：振幅极小 → 非变星/噪声 + if features.amplitude < 0.01 { + classification = VariableStarType::Stochastic; + confidence = 0.8; + explanation = "振幅极小（< 0.01 mag），可能是恒星活动或仪器噪声".to_string(); + scores.insert("stochastic".to_string(), 0.8); + } + // 规则 2：周期性 + 振幅 + 不对称光变 → RR Lyrae + else if let Some(period) = features.period { + if (0.2..=1.0).contains(&period) + && features.amplitude > 0.3 + && features.skewness.abs() > 0.3 + { + classification = VariableStarType::RrLyrae; + confidence = 0.85; + explanation = format!( + "周期 {:.4} 天在 RR Lyrae 范围（0.2-1.0 天），振幅 {:.3} mag，偏度 {:.2} 表明不对称光变", + period, features.amplitude, features.skewness + ); + scores.insert("rr_lyrae".to_string(), 0.85); + } + // 规则 3：长周期 + 大振幅 → Cepheid + else if (1.0..=100.0).contains(&period) && features.amplitude > 0.3 { + classification = VariableStarType::Cepheid; + confidence = 0.75; + explanation = format!( + "周期 {:.2} 天在 Cepheid 范围（1-100 天），振幅 {:.3} mag", + period, features.amplitude + ); + scores.insert("cepheid".to_string(), 0.75); + } + // 规则 4：规则周期 + 低偏度 → 食双星 + else if features.amplitude > 0.05 + && features.skewness.abs() < 0.5 + && features.kurtosis > 0.5 + && features.zero_crossing_rate < 0.3 + { + classification = VariableStarType::EclipsingBinary; + confidence = 0.7; + explanation = format!( + "周期 {:.4} 天，振幅 {:.3} mag，低偏度+高峰度提示方波状食信号", + period, features.amplitude + ); + scores.insert("eclipsing_binary".to_string(), 0.7); + } + // 规则 5：有周期但不符合以上 + else { + classification = VariableStarType::Pulsating; + confidence = 0.5; + explanation = format!( + "周期 {:.4} 天，振幅 {:.3} mag，但不符合已知变星类型的特征模式", + period, features.amplitude + ); + scores.insert("pulsating".to_string(), 0.5); + } + } + // 规则 6：无周期 + 大振幅 + 高过零率 → AGN + else if features.amplitude > 0.1 && features.zero_crossing_rate > 0.4 && features.eta > 0.5 { + classification = VariableStarType::Agn; + confidence = 0.6; + explanation = "无周期性，快速不规则变化（高过零率+高η指标）提示 AGN".to_string(); + scores.insert("agn".to_string(), 0.6); + } + // 规则 7：无周期 + 低变化 → 随机变化 + else if features.amplitude < 0.05 && features.std / features.amplitude.max(1e-10) < 0.5 { + classification = VariableStarType::Stochastic; + confidence = 0.7; + explanation = "变化幅度小且无明显结构，可能是随机光变".to_string(); + scores.insert("stochastic".to_string(), 0.7); + } + // 默认：未知 + else { + classification = VariableStarType::Unknown; + confidence = 0.3; + explanation = "特征不足以做出可靠分类".to_string(); + scores.insert("unknown".to_string(), 0.3); + } + + ClassificationResult { + classification, + confidence, + features: features.clone(), + scores, + explanation, + } +} + +/// 完整流程：从时域数据提取特征 + 分类 +pub fn classify_variable_star( + time: &[f64], + flux: &[f64], + period: Option, +) -> Result { + let time_span = time.iter().cloned().fold(f64::NEG_INFINITY, f64::max) + - time.iter().cloned().fold(f64::INFINITY, f64::min); + + let mut features = extract_features(time, flux)?; + features.period = period; + + Ok(classify(&mut features, time_span)) +} + +/// 格式化分类结果为 Markdown +pub fn format_classification(result: &ClassificationResult) -> String { + let mut out = "## 变星分类结果\n\n".to_string(); + out.push_str(&format!( + "- **分类**: {} ({})\n", + result.classification.display(), + result.classification.description() + )); + out.push_str(&format!( + "- **置信度**: {:.1}%\n", + result.confidence * 100.0 + )); + out.push_str(&format!("- **解释**: {}\n\n", result.explanation)); + + out.push_str("### 特征\n\n"); + out.push_str(&format!("- 振幅: {:.4} mag\n", result.features.amplitude)); + out.push_str(&format!("- MAD: {:.4}\n", result.features.mad)); + out.push_str(&format!("- 标准差: {:.4}\n", result.features.std)); + out.push_str(&format!("- 偏度: {:.4}\n", result.features.skewness)); + out.push_str(&format!("- 峰度: {:.4}\n", result.features.kurtosis)); + out.push_str(&format!( + "- 过零率: {:.4}\n", + result.features.zero_crossing_rate + )); + out.push_str(&format!("- η 指标: {:.4}\n", result.features.eta)); + out.push_str(&format!( + "- beyond1std: {:.4}\n", + result.features.beyond_1std + )); + out.push_str(&format!("- 最大斜率: {:.4}\n", result.features.max_slope)); + if let Some(p) = result.features.period { + out.push_str(&format!("- 周期: {:.6} 天\n", p)); + } + + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extract_features_sine() { + let n = 200; + let mut time = Vec::new(); + let mut flux = Vec::new(); + for i in 0..n { + let t = i as f64 * 0.01; + time.push(t); + flux.push((2.0 * std::f64::consts::PI * t / 1.0).sin() + 5.0); + } + + let features = extract_features(&time, &flux).unwrap(); + assert!(features.amplitude > 1.5); // 正弦波振幅约 2 + assert!(features.std > 0.5); + } + + #[test] + fn test_classify_rr_lyrae() { + let features = VariabilityFeatures { + amplitude: 0.8, + mad: 0.2, + std: 0.25, + skewness: 0.5, + kurtosis: 0.1, + zero_crossing_rate: 0.15, + eta: 0.3, + beyond_1std: 0.3, + max_slope: 5.0, + max_monotonic: 0.3, + period: Some(0.5), + }; + let result = classify(&mut features.clone(), 100.0); + assert_eq!(result.classification, VariableStarType::RrLyrae); + } + + #[test] + fn test_classify_eclipsing_binary() { + let features = VariabilityFeatures { + amplitude: 0.2, + mad: 0.05, + std: 0.06, + skewness: 0.1, + kurtosis: 1.0, + zero_crossing_rate: 0.1, + eta: 0.2, + beyond_1std: 0.2, + max_slope: 3.0, + max_monotonic: 0.4, + period: Some(2.0), + }; + let result = classify(&mut features.clone(), 100.0); + assert_eq!(result.classification, VariableStarType::EclipsingBinary); + } +} diff --git a/src/services/timeseries/mod.rs b/src/services/timeseries/mod.rs new file mode 100644 index 0000000..75b78e7 --- /dev/null +++ b/src/services/timeseries/mod.rs @@ -0,0 +1,10 @@ +// src/services/timeseries/mod.rs +// +// 时序分析服务层 —— 周期搜索、变星分类 +// +// 模块布局： +// periodicity.rs —— Lomb-Scargle 周期图 + BLS（Box Least Squares）凌星检测 +// classification.rs —— 基于统计特征的变星类型自动分类 + +pub mod classification; +pub mod periodicity; diff --git a/src/services/timeseries/periodicity.rs b/src/services/timeseries/periodicity.rs new file mode 100644 index 0000000..8e6e0cf --- /dev/null +++ b/src/services/timeseries/periodicity.rs @@ -0,0 +1,469 @@ +// src/services/timeseries/periodicity.rs +// +// 光变周期搜索 —— Lomb-Scargle 周期图 + BLS（Box Least Squares） +// +// Lomb-Scargle：适用于非均匀采样的时域数据，自动报告最优周期和 FAP +// BLS：凌星检测专用，搜索方波信号 + +use anyhow::{anyhow, Result}; +use rayon::prelude::*; +use serde::{Deserialize, Serialize}; + +/// 周期搜索结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PeriodSearchResult { + /// 最优周期 (天) + pub best_period: f64, + /// 最优周期对应的功率（0-1 归一化） + pub best_power: f64, + /// 误报概率（FAP），越小越显著 + pub fap: Option, + /// 次优周期列表（按功率降序） + pub secondary_periods: Vec, + /// 周期图数据点 (period, power)，供前端绘图 + pub periodogram: Vec<(f64, f64)>, + /// 折叠后的光变曲线（time/phase, flux），供前端绘图 + pub folded_lightcurve: Option>, + /// 搜索方法 + pub method: String, +} + +/// 候选周期 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PeriodCandidate { + pub period: f64, + pub power: f64, +} + +/// BLS 搜索结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BlsResult { + /// 最优周期 (天) + pub best_period: f64, + /// 最优凌星深度 + pub depth: f64, + /// 最优凌星持续时间 (天) + pub duration: f64, + /// 最优凌星相位 + pub phase: f64, + /// SDE（信号检测效率） + pub sde: f64, + /// 周期图 + pub periodogram: Vec<(f64, f64)>, +} + +/// Lomb-Scargle 周期图 +/// +/// 输入：时间数组 (天)、通量数组、最小周期 (天)、最大周期 (天)、频率点数 +pub fn lomb_scargle( + time: &[f64], + flux: &[f64], + min_period: Option, + max_period: Option, + n_freq: usize, +) -> Result { + if time.len() != flux.len() { + return Err(anyhow!("时间和通量数组长度不一致")); + } + if time.len() < 4 { + return Err(anyhow!("数据点太少（需要 ≥4）")); + } + + let n = time.len(); + + // 默认周期范围 + let t_min = time.iter().cloned().fold(f64::INFINITY, f64::min); + let t_max = time.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + let baseline = t_max - t_min; + let min_p = min_period.unwrap_or(baseline / 10.0).max(0.1); + let max_p = max_period.unwrap_or(baseline / 2.0).min(baseline * 2.0); + + // 频率网格（均匀频率采样） + let freq_min = 1.0 / max_p; + let freq_max = 1.0 / min_p; + let df = (freq_max - freq_min) / n_freq as f64; + + // 均值和方差归一化 + let mean_flux: f64 = flux.iter().sum::() / n as f64; + let var_flux: f64 = flux.iter().map(|f| (f - mean_flux).powi(2)).sum::() / n as f64; + + if var_flux < 1e-20 { + return Err(anyhow!("通量方差为零，无法搜索周期")); + } + + // 计算周期图 + let periodogram: Vec<(f64, f64)> = (0..n_freq) + .into_par_iter() + .map(|i| { + let freq = freq_min + df * i as f64; + let omega = 2.0 * std::f64::consts::PI * freq; + let tau = estimate_tau(time, omega); + + // Lomb-Scargle 功率 + let (s, c, sc, ss) = { + let mut s = 0.0; + let mut c = 0.0; + let mut sc = 0.0; + let mut ss = 0.0; + for &t in time { + let wt = omega * (t - tau); + let st = wt.sin(); + let ct = wt.cos(); + s += st; + c += ct; + sc += st * ct; + ss += st * st; + } + (s, c, sc, ss) + }; + + let (h_c, h_s) = { + let mut hc = 0.0; + let mut hs = 0.0; + for (t, f) in time.iter().zip(flux.iter()) { + let wt = omega * (t - tau); + let y = f - mean_flux; + hc += y * wt.cos(); + hs += y * wt.sin(); + } + (hc, hs) + }; + + let denom = 2.0 * (ss * (1.0 - sc / n as f64) - (s * s - c * c) / (4.0 * n as f64)); + let power = if denom.abs() > 1e-20 { + (h_c * h_c + / (c * c / n as f64 + 2.0 * ss * (1.0 - sc / n as f64) + - (s * s + c * c) / n as f64) + + h_s * h_s + / (s * s / n as f64 + 2.0 * ss * (1.0 - sc / n as f64) + - (s * s + c * c) / n as f64)) + / (2.0 * var_flux) + } else { + // 简化 LS 公式 + let omega_tau_shifted: Vec<(f64, f64)> = time + .iter() + .map(|&t| { + let wt = omega * (t - tau); + (wt.cos(), wt.sin()) + }) + .collect(); + let y: Vec = flux.iter().map(|f| f - mean_flux).collect(); + let yc: f64 = y + .iter() + .zip(omega_tau_shifted.iter()) + .map(|(y, (c, _))| y * c) + .sum(); + let ys: f64 = y + .iter() + .zip(omega_tau_shifted.iter()) + .map(|(y, (_, s))| y * s) + .sum(); + let cc: f64 = omega_tau_shifted.iter().map(|(_, c)| c * c).sum(); + let ss2: f64 = omega_tau_shifted.iter().map(|(_, s)| s * s).sum(); + let cs: f64 = omega_tau_shifted.iter().map(|(c, s)| c * s).sum(); + let denom2 = cc * ss2 - cs * cs; + if denom2.abs() > 1e-20 { + (ss2 * yc * yc + cc * ys * ys - 2.0 * cs * yc * ys) / denom2 / var_flux + } else { + 0.0 + } + }; + + let period = 1.0 / freq; + (period, power) + }) + .collect(); + + // 找到最大功率及其对应频率 + let (best_period, max_power) = + periodogram + .iter() + .fold((0.0, 0.0), |(best_p, max_p), &(p, power)| { + if power > max_p { + (p, power) + } else { + (best_p, max_p) + } + }); + + // 估算 FAP（基于峰值数量的简化方法） + let fap = estimate_fap_scargle(max_power, n); + + // 找次优周期（排除最优周期附近的频率） + let mut secondary = Vec::new(); + let period_window = best_period * 0.05; // ±5% 排除窗口 + for &(p, power) in &periodogram { + if (p - best_period).abs() > period_window && power > max_power * 0.3 { + secondary.push(PeriodCandidate { period: p, power }); + } + } + secondary.sort_by(|a, b| { + b.power + .partial_cmp(&a.power) + .unwrap_or(std::cmp::Ordering::Equal) + }); + secondary.truncate(5); + + // 折叠光变曲线 + let folded = fold_lightcurve(time, flux, best_period); + + Ok(PeriodSearchResult { + best_period, + best_power: max_power, + fap: Some(fap), + secondary_periods: secondary, + periodogram, + folded_lightcurve: Some(folded), + method: "Lomb-Scargle".to_string(), + }) +} + +/// 估计 τ 参数（Lomb-Scargle 的时间偏移） +fn estimate_tau(time: &[f64], omega: f64) -> f64 { + let (s, c) = time.iter().fold((0.0, 0.0), |(s, c), &t| { + (s + (2.0 * omega * t).sin(), c + (2.0 * omega * t).cos()) + }); + 0.5 * s.atan2(c) / omega +} + +/// 估算 FAP（False Alarm Probability） +fn estimate_fap_scargle(power: f64, n: usize) -> f64 { + // 简化公式：FAP ≈ (1 - P)^(N/2)（Horvath et al. 2022 近似） + // 更精确的方法需要有效自由度，这里用简化版本 + let neff = n as f64; + let fap = (1.0 - power).powf(neff / 2.0); + fap.clamp(0.0, 1.0) +} + +/// 折叠光变曲线 +pub fn fold_lightcurve(time: &[f64], flux: &[f64], period: f64) -> Vec<(f64, f64)> { + time.iter() + .zip(flux.iter()) + .map(|(&t, &f)| { + let phase = (t % period) / period; + let phase = if phase < 0.0 { phase + 1.0 } else { phase }; + (phase, f) + }) + .collect() +} + +/// BLS（Box Least Squares）周期搜索 +/// +/// 搜索方波（凌星）信号，返回最优周期、深度、持续时间 +pub fn bls( + time: &[f64], + flux: &[f64], + min_period: Option, + max_period: Option, + n_freq: usize, + min_duration_frac: f64, + max_duration_frac: f64, +) -> Result { + if time.len() != flux.len() { + return Err(anyhow!("时间和通量数组长度不一致")); + } + if time.len() < 10 { + return Err(anyhow!("数据点太少（需要 ≥10）")); + } + + let t_min = time.iter().cloned().fold(f64::INFINITY, f64::min); + let t_max = time.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + let baseline = t_max - t_min; + let min_p = min_period.unwrap_or(baseline / 20.0).max(0.5); + let max_p = max_period.unwrap_or(baseline / 2.0).min(baseline); + + let freq_min = 1.0 / max_p; + let freq_max = 1.0 / min_p; + let df = (freq_max - freq_min) / n_freq as f64; + + let mean_flux: f64 = flux.iter().sum::() / flux.len() as f64; + let var_flux: f64 = + flux.iter().map(|f| (f - mean_flux).powi(2)).sum::() / flux.len() as f64; + + let n_duration = 20; + let results: Vec<(f64, f64, f64, f64, f64)> = (0..n_freq) + .into_par_iter() + .map(|i| { + let freq = freq_min + df * i as f64; + let period = 1.0 / freq; + + // 扫描不同凌星持续时间 + let mut max_bls_power = 0.0; + let mut best_d_in_period = 0.0; + let mut best_offset = 0.0; + + for j in 0..n_duration { + let d_frac = min_duration_frac + + (max_duration_frac - min_duration_frac) * j as f64 / (n_duration - 1) as f64; + let duration = period * d_frac; + + // 将数据折叠到该周期 + let mut folded: Vec<(f64, f64)> = time + .iter() + .zip(flux.iter()) + .map(|(&t, &f)| { + let phase = (t % period) / period; + let phase = if phase < 0.0 { phase + 1.0 } else { phase }; + (phase, f) + }) + .collect(); + folded.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)); + + let d_in_phase = duration / period; + + // BLS：扫描凌星窗口起始相位 + let n_offsets = 50; + for k in 0..n_offsets { + let offset = k as f64 / n_offsets as f64; + + // 计算窗口内外的均值 + let mut in_sum = 0.0; + let mut in_count = 0; + let mut out_sum = 0.0; + let mut out_count = 0; + + for &(phase, flux_val) in &folded { + let in_window = if offset + d_in_phase <= 1.0 { + phase >= offset && phase < offset + d_in_phase + } else { + phase >= offset || phase < (offset + d_in_phase) - 1.0 + }; + + if in_window { + in_sum += flux_val; + in_count += 1; + } else { + out_sum += flux_val; + out_count += 1; + } + } + + if in_count == 0 || out_count == 0 { + continue; + } + + let in_mean = in_sum / in_count as f64; + let out_mean = out_sum / out_count as f64; + + // BLS 统计量（归一化） + let bls_power = + (out_mean - in_mean).powi(2) * in_count as f64 * out_count as f64 + / flux.len() as f64 + / var_flux; + + if bls_power > max_bls_power { + max_bls_power = bls_power; + best_d_in_period = d_in_phase; + best_offset = offset; + } + } + } + + let sde = if max_bls_power > 0.0 { + max_bls_power / var_flux * mean_flux.powi(2) + } else { + 0.0 + }; + + (period, max_bls_power, sde, best_d_in_period, best_offset) + }) + .collect(); + + let mut periodogram = Vec::with_capacity(n_freq); + let mut max_sde = 0.0; + let mut best_period = 0.0; + let mut best_depth = 0.0; + let mut best_duration = 0.0; + let mut best_phase = 0.0; + + for (period, max_bls_power, sde, best_d_in_period, best_offset) in results { + periodogram.push((period, max_bls_power)); + + if sde > max_sde { + max_sde = sde; + best_period = period; + best_depth = 1.0 - (mean_flux - max_bls_power.sqrt() * mean_flux) / mean_flux; + best_duration = period * best_d_in_period; + best_phase = best_offset; + } + } + + Ok(BlsResult { + best_period, + depth: best_depth.abs().min(1.0), + duration: best_duration, + phase: best_phase, + sde: max_sde, + periodogram, + }) +} + +/// 格式化周期搜索结果为 Markdown +pub fn format_period_result(result: &PeriodSearchResult) -> String { + let mut out = format!("## 周期搜索结果（{}）\n\n", result.method); + out.push_str(&format!("- **最优周期**: {:.6} 天\n", result.best_period)); + out.push_str(&format!("- **周期功率**: {:.4}\n", result.best_power)); + if let Some(fap) = result.fap { + out.push_str(&format!("- **FAP**: {:.6e}\n", fap)); + if fap < 0.001 { + out.push_str(" - **高度显著** (FAP < 0.1%)\n"); + } else if fap < 0.01 { + out.push_str(" - **显著** (FAP < 1%)\n"); + } else if fap < 0.05 { + out.push_str(" - **可能显著** (FAP < 5%)\n"); + } else { + out.push_str(" - **不显著** (FAP ≥ 5%)\n"); + } + } + + if !result.secondary_periods.is_empty() { + out.push_str("\n### 次优周期\n\n"); + out.push_str("| 周期 (天) | 功率 |\n|----------|------|\n"); + for p in &result.secondary_periods { + out.push_str(&format!("| {:.6} | {:.4} |\n", p.period, p.power)); + } + } + + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_lomb_scargle_sine_wave() { + // 生成一个周期为 2.0 天的正弦波 + let n = 100; + let period = 2.0; + let mut time = Vec::new(); + let mut flux = Vec::new(); + for i in 0..n { + let t = i as f64 * 0.1 + 0.05 * (i as f64 * 0.3).sin(); // 非均匀采样 + time.push(t); + flux.push( + (2.0 * std::f64::consts::PI * t / period).sin() + 0.1 * (i as f64 * 0.7).sin(), + ); + } + + let result = lomb_scargle(&time, &flux, Some(0.5), Some(5.0), 1000).unwrap(); + assert!( + (result.best_period - period).abs() < 0.1, + "检测到的周期 {} 应接近 {}", + result.best_period, + period + ); + } + + #[test] + fn test_fold_lightcurve() { + let time = vec![0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]; + let flux = vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0]; + let folded = fold_lightcurve(&time, &flux, 1.0); + assert_eq!(folded.len(), 7); + // 相位应在 [0, 1) + for (phase, _) in &folded { + assert!(*phase >= 0.0 && *phase < 1.0); + } + } +}