fix(agent): supervise background tasks atomically

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<ToolRegistry>,
     storage: Option<Arc<crate::storage::Storage>>,
     active_tasks: Arc<DashMap<String, CancellationToken>>,
+    background_permits: Arc<Semaphore>,
     notify_tx: tokio::sync::mpsc::UnboundedSender<TaskNotification>,
     max_concurrent_background_tasks: usize,
     skills_loader: Option<Arc<SkillsLoader>>,
+    task_supervisor: crate::task_supervisor::TaskSupervisor,
 }
 
 impl SubAgentManager {
@@ -130,15 +133,18 @@ impl SubAgentManager {
         notify_tx: tokio::sync::mpsc::UnboundedSender<TaskNotification>,
         max_concurrent_background_tasks: usize,
         skills_loader: Option<Arc<SkillsLoader>>,
+        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<String, SubAgentError> {
-        if self.active_tasks.len() >= self.max_concurrent_background_tasks {
+        let permit = self
-            return Err(SubAgentError::TooManyTasks(
+            .background_permits
-                self.max_concurrent_background_tasks,
+            .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 {
+                                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: 0,
+                                        tool_calls_count,
-                                    iterations: 0,
+                                        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 {

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()));
 

src/storage/mod.rs

@@ -1099,12 +1099,16 @@ impl Storage {
         error: Option<&str>,
         started_at: Option<i64>,
         finished_at: Option<i64>,
+        tool_calls_count: Option<i64>,
+        iterations: Option<i64>,
     ) -> 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();

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<F>(&self, name: impl Into<String>, future: F) -> bool
+    where
+        F: Future<Output = ()> + 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));
+    }
 }