AstroResearch/src/api/agent.rs

// 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<String>,
    /// Agent 运行模式： "default" / "deep-research" / "literature-reader"
    #[serde(default = "default_mode")]
    pub mode: String,
    /// 是否启用 LLM 思考模式。None = 由 mode 决定，Some(true/false) = 用户显式覆盖。
    #[serde(default)]
    pub thinking: Option<bool>,
    /// 可选的图片附件（base64 编码 + MIME 类型）
    #[serde(default)]
    pub image: Option<AttachedImage>,
}

/// 用户附带的图片，用于多模态 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<String>,
}

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<Vec<AgentModeDto>> {
    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<Arc<AppState>>) -> Json<Vec<AgentToolDto>> {
    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<Arc<AppState>>,
    Json(req): Json<AgentChatRequest>,
) -> ApiResult<Sse<impl Stream<Item = Result<Event, Infallible>>>> {
    // 截断日志中的问题内容，避免打印敏感信息
    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<String>, Option<String>) = 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::<AgentStreamEvent>();

    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<i64>,
    pub offset: Option<i64>,
}

pub async fn list_sessions(
    State(state): State<Arc<AppState>>,
    Query(params): Query<SessionListParams>,
) -> ApiResult<Json<Vec<crate::services::session::SessionSummary>>> {
    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<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<crate::services::session::SessionDetail>> {
    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<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<serde_json::Value>> {
    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<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<serde_json::Value>> {
    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<Arc<AppState>>,
) -> ApiResult<Json<crate::services::session::AgentMetrics>> {
    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<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<Vec<crate::services::session::AuditLogEntry>>> {
    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<String>,
    pub free_text: Option<String>,
}

pub async fn answer_question(
    State(state): State<Arc<AppState>>,
    Json(req): Json<AnswerQuestionRequest>,
) -> ApiResult<Json<serde_json::Value>> {
    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<Arc<AppState>>,
) -> Json<Vec<serde_json::Value>> {
    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<serde_json::Value> = pending
        .iter()
        .map(|(id, pq)| {
            serde_json::from_str::<serde_json::Value>(&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<Arc<AppState>>,
    Path(session_id): Path<String>,
    Json(req): Json<super::PermissionResponse>,
) -> ApiResult<Json<serde_json::Value>> {
    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<Arc<AppState>>,
) -> Json<Vec<serde_json::Value>> {
    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<serde_json::Value> = 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<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<BranchResponse>> {
    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<String>,
}

pub async fn retry_session(
    State(state): State<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<RetryResponse>> {
    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<usize>,
    /// 或者指定回退到的消息 ID
    pub message_id: Option<i64>,
}

#[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<Arc<AppState>>,
    Path(session_id): Path<String>,
    Json(req): Json<RewindRequest>,
) -> ApiResult<Json<RewindResponse>> {
    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<Arc<AppState>>,
    Path(session_id): Path<String>,
) -> ApiResult<Json<RestoreResponse>> {
    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,
    }))
}