From f691c1aad6d227be704c9041299fa761a9e6b1ae Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Tue, 14 Jul 2026 11:30:32 +0800 Subject: [PATCH] fix(agent): supervise background tasks atomically --- src/agent/sub_agent.rs | 142 +++++++++++++++++++++++++++++++++++++---- src/session/session.rs | 1 + src/storage/mod.rs | 48 +++++++++++++- src/task_supervisor.rs | 42 ++++++++++++ 4 files changed, 218 insertions(+), 15 deletions(-) diff --git a/src/agent/sub_agent.rs b/src/agent/sub_agent.rs index b42e2fa..31bd492 100644 --- a/src/agent/sub_agent.rs +++ b/src/agent/sub_agent.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use std::time::Instant; use dashmap::DashMap; +use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; use uuid::Uuid; @@ -117,9 +118,11 @@ pub struct SubAgentManager { full_tools: Arc, storage: Option>, active_tasks: Arc>, + background_permits: Arc, notify_tx: tokio::sync::mpsc::UnboundedSender, max_concurrent_background_tasks: usize, skills_loader: Option>, + task_supervisor: crate::task_supervisor::TaskSupervisor, } impl SubAgentManager { @@ -130,15 +133,18 @@ impl SubAgentManager { notify_tx: tokio::sync::mpsc::UnboundedSender, max_concurrent_background_tasks: usize, skills_loader: Option>, + task_supervisor: crate::task_supervisor::TaskSupervisor, ) -> Self { Self { provider_config, full_tools, storage, active_tasks: Arc::new(DashMap::new()), + background_permits: Arc::new(Semaphore::new(max_concurrent_background_tasks)), notify_tx, max_concurrent_background_tasks, skills_loader, + task_supervisor, } } @@ -306,11 +312,11 @@ impl SubAgentManager { config: SubAgentConfig, ctx: DelegateContext, ) -> Result { - if self.active_tasks.len() >= self.max_concurrent_background_tasks { - return Err(SubAgentError::TooManyTasks( - self.max_concurrent_background_tasks, - )); - } + let permit = self + .background_permits + .clone() + .try_acquire_owned() + .map_err(|_| SubAgentError::TooManyTasks(self.max_concurrent_background_tasks))?; let task_id = generate_task_id(); let cancel_token = CancellationToken::new(); @@ -368,6 +374,7 @@ impl SubAgentManager { let storage = self.storage.clone(); let notify_tx = self.notify_tx.clone(); let active_tasks = Arc::clone(&self.active_tasks); + let shutdown = self.task_supervisor.cancellation_token(); let tid = task_id.clone(); let sess_id = ctx.session_id.clone(); @@ -375,7 +382,10 @@ impl SubAgentManager { let cid = ctx.chat_id.clone(); let prompt = config.prompt.clone(); - tokio::spawn(async move { + let spawned = self.task_supervisor.spawn_graceful( + format!("sub-agent:{task_id}"), + async move { + let _permit = permit; let started_at = chrono::Utc::now().timestamp_millis(); // Update DB: running @@ -388,6 +398,8 @@ impl SubAgentManager { None, Some(started_at), None, + None, + None, ) .await; } @@ -423,14 +435,20 @@ impl SubAgentManager { agent.process(history), ) => { match r { - Ok(Ok(agent_result)) => SubAgentResult { - task_id: tid.clone(), - content: agent_result.final_response.content, - content_truncated: false, - status: TaskStatus::Completed, - tool_calls_count: 0, - iterations: 0, - duration_ms: 0, + Ok(Ok(agent_result)) => { + let tool_calls_count = agent_result.emitted_messages + .iter().filter(|m| m.tool_calls.is_some()).count(); + let iterations = agent_result.emitted_messages + .iter().filter(|m| m.role == "assistant" && m.tool_calls.is_some()).count(); + SubAgentResult { + task_id: tid.clone(), + content: agent_result.final_response.content, + content_truncated: false, + status: TaskStatus::Completed, + tool_calls_count, + iterations, + duration_ms: 0, + } }, Ok(Err(e)) => SubAgentResult { task_id: tid.clone(), @@ -461,6 +479,15 @@ impl SubAgentManager { iterations: 0, duration_ms: 0, }, + _ = shutdown.cancelled() => SubAgentResult { + task_id: tid.clone(), + content: String::new(), + content_truncated: false, + status: TaskStatus::Cancelled, + tool_calls_count: 0, + iterations: 0, + duration_ms: 0, + }, } } None => SubAgentResult { @@ -493,6 +520,8 @@ impl SubAgentManager { error_val.as_deref(), Some(started_at), Some(finished_at), + Some(result.tool_calls_count as i64), + Some(result.iterations as i64), ) .await; } @@ -509,6 +538,27 @@ impl SubAgentManager { active_tasks.remove(&tid); }); + if !spawned { + self.active_tasks.remove(&task_id); + if let Some(ref storage) = self.storage { + let _ = storage + .update_background_task_status( + &task_id, + "cancelled", + None, + Some("gateway shutdown"), + None, + Some(chrono::Utc::now().timestamp_millis()), + None, + None, + ) + .await; + } + return Err(SubAgentError::Other( + "gateway is shutting down and cannot accept background tasks".to_string(), + )); + } + Ok(task_id) } @@ -523,6 +573,8 @@ impl SubAgentManager { None, None, Some(chrono::Utc::now().timestamp_millis()), + None, + None, ) .await .map_err(|e| SubAgentError::Storage(e.to_string()))?; @@ -610,6 +662,68 @@ fn truncate_sub_agent_result(content: &str) -> (String, bool) { } } +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + fn manager(max_tasks: usize) -> SubAgentManager { + let (notify_tx, _notify_rx) = tokio::sync::mpsc::unbounded_channel(); + SubAgentManager::new( + LLMProviderConfig { + provider_type: "openai".into(), + name: "test".into(), + base_url: "http://localhost".into(), + api_key: "test".into(), + extra_headers: HashMap::new(), + model_id: "test".into(), + temperature: None, + max_tokens: None, + model_extra: HashMap::new(), + max_tool_iterations: 1, + token_limit: 4096, + workspace_dir: std::env::temp_dir(), + input_types: vec!["text".into()], + }, + Arc::new(ToolRegistry::new()), + None, + notify_tx, + max_tasks, + None, + crate::task_supervisor::TaskSupervisor::new(), + ) + } + + #[tokio::test] + async fn background_limit_is_enforced_by_atomic_permit() { + let manager = manager(1); + let _permit = manager + .background_permits + .clone() + .try_acquire_owned() + .unwrap(); + let error = manager + .run_background( + SubAgentConfig { + prompt: "test".into(), + mode: ExecutionMode::Background, + allowed_tools: None, + max_iterations: None, + timeout_secs: Some(1), + }, + DelegateContext { + session_id: "cli:test:dialog".into(), + channel: "cli".into(), + chat_id: "test".into(), + }, + ) + .await + .unwrap_err(); + + assert!(matches!(error, SubAgentError::TooManyTasks(1))); + } +} + fn summarize_for_notification(content: &str, _duration_ms: u64) -> String { const MAX_SUMMARY_BYTES: usize = 500; if content.len() <= MAX_SUMMARY_BYTES { diff --git a/src/session/session.rs b/src/session/session.rs index a17f56c..7104eee 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -1063,6 +1063,7 @@ impl SessionManager { notify_tx, max_concurrent_background_tasks, Some(skills_loader.clone()), + task_supervisor.clone(), )); tools.register(crate::tools::DelegateTool::new(sub_agent_manager.clone())); diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 076a574..59741f1 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -1099,12 +1099,16 @@ impl Storage { error: Option<&str>, started_at: Option, finished_at: Option, + tool_calls_count: Option, + iterations: Option, ) -> Result<(), StorageError> { sqlx::query( r#" UPDATE background_tasks SET status = ?, result = COALESCE(?, result), error = COALESCE(?, error), - started_at = COALESCE(?, started_at), finished_at = COALESCE(?, finished_at) + started_at = COALESCE(?, started_at), finished_at = COALESCE(?, finished_at), + tool_calls_count = COALESCE(?, tool_calls_count), + iterations = COALESCE(?, iterations) WHERE id = ? "#, ) @@ -1113,6 +1117,8 @@ impl Storage { .bind(error) .bind(started_at) .bind(finished_at) + .bind(tool_calls_count) + .bind(iterations) .bind(id) .execute(self.pool()) .await?; @@ -1251,6 +1257,46 @@ mod tests { assert!(orphan.is_err()); } + #[tokio::test] + async fn background_task_completion_persists_execution_metrics() { + let (storage, _dir) = create_test_storage().await; + let task = crate::storage::BackgroundTask { + id: "task-metrics".into(), + session_id: "cli:test:dialog".into(), + channel: "cli".into(), + chat_id: "test".into(), + prompt: "measure".into(), + allowed_tools: None, + status: "pending".into(), + result: None, + error: None, + tool_calls_count: 0, + iterations: 0, + started_at: None, + finished_at: None, + created_at: 1, + }; + storage.create_background_task(&task).await.unwrap(); + + storage + .update_background_task_status( + &task.id, + "completed", + Some("done"), + None, + Some(2), + Some(3), + Some(4), + Some(5), + ) + .await + .unwrap(); + + let persisted = storage.get_background_task(&task.id).await.unwrap(); + assert_eq!(persisted.tool_calls_count, 4); + assert_eq!(persisted.iterations, 5); + } + #[tokio::test] async fn legacy_schema_is_migrated_without_rebuild() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/task_supervisor.rs b/src/task_supervisor.rs index d890a07..c07c614 100644 --- a/src/task_supervisor.rs +++ b/src/task_supervisor.rs @@ -88,6 +88,32 @@ impl TaskSupervisor { true } + /// Register a task that performs its own cooperative cancellation and + /// cleanup. The supervisor broadcasts cancellation during shutdown, but + /// does not drop this future until the grace period expires. + pub fn spawn_graceful(&self, name: impl Into, future: F) -> bool + where + F: Future + Send + 'static, + { + let name = name.into(); + let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner()); + if state.stopping { + return false; + } + + state.tasks.retain(|task| !task.handle.is_finished()); + let task_name = name.clone(); + let handle = tokio::spawn(async move { + tracing::debug!(task = %task_name, "Graceful background task started"); + match AssertUnwindSafe(future).catch_unwind().await { + Ok(()) => tracing::debug!(task = %task_name, "Graceful background task finished"), + Err(_) => tracing::error!(task = %task_name, "Graceful background task panicked"), + } + }); + state.tasks.push(ManagedTask { name, handle }); + true + } + pub fn cancel(&self) { let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner()); state.stopping = true; @@ -157,4 +183,20 @@ mod tests { supervisor.shutdown(Duration::from_secs(1)).await; assert!(dropped.load(Ordering::SeqCst)); } + + #[tokio::test] + async fn graceful_task_observes_cancellation_before_shutdown_returns() { + let supervisor = TaskSupervisor::new(); + let cancellation = supervisor.cancellation_token(); + let cleaned_up = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let task_cleaned_up = cleaned_up.clone(); + assert!(supervisor.spawn_graceful("graceful", async move { + cancellation.cancelled().await; + task_cleaned_up.store(true, std::sync::atomic::Ordering::SeqCst); + })); + + supervisor.shutdown(Duration::from_secs(1)).await; + + assert!(cleaned_up.load(std::sync::atomic::Ordering::SeqCst)); + } }