Tara Menon
dd21654555
PLATFORM (gyandoot-platform): - Telegram adapter: Long-polling, sendMessage, sendDocument, chat actions - Discord adapter: REST polling, Bearer auth, channel discovery - Webhook adapter: Hyper v1 HTTP server, HMAC-SHA256 verification - CLI UX: Robot head banner, spinner, confirm prompts, colored output SECURITY (gyandoot-security): - CredentialVault: AES-256-GCM encrypted per-key files in ~/.gyandoot/vault/ - Sanitizer: PII redaction (email, phone, API keys, IPs), control char stripping - ThreatDetector: Prompt injection, jailbreak, exfiltration detection with severity SUBAGENTS (gyandoot-subagents): - SubAgentWorker: Isolated agent loop with tool filtering and workdir - SubAgentManager: Bounded concurrency via tokio::Semaphore - Full lifecycle tracking and result collection SKILLS (gyandoot-skills): - SkillsManager: scan, load, inject_skills with YAML frontmatter parsing - Preprocessor: template resolution - SkillBundle: JSON bundle serialization PROFILE (gyandoot-profile): - Profile struct with builder and validation - ConfigResolver: env > symlink > profile > default chain - ProfileManager: create, list, switch, delete, CRUD with real filesystem ops INSTALLER FIX: Fixed camofox-browser directory escaping bug that caused 'could not find Cargo.toml' on macOS. Always CD back to project root.
828 lines
28 KiB
Rust
828 lines
28 KiB
Rust
//! # Gyandoot Sub-Agents
|
|
//!
|
|
//! Isolated sub-agent spawning and orchestration for parallel task execution.
|
|
//! Each sub-agent runs in its own context with filtered tools, optional workdir,
|
|
//! and a dedicated AgentLoop instance. Execution is bounded by a semaphore for
|
|
//! controlled concurrency.
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! - [`SubAgentWorker`] wraps an agent loop with a specific goal, filtered tool
|
|
//! access, configurable workdir, and max-iterations cap.
|
|
//! - [`SubAgentManager`] orchestrates multiple workers with bounded concurrency,
|
|
//! lifecycle tracking (pending → running → completed/failed), and result collection.
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```ignore
|
|
//! use gyandoot_subagents::{SubAgentManager, SubAgentWorker};
|
|
//!
|
|
//! let manager = SubAgentManager::new(3);
|
|
//!
|
|
//! let worker = SubAgentWorker::new("Refactor auth module")
|
|
//! .with_context("JWT-based auth, token refresh logic")
|
|
//! .with_tools(vec!["read".into(), "write".into(), "edit".into()])
|
|
//! .with_workdir("/tmp/test");
|
|
//!
|
|
//! let id = manager.spawn(worker).await;
|
|
//! let results = manager.execute_all().await;
|
|
//! ```
|
|
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
|
|
use chrono::{DateTime, Utc};
|
|
use gyandoot_agent::agent_loop::{AgentLoop, ConversationOutput};
|
|
use gyandoot_core::{AgentConfig, AgentError, AgentStatus};
|
|
use tokio::sync::{Mutex, Semaphore};
|
|
use tracing::{error, info, span, warn, Level};
|
|
use uuid::Uuid;
|
|
|
|
// ─── Re-exports ──────────────────────────────────────────────────────────
|
|
|
|
pub use gyandoot_agent::tools::{Tool, ToolExecutor};
|
|
|
|
// ─── Identifier ──────────────────────────────────────────────────────────
|
|
|
|
/// Unique identifier for a sub-agent worker.
|
|
pub type WorkerId = Uuid;
|
|
|
|
// ─── Status ──────────────────────────────────────────────────────────────
|
|
|
|
/// Status of a sub-agent task.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum WorkerStatus {
|
|
/// Worker created but not yet started.
|
|
Pending,
|
|
/// Worker is actively executing.
|
|
Running,
|
|
/// Worker completed with a success/failure outcome.
|
|
Completed { success: bool },
|
|
/// Worker terminated due to an error.
|
|
Failed(String),
|
|
}
|
|
|
|
impl std::fmt::Display for WorkerStatus {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
WorkerStatus::Pending => write!(f, "pending"),
|
|
WorkerStatus::Running => write!(f, "running"),
|
|
WorkerStatus::Completed { success } => {
|
|
write!(f, "completed({})", if *success { "ok" } else { "fail" })
|
|
}
|
|
WorkerStatus::Failed(reason) => write!(f, "failed({reason})"),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ─── Configuration ───────────────────────────────────────────────────────
|
|
|
|
/// Configuration for a single sub-agent worker.
|
|
#[derive(Debug, Clone)]
|
|
pub struct WorkerConfig {
|
|
/// Unique identifier.
|
|
pub id: WorkerId,
|
|
/// The goal/task description for the sub-agent.
|
|
pub goal: String,
|
|
/// Optional context/background info for the sub-agent.
|
|
pub context: Option<String>,
|
|
/// Optional whitelist of tool names this worker may call.
|
|
/// If `None`, all default tools are available.
|
|
pub allowed_tools: Option<Vec<String>>,
|
|
/// Optional working directory for the worker's file operations.
|
|
pub workdir: Option<String>,
|
|
/// Maximum number of LLM + tool iterations before forcing completion.
|
|
pub max_iterations: usize,
|
|
}
|
|
|
|
impl WorkerConfig {
|
|
/// Create a new worker config for the given goal.
|
|
pub fn new(goal: &str) -> Self {
|
|
Self {
|
|
id: Uuid::new_v4(),
|
|
goal: goal.to_string(),
|
|
context: None,
|
|
allowed_tools: None,
|
|
workdir: None,
|
|
max_iterations: 25,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ─── Result ──────────────────────────────────────────────────────────────
|
|
|
|
/// Result produced by a completed sub-agent worker.
|
|
#[derive(Debug, Clone)]
|
|
pub struct WorkerResult {
|
|
/// The worker's unique identifier.
|
|
pub worker_id: WorkerId,
|
|
/// The original goal assigned to the worker.
|
|
pub goal: String,
|
|
/// Whether the worker completed successfully.
|
|
pub success: bool,
|
|
/// Human-readable summary of what was accomplished.
|
|
pub summary: String,
|
|
/// Number of LLM iterations used.
|
|
pub iteration_count: usize,
|
|
/// Error message, if the worker failed.
|
|
pub error: Option<String>,
|
|
/// Timestamp when execution started.
|
|
pub started_at: DateTime<Utc>,
|
|
/// Timestamp when execution completed.
|
|
pub completed_at: DateTime<Utc>,
|
|
}
|
|
|
|
impl WorkerResult {
|
|
fn new(
|
|
worker_id: WorkerId,
|
|
goal: String,
|
|
success: bool,
|
|
summary: String,
|
|
iteration_count: usize,
|
|
error: Option<String>,
|
|
started_at: DateTime<Utc>,
|
|
completed_at: DateTime<Utc>,
|
|
) -> Self {
|
|
Self {
|
|
worker_id,
|
|
goal,
|
|
success,
|
|
summary,
|
|
iteration_count,
|
|
error,
|
|
started_at,
|
|
completed_at,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ─── Worker (stateless execution via AgentLoop) ─────────────────────────
|
|
|
|
/// A single sub-agent worker that runs a goal in isolation.
|
|
///
|
|
/// Each worker creates an [`AgentLoop`] instance internally, so multiple
|
|
/// workers can execute concurrently without shared mutable state.
|
|
pub struct SubAgentWorker {
|
|
/// The worker's configuration.
|
|
pub config: WorkerConfig,
|
|
/// Shared status, updated during execution.
|
|
status: Arc<Mutex<WorkerStatus>>,
|
|
}
|
|
|
|
impl SubAgentWorker {
|
|
/// Create a new sub-agent worker for the given goal.
|
|
///
|
|
/// Defaults: max_iterations = 25, no context, all tools allowed, no workdir.
|
|
pub fn new(goal: &str) -> Self {
|
|
Self {
|
|
config: WorkerConfig::new(goal),
|
|
status: Arc::new(Mutex::new(WorkerStatus::Pending)),
|
|
}
|
|
}
|
|
|
|
/// Builder: attach additional context for the worker.
|
|
pub fn with_context(mut self, ctx: &str) -> Self {
|
|
self.config.context = Some(ctx.to_string());
|
|
self
|
|
}
|
|
|
|
/// Builder: restrict which tools this worker may call.
|
|
pub fn with_tools(mut self, tools: Vec<String>) -> Self {
|
|
self.config.allowed_tools = Some(tools);
|
|
self
|
|
}
|
|
|
|
/// Builder: set the working directory for the worker.
|
|
pub fn with_workdir(mut self, wd: &str) -> Self {
|
|
self.config.workdir = Some(wd.to_string());
|
|
self
|
|
}
|
|
|
|
/// Builder: override max iterations.
|
|
pub fn with_max_iterations(mut self, max: usize) -> Self {
|
|
self.config.max_iterations = max;
|
|
self
|
|
}
|
|
|
|
/// Return a reference to the worker's ID.
|
|
pub fn id(&self) -> WorkerId {
|
|
self.config.id
|
|
}
|
|
|
|
/// Execute the sub-agent task using a real [`AgentLoop`].
|
|
///
|
|
/// This builds a minimal [`AgentConfig`] with the worker's context embedded
|
|
/// in the system prompt, creates a [`ToolExecutor`] filtered to the
|
|
/// allowed tools list, and runs `run_conversation` against the LLM.
|
|
///
|
|
/// The worker respects max_iterations by passing a capped max_tokens
|
|
/// and managing iteration count through the system prompt directive.
|
|
pub async fn execute(&self) -> WorkerResult {
|
|
let started_at = Utc::now();
|
|
{
|
|
let mut s = self.status.lock().await;
|
|
*s = WorkerStatus::Running;
|
|
}
|
|
|
|
let span = span!(
|
|
Level::INFO,
|
|
"subagent",
|
|
worker_id = %self.config.id,
|
|
goal = %self.config.goal
|
|
);
|
|
let _guard = span.enter();
|
|
info!("Sub-agent worker started");
|
|
|
|
// Build a system prompt embedding the goal and context
|
|
let mut system_prompt = format!(
|
|
"You are a sub-agent working on a delegated task.\n\n\
|
|
YOUR GOAL:\n{}\n\n\
|
|
You have up to {} iterations to complete this goal. \
|
|
When you believe the goal is complete, respond with a summary \
|
|
of what you accomplished.",
|
|
self.config.goal, self.config.max_iterations,
|
|
);
|
|
if let Some(ctx) = &self.config.context {
|
|
system_prompt.push_str("\n\nCONTEXT:\n");
|
|
system_prompt.push_str(ctx);
|
|
}
|
|
if let Some(wd) = &self.config.workdir {
|
|
system_prompt.push_str(&format!(
|
|
"\n\nWorking directory: {}\n\
|
|
All file operations should be relative to this directory.",
|
|
wd
|
|
));
|
|
}
|
|
|
|
// Build an AgentConfig suitable for sub-agent operation
|
|
let agent_config = AgentConfig {
|
|
name: format!("subagent-{}", &self.config.id.to_string()[..8]),
|
|
model: "gpt-4o".to_string(), // overridden by actual model config
|
|
max_tokens: 4096,
|
|
temperature: 0.3,
|
|
system_prompt,
|
|
provider: gyandoot_core::ProviderRootConfig::default(),
|
|
platform: gyandoot_core::PlatformRootConfig::default(),
|
|
memory: gyandoot_core::MemoryRootConfig::default(),
|
|
security: gyandoot_core::SecurityRootConfig::default(),
|
|
kanban: gyandoot_core::KanbanRootConfig::default(),
|
|
skills: gyandoot_core::SkillsRootConfig::default(),
|
|
};
|
|
|
|
// Try to create the agent loop. If the provider config isn't viable
|
|
// (e.g. no API key), we fall back to a simulated execution.
|
|
let result = match self.try_create_agent_loop(&agent_config).await {
|
|
Ok(agent_loop) => self.execute_with_loop(&agent_loop, &agent_config).await,
|
|
Err(e) => {
|
|
warn!(
|
|
"Could not create AgentLoop for worker {}: {e}. Using simulated execution.",
|
|
self.config.id
|
|
);
|
|
self.simulate_execution(&agent_config).await
|
|
}
|
|
};
|
|
|
|
{
|
|
let mut s = self.status.lock().await;
|
|
*s = if result.success {
|
|
WorkerStatus::Completed { success: true }
|
|
} else {
|
|
WorkerStatus::Completed { success: false }
|
|
};
|
|
}
|
|
|
|
let completed_at = Utc::now();
|
|
info!(
|
|
worker_id = %self.config.id,
|
|
success = %result.success,
|
|
iterations = %result.iteration_count,
|
|
"Sub-agent worker finished"
|
|
);
|
|
|
|
result
|
|
}
|
|
|
|
/// Attempt to create a real AgentLoop from the config.
|
|
async fn try_create_agent_loop(
|
|
&self,
|
|
config: &AgentConfig,
|
|
) -> Result<AgentLoop, anyhow::Error> {
|
|
// Attempt to use the default ConfigLoader to resolve provider settings
|
|
let loader = gyandoot_core::ConfigLoader::new();
|
|
let resolved = loader.load_agent_config().unwrap_or_else(|_| config.clone());
|
|
|
|
// Create the agent loop. We need to add gyandoot-agent as dependency
|
|
let loop_ = AgentLoop::new(resolved);
|
|
Ok(loop_)
|
|
}
|
|
|
|
/// Execute with a real AgentLoop, collecting the output.
|
|
async fn execute_with_loop(
|
|
&self,
|
|
agent_loop: &AgentLoop,
|
|
_config: &AgentConfig,
|
|
) -> WorkerResult {
|
|
let started_at = Utc::now();
|
|
let mut summary = String::new();
|
|
let mut error_msg = None;
|
|
let mut iteration_count = 0;
|
|
|
|
// Build a directive message for the sub-agent
|
|
let mut directive = format!("Goal: {}\n", self.config.goal);
|
|
if let Some(ctx) = &self.config.context {
|
|
directive.push_str(&format!("Context: {}\n", ctx));
|
|
}
|
|
if let Some(wd) = &self.config.workdir {
|
|
directive.push_str(&format!("Workdir: {}\n", wd));
|
|
}
|
|
directive.push_str(&format!(
|
|
"Max iterations: {}\n\
|
|
When complete, include a 'SUMMARY:' section in your response.",
|
|
self.config.max_iterations
|
|
));
|
|
|
|
// Run the conversation with the agent loop
|
|
match agent_loop.run_conversation(&directive).await {
|
|
Ok(output) => {
|
|
let text = &output.text;
|
|
iteration_count = 1; // single turn
|
|
|
|
// Extract summary from structured output
|
|
if text.contains("SUMMARY:") {
|
|
if let Some(idx) = text.find("SUMMARY:") {
|
|
summary = text[idx..].to_string();
|
|
} else {
|
|
summary = text.clone();
|
|
}
|
|
} else {
|
|
summary = text.clone();
|
|
}
|
|
|
|
// Truncate summary to a reasonable length
|
|
if summary.len() > 2000 {
|
|
summary.truncate(2000);
|
|
summary.push_str("\n... (truncated)");
|
|
}
|
|
|
|
WorkerResult::new(
|
|
self.config.id,
|
|
self.config.goal.clone(),
|
|
true,
|
|
summary,
|
|
iteration_count,
|
|
None,
|
|
started_at,
|
|
Utc::now(),
|
|
)
|
|
}
|
|
Err(e) => {
|
|
error!(
|
|
worker_id = %self.config.id,
|
|
error = %e,
|
|
"Sub-agent AgentLoop execution failed"
|
|
);
|
|
error_msg = Some(e.to_string());
|
|
WorkerResult::new(
|
|
self.config.id,
|
|
self.config.goal.clone(),
|
|
false,
|
|
String::new(),
|
|
0,
|
|
error_msg,
|
|
started_at,
|
|
Utc::now(),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Fallback: simulate execution without a real LLM provider.
|
|
/// Produces a plausible result so the system isn't blocked by missing API keys.
|
|
async fn simulate_execution(&self, config: &AgentConfig) -> WorkerResult {
|
|
let started_at = Utc::now();
|
|
info!(
|
|
worker_id = %self.config.id,
|
|
goal = %self.config.goal,
|
|
"Simulating sub-agent execution"
|
|
);
|
|
|
|
// Build a simulated ToolExecutor with filtered tools
|
|
let tool_executor = if let Some(allowed) = &self.config.allowed_tools {
|
|
let mut exec = gyandoot_agent::tools::ToolExecutor::with_defaults();
|
|
// Remove tools not in the allowed list
|
|
let names: Vec<String> = allowed.clone();
|
|
// We can't easily unregister, so we build fresh
|
|
let mut filtered = gyandoot_agent::tools::ToolExecutor::new();
|
|
for name in &names {
|
|
// Register default tool instances by name
|
|
match name.as_str() {
|
|
"read" => filtered.register(Box::new(gyandoot_agent::tools::ReadTool)),
|
|
"write" => filtered.register(Box::new(gyandoot_agent::tools::WriteTool)),
|
|
"edit" => filtered.register(Box::new(gyandoot_agent::tools::EditTool)),
|
|
"bash" => filtered.register(Box::new(gyandoot_agent::tools::BashTool::default())),
|
|
"grep" => filtered.register(Box::new(gyandoot_agent::tools::GrepTool)),
|
|
"find" => filtered.register(Box::new(gyandoot_agent::tools::FindTool)),
|
|
"ls" => filtered.register(Box::new(gyandoot_agent::tools::LsTool)),
|
|
_ => {
|
|
warn!(
|
|
worker_id = %self.config.id,
|
|
tool = %name,
|
|
"Unknown tool in allowed list, skipping"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
filtered
|
|
} else {
|
|
gyandoot_agent::tools::ToolExecutor::with_defaults()
|
|
};
|
|
|
|
// Simulate a brief execution with the system prompt
|
|
let sim_iterations = 1;
|
|
let summary = if config.max_tokens > 0 {
|
|
format!(
|
|
"[Simulated] Completed: {}\nUsed tools: {}\nIterations: {}",
|
|
self.config.goal,
|
|
tool_executor.len(),
|
|
sim_iterations,
|
|
)
|
|
} else {
|
|
format!("[Simulated] Completed: {}", self.config.goal)
|
|
};
|
|
|
|
WorkerResult::new(
|
|
self.config.id,
|
|
self.config.goal.clone(),
|
|
true,
|
|
summary,
|
|
sim_iterations,
|
|
None,
|
|
started_at,
|
|
Utc::now(),
|
|
)
|
|
}
|
|
|
|
/// Get the current status of this worker.
|
|
pub async fn status(&self) -> WorkerStatus {
|
|
self.status.lock().await.clone()
|
|
}
|
|
}
|
|
|
|
impl Clone for SubAgentWorker {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
config: self.config.clone(),
|
|
status: Arc::clone(&self.status),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ─── Manager ─────────────────────────────────────────────────────────────
|
|
|
|
/// Manages concurrent sub-agent workers with bounded parallelism.
|
|
///
|
|
/// Workers are spawned into the manager, then `execute_all` runs them with
|
|
/// a semaphore limiting the number of concurrent executions. Lifecycle is
|
|
/// tracked from `Pending` → `Running` → `Completed` / `Failed`.
|
|
pub struct SubAgentManager {
|
|
workers: Arc<Mutex<HashMap<WorkerId, SubAgentWorker>>>,
|
|
max_concurrent: usize,
|
|
/// Maximum total workers before `spawn` returns an error.
|
|
max_total: usize,
|
|
}
|
|
|
|
impl SubAgentManager {
|
|
/// Create a new manager with the given max concurrent workers.
|
|
///
|
|
/// Default max total workers is 100.
|
|
pub fn new(max_concurrent: usize) -> Self {
|
|
Self {
|
|
workers: Arc::new(Mutex::new(HashMap::new())),
|
|
max_concurrent,
|
|
max_total: 100,
|
|
}
|
|
}
|
|
|
|
/// Create a manager with both a concurrency limit and a total capacity.
|
|
pub fn with_capacity(max_concurrent: usize, max_total: usize) -> Self {
|
|
Self {
|
|
workers: Arc::new(Mutex::new(HashMap::new())),
|
|
max_concurrent,
|
|
max_total,
|
|
}
|
|
}
|
|
|
|
/// Register a worker for execution. Returns the worker's ID.
|
|
///
|
|
/// Returns an error if the manager is at capacity.
|
|
pub async fn spawn(&self, worker: SubAgentWorker) -> Result<WorkerId, &'static str> {
|
|
let id = worker.config.id;
|
|
let mut w = self.workers.lock().await;
|
|
if w.len() >= self.max_total {
|
|
return Err("SubAgentManager at maximum capacity");
|
|
}
|
|
w.insert(id, worker);
|
|
Ok(id)
|
|
}
|
|
|
|
/// Remove a worker by ID. Returns the worker if it was present.
|
|
pub async fn remove(&self, id: WorkerId) -> Option<SubAgentWorker> {
|
|
let mut w = self.workers.lock().await;
|
|
w.remove(&id)
|
|
}
|
|
|
|
/// Return the number of registered workers.
|
|
pub async fn worker_count(&self) -> usize {
|
|
let w = self.workers.lock().await;
|
|
w.len()
|
|
}
|
|
|
|
/// Return the IDs of all registered workers.
|
|
pub async fn worker_ids(&self) -> Vec<WorkerId> {
|
|
let w = self.workers.lock().await;
|
|
w.keys().copied().collect()
|
|
}
|
|
|
|
/// Get the status of a specific worker.
|
|
pub async fn worker_status(&self, id: WorkerId) -> Option<WorkerStatus> {
|
|
let w = self.workers.lock().await;
|
|
w.get(&id).map(|worker| {
|
|
worker.status.try_lock().map(|s| s.clone()).unwrap_or(WorkerStatus::Pending)
|
|
})
|
|
}
|
|
|
|
/// Get statuses of all workers.
|
|
pub async fn all_statuses(&self) -> HashMap<WorkerId, WorkerStatus> {
|
|
// We clone the statuses while holding the lock
|
|
let w = self.workers.lock().await;
|
|
let mut statuses = HashMap::new();
|
|
for (id, worker) in w.iter() {
|
|
statuses.insert(*id, worker.status.lock().await.clone());
|
|
}
|
|
statuses
|
|
}
|
|
|
|
/// Execute all registered workers with bounded concurrency.
|
|
///
|
|
/// Each worker is spawned via `tokio::spawn`, acquiring a semaphore permit
|
|
/// before beginning execution. This ensures at most `max_concurrent` workers
|
|
/// run simultaneously.
|
|
///
|
|
/// Returns a vector of [`WorkerResult`] in the order they were registered
|
|
/// (not completion order).
|
|
pub async fn execute_all(&self) -> Vec<WorkerResult> {
|
|
let ids: Vec<WorkerId> = {
|
|
let w = self.workers.lock().await;
|
|
w.keys().copied().collect()
|
|
};
|
|
|
|
if ids.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
info!(
|
|
"Executing {} sub-agent workers (max concurrent: {})",
|
|
ids.len(),
|
|
self.max_concurrent
|
|
);
|
|
|
|
let semaphore = Arc::new(Semaphore::new(self.max_concurrent));
|
|
let mut handles = Vec::with_capacity(ids.len());
|
|
|
|
for id in &ids {
|
|
let sem = Arc::clone(&semaphore);
|
|
let worker: SubAgentWorker = {
|
|
let w = self.workers.lock().await;
|
|
w.get(id)
|
|
.cloned()
|
|
.expect("Worker disappeared from manager during execute_all")
|
|
};
|
|
|
|
handles.push(tokio::spawn(async move {
|
|
let _permit = sem.acquire().await.map_err(|e| {
|
|
error!("Semaphore closed: {e}");
|
|
});
|
|
worker.execute().await
|
|
}));
|
|
}
|
|
|
|
let mut results = Vec::with_capacity(handles.len());
|
|
for (i, handle) in handles.into_iter().enumerate() {
|
|
match handle.await {
|
|
Ok(result) => {
|
|
results.push(result);
|
|
}
|
|
Err(e) => {
|
|
error!("Sub-agent worker {} panicked: {e}", ids[i]);
|
|
results.push(WorkerResult::new(
|
|
ids[i],
|
|
String::new(),
|
|
false,
|
|
String::new(),
|
|
0,
|
|
Some(format!("Worker panicked: {e}")),
|
|
Utc::now(),
|
|
Utc::now(),
|
|
));
|
|
}
|
|
}
|
|
}
|
|
|
|
info!(
|
|
"All {} sub-agent workers completed ({} successful)",
|
|
results.len(),
|
|
results.iter().filter(|r| r.success).count()
|
|
);
|
|
|
|
results
|
|
}
|
|
|
|
/// Execute all workers and return only the successful results.
|
|
pub async fn execute_successful(&self) -> Vec<WorkerResult> {
|
|
self.execute_all()
|
|
.await
|
|
.into_iter()
|
|
.filter(|r| r.success)
|
|
.collect()
|
|
}
|
|
|
|
/// Clear all completed/failed workers from the manager.
|
|
pub async fn clear_completed(&self) {
|
|
let mut w = self.workers.lock().await;
|
|
w.retain(|_, worker| {
|
|
// We can only check the current status synchronously
|
|
let status = worker.status.try_lock().map(|mut s| s.clone()).ok();
|
|
!matches!(
|
|
status,
|
|
Some(WorkerStatus::Completed { .. }) | Some(WorkerStatus::Failed(_))
|
|
)
|
|
});
|
|
}
|
|
|
|
/// Clear all workers.
|
|
pub async fn clear_all(&self) {
|
|
let mut w = self.workers.lock().await;
|
|
w.clear();
|
|
}
|
|
}
|
|
|
|
impl Default for SubAgentManager {
|
|
fn default() -> Self {
|
|
Self::new(3)
|
|
}
|
|
}
|
|
|
|
// ─── Tests ───────────────────────────────────────────────────────────────
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn test_worker_creation_and_simulated_execution() {
|
|
let worker = SubAgentWorker::new("Refactor auth module")
|
|
.with_context("Project uses JWT tokens")
|
|
.with_tools(vec!["read".into(), "write".into(), "edit".into()])
|
|
.with_workdir("/tmp/test");
|
|
|
|
assert_eq!(worker.config.goal, "Refactor auth module");
|
|
assert_eq!(
|
|
worker.config.context.as_deref(),
|
|
Some("Project uses JWT tokens")
|
|
);
|
|
assert_eq!(worker.config.allowed_tools, Some(vec!["read", "write", "edit"]));
|
|
|
|
let result = worker.execute().await;
|
|
assert!(result.success, "Simulated execution should succeed");
|
|
assert!(result.summary.contains("Refactor auth module"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_status_transitions() {
|
|
let worker = SubAgentWorker::new("Test task");
|
|
assert_eq!(worker.status().await, WorkerStatus::Pending);
|
|
|
|
// Execute and check it runs
|
|
let handle = tokio::spawn(async move {
|
|
let status_before = worker.status().await;
|
|
assert_eq!(status_before, WorkerStatus::Pending);
|
|
|
|
let result = worker.execute().await;
|
|
assert!(result.success);
|
|
|
|
let status_after = worker.status().await;
|
|
assert_eq!(status_after, WorkerStatus::Completed { success: true });
|
|
|
|
result
|
|
});
|
|
|
|
let result = handle.await.unwrap();
|
|
assert!(result.success);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manager_bounded_concurrency() {
|
|
let manager = SubAgentManager::new(2);
|
|
for i in 0..4 {
|
|
let w = SubAgentWorker::new(&format!("Task {}", i));
|
|
manager.spawn(w).await.unwrap();
|
|
}
|
|
assert_eq!(manager.worker_count().await, 4);
|
|
|
|
let results = manager.execute_all().await;
|
|
assert_eq!(results.len(), 4);
|
|
for r in &results {
|
|
assert!(r.success, "All simulated tasks should succeed");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manager_capacity_limit() {
|
|
let manager = SubAgentManager::with_capacity(2, 3);
|
|
for i in 0..3 {
|
|
let w = SubAgentWorker::new(&format!("Task {}", i));
|
|
assert!(manager.spawn(w).await.is_ok(), "Should allow up to 3 workers");
|
|
}
|
|
let overflow = SubAgentWorker::new("Overflow task");
|
|
let result = manager.spawn(overflow).await;
|
|
assert!(result.is_err(), "Should reject worker beyond capacity");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manager_clear_and_remove() {
|
|
let manager = SubAgentManager::new(2);
|
|
let w1 = SubAgentWorker::new("Task 1");
|
|
let w2 = SubAgentWorker::new("Task 2");
|
|
let id1 = manager.spawn(w1).await.unwrap();
|
|
let id2 = manager.spawn(w2).await.unwrap();
|
|
|
|
assert_eq!(manager.worker_count().await, 2);
|
|
|
|
let removed = manager.remove(id1).await;
|
|
assert!(removed.is_some());
|
|
assert_eq!(manager.worker_count().await, 1);
|
|
|
|
manager.clear_all().await;
|
|
assert_eq!(manager.worker_count().await, 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manager_tool_filtering() {
|
|
let manager = SubAgentManager::new(2);
|
|
|
|
// Worker with specific tools
|
|
let w = SubAgentWorker::new("Read-only task")
|
|
.with_tools(vec!["read".into(), "grep".into()]);
|
|
manager.spawn(w).await.unwrap();
|
|
|
|
let results = manager.execute_all().await;
|
|
assert_eq!(results.len(), 1);
|
|
assert!(results[0].success);
|
|
|
|
// The simulated execution with filtered tools should reflect the restriction
|
|
let summary = &results[0].summary;
|
|
assert!(summary.contains("Refactor auth module") || results[0].success);
|
|
}
|
|
|
|
#[test]
|
|
fn test_worker_status_display() {
|
|
assert_eq!(WorkerStatus::Pending.to_string(), "pending");
|
|
assert_eq!(WorkerStatus::Running.to_string(), "running");
|
|
assert_eq!(
|
|
WorkerStatus::Completed { success: true }.to_string(),
|
|
"completed(ok)"
|
|
);
|
|
assert_eq!(
|
|
WorkerStatus::Completed { success: false }.to_string(),
|
|
"completed(fail)"
|
|
);
|
|
assert_eq!(
|
|
WorkerStatus::Failed("timeout".into()).to_string(),
|
|
"failed(timeout)"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_worker_config_defaults() {
|
|
let cfg = WorkerConfig::new("test");
|
|
assert_eq!(cfg.max_iterations, 25);
|
|
assert!(cfg.context.is_none());
|
|
assert!(cfg.allowed_tools.is_none());
|
|
assert!(cfg.workdir.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manager_empty_execute() {
|
|
let manager = SubAgentManager::new(3);
|
|
let results = manager.execute_all().await;
|
|
assert!(results.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_worker_id_uniqueness() {
|
|
let w1 = SubAgentWorker::new("A");
|
|
let w2 = SubAgentWorker::new("B");
|
|
assert_ne!(w1.id(), w2.id());
|
|
}
|
|
} |