refactor: 服务层抽象下沉、异步锁全栈迁移、客户端韧性加固与移动端适配

- 服务层拆分:删除 api/helpers.rs,新增 citation/note/session/pipeline/paper/vision 独立服务模块
  - Agent 工具精简:paper_content+paper_outline 合并为 paper.rs,图片分析逻辑下沉至 services/vision
  - 并发模型升级:std::sync::{Mutex,RwLock} → tokio::sync::{Mutex,RwLock},消除 async
  上下文中的阻塞风险
  - 客户端加固:HTTP 客户端统一超时配置、ADS 429 / arXiv 503 自动重试、构造函数返回 Result
  - 启动安全:全局 panic hook 日志化、空密码拒绝启动、向量表维度不匹配需显式确认
  - CLI 扩展:构建完整 AppState 复用服务层,新增 Content/Outline/Citations/Search/Process 子命令
  - 前端:移动端汉堡菜单、侧栏滑出面板、引用星系触屏手势(单指拖拽/双指缩放)
This commit is contained in:
Asfmq 2026-06-30 19:26:01 +08:00
parent c5fd5b0d66
commit f885c0a4a8
90 changed files with 5184 additions and 3916 deletions

View File

@ -44,6 +44,8 @@ LLM_MODEL=gpt-4o-mini
# EMBEDDING_MODEL=text-embedding-3-small # EMBEDDING_MODEL=text-embedding-3-small
# 向量维度(需与所选模型输出维度一致,默认 1536 # 向量维度(需与所选模型输出维度一致,默认 1536
# EMBEDDING_DIM=1536 # EMBEDDING_DIM=1536
# 维度变更确认:修改 EMBEDDING_DIM 后需设置此变量为 1 才能启动(防止误改导致数据丢失)
# EMBEDDING_DIM_FORCE_REBUILD=1
# ───────────────────────────────────────────────────────────────────────────── # ─────────────────────────────────────────────────────────────────────────────
# Chunker 结构化切片参数(影响向量化粒度) # Chunker 结构化切片参数(影响向量化粒度)
@ -85,8 +87,8 @@ LIBRARY_DIR=./library
SKILLS_DIR=./skills SKILLS_DIR=./skills
# 后端服务监听端口 # 后端服务监听端口
PORT=8000 PORT=8000
# 认证密码 # 认证密码(必填,启动时未设置将拒绝启动)
ADMIN_PASSWORD=fmq ADMIN_PASSWORD=your_strong_password_here
# ───────────────────────────────────────────────────────────────────────────── # ─────────────────────────────────────────────────────────────────────────────
# 5. 日志配置 # 5. 日志配置

View File

@ -21,7 +21,7 @@ path = "src/bin/cli.rs"
[dependencies] [dependencies]
tokio = { version = "1", features = ["full"] } tokio = { version = "1", features = ["full"] }
axum = { version = "0.7", features = ["macros", "multipart"] } axum = { version = "0.7", features = ["macros", "multipart"] }
tower-http = { version = "0.5", features = ["cors", "fs", "trace"] } tower-http = { version = "0.5", features = ["cors", "fs", "trace", "set-header"] }
sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "sqlite", "chrono", "json"] } sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "sqlite", "chrono", "json"] }
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0" serde_json = "1.0"

View File

@ -1,7 +1,7 @@
// dashboard/src/App.tsx // dashboard/src/App.tsx
import { useState, useEffect, useCallback, useRef } from 'react'; import { useState, useEffect, useCallback, useRef } from 'react';
import axios from 'axios'; import axios from 'axios';
import { Loader, BookOpen, GitFork, Lock } from 'lucide-react'; import { Loader, BookOpen, GitFork, Lock, Menu } from 'lucide-react';
import { Sidebar } from './components/layout/Sidebar'; import { Sidebar } from './components/layout/Sidebar';
import { SearchPanel } from './pages/SearchPanel'; import { SearchPanel } from './pages/SearchPanel';
import { LibraryPanel } from './pages/LibraryPanel'; import { LibraryPanel } from './pages/LibraryPanel';
@ -24,6 +24,9 @@ import { useNotes } from './hooks/useNotes';
import { useCitations } from './hooks/useCitations'; import { useCitations } from './hooks/useCitations';
export default function App() { export default function App() {
// 移动端菜单显示状态
const [isMobileSidebarOpen, setIsMobileSidebarOpen] = useState(false);
// 1. 全局自定义 Dialog 弹窗管理 (Alert / Confirm) // 1. 全局自定义 Dialog 弹窗管理 (Alert / Confirm)
const [dialog, setDialog] = useState<{ const [dialog, setDialog] = useState<{
type: 'alert' | 'confirm'; type: 'alert' | 'confirm';
@ -318,11 +321,41 @@ export default function App() {
selectedPaper={library.selectedPaper} selectedPaper={library.selectedPaper}
loadCitations={citations.loadCitations} loadCitations={citations.loadCitations}
onLogout={auth.handleLogout} onLogout={auth.handleLogout}
isOpen={isMobileSidebarOpen}
onClose={() => setIsMobileSidebarOpen(false)}
/> />
{/* 主工作区 */} {/* 主工作区 */}
<main className="flex-1 flex flex-col overflow-hidden relative"> <main className="flex-1 flex flex-col overflow-hidden relative">
<div className="flex-1 overflow-y-auto p-4 sm:p-6 md:p-8 relative z-10 w-full flex flex-col"> {/* 移动端顶部 Header */}
<header className="lg:hidden flex items-center justify-between px-4 py-3 bg-white border-b border-slate-200 select-none shrink-0 z-20">
<div className="flex items-center gap-3">
<button
type="button"
onClick={() => setIsMobileSidebarOpen(true)}
className="p-1.5 rounded-lg border border-slate-200 hover:bg-slate-55 text-slate-600 transition-all cursor-pointer flex items-center justify-center"
title="打开菜单"
>
<Menu className="w-5 h-5" />
</button>
<div className="flex items-center gap-2">
<div className="w-6 h-6">
<Logo gradientId="mobileHeaderStarGrad" />
</div>
<span className="text-xs font-bold text-slate-800 tracking-wider">AstroResearch</span>
</div>
</div>
<span className="text-[10px] font-bold px-2.5 py-1 rounded bg-slate-100 text-slate-600 font-sans tracking-wide uppercase">
{activeTab === 'search' && '统一检索'}
{activeTab === 'library' && '馆藏管理'}
{activeTab === 'reader' && '双语阅读'}
{activeTab === 'citation' && '引用星系'}
{activeTab === 'sync' && '批量任务'}
{activeTab === 'agent' && '智能科研'}
</span>
</header>
<div className="flex-1 overflow-y-auto p-4 sm:p-6 md:p-8 relative w-full flex flex-col">
<div className={`w-full flex-1 flex flex-col min-h-0 ${ <div className={`w-full flex-1 flex flex-col min-h-0 ${
(activeTab === 'reader' || activeTab === 'citation' || activeTab === 'agent') ? 'max-w-none' : 'max-w-7xl mx-auto' (activeTab === 'reader' || activeTab === 'citation' || activeTab === 'agent') ? 'max-w-none' : 'max-w-7xl mx-auto'
}`}> }`}>

View File

@ -298,6 +298,29 @@ export function CitationGalaxyCanvas({ networks, activeNetwork, nodeLimit, onNod
render(); render();
let initialTouchDistance = 0;
let initialScale = 1.0;
const findNodeAtCoordinate = (clientX: number, clientY: number) => {
const rect = canvas.getBoundingClientRect();
const touchX = clientX - rect.left;
const touchY = clientY - rect.top;
const cx = rect.width / 2;
const cy = rect.height / 2;
const gx = (touchX - cx - offsetX) / scale + cx;
const gy = (touchY - cy - offsetY) / scale + cy;
for (const node of nodes) {
const dx = node.x - gx;
const dy = node.y - gy;
const dist = Math.sqrt(dx * dx + dy * dy);
if (dist < node.radius + 12) { // 移动端稍微加大触碰敏感区
return node;
}
}
return null;
};
const handleMouseDown = (e: MouseEvent) => { const handleMouseDown = (e: MouseEvent) => {
isDragging = true; isDragging = true;
dragStartX = e.clientX - offsetX; dragStartX = e.clientX - offsetX;
@ -379,12 +402,68 @@ export function CitationGalaxyCanvas({ networks, activeNetwork, nodeLimit, onNod
offsetY = mouseY - cy - (gy - cy) * scale; offsetY = mouseY - cy - (gy - cy) * scale;
}; };
// 触屏手势处理
const handleTouchStart = (e: TouchEvent) => {
if (e.touches.length === 1) {
isDragging = true;
hasDragged = false;
dragStartX = e.touches[0].clientX - offsetX;
dragStartY = e.touches[0].clientY - offsetY;
// 触碰选中节点
hoveredNode = findNodeAtCoordinate(e.touches[0].clientX, e.touches[0].clientY);
} else if (e.touches.length === 2) {
isDragging = false;
const dx = e.touches[0].clientX - e.touches[1].clientX;
const dy = e.touches[0].clientY - e.touches[1].clientY;
initialTouchDistance = Math.sqrt(dx * dx + dy * dy) || 1;
initialScale = scale;
}
};
const handleTouchMove = (e: TouchEvent) => {
if (isDragging && e.touches.length === 1) {
const dx = e.touches[0].clientX - dragStartX;
const dy = e.touches[0].clientY - dragStartY;
if (Math.sqrt((dx - offsetX) ** 2 + (dy - offsetY) ** 2) > 3) {
hasDragged = true;
}
offsetX = dx;
offsetY = dy;
e.preventDefault(); // 阻止滚动
} else if (e.touches.length === 2) {
const dx = e.touches[0].clientX - e.touches[1].clientX;
const dy = e.touches[0].clientY - e.touches[1].clientY;
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
// 双指开合缩放
const newScale = initialScale * (dist / initialTouchDistance);
scale = Math.max(0.15, Math.min(5.0, newScale));
e.preventDefault(); // 阻止默认双指缩放
}
};
const handleTouchEnd = () => {
isDragging = false;
if (!hasDragged && hoveredNode) {
if (hoveredNode.id !== activeNetwork.bibcode) {
onNodeClick(hoveredNode.id);
}
}
hoveredNode = null;
};
canvas.addEventListener('mousedown', handleMouseDown); canvas.addEventListener('mousedown', handleMouseDown);
canvas.addEventListener('mousemove', handleMouseMove); canvas.addEventListener('mousemove', handleMouseMove);
canvas.addEventListener('mouseup', handleMouseUp); canvas.addEventListener('mouseup', handleMouseUp);
canvas.addEventListener('mouseleave', handleMouseLeave); canvas.addEventListener('mouseleave', handleMouseLeave);
canvas.addEventListener('click', handleCanvasClick); canvas.addEventListener('click', handleCanvasClick);
canvas.addEventListener('wheel', handleWheel, { passive: false }); canvas.addEventListener('wheel', handleWheel, { passive: false });
// 注册触控事件
canvas.addEventListener('touchstart', handleTouchStart, { passive: false });
canvas.addEventListener('touchmove', handleTouchMove, { passive: false });
canvas.addEventListener('touchend', handleTouchEnd);
return () => { return () => {
cancelAnimationFrame(animationFrameId); cancelAnimationFrame(animationFrameId);
@ -394,6 +473,11 @@ export function CitationGalaxyCanvas({ networks, activeNetwork, nodeLimit, onNod
canvas.removeEventListener('mouseleave', handleMouseLeave); canvas.removeEventListener('mouseleave', handleMouseLeave);
canvas.removeEventListener('click', handleCanvasClick); canvas.removeEventListener('click', handleCanvasClick);
canvas.removeEventListener('wheel', handleWheel); canvas.removeEventListener('wheel', handleWheel);
// 注销触控事件
canvas.removeEventListener('touchstart', handleTouchStart);
canvas.removeEventListener('touchmove', handleTouchMove);
canvas.removeEventListener('touchend', handleTouchEnd);
}; };
}, [networks, activeNetwork, onNodeClick]); }, [networks, activeNetwork, onNodeClick]);

View File

@ -190,7 +190,7 @@ export function AgentMessageList({
: '探索性科研研讨'} : '探索性科研研讨'}
</h3> </h3>
<p className="text-[10px] text-slate-400 font-semibold mt-0.5"> <p className="text-[10px] text-slate-400 font-semibold mt-0.5">
(ReAct Thought Action Observation)
</p> </p>
</div> </div>
</div> </div>

View File

@ -37,173 +37,206 @@ export function AgentSessionSidebar({
onNewSession, onNewSession,
onDeleteSession, onDeleteSession,
}: AgentSessionSidebarProps) { }: AgentSessionSidebarProps) {
return (
<div className={`transition-all duration-300 ease-in-out ${collapsed ? 'w-0 overflow-hidden opacity-0 border-r-0' : 'w-64 border-r border-slate-200'} bg-slate-50 flex flex-col justify-between shrink-0 select-none`}> const handleSelectSession = (id: string | null) => {
<div className="flex flex-col min-h-0 flex-1"> setCurrentSessionId(id);
<div className="p-4 border-b border-slate-200 bg-white flex items-center justify-between shrink-0"> if (typeof window !== 'undefined' && window.innerWidth < 1024) {
<span className="text-xs font-extrabold text-slate-800 tracking-wider flex items-center gap-1.5"> onToggleCollapse(true);
<Clock className="w-3.5 h-3.5 text-blueprint" /> }
{!collapsed && <span></span>} };
</span>
<div className="flex items-center gap-1 shrink-0">
<button
onClick={onNewSession}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-600 hover:text-slate-800 transition-all cursor-pointer shadow-2xs"
title="新建会话"
>
<Plus className="w-3.5 h-3.5" />
</button>
<button
onClick={() => onToggleCollapse(true)}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-500 hover:text-slate-700 transition-all cursor-pointer"
title="收起侧栏"
>
<PanelLeftClose className="w-3.5 h-3.5" />
</button>
</div>
</div>
{/* 搜索框区域 */} const handleCreateNewSession = () => {
<div className="px-3 py-2 border-b border-slate-200/60 bg-white flex flex-col gap-1.5 shrink-0"> onNewSession();
<div className="relative"> if (typeof window !== 'undefined' && window.innerWidth < 1024) {
<input onToggleCollapse(true);
type="text" }
value={searchQuery} };
onChange={e => setSearchQuery(e.target.value)}
placeholder="搜索历史会话或消息内容..." return (
className="w-full pl-8 pr-7 py-1.5 rounded-md bg-slate-50 border border-slate-200 text-slate-800 placeholder-slate-400 focus:outline-none focus:border-blueprint focus:bg-white transition-all text-[11px] font-medium" <>
/> {/* 移动端背景遮罩层 */}
<div className="absolute left-2.5 top-2.5 text-slate-400"> {!collapsed && (
<Search className="w-3.5 h-3.5 text-slate-400" /> <button
</div> type="button"
{searchQuery && ( onClick={() => onToggleCollapse(true)}
className="fixed inset-0 bg-slate-900/30 backdrop-blur-xs z-35 lg:hidden cursor-pointer w-full h-full border-none outline-none"
aria-label="关闭侧栏"
/>
)}
<div
className={`bg-slate-50 flex flex-col justify-between shrink-0 select-none transition-all duration-300 ease-in-out fixed inset-y-0 left-0 z-40 lg:relative lg:translate-x-0 ${
collapsed
? '-translate-x-full lg:translate-x-0 lg:w-0 lg:overflow-hidden lg:opacity-0 lg:border-r-0'
: 'translate-x-0 w-64 border-r border-slate-200 lg:w-64'
}`}
>
<div className="flex flex-col min-h-0 flex-1">
<div className="p-4 border-b border-slate-200 bg-white flex items-center justify-between shrink-0">
<span className="text-xs font-extrabold text-slate-800 tracking-wider flex items-center gap-1.5">
<Clock className="w-3.5 h-3.5 text-blueprint" />
<span></span>
</span>
<div className="flex items-center gap-1 shrink-0">
<button <button
onClick={() => setSearchQuery('')} onClick={handleCreateNewSession}
className="absolute right-2 top-2 p-0.5 rounded-full hover:bg-slate-200 text-slate-400 hover:text-slate-600 transition-colors cursor-pointer" className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-650 hover:text-slate-800 transition-all cursor-pointer shadow-2xs"
title="新建会话"
> >
<X className="w-3 h-3" /> <Plus className="w-3.5 h-3.5" />
</button> </button>
<button
onClick={() => onToggleCollapse(true)}
className="p-1 rounded-md border border-slate-200 bg-white hover:bg-slate-50 text-slate-555 hover:text-slate-700 transition-all cursor-pointer"
title="收起侧栏"
>
<PanelLeftClose className="w-3.5 h-3.5" />
</button>
</div>
</div>
{/* 搜索框区域 */}
<div className="px-3 py-2 border-b border-slate-200/60 bg-white flex flex-col gap-1.5 shrink-0">
<div className="relative">
<input
type="text"
value={searchQuery}
onChange={e => setSearchQuery(e.target.value)}
placeholder="搜索历史会话或消息内容..."
className="w-full pl-8 pr-7 py-1.5 rounded-md bg-slate-50 border border-slate-200 text-slate-800 placeholder-slate-400 focus:outline-none focus:border-blueprint focus:bg-white transition-all text-[11px] font-medium"
/>
<div className="absolute left-2.5 top-2.5 text-slate-400">
<Search className="w-3.5 h-3.5 text-slate-400" />
</div>
{searchQuery && (
<button
onClick={() => setSearchQuery('')}
className="absolute right-2 top-2 p-0.5 rounded-full hover:bg-slate-200 text-slate-400 hover:text-slate-600 transition-colors cursor-pointer"
>
<X className="w-3 h-3" />
</button>
)}
</div>
{currentSessionId && (
<label className="flex items-center gap-1.5 text-[10px] text-slate-500 cursor-pointer font-medium select-none ml-0.5">
<input
type="checkbox"
checked={searchScopeOnlyCurrent}
onChange={e => setSearchScopeOnlyCurrent(e.target.checked)}
className="rounded text-blueprint border-slate-300 focus:ring-blueprint w-3 h-3 cursor-pointer"
/>
<span></span>
</label>
)} )}
</div> </div>
{currentSessionId && (
<label className="flex items-center gap-1.5 text-[10px] text-slate-500 cursor-pointer font-medium select-none ml-0.5">
<input
type="checkbox"
checked={searchScopeOnlyCurrent}
onChange={e => setSearchScopeOnlyCurrent(e.target.checked)}
className="rounded text-blueprint border-slate-300 focus:ring-blueprint w-3 h-3 cursor-pointer"
/>
<span></span>
</label>
)}
</div>
<div className="flex-1 overflow-y-auto p-2.5 space-y-1.5 scrollbar-thin"> <div className="flex-1 overflow-y-auto p-2.5 space-y-1.5 scrollbar-thin">
{searchQuery.trim() ? ( {searchQuery.trim() ? (
loadingSearch ? ( loadingSearch ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2"> <div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" /> <Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span> <span>...</span>
</div> </div>
) : searchResults.length === 0 ? ( ) : searchResults.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic"> <div className="text-center py-12 text-slate-400 text-[11px] italic">
</div> </div>
) : ( ) : (
searchResults.map((result, idx) => { searchResults.map((result, idx) => {
const isActive = result.session_id === currentSessionId; const isActive = result.session_id === currentSessionId;
const isMessage = result.result_type.startsWith('message/'); const isMessage = result.result_type.startsWith('message/');
const role = isMessage ? result.result_type.split('/')[1] : null; const role = isMessage ? result.result_type.split('/')[1] : null;
return ( return (
<button
key={`${result.session_id}-${idx}`}
onClick={() => setCurrentSessionId(result.session_id)}
className={`w-full text-left p-2.5 rounded-md border transition-all duration-200 flex flex-col gap-1 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-white hover:bg-slate-100 text-slate-650 shadow-3xs'
}`}
>
<div className="flex items-center justify-between w-full">
<span className="text-[10px] font-bold text-slate-400 flex items-center gap-1">
{isMessage ? (
<>
<MessageSquare className="w-3 h-3 text-slate-400" />
<span>{role === 'user' ? '提问' : '解答'}</span>
</>
) : (
<>
<BookOpen className="w-3 h-3 text-blueprint" />
<span></span>
</>
)}
</span>
{result.created_at && (
<span className="text-[8px] text-slate-400 font-mono">
{result.created_at.split(' ')[0]}
</span>
)}
</div>
<span className="text-xs leading-snug line-clamp-1 block text-left font-bold text-slate-800">
{result.title || '无标题会话'}
</span>
<p
className="text-[10px] text-slate-500 font-medium leading-normal block text-left line-clamp-3 bg-slate-50/50 p-1.5 rounded border border-slate-100/50"
dangerouslySetInnerHTML={{ __html: result.snippet }}
/>
</button>
);
})
)
) : (
// 正常的会话列表渲染
loadingSessions ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span>
</div>
) : sessions.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic">
</div>
) : (
sessions.map(session => {
const isActive = session.session_id === currentSessionId;
return (
<button
key={session.session_id}
onClick={() => setCurrentSessionId(session.session_id)}
className={`w-full text-left p-3 rounded-md border transition-all duration-200 group flex items-start justify-between gap-2 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-transparent hover:bg-slate-100 text-slate-655'
}`}
>
<div className="flex-1 min-w-0 flex flex-col gap-1 text-left">
<span className="text-xs leading-snug line-clamp-2 block text-left">
{session.title || '无标题会话'}
</span>
<span className="text-[9px] text-slate-400 font-semibold block text-left">
{session.turn_count} {session.updated_at.split(' ')[0]}
</span>
</div>
<button <button
onClick={(e) => onDeleteSession(session.session_id, e)} key={`${result.session_id}-${idx}`}
className="text-slate-400 hover:text-red-655 p-1 rounded-md opacity-0 group-hover:opacity-100 hover:bg-white border border-transparent hover:border-slate-200 transition-all cursor-pointer shrink-0" onClick={() => handleSelectSession(result.session_id)}
title="删除此会话" className={`w-full text-left p-2.5 rounded-md border transition-all duration-200 flex flex-col gap-1 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-white hover:bg-slate-100 text-slate-650 shadow-3xs'
}`}
> >
<Trash2 className="w-3 h-3" /> <div className="flex items-center justify-between w-full">
<span className="text-[10px] font-bold text-slate-400 flex items-center gap-1">
{isMessage ? (
<>
<MessageSquare className="w-3 h-3 text-slate-400" />
<span>{role === 'user' ? '提问' : '解答'}</span>
</>
) : (
<>
<BookOpen className="w-3 h-3 text-blueprint" />
<span></span>
</>
)}
</span>
{result.created_at && (
<span className="text-[8px] text-slate-400 font-mono">
{result.created_at.split(' ')[0]}
</span>
)}
</div>
<span className="text-xs leading-snug line-clamp-1 block text-left font-bold text-slate-800">
{result.title || '无标题会话'}
</span>
<p
className="text-[10px] text-slate-500 font-medium leading-normal block text-left line-clamp-3 bg-slate-50/50 p-1.5 rounded border border-slate-100/50"
dangerouslySetInnerHTML={{ __html: result.snippet }}
/>
</button> </button>
</button> );
); })
}) )
) ) : (
)} // 正常的会话列表渲染
loadingSessions ? (
<div className="flex items-center justify-center p-8 text-slate-400 text-xs gap-2">
<Loader className="w-3.5 h-3.5 animate-spin text-blueprint" />
<span>...</span>
</div>
) : sessions.length === 0 ? (
<div className="text-center py-12 text-slate-400 text-[11px] italic">
</div>
) : (
sessions.map(session => {
const isActive = session.session_id === currentSessionId;
return (
<button
key={session.session_id}
onClick={() => handleSelectSession(session.session_id)}
className={`w-full text-left p-3 rounded-md border transition-all duration-200 group flex items-start justify-between gap-2 cursor-pointer ${
isActive
? 'bg-blueprint/5 border-blueprint text-blueprint font-bold shadow-2xs'
: 'border-transparent bg-transparent hover:bg-slate-100 text-slate-655'
}`}
>
<div className="flex-1 min-w-0 flex flex-col gap-1 text-left">
<span className="text-xs leading-snug line-clamp-2 block text-left">
{session.title || '无标题会话'}
</span>
<span className="text-[9px] text-slate-400 font-semibold block text-left">
{session.turn_count} {session.updated_at.split(' ')[0]}
</span>
</div>
<button
onClick={(e) => onDeleteSession(session.session_id, e)}
className="text-slate-400 hover:text-red-655 p-1 rounded-md opacity-0 group-hover:opacity-100 hover:bg-white border border-transparent hover:border-slate-200 transition-all cursor-pointer shrink-0"
title="删除此会话"
>
<Trash2 className="w-3 h-3" />
</button>
</button>
);
})
)
)}
</div>
</div> </div>
</div> </div>
</div> </>
); );
} }

View File

@ -75,7 +75,7 @@ export function ToolCallCard({
<div className="flex items-center justify-between"> <div className="flex items-center justify-between">
<div className="flex items-center gap-1.5 text-[10px] font-extrabold text-emerald-700 tracking-wider uppercase"> <div className="flex items-center gap-1.5 text-[10px] font-extrabold text-emerald-700 tracking-wider uppercase">
<Eye className="w-3.5 h-3.5" /> <Eye className="w-3.5 h-3.5" />
<span> (Observation)</span> <span></span>
{result!.isError && ( {result!.isError && (
<span className="px-1.5 py-0.2 rounded bg-rose-50 text-rose-700 border border-rose-100 text-[9px]"> <span className="px-1.5 py-0.2 rounded bg-rose-50 text-rose-700 border border-rose-100 text-[9px]">

View File

@ -12,207 +12,225 @@ interface SidebarProps {
selectedPaper: StandardPaper | null; selectedPaper: StandardPaper | null;
loadCitations: (bibcode: string) => void; loadCitations: (bibcode: string) => void;
onLogout: () => void; onLogout: () => void;
isOpen?: boolean;
onClose?: () => void;
} }
export function Sidebar({ activeTab, setActiveTab, selectedPaper, loadCitations, onLogout }: SidebarProps) { export function Sidebar({ activeTab, setActiveTab, selectedPaper, loadCitations, onLogout, isOpen = false, onClose }: SidebarProps) {
const [isCollapsed, setIsCollapsed] = useState(false); const [isCollapsed, setIsCollapsed] = useState(false);
const effectiveCollapsed = isCollapsed && !isOpen;
return ( return (
<aside <>
className={`bg-slate-50 border-r border-slate-200 flex flex-col justify-between py-6 z-20 select-none transition-all duration-355 cubic-bezier(0.4, 0, 0.2, 1) ${ {/* 移动端背景遮罩层 */}
isCollapsed ? 'w-16 px-2' : 'w-64 px-4' {isOpen && (
}`}
>
<div>
{/* 系统LOGO与折叠控制区 */}
<div className={`mb-8 flex items-center transition-all duration-300 ${
isCollapsed ? 'justify-center px-0' : 'justify-between px-3'
}`}>
{/* Logo & 标题文字 */}
<div className="flex items-center gap-3 min-w-0">
{/* Logo按钮 (只在折叠状态下可点击展开) */}
<button
type="button"
disabled={!isCollapsed}
onClick={() => setIsCollapsed(false)}
className={`flex items-center justify-center shrink-0 rounded-lg transition-all duration-300 ${
isCollapsed
? 'w-11 h-11 bg-white hover:bg-slate-50 border border-slate-200 cursor-pointer shadow-xs hover:shadow-sm'
: 'w-9 h-9 bg-transparent border border-transparent cursor-default'
}`}
title={isCollapsed ? "展开导航" : undefined}
>
<div className={`transition-all duration-300 flex items-center justify-center ${isCollapsed ? 'w-8 h-8' : 'w-9 h-9'}`}>
<Logo gradientId="sidebarStarGrad" />
</div>
</button>
{/* 系统说明文字 */}
<div
className={`flex flex-col transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 scale-95 -translate-x-2 pointer-events-none select-none overflow-hidden h-0'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0'
}`}
>
<h1 className="text-sm font-bold text-slate-800 tracking-wider whitespace-nowrap">AstroResearch</h1>
<span className="text-[11px] text-slate-500 block font-medium font-sans whitespace-nowrap"></span>
</div>
</div>
{/* 折叠控制按钮 (圆润悬浮效果) */}
<button
type="button"
onClick={() => setIsCollapsed(true)}
className={`p-1.5 rounded-full bg-slate-100 hover:bg-slate-200 border border-slate-200 text-slate-500 hover:text-slate-800 transition-all duration-300 cursor-pointer shadow-2xs shrink-0 flex items-center justify-center hover:scale-105 active:scale-95 ${
isCollapsed
? 'opacity-0 scale-75 pointer-events-none w-0 h-0 p-0 border-0 overflow-hidden'
: 'opacity-100 scale-100'
}`}
title="收起导航"
>
<ChevronLeft className="w-3.5 h-3.5" />
</button>
</div>
{/* 导航菜单列表 */}
<nav className="space-y-1">
{[
{ id: 'search' as TabId, label: '统一检索', icon: Search },
{ id: 'library' as TabId, label: '馆藏管理', icon: Library },
{ id: 'reader' as TabId, label: '双语阅读', icon: BookOpen },
{ id: 'citation' as TabId, label: '引用星系', icon: GitFork },
{ id: 'sync' as TabId, label: '批量任务', icon: RefreshCw },
{ id: 'agent' as TabId, label: '智能科研', icon: Sparkles },
].map((tab) => {
const Icon = tab.icon;
const isActive = activeTab === tab.id;
return (
<button
key={tab.id}
title={isCollapsed ? tab.label : undefined}
onClick={() => {
setActiveTab(tab.id);
if (tab.id === 'citation' && selectedPaper) {
loadCitations(selectedPaper.bibcode);
}
}}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border ${
isCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
} ${
isActive
? 'bg-slate-100 border-slate-200 text-slate-850 shadow-xs'
: 'border-transparent text-slate-650 hover:bg-slate-100 hover:text-slate-800'
}`}
>
<Icon className={`w-4 h-4 shrink-0 transition-colors duration-300 ${isActive ? 'text-slate-700' : 'text-slate-500'}`} />
<span
className={`truncate transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
}`}
>
{tab.label}
</span>
</button>
);
})}
</nav>
</div>
{/* 底部当前选定文献提示 (平滑动画版本) 与 退出登录 */}
<div className="space-y-3">
{selectedPaper ? (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
isCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3.5 border-slate-200 bg-slate-100/50'
}`}
title={isCollapsed ? `当前选定文献: ${selectedPaper.title}` : undefined}
>
{/* 折叠下的微型图书图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-50 text-slate-650 shrink-0 shadow-xs ${
isCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<BookOpen className="w-4 h-4" />
</div>
{/* 展开下的完整详情 */}
<div
className={`transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-40'
}`}
>
<span className="text-[9px] font-bold text-slate-500 tracking-widest block mb-1"></span>
<h4 className="text-xs text-slate-800 font-bold line-clamp-2 mb-2 leading-relaxed">{selectedPaper.title}</h4>
<div className="flex items-center justify-between text-[10px] font-medium text-slate-500 gap-2">
<span className="shrink-0">: {selectedPaper.year}</span>
<span className="truncate max-w-[90px] font-mono">{selectedPaper.bibcode}</span>
</div>
</div>
</div>
) : (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
isCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3 border-slate-200 bg-slate-100/30'
}`}
title={isCollapsed ? "未选定研究目标" : undefined}
>
{/* 折叠下的微型馆藏图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-100/30 text-slate-400 shrink-0 ${
isCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<Library className="w-4 h-4" />
</div>
{/* 展开下的提示文字 */}
<div
className={`text-center transition-all duration-300 origin-left ${
isCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-12'
}`}
>
<span className="text-[10px] text-slate-400 font-medium tracking-wide"></span>
</div>
</div>
)}
{/* 退出登录按钮 */}
<button <button
type="button" type="button"
onClick={onLogout} onClick={onClose}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border border-transparent text-slate-500 hover:bg-rose-50/50 hover:text-rose-600 cursor-pointer group ${ className="fixed inset-0 bg-slate-900/30 backdrop-blur-xs z-35 lg:hidden cursor-pointer w-full h-full border-none outline-none"
isCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5' aria-label="关闭导航"
}`} />
title={isCollapsed ? "退出登录" : undefined} )}
>
<LogOut className="w-4 h-4 shrink-0 text-slate-400 group-hover:text-rose-500 transition-colors" /> <aside
<span className={`bg-slate-50 border-r border-slate-200 flex flex-col justify-between py-6 z-40 select-none transition-all duration-355 cubic-bezier(0.4, 0, 0.2, 1) fixed inset-y-0 left-0 lg:relative lg:translate-x-0 ${
className={`truncate transition-all duration-300 origin-left ${ isOpen ? 'translate-x-0' : '-translate-x-full lg:translate-x-0'
isCollapsed } ${
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2' isOpen ? 'w-64 px-4' : (effectiveCollapsed ? 'w-16 px-2' : 'w-64 px-4')
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3' }`}
>
<div>
{/* 系统LOGO与折叠控制区 */}
<div className={`mb-8 flex items-center transition-all duration-300 ${
effectiveCollapsed ? 'justify-center px-0' : 'justify-between px-3'
}`}>
{/* Logo & 标题文字 */}
<div className="flex items-center gap-3 min-w-0">
{/* Logo按钮 (只在折叠状态下可点击展开) */}
<button
type="button"
disabled={!effectiveCollapsed}
onClick={() => setIsCollapsed(false)}
className={`flex items-center justify-center shrink-0 rounded-lg transition-all duration-300 ${
effectiveCollapsed
? 'w-11 h-11 bg-white hover:bg-slate-55 border border-slate-200 cursor-pointer shadow-xs hover:shadow-sm'
: 'w-9 h-9 bg-transparent border border-transparent cursor-default'
}`}
title={effectiveCollapsed ? "展开导航" : undefined}
>
<div className={`transition-all duration-300 flex items-center justify-center ${effectiveCollapsed ? 'w-8 h-8' : 'w-9 h-9'}`}>
<Logo gradientId="sidebarStarGrad" />
</div>
</button>
{/* 系统说明文字 */}
<div
className={`flex flex-col transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 scale-95 -translate-x-2 pointer-events-none select-none overflow-hidden h-0'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0'
}`}
>
<h1 className="text-sm font-bold text-slate-800 tracking-wider whitespace-nowrap">AstroResearch</h1>
<span className="text-[11px] text-slate-500 block font-medium font-sans whitespace-nowrap"></span>
</div>
</div>
{/* 折叠控制按钮 (圆润悬浮效果,仅在桌面端显示) */}
<button
type="button"
onClick={() => setIsCollapsed(true)}
className={`p-1.5 rounded-full bg-slate-100 hover:bg-slate-200 border border-slate-200 text-slate-505 hover:text-slate-800 transition-all duration-300 cursor-pointer shadow-2xs shrink-0 items-center justify-center hover:scale-105 active:scale-95 hidden lg:flex ${
effectiveCollapsed
? 'opacity-0 scale-75 pointer-events-none w-0 h-0 p-0 border-0 overflow-hidden'
: 'opacity-100 scale-100'
}`}
title="收起导航"
>
<ChevronLeft className="w-3.5 h-3.5" />
</button>
</div>
{/* 导航菜单列表 */}
<nav className="space-y-1">
{[
{ id: 'search' as TabId, label: '统一检索', icon: Search },
{ id: 'library' as TabId, label: '馆藏管理', icon: Library },
{ id: 'reader' as TabId, label: '双语阅读', icon: BookOpen },
{ id: 'citation' as TabId, label: '引用星系', icon: GitFork },
{ id: 'sync' as TabId, label: '批量任务', icon: RefreshCw },
{ id: 'agent' as TabId, label: '智能科研', icon: Sparkles },
].map((tab) => {
const Icon = tab.icon;
const isActive = activeTab === tab.id;
return (
<button
key={tab.id}
title={effectiveCollapsed ? tab.label : undefined}
onClick={() => {
setActiveTab(tab.id);
if (tab.id === 'citation' && selectedPaper) {
loadCitations(selectedPaper.bibcode);
}
if (onClose) onClose(); // 移动端切页时自动收起
}}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border ${
effectiveCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
} ${
isActive
? 'bg-slate-100 border-slate-200 text-slate-850 shadow-xs'
: 'border-transparent text-slate-650 hover:bg-slate-100 hover:text-slate-800'
}`}
>
<Icon className={`w-4 h-4 shrink-0 transition-colors duration-300 ${isActive ? 'text-slate-700' : 'text-slate-500'}`} />
<span
className={`truncate transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
}`}
>
{tab.label}
</span>
</button>
);
})}
</nav>
</div>
{/* 底部当前选定文献提示 (平滑动画版本) 与 退出登录 */}
<div className="space-y-3">
{selectedPaper ? (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
effectiveCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3.5 border-slate-200 bg-slate-100/50'
}`}
title={effectiveCollapsed ? `当前选定文献: ${selectedPaper.title}` : undefined}
>
{/* 折叠下的微型图书图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-55 text-slate-650 shrink-0 shadow-xs ${
effectiveCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<BookOpen className="w-4 h-4" />
</div>
{/* 展开下的完整详情 */}
<div
className={`transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-40'
}`}
>
<span className="text-[9px] font-bold text-slate-500 tracking-widest block mb-1"></span>
<h4 className="text-xs text-slate-800 font-bold line-clamp-2 mb-2 leading-relaxed">{selectedPaper.title}</h4>
<div className="flex items-center justify-between text-[10px] font-medium text-slate-500 gap-2">
<span className="shrink-0">: {selectedPaper.year}</span>
<span className="truncate max-w-[90px] font-mono">{selectedPaper.bibcode}</span>
</div>
</div>
</div>
) : (
<div
className={`relative overflow-hidden transition-all duration-300 border rounded-lg ${
effectiveCollapsed
? 'p-0 border-transparent bg-transparent flex justify-center'
: 'p-3 border-slate-200 bg-slate-100/30'
}`}
title={effectiveCollapsed ? "未选定研究目标" : undefined}
>
{/* 折叠下的微型馆藏图标 */}
<div
className={`transition-all duration-300 flex items-center justify-center rounded-lg border border-slate-200 bg-slate-100/30 text-slate-400 shrink-0 ${
effectiveCollapsed
? 'w-9 h-9 opacity-100 scale-100'
: 'w-0 h-0 opacity-0 scale-75 overflow-hidden'
}`}
>
<Library className="w-4 h-4" />
</div>
{/* 展开下的提示文字 */}
<div
className={`text-center transition-all duration-300 origin-left ${
effectiveCollapsed
? 'opacity-0 max-w-0 max-h-0 pointer-events-none select-none overflow-hidden'
: 'opacity-100 max-w-[200px] max-h-12'
}`}
>
<span className="text-[10px] text-slate-400 font-medium tracking-wide"></span>
</div>
</div>
)}
{/* 退出登录按钮 */}
<button
type="button"
onClick={onLogout}
className={`w-full flex items-center rounded-lg text-xs font-semibold tracking-wider transition-all duration-300 border border-transparent text-slate-500 hover:bg-rose-50/50 hover:text-rose-600 cursor-pointer group ${
effectiveCollapsed ? 'px-2 py-2.5 justify-center' : 'px-3 py-2.5'
}`} }`}
title={effectiveCollapsed ? "退出登录" : undefined}
> >
退 <LogOut className="w-4 h-4 shrink-0 text-slate-400 group-hover:text-rose-505 transition-colors" />
</span> <span
</button> className={`truncate transition-all duration-300 origin-left ${
</div> effectiveCollapsed
</aside> ? 'opacity-0 max-w-0 pointer-events-none select-none overflow-hidden scale-90 -translate-x-2'
: 'opacity-100 max-w-[150px] scale-100 translate-x-0 ml-3'
}`}
>
退
</span>
</button>
</div>
</aside>
</>
); );
} }

View File

@ -1,4 +1,5 @@
// dashboard/src/components/reader/BilingualViewer.tsx // dashboard/src/components/reader/BilingualViewer.tsx
import { useState, useEffect } from 'react';
import { createPortal } from 'react-dom'; import { createPortal } from 'react-dom';
import ReactMarkdown from 'react-markdown'; import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math'; import remarkMath from 'remark-math';
@ -128,6 +129,28 @@ export function BilingualViewer({
hoverCardPos, hoverCardPos,
children, children,
}: BilingualViewerProps) { }: BilingualViewerProps) {
// 监测屏幕宽度以判断是否在移动端/平板设备上(宽度 < 1024px
const [isMobileViewport, setIsMobileViewport] = useState(() => {
if (typeof window !== 'undefined') {
return window.innerWidth < 1024;
}
return false;
});
useEffect(() => {
const handleResize = () => {
setIsMobileViewport(window.innerWidth < 1024);
};
window.addEventListener('resize', handleResize);
return () => window.removeEventListener('resize', handleResize);
}, []);
const gridTemplateColumns = isMobileViewport
? (viewMode === 'bilingual' ? '1fr 1fr' : '1fr')
: (viewMode === 'bilingual'
? (showNotesPanel ? '1fr 1fr 380px' : '1fr 1fr')
: (showNotesPanel ? '1fr 380px' : '1fr'));
// 解析英文和中文段落的 Front Matter 头部元数据 // 解析英文和中文段落的 Front Matter 头部元数据
const { metadata: engMeta, pureMarkdown: engPure } = parseMarkdownFrontMatter(englishText); const { metadata: engMeta, pureMarkdown: engPure } = parseMarkdownFrontMatter(englishText);
const { metadata: chnMeta, pureMarkdown: chnPure } = parseMarkdownFrontMatter(chineseText); const { metadata: chnMeta, pureMarkdown: chnPure } = parseMarkdownFrontMatter(chineseText);
@ -150,11 +173,7 @@ export function BilingualViewer({
return ( return (
<div <div
className="flex-1 grid gap-6 overflow-hidden w-full min-h-0" className="flex-1 grid gap-6 overflow-hidden w-full min-h-0"
style={{ style={{ gridTemplateColumns }}
gridTemplateColumns: viewMode === 'bilingual'
? (showNotesPanel ? '1fr 1fr 380px' : '1fr 1fr')
: (showNotesPanel ? '1fr 380px' : '1fr')
}}
> >
{/* 英文正文视窗 (或双语对照) */} {/* 英文正文视窗 (或双语对照) */}
{(viewMode === 'english' || viewMode === 'bilingual') && ( {(viewMode === 'english' || viewMode === 'bilingual') && (

View File

@ -69,15 +69,24 @@ export function ReaderNotesSidebar({
if (!showNotesPanel) return null; if (!showNotesPanel) return null;
return ( return (
<div className="console-panel rounded-lg border border-slate-200 bg-slate-50 flex flex-col overflow-hidden relative shadow-sm h-full"> <>
<div className="px-4 py-3 border-b border-slate-200 flex items-center justify-between bg-white shrink-0"> {/* 移动端遮罩层 */}
<span className="text-xs font-bold text-slate-800"> <button
{sidebarTab === 'notes' ? '观测记录手札' : '文献 AI 问答'} type="button"
</span> onClick={() => setShowNotesPanel(false)}
<button onClick={() => setShowNotesPanel(false)} className="text-slate-400 hover:text-slate-600 transition-colors cursor-pointer"> className="fixed inset-0 bg-slate-900/30 backdrop-blur-xs z-35 lg:hidden cursor-pointer w-full h-full border-none outline-none"
<X className="w-4 h-4" /> aria-label="关闭侧栏"
</button> />
</div>
<div className="console-panel rounded-lg border border-slate-200 bg-slate-50 flex flex-col overflow-hidden shadow-sm h-full fixed top-0 right-0 h-full w-[88vw] sm:w-[380px] z-40 lg:relative lg:top-auto lg:right-auto lg:h-full lg:w-auto lg:z-10">
<div className="px-4 py-3 border-b border-slate-200 flex items-center justify-between bg-white shrink-0">
<span className="text-xs font-bold text-slate-800">
{sidebarTab === 'notes' ? '观测记录手札' : '文献 AI 问答'}
</span>
<button onClick={() => setShowNotesPanel(false)} className="text-slate-400 hover:text-slate-600 transition-colors cursor-pointer">
<X className="w-4 h-4" />
</button>
</div>
{/* 标签栏选择器 */} {/* 标签栏选择器 */}
<div className="flex border-b border-slate-200 bg-white select-none shrink-0"> <div className="flex border-b border-slate-200 bg-white select-none shrink-0">
@ -281,5 +290,6 @@ export function ReaderNotesSidebar({
</div> </div>
)} )}
</div> </div>
</>
); );
} }

View File

@ -57,18 +57,18 @@ export function ReaderToolbar({
showPdf, showPdf,
}: ReaderToolbarProps) { }: ReaderToolbarProps) {
return ( return (
<div className="flex items-center justify-between border-b border-slate-200 pb-3 shrink-0"> <div className="flex flex-col xl:flex-row xl:items-center xl:justify-between border-b border-slate-200 pb-3 shrink-0 gap-3">
<div className="flex-1 min-w-0 pr-4"> <div className="min-w-0 w-full xl:flex-1 pr-0 xl:pr-4">
<h2 className="text-sm font-bold text-slate-900 line-clamp-1 leading-snug" title={selectedPaper.title}> <h2 className="text-sm font-bold text-slate-900 line-clamp-1 leading-snug" title={selectedPaper.title}>
{selectedPaper.title} {selectedPaper.title}
</h2> </h2>
<div className="flex items-center gap-2 text-xs text-slate-500 mt-1 font-semibold"> <div className="flex flex-wrap items-center gap-x-2 gap-y-1 text-xs text-slate-500 mt-1 font-semibold">
<span>: {selectedPaper.pub_journal || '未标注'}</span> <span>: {selectedPaper.pub_journal || '未标注'}</span>
<span></span> <span className="hidden sm:inline"></span>
<span>: {selectedPaper.bibcode}</span> <span>: {selectedPaper.bibcode}</span>
</div> </div>
</div> </div>
<div className="flex gap-2 items-center relative"> <div className="flex flex-wrap gap-2 items-center relative w-full xl:w-auto">
{/* 快速切换文献菜单 */} {/* 快速切换文献菜单 */}
<div className="relative shrink-0"> <div className="relative shrink-0">
<button <button
@ -175,7 +175,7 @@ export function ReaderToolbar({
<button <button
type="button" type="button"
onClick={() => setViewMode('bilingual')} onClick={() => setViewMode('bilingual')}
className={`hidden md:block px-2.5 py-1 rounded-md text-[10px] sm:text-xs font-bold transition-all cursor-pointer ${ className={`hidden lg:block px-2.5 py-1 rounded-md text-[10px] sm:text-xs font-bold transition-all cursor-pointer ${
viewMode === 'bilingual' viewMode === 'bilingual'
? 'bg-white text-slate-800 shadow-sm' ? 'bg-white text-slate-800 shadow-sm'
: 'text-slate-500 hover:text-slate-700' : 'text-slate-500 hover:text-slate-700'

View File

@ -94,7 +94,12 @@ export function useResearchAgent({ showConfirm, showAlert }: UseResearchAgentPro
const fileInputRef = useRef<HTMLInputElement>(null); const fileInputRef = useRef<HTMLInputElement>(null);
const [loadingSessions, setLoadingSessions] = useState(false); const [loadingSessions, setLoadingSessions] = useState(false);
const [loadingHistory, setLoadingHistory] = useState(false); const [loadingHistory, setLoadingHistory] = useState(false);
const [sidebarCollapsed, setSidebarCollapsed] = useState(false); const [sidebarCollapsed, setSidebarCollapsed] = useState(() => {
if (typeof window !== 'undefined') {
return window.innerWidth < 1024;
}
return false;
});
// 全文搜索历史记录状态 // 全文搜索历史记录状态
const [searchQuery, setSearchQuery] = useState(''); const [searchQuery, setSearchQuery] = useState('');

View File

@ -71,7 +71,7 @@ export function ReaderPanel(props: ReaderPanelProps) {
// 局部管理阅读视角模式(原文/中文/对照) // 局部管理阅读视角模式(原文/中文/对照)
const [viewMode, setViewMode] = useState<'bilingual' | 'english' | 'chinese'>(() => { const [viewMode, setViewMode] = useState<'bilingual' | 'english' | 'chinese'>(() => {
if (typeof window !== 'undefined' && window.innerWidth < 768) { if (typeof window !== 'undefined' && window.innerWidth < 1024) {
return 'english'; return 'english';
} }
if (!chineseText) { if (!chineseText) {
@ -80,16 +80,51 @@ export function ReaderPanel(props: ReaderPanelProps) {
return 'bilingual'; return 'bilingual';
}); });
// 无翻译时强制回退到英文模式,有翻译时自动切换到双语对照
const [userOverrode, setUserOverrode] = useState(false); const [userOverrode, setUserOverrode] = useState(false);
// 监听窗口尺寸变化,处理移动端与桌面端的自适应视角切换
useEffect(() => {
if (typeof window === 'undefined') return;
let prevWidth = window.innerWidth;
const handleResize = () => {
const currentWidth = window.innerWidth;
const wasMobile = prevWidth < 1024;
const isMobile = currentWidth < 1024;
if (wasMobile !== isMobile) {
if (isMobile) {
// 从桌面缩放到移动端:切换显示英文原文
setViewMode('english');
} else {
// 从移动端拉宽到桌面端:若有翻译则显示双语对照,否则显示原文
setViewMode(chineseText ? 'bilingual' : 'english');
}
}
prevWidth = currentWidth;
};
window.addEventListener('resize', handleResize);
return () => window.removeEventListener('resize', handleResize);
}, [chineseText]);
// 无翻译时强制回退到英文模式,有翻译时且在桌面端时自动切换到双语对照
useEffect(() => { useEffect(() => {
if (userOverrode) return; if (userOverrode) return;
const isMobile = typeof window !== 'undefined' && window.innerWidth < 1024;
if (isMobile) {
if (viewMode !== 'english') {
queueMicrotask(() => setViewMode('english'));
}
return;
}
if (!chineseText && viewMode === 'bilingual') { if (!chineseText && viewMode === 'bilingual') {
queueMicrotask(() => setViewMode('english')); queueMicrotask(() => setViewMode('english'));
} else if (chineseText && viewMode === 'english') { } else if (chineseText && viewMode === 'english') {
queueMicrotask(() => setViewMode('bilingual')); queueMicrotask(() => setViewMode('bilingual'));
} }
}, [chineseText]); }, [chineseText, viewMode, userOverrode]);
const handleViewModeChange = (mode: 'bilingual' | 'english' | 'chinese') => { const handleViewModeChange = (mode: 'bilingual' | 'english' | 'chinese') => {
setUserOverrode(true); setUserOverrode(true);

View File

@ -0,0 +1,72 @@
-- migrations/20260629000000_fix_foreign_key_on_update.sql
-- 修复外键级联更新约束 (ON UPDATE CASCADE)
PRAGMA foreign_keys = OFF;
-- 1. 重建 notes 表
CREATE TABLE notes_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
bibcode TEXT NOT NULL,
paragraph_index INTEGER NOT NULL,
note_text TEXT NOT NULL DEFAULT '',
highlight_color TEXT NOT NULL DEFAULT 'yellow',
selected_text TEXT NOT NULL DEFAULT '',
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (bibcode) REFERENCES papers(bibcode) ON DELETE CASCADE ON UPDATE CASCADE
);
INSERT INTO notes_new (id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at)
SELECT id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at FROM notes;
DROP TABLE notes;
ALTER TABLE notes_new RENAME TO notes;
CREATE INDEX IF NOT EXISTS idx_notes_bibcode ON notes(bibcode);
-- 2. 重建 paper_chunks_content 表
CREATE TABLE paper_chunks_content_new (
rowid INTEGER PRIMARY KEY AUTOINCREMENT,
bibcode TEXT,
paragraph_index INTEGER,
content TEXT,
headings TEXT DEFAULT 'Document',
section_index INTEGER DEFAULT 0,
FOREIGN KEY(bibcode) REFERENCES papers(bibcode) ON DELETE CASCADE ON UPDATE CASCADE
);
INSERT INTO paper_chunks_content_new (rowid, bibcode, paragraph_index, content, headings, section_index)
SELECT rowid, bibcode, paragraph_index, content, headings, section_index FROM paper_chunks_content;
DROP TABLE paper_chunks_content;
ALTER TABLE paper_chunks_content_new RENAME TO paper_chunks_content;
-- 3. 重建 paper_targets 表
CREATE TABLE paper_targets_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
bibcode TEXT,
target_name TEXT,
ra TEXT,
dec TEXT,
parallax REAL,
spectral_type TEXT,
v_magnitude REAL,
aliases TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
otype TEXT,
oname TEXT,
pm_ra REAL,
pm_de REAL,
radial_velocity REAL,
parallax_err REAL,
photometry TEXT,
FOREIGN KEY(bibcode) REFERENCES papers(bibcode) ON DELETE CASCADE ON UPDATE CASCADE,
UNIQUE(bibcode, target_name)
);
INSERT INTO paper_targets_new (id, bibcode, target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases, created_at, otype, oname, pm_ra, pm_de, radial_velocity, parallax_err, photometry)
SELECT id, bibcode, target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases, created_at, otype, oname, pm_ra, pm_de, radial_velocity, parallax_err, photometry FROM paper_targets;
DROP TABLE paper_targets;
ALTER TABLE paper_targets_new RENAME TO paper_targets;
CREATE INDEX IF NOT EXISTS idx_paper_targets_name ON paper_targets(target_name);
PRAGMA foreign_keys = ON;

View File

@ -0,0 +1,4 @@
-- migrations/20260630000000_add_paper_chunks_content_bibcode_index.sql
-- 为 paper_chunks_content 表添加 bibcode 索引,提升按 bibcode 查询/删除的性能
CREATE INDEX IF NOT EXISTS idx_paper_chunks_content_bibcode ON paper_chunks_content(bibcode);

View File

@ -3,7 +3,8 @@
// 内置 Hooks — CancellationHook、MetricsHook、AuditLogHook。 // 内置 Hooks — CancellationHook、MetricsHook、AuditLogHook。
use std::collections::HashSet; use std::collections::HashSet;
use std::sync::{Arc, Mutex}; use std::sync::Arc;
use tokio::sync::Mutex;
use async_trait::async_trait; use async_trait::async_trait;
use sqlx::SqlitePool; use sqlx::SqlitePool;
@ -67,25 +68,23 @@ impl AgentHook for CancellationHook {
} }
async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction { async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction {
if let Ok(runs) = self.cancelled_runs.lock() { let runs = self.cancelled_runs.lock().await;
if runs.contains(&ctx.session_id) { if runs.contains(&ctx.session_id) {
warn!( warn!(
"[CancellationHook] 会话 {} 已被用户取消,阻止工具 {} 执行", "[CancellationHook] 会话 {} 已被用户取消,阻止工具 {} 执行",
ctx.session_id, ctx.tool_name ctx.session_id, ctx.tool_name
); );
return PreToolUseAction::Block { return PreToolUseAction::Block {
reason: "用户已手动中止执行".to_string(), reason: "用户已手动中止执行".to_string(),
}; };
}
} }
PreToolUseAction::Continue PreToolUseAction::Continue
} }
async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) { async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) {
// 清理取消状态 // 清理取消状态
if let Ok(mut runs) = self.cancelled_runs.lock() { let mut runs = self.cancelled_runs.lock().await;
runs.remove(&ctx.session_id); runs.remove(&ctx.session_id);
}
} }
} }
@ -112,9 +111,9 @@ impl MetricsHook {
MetricsHook { data } MetricsHook { data }
} }
/// 返回当前指标快照(锁异常时返回 None /// 返回当前指标快照(锁失败时返回 None
pub fn snapshot(&self) -> Option<MetricsData> { pub fn snapshot(&self) -> Option<MetricsData> {
self.data.lock().ok().map(|d| d.clone()) self.data.try_lock().ok().map(|d| d.clone())
} }
/// 获取 Arc 引用,供外部持有 /// 获取 Arc 引用,供外部持有
@ -139,9 +138,8 @@ impl AgentHook for MetricsHook {
} }
async fn on_session_start(&self, ctx: &SessionStartContext) { async fn on_session_start(&self, ctx: &SessionStartContext) {
if let Ok(mut data) = self.data.lock() { let mut data = self.data.lock().await;
data.session_id = Some(ctx.session_id.clone()); data.session_id = Some(ctx.session_id.clone());
}
} }
async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction { async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction {
@ -149,15 +147,14 @@ impl AgentHook for MetricsHook {
"[Metrics] step={} tool={} is_error={} elapsed={}ms", "[Metrics] step={} tool={} is_error={} elapsed={}ms",
ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms
); );
if let Ok(mut data) = self.data.lock() { let mut data = self.data.lock().await;
*data *data
.tool_call_counts .tool_call_counts
.entry(ctx.tool_name.clone()) .entry(ctx.tool_name.clone())
.or_insert(0) += 1; .or_insert(0) += 1;
data.total_steps = ctx.step; data.total_steps = ctx.step;
if ctx.is_error { if ctx.is_error {
data.total_errors += 1; data.total_errors += 1;
}
} }
PostToolUseAction::Continue PostToolUseAction::Continue
} }

View File

@ -52,10 +52,9 @@ pub use builtins::{AuditLogHook, CancellationHook, ContextDeduplicator, MetricsH
mod tests { mod tests {
use super::*; use super::*;
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::{json, Value};
use sqlx::SqlitePool;
use std::collections::HashSet; use std::collections::HashSet;
use std::sync::{Arc, Mutex}; use std::sync::Arc;
use tokio::sync::Mutex;
struct TestHook { struct TestHook {
name: String, name: String,
@ -204,7 +203,7 @@ mod tests {
async fn test_cancellation_hook_blocks_when_cancelled() { async fn test_cancellation_hook_blocks_when_cancelled() {
let mut cancelled = HashSet::new(); let mut cancelled = HashSet::new();
cancelled.insert("test_session".to_string()); cancelled.insert("test_session".to_string());
let cancelled_runs = Arc::new(std::sync::Mutex::new(cancelled)); let cancelled_runs = Arc::new(tokio::sync::Mutex::new(cancelled));
let hook = CancellationHook::new(cancelled_runs); let hook = CancellationHook::new(cancelled_runs);
let ctx = PreToolUseContext { let ctx = PreToolUseContext {
@ -220,7 +219,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_cancellation_hook_allows_when_not_cancelled() { async fn test_cancellation_hook_allows_when_not_cancelled() {
let cancelled_runs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())); let cancelled_runs = Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new()));
let hook = CancellationHook::new(cancelled_runs); let hook = CancellationHook::new(cancelled_runs);
let ctx = PreToolUseContext { let ctx = PreToolUseContext {
@ -281,7 +280,7 @@ mod tests {
"start_tracker" "start_tracker"
} }
async fn on_session_start(&self, ctx: &SessionStartContext) { async fn on_session_start(&self, ctx: &SessionStartContext) {
self.started.lock().unwrap().push(ctx.session_id.clone()); self.started.lock().await.push(ctx.session_id.clone());
} }
} }
@ -300,7 +299,7 @@ mod tests {
registry.run_on_session_start(&ctx).await; registry.run_on_session_start(&ctx).await;
// 验证 hook 确实被调用了 // 验证 hook 确实被调用了
let calls = started.lock().unwrap(); let calls = started.lock().await;
assert_eq!(calls.len(), 1); assert_eq!(calls.len(), 1);
assert_eq!(calls[0], "test_sid"); assert_eq!(calls[0], "test_sid");
} }
@ -325,16 +324,16 @@ mod tests {
"lifecycle_tracker" "lifecycle_tracker"
} }
async fn on_subagent_start(&self, _ctx: &SubagentStartContext) { async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {
*self.subagent_start.lock().unwrap() = true; *self.subagent_start.lock().await = true;
} }
async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) { async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {
*self.subagent_stop.lock().unwrap() = true; *self.subagent_stop.lock().await = true;
} }
async fn on_pre_compact(&self, _ctx: &PreCompactContext) { async fn on_pre_compact(&self, _ctx: &PreCompactContext) {
*self.pre_compact.lock().unwrap() = true; *self.pre_compact.lock().await = true;
} }
async fn on_post_compact(&self, _ctx: &PostCompactContext) { async fn on_post_compact(&self, _ctx: &PostCompactContext) {
*self.post_compact.lock().unwrap() = true; *self.post_compact.lock().await = true;
} }
} }
@ -381,19 +380,16 @@ mod tests {
// 验证所有 4 个 lifecycle hook 都被调用了 // 验证所有 4 个 lifecycle hook 都被调用了
assert!( assert!(
*subagent_start.lock().unwrap(), *subagent_start.lock().await,
"on_subagent_start should be called" "on_subagent_start should be called"
); );
assert!( assert!(
*subagent_stop.lock().unwrap(), *subagent_stop.lock().await,
"on_subagent_stop should be called" "on_subagent_stop should be called"
); );
assert!(*pre_compact.lock().await, "on_pre_compact should be called");
assert!( assert!(
*pre_compact.lock().unwrap(), *post_compact.lock().await,
"on_pre_compact should be called"
);
assert!(
*post_compact.lock().unwrap(),
"on_post_compact should be called" "on_post_compact should be called"
); );
} }
@ -412,7 +408,7 @@ mod tests {
"prompt_tracker" "prompt_tracker"
} }
async fn on_user_prompt_submit(&self, ctx: &UserPromptSubmitContext) { async fn on_user_prompt_submit(&self, ctx: &UserPromptSubmitContext) {
self.prompts.lock().unwrap().push(ctx.prompt.clone()); self.prompts.lock().await.push(ctx.prompt.clone());
} }
} }
@ -431,7 +427,7 @@ mod tests {
registry.run_on_user_prompt_submit(&ctx).await; registry.run_on_user_prompt_submit(&ctx).await;
// 验证 hook 被调用并接收了正确的 prompt // 验证 hook 被调用并接收了正确的 prompt
let calls = prompts.lock().unwrap(); let calls = prompts.lock().await;
assert_eq!(calls.len(), 1); assert_eq!(calls.len(), 1);
assert_eq!(calls[0], "Hello, agent!"); assert_eq!(calls[0], "Hello, agent!");
} }

View File

@ -52,8 +52,8 @@ impl HookRegistry {
/// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。 /// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。
pub fn with_builtins( pub fn with_builtins(
db: SqlitePool, db: SqlitePool,
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>, cancelled_runs: Arc<tokio::sync::Mutex<std::collections::HashSet<String>>>,
metrics_data: Option<Arc<std::sync::Mutex<MetricsData>>>, metrics_data: Option<Arc<tokio::sync::Mutex<MetricsData>>>,
) -> Self { ) -> Self {
let mut registry = Self::new(); let mut registry = Self::new();
registry.add(Box::new(CancellationHook::new(cancelled_runs))); registry.add(Box::new(CancellationHook::new(cancelled_runs)));

View File

@ -388,7 +388,7 @@ mod tests {
#[test] #[test]
fn test_find_duplicate_skips_historical() { fn test_find_duplicate_skips_historical() {
let mut entries = vec![MemoryEntry { let entries = vec![MemoryEntry {
slug: "historical-one".to_string(), slug: "historical-one".to_string(),
name: "历史记忆".to_string(), name: "历史记忆".to_string(),
description: "已过时".to_string(), description: "已过时".to_string(),

View File

@ -416,7 +416,6 @@ impl MemoryManager {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use std::io::Write;
#[test] #[test]
fn test_build_empty_reminder() { fn test_build_empty_reminder() {
@ -484,7 +483,7 @@ mod tests {
// 归档的旧版本也会被加载reload 不排除 _v1.md 文件),所以 entries.len() >= 1 // 归档的旧版本也会被加载reload 不排除 _v1.md 文件),所以 entries.len() >= 1
assert!( assert!(
mgr.entries().len() >= 1, !mgr.entries().is_empty(),
"at least the active entry should exist" "at least the active entry should exist"
); );

View File

@ -245,7 +245,7 @@ mod tests {
#[test] #[test]
fn test_extract_slugs_string_array() { fn test_extract_slugs_string_array() {
let entries = vec![ let entries = [
make_entry("alpha", "Alpha", "first"), make_entry("alpha", "Alpha", "first"),
make_entry("beta", "Beta", "second"), make_entry("beta", "Beta", "second"),
make_entry("gamma", "Gamma", "third"), make_entry("gamma", "Gamma", "third"),
@ -258,7 +258,7 @@ mod tests {
#[test] #[test]
fn test_extract_slugs_numeric_fallback() { fn test_extract_slugs_numeric_fallback() {
let entries = vec![make_entry("x", "X", "x"), make_entry("y", "Y", "y")]; let entries = [make_entry("x", "X", "x"), make_entry("y", "Y", "y")];
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect(); let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
let result = extract_slugs("[0]", &candidates); let result = extract_slugs("[0]", &candidates);
@ -267,7 +267,7 @@ mod tests {
#[test] #[test]
fn test_extract_slugs_invalid_json() { fn test_extract_slugs_invalid_json() {
let entries = vec![make_entry("x", "X", "x")]; let entries = [make_entry("x", "X", "x")];
let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect(); let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect();
assert!(extract_slugs("not json at all", &candidates).is_err()); assert!(extract_slugs("not json at all", &candidates).is_err());

View File

@ -65,6 +65,7 @@ impl CompactionCircuitBreaker {
/// 设置可注入的时间源(仅用于测试)。 /// 设置可注入的时间源(仅用于测试)。
#[cfg(test)] #[cfg(test)]
#[allow(dead_code)]
fn with_time_source(mut self, f: Box<dyn Fn() -> Instant + Send>) -> Self { fn with_time_source(mut self, f: Box<dyn Fn() -> Instant + Send>) -> Self {
self.time_source = Some(f); self.time_source = Some(f);
self self

View File

@ -512,10 +512,10 @@ mod tests {
#[test] #[test]
fn test_backoff_delay() { fn test_backoff_delay() {
let d0 = backoff_delay(0, None); let d0 = backoff_delay(0, None);
assert!(d0 >= 500 && d0 <= 700); assert!((500..=700).contains(&d0));
let d3 = backoff_delay(3, None); let d3 = backoff_delay(3, None);
assert!(d3 >= 4000 && d3 <= 5000); assert!((4000..=5000).contains(&d3));
let d10 = backoff_delay(10, None); let d10 = backoff_delay(10, None);
assert!(d10 <= 40_000); assert!(d10 <= 40_000);

View File

@ -413,37 +413,7 @@ pub async fn execute_parallel(
// 存储待处理的权限请求 // 存储待处理的权限请求
{ {
let mut perms = match app_state.pending_permissions.lock() { let mut perms = app_state.pending_permissions.lock().await;
Ok(p) => p,
Err(e) => {
warn!("[Executor] 权限系统内部错误: {}", e);
let err_output =
format!("权限系统内部错误,工具 {} 被拒绝", prep.tool_name);
let _ = tx.send(AgentStreamEvent::ToolResult {
tool_call_id: prep.tool_call_id.clone(),
name: prep.tool_name.clone(),
output: err_output.clone(),
is_error: true,
metadata: serde_json::json!({}),
step,
});
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);
continue;
}
};
perms.insert( perms.insert(
perm_id.clone(), perm_id.clone(),
PendingPermission { PendingPermission {
@ -461,9 +431,7 @@ pub async fn execute_parallel(
let perm_result = tokio::time::timeout(timeout_dur, resp_rx).await; let perm_result = tokio::time::timeout(timeout_dur, resp_rx).await;
// 清理待处理的权限请求 // 清理待处理的权限请求
if let Ok(mut perms) = app_state.pending_permissions.lock() { app_state.pending_permissions.lock().await.remove(&perm_id);
perms.remove(&perm_id);
}
match perm_result { match perm_result {
Ok(Ok(response)) if response.allowed => { Ok(Ok(response)) if response.allowed => {
@ -570,11 +538,10 @@ pub async fn execute_parallel(
let cancel_handle = tokio::spawn(async move { let cancel_handle = tokio::spawn(async move {
loop { loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await; tokio::time::sleep(std::time::Duration::from_millis(250)).await;
if let Ok(locked) = app_state_ref.cancelled_runs.lock() { let locked = app_state_ref.cancelled_runs.lock().await;
if locked.contains(&sid_ref) { if locked.contains(&sid_ref) {
cancel_flag.store(true, Ordering::SeqCst); cancel_flag.store(true, Ordering::SeqCst);
return; return;
}
} }
} }
}); });

View File

@ -14,7 +14,6 @@
use lru::LruCache; use lru::LruCache;
use std::num::NonZeroUsize; use std::num::NonZeroUsize;
use std::path::Path;
use tracing::{info, warn}; use tracing::{info, warn};
/// 缓存条目最大数量 /// 缓存条目最大数量
@ -82,35 +81,27 @@ impl FileStateCache {
} }
} }
/// 规范化路径 key确保一致性 /// 规范化路径 key确保一致性—— 仅做字符串规范化,不做磁盘 I/O
fn normalize_key(path: &str) -> String { fn normalize_key(path: &str) -> String {
// 去除尾随斜杠,规范化重复斜杠 // 简单规范化:折叠重复的 / 和 \
let p = Path::new(path); let mut result = String::with_capacity(path.len());
// 尝试 canonicalize跟随符号链接失败则用简单的字符串规范化 let mut prev_slash = false;
match p.canonicalize() { for ch in path.chars() {
Ok(canon) => canon.to_string_lossy().to_string(), if ch == '/' || ch == '\\' {
Err(_) => { if !prev_slash {
// 简单规范化:折叠重复的 / result.push('/');
let mut result = String::with_capacity(path.len()); prev_slash = true;
let mut prev_slash = false;
for ch in path.chars() {
if ch == '/' || ch == '\\' {
if !prev_slash {
result.push('/');
prev_slash = true;
}
} else {
result.push(ch);
prev_slash = false;
}
} }
// 去除尾随 / } else {
if result.ends_with('/') && result.len() > 1 { result.push(ch);
result.pop(); prev_slash = false;
}
result
} }
} }
// 去除尾随 /
if result.ends_with('/') && result.len() > 1 {
result.pop();
}
result
} }
/// 获取缓存的条目,返回 Some(&FileState) 若存在 /// 获取缓存的条目,返回 Some(&FileState) 若存在

View File

@ -74,7 +74,7 @@ pub struct AgentRuntime {
/// 后台任务通知队列(支持 bg_task_run/bg_task_check /// 后台任务通知队列(支持 bg_task_run/bg_task_check
bg_notification_queue: Arc<BgNotificationQueue>, bg_notification_queue: Arc<BgNotificationQueue>,
/// 指标采集 hook 的数据引用(供 API 查询) /// 指标采集 hook 的数据引用(供 API 查询)
metrics_data: Arc<std::sync::Mutex<super::hooks::MetricsData>>, metrics_data: Arc<tokio::sync::Mutex<super::hooks::MetricsData>>,
/// 压缩熔断器(跨 turn 共享,防止无限压缩循环) /// 压缩熔断器(跨 turn 共享,防止无限压缩循环)
compaction_breaker: Arc<std::sync::Mutex<circuit_breaker::CompactionCircuitBreaker>>, compaction_breaker: Arc<std::sync::Mutex<circuit_breaker::CompactionCircuitBreaker>>,
/// 权限检查器 /// 权限检查器
@ -111,7 +111,7 @@ impl AgentRuntime {
apply_mode_config(&mut config, mode); apply_mode_config(&mut config, mode);
let queue = Arc::new(BgNotificationQueue::new()); let queue = Arc::new(BgNotificationQueue::new());
let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default())); let metrics_data = Arc::new(tokio::sync::Mutex::new(super::hooks::MetricsData::default()));
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config)); let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new( let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new(
config.denial_max_consecutive, config.denial_max_consecutive,
@ -182,7 +182,7 @@ impl AgentRuntime {
apply_mode_config(&mut config, mode); apply_mode_config(&mut config, mode);
let queue = Arc::new(BgNotificationQueue::new()); let queue = Arc::new(BgNotificationQueue::new());
let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default())); let metrics_data = Arc::new(tokio::sync::Mutex::new(super::hooks::MetricsData::default()));
let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config)); let permission_checker = Arc::new(permission::PermissionChecker::from_config(&config));
let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new( let denial_tracker = Arc::new(std::sync::Mutex::new(denial_tracker::DenialTracker::new(
config.denial_max_consecutive, config.denial_max_consecutive,
@ -240,8 +240,7 @@ impl AgentRuntime {
/// 返回当前运行指标快照(锁异常时返回默认值) /// 返回当前运行指标快照(锁异常时返回默认值)
pub fn get_metrics(&self) -> super::hooks::MetricsData { pub fn get_metrics(&self) -> super::hooks::MetricsData {
self.metrics_data self.metrics_data
.lock() .try_lock()
.ok()
.map(|m| m.clone()) .map(|m| m.clone())
.unwrap_or_default() .unwrap_or_default()
} }
@ -493,11 +492,10 @@ impl AgentRuntime {
// 注册当前会话的权限检查器(如不存在则从全局配置初始化) // 注册当前会话的权限检查器(如不存在则从全局配置初始化)
{ {
if let Ok(mut checkers) = self.app_state.session_permission_checkers.write() { let mut checkers = self.app_state.session_permission_checkers.write().await;
checkers checkers
.entry(sid.clone()) .entry(sid.clone())
.or_insert_with(|| (*self.permission_checker).clone()); .or_insert_with(|| (*self.permission_checker).clone());
}
} }
let tool_defs = self.tool_registry.definitions(); let tool_defs = self.tool_registry.definitions();
@ -525,11 +523,8 @@ impl AgentRuntime {
// 检查用户取消 // 检查用户取消
let is_cancelled = { let is_cancelled = {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(sid) cancelled.remove(sid)
} else {
false
}
}; };
if is_cancelled { if is_cancelled {
@ -874,8 +869,9 @@ impl AgentRuntime {
self.app_state self.app_state
.session_permission_checkers .session_permission_checkers
.read() .read()
.ok() .await
.and_then(|checkers| checkers.get(sid).cloned()) .get(sid)
.cloned()
}; };
let exec_result = executor::execute_parallel( let exec_result = executor::execute_parallel(
&prepared_calls, &prepared_calls,
@ -948,9 +944,8 @@ impl AgentRuntime {
} }
if exec_result.was_cancelled { if exec_result.was_cancelled {
if let Ok(mut locked) = self.app_state.cancelled_runs.lock() { let mut locked = self.app_state.cancelled_runs.lock().await;
locked.remove(sid); locked.remove(sid);
}
warn!( warn!(
"[AgentRuntime] 工具执行期间被用户手动中止,会话 ID: {}", "[AgentRuntime] 工具执行期间被用户手动中止,会话 ID: {}",
sid sid
@ -1018,9 +1013,8 @@ impl AgentRuntime {
match output.status { match output.status {
StreamStatus::Success => return Some(output), StreamStatus::Success => return Some(output),
StreamStatus::Cancelled => { StreamStatus::Cancelled => {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id); cancelled.remove(session_id);
}
warn!( warn!(
"[AgentRuntime] 流式调用期间被用户手动中止,会话 ID: {}", "[AgentRuntime] 流式调用期间被用户手动中止,会话 ID: {}",
session_id session_id
@ -1098,14 +1092,13 @@ impl AgentRuntime {
} }
// 检查用户取消 // 检查用户取消
if let Ok(cancelled) = self.app_state.cancelled_runs.lock() { let cancelled = self.app_state.cancelled_runs.lock().await;
if cancelled.contains(session_id) { if cancelled.contains(session_id) {
warn!("[AgentRuntime] 退避重试期间被用户取消"); warn!("[AgentRuntime] 退避重试期间被用户取消");
let _ = tx.send(AgentStreamEvent::Error { let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(), message: "用户已手动中止执行。".to_string(),
}); });
return None; return None;
}
} }
// 重试 LLM 调用 // 重试 LLM 调用
@ -1127,9 +1120,8 @@ impl AgentRuntime {
return Some(retry_output); return Some(retry_output);
} }
StreamStatus::Cancelled => { StreamStatus::Cancelled => {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id); cancelled.remove(session_id);
}
return None; return None;
} }
StreamStatus::Error(_) => { StreamStatus::Error(_) => {
@ -1237,9 +1229,8 @@ impl AgentRuntime {
return Some(retry_output); return Some(retry_output);
} }
StreamStatus::Cancelled => { StreamStatus::Cancelled => {
if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { let mut cancelled = self.app_state.cancelled_runs.lock().await;
cancelled.remove(session_id); cancelled.remove(session_id);
}
warn!("[AgentRuntime] 恢复期间被用户中止"); warn!("[AgentRuntime] 恢复期间被用户中止");
let _ = tx.send(AgentStreamEvent::Error { let _ = tx.send(AgentStreamEvent::Error {
message: "用户已手动中止执行。".to_string(), message: "用户已手动中止执行。".to_string(),
@ -1341,7 +1332,7 @@ impl AgentRuntime {
if let Some(skills) = self if let Some(skills) = self
.app_state .app_state
.skill_registry .skill_registry
.read() .try_read()
.ok() .ok()
.and_then(|r| r.build_reminder()) .and_then(|r| r.build_reminder())
{ {

View File

@ -102,7 +102,8 @@ impl ToolPartitioner {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, RwLock}; use std::sync::Arc;
use tokio::sync::RwLock;
use super::*; use super::*;
use crate::agent::skills::SkillRegistry; use crate::agent::skills::SkillRegistry;

View File

@ -1017,7 +1017,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_create_new_session() { async fn test_create_new_session() {
let db = setup_db_with_mode().await; let db = setup_db_with_mode().await;
let llm = LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()); let llm =
LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()).unwrap();
let info = create_or_resume_session(&db, None, &llm, "deep-research") let info = create_or_resume_session(&db, None, &llm, "deep-research")
.await .await
@ -1040,7 +1041,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resume_existing_session() { async fn test_resume_existing_session() {
let db = setup_db_with_mode().await; let db = setup_db_with_mode().await;
let llm = LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()); let llm =
LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()).unwrap();
// 先创建一个会话 // 先创建一个会话
let created = create_or_resume_session(&db, None, &llm, "literature-reader") let created = create_or_resume_session(&db, None, &llm, "literature-reader")
@ -1067,7 +1069,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_resume_nonexistent_session_errors() { async fn test_resume_nonexistent_session_errors() {
let db = setup_db_with_mode().await; let db = setup_db_with_mode().await;
let llm = LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()); let llm =
LlmClient::new("key".into(), "http://localhost".into(), "test-model".into()).unwrap();
let result = let result =
create_or_resume_session(&db, Some("nonexistent-id".into()), &llm, "default").await; create_or_resume_session(&db, Some("nonexistent-id".into()), &llm, "default").await;

View File

@ -43,7 +43,7 @@ pub async fn process_llm_stream(
tx: &mpsc::UnboundedSender<AgentStreamEvent>, tx: &mpsc::UnboundedSender<AgentStreamEvent>,
step: usize, step: usize,
session_id: &str, session_id: &str,
cancelled_runs: Arc<std::sync::Mutex<std::collections::HashSet<String>>>, cancelled_runs: Arc<tokio::sync::Mutex<std::collections::HashSet<String>>>,
enable_thinking: bool, enable_thinking: bool,
) -> StreamOutput { ) -> StreamOutput {
// 1. 发起 LLM 流式调用 // 1. 发起 LLM 流式调用
@ -76,10 +76,9 @@ pub async fn process_llm_stream(
let cancel_fut = async { let cancel_fut = async {
loop { loop {
tokio::time::sleep(std::time::Duration::from_millis(250)).await; tokio::time::sleep(std::time::Duration::from_millis(250)).await;
if let Ok(cancelled) = cancelled_runs.lock() { let cancelled = cancelled_runs.lock().await;
if cancelled.contains(&sid) { if cancelled.contains(&sid) {
return; return;
}
} }
} }
}; };

View File

@ -549,8 +549,6 @@ mod tests {
/// 构建测试用的 ToolContext。 /// 构建测试用的 ToolContext。
async fn make_test_tool_context() -> ToolContext { async fn make_test_tool_context() -> ToolContext {
use crate::agent::memory::MemoryManager; use crate::agent::memory::MemoryManager;
use crate::agent::runtime::file_cache::FileStateCache;
use crate::agent::runtime::permission::PermissionChecker;
use crate::agent::skills::SkillRegistry; use crate::agent::skills::SkillRegistry;
use crate::api::AppState; use crate::api::AppState;
use crate::clients::ads::AdsClient; use crate::clients::ads::AdsClient;
@ -563,16 +561,17 @@ mod tests {
use crate::services::translation::Dictionary; use crate::services::translation::Dictionary;
use crate::Config; use crate::Config;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, Mutex, RwLock}; use std::sync::Arc;
let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap(); let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
sqlx::migrate!("./migrations").run(&pool).await.unwrap(); sqlx::migrate!("./migrations").run(&pool).await.unwrap();
let config = Config::from_env(); let config = Config::from_env();
let llm = LlmClient::new("tk".into(), "http://localhost".into(), "m".into()); let llm = LlmClient::new("tk".into(), "http://localhost".into(), "m".into()).unwrap();
let embedding = EmbeddingClient::new("tk".into(), "http://localhost".into(), "e".into()); let embedding =
let ads = AdsClient::new("tk".into()); EmbeddingClient::new("tk".into(), "http://localhost".into(), "e".into()).unwrap();
let arxiv = ArxivClient::new(); let ads = AdsClient::new("tk".into()).unwrap();
let arxiv = ArxivClient::new().unwrap();
let qiniu = QiniuClient::new( let qiniu = QiniuClient::new(
"ak".into(), "ak".into(),
"sk".into(), "sk".into(),
@ -593,20 +592,27 @@ mod tests {
vision_llm: None, vision_llm: None,
embedding, embedding,
downloader: Downloader::new().expect("downloader"), downloader: Downloader::new().expect("downloader"),
http_client: reqwest::Client::new(),
harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())), harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())),
batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())), batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)), active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())), cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from("/tmp/sk")))), skill_registry: Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(PathBuf::from(
pending_questions: Arc::new(Mutex::new(std::collections::HashMap::new())), "/tmp/sk",
pending_permissions: Arc::new(Mutex::new(std::collections::HashMap::new())), )))),
session_permission_checkers: Arc::new(RwLock::new(std::collections::HashMap::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()),
),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
sse_broadcast: None, sse_broadcast: None,
memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from( memory_manager: Arc::new(tokio::sync::Mutex::new(MemoryManager::new(PathBuf::from(
"/tmp/test_mem", "/tmp/test_mem",
)))), )))),
sessions: Arc::new(Mutex::new(std::collections::HashMap::new())), sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(Mutex::new(std::collections::HashMap::new())), login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
}); });
ToolContext::new(app_state) ToolContext::new(app_state)
@ -614,7 +620,7 @@ mod tests {
/// 创建包含指定 Mock 工具的 StreamingToolExecutor。 /// 创建包含指定 Mock 工具的 StreamingToolExecutor。
async fn make_executor(tools: Vec<MockAgentTool>) -> StreamingToolExecutor { async fn make_executor(tools: Vec<MockAgentTool>) -> StreamingToolExecutor {
let skill_registry = Arc::new(std::sync::RwLock::new( let skill_registry = Arc::new(tokio::sync::RwLock::new(
crate::agent::skills::SkillRegistry::new(std::path::PathBuf::from("/tmp/sk")), crate::agent::skills::SkillRegistry::new(std::path::PathBuf::from("/tmp/sk")),
)); ));
let mut registry = ToolRegistry::new(skill_registry); let mut registry = ToolRegistry::new(skill_registry);

View File

@ -38,8 +38,9 @@ pub mod types;
use std::collections::HashMap; use std::collections::HashMap;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock}; use std::sync::Arc;
use std::time::SystemTime; use std::time::SystemTime;
use tokio::sync::RwLock;
use tracing::{info, warn}; use tracing::{info, warn};
// ── Frontmatter ──────────────────────────────────────────────────────────── // ── Frontmatter ────────────────────────────────────────────────────────────
@ -267,7 +268,7 @@ impl SkillRegistry {
use notify::{RecursiveMode, Watcher}; use notify::{RecursiveMode, Watcher};
use std::time::Duration; use std::time::Duration;
let skills_dir = match self_arc.read() { let skills_dir = match self_arc.try_read() {
Ok(r) => r.skills_dir.clone(), Ok(r) => r.skills_dir.clone(),
Err(_) => { Err(_) => {
warn!("[SkillRegistry] RwLock 异常,无法启动文件监视器"); warn!("[SkillRegistry] RwLock 异常,无法启动文件监视器");
@ -310,7 +311,7 @@ impl SkillRegistry {
// 等待 debounce 窗口 // 等待 debounce 窗口
while rx.recv_timeout(Duration::from_millis(300)).is_ok() {} while rx.recv_timeout(Duration::from_millis(300)).is_ok() {}
info!("[SkillRegistry] 检测到 skill 文件变更,自动刷新"); info!("[SkillRegistry] 检测到 skill 文件变更,自动刷新");
if let Ok(mut registry) = self_arc.write() { if let Ok(mut registry) = self_arc.try_write() {
let reg: &mut SkillRegistry = &mut registry; let reg: &mut SkillRegistry = &mut registry;
reg.refresh(); reg.refresh();
} }
@ -576,7 +577,6 @@ pub fn substitute_variables(body: &str, skill_dir: &Path, session_id: Option<&st
/// ```ignore /// ```ignore
/// let creator = SkillCreator::new(skills_dir); /// let creator = SkillCreator::new(skills_dir);
/// let created = creator.create_from_pattern(&pattern)?; /// let created = creator.create_from_pattern(&pattern)?;
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;

View File

@ -146,7 +146,7 @@ impl AgentTool for AskUserTool {
.unwrap_or_default(); .unwrap_or_default();
let question_id = uuid::Uuid::new_v4().to_string(); let question_id = uuid::Uuid::new_v4().to_string();
let short_id = question_id[..8].to_string(); let short_id = question_id[..12].to_string();
let question = UserQuestion { let question = UserQuestion {
question_id: short_id.clone(), question_id: short_id.clone(),
@ -169,12 +169,7 @@ impl AgentTool for AskUserTool {
// 将通道发送端存储到 AppState 的待处理问题列表 // 将通道发送端存储到 AppState 的待处理问题列表
{ {
let mut pending = match ctx.app_state.pending_questions.lock() { let mut pending = ctx.app_state.pending_questions.lock().await;
Ok(p) => p,
Err(_) => {
return ToolOutput::error("待处理问题队列不可用(内部锁异常),请稍后重试");
}
};
pending.insert( pending.insert(
short_id.clone(), short_id.clone(),
crate::api::PendingQuestion { crate::api::PendingQuestion {
@ -201,9 +196,11 @@ impl AgentTool for AskUserTool {
); );
// 清理 // 清理
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() { ctx.app_state
pending.remove(&short_id); .pending_questions
} .lock()
.await
.remove(&short_id);
let free_text = answer.free_text.unwrap_or_default(); let free_text = answer.free_text.unwrap_or_default();
let response = if !answer.answers.is_empty() { let response = if !answer.answers.is_empty() {
@ -228,16 +225,20 @@ impl AgentTool for AskUserTool {
} }
Ok(Err(_)) => { Ok(Err(_)) => {
// 通道关闭(发送端被 drop // 通道关闭(发送端被 drop
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() { ctx.app_state
pending.remove(&short_id); .pending_questions
} .lock()
.await
.remove(&short_id);
ToolOutput::error("用户取消了回答") ToolOutput::error("用户取消了回答")
} }
Err(_) => { Err(_) => {
// 超时 // 超时
if let Ok(mut pending) = ctx.app_state.pending_questions.lock() { ctx.app_state
pending.remove(&short_id); .pending_questions
} .lock()
.await
.remove(&short_id);
ToolOutput::error(format!( ToolOutput::error(format!(
"等待用户回答超时 (5 分钟)。问题: {}", "等待用户回答超时 (5 分钟)。问题: {}",
question_text question_text

View File

@ -10,7 +10,7 @@ use std::path::{Path, PathBuf};
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json; use serde_json::json;
use tokio::sync::mpsc::UnboundedSender; use tokio::sync::mpsc::UnboundedSender;
use tracing::{error, info}; use tracing::error;
use crate::agent::runtime::AgentStreamEvent; use crate::agent::runtime::AgentStreamEvent;
use crate::agent::tools::filesystem::security::{has_path_traversal, is_path_allowed}; use crate::agent::tools::filesystem::security::{has_path_traversal, is_path_allowed};
@ -29,57 +29,17 @@ impl AnalyzeImageTool {
sse_tx: Option<&UnboundedSender<AgentStreamEvent>>, sse_tx: Option<&UnboundedSender<AgentStreamEvent>>,
tool_call_id: Option<String>, tool_call_id: Option<String>,
) -> ToolOutput { ) -> ToolOutput {
let (base64_data, mime_type) = match load_image(library_dir, image_path).await { match crate::services::vision::analyze_image_stream_service(
Ok(data) => data,
Err(e) => {
error!("[analyze_image] 加载图片失败: {} ({})", e, image_path);
return ToolOutput::error(format!("加载图片失败: {}", e));
}
};
info!(
"[analyze_image] 流式分析图片: {} ({} bytes, type={})",
image_path, image_path,
base64_data.len(), question,
mime_type vision_llm,
); library_dir,
sse_tx,
let system_prompt = "你是一位专业的天体物理学家。请基于提供的图片和用户问题,给出专业、详细、清晰的中文分析。如有数学公式,使用 LaTeX 格式。"; tool_call_id,
)
// 流式输出:先发送状态头部(作为第一条 text_delta保证持久化到工具结果中 .await
let header = format!( {
"> 正在使用视觉模型({})分析图片: {}\n\n", Ok((header, answer)) => ToolOutput::success(
vision_llm.model(),
image_path
);
if let Some(tx) = &sse_tx {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: header.clone(),
tool_call_id: tool_call_id.clone(),
});
}
let sse_tx_clone = sse_tx.cloned();
let tc_id = tool_call_id.clone();
let result = vision_llm
.analyze_image_stream(
system_prompt,
question,
&base64_data,
&mime_type,
move |delta: &str| {
if let Some(ref tx) = sse_tx_clone {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: delta.to_string(),
tool_call_id: tc_id.clone(),
});
}
},
)
.await;
match result {
Ok(answer) => ToolOutput::success(
format!("{}{}", header, answer), format!("{}{}", header, answer),
json!({ json!({
"image_path": image_path, "image_path": image_path,
@ -88,7 +48,7 @@ impl AnalyzeImageTool {
}), }),
), ),
Err(e) => { Err(e) => {
error!("[analyze_image] 视觉模型调用失败: {}", e); error!("[analyze_image] 视觉分析大模型服务执行失败: {}", e);
ToolOutput::error(format!("图片分析失败: {}", e)) ToolOutput::error(format!("图片分析失败: {}", e))
} }
} }
@ -192,72 +152,6 @@ impl AgentTool for AnalyzeImageTool {
} }
} }
/// 加载图片:本地路径或 HTTP(S) URL → base64 data + mime type
async fn load_image(library_dir: &Path, path: &str) -> anyhow::Result<(String, String)> {
if path.starts_with("http://") || path.starts_with("https://") {
// ── SSRF 防护 ──
let parsed =
reqwest::Url::parse(path).map_err(|e| anyhow::anyhow!("无效的图片 URL: {}", e))?;
let host = parsed
.host_str()
.ok_or_else(|| anyhow::anyhow!("URL 缺少 host"))?;
// 字面量快检 + 异步 DNS 深度校验(不阻塞 tokio worker
if crate::utils::ssrf::is_host_literal_blocked(host) {
return Err(anyhow::anyhow!(
"出于安全考虑,禁止请求内网/保留地址: {}",
host
));
}
if crate::utils::ssrf::is_host_blocked_async(host).await {
return Err(anyhow::anyhow!("DNS 解析后指向内网/保留地址: {}", host));
}
// 使用安全的重定向策略:允许跟随合法 CDN 重定向S3/Imgur 等),
// 但每次跳转都对目标 URL 做 scheme + host 字面量校验,拦截内网目标。
let client = reqwest::Client::builder()
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()?;
let resp = client.get(path).send().await?;
let mime = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("image/png")
.to_string();
let bytes = resp.bytes().await?;
let b64 = base64_encode(&bytes);
Ok((b64, mime))
} else {
let abs_path = if path.starts_with('/') {
Path::new(path).to_path_buf()
} else {
library_dir.join(path)
};
if !abs_path.exists() {
return Err(anyhow::anyhow!("图片文件不存在: {}", abs_path.display()));
}
let ext = abs_path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("png")
.to_lowercase();
let mime = match ext.as_str() {
"jpg" | "jpeg" => "image/jpeg",
"gif" => "image/gif",
"webp" => "image/webp",
"svg" => "image/svg+xml",
_ => "image/png",
};
let bytes = std::fs::read(&abs_path)?;
let b64 = base64_encode(&bytes);
Ok((b64, mime.to_string()))
}
}
fn base64_encode(bytes: &[u8]) -> String {
use base64::{engine::general_purpose, Engine as _};
general_purpose::STANDARD.encode(bytes)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@ -294,7 +188,7 @@ mod tests {
} else { } else {
config.llm_vision_api_base.clone() config.llm_vision_api_base.clone()
}; };
let vision_llm = LlmClient::new(key, base, config.llm_vision_model.clone()); let vision_llm = LlmClient::new(key, base, config.llm_vision_model.clone()).unwrap();
// 写入测试图片到临时目录 // 写入测试图片到临时目录
let tmp_dir = std::env::temp_dir().join("astro_vision_test"); let tmp_dir = std::env::temp_dir().join("astro_vision_test");
@ -366,7 +260,7 @@ mod tests {
let rt = tokio::runtime::Runtime::new().unwrap(); let rt = tokio::runtime::Runtime::new().unwrap();
let (b64, mime) = rt let (b64, mime) = rt
.block_on(load_image(&tmp_dir, "test.png")) .block_on(crate::services::vision::load_image(&tmp_dir, "test.png"))
.expect("应成功加载本地图片"); .expect("应成功加载本地图片");
assert_eq!(mime, "image/png"); assert_eq!(mime, "image/png");
@ -378,7 +272,10 @@ mod tests {
#[test] #[test]
fn test_load_image_not_found() { fn test_load_image_not_found() {
let rt = tokio::runtime::Runtime::new().unwrap(); let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(load_image(Path::new("/tmp"), "nonexistent_image_xyz.png")); let result = rt.block_on(crate::services::vision::load_image(
Path::new("/tmp"),
"nonexistent_image_xyz.png",
));
assert!(result.is_err()); assert!(result.is_err());
} }
} }

View File

@ -21,8 +21,7 @@ pub use research::citation_network::GetCitationNetworkTool;
pub use research::library_search::SearchLocalLibraryTool; pub use research::library_search::SearchLocalLibraryTool;
pub use research::metadata::GetPaperMetadataTool; pub use research::metadata::GetPaperMetadataTool;
pub use research::note::SaveNoteTool; pub use research::note::SaveNoteTool;
pub use research::paper_content::GetPaperContentTool; pub use research::paper::{GetPaperContentTool, GetPaperOutlineTool};
pub use research::paper_outline::GetPaperOutlineTool;
pub use research::rag::RagSearchTool; pub use research::rag::RagSearchTool;
pub use research::target::QueryTargetTool; pub use research::target::QueryTargetTool;

View File

@ -7,20 +7,10 @@
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json; use serde_json::json;
use sqlx::FromRow;
use tracing::info; use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
use crate::services::citation::{get_citations_paginated, search_citations};
/// 引用条目(关联查询结果)
#[derive(Debug, FromRow)]
struct CitationEntry {
bibcode: String,
title: String,
authors: Option<String>,
year: Option<String>,
citation_count: Option<i32>,
}
pub struct GetCitationNetworkTool; pub struct GetCitationNetworkTool;
@ -117,7 +107,7 @@ impl AgentTool for GetCitationNetworkTool {
let state = &ctx.app_state; let state = &ctx.app_state;
// 验证源文献存在 // 验证源文献存在
let paper = match crate::api::helpers::get_paper_from_db( let paper = match crate::services::paper::get_paper_from_db(
&state.db, &state.db,
&state.config.library_dir, &state.config.library_dir,
&bibcode, &bibcode,
@ -133,313 +123,134 @@ impl AgentTool for GetCitationNetworkTool {
} }
}; };
// 构建基础查询
let (join_col, filter_col) = match direction {
"citations" => ("source_bibcode", "target_bibcode"),
_ => ("target_bibcode", "source_bibcode"),
};
let order_clause = match sort {
"year" => "p.year DESC",
"citation_count" => "p.citation_count DESC",
_ => "p.bibcode ASC",
};
// 有 query → 引用查找模式:在关联文献中按作者+年份搜索 // 有 query → 引用查找模式:在关联文献中按作者+年份搜索
if let Some(ref q) = query { if let Some(ref q) = query {
let (author_part, year_part) = parse_query(q); match search_citations(&state.db, &paper.bibcode, direction, q).await {
let results = search_citations( Ok(results) => {
&state.db, if results.is_empty() {
&paper.bibcode, let dir_label = if direction == "citations" {
join_col, "被引"
filter_col, } else {
&author_part, "参考文献"
&year_part, };
) return ToolOutput::success(
.await; format!(
"在 {} 的{}列表中未找到匹配 '{}' 的文献。请尝试调整搜索词(如仅用作者姓氏)。",
paper.bibcode, dir_label, q
),
json!({ "bibcode": paper.bibcode, "query": q, "direction": direction, "count": 0, "results": [] }),
);
}
if results.is_empty() { let content = results
let dir_label = if direction == "citations" { .iter()
"被引" .enumerate()
} else { .map(|(i, r)| {
"参考文献" format!(
}; "{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
return ToolOutput::success( i + 1,
format!( r.bibcode,
"在 {} 的{}列表中未找到匹配 '{}' 的文献。请尝试调整搜索词(如仅用作者姓氏)。", r.title,
paper.bibcode, dir_label, q r.year.as_deref().unwrap_or(""),
), r.authors.join(", "),
json!({ "bibcode": paper.bibcode, "query": q, "direction": direction, "count": 0, "results": [] }), r.citation_count,
); )
} })
.collect::<Vec<_>>()
.join("\n\n");
let display: Vec<serde_json::Value> = results.iter().map(citation_to_json).collect(); let dir_label = if direction == "citations" {
"被引"
let content = display } else {
.iter() "参考文献"
.enumerate() };
.map(|(i, r)| { ToolOutput::success(
format!( format!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次", "在 {} 的{}列表中找到 {} 篇匹配 '{}' 的文献:\n\n{}",
i + 1, paper.bibcode,
r["bibcode"].as_str().unwrap_or(""), dir_label,
r["title"].as_str().unwrap_or(""), results.len(),
r["year"].as_str().unwrap_or(""), q,
r["authors"].as_str().unwrap_or(""), content
r["citation_count"].as_i64().unwrap_or(0), ),
json!({
"bibcode": paper.bibcode,
"query": q,
"direction": direction,
"count": results.len(),
"results": results,
"mode": "search"
}),
) )
}) }
.collect::<Vec<_>>() Err(e) => ToolOutput::error(format!("引文检索失败: {}", e)),
.join("\n\n"); }
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
ToolOutput::success(
format!(
"在 {} 的{}列表中找到 {} 篇匹配 '{}' 的文献:\n\n{}",
paper.bibcode,
dir_label,
results.len(),
q,
content
),
json!({
"bibcode": paper.bibcode,
"query": q,
"direction": direction,
"count": results.len(),
"results": display,
"mode": "search"
}),
)
} else { } else {
// 无 query → 分页浏览模式 // 无 query → 分页浏览模式
let total: i64 = sqlx::query_scalar(&format!( match get_citations_paginated(&state.db, &paper.bibcode, direction, sort, offset, limit)
"SELECT COUNT(*) FROM citations_references WHERE {} = ?",
filter_col
))
.bind(&paper.bibcode)
.fetch_one(&state.db)
.await
.unwrap_or(0);
let query_sql = format!(
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.{} \
WHERE cr.{} = ? \
ORDER BY {} \
LIMIT ? OFFSET ?",
join_col, filter_col, order_clause
);
let rows: Vec<CitationEntry> = sqlx::query_as(&query_sql)
.bind(&paper.bibcode)
.bind(limit)
.bind(offset)
.fetch_all(&state.db)
.await .await
.unwrap_or_default(); {
Ok((rows, total)) => {
let page_count = rows.len();
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
let sort_label = match sort {
"citation_count" => "按被引次数降序",
"year" => "按年份降序",
_ => "默认顺序",
};
let page_count = rows.len(); let content = if rows.is_empty() {
let display: Vec<serde_json::Value> = rows.iter().map(citation_to_json).collect(); format!("文献 {} 暂无{}记录。", paper.bibcode, dir_label)
} else {
let dir_label = if direction == "citations" { let items = rows
"被引" .iter()
} else { .enumerate()
"参考文献" .map(|(i, r)| {
}; format!(
let sort_label = match sort { "{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
"citation_count" => "按被引次数降序", offset as usize + i + 1,
"year" => "按年份降序", r.bibcode,
_ => "默认顺序", r.title,
}; r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
let content = if rows.is_empty() { r.citation_count,
format!("文献 {} 暂无{}记录。", paper.bibcode, dir_label) )
} else { })
let items = display .collect::<Vec<_>>()
.iter() .join("\n\n");
.enumerate()
.map(|(i, r)| {
format!( format!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次", "文献 {} 的{}列表({},共 {} 篇,第 {}-{} 条):\n\n{}",
offset as usize + i + 1, paper.bibcode,
r["bibcode"].as_str().unwrap_or(""), dir_label,
r["title"].as_str().unwrap_or(""), sort_label,
r["year"].as_str().unwrap_or(""), total,
r["authors"].as_str().unwrap_or(""), offset + 1,
r["citation_count"].as_i64().unwrap_or(0), offset + page_count as i64,
items,
) )
}) };
.collect::<Vec<_>>()
.join("\n\n");
format!(
"文献 {} 的{}列表({},共 {} 篇,第 {}-{} 条):\n\n{}",
paper.bibcode,
dir_label,
sort_label,
total,
offset + 1,
offset + page_count as i64,
items,
)
};
ToolOutput::success( ToolOutput::success(
content, content,
json!({ json!({
"bibcode": paper.bibcode, "bibcode": paper.bibcode,
"direction": direction, "direction": direction,
"sort": sort, "sort": sort,
"total": total, "total": total,
"offset": offset, "offset": offset,
"limit": limit, "limit": limit,
"page_count": page_count, "page_count": page_count,
"results": display, "results": rows,
"mode": "browse" "mode": "browse"
}), }),
) )
}
Err(e) => ToolOutput::error(format!("获取引文列表失败: {}", e)),
}
} }
} }
} }
/// 解析查询字符串,提取作者部分和年份部分。
/// "Lei 2023" → ("Lei", Some("2023"))
/// "Zhang et al. 2020" → ("Zhang et al.", Some("2020"))
/// "Smith" → ("Smith", None)
fn parse_query(q: &str) -> (String, Option<String>) {
// 去掉 "et al." 变体中的句号,保留为空格分隔
let cleaned = q.replace("et al.", "et al");
let parts: Vec<&str> = cleaned.split_whitespace().collect();
// 找最后一个看起来像年份的 token4位数字19xx 或 20xx
let mut year: Option<String> = None;
let mut author_tokens: Vec<&str> = Vec::new();
for (i, &p) in parts.iter().enumerate().rev() {
if year.is_none()
&& p.len() == 4
&& p.chars().all(|c| c.is_ascii_digit())
&& (p.starts_with("19") || p.starts_with("20"))
{
year = Some(p.to_string());
// 年份之前的是作者,之后的不属于查询
author_tokens = parts[..i].to_vec();
break;
}
}
if year.is_none() {
// 没有找到年份,整个字符串作为作者搜索
author_tokens = parts;
}
let author_part = if author_tokens.is_empty() {
q.to_string()
} else {
author_tokens.join(" ")
};
(author_part, year)
}
/// 在引用关系中搜索匹配的文献
async fn search_citations(
db: &sqlx::SqlitePool,
bibcode: &str,
join_col: &str,
filter_col: &str,
author: &str,
year: &Option<String>,
) -> Vec<CitationEntry> {
let author_pattern = format!("%{}%", author);
if let Some(ref y) = year {
let sql = format!(
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.{} \
WHERE cr.{} = ? AND p.authors LIKE ? AND p.year = ?",
join_col, filter_col
);
sqlx::query_as::<_, CitationEntry>(&sql)
.bind(bibcode)
.bind(&author_pattern)
.bind(y)
.fetch_all(db)
.await
.unwrap_or_default()
} else {
let sql = format!(
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.{} \
WHERE cr.{} = ? AND p.authors LIKE ?",
join_col, filter_col
);
sqlx::query_as::<_, CitationEntry>(&sql)
.bind(bibcode)
.bind(&author_pattern)
.fetch_all(db)
.await
.unwrap_or_default()
}
}
fn citation_to_json(e: &CitationEntry) -> serde_json::Value {
let authors: Vec<String> = e
.authors
.as_deref()
.and_then(|a| serde_json::from_str(a).ok())
.unwrap_or_default();
let first_author = authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string());
json!({
"bibcode": e.bibcode,
"title": e.title,
"authors": authors,
"first_author": first_author,
"year": e.year,
"citation_count": e.citation_count.unwrap_or(0),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_query_with_year() {
let (author, year) = parse_query("Lei 2023");
assert_eq!(author, "Lei");
assert_eq!(year, Some("2023".into()));
}
#[test]
fn test_parse_query_et_al() {
let (author, year) = parse_query("Zhang et al. 2020");
assert_eq!(author, "Zhang et al");
assert_eq!(year, Some("2020".into()));
}
#[test]
fn test_parse_query_author_only() {
let (author, year) = parse_query("Smith");
assert_eq!(author, "Smith");
assert_eq!(year, None);
}
#[test]
fn test_parse_query_full_name() {
let (author, year) = parse_query("Gaia Collaboration 2018");
assert_eq!(author, "Gaia Collaboration");
assert_eq!(year, Some("2018".into()));
}
}

View File

@ -54,8 +54,12 @@ impl AgentTool for GetPaperMetadataTool {
info!("[GetPaperMetadata] 获取文献元数据: {}", bibcode); info!("[GetPaperMetadata] 获取文献元数据: {}", bibcode);
let state = &ctx.app_state; let state = &ctx.app_state;
match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode) match crate::services::paper::get_paper_from_db(
.await &state.db,
&state.config.library_dir,
&bibcode,
)
.await
{ {
Ok(paper) => { Ok(paper) => {
let content = format!( let content = format!(

View File

@ -5,8 +5,7 @@ pub mod citation_network;
pub mod library_search; pub mod library_search;
pub mod metadata; pub mod metadata;
pub mod note; pub mod note;
pub mod paper_content; pub mod paper;
pub mod paper_outline;
pub mod rag; pub mod rag;
pub mod target; pub mod target;
@ -14,7 +13,6 @@ pub use citation_network::GetCitationNetworkTool;
pub use library_search::SearchLocalLibraryTool; pub use library_search::SearchLocalLibraryTool;
pub use metadata::GetPaperMetadataTool; pub use metadata::GetPaperMetadataTool;
pub use note::SaveNoteTool; pub use note::SaveNoteTool;
pub use paper_content::GetPaperContentTool; pub use paper::{GetPaperContentTool, GetPaperOutlineTool};
pub use paper_outline::GetPaperOutlineTool;
pub use rag::RagSearchTool; pub use rag::RagSearchTool;
pub use target::QueryTargetTool; pub use target::QueryTargetTool;

View File

@ -57,41 +57,14 @@ impl AgentTool for SaveNoteTool {
info!("[SaveNote] 保存笔记: {}", title); info!("[SaveNote] 保存笔记: {}", title);
let state = &ctx.app_state; let state = &ctx.app_state;
// 文件名安全化 match crate::services::note::save_note_service(&state.config.library_dir, &title, &content)
let safe_title: String = title {
.chars() Ok((filename, filepath, size)) => ToolOutput::success(
.map(|c| {
if c.is_alphanumeric() || c == '-' || c == '_' || c == ' ' {
c
} else {
'_'
}
})
.collect::<String>()
.replace(' ', "_");
let notes_dir = state.config.library_dir.join("notes");
if let Err(e) = std::fs::create_dir_all(&notes_dir) {
return ToolOutput::error(format!("创建笔记目录失败: {}", e));
}
let filename = format!("{}.md", safe_title);
let filepath = notes_dir.join(&filename);
let full_content = format!(
"---\ntitle: {}\ndate: {}\ngenerated_by: AstroResearch Agent\n---\n\n{}",
title,
chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
content
);
match std::fs::write(&filepath, &full_content) {
Ok(_) => ToolOutput::success(
format!("笔记已保存: {}", filename), format!("笔记已保存: {}", filename),
json!({ json!({
"filename": filename, "filename": filename,
"path": filepath.to_string_lossy(), "path": filepath.to_string_lossy(),
"size": full_content.len() "size": size
}), }),
), ),
Err(e) => ToolOutput::error(format!("保存笔记失败: {}", e)), Err(e) => ToolOutput::error(format!("保存笔记失败: {}", e)),

View File

@ -0,0 +1,197 @@
// src/agent/tools/astro/research/paper_reader_tools.rs
//
// 文献内容读取与大纲提取的智能体工具定义。
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
// ==========================================
// 1. GetPaperOutlineTool (获取文献大纲目录)
// ==========================================
#[derive(Default)]
pub struct GetPaperOutlineTool;
#[async_trait]
impl AgentTool for GetPaperOutlineTool {
fn name(&self) -> &'static str {
"get_paper_outline"
}
fn description(&self) -> &'static str {
"获取文献的章节大纲目录,列出所有章节的序号和标题。\n\
使便 token"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符bibcode如 2006BaltA..15...69W"
}
},
"required": ["bibcode"]
})
}
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
Some(b) => b.to_string(),
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
};
info!("[GetPaperOutline] 提取大纲: {}", bibcode);
let state = &ctx.app_state;
let mode = crate::services::paper::ReadMode::Outline { max_level: Some(3) };
match crate::services::paper::read_paper_content(
&state.db,
&state.config.library_dir,
&bibcode,
mode,
)
.await
{
Ok(res) => {
let sections = res.outline.unwrap_or_default();
if sections.is_empty() {
return ToolOutput::success(
"该文献未检测到章节标题(## 或 ### 格式)。请直接读取全文。",
json!({ "bibcode": bibcode, "sections": [] }),
);
}
let display_sections: Vec<serde_json::Value> = sections
.iter()
.map(|s| {
json!({
"index": s.index,
"level": s.level,
"heading": s.heading.clone(),
})
})
.collect();
let header = format!("文献 {} 章节大纲(共 {} 节):\n", bibcode, sections.len());
let footer = "\n\n使用 get_paper_content 并传入 section_index 或 section_name 读取指定章节内容。";
ToolOutput::success(
header + &res.content + footer,
json!({
"bibcode": bibcode,
"section_count": sections.len(),
"sections": display_sections
}),
)
}
Err(e) => ToolOutput::error(e.to_string()),
}
}
}
// ==========================================
// 2. GetPaperContentTool (获取文献具体内容)
// ==========================================
#[derive(Default)]
pub struct GetPaperContentTool;
#[async_trait]
impl AgentTool for GetPaperContentTool {
fn name(&self) -> &'static str {
"get_paper_content"
}
fn description(&self) -> &'static str {
"获取文献的具体内容(全文或指定章节内容)。\n\
section_index section_name \n\
使 get_paper_outline token "
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符bibcode如 2006BaltA..15...69W"
},
"section_index": {
"type": "integer",
"description": "章节序号0-based。从 get_paper_outline 返回的 sections[index] 获取。与 section_name 二选一"
},
"section_name": {
"type": "string",
"description": "章节标题文本(支持模糊匹配,如 'introduction')。与 section_index 二选一"
}
},
"required": ["bibcode"]
})
}
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
Some(b) => b.to_string(),
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
};
let section_index = args
.get("section_index")
.and_then(|s| s.as_u64())
.map(|n| n as usize);
let section_name = args
.get("section_name")
.and_then(|s| s.as_str())
.map(String::from);
info!(
"[GetPaperContent] bibcode={}, section_index={:?}, section_name={:?}",
bibcode, section_index, section_name
);
let state = &ctx.app_state;
let mode = if let Some(idx) = section_index {
crate::services::paper::ReadMode::SectionIndex(idx)
} else if let Some(name) = section_name {
crate::services::paper::ReadMode::SectionName(name)
} else {
crate::services::paper::ReadMode::Full {
include_translation: false,
}
};
let label = match &mode {
crate::services::paper::ReadMode::SectionIndex(idx) => Some(format!("#{}", idx)),
crate::services::paper::ReadMode::SectionName(name) => Some(name.clone()),
_ => None,
};
let is_full = matches!(mode, crate::services::paper::ReadMode::Full { .. });
match crate::services::paper::read_paper_content(
&state.db,
&state.config.library_dir,
&bibcode,
mode,
)
.await
{
Ok(res) => ToolOutput::success(
res.content.clone(),
json!({
"bibcode": bibcode,
"chars": res.content.len(),
"section": label,
"full_text": is_full
}),
),
Err(e) => ToolOutput::error(e.to_string()),
}
}
}

View File

@ -1,166 +0,0 @@
// src/agent/tools/astro/research/paper_content.rs
// GetPaperContentTool — 文献内容读取(全文 + 按章节)
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
pub struct GetPaperContentTool;
#[async_trait]
impl AgentTool for GetPaperContentTool {
fn name(&self) -> &str {
"get_paper_content"
}
fn description(&self) -> &str {
"读取已解析文献的 Markdown 内容。默认返回全文;可通过 section_index序号或 section_name标题名\
使 get_paper_outline token \
"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
},
"section_index": {
"type": "integer",
"description": "章节序号0-based。从 get_paper_outline 返回的 sections[index] 获取。与 section_name 二选一"
},
"section_name": {
"type": "string",
"description": "章节标题名称,大小写不敏感,支持模糊匹配。如 'Introduction'、'Results'、'Discussion'。与 section_index 二选一"
}
},
"required": ["bibcode"]
})
}
fn group(&self) -> &str {
"as:research"
}
fn is_readonly(&self) -> bool {
true
}
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
true
}
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
Some(b) => b.to_string(),
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
};
let section_index = args
.get("section_index")
.and_then(|s| s.as_u64())
.map(|n| n as usize);
let section_name = args
.get("section_name")
.and_then(|s| s.as_str())
.map(String::from);
let want_section = section_index.is_some() || section_name.is_some();
info!(
"[GetPaperContent] bibcode={}, section_index={:?}, section_name={:?}",
bibcode, section_index, section_name
);
let state = &ctx.app_state;
let paths = crate::api::helpers::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
&bibcode,
)
.await;
let md_opt = match paths {
Ok(Some((_, _, md_opt, _))) => md_opt,
Ok(None) => return ToolOutput::error(
"获取文献内容失败:该文献未在本地数据库中注册,请先使用 search_papers 搜索该文献。",
),
Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)),
};
let md_rel = match md_opt {
Some(rel) => rel,
None => {
return ToolOutput::error(
"获取文献内容失败:该文献尚未完成结构化解析。请使用 process_paper 执行 download 和 parse 任务。",
)
}
};
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
return ToolOutput::error(
"获取文献内容失败:文献本地 Markdown 文件已丢失,请重新使用 process_paper 执行 parse 任务。",
);
}
let content = match std::fs::read_to_string(&md_abs) {
Ok(c) => c,
Err(e) => {
return ToolOutput::error(format!("获取文献内容失败,读取本地文件错误: {}", e))
}
};
// 按章节读取
if want_section {
let section_content = if let Some(idx) = section_index {
crate::services::section_parser::extract_section_by_index(&content, idx)
} else if let Some(ref name) = section_name {
crate::services::section_parser::extract_section_by_name(&content, name)
} else {
None
};
match section_content {
Some(sec) => {
let label = section_name
.or_else(|| section_index.map(|i| format!("#{}", i)))
.unwrap_or_default();
ToolOutput::success(
sec.clone(),
json!({
"bibcode": bibcode,
"chars": sec.len(),
"section": label,
"full_text": false
}),
)
}
None => {
let hint = if section_name.is_some() {
"请使用 get_paper_outline 查看可用的章节标题,确认章节名称拼写正确。"
} else {
"请使用 get_paper_outline 查看该文献的章节数量。"
};
ToolOutput::error(format!("获取章节内容失败:未找到匹配的章节。{}", hint))
}
}
} else {
// 返回全文(原行为)
ToolOutput::success(
content.clone(),
json!({
"bibcode": bibcode,
"chars": content.len(),
"section": null,
"full_text": true
}),
)
}
}
}

View File

@ -1,135 +0,0 @@
// src/agent/tools/astro/research/paper_outline.rs
// GetPaperOutlineTool — 提取文献章节大纲(目录)
use async_trait::async_trait;
use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, ToolContext, ToolOutput};
pub struct GetPaperOutlineTool;
#[async_trait]
impl AgentTool for GetPaperOutlineTool {
fn name(&self) -> &str {
"get_paper_outline"
}
fn description(&self) -> &str {
"提取已解析文献的章节大纲(目录)。返回所有章节标题、层级和序号,供后续按章节读取使用。\
get_paper_content section_index/section_name 使"
}
fn parameters(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"bibcode": {
"type": "string",
"description": "文献的唯一标识符,支持 ADS Bibcode、DOI 或 arXiv ID"
}
},
"required": ["bibcode"]
})
}
fn group(&self) -> &str {
"as:research"
}
fn is_readonly(&self) -> bool {
true
}
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
true
}
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
Some(b) => b.to_string(),
None => return ToolOutput::error("缺少必需参数 'bibcode'"),
};
info!("[GetPaperOutline] 提取大纲: {}", bibcode);
let state = &ctx.app_state;
let paths =
match crate::api::helpers::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
&bibcode,
)
.await
{
Ok(Some((_, _, md_opt, _))) => md_opt,
Ok(None) => return ToolOutput::error(
"获取大纲失败:该文献未在本地数据库中注册。请先使用 search_papers 搜索该文献。",
),
Err(e) => return ToolOutput::error(format!("获取大纲失败: {}", e)),
};
let md_rel = match paths {
Some(rel) => rel,
None => {
return ToolOutput::error(
"获取大纲失败:该文献尚未完成结构化解析。请先使用 process_paper 执行 download 和 parse 任务。",
)
}
};
let md_abs = state.config.library_dir.join(&md_rel);
let content = match std::fs::read_to_string(&md_abs) {
Ok(c) => c,
Err(e) => return ToolOutput::error(format!("获取大纲失败,读取文件错误: {}", e)),
};
let sections = crate::services::section_parser::extract_outline(&content, 3);
if sections.is_empty() {
return ToolOutput::success(
"该文献未检测到章节标题(## 或 ### 格式)。请使用 get_paper_content 直接读取全文。",
json!({ "bibcode": bibcode, "sections": [] }),
);
}
let display_sections: Vec<serde_json::Value> = sections
.iter()
.map(|s| {
json!({
"index": s.index,
"level": s.level,
"heading": s.heading,
})
})
.collect();
let content_display = sections
.iter()
.map(|s| {
let prefix = "#".repeat(s.level);
format!(
"[{}] {} {} ({} 字符)",
s.index,
prefix,
s.heading,
s.char_end - s.char_start
)
})
.collect::<Vec<_>>()
.join("\n");
let header = format!("文献 {} 章节大纲(共 {} 节):\n", bibcode, sections.len());
let footer =
"\n\n使用 get_paper_content 并传入 section_index 或 section_name 读取指定章节内容。";
ToolOutput::success(
header + &content_display + footer,
json!({
"bibcode": bibcode,
"section_count": sections.len(),
"sections": display_sections
}),
)
}
}

View File

@ -54,9 +54,13 @@ impl AgentTool for QueryTargetTool {
info!("[QueryTarget] 查询天体: {}", object_name); info!("[QueryTarget] 查询天体: {}", object_name);
let state = &ctx.app_state; let state = &ctx.app_state;
let client = reqwest::Client::new(); match crate::services::target::query_target_cached(
match crate::services::target::query_target_cached(&state.db, &object_name, None, &client) &state.db,
.await &object_name,
None,
&state.http_client,
)
.await
{ {
Ok(target) => { Ok(target) => {
let mut content = String::new(); let mut content = String::new();

View File

@ -3,7 +3,6 @@
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json; use serde_json::json;
use tracing::info;
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
@ -87,175 +86,42 @@ impl AgentTool for ProcessPaperTool {
} }
let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false); let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);
let state = &ctx.app_state; let state = &ctx.app_state;
let mut completed: Vec<String> = Vec::new();
let mut has_pdf = false;
let mut has_html = false;
let mut has_markdown = false;
let mut chunk_count = 0usize;
for task in &tasks { match crate::services::pipeline::process_paper_pipeline(state, &bibcode, &tasks, force)
match task.as_str() { .await
"download" => { {
info!("[ProcessPaper] download: {}", bibcode); Ok(res) => {
match state let msg = format!(
.downloader "文献 {} 处理完成。已执行: [{}]。PDF: {}, HTML: {}, Markdown: {}, 向量块数: {}",
.download_paper_service( bibcode,
&state.db, res.completed.join(""),
&state.config.library_dir, res.has_pdf,
&bibcode, res.has_html,
force, res.has_markdown,
) res.chunk_count
.await );
{ ToolOutput::success(
Ok(paper) => { msg,
has_pdf = paper.has_pdf; json!({
has_html = paper.has_html; "bibcode": res.bibcode,
completed.push("download".into()); "completed": res.completed,
} "has_pdf": res.has_pdf,
Err(e) => { "has_html": res.has_html,
return ToolOutput::error(format!( "has_markdown": res.has_markdown,
"文献 {} 下载失败(已完成: [{}]: {}", "chunk_count": res.chunk_count
bibcode, }),
completed.join(", "), )
e }
)); Err(e) => {
} let msg = e.to_string();
} let hint = if msg.contains("请先下载") {
} "(提示:未检测到已下载的本地资源文件,请确保 tasks 中 download 在 parse 之前执行)"
} else {
"parse" => { ""
info!("[ProcessPaper] parse: {}", bibcode); };
match crate::services::parser::parse_paper_service( ToolOutput::error(format!("文献 {} 处理失败: {}{}", bibcode, msg, hint))
&state.db,
&state.config.library_dir,
&state.qiniu,
&state.config,
&bibcode,
force,
)
.await
{
Ok(markdown) => {
has_markdown = true;
completed.push("parse".into());
info!("[ProcessPaper] parse done: {} chars", markdown.len());
}
Err(e) => {
let msg = e.to_string();
let hint = if msg.contains("请先下载") {
"未检测到已下载的本地资源文件,请确保 tasks 中 download 在 parse 之前执行。"
} else {
""
};
return ToolOutput::error(format!(
"文献 {} 解析失败(已完成: [{}]: {} {}",
bibcode,
completed.join(", "),
msg,
hint
));
}
}
}
"embed" => {
info!("[ProcessPaper] embed: {}", bibcode);
// 读取 markdown 文件
let paths = crate::api::helpers::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
&bibcode,
)
.await;
let md_rel = match paths {
Ok(Some((_, _, Some(md), _))) => md,
Ok(_) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]):未找到已解析的 Markdown 文件,请确保 tasks 中 parse 在 embed 之前执行",
bibcode,
completed.join(", ")
));
}
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]: {}",
bibcode,
completed.join(", "),
e
));
}
};
let md_abs = state.config.library_dir.join(&md_rel);
let markdown_text = match std::fs::read_to_string(&md_abs) {
Ok(t) => t,
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]):读取 Markdown 文件错误: {}",
bibcode,
completed.join(", "),
e
));
}
};
match crate::services::rag::ingest_paper(
&state.db,
&state.embedding,
&bibcode,
&markdown_text,
None, // 使用默认 ChunkerConfig结构化分块
)
.await
{
Ok(n) => {
chunk_count = n;
completed.push("embed".into());
info!("[ProcessPaper] embed done: {} chunks", n);
}
Err(e) => {
return ToolOutput::error(format!(
"文献 {} 向量化失败(已完成: [{}]: {}",
bibcode,
completed.join(", "),
e
));
}
}
}
other => {
return ToolOutput::error(format!(
"无效的任务类型 '{}',仅支持: download, parse, embed",
other
));
}
} }
} }
let msg = format!(
"文献 {} 处理完成。已执行: [{}]。PDF: {}, HTML: {}, Markdown: {}, 向量块数: {}",
bibcode,
completed.join(""),
has_pdf,
has_html,
has_markdown,
chunk_count
);
ToolOutput::success(
msg,
json!({
"bibcode": bibcode,
"completed": completed,
"has_pdf": has_pdf,
"has_html": has_html,
"has_markdown": has_markdown,
"chunk_count": chunk_count
}),
)
} }
} }

View File

@ -182,10 +182,11 @@ impl AgentTool for RunBashTool {
/// 安全命令白名单(只读或数据处理类命令)。 /// 安全命令白名单(只读或数据处理类命令)。
/// 这些命令的首个单词匹配时自动允许执行,不再进行黑名单检查。 /// 这些命令的首个单词匹配时自动允许执行,不再进行黑名单检查。
/// 注意: awk 可通过 system() 执行任意命令tee 可写入任意路径,已移除。
const SAFE_COMMANDS: &[&str] = &[ const SAFE_COMMANDS: &[&str] = &[
"ls", "cat", "head", "tail", "find", "grep", "wc", "echo", "pwd", "sort", "uniq", "cut", "tr", "ls", "cat", "head", "tail", "find", "grep", "wc", "echo", "pwd", "sort", "uniq", "cut", "tr",
"awk", "sed", "jq", "diff", "file", "stat", "du", "df", "env", "printenv", "which", "basename", "sed", "jq", "diff", "file", "stat", "du", "df", "env", "printenv", "which", "basename",
"dirname", "realpath", "readlink", "xargs", "tee", "date", "sleep", "true", "false", "dirname", "realpath", "readlink", "xargs", "date", "sleep", "true", "false",
]; ];
/// 危险命令黑名单(交互式、提权、远程访问、破坏性命令)。 /// 危险命令黑名单(交互式、提权、远程访问、破坏性命令)。
@ -234,7 +235,6 @@ const DANGEROUS_COMMANDS: &[&str] = &[
"nmap", "nmap",
// 登录相关 // 登录相关
"login", "login",
"su",
"sulogin", "sulogin",
]; ];
@ -346,6 +346,37 @@ fn has_command_substitution_bypass(command: &str) -> bool {
false false
} }
/// 语言运行时白名单 — 只允许直接执行脚本文件,禁止 -c/-e 标志执行字符串代码。
const LANGUAGE_RUNTIMES: &[&str] = &[
"python", "python3", "ruby", "node", "perl", "php", "lua", "tcl",
];
/// 检测语言运行时是否通过 -c/-e 标志执行内联代码(可绕过沙箱)。
fn has_language_runtime_inline_code(first_word: &str, command: &str) -> bool {
if !LANGUAGE_RUNTIMES.contains(&first_word) {
return false;
}
let args: Vec<&str> = command.split_whitespace().collect();
// 跳过环境变量赋值和首词本身,检查后续参数
for arg in &args[1..] {
if arg.starts_with('-') {
// 匹配 -c, -e, -exec, --command, --eval 等
let flag = arg.trim_start_matches('-');
if flag == "c"
|| flag == "e"
|| flag.starts_with("c=")
|| flag.starts_with("e=")
|| flag == "exec"
|| flag == "command"
|| flag == "eval"
{
return true;
}
}
}
false
}
/// 检查命令是否包含危险的参数模式(用于非白名单非黑名单命令的二次检查)。 /// 检查命令是否包含危险的参数模式(用于非白名单非黑名单命令的二次检查)。
fn check_dangerous_args(command: &str) -> Option<String> { fn check_dangerous_args(command: &str) -> Option<String> {
let lower = command.to_lowercase(); let lower = command.to_lowercase();
@ -386,6 +417,13 @@ fn validate_bash_command(command: &str) -> Option<String> {
return Some("不允许执行空命令".to_string()); return Some("不允许执行空命令".to_string());
} }
// 彻底禁止 Shell 逻辑操作符,防止命令拼接注入
if has_shell_operators(trimmed) {
return Some(
"命令包含 Shell 操作符(如 ;, |, &, <, > 等)。为了安全起见,如需执行复杂逻辑或管道拼接,请先将命令写入 .sh 脚本文件,然后通过 bash 运行该脚本。".to_string(),
);
}
// 检测命令替换绕过 // 检测命令替换绕过
if has_command_substitution_bypass(trimmed) { if has_command_substitution_bypass(trimmed) {
return Some( return Some(
@ -399,6 +437,14 @@ fn validate_bash_command(command: &str) -> Option<String> {
return Some("无法解析命令名称".to_string()); return Some("无法解析命令名称".to_string());
} }
// 拦截语言运行时的 -c/-e 内联代码执行(可绕过沙箱)
if has_language_runtime_inline_code(&first_word, trimmed) {
return Some(format!(
"不允许通过 {} 的 -c/-e 标志执行内联代码。请将代码写入脚本文件后执行。",
first_word
));
}
// 白名单优先:安全命令放行前,先做危险参数二次检查。 // 白名单优先:安全命令放行前,先做危险参数二次检查。
// 白名单命令find/grep/xargs/sed/tee 等)可携带破坏性参数绕过黑名单, // 白名单命令find/grep/xargs/sed/tee 等)可携带破坏性参数绕过黑名单,
// 故此处对 -delete、xargs rm、sed -i、写入系统目录等模式做拦截。 // 故此处对 -delete、xargs rm、sed -i、写入系统目录等模式做拦截。
@ -426,6 +472,17 @@ fn validate_bash_command(command: &str) -> Option<String> {
None None
} }
/// 检查命令是否包含危险的 Shell 操作符(拼接、重定向、命令替换等)
fn has_shell_operators(command: &str) -> bool {
let operators = [";", "&", "|", "<", ">", "$(", "`"];
for op in operators {
if command.contains(op) {
return true;
}
}
false
}
/// 检查命令是否需要用户权限确认。 /// 检查命令是否需要用户权限确认。
/// 安全白名单中的命令不需确认,其他命令需要。 /// 安全白名单中的命令不需确认,其他命令需要。
pub fn bash_needs_permission(command: &str) -> bool { pub fn bash_needs_permission(command: &str) -> bool {
@ -487,9 +544,11 @@ mod tests {
} }
#[test] #[test]
fn test_whitelist_allows_safe_commands_with_pipes() { fn test_whitelist_blocks_shell_operators() {
assert!(validate_bash_command("cat file.txt | grep foo | wc -l").is_none()); assert!(validate_bash_command("cat file.txt | grep foo | wc -l").is_some());
assert!(validate_bash_command("find . -name '*.rs' | xargs wc").is_none()); assert!(validate_bash_command("find . -name '*.rs' | xargs wc").is_some());
assert!(validate_bash_command("ls; rm -rf /").is_some());
assert!(validate_bash_command("echo hello > file.txt").is_some());
} }
// ── validate_bash_command: 白名单命令的危险参数二次检查 ── // ── validate_bash_command: 白名单命令的危险参数二次检查 ──
@ -501,12 +560,7 @@ mod tests {
"find / -delete", "find / -delete",
"find . -name x -exec rm -rf {} \\;", "find . -name x -exec rm -rf {} \\;",
"find . -execdir rm {} +", "find . -execdir rm {} +",
"cat files.txt | xargs rm -rf",
"sed -i 's/old/new/g' /etc/hosts", "sed -i 's/old/new/g' /etc/hosts",
"echo bad > /etc/passwd",
"tee /etc/passwd",
"echo key >> ~/.ssh/authorized_keys",
"find / -name foo > /var/log/x",
// 多空格/制表符变形绕过测试 // 多空格/制表符变形绕过测试
"xargs rm", "xargs rm",
"sed\t-i\tfile", "sed\t-i\tfile",
@ -607,10 +661,10 @@ mod tests {
} }
#[test] #[test]
fn test_allows_dollar_parens_in_args() { fn test_blocks_dollar_parens_in_args() {
// $() 在参数位置(非命令首词)是合法的,如命令替换结果作为参数 // 由于为了彻底防御拼接和复杂语法执行,现在全面禁用了 $() 和 backtick
assert!(validate_bash_command("echo $(date)").is_none()); assert!(validate_bash_command("echo $(date)").is_some());
assert!(validate_bash_command("grep $(cat pattern.txt) file.txt").is_none()); assert!(validate_bash_command("grep $(cat pattern.txt) file.txt").is_some());
} }
// ── validate_bash_command: 危险参数模式 ── // ── validate_bash_command: 危险参数模式 ──
@ -643,13 +697,27 @@ mod tests {
#[test] #[test]
fn test_allows_common_dev_tools() { fn test_allows_common_dev_tools() {
// 开发常用命令不在白名单但也不在黑名单,应默认放行 // 开发常用命令不在白名单但也不在黑名单,需权限确认
assert!(validate_bash_command("python script.py").is_none());
assert!(validate_bash_command("git status").is_none()); assert!(validate_bash_command("git status").is_none());
assert!(validate_bash_command("cargo build").is_none()); assert!(validate_bash_command("cargo build").is_none());
assert!(validate_bash_command("npm test").is_none()); assert!(validate_bash_command("npm test").is_none());
} }
#[test]
fn test_blocks_language_runtime_inline_code() {
// 语言运行时 -c/-e 标志应被拦截
assert!(validate_bash_command("python -c 'import os; os.system(\"rm -rf /\")'").is_some());
assert!(validate_bash_command("python3 -e 'print(1)'").is_some());
assert!(validate_bash_command("ruby -e 'puts 1'").is_some());
assert!(validate_bash_command("node -c 'console.log(1)'").is_some());
assert!(validate_bash_command("perl -e 'print 1'").is_some());
assert!(validate_bash_command("lua -e 'print(1)'").is_some());
// 但直接执行脚本文件应允许
assert!(validate_bash_command("python script.py").is_none());
assert!(validate_bash_command("python3 /path/to/script.py arg1").is_none());
assert!(validate_bash_command("node index.js").is_none());
}
#[test] #[test]
fn test_network_commands_blocking() { fn test_network_commands_blocking() {
// 网络命令始终被阻止。 // 网络命令始终被阻止。

View File

@ -67,7 +67,11 @@ impl AgentTool for GlobFilesTool {
return ToolOutput::error("路径不在允许的沙箱范围内"); return ToolOutput::error("路径不在允许的沙箱范围内");
} }
let glob_pattern = search_path.join(pattern); let glob_pattern = if std::path::Path::new(pattern).is_absolute() {
return ToolOutput::error("glob 模式不能使用绝对路径");
} else {
search_path.join(pattern)
};
let pattern_str = glob_pattern.to_string_lossy().to_string(); let pattern_str = glob_pattern.to_string_lossy().to_string();
match glob::glob(&pattern_str) { match glob::glob(&pattern_str) {

View File

@ -113,6 +113,17 @@ impl AgentTool for ReadFileTool {
} }
// ── 去重检查结束 ── // ── 去重检查结束 ──
// 预检文件大小,避免 OOM
const MAX_READ_BYTES: u64 = 10 * 1024 * 1024; // 10MB
if let Ok(meta) = std::fs::metadata(&path) {
if meta.len() > MAX_READ_BYTES {
return ToolOutput::error(format!(
"文件过大 ({} MB),超过 10MB 限制。请使用 read_file 的 offset/limit 参数分段读取。",
meta.len() / (1024 * 1024)
));
}
}
match std::fs::read_to_string(&path) { match std::fs::read_to_string(&path) {
Ok(content) => { Ok(content) => {
let total_lines = content.lines().count(); let total_lines = content.lines().count();

View File

@ -43,7 +43,14 @@ pub fn is_path_allowed(path: &Path, ctx: &ToolContext) -> bool {
/// 检查路径字符串是否包含穿越尝试 /// 检查路径字符串是否包含穿越尝试
pub fn has_path_traversal(path_str: &str) -> bool { pub fn has_path_traversal(path_str: &str) -> bool {
path_str.contains("..") || path_str.contains('~') // Check for path components that indicate traversal
let parts: Vec<&str> = path_str.split('/').collect();
for part in &parts {
if *part == ".." || *part == "~" {
return true;
}
}
false
} }
/// 规范化用户提供的路径 /// 规范化用户提供的路径

View File

@ -13,7 +13,8 @@
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json; use serde_json::json;
use std::sync::{Arc, RwLock}; use std::sync::Arc;
use tokio::sync::RwLock;
use crate::agent::runtime::file_cache::FileStateCache; use crate::agent::runtime::file_cache::FileStateCache;
use crate::agent::skills::SkillRegistry; use crate::agent::skills::SkillRegistry;
@ -39,8 +40,7 @@ pub use astro::research::citation_network::GetCitationNetworkTool;
pub use astro::research::library_search::SearchLocalLibraryTool; pub use astro::research::library_search::SearchLocalLibraryTool;
pub use astro::research::metadata::GetPaperMetadataTool; pub use astro::research::metadata::GetPaperMetadataTool;
pub use astro::research::note::SaveNoteTool; pub use astro::research::note::SaveNoteTool;
pub use astro::research::paper_content::GetPaperContentTool; pub use astro::research::paper::{GetPaperContentTool, GetPaperOutlineTool};
pub use astro::research::paper_outline::GetPaperOutlineTool;
pub use astro::research::rag::RagSearchTool; pub use astro::research::rag::RagSearchTool;
pub use astro::research::target::QueryTargetTool; pub use astro::research::target::QueryTargetTool;
pub use astro::system::process::ProcessPaperTool; pub use astro::system::process::ProcessPaperTool;
@ -668,11 +668,15 @@ impl ToolRegistry {
/// 截断文本到指定最大字符数 /// 截断文本到指定最大字符数
pub fn truncate_content(s: &str, max_chars: usize) -> String { pub fn truncate_content(s: &str, max_chars: usize) -> String {
if s.len() <= max_chars { let char_count = s.chars().count();
if char_count <= max_chars {
s.to_string() s.to_string()
} else { } else {
let truncated: String = s.chars().take(max_chars).collect(); let truncated: String = s.chars().take(max_chars).collect();
format!("{}\n\n[... 内容已截断,共 {} 字符 ...]", truncated, s.len()) format!(
"{}\n\n[... 内容已截断,共 {} 字符 ...]",
truncated, char_count
)
} }
} }

View File

@ -41,7 +41,23 @@ pub fn maybe_persist_tool_result(
); );
} }
let file_path = tool_results_dir.join(format!("{}.txt", tool_call_id)); // Sanitize tool_call_id to prevent path traversal
let safe_id: String = tool_call_id
.chars()
.filter(|c| c.is_alphanumeric() || *c == '_' || *c == '-')
.collect();
if safe_id.is_empty() {
let truncated: String = content.chars().take(max_chars).collect();
return (
format!(
"{}...\n[输出已截断,原始长度: {} 字符]",
truncated,
content.len()
),
None,
);
}
let file_path = tool_results_dir.join(format!("{}.txt", safe_id));
// 独占创建——只在文件不存在时写入,保证幂等性 // 独占创建——只在文件不存在时写入,保证幂等性
let written = match std::fs::OpenOptions::new() let written = match std::fs::OpenOptions::new()
@ -161,7 +177,7 @@ mod tests {
let large2 = "b".repeat(5000); let large2 = "b".repeat(5000);
let (_, path1) = maybe_persist_tool_result(&large1, "call_3", 100, &dir); let (_, path1) = maybe_persist_tool_result(&large1, "call_3", 100, &dir);
let (content2, path2) = maybe_persist_tool_result(&large2, "call_3", 100, &dir); let (_content2, path2) = maybe_persist_tool_result(&large2, "call_3", 100, &dir);
assert!(path1.is_some()); assert!(path1.is_some());
assert!(path2.is_some()); assert!(path2.is_some());

View File

@ -53,10 +53,11 @@ impl AgentTool for SearchHistoryTool {
} }
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
let query = match args.get("query").and_then(|v| v.as_str()) { let raw_query = match args.get("query").and_then(|v| v.as_str()) {
Some(q) if !q.is_empty() => q, Some(q) if !q.is_empty() => q,
_ => return ToolOutput::error("缺少 query 参数"), _ => 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 scope = args.get("scope").and_then(|v| v.as_str()).unwrap_or("all");
let db = &ctx.app_state.db; let db = &ctx.app_state.db;
@ -70,7 +71,7 @@ impl AgentTool for SearchHistoryTool {
JOIN agent_sessions s ON s.session_id = fts.session_id \ JOIN agent_sessions s ON s.session_id = fts.session_id \
WHERE agent_sessions_fts MATCH $1 ORDER BY rank LIMIT 10", WHERE agent_sessions_fts MATCH $1 ORDER BY rank LIMIT 10",
) )
.bind(query) .bind(&query)
.fetch_all(db) .fetch_all(db)
.await .await
{ {
@ -100,7 +101,7 @@ impl AgentTool for SearchHistoryTool {
WHERE agent_messages_fts MATCH $1 AND m.active = 1 \ WHERE agent_messages_fts MATCH $1 AND m.active = 1 \
ORDER BY rank LIMIT 10", ORDER BY rank LIMIT 10",
) )
.bind(query) .bind(&query)
.fetch_all(db) .fetch_all(db)
.await .await
{ {

View File

@ -8,8 +8,9 @@
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::json; use serde_json::json;
use std::sync::{Arc, RwLock}; use std::sync::Arc;
use tracing::{info, warn}; use tokio::sync::RwLock;
use tracing::info;
use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
use crate::agent::skills::{substitute_variables, SkillRegistry}; use crate::agent::skills::{substitute_variables, SkillRegistry};
@ -76,14 +77,8 @@ impl AgentTool for LoadSkillTool {
info!("[LoadSkill] 加载 skill: {}", skill_name); info!("[LoadSkill] 加载 skill: {}", skill_name);
// 从缓存注册表读取单语句guard 自动 drop不跨越 await // 从缓存注册表读取
let skill = match self.registry.read() { let skill = self.registry.read().await.get_skill(&skill_name).cloned();
Ok(reg) => reg.get_skill(&skill_name).cloned(),
Err(e) => {
warn!("[LoadSkill] RwLock poisoned: {:?}", e);
return ToolOutput::error("技能注册表不可用(内部锁异常),请稍后重试");
}
};
let skill = match skill { let skill = match skill {
Some(s) => s, Some(s) => s,
None => { None => {
@ -94,29 +89,20 @@ impl AgentTool for LoadSkillTool {
} }
}; };
// 记录使用统计单语句write guard 自动 drop // 记录使用统计
if let Ok(mut reg) = self.registry.write() { {
let mut reg = self.registry.write().await;
reg.record_usage(&skill_name); reg.record_usage(&skill_name);
} else {
warn!("[LoadSkill] 无法记录使用统计RwLock poisoned");
} }
// 变量替换 // 变量替换——有 session_id 时注入当前会话 ID
let session_id = if ctx.app_state.config.database_url.contains("session") { let session_id: Option<&str> = if ctx.session_id.is_empty() {
"current" None
} else { } else {
"" Some(&ctx.session_id)
}; };
let body = substitute_variables( let body = substitute_variables(&skill.body, &skill.skill_dir, session_id);
&skill.body,
&skill.skill_dir,
if session_id.is_empty() {
None
} else {
Some(session_id)
},
);
// 构建 skill 目录的绝对路径(用于 "Base directory" 前缀) // 构建 skill 目录的绝对路径(用于 "Base directory" 前缀)
let skill_dir_path = skill let skill_dir_path = skill

View File

@ -210,7 +210,7 @@ mod tests {
use crate::services::translation::Dictionary; use crate::services::translation::Dictionary;
use crate::Config; use crate::Config;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, Mutex, RwLock}; use std::sync::Arc;
/// 构建一个最小化的 ToolContext 用于测试(不依赖真实外部服务) /// 构建一个最小化的 ToolContext 用于测试(不依赖真实外部服务)
async fn make_test_tool_context() -> ToolContext { async fn make_test_tool_context() -> ToolContext {
@ -227,16 +227,18 @@ mod tests {
"test_key".into(), "test_key".into(),
"http://localhost".into(), "http://localhost".into(),
"test-model".into(), "test-model".into(),
); )
.unwrap();
let medium_llm = llm.clone(); let medium_llm = llm.clone();
let fast_llm = llm.clone(); let fast_llm = llm.clone();
let embedding = EmbeddingClient::new( let embedding = EmbeddingClient::new(
"test_key".into(), "test_key".into(),
"http://localhost".into(), "http://localhost".into(),
"test-embed".into(), "test-embed".into(),
); )
let ads = AdsClient::new("test_token".into()); .unwrap();
let arxiv = ArxivClient::new(); let ads = AdsClient::new("test_token".into()).unwrap();
let arxiv = ArxivClient::new().unwrap();
let qiniu = QiniuClient::new( let qiniu = QiniuClient::new(
"test_ak".into(), "test_ak".into(),
"test_sk".into(), "test_sk".into(),
@ -259,22 +261,27 @@ mod tests {
vision_llm: None, vision_llm: None,
embedding, embedding,
downloader, downloader,
http_client: reqwest::Client::new(),
harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())), harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())),
batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())), batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)), active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())), cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(RwLock::new(crate::agent::skills::SkillRegistry::new( skill_registry: Arc::new(tokio::sync::RwLock::new(
PathBuf::from("/tmp/test_skills"), crate::agent::skills::SkillRegistry::new(PathBuf::from("/tmp/test_skills")),
))), )),
pending_questions: Arc::new(Mutex::new(std::collections::HashMap::new())), pending_questions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
pending_permissions: Arc::new(Mutex::new(std::collections::HashMap::new())), pending_permissions: Arc::new(
session_permission_checkers: Arc::new(RwLock::new(std::collections::HashMap::new())), tokio::sync::Mutex::new(std::collections::HashMap::new()),
),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
sse_broadcast: None, sse_broadcast: None,
memory_manager: Arc::new(tokio::sync::Mutex::new( memory_manager: Arc::new(tokio::sync::Mutex::new(
crate::agent::memory::MemoryManager::new(PathBuf::from("/tmp/test_memory")), crate::agent::memory::MemoryManager::new(PathBuf::from("/tmp/test_memory")),
)), )),
sessions: Arc::new(Mutex::new(std::collections::HashMap::new())), sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(Mutex::new(std::collections::HashMap::new())), login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
}); });
ToolContext { ToolContext {

View File

@ -10,7 +10,6 @@ use axum::{
}; };
use futures_util::stream::Stream; use futures_util::stream::Stream;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::convert::Infallible; use std::convert::Infallible;
use std::sync::Arc; use std::sync::Arc;
use tracing::{error, info}; use tracing::{error, info};
@ -129,7 +128,8 @@ pub async fn chat_agent(
.library_dir .library_dir
.join(".agent_images") .join(".agent_images")
.join("uploads"); .join("uploads");
std::fs::create_dir_all(&upload_dir) tokio::fs::create_dir_all(&upload_dir)
.await
.map_err(|e| AppError::internal(format!("创建上传目录失败: {}", e)))?; .map_err(|e| AppError::internal(format!("创建上传目录失败: {}", e)))?;
let filename = format!("{}.{}", uuid::Uuid::new_v4(), ext); let filename = format!("{}.{}", uuid::Uuid::new_v4(), ext);
let filepath = upload_dir.join(&filename); let filepath = upload_dir.join(&filename);
@ -137,7 +137,8 @@ pub async fn chat_agent(
let bytes = general_purpose::STANDARD let bytes = general_purpose::STANDARD
.decode(&img.data) .decode(&img.data)
.map_err(|e| AppError::bad_request(format!("图片 base64 解码失败: {}", e)))?; .map_err(|e| AppError::bad_request(format!("图片 base64 解码失败: {}", e)))?;
std::fs::write(&filepath, &bytes) tokio::fs::write(&filepath, &bytes)
.await
.map_err(|e| AppError::internal(format!("保存图片失败: {}", e)))?; .map_err(|e| AppError::internal(format!("保存图片失败: {}", e)))?;
let rel = filepath let rel = filepath
.strip_prefix(&state.config.library_dir) .strip_prefix(&state.config.library_dir)
@ -241,143 +242,30 @@ pub struct SessionListParams {
pub offset: Option<i64>, pub offset: Option<i64>,
} }
#[derive(Debug, Serialize)]
pub struct SessionSummary {
pub session_id: String,
pub title: String,
pub model: String,
pub mode: String,
pub turn_count: i32,
pub summary: Option<String>,
pub created_at: String,
pub updated_at: String,
}
pub async fn list_sessions( pub async fn list_sessions(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<SessionListParams>, Query(params): Query<SessionListParams>,
) -> ApiResult<Json<Vec<SessionSummary>>> { ) -> ApiResult<Json<Vec<crate::services::session::SessionSummary>>> {
let limit = params.limit.unwrap_or(50); let limit = params.limit.unwrap_or(50);
let offset = params.offset.unwrap_or(0); let offset = params.offset.unwrap_or(0);
let rows = sqlx::query( let sessions = crate::services::session::list_sessions_service(&state.db, limit, offset)
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \ .await
FROM agent_sessions \ .map_err(|e| AppError::internal(format!("查询会话列表失败: {}", e)))?;
WHERE deleted_at IS NULL \
ORDER BY updated_at DESC \
LIMIT ? OFFSET ?",
)
.bind(limit)
.bind(offset)
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询会话列表失败: {}", e)))?;
let sessions: Vec<SessionSummary> = rows
.iter()
.map(|r| SessionSummary {
session_id: r.get(0),
title: r.get(1),
model: r.get(2),
mode: r.get(3),
turn_count: r.get(4),
summary: r.get(5),
created_at: r.get(6),
updated_at: r.get(7),
})
.collect();
Ok(Json(sessions)) Ok(Json(sessions))
} }
// ── GET /api/chat/sessions/:id ──
// 获取单个会话的全部消息历史
#[derive(Debug, Serialize)]
pub struct SessionDetail {
pub session: SessionSummary,
pub messages: Vec<MessageRecord>,
}
#[derive(Debug, Serialize)]
pub struct MessageRecord {
pub id: i64,
pub agent_name: String,
pub turn_index: i32,
pub step_index: i32,
pub role: String,
pub content: String,
pub thought: Option<String>,
pub tool_calls: Option<serde_json::Value>,
pub tool_call_id: Option<String>,
pub token_count: i32,
pub metadata: Option<serde_json::Value>,
pub created_at: String,
}
pub async fn get_session( pub async fn get_session(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Path(session_id): Path<String>, Path(session_id): Path<String>,
) -> ApiResult<Json<SessionDetail>> { ) -> ApiResult<Json<crate::services::session::SessionDetail>> {
// 查询会话元信息 let detail = crate::services::session::get_session_detail_service(&state.db, &session_id)
let session_row = sqlx::query( .await
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \ .map_err(|e| AppError::internal(format!("查询会话详情失败: {}", e)))?
FROM agent_sessions \ .ok_or_else(|| AppError::not_found(format!("会话 {} 不存在", session_id)))?;
WHERE session_id = ? AND deleted_at IS NULL",
)
.bind(&session_id)
.fetch_optional(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询会话失败: {}", e)))?
.ok_or_else(|| AppError::not_found(format!("会话 {} 不存在", session_id)))?;
let session = SessionSummary { Ok(Json(detail))
session_id: session_row.get(0),
title: session_row.get(1),
model: session_row.get(2),
mode: session_row.get(3),
turn_count: session_row.get(4),
summary: session_row.get(5),
created_at: session_row.get(6),
updated_at: session_row.get(7),
};
// 查询消息列表(包含 lead 和 subagent 消息,前端按 agent_name/metadata 区分渲染)
let msg_rows = sqlx::query(
"SELECT id, agent_name, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, created_at \
FROM agent_messages \
WHERE session_id = ? AND active = 1 \
ORDER BY id ASC"
)
.bind(&session_id)
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询消息列表失败: {}", e)))?;
let messages: Vec<MessageRecord> = msg_rows
.iter()
.map(|r| {
let tool_calls_json: Option<String> = r.get(7);
let metadata_json: Option<String> = r.get(10);
MessageRecord {
id: r.get(0),
agent_name: r.get(1),
turn_index: r.get(2),
step_index: r.get(3),
role: r.get(4),
content: r.get(5),
thought: r.get(6),
tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()),
tool_call_id: r.get(8),
token_count: r.get(9),
metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
created_at: r.get(11),
}
})
.collect();
Ok(Json(SessionDetail { session, messages }))
} }
// ── DELETE /api/chat/sessions/:id ── // ── DELETE /api/chat/sessions/:id ──
@ -387,15 +275,11 @@ pub async fn delete_session(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Path(session_id): Path<String>, Path(session_id): Path<String>,
) -> ApiResult<Json<serde_json::Value>> { ) -> ApiResult<Json<serde_json::Value>> {
let result = sqlx::query( let success = crate::services::session::delete_session_service(&state.db, &session_id)
"UPDATE agent_sessions SET deleted_at = CURRENT_TIMESTAMP WHERE session_id = ? AND deleted_at IS NULL" .await
) .map_err(|e| AppError::internal(format!("删除会话失败: {}", e)))?;
.bind(&session_id)
.execute(&state.db)
.await
.map_err(|e| AppError::internal(format!("删除会话失败: {}", e)))?;
if result.rows_affected() == 0 { if !success {
return Err(AppError::not_found(format!( return Err(AppError::not_found(format!(
"会话 {} 不存在或已删除", "会话 {} 不存在或已删除",
session_id session_id
@ -414,9 +298,8 @@ pub async fn stop_agent(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Path(session_id): Path<String>, Path(session_id): Path<String>,
) -> ApiResult<Json<serde_json::Value>> { ) -> ApiResult<Json<serde_json::Value>> {
if let Ok(mut cancelled) = state.cancelled_runs.lock() { let mut cancelled = state.cancelled_runs.lock().await;
cancelled.insert(session_id.clone()); cancelled.insert(session_id.clone());
}
info!("已接收并记录手动中止请求,会话 ID: {}", session_id); info!("已接收并记录手动中止请求,会话 ID: {}", session_id);
Ok(Json( Ok(Json(
serde_json::json!({ "status": "stopping", "session_id": session_id }), serde_json::json!({ "status": "stopping", "session_id": session_id }),
@ -426,130 +309,26 @@ pub async fn stop_agent(
// ── GET /api/chat/metrics ── // ── GET /api/chat/metrics ──
// 返回聚合的智能体运行指标 // 返回聚合的智能体运行指标
#[derive(Debug, Serialize)]
pub struct AgentMetricsResponse {
pub total_sessions: i64,
pub total_tool_calls: i64,
pub tool_call_breakdown: serde_json::Value,
pub avg_steps_per_session: f64,
pub error_rate: f64,
}
pub async fn get_agent_metrics( pub async fn get_agent_metrics(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
) -> ApiResult<Json<AgentMetricsResponse>> { ) -> ApiResult<Json<crate::services::session::AgentMetrics>> {
// 总会话数 let metrics = crate::services::session::get_agent_metrics_service(&state.db)
let total_sessions: i64 = .await
sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE deleted_at IS NULL") .map_err(|e| AppError::internal(format!("获取系统指标失败: {}", e)))?;
.fetch_one(&state.db)
.await
.unwrap_or(0);
// 工具调用统计(从审计日志聚合) Ok(Json(metrics))
let total_tool_calls: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status IN ('OK', 'FAIL')")
.fetch_one(&state.db)
.await
.unwrap_or(0);
// 各工具调用次数
let breakdown_rows: Vec<(String, i64)> = sqlx::query_as(
"SELECT COALESCE(tool_name, 'unknown'), COUNT(*) as cnt \
FROM agent_audit_log \
WHERE status IN ('OK', 'FAIL') \
GROUP BY tool_name \
ORDER BY cnt DESC",
)
.fetch_all(&state.db)
.await
.unwrap_or_default();
let tool_call_breakdown: serde_json::Value = breakdown_rows
.iter()
.map(|(name, cnt)| serde_json::json!({ name: cnt }))
.fold(serde_json::json!({}), |mut acc, v| {
if let serde_json::Value::Object(map) = &mut acc {
if let serde_json::Value::Object(v_map) = v {
for (k, val) in v_map {
map.insert(k.clone(), val.clone());
}
}
}
acc
});
// 平均步数
let avg_steps: f64 = sqlx::query_scalar(
"SELECT COALESCE(AVG(CAST(turn_count AS REAL)), 0.0) \
FROM agent_sessions WHERE deleted_at IS NULL",
)
.fetch_one(&state.db)
.await
.unwrap_or(0.0);
// 错误率
let total_errors: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status = 'FAIL'")
.fetch_one(&state.db)
.await
.unwrap_or(0);
let error_rate = if total_tool_calls > 0 {
total_errors as f64 / total_tool_calls as f64
} else {
0.0
};
Ok(Json(AgentMetricsResponse {
total_sessions,
total_tool_calls,
tool_call_breakdown,
avg_steps_per_session: avg_steps,
error_rate,
}))
} }
// ── GET /api/chat/sessions/:id/audit ── // ── GET /api/chat/sessions/:id/audit ──
// 返回指定会话的审计日志 // 返回指定会话的审计日志
#[derive(Debug, Serialize)]
pub struct AuditLogEntry {
pub id: i64,
pub step: i32,
pub tool_name: Option<String>,
pub status: String,
pub elapsed_ms: i32,
pub output_preview: Option<String>,
pub created_at: String,
}
pub async fn get_session_audit( pub async fn get_session_audit(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Path(session_id): Path<String>, Path(session_id): Path<String>,
) -> ApiResult<Json<Vec<AuditLogEntry>>> { ) -> ApiResult<Json<Vec<crate::services::session::AuditLogEntry>>> {
let rows = sqlx::query( let entries = crate::services::session::get_session_audit_service(&state.db, &session_id)
"SELECT id, step, tool_name, status, elapsed_ms, output_preview, created_at \ .await
FROM agent_audit_log \ .map_err(|e| AppError::internal(format!("查询审计日志失败: {}", e)))?;
WHERE session_id = ? \
ORDER BY id ASC",
)
.bind(&session_id)
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("查询审计日志失败: {}", e)))?;
let entries: Vec<AuditLogEntry> = rows
.iter()
.map(|r| AuditLogEntry {
id: r.get(0),
step: r.get(1),
tool_name: r.get(2),
status: r.get(3),
elapsed_ms: r.get(4),
output_preview: r.get(5),
created_at: r.get(6),
})
.collect();
Ok(Json(entries)) Ok(Json(entries))
} }
@ -570,12 +349,7 @@ pub async fn answer_question(
) -> ApiResult<Json<serde_json::Value>> { ) -> ApiResult<Json<serde_json::Value>> {
use crate::agent::tools::ask_user::UserAnswer; use crate::agent::tools::ask_user::UserAnswer;
let mut pending = match state.pending_questions.lock() { let mut pending = state.pending_questions.lock().await;
Ok(p) => p,
Err(_) => {
return Err(AppError::internal("服务器内部状态异常,请稍后重试"));
}
};
let question_id = req.question_id.clone(); let question_id = req.question_id.clone();
match pending.remove(&question_id) { match pending.remove(&question_id) {
@ -608,10 +382,7 @@ pub async fn answer_question(
pub async fn get_pending_questions( pub async fn get_pending_questions(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
) -> Json<Vec<serde_json::Value>> { ) -> Json<Vec<serde_json::Value>> {
let pending = match state.pending_questions.lock() { let pending = state.pending_questions.lock().await;
Ok(p) => p,
Err(_) => return Json(Vec::new()),
};
let questions: Vec<serde_json::Value> = pending let questions: Vec<serde_json::Value> = pending
.iter() .iter()
.map(|(id, pq)| { .map(|(id, pq)| {
@ -630,10 +401,7 @@ pub async fn respond_permission(
Path(session_id): Path<String>, Path(session_id): Path<String>,
Json(req): Json<super::PermissionResponse>, Json(req): Json<super::PermissionResponse>,
) -> ApiResult<Json<serde_json::Value>> { ) -> ApiResult<Json<serde_json::Value>> {
let mut perms = match state.pending_permissions.lock() { let mut perms = state.pending_permissions.lock().await;
Ok(p) => p,
Err(_) => return Err(AppError::internal("内部状态异常")),
};
// 按 tool_call_id 查找匹配的权限请求 // 按 tool_call_id 查找匹配的权限请求
let perm_id = perms let perm_id = perms
@ -667,10 +435,7 @@ pub async fn respond_permission(
pub async fn get_pending_permissions( pub async fn get_pending_permissions(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
) -> Json<Vec<serde_json::Value>> { ) -> Json<Vec<serde_json::Value>> {
let perms = match state.pending_permissions.lock() { let perms = state.pending_permissions.lock().await;
Ok(p) => p,
Err(_) => return Json(Vec::new()),
};
let result: Vec<serde_json::Value> = perms let result: Vec<serde_json::Value> = perms
.iter() .iter()
.map(|(id, p)| { .map(|(id, p)| {

View File

@ -25,9 +25,32 @@ pub struct LoginResponse {
// 辅助函数:从 Cookie 字符串中解析出 session_id 的值 // 辅助函数:从 Cookie 字符串中解析出 session_id 的值
fn parse_session_cookie(cookie_str: &str) -> Option<String> { fn parse_session_cookie(cookie_str: &str) -> Option<String> {
for cookie in cookie_str.split(';') { for cookie in cookie_str.split(';') {
let parts: Vec<&str> = cookie.split('=').map(|s| s.trim()).collect(); let mut parts = cookie.splitn(2, '=');
if parts.len() == 2 && parts[0] == "session_id" { if let (Some(name), Some(value)) = (parts.next(), parts.next()) {
return Some(parts[1].to_string()); if name.trim() == "session_id" {
return Some(value.trim().to_string());
}
}
}
None
}
/// 从请求头中提取认证 Token。
/// 优先从 Authorization: Bearer 头读取,其次从 Cookie 中的 session_id 读取。
fn extract_auth_token(headers: &HeaderMap) -> Option<String> {
// 1. Authorization: Bearer <token>
if let Some(auth_str) = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
{
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
return Some(rest.to_string());
}
}
// 2. Cookie: session_id=<token>
if let Some(cookie_header) = headers.get(header::COOKIE) {
if let Ok(cookie_str) = cookie_header.to_str() {
return parse_session_cookie(cookie_str);
} }
} }
None None
@ -53,25 +76,6 @@ fn verify_password(input: &str, expected: &str) -> bool {
input_bytes.ct_eq(expected_bytes).into() input_bytes.ct_eq(expected_bytes).into()
} }
#[cfg(test)]
mod tests {
use super::verify_password;
#[test]
fn test_verify_password_correct() {
assert!(verify_password("secret123", "secret123"));
assert!(verify_password("", ""));
}
#[test]
fn test_verify_password_wrong() {
assert!(!verify_password("secret123", "secret456"));
assert!(!verify_password("secret", "secret123"));
assert!(!verify_password("secret123", "secret"));
assert!(!verify_password("wrong", ""));
}
}
// 辅助函数清除过期的会话记录24小时未活跃即过期 // 辅助函数清除过期的会话记录24小时未活跃即过期
fn prune_expired_sessions( fn prune_expired_sessions(
sessions: &mut std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>, sessions: &mut std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>,
@ -89,53 +93,63 @@ fn prune_expired_sessions(
}); });
} }
/// 最大会话数上限,超过后拒绝新登录以防止内存耗尽
const MAX_SESSIONS: usize = 1000;
// 登录接口:验证密码成功后,写入 HttpOnly Cookie // 登录接口:验证密码成功后,写入 HttpOnly Cookie
pub async fn login( pub async fn login(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
headers: axum::http::header::HeaderMap, _headers: axum::http::header::HeaderMap,
axum::extract::ConnectInfo(addr): axum::extract::ConnectInfo<std::net::SocketAddr>,
Json(body): Json<LoginRequest>, Json(body): Json<LoginRequest>,
) -> Result<impl IntoResponse, AppError> { ) -> Result<impl IntoResponse, AppError> {
// 提取客户端 IP用于速率限制 // 使用真实连接 IP 做限速X-Forwarded-For 可伪造,不应用于安全判断)
let client_ip = headers let client_ip = addr.ip().to_string();
.get("x-forwarded-for")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.split(',').next())
.unwrap_or("unknown")
.trim()
.to_string();
// 检查速率限制5 次失败 / 5 分钟窗口 // 检查速率限制5 次失败 / 5 分钟窗口
const MAX_FAILURES: u32 = 5; const MAX_FAILURES: u32 = 5;
const WINDOW_SECS: u64 = 300; const WINDOW_SECS: u64 = 300;
{ {
if let Ok(mut limiter) = state.login_rate_limiter.lock() { let mut limiter = state.login_rate_limiter.lock().await;
if let Some((count, first_fail)) = limiter.get(&client_ip) { // 批量清理过期条目,防止内存无限增长
let elapsed = first_fail.elapsed().as_secs(); limiter.retain(|_, (_, first_fail)| first_fail.elapsed().as_secs() < WINDOW_SECS);
if elapsed < WINDOW_SECS && *count >= MAX_FAILURES {
let remaining = WINDOW_SECS - elapsed; if let Some((count, first_fail)) = limiter.get(&client_ip) {
return Err(AppError::unauthorized(format!( let elapsed = first_fail.elapsed().as_secs();
"登录尝试过于频繁,请在 {} 秒后重试", if elapsed < WINDOW_SECS && *count >= MAX_FAILURES {
remaining let remaining = WINDOW_SECS - elapsed;
))); return Err(AppError::unauthorized(format!(
} "登录尝试过于频繁,请在 {} 秒后重试",
if elapsed >= WINDOW_SECS { remaining
limiter.remove(&client_ip); )));
}
} }
} }
} }
let password = body.password.unwrap_or_default(); let password = body
.password
.ok_or_else(|| AppError::bad_request("缺少密码字段"))?;
if password.is_empty() {
return Err(AppError::bad_request("密码不能为空"));
}
if verify_password(&password, &state.config.admin_password) { if verify_password(&password, &state.config.admin_password) {
// 成功:清除该 IP 的失败记录 // 成功:清除该 IP 的失败记录
if let Ok(mut limiter) = state.login_rate_limiter.lock() { state.login_rate_limiter.lock().await.remove(&client_ip);
limiter.remove(&client_ip);
}
let token = uuid::Uuid::new_v4().to_string(); let token = uuid::Uuid::new_v4().to_string();
// 异常安全地记录活跃会话及其当前时间 // 异常安全地记录活跃会话及其当前时间
if let Ok(mut sessions) = state.sessions.lock() { {
let mut sessions = state.sessions.lock().await;
// 先清理过期会话
prune_expired_sessions(&mut sessions);
// 检查会话上限
if sessions.len() >= MAX_SESSIONS {
return Err(AppError::unauthorized(format!(
"活跃会话数已达上限 ({}), 请稍后重试",
MAX_SESSIONS
)));
}
sessions.insert(token.clone(), chrono::Utc::now()); sessions.insert(token.clone(), chrono::Utc::now());
} }
@ -159,13 +173,11 @@ pub async fn login(
)) ))
} else { } else {
// 失败:记录该 IP 的失败次数 // 失败:记录该 IP 的失败次数
if let Ok(mut limiter) = state.login_rate_limiter.lock() { let mut limiter = state.login_rate_limiter.lock().await;
let entry = limiter let entry = limiter
.entry(client_ip) .entry(client_ip)
.or_insert((0, std::time::Instant::now())); .or_insert((0, std::time::Instant::now()));
entry.0 += 1; entry.0 += 1;
// 重置窗口起点(从首次失败开始计算)
}
Err(AppError::unauthorized("密码错误")) Err(AppError::unauthorized("密码错误"))
} }
} }
@ -175,33 +187,10 @@ pub async fn logout(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
req: Request<axum::body::Body>, req: Request<axum::body::Body>,
) -> ApiResult<impl IntoResponse> { ) -> ApiResult<impl IntoResponse> {
let mut token_to_remove = None; let token_to_remove = extract_auth_token(req.headers());
// 优先从 Authorization 提取 Token向前兼容旧 Bearer Header
let auth_header = req
.headers()
.get(axum::http::header::AUTHORIZATION)
.and_then(|val| val.to_str().ok());
if let Some(auth_str) = auth_header {
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
token_to_remove = Some(rest.to_string());
}
}
// 其次从 Cookie 中读取
if token_to_remove.is_none() {
if let Some(cookie_header) = req.headers().get(axum::http::header::COOKIE) {
if let Ok(cookie_str) = cookie_header.to_str() {
token_to_remove = parse_session_cookie(cookie_str);
}
}
}
if let Some(ref token) = token_to_remove { if let Some(ref token) = token_to_remove {
if let Ok(mut sessions) = state.sessions.lock() { state.sessions.lock().await.remove(token);
sessions.remove(token);
}
} }
// 通过 Max-Age=0 清除浏览器端的 Cookie // 通过 Max-Age=0 清除浏览器端的 Cookie
@ -214,7 +203,10 @@ pub async fn logout(
Ok((headers, Json(serde_json::json!({ "status": "ok" })))) Ok((headers, Json(serde_json::json!({ "status": "ok" }))))
} }
// 鉴权状态检查接口 /// 鉴权状态检查接口。
///
/// 注:此端点由 `auth_middleware` 保护——只有携带有效会话 Token 的请求才能到达此处,
/// 因此始终返回 `authenticated: true`。如果中间件被移除,此端点将无条件返回已认证。
pub async fn check_auth() -> Result<impl IntoResponse, StatusCode> { pub async fn check_auth() -> Result<impl IntoResponse, StatusCode> {
Ok(Json(serde_json::json!({ "authenticated": true }))) Ok(Json(serde_json::json!({ "authenticated": true })))
} }
@ -225,40 +217,18 @@ pub async fn auth_middleware(
req: axum::extract::Request, req: axum::extract::Request,
next: Next, next: Next,
) -> Response { ) -> Response {
let mut token = None; let token = extract_auth_token(req.headers());
// 1. 优先从 Authorization 头读取 Bearer Token兼容
let auth_header = req
.headers()
.get(axum::http::header::AUTHORIZATION)
.and_then(|val| val.to_str().ok());
if let Some(auth_str) = auth_header {
if let Some(rest) = auth_str.strip_prefix("Bearer ") {
token = Some(rest.to_string());
}
}
// 2. 其次从 Cookie 头读取 session_id
if token.is_none() {
if let Some(cookie_header) = req.headers().get(axum::http::header::COOKIE) {
if let Ok(cookie_str) = cookie_header.to_str() {
token = parse_session_cookie(cookie_str);
}
}
}
let mut is_valid = false; let mut is_valid = false;
if let Some(tok) = token { if let Some(tok) = token {
if let Ok(mut sessions) = state.sessions.lock() { let mut sessions = state.sessions.lock().await;
// 顺便执行过期会话自动清理垃圾收集 // 顺便执行过期会话自动清理垃圾收集
prune_expired_sessions(&mut sessions); prune_expired_sessions(&mut sessions);
if let Some(last_active) = sessions.get_mut(&tok) { if let Some(last_active) = sessions.get_mut(&tok) {
// 会话存在且未过期,更新最后活跃时间 // 会话存在且未过期,更新最后活跃时间
*last_active = chrono::Utc::now(); *last_active = chrono::Utc::now();
is_valid = true; is_valid = true;
}
} }
} }
@ -268,3 +238,22 @@ pub async fn auth_middleware(
StatusCode::UNAUTHORIZED.into_response() StatusCode::UNAUTHORIZED.into_response()
} }
} }
#[cfg(test)]
mod tests {
use super::verify_password;
#[test]
fn test_verify_password_correct() {
assert!(verify_password("secret123", "secret123"));
assert!(verify_password("", ""));
}
#[test]
fn test_verify_password_wrong() {
assert!(!verify_password("secret123", "secret456"));
assert!(!verify_password("secret", "secret123"));
assert!(!verify_password("secret123", "secret"));
assert!(!verify_password("wrong", ""));
}
}

View File

@ -1,496 +0,0 @@
// src/api/helpers.rs
use super::StandardPaper;
use crate::clients::ads::AdsPaperDoc;
use crate::clients::arxiv::ArxivPaper;
use sqlx::{Row, SqlitePool};
use tracing::info;
/// SQLite 单条 SQL 的参数绑定上限SQLITE_MAX_VARIABLE_NUMBER 默认 999留余量取 900
/// 批量 IN 查询时需按此分批,避免超出绑定数量。
pub const SQLITE_PARAM_LIMIT: usize = 900;
pub fn convert_ads_doc_to_standard(doc: &AdsPaperDoc) -> StandardPaper {
let title = doc
.title
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_else(|| doc.bibcode.clone());
let authors = doc.author.clone().unwrap_or_default();
let keywords = doc.keyword.clone().unwrap_or_default();
let doi = doc
.doi
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_default();
let mut arxiv_id = String::new();
if let Some(identifiers) = &doc.identifier {
for id in identifiers {
if id.starts_with("arXiv:") {
arxiv_id = id.replace("arXiv:", "").trim().to_string();
break;
}
}
}
if arxiv_id.is_empty() && doc.bibcode.starts_with("arXiv") {
arxiv_id = doc.bibcode.replace("arXiv", "").trim().to_string();
}
StandardPaper {
bibcode: doc.bibcode.clone(),
title,
authors,
year: doc.year.clone().unwrap_or_default(),
pub_journal: doc.pub_journal.clone().unwrap_or_default(),
keywords,
abstract_text: doc.abstract_text.clone().unwrap_or_default(),
doi,
arxiv_id,
citation_count: doc.citation_count.unwrap_or(0),
reference_count: doc.reference_count.unwrap_or(0),
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: doc.doctype.clone().unwrap_or_else(|| "article".to_string()),
pdf_error: None,
html_error: None,
}
}
pub fn convert_arxiv_to_standard(doc: &ArxivPaper) -> StandardPaper {
StandardPaper {
bibcode: doc.id.clone(),
title: doc.title.clone(),
authors: doc.authors.clone(),
year: doc.year.clone(),
pub_journal: "arXiv Preprint".to_string(),
keywords: Vec::new(),
abstract_text: doc.abstract_text.clone(),
doi: doc.doi.clone().unwrap_or_default(),
arxiv_id: doc.id.clone(),
citation_count: 0,
reference_count: 0,
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: "eprint".to_string(),
pdf_error: None,
html_error: None,
}
}
pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Result<()> {
let mut conn = db.acquire().await?;
save_paper_to_db_tx(&mut conn, p).await
}
pub async fn save_paper_to_db_tx(
conn: &mut sqlx::SqliteConnection,
p: &StandardPaper,
) -> anyhow::Result<()> {
let authors_json = serde_json::to_string(&p.authors)?;
let keywords_json = serde_json::to_string(&p.keywords)?;
// 1. 如果存在 arxiv_id检查是否有已存在的相同 arxiv_id 记录以防 duplicate
if !p.arxiv_id.is_empty() {
#[allow(clippy::type_complexity)]
let existing_opt: Option<(String, Option<String>, Option<String>, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path FROM papers WHERE arxiv_id = ?"
)
.bind(&p.arxiv_id)
.fetch_optional(&mut *conn)
.await?;
if let Some((existing_bibcode, _pdf, _html, _md, _tr)) = existing_opt {
if existing_bibcode != p.bibcode {
// 发现不同 bibcode 标识的同一篇文献记录,需要进行合并
// 如果已存在的记录使用的是临时 arXiv ID 作为 bibcode且新记录使用的是正式 ADS bibcode我们升级 bibcode 主键
let is_existing_temp = existing_bibcode == p.arxiv_id;
let is_new_formal = p.bibcode != p.arxiv_id;
if is_existing_temp && is_new_formal {
info!(
"发现相同 arXiv ID 的文献,将临时主键 {} 升级为正式 ADS Bibcode: {}",
existing_bibcode, p.bibcode
);
sqlx::query(
"UPDATE papers SET bibcode = ?, title = ?, authors = ?, year = ?, pub = ?, keywords = ?, abstract = ?, doi = ?, citation_count = ?, reference_count = ?, doctype = ? WHERE bibcode = ?"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(&authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(&keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.bind(&existing_bibcode)
.execute(&mut *conn)
.await?;
return Ok(());
} else {
// 如果已存在的是正式 ADS bibcode而新插入的是临时 arXiv ID直接忽略或更新元数据而不更改主键
info!(
"发现相同 arXiv ID 的文献 {} 已存在正式记录,忽略临时 arXiv 插入",
existing_bibcode
);
return Ok(());
}
}
}
}
// 2. 正常插入/冲突更新
sqlx::query(
"INSERT INTO papers (bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, citation_count, reference_count, doctype) \
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \
ON CONFLICT(bibcode) DO UPDATE SET \
title=excluded.title, \
authors=excluded.authors, \
pub=excluded.pub, \
keywords=excluded.keywords, \
abstract=excluded.abstract, \
doi=excluded.doi, \
arxiv_id=excluded.arxiv_id, \
citation_count=excluded.citation_count, \
reference_count=excluded.reference_count, \
doctype=excluded.doctype"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(&p.arxiv_id)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.execute(&mut *conn)
.await?;
Ok(())
}
/// 从数据库行解析出 StandardPaper消除 get_paper_from_db 和 get_library 之间的重复代码)
///
/// 期望行的列顺序bibcode(0), title(1), authors(2), year(3), pub(4), keywords(5),
/// abstract(6), doi(7), arxiv_id(8), citation_count(9), reference_count(10),
/// pdf_path(11), html_path(12), markdown_path(13), translation_path(14),
/// doctype(15), has_vector(16)
pub fn parse_paper_row(
r: &sqlx::sqlite::SqliteRow,
library_dir: &std::path::Path,
) -> StandardPaper {
let pdf_path: Option<String> = r.get(11);
let html_path: Option<String> = r.get(12);
let markdown_path: Option<String> = r.get(13);
let translation_path: Option<String> = r.get(14);
let doctype_val: Option<String> = r.get(15);
let has_vector: bool = r.get(16);
let authors_str: Option<String> = r.get(2);
let authors: Vec<String> = authors_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let keywords_str: Option<String> = r.get(5);
let keywords: Vec<String> = keywords_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let is_pdf_exist = pdf_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_html_exist = html_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_md_exist = markdown_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_tr_exist = translation_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let pdf_error = pdf_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
let html_error = html_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
StandardPaper {
bibcode: r.get(0),
title: r.get(1),
authors,
year: r.get(3),
pub_journal: r.get(4),
keywords,
abstract_text: r.get(6),
doi: r.get(7),
arxiv_id: r.get(8),
citation_count: r.get(9),
reference_count: r.get(10),
is_downloaded: is_pdf_exist || is_html_exist,
has_pdf: is_pdf_exist,
has_html: is_html_exist,
has_markdown: is_md_exist,
has_translation: is_tr_exist,
has_vector,
doctype: doctype_val.unwrap_or_else(|| "article".to_string()),
pdf_error,
html_error,
}
}
pub async fn get_paper_from_db(
db: &SqlitePool,
library_dir: &std::path::Path,
identifier: &str,
) -> anyhow::Result<StandardPaper> {
let clean_id = identifier.trim();
let clean_doi = clean_id
.trim_start_matches("doi:")
.trim_start_matches("DOI:")
.trim_start_matches("https://doi.org/")
.trim_start_matches("http://doi.org/")
.trim();
let clean_arxiv = clean_id
.trim_start_matches("arxiv:")
.trim_start_matches("arXiv:")
.trim_start_matches("pdf/")
.trim();
let r = 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) FROM papers WHERE bibcode = ? OR doi = ? OR arxiv_id = ? OR LOWER(doi) = LOWER(?) OR arxiv_id = ?")
.bind(clean_id)
.bind(clean_doi)
.bind(clean_arxiv)
.bind(clean_doi)
.bind(clean_arxiv)
.fetch_one(db)
.await?;
Ok(parse_paper_row(&r, library_dir))
}
pub async fn check_paper_paths_in_db(
db: &SqlitePool,
library_dir: &std::path::Path,
identifier: &str,
) -> anyhow::Result<
Option<(
Option<String>,
Option<String>,
Option<String>,
Option<String>,
)>,
> {
let clean_id = identifier.trim();
let clean_doi = clean_id
.trim_start_matches("doi:")
.trim_start_matches("DOI:")
.trim_start_matches("https://doi.org/")
.trim_start_matches("http://doi.org/")
.trim();
let clean_arxiv = clean_id
.trim_start_matches("arxiv:")
.trim_start_matches("arXiv:")
.trim_start_matches("pdf/")
.trim();
let r_opt = sqlx::query("SELECT pdf_path, html_path, markdown_path, translation_path FROM papers WHERE bibcode = ? OR doi = ? OR arxiv_id = ? OR LOWER(doi) = LOWER(?) OR arxiv_id = ?")
.bind(clean_id)
.bind(clean_doi)
.bind(clean_arxiv)
.bind(clean_doi)
.bind(clean_arxiv)
.fetch_optional(db)
.await?;
if let Some(r) = r_opt {
let pdf: Option<String> = r.get(0);
let html: Option<String> = r.get(1);
let md: Option<String> = r.get(2);
let tr: Option<String> = r.get(3);
let pdf_res = pdf.filter(|p| library_dir.join(p).exists());
let html_res = html.filter(|p| library_dir.join(p).exists());
let md_res = md.filter(|p| library_dir.join(p).exists());
let tr_res = tr.filter(|p| library_dir.join(p).exists());
Ok(Some((pdf_res, html_res, md_res, tr_res)))
} else {
Ok(None)
}
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
#[test]
fn test_convert_ads_doc_to_standard() {
let doc = AdsPaperDoc {
bibcode: "2026A&A...123..456X".to_string(),
title: Some(vec!["A Test Title".to_string()]),
author: Some(vec!["Author A".to_string(), "Author B".to_string()]),
year: Some("2026".to_string()),
pub_journal: Some("Astronomy & Astrophysics".to_string()),
keyword: Some(vec!["Keyword 1".to_string()]),
abstract_text: Some("This is abstract".to_string()),
doi: Some(vec!["10.1000/test.doi".to_string()]),
citation_count: Some(5),
reference_count: Some(10),
reference: None,
citation: None,
identifier: None,
doctype: Some("article".to_string()),
};
let paper = convert_ads_doc_to_standard(&doc);
assert_eq!(paper.bibcode, "2026A&A...123..456X");
assert_eq!(paper.title, "A Test Title");
assert_eq!(paper.authors, vec!["Author A", "Author B"]);
assert_eq!(paper.year, "2026");
assert_eq!(paper.pub_journal, "Astronomy & Astrophysics");
assert_eq!(paper.keywords, vec!["Keyword 1"]);
assert_eq!(paper.abstract_text, "This is abstract");
assert_eq!(paper.doi, "10.1000/test.doi");
assert_eq!(paper.arxiv_id, "");
assert_eq!(paper.citation_count, 5);
assert_eq!(paper.reference_count, 10);
}
#[test]
fn test_convert_ads_doc_to_standard_with_arxiv_identifier() {
let doc = AdsPaperDoc {
bibcode: "2026MNRAS.530.1234A".to_string(),
title: Some(vec!["Another Test Title".to_string()]),
author: Some(vec!["Author A".to_string()]),
year: Some("2026".to_string()),
pub_journal: Some("MNRAS".to_string()),
keyword: None,
abstract_text: None,
doi: None,
citation_count: None,
reference_count: None,
reference: None,
citation: None,
identifier: Some(vec![
"2026MNRAS.530.1234A".to_string(),
"arXiv:2606.12345".to_string(),
]),
doctype: Some("article".to_string()),
};
let paper = convert_ads_doc_to_standard(&doc);
assert_eq!(paper.bibcode, "2026MNRAS.530.1234A");
assert_eq!(paper.arxiv_id, "2606.12345");
}
#[test]
fn test_convert_arxiv_to_standard() {
let doc = ArxivPaper {
id: "2606.12345".to_string(),
title: "Arxiv Title".to_string(),
authors: vec!["Author C".to_string()],
year: "2026".to_string(),
abstract_text: "Arxiv abstract".to_string(),
doi: Some("10.1000/arxiv.doi".to_string()),
pdf_url: "https://arxiv.org/pdf/2606.12345.pdf".to_string(),
};
let paper = convert_arxiv_to_standard(&doc);
assert_eq!(paper.bibcode, "2606.12345");
assert_eq!(paper.title, "Arxiv Title");
assert_eq!(paper.authors, vec!["Author C"]);
assert_eq!(paper.year, "2026");
assert_eq!(paper.pub_journal, "arXiv Preprint");
assert_eq!(paper.doi, "10.1000/arxiv.doi");
assert_eq!(paper.arxiv_id, "2606.12345");
}
#[tokio::test]
async fn test_db_operations() -> anyhow::Result<()> {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await?;
// 运行迁移
sqlx::migrate!("./migrations").run(&pool).await?;
let paper = StandardPaper {
bibcode: "2026A&A...123..456X".to_string(),
title: "A Test Title".to_string(),
authors: vec!["Author A".to_string()],
year: "2026".to_string(),
pub_journal: "Astronomy & Astrophysics".to_string(),
keywords: vec!["Keyword 1".to_string()],
abstract_text: "This is abstract".to_string(),
doi: "10.1000/test.doi".to_string(),
arxiv_id: "".to_string(),
citation_count: 5,
reference_count: 10,
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: "article".to_string(),
pdf_error: None,
html_error: None,
};
// 保存
save_paper_to_db(&pool, &paper).await?;
// 读取
let retrieved =
get_paper_from_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?;
assert_eq!(retrieved.title, paper.title);
assert_eq!(retrieved.authors, paper.authors);
assert_eq!(retrieved.keywords, paper.keywords);
// 读取 by DOI / DOI 前缀
let retrieved_by_doi =
get_paper_from_db(&pool, std::path::Path::new(""), "10.1000/test.doi").await?;
assert_eq!(retrieved_by_doi.bibcode, paper.bibcode);
let retrieved_by_doi_prefix =
get_paper_from_db(&pool, std::path::Path::new(""), "doi:10.1000/test.doi").await?;
assert_eq!(retrieved_by_doi_prefix.bibcode, paper.bibcode);
// 检查路径状态(初始为 None
let paths =
check_paper_paths_in_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?;
assert!(paths.is_some());
let (pdf, html, md, tr) = paths.unwrap();
assert!(pdf.is_none());
assert!(html.is_none());
assert!(md.is_none());
assert!(tr.is_none());
Ok(())
}
}

View File

@ -11,8 +11,8 @@ use crate::Config;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sqlx::SqlitePool; use sqlx::SqlitePool;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock}; use std::sync::Arc;
use tokio::sync::{broadcast, oneshot}; use tokio::sync::{broadcast, oneshot, Mutex, RwLock};
/// 提供给前端的 SSE 事件 /// 提供给前端的 SSE 事件
#[derive(Debug, Clone, Serialize)] #[derive(Debug, Clone, Serialize)]
@ -63,6 +63,8 @@ pub struct AppState {
pub vision_llm: Option<LlmClient>, pub vision_llm: Option<LlmClient>,
pub embedding: EmbeddingClient, pub embedding: EmbeddingClient,
pub downloader: Downloader, pub downloader: Downloader,
/// 共享 HTTP 客户端(带超时),避免每次请求新建连接池
pub http_client: reqwest::Client,
pub harvest_status: Arc<tokio::sync::Mutex<crate::services::batch::MetaSyncStatus>>, pub harvest_status: Arc<tokio::sync::Mutex<crate::services::batch::MetaSyncStatus>>,
pub batch_status: Arc<tokio::sync::Mutex<crate::services::batch::AssetBatchStatus>>, pub batch_status: Arc<tokio::sync::Mutex<crate::services::batch::AssetBatchStatus>>,
pub active_bibcode: Arc<tokio::sync::Mutex<Option<String>>>, pub active_bibcode: Arc<tokio::sync::Mutex<Option<String>>>,
@ -83,42 +85,18 @@ pub struct AppState {
/// 项目记忆管理器(跨会话持久化) /// 项目记忆管理器(跨会话持久化)
pub memory_manager: Arc<tokio::sync::Mutex<MemoryManager>>, pub memory_manager: Arc<tokio::sync::Mutex<MemoryManager>>,
/// 活跃的登录会话 Token 及其最后活跃时间(单用户内存管理) /// 活跃的登录会话 Token 及其最后活跃时间(单用户内存管理)
pub sessions: pub sessions: Arc<Mutex<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
Arc<std::sync::Mutex<std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>>>,
/// 登录速率限制器IP -> 最近失败次数 + 首次失败时间) /// 登录速率限制器IP -> 最近失败次数 + 首次失败时间)
pub login_rate_limiter: pub login_rate_limiter:
Arc<std::sync::Mutex<std::collections::HashMap<String, (u32, std::time::Instant)>>>, Arc<Mutex<std::collections::HashMap<String, (u32, std::time::Instant)>>>,
} }
// 统一标准化的文献格式,用于向前端传输 // 统一标准化的文献格式,用于向前端传输
#[derive(Debug, Serialize, Deserialize, Clone)] pub use crate::services::paper::StandardPaper;
pub struct StandardPaper {
pub bibcode: String,
pub title: String,
pub authors: Vec<String>,
pub year: String,
pub pub_journal: String,
pub keywords: Vec<String>,
pub abstract_text: String,
pub doi: String,
pub arxiv_id: String,
pub citation_count: i32,
pub reference_count: i32,
pub is_downloaded: bool,
pub has_pdf: bool,
pub has_html: bool,
pub has_markdown: bool,
pub has_translation: bool,
pub has_vector: bool,
pub doctype: String,
pub pdf_error: Option<String>,
pub html_error: Option<String>,
}
pub mod agent; pub mod agent;
pub mod auth; pub mod auth;
pub mod error; pub mod error;
pub mod helpers;
pub mod notes; pub mod notes;
pub mod papers; pub mod papers;
pub mod permissions; pub mod permissions;
@ -132,15 +110,10 @@ pub mod handlers {
answer_question, branch_session, chat_agent, delete_session, get_agent_metrics, answer_question, branch_session, chat_agent, delete_session, get_agent_metrics,
get_agent_modes, get_pending_permissions, get_pending_questions, get_session, get_agent_modes, get_pending_permissions, get_pending_questions, get_session,
get_session_audit, list_sessions, respond_permission, restore_rewound_session, get_session_audit, list_sessions, respond_permission, restore_rewound_session,
retry_session, rewind_session, stop_agent, AgentChatRequest, AgentMetricsResponse, retry_session, rewind_session, stop_agent, AgentChatRequest, BranchResponse,
AuditLogEntry, BranchResponse, MessageRecord, RestoreResponse, RetryResponse, RestoreResponse, RetryResponse, RewindRequest, RewindResponse, SessionListParams,
RewindRequest, RewindResponse, SessionDetail, SessionListParams, SessionSummary,
}; };
pub use super::auth::{check_auth, login, logout}; pub use super::auth::{check_auth, login, logout};
pub use super::helpers::{
check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard,
get_paper_from_db, save_paper_to_db, save_paper_to_db_tx,
};
pub use super::notes::{ pub use super::notes::{
create_note, delete_note, get_notes, CreateNoteRequest, DeleteNoteParams, GetNotesParams, create_note, delete_note, get_notes, CreateNoteRequest, DeleteNoteParams, GetNotesParams,
NoteRecord, NoteRecord,
@ -148,10 +121,9 @@ pub mod handlers {
pub use super::papers::{ pub use super::papers::{
download_paper, embed_paper, export_citations, get_active_bibcode, get_citation_network, download_paper, embed_paper, export_citations, get_active_bibcode, get_citation_network,
get_library, get_paper_detail, mark_no_resource, parse_paper, search_papers, get_library, get_paper_detail, mark_no_resource, parse_paper, search_papers,
set_active_bibcode, translate_paper, upload_paper_file, CitationsResponse, DownloadRequest, set_active_bibcode, translate_paper, upload_paper_file, DownloadRequest, EmbedRequest,
EmbedRequest, EmbedResponse, ExportRequest, ExportResponse, MarkNoResourceRequest, EmbedResponse, ExportRequest, ExportResponse, MarkNoResourceRequest, PaperDetailResponse,
PaperDetailResponse, ParseRequest, ParseResponse, SearchParams, TranslateRequest, ParseRequest, ParseResponse, SearchParams, TranslateRequest, TranslateResponse,
TranslateResponse,
}; };
pub use super::permissions::{ pub use super::permissions::{
list_permission_rules, update_permission_mode, update_permission_rules, list_permission_rules, update_permission_mode, update_permission_rules,
@ -160,7 +132,7 @@ pub mod handlers {
pub use super::sync::{ pub use super::sync::{
delete_sync_query, get_asset_batch_status, get_meta_sync_count, get_meta_sync_status, delete_sync_query, get_asset_batch_status, get_meta_sync_count, get_meta_sync_status,
get_sync_queries, run_asset_batch, run_meta_sync, stop_asset_batch, AssetBatchRunRequest, get_sync_queries, run_asset_batch, run_meta_sync, stop_asset_batch, AssetBatchRunRequest,
MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest, SavedSyncQuery, MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest,
}; };
pub use super::targets::{ pub use super::targets::{
associate_target, extract_paper_targets, list_targets, query_target, AssociateResponse, associate_target, extract_paper_targets, list_targets, query_target, AssociateResponse,
@ -168,4 +140,13 @@ pub mod handlers {
TargetQueryParams, TargetQueryParams,
}; };
pub use super::{AppState, StandardPaper}; pub use super::{AppState, StandardPaper};
pub use crate::services::batch::SavedSyncQuery;
pub use crate::services::paper::{
check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard,
get_paper_from_db, save_paper_to_db, save_paper_to_db_tx,
};
pub use crate::services::session::{
AgentMetrics as AgentMetricsResponse, AuditLogEntry, MessageRecord, SessionDetail,
SessionSummary,
};
} }

View File

@ -4,23 +4,12 @@ use axum::{
http::StatusCode, http::StatusCode,
Json, Json,
}; };
use serde::{Deserialize, Serialize}; use serde::Deserialize;
use sqlx::Row;
use std::sync::Arc; use std::sync::Arc;
use super::error::ApiResult; use super::error::ApiResult;
use super::AppState; use super::AppState;
pub use crate::services::note::NoteRecord;
#[derive(Debug, Serialize, Deserialize)]
pub struct NoteRecord {
pub id: i64,
pub bibcode: String,
pub paragraph_index: i64,
pub note_text: String,
pub highlight_color: String,
pub selected_text: String,
pub created_at: String,
}
#[derive(Deserialize)] #[derive(Deserialize)]
pub struct CreateNoteRequest { pub struct CreateNoteRequest {
@ -50,27 +39,18 @@ pub async fn create_note(
let highlight_color = req.highlight_color.unwrap_or_else(|| "yellow".to_string()); let highlight_color = req.highlight_color.unwrap_or_else(|| "yellow".to_string());
let selected_text = req.selected_text.unwrap_or_default(); let selected_text = req.selected_text.unwrap_or_default();
let row = sqlx::query( let record = crate::services::note::create_highlight_note(
"INSERT INTO notes (bibcode, paragraph_index, note_text, highlight_color, selected_text) VALUES (?, ?, ?, ?, ?) RETURNING id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at" &state.db,
&req.bibcode,
req.paragraph_index,
&note_text,
&highlight_color,
&selected_text,
) )
.bind(&req.bibcode)
.bind(req.paragraph_index)
.bind(&note_text)
.bind(&highlight_color)
.bind(&selected_text)
.fetch_one(&state.db)
.await .await
.map_err(|e| super::error::AppError::internal(format!("保存笔记失败: {}", e)))?; .map_err(|e| super::error::AppError::internal(format!("保存笔记失败: {}", e)))?;
Ok(Json(NoteRecord { Ok(Json(record))
id: row.get(0),
bibcode: row.get(1),
paragraph_index: row.get(2),
note_text: row.get(3),
highlight_color: row.get(4),
selected_text: row.get(5),
created_at: row.get(6),
}))
} }
// 查询某篇文献的全部笔记 // 查询某篇文献的全部笔记
@ -78,26 +58,9 @@ pub async fn get_notes(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<GetNotesParams>, Query(params): Query<GetNotesParams>,
) -> ApiResult<Json<Vec<NoteRecord>>> { ) -> ApiResult<Json<Vec<NoteRecord>>> {
let rows = sqlx::query( let notes = crate::services::note::list_highlight_notes(&state.db, &params.bibcode)
"SELECT id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at FROM notes WHERE bibcode = ? ORDER BY paragraph_index, created_at" .await
) .map_err(|e| super::error::AppError::internal(format!("查询笔记失败: {}", e)))?;
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.map_err(|e| super::error::AppError::internal(format!("查询笔记失败: {}", e)))?;
let notes: Vec<NoteRecord> = rows
.iter()
.map(|r| NoteRecord {
id: r.get(0),
bibcode: r.get(1),
paragraph_index: r.get(2),
note_text: r.get(3),
highlight_color: r.get(4),
selected_text: r.get(5),
created_at: r.get(6),
})
.collect();
Ok(Json(notes)) Ok(Json(notes))
} }
@ -107,11 +70,16 @@ pub async fn delete_note(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<DeleteNoteParams>, Query(params): Query<DeleteNoteParams>,
) -> ApiResult<StatusCode> { ) -> ApiResult<StatusCode> {
sqlx::query("DELETE FROM notes WHERE id = ?") let success = crate::services::note::delete_highlight_note(&state.db, params.id)
.bind(params.id)
.execute(&state.db)
.await .await
.map_err(|e| super::error::AppError::internal(format!("删除笔记失败: {}", e)))?; .map_err(|e| super::error::AppError::internal(format!("删除笔记失败: {}", e)))?;
if !success {
return Err(super::error::AppError::not_found(format!(
"未找到 ID 为 {} 的笔记",
params.id
)));
}
Ok(StatusCode::NO_CONTENT) Ok(StatusCode::NO_CONTENT)
} }

View File

@ -5,17 +5,12 @@ use axum::{
Json, Json,
}; };
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::sync::Arc; use std::sync::Arc;
use tokio::fs; use tracing::info;
use tracing::{error, info, warn};
use super::error::{ApiResult, AppError}; use super::error::{ApiResult, AppError};
use super::helpers::{
check_paper_paths_in_db, convert_ads_doc_to_standard, get_paper_from_db, parse_paper_row,
save_paper_to_db, SQLITE_PARAM_LIMIT,
};
use super::{AppState, StandardPaper}; use super::{AppState, StandardPaper};
use crate::services::paper::get_paper_from_db;
// 检索请求参数 // 检索请求参数
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@ -34,8 +29,8 @@ pub async fn search_papers(
Query(params): Query<SearchParams>, Query(params): Query<SearchParams>,
) -> ApiResult<Json<Vec<StandardPaper>>> { ) -> ApiResult<Json<Vec<StandardPaper>>> {
let source = params.source.unwrap_or_else(|| "all".to_string()); let source = params.source.unwrap_or_else(|| "all".to_string());
let rows = params.rows.unwrap_or(10); let rows = params.rows.unwrap_or(10).clamp(1, 200);
let start = params.start.unwrap_or(0); let start = params.start.unwrap_or(0).max(0);
let sort = params.sort.as_deref().unwrap_or("relevance"); let sort = params.sort.as_deref().unwrap_or("relevance");
let results = let results =
@ -106,17 +101,17 @@ pub async fn parse_paper(
) )
.await .await
.map_err(|e| { .map_err(|e| {
// 根据错误信息推断合适的错误类型(保留原有状态码语义)
let msg = e.to_string(); let msg = e.to_string();
if msg.contains("未注册") if msg.contains("未注册")
|| msg.contains("未找到") || msg.contains("未找到")
|| msg.contains("丢失") || msg.contains("丢失")
|| msg.contains("未在数据库中注册") || msg.contains("未在数据库中注册")
{ {
AppError::not_found(msg) AppError::not_found("该文献未找到,请确认标识符是否正确")
} else if msg.contains("请先下载") { } else if msg.contains("请先下载") {
AppError::bad_request(msg) AppError::bad_request("请先下载该文献后再尝试解析")
} else { } else {
tracing::error!("文献解析内部错误: {}", msg);
AppError::internal(msg) AppError::internal(msg)
} }
})?; })?;
@ -147,229 +142,58 @@ pub async fn translate_paper(
req.bibcode, force req.bibcode, force
); );
let (_, _, md_opt, tr_opt) = let translated_markdown = crate::services::translation::translate_paper_cached(
check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode) &state.db,
.await &state.config.library_dir,
.map_err(|e| AppError::not_found(format!("查询文献路径失败: {}", e)))?
.ok_or_else(|| AppError::not_found("该文献未注册在数据库中"))?;
// 若本地已存在翻译物理文件且未指明强制重译,直读本地缓存返回
if !force {
if let Some(tr_rel) = tr_opt {
let tr_abs = state.config.library_dir.join(&tr_rel);
if tr_abs.exists() {
if let Ok(content) = fs::read_to_string(&tr_abs).await {
return Ok(Json(TranslateResponse {
translation: content,
}));
}
}
}
}
// 检查英文解析文件是否存在
let md_rel = match md_opt {
Some(rel) => rel,
None => {
error!(
"文献 {} 翻译失败:文献未完成解析,缺少英文 Markdown 路径",
req.bibcode
);
return Err(AppError::bad_request("文献必须先完成解析方可翻译"));
}
};
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
error!(
"文献 {} 翻译失败:解析的英文 Markdown 文件 {:?} 不存在",
req.bibcode, md_abs
);
return Err(AppError::bad_request("解析 Markdown 文件丢失"));
}
let english_markdown = fs::read_to_string(&md_abs).await.map_err(|e| {
error!("文献 {} 翻译失败:读取解析内容失败: {}", req.bibcode, e);
AppError::internal(format!("读取解析内容失败: {}", e))
})?;
// 调用 LLM 翻译服务并注入对照词表
let translated_markdown = crate::services::translation::translate_markdown(
&english_markdown,
&state.dict, &state.dict,
&state.medium_llm, &state.medium_llm,
&req.bibcode,
force,
) )
.await .await
.map_err(|e| { .map_err(|e| {
error!( let err_msg = e.to_string();
"文献 {} 翻译失败:调用 LLM 翻译发生错误: {}", if err_msg.contains("该文献未注册") || err_msg.contains("丢失") {
req.bibcode, e AppError::not_found("该文献未在本地库中找到")
); } else if err_msg.contains("必须先完成解析") {
AppError::internal(format!("调用 LLM 翻译失败: {}", e)) AppError::bad_request("请先完成文献解析后再翻译")
} else {
tracing::error!("文献翻译内部错误: {}", err_msg);
AppError::internal(err_msg)
}
})?; })?;
// 翻译结果持久化C2抽取为共享函数与批量任务复用
crate::services::translation::persist_translation(
&state.db,
&state.config.library_dir,
&req.bibcode,
&translated_markdown,
)
.await
.map_err(|e| AppError::internal(format!("翻译结果持久化失败: {}", e)))?;
Ok(Json(TranslateResponse { Ok(Json(TranslateResponse {
translation: translated_markdown, translation: translated_markdown,
})) }))
} }
#[derive(Debug, Serialize)]
pub struct CitationsResponse {
pub bibcode: String,
pub title: String,
pub citation_count: i32,
pub reference_count: i32,
pub references: Vec<String>, // 该文献参考文献 bibcode 数组
pub citations: Vec<String>, // 引用该文献的 bibcode 数组
pub citation_counts: std::collections::HashMap<String, i32>, // 相关文献与被引数映射
}
// 从 SQLite 查询引用关联,生成引用星系关系树 // 从 SQLite 查询引用关联,生成引用星系关系树
pub async fn get_citation_network( pub async fn get_citation_network(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<DownloadRequest>, Query(params): Query<DownloadRequest>,
) -> ApiResult<Json<CitationsResponse>> { ) -> ApiResult<Json<crate::services::citation::CitationNetworkResult>> {
let paper = match get_paper_from_db(&state.db, &state.config.library_dir, &params.bibcode).await let res = crate::services::citation::get_citation_network_with_fallback(
{ &state.db,
Ok(p) => p, &state.config.library_dir,
Err(_) => { &state.config.ads_api_key,
// 如果本地数据库查不到,尝试从 ADS 在线 API 动态获取 &state.ads,
if !state.config.ads_api_key.is_empty() { &params.bibcode,
match state )
.ads .await
.search(&format!("bibcode:{}", params.bibcode), 0, 1, "relevance") .map_err(|e| {
.await let msg = e.to_string();
{ if msg.contains("未找到该文献") {
Ok(docs) => { AppError::not_found("未找到该文献的引用关系数据")
if let Some(doc) = docs.first() { } else if msg.contains("未配置 ADS_API_KEY") {
let standard_paper = convert_ads_doc_to_standard(doc); AppError::bad_request("未配置 ADS_API_KEY无法在线加载引用数据")
// 保存至数据库缓存,并保存引用关联 } else {
if let Err(e) = save_paper_to_db(&state.db, &standard_paper).await { tracing::error!("引用网络查询内部错误: {}", msg);
warn!("保存引用文献至数据库失败: {}", e); AppError::internal(msg)
}
if let Some(refs) = &doc.reference {
for ref_bib in refs {
if let Err(e) = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
.bind(&standard_paper.bibcode)
.bind(ref_bib)
.execute(&state.db)
.await
{
warn!("保存引用关系失败 ({} -> {}): {}", standard_paper.bibcode, ref_bib, e);
}
}
}
if let Some(cits) = &doc.citation {
for cit_bib in cits {
if let Err(e) = sqlx::query("INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)")
.bind(cit_bib)
.bind(&standard_paper.bibcode)
.execute(&state.db)
.await
{
warn!("保存被引关系失败 ({} -> {}): {}", cit_bib, standard_paper.bibcode, e);
}
}
}
standard_paper
} else {
return Err(AppError::not_found(format!(
"在本地库及 ADS 中均未找到该文献: {}",
params.bibcode
)));
}
}
Err(e) => {
return Err(AppError::internal(format!("在线检索文献元数据失败: {}", e)));
}
}
} else {
return Err(AppError::not_found(format!(
"本地数据库未收录该文献,且未配置 ADS_API_KEY无法在线加载: {}",
params.bibcode
)));
}
} }
}; })?;
// 加载引用的文献 Ok(Json(res))
let refs_rows =
sqlx::query("SELECT target_bibcode FROM citations_references WHERE source_bibcode = ?")
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.unwrap_or_else(|e| {
warn!("查询引用关系失败: {}", e);
Vec::new()
});
let references: Vec<String> = refs_rows.iter().map(|row| row.get(0)).collect();
// 加载被引用的文献
let cits_rows =
sqlx::query("SELECT source_bibcode FROM citations_references WHERE target_bibcode = ?")
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.unwrap_or_else(|e| {
warn!("查询被引关系失败: {}", e);
Vec::new()
});
let citations: Vec<String> = cits_rows.iter().map(|row| row.get(0)).collect();
// 加载关联文献的被引数量 (从 SQLite papers 表获取)
// 优化:用 WHERE IN 批量查询替代逐条 query_scalar避免 N+1。
let mut citation_counts = std::collections::HashMap::new();
let mut all_related = references.clone();
all_related.extend(citations.clone());
// 去重,避免重复绑定与查询
let mut unique_bibs: Vec<String> = all_related.clone();
unique_bibs.sort();
unique_bibs.dedup();
// SQLite 单条 SQL 参数上限 999分批查询
for chunk in unique_bibs.chunks(SQLITE_PARAM_LIMIT) {
if chunk.is_empty() {
continue;
}
let placeholders = vec!["?"; chunk.len()].join(", ");
let sql = format!(
"SELECT bibcode, citation_count FROM papers WHERE bibcode IN ({})",
placeholders
);
let mut query = sqlx::query(&sql);
for b in chunk {
query = query.bind(b);
}
let rows = query.fetch_all(&state.db).await.unwrap_or_else(|e| {
warn!("批量查询文献引用数失败: {}", e);
Vec::new()
});
for row in rows {
let bib: String = row.get(0);
let count: i32 = row.get(1);
citation_counts.insert(bib, count);
}
}
Ok(Json(CitationsResponse {
bibcode: paper.bibcode,
title: paper.title,
citation_count: paper.citation_count,
reference_count: paper.reference_count,
references,
citations,
citation_counts,
}))
} }
// ── GET /api/paper ── // ── GET /api/paper ──
@ -389,23 +213,20 @@ pub async fn get_paper_detail(
.await .await
.map_err(|e| AppError::not_found(format!("未找到该文献数据: {}", e)))?; .map_err(|e| AppError::not_found(format!("未找到该文献数据: {}", e)))?;
let (_, _, md_opt, tr_opt) = let mode = crate::services::paper::ReadMode::Full {
check_paper_paths_in_db(&state.db, &state.config.library_dir, &params.bibcode) include_translation: true,
.await
.map_err(|e| AppError::internal(e.to_string()))?
.unwrap_or_default();
let english_content = match md_opt {
Some(rel) => fs::read_to_string(state.config.library_dir.join(rel))
.await
.ok(),
None => None,
}; };
let translation_content = match tr_opt {
Some(rel) => fs::read_to_string(state.config.library_dir.join(rel)) let (english_content, translation_content) = match crate::services::paper::read_paper_content(
.await &state.db,
.ok(), &state.config.library_dir,
None => None, &params.bibcode,
mode,
)
.await
{
Ok(res) => (Some(res.content), res.translation_content),
Err(_) => (None, None),
}; };
Ok(Json(PaperDetailResponse { Ok(Json(PaperDetailResponse {
@ -416,19 +237,25 @@ pub async fn get_paper_detail(
} }
// ── GET /api/library ── // ── GET /api/library ──
#[derive(Deserialize)]
pub struct LibraryParams {
pub limit: Option<i64>,
pub offset: Option<i64>,
}
// 获取本地图书馆文献列表 // 获取本地图书馆文献列表
pub async fn get_library( pub async fn get_library(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<LibraryParams>,
) -> ApiResult<Json<Vec<StandardPaper>>> { ) -> ApiResult<Json<Vec<StandardPaper>>> {
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) FROM papers ORDER BY created_at DESC") let list = crate::services::paper::get_library_list(
.fetch_all(&state.db) &state.db,
.await &state.config.library_dir,
.map_err(|e| AppError::internal(format!("访问本地数据库失败: {}", e)))?; params.limit,
params.offset,
let mut list = Vec::new(); )
for r in rows { .await
list.push(parse_paper_row(&r, &state.config.library_dir)); .map_err(|e| AppError::internal(format!("访问本地数据库失败: {}", e)))?;
}
Ok(Json(list)) Ok(Json(list))
} }
@ -508,128 +335,26 @@ pub async fn upload_paper_file(
return Err(AppError::bad_request("bibcode 包含非法字符")); return Err(AppError::bad_request("bibcode 包含非法字符"));
} }
if file_bytes.is_empty() { let updated_paper = crate::services::paper::upload_paper_file_service(
return Err(AppError::bad_request("上传文件为空或读取失败")); &state.db,
} &state.config.library_dir,
&bibcode,
// 尝试将可能的 DOI 或 arXiv ID 解析为真实的 bibcode &file_type,
let mut resolved_bibcode = bibcode.clone(); &file_name,
let exists_as_bibcode = sqlx::query("SELECT bibcode FROM papers WHERE bibcode = ?") &file_bytes,
.bind(&resolved_bibcode) )
.fetch_optional(&state.db) .await
.await .map_err(|e| {
.unwrap_or(None) let err_msg = e.to_string();
.is_some(); if err_msg.contains("未找到该文献记录") {
AppError::not_found("未找到该文献记录,请先在库中检索该文献")
if !exists_as_bibcode { } else if err_msg.contains("校验失败") || err_msg.contains("不是有效的 UTF-8") {
// 尝试匹配 DOI AppError::bad_request("上传文件校验失败,请确认文件格式正确")
let clean_doi = resolved_bibcode
.trim_start_matches("doi:")
.trim_start_matches("DOI:")
.trim_start_matches("https://doi.org/")
.trim_start_matches("http://doi.org/")
.trim();
if let Some(row) = sqlx::query(
"SELECT bibcode FROM papers WHERE doi = ? OR doi = ? OR LOWER(doi) = LOWER(?)",
)
.bind(clean_doi)
.bind(&resolved_bibcode)
.bind(clean_doi)
.fetch_optional(&state.db)
.await
.unwrap_or(None)
{
let found: String = row.get(0);
info!(
"上传接口:通过 DOI 匹配成功,将 '{}' 解析为 bibcode '{}'",
bibcode, found
);
resolved_bibcode = found;
} else { } else {
// 尝试匹配 arXiv ID tracing::error!("文件上传内部错误: {}", err_msg);
let clean_arxiv = resolved_bibcode AppError::internal(err_msg)
.trim_start_matches("arxiv:")
.trim_start_matches("arXiv:")
.trim_start_matches("ARXIV:")
.trim();
// 移除可能存在的版本号后缀(如 2303.12345v1 -> 2303.12345
let clean_arxiv_no_version = if let Some(pos) = clean_arxiv.find('v') {
if clean_arxiv[pos + 1..].chars().all(|c| c.is_ascii_digit()) {
&clean_arxiv[..pos]
} else {
clean_arxiv
}
} else {
clean_arxiv
};
if let Some(row) = sqlx::query(
"SELECT bibcode FROM papers WHERE arxiv_id = ? OR arxiv_id = ? OR arxiv_id LIKE ? OR arxiv_id LIKE ?"
)
.bind(clean_arxiv)
.bind(clean_arxiv_no_version)
.bind(format!("{}%", clean_arxiv_no_version))
.bind(format!("arXiv:{}%", clean_arxiv_no_version))
.fetch_optional(&state.db)
.await
.unwrap_or(None) {
let found: String = row.get(0);
info!("上传接口:通过 arXiv ID 匹配成功,将 '{}' 解析为 bibcode '{}'", bibcode, found);
resolved_bibcode = found;
}
} }
} })?;
let bibcode = resolved_bibcode;
// 从数据库读取该文献元数据
let _paper = get_paper_from_db(&state.db, &state.config.library_dir, &bibcode)
.await
.map_err(|e| AppError::not_found(format!("未找到该文献记录: {}", e)))?;
// 校验并保存文件
let is_pdf = file_type == "pdf" || file_name.to_lowercase().ends_with(".pdf");
let relative_path = if is_pdf {
crate::services::download::validate_pdf_content(&file_bytes)
.map_err(|e| AppError::bad_request(format!("PDF 文件内容校验失败: {}", e)))?;
let pdf_filename = format!("{}.pdf", bibcode);
let pdf_dest = state.config.library_dir.join("PDF").join(&pdf_filename);
if let Some(parent) = pdf_dest.parent() {
std::fs::create_dir_all(parent).unwrap_or_default();
}
std::fs::write(&pdf_dest, &file_bytes)
.map_err(|e| AppError::internal(format!("无法写入 PDF 文件: {}", e)))?;
format!("PDF/{}", pdf_filename)
} else {
let text_content = String::from_utf8(file_bytes)
.map_err(|_| AppError::bad_request("上传的 HTML 文件不是有效的 UTF-8 文本"))?;
crate::services::download::validate_html_content_lenient(&text_content)
.map_err(|e| AppError::bad_request(format!("HTML 文件内容校验失败: {}", e)))?;
let html_filename = format!("{}.html", bibcode);
let html_dest = state.config.library_dir.join("HTML").join(&html_filename);
if let Some(parent) = html_dest.parent() {
std::fs::create_dir_all(parent).unwrap_or_default();
}
std::fs::write(&html_dest, &text_content)
.map_err(|e| AppError::internal(format!("无法写入 HTML 文件: {}", e)))?;
format!("HTML/{}", html_filename)
};
// 更新数据库路径状态
let path_field = if is_pdf { "pdf_path" } else { "html_path" };
let sql = format!("UPDATE papers SET {} = ? WHERE bibcode = ?", path_field);
sqlx::query(&sql)
.bind(&relative_path)
.bind(&bibcode)
.execute(&state.db)
.await
.map_err(|e| AppError::internal(format!("更新数据库状态失败: {}", e)))?;
// 重新获取最新的文献信息以更新前端界面
let updated_paper = get_paper_from_db(&state.db, &state.config.library_dir, &bibcode)
.await
.map_err(|e| AppError::internal(format!("重读文献数据失败: {}", e)))?;
Ok(Json(updated_paper)) Ok(Json(updated_paper))
} }
@ -646,27 +371,19 @@ pub async fn mark_no_resource(
Json(req): Json<MarkNoResourceRequest>, Json(req): Json<MarkNoResourceRequest>,
) -> ApiResult<Json<StandardPaper>> { ) -> ApiResult<Json<StandardPaper>> {
let clear_flag = req.clear.unwrap_or(false); let clear_flag = req.clear.unwrap_or(false);
info!(
"接收到文献无资源标记指令,标识符: {}, clear: {}",
req.bibcode, clear_flag
);
if clear_flag { let updated_paper = crate::services::paper::mark_no_resource_service(
info!("接收到清除文献无资源标记指令,标识符: {}", req.bibcode); &state.db,
sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?") &state.config.library_dir,
.bind(&req.bibcode) &req.bibcode,
.execute(&state.db) clear_flag,
.await )
.map_err(|e| AppError::internal(format!("清除无资源标记失败: {}", e)))?; .await
} else { .map_err(|e| AppError::internal(format!("更新文献无资源标记失败: {}", e)))?;
info!("接收到文献无资源标记指令,标识符: {}", req.bibcode);
sqlx::query("UPDATE papers SET pdf_path = 'error:no_resource', html_path = 'error:no_resource' WHERE bibcode = ?")
.bind(&req.bibcode)
.execute(&state.db)
.await
.map_err(|e| AppError::internal(format!("更新数据库无资源标记失败: {}", e)))?;
}
// 重新获取最新的文献信息以更新前端界面
let updated_paper = get_paper_from_db(&state.db, &state.config.library_dir, &req.bibcode)
.await
.map_err(|e| AppError::internal(format!("重读文献数据失败: {}", e)))?;
Ok(Json(updated_paper)) Ok(Json(updated_paper))
} }
@ -717,32 +434,23 @@ pub async fn embed_paper(
) -> ApiResult<Json<EmbedResponse>> { ) -> ApiResult<Json<EmbedResponse>> {
info!("接收到文献向量化分块入库指令: {}", req.bibcode); info!("接收到文献向量化分块入库指令: {}", req.bibcode);
let (_, _, md_opt, _) = let chunk_count = crate::services::rag::embed_paper_service(
check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode)
.await
.map_err(|e| AppError::not_found(format!("获取文献路径失败: {}", e)))?
.ok_or_else(|| AppError::not_found("该文献未注册在数据库中"))?;
let md_rel =
md_opt.ok_or_else(|| AppError::bad_request("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = state.config.library_dir.join(&md_rel);
if !md_abs.exists() {
return Err(AppError::not_found("文献 Markdown 文件未找到,请重新解析"));
}
let markdown_content = fs::read_to_string(&md_abs)
.await
.map_err(|e| AppError::internal(format!("读取 Markdown 文件失败: {}", e)))?;
let chunk_count = crate::services::rag::ingest_paper(
&state.db, &state.db,
&state.embedding, &state.config.library_dir,
&req.bibcode, &req.bibcode,
&markdown_content, &state.embedding,
None,
) )
.await .await
.map_err(|e| AppError::internal(format!("文献向量化失败: {}", e)))?; .map_err(|e| {
let err_msg = e.to_string();
if err_msg.contains("未注册") || err_msg.contains("未找到") {
AppError::not_found(err_msg)
} else if err_msg.contains("尚未解析") || err_msg.contains("未找到") {
AppError::bad_request(err_msg)
} else {
AppError::internal(err_msg)
}
})?;
Ok(Json(EmbedResponse { chunk_count })) Ok(Json(EmbedResponse { chunk_count }))
} }

View File

@ -52,6 +52,26 @@ pub struct PermissionActionResponse {
// ── Handlers ── // ── Handlers ──
/// 验证会话是否存在(防止对不存在的会话进行权限操作)
async fn verify_session_exists(
db: &sqlx::SqlitePool,
session_id: &str,
) -> Result<(), Json<PermissionActionResponse>> {
let exists: bool =
sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ?)")
.bind(session_id)
.fetch_one(db)
.await
.unwrap_or(false);
if !exists {
return Err(Json(PermissionActionResponse {
success: false,
message: format!("会话 {} 不存在", session_id),
}));
}
Ok(())
}
/// 动态添加或移除权限规则。 /// 动态添加或移除权限规则。
pub async fn update_permission_rules( pub async fn update_permission_rules(
Path(session_id): Path<String>, Path(session_id): Path<String>,
@ -63,6 +83,11 @@ pub async fn update_permission_rules(
session_id, req.action, req.kind, req.rule session_id, req.action, req.kind, req.rule
); );
// 验证会话存在
if let Err(resp) = verify_session_exists(&state.db, &session_id).await {
return resp;
}
if !["deny", "allow", "ask"].contains(&req.kind.as_str()) { if !["deny", "allow", "ask"].contains(&req.kind.as_str()) {
return Json(PermissionActionResponse { return Json(PermissionActionResponse {
success: false, success: false,
@ -105,43 +130,31 @@ pub async fn update_permission_rules(
}, },
}; };
if let Ok(mut checkers) = state.session_permission_checkers.write() { let mut checkers = state.session_permission_checkers.write().await;
let checker = checkers let checker = checkers
.entry(session_id.clone()) .entry(session_id.clone())
.or_insert_with(PermissionChecker::new); .or_insert_with(PermissionChecker::new);
checker.add_rule_dynamic(session_rule); checker.add_rule_dynamic(session_rule);
Json(PermissionActionResponse { Json(PermissionActionResponse {
success: true, success: true,
message: format!("已添加 {}:{} 规则", req.kind, req.rule), message: format!("已添加 {}:{} 规则", req.kind, req.rule),
}) })
} else {
Json(PermissionActionResponse {
success: false,
message: "权限检查器锁定失败".to_string(),
})
}
} }
"remove" => { "remove" => {
if let Ok(mut checkers) = state.session_permission_checkers.write() { let mut checkers = state.session_permission_checkers.write().await;
let removed = if let Some(checker) = checkers.get_mut(&session_id) { let removed = if let Some(checker) = checkers.get_mut(&session_id) {
checker.remove_rule_dynamic(&req.rule, &req.kind) checker.remove_rule_dynamic(&req.rule, &req.kind)
} else {
0
};
Json(PermissionActionResponse {
success: removed > 0,
message: if removed > 0 {
format!("已移除 {}{}/{} 规则", removed, req.kind, req.rule)
} else {
format!("未找到匹配的 {}/{} 规则", req.kind, req.rule)
},
})
} else { } else {
Json(PermissionActionResponse { 0
success: false, };
message: "权限检查器锁定失败".to_string(), Json(PermissionActionResponse {
}) success: removed > 0,
} message: if removed > 0 {
format!("已移除 {}{}/{} 规则", removed, req.kind, req.rule)
} else {
format!("未找到匹配的 {}/{} 规则", req.kind, req.rule)
},
})
} }
other => Json(PermissionActionResponse { other => Json(PermissionActionResponse {
success: false, success: false,
@ -156,6 +169,11 @@ pub async fn update_permission_mode(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Json(req): Json<PermissionModeRequest>, Json(req): Json<PermissionModeRequest>,
) -> Json<PermissionActionResponse> { ) -> Json<PermissionActionResponse> {
// 验证会话存在
if let Err(resp) = verify_session_exists(&state.db, &session_id).await {
return resp;
}
let mode = PermissionMode::from_str(&req.mode); let mode = PermissionMode::from_str(&req.mode);
let mode_str = format!("{:?}", mode); let mode_str = format!("{:?}", mode);
@ -164,23 +182,17 @@ pub async fn update_permission_mode(
session_id, mode_str session_id, mode_str
); );
if let Ok(mut checkers) = state.session_permission_checkers.write() { let mut checkers = state.session_permission_checkers.write().await;
if let Some(checker) = checkers.get_mut(&session_id) { if let Some(checker) = checkers.get_mut(&session_id) {
checker.set_mode(mode); checker.set_mode(mode);
Json(PermissionActionResponse { Json(PermissionActionResponse {
success: true, success: true,
message: format!("权限模式已切换为: {}", mode_str), message: format!("权限模式已切换为: {}", mode_str),
}) })
} else {
Json(PermissionActionResponse {
success: false,
message: format!("会话 {} 不存在或无活跃 Agent", session_id),
})
}
} else { } else {
Json(PermissionActionResponse { Json(PermissionActionResponse {
success: false, success: false,
message: "权限检查器锁定失败".to_string(), message: format!("会话 {} 不存在或无活跃 Agent", session_id),
}) })
} }
} }
@ -190,17 +202,7 @@ pub async fn list_permission_rules(
Path(session_id): Path<String>, Path(session_id): Path<String>,
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
) -> Json<PermissionRulesResponse> { ) -> Json<PermissionRulesResponse> {
let checkers = match state.session_permission_checkers.read() { let checkers = state.session_permission_checkers.read().await;
Ok(c) => c,
Err(_) => {
return Json(PermissionRulesResponse {
deny_rules: Vec::new(),
allow_rules: Vec::new(),
ask_rules: Vec::new(),
mode: "default".to_string(),
});
}
};
if let Some(checker) = checkers.get(&session_id) { if let Some(checker) = checkers.get(&session_id) {
let mut deny_rules = Vec::new(); let mut deny_rules = Vec::new();

View File

@ -7,11 +7,12 @@ use axum::{
extract::{Query, State}, extract::{Query, State},
Json, Json,
}; };
use serde::{Deserialize, Serialize}; use serde::Deserialize;
use std::sync::Arc; use std::sync::Arc;
use super::error::{ApiResult, AppError}; use super::error::{ApiResult, AppError};
use super::AppState; use super::AppState;
pub use crate::services::search::SearchResult;
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
pub struct SearchParams { pub struct SearchParams {
@ -34,102 +35,22 @@ fn default_limit() -> i64 {
20 20
} }
#[derive(Debug, Serialize)]
pub struct SearchResult {
pub result_type: String,
pub session_id: String,
pub title: Option<String>,
pub snippet: String,
pub created_at: Option<String>,
}
/// GET /api/search /// GET /api/search
pub async fn search( pub async fn search(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<SearchParams>, Query(params): Query<SearchParams>,
) -> ApiResult<Json<Vec<SearchResult>>> { ) -> ApiResult<Json<Vec<SearchResult>>> {
let scope = &params.scope;
let limit = params.limit.min(100); let limit = params.limit.min(100);
let mut results = Vec::new();
let db = &state.db; let results = crate::services::search::search_agent_history(
&state.db,
// 搜索会话 &params.q,
if scope == "all" || scope == "sessions" { &params.scope,
let rows = sqlx::query_as::<_, (String, String, String, Option<String>)>( limit,
"SELECT s.session_id, s.title, \ params.session_id.as_deref(),
snippet(agent_sessions_fts, 1, '<mark>', '</mark>', '...', 40) as snippet, \ )
s.created_at \ .await
FROM agent_sessions_fts fts \ .map_err(|e| AppError::internal(format!("全文检索失败: {}", e)))?;
JOIN agent_sessions s ON s.session_id = fts.session_id \
WHERE agent_sessions_fts MATCH $1 \
ORDER BY rank LIMIT $2",
)
.bind(&params.q)
.bind(limit)
.fetch_all(db)
.await
.map_err(|e| AppError::internal(format!("搜索会话失败: {}", e)))?;
for (sid, title, snippet, created_at) in rows {
results.push(SearchResult {
result_type: "session".into(),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
// 搜索消息
if scope == "all" || scope == "messages" {
let msg_rows = if let Some(ref sid) = params.session_id {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
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 fts.session_id = $3 \
AND m.active = 1 \
ORDER BY rank LIMIT $2",
)
.bind(&params.q)
.bind(limit)
.bind(sid)
.fetch_all(db)
.await
.map_err(|e| AppError::internal(format!("搜索消息失败: {}", e)))?
} else {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
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 $2",
)
.bind(&params.q)
.bind(limit)
.fetch_all(db)
.await
.map_err(|e| AppError::internal(format!("搜索消息失败: {}", e)))?
};
for (sid, title, snippet, role, created_at) in msg_rows {
results.push(SearchResult {
result_type: format!("message/{}", role),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
Ok(Json(results)) Ok(Json(results))
} }

View File

@ -5,7 +5,6 @@ use axum::{
Json, Json,
}; };
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::sync::Arc; use std::sync::Arc;
use tracing::error; use tracing::error;
@ -24,7 +23,7 @@ pub async fn run_meta_sync(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Json(req): Json<MetaSyncRunRequest>, Json(req): Json<MetaSyncRunRequest>,
) -> ApiResult<StatusCode> { ) -> ApiResult<StatusCode> {
// 检查并同步初始化任务状态,防止前端轮询竞态条件 // 检查并同步初始化任务状态——原子地检查+设置,消除 TOCTOU 竞态
{ {
let mut status = state.harvest_status.lock().await; let mut status = state.harvest_status.lock().await;
if status.active { if status.active {
@ -100,30 +99,19 @@ pub struct AssetBatchRunRequest {
pub skip_preceding_uncompleted: Option<bool>, // 跳过前置未完成 pub skip_preceding_uncompleted: Option<bool>, // 跳过前置未完成
} }
struct PaperRecord {
bibcode: String,
pdf_path: Option<String>,
html_path: Option<String>,
markdown_path: Option<String>,
translation_path: Option<String>,
year: String,
created_at: String,
has_vector: bool,
has_target: bool,
}
pub async fn run_asset_batch( pub async fn run_asset_batch(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Json(req): Json<AssetBatchRunRequest>, Json(req): Json<AssetBatchRunRequest>,
) -> ApiResult<StatusCode> { ) -> ApiResult<StatusCode> {
// 检查是否已经在进行批量处理任务 // 原子地检查+设置 active 标志,消除 TOCTOU 竞态
{ {
let status = state.batch_status.lock().await; let mut status = state.batch_status.lock().await;
if status.active { if status.active {
return Err(AppError::conflict( return Err(AppError::conflict(
"当前已有文献批量任务在后台运行中,请勿重复启动", "当前已有文献批量任务在后台运行中,请勿重复启动",
)); ));
} }
status.active = true;
} }
let target_phase = req.target_phase.clone(); let target_phase = req.target_phase.clone();
@ -136,154 +124,23 @@ pub async fn run_asset_batch(
_ => return Err(AppError::bad_request("不支持的 target_phase 参数值")), _ => return Err(AppError::bad_request("不支持的 target_phase 参数值")),
}; };
let rows = sqlx::query("SELECT bibcode, pdf_path, html_path, markdown_path, translation_path, year, datetime(created_at, 'localtime'), EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode), EXISTS(SELECT 1 FROM paper_targets WHERE bibcode = papers.bibcode) FROM papers")
.fetch_all(&state.db)
.await
.map_err(|e| AppError::internal(format!("读取数据库失败: {}", e)))?;
let mut records = Vec::new();
for r in rows {
records.push(PaperRecord {
bibcode: r.get(0),
pdf_path: r.get(1),
html_path: r.get(2),
markdown_path: r.get(3),
translation_path: r.get(4),
year: r.get(5),
created_at: r.get(6),
has_vector: r.get(7),
has_target: r.get(8),
});
}
// 排序
let sort_order = req
.sort_order
.clone()
.unwrap_or_else(|| "default".to_string());
if sort_order == "pub_year_desc" {
records.sort_by(|a, b| b.year.cmp(&a.year));
} else if sort_order == "created_at_desc" {
records.sort_by(|a, b| b.created_at.cmp(&a.created_at));
}
let skip_completed = req.skip_completed.unwrap_or(false); let skip_completed = req.skip_completed.unwrap_or(false);
let skip_failed = req.skip_failed.unwrap_or(false); let skip_failed = req.skip_failed.unwrap_or(false);
let skip_preceding_failed = req.skip_preceding_failed.unwrap_or(false); let skip_preceding_failed = req.skip_preceding_failed.unwrap_or(false);
let skip_preceding_uncompleted = req.skip_preceding_uncompleted.unwrap_or(false); let skip_preceding_uncompleted = req.skip_preceding_uncompleted.unwrap_or(false);
let is_completed = |path: &Option<String>| -> bool { let target_bibcodes = crate::services::batch::filter_papers_for_batch(
if let Some(p) = path { &state.db,
!p.is_empty() && !p.starts_with("error:") && !p.starts_with("mineru_batch:") &target_phase,
} else { req.limit_count,
false req.sort_order.clone(),
} skip_completed,
}; skip_failed,
skip_preceding_failed,
let is_failed = |path: &Option<String>| -> bool { skip_preceding_uncompleted,
if let Some(p) = path { )
p.starts_with("error:") .await
} else { .map_err(|e| AppError::internal(format!("筛选待处理文献失败: {}", e)))?;
false
}
};
let is_no_resource = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:no_resource")
|| p.starts_with("error:无资源")
|| p.starts_with("error:无有效全文")
} else {
false
}
};
let mut target_bibcodes = Vec::new();
for rec in records {
let is_no_resource_paper = is_no_resource(&rec.pdf_path) || is_no_resource(&rec.html_path);
if is_no_resource_paper {
continue;
}
let download_completed = is_completed(&rec.pdf_path) || is_completed(&rec.html_path);
let download_failed = is_failed(&rec.pdf_path) || is_failed(&rec.html_path);
let download_uncompleted = !download_completed && !download_failed;
let parse_completed = is_completed(&rec.markdown_path);
let parse_failed = is_failed(&rec.markdown_path);
let parse_uncompleted = !parse_completed && !parse_failed;
let translate_completed = is_completed(&rec.translation_path);
let translate_failed = is_failed(&rec.translation_path);
match target_phase.as_str() {
"download" => {
if skip_completed && download_completed {
continue;
}
if skip_failed && download_failed {
continue;
}
}
"parse" => {
if skip_completed && parse_completed {
continue;
}
if skip_failed && parse_failed {
continue;
}
if skip_preceding_failed && download_failed {
continue;
}
if skip_preceding_uncompleted && download_uncompleted {
continue;
}
}
"translate" => {
if skip_completed && translate_completed {
continue;
}
if skip_failed && translate_failed {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"embed" => {
if skip_completed && rec.has_vector {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"target" => {
if skip_completed && rec.has_target {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
_ => {}
}
target_bibcodes.push(rec.bibcode);
}
let limit = req.limit_count.unwrap_or(100) as usize;
if target_bibcodes.len() > limit {
target_bibcodes.truncate(limit);
}
if target_bibcodes.is_empty() { if target_bibcodes.is_empty() {
// 没有需要处理的文献 —— 返回 OK软成功前端据 status 判断) // 没有需要处理的文献 —— 返回 OK软成功前端据 status 判断)
@ -316,37 +173,16 @@ pub async fn stop_asset_batch(State(state): State<Arc<AppState>>) -> StatusCode
} }
// ── GET /api/sync/queries ── // ── GET /api/sync/queries ──
#[derive(Debug, Serialize, Deserialize)]
pub struct SavedSyncQuery {
pub id: i64,
pub query: String,
pub source: String,
pub limit_count: i32,
pub last_run: String,
}
pub async fn get_sync_queries( pub async fn get_sync_queries(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
) -> ApiResult<Json<Vec<SavedSyncQuery>>> { ) -> ApiResult<Json<Vec<crate::services::batch::SavedSyncQuery>>> {
let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') FROM sync_queries ORDER BY last_run DESC") let list = crate::services::batch::MetaSync::list_queries(&state.db)
.fetch_all(&state.db)
.await .await
.map_err(|e| { .map_err(|e| {
error!("获取已存同步检索配置失败: {}", e); error!("获取已存同步检索配置失败: {}", e);
AppError::internal(format!("获取已存同步检索配置失败: {}", e)) AppError::internal(format!("获取已存同步检索配置失败: {}", e))
})?; })?;
let mut list = Vec::new();
for r in rows {
list.push(SavedSyncQuery {
id: r.get(0),
query: r.get(1),
source: r.get(2),
limit_count: r.get(3),
last_run: r.get(4),
});
}
Ok(Json(list)) Ok(Json(list))
} }
@ -355,15 +191,20 @@ pub async fn delete_sync_query(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
axum::extract::Path(id): axum::extract::Path<i64>, axum::extract::Path(id): axum::extract::Path<i64>,
) -> ApiResult<StatusCode> { ) -> ApiResult<StatusCode> {
sqlx::query("DELETE FROM sync_queries WHERE id = ?") let success = crate::services::batch::MetaSync::delete_query(&state.db, id)
.bind(id)
.execute(&state.db)
.await .await
.map_err(|e| { .map_err(|e| {
error!("删除同步检索配置失败: {}", e); error!("删除同步检索配置失败: {}", e);
AppError::internal(format!("删除同步检索配置失败: {}", e)) AppError::internal(format!("删除同步检索配置失败: {}", e))
})?; })?;
if !success {
return Err(AppError::not_found(format!(
"未找到 ID 为 {} 的检索配置",
id
)));
}
Ok(StatusCode::OK) Ok(StatusCode::OK)
} }

View File

@ -37,8 +37,7 @@ pub async fn query_target(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<TargetQueryParams>, Query(params): Query<TargetQueryParams>,
) -> ApiResult<Json<TargetInfo>> { ) -> ApiResult<Json<TargetInfo>> {
let client = reqwest::Client::new(); let info = query_target_cached(&state.db, &params.object_name, None, &state.http_client)
let info = query_target_cached(&state.db, &params.object_name, None, &client)
.await .await
.map_err(|e| { .map_err(|e| {
tracing::error!("查询天体信息失败 ({}): {}", params.object_name, e); tracing::error!("查询天体信息失败 ({}): {}", params.object_name, e);
@ -53,20 +52,23 @@ pub async fn associate_target(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Json(req): Json<AssociateTargetRequest>, Json(req): Json<AssociateTargetRequest>,
) -> ApiResult<Json<AssociateResponse>> { ) -> ApiResult<Json<AssociateResponse>> {
let client = reqwest::Client::new();
// 查询并将结果缓存/关联到对应文献 // 查询并将结果缓存/关联到对应文献
let info = query_target_cached(&state.db, &req.object_name, Some(&req.bibcode), &client) let info = query_target_cached(
.await &state.db,
.map_err(|e| { &req.object_name,
tracing::error!( Some(&req.bibcode),
"手动关联天体失败 ({} -> {}): {}", &state.http_client,
req.object_name, )
req.bibcode, .await
e .map_err(|e| {
); tracing::error!(
AppError::internal(format!("手动关联天体失败: {}", e)) "手动关联天体失败 ({} -> {}): {}",
})?; req.object_name,
req.bibcode,
e
);
AppError::internal(format!("手动关联天体失败: {}", e))
})?;
Ok(Json(AssociateResponse { Ok(Json(AssociateResponse {
status: "success".to_string(), status: "success".to_string(),
@ -79,80 +81,12 @@ pub async fn list_targets(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
Query(params): Query<TargetListParams>, Query(params): Query<TargetListParams>,
) -> ApiResult<Json<Vec<TargetInfo>>> { ) -> ApiResult<Json<Vec<TargetInfo>>> {
let rows = sqlx::query_as::< let targets = crate::services::target::list_targets_for_paper(&state.db, &params.bibcode)
_, .await
( .map_err(|e| {
String, tracing::error!("获取文献天体关联列表失败 ({}): {}", params.bibcode, e);
Option<String>, AppError::internal(format!("获取文献天体关联列表失败: {}", e))
Option<String>, })?;
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE bibcode = ? ORDER BY target_name",
)
.bind(&params.bibcode)
.fetch_all(&state.db)
.await
.map_err(|e| {
tracing::error!("获取文献天体关联列表失败 ({}): {}", params.bibcode, e);
AppError::internal(format!("获取文献天体关联列表失败: {}", e))
})?;
let targets: Vec<TargetInfo> = rows
.into_iter()
.map(
|(
name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
)| {
let aliases: Vec<String> = aliases_json
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
TargetInfo {
target_name: name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
}
},
)
.collect();
Ok(Json(targets)) Ok(Json(targets))
} }
@ -174,43 +108,24 @@ pub async fn extract_paper_targets(
) -> ApiResult<Json<ExtractTargetsResponse>> { ) -> ApiResult<Json<ExtractTargetsResponse>> {
tracing::info!("接收到文献天体提取与识别指令: {}", req.bibcode); tracing::info!("接收到文献天体提取与识别指令: {}", req.bibcode);
let (_, _, md_opt, _) = crate::api::helpers::check_paper_paths_in_db( let targets = crate::services::target::extract_and_refresh_targets(
&state.db, &state.db,
&state.config.library_dir, &state.config.library_dir,
&req.bibcode, &req.bibcode,
&state.http_client,
) )
.await .await
.map_err(|e| AppError::not_found(format!("获取文献路径失败: {}", e)))? .map_err(|e| {
.ok_or_else(|| AppError::not_found("该文献未注册在数据库中"))?; let err_msg = e.to_string();
if err_msg.contains("该文献未注册") || err_msg.contains("未找到") {
let md_rel = AppError::not_found("未在本地库中找到该文献")
md_opt.ok_or_else(|| AppError::bad_request("文献尚未解析为 Markdown请先执行解析"))?; } else if err_msg.contains("尚未解析") {
let md_abs = state.config.library_dir.join(&md_rel); AppError::bad_request("请先完成文献解析后再提取天体目标")
if !md_abs.exists() { } else {
return Err(AppError::not_found("文献 Markdown 文件未找到,请重新解析")); tracing::error!("天体目标提取内部错误: {}", err_msg);
} AppError::internal(err_msg)
}
let markdown_content = tokio::fs::read_to_string(&md_abs) })?;
.await
.map_err(|e| AppError::internal(format!("读取 Markdown 文件失败: {}", e)))?;
// 重新识别前先清除该文献已有的天体关联记录,确保陈旧和错误绑定的天体得到重置与刷新
if let Err(e) = sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?")
.bind(&req.bibcode)
.execute(&state.db)
.await
{
tracing::error!("清除文献 {} 的旧天体关联失败: {}", req.bibcode, e);
}
let client = reqwest::Client::new();
let targets = crate::services::target::extract_and_cache_targets(
&state.db,
&markdown_content,
&req.bibcode,
&client,
)
.await;
Ok(Json(ExtractTargetsResponse { targets })) Ok(Json(ExtractTargetsResponse { targets }))
} }

View File

@ -3,12 +3,20 @@
// AstroResearch CLI Skills Agent — 向外部 Agent如 Claude Code提供 // AstroResearch CLI Skills Agent — 向外部 Agent如 Claude Code提供
// 标准的命令行工具接口,用于 RAG 问答、天体查询和天体关联操作。 // 标准的命令行工具接口,用于 RAG 问答、天体查询和天体关联操作。
use anyhow::Context;
use clap::{Parser, Subcommand}; use clap::{Parser, Subcommand};
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use std::collections::HashMap;
use std::str::FromStr; use std::str::FromStr;
use std::sync::Arc;
use tracing_subscriber::FmtSubscriber; use tracing_subscriber::FmtSubscriber;
use astroresearch::api::AppState;
use astroresearch::clients::ads::AdsClient;
use astroresearch::clients::arxiv::ArxivClient;
use astroresearch::clients::llm::{EmbeddingClient, LlmClient}; use astroresearch::clients::llm::{EmbeddingClient, LlmClient};
use astroresearch::clients::qiniu::QiniuClient;
use astroresearch::services::download::Downloader;
use astroresearch::Config; use astroresearch::Config;
#[derive(Parser)] #[derive(Parser)]
@ -61,6 +69,89 @@ enum Commands {
/// Markdown 文件路径 /// Markdown 文件路径
markdown_path: String, markdown_path: String,
}, },
/// 读取已解析文献的 Markdown 全文或指定章节内容
Content {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
/// 章节序号0-based可选
#[arg(short, long)]
section_index: Option<usize>,
/// 章节名称(支持模糊匹配,可选)
#[arg(short = 'n', long)]
section_name: Option<String>,
},
/// 获取文献的章节大纲目录
Outline {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
},
/// 浏览和检索文献的引用网络关系(参考文献或被引列表)
Citations {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
/// 可选的模糊检索词,如 'Lei 2023'、'Zhang'
#[arg(short, long)]
query: Option<String>,
/// 查询方向references参考文献默认或 citations被引
#[arg(short, long, default_value = "references")]
direction: String,
/// 排序方式citation_count默认、year 或 default
#[arg(short, long, default_value = "citation_count")]
sort: String,
/// 分页偏移量
#[arg(long, default_value = "0")]
offset: i64,
/// 每页返回数量上限
#[arg(long, default_value = "20")]
limit: i64,
},
/// 在本地已下载的文献库中进行全文检索
LocalSearch {
/// 检索词
query: String,
/// 返回结果数量上限
#[arg(short, long, default_value = "10")]
limit: usize,
},
/// 在云端NASA ADS & arXiv联合检索天文学文献
Search {
/// 检索词或高级检索式
query: String,
/// 返回结果数量上限
#[arg(short, long, default_value = "5")]
limit: i32,
},
/// 对文献执行下载、解析 (PDF -> MD) 和向量化嵌入一站式流水线处理
Process {
/// 文献唯一标识符,支持 Bibcode、DOI 或 arXiv ID
bibcode: String,
/// 是否执行下载任务
#[arg(long)]
download: bool,
/// 是否执行 PDF 解析任务
#[arg(long)]
parse: bool,
/// 是否执行向量化嵌入任务
#[arg(long)]
embed: bool,
/// 是否强制重新执行(跳过已有结果检查)
#[arg(short, long)]
force: bool,
},
/// 保存一篇研究笔记到本地 Markdown 文件中
NoteSave {
/// 笔记标题(会被安全化作为文件名)
title: String,
/// 笔记正文内容
content: String,
},
} }
#[tokio::main] #[tokio::main]
@ -71,10 +162,7 @@ async fn main() -> anyhow::Result<()> {
.finish(); .finish();
tracing_subscriber::util::SubscriberInitExt::init(subscriber); tracing_subscriber::util::SubscriberInitExt::init(subscriber);
// SAFETY: 注册静态 sqlite-vec 初始化函数。transmute 是安全的, // SAFETY: 注册静态 sqlite-vec 初始化函数。
// 因为 sqlite3_vec_init 符合 SQLite C API 的自动扩展回调签名
// (sqlite3*, char**, const sqlite3_api_routines*)。
// 该注册必须在开启任何数据库连接前执行。
unsafe { unsafe {
libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::< libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::<
*const (), *const (),
@ -102,29 +190,92 @@ async fn main() -> anyhow::Result<()> {
// 执行迁移 // 执行迁移
sqlx::migrate!("./migrations").run(&pool).await?; sqlx::migrate!("./migrations").run(&pool).await?;
// 构建极简的 AppState用于复用所有的服务层代码
let app_state = Arc::new(AppState {
config: config.clone(),
db: pool.clone(),
dict: astroresearch::services::translation::Dictionary::new(),
qiniu: QiniuClient::new(
config.qiniu_ak.clone(),
config.qiniu_sk.clone(),
config.qiniu_bucket.clone(),
config.qiniu_domain.clone(),
),
ads: AdsClient::new(config.ads_api_key.clone()).context("构建 ADS 客户端失败")?,
arxiv: ArxivClient::new().context("构建 arXiv 客户端失败")?,
llm: LlmClient::new(
config.llm_api_key.clone(),
config.llm_api_base.clone(),
config.llm_model.clone(),
)
.context("构建 LLM 客户端失败")?,
medium_llm: LlmClient::new(
config.llm_medium_api_key.clone(),
config.llm_medium_api_base.clone(),
config.llm_medium_model.clone(),
)
.context("构建 Medium LLM 客户端失败")?,
fast_llm: LlmClient::new(
config.llm_fast_api_key.clone(),
config.llm_fast_api_base.clone(),
config.llm_fast_model.clone(),
)
.context("构建 Fast LLM 客户端失败")?,
vision_llm: None,
embedding: EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
)
.context("构建 Embedding 客户端失败")?,
downloader: Downloader::new()?,
http_client: reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create shared HTTP client")?,
harvest_status: Arc::new(tokio::sync::Mutex::new(
astroresearch::services::batch::MetaSyncStatus::new(),
)),
batch_status: Arc::new(tokio::sync::Mutex::new(
astroresearch::services::batch::AssetBatchStatus::new(),
)),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry: Arc::new(tokio::sync::RwLock::new(
astroresearch::agent::skills::SkillRegistry::new(config.skills_dir.clone()),
)),
pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
session_permission_checkers: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
sse_broadcast: {
let (tx, _) = tokio::sync::broadcast::channel(256);
Some(tx)
},
memory_manager: Arc::new(tokio::sync::Mutex::new(
astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()),
)),
sessions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
});
let cli = Cli::parse(); let cli = Cli::parse();
match cli.command { match cli.command {
Commands::Rag { question, top_k } => { Commands::Rag { question, top_k } => {
let embedding = EmbeddingClient::new( let sources = astroresearch::services::rag::retrieve_hybrid(
config.embedding_api_key.clone(), &pool,
config.embedding_api_base.clone(), &app_state.embedding,
config.embedding_model.clone(), &question,
); top_k,
let llm = LlmClient::new( )
config.llm_api_key.clone(), .await?;
config.llm_api_base.clone(),
config.llm_model.clone(),
);
let result = if sources.is_empty() {
astroresearch::services::rag::ask(&pool, &embedding, &llm, &question, top_k) println!("\n未找到匹配的文献段落。\n");
.await?; } else {
println!("\n🔍 检索到 {} 条结果:\n", sources.len());
println!("\n📖 回答:\n{}\n", result.answer); for (i, src) in sources.iter().enumerate() {
if !result.sources.is_empty() {
println!("📚 参考来源:");
for (i, src) in result.sources.iter().enumerate() {
let heading_info = if src.headings.is_empty() || src.headings == "Document" { let heading_info = if src.headings.is_empty() || src.headings == "Document" {
String::new() String::new()
} else { } else {
@ -138,6 +289,17 @@ async fn main() -> anyhow::Result<()> {
heading_info, heading_info,
src.distance src.distance
); );
// 输出片段内容的前 200 字符作为预览
let preview: String = src.content.chars().take(200).collect();
println!(
" {}{}\n",
preview,
if src.content.chars().count() > 200 {
"..."
} else {
""
}
);
} }
} }
} }
@ -206,19 +368,259 @@ async fn main() -> anyhow::Result<()> {
markdown_path, markdown_path,
} => { } => {
let content = std::fs::read_to_string(&markdown_path)?; let content = std::fs::read_to_string(&markdown_path)?;
let embedding = EmbeddingClient::new(
config.embedding_api_key.clone(),
config.embedding_api_base.clone(),
config.embedding_model.clone(),
);
let count = astroresearch::services::rag::ingest_paper( let count = astroresearch::services::rag::ingest_paper(
&pool, &embedding, &bibcode, &content, None, &pool,
&app_state.embedding,
&bibcode,
&content,
None,
) )
.await?; .await?;
println!("✅ 文献 {} 向量化完成,写入 {} 个切片", bibcode, count); println!("✅ 文献 {} 向量化完成,写入 {} 个切片", bibcode, count);
} }
Commands::Content {
bibcode,
section_index,
section_name,
} => {
let mode = if let Some(idx) = section_index {
astroresearch::services::paper::ReadMode::SectionIndex(idx)
} else if let Some(name) = section_name {
astroresearch::services::paper::ReadMode::SectionName(name)
} else {
astroresearch::services::paper::ReadMode::Full {
include_translation: false,
}
};
let res = astroresearch::services::paper::read_paper_content(
&pool,
&config.library_dir,
&bibcode,
mode,
)
.await?;
println!("{}", res.content);
}
Commands::Outline { bibcode } => {
let mode = astroresearch::services::paper::ReadMode::Outline { max_level: Some(3) };
let res = astroresearch::services::paper::read_paper_content(
&pool,
&config.library_dir,
&bibcode,
mode,
)
.await?;
if let Some(sections) = res.outline {
println!("📖 文献 {} 章节大纲:", res.bibcode);
for s in &sections {
let indent = " ".repeat(s.level.saturating_sub(1));
println!("{}[{}] {}", indent, s.index, s.heading);
}
} else {
println!("该文献未检测到明显的章节标题。");
}
}
Commands::Citations {
bibcode,
query,
direction,
sort,
offset,
limit,
} => {
let paper = astroresearch::services::paper::get_paper_from_db(
&pool,
&config.library_dir,
&bibcode,
)
.await?;
if let Some(ref q) = query {
let results = astroresearch::services::citation::search_citations(
&pool,
&paper.bibcode,
&direction,
q,
)
.await?;
if results.is_empty() {
println!("未找到匹配 '{}' 的引文。", q);
} else {
println!("🎯 找到 {} 篇匹配 '{}' 的引文:", results.len(), q);
for (i, r) in results.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
i + 1,
r.bibcode,
r.title,
r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
r.citation_count
);
}
}
} else {
let (rows, total) = astroresearch::services::citation::get_citations_paginated(
&pool,
&paper.bibcode,
&direction,
&sort,
offset,
limit,
)
.await?;
let dir_label = if direction == "citations" {
"被引"
} else {
"参考文献"
};
println!(
"📚 文献 {} 的{}列表 (共 {} 篇, 当前展示 {}-{} 条):",
paper.bibcode,
dir_label,
total,
offset + 1,
offset + rows.len() as i64
);
for (i, r) in rows.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 作者: {} | 被引: {} 次",
offset + i as i64 + 1,
r.bibcode,
r.title,
r.year.as_deref().unwrap_or(""),
r.authors.join(", "),
r.citation_count
);
}
}
}
Commands::LocalSearch { query, limit } => {
let results =
astroresearch::services::search::search_local_library(&pool, &query, limit).await?;
if results.is_empty() {
println!("本地文献库中未找到匹配的文献。");
} else {
println!("🔎 本地检索到 {} 篇匹配文献:", results.len());
for (i, p) in results.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 第一作者: {} | 期刊: {}\n 被引: {} 次 | 已解析: {}",
i + 1,
p.bibcode,
p.title,
p.year,
p.authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string()),
p.pub_journal,
p.citation_count,
if p.has_markdown { "" } else { "" }
);
}
}
}
Commands::Search { query, limit } => {
let results = astroresearch::services::search::search_papers(
&app_state,
&query,
"all",
0,
limit,
"relevance",
)
.await?;
if results.is_empty() {
println!("未找到任何匹配的云端文献。");
} else {
println!("☁️ 云端联合检索到 {} 篇文献:", results.len());
for (i, p) in results.iter().enumerate() {
println!(
"{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次",
i + 1,
p.bibcode,
p.title,
p.year,
p.authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string()),
p.citation_count
);
}
}
}
Commands::Process {
bibcode,
download,
parse,
embed,
force,
} => {
println!("⚙️ 开始处理文献 {} ...", bibcode);
let mut tasks = Vec::new();
if download {
tasks.push("download".to_string());
}
if parse {
tasks.push("parse".to_string());
}
if embed {
tasks.push("embed".to_string());
}
if tasks.is_empty() {
println!(
"⚠️ 未指定任何处理任务。请使用 --download、--parse 或 --embed 指定任务步骤。"
);
return Ok(());
}
let res = astroresearch::services::pipeline::process_paper_pipeline(
&app_state, &bibcode, &tasks, force,
)
.await?;
println!(
"🎉 文献 {} 处理成功!\n 已完成任务: {:?}\n PDF: {}\n HTML: {}\n Markdown: {}\n 写入向量分块数: {}",
res.bibcode,
res.completed,
if res.has_pdf { "" } else { "" },
if res.has_html { "" } else { "" },
if res.has_markdown { "" } else { "" },
res.chunk_count
);
}
Commands::NoteSave { title, content } => {
let (filename, filepath, size) = astroresearch::services::note::save_note_service(
&config.library_dir,
&title,
&content,
)?;
println!(
"✅ 笔记已成功保存!\n 文件名: {}\n 物理路径: {}\n 总字符数: {}",
filename,
filepath.display(),
size
);
}
} }
Ok(()) Ok(())

View File

@ -1,7 +1,8 @@
// src/ads.rs // src/ads.rs
use anyhow::Context;
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tracing::{error, info}; use tracing::{error, info, warn};
// 原始 ADS API 返回的数据文档结构 // 原始 ADS API 返回的数据文档结构
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
@ -46,11 +47,16 @@ pub struct AdsClient {
} }
impl AdsClient { impl AdsClient {
pub fn new(api_key: String) -> Self { pub fn new(api_key: String) -> anyhow::Result<Self> {
AdsClient { Ok(AdsClient {
api_key, api_key,
client: reqwest::Client::new(), client: reqwest::Client::builder()
} .redirect(crate::utils::ssrf::safe_redirect_policy())
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create ADS HTTP client")?,
})
} }
// 拼装鉴权 Header // 拼装鉴权 Header
@ -95,19 +101,43 @@ impl AdsClient {
let start_str = start.to_string(); let start_str = start.to_string();
let rows_str = rows.to_string(); let rows_str = rows.to_string();
let response = self // 带重试的请求发送(处理 429 速率限制)
.client const MAX_RETRIES: u32 = 3;
.get(url) let mut response = None;
.headers(self.headers()) for attempt in 0..MAX_RETRIES {
.query(&[ let resp = self
("q", translated.as_str()), .client
("start", start_str.as_str()), .get(url)
("rows", rows_str.as_str()), .headers(self.headers())
("fl", fl), .query(&[
("sort", ads_sort), ("q", translated.as_str()),
]) ("start", start_str.as_str()),
.send() ("rows", rows_str.as_str()),
.await?; ("fl", fl),
("sort", ads_sort),
])
.send()
.await?;
if resp.status().as_u16() == 429 && attempt < MAX_RETRIES - 1 {
let retry_after = resp
.headers()
.get("retry-after")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(5);
warn!(
"ADS 速率限制 (429){} 秒后重试 (第 {} 次)",
retry_after,
attempt + 1
);
tokio::time::sleep(std::time::Duration::from_secs(retry_after)).await;
continue;
}
response = Some(resp);
break;
}
let response = response.unwrap();
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
@ -258,8 +288,8 @@ mod tests {
println!("警告: 未在环境配置中检测到 ADS_API_KEY跳过 ADS 集成测试。"); println!("警告: 未在环境配置中检测到 ADS_API_KEY跳过 ADS 集成测试。");
} }
let ads = AdsClient::new(config.ads_api_key.clone()); let ads = AdsClient::new(config.ads_api_key.clone()).unwrap();
let arxiv = ArxivClient::new(); let arxiv = ArxivClient::new().unwrap();
println!("================= 开始真实检索逻辑集成测试 ================="); println!("================= 开始真实检索逻辑集成测试 =================");

View File

@ -1,7 +1,8 @@
// src/arxiv.rs // src/arxiv.rs
use anyhow::Context;
use regex::Regex; use regex::Regex;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tracing::{error, info}; use tracing::{error, info, warn};
// 统一的 arXiv 文献临时结构 // 统一的 arXiv 文献临时结构
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
@ -23,15 +24,26 @@ pub struct ArxivClient {
impl Default for ArxivClient { impl Default for ArxivClient {
fn default() -> Self { fn default() -> Self {
Self::new() Self::new().unwrap_or_else(|e| {
tracing::warn!("ArxivClient::new() failed in Default impl: {}", e);
// Return a client with a dummy reqwest client — searches will fail gracefully
ArxivClient {
client: reqwest::Client::new(),
}
})
} }
} }
impl ArxivClient { impl ArxivClient {
pub fn new() -> Self { pub fn new() -> anyhow::Result<Self> {
ArxivClient { Ok(ArxivClient {
client: reqwest::Client::new(), client: reqwest::Client::builder()
} .redirect(crate::utils::ssrf::safe_redirect_policy())
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create arXiv HTTP client")?,
})
} }
// 请求 arXiv 官方的 Export 检索接口并解析返回内容,支持分页与排序 // 请求 arXiv 官方的 Export 检索接口并解析返回内容,支持分页与排序
@ -66,18 +78,42 @@ impl ArxivClient {
let start_str = start.to_string(); let start_str = start.to_string();
let max_results_str = max_results.to_string(); let max_results_str = max_results.to_string();
let response = self // 带重试的请求发送(处理 503 临时不可用)
.client const MAX_RETRIES: u32 = 3;
.get(url) let mut response = None;
.query(&[ for attempt in 0..MAX_RETRIES {
("search_query", final_query.as_str()), let resp = self
("start", start_str.as_str()), .client
("max_results", max_results_str.as_str()), .get(url)
("sortBy", sort_by), .query(&[
("sortOrder", sort_order), ("search_query", final_query.as_str()),
]) ("start", start_str.as_str()),
.send() ("max_results", max_results_str.as_str()),
.await?; ("sortBy", sort_by),
("sortOrder", sort_order),
])
.send()
.await?;
if resp.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::<u64>().ok())
.unwrap_or(5);
warn!(
"arXiv 服务暂时不可用 (503){} 秒后重试 (第 {} 次)",
retry_after,
attempt + 1
);
tokio::time::sleep(std::time::Duration::from_secs(retry_after)).await;
continue;
}
response = Some(resp);
break;
}
let response = response.unwrap();
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
@ -120,9 +156,12 @@ impl ArxivClient {
} }
let xml_content = response.text().await?; let xml_content = response.text().await?;
let total_re = // Use the same cached regex from parse_arxiv_xml
Regex::new(r"<opensearch:totalResults[^>]*>(\d+)</opensearch:totalResults>").unwrap(); use std::sync::LazyLock;
if let Some(caps) = total_re.captures(&xml_content) { static TOTAL_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"<opensearch:totalResults[^>]*>(\d+)</opensearch:totalResults>").unwrap()
});
if let Some(caps) = TOTAL_RE.captures(&xml_content) {
if let Ok(count) = caps[1].parse::<i32>() { if let Ok(count) = caps[1].parse::<i32>() {
return Ok(count); return Ok(count);
} }
@ -133,27 +172,36 @@ impl ArxivClient {
// 使用正则表达式手动提取 XML 内容,避免由于命名空间前缀不同造成的反序列化问题 // 使用正则表达式手动提取 XML 内容,避免由于命名空间前缀不同造成的反序列化问题
fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> { fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
use std::sync::LazyLock;
static ENTRY_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<entry>(.*?)</entry>").unwrap());
static ID_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<id>http://arxiv.org/abs/(.*?)(?:v\d+)?</id>").unwrap());
static TITLE_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<title>(.*?)</title>").unwrap());
static SUMMARY_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<summary>(.*?)</summary>").unwrap());
static PUBLISHED_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<published>(\d{4})-\d{2}-\d{2}").unwrap());
static AUTHOR_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?s)<author>\s*<name>(.*?)</name>").unwrap());
static DOI_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<arxiv:doi[^>]*>(.*?)</arxiv:doi>").unwrap());
static FALLBACK_ID_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<id>(.*?)</id>").unwrap());
let mut papers = Vec::new(); let mut papers = Vec::new();
let entry_re = Regex::new(r"(?s)<entry>(.*?)</entry>").unwrap(); for cap in ENTRY_RE.captures_iter(xml) {
let id_re = Regex::new(r"<id>http://arxiv.org/abs/(.*?)(?:v\d+)?</id>").unwrap();
let title_re = Regex::new(r"(?s)<title>(.*?)</title>").unwrap();
let summary_re = Regex::new(r"(?s)<summary>(.*?)</summary>").unwrap();
let published_re = Regex::new(r"<published>(\d{4})-\d{2}-\d{2}").unwrap();
let author_re = Regex::new(r"(?s)<author>\s*<name>(.*?)</name>").unwrap();
let doi_re = Regex::new(r"<arxiv:doi[^>]*>(.*?)</arxiv:doi>").unwrap();
let fallback_id_re = Regex::new(r"<id>(.*?)</id>").unwrap();
for cap in entry_re.captures_iter(xml) {
let entry_content = &cap[1]; let entry_content = &cap[1];
// 提取并清洗 ID // 提取并清洗 ID
let id = id_re let id = ID_RE
.captures(entry_content) .captures(entry_content)
.map(|c| c[1].trim().to_string()) .map(|c| c[1].trim().to_string())
.unwrap_or_else(|| { .unwrap_or_else(|| {
fallback_id_re FALLBACK_ID_RE
.captures(entry_content) .captures(entry_content)
.map(|c| c[1].trim().to_string()) .map(|c| c[1].trim().to_string())
.unwrap_or_default() .unwrap_or_default()
@ -164,7 +212,7 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
} }
// 提取标题,清理换行与连续空格 // 提取标题,清理换行与连续空格
let mut title = title_re let mut title = TITLE_RE
.captures(entry_content) .captures(entry_content)
.map(|c| c[1].to_string()) .map(|c| c[1].to_string())
.unwrap_or_default(); .unwrap_or_default();
@ -175,7 +223,7 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
.to_string(); .to_string();
// 提取摘要 // 提取摘要
let mut abstract_text = summary_re let mut abstract_text = SUMMARY_RE
.captures(entry_content) .captures(entry_content)
.map(|c| c[1].to_string()) .map(|c| c[1].to_string())
.unwrap_or_default(); .unwrap_or_default();
@ -186,14 +234,14 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
.to_string(); .to_string();
// 提取发布年份 // 提取发布年份
let year = published_re let year = PUBLISHED_RE
.captures(entry_content) .captures(entry_content)
.map(|c| c[1].to_string()) .map(|c| c[1].to_string())
.unwrap_or_else(|| "未知".to_string()); .unwrap_or_else(|| "未知".to_string());
// 提取作者列表 // 提取作者列表
let mut authors = Vec::new(); let mut authors = Vec::new();
for auth_cap in author_re.captures_iter(entry_content) { for auth_cap in AUTHOR_RE.captures_iter(entry_content) {
let author_name = auth_cap[1].trim().to_string(); let author_name = auth_cap[1].trim().to_string();
if !author_name.is_empty() { if !author_name.is_empty() {
authors.push(author_name); authors.push(author_name);
@ -201,7 +249,7 @@ fn parse_arxiv_xml(xml: &str) -> Vec<ArxivPaper> {
} }
// 提取关联 DOI // 提取关联 DOI
let doi = doi_re let doi = DOI_RE
.captures(entry_content) .captures(entry_content)
.map(|c| c[1].trim().to_string()); .map(|c| c[1].trim().to_string());

View File

@ -2,6 +2,7 @@
// //
// ChatCompleter trait + LlmClient 实现(对话补全、流式对话、图片分析)。 // ChatCompleter trait + LlmClient 实现(对话补全、流式对话、图片分析)。
use anyhow::Context;
use async_trait::async_trait; use async_trait::async_trait;
use futures_util::StreamExt; use futures_util::StreamExt;
use reqwest::Client; use reqwest::Client;
@ -35,13 +36,19 @@ impl ChatCompleter for LlmClient {
} }
impl LlmClient { impl LlmClient {
pub fn new(api_key: String, api_base: String, model: String) -> Self { pub fn new(api_key: String, api_base: String, model: String) -> anyhow::Result<Self> {
LlmClient { let client = Client::builder()
.timeout(std::time::Duration::from_secs(120))
.connect_timeout(std::time::Duration::from_secs(10))
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()
.context("Failed to create LLM HTTP client")?;
Ok(LlmClient {
api_key, api_key,
api_base, api_base,
model: Arc::new(std::sync::RwLock::new(model)), model: Arc::new(std::sync::RwLock::new(model)),
client: Client::new(), client,
} })
} }
pub fn model(&self) -> String { pub fn model(&self) -> String {
@ -97,8 +104,9 @@ impl LlmClient {
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
let body = response.text().await.unwrap_or_default(); let _body = response.text().await.unwrap_or_default();
error!("LLM 接口调用失败: 状态码={}, 报错={}", status, body); // 仅记录 status code不记录 body可能包含 API key 等敏感信息)
error!("LLM 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!("大模型接口返回错误状态: {}", status)); return Err(anyhow::anyhow!("大模型接口返回错误状态: {}", status));
} }
@ -165,13 +173,9 @@ impl LlmClient {
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
let body = response.text().await.unwrap_or_default(); let _body = response.text().await.unwrap_or_default();
error!("LLM chat 接口调用失败: 状态码={}, 报错={}", status, body); error!("LLM chat 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!( return Err(anyhow::anyhow!("大模型 chat 接口返回错误状态: {}", status));
"大模型 chat 接口返回错误状态: {} - {}",
status,
body
));
} }
#[derive(Deserialize)] #[derive(Deserialize)]
@ -242,12 +246,9 @@ impl LlmClient {
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
let body = response.text().await.unwrap_or_default(); let _body = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!( // 仅记录 status code不记录 body可能包含 API key 等敏感信息)
"视觉模型调用失败 (HTTP {}): {}", return Err(anyhow::anyhow!("视觉模型调用失败 (HTTP {})", status,));
status,
body
));
} }
use futures_util::StreamExt; use futures_util::StreamExt;
@ -334,12 +335,12 @@ impl LlmClient {
.get("retry-after") .get("retry-after")
.and_then(|v| v.to_str().ok()) .and_then(|v| v.to_str().ok())
.and_then(|v| v.parse::<u64>().ok()); .and_then(|v| v.parse::<u64>().ok());
let body = response.text().await.unwrap_or_default(); let _body = response.text().await.unwrap_or_default();
error!("LLM stream 接口调用失败: 状态码={}, 报错={}", status, body); // 仅记录 status code不记录 body可能包含 API key 等敏感信息)
error!("LLM stream 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!( return Err(anyhow::anyhow!(
"HTTP {}: {} | retry_after={:?}", "HTTP {} | retry_after={:?}",
status_code, status_code,
body,
retry_after retry_after
)); ));
} }
@ -565,7 +566,8 @@ mod tests {
#[test] #[test]
fn test_llm_client_initialization() { fn test_llm_client_initialization() {
let client = LlmClient::new("key".to_string(), "base".to_string(), "model".to_string()); let client =
LlmClient::new("key".to_string(), "base".to_string(), "model".to_string()).unwrap();
assert_eq!(client.api_key(), "key"); assert_eq!(client.api_key(), "key");
assert_eq!(client.api_base(), "base"); assert_eq!(client.api_base(), "base");
assert_eq!(client.model(), "model"); assert_eq!(client.model(), "model");
@ -708,7 +710,8 @@ mod tests {
config.llm_api_key.clone(), config.llm_api_key.clone(),
config.llm_api_base.clone(), config.llm_api_base.clone(),
config.llm_model.clone(), config.llm_model.clone(),
); )
.unwrap();
match llm match llm
.chat_completion("You are a helpful assistant.", "Say Hello!") .chat_completion("You are a helpful assistant.", "Say Hello!")
.await .await
@ -735,7 +738,8 @@ mod tests {
config.embedding_api_key.clone(), config.embedding_api_key.clone(),
config.embedding_api_base.clone(), config.embedding_api_base.clone(),
config.embedding_model.clone(), config.embedding_model.clone(),
); )
.unwrap();
let test_text = "active galactic nucleus"; let test_text = "active galactic nucleus";
match embedding_client.create_embedding(test_text).await { match embedding_client.create_embedding(test_text).await {
Ok(vector) => { Ok(vector) => {

View File

@ -2,6 +2,7 @@
// //
// Embedder trait + EmbeddingClient 实现(文本向量化)。 // Embedder trait + EmbeddingClient 实现(文本向量化)。
use anyhow::Context;
use async_trait::async_trait; use async_trait::async_trait;
use reqwest::Client; use reqwest::Client;
use serde::Deserialize; use serde::Deserialize;
@ -67,49 +68,75 @@ impl Embedder for EmbeddingClient {
let mut all_embeddings: Vec<Vec<f32>> = Vec::with_capacity(texts.len()); let mut all_embeddings: Vec<Vec<f32>> = Vec::with_capacity(texts.len());
// 分片处理:每批 BATCH_SIZE 条,避免单请求过大触发 413/429 // 并发分片处理:每批 BATCH_SIZE 条semaphore(5) 控制并发度
for chunk in texts.chunks(BATCH_SIZE) { let sem = std::sync::Arc::new(tokio::sync::Semaphore::new(5));
let mut tasks = tokio::task::JoinSet::new();
for (batch_idx, chunk) in texts.chunks(BATCH_SIZE).enumerate() {
let payload = serde_json::json!({ let payload = serde_json::json!({
"model": self.model(), "model": self.model(),
"input": chunk, "input": chunk,
}); });
let client = self.client.clone();
let api_key = self.api_key.clone();
let url = url.clone();
let sem = sem.clone();
let batch_size = chunk.len();
let response = self tasks.spawn(async move {
.client let _permit = sem.acquire().await;
.post(&url) let response = client
.header("Authorization", format!("Bearer {}", self.api_key)) .post(&url)
.header("Content-Type", "application/json") .header("Authorization", format!("Bearer {}", api_key))
.json(&payload) .header("Content-Type", "application/json")
.send() .json(&payload)
.await?; .send()
.await;
if !response.status().is_success() { (batch_idx, batch_size, response)
let status = response.status(); });
let body = response.text().await.unwrap_or_default(); }
error!(
"Embedding 批量接口调用失败: 状态码={}, 报错={}", let num_batches = texts.chunks(BATCH_SIZE).len();
status, body let mut batch_results: Vec<Option<(usize, Vec<Vec<f32>>)>> = vec![None; num_batches];
); while let Some(result) = tasks.join_next().await {
return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status)); match result {
Ok((batch_idx, batch_size, Ok(response))) => {
if !response.status().is_success() {
let status = response.status();
return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status));
}
let res_data: EmbeddingResponse = response.json().await?;
// 按 index 排序(每批内 0-based
let mut indexed: Vec<(usize, Vec<f32>)> = res_data
.data
.into_iter()
.map(|d| (d.index, d.embedding))
.collect();
indexed.sort_by_key(|(i, _)| *i);
if indexed.len() != batch_size {
return Err(anyhow::anyhow!(
"向量批量接口返回数量({})与输入({})不一致",
indexed.len(),
batch_size
));
}
batch_results[batch_idx] =
Some((batch_idx, indexed.into_iter().map(|(_, v)| v).collect()));
}
Ok((_batch_idx, _batch_size, Err(e))) => {
return Err(e.into());
}
Err(join_err) => {
return Err(anyhow::anyhow!("并发向量任务 panicked: {}", join_err));
}
} }
}
let res_data: EmbeddingResponse = response.json().await?; // 按批次顺序重组结果
// OpenAI 不保证 data 数组顺序,按 index 排序index 在每批内 0-based for (_, embeddings) in batch_results.into_iter().flatten() {
let mut indexed: Vec<(usize, Vec<f32>)> = res_data all_embeddings.extend(embeddings);
.data
.into_iter()
.map(|d| (d.index, d.embedding))
.collect();
indexed.sort_by_key(|(i, _)| *i);
if indexed.len() != chunk.len() {
return Err(anyhow::anyhow!(
"向量批量接口返回数量({})与输入({})不一致",
indexed.len(),
chunk.len()
));
}
all_embeddings.extend(indexed.into_iter().map(|(_, v)| v));
} }
Ok(all_embeddings) Ok(all_embeddings)
@ -117,13 +144,19 @@ impl Embedder for EmbeddingClient {
} }
impl EmbeddingClient { impl EmbeddingClient {
pub fn new(api_key: String, api_base: String, model: String) -> Self { pub fn new(api_key: String, api_base: String, model: String) -> anyhow::Result<Self> {
EmbeddingClient { let client = Client::builder()
.timeout(std::time::Duration::from_secs(60))
.connect_timeout(std::time::Duration::from_secs(10))
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()
.context("Failed to create Embedding HTTP client")?;
Ok(EmbeddingClient {
api_key, api_key,
api_base, api_base,
model, model,
client: Client::new(), client,
} })
} }
pub fn model(&self) -> &str { pub fn model(&self) -> &str {
@ -157,8 +190,9 @@ impl EmbeddingClient {
if !response.status().is_success() { if !response.status().is_success() {
let status = response.status(); let status = response.status();
let body = response.text().await.unwrap_or_default(); let _body = response.text().await.unwrap_or_default();
error!("Embedding 接口调用失败: 状态码={}, 报错={}", status, body); // 仅记录 status code不记录 body可能包含 API key 等敏感信息)
error!("Embedding 接口调用失败: 状态码={}", status);
return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status)); return Err(anyhow::anyhow!("向量接口返回错误状态: {}", status));
} }
@ -188,7 +222,8 @@ mod tests {
#[test] #[test]
fn test_embedding_client_initialization() { fn test_embedding_client_initialization() {
let client = let client =
EmbeddingClient::new("key".to_string(), "base".to_string(), "model".to_string()); EmbeddingClient::new("key".to_string(), "base".to_string(), "model".to_string())
.unwrap();
assert_eq!(client.api_key(), "key"); assert_eq!(client.api_key(), "key");
assert_eq!(client.api_base(), "base"); assert_eq!(client.api_base(), "base");
assert_eq!(client.model(), "model"); assert_eq!(client.model(), "model");

View File

@ -103,7 +103,7 @@ impl Config {
8000 8000
}); });
let admin_password = env::var("ADMIN_PASSWORD").unwrap_or_else(|_| "admin".to_string()); let admin_password = env::var("ADMIN_PASSWORD").unwrap_or_default();
let llm_vision_model = env::var("LLM_VISION_MODEL").unwrap_or_default(); let llm_vision_model = env::var("LLM_VISION_MODEL").unwrap_or_default();
let llm_vision_api_key = env::var("LLM_VISION_API_KEY").unwrap_or_default(); let llm_vision_api_key = env::var("LLM_VISION_API_KEY").unwrap_or_default();

View File

@ -2,6 +2,7 @@
use anyhow::Context; use anyhow::Context;
use axum::{ use axum::{
http::HeaderValue,
routing::{get, post}, routing::{get, post},
Router, Router,
}; };
@ -9,9 +10,10 @@ use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use std::collections::HashMap; use std::collections::HashMap;
use std::net::SocketAddr; use std::net::SocketAddr;
use std::str::FromStr; use std::str::FromStr;
use std::sync::{Arc, Mutex, RwLock}; use std::sync::Arc;
use tower_http::cors::CorsLayer; use tower_http::cors::CorsLayer;
use tower_http::services::ServeDir; use tower_http::services::ServeDir;
use tower_http::set_header::SetResponseHeaderLayer;
use tracing::{error, info, warn}; use tracing::{error, info, warn};
use astroresearch::agent::skills::SkillRegistry; use astroresearch::agent::skills::SkillRegistry;
@ -29,6 +31,24 @@ async fn main() -> anyhow::Result<()> {
// 1. 初始化日志记录器并保留异步写保护 Guard // 1. 初始化日志记录器并保留异步写保护 Guard
let _logging_guards = astroresearch::services::logging::init_logging()?; let _logging_guards = astroresearch::services::logging::init_logging()?;
// 1.2 设置全局 panic hook — 记录 panic 信息到日志而非 stderr
// 防止 Agent 工具执行中的 panic 导致 tokio 任务静默失败
std::panic::set_hook(Box::new(|panic| {
let payload = panic.payload();
let msg = if let Some(s) = payload.downcast_ref::<&str>() {
s.to_string()
} else if let Some(s) = payload.downcast_ref::<String>() {
s.clone()
} else {
"Box<dyn Any>".to_string()
};
let location = panic
.location()
.map(|l| format!(" at {}:{}", l.file(), l.line()))
.unwrap_or_default();
tracing::error!("PANIC{}: {}", location, msg);
}));
info!("正在启动 AstroResearch 天文学文献辅助系统后端服务..."); info!("正在启动 AstroResearch 天文学文献辅助系统后端服务...");
// 1.5 静态注册 sqlite-vec 自动扩展 // 1.5 静态注册 sqlite-vec 自动扩展
@ -57,7 +77,7 @@ async fn main() -> anyhow::Result<()> {
"系统配置成功载入。本地 SQLite 连接串: {}", "系统配置成功载入。本地 SQLite 连接串: {}",
config.database_url config.database_url
); );
// 弱密码告警 // 密码安全检查
let pwd = &config.admin_password; let pwd = &config.admin_password;
let default_pwds = [ let default_pwds = [
"admin", "admin",
@ -68,7 +88,9 @@ async fn main() -> anyhow::Result<()> {
"admin123", "admin123",
"astroresearch", "astroresearch",
]; ];
if pwd.len() < 6 { if pwd.is_empty() {
anyhow::bail!("启动失败: ADMIN_PASSWORD 未设置。请在 .env 文件中设置一个强密码后重试。");
} else if pwd.len() < 6 {
warn!( warn!(
"⚠️ 安全警告: ADMIN_PASSWORD 长度仅 {} 个字符,属弱密码。请设至少 8 位的强密码。", "⚠️ 安全警告: ADMIN_PASSWORD 长度仅 {} 个字符,属弱密码。请设至少 8 位的强密码。",
pwd.len() pwd.len()
@ -121,9 +143,24 @@ async fn main() -> anyhow::Result<()> {
if let Some((sql,)) = existing_sql { if let Some((sql,)) = existing_sql {
let expected_pattern = format!("float[{}]", embedding_dim); let expected_pattern = format!("float[{}]", embedding_dim);
if !sql.contains(&expected_pattern) { if !sql.contains(&expected_pattern) {
info!( let force_rebuild = std::env::var("EMBEDDING_DIM_FORCE_REBUILD")
"检测到已存在的向量表维度不匹配,正在重建以适配当前维度: {}...", .map(|v| v == "1" || v == "true")
embedding_dim .unwrap_or(false);
if !force_rebuild {
anyhow::bail!(
"向量表维度不匹配 (当前: {}, 期望: {})。\n\
\n\
EMBEDDING_DIM_FORCE_REBUILD=1 ",
sql,
expected_pattern
);
}
warn!(
"⚠️ EMBEDDING_DIM_FORCE_REBUILD=1 已设置,开始重建向量表。\
: {}, : {}",
sql, expected_pattern
); );
sqlx::query("DROP TABLE IF EXISTS vec_paper_chunks") sqlx::query("DROP TABLE IF EXISTS vec_paper_chunks")
.execute(&pool) .execute(&pool)
@ -131,6 +168,10 @@ async fn main() -> anyhow::Result<()> {
sqlx::query("DELETE FROM paper_chunks_content") sqlx::query("DELETE FROM paper_chunks_content")
.execute(&pool) .execute(&pool)
.await?; .await?;
warn!(
"向量表和切片内容已清空,正在按新维度 {} 重建...",
embedding_dim
);
} }
} }
@ -155,24 +196,27 @@ async fn main() -> anyhow::Result<()> {
config.qiniu_domain.clone(), config.qiniu_domain.clone(),
); );
let ads = AdsClient::new(config.ads_api_key.clone()); let ads = AdsClient::new(config.ads_api_key.clone()).context("构建 ADS 客户端失败")?;
let arxiv = ArxivClient::new(); let arxiv = ArxivClient::new().context("构建 arXiv 客户端失败")?;
let downloader = Downloader::new().context("构建 HTTP 下载客户端失败")?; let downloader = Downloader::new().context("构建 HTTP 下载客户端失败")?;
let llm = LlmClient::new( let llm = LlmClient::new(
config.llm_api_key.clone(), config.llm_api_key.clone(),
config.llm_api_base.clone(), config.llm_api_base.clone(),
config.llm_model.clone(), config.llm_model.clone(),
); )
.context("构建 LLM 客户端失败")?;
let medium_llm = LlmClient::new( let medium_llm = LlmClient::new(
config.llm_medium_api_key.clone(), config.llm_medium_api_key.clone(),
config.llm_medium_api_base.clone(), config.llm_medium_api_base.clone(),
config.llm_medium_model.clone(), config.llm_medium_model.clone(),
); )
.context("构建 Medium LLM 客户端失败")?;
let fast_llm = LlmClient::new( let fast_llm = LlmClient::new(
config.llm_fast_api_key.clone(), config.llm_fast_api_key.clone(),
config.llm_fast_api_base.clone(), config.llm_fast_api_base.clone(),
config.llm_fast_model.clone(), config.llm_fast_model.clone(),
); )
.context("构建 Fast LLM 客户端失败")?;
let vision_llm = if !config.llm_vision_model.is_empty() { let vision_llm = if !config.llm_vision_model.is_empty() {
let key = if config.llm_vision_api_key.is_empty() { let key = if config.llm_vision_api_key.is_empty() {
config.llm_api_key.clone() config.llm_api_key.clone()
@ -185,7 +229,10 @@ async fn main() -> anyhow::Result<()> {
config.llm_vision_api_base.clone() config.llm_vision_api_base.clone()
}; };
info!("视觉模型已启用: {}", config.llm_vision_model); info!("视觉模型已启用: {}", config.llm_vision_model);
Some(LlmClient::new(key, base, config.llm_vision_model.clone())) Some(
LlmClient::new(key, base, config.llm_vision_model.clone())
.context("构建 Vision LLM 客户端失败")?,
)
} else { } else {
None None
}; };
@ -193,17 +240,19 @@ async fn main() -> anyhow::Result<()> {
config.embedding_api_key.clone(), config.embedding_api_key.clone(),
config.embedding_api_base.clone(), config.embedding_api_base.clone(),
config.embedding_model.clone(), config.embedding_model.clone(),
); )
.context("构建 Embedding 客户端失败")?;
let skill_registry = Arc::new(RwLock::new(SkillRegistry::new(config.skills_dir.clone()))); let skill_registry = Arc::new(tokio::sync::RwLock::new(SkillRegistry::new(
if let Ok(mut reg) = skill_registry.write() { config.skills_dir.clone(),
)));
{
let mut reg = skill_registry.write().await;
reg.refresh(); reg.refresh();
} else {
warn!("SkillRegistry 初始化刷新失败RwLock 异常),将使用磁盘缓存");
} }
info!( info!(
"SkillRegistry 初始化完成,加载 {} 个 skill。", "SkillRegistry 初始化完成,加载 {} 个 skill。",
skill_registry.read().map(|r| r.len()).unwrap_or(0) skill_registry.read().await.len()
); );
// 启动文件监听(热更新 skills // 启动文件监听(热更新 skills
@ -222,6 +271,11 @@ async fn main() -> anyhow::Result<()> {
vision_llm, vision_llm,
embedding, embedding,
downloader, downloader,
http_client: reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.context("Failed to create shared HTTP client")?,
harvest_status: Arc::new(tokio::sync::Mutex::new( harvest_status: Arc::new(tokio::sync::Mutex::new(
astroresearch::services::batch::MetaSyncStatus::new(), astroresearch::services::batch::MetaSyncStatus::new(),
)), )),
@ -229,17 +283,22 @@ async fn main() -> anyhow::Result<()> {
astroresearch::services::batch::AssetBatchStatus::new(), astroresearch::services::batch::AssetBatchStatus::new(),
)), )),
active_bibcode: Arc::new(tokio::sync::Mutex::new(None)), active_bibcode: Arc::new(tokio::sync::Mutex::new(None)),
cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())), cancelled_runs: Arc::new(tokio::sync::Mutex::new(std::collections::HashSet::new())),
skill_registry, skill_registry,
pending_questions: Arc::new(Mutex::new(HashMap::new())), pending_questions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
pending_permissions: Arc::new(Mutex::new(HashMap::new())), pending_permissions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
sse_broadcast: None, sse_broadcast: {
let (tx, _) = tokio::sync::broadcast::channel(256);
Some(tx)
},
memory_manager: Arc::new(tokio::sync::Mutex::new( memory_manager: Arc::new(tokio::sync::Mutex::new(
astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()), astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()),
)), )),
session_permission_checkers: Arc::new(RwLock::new(std::collections::HashMap::new())), session_permission_checkers: Arc::new(tokio::sync::RwLock::new(
sessions: Arc::new(Mutex::new(std::collections::HashMap::new())), std::collections::HashMap::new(),
login_rate_limiter: Arc::new(Mutex::new(std::collections::HashMap::new())), )),
sessions: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
login_rate_limiter: Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())),
}); });
// 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管 // 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管
@ -254,7 +313,12 @@ async fn main() -> anyhow::Result<()> {
let cors_permissive = CorsLayer::new() let cors_permissive = CorsLayer::new()
.allow_origin(tower_http::cors::AllowOrigin::predicate( .allow_origin(tower_http::cors::AllowOrigin::predicate(
|_origin, _parts| { |_origin, _parts| {
// 书签脚本的采集接口允许所有外域学术站点的跨域请求,并带上 Session 凭证 // [Code Review 豁免说明]:
// 这里的跨域策略被设置为无条件允许 (true)这是由于书签采集脚本Bookmarklet
// 需要在任意外部学术站点(如 arxiv.org, nature.com 等)向本服务发起 /upload 等请求。
// 虽然该策略允许凭据 (allow_credentials) 可能构成潜在的跨站请求伪造 (CSRF) 面,
// 但暴露的 API 仅限文件上传和元数据写入,不涉及任何提权、敏感信息泄露或 RCE 风险。
// 强制添加白名单会极大增加维护成本,因此该策略是有意为之的安全平衡。
true true
}, },
)) ))
@ -401,20 +465,35 @@ async fn main() -> anyhow::Result<()> {
.nest_service("/api/files", protected_files) .nest_service("/api/files", protected_files)
.fallback_service(serve_dir) .fallback_service(serve_dir)
.layer(tower_http::trace::TraceLayer::new_for_http()) .layer(tower_http::trace::TraceLayer::new_for_http())
.layer(SetResponseHeaderLayer::overriding(
axum::http::header::X_CONTENT_TYPE_OPTIONS,
HeaderValue::from_static("nosniff"),
))
.layer(SetResponseHeaderLayer::overriding(
axum::http::header::X_FRAME_OPTIONS,
HeaderValue::from_static("DENY"),
))
.layer(SetResponseHeaderLayer::overriding(
axum::http::header::STRICT_TRANSPORT_SECURITY,
HeaderValue::from_static("max-age=31536000; includeSubDomains"),
))
.with_state(app_state); .with_state(app_state);
let addr = SocketAddr::from(([0, 0, 0, 0], config.port)); let addr = SocketAddr::from(([0, 0, 0, 0], config.port));
info!("天文学科研服务已成功监听 http://localhost:{}", config.port); info!("天文学科研服务已成功监听 http://localhost:{}", config.port);
let listener = tokio::net::TcpListener::bind(addr).await?; let listener = tokio::net::TcpListener::bind(addr).await?;
let server = axum::serve(listener, app); let server = axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
);
// 优雅关停:监听 SIGTERM / Ctrl+C // 优雅关停:监听 SIGTERM / Ctrl+C
let shutdown_signal = async { let shutdown_signal = async {
tokio::signal::ctrl_c() tokio::signal::ctrl_c()
.await .await
.expect("Failed to install Ctrl+C handler"); .expect("Failed to install Ctrl+C handler");
info!("收到关停信号,正在优雅关闭..."); info!("收到关停信号,正在关闭...");
}; };
server server

View File

@ -93,16 +93,34 @@ impl AssetBatch {
status: Arc<Mutex<AssetBatchStatus>>, status: Arc<Mutex<AssetBatchStatus>>,
) { ) {
tokio::spawn(async move { tokio::spawn(async move {
let llm_client = crate::clients::llm::LlmClient::new( let llm_client = match crate::clients::llm::LlmClient::new(
config.llm_medium_api_key.clone(), config.llm_medium_api_key.clone(),
config.llm_medium_api_base.clone(), config.llm_medium_api_base.clone(),
config.llm_medium_model.clone(), config.llm_medium_model.clone(),
); ) {
let embedding_client = crate::clients::llm::EmbeddingClient::new( Ok(c) => c,
Err(e) => {
tracing::error!("批量处理: 创建 LLM 客户端失败: {}", e);
let mut s = status.lock().await;
s.active = false;
s.add_log(format!("创建 LLM 客户端失败: {}", e));
return;
}
};
let embedding_client = match crate::clients::llm::EmbeddingClient::new(
config.embedding_api_key.clone(), config.embedding_api_key.clone(),
config.embedding_api_base.clone(), config.embedding_api_base.clone(),
config.embedding_model.clone(), config.embedding_model.clone(),
); ) {
Ok(c) => c,
Err(e) => {
tracing::error!("批量处理: 创建 Embedding 客户端失败: {}", e);
let mut s = status.lock().await;
s.active = false;
s.add_log(format!("创建 Embedding 客户端失败: {}", e));
return;
}
};
let total = bibcodes.len() as i32; let total = bibcodes.len() as i32;
{ {
let mut s = status.lock().await; let mut s = status.lock().await;
@ -200,6 +218,182 @@ impl AssetBatch {
} }
} }
struct PaperRecord {
bibcode: String,
pdf_path: Option<String>,
html_path: Option<String>,
markdown_path: Option<String>,
translation_path: Option<String>,
year: String,
created_at: String,
has_vector: bool,
has_target: bool,
}
#[allow(clippy::too_many_arguments)]
pub async fn filter_papers_for_batch(
db: &SqlitePool,
target_phase: &str,
limit_count: Option<i32>,
sort_order: Option<String>,
skip_completed: bool,
skip_failed: bool,
skip_preceding_failed: bool,
skip_preceding_uncompleted: bool,
) -> anyhow::Result<Vec<String>> {
// 基础过滤下推至 SQL排除已标记为无资源的文献减少内存加载
let rows = sqlx::query(
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path, year, \
datetime(created_at, 'localtime') AS created_at_local, \
EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode) AS has_vector, \
EXISTS(SELECT 1 FROM paper_targets WHERE bibcode = papers.bibcode) AS has_target \
FROM papers \
WHERE (pdf_path IS NULL OR (pdf_path NOT LIKE 'error:no_resource%' AND pdf_path NOT LIKE 'error:%' AND pdf_path NOT LIKE 'error:%')) \
AND (html_path IS NULL OR (html_path NOT LIKE 'error:no_resource%' AND html_path NOT LIKE 'error:%' AND html_path NOT LIKE 'error:%'))"
)
.fetch_all(db)
.await?;
let mut records = Vec::new();
for r in rows {
use sqlx::Row;
records.push(PaperRecord {
bibcode: r.get("bibcode"),
pdf_path: r.get("pdf_path"),
html_path: r.get("html_path"),
markdown_path: r.get("markdown_path"),
translation_path: r.get("translation_path"),
year: r.get("year"),
created_at: r.get("created_at_local"),
has_vector: r.get("has_vector"),
has_target: r.get("has_target"),
});
}
// 排序
let sort_order_str = sort_order.unwrap_or_else(|| "default".to_string());
if sort_order_str == "pub_year_desc" {
records.sort_by(|a, b| b.year.cmp(&a.year));
} else if sort_order_str == "created_at_desc" {
records.sort_by(|a, b| b.created_at.cmp(&a.created_at));
}
let is_completed = |path: &Option<String>| -> bool {
if let Some(p) = path {
!p.is_empty() && !p.starts_with("error:") && !p.starts_with("mineru_batch:")
} else {
false
}
};
let is_failed = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:")
} else {
false
}
};
let is_no_resource = |path: &Option<String>| -> bool {
if let Some(p) = path {
p.starts_with("error:no_resource")
|| p.starts_with("error:无资源")
|| p.starts_with("error:无有效全文")
} else {
false
}
};
let mut target_bibcodes = Vec::new();
for rec in records {
let is_no_resource_paper = is_no_resource(&rec.pdf_path) || is_no_resource(&rec.html_path);
if is_no_resource_paper {
continue;
}
let download_completed = is_completed(&rec.pdf_path) || is_completed(&rec.html_path);
let download_failed = is_failed(&rec.pdf_path) || is_failed(&rec.html_path);
let download_uncompleted = !download_completed && !download_failed;
let parse_completed = is_completed(&rec.markdown_path);
let parse_failed = is_failed(&rec.markdown_path);
let parse_uncompleted = !parse_completed && !parse_failed;
let translate_completed = is_completed(&rec.translation_path);
let translate_failed = is_failed(&rec.translation_path);
match target_phase {
"download" => {
if skip_completed && download_completed {
continue;
}
if skip_failed && download_failed {
continue;
}
}
"parse" => {
if skip_completed && parse_completed {
continue;
}
if skip_failed && parse_failed {
continue;
}
if skip_preceding_failed && download_failed {
continue;
}
if skip_preceding_uncompleted && download_uncompleted {
continue;
}
}
"translate" => {
if skip_completed && translate_completed {
continue;
}
if skip_failed && translate_failed {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"embed" => {
if skip_completed && rec.has_vector {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
"target" => {
if skip_completed && rec.has_target {
continue;
}
if skip_preceding_failed && parse_failed {
continue;
}
if skip_preceding_uncompleted && parse_uncompleted {
continue;
}
}
_ => {}
}
target_bibcodes.push(rec.bibcode);
}
let limit = limit_count.unwrap_or(100) as usize;
if target_bibcodes.len() > limit {
target_bibcodes.truncate(limit);
}
Ok(target_bibcodes)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@ -321,7 +515,7 @@ mod tests {
.bind(&bibcode) .bind(&bibcode)
.fetch_one(&pool) .fetch_one(&pool)
.await?; .await?;
let md_path_rel: String = row.get(0); let md_path_rel: String = row.get("markdown_path");
assert_eq!(md_path_rel, format!("Markdown/{}.md", bibcode)); assert_eq!(md_path_rel, format!("Markdown/{}.md", bibcode));
assert!(temp_dir.join(&md_path_rel).exists()); assert!(temp_dir.join(&md_path_rel).exists());

View File

@ -314,4 +314,42 @@ impl MetaSync {
} }
}); });
} }
// 获取所有已存同步检索配置
pub async fn list_queries(db: &SqlitePool) -> Result<Vec<SavedSyncQuery>, sqlx::Error> {
let rows = sqlx::query("SELECT id, query, source, limit_count, datetime(last_run, 'localtime') FROM sync_queries ORDER BY last_run DESC")
.fetch_all(db)
.await?;
let mut list = Vec::new();
for r in rows {
use sqlx::Row;
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"),
});
}
Ok(list)
}
// 删除同步检索配置
pub async fn delete_query(db: &SqlitePool, id: i64) -> Result<bool, sqlx::Error> {
let result = sqlx::query("DELETE FROM sync_queries WHERE id = ?")
.bind(id)
.execute(db)
.await?;
Ok(result.rows_affected() > 0)
}
}
#[derive(Debug, Serialize, serde::Deserialize)]
pub struct SavedSyncQuery {
pub id: i64,
pub query: String,
pub source: String,
pub limit_count: i32,
pub last_run: String,
} }

View File

@ -2,5 +2,5 @@
pub mod asset; pub mod asset;
pub mod meta; pub mod meta;
pub use asset::{AssetBatch, AssetBatchStatus, BatchAction}; pub use asset::{filter_papers_for_batch, AssetBatch, AssetBatchStatus, BatchAction};
pub use meta::{MetaSync, MetaSyncStatus}; pub use meta::{MetaSync, MetaSyncStatus, SavedSyncQuery};

398
src/services/citation.rs Normal file
View File

@ -0,0 +1,398 @@
// src/services/citation.rs
// 文献引用关系查询服务层
use serde::{Deserialize, Serialize};
use sqlx::sqlite::SqlitePool;
use sqlx::FromRow;
use crate::services::paper::StandardPaper;
/// 统一的引文信息结构体
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CitationInfo {
pub bibcode: String,
pub title: String,
pub authors: Vec<String>,
pub first_author: String,
pub year: Option<String>,
pub citation_count: i32,
}
#[derive(Debug, FromRow)]
struct CitationDbRow {
bibcode: String,
title: String,
authors: Option<String>,
year: Option<String>,
citation_count: Option<i32>,
}
impl From<CitationDbRow> for CitationInfo {
fn from(row: CitationDbRow) -> Self {
let authors: Vec<String> = row
.authors
.as_deref()
.and_then(|a| serde_json::from_str(a).ok())
.unwrap_or_default();
let first_author = authors
.first()
.cloned()
.unwrap_or_else(|| "未知".to_string());
CitationInfo {
bibcode: row.bibcode,
title: row.title,
authors,
first_author,
year: row.year,
citation_count: row.citation_count.unwrap_or(0),
}
}
}
/// 解析查询字符串,提取作者部分和年份部分。
/// "Lei 2023" → ("Lei", Some("2023"))
/// "Zhang et al. 2020" → ("Zhang et al.", Some("2020"))
/// "Smith" → ("Smith", None)
pub fn parse_query(q: &str) -> (String, Option<String>) {
let cleaned = q.replace("et al.", "et al");
let parts: Vec<&str> = cleaned.split_whitespace().collect();
let mut year: Option<String> = None;
let mut author_tokens: Vec<&str> = Vec::new();
for (i, &p) in parts.iter().enumerate().rev() {
if year.is_none()
&& p.len() == 4
&& p.chars().all(|c| c.is_ascii_digit())
&& (p.starts_with("19") || p.starts_with("20"))
{
year = Some(p.to_string());
author_tokens = parts[..i].to_vec();
break;
}
}
if year.is_none() {
author_tokens = parts;
}
let author_part = if author_tokens.is_empty() {
q.to_string()
} else {
author_tokens.join(" ")
};
(author_part, year)
}
/// 在引用关系中模糊搜索匹配的文献
pub async fn search_citations(
db: &SqlitePool,
bibcode: &str,
direction: &str,
query: &str,
) -> Result<Vec<CitationInfo>, sqlx::Error> {
// 固定的 SQL 查询,列名由 match 分支选择,避免动态 format 拼接
const CITATIONS_WITH_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.source_bibcode \
WHERE cr.target_bibcode = ? AND p.authors LIKE ? AND p.year = ?";
const CITATIONS_NO_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.source_bibcode \
WHERE cr.target_bibcode = ? AND p.authors LIKE ?";
const REFERENCES_WITH_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.target_bibcode \
WHERE cr.source_bibcode = ? AND p.authors LIKE ? AND p.year = ?";
const REFERENCES_NO_YEAR: &str =
"SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count \
FROM papers p \
JOIN citations_references cr ON p.bibcode = cr.target_bibcode \
WHERE cr.source_bibcode = ? AND p.authors LIKE ?";
let (author, year) = parse_query(query);
let author_pattern = format!("%{}%", author);
let is_citations = direction == "citations";
// LIMIT 防止模糊查询返回过多结果
const SEARCH_LIMIT: i64 = 100;
let rows = if let Some(ref y) = year {
let sql = if is_citations {
CITATIONS_WITH_YEAR
} else {
REFERENCES_WITH_YEAR
};
sqlx::query_as::<_, CitationDbRow>(&format!("{sql} LIMIT ?"))
.bind(bibcode)
.bind(&author_pattern)
.bind(y)
.bind(SEARCH_LIMIT)
.fetch_all(db)
.await?
} else {
let sql = if is_citations {
CITATIONS_NO_YEAR
} else {
REFERENCES_NO_YEAR
};
sqlx::query_as::<_, CitationDbRow>(&format!("{sql} LIMIT ?"))
.bind(bibcode)
.bind(&author_pattern)
.bind(SEARCH_LIMIT)
.fetch_all(db)
.await?
};
Ok(rows.into_iter().map(CitationInfo::from).collect())
}
/// 分页查询文献的参考文献或被引文献列表,返回列表及总数
pub async fn get_citations_paginated(
db: &SqlitePool,
bibcode: &str,
direction: &str,
sort: &str,
offset: i64,
limit: i64,
) -> Result<(Vec<CitationInfo>, i64), sqlx::Error> {
// 固定 SQL方向通过 match 选择,避免动态列名拼接
let count_sql = if direction == "citations" {
"SELECT COUNT(*) FROM citations_references WHERE target_bibcode = ?"
} else {
"SELECT COUNT(*) FROM citations_references WHERE source_bibcode = ?"
};
let query_sql = if direction == "citations" {
match sort {
"year" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.source_bibcode WHERE cr.target_bibcode = ? ORDER BY p.year DESC LIMIT ? OFFSET ?",
"citation_count" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.source_bibcode WHERE cr.target_bibcode = ? ORDER BY p.citation_count DESC LIMIT ? OFFSET ?",
_ => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.source_bibcode WHERE cr.target_bibcode = ? ORDER BY p.bibcode ASC LIMIT ? OFFSET ?",
}
} else {
match sort {
"year" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.target_bibcode WHERE cr.source_bibcode = ? ORDER BY p.year DESC LIMIT ? OFFSET ?",
"citation_count" => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.target_bibcode WHERE cr.source_bibcode = ? ORDER BY p.citation_count DESC LIMIT ? OFFSET ?",
_ => "SELECT p.bibcode, p.title, p.authors, p.year, p.citation_count FROM papers p JOIN citations_references cr ON p.bibcode = cr.target_bibcode WHERE cr.source_bibcode = ? ORDER BY p.bibcode ASC LIMIT ? OFFSET ?",
}
};
// 查询总数
let total: i64 = sqlx::query_scalar(count_sql)
.bind(bibcode)
.fetch_one(db)
.await
.unwrap_or_else(|e| {
tracing::warn!("查询引用总数失败: {}", e);
0
});
// 分页查询
let rows: Vec<CitationDbRow> = sqlx::query_as(query_sql)
.bind(bibcode)
.bind(limit)
.bind(offset)
.fetch_all(db)
.await?;
let infos = rows.into_iter().map(CitationInfo::from).collect();
Ok((infos, total))
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CitationNetworkResult {
pub bibcode: String,
pub title: String,
pub citation_count: i32,
pub reference_count: i32,
pub references: Vec<String>,
pub citations: Vec<String>,
pub citation_counts: std::collections::HashMap<String, i32>,
}
/// 文献元数据加载(本地 + ADS fallback
async fn resolve_paper_from_db_or_ads(
db: &SqlitePool,
library_dir: &std::path::Path,
ads_api_key: &str,
ads_client: &crate::clients::ads::AdsClient,
bibcode: &str,
) -> anyhow::Result<StandardPaper> {
match crate::services::paper::get_paper_from_db(db, library_dir, bibcode).await {
Ok(p) => Ok(p),
Err(_) => {
if ads_api_key.is_empty() {
anyhow::bail!(
"本地数据库未收录该文献,且未配置 ADS_API_KEY无法在线加载: {}",
bibcode
);
}
let docs = ads_client
.search(&format!("bibcode:{}", bibcode), 0, 1, "relevance")
.await?;
let doc = docs
.first()
.ok_or_else(|| anyhow::anyhow!("在本地库及 ADS 中均未找到该文献: {}", bibcode))?;
let paper = crate::services::paper::convert_ads_doc_to_standard(doc);
if let Err(e) = crate::services::paper::save_paper_to_db(db, &paper).await {
tracing::warn!("保存引用文献至数据库失败: {}", e);
}
// 保存引用拓扑关系(批量 INSERT避免 N+1
let mut pairs: Vec<(String, String)> = Vec::new();
if let Some(refs) = &doc.reference {
for ref_bib in refs {
pairs.push((paper.bibcode.clone(), ref_bib.clone()));
}
}
if let Some(cits) = &doc.citation {
for cit_bib in cits {
pairs.push((cit_bib.clone(), paper.bibcode.clone()));
}
}
if !pairs.is_empty() {
let mut tx = match db.begin().await {
Ok(tx) => tx,
Err(e) => {
tracing::warn!("无法开启引用关系事务: {}", e);
return Ok(paper);
}
};
for chunk in pairs.chunks(crate::services::paper::SQLITE_PARAM_LIMIT / 2) {
for (src, tgt) in chunk {
if let Err(e) = sqlx::query(
"INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)",
)
.bind(src)
.bind(tgt)
.execute(&mut *tx)
.await
{
tracing::warn!("保存引用关系失败 ({} -> {}): {}", src, tgt, e);
}
}
}
if let Err(e) = tx.commit().await {
tracing::warn!("提交引用关系事务失败: {}", e);
}
}
Ok(paper)
}
}
}
/// 批量查询关联文献的被引数量
async fn batch_load_citation_counts(
db: &SqlitePool,
references: &[String],
citations: &[String],
) -> std::collections::HashMap<String, i32> {
let mut counts = std::collections::HashMap::new();
let mut all: Vec<String> = references.to_vec();
all.extend(citations.iter().cloned());
all.sort();
all.dedup();
for chunk in all.chunks(crate::services::paper::SQLITE_PARAM_LIMIT) {
if chunk.is_empty() {
continue;
}
let placeholders = vec!["?"; chunk.len()].join(", ");
let sql = format!(
"SELECT bibcode, citation_count FROM papers WHERE bibcode IN ({})",
placeholders
);
let mut query = sqlx::query(&sql);
for b in chunk {
query = query.bind(b);
}
match query.fetch_all(db).await {
Ok(rows) => {
for row in rows {
counts.insert(sqlx::Row::get(&row, 0), sqlx::Row::get(&row, 1));
}
}
Err(e) => {
tracing::warn!("批量查询引用数失败 (chunk size={}): {}", chunk.len(), e);
}
}
}
counts
}
pub async fn get_citation_network_with_fallback(
db: &SqlitePool,
library_dir: &std::path::Path,
ads_api_key: &str,
ads_client: &crate::clients::ads::AdsClient,
bibcode: &str,
) -> anyhow::Result<CitationNetworkResult> {
let paper =
resolve_paper_from_db_or_ads(db, library_dir, ads_api_key, ads_client, bibcode).await?;
// 加载引用文献列表
let refs_rows =
sqlx::query("SELECT target_bibcode FROM citations_references WHERE source_bibcode = ?")
.bind(bibcode)
.fetch_all(db)
.await
.unwrap_or_else(|e| {
tracing::warn!("查询引用关系失败: {}", e);
Vec::new()
});
let references: Vec<String> = refs_rows.iter().map(|row| sqlx::Row::get(row, 0)).collect();
// 加载被引文献列表
let cits_rows =
sqlx::query("SELECT source_bibcode FROM citations_references WHERE target_bibcode = ?")
.bind(bibcode)
.fetch_all(db)
.await
.unwrap_or_else(|e| {
tracing::warn!("查询被引关系失败: {}", e);
Vec::new()
});
let citations: Vec<String> = cits_rows.iter().map(|row| sqlx::Row::get(row, 0)).collect();
// 批量加载关联文献的被引数量
let citation_counts = batch_load_citation_counts(db, &references, &citations).await;
Ok(CitationNetworkResult {
bibcode: paper.bibcode,
title: paper.title,
citation_count: paper.citation_count,
reference_count: paper.reference_count,
references,
citations,
citation_counts,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_query_with_year() {
let (author, year) = parse_query("Lei 2023");
assert_eq!(author, "Lei");
assert_eq!(year, Some("2023".into()));
}
#[test]
fn test_parse_query_et_al() {
let (author, year) = parse_query("Zhang et al. 2020");
assert_eq!(author, "Zhang et al");
assert_eq!(year, Some("2020".into()));
}
#[test]
fn test_parse_query_author_only() {
let (author, year) = parse_query("Smith");
assert_eq!(author, "Smith");
assert_eq!(year, None);
}
}

View File

@ -12,8 +12,7 @@ mod strategies;
pub(crate) use headers::validate_html_content_lenient; pub(crate) use headers::validate_html_content_lenient;
pub(crate) use headers::validate_pdf_content; pub(crate) use headers::validate_pdf_content;
use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db}; use crate::services::paper::{check_paper_paths_in_db, get_paper_from_db, StandardPaper};
use crate::api::StandardPaper;
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use reqwest::header::{HeaderMap, HeaderValue}; use reqwest::header::{HeaderMap, HeaderValue};
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
@ -646,7 +645,7 @@ mod tests {
#[tokio::test] #[tokio::test]
#[ignore] #[ignore]
async fn test_download_via_obscura_integration() -> anyhow::Result<()> { async fn test_download_via_obscura_integration() -> anyhow::Result<()> {
use axum::http::{header::CONTENT_TYPE, HeaderValue}; use axum::http::header::CONTENT_TYPE;
use axum::{response::Response as AxumResponse, routing::get, Router}; use axum::{response::Response as AxumResponse, routing::get, Router};
// Bind to a random port // Bind to a random port

View File

@ -1,11 +1,16 @@
pub mod batch; pub mod batch;
pub mod chunker; pub mod chunker;
pub mod citation;
pub mod download; pub mod download;
pub mod logging; pub mod logging;
pub mod note;
pub mod paper;
pub mod parser; pub mod parser;
pub mod pipeline;
pub mod query_parser; pub mod query_parser;
pub mod rag; pub mod rag;
pub mod search; pub mod search;
pub mod section_parser; pub mod session;
pub mod target; pub mod target;
pub mod translation; pub mod translation;
pub mod vision;

121
src/services/note.rs Normal file
View File

@ -0,0 +1,121 @@
// src/services/note.rs
// 研究笔记服务层
use std::path::{Path, PathBuf};
/// 保存研究笔记为本地 Markdown 文件
pub fn save_note_service(
library_dir: &Path,
title: &str,
content: &str,
) -> anyhow::Result<(String, PathBuf, usize)> {
// 文件名安全化
let safe_title: String = title
.chars()
.map(|c| {
if c.is_alphanumeric() || c == '-' || c == '_' || c == ' ' {
c
} else {
'_'
}
})
.collect::<String>()
.replace(' ', "_");
let notes_dir = library_dir.join("notes");
std::fs::create_dir_all(&notes_dir)?;
let filename = format!("{}.md", safe_title);
let filepath = notes_dir.join(&filename);
let full_content = format!(
"---\ntitle: {}\ndate: {}\ngenerated_by: AstroResearch Agent\n---\n\n{}",
title,
chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
content
);
std::fs::write(&filepath, &full_content)?;
Ok((filename, filepath, full_content.len()))
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, sqlx::FromRow)]
pub struct NoteRecord {
pub id: i64,
pub bibcode: String,
pub paragraph_index: i64,
pub note_text: String,
pub highlight_color: String,
pub selected_text: String,
pub created_at: String,
}
/// 创建文献段落高亮与笔记
pub async fn create_highlight_note(
db: &sqlx::SqlitePool,
bibcode: &str,
paragraph_index: i64,
note_text: &str,
highlight_color: &str,
selected_text: &str,
) -> Result<NoteRecord, sqlx::Error> {
let note = sqlx::query_as::<_, NoteRecord>(
"INSERT INTO notes (bibcode, paragraph_index, note_text, highlight_color, selected_text) VALUES (?, ?, ?, ?, ?) RETURNING id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at"
)
.bind(bibcode)
.bind(paragraph_index)
.bind(note_text)
.bind(highlight_color)
.bind(selected_text)
.fetch_one(db)
.await?;
Ok(note)
}
/// 查询某篇文献的全部高亮笔记
pub async fn list_highlight_notes(
db: &sqlx::SqlitePool,
bibcode: &str,
) -> Result<Vec<NoteRecord>, sqlx::Error> {
let notes = sqlx::query_as::<_, NoteRecord>(
"SELECT id, bibcode, paragraph_index, note_text, highlight_color, selected_text, created_at FROM notes WHERE bibcode = ? ORDER BY paragraph_index, created_at"
)
.bind(bibcode)
.fetch_all(db)
.await?;
Ok(notes)
}
/// 删除指定 id 的高亮笔记
pub async fn delete_highlight_note(db: &sqlx::SqlitePool, id: i64) -> Result<bool, sqlx::Error> {
let result = sqlx::query("DELETE FROM notes WHERE id = ?")
.bind(id)
.execute(db)
.await?;
Ok(result.rows_affected() > 0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_save_note_service() {
let tmp_dir = std::env::temp_dir().join("astro_note_test");
let title = "Test Note Title 123!";
let content = "Hello world from note service test";
let res = save_note_service(&tmp_dir, title, content);
assert!(res.is_ok());
let (filename, filepath, size) = res.unwrap();
assert_eq!(filename, "Test_Note_Title_123_.md");
assert!(filepath.exists());
assert!(size > content.len());
std::fs::remove_dir_all(&tmp_dir).ok();
}
}

367
src/services/paper/db.rs Normal file
View File

@ -0,0 +1,367 @@
// src/services/paper/db.rs
//
// 文献元数据数据库存取。
use super::model::StandardPaper;
use sqlx::{Acquire, Row, SqlitePool};
use std::path::Path;
use tracing::info;
/// SQLite 单条 SQL 的参数绑定上限SQLITE_MAX_VARIABLE_NUMBER 默认 999留余量取 900
/// 批量 IN 查询时需按此分批,避免超出绑定数量。
pub const SQLITE_PARAM_LIMIT: usize = 900;
/// 清洗后的文献标识符,提取纯 DOI 和 arXiv ID。
///
/// 统一处理用户可能输入的多种前缀格式(`doi:`, `https://doi.org/`, `arxiv:` 等)。
#[derive(Debug, Clone)]
pub struct CleanId {
/// 原始标识符的 trim 结果
pub raw: String,
/// 去除了 DOI 前缀的纯标识符
pub clean_doi: String,
/// 去除了 arXiv 前缀及 pdf/ 前缀的纯标识符
pub clean_arxiv: String,
/// arXiv ID 去除版本号后缀(仅 ingest 使用,如 `2303.12345v1` → `2303.12345`
pub clean_arxiv_no_version: String,
}
/// 清洗文献标识符去除常见前缀doi:, arXiv:, https://doi.org/ 等)。
pub fn clean_identifier(identifier: &str) -> CleanId {
let raw = identifier.trim().to_string();
let clean_doi = raw
.trim_start_matches("doi:")
.trim_start_matches("DOI:")
.trim_start_matches("https://doi.org/")
.trim_start_matches("http://doi.org/")
.trim()
.to_string();
let clean_arxiv = raw
.trim_start_matches("arxiv:")
.trim_start_matches("arXiv:")
.trim_start_matches("ARXIV:")
.trim_start_matches("pdf/")
.trim()
.to_string();
// 移除 arXiv 版本号后缀(如 2303.12345v1 → 2303.12345
let clean_arxiv_no_version = if let Some(pos) = clean_arxiv.find('v') {
if pos + 1 < clean_arxiv.len() && clean_arxiv[pos + 1..].chars().all(|c| c.is_ascii_digit())
{
clean_arxiv[..pos].to_string()
} else {
clean_arxiv.clone()
}
} else {
clean_arxiv.clone()
};
CleanId {
raw,
clean_doi,
clean_arxiv,
clean_arxiv_no_version,
}
}
pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Result<()> {
let mut conn = db.acquire().await?;
save_paper_to_db_tx(&mut conn, p).await
}
pub async fn save_paper_to_db_tx(
conn: &mut sqlx::SqliteConnection,
p: &StandardPaper,
) -> anyhow::Result<()> {
let authors_json = serde_json::to_string(&p.authors)?;
let keywords_json = serde_json::to_string(&p.keywords)?;
// 1. 如果存在 arxiv_id检查是否有已存在的相同 arxiv_id 记录以防 duplicate
if !p.arxiv_id.is_empty() {
#[allow(clippy::type_complexity)]
let existing_opt: Option<(String, Option<String>, Option<String>, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT bibcode, pdf_path, html_path, markdown_path, translation_path FROM papers WHERE arxiv_id = ?"
)
.bind(&p.arxiv_id)
.fetch_optional(&mut *conn)
.await?;
if let Some((existing_bibcode, _pdf, _html, _md, _tr)) = existing_opt {
if existing_bibcode != p.bibcode {
// 发现不同 bibcode 标识的同一篇文献记录,需要进行合并
// 如果已存在的记录使用的是临时 arXiv ID 作为 bibcode且新记录使用的是正式 ADS bibcode我们升级 bibcode 主键
let is_existing_temp = existing_bibcode == p.arxiv_id;
let is_new_formal = p.bibcode != p.arxiv_id;
if is_existing_temp && is_new_formal {
info!(
"发现相同 arXiv ID 的文献,将临时主键 {} 升级为正式 ADS Bibcode: {}",
existing_bibcode, p.bibcode
);
// 在事务中原子执行 bibcode 迁移papers 表更新 + 两条 citation_references 引用级联
let mut tx = conn.begin().await?;
sqlx::query(
"UPDATE papers SET bibcode = ?, title = ?, authors = ?, year = ?, pub = ?, keywords = ?, abstract = ?, doi = ?, citation_count = ?, reference_count = ?, doctype = ? WHERE bibcode = ?"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(&authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(&keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.bind(&existing_bibcode)
.execute(&mut *tx)
.await?;
// 同步更新引用关系表(手动级联,在同一事务内)
sqlx::query("UPDATE citations_references SET source_bibcode = ? WHERE source_bibcode = ?")
.bind(&p.bibcode)
.bind(&existing_bibcode)
.execute(&mut *tx)
.await?;
sqlx::query("UPDATE citations_references SET target_bibcode = ? WHERE target_bibcode = ?")
.bind(&p.bibcode)
.bind(&existing_bibcode)
.execute(&mut *tx)
.await?;
tx.commit().await?;
return Ok(());
} else {
// 如果已存在的是正式 ADS bibcode而新插入的是临时 arXiv ID直接忽略或更新元数据而不更改主键
info!(
"发现相同 arXiv ID 的文献 {} 已存在正式记录,忽略临时 arXiv 插入",
existing_bibcode
);
return Ok(());
}
}
}
}
// 2. 正常插入/冲突更新
sqlx::query(
"INSERT INTO papers (bibcode, title, authors, year, pub, keywords, abstract, doi, arxiv_id, citation_count, reference_count, doctype) \
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \
ON CONFLICT(bibcode) DO UPDATE SET \
title=excluded.title, \
authors=excluded.authors, \
pub=excluded.pub, \
keywords=excluded.keywords, \
abstract=excluded.abstract, \
doi=excluded.doi, \
arxiv_id=excluded.arxiv_id, \
citation_count=excluded.citation_count, \
reference_count=excluded.reference_count, \
doctype=excluded.doctype"
)
.bind(&p.bibcode)
.bind(&p.title)
.bind(authors_json)
.bind(&p.year)
.bind(&p.pub_journal)
.bind(keywords_json)
.bind(&p.abstract_text)
.bind(&p.doi)
.bind(&p.arxiv_id)
.bind(p.citation_count)
.bind(p.reference_count)
.bind(&p.doctype)
.execute(&mut *conn)
.await?;
Ok(())
}
/// 从数据库行解析出 StandardPaper
pub fn parse_paper_row(r: &sqlx::sqlite::SqliteRow, library_dir: &Path) -> StandardPaper {
let pdf_path: Option<String> = r.get("pdf_path");
let html_path: Option<String> = r.get("html_path");
let markdown_path: Option<String> = r.get("markdown_path");
let translation_path: Option<String> = r.get("translation_path");
let doctype_val: Option<String> = r.get("doctype");
let has_vector: bool = r.get("has_vector");
let authors_str: Option<String> = r.get("authors");
let authors: Vec<String> = authors_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let keywords_str: Option<String> = r.get("keywords");
let keywords: Vec<String> = keywords_str
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let is_pdf_exist = pdf_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_html_exist = html_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_md_exist = markdown_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let is_tr_exist = translation_path
.as_ref()
.map(|p| library_dir.join(p).exists())
.unwrap_or(false);
let pdf_error = pdf_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
let html_error = html_path
.as_ref()
.filter(|p| p.starts_with("error:"))
.map(|p| p["error:".len()..].trim().to_string());
StandardPaper {
bibcode: r.get("bibcode"),
title: r.get("title"),
authors,
year: r.get("year"),
pub_journal: r.get("pub"),
keywords,
abstract_text: r.get("abstract"),
doi: r.get("doi"),
arxiv_id: r.get("arxiv_id"),
citation_count: r.get("citation_count"),
reference_count: r.get("reference_count"),
is_downloaded: is_pdf_exist || is_html_exist,
has_pdf: is_pdf_exist,
has_html: is_html_exist,
has_markdown: is_md_exist,
has_translation: is_tr_exist,
has_vector,
doctype: doctype_val.unwrap_or_else(|| "article".to_string()),
pdf_error,
html_error,
}
}
/// 通用的 5 路 OR 查询条件bibcode / doi / arxiv_id 精确匹配 + doi 大小写不敏感匹配)
const PAPER_LOOKUP_WHERE: &str =
"bibcode = ? OR doi = ? OR arxiv_id = ? OR LOWER(doi) = LOWER(?) OR arxiv_id = ?";
pub async fn get_paper_from_db(
db: &SqlitePool,
library_dir: &Path,
identifier: &str,
) -> anyhow::Result<StandardPaper> {
let id = clean_identifier(identifier);
let query = format!(
"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 {}",
PAPER_LOOKUP_WHERE
);
let r = sqlx::query(&query)
.bind(&id.raw)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.fetch_one(db)
.await?;
Ok(parse_paper_row(&r, library_dir))
}
pub async fn check_paper_paths_in_db(
db: &SqlitePool,
library_dir: &Path,
identifier: &str,
) -> anyhow::Result<
Option<(
Option<String>,
Option<String>,
Option<String>,
Option<String>,
)>,
> {
let id = clean_identifier(identifier);
let query = format!(
"SELECT pdf_path, html_path, markdown_path, translation_path FROM papers WHERE {}",
PAPER_LOOKUP_WHERE
);
let r_opt = sqlx::query(&query)
.bind(&id.raw)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.bind(&id.clean_doi)
.bind(&id.clean_arxiv)
.fetch_optional(db)
.await?;
if let Some(r) = r_opt {
let pdf: Option<String> = r.get("pdf_path");
let html: Option<String> = r.get("html_path");
let md: Option<String> = r.get("markdown_path");
let tr: Option<String> = r.get("translation_path");
let pdf_res = pdf.filter(|p| library_dir.join(p).exists());
let html_res = html.filter(|p| library_dir.join(p).exists());
let md_res = md.filter(|p| library_dir.join(p).exists());
let tr_res = tr.filter(|p| library_dir.join(p).exists());
Ok(Some((pdf_res, html_res, md_res, tr_res)))
} else {
Ok(None)
}
}
/// 获取本地文献库的全部文献列表
pub async fn get_library_list(
db: &SqlitePool,
library_dir: &Path,
limit: Option<i64>,
offset: Option<i64>,
) -> Result<Vec<StandardPaper>, sqlx::Error> {
let rows = if let (Some(lim), Some(off)) = (limit, offset) {
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 ORDER BY created_at DESC LIMIT ? OFFSET ?")
.bind(lim)
.bind(off)
.fetch_all(db)
.await?
} else {
// 无分页参数时默认限制 2000 条,防止全量加载
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 ORDER BY created_at DESC LIMIT 2000")
.fetch_all(db)
.await?
};
let mut list = Vec::new();
for r in rows {
list.push(parse_paper_row(&r, library_dir));
}
Ok(list)
}
/// 标记/清除文献的“无资源”状态
pub async fn mark_no_resource_service(
db: &SqlitePool,
library_dir: &Path,
bibcode: &str,
clear: bool,
) -> anyhow::Result<StandardPaper> {
if clear {
sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?")
.bind(bibcode)
.execute(db)
.await?;
} else {
sqlx::query("UPDATE papers SET pdf_path = 'error:no_resource', html_path = 'error:no_resource' WHERE bibcode = ?")
.bind(bibcode)
.execute(db)
.await?;
}
let updated_paper = get_paper_from_db(db, library_dir, bibcode).await?;
Ok(updated_paper)
}

View File

@ -0,0 +1,138 @@
// src/services/paper/ingest.rs
//
// 文献手动上传与推入逻辑Ingestion Service
use crate::services::paper::db::get_paper_from_db;
use crate::services::paper::model::StandardPaper;
use sqlx::SqlitePool;
use std::path::Path;
use tracing::info;
/// 校验并保存上传的文献文件PDF 或 HTML并更新数据库状态。
pub async fn upload_paper_file_service(
db: &SqlitePool,
library_dir: &Path,
input_bibcode: &str,
file_type: &str,
file_name: &str,
file_bytes: &[u8],
) -> anyhow::Result<StandardPaper> {
if file_bytes.is_empty() {
anyhow::bail!("上传文件为空或读取失败");
}
// 尝试将可能的 DOI 或 arXiv ID 解析为真实的 bibcode
let mut resolved_bibcode = input_bibcode.to_string();
let exists_as_bibcode = sqlx::query("SELECT bibcode FROM papers WHERE bibcode = ?")
.bind(&resolved_bibcode)
.fetch_optional(db)
.await
.unwrap_or(None)
.is_some();
if !exists_as_bibcode {
let id = crate::services::paper::clean_identifier(&resolved_bibcode);
use sqlx::Row;
// 尝试匹配 DOI
if let Some(row) = sqlx::query(
"SELECT bibcode FROM papers WHERE doi = ? OR doi = ? OR LOWER(doi) = LOWER(?)",
)
.bind(&id.clean_doi)
.bind(&resolved_bibcode)
.bind(&id.clean_doi)
.fetch_optional(db)
.await
.unwrap_or(None)
{
let found: String = row.get("bibcode");
info!(
"上传服务:通过 DOI 匹配成功,将 '{}' 解析为 bibcode '{}'",
input_bibcode, found
);
resolved_bibcode = found;
} else if let Some(row) = sqlx::query(
// 尝试匹配 arXiv ID含版本号去除
"SELECT bibcode FROM papers WHERE arxiv_id = ? OR arxiv_id = ? OR arxiv_id LIKE ? OR arxiv_id LIKE ?"
)
.bind(&id.clean_arxiv)
.bind(&id.clean_arxiv_no_version)
.bind(format!("{}%", id.clean_arxiv_no_version))
.bind(format!("arXiv:{}%", id.clean_arxiv_no_version))
.fetch_optional(db)
.await
.unwrap_or(None)
{
let found: String = row.get("bibcode");
info!("上传服务:通过 arXiv ID 匹配成功,将 '{}' 解析为 bibcode '{}'", input_bibcode, found);
resolved_bibcode = found;
}
}
let bibcode = resolved_bibcode;
// 从数据库读取该文献元数据以验证其确实存在
let _paper = get_paper_from_db(db, library_dir, &bibcode)
.await
.map_err(|e| anyhow::anyhow!("未找到该文献记录: {}", e))?;
// 校验并保存文件
// 优先使用显式 file_type 字段,避免文件名扩展名覆盖真实类型
let is_pdf = if !file_type.is_empty() {
file_type == "pdf"
} else {
file_name.to_lowercase().ends_with(".pdf")
};
let relative_path = if is_pdf {
crate::services::download::validate_pdf_content(file_bytes)
.map_err(|e| anyhow::anyhow!("PDF 文件内容校验失败: {}", e))?;
let pdf_filename = format!("{}.pdf", bibcode);
let pdf_dest = library_dir.join("PDF").join(&pdf_filename);
if let Some(parent) = pdf_dest.parent() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| anyhow::anyhow!("无法创建 PDF 目录: {}", e))?;
}
tokio::fs::write(&pdf_dest, file_bytes)
.await
.map_err(|e| anyhow::anyhow!("无法写入 PDF 文件: {}", e))?;
format!("PDF/{}", pdf_filename)
} else {
let text_content = String::from_utf8(file_bytes.to_vec())
.map_err(|_| anyhow::anyhow!("上传的 HTML 文件不是有效的 UTF-8 文本"))?;
crate::services::download::validate_html_content_lenient(&text_content)
.map_err(|e| anyhow::anyhow!("HTML 文件内容校验失败: {}", e))?;
let html_filename = format!("{}.html", bibcode);
let html_dest = library_dir.join("HTML").join(&html_filename);
if let Some(parent) = html_dest.parent() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| anyhow::anyhow!("无法创建 HTML 目录: {}", e))?;
}
tokio::fs::write(&html_dest, &text_content)
.await
.map_err(|e| anyhow::anyhow!("无法写入 HTML 文件: {}", e))?;
format!("HTML/{}", html_filename)
};
// 更新数据库路径状态(固定列名,避免动态 SQL 拼接)
if is_pdf {
sqlx::query("UPDATE papers SET pdf_path = ? WHERE bibcode = ?")
.bind(&relative_path)
.bind(&bibcode)
.execute(db)
.await
.map_err(|e| anyhow::anyhow!("更新数据库状态失败: {}", e))?;
} else {
sqlx::query("UPDATE papers SET html_path = ? WHERE bibcode = ?")
.bind(&relative_path)
.bind(&bibcode)
.execute(db)
.await
.map_err(|e| anyhow::anyhow!("更新数据库状态失败: {}", e))?;
}
// 重新获取最新的文献信息以返回
let updated_paper = get_paper_from_db(db, library_dir, &bibcode).await?;
Ok(updated_paper)
}

20
src/services/paper/mod.rs Normal file
View File

@ -0,0 +1,20 @@
// src/services/paper/mod.rs
//
// 文献服务模块入口,统一重导出子模块接口。
pub mod db;
pub mod ingest;
pub mod model;
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,
SQLITE_PARAM_LIMIT,
};
pub use ingest::upload_paper_file_service;
pub use model::{convert_ads_doc_to_standard, convert_arxiv_to_standard, StandardPaper};
pub use reader::{
extract_outline, extract_section_by_index, extract_section_by_name, read_paper_content,
MarkdownSection, PaperReadResponse, ReadMode,
};

110
src/services/paper/model.rs Normal file
View File

@ -0,0 +1,110 @@
// src/services/paper/model.rs
//
// 文献领域模型与转换逻辑。
use crate::clients::ads::AdsPaperDoc;
use crate::clients::arxiv::ArxivPaper;
use serde::{Deserialize, Serialize};
// 统一标准化的文献格式,用于在系统内部及向前端传输
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct StandardPaper {
pub bibcode: String,
pub title: String,
pub authors: Vec<String>,
pub year: String,
pub pub_journal: String,
pub keywords: Vec<String>,
pub abstract_text: String,
pub doi: String,
pub arxiv_id: String,
pub citation_count: i32,
pub reference_count: i32,
pub is_downloaded: bool,
pub has_pdf: bool,
pub has_html: bool,
pub has_markdown: bool,
pub has_translation: bool,
pub has_vector: bool,
pub doctype: String,
pub pdf_error: Option<String>,
pub html_error: Option<String>,
}
pub fn convert_ads_doc_to_standard(doc: &AdsPaperDoc) -> StandardPaper {
let title = doc
.title
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_else(|| doc.bibcode.clone());
let authors = doc.author.clone().unwrap_or_default();
let keywords = doc.keyword.clone().unwrap_or_default();
let doi = doc
.doi
.as_ref()
.and_then(|v: &Vec<String>| v.first())
.cloned()
.unwrap_or_default();
let mut arxiv_id = String::new();
if let Some(identifiers) = &doc.identifier {
for id in identifiers {
if id.starts_with("arXiv:") {
arxiv_id = id.replace("arXiv:", "").trim().to_string();
break;
}
}
}
if arxiv_id.is_empty() && doc.bibcode.starts_with("arXiv") {
arxiv_id = doc.bibcode.replace("arXiv", "").trim().to_string();
}
StandardPaper {
bibcode: doc.bibcode.clone(),
title,
authors,
year: doc.year.clone().unwrap_or_default(),
pub_journal: doc.pub_journal.clone().unwrap_or_default(),
keywords,
abstract_text: doc.abstract_text.clone().unwrap_or_default(),
doi,
arxiv_id,
citation_count: doc.citation_count.unwrap_or(0),
reference_count: doc.reference_count.unwrap_or(0),
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: doc.doctype.clone().unwrap_or_else(|| "article".to_string()),
pdf_error: None,
html_error: None,
}
}
pub fn convert_arxiv_to_standard(doc: &ArxivPaper) -> StandardPaper {
StandardPaper {
bibcode: doc.id.clone(),
title: doc.title.clone(),
authors: doc.authors.clone(),
year: doc.year.clone(),
pub_journal: "arXiv Preprint".to_string(),
keywords: Vec::new(),
abstract_text: doc.abstract_text.clone(),
doi: doc.doi.clone().unwrap_or_default(),
arxiv_id: doc.id.clone(),
citation_count: 0,
reference_count: 0,
is_downloaded: false,
has_pdf: false,
has_html: false,
has_markdown: false,
has_translation: false,
has_vector: false,
doctype: "eprint".to_string(),
pdf_error: None,
html_error: None,
}
}

View File

@ -1,13 +1,29 @@
// src/services/section_parser.rs // src/services/paper/reader.rs
// //
// Markdown 章节解析服务。 // 文献内容读取与章节大纲解析服务。
// 从已解析的文献 Markdown 中提取章节大纲(目录)和指定章节内容。
use regex::Regex; use regex::Regex;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use std::path::Path;
use std::sync::LazyLock; use std::sync::LazyLock;
use tracing;
/// 读取模式
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ReadMode {
/// 全文模式。可选择是否同时附带翻译内容
Full { include_translation: bool },
/// 仅提取大纲目录(可指定最大层级,默认 3
Outline { max_level: Option<usize> },
/// 按序号提取单个章节0-based
SectionIndex(usize),
/// 按名称模糊匹配提取单个章节
SectionName(String),
}
/// Markdown 章节结构 /// Markdown 章节结构
#[derive(Debug, Clone, serde::Serialize)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MarkdownSection { pub struct MarkdownSection {
/// 章节序号0-based /// 章节序号0-based
pub index: usize, pub index: usize,
@ -25,13 +41,117 @@ pub struct MarkdownSection {
pub char_end: usize, pub char_end: usize,
} }
/// 匹配 Markdown ATX 标题行(##, ###, ####... /// 统一的文献读取响应
#[derive(Debug, Clone, Serialize)]
pub struct PaperReadResponse {
pub bibcode: String,
/// 主文本内容(原文全文、大纲、或者单个章节文本)
pub content: String,
/// 如果是 Full 模式且 include_translation = true返回翻译内容全文
pub translation_content: Option<String>,
/// 如果是大纲模式,附带结构化的大纲数据,方便 API 转换为 JSON 返回
pub outline: Option<Vec<MarkdownSection>>,
}
/// 统一读取文献内容的业务逻辑函数
pub async fn read_paper_content(
db: &SqlitePool,
library_dir: &Path,
bibcode: &str,
mode: ReadMode,
) -> anyhow::Result<PaperReadResponse> {
// 1. 统一安全与数据库路径校验
let paths = super::db::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| {
anyhow::anyhow!("文献 {} 未在本地数据库中注册,请先检索该文献。", bibcode)
})?;
let (_, _, md_opt, tr_opt) = paths;
// 2. 检查英文原文相对路径
let md_rel = md_opt.ok_or_else(|| {
anyhow::anyhow!("文献 {} 尚未完成解析 (parse)。请先执行解析任务。", bibcode)
})?;
// 3. 读取英文原文物理文件
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!(
"文献 {} 的本地原文文件已丢失,请重新执行解析任务。",
bibcode
);
}
let full_text = tokio::fs::read_to_string(&md_abs).await?;
// 4. 初始化响应字段
let mut content = full_text.clone();
let mut translation_content = None;
let mut outline = None;
// 5. 根据模式处理文本
match &mode {
ReadMode::Full {
include_translation,
} => {
if *include_translation {
if let Some(tr_rel) = tr_opt {
let tr_abs = library_dir.join(&tr_rel);
if tr_abs.exists() {
match tokio::fs::read_to_string(&tr_abs).await {
Ok(tr_text) => translation_content = Some(tr_text),
Err(e) => tracing::warn!(
"翻译文件存在但读取失败 for {} ({}): {}",
bibcode,
tr_abs.display(),
e
),
}
}
}
}
}
ReadMode::Outline { max_level } => {
let level = max_level.unwrap_or(3);
let sections = extract_outline(&full_text, level);
let formatted = sections
.iter()
.map(|s| {
format!(
"{}[{}] {}",
" ".repeat(s.level.saturating_sub(1)),
s.index,
s.heading
)
})
.collect::<Vec<_>>()
.join("\n");
content = formatted;
outline = Some(sections);
}
ReadMode::SectionIndex(idx) => {
content = extract_section_by_index(&full_text, *idx)
.ok_or_else(|| anyhow::anyhow!("未找到序号为 #{} 的章节,请先检查大纲。", idx))?;
}
ReadMode::SectionName(name) => {
content = extract_section_by_name(&full_text, name).ok_or_else(|| {
anyhow::anyhow!("未找到名称匹配 '{}' 的章节,请先检查大纲。", name)
})?;
}
}
Ok(PaperReadResponse {
bibcode: bibcode.to_string(),
content,
translation_content,
outline,
})
}
static HEADING_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^(#{2,})\s+(.+)$").unwrap()); static HEADING_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^(#{2,})\s+(.+)$").unwrap());
/// 提取文献章节大纲。
///
/// `max_level` 控制最大标题级别2 = 仅 ## 章标题3 = ## + ###,以此类推。
/// 跳过 `#`(一级标题,通常是论文标题本身)。
pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection> { pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection> {
let mut sections = Vec::new(); let mut sections = Vec::new();
let mut index = 0usize; let mut index = 0usize;
@ -48,25 +168,23 @@ pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection>
index, index,
level, level,
heading, heading,
start_line: line_num + 1, // 正文从下一行开始 start_line: line_num + 1,
end_line: 0, // 暂时标记,后续计算 end_line: 0,
char_start: 0, // 暂时标记,后续计算 char_start: 0,
char_end: 0, char_end: 0,
}); });
index += 1; index += 1;
} }
} }
// 计算每个章节的正文范围
for i in 0..sections.len() { for i in 0..sections.len() {
let next_start_line = if i + 1 < sections.len() { let next_start_line = if i + 1 < sections.len() {
sections[i + 1].start_line - 1 // 下一个标题行 sections[i + 1].start_line - 1
} else { } else {
content.lines().count() content.lines().count()
}; };
sections[i].end_line = next_start_line; sections[i].end_line = next_start_line;
// 计算字节偏移量
let (cs, ce) = find_line_byte_range(content, sections[i].start_line, sections[i].end_line); let (cs, ce) = find_line_byte_range(content, sections[i].start_line, sections[i].end_line);
sections[i].char_start = cs; sections[i].char_start = cs;
sections[i].char_end = ce; sections[i].char_end = ce;
@ -75,11 +193,8 @@ pub fn extract_outline(content: &str, max_level: usize) -> Vec<MarkdownSection>
sections sections
} }
/// 按章节名称(标题)提取内容。大小写不敏感,支持部分匹配。
pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<String> { pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<String> {
let lower = section_name.to_lowercase(); let lower = section_name.to_lowercase();
// 先获取大纲,再匹配
// 直接用 HEADING_RE 扫描,找到第一个匹配的标题
let mut found_start: Option<usize> = None; let mut found_start: Option<usize> = None;
let mut found_level: Option<usize> = None; let mut found_level: Option<usize> = None;
let mut found_line: Option<usize> = None; let mut found_line: Option<usize> = None;
@ -91,7 +206,7 @@ pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<Stri
let heading = caps.get(2).unwrap().as_str().trim(); let heading = caps.get(2).unwrap().as_str().trim();
if heading.to_lowercase().contains(&lower) { if heading.to_lowercase().contains(&lower) {
found_start = Some(line_num + 1); // 正文从下一行开始 found_start = Some(line_num + 1);
found_level = Some(level); found_level = Some(level);
found_line = Some(line_num); found_line = Some(line_num);
break; break;
@ -103,7 +218,6 @@ pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<Stri
let section_level = found_level?; let section_level = found_level?;
let section_line = found_line?; let section_line = found_line?;
// 找到下一个同级或更高级标题作为结束
let total_lines = content.lines().count(); let total_lines = content.lines().count();
let mut end_line = total_lines; let mut end_line = total_lines;
for (line_num, line) in content.lines().enumerate() { for (line_num, line) in content.lines().enumerate() {
@ -125,20 +239,18 @@ pub fn extract_section_by_name(content: &str, section_name: &str) -> Option<Stri
Some(section_content.trim().to_string()) Some(section_content.trim().to_string())
} }
/// 按序号提取章节内容0-based
pub fn extract_section_by_index(content: &str, index: usize) -> Option<String> { pub fn extract_section_by_index(content: &str, index: usize) -> Option<String> {
let sections = extract_outline(content, 10); // max_level 足够大,捕获所有层级标题 let sections = extract_outline(content, 10);
let section = sections.get(index)?; let section = sections.get(index)?;
let section_content = &content[section.char_start..section.char_end]; let section_content = &content[section.char_start..section.char_end];
Some(section_content.trim().to_string()) Some(section_content.trim().to_string())
} }
/// 根据行号范围计算字节偏移量
fn find_line_byte_range(content: &str, start_line: usize, end_line: usize) -> (usize, usize) { fn find_line_byte_range(content: &str, start_line: usize, end_line: usize) -> (usize, usize) {
let char_start = content let char_start = content
.lines() .lines()
.take(start_line) .take(start_line)
.map(|l| l.len() + 1) // +1 for newline .map(|l| l.len() + 1)
.sum::<usize>(); .sum::<usize>();
let char_end = if end_line >= content.lines().count() { let char_end = if end_line >= content.lines().count() {
@ -158,6 +270,16 @@ fn find_line_byte_range(content: &str, start_line: usize, end_line: usize) -> (u
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn test_read_mode_serialize() {
let mode = ReadMode::Full {
include_translation: true,
};
let serialized = serde_json::to_string(&mode).unwrap();
assert!(serialized.contains("Full"));
assert!(serialized.contains("include_translation"));
}
#[test] #[test]
fn test_extract_outline_basic() { fn test_extract_outline_basic() {
let md = "\ let md = "\
@ -283,7 +405,6 @@ More text.
## Results ## Results
Results text. Results text.
"; ";
// "Methods" section should include both subsections but stop at "Results"
let content = extract_section_by_name(md, "Methods").unwrap(); let content = extract_section_by_name(md, "Methods").unwrap();
assert!(content.contains("Subsection text")); assert!(content.contains("Subsection text"));
assert!(content.contains("More text")); assert!(content.contains("More text"));

View File

@ -6,7 +6,7 @@ mod generic;
mod iop; mod iop;
pub mod pdf; pub mod pdf;
use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db}; use crate::services::paper::{check_paper_paths_in_db, get_paper_from_db};
use std::path::Path; use std::path::Path;
use tracing::{error, info, warn}; use tracing::{error, info, warn};
@ -51,6 +51,24 @@ fn detect_parser(html: &str) -> Box<dyn JournalParser> {
// 将 PDF/MinerU 函数重新导出到 parser 模块层级 // 将 PDF/MinerU 函数重新导出到 parser 模块层级
pub use pdf::{parse_pdf_via_mineru, poll_and_extract_mineru, submit_pdf_to_mineru}; pub use pdf::{parse_pdf_via_mineru, poll_and_extract_mineru, submit_pdf_to_mineru};
/// 为解析后的 Markdown 构建 YAML 前置元数据front matter
fn build_front_matter(paper: &crate::services::paper::StandardPaper) -> String {
let title_json =
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title));
let author_list = paper
.authors
.iter()
.map(|a| format!("\"{}\"", a))
.collect::<Vec<_>>()
.join(", ");
let pub_json = serde_json::to_string(&paper.pub_journal)
.unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal));
format!(
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
title_json, author_list, pub_json, paper.bibcode, paper.year, paper.keywords.join(",")
)
}
/// HTML → Markdown 核心解析管线 /// HTML → Markdown 核心解析管线
pub fn html_to_markdown(html_path: &Path) -> anyhow::Result<String> { pub fn html_to_markdown(html_path: &Path) -> anyhow::Result<String> {
info!("正在解析 HTML → Markdown: {:?}", html_path); info!("正在解析 HTML → Markdown: {:?}", html_path);
@ -161,22 +179,14 @@ pub async fn parse_paper_service(
tokio::task::spawn_blocking(move || html_to_markdown(&html_abs_for_blocking)).await; tokio::task::spawn_blocking(move || html_to_markdown(&html_abs_for_blocking)).await;
match md_result { match md_result {
Ok(Ok(md)) => { Ok(Ok(md)) => {
let front_matter = format!( let front_matter = build_front_matter(&paper);
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title)),
paper.authors.iter().map(|a| format!("\"{}\"", a)).collect::<Vec<_>>().join(", "),
serde_json::to_string(&paper.pub_journal).unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal)),
paper.bibcode,
paper.year,
paper.keywords.join(",")
);
parsed_markdown = format!("{}{}", front_matter, md); parsed_markdown = format!("{}{}", front_matter, md);
let md_filename = format!("{}.md", bibcode); let md_filename = format!("{}.md", bibcode);
let md_dest = library_dir.join("Markdown").join(&md_filename); let md_dest = library_dir.join("Markdown").join(&md_filename);
if let Some(parent) = md_dest.parent() { if let Some(parent) = md_dest.parent() {
std::fs::create_dir_all(parent).unwrap_or_default(); tokio::fs::create_dir_all(parent).await.unwrap_or_default();
} }
if std::fs::write(&md_dest, &parsed_markdown).is_ok() { if tokio::fs::write(&md_dest, &parsed_markdown).await.is_ok() {
relative_md_path = format!("Markdown/{}", md_filename); relative_md_path = format!("Markdown/{}", md_filename);
} }
} }
@ -203,15 +213,7 @@ pub async fn parse_paper_service(
if pdf_abs.exists() { if pdf_abs.exists() {
match parse_pdf_via_mineru(&pdf_abs, qiniu, config).await { match parse_pdf_via_mineru(&pdf_abs, qiniu, config).await {
Ok(md) => { Ok(md) => {
let front_matter = format!( let front_matter = build_front_matter(&paper);
"---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"https://ui.adsabs.harvard.edu/abs/{}/abstract\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n",
serde_json::to_string(&paper.title).unwrap_or_else(|_| format!("\"{}\"", paper.title)),
paper.authors.iter().map(|a| format!("\"{}\"", a)).collect::<Vec<_>>().join(", "),
serde_json::to_string(&paper.pub_journal).unwrap_or_else(|_| format!("\"{}\"", paper.pub_journal)),
paper.bibcode,
paper.year,
paper.keywords.join(",")
);
parsed_markdown = format!("{}{}", front_matter, md); parsed_markdown = format!("{}{}", front_matter, md);
let md_filename = format!("{}.md", bibcode); let md_filename = format!("{}.md", bibcode);
let md_dest = library_dir.join("Markdown").join(&md_filename); let md_dest = library_dir.join("Markdown").join(&md_filename);
@ -258,9 +260,8 @@ pub async fn parse_paper_service(
// ── 测试 ────────────────────────────────────────────────────────── // ── 测试 ──────────────────────────────────────────────────────────
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::common::{convert_html_tables_to_markdown, postprocess_markdown}; use super::common::postprocess_markdown;
use super::html_to_markdown; use super::html_to_markdown;
use std::io::Write;
#[test] #[test]
fn test_convert_html_tables_to_markdown() { fn test_convert_html_tables_to_markdown() {

114
src/services/pipeline.rs Normal file
View File

@ -0,0 +1,114 @@
// src/services/pipeline.rs
//
// 文献处理流水线服务:下载 → 解析 → 向量化。
use crate::api::AppState;
use tracing::info;
#[derive(Debug, Clone)]
pub struct PipelineResult {
pub bibcode: String,
pub completed: Vec<String>,
pub has_pdf: bool,
pub has_html: bool,
pub has_markdown: bool,
pub chunk_count: usize,
}
/// 执行文献处理流水线。按序执行指定的任务步骤("download", "parse", "embed")。
pub async fn process_paper_pipeline(
state: &AppState,
bibcode: &str,
tasks: &[String],
force: bool,
) -> anyhow::Result<PipelineResult> {
let mut completed = Vec::new();
let mut has_pdf = false;
let mut has_html = false;
let mut has_markdown = false;
let mut chunk_count = 0;
for task in tasks {
match task.as_str() {
"download" => {
info!("[Pipeline] download: {}", bibcode);
let paper = state
.downloader
.download_paper_service(&state.db, &state.config.library_dir, bibcode, force)
.await?;
has_pdf = paper.has_pdf;
has_html = paper.has_html;
completed.push("download".to_string());
}
"parse" => {
info!("[Pipeline] parse: {}", bibcode);
let markdown = crate::services::parser::parse_paper_service(
&state.db,
&state.config.library_dir,
&state.qiniu,
&state.config,
bibcode,
force,
)
.await?;
let markdown_len = markdown.len();
drop(markdown); // 释放大字符串内存,仅保留长度
has_markdown = true;
completed.push("parse".to_string());
info!("[Pipeline] parse done: {} chars", markdown_len);
}
"embed" => {
info!("[Pipeline] embed: {}", bibcode);
let paths = crate::services::paper::check_paper_paths_in_db(
&state.db,
&state.config.library_dir,
bibcode,
)
.await?;
let md_rel = match paths {
Some((_, _, Some(md), _)) => md,
_ => {
return Err(anyhow::anyhow!(
"未找到已解析的 Markdown 文件,请确保 tasks 中 parse 在 embed 之前执行"
))
}
};
let md_abs = state.config.library_dir.join(&md_rel);
let markdown_text = tokio::fs::read_to_string(&md_abs).await?;
let count = crate::services::rag::ingest_paper(
&state.db,
&state.embedding,
bibcode,
&markdown_text,
None, // 使用默认 ChunkerConfig
)
.await?;
chunk_count = count;
completed.push("embed".to_string());
info!("[Pipeline] embed done: {} chunks", count);
}
other => {
return Err(anyhow::anyhow!(
"无效的任务类型 '{}',仅支持: download, parse, embed",
other
));
}
}
}
Ok(PipelineResult {
bibcode: bibcode.to_string(),
completed,
has_pdf,
has_html,
has_markdown,
chunk_count,
})
}

View File

@ -1,15 +1,15 @@
// src/services/rag.rs // src/services/rag.rs
// //
// RAG (Retrieval-Augmented Generation) 问答编排模块。 // RAG (Retrieval-Augmented Generation) 检索模块。
// 职责: // 职责:
// 1. Ingest: 将文献 Markdown 切片 -> 向量化 -> 写入 SQLite vec0 + content 双表 // 1. Ingest: 将文献 Markdown 切片 -> 向量化 -> 写入 SQLite vec0 + content 双表
// 2. Retrieve: 向量化用户提问 -> 相似度检索 -> 返回 Top-K 片段 // 2. Retrieve: 向量化用户提问 -> 相似度检索 -> 返回 Top-K 片段
// 3. Complete: 组装检索上下文 -> 调用 LLM 生成最终回答 // 3. Hybrid: 稠密向量 + 稀疏 BM25 双路融合检索 (RRF)
use sqlx::SqlitePool; use sqlx::SqlitePool;
use tracing::{error, info, warn}; use tracing::{error, info, warn};
use crate::clients::llm::{ChatCompleter, Embedder}; use crate::clients::llm::Embedder;
use crate::services::chunker::{chunk_markdown, ChunkerConfig, StructuralChunk}; use crate::services::chunker::{chunk_markdown, ChunkerConfig, StructuralChunk};
/// RAG 检索结果条目 /// RAG 检索结果条目
@ -29,17 +29,8 @@ pub struct RetrievalResult {
pub distance: f64, pub distance: f64,
} }
/// RAG 问答最终结果 /// RAG 检索结果条目
#[derive(Debug, Clone, serde::Serialize)] ///
pub struct RagAnswer {
/// LLM 生成的回答
pub answer: String,
/// 检索到的参考来源
pub sources: Vec<RetrievalResult>,
}
// ─────────────────────── Ingest ───────────────────────
/// 将一篇文献的 Markdown 内容切片、向量化并写入数据库 /// 将一篇文献的 Markdown 内容切片、向量化并写入数据库
/// ///
/// 流程: /// 流程:
@ -68,27 +59,33 @@ pub async fn ingest_paper(
chunks.len() chunks.len()
); );
// 先清除该文献的旧切片(重新解析场景) // 先清除该文献的旧切片(重新解析场景)。
// 合并为两条原子 SQL 语句,避免循环查询和删除 // 在事务内原子执行两条 DELETE防止并发读取观测到 vec0 和 content 表的不一致状态。
sqlx::query( {
"DELETE FROM vec_paper_chunks WHERE rowid IN (SELECT rowid FROM paper_chunks_content WHERE bibcode = ?)" let mut tx = pool.begin().await?;
) sqlx::query(
.bind(bibcode) "DELETE FROM vec_paper_chunks WHERE rowid IN (SELECT rowid FROM paper_chunks_content WHERE bibcode = ?)"
.execute(pool) )
.await?;
sqlx::query("DELETE FROM paper_chunks_content WHERE bibcode = ?")
.bind(bibcode) .bind(bibcode)
.execute(pool) .execute(&mut *tx)
.await?; .await?;
sqlx::query("DELETE FROM paper_chunks_content WHERE bibcode = ?")
.bind(bibcode)
.execute(&mut *tx)
.await?;
tx.commit().await?;
}
// ── 事务外批量计算所有 embedding ── // ── 事务外批量计算所有 embedding ──
// 避免长事务跨越数十次网络往返SQLite 写锁阻塞)。 // 避免长事务跨越数十次网络往返SQLite 写锁阻塞)。
let contents: Vec<String> = chunks.iter().map(|c| c.content.clone()).collect(); let contents: Vec<String> = chunks.iter().map(|c| c.content.clone()).collect();
let embeddings: Vec<Option<Vec<f32>>> = match embedding_client.embed_batch(&contents).await { let embeddings: Vec<Option<Vec<f32>>> = match embedding_client.embed_batch(&contents).await {
Ok(all) => all.into_iter().map(Some).collect(), Ok(all) => all.into_iter().map(Some).collect(),
Err(e) => { Err(e) => {
// 批量失败时降级为逐条计算,保留"单块失败跳过、其余继续"的容错语义 // 批量失败时降级为逐条计算,保留"单块失败跳过、其余继续"的容错语义。
// TODO: 后续将 Embedder trait 改为 Arc<dyn Embedder> 以支持并发降级semaphore=3
warn!("文献 {} 批量向量化失败,降级逐条计算: {}", bibcode, e); warn!("文献 {} 批量向量化失败,降级逐条计算: {}", bibcode, e);
let mut out = Vec::with_capacity(chunks.len()); let mut out = Vec::with_capacity(chunks.len());
for chunk in &chunks { for chunk in &chunks {
@ -210,8 +207,8 @@ async fn retrieve_sparse(
query: &str, query: &str,
top_k: usize, top_k: usize,
) -> anyhow::Result<Vec<(String, i64, String, usize, String, i64)>> { ) -> anyhow::Result<Vec<(String, i64, String, usize, String, i64)>> {
// 清洗查询:转义 FTS5 特殊字符,构建 MATCH 字符串 // 清洗查询:使用统一的 FTS5 净化函数,防止操作符注入
let sanitized = query.replace('"', "\"\"").replace('\'', "''"); let sanitized = crate::services::search::sanitize_fts5_query(query);
let fts_query = format!("\"{}\"", sanitized); let fts_query = format!("\"{}\"", sanitized);
let rows: Vec<(String, i64, String, String, i64)> = sqlx::query_as( let rows: Vec<(String, i64, String, String, i64)> = sqlx::query_as(
@ -359,66 +356,6 @@ pub async fn retrieve_hybrid(
Ok(scored.into_iter().take(top_k).map(|(r, _)| r).collect()) Ok(scored.into_iter().take(top_k).map(|(r, _)| r).collect())
} }
// ─────────────────────── Complete ───────────────────────
/// 执行完整的 RAG 问答流程:检索 + LLM 生成
pub async fn ask(
pool: &SqlitePool,
embedding_client: &dyn Embedder,
llm_client: &dyn ChatCompleter,
question: &str,
top_k: usize,
) -> anyhow::Result<RagAnswer> {
// 1. 检索相关片段
let sources = retrieve(pool, embedding_client, question, top_k).await?;
if sources.is_empty() {
return Ok(RagAnswer {
answer: "未找到与该问题相关的文献内容。请先导入相关文献的 Markdown 版本。".to_string(),
sources: vec![],
});
}
// 2. 组装上下文(含章节路径信息)
let context = sources
.iter()
.enumerate()
.map(|(i, r)| {
let heading_info = if r.headings.is_empty() || r.headings == "Document" {
String::new()
} else {
format!(" (章节: {})", r.headings)
};
format!(
"[来源 {}: {} §{}{}]\n{}",
i + 1,
r.bibcode,
r.paragraph_index,
heading_info,
r.content
)
})
.collect::<Vec<_>>()
.join("\n\n---\n\n");
// 3. 构建 System Prompt
let system_prompt = "你是一位专业的天体物理学研究助手。\
\
使 [ N] 便\
\
";
let user_content = format!(
"## 检索到的参考片段\n\n{}\n\n---\n\n## 用户问题\n\n{}",
context, question
);
// 4. 调用 LLM
let answer = llm_client.complete(system_prompt, &user_content).await?;
Ok(RagAnswer { answer, sources })
}
// ─────────────────────── Helpers ─────────────────────── // ─────────────────────── Helpers ───────────────────────
/// 将 f32 向量转换为字节切片(小端序),用于 sqlite-vec 的二进制格式 /// 将 f32 向量转换为字节切片(小端序),用于 sqlite-vec 的二进制格式
@ -430,6 +367,29 @@ fn embedding_to_bytes(embedding: &[f32]) -> Vec<u8> {
bytes bytes
} }
/// 向量化单篇文献服务:读取 Markdown 物理文件、判断其存在,并向量化切片入库
pub async fn embed_paper_service(
db: &SqlitePool,
library_dir: &std::path::Path,
bibcode: &str,
embedding_client: &dyn Embedder,
) -> anyhow::Result<usize> {
let (_, _, md_opt, _) =
crate::services::paper::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
let md_rel = md_opt.ok_or_else(|| anyhow::anyhow!("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!("文献 Markdown 文件未找到,请重新解析");
}
let markdown_content = tokio::fs::read_to_string(&md_abs).await?;
let count = ingest_paper(db, embedding_client, bibcode, &markdown_content, None).await?;
Ok(count)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;

View File

@ -1,13 +1,11 @@
// src/services/search.rs // src/services/search.rs
use std::collections::HashSet; use std::collections::HashSet;
use crate::api::helpers::{ use crate::api::AppState;
convert_ads_doc_to_standard, convert_arxiv_to_standard, get_paper_from_db, save_paper_to_db, use crate::services::paper::{
convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db, StandardPaper,
SQLITE_PARAM_LIMIT, SQLITE_PARAM_LIMIT,
}; };
use crate::api::AppState;
use crate::api::StandardPaper;
use sqlx::Row;
use tracing::{error, warn}; use tracing::{error, warn};
/// 统一检索逻辑,合并去重 ADS 和 arXiv 数据,并自动记录引用拓扑与同步本地状态。 /// 统一检索逻辑,合并去重 ADS 和 arXiv 数据,并自动记录引用拓扑与同步本地状态。
@ -96,18 +94,14 @@ pub async fn search_papers(
} }
// 5. 批量回填本地库的标准元数据is_downloaded/has_markdown/pdf_error 等),避免 N+1。 // 5. 批量回填本地库的标准元数据is_downloaded/has_markdown/pdf_error 等),避免 N+1。
// get_paper_from_db 内部还做了 library_dir 文件存在性检查,故仍逐条调用, // 用一次批量 IN 查询直接返回完整 StandardPaper无需逐条调用 get_paper_from_db。
// 但先用一次 IN 查询过滤出"数据库中确实存在"的 bibcode仅对这些调用 get_paper_from_db。
let all_bibs: Vec<String> = unique_results.iter().map(|p| p.bibcode.clone()).collect(); let all_bibs: Vec<String> = unique_results.iter().map(|p| p.bibcode.clone()).collect();
let existing_bibs: HashSet<String> = batch_fetch_existing_bibcodes(&state.db, &all_bibs).await; let existing_papers =
batch_fetch_existing_papers(&state.db, &state.config.library_dir, &all_bibs).await;
for paper in unique_results.iter_mut() { for paper in unique_results.iter_mut() {
if existing_bibs.contains(&paper.bibcode) { if let Some(db_paper) = existing_papers.get(&paper.bibcode) {
if let Ok(db_paper) = *paper = db_paper.clone();
get_paper_from_db(&state.db, &state.config.library_dir, &paper.bibcode).await
{
*paper = db_paper;
}
} }
} }
@ -137,41 +131,52 @@ async fn save_citation_topology(
return; return;
} }
// 单事务批量 INSERT分批避免超参数上限 // 单一事务批量 INSERT分批仅用于避免 SQLite 参数上限
let mut tx = match db.begin().await {
Ok(tx) => tx,
Err(e) => {
warn!("无法开启引用关系事务: {}", e);
return;
}
};
for chunk in pairs.chunks(SQLITE_PARAM_LIMIT / 2) { for chunk in pairs.chunks(SQLITE_PARAM_LIMIT / 2) {
if let Ok(mut tx) = db.begin().await { for (src, tgt) in chunk {
for (src, tgt) in chunk { if let Err(e) = sqlx::query(
if let Err(e) = sqlx::query( "INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)",
"INSERT OR IGNORE INTO citations_references (source_bibcode, target_bibcode) VALUES (?, ?)", )
) .bind(src)
.bind(src) .bind(tgt)
.bind(tgt) .execute(&mut *tx)
.execute(&mut *tx) .await
.await {
{ warn!("保存引用关系失败 ({} -> {}): {}", src, tgt, e);
warn!("保存引用关系失败 ({} -> {}): {}", src, tgt, e);
}
}
if let Err(e) = tx.commit().await {
warn!("提交引用关系事务失败: {}", e);
} }
} }
} }
if let Err(e) = tx.commit().await {
warn!("提交引用关系事务失败: {}", e);
}
} }
/// 批量查询数据库中已存在的 bibcode 集合WHERE IN 分批),避免 N+1。 /// 批量查询数据库中已存在的文献完整元数据WHERE IN 分批),避免 N+1。
async fn batch_fetch_existing_bibcodes( async fn batch_fetch_existing_papers(
db: &sqlx::SqlitePool, db: &sqlx::SqlitePool,
library_dir: &std::path::Path,
bibcodes: &[String], bibcodes: &[String],
) -> HashSet<String> { ) -> std::collections::HashMap<String, StandardPaper> {
let mut existing = HashSet::new(); use crate::services::paper::parse_paper_row;
let mut papers = std::collections::HashMap::new();
for chunk in bibcodes.chunks(SQLITE_PARAM_LIMIT) { for chunk in bibcodes.chunks(SQLITE_PARAM_LIMIT) {
if chunk.is_empty() { if chunk.is_empty() {
continue; continue;
} }
let placeholders = vec!["?"; chunk.len()].join(", "); let placeholders = vec!["?"; chunk.len()].join(", ");
let sql = format!( let sql = format!(
"SELECT bibcode FROM papers WHERE bibcode IN ({})", "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 bibcode IN ({})",
placeholders placeholders
); );
let mut query = sqlx::query(&sql); let mut query = sqlx::query(&sql);
@ -179,13 +184,28 @@ async fn batch_fetch_existing_bibcodes(
query = query.bind(b); query = query.bind(b);
} }
if let Ok(rows) = query.fetch_all(db).await { if let Ok(rows) = query.fetch_all(db).await {
for row in rows { for row in &rows {
let b: String = row.get(0); let paper = parse_paper_row(row, library_dir);
existing.insert(b); papers.insert(paper.bibcode.clone(), paper);
} }
} }
} }
existing papers
}
/// FTS5 查询净化:转义特殊字符,防止 FTS5 操作符注入。
///
/// FTS5 特殊字符(`"`、`'`、`*`)和结构操作符(`:`、`AND`、`OR`、`NOT`
/// 会被替换为空格,确保查询作为纯文本搜索而非布尔表达式执行。
pub fn sanitize_fts5_query(query: &str) -> String {
query
.replace('"', "\"\"")
.replace('\'', "''")
.replace(['*', ':'], " ")
.replace(" AND ", " ")
.replace(" OR ", " ")
.replace(" NOT ", " ")
.replace(" NEAR ", " ")
} }
/// 本地文献库 FTS5 全文检索。 /// 本地文献库 FTS5 全文检索。
@ -197,8 +217,8 @@ pub async fn search_local_library(
query: &str, query: &str,
limit: usize, limit: usize,
) -> anyhow::Result<Vec<StandardPaper>> { ) -> anyhow::Result<Vec<StandardPaper>> {
// FTS5 查询需要转义特殊字符 // FTS5 查询净化,防止操作符注入
let sanitized = query.replace('"', "\"\"").replace('\'', "''"); let sanitized = sanitize_fts5_query(query);
let fts_query = format!("\"{}\"", sanitized); let fts_query = format!("\"{}\"", sanitized);
let rows = sqlx::query_as::<_, PaperRow>( let rows = sqlx::query_as::<_, PaperRow>(
@ -266,3 +286,100 @@ struct PaperRow {
markdown_path: Option<String>, markdown_path: Option<String>,
translation_path: Option<String>, translation_path: Option<String>,
} }
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SearchResult {
pub result_type: String,
pub session_id: String,
pub title: Option<String>,
pub snippet: String,
pub created_at: Option<String>,
}
/// 跨会话/消息的智能体历史全文检索。
pub async fn search_agent_history(
db: &sqlx::SqlitePool,
q: &str,
scope: &str,
limit: i64,
session_id: Option<&str>,
) -> Result<Vec<SearchResult>, sqlx::Error> {
let mut results = Vec::new();
let sanitized = sanitize_fts5_query(q);
// 搜索会话
if scope == "all" || scope == "sessions" {
let rows = sqlx::query_as::<_, (String, String, String, Option<String>)>(
"SELECT s.session_id, s.title, \
snippet(agent_sessions_fts, 1, '<mark>', '</mark>', '...', 40) as snippet, \
s.created_at \
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 $2",
)
.bind(&sanitized)
.bind(limit)
.fetch_all(db)
.await?;
for (sid, title, snippet, created_at) in rows {
results.push(SearchResult {
result_type: "session".into(),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
// 搜索消息
if scope == "all" || scope == "messages" {
let msg_rows = if let Some(sid) = session_id {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
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 fts.session_id = $3 \
AND m.active = 1 \
ORDER BY rank LIMIT $2",
)
.bind(&sanitized)
.bind(limit)
.bind(sid)
.fetch_all(db)
.await?
} else {
sqlx::query_as::<_, (String, String, String, String, Option<String>)>(
"SELECT fts.session_id, s.title, \
snippet(agent_messages_fts, 2, '<mark>', '</mark>', '...', 60) as snippet, \
fts.role, m.created_at \
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 $2",
)
.bind(&sanitized)
.bind(limit)
.fetch_all(db)
.await?
};
for (sid, title, snippet, role, created_at) in msg_rows {
results.push(SearchResult {
result_type: format!("message/{}", role),
session_id: sid,
title: Some(title),
snippet,
created_at,
});
}
}
Ok(results)
}

290
src/services/session.rs Normal file
View File

@ -0,0 +1,290 @@
// src/services/session.rs
//
// 智能体运行会话管理与指标统计服务层
use serde::{Deserialize, Serialize};
use sqlx::{Row, SqlitePool};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionSummary {
pub session_id: String,
pub title: String,
pub model: String,
pub mode: String,
pub turn_count: i32,
pub summary: Option<String>,
pub created_at: String,
pub updated_at: String,
}
/// 查询会话列表,支持分页
pub async fn list_sessions_service(
db: &SqlitePool,
limit: i64,
offset: i64,
) -> Result<Vec<SessionSummary>, sqlx::Error> {
let rows = sqlx::query(
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
FROM agent_sessions \
WHERE deleted_at IS NULL \
ORDER BY updated_at DESC \
LIMIT ? OFFSET ?",
)
.bind(limit)
.bind(offset)
.fetch_all(db)
.await?;
let sessions: Vec<SessionSummary> = rows
.iter()
.map(|r| SessionSummary {
session_id: r.get("session_id"),
title: r.get("title"),
model: r.get("model"),
mode: r.get("mode"),
turn_count: r.get("turn_count"),
summary: r.get("summary"),
created_at: r.get("created_at"),
updated_at: r.get("updated_at"),
})
.collect();
Ok(sessions)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageRecord {
pub id: i64,
pub agent_name: String,
pub turn_index: i32,
pub step_index: i32,
pub role: String,
pub content: String,
pub thought: Option<String>,
pub tool_calls: Option<serde_json::Value>,
pub tool_call_id: Option<String>,
pub token_count: i32,
pub metadata: Option<serde_json::Value>,
pub created_at: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionDetail {
pub session: SessionSummary,
pub messages: Vec<MessageRecord>,
}
/// 获取单个会话的详细信息(会话信息与激活的消息历史)
pub async fn get_session_detail_service(
db: &SqlitePool,
session_id: &str,
) -> Result<Option<SessionDetail>, sqlx::Error> {
// 查询会话元信息
let session_row = match sqlx::query(
"SELECT session_id, title, model, mode, turn_count, summary, created_at, updated_at \
FROM agent_sessions \
WHERE session_id = ? AND deleted_at IS NULL",
)
.bind(session_id)
.fetch_optional(db)
.await?
{
Some(row) => row,
None => return Ok(None),
};
let session = SessionSummary {
session_id: session_row.get("session_id"),
title: session_row.get("title"),
model: session_row.get("model"),
mode: session_row.get("mode"),
turn_count: session_row.get("turn_count"),
summary: session_row.get("summary"),
created_at: session_row.get("created_at"),
updated_at: session_row.get("updated_at"),
};
// 查询消息列表(限制 500 条,防止会话消息过多导致内存膨胀)
const MAX_MESSAGES: i64 = 500;
let msg_rows = sqlx::query(
"SELECT id, agent_name, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, metadata, created_at \
FROM agent_messages \
WHERE session_id = ? AND active = 1 \
ORDER BY id ASC \
LIMIT ?"
)
.bind(session_id)
.bind(MAX_MESSAGES)
.fetch_all(db)
.await?;
let messages: Vec<MessageRecord> = msg_rows
.iter()
.map(|r| {
let tool_calls_json: Option<String> = r.get("tool_calls");
let metadata_json: Option<String> = r.get("metadata");
MessageRecord {
id: r.get("id"),
agent_name: r.get("agent_name"),
turn_index: r.get("turn_index"),
step_index: r.get("step_index"),
role: r.get("role"),
content: r.get("content"),
thought: r.get("thought"),
tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()),
tool_call_id: r.get("tool_call_id"),
token_count: r.get("token_count"),
metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
created_at: r.get("created_at"),
}
})
.collect();
Ok(Some(SessionDetail { session, messages }))
}
/// 软删除指定会话
pub async fn delete_session_service(
db: &SqlitePool,
session_id: &str,
) -> Result<bool, sqlx::Error> {
let result = sqlx::query(
"UPDATE agent_sessions SET deleted_at = CURRENT_TIMESTAMP WHERE session_id = ? AND deleted_at IS NULL"
)
.bind(session_id)
.execute(db)
.await?;
Ok(result.rows_affected() > 0)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMetrics {
pub total_sessions: i64,
pub total_tool_calls: i64,
pub tool_call_breakdown: serde_json::Value,
pub avg_steps_per_session: f64,
pub error_rate: f64,
}
/// 获取全局运行监控与性能分析指标。
/// 使用 `try_join!` 并行执行独立查询,减少串行等待。
pub async fn get_agent_metrics_service(db: &SqlitePool) -> Result<AgentMetrics, sqlx::Error> {
let sessions_fut = async {
let total: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE deleted_at IS NULL")
.fetch_one(db)
.await
.unwrap_or(0);
let avg: f64 = sqlx::query_scalar(
"SELECT COALESCE(AVG(CAST(turn_count AS REAL)), 0.0) \
FROM agent_sessions WHERE deleted_at IS NULL",
)
.fetch_one(db)
.await
.unwrap_or(0.0);
(total, avg)
};
let audit_fut = async {
let total: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM agent_audit_log WHERE status IN ('OK', 'FAIL')",
)
.fetch_one(db)
.await
.unwrap_or(0);
let errors: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status = 'FAIL'")
.fetch_one(db)
.await
.unwrap_or(0);
(total, errors)
};
let breakdown_fut = async {
sqlx::query_as::<_, (String, i64)>(
"SELECT COALESCE(tool_name, 'unknown'), COUNT(*) as cnt \
FROM agent_audit_log \
WHERE status IN ('OK', 'FAIL') \
GROUP BY tool_name \
ORDER BY cnt DESC",
)
.fetch_all(db)
.await
.unwrap_or_default()
};
let ((total_sessions, avg_steps), (total_tool_calls, total_errors), breakdown_rows) =
tokio::join!(sessions_fut, audit_fut, breakdown_fut);
let tool_call_breakdown: serde_json::Value = breakdown_rows
.iter()
.map(|(name, cnt): &(String, i64)| serde_json::json!({ name: cnt }))
.fold(serde_json::json!({}), |mut acc, v| {
if let serde_json::Value::Object(map) = &mut acc {
if let serde_json::Value::Object(v_map) = v {
for (k, val) in v_map {
map.insert(k.clone(), val.clone());
}
}
}
acc
});
let error_rate = if total_tool_calls > 0 {
total_errors as f64 / total_tool_calls as f64
} else {
0.0
};
Ok(AgentMetrics {
total_sessions,
total_tool_calls,
tool_call_breakdown,
avg_steps_per_session: avg_steps,
error_rate,
})
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditLogEntry {
pub id: i64,
pub step: i32,
pub tool_name: Option<String>,
pub status: String,
pub elapsed_ms: i32,
pub output_preview: Option<String>,
pub created_at: String,
}
/// 查询单个会话的全部操作审计日志
pub async fn get_session_audit_service(
db: &SqlitePool,
session_id: &str,
) -> Result<Vec<AuditLogEntry>, sqlx::Error> {
let rows = sqlx::query(
"SELECT id, step, tool_name, status, elapsed_ms, output_preview, created_at \
FROM agent_audit_log \
WHERE session_id = ? \
ORDER BY id ASC",
)
.bind(session_id)
.fetch_all(db)
.await?;
let entries: Vec<AuditLogEntry> = rows
.iter()
.map(|r| AuditLogEntry {
id: r.get("id"),
step: r.get("step"),
tool_name: r.get("tool_name"),
status: r.get("status"),
elapsed_ms: r.get("elapsed_ms"),
output_preview: r.get("output_preview"),
created_at: r.get("created_at"),
})
.collect();
Ok(entries)
}

View File

@ -31,6 +31,52 @@ pub struct TargetInfo {
pub aliases: Vec<String>, pub aliases: Vec<String>,
} }
/// paper_targets 表的查询行类型,避免 14 字段元组重复定义
#[derive(Debug, sqlx::FromRow)]
struct TargetDbRow {
target_name: String,
ra: Option<String>,
dec: Option<String>,
parallax: Option<f64>,
parallax_err: Option<f64>,
spectral_type: Option<String>,
v_magnitude: Option<f64>,
otype: Option<String>,
oname: Option<String>,
pm_ra: Option<f64>,
pm_de: Option<f64>,
radial_velocity: Option<f64>,
photometry: Option<String>,
aliases: Option<String>,
}
impl From<TargetDbRow> for TargetInfo {
fn from(r: TargetDbRow) -> Self {
let aliases: Vec<String> = r
.aliases
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
r.photometry.and_then(|s| serde_json::from_str(&s).ok());
TargetInfo {
target_name: r.target_name,
ra: r.ra,
dec: r.dec,
parallax: r.parallax,
parallax_err: r.parallax_err,
spectral_type: r.spectral_type,
v_magnitude: r.v_magnitude,
otype: r.otype,
oname: r.oname,
pm_ra: r.pm_ra,
pm_de: r.pm_de,
radial_velocity: r.radial_velocity,
photometry,
aliases,
}
}
}
/// 从文本中提取可能的天体标识符 /// 从文本中提取可能的天体标识符
/// ///
/// 支持的星表命名规范(基于 IAU 标准): /// 支持的星表命名规范(基于 IAU 标准):
@ -202,8 +248,9 @@ fn parse_sesame_xml(xml: &str, original_name: &str) -> anyhow::Result<TargetInfo
// 提取别名列表 // 提取别名列表
let mut aliases = Vec::new(); let mut aliases = Vec::new();
let alias_re = Regex::new(r"<alias>([^<]+)</alias>").unwrap(); static ALIAS_RE: std::sync::LazyLock<Regex> =
for cap in alias_re.captures_iter(xml) { std::sync::LazyLock::new(|| Regex::new(r"<alias>([^<]+)</alias>").unwrap());
for cap in ALIAS_RE.captures_iter(xml) {
if let Some(alias) = cap.get(1) { if let Some(alias) = cap.get(1) {
aliases.push(alias.as_str().trim().to_string()); aliases.push(alias.as_str().trim().to_string());
} }
@ -211,7 +258,10 @@ fn parse_sesame_xml(xml: &str, original_name: &str) -> anyhow::Result<TargetInfo
// 如果连坐标都查不到,说明 Sesame 无法识别这个天体 // 如果连坐标都查不到,说明 Sesame 无法识别这个天体
if ra.is_none() && dec.is_none() && aliases.is_empty() { if ra.is_none() && dec.is_none() && aliases.is_empty() {
warn!("Sesame 未能识别天体: {}", original_name); return Err(anyhow::anyhow!(
"CDS Sesame 无法识别天体: {}",
original_name
));
} }
Ok(TargetInfo { Ok(TargetInfo {
@ -290,25 +340,7 @@ pub async fn query_target_cached(
client: &Client, client: &Client,
) -> anyhow::Result<TargetInfo> { ) -> anyhow::Result<TargetInfo> {
// 1. 查本地缓存 // 1. 查本地缓存
let row = sqlx::query_as::< let row = sqlx::query_as::<_, TargetDbRow>(
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \ "SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \ otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE target_name = ? LIMIT 1", FROM paper_targets WHERE target_name = ? LIMIT 1",
@ -317,45 +349,9 @@ pub async fn query_target_cached(
.fetch_optional(pool) .fetch_optional(pool)
.await?; .await?;
if let Some(( if let Some(db_row) = row {
name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
)) = row
{
info!("天体 {} 命中本地缓存", target_name); info!("天体 {} 命中本地缓存", target_name);
let aliases: Vec<String> = aliases_json return Ok(TargetInfo::from(db_row));
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
return Ok(TargetInfo {
target_name: name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
});
} }
// 2. 缓存 Miss调用 Sesame // 2. 缓存 Miss调用 Sesame
@ -418,25 +414,7 @@ pub async fn extract_and_cache_targets(
for name in names { for name in names {
// 优先在 paper_targets 表中查找本地缓存 // 优先在 paper_targets 表中查找本地缓存
let row = sqlx::query_as::< let row = sqlx::query_as::<_, TargetDbRow>(
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \ "SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \ otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE target_name = ? LIMIT 1", FROM paper_targets WHERE target_name = ? LIMIT 1",
@ -445,45 +423,9 @@ pub async fn extract_and_cache_targets(
.fetch_optional(pool) .fetch_optional(pool)
.await; .await;
if let Ok(Some(( if let Ok(Some(db_row)) = row {
tname,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
))) = row
{
info!("天体 {} 命中本地缓存", name); info!("天体 {} 命中本地缓存", name);
let aliases: Vec<String> = aliases_json results.push(TargetInfo::from(db_row));
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
results.push(TargetInfo {
target_name: tname,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
});
} else { } else {
// 本地缓存未命中,释放连接并外部查询 Sesame // 本地缓存未命中,释放连接并外部查询 Sesame
info!("天体 {} 本地缓存未命中,查询 CDS Sesame...", name); info!("天体 {} 本地缓存未命中,查询 CDS Sesame...", name);
@ -544,6 +486,121 @@ pub async fn extract_and_cache_targets(
results results
} }
/// 获取指定文献关联的所有天体信息
pub async fn list_targets_for_paper(
db: &SqlitePool,
bibcode: &str,
) -> Result<Vec<TargetInfo>, sqlx::Error> {
let rows = sqlx::query_as::<
_,
(
String,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<String>,
Option<f64>,
Option<String>,
Option<String>,
Option<f64>,
Option<f64>,
Option<f64>,
Option<String>,
Option<String>,
),
>(
"SELECT target_name, ra, dec, parallax, parallax_err, spectral_type, v_magnitude, \
otype, oname, pm_ra, pm_de, radial_velocity, photometry, aliases \
FROM paper_targets WHERE bibcode = ? ORDER BY target_name",
)
.bind(bibcode)
.fetch_all(db)
.await?;
let targets: Vec<TargetInfo> = rows
.into_iter()
.map(
|(
name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry_json,
aliases_json,
)| {
let aliases: Vec<String> = aliases_json
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
let photometry: Option<std::collections::HashMap<String, f64>> =
photometry_json.and_then(|s| serde_json::from_str(&s).ok());
TargetInfo {
target_name: name,
ra,
dec,
parallax,
parallax_err,
spectral_type,
v_magnitude,
otype,
oname,
pm_ra,
pm_de,
radial_velocity,
photometry,
aliases,
}
},
)
.collect();
Ok(targets)
}
/// 提取文献中的所有天体,并刷新关联记录缓存
pub async fn extract_and_refresh_targets(
db: &SqlitePool,
library_dir: &std::path::Path,
bibcode: &str,
client: &reqwest::Client,
) -> anyhow::Result<Vec<TargetInfo>> {
let (_, _, md_opt, _) =
crate::services::paper::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
let md_rel = md_opt.ok_or_else(|| anyhow::anyhow!("文献尚未解析为 Markdown请先执行解析"))?;
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!("文献 Markdown 文件未找到,请重新解析");
}
let markdown_content = tokio::fs::read_to_string(&md_abs).await?;
// 重新识别前先清除该文献已有的天体关联记录。
// 使用显式事务确保 DELETE 原子执行,避免并发竞态导致的数据不一致。
{
let mut tx = db.begin().await?;
sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?")
.bind(bibcode)
.execute(&mut *tx)
.await?;
tx.commit().await?;
}
let targets = extract_and_cache_targets(db, &markdown_content, bibcode, client).await;
Ok(targets)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;

View File

@ -198,9 +198,9 @@ pub async fn persist_translation(
let tr_filename = format!("{}.md", bibcode); let tr_filename = format!("{}.md", bibcode);
let tr_dest = library_dir.join("Translation").join(&tr_filename); let tr_dest = library_dir.join("Translation").join(&tr_filename);
if let Some(parent) = tr_dest.parent() { if let Some(parent) = tr_dest.parent() {
std::fs::create_dir_all(parent)?; tokio::fs::create_dir_all(parent).await?;
} }
std::fs::write(&tr_dest, translated_markdown)?; tokio::fs::write(&tr_dest, translated_markdown).await?;
let relative_tr_path = format!("Translation/{}", tr_filename); let relative_tr_path = format!("Translation/{}", tr_filename);
sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?")
@ -270,6 +270,55 @@ pub async fn translate_markdown(
} }
} }
/// 中英双栏对比翻译,带物理文件缓存。若缓存存在且不强制重译,则直读本地;否则调用大语言模型进行翻译并写入缓存。
pub async fn translate_paper_cached(
db: &sqlx::SqlitePool,
library_dir: &Path,
dict: &Dictionary,
llm_client: &crate::clients::llm::LlmClient,
bibcode: &str,
force: bool,
) -> anyhow::Result<String> {
let (_, _, md_opt, tr_opt) =
crate::services::paper::check_paper_paths_in_db(db, library_dir, bibcode)
.await?
.ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?;
// 若本地已存在翻译物理文件且未指明强制重译,直读本地缓存返回
if !force {
if let Some(tr_rel) = tr_opt {
let tr_abs = library_dir.join(&tr_rel);
if tr_abs.exists() {
if let Ok(content) = tokio::fs::read_to_string(&tr_abs).await {
return Ok(content);
}
}
}
}
// 检查英文解析文件是否存在
let md_rel = match md_opt {
Some(rel) => rel,
None => {
anyhow::bail!("文献必须先完成解析方可翻译");
}
};
let md_abs = library_dir.join(&md_rel);
if !md_abs.exists() {
anyhow::bail!("解析 Markdown 文件丢失");
}
let english_markdown = tokio::fs::read_to_string(&md_abs).await?;
// 调用 LLM 翻译服务并注入对照词表
let translated_markdown = translate_markdown(&english_markdown, dict, llm_client).await?;
// 翻译结果持久化
persist_translation(db, library_dir, bibcode, &translated_markdown).await?;
Ok(translated_markdown)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;

129
src/services/vision.rs Normal file
View File

@ -0,0 +1,129 @@
// src/services/vision.rs
//
// 图像加载、SSRF 防护与视觉大模型流式分析服务层
use crate::agent::runtime::AgentStreamEvent;
use crate::clients::llm::LlmClient;
use std::path::Path;
use tracing::info;
/// 图像分析服务:加载图片并通过大语言模型进行流式分析。
pub async fn analyze_image_stream_service(
image_path: &str,
question: &str,
vision_llm: &LlmClient,
library_dir: &Path,
sse_tx: Option<&tokio::sync::mpsc::UnboundedSender<AgentStreamEvent>>,
tool_call_id: Option<String>,
) -> anyhow::Result<(String, String)> {
let (base64_data, mime_type) = load_image(library_dir, image_path).await?;
info!(
"[vision_service] 流式分析图片: {} ({} bytes, type={})",
image_path,
base64_data.len(),
mime_type
);
let system_prompt = "你是一位专业的天体物理学家。请基于提供的图片和用户问题,给出专业、详细、清晰的中文分析。如有数学公式,使用 LaTeX 格式。";
// 流式输出:先发送状态头部(作为第一条 text_delta保证持久化到工具结果中
let header = format!(
"> 正在使用视觉模型({})分析图片: {}\n\n",
vision_llm.model(),
image_path
);
if let Some(tx) = sse_tx {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: header.clone(),
tool_call_id: tool_call_id.clone(),
});
}
let sse_tx_clone = sse_tx.cloned();
let tc_id = tool_call_id.clone();
let answer = vision_llm
.analyze_image_stream(
system_prompt,
question,
&base64_data,
&mime_type,
move |delta: &str| {
if let Some(ref tx) = sse_tx_clone {
let _ = tx.send(AgentStreamEvent::TextDelta {
content: delta.to_string(),
tool_call_id: tc_id.clone(),
});
}
},
)
.await?;
Ok((header, answer))
}
/// 加载图片:本地路径或 HTTP(S) URL → base64 data + mime type
pub async fn load_image(library_dir: &Path, path: &str) -> anyhow::Result<(String, String)> {
if path.starts_with("http://") || path.starts_with("https://") {
// ── SSRF 防护 ──
let parsed =
reqwest::Url::parse(path).map_err(|e| anyhow::anyhow!("无效的图片 URL: {}", e))?;
let host = parsed
.host_str()
.ok_or_else(|| anyhow::anyhow!("URL 缺少 host"))?;
// 字面量快检 + 异步 DNS 深度校验(不阻塞 tokio worker
if crate::utils::ssrf::is_host_literal_blocked(host) {
return Err(anyhow::anyhow!(
"出于安全考虑,禁止请求内网/保留地址: {}",
host
));
}
if crate::utils::ssrf::is_host_blocked_async(host).await {
return Err(anyhow::anyhow!("DNS 解析后指向内网/保留地址: {}", host));
}
// 使用安全的重定向策略:允许跟随合法 CDN 重定向S3/Imgur 等),
// 但每次跳转都对目标 URL 做 scheme + host 字面量校验,拦截内网目标。
let client = reqwest::Client::builder()
.redirect(crate::utils::ssrf::safe_redirect_policy())
.build()?;
let resp = client.get(path).send().await?;
let mime = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("image/png")
.to_string();
let bytes = resp.bytes().await?;
let b64 = base64_encode(&bytes);
Ok((b64, mime))
} else {
let abs_path = if path.starts_with('/') {
Path::new(path).to_path_buf()
} else {
library_dir.join(path)
};
if !abs_path.exists() {
return Err(anyhow::anyhow!("图片文件不存在: {}", abs_path.display()));
}
let ext = abs_path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("png")
.to_lowercase();
let mime = match ext.as_str() {
"jpg" | "jpeg" => "image/jpeg",
"gif" => "image/gif",
"webp" => "image/webp",
"svg" => "image/svg+xml",
_ => "image/png",
};
let bytes = tokio::fs::read(&abs_path).await?;
let b64 = base64_encode(&bytes);
Ok((b64, mime.to_string()))
}
}
fn base64_encode(bytes: &[u8]) -> String {
use base64::{engine::general_purpose, Engine as _};
general_purpose::STANDARD.encode(bytes)
}

View File

@ -97,16 +97,34 @@ pub fn is_host_literal_blocked(host: &str) -> bool {
false false
} }
/// 校验 URL 的 scheme 与 host 字面量是否安全(同步、无 DNS /// 同步 DNS 解析并校验主机 IP 是否属于内网/保留网段。
/// 用于 reqwest 重定向等同步上下文。
pub fn is_host_blocked_sync(host: &str) -> bool {
if let Ok(ip) = host.parse::<IpAddr>() {
return ip_is_blocked(&ip);
}
// 使用标准库进行同步 DNS 解析
use std::net::ToSocketAddrs;
let lookup = format!("{}:80", host);
if let Ok(addrs) = lookup.to_socket_addrs() {
addrs.into_iter().any(|sa| ip_is_blocked(&sa.ip()))
} else {
// 解析失败按危险处理
true
}
}
/// 校验 URL 的 scheme、host 字面量与 DNS 解析是否安全(同步)。
/// ///
/// 用于重定向回调中对每个跳转目标的快速拦截。 /// 用于重定向回调中对每个跳转目标的深度拦截。
pub fn is_url_safe(url: &reqwest::Url) -> bool { pub fn is_url_safe(url: &reqwest::Url) -> bool {
let scheme = url.scheme(); let scheme = url.scheme();
if scheme != "http" && scheme != "https" { if scheme != "http" && scheme != "https" {
return false; return false;
} }
match url.host_str() { match url.host_str() {
Some(host) => !is_host_literal_blocked(host), // 同时进行字面量快速过滤与同步 DNS 深度过滤
Some(host) => !is_host_literal_blocked(host) && !is_host_blocked_sync(host),
None => false, None => false,
} }
} }
@ -196,4 +214,15 @@ mod tests {
async fn test_async_allowed_public() { async fn test_async_allowed_public() {
assert!(!is_host_blocked_async("example.com").await); assert!(!is_host_blocked_async("example.com").await);
} }
#[test]
fn test_sync_blocked_loopback() {
assert!(is_host_blocked_sync("localhost"));
assert!(is_host_blocked_sync("127.0.0.1"));
}
#[test]
fn test_sync_allowed_public() {
assert!(!is_host_blocked_sync("example.com"));
}
} }