xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): supervise background tasks atomically

Browse Source
This commit is contained in:
xiaoxixi 2026-07-14 11:30:32 +08:00
parent a9c297764e
commit f691c1aad6
4 changed files with 218 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

132
src/agent/sub_agent.rs
View File

@ -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 {
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 {
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,
iterations: 0,
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 {

1
src/session/session.rs
View File

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

48
src/storage/mod.rs
View File

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

42
src/task_supervisor.rs
View File

@ -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));
}
}