diff --git a/src/main.rs b/src/main.rs index dc0ee79..fcbc670 100644 --- a/src/main.rs +++ b/src/main.rs @@ -438,6 +438,14 @@ async fn cmd_listen( .await; }); info!("定时任务调度器已启动"); + // 审批定期清理 + let approval_clean = approval_manager.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + approval_clean.clean_expired().await; + } + }); } info!("开始监听消息 (llm={}, echo={})...", enable_llm, echo); diff --git a/src/scheduler.rs b/src/scheduler.rs index cbeff9f..423a0a8 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -70,6 +70,7 @@ struct DueTask { } async fn fetch_due_tasks(pool: &PgPool) -> Result, String> { + let mut tx = pool.begin().await.map_err(|e| format!("开始事务失败: {}", e))?; let rows = sqlx::query_as::<_, (uuid::Uuid, String, String, String, String, String)>( r#" SELECT id, name, user_id, command, shell, cwd @@ -80,10 +81,20 @@ async fn fetch_due_tasks(pool: &PgPool) -> Result, String> { FOR UPDATE SKIP LOCKED "#, ) - .fetch_all(pool) + .fetch_all(&mut *tx) .await .map_err(|e| format!("查询到期任务失败: {}", e))?; + // 立即标记为已调度（防止其他实例重复获取） + for (id, _, _, _, _, _) in &rows { + sqlx::query("UPDATE scheduled_tasks SET next_run_at = NOW() + (interval_seconds * INTERVAL '1 second') WHERE id = $1") + .bind(id) + .execute(&mut *tx) + .await + .ok(); + } + tx.commit().await.map_err(|e| format!("提交事务失败: {}", e))?; + Ok(rows .into_iter() .map(|(id, name, user_id, command, shell, cwd)| DueTask { diff --git a/src/tools/approval.rs b/src/tools/approval.rs index 6045051..48c5b63 100644 --- a/src/tools/approval.rs +++ b/src/tools/approval.rs @@ -89,61 +89,91 @@ impl ApprovalManager { /// 处理用户回复（由消息循环调用） /// 匹配成功时返回 Some(技能名)，匹配失败返回 None pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option { - let mut map = self.pending.lock().await; - let entries = map.get_mut(user_id)?; - if entries.is_empty() { - return None; - } + let (result, name) = { + let mut map = self.pending.lock().await; + let entries = map.get_mut(user_id)?; + if entries.is_empty() { return None; } - let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes())); - let reply_trimmed = reply.trim(); + let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes())); + let reply_trimmed = reply.trim(); - // 检查取消 - if reply_trimmed == "0" || reply_trimmed == "取消" { - let entry = entries.remove(0); - let name = entry.skill_name.clone(); - let _ = entry.sender.send(ApprovalDecision::Rejected); - return Some(name); - } - - // 检查确认码 - let pos = entries.iter().position(|e| e.code_hash == reply_hash); - if let Some(idx) = pos { - let entry = entries.remove(idx); - let name = entry.skill_name.clone(); - let _ = entry.sender.send(ApprovalDecision::Approved); - return Some(name); - } - - // 确认码错误 - if let Some(entry) = entries.first_mut() { - entry.attempts_left -= 1; - if entry.attempts_left <= 0 { + if reply_trimmed == "0" || reply_trimmed == "取消" { let entry = entries.remove(0); - let name = entry.skill_name.clone(); let _ = entry.sender.send(ApprovalDecision::Rejected); - return Some(name); + (ApprovalDecision::Rejected, entry.skill_name.clone()) + } else if let Some(idx) = entries.iter().position(|e| e.code_hash == reply_hash) { + let entry = entries.remove(idx); + let _ = entry.sender.send(ApprovalDecision::Approved); + (ApprovalDecision::Approved, entry.skill_name.clone()) + } else if let Some(entry) = entries.first_mut() { + entry.attempts_left -= 1; + if entry.attempts_left <= 0 { + let entry = entries.remove(0); + let n = entry.skill_name.clone(); + let _ = entry.sender.send(ApprovalDecision::Rejected); + (ApprovalDecision::Rejected, n) + } else { + return None; + } + } else { + return None; } - } - - None - } - - /// 清理过期（由定时任务调用） - pub async fn clean_expired(&self) { - let mut map = self.pending.lock().await; - for entries in map.values_mut() { - entries.retain(|e| !e.sender.is_closed()); - } - map.retain(|_, v| !v.is_empty()); + }; + // 更新数据库状态 if let Some(ref pool) = self.pool { + let status = match result { + ApprovalDecision::Approved => "approved", + _ => "rejected", + }; let _ = sqlx::query( - "UPDATE pending_approvals SET status = 'expired' WHERE status = 'pending' AND expires_at < NOW()", + "UPDATE pending_approvals SET status = $1, consumed_at = NOW() WHERE user_id = $2 AND code_hash = $3 AND status = 'pending'", ) + .bind(status) + .bind(user_id) + .bind(&format!("{:x}", Sha256::digest(reply.as_bytes()))) .execute(pool.as_ref()) .await; } + + Some(name) + } + + /// 清理过期审批（通知等待方 + 更新 DB） + pub async fn clean_expired(&self) { + let mut expired = Vec::new(); + { + let map = self.pending.lock().await; + for (uid, entries) in map.iter() { + for (i, e) in entries.iter().enumerate() { + if e.sender.is_closed() { + expired.push((uid.clone(), i, e.code_hash.clone())); + } + } + } + } + + if let Some(ref pool) = self.pool { + for (_uid, _i, _code_hash) in &expired { + let _ = sqlx::query( + "UPDATE pending_approvals SET status = 'expired' WHERE code_hash = $1 AND status = 'pending'", + ) + .bind(_code_hash) + .execute(pool.as_ref()) + .await; + } + } + + let mut map = self.pending.lock().await; + for (uid, i, _) in expired.iter().rev() { + if let Some(entries) = map.get_mut(uid) { + if *i < entries.len() { + let entry = entries.remove(*i); + let _ = entry.sender.send(ApprovalDecision::Expired); + } + } + } + map.retain(|_, v| !v.is_empty()); } /// 获取某个用户的待审批条数 diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 982931e..b934eeb 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -267,6 +267,8 @@ async fn run_skill(skill: &Skill, params: serde_json::Value, _ctx: &ExecutionCon } else { let msg = if !stderr.is_empty() { format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stderr) + } else if !stdout.is_empty() { + format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stdout) } else { format!("退出码 {}", output.status.code().unwrap_or(-1)) };