use crate::{ error::AppError, state::{AppState, TaskReply}, }; use chrono::Utc; use lightops_common::protocol::ServerMessage; use serde_json::Value; use tokio::sync::oneshot; use uuid::Uuid; pub async fn request_agent_task( state: &AppState, agent_id: &str, user_id: Option, action: &str, params: Value, ) -> Result { request_agent_task_with_timeout( state, agent_id, user_id, action, params, state.cfg.task_timeout_secs, ) .await } pub async fn request_agent_task_with_timeout( state: &AppState, agent_id: &str, user_id: Option, action: &str, params: Value, timeout_secs: u64, ) -> Result { request_agent_task_with_timeout_and_stored_params( state, agent_id, user_id, action, params.clone(), params, timeout_secs, ) .await } pub async fn request_agent_task_with_timeout_and_stored_params( state: &AppState, agent_id: &str, user_id: Option, action: &str, params: Value, stored_params: Value, timeout_secs: u64, ) -> Result { let task_id = Uuid::new_v4().to_string(); sqlx::query( "INSERT INTO tasks (id, agent_id, user_id, action, params_json, status, started_at) VALUES (?, ?, ?, ?, ?, 'running', ?)", ) .bind(&task_id) .bind(agent_id) .bind(user_id) .bind(action) .bind(stored_params.to_string()) .bind(Utc::now().to_rfc3339()) .execute(&state.pool) .await?; append_task_event( state, &task_id, "info", "任务已创建", Some(&serde_json::json!({ "action": action })), ) .await?; let handle = { let agents = state.agents.read().await; agents.get(agent_id).cloned() } .ok_or(AppError::AgentOffline)?; let (tx, rx) = oneshot::channel(); state.pending.insert(task_id.clone(), tx); state .pending_agents .insert(task_id.clone(), agent_id.to_string()); let msg = ServerMessage::TaskRequest { task_id: task_id.clone(), action: action.to_string(), params, }; if handle.tx.send(msg).is_err() { state.pending.remove(&task_id); state.pending_agents.remove(&task_id); mark_task(state, &task_id, "failed", None, Some("Agent 已断开连接")).await?; return Err(AppError::AgentOffline); } append_task_event(state, &task_id, "info", "任务已下发到 Agent", None).await?; let reply = match tokio::time::timeout(std::time::Duration::from_secs(timeout_secs), rx).await { Ok(Ok(reply)) => reply, Ok(Err(_)) => { state.pending_agents.remove(&task_id); mark_task(state, &task_id, "failed", None, Some("Agent 已断开连接")).await?; return Err(AppError::AgentOffline); } Err(_) => { state.pending.remove(&task_id); state.pending_agents.remove(&task_id); mark_task(state, &task_id, "timeout", None, Some("任务超时")).await?; return Err(AppError::Timeout); } }; let status = if reply.success { "success" } else { "failed" }; mark_task( state, &task_id, status, Some(&reply.data), reply.error.as_deref(), ) .await?; Ok(reply) } pub async fn append_task_event( state: &AppState, task_id: &str, level: &str, message: &str, data: Option<&Value>, ) -> Result<(), AppError> { sqlx::query( "INSERT INTO task_events(task_id, level, message, data_json, created_at) VALUES(?, ?, ?, ?, ?)", ) .bind(task_id) .bind(level) .bind(message.chars().take(500).collect::()) .bind(data.map(|value| value.to_string())) .bind(Utc::now().to_rfc3339()) .execute(&state.pool) .await?; Ok(()) } pub async fn fail_agent_pending_tasks(state: &AppState, agent_id: &str, reason: &str) { let task_ids: Vec = state .pending_agents .iter() .filter(|item| item.value() == agent_id) .map(|item| item.key().clone()) .collect(); for task_id in task_ids { state.pending_agents.remove(&task_id); if let Some((_, tx)) = state.pending.remove(&task_id) { let _ = tx.send(TaskReply { success: false, data: serde_json::json!({}), error: Some(reason.to_string()), }); } let _ = mark_task(state, &task_id, "failed", None, Some(reason)).await; } } pub async fn mark_task( state: &AppState, task_id: &str, status: &str, result: Option<&Value>, error: Option<&str>, ) -> Result<(), AppError> { sqlx::query( "UPDATE tasks SET status = ?, result_json = ?, error = ?, finished_at = ? WHERE id = ?", ) .bind(status) .bind(result.map(|v| v.to_string())) .bind(error) .bind(Utc::now().to_rfc3339()) .bind(task_id) .execute(&state.pool) .await?; let level = if status == "success" { "info" } else { "error" }; let message = match status { "success" => "任务成功完成", "timeout" => "任务执行超时", "cancelled" => "任务已取消", _ => "任务执行失败", }; let data = serde_json::json!({ "status": status, "error": error, "result": result }); append_task_event(state, task_id, level, message, Some(&data)).await?; Ok(()) } pub async fn audit( state: &AppState, user_id: Option, agent_id: Option<&str>, action: &str, target: Option<&str>, params_summary: Option, success: bool, error: Option<&str>, ) { let _ = sqlx::query( "INSERT INTO audit_logs (user_id, agent_id, action, target, params_summary, success, error) VALUES (?, ?, ?, ?, ?, ?, ?)", ) .bind(user_id) .bind(agent_id) .bind(action) .bind(target) .bind(params_summary) .bind(if success { 1 } else { 0 }) .bind(error) .execute(&state.pool) .await; }