855 lines
30 KiB
Rust
855 lines
30 KiB
Rust
use std::collections::HashSet;
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
|
|
use dashmap::DashMap;
|
|
use tokio::sync::Semaphore;
|
|
use tokio_util::sync::CancellationToken;
|
|
use uuid::Uuid;
|
|
|
|
use crate::agent::AgentError;
|
|
use crate::agent::AgentLoop;
|
|
use crate::agent::system_prompt::build_sub_agent_system_prompt;
|
|
use crate::bus::ChatMessage;
|
|
use crate::config::LLMProviderConfig;
|
|
use crate::providers::{LLMProvider, create_provider};
|
|
use crate::skills::SkillsLoader;
|
|
use crate::tools::ToolRegistry;
|
|
|
|
tokio::task_local! {
|
|
pub(crate) static DELEGATE_CONTEXT: DelegateContext;
|
|
}
|
|
|
|
/// Read the delegate context from the current task. Returns an error if not set.
|
|
pub fn get_delegate_context() -> Result<DelegateContext, String> {
|
|
DELEGATE_CONTEXT
|
|
.try_with(|ctx| ctx.clone())
|
|
.map_err(|_| "DELEGATE_CONTEXT not set".to_string())
|
|
}
|
|
|
|
const DEFAULT_MAX_ITERATIONS: usize = 99;
|
|
const DEFAULT_TIMEOUT_SECS: u64 = 3600;
|
|
const MAX_INLINE_RESULT_CHARS: usize = 8000;
|
|
|
|
const DEFAULT_READONLY_TOOLS: &[&str] = &[
|
|
"file_read",
|
|
"file_search",
|
|
"content_search",
|
|
"web_fetch",
|
|
"http_request",
|
|
"calculator",
|
|
];
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct SubAgentConfig {
|
|
pub prompt: String,
|
|
pub mode: ExecutionMode,
|
|
pub allowed_tools: Option<Vec<String>>,
|
|
pub max_iterations: Option<usize>,
|
|
pub timeout_secs: Option<u64>,
|
|
pub plan_item_id: Option<String>,
|
|
pub session_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum ExecutionMode {
|
|
Inline,
|
|
Background,
|
|
Parallel,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct SubAgentResult {
|
|
pub task_id: String,
|
|
pub content: String,
|
|
pub content_truncated: bool,
|
|
pub status: TaskStatus,
|
|
pub tool_calls_count: usize,
|
|
pub iterations: usize,
|
|
pub duration_ms: u64,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum TaskStatus {
|
|
Completed,
|
|
Failed(String),
|
|
Cancelled,
|
|
TimedOut,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct TaskNotification {
|
|
pub task_id: String,
|
|
pub session_id: String,
|
|
pub channel: String,
|
|
pub chat_id: String,
|
|
pub status: TaskStatus,
|
|
pub result_summary: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct DelegateContext {
|
|
pub session_id: String,
|
|
pub channel: String,
|
|
pub chat_id: String,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum SubAgentError {
|
|
TooManyTasks(usize),
|
|
ProviderCreation(String),
|
|
Storage(String),
|
|
Other(String),
|
|
}
|
|
|
|
impl std::fmt::Display for SubAgentError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
Self::TooManyTasks(max) => write!(f, "后台任务已达上限({})，请稍后重试", max),
|
|
Self::ProviderCreation(e) => write!(f, "provider creation failed: {}", e),
|
|
Self::Storage(e) => write!(f, "storage error: {}", e),
|
|
Self::Other(e) => write!(f, "{}", e),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for SubAgentError {}
|
|
|
|
pub struct SubAgentManager {
|
|
provider_config: LLMProviderConfig,
|
|
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>>,
|
|
work_manager: Option<Arc<crate::work::WorkManager>>,
|
|
task_supervisor: crate::task_supervisor::TaskSupervisor,
|
|
}
|
|
|
|
impl SubAgentManager {
|
|
pub fn new(
|
|
provider_config: LLMProviderConfig,
|
|
full_tools: Arc<ToolRegistry>,
|
|
storage: Option<Arc<crate::storage::Storage>>,
|
|
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,
|
|
work_manager: None,
|
|
task_supervisor,
|
|
}
|
|
}
|
|
|
|
pub fn with_work_manager(mut self, work_manager: Arc<crate::work::WorkManager>) -> Self {
|
|
self.work_manager = Some(work_manager);
|
|
self
|
|
}
|
|
|
|
pub fn filter_tools(&self, allowed: &Option<Vec<String>>) -> Arc<ToolRegistry> {
|
|
let allowed_set: HashSet<&str> = match allowed {
|
|
Some(list) => list.iter().map(|s| s.as_str()).collect(),
|
|
None => DEFAULT_READONLY_TOOLS.iter().copied().collect(),
|
|
};
|
|
let filtered = ToolRegistry::new();
|
|
for (name, tool) in self.full_tools.iter() {
|
|
if allowed_set.contains(name.as_str()) && name != "delegate" && name != "todo" {
|
|
filtered.register_raw(name, tool);
|
|
}
|
|
}
|
|
Arc::new(filtered)
|
|
}
|
|
|
|
fn get_skills_prompt(&self, tools: &ToolRegistry) -> Option<String> {
|
|
let has_get_skill = tools.iter().iter().any(|(name, _)| name == "get_skill");
|
|
if has_get_skill && let Some(ref loader) = self.skills_loader {
|
|
let prompt = loader.build_skills_prompt();
|
|
if !prompt.is_empty() {
|
|
return Some(prompt);
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
pub fn build_sub_agent(
|
|
&self,
|
|
config: &SubAgentConfig,
|
|
tools: Arc<ToolRegistry>,
|
|
) -> Result<AgentLoop, AgentError> {
|
|
let mut provider = create_provider(self.provider_config.clone())
|
|
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
|
if let Some(ref s) = self.storage {
|
|
provider.set_storage(s.clone());
|
|
}
|
|
let provider: Arc<dyn LLMProvider> = Arc::from(provider);
|
|
|
|
let max_iterations = config.max_iterations.unwrap_or(DEFAULT_MAX_ITERATIONS);
|
|
let workspace_dir = self.provider_config.workspace_dir.clone();
|
|
let model_name = self.provider_config.model_id.clone();
|
|
let input_types = self.provider_config.input_types.clone();
|
|
|
|
let agent = AgentLoop::with_provider_and_tools(
|
|
provider,
|
|
tools,
|
|
max_iterations,
|
|
model_name,
|
|
workspace_dir,
|
|
input_types,
|
|
)
|
|
.with_context_window(self.provider_config.token_limit);
|
|
|
|
Ok(agent)
|
|
}
|
|
|
|
pub async fn run_inline(
|
|
&self,
|
|
config: SubAgentConfig,
|
|
) -> Result<SubAgentResult, SubAgentError> {
|
|
let task_id = generate_task_id();
|
|
let tools = self.filter_tools(&config.allowed_tools);
|
|
let timeout_secs = config.timeout_secs.unwrap_or(DEFAULT_TIMEOUT_SECS);
|
|
let timeout_human = format_duration(timeout_secs);
|
|
let http_get_only = config.allowed_tools.is_none()
|
|
|| config
|
|
.allowed_tools
|
|
.as_ref()
|
|
.is_some_and(|v| v.iter().any(|t| t == "http_request"));
|
|
let skills_prompt = self.get_skills_prompt(&tools);
|
|
let system_prompt = build_sub_agent_system_prompt(
|
|
&config.prompt,
|
|
&timeout_human,
|
|
&tools,
|
|
&self.provider_config.workspace_dir,
|
|
&self.provider_config.model_id,
|
|
skills_prompt,
|
|
http_get_only,
|
|
);
|
|
|
|
let agent = self
|
|
.build_sub_agent(&config, tools)
|
|
.map_err(|e| SubAgentError::ProviderCreation(e.to_string()))?;
|
|
self.assign_work_item(&config, &task_id).await?;
|
|
|
|
let history = vec![
|
|
ChatMessage::system(system_prompt),
|
|
ChatMessage::user(&config.prompt),
|
|
];
|
|
|
|
let start = Instant::now();
|
|
|
|
let result = tokio::time::timeout(
|
|
std::time::Duration::from_secs(timeout_secs),
|
|
agent.process(history),
|
|
)
|
|
.await;
|
|
|
|
let duration_ms = start.elapsed().as_millis() as u64;
|
|
|
|
let result = match result {
|
|
Ok(Ok(agent_result)) => {
|
|
let (content, truncated) =
|
|
truncate_sub_agent_result(&agent_result.final_response.content);
|
|
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: task_id.clone(),
|
|
content,
|
|
content_truncated: truncated,
|
|
status: TaskStatus::Completed,
|
|
tool_calls_count,
|
|
iterations,
|
|
duration_ms,
|
|
}
|
|
}
|
|
Ok(Err(e)) => SubAgentResult {
|
|
task_id: task_id.clone(),
|
|
content: String::new(),
|
|
content_truncated: false,
|
|
status: TaskStatus::Failed(e.to_string()),
|
|
tool_calls_count: 0,
|
|
iterations: 0,
|
|
duration_ms,
|
|
},
|
|
Err(_elapsed) => SubAgentResult {
|
|
task_id: task_id.clone(),
|
|
content: String::new(),
|
|
content_truncated: false,
|
|
status: TaskStatus::TimedOut,
|
|
tool_calls_count: 0,
|
|
iterations: 0,
|
|
duration_ms,
|
|
},
|
|
};
|
|
self.finish_work_item(&config, &result).await;
|
|
Ok(result)
|
|
}
|
|
|
|
pub async fn run_parallel(
|
|
&self,
|
|
configs: Vec<SubAgentConfig>,
|
|
) -> Result<Vec<SubAgentResult>, SubAgentError> {
|
|
let futures: Vec<_> = configs
|
|
.into_iter()
|
|
.map(|config| {
|
|
let mgr = self; // &self borrow, all tasks share the same manager
|
|
async move { mgr.run_inline(config).await }
|
|
})
|
|
.collect();
|
|
|
|
let results = futures_util::future::join_all(futures).await;
|
|
results.into_iter().collect::<Result<Vec<_>, _>>()
|
|
}
|
|
|
|
pub async fn run_background(
|
|
&self,
|
|
config: SubAgentConfig,
|
|
ctx: DelegateContext,
|
|
) -> Result<String, SubAgentError> {
|
|
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 mut work_config = config.clone();
|
|
if work_config.session_id.is_none() {
|
|
work_config.session_id = Some(ctx.session_id.clone());
|
|
}
|
|
let cancel_token = CancellationToken::new();
|
|
|
|
// Write DB: pending
|
|
if let Some(ref storage) = self.storage {
|
|
let allowed_tools_json = config
|
|
.allowed_tools
|
|
.as_ref()
|
|
.and_then(|v| serde_json::to_string(v).ok());
|
|
let record = crate::storage::BackgroundTask {
|
|
id: task_id.clone(),
|
|
session_id: ctx.session_id.clone(),
|
|
channel: ctx.channel.clone(),
|
|
chat_id: ctx.chat_id.clone(),
|
|
prompt: config.prompt.clone(),
|
|
allowed_tools: allowed_tools_json,
|
|
status: "pending".to_string(),
|
|
result: None,
|
|
error: None,
|
|
tool_calls_count: 0,
|
|
iterations: 0,
|
|
started_at: None,
|
|
finished_at: None,
|
|
created_at: chrono::Utc::now().timestamp_millis(),
|
|
};
|
|
storage
|
|
.create_background_task(&record)
|
|
.await
|
|
.map_err(|e| SubAgentError::Storage(e.to_string()))?;
|
|
}
|
|
if let Err(error) = self.assign_work_item(&work_config, &task_id).await {
|
|
if let Some(ref storage) = self.storage {
|
|
let _ = storage
|
|
.update_background_task_status(
|
|
&task_id,
|
|
crate::storage::background_task::BackgroundTaskUpdate {
|
|
status: "cancelled",
|
|
result: None,
|
|
error: Some("plan item assignment failed"),
|
|
started_at: None,
|
|
finished_at: Some(chrono::Utc::now().timestamp_millis()),
|
|
tool_calls_count: None,
|
|
iterations: None,
|
|
},
|
|
)
|
|
.await;
|
|
}
|
|
return Err(error);
|
|
}
|
|
|
|
self.active_tasks
|
|
.insert(task_id.clone(), cancel_token.clone());
|
|
|
|
let tools = self.filter_tools(&config.allowed_tools);
|
|
let timeout_secs = config.timeout_secs.unwrap_or(DEFAULT_TIMEOUT_SECS);
|
|
let timeout_human = format_duration(timeout_secs);
|
|
let http_get_only = config.allowed_tools.is_none()
|
|
|| config
|
|
.allowed_tools
|
|
.as_ref()
|
|
.is_some_and(|v| v.iter().any(|t| t == "http_request"));
|
|
let skills_prompt = self.get_skills_prompt(&tools);
|
|
let system_prompt = build_sub_agent_system_prompt(
|
|
&config.prompt,
|
|
&timeout_human,
|
|
&tools,
|
|
&self.provider_config.workspace_dir,
|
|
&self.provider_config.model_id,
|
|
skills_prompt,
|
|
http_get_only,
|
|
);
|
|
let provider_config = self.provider_config.clone();
|
|
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();
|
|
let ch = ctx.channel.clone();
|
|
let cid = ctx.chat_id.clone();
|
|
let prompt = config.prompt.clone();
|
|
let work_manager = self.work_manager.clone();
|
|
let work_item_id = work_config.plan_item_id.clone();
|
|
let work_session_id = work_config.session_id.clone();
|
|
|
|
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
|
|
if let Some(ref s) = storage {
|
|
let _ = s
|
|
.update_background_task_status(&tid, crate::storage::background_task::BackgroundTaskUpdate {
|
|
status: "running",
|
|
result: None,
|
|
error: None,
|
|
started_at: Some(started_at),
|
|
finished_at: None,
|
|
tool_calls_count: None,
|
|
iterations: None,
|
|
})
|
|
.await;
|
|
}
|
|
|
|
let mut provider = create_provider(provider_config.clone()).ok();
|
|
if let Some(ref mut p) = provider
|
|
&& let Some(ref s) = storage
|
|
{
|
|
p.set_storage(s.clone());
|
|
}
|
|
let provider_result: Option<Arc<dyn LLMProvider>> = provider.map(Arc::from);
|
|
|
|
let result = match provider_result {
|
|
Some(provider) => {
|
|
let agent = AgentLoop::with_provider_and_tools(
|
|
provider,
|
|
tools,
|
|
DEFAULT_MAX_ITERATIONS,
|
|
provider_config.model_id.clone(),
|
|
provider_config.workspace_dir.clone(),
|
|
provider_config.input_types.clone(),
|
|
)
|
|
.with_context_window(provider_config.token_limit);
|
|
|
|
let history = vec![
|
|
ChatMessage::system(system_prompt),
|
|
ChatMessage::user(&prompt),
|
|
];
|
|
|
|
tokio::select! {
|
|
r = tokio::time::timeout(
|
|
std::time::Duration::from_secs(timeout_secs),
|
|
agent.process(history),
|
|
) => {
|
|
match r {
|
|
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(),
|
|
content: String::new(),
|
|
content_truncated: false,
|
|
status: TaskStatus::Failed(e.to_string()),
|
|
tool_calls_count: 0,
|
|
iterations: 0,
|
|
duration_ms: 0,
|
|
},
|
|
Err(_) => SubAgentResult {
|
|
task_id: tid.clone(),
|
|
content: String::new(),
|
|
content_truncated: false,
|
|
status: TaskStatus::TimedOut,
|
|
tool_calls_count: 0,
|
|
iterations: 0,
|
|
duration_ms: 0,
|
|
},
|
|
}
|
|
}
|
|
_ = cancel_token.cancelled() => SubAgentResult {
|
|
task_id: tid.clone(),
|
|
content: String::new(),
|
|
content_truncated: false,
|
|
status: TaskStatus::Cancelled,
|
|
tool_calls_count: 0,
|
|
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 {
|
|
task_id: tid.clone(),
|
|
content: String::new(),
|
|
content_truncated: false,
|
|
status: TaskStatus::Failed("provider creation failed".into()),
|
|
tool_calls_count: 0,
|
|
iterations: 0,
|
|
duration_ms: 0,
|
|
},
|
|
};
|
|
|
|
let finished_at = chrono::Utc::now().timestamp_millis();
|
|
let duration_ms = (finished_at - started_at) as u64;
|
|
|
|
let (status_str, error_val) = match &result.status {
|
|
TaskStatus::Completed => ("completed".to_string(), None),
|
|
TaskStatus::Failed(e) => ("failed".to_string(), Some(e.clone())),
|
|
TaskStatus::Cancelled => ("cancelled".to_string(), None),
|
|
TaskStatus::TimedOut => ("failed".to_string(), Some("timeout".to_string())),
|
|
};
|
|
|
|
if let (Some(manager), Some(session_id), Some(item_id)) =
|
|
(work_manager.as_ref(), work_session_id.as_deref(), work_item_id.as_deref())
|
|
{
|
|
let completed = matches!(result.status, TaskStatus::Completed);
|
|
let summary = if completed {
|
|
Some(result.content.as_str())
|
|
} else {
|
|
error_val.as_deref().or(Some("子 Agent 未完成任务"))
|
|
};
|
|
if let Err(error) = manager
|
|
.finish_sub_agent(session_id, item_id, &tid, completed, summary)
|
|
.await
|
|
{
|
|
tracing::warn!(task_id = %tid, item_id, error = %error, "Failed to update plan item");
|
|
}
|
|
}
|
|
|
|
if let Some(ref s) = storage {
|
|
let _ = s
|
|
.update_background_task_status(&tid, crate::storage::background_task::BackgroundTaskUpdate {
|
|
status: &status_str,
|
|
result: Some(&result.content),
|
|
error: error_val.as_deref(),
|
|
started_at: Some(started_at),
|
|
finished_at: Some(finished_at),
|
|
tool_calls_count: Some(result.tool_calls_count as i64),
|
|
iterations: Some(result.iterations as i64),
|
|
})
|
|
.await;
|
|
}
|
|
|
|
let _ = notify_tx.send(TaskNotification {
|
|
task_id: tid.clone(),
|
|
session_id: sess_id,
|
|
channel: ch,
|
|
chat_id: cid,
|
|
status: result.status,
|
|
result_summary: summarize_for_notification(&result.content, duration_ms),
|
|
});
|
|
|
|
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,
|
|
crate::storage::background_task::BackgroundTaskUpdate {
|
|
status: "cancelled",
|
|
result: None,
|
|
error: Some("gateway shutdown"),
|
|
started_at: None,
|
|
finished_at: Some(chrono::Utc::now().timestamp_millis()),
|
|
tool_calls_count: None,
|
|
iterations: None,
|
|
},
|
|
)
|
|
.await;
|
|
}
|
|
if let (Some(manager), Some(session_id), Some(item_id)) = (
|
|
self.work_manager.as_ref(),
|
|
work_config.session_id.as_deref(),
|
|
work_config.plan_item_id.as_deref(),
|
|
) {
|
|
let _ = manager
|
|
.finish_sub_agent(
|
|
session_id,
|
|
item_id,
|
|
&task_id,
|
|
false,
|
|
Some("gateway shutdown"),
|
|
)
|
|
.await;
|
|
}
|
|
return Err(SubAgentError::Other(
|
|
"gateway is shutting down and cannot accept background tasks".to_string(),
|
|
));
|
|
}
|
|
|
|
Ok(task_id)
|
|
}
|
|
|
|
pub async fn cancel_task(&self, task_id: &str) -> Result<bool, SubAgentError> {
|
|
if let Some((_, token)) = self.active_tasks.remove(task_id) {
|
|
token.cancel();
|
|
if let Some(ref s) = self.storage {
|
|
s.update_background_task_status(
|
|
task_id,
|
|
crate::storage::background_task::BackgroundTaskUpdate {
|
|
status: "cancelled",
|
|
result: None,
|
|
error: None,
|
|
started_at: None,
|
|
finished_at: Some(chrono::Utc::now().timestamp_millis()),
|
|
tool_calls_count: None,
|
|
iterations: None,
|
|
},
|
|
)
|
|
.await
|
|
.map_err(|e| SubAgentError::Storage(e.to_string()))?;
|
|
}
|
|
Ok(true)
|
|
} else if let Some(ref s) = self.storage {
|
|
match s.get_background_task(task_id).await {
|
|
Ok(task) => match task.status.as_str() {
|
|
"pending" | "running" => {
|
|
tracing::warn!(task_id, "task in DB but not in active_tasks");
|
|
Ok(false)
|
|
}
|
|
_ => Ok(false),
|
|
},
|
|
Err(_) => Ok(false),
|
|
}
|
|
} else {
|
|
Ok(false)
|
|
}
|
|
}
|
|
|
|
pub async fn check_task(&self, task_id: &str) -> Option<crate::storage::BackgroundTask> {
|
|
if let Some(ref s) = self.storage {
|
|
s.get_background_task(task_id).await.ok()
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
pub async fn list_tasks(&self, session_id: &str) -> Vec<crate::storage::BackgroundTask> {
|
|
if let Some(ref s) = self.storage {
|
|
s.list_background_tasks(session_id)
|
|
.await
|
|
.unwrap_or_default()
|
|
} else {
|
|
vec![]
|
|
}
|
|
}
|
|
|
|
pub async fn cancel_by_session(&self, session_id: &str) {
|
|
// Cancel all running tasks for a session by checking DB
|
|
if let Some(ref s) = self.storage
|
|
&& let Ok(tasks) = s.list_background_tasks(session_id).await
|
|
{
|
|
for task in &tasks {
|
|
if task.status == "pending" || task.status == "running" {
|
|
let _ = self.cancel_task(&task.id).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn active_task_count(&self) -> usize {
|
|
self.active_tasks.len()
|
|
}
|
|
|
|
async fn assign_work_item(
|
|
&self,
|
|
config: &SubAgentConfig,
|
|
task_id: &str,
|
|
) -> Result<(), SubAgentError> {
|
|
if let (Some(manager), Some(session_id), Some(item_id)) = (
|
|
self.work_manager.as_ref(),
|
|
config.session_id.as_deref(),
|
|
config.plan_item_id.as_deref(),
|
|
) {
|
|
manager
|
|
.assign_sub_agent(session_id, item_id, task_id)
|
|
.await
|
|
.map_err(|error| SubAgentError::Storage(error.to_string()))?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn finish_work_item(&self, config: &SubAgentConfig, result: &SubAgentResult) {
|
|
let (Some(manager), Some(session_id), Some(item_id)) = (
|
|
self.work_manager.as_ref(),
|
|
config.session_id.as_deref(),
|
|
config.plan_item_id.as_deref(),
|
|
) else {
|
|
return;
|
|
};
|
|
let completed = matches!(result.status, TaskStatus::Completed);
|
|
let summary = if completed {
|
|
Some(result.content.as_str())
|
|
} else {
|
|
match &result.status {
|
|
TaskStatus::Failed(error) => Some(error.as_str()),
|
|
TaskStatus::TimedOut => Some("子 Agent 执行超时"),
|
|
TaskStatus::Cancelled => Some("子 Agent 已取消"),
|
|
TaskStatus::Completed => None,
|
|
}
|
|
};
|
|
if let Err(error) = manager
|
|
.finish_sub_agent(session_id, item_id, &result.task_id, completed, summary)
|
|
.await
|
|
{
|
|
tracing::warn!(task_id = %result.task_id, item_id, error = %error, "Failed to update plan item");
|
|
}
|
|
}
|
|
}
|
|
|
|
fn generate_task_id() -> String {
|
|
Uuid::new_v4().to_string()[..8].to_string()
|
|
}
|
|
|
|
fn format_duration(seconds: u64) -> String {
|
|
if seconds < 60 {
|
|
format!("{}s", seconds)
|
|
} else if seconds < 3600 {
|
|
format!("{}m", seconds / 60)
|
|
} else {
|
|
format!("{}h", seconds / 3600)
|
|
}
|
|
}
|
|
|
|
fn truncate_sub_agent_result(content: &str) -> (String, bool) {
|
|
if content.len() <= MAX_INLINE_RESULT_CHARS {
|
|
(content.to_string(), false)
|
|
} else {
|
|
let truncate_at = content.floor_char_boundary(MAX_INLINE_RESULT_CHARS);
|
|
(
|
|
format!(
|
|
"{}\n\n[... 结果已截断，共 {} 字符，完整结果请使用 check_task 查看 ...]",
|
|
&content[..truncate_at],
|
|
content.len()
|
|
),
|
|
true,
|
|
)
|
|
}
|
|
}
|
|
|
|
fn summarize_for_notification(content: &str, _duration_ms: u64) -> String {
|
|
const MAX_SUMMARY_BYTES: usize = 500;
|
|
if content.len() <= MAX_SUMMARY_BYTES {
|
|
content.to_string()
|
|
} else {
|
|
let truncate_at = content.floor_char_boundary(MAX_SUMMARY_BYTES);
|
|
format!("{}...", &content[..truncate_at])
|
|
}
|
|
}
|
|
|
|
#[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),
|
|
plan_item_id: None,
|
|
session_id: None,
|
|
},
|
|
DelegateContext {
|
|
session_id: "cli:test:dialog".into(),
|
|
channel: "cli".into(),
|
|
chat_id: "test".into(),
|
|
},
|
|
)
|
|
.await
|
|
.unwrap_err();
|
|
|
|
assert!(matches!(error, SubAgentError::TooManyTasks(1)));
|
|
}
|
|
}
|