use super::AgentChain;
use super::AgentInvocation;
use super::ExecutionPlan;
use super::ExecutionStep;
use super::InvocationRequest;
use super::SubagentContext;
use super::SubagentError;
use super::SubagentExecution;
use super::SubagentRegistry;
use super::SubagentStatus;
use crate::modes::OperatingMode;
use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::SystemTime;
use tokio::sync::Mutex;
use tokio::sync::RwLock;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::time::sleep;
use tokio::time::timeout;
use tracing::debug;
use tracing::info;
use tracing::warn;
use uuid::Uuid;
const DEFAULT_MAX_CONCURRENCY: usize = 8;
const DEFAULT_AGENT_TIMEOUT: Duration = Duration::from_secs(300);
const MAX_RETRIES: u32 = 3;
const RETRY_BACKOFF: Duration = Duration::from_secs(2);
const CIRCUIT_BREAKER_THRESHOLD: u32 = 5;
const CIRCUIT_BREAKER_RESET: Duration = Duration::from_secs(60);
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OrchestratorConfig {
pub max_concurrency: usize,
pub agent_timeout: Duration,
pub enable_retries: bool,
pub max_retries: u32,
pub retry_backoff: Duration,
pub enable_circuit_breaker: bool,
pub circuit_breaker_threshold: u32,
pub circuit_breaker_reset: Duration,
pub monitor_memory: bool,
pub memory_threshold_mb: usize,
}
impl Default for OrchestratorConfig {
fn default() -> Self {
Self {
max_concurrency: DEFAULT_MAX_CONCURRENCY,
agent_timeout: DEFAULT_AGENT_TIMEOUT,
enable_retries: true,
max_retries: MAX_RETRIES,
retry_backoff: RETRY_BACKOFF,
enable_circuit_breaker: true,
circuit_breaker_threshold: CIRCUIT_BREAKER_THRESHOLD,
circuit_breaker_reset: CIRCUIT_BREAKER_RESET,
monitor_memory: true,
memory_threshold_mb: 2048, }
}
}
#[derive(Debug, Clone)]
pub struct SharedContext {
data: Arc<RwLock<HashMap<String, serde_json::Value>>>,
previous_outputs: Arc<RwLock<Vec<String>>>,
modified_files: Arc<RwLock<Vec<PathBuf>>>,
errors: Arc<RwLock<Vec<String>>>,
}
impl Default for SharedContext {
fn default() -> Self {
Self::new()
}
}
impl SharedContext {
pub fn new() -> Self {
Self {
data: Arc::new(RwLock::new(HashMap::new())),
previous_outputs: Arc::new(RwLock::new(Vec::new())),
modified_files: Arc::new(RwLock::new(Vec::new())),
errors: Arc::new(RwLock::new(Vec::new())),
}
}
pub async fn set(&self, key: String, value: serde_json::Value) {
self.data.write().await.insert(key, value);
}
pub async fn get(&self, key: &str) -> Option<serde_json::Value> {
self.data.read().await.get(key).cloned()
}
pub async fn add_output(&self, output: String) {
self.previous_outputs.write().await.push(output);
}
pub async fn last_output(&self) -> Option<String> {
self.previous_outputs.read().await.last().cloned()
}
pub async fn all_outputs(&self) -> Vec<String> {
self.previous_outputs.read().await.clone()
}
pub async fn add_modified_files(&self, files: Vec<PathBuf>) {
self.modified_files.write().await.extend(files);
}
pub async fn modified_files(&self) -> Vec<PathBuf> {
self.modified_files.read().await.clone()
}
pub async fn add_error(&self, error: String) {
self.errors.write().await.push(error);
}
pub async fn errors(&self) -> Vec<String> {
self.errors.read().await.clone()
}
pub async fn snapshot(&self) -> ContextSnapshot {
ContextSnapshot {
data: self.data.read().await.clone(),
previous_outputs: self.previous_outputs.read().await.clone(),
modified_files: self.modified_files.read().await.clone(),
errors: self.errors.read().await.clone(),
timestamp: SystemTime::now(),
}
}
pub async fn restore(&self, snapshot: ContextSnapshot) {
*self.data.write().await = snapshot.data;
*self.previous_outputs.write().await = snapshot.previous_outputs;
*self.modified_files.write().await = snapshot.modified_files;
*self.errors.write().await = snapshot.errors;
}
pub async fn merge(&self, other: &SharedContext) {
let other_data = other.data.read().await;
for (key, value) in other_data.iter() {
self.data.write().await.insert(key.clone(), value.clone());
}
let other_outputs = other.previous_outputs.read().await;
self.previous_outputs
.write()
.await
.extend(other_outputs.clone());
let other_files = other.modified_files.read().await;
self.modified_files
.write()
.await
.extend(other_files.clone());
let other_errors = other.errors.read().await;
self.errors.write().await.extend(other_errors.clone());
}
}
#[derive(Debug, Clone)]
pub struct ContextSnapshot {
pub data: HashMap<String, serde_json::Value>,
pub previous_outputs: Vec<String>,
pub modified_files: Vec<PathBuf>,
pub errors: Vec<String>,
pub timestamp: SystemTime,
}
#[derive(Debug, Clone)]
pub struct ProgressUpdate {
pub execution_id: Uuid,
pub agent_name: String,
pub status: SubagentStatus,
pub message: Option<String>,
pub progress_percentage: Option<u8>,
pub timestamp: SystemTime,
}
#[derive(Debug)]
pub struct OrchestratorResult {
pub request: InvocationRequest,
pub executions: Vec<SubagentExecution>,
pub context: SharedContext,
pub total_duration: Duration,
pub success: bool,
pub partial_success: bool,
}
#[derive(Debug)]
struct CircuitBreaker {
failure_count: AtomicU32,
is_open: AtomicBool,
last_failure_time: Arc<Mutex<Option<SystemTime>>>,
threshold: u32,
reset_duration: Duration,
}
impl CircuitBreaker {
fn new(threshold: u32, reset_duration: Duration) -> Self {
Self {
failure_count: AtomicU32::new(0),
is_open: AtomicBool::new(false),
last_failure_time: Arc::new(Mutex::new(None)),
threshold,
reset_duration,
}
}
async fn record_success(&self) {
self.failure_count.store(0, Ordering::SeqCst);
self.is_open.store(false, Ordering::SeqCst);
*self.last_failure_time.lock().await = None;
}
async fn record_failure(&self) {
let count = self.failure_count.fetch_add(1, Ordering::SeqCst) + 1;
*self.last_failure_time.lock().await = Some(SystemTime::now());
if count >= self.threshold {
self.is_open.store(true, Ordering::SeqCst);
warn!(
"Circuit breaker opened after {} consecutive failures",
count
);
}
}
async fn is_open(&self) -> bool {
if !self.is_open.load(Ordering::SeqCst) {
return false;
}
if let Some(last_failure) = *self.last_failure_time.lock().await
&& let Ok(elapsed) = SystemTime::now().duration_since(last_failure)
&& elapsed > self.reset_duration
{
self.is_open.store(false, Ordering::SeqCst);
self.failure_count.store(0, Ordering::SeqCst);
info!("Circuit breaker reset after {:?}", elapsed);
return false;
}
true
}
}
#[derive(Debug)]
pub struct AgentOrchestrator {
config: OrchestratorConfig,
registry: Arc<SubagentRegistry>,
concurrency_limiter: Arc<Semaphore>,
circuit_breakers: Arc<RwLock<HashMap<String, Arc<CircuitBreaker>>>>,
progress_tx: mpsc::UnboundedSender<ProgressUpdate>,
progress_rx: Arc<Mutex<mpsc::UnboundedReceiver<ProgressUpdate>>>,
cancelled: Arc<AtomicBool>,
active_executions: Arc<AtomicUsize>,
operating_mode: OperatingMode,
}
impl AgentOrchestrator {
pub fn new(
registry: Arc<SubagentRegistry>,
config: OrchestratorConfig,
operating_mode: OperatingMode,
) -> Self {
let (progress_tx, progress_rx) = mpsc::unbounded_channel();
Self {
config: config.clone(),
registry,
concurrency_limiter: Arc::new(Semaphore::new(config.max_concurrency)),
circuit_breakers: Arc::new(RwLock::new(HashMap::new())),
progress_tx,
progress_rx: Arc::new(Mutex::new(progress_rx)),
cancelled: Arc::new(AtomicBool::new(false)),
active_executions: Arc::new(AtomicUsize::new(0)),
operating_mode,
}
}
pub async fn execute_plan(
&self,
request: InvocationRequest,
) -> Result<OrchestratorResult, SubagentError> {
let start_time = SystemTime::now();
info!("Starting orchestrator execution for request {}", request.id);
if self.cancelled.load(Ordering::SeqCst) {
return Err(SubagentError::ExecutionFailed(
"Execution cancelled".to_string(),
));
}
if self.config.monitor_memory {
self.check_memory_pressure().await?;
}
let context = SharedContext::new();
let executions = match &request.execution_plan {
ExecutionPlan::Single(invocation) => {
vec![self.execute_single(invocation.clone(), &context).await?]
}
ExecutionPlan::Sequential(chain) => {
self.execute_sequential(chain.clone(), &context).await?
}
ExecutionPlan::Parallel(invocations) => {
self.execute_parallel(invocations.clone(), &context).await?
}
ExecutionPlan::Mixed(steps) => self.execute_mixed(steps.clone(), &context).await?,
ExecutionPlan::Conditional(cond) => {
let mut results = Vec::new();
for agent in &cond.agents {
results.push(self.execute_single(agent.clone(), &context).await?);
}
results
}
};
let total_duration = SystemTime::now()
.duration_since(start_time)
.unwrap_or_default();
let success = executions
.iter()
.all(|e| e.status == SubagentStatus::Completed);
let partial_success = executions
.iter()
.any(|e| e.status == SubagentStatus::Completed);
info!(
"Orchestrator execution completed for request {} in {:?} (success: {}, partial: {})",
request.id, total_duration, success, partial_success
);
Ok(OrchestratorResult {
request,
executions,
context,
total_duration,
success,
partial_success,
})
}
pub async fn execute_single(
&self,
invocation: AgentInvocation,
shared_context: &SharedContext,
) -> Result<SubagentExecution, SubagentError> {
if self.config.enable_circuit_breaker {
let breaker = self
.get_or_create_circuit_breaker(&invocation.agent_name)
.await;
if breaker.is_open().await {
return Err(SubagentError::ExecutionFailed(format!(
"Circuit breaker open for agent {}",
invocation.agent_name
)));
}
}
let _permit = self.concurrency_limiter.acquire().await.map_err(|e| {
SubagentError::ExecutionFailed(format!("Failed to acquire permit: {}", e))
})?;
self.active_executions.fetch_add(1, Ordering::SeqCst);
let result = self.execute_with_retry(invocation, shared_context).await;
self.active_executions.fetch_sub(1, Ordering::SeqCst);
if self.config.enable_circuit_breaker {
let breaker = self
.get_or_create_circuit_breaker(
&result
.as_ref()
.map(|e| e.agent_name.clone())
.unwrap_or_default(),
)
.await;
match &result {
Ok(execution) if execution.status == SubagentStatus::Completed => {
breaker.record_success().await;
}
_ => {
breaker.record_failure().await;
}
}
}
result
}
pub async fn execute_sequential(
&self,
chain: AgentChain,
shared_context: &SharedContext,
) -> Result<Vec<SubagentExecution>, SubagentError> {
let mut executions = Vec::new();
for invocation in chain.agents {
if self.cancelled.load(Ordering::SeqCst) {
warn!("Sequential execution cancelled");
break;
}
let execution = self.execute_single(invocation, shared_context).await?;
if chain.pass_output
&& let Some(output) = &execution.output
{
shared_context.add_output(output.clone()).await;
}
shared_context
.add_modified_files(execution.modified_files.clone())
.await;
executions.push(execution);
}
Ok(executions)
}
pub async fn execute_parallel(
&self,
invocations: Vec<AgentInvocation>,
shared_context: &SharedContext,
) -> Result<Vec<SubagentExecution>, SubagentError> {
let mut tasks = Vec::new();
for invocation in invocations {
let self_clone = self.clone_for_task();
let context_clone = shared_context.clone();
let task =
tokio::spawn(
async move { self_clone.execute_single(invocation, &context_clone).await },
);
tasks.push(task);
}
let mut executions = Vec::new();
let mut errors = Vec::new();
for task in tasks {
match task.await {
Ok(Ok(execution)) => {
shared_context
.add_modified_files(execution.modified_files.clone())
.await;
executions.push(execution);
}
Ok(Err(e)) => {
errors.push(e.to_string());
shared_context.add_error(e.to_string()).await;
}
Err(e) => {
let error = format!("Task join error: {}", e);
errors.push(error.clone());
shared_context.add_error(error).await;
}
}
}
if !executions.is_empty() {
Ok(executions)
} else if !errors.is_empty() {
Err(SubagentError::ExecutionFailed(format!(
"All parallel executions failed: {}",
errors.join(", ")
)))
} else {
Err(SubagentError::ExecutionFailed(
"No executions completed".to_string(),
))
}
}
pub async fn execute_mixed(
&self,
steps: Vec<ExecutionStep>,
shared_context: &SharedContext,
) -> Result<Vec<SubagentExecution>, SubagentError> {
let mut all_executions = Vec::new();
for step in steps {
if self.cancelled.load(Ordering::SeqCst) {
warn!("Mixed execution cancelled");
break;
}
match step {
ExecutionStep::Single(invocation) => {
let execution = self.execute_single(invocation, shared_context).await?;
all_executions.push(execution);
}
ExecutionStep::Parallel(invocations) => {
let executions = self.execute_parallel(invocations, shared_context).await?;
all_executions.extend(executions);
}
ExecutionStep::Conditional(cond) => {
for agent in cond.agents {
let execution = self.execute_single(agent, shared_context).await?;
all_executions.push(execution);
}
}
ExecutionStep::Barrier => {
while self.active_executions.load(Ordering::SeqCst) > 0 {
sleep(Duration::from_millis(100)).await;
}
debug!("Barrier: All previous executions completed");
}
}
}
Ok(all_executions)
}
async fn execute_with_retry(
&self,
invocation: AgentInvocation,
shared_context: &SharedContext,
) -> Result<SubagentExecution, SubagentError> {
let mut attempts = 0;
let max_attempts = if self.config.enable_retries {
self.config.max_retries + 1
} else {
1
};
loop {
attempts += 1;
let mut execution = SubagentExecution::new(invocation.agent_name.clone());
self.send_progress(ProgressUpdate {
execution_id: execution.id,
agent_name: invocation.agent_name.clone(),
status: SubagentStatus::Running,
message: Some(format!(
"Starting execution (attempt {}/{})",
attempts, max_attempts
)),
progress_percentage: Some(0),
timestamp: SystemTime::now(),
})
.await;
let result = timeout(
self.config.agent_timeout,
self.execute_agent_internal(&invocation, shared_context, &mut execution),
)
.await;
match result {
Ok(Ok(())) => {
self.send_progress(ProgressUpdate {
execution_id: execution.id,
agent_name: invocation.agent_name.clone(),
status: SubagentStatus::Completed,
message: Some("Execution completed successfully".to_string()),
progress_percentage: Some(100),
timestamp: SystemTime::now(),
})
.await;
return Ok(execution);
}
Ok(Err(e)) if attempts < max_attempts && self.is_retriable_error(&e) => {
warn!(
"Agent {} failed with retriable error (attempt {}/{}): {}",
invocation.agent_name, attempts, max_attempts, e
);
self.send_progress(ProgressUpdate {
execution_id: execution.id,
agent_name: invocation.agent_name.clone(),
status: SubagentStatus::Running,
message: Some(format!("Retrying after error: {}", e)),
progress_percentage: None,
timestamp: SystemTime::now(),
})
.await;
sleep(self.config.retry_backoff * attempts).await;
continue;
}
Ok(Err(e)) => {
execution.fail(e.to_string());
self.send_progress(ProgressUpdate {
execution_id: execution.id,
agent_name: invocation.agent_name.clone(),
status: SubagentStatus::Failed(e.to_string()),
message: Some(format!("Execution failed: {}", e)),
progress_percentage: None,
timestamp: SystemTime::now(),
})
.await;
return Err(e);
}
Err(_) => {
let error = format!(
"Agent {} timed out after {:?}",
invocation.agent_name, self.config.agent_timeout
);
execution.fail(error.clone());
self.send_progress(ProgressUpdate {
execution_id: execution.id,
agent_name: invocation.agent_name.clone(),
status: SubagentStatus::Failed(error.clone()),
message: Some(error.clone()),
progress_percentage: None,
timestamp: SystemTime::now(),
})
.await;
if attempts < max_attempts {
warn!(
"Agent {} timed out (attempt {}/{})",
invocation.agent_name, attempts, max_attempts
);
sleep(self.config.retry_backoff * attempts).await;
continue;
}
return Err(SubagentError::Timeout {
name: invocation.agent_name,
});
}
}
}
}
async fn execute_agent_internal(
&self,
invocation: &AgentInvocation,
shared_context: &SharedContext,
execution: &mut SubagentExecution,
) -> Result<(), SubagentError> {
execution.start();
let _agent = self
.registry
.get_agent(&invocation.agent_name)
.ok_or_else(|| SubagentError::AgentNotFound {
name: invocation.agent_name.clone(),
})?;
let agent_context = SubagentContext {
execution_id: execution.id,
mode: self.operating_mode, available_tools: vec![], conversation_context: shared_context.last_output().await.unwrap_or_default(),
working_directory: std::env::current_dir().unwrap_or_default(),
parameters: invocation.parameters.clone(),
metadata: HashMap::new(),
};
info!(
"Executing agent {} with context: {:?}",
invocation.agent_name, agent_context
);
for i in 1..=10 {
if self.cancelled.load(Ordering::SeqCst) {
return Err(SubagentError::ExecutionFailed(
"Execution cancelled".to_string(),
));
}
self.send_progress(ProgressUpdate {
execution_id: execution.id,
agent_name: invocation.agent_name.clone(),
status: SubagentStatus::Running,
message: Some(format!("Processing step {}/10", i)),
progress_percentage: Some((i * 10) as u8),
timestamp: SystemTime::now(),
})
.await;
sleep(Duration::from_millis(100)).await;
}
let output = format!(
"Agent {} completed successfully with parameters: {:?}",
invocation.agent_name, invocation.parameters
);
execution.complete(output.clone(), vec![]);
Ok(())
}
const fn is_retriable_error(&self, error: &SubagentError) -> bool {
matches!(
error,
SubagentError::Timeout { .. }
| SubagentError::ExecutionFailed(_)
| SubagentError::Io(_)
)
}
async fn get_or_create_circuit_breaker(&self, agent_name: &str) -> Arc<CircuitBreaker> {
let mut breakers = self.circuit_breakers.write().await;
breakers
.entry(agent_name.to_string())
.or_insert_with(|| {
Arc::new(CircuitBreaker::new(
self.config.circuit_breaker_threshold,
self.config.circuit_breaker_reset,
))
})
.clone()
}
async fn check_memory_pressure(&self) -> Result<(), SubagentError> {
let memory_usage_mb = 500;
if memory_usage_mb > self.config.memory_threshold_mb {
Err(SubagentError::ExecutionFailed(format!(
"Memory pressure too high: {}MB > {}MB threshold",
memory_usage_mb, self.config.memory_threshold_mb
)))
} else {
Ok(())
}
}
async fn send_progress(&self, update: ProgressUpdate) {
if let Err(e) = self.progress_tx.send(update) {
warn!("Failed to send progress update: {}", e);
}
}
pub async fn progress_receiver(&self) -> mpsc::UnboundedReceiver<ProgressUpdate> {
let (_tx, rx) = mpsc::unbounded_channel();
rx
}
pub fn cancel(&self) {
self.cancelled.store(true, Ordering::SeqCst);
info!("Orchestrator execution cancelled");
}
pub fn reset_cancellation(&self) {
self.cancelled.store(false, Ordering::SeqCst);
}
pub fn active_count(&self) -> usize {
self.active_executions.load(Ordering::SeqCst)
}
fn clone_for_task(&self) -> Arc<Self> {
Arc::new(Self {
config: self.config.clone(),
registry: self.registry.clone(),
concurrency_limiter: self.concurrency_limiter.clone(),
circuit_breakers: self.circuit_breakers.clone(),
progress_tx: self.progress_tx.clone(),
progress_rx: self.progress_rx.clone(),
cancelled: self.cancelled.clone(),
active_executions: self.active_executions.clone(),
operating_mode: self.operating_mode,
})
}
}
impl AgentOrchestrator {
pub async fn execute_conditional<F>(
&self,
invocation: AgentInvocation,
shared_context: &SharedContext,
condition: F,
) -> Result<Option<SubagentExecution>, SubagentError>
where
F: Fn(&SharedContext) -> futures::future::BoxFuture<'_, bool> + Send + Sync,
{
if condition(shared_context).await {
Ok(Some(self.execute_single(invocation, shared_context).await?))
} else {
info!("Skipping agent {} due to condition", invocation.agent_name);
Ok(None)
}
}
pub async fn execute_with_dependencies(
&self,
invocation: AgentInvocation,
dependencies: Vec<String>,
shared_context: &SharedContext,
) -> Result<SubagentExecution, SubagentError> {
let outputs = shared_context.all_outputs().await;
for dep in dependencies {
if !outputs.iter().any(|o| o.contains(&dep)) {
return Err(SubagentError::ExecutionFailed(format!(
"Dependency {} not satisfied for agent {}",
dep, invocation.agent_name
)));
}
}
self.execute_single(invocation, shared_context).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::subagents::config::SubagentConfig;
async fn create_test_orchestrator() -> AgentOrchestrator {
use tempfile::TempDir;
let temp_dir = TempDir::new().unwrap();
let home_dir = temp_dir.path().join(".agcodex");
std::fs::create_dir_all(&home_dir).unwrap();
unsafe {
std::env::set_var("HOME", temp_dir.path());
}
let registry = Arc::new(SubagentRegistry::new().unwrap());
let test_agent = SubagentConfig {
name: "test-agent".to_string(),
description: "Test agent".to_string(),
mode_override: None,
intelligence: crate::subagents::config::IntelligenceLevel::Medium,
tools: std::collections::HashMap::new(),
prompt: "Test prompt".to_string(),
parameters: vec![],
template: None,
timeout_seconds: 10,
chainable: true,
parallelizable: true,
metadata: std::collections::HashMap::new(),
file_patterns: vec![],
tags: vec![], };
let global_agents_dir = temp_dir
.path()
.join(".agcodex")
.join("agents")
.join("global");
std::fs::create_dir_all(&global_agents_dir).unwrap();
let config_path = global_agents_dir.join("test-agent.toml");
test_agent.to_file(&config_path).unwrap();
registry.load_all().unwrap();
AgentOrchestrator::new(
registry,
OrchestratorConfig::default(),
OperatingMode::Build,
)
}
#[tokio::test]
async fn test_single_agent_execution() {
let orchestrator = create_test_orchestrator().await;
let invocation = AgentInvocation {
agent_name: "test-agent".to_string(),
parameters: HashMap::new(),
raw_parameters: String::new(),
position: 0,
intelligence_override: None,
mode_override: None,
};
let context = SharedContext::new();
let result = orchestrator.execute_single(invocation, &context).await;
assert!(result.is_ok());
let execution = result.unwrap();
assert_eq!(execution.status, SubagentStatus::Completed);
}
#[tokio::test]
async fn test_sequential_execution() {
let orchestrator = create_test_orchestrator().await;
let chain = AgentChain {
agents: vec![
AgentInvocation {
agent_name: "test-agent".to_string(),
parameters: HashMap::new(),
raw_parameters: String::new(),
position: 0,
intelligence_override: None,
mode_override: None,
},
AgentInvocation {
agent_name: "test-agent".to_string(),
parameters: HashMap::new(),
raw_parameters: String::new(),
position: 1,
intelligence_override: None,
mode_override: None,
},
],
pass_output: true,
};
let context = SharedContext::new();
let result = orchestrator.execute_sequential(chain, &context).await;
assert!(result.is_ok());
let executions = result.unwrap();
assert_eq!(executions.len(), 2);
assert!(
executions
.iter()
.all(|e| e.status == SubagentStatus::Completed)
);
}
#[tokio::test]
async fn test_parallel_execution() {
let orchestrator = create_test_orchestrator().await;
let invocations = vec![
AgentInvocation {
agent_name: "test-agent".to_string(),
parameters: HashMap::new(),
raw_parameters: String::new(),
position: 0,
intelligence_override: None,
mode_override: None,
},
AgentInvocation {
agent_name: "test-agent".to_string(),
parameters: HashMap::new(),
raw_parameters: String::new(),
position: 1,
intelligence_override: None,
mode_override: None,
},
];
let context = SharedContext::new();
let result = orchestrator.execute_parallel(invocations, &context).await;
assert!(result.is_ok());
let executions = result.unwrap();
assert_eq!(executions.len(), 2);
}
#[tokio::test]
async fn test_context_sharing() {
let context = SharedContext::new();
context
.set("key1".to_string(), serde_json::json!("value1"))
.await;
let value = context.get("key1").await;
assert_eq!(value, Some(serde_json::json!("value1")));
context.add_output("output1".to_string()).await;
context.add_output("output2".to_string()).await;
assert_eq!(context.last_output().await, Some("output2".to_string()));
assert_eq!(context.all_outputs().await.len(), 2);
let snapshot = context.snapshot().await;
context
.set("key2".to_string(), serde_json::json!("value2"))
.await;
assert!(context.get("key2").await.is_some());
context.restore(snapshot).await;
assert!(context.get("key2").await.is_none());
assert_eq!(context.get("key1").await, Some(serde_json::json!("value1")));
}
#[tokio::test]
async fn test_circuit_breaker() {
let breaker = CircuitBreaker::new(3, Duration::from_secs(1));
for _ in 0..3 {
breaker.record_failure().await;
}
assert!(breaker.is_open().await);
tokio::time::sleep(Duration::from_secs(2)).await;
assert!(!breaker.is_open().await);
breaker.record_success().await;
assert!(!breaker.is_open().await);
}
#[tokio::test]
async fn test_cancellation() {
let orchestrator = create_test_orchestrator().await;
orchestrator.cancel();
let invocation = AgentInvocation {
agent_name: "test-agent".to_string(),
parameters: HashMap::new(),
raw_parameters: String::new(),
position: 0,
intelligence_override: None,
mode_override: None,
};
let context = SharedContext::new();
let result = orchestrator.execute_single(invocation, &context).await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SubagentError::ExecutionFailed(_)
));
}
}