// src/api/agent.rs // // 科研智能体 API 控制器。 // 提供 SSE 流式对话接口和会话管理 CRUD 接口。 use axum::{ extract::{Path, Query, State}, response::sse::{Event, Sse}, Json, }; use futures_util::stream::Stream; use serde::{Deserialize, Serialize}; use std::convert::Infallible; use std::sync::Arc; use tracing::{error, info}; use super::error::{ApiResult, AppError}; use super::AppState; use crate::agent::runtime::AgentStreamEvent; // ── POST /api/chat/agent ── // SSE 流式智能体对话接口 #[derive(Debug, Deserialize)] pub struct AgentChatRequest { pub question: String, pub session_id: Option, /// Agent 运行模式: "default" / "deep-research" / "literature-reader" #[serde(default = "default_mode")] pub mode: String, /// 是否启用 LLM 思考模式。None = 由 mode 决定,Some(true/false) = 用户显式覆盖。 #[serde(default)] pub thinking: Option, /// 可选的图片附件(base64 编码 + MIME 类型) #[serde(default)] pub image: Option, } /// 用户附带的图片,用于多模态 Agent 提问。 #[derive(Debug, Deserialize)] pub struct AttachedImage { /// base64 编码的图片数据(不含 data:xxx;base64, 前缀)。与 path 互斥。 #[serde(default)] pub data: String, /// MIME 类型,如 "image/png"、"image/jpeg" #[serde(default)] pub mime_type: String, /// 已有图片的相对路径(重试时复用已有文件,不再重新 base64 解码存盘) #[serde(default)] pub path: Option, } fn default_mode() -> String { "default".to_string() } #[derive(Debug, Serialize)] pub struct AgentModeDto { pub id: &'static str, pub name: &'static str, pub description: &'static str, pub icon: &'static str, } // ── GET /api/chat/modes ── // 获取可用的智能体运行模式列表 pub async fn get_agent_modes() -> Json> { use crate::agent::modes::ModeRegistry; let registry = ModeRegistry::builtins(); let modes = registry .list() .iter() .map(|m| AgentModeDto { id: m.id, name: m.name, description: m.description, icon: m.icon, }) .collect(); Json(modes) } #[derive(Debug, Serialize)] pub struct AgentToolDto { pub name: String, pub display_name: String, pub is_internal: bool, } // ── GET /api/chat/tools ── // 获取系统注册的工具元数据列表 pub async fn get_agent_tools(State(state): State>) -> Json> { use crate::agent::tools::ToolRegistry; let skill_registry = state.skill_registry.clone(); let registry = ToolRegistry::new(skill_registry); let tools = registry .list() .iter() .map(|t| AgentToolDto { name: t.name().to_string(), display_name: t.display_name().to_string(), is_internal: t.is_internal(), }) .collect(); Json(tools) } pub async fn chat_agent( State(state): State>, Json(req): Json, ) -> ApiResult>>> { // 截断日志中的问题内容,避免打印敏感信息 let question_preview = if req.question.len() > 50 { format!("{}...", &req.question[..50]) } else { req.question.clone() }; info!( "接收到智能体对话请求: question='{}', session_id={:?}, has_image={}", question_preview, req.session_id, req.image.is_some() ); // 处理图片附件:保存到磁盘,路径注入 Agent 上下文,前端和 DB 保留原始问题 let question = req.question.clone(); let (image_context, image_path_for_db): (Option, Option) = match req.image { Some(ref img) => { // 如果带有 path 字段(重试场景),复用已有文件,不重新保存 let relative_path: String = if let Some(ref existing_path) = img.path { let full = state.config.storage.library_dir.join(existing_path); if full.exists() { info!("重试复用已有图片: {}", existing_path); existing_path.clone() } else { return Err(AppError::bad_request(format!( "图片文件不存在: {}", existing_path ))); } } else { if img.data.is_empty() { return Err(AppError::bad_request("图片数据为空")); } if !img.mime_type.starts_with("image/") { return Err(AppError::bad_request(format!( "不支持的图片类型: {}", img.mime_type ))); } if state.llm.vision.is_none() { return Err(AppError::bad_request( "图片分析功能未启用。请配置 LLM_VISION_MODEL 环境变量后重试。", )); } let ext = img.mime_type.strip_prefix("image/").unwrap_or("png"); let upload_dir = state .config .storage .library_dir .join(".agent") .join("images") .join("uploads"); tokio::fs::create_dir_all(&upload_dir) .await .map_err(|e| AppError::internal(format!("创建上传目录失败: {}", e)))?; let filename = format!("{}.{}", uuid::Uuid::new_v4(), ext); let filepath = upload_dir.join(&filename); use base64::{engine::general_purpose, Engine as _}; let bytes = general_purpose::STANDARD .decode(&img.data) .map_err(|e| AppError::bad_request(format!("图片 base64 解码失败: {}", e)))?; tokio::fs::write(&filepath, &bytes) .await .map_err(|e| AppError::internal(format!("保存图片失败: {}", e)))?; let rel = filepath .strip_prefix(&state.config.storage.library_dir) .unwrap_or(&filepath) .display() .to_string(); info!("用户图片已保存: {}", rel); rel }; let ctx = format!( "用户上传了一张图片,已保存到: {}\n如需分析此图片,请使用 analyze_image 工具,传入 image_path=\"{}\"。", relative_path, relative_path ); (Some(ctx), Some(relative_path)) } None => (None, None), }; // ── 会话解析与运行时复用 ── // 预分配会话 ID:新会话的首个请求也能写入取消标记并命中运行时缓存 // (历史上新会话首请求超时只能 abort 任务,无法写入取消标记) let session_key = req .session_id .clone() .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); // 模式回放:已存在的会话从 DB 读取创建时的模式。 // 修复:恢复会话时传不同 mode 会静默改变行为——现在会话模式一次创建后保持稳定。 let mode_id = crate::agent::runtime::session::load_session_mode(&state.db, &session_key) .await .unwrap_or_else(|| req.mode.clone()); let runtime = state .agent_runtimes .get_or_create(Arc::clone(&state), &session_key, &mode_id); // 只有 mode 未强制固定 thinking 时,用户才可以覆盖 if runtime.mode_fixed_thinking().is_none() { if let Some(thinking) = req.thinking { runtime.set_thinking(thinking); } } // 同会话 turn 串行化:并发请求 fail-loud 拒绝,防止 turn_index/消息顺序被破坏 let turn_lock = state .agent_runtimes .turn_lock(&session_key) .unwrap_or_else(|| std::sync::Arc::new(tokio::sync::Mutex::new(()))); // 快速检测:已占用直接 409(guard 立即释放,任务内部会重新 try_lock 兜底竞态) if turn_lock.try_lock().is_err() { return Err(AppError::conflict( "该会话正在执行中,请等待当前回合完成后再发送新消息", )); } let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); let question_owned = question.clone(); let image_context_owned = image_context.clone(); let image_path_owned = image_path_for_db.clone(); let cancel_session_key = session_key.clone(); // 在后台 tokio 任务中执行 Agent 循环 let agent_handle = tokio::spawn(async move { // 持有 turn 串行锁直到回合结束(内部重新 try_lock 兜底竞态窗口) let _turn_guard = match turn_lock.try_lock() { Ok(guard) => guard, Err(_) => { let _ = tx.send(AgentStreamEvent::Error { message: "该会话正在执行中,请稍后重试。".to_string(), }); let _ = tx.send(AgentStreamEvent::Done); return; } }; match runtime .run_turn_with_image_context( Some(session_key.clone()), &question_owned, image_context_owned, image_path_owned, tx.clone(), ) .await { Ok(sid) => { info!("智能体对话完成: session_id={}", sid); } Err(e) => { error!("智能体对话执行出错: {}", e); let _ = tx.send(AgentStreamEvent::Error { message: format!("智能体执行错误: {}", e), }); let _ = tx.send(AgentStreamEvent::Done); } } }); // 将 mpsc 通道转换为 SSE 事件流(带 10 分钟超时) const SSE_TIMEOUT_SECS: u64 = 600; let cancelled_runs = state.session.cancelled_runs.clone(); let stream = async_stream::stream! { let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(SSE_TIMEOUT_SECS); loop { let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); if remaining.is_zero() { // 超时:通知 CancellationHook 停止 Agent,并中止后台任务 cancelled_runs.insert(cancel_session_key.clone(), ()); agent_handle.abort(); let timeout_event = AgentStreamEvent::Error { message: "Agent 执行超时(10 分钟),请重试。".to_string(), }; let data = serde_json::to_string(&timeout_event).unwrap_or_default(); yield Ok(Event::default().data(data)); break; } match tokio::time::timeout(remaining, rx.recv()).await { Ok(Some(event)) => { let data = serde_json::to_string(&event).unwrap_or_default(); let is_done = matches!(event, AgentStreamEvent::Done); yield Ok(Event::default().data(data)); if is_done { break; } } Ok(None) => break, // channel closed Err(_) => { // 超时:通知 CancellationHook 停止 Agent,并中止后台任务 cancelled_runs.insert(cancel_session_key.clone(), ()); agent_handle.abort(); let timeout_event = AgentStreamEvent::Error { message: "Agent 执行超时(10 分钟),请重试。".to_string(), }; let data = serde_json::to_string(&timeout_event).unwrap_or_default(); yield Ok(Event::default().data(data)); break; } } } }; Ok(Sse::new(stream)) } // ── GET /api/chat/sessions ── // 获取会话列表 #[derive(Debug, Deserialize)] pub struct SessionListParams { pub limit: Option, pub offset: Option, } pub async fn list_sessions( State(state): State>, Query(params): Query, ) -> ApiResult>> { let limit = params.limit.unwrap_or(50).clamp(1, 200); let offset = params.offset.unwrap_or(0).max(0); let sessions = crate::services::session::list_sessions_service(&state.db, limit, offset) .await .map_err(|e| AppError::internal(format!("查询会话列表失败: {}", e)))?; Ok(Json(sessions)) } pub async fn get_session( State(state): State>, Path(session_id): Path, ) -> ApiResult> { let detail = crate::services::session::get_session_detail_service(&state.db, &session_id) .await .map_err(|e| AppError::internal(format!("查询会话详情失败: {}", e)))? .ok_or_else(|| AppError::not_found(format!("会话 {} 不存在", session_id)))?; Ok(Json(detail)) } // ── DELETE /api/chat/sessions/:id ── // 软删除会话 pub async fn delete_session( State(state): State>, Path(session_id): Path, ) -> ApiResult> { let success = crate::services::session::delete_session_service(&state.db, &session_id) .await .map_err(|e| AppError::internal(format!("删除会话失败: {}", e)))?; if success { // 同步移除会话级运行时缓存(后台队列/压缩日志等状态随之释放) state.agent_runtimes.remove(&session_id); } if !success { return Err(AppError::not_found(format!( "会话 {} 不存在或已删除", session_id ))); } info!("会话已软删除: {}", session_id); Ok(Json( serde_json::json!({ "status": "deleted", "session_id": session_id }), )) } // ── POST /api/chat/sessions/:id/stop ── // 手动停止智能体执行接口 pub async fn stop_agent( State(state): State>, Path(session_id): Path, ) -> ApiResult> { state.session.cancelled_runs.insert(session_id.clone(), ()); info!("已接收并记录手动中止请求,会话 ID: {}", session_id); Ok(Json( serde_json::json!({ "status": "stopping", "session_id": session_id }), )) } // ── GET /api/chat/metrics ── // 返回聚合的智能体运行指标 pub async fn get_agent_metrics( State(state): State>, ) -> ApiResult> { let metrics = crate::services::session::get_agent_metrics_service(&state.db) .await .map_err(|e| AppError::internal(format!("获取系统指标失败: {}", e)))?; Ok(Json(metrics)) } // ── GET /api/chat/sessions/:id/audit ── // 返回指定会话的审计日志 pub async fn get_session_audit( State(state): State>, Path(session_id): Path, ) -> ApiResult>> { let entries = crate::services::session::get_session_audit_service(&state.db, &session_id) .await .map_err(|e| AppError::internal(format!("查询审计日志失败: {}", e)))?; Ok(Json(entries)) } // ── POST /api/chat/answer_question ── // 用户回答 Agent 的提问(ask_user 工具配合使用) #[derive(Debug, Deserialize)] pub struct AnswerQuestionRequest { pub question_id: String, pub answers: Vec, pub free_text: Option, } pub async fn answer_question( State(state): State>, Json(req): Json, ) -> ApiResult> { use crate::agent::tools::ask_user::UserAnswer; let mut pending = state.session.pending_questions.lock().await; let question_id = req.question_id.clone(); match pending.remove(&question_id) { Some(pq) => { let answer = UserAnswer { question_id: question_id.clone(), answers: req.answers.clone(), free_text: req.free_text.clone(), }; match pq.answer_tx.send(answer) { Ok(()) => { info!("[API] 用户回答了问题: id={}", question_id); Ok(Json( serde_json::json!({"status": "ok", "question_id": question_id}), )) } Err(_) => Err(AppError::gone("问题已超时或已被回答")), } } None => Err(AppError::not_found(format!( "未找到待回答问题: {}", question_id ))), } } // ── GET /api/chat/pending_questions ── // 获取当前待回答的问题(前端轮询或初始化) /// 待回答问题的 TTL(10 分钟),超过此时间自动清理 const PENDING_TTL_SECS: u64 = 600; pub async fn get_pending_questions( State(state): State>, ) -> Json> { let mut pending = state.session.pending_questions.lock().await; // 清理过期条目(agent 崩溃后不会被 answer_question 清理) pending.retain(|_, pq| pq.created_at.elapsed().as_secs() < PENDING_TTL_SECS); let questions: Vec = pending .iter() .map(|(id, pq)| { serde_json::from_str::(&pq.question_json) .unwrap_or(serde_json::json!({"question_id": id})) }) .collect(); Json(questions) } // ── POST /api/chat/sessions/:id/permissions/respond ── // 用户响应权限请求 pub async fn respond_permission( State(state): State>, Path(session_id): Path, Json(req): Json, ) -> ApiResult> { let mut perms = state.session.pending_permissions.lock().await; // 按 tool_call_id 查找匹配的权限请求 let perm_id = perms .iter() .find(|(_, p)| p.tool_call_id == req.tool_call_id) .map(|(id, _)| id.clone()); match perm_id { Some(id) => { let perm = match perms.remove(&id) { Some(p) => p, None => { return Err(AppError::internal("权限请求在查找后消失")); } }; match perm.response_tx.send(req) { Ok(()) => { info!( "[API] 用户响应了权限请求: session={} tool_call_id={}", session_id, perm.tool_call_id ); Ok(Json(serde_json::json!({"status": "ok"}))) } Err(_) => Err(AppError::gone("权限请求已超时或已处理")), } } None => Err(AppError::not_found( "未找到该权限请求(可能已超时或已处理)", )), } } // ── GET /api/chat/sessions/:id/permissions ── // 获取当前待处理的权限请求(前端轮询) pub async fn get_pending_permissions( State(state): State>, ) -> Json> { let mut perms = state.session.pending_permissions.lock().await; // 清理过期条目(agent 崩溃后不会被 respond_permission 清理) perms.retain(|_, p| p.created_at.elapsed().as_secs() < PENDING_TTL_SECS); let result: Vec = perms .iter() .map(|(id, p)| { serde_json::json!({ "permission_id": id, "tool_call_id": p.tool_call_id, "tool_name": p.tool_name, "message": p.message, "arguments": p.arguments, }) }) .collect(); Json(result) } // ── POST /api/chat/sessions/:id/branch ── // 创建会话分叉(复制所有 active=1 消息到新会话) #[derive(Debug, Serialize)] pub struct BranchResponse { pub branch_session_id: String, pub forked_at_message_id: i64, pub copied_count: usize, } pub async fn branch_session( State(state): State>, Path(session_id): Path, ) -> ApiResult> { let result = crate::agent::runtime::session::branch_session(&state.db, &session_id) .await .map_err(|e| AppError::bad_request(e.to_string()))?; Ok(Json(BranchResponse { branch_session_id: result.branch_session_id, forked_at_message_id: result.forked_at_message_id, copied_count: result.copied_count, })) } // ── POST /api/chat/sessions/:id/retry ── // 重试最后一次对话(硬删除 + 返回消息文本供前端重提交) #[derive(Debug, Serialize)] pub struct RetryResponse { /// 被删除的用户消息文本(前端可自动重提交) pub retried_message: String, pub new_turn_index: i32, pub deleted_count: i64, pub session_id: String, /// 原消息附带的图片路径(如果有) pub image_path: Option, } pub async fn retry_session( State(state): State>, Path(session_id): Path, ) -> ApiResult> { let (retried_message, new_turn_index, image_path) = crate::agent::runtime::session::retry_last_turn(&state.db, &session_id) .await .map_err(|e| AppError::bad_request(e.to_string()))?; Ok(Json(RetryResponse { retried_message, new_turn_index, deleted_count: 0, session_id, image_path, })) } // ── POST /api/chat/sessions/:id/rewind ── // 回退会话到指定的消息之前(软删除) // // 回退后的消息标记为 active=0(审计保留)。 // 在未产生新对话前可通过 /rewind/restore 恢复。 // 如果已产生新对话,使用 /branch 分叉探索替代路径。 #[derive(Debug, Deserialize)] pub struct RewindRequest { /// 回退 N 个用户轮次(默认 1) pub n: Option, /// 或者指定回退到的消息 ID pub message_id: Option, } #[derive(Debug, Serialize)] pub struct RewindResponse { pub rewound_count: usize, pub target_preview: String, pub new_turn_index: i32, pub session_id: String, } pub async fn rewind_session( State(state): State>, Path(session_id): Path, Json(req): Json, ) -> ApiResult> { let result = if let Some(msg_id) = req.message_id { crate::agent::runtime::session::rewind_to_message(&state.db, &session_id, msg_id) .await .map_err(|e| AppError::bad_request(e.to_string()))? } else { let n = req.n.unwrap_or(1); crate::agent::runtime::session::rewind_n_turns(&state.db, &session_id, n) .await .map_err(|e| AppError::bad_request(e.to_string()))? }; Ok(Json(RewindResponse { rewound_count: result.rewound_count, target_preview: result.target_preview, new_turn_index: result.new_turn_index, session_id: session_id.clone(), })) } // ── POST /api/chat/sessions/:id/rewind/restore ── // 恢复最近一次回退(undo-of-undo)。 // 仅当回退后未产生新对话时才可恢复。 #[derive(Debug, Serialize)] pub struct RestoreResponse { pub restored_count: usize, pub session_id: String, } pub async fn restore_rewound_session( State(state): State>, Path(session_id): Path, ) -> ApiResult> { let count = crate::agent::runtime::session::restore_rewound(&state.db, &session_id) .await .map_err(|e| AppError::conflict(e.to_string()))?; Ok(Json(RestoreResponse { restored_count: count, session_id, })) }